diff --git a/vectorstore.py b/vectorstore.py index 2c6a67e..3d9c2fe 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,53 +1,78 @@ -""" -Vector store utilities for ChromaDB. +"""Vector store utilities using ChromaDB and Ollama embeddings. + +This module provides functions to create a persistent ChromaDB vector store and +load documents from a directory into it. The documents are split into +manageable chunks using a recursive character text splitter. """ from pathlib import Path +from typing import Iterable -from langchain_ollama import OllamaEmbeddings from langchain_chroma import Chroma +from langchain_ollama import OllamaEmbeddings from langchain_text_splitters import RecursiveCharacterTextSplitter +# Default persistence directory +DEFAULT_PERSIST_DIR = "./chroma_db" -def create_vectorstore(persist_directory: str = "./chroma_db"): +# Default chunking parameters – these can be tuned +DEFAULT_CHUNK_SIZE = 1000 +DEFAULT_CHUNK_OVERLAP = 200 + + +def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma: """Create or load a Chroma vector store. Parameters ---------- persist_directory: str - Directory where the Chroma database is persisted. + Directory where the ChromaDB files are stored. Returns ------- Chroma - The Chroma vector store instance. + An instance of the Chroma vector store backed by the given directory. """ embeddings = OllamaEmbeddings(model="nomic-embed-text") - return Chroma(persist_directory=persist_directory, embedding_function=embeddings) + vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings) + return vectorstore -def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200): - """Read all .txt and .md files from a directory, split them into chunks and add to the vectorstore. +def _load_text_files(directory: str | Path) -> Iterable[str]: + """Yield the contents of all .txt and .md files in *directory*. Parameters ---------- - directory: str - Path to the folder containing documents. - vectorstore: Chroma - The vector store to which documents will be added. - chunk_size: int - Maximum number of characters per chunk. - chunk_overlap: int - Number of overlapping characters between consecutive chunks. + directory: str | Path + Path to the directory containing the documents. """ - splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) - docs = [] - for file_path in Path(directory).rglob("*.txt"): - docs.append(file_path.read_text(encoding="utf-8")) - for file_path in Path(directory).rglob("*.md"): - docs.append(file_path.read_text(encoding="utf-8")) - if not docs: - return - texts = splitter.split_text("\n\n".join(docs)) - vectorstore.add_texts(texts) + directory = Path(directory) + for file_path in directory.rglob("*.txt"): + yield file_path.read_text(encoding="utf-8") + for file_path in directory.rglob("*.md"): + yield file_path.read_text(encoding="utf-8") + + +def load_documents(directory: str | Path, vectorstore: Chroma) -> None: + """Load documents from *directory* into *vectorstore*. + + The documents are split into chunks using a recursive character splitter + and then added to the vector store. Existing documents are not removed – + this function simply appends new data. + + Parameters + ---------- + directory: str | Path + Directory containing the source documents. + vectorstore: Chroma + The vector store to populate. + """ + splitter = RecursiveCharacterTextSplitter(chunk_size=DEFAULT_CHUNK_SIZE, chunk_overlap=DEFAULT_CHUNK_OVERLAP) + for text in _load_text_files(directory): + chunks = splitter.split_text(text) + vectorstore.add_texts(chunks) + + # Persist changes vectorstore.persist() + +# End of vectorstore.py