From d6e39c7557f7ba0483eda78c24a8c65183ff7612 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A0=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= Date: Tue, 2 Jun 2026 07:47:23 +0000 Subject: [PATCH] Update vectorstore.py --- vectorstore.py | 104 +++++++++++++++++++++---------------------------- 1 file changed, 44 insertions(+), 60 deletions(-) diff --git a/vectorstore.py b/vectorstore.py index d45593d..d140f20 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,93 +1,77 @@ -"""Utilities for creating and populating a ChromaDB vector store. +""" +Vector store utilities for the RAG agent. -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. +Provides functions to create a ChromaDB vector store backed by Ollama embeddings +and to load documents from a directory into the store. """ from pathlib import Path -from typing import Iterable +from typing import List -from langchain_chroma import Chroma from langchain_ollama import OllamaEmbeddings +from langchain_chroma import Chroma from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain.docstore.document import Document # --------------------------------------------------------------------------- -# Vector store creation +# Configuration constants +# --------------------------------------------------------------------------- +DEFAULT_EMBEDDING_MODEL = "nomic-embed-text" +DEFAULT_PERSIST_DIR = "./chroma_db" + +# --------------------------------------------------------------------------- +# Public API # --------------------------------------------------------------------------- -def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: - """Create a Chroma vector store backed by Ollama embeddings. +def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma: + """Create (or load) a Chroma vector store. Parameters ---------- persist_directory: str - Path to the directory where the Chroma DB will be stored. + Directory where the Chroma database will be persisted. Returns ------- Chroma - A Chroma vector store instance. + An instance of the Chroma vector store. """ - embeddings = OllamaEmbeddings(model="nomic-embed-text") - return Chroma( - persist_directory=persist_directory, - embedding_function=embeddings, - ) + embeddings = OllamaEmbeddings(model=DEFAULT_EMBEDDING_MODEL) + return Chroma(persist_directory=persist_directory, embedding_function=embeddings) -# --------------------------------------------------------------------------- -# Document ingestion -# --------------------------------------------------------------------------- +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*. -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. + 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. Parameters ---------- - directory: str | Path - Directory containing the documents. + directory: str + Path to the folder containing the 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. """ - path = Path(directory) - if not path.is_dir(): - raise ValueError(f"{directory!r} is not a directory") + splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) + docs: List[Document] = [] - 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")) + 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)) - if not docs: - print("No documents found in", directory) - return + # Convert list of strings to list of Documents + documents = [Document(page_content=chunk) for chunk in docs] - # 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.") + if documents: + vectorstore.add_documents(documents) + vectorstore.persist() # --------------------------------------------------------------------------- -# Example usage (uncomment to run manually) -# --------------------------------------------------------------------------- -# if __name__ == "__main__": -# store = create_vectorstore() -# load_documents("documents", store) +# End of module +# --------------------------------------------------------------------------- \ No newline at end of file