Update vectorstore.py

vectorstore.py

@@ -1,78 +1,81 @@
-"""Vector store utilities using ChromaDB and Ollama embeddings.
+"""
-
+Vector store utilities using ChromaDB and Ollama embeddings.
-This module provides functions to create a persistent ChromaDB vector store and
-load documents from a directory into it. The documents are split into
-manageable chunks using a recursive character text splitter.
 """
 
+import os
 from pathlib import Path
-from typing import Iterable
+from typing import List
 
 from langchain_chroma import Chroma
 from langchain_ollama import OllamaEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_ollama import ChatOllama
 
-# Default persistence directory
+# Create the vector store with persistence
-DEFAULT_PERSIST_DIR = "./chroma_db"
 
-# Default chunking parameters – these can be tuned
+def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
-DEFAULT_CHUNK_SIZE = 1000
-DEFAULT_CHUNK_OVERLAP = 200
-
-
-def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma:
     """Create or load a Chroma vector store.
 
     Parameters
     ----------
     persist_directory: str
-        Directory where the ChromaDB files are stored.
+        Directory where the Chroma DB will be persisted.
 
     Returns
     -------
     Chroma
-        An instance of the Chroma vector store backed by the given directory.
+        The Chroma vector store instance.
     """
     embeddings = OllamaEmbeddings(model="nomic-embed-text")
-    vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
-    return vectorstore
 
+# Load documents from a directory and add them to the vector store
 
-def _load_text_files(directory: str | Path) -> Iterable[str]:
+def load_documents(directory: str, vectorstore: Chroma) -> None:
-    """Yield the contents of all .txt and .md files in *directory*.
+    """Load .txt and .md files from *directory*, chunk them, and add to *vectorstore*.
 
     Parameters
     ----------
-    directory: str | Path
+    directory: str
         Path to the directory containing the documents.
-    """
-    directory = Path(directory)
-    for file_path in directory.rglob("*.txt"):
-        yield file_path.read_text(encoding="utf-8")
-    for file_path in directory.rglob("*.md"):
-        yield file_path.read_text(encoding="utf-8")
-
-
-def load_documents(directory: str | Path, vectorstore: Chroma) -> None:
-    """Load documents from *directory* into *vectorstore*.
-
-    The documents are split into chunks using a recursive character splitter
-    and then added to the vector store. Existing documents are not removed –
-    this function simply appends new data.
-
-    Parameters
-    ----------
-    directory: str | Path
-        Directory containing the source documents.
     vectorstore: Chroma
-        The vector store to populate.
+        The vector store to which the documents will be added.
     """
-    splitter = RecursiveCharacterTextSplitter(chunk_size=DEFAULT_CHUNK_SIZE, chunk_overlap=DEFAULT_CHUNK_OVERLAP)
+    # Ensure the directory exists
-    for text in _load_text_files(directory):
+    path = Path(directory)
-        chunks = splitter.split_text(text)
+    if not path.is_dir():
-        vectorstore.add_texts(chunks)
+        raise FileNotFoundError(f"Directory {directory} does not exist")
 
-    # Persist changes
+    # Collect all .txt and .md files
-    vectorstore.persist()
+    files = list(path.rglob("*.txt")) + list(path.rglob("*.md"))
+    if not files:
+        print(f"No .txt or .md files found in {directory}")
+        return
 
-# End of vectorstore.py
+    # 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},
+            })
+
+    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)