Add vector_store.py

vector_store.py

@@ -0,0 +1,65 @@
+"""
+Vector store implementation using ChromaDB.
+
+Provides functions to add documents and perform semantic search.
+"""
+
+import os
+from typing import List, Dict
+from langchain_ollama import OllamaEmbeddings
+from chromadb import Client
+from chromadb.config import Settings
+
+# Initialize embeddings model (Ollama)
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
+
+# ChromaDB client – in‑memory by default, persistent folder "chromadb"
+CHROMA_DIR = os.path.join(os.getcwd(), "chromadb")
+client = Client(Settings(chroma_db_impl="duckdb+parquet", persist_directory=CHROMA_DIR))
+collection_name = "rag_collection"
+
+# Ensure collection exists
+if collection_name not in client.list_collections():
+    client.create_collection(name=collection_name)
+col = client.get_or_create_collection(name=collection_name)
+
+class VectorStore:
+    def __init__(self, collection):
+        self.collection = collection
+
+    def add_document(self, doc_id: str, text: str, metadata: Dict | None = None) -> None:
+        """Add a single document to the collection.
+
+        Parameters
+        ----------
+        doc_id: str
+            Unique identifier for the document.
+        text: str
+            Raw text content.
+        metadata: dict, optional
+            Additional key/value pairs stored with the vector.
+        """
+        vec = embeddings.embed_query(text)
+        self.collection.add(ids=[doc_id], documents=[text], metadatas=[metadata or {}])
+
+    def search(self, query: str, k: int = 5) -> List[Dict]:
+        """Semantic search over the collection.
+
+        Returns a list of dicts with keys: id, document, score, metadata.
+        """
+        results = self.collection.query(
+            query_texts=[query], n_results=k, include=['documents', 'distances', 'metadatas']
+        )
+        hits = []
+        for i in range(len(results["ids"][0])):
+            hit = {
+                "id": results["ids"][0][i],
+                "document": results["documents"][0][i],
+                "score": 1 - results["distances"][0][i],  # distance to similarity
+                "metadata": results["metadatas"][0][i],
+            }
+            hits.append(hit)
+        return hits
+
+# Singleton instance for easy import
+vector_store = VectorStore(col)