Update vectorstore.py

vectorstore.py

@@ -1,19 +1,20 @@
 """Vector store utilities using ChromaDB and Ollama embeddings.
 
-This module provides functions to create a persistent Chroma vector store
+This module provides functions to create a persistent Chroma vector store and load
-and to load documents from a directory into the store.
+text documents from a directory into it.  The store is exposed via the global
+``store`` variable so that other modules (e.g. tools) can access it.
 """
 
 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_core.documents import Document
 
-# ---------------------------------------------------------------------------
+# Global store that will be initialised in ``create_vectorstore``.
-# Create a persistent Chroma vector store.
+store: Chroma | None = None
-# ---------------------------------------------------------------------------
 
 def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
     """Create a Chroma vector store with Ollama embeddings.
@@ -21,41 +22,45 @@ def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
     Parameters
     ----------
     persist_directory: str
-        Directory where the Chroma database will be stored.
+        Directory where the vector data will be persisted.
-
-    Returns
-    -------
-    Chroma
-        A Chroma vector store instance.
     """
+    global store
     embeddings = OllamaEmbeddings(model="nomic-embed-text")
-    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+    store = Chroma(
+        collection_name="rag_collection",
+        embedding_function=embeddings,
+        persist_directory=persist_directory,
+    )
+    return store
 
-# ---------------------------------------------------------------------------
+def load_documents(directory: str | Path, vectorstore: Chroma) -> None:
-# Load documents from a directory into the vector store.
-# ---------------------------------------------------------------------------
-
-def load_documents(directory: str, vectorstore: Chroma) -> None:
     """Load all .txt and .md files from *directory* into *vectorstore*.
 
-    The documents are split into chunks using a RecursiveCharacterTextSplitter
+    The documents are split using ``RecursiveCharacterTextSplitter`` before
-    before being added to the vector store.
+    being added to the collection.
     """
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    dir_path = Path(directory)
-    docs = []
+    txt_files = list(dir_path.rglob("*.txt")) + list(dir_path.rglob("*.md"))
-    for path in Path(directory).rglob("*.txt"):
+    if not txt_files:
-        docs.append(path.read_text(encoding="utf-8"))
+        print(f"No .txt or .md files found in {dir_path}")
-    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)
 
-# ---------------------------------------------------------------------------
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-# End of vectorstore.py
+    docs: List[Document] = []
-# ---------------------------------------------------------------------------
+    for file_path in txt_files:
+        text = file_path.read_text(encoding="utf-8")
+        docs.extend(splitter.create_documents([text], metadata={"source": str(file_path)}))
+
+    vectorstore.add_documents(docs)
+    # Persist the collection to disk.
+    vectorstore.persist()
+    print(f"Loaded {len(docs)} documents from {dir_path} into Chroma.")
+
+# Helper to get the global store.
+
+def get_vectorstore() -> Chroma:
+    if store is None:
+        raise RuntimeError("Vector store has not been initialised. Call create_vectorstore() first.")
+    return store
+
+"""End of vectorstore.py"""