add vector_store.py

vector_store.py

@@ -1,54 +1,64 @@
 """
-Vector store abstraction using ChromaDB.
+Vector store implementation using ChromaDB.
 
-Provides methods to add documents with embeddings and perform similarity search.
+Provides methods to add documents and perform similarity search.
 """
 import os
 from pathlib import Path
-from typing import List, Dict, Any
+from typing import List, Dict
 
-from chromadb import Client as ChromaClient
+from chromadb import Client as ChromadbClient
 from chromadb.config import Settings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_ollama import OllamaEmbeddings
 
-# Initialize global client (in-memory for simplicity)
-client = ChromaClient(Settings(chroma_db_impl="duckdb+parquet", persist_directory=None))
-collection_name = "rag_collection"
+# Ensure persistent directory exists
+CHROMA_DIR = Path("./chroma_db")
+CHROMA_DIR.mkdir(parents=True, exist_ok=True)
 
-# Ensure collection exists
+# Embedding model via Ollama
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
+
+# Chroma client with persistence
+client = ChromadbClient(Settings(persist_directory=str(CHROMA_DIR)))
+collection_name = "knowledge"
 if collection_name not in client.list_collections():
     client.create_collection(name=collection_name)
 col = client.get_or_create_collection(name=collection_name)
 
-embeddings = OllamaEmbeddings(model="nomic-embed-text")
+# Chunker from chunker.py
+from chunker import chunker
 
-class ChromaVectorStore:
-    """Wrapper around a Chroma collection."""
+class VectorStore:
+    def add_documents(self, docs: List[str], metadatas: List[Dict]):
+        """Add documents to the collection.
 
-    def __init__(self, collection):
-        self.collection = collection
+        Parameters
+        ----------
+        docs: list of raw text strings.
+        metadatas: list of metadata dicts corresponding to each doc.
+        """
+        # Split into chunks and embed
+        all_chunks = []
+        all_metadatas = []
+        for doc, meta in zip(docs, metadatas):
+            chunks = chunker.split_text(doc)
+            all_chunks.extend(chunks)
+            all_metadatas.extend([meta] * len(chunks))
 
-    def add_documents(self, documents: List[Dict[str, Any]]):
-        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)
+        # Generate embeddings via Ollama
+        embeds = embeddings.embed_documents(all_chunks)
+        ids = [f"chunk_{i}" for i in range(len(all_chunks))]
+        col.add(ids=ids, documents=all_chunks, metadatas=all_metadatas, embeddings=embeds)
 
-    def similarity_search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
-        results = self.collection.query(
+    def similarity_search(self, query: str, k: int = 5) -> List[Dict]:
+        """Return top-k similar chunks with metadata."""
+        results = col.query(
             query_texts=[query],
             n_results=k,
-            include=['documents', 'distances', 'metadatas'],
+            include=['documents', 'metadatas'],
         )
-        docs = []
-        for doc, dist, meta in zip(results["documents"][0], results["distances"][0], results["metadatas"][0]):
-            docs.append({"content": doc, "distance": dist, "metadata": meta})
-        return docs
-
-# Singleton instance
-vector_store = ChromaVectorStore(col)
+        # results is dict with keys documents, metadatas
+        docs = results["documents"][0]
+        metas = results["metadatas"][0]
+        return [{"content": d, "metadata": m} for d, m in zip(docs, metas)]