Update vectorstore.py

vectorstore.py

@@ -1,25 +1,29 @@
-"""
-Vector store utilities using ChromaDB and Ollama embeddings.
+"""Vector store utilities using ChromaDB and Ollama embeddings.
+
+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_ollama import ChatOllama
 
-# Create the vector store with persistence
+# Default persistence directory
+DEFAULT_PERSIST_DIR = "./chroma_db"
 
-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 persisted.
+        Directory where the Chroma DB will be stored.
 
     Returns
     -------
@@ -29,53 +33,29 @@ def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
     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) -> None:
-    """Load .txt and .md files from *directory*, chunk them, and add to *vectorstore*.
+def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
+    """Load all .txt and .md files from *directory* into *vectorstore*.
 
-    Parameters
-    ----------
-    directory: str
-        Path to the directory containing the documents.
-    vectorstore: Chroma
-        The vector store to which the documents will be added.
+    The files are read, split into chunks with a recursive character splitter and
+    added to the Chroma collection.
     """
-    # Ensure the directory exists
-    path = Path(directory)
-    if not path.is_dir():
-        raise FileNotFoundError(f"Directory {directory} does not exist")
+    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    docs: List[str] = []
+    for file_path in Path(directory).rglob("*.txt"):
+        docs.append(file_path.read_text(encoding="utf-8"))
+    for file_path in Path(directory).rglob("*.md"):
+        docs.append(file_path.read_text(encoding="utf-8"))
 
-    # Collect all .txt and .md files
-    files = list(path.rglob("*.txt")) + list(path.rglob("*.md"))
-    if not files:
-        print(f"No .txt or .md files found in {directory}")
+    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)
 
-    # Read and chunk the documents
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    documents = []
-    for file_path in files:
-        try:
-            content = file_path.read_text(encoding="utf-8")
-        except Exception as e:
-            print(f"Failed to read {file_path}: {e}")
-            continue
-        # Split into chunks
-        chunks = text_splitter.split_text(content)
-        for i, chunk in enumerate(chunks):
-            documents.append({
-                "page_content": chunk,
-                "metadata": {"source": str(file_path), "chunk_index": i},
-            })
+    # Persist changes
+    vectorstore.persist()
 
-    if documents:
-        vectorstore.add_documents(documents)
-        print(f"Added {len(documents)} chunks to the vector store from {directory}")
-    else:
-        print("No documents were processed.")
-
-# Example usage:
-# if __name__ == "__main__":
-#     store = create_vectorstore()
-#     load_documents("./documents", store)
+    print(f"Loaded {len(texts)} chunks into ChromaDB.")