mcp-sqlite-memory-bank 1.5.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

mcp_sqlite_memory_bank/semantic.py

@@ -21,7 +21,9 @@ except ImportError:
     SENTENCE_TRANSFORMERS_AVAILABLE = False
     SentenceTransformer = None  # type: ignore
     util = None  # type: ignore
-    logging.warning("sentence-transformers not available. Install with: pip install sentence-transformers")
+    logging.warning(
+        "sentence-transformers not available. Install with: pip install sentence-transformers"
+    )
 
 try:
     import torch
@@ -58,7 +60,7 @@ class SemanticSearchEngine:
             )
 
     @property
-    def model(self):
+    def model(self) -> Any:
         """Lazy load the sentence transformer model."""
         if self._model is None:
             if not SENTENCE_TRANSFORMERS_AVAILABLE or SentenceTransformer is None:
@@ -67,7 +69,9 @@ class SemanticSearchEngine:
                 self._model = SentenceTransformer(self.model_name)
                 logging.info(f"Loaded semantic search model: {self.model_name}")
             except Exception as e:
-                raise DatabaseError(f"Failed to load semantic search model {self.model_name}: {e}")
+                raise DatabaseError(
+                    f"Failed to load semantic search model {self.model_name}: {e}"
+                )
         return self._model
 
     def get_embedding_dimensions(self) -> Optional[int]:
@@ -124,13 +128,17 @@ class SemanticSearchEngine:
 
         try:
             embeddings = self.model.encode(
-                valid_texts, convert_to_tensor=False, show_progress_bar=len(valid_texts) > 10
+                valid_texts,
+                convert_to_tensor=False,
+                show_progress_bar=len(valid_texts) > 10,
             )
             return [emb.tolist() for emb in embeddings]
         except Exception as e:
             raise DatabaseError(f"Failed to generate batch embeddings: {e}")
 
-    def calculate_similarity(self, embedding1: List[float], embedding2: List[float]) -> float:
+    def calculate_similarity(
+        self, embedding1: List[float], embedding2: List[float]
+    ) -> float:
         """
         Calculate cosine similarity between two embeddings.
 
@@ -193,7 +201,12 @@ class SemanticSearchEngine:
             return []
 
         # Use efficient torch/sentence-transformers if available
-        if TORCH_AVAILABLE and torch is not None and SENTENCE_TRANSFORMERS_AVAILABLE and util is not None:
+        if (
+            TORCH_AVAILABLE
+            and torch is not None
+            and SENTENCE_TRANSFORMERS_AVAILABLE
+            and util is not None
+        ):
             try:
                 # Convert to tensors
                 query_tensor = torch.tensor(query_embedding).unsqueeze(0)
@@ -213,7 +226,9 @@ class SemanticSearchEngine:
                 results.sort(key=lambda x: x[1], reverse=True)
                 return results[:top_k]
             except Exception as e:
-                logging.warning(f"Torch similarity search failed, using numpy fallback: {e}")
+                logging.warning(
+                    f"Torch similarity search failed, using numpy fallback: {e}"
+                )
 
         # Fallback to numpy implementation
         results = []
@@ -291,6 +306,10 @@ class SemanticSearchEngine:
                 original_idx = valid_indices[candidate_idx]
                 row = content_data[original_idx].copy()
 
+                # Remove embedding data to avoid polluting LLM responses
+                if embedding_column in row:
+                    del row[embedding_column]
+
                 # Add similarity score
                 row["similarity_score"] = round(similarity_score, 3)
 
@@ -298,7 +317,11 @@ class SemanticSearchEngine:
                 if content_columns:
                     matched_content = []
                     for col in content_columns:
-                        if col in row and row[col] and query.lower() in str(row[col]).lower():
+                        if (
+                            col in row
+                            and row[col]
+                            and query.lower() in str(row[col]).lower()
+                        ):
                             matched_content.append(f"{col}: {row[col]}")
                     if matched_content:
                         row["matched_content"] = matched_content
@@ -346,7 +369,11 @@ class SemanticSearchEngine:
 
         # Get semantic search results
         semantic_results = self.semantic_search(
-            query, content_data, embedding_column, similarity_threshold=0.3, top_k=top_k * 2  # Get more for reranking
+            query,
+            content_data,
+            embedding_column,
+            similarity_threshold=0.3,
+            top_k=top_k * 2,  # Get more for reranking
         )
 
         # Add text matching scores
@@ -360,11 +387,15 @@ class SemanticSearchEngine:
                         content = str(result[col]).lower()
                         if query_lower in content:
                             # Simple frequency-based text scoring
-                            text_score += content.count(query_lower) / len(content.split())
+                            text_score += content.count(query_lower) / len(
+                                content.split()
+                            )
 
             # Combine scores
             semantic_score = result.get("similarity_score", 0.0)
-            combined_score = (semantic_score * semantic_weight) + (text_score * text_weight)
+            combined_score = (semantic_score * semantic_weight) + (
+                text_score * text_weight
+            )
             result["combined_score"] = round(combined_score, 3)
             result["text_score"] = round(text_score, 3)
 
@@ -373,7 +404,7 @@ class SemanticSearchEngine:
 
         return semantic_results[:top_k]
 
-    def clear_cache(self):
+    def clear_cache(self) -> None:
         """Clear the embedding cache."""
         self._embedding_cache.clear()
         logging.info("Semantic search cache cleared")
@@ -390,15 +421,17 @@ def get_semantic_engine(model_name: str = "all-MiniLM-L6-v2") -> SemanticSearchE
     try:
         if _semantic_engine is None or _semantic_engine.model_name != model_name:
             if not SENTENCE_TRANSFORMERS_AVAILABLE:
-                raise ValueError("Sentence transformers not available for semantic search")
+                raise ValueError(
+                    "Sentence transformers not available for semantic search"
+                )
             _semantic_engine = SemanticSearchEngine(model_name)
-        
+
         # Verify the engine is properly initialized
-        if not hasattr(_semantic_engine, 'hybrid_search'):
+        if not hasattr(_semantic_engine, "hybrid_search"):
             raise ValueError("Semantic engine missing hybrid_search method")
-            
+
         return _semantic_engine
-        
+
     except Exception as e:
         raise DatabaseError(f"Failed to initialize semantic engine: {e}")