diff --git a/vectorstore.py b/vectorstore.py index d140f20..2d83e06 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,77 +1,79 @@ -""" -Vector store utilities for the RAG agent. +"""Utilities for creating and populating a Chroma vector store. -Provides functions to create a ChromaDB vector store backed by Ollama embeddings -and to load documents from a directory into the store. +This module provides two helper functions: + +* ``create_vectorstore`` – creates a Chroma collection backed by Ollama embeddings. +* ``load_documents`` – reads ``.txt``/``.md`` files, splits them into chunks and adds them to the collection. + +The vector store is persisted in ``./chroma_db`` by default. """ from pathlib import Path from typing import List -from langchain_ollama import OllamaEmbeddings from langchain_chroma import Chroma +from langchain_ollama import OllamaEmbeddings from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain.docstore.document import Document # --------------------------------------------------------------------------- -# Configuration constants +# Configuration # --------------------------------------------------------------------------- -DEFAULT_EMBEDDING_MODEL = "nomic-embed-text" DEFAULT_PERSIST_DIR = "./chroma_db" +EMBEDDING_MODEL = "nomic-embed-text" # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma: - """Create (or load) a Chroma vector store. + """Create a Chroma vector store with Ollama embeddings. Parameters ---------- persist_directory: str - Directory where the Chroma database will be persisted. + Directory where the vector store will be persisted. Returns ------- Chroma - An instance of the Chroma vector store. + A Chroma collection ready for adding documents. """ - embeddings = OllamaEmbeddings(model=DEFAULT_EMBEDDING_MODEL) + embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) return Chroma(persist_directory=persist_directory, embedding_function=embeddings) -def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None: - """Load all .txt and .md files from *directory*, chunk them and add to *vectorstore*. - The function is idempotent – if the same files are loaded again, duplicates will - not be created because Chroma will deduplicate based on the content hash. +def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None: + """Load all ``.txt`` and ``.md`` files from *directory*, split them into chunks and add to *vectorstore*. Parameters ---------- directory: str - Path to the folder containing the documents. + Path to the folder containing documents. vectorstore: Chroma - The vector store to populate. - chunk_size: int, optional - Maximum number of characters per chunk. - chunk_overlap: int, optional - Number of characters to overlap between consecutive chunks. + The Chroma collection to populate. + chunk_size: int + Number of characters per chunk. + chunk_overlap: int + Number of characters that overlap between consecutive chunks. """ splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) - docs: List[Document] = [] + docs: List[str] = [] - for file_path in Path(directory).glob("**/*"): - if file_path.suffix.lower() not in {".txt", ".md"}: - continue - text = file_path.read_text(encoding="utf-8") - docs.extend(splitter.split_text(text)) + for path in Path(directory).glob("**/*"): + if path.is_file() and path.suffix.lower() in {".txt", ".md"}: + text = path.read_text(encoding="utf-8") + docs.extend(splitter.split_text(text)) - # Convert list of strings to list of Documents - documents = [Document(page_content=chunk) for chunk in docs] - - if documents: - vectorstore.add_documents(documents) - vectorstore.persist() + if docs: + vectorstore.add_texts(docs) + else: + print("[vectorstore] No documents found in", directory) # --------------------------------------------------------------------------- -# End of module -# --------------------------------------------------------------------------- \ No newline at end of file +# Example usage (uncomment to run as a script) +# --------------------------------------------------------------------------- +# if __name__ == "__main__": +# store = create_vectorstore() +# load_documents("documents", store) +# print("Vector store populated.") +"""