add vectorstore.py

vectorstore.py

@@ -0,0 +1,61 @@
+"""ChromaDB vector store with Ollama embeddings and document chunking.""
+import os
+from typing import List
+
+from langchain_chroma import Chroma
+from langchain_ollama import OllamaEmbeddings
+from langchain_core.documents import Document
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+COLLECTION_NAME = "knowledge_base"
+PERSIST_DIR = os.getenv("CHROMA_PERSIST_DIR", "./chroma_db")
+
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
+
+_store: Chroma | None = None
+
+
+def get_vector_store() -> Chroma:
+    """Return singleton ChromaDB instance (persisted to disk)."""
+    global _store
+    if _store is None:
+        _store = Chroma(
+            collection_name=COLLECTION_NAME,
+            embedding_function=embeddings,
+            persist_directory=PERSIST_DIR,
+        )
+    return _store
+
+
+def add_documents(texts: List[str], chunk_size: int = 500, chunk_overlap: int = 100) -> int:
+    """Split texts into chunks and index them in ChromaDB.
+
+    Args:
+        texts: list of raw text strings
+        chunk_size: maximum chunk length in characters
+        chunk_overlap: overlap between consecutive chunks
+
+    Returns:
+        number of chunks added to the store
+    """
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+    )
+    docs: List[Document] = splitter.create_documents(texts)
+    if docs:
+        get_vector_store().add_documents(docs)
+    return len(docs)
+
+
+def similarity_search(query: str, k: int = 5) -> List[Document]:
+    """Semantic search in ChromaDB.
+
+    Args:
+        query: natural language search query
+        k: number of top results to return
+
+    Returns:
+        list of the most relevant Document objects
+    """
+    return get_vector_store().similarity_search(query, k=k)