diff --git a/vectorstore.py b/vectorstore.py index 510bef0..d45593d 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,14 +1,15 @@ """Utilities for creating and populating a ChromaDB vector store. -This module provides two functions: +This module contains two helper functions: -- :func:`create_vectorstore` – creates a Chroma vector store backed by a local directory. -- :func:`load_documents` – reads all ``.txt`` and ``.md`` files from a directory, splits them into chunks using - :class:`langchain_text_splitters.RecursiveCharacterTextSplitter`, and upserts the chunks into the - provided vector store. +* :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 persistent across runs – the ``persist_directory`` argument defaults to -``"./chroma_db"``. +The vector store is persisted in ``./chroma_db`` by default, so it survives program +restarts. """ from pathlib import Path @@ -18,46 +19,75 @@ from langchain_chroma import Chroma from langchain_ollama import OllamaEmbeddings from langchain_text_splitters import RecursiveCharacterTextSplitter -# Default embedding model used by the vector store -EMBEDDING_MODEL = "nomic-embed-text" - +# --------------------------------------------------------------------------- +# Vector store creation +# --------------------------------------------------------------------------- def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: - """Create a Chroma vector store. + """Create a Chroma vector store backed by Ollama embeddings. Parameters ---------- persist_directory: str - Directory where the vector store will be persisted. + Path to the directory where the Chroma DB will be stored. Returns ------- Chroma A Chroma vector store instance. """ - embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) - return Chroma(persist_directory=persist_directory, embedding_function=embeddings) + embeddings = OllamaEmbeddings(model="nomic-embed-text") + return Chroma( + persist_directory=persist_directory, + embedding_function=embeddings, + ) +# --------------------------------------------------------------------------- +# Document ingestion +# --------------------------------------------------------------------------- -def _read_text_files(directory: str) -> Iterable[str]: - """Yield the content of all ``.txt`` and ``.md`` files in *directory*. +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) - for file_path in path.rglob("*.txt"): - yield file_path.read_text(encoding="utf-8") - for file_path in path.rglob("*.md"): - yield file_path.read_text(encoding="utf-8") + 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")) -def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None: - """Load documents from *directory* into *vectorstore*. + if not docs: + print("No documents found in", directory) + return - The documents are split into chunks using :class:`RecursiveCharacterTextSplitter` and then - upserted into the vector store. - """ - splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) - for text in _read_text_files(directory): - chunks = splitter.split_text(text) - vectorstore.add_texts(chunks) + # Split all documents into chunks + chunks = splitter.split_text("\n\n".join(docs)) + # Create LangChain Document objects + from langchain.docstore.document import Document -# End of vectorstore.py + 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)