Update vectorstore.py

vectorstore.py

@@ -1,14 +1,15 @@
 """Utilities for creating and populating a ChromaDB vector store.
 
-This module provides two functions:
+This module contains two helper functions:
 
-- :func:`create_vectorstore` – creates a Chroma vector store backed by a local directory.
+* :func:`create_vectorstore` – returns a :class:`langchain_chroma.Chroma` instance backed by
-- :func:`load_documents` – reads all ``.txt`` and ``.md`` files from a directory, splits them into chunks using
+  an ``OllamaEmbeddings`` model.
-  :class:`langchain_text_splitters.RecursiveCharacterTextSplitter`, and upserts the chunks into the
+* :func:`load_documents` – reads ``.txt``/``.md`` files from a directory, splits them into
-  provided vector store.
+  chunks using :class:`langchain_text_splitters.RecursiveCharacterTextSplitter`, and adds
+  the chunks to the vector store.
 
-The vector store is persistent across runs – the ``persist_directory`` argument defaults to
+The vector store is persisted in ``./chroma_db`` by default, so it survives program
-``"./chroma_db"``.
+restarts.
 """
 
 from pathlib import Path
@@ -18,46 +19,75 @@ from langchain_chroma import Chroma
 from langchain_ollama import OllamaEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
-# Default embedding model used by the vector store
+# ---------------------------------------------------------------------------
-EMBEDDING_MODEL = "nomic-embed-text"
+# Vector store creation
-
+# ---------------------------------------------------------------------------
 
 def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
-    """Create a Chroma vector store.
+    """Create a Chroma vector store backed by Ollama embeddings.
 
     Parameters
     ----------
     persist_directory: str
-        Directory where the vector store will be persisted.
+        Path to the directory where the Chroma DB will be stored.
 
     Returns
     -------
     Chroma
         A Chroma vector store instance.
     """
-    embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
-    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+    return Chroma(
+        persist_directory=persist_directory,
+        embedding_function=embeddings,
+    )
 
+# ---------------------------------------------------------------------------
+# Document ingestion
+# ---------------------------------------------------------------------------
 
-def _read_text_files(directory: str) -> Iterable[str]:
+def load_documents(directory: str | Path, vectorstore: Chroma) -> None:
-    """Yield the content of all ``.txt`` and ``.md`` files in *directory*.
+    """Load all ``.txt`` and ``.md`` files from *directory* into *vectorstore*.
+
+    The files are split into chunks using
+    :class:`langchain_text_splitters.RecursiveCharacterTextSplitter` before being
+    added to the vector store.
+
+    Parameters
+    ----------
+    directory: str | Path
+        Directory containing the documents.
+    vectorstore: Chroma
+        The vector store to populate.
     """
     path = Path(directory)
-    for file_path in path.rglob("*.txt"):
+    if not path.is_dir():
-        yield file_path.read_text(encoding="utf-8")
+        raise ValueError(f"{directory!r} is not a directory")
-    for file_path in path.rglob("*.md"):
-        yield file_path.read_text(encoding="utf-8")
 
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    docs = []
+    for file in path.rglob("*.txt"):
+        docs.append(file.read_text(encoding="utf-8"))
+    for file in path.rglob("*.md"):
+        docs.append(file.read_text(encoding="utf-8"))
 
-def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
+    if not docs:
-    """Load documents from *directory* into *vectorstore*.
+        print("No documents found in", directory)
+        return
 
-    The documents are split into chunks using :class:`RecursiveCharacterTextSplitter` and then
+    # Split all documents into chunks
-    upserted into the vector store.
+    chunks = splitter.split_text("\n\n".join(docs))
-    """
+    # Create LangChain Document objects
-    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    from langchain.docstore.document import Document
-    for text in _read_text_files(directory):
-        chunks = splitter.split_text(text)
-        vectorstore.add_texts(chunks)
 
-# End of vectorstore.py
+    documents = [Document(page_content=chunk) for chunk in chunks]
+    vectorstore.add_documents(documents)
+    vectorstore.persist()
+    print(f"Added {len(documents)} chunks to the vector store.")
+
+# ---------------------------------------------------------------------------
+# Example usage (uncomment to run manually)
+# ---------------------------------------------------------------------------
+# if __name__ == "__main__":
+#     store = create_vectorstore()
+#     load_documents("documents", store)