Create utils/vector_store.py

utils/vector_store.py

@@ -0,0 +1,76 @@
+# utils/vector_store.py
+import faiss
+import numpy as np
+from typing import List, Dict, Optional
+import pickle
+import os
+from sentence_transformers import SentenceTransformer
+import streamlit as st
+
+class VectorStore:
+    def __init__(self, case_id: str):
+        self.case_id = case_id
+        self.base_path = f"data/cases/{case_id}/vector_store"
+        os.makedirs(self.base_path, exist_ok=True)
+        
+        if 'embedder' not in st.session_state:
+            st.session_state.embedder = SentenceTransformer('all-MiniLM-L6-v2')
+        
+        self.embedder = st.session_state.embedder
+        self.dimension = 384  # Dimension of embeddings
+        
+        self._load_or_create_index()
+
+    def _load_or_create_index(self):
+        """Load existing index or create new one"""
+        index_path = os.path.join(self.base_path, "faiss.index")
+        metadata_path = os.path.join(self.base_path, "metadata.pkl")
+        
+        if os.path.exists(index_path) and os.path.exists(metadata_path):
+            self.index = faiss.read_index(index_path)
+            with open(metadata_path, 'rb') as f:
+                self.metadata = pickle.load(f)
+        else:
+            self.index = faiss.IndexFlatL2(self.dimension)
+            self.metadata = []
+
+    def add_texts(self, texts: List[str], metadatas: Optional[List[Dict]] = None):
+        """Add texts to the vector store"""
+        if not texts:
+            return
+            
+        embeddings = self.embedder.encode(texts)
+        
+        self.index.add(np.array(embeddings).astype('float32'))
+        
+        if metadatas:
+            self.metadata.extend(metadatas)
+        else:
+            self.metadata.extend([{} for _ in texts])
+            
+        self._save_index()
+
+    def similarity_search(self, query: str, k: int = 5) -> List[Dict]:
+        """Search for similar texts"""
+        query_embedding = self.embedder.encode([query])
+        D, I = self.index.search(np.array(query_embedding).astype('float32'), k)
+        
+        results = []
+        for score, idx in zip(D[0], I[0]):
+            if idx < len(self.metadata):
+                result = {
+                    "metadata": self.metadata[idx],
+                    "score": float(score)
+                }
+                results.append(result)
+                
+        return results
+
+    def _save_index(self):
+        """Save index and metadata"""
+        index_path = os.path.join(self.base_path, "faiss.index")
+        metadata_path = os.path.join(self.base_path, "metadata.pkl")
+        
+        faiss.write_index(self.index, index_path)
+        with open(metadata_path, 'wb') as f:
+            pickle.dump(self.metadata, f)