Update vectorstore.py

vectorstore.py

@@ -1,31 +1,18 @@
-"""Utilities for creating and populating a Chroma vector store.
+"""Vector store utilities for ChromaDB.
 
-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.
+This module provides functions to create a persistent ChromaDB vector store using
+Ollama embeddings and to load documents from a directory into the store.
 """
 
 from pathlib import Path
-from typing import List
 
 from langchain_chroma import Chroma
 from langchain_ollama import OllamaEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
 
-# ---------------------------------------------------------------------------
-# Configuration
-# ---------------------------------------------------------------------------
-DEFAULT_PERSIST_DIR = "./chroma_db"
-EMBEDDING_MODEL = "nomic-embed-text"
 
-# ---------------------------------------------------------------------------
-# Public API
-# ---------------------------------------------------------------------------
-
-def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma:
+def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
     """Create a Chroma vector store with Ollama embeddings.
 
     Parameters
@@ -36,44 +23,29 @@ def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma:
     Returns
     -------
     Chroma
-        A Chroma collection ready for adding documents.
+        A Chroma vector store instance.
     """
-    embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
     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*, split them into chunks and add to *vectorstore*.
+def load_documents(directory: str, vectorstore: Chroma) -> None:
+    """Load .txt and .md files from *directory* into *vectorstore*.
 
-    Parameters
-    ----------
-    directory: str
-        Path to the folder containing documents.
-    vectorstore: Chroma
-        The Chroma collection to populate.
-    chunk_size: int
-        Number of characters per chunk.
-    chunk_overlap: int
-        Number of characters that overlap between consecutive chunks.
+    The documents are split into chunks using ``RecursiveCharacterTextSplitter``
+    before being added to the vector store.
     """
-    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
-    docs: List[str] = []
-
-    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))
-
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    docs = []
+    for file_path in Path(directory).glob("*"):
+        if file_path.suffix.lower() not in {".txt", ".md"}:
+            continue
+        with open(file_path, "r", encoding="utf-8") as f:
+            text = f.read()
+        chunks = splitter.split_text(text)
+        docs.extend([Document(page_content=c, metadata={"source": str(file_path)}) for c in chunks])
     if docs:
-        vectorstore.add_texts(docs)
+        vectorstore.add_documents(docs)
+        print(f"Loaded {len(docs)} chunks from {directory} into ChromaDB.")
     else:
-        print("[vectorstore] No documents found in", directory)
-
-# ---------------------------------------------------------------------------
-# Example usage (uncomment to run as a script)
-# ---------------------------------------------------------------------------
-# if __name__ == "__main__":
-#     store = create_vectorstore()
-#     load_documents("documents", store)
-#     print("Vector store populated.")
-"""
+        print(f"No .txt/.md files found in {directory}.")