add vector_store.py

vector_store.py

@@ -1,21 +1,18 @@
 """
-Vector store implementation using ChromaDB.
+Vector store abstraction using ChromaDB.
 
-Provides functions to add documents and perform semantic search.
+Provides methods to add documents with embeddings and perform similarity search.
 """
-
 import os
-from typing import List, Dict
+from pathlib import Path
-from langchain_ollama import OllamaEmbeddings
+from typing import List, Dict, Any
-from chromadb import Client
+
+from chromadb import Client as ChromaClient
 from chromadb.config import Settings
+from langchain_ollama import OllamaEmbeddings
 
-# Initialize embeddings model (Ollama)
+# Initialize global client (in-memory for simplicity)
-embeddings = OllamaEmbeddings(model="nomic-embed-text")
+client = ChromaClient(Settings(chroma_db_impl="duckdb+parquet", persist_directory=None))
-
-# 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
@@ -23,43 +20,35 @@ 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:
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
+
+class ChromaVectorStore:
+    """Wrapper around a Chroma collection."""
+
     def __init__(self, collection):
         self.collection = collection
 
-    def add_document(self, doc_id: str, text: str, metadata: Dict | None = None) -> None:
+    def add_documents(self, documents: List[Dict[str, Any]]):
-        """Add a single document to the collection.
+        ids = []
+        texts = []
+        metadatas = []
+        for doc in documents:
+            ids.append(doc.get("id", os.urandom(8).hex()))
+            texts.append(doc["content"])
+            metadatas.append(doc.get("metadata", {}))
+        embeddings_list = embeddings.embed_documents(texts)
+        self.collection.add(ids=ids, documents=texts, embeddings=embeddings_list, metadatas=metadatas)
 
-        Parameters
+    def similarity_search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
-        ----------
-        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']
+            query_texts=[query],
+            n_results=k,
+            include=['documents', 'distances', 'metadatas'],
         )
-        hits = []
+        docs = []
-        for i in range(len(results["ids"][0])):
+        for doc, dist, meta in zip(results["documents"][0], results["distances"][0], results["metadatas"][0]):
-            hit = {
+            docs.append({"content": doc, "distance": dist, "metadata": meta})
-                "id": results["ids"][0][i],
+        return docs
-                "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
+# Singleton instance
-vector_store = VectorStore(col)
+vector_store = ChromaVectorStore(col)