Update vectorstore.py

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
+# ---------------------------------------------------------------------------