Update utils/vector_store.py

utils/vector_store.py

@@ -4,141 +4,123 @@ from typing import List, Dict, Any
 from sentence_transformers import SentenceTransformer, util
 import numpy as np
 from datetime import datetime
-
+import streamlit as st
 
 class VectorStore:
-    def __init__(self, storage_path: str = "data/vector_store", model_name: str = 'all-MiniLM-L6-v2'):
-        """Initialize VectorStore with improved chunk handling."""
+    def __init__(self, storage_path: str = "data/vector_store"):
+        """Initialize VectorStore with storage management."""
         self.storage_path = storage_path
         os.makedirs(storage_path, exist_ok=True)
+        
+        self.model = SentenceTransformer('all-MiniLM-L6-v2')
+        self.vectors = []  # Initialize empty list
+        self._load_vectors()  # Load any existing vectors
 
-        self.model = SentenceTransformer(model_name)
-        self.vectors = self._load_vectors()
-        self.chunk_size = 512  # Optimal size for most transformer models
-        self.chunk_overlap = 50  # Overlap to maintain context
-
-    def _load_vectors(self) -> List[Dict]:
-        """Load vectors with error handling and versioning."""
+    def _load_vectors(self):
+        """Load stored vectors with error handling."""
         vector_file = os.path.join(self.storage_path, "vectors.pkl")
         try:
             if os.path.exists(vector_file):
                 with open(vector_file, "rb") as f:
-                    vectors = pickle.load(f)
-                return vectors if isinstance(vectors, list) else []
+                    self.vectors = pickle.load(f)
+            if not isinstance(self.vectors, list):
+                self.vectors = []
         except Exception as e:
-            print(f"Error loading vectors: {e}")
-            return []
+            st.error(f"Error loading vectors: {str(e)}")
+            self.vectors = []
 
     def _save_vectors(self):
-        """Save vectors with backup and atomic write."""
+        """Save vectors with error handling."""
         vector_file = os.path.join(self.storage_path, "vectors.pkl")
-        backup_file = vector_file + ".backup"
-
-        # Create backup of existing vectors
-        if os.path.exists(vector_file):
-            os.replace(vector_file, backup_file)
-
         try:
             with open(vector_file, "wb") as f:
                 pickle.dump(self.vectors, f)
-            # Remove backup after successful save
-            if os.path.exists(backup_file):
-                os.remove(backup_file)
         except Exception as e:
-            print(f"Error saving vectors: {e}")
-            # Restore from backup if save failed
-            if os.path.exists(backup_file):
-                os.replace(backup_file, vector_file)
-
-    def add_document(self, doc_id: str, text: str, metadata: Dict[str, Any]):
-        """Add document with improved chunking and metadata."""
-        # Create chunks with overlap
-        chunks = self._create_chunks(text)
-
-        # Add timestamp and chunk info to metadata
-        base_metadata = {
-            **metadata,
-            "added_at": datetime.now().isoformat(),
-            "doc_id": doc_id,
-            "total_chunks": len(chunks)
-        }
-
-        # Process and store chunks
-        for chunk_idx, chunk in enumerate(chunks):
-            chunk_metadata = {
-                **base_metadata,
-                "chunk_idx": chunk_idx,
-                "chunk_text": chunk[:200]  # Store preview of chunk text
-            }
-
-            # Encode chunk
-            vector = self.model.encode(chunk, convert_to_tensor=True)
+            st.error(f"Error saving vectors: {str(e)}")
 
-            # Store chunk with metadata
-            self.vectors.append({
-                "doc_id": f"{doc_id}_chunk_{chunk_idx}",
+    def add_document(self, doc_id: str, text: str, metadata: Dict[str, Any] = None):
+        """Add a document to the vector store."""
+        try:
+            # Create vector embedding
+            vector = self.model.encode(text, convert_to_tensor=True)
+            
+            # Create document record
+            doc_record = {
+                "doc_id": doc_id,
                 "vector": vector,
-                "text": chunk,
-                "metadata": chunk_metadata
-            })
-
-        self._save_vectors()
-
-    def _create_chunks(self, text: str) -> List[str]:
-        """Create overlapping chunks with improved sentence boundary handling."""
-        # Split into sentences first
-        sentences = [s.strip() for s in text.split('.') if s.strip()]
-        chunks = []
-        current_chunk = []
-        current_size = 0
-
-        for sentence in sentences:
-            sentence_size = len(sentence.split())
-
-            if current_size + sentence_size > self.chunk_size:
-                # Save current chunk
-                if current_chunk:
-                    chunks.append(' '.join(current_chunk))
-                # Start new chunk with overlap
-                overlap_start = max(0, len(current_chunk) - self.chunk_overlap)
-                current_chunk = current_chunk[overlap_start:] + [sentence]
-                current_size = sum(len(s.split()) for s in current_chunk)
-            else:
-                current_chunk.append(sentence)
-                current_size += sentence_size
-
-        # Add final chunk
-        if current_chunk:
-            chunks.append(' '.join(current_chunk))
-
-        return chunks
+                "text": text,
+                "metadata": metadata or {}
+            }
+            
+            # Add to vectors list
+            if not isinstance(self.vectors, list):
+                self.vectors = []
+            self.vectors.append(doc_record)
+            
+            # Save updated vectors
+            self._save_vectors()
+            
+        except Exception as e:
+            st.error(f"Error adding document to vector store: {str(e)}")
+            raise
 
     def similarity_search(self, query: str, k: int = 3) -> List[Dict]:
         """Perform similarity search with error handling."""
         try:
-            # If no vectors are stored yet, return empty list
+            # Handle empty vectors
             if not self.vectors:
                 return []
 
+            # Encode query
             query_vector = self.model.encode(query, convert_to_tensor=True)
+            
+            # Calculate similarities
             results = []
-
             for doc in self.vectors:
-                similarity = util.pytorch_cos_sim(query_vector, doc["vector"]).item()
-                results.append({
-                    "text": doc["text"],
-                    "metadata": doc["metadata"],
-                    "score": similarity
-                })
-
-            # Sort by similarity and return top k
+                try:
+                    similarity = util.pytorch_cos_sim(query_vector, doc["vector"]).item()
+                    results.append({
+                        "doc_id": doc["doc_id"],
+                        "text": doc["text"],
+                        "metadata": doc["metadata"],
+                        "score": float(similarity)  # Convert to float for serialization
+                    })
+                except Exception as e:
+                    st.warning(f"Skipping document due to error: {str(e)}")
+                    continue
+
+            # Sort by similarity
             results.sort(key=lambda x: x["score"], reverse=True)
             return results[:k]
 
         except Exception as e:
-            print(f"Error in similarity search: {str(e)}")
+            st.error(f"Error in similarity search: {str(e)}")
             return []
 
+    def get_document(self, doc_id: str) -> Dict:
+        """Retrieve a document by ID."""
+        try:
+            for doc in self.vectors:
+                if doc["doc_id"] == doc_id:
+                    return {
+                        "doc_id": doc["doc_id"],
+                        "text": doc["text"],
+                        "metadata": doc["metadata"]
+                    }
+            return None
+        except Exception as e:
+            st.error(f"Error retrieving document: {str(e)}")
+            return None
+
+    def clear(self):
+        """Clear all vectors."""
+        self.vectors = []
+        self._save_vectors()
+
+    def __len__(self):
+        """Get number of documents in store."""
+        return len(self.vectors) if self.vectors is not None else 0
+        
     def _rerank_results(self, results: List[Dict], query: str) -> List[Dict]:
         """Re-rank results considering chunk position and metadata relevance."""
         for result in results: