Create utils/vector_store.py

utils/vector_store.py

@@ -0,0 +1,58 @@
+# utils/vector_store.py
+import chromadb
+from chromadb.config import Settings
+from typing import List, Dict, Optional
+import numpy as np
+
+class VectorStore:
+    def __init__(self, persist_directory: str = "./data/chroma"):
+        self.client = chromadb.Client(Settings(
+            chroma_db_impl="duckdb+parquet",
+            persist_directory=persist_directory
+        ))
+        self.collection = self.client.get_or_create_collection("legal_docs")
+
+    def add_documents(self, chunks: List[Dict], metadata: Optional[Dict] = None):
+        """Add document chunks to vector store"""
+        for i, chunk in enumerate(chunks):
+            chunk_metadata = {
+                "chunk_id": i,
+                "length": chunk["metadata"]["length"],
+                "token_count": chunk["metadata"]["token_count"]
+            }
+            if metadata:
+                chunk_metadata.update(metadata)
+
+            self.collection.add(
+                embeddings=[chunk["embeddings"].tolist()],
+                documents=[chunk["text"]],
+                metadatas=[chunk_metadata],
+                ids=[f"{metadata.get('filename', 'doc')}_{i}"]
+            )
+
+    def search(self, query: str, n_results: int = 5) -> List[Dict]:
+        """Search for similar documents"""
+        results = self.collection.query(
+            query_texts=[query],
+            n_results=n_results,
+            include=["documents", "metadatas", "distances"]
+        )
+        
+        return [
+            {
+                "text": doc,
+                "metadata": meta,
+                "distance": dist
+            }
+            for doc, meta, dist in zip(
+                results["documents"][0],
+                results["metadatas"][0],
+                results["distances"][0]
+            )
+        ]
+
+    def get_all_documents(self) -> List[Dict]:
+        """Get all stored documents"""
+        return self.collection.get(
+            include=["documents", "metadatas", "embeddings"]
+        )