Update vectorstore.py

vectorstore.py

@@ -1,51 +1,94 @@
-"""Vector store utilities for ChromaDB.
+"""Vector store utilities for ChromaDB with Ollama embeddings.
 
-This module provides functions to create a persistent ChromaDB vector store using
-Ollama embeddings and to load documents from a directory into the store.
+This module provides functions to create a persistent Chroma vector store and
+load documents from a directory into it. Documents are split into chunks using
+`RecursiveCharacterTextSplitter` and stored in the Chroma collection.
 """
 
+import os
 from pathlib import Path
+from typing import List
 
+from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_chroma import Chroma
 from langchain_ollama import OllamaEmbeddings
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_core.documents import Document
 
+# ---------------------------------------------------------------------------
+# Vector store creation
+# ---------------------------------------------------------------------------
 
 def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
-    """Create a Chroma vector store with Ollama embeddings.
+    """Create or load a Chroma vector store.
 
     Parameters
     ----------
     persist_directory: str
-        Directory where the vector store will be persisted.
+        Directory where the Chroma DB files are stored.
 
     Returns
     -------
     Chroma
         A Chroma vector store instance.
     """
+    # Ensure directory exists
+    Path(persist_directory).mkdir(parents=True, exist_ok=True)
+    # Use Ollama embeddings
     embeddings = OllamaEmbeddings(model="nomic-embed-text")
-    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+    # Create Chroma store
+    vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+    return vectorstore
+
+# ---------------------------------------------------------------------------
+# Document loading
+# ---------------------------------------------------------------------------
+
+def _load_text_files(directory: str) -> List[str]:
+    """Load all .txt and .md files from a directory into a list of strings."""
+    texts = []
+    for root, _, files in os.walk(directory):
+        for file in files:
+            if file.lower().endswith(('.txt', '.md')):
+                path = Path(root) / file
+                try:
+                    content = path.read_text(encoding="utf-8")
+                    texts.append(content)
+                except Exception as e:
+                    print(f"Failed to read {path}: {e}")
+    return texts
 
 
-def load_documents(directory: str, vectorstore: Chroma) -> None:
-    """Load .txt and .md files from *directory* into *vectorstore*.
+def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
+    """Load documents from a directory into the provided vector store.
 
-    The documents are split into chunks using ``RecursiveCharacterTextSplitter``
-    before being added to the vector store.
+    Parameters
+    ----------
+    directory: str
+        Path to the directory containing .txt/.md files.
+    vectorstore: Chroma
+        The vector store to add documents to.
+    chunk_size: int, optional
+        Maximum size of each chunk.
+    chunk_overlap: int, optional
+        Number of characters to overlap between chunks.
     """
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    texts = _load_text_files(directory)
+    if not texts:
+        print("No text files found in the directory.")
+        return
+
+    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
     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_documents(docs)
-        print(f"Loaded {len(docs)} chunks from {directory} into ChromaDB.")
-    else:
-        print(f"No .txt/.md files found in {directory}.")
+    for text in texts:
+        docs.extend(splitter.split_text(text))
+
+    # Add documents to Chroma
+    vectorstore.add_texts(docs)
+    print(f"Loaded {len(docs)} chunks into the vector store.")
+
+# ---------------------------------------------------------------------------
+# Example usage (uncomment to run directly)
+# ---------------------------------------------------------------------------
+# if __name__ == "__main__":
+#     store = create_vectorstore()
+#     load_documents("documents", store)
+""