Update vectorstore.py

vectorstore.py

@@ -1,53 +1,78 @@
-"""
+"""Vector store utilities using ChromaDB and Ollama embeddings.
-Vector store utilities for ChromaDB.
+
+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.
 """
 
 from pathlib import Path
+from typing import Iterable
 
-from langchain_ollama import OllamaEmbeddings
 from langchain_chroma import Chroma
+from langchain_ollama import OllamaEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
+# Default persistence directory
+DEFAULT_PERSIST_DIR = "./chroma_db"
 
-def create_vectorstore(persist_directory: str = "./chroma_db"):
+# Default chunking parameters – these can be tuned
+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 Chroma database is persisted.
+        Directory where the ChromaDB files are stored.
 
     Returns
     -------
     Chroma
-        The Chroma vector store instance.
+        An instance of the Chroma vector store backed by the given directory.
     """
     embeddings = OllamaEmbeddings(model="nomic-embed-text")
-    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+    vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+    return vectorstore
 
 
-def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200):
+def _load_text_files(directory: str | Path) -> Iterable[str]:
-    """Read all .txt and .md files from a directory, split them into chunks and add to the vectorstore.
+    """Yield the contents of all .txt and .md files in *directory*.
 
     Parameters
     ----------
-    directory: str
+    directory: str | Path
-        Path to the folder containing documents.
+        Path to the directory containing the documents.
-    vectorstore: Chroma
-        The vector store to which documents will be added.
-    chunk_size: int
-        Maximum number of characters per chunk.
-    chunk_overlap: int
-        Number of overlapping characters between consecutive chunks.
     """
-    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    directory = Path(directory)
-    docs = []
+    for file_path in directory.rglob("*.txt"):
-    for file_path in Path(directory).rglob("*.txt"):
+        yield file_path.read_text(encoding="utf-8")
-        docs.append(file_path.read_text(encoding="utf-8"))
+    for file_path in directory.rglob("*.md"):
-    for file_path in Path(directory).rglob("*.md"):
+        yield file_path.read_text(encoding="utf-8")
-        docs.append(file_path.read_text(encoding="utf-8"))
+
-    if not docs:
+
-        return
+def load_documents(directory: str | Path, vectorstore: Chroma) -> None:
-    texts = splitter.split_text("\n\n".join(docs))
+    """Load documents from *directory* into *vectorstore*.
-    vectorstore.add_texts(texts)
+
+    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.
+    """
+    splitter = RecursiveCharacterTextSplitter(chunk_size=DEFAULT_CHUNK_SIZE, chunk_overlap=DEFAULT_CHUNK_OVERLAP)
+    for text in _load_text_files(directory):
+        chunks = splitter.split_text(text)
+        vectorstore.add_texts(chunks)
+
+    # Persist changes
     vectorstore.persist()
+
+# End of vectorstore.py