Update vectorstore.py

vectorstore.py

@@ -1,77 +1,79 @@
-"""
+"""Utilities for creating and populating a Chroma vector store.
-Vector store utilities for the RAG agent.
 
-Provides functions to create a ChromaDB vector store backed by Ollama embeddings
+This module provides two helper functions:
-and to load documents from a directory into the store.
+
+* ``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
+def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
-    not be created because Chroma will deduplicate based on the content hash.
+    """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.
+        The Chroma collection to populate.
-    chunk_size: int, optional
+    chunk_size: int
-        Maximum number of characters per chunk.
+        Number of characters per chunk.
-    chunk_overlap: int, optional
+    chunk_overlap: int
-        Number of characters to overlap between consecutive chunks.
+        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("**/*"):
+    for path in Path(directory).glob("**/*"):
-        if file_path.suffix.lower() not in {".txt", ".md"}:
+        if path.is_file() and path.suffix.lower() in {".txt", ".md"}:
-            continue
+            text = path.read_text(encoding="utf-8")
-        text = file_path.read_text(encoding="utf-8")
+            docs.extend(splitter.split_text(text))
-        docs.extend(splitter.split_text(text))
 
-    # Convert list of strings to list of Documents
+    if docs:
-    documents = [Document(page_content=chunk) for chunk in docs]
+        vectorstore.add_texts(docs)
-
+    else:
-    if documents:
+        print("[vectorstore] No documents found in", directory)
-        vectorstore.add_documents(documents)
-        vectorstore.persist()
 
 # ---------------------------------------------------------------------------
-# End of module
+# Example usage (uncomment to run as a script)
 # ---------------------------------------------------------------------------
+# if __name__ == "__main__":
+#     store = create_vectorstore()
+#     load_documents("documents", store)
+#     print("Vector store populated.")
+"""