Update vectorstore.py

vectorstore.py

@@ -1,94 +1,61 @@
-"""Vector store utilities for ChromaDB with Ollama embeddings.
+"""Vector store utilities using ChromaDB and Ollama embeddings.
 
-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.
+This module provides functions to create a persistent Chroma vector store
+and to load documents from a directory into the store.
 """
 
-import os
 from pathlib import Path
-from typing import List
+from typing import Iterable
 
-from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_chroma import Chroma
 from langchain_ollama import OllamaEmbeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
 
 # ---------------------------------------------------------------------------
-# Vector store creation
+# Create a persistent Chroma vector store.
 # ---------------------------------------------------------------------------
 
 def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
-    """Create or load a Chroma vector store.
+    """Create a Chroma vector store with Ollama embeddings.
 
     Parameters
     ----------
     persist_directory: str
-        Directory where the Chroma DB files are stored.
+        Directory where the Chroma database will be 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")
-    # Create Chroma store
-    vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
-    return vectorstore
+    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
 
 # ---------------------------------------------------------------------------
-# Document loading
+# Load documents from a directory into the vector store.
 # ---------------------------------------------------------------------------
 
-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 all .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.
-
-    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.
+    The documents are split into chunks using a RecursiveCharacterTextSplitter
+    before being added to the vector store.
     """
-    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)
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
     docs = []
-    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.")
+    for path in Path(directory).rglob("*.txt"):
+        docs.append(path.read_text(encoding="utf-8"))
+    for path in Path(directory).rglob("*.md"):
+        docs.append(path.read_text(encoding="utf-8"))
+    if not docs:
+        return
+    # Split the documents into chunks.
+    chunks = splitter.split_text("\n\n".join(docs))
+    # Create LangChain Document objects.
+    from langchain.docstore.document import Document
+    documents = [Document(page_content=chunk) for chunk in chunks]
+    vectorstore.add_documents(documents)
 
 # ---------------------------------------------------------------------------
-# Example usage (uncomment to run directly)
-# ---------------------------------------------------------------------------
-# if __name__ == "__main__":
-#     store = create_vectorstore()
-#     load_documents("documents", store)
-""
+# End of vectorstore.py
+# ---------------------------------------------------------------------------