94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
"""Utilities for creating and populating a ChromaDB vector store.
|
||
|
||
This module contains two helper functions:
|
||
|
||
* :func:`create_vectorstore` – returns a :class:`langchain_chroma.Chroma` instance backed by
|
||
an ``OllamaEmbeddings`` model.
|
||
* :func:`load_documents` – reads ``.txt``/``.md`` files from a directory, splits them into
|
||
chunks using :class:`langchain_text_splitters.RecursiveCharacterTextSplitter`, and adds
|
||
the chunks to the vector store.
|
||
|
||
The vector store is persisted in ``./chroma_db`` by default, so it survives program
|
||
restarts.
|
||
"""
|
||
|
||
from pathlib import Path
|
||
from typing import Iterable
|
||
|
||
from langchain_chroma import Chroma
|
||
from langchain_ollama import OllamaEmbeddings
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Vector store creation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
|
||
"""Create a Chroma vector store backed by Ollama embeddings.
|
||
|
||
Parameters
|
||
----------
|
||
persist_directory: str
|
||
Path to the directory where the Chroma DB will be stored.
|
||
|
||
Returns
|
||
-------
|
||
Chroma
|
||
A Chroma vector store instance.
|
||
"""
|
||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||
return Chroma(
|
||
persist_directory=persist_directory,
|
||
embedding_function=embeddings,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Document ingestion
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def load_documents(directory: str | Path, vectorstore: Chroma) -> None:
|
||
"""Load all ``.txt`` and ``.md`` files from *directory* into *vectorstore*.
|
||
|
||
The files are split into chunks using
|
||
:class:`langchain_text_splitters.RecursiveCharacterTextSplitter` before being
|
||
added to the vector store.
|
||
|
||
Parameters
|
||
----------
|
||
directory: str | Path
|
||
Directory containing the documents.
|
||
vectorstore: Chroma
|
||
The vector store to populate.
|
||
"""
|
||
path = Path(directory)
|
||
if not path.is_dir():
|
||
raise ValueError(f"{directory!r} is not a directory")
|
||
|
||
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
||
docs = []
|
||
for file in path.rglob("*.txt"):
|
||
docs.append(file.read_text(encoding="utf-8"))
|
||
for file in path.rglob("*.md"):
|
||
docs.append(file.read_text(encoding="utf-8"))
|
||
|
||
if not docs:
|
||
print("No documents found in", directory)
|
||
return
|
||
|
||
# Split all documents into chunks
|
||
chunks = splitter.split_text("\n\n".join(docs))
|
||
# Create LangChain Document objects
|
||
from langchain.docstore.document import Document
|
||
|
||
documents = [Document(page_content=chunk) for chunk in chunks]
|
||
vectorstore.add_documents(documents)
|
||
vectorstore.persist()
|
||
print(f"Added {len(documents)} chunks to the vector store.")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Example usage (uncomment to run manually)
|
||
# ---------------------------------------------------------------------------
|
||
# if __name__ == "__main__":
|
||
# store = create_vectorstore()
|
||
# load_documents("documents", store)
|