Update utils/vector_store.py

utils/vector_store.py

@@ -7,96 +7,104 @@ from datetime import datetime
 import streamlit as st
 
 class VectorStore:
-    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 = []
-        self._load_vectors()
-
-    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:
-                    self.vectors = pickle.load(f)
-            if not isinstance(self.vectors, list):
-                self.vectors = []
-        except Exception as e:
-            self.vectors = []
-
-    def _save_vectors(self):
-        """Save vectors with error handling."""
-        vector_file = os.path.join(self.storage_path, "vectors.pkl")
-        try:
-            with open(vector_file, "wb") as f:
-                pickle.dump(self.vectors, f)
-        except Exception as e:
-            raise Exception(f"Error saving vectors: {str(e)}")
-
-    def add_document(self, doc_id: str, text: str, metadata: Dict[str, Any] = None):
-        """Add a document to the vector store."""
-        try:
-            vector = self.model.encode(text, convert_to_tensor=True)
-            doc_record = {
-                "doc_id": doc_id,
-                "vector": vector,
-                "text": text,
-                "metadata": metadata or {}
-            }
-            if not isinstance(self.vectors, list):
-                self.vectors = []
-            self.vectors.append(doc_record)
-            self._save_vectors()
-        except Exception as e:
-            raise Exception(f"Error adding document: {str(e)}")
-
     def similarity_search(self, query: str, k: int = 3, filter_docs: Optional[List[str]] = None) -> List[Dict]:
-        """Perform similarity search with document filtering."""
+        """Enhanced similarity search with ontology awareness."""
         try:
             if not self.vectors:
                 return []
 
+            # Encode query
             query_vector = self.model.encode(query, convert_to_tensor=True)
+            
+            # Calculate enhanced similarities
             results = []
-
             for doc in self.vectors:
-                # Apply document filter if provided
+                # Skip if document is filtered out
                 if filter_docs and doc["doc_id"] not in filter_docs:
                     continue
+                
+                try:
+                    # Base similarity score
+                    base_similarity = util.pytorch_cos_sim(query_vector, doc["vector"]).item()
+                    
+                    # Calculate ontology boost
+                    ontology_boost = self._calculate_ontology_relevance(
+                        query, 
+                        doc.get('metadata', {}).get('ontology_links', [])
+                    )
+                    
+                    # Final score combining vector similarity and ontology relevance
+                    final_score = (base_similarity * 0.7) + (ontology_boost * 0.3)
+                    
+                    results.append({
+                        "doc_id": doc["doc_id"],
+                        "text": doc["text"],
+                        "metadata": doc["metadata"],
+                        "score": float(final_score),
+                        "base_similarity": float(base_similarity),
+                        "ontology_boost": float(ontology_boost)
+                    })
+                    
+                except Exception as e:
+                    st.warning(f"Error processing document: {str(e)}")
+                    continue
 
-                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)
-                })
-
+            # Sort by final score
             results.sort(key=lambda x: x["score"], reverse=True)
             return results[:k]
 
         except Exception as e:
-            raise Exception(f"Error in similarity search: {str(e)}")
+            st.error(f"Error in similarity search: {str(e)}")
+            return []
 
-    def delete_document(self, doc_id: str) -> bool:
-        """Delete a document from the vector store."""
+    def _calculate_ontology_relevance(self, query: str, ontology_links: List[Dict]) -> float:
+        """Calculate ontology-based relevance score."""
+        if not ontology_links:
+            return 0.0
+            
+        query_lower = query.lower()
+        relevance_score = 0.0
+        
+        for link in ontology_links:
+            # Direct concept match
+            if link['concept'].lower() in query_lower:
+                relevance_score += 0.3
+                
+            # Description match
+            if 'description' in link and any(term in query_lower 
+                                           for term in link['description'].lower().split()):
+                relevance_score += 0.2
+                
+            # Related concepts match
+            if 'relationships' in link:
+                for related in link['relationships']:
+                    if related.lower() in query_lower:
+                        relevance_score += 0.1
+
+        # Normalize score to [0, 1]
+        return min(1.0, relevance_score)
+
+    def add_document(self, doc_id: str, text: str, metadata: Dict[str, Any] = None):
+        """Add document with enhanced metadata processing."""
         try:
-            initial_length = len(self.vectors)
-            self.vectors = [doc for doc in self.vectors if doc["doc_id"] != doc_id]
+            # Create vector embedding
+            vector = self.model.encode(text, convert_to_tensor=True)
+            
+            # Ensure metadata includes ontology links
+            if metadata and 'ontology_links' not in metadata:
+                metadata['ontology_links'] = []
+            
+            doc_record = {
+                "doc_id": doc_id,
+                "vector": vector,
+                "text": text,
+                "metadata": metadata or {}
+            }
+            
+            if not isinstance(self.vectors, list):
+                self.vectors = []
+            self.vectors.append(doc_record)
             self._save_vectors()
-            return len(self.vectors) < initial_length
+            
         except Exception as e:
-            raise Exception(f"Error deleting document: {str(e)}")
-
-    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
+            raise Exception(f"Error adding document: {str(e)}")