diff --git a/vectorstore.py b/vectorstore.py index 3988cf4..510bef0 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,82 +1,63 @@ -""" -Vector store utilities for ChromaDB. +"""Utilities for creating and populating a ChromaDB vector store. + +This module provides two 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. + +The vector store is persistent across runs – the ``persist_directory`` argument defaults to +``"./chroma_db"``. """ -import os from pathlib import Path +from typing import Iterable from langchain_chroma import Chroma from langchain_ollama import OllamaEmbeddings from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_community.document_loaders import TextLoader, UnstructuredMarkdownLoader -# --------------------------------------------------------------------------- -# Create a persistent Chroma vector store -# --------------------------------------------------------------------------- +# Default embedding model used by the vector store +EMBEDDING_MODEL = "nomic-embed-text" + def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: - """Create or load a Chroma vector store. + """Create a Chroma vector store. Parameters ---------- persist_directory: str - Directory where the Chroma DB will be persisted. + Directory where the vector store will be persisted. Returns ------- Chroma A Chroma vector store instance. """ - os.makedirs(persist_directory, exist_ok=True) - embeddings = OllamaEmbeddings(model="nomic-embed-text") + embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) return Chroma(persist_directory=persist_directory, embedding_function=embeddings) -# --------------------------------------------------------------------------- -# Load documents from a directory and add them to the vector store -# --------------------------------------------------------------------------- -def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 500, chunk_overlap: int = 50) -> None: - """Load `.txt` and `.md` files from *directory*, split them into chunks, and add to *vectorstore*. - - Parameters - ---------- - directory: str - Directory containing the source documents. - vectorstore: Chroma - The vector store to which documents will be added. - chunk_size: int, optional - Maximum chunk size in characters. - chunk_overlap: int, optional - Number of overlapping characters between consecutive chunks. +def _read_text_files(directory: str) -> Iterable[str]: + """Yield the content of all ``.txt`` and ``.md`` files in *directory*. + """ + 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") + + +def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None: + """Load documents from *directory* into *vectorstore*. + + The documents are split into chunks using :class:`RecursiveCharacterTextSplitter` and then + upserted into the vector store. """ - loader_classes = { - ".txt": TextLoader, - ".md": UnstructuredMarkdownLoader, - } 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) - docs = [] - for root, _, files in os.walk(directory): - for file in files: - ext = Path(file).suffix.lower() - if ext not in loader_classes: - continue - loader = loader_classes[ext](os.path.join(root, file)) - loaded_docs = loader.load() - docs.extend(loaded_docs) - - if not docs: - return - - # Split documents into smaller chunks - split_docs = splitter.split_documents(docs) - vectorstore.add_documents(split_docs) - -# --------------------------------------------------------------------------- -# Example usage -# --------------------------------------------------------------------------- -if __name__ == "__main__": - # This block is only executed when running the module directly. - store = create_vectorstore() - load_documents("documents", store) - print("Vector store populated.") +# End of vectorstore.py