Update vectorstore.py

vectorstore.py

@@ -1,82 +1,63 @@
-"""
-Vector store utilities for ChromaDB.
+"""Utilities for creating and populating a ChromaDB vector store.
+
+This module provides two functions:
+
+- :func:`create_vectorstore` – creates a Chroma vector store backed by a local directory.
+- :func:`load_documents` – reads all ``.txt`` and ``.md`` files from a directory, splits them into chunks using
+  :class:`langchain_text_splitters.RecursiveCharacterTextSplitter`, and upserts the chunks into the
+  provided vector store.
+
+The vector store is persistent across runs – the ``persist_directory`` argument defaults to
+``"./chroma_db"``.
 """
 
-import os
 from pathlib import Path
+from typing import Iterable
 
 from langchain_chroma import Chroma
 from langchain_ollama import OllamaEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_community.document_loaders import TextLoader, UnstructuredMarkdownLoader
 
-# ---------------------------------------------------------------------------
-# Create a persistent Chroma vector store
-# ---------------------------------------------------------------------------
+# Default embedding model used by the vector store
+EMBEDDING_MODEL = "nomic-embed-text"
+
 
 def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
-    """Create or load a Chroma vector store.
+    """Create a Chroma vector store.
 
     Parameters
     ----------
     persist_directory: str
-        Directory where the Chroma DB will be persisted.
+        Directory where the vector store will be persisted.
 
     Returns
     -------
     Chroma
         A Chroma vector store instance.
     """
-    os.makedirs(persist_directory, exist_ok=True)
-    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+    embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
     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 = 500, chunk_overlap: int = 50) -> None:
-    """Load `.txt` and `.md` files from *directory*, split them into chunks, and add to *vectorstore*.
-
-    Parameters
-    ----------
-    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.
+def _read_text_files(directory: str) -> Iterable[str]:
+    """Yield the content of all ``.txt`` and ``.md`` files in *directory*.
+    """
+    path = Path(directory)
+    for file_path in path.rglob("*.txt"):
+        yield file_path.read_text(encoding="utf-8")
+    for file_path in path.rglob("*.md"):
+        yield file_path.read_text(encoding="utf-8")
+
+
+def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
+    """Load documents from *directory* into *vectorstore*.
+
+    The documents are split into chunks using :class:`RecursiveCharacterTextSplitter` and then
+    upserted into the vector store.
     """
-    loader_classes = {
-        ".txt": TextLoader,
-        ".md": UnstructuredMarkdownLoader,
-    }
     splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    for text in _read_text_files(directory):
+        chunks = splitter.split_text(text)
+        vectorstore.add_texts(chunks)
 
-    docs = []
-    for root, _, files in os.walk(directory):
-        for file in files:
-            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 smaller chunks
-    split_docs = splitter.split_documents(docs)
-    vectorstore.add_documents(split_docs)
-
-# ---------------------------------------------------------------------------
-# 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.")
+# End of vectorstore.py