Update vectorstore.py

vectorstore.py

@@ -1,61 +1,82 @@
-"""Vector store utilities using ChromaDB and Ollama embeddings.
+"""
-
+Vector store utilities for ChromaDB.
-This module provides functions to create a persistent Chroma vector store and to load documents
-from a directory into the store. Documents are split into chunks using
-`RecursiveCharacterTextSplitter`.
 """
 
 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_community.document_loaders import TextLoader, UnstructuredMarkdownLoader
 
-# Default persistence directory
+# ---------------------------------------------------------------------------
-DEFAULT_PERSIST_DIR = "./chroma_db"
+# Create a persistent Chroma vector store
+# ---------------------------------------------------------------------------
 
-
+def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
-def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma:
     """Create or load a Chroma vector store.
 
     Parameters
     ----------
     persist_directory: str
-        Directory where the Chroma DB will be stored.
+        Directory where the Chroma DB will be persisted.
 
     Returns
     -------
     Chroma
-        The Chroma vector store instance.
+        A Chroma vector store instance.
     """
+    os.makedirs(persist_directory, exist_ok=True)
     embeddings = OllamaEmbeddings(model="nomic-embed-text")
     return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
 
+# ---------------------------------------------------------------------------
+# Load documents from a directory and add them to the vector store
+# ---------------------------------------------------------------------------
 
-def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
+def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 500, chunk_overlap: int = 50) -> None:
-    """Load all .txt and .md files from *directory* into *vectorstore*.
+    """Load `.txt` and `.md` files from *directory*, split them into chunks, and add to *vectorstore*.
 
-    The files are read, split into chunks with a recursive character splitter and
+    Parameters
-    added to the Chroma collection.
+    ----------
+    directory: str
+        Directory containing the source documents.
+    vectorstore: Chroma
+        The vector store to which documents will be added.
+    chunk_size: int, optional
+        Maximum chunk size in characters.
+    chunk_overlap: int, optional
+        Number of overlapping characters between consecutive chunks.
     """
+    loader_classes = {
+        ".txt": TextLoader,
+        ".md": UnstructuredMarkdownLoader,
+    }
     splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
-    docs: List[str] = []
+
-    for file_path in Path(directory).rglob("*.txt"):
+    docs = []
-        docs.append(file_path.read_text(encoding="utf-8"))
+    for root, _, files in os.walk(directory):
-    for file_path in Path(directory).rglob("*.md"):
+        for file in files:
-        docs.append(file_path.read_text(encoding="utf-8"))
+            ext = Path(file).suffix.lower()
+            if ext not in loader_classes:
+                continue
+            loader = loader_classes[ext](os.path.join(root, file))
+            loaded_docs = loader.load()
+            docs.extend(loaded_docs)
 
     if not docs:
         return
-    # Split documents into chunks
-    texts = splitter.split_text("\n\n".join(docs))
-    # Create a list of dicts with metadata (optional)
-    metadatas = [{"source": "local"} for _ in texts]
-    vectorstore.add_texts(texts, metadatas=metadatas)
 
-    # Persist changes
+    # Split documents into smaller chunks
-    vectorstore.persist()
+    split_docs = splitter.split_documents(docs)
+    vectorstore.add_documents(split_docs)
 
-    print(f"Loaded {len(texts)} chunks into ChromaDB.")
+# ---------------------------------------------------------------------------
+# Example usage
+# ---------------------------------------------------------------------------
+if __name__ == "__main__":
+    # This block is only executed when running the module directly.
+    store = create_vectorstore()
+    load_documents("documents", store)
+    print("Vector store populated.")