Add vectorstore.py

vectorstore.py

@@ -0,0 +1,53 @@
+"""
+Vector store utilities for the RAG agent.
+
+Provides:
+- create_vectorstore(persist_directory)
+- load_documents(directory, vectorstore)
+"""
+import os
+from pathlib import Path
+from typing import List
+
+from langchain_ollama import OllamaEmbeddings
+from langchain_chroma import Chroma
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+
+def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
+    """Create or load a Chroma vector store.
+
+    Parameters
+    ----------
+    persist_directory: str
+        Directory where the Chroma DB will be stored. If it does not exist, it is created.
+    """
+    Path(persist_directory).mkdir(parents=True, exist_ok=True)
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+    return Chroma(
+        persist_directory=persist_directory,
+        embedding_function=embeddings
+    )
+
+
+def load_documents(directory: str, vectorstore: Chroma) -> None:
+    """Load all .txt/.md files from *directory* into the vector store.
+
+    The function reads each file, splits it with a RecursiveCharacterTextSplitter and adds the chunks to the collection.
+    Existing documents are overwritten – this is fine for an init script.
+    """
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    docs: List[str] = []
+    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:
+        print(f"No documents found in {directory}")
+        return
+
+    # Split and add to vectorstore
+    chunks = text_splitter.split_documents([{"content": d} for d in docs])
+    vectorstore.add_texts([c["content"] for c in chunks])
+    print(f"Loaded {len(chunks)} chunks into ChromaDB.")