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

mcp_sqlite_memory_bank/database.py

@@ -12,7 +12,19 @@ import json
 import logging
 from functools import wraps
 from typing import Dict, List, Any, Optional, Callable, cast
-from sqlalchemy import create_engine, MetaData, Table, select, insert, update, delete, text, inspect, and_, or_
+from sqlalchemy import (
+    create_engine,
+    MetaData,
+    Table,
+    select,
+    insert,
+    update,
+    delete,
+    text,
+    inspect,
+    and_,
+    or_,
+)
 from sqlalchemy.engine import Engine
 from sqlalchemy.exc import SQLAlchemyError
 from contextlib import contextmanager
@@ -63,7 +75,7 @@ class SQLiteMemoryDatabase:
         except Exception as e:
             logging.warning(f"Error closing database: {e}")
 
-    def __del__(self):
+    def __del__(self) -> None:
         """Ensure cleanup when object is garbage collected."""
         self.close()
 
@@ -76,7 +88,7 @@ class SQLiteMemoryDatabase:
             logging.warning(f"Failed to refresh metadata: {e}")
 
     @contextmanager
-    def get_connection(self):
+    def get_connection(self) -> Any:
         """Get a database connection with automatic cleanup."""
         conn = self.engine.connect()
         try:
@@ -96,12 +108,17 @@ class SQLiteMemoryDatabase:
 
         return self.metadata.tables[table_name]
 
-    def _validate_columns(self, table: Table, column_names: List[str], context: str = "operation") -> None:
+    def _validate_columns(
+        self, table: Table, column_names: List[str], context: str = "operation"
+    ) -> None:
         """Validate that all column names exist in the table."""
         valid_columns = set(col.name for col in table.columns)
         for col_name in column_names:
             if col_name not in valid_columns:
-                raise ValidationError(f"Invalid column '{col_name}' for table " f"'{table.name}' in {context}")
+                raise ValidationError(
+                    f"Invalid column '{col_name}' for table "
+                    f"'{table.name}' in {context}"
+                )
 
     def _build_where_conditions(self, table: Table, where: Dict[str, Any]) -> List:
         """Build SQLAlchemy WHERE conditions from a dictionary."""
@@ -135,7 +152,9 @@ class SQLiteMemoryDatabase:
 
         return decorator
 
-    def create_table(self, table_name: str, columns: List[Dict[str, str]]) -> ToolResponse:
+    def create_table(
+        self, table_name: str, columns: List[Dict[str, str]]
+    ) -> ToolResponse:
         """Create a new table with the specified columns."""
         # Input validation
         if not table_name or not table_name.isidentifier():
@@ -170,7 +189,11 @@ class SQLiteMemoryDatabase:
         try:
             with self.get_connection() as conn:
                 inspector = inspect(conn)
-                tables = [name for name in inspector.get_table_names() if not name.startswith("sqlite_")]
+                tables = [
+                    name
+                    for name in inspector.get_table_names()
+                    if not name.startswith("sqlite_")
+                ]
             return {"success": True, "tables": tables}
         except SQLAlchemyError as e:
             raise DatabaseError(f"Failed to list tables: {str(e)}")
@@ -229,7 +252,9 @@ class SQLiteMemoryDatabase:
         except (ValidationError, SQLAlchemyError) as e:
             if isinstance(e, ValidationError):
                 raise e
-            raise DatabaseError(f"Failed to rename table from {old_name} to {new_name}: {str(e)}")
+            raise DatabaseError(
+                f"Failed to rename table from {old_name} to {new_name}: {str(e)}"
+            )
 
     def insert_row(self, table_name: str, data: Dict[str, Any]) -> ToolResponse:
         """Insert a row into a table."""
@@ -248,7 +273,10 @@ class SQLiteMemoryDatabase:
             raise DatabaseError(f"Failed to insert into table {table_name}: {str(e)}")
 
     def read_rows(
-        self, table_name: str, where: Optional[Dict[str, Any]] = None, limit: Optional[int] = None
+        self,
+        table_name: str,
+        where: Optional[Dict[str, Any]] = None,
+        limit: Optional[int] = None,
     ) -> ToolResponse:
         """Read rows from a table with optional filtering."""
         try:
@@ -275,7 +303,10 @@ class SQLiteMemoryDatabase:
             raise DatabaseError(f"Failed to read from table {table_name}: {str(e)}")
 
     def update_rows(
-        self, table_name: str, data: Dict[str, Any], where: Optional[Dict[str, Any]] = None
+        self,
+        table_name: str,
+        data: Dict[str, Any],
+        where: Optional[Dict[str, Any]] = None,
     ) -> ToolResponse:
         """Update rows in a table."""
         if not data:
@@ -299,7 +330,9 @@ class SQLiteMemoryDatabase:
                 raise e
             raise DatabaseError(f"Failed to update table {table_name}: {str(e)}")
 
-    def delete_rows(self, table_name: str, where: Optional[Dict[str, Any]] = None) -> ToolResponse:
+    def delete_rows(
+        self, table_name: str, where: Optional[Dict[str, Any]] = None
+    ) -> ToolResponse:
         """Delete rows from a table."""
         try:
             table = self._ensure_table_exists(table_name)
@@ -310,7 +343,9 @@ class SQLiteMemoryDatabase:
             if conditions:
                 stmt = stmt.where(and_(*conditions))
             else:
-                logging.warning(f"delete_rows called without WHERE clause on table {table_name}")
+                logging.warning(
+                    f"delete_rows called without WHERE clause on table {table_name}"
+                )
 
             result = self._execute_with_commit(stmt)
             return {"success": True, "rows_affected": result.rowcount}
@@ -363,13 +398,16 @@ class SQLiteMemoryDatabase:
         try:
             self._refresh_metadata()
             schemas = {
-                table_name: [col.name for col in table.columns] for table_name, table in self.metadata.tables.items()
+                table_name: [col.name for col in table.columns]
+                for table_name, table in self.metadata.tables.items()
             }
             return {"success": True, "schemas": schemas}
         except SQLAlchemyError as e:
             raise DatabaseError(f"Failed to list all columns: {str(e)}")
 
-    def search_content(self, query: str, tables: Optional[List[str]] = None, limit: int = 50) -> ToolResponse:
+    def search_content(
+        self, query: str, tables: Optional[List[str]] = None, limit: int = 50
+    ) -> ToolResponse:
         """Perform full-text search across table content."""
         if not query or not query.strip():
             raise ValidationError("Search query cannot be empty")
@@ -390,7 +428,8 @@ class SQLiteMemoryDatabase:
                     text_columns = [
                         col
                         for col in table.columns
-                        if "TEXT" in str(col.type).upper() or "VARCHAR" in str(col.type).upper()
+                        if "TEXT" in str(col.type).upper()
+                        or "VARCHAR" in str(col.type).upper()
                     ]
 
                     if not text_columns:
@@ -413,42 +452,72 @@ class SQLiteMemoryDatabase:
                             if col.name in row_dict and row_dict[col.name]:
                                 content = str(row_dict[col.name]).lower()
                                 content_length = len(content)
-                                
+
                                 if query_lower in content:
                                     # Factor 1: Exact phrase frequency (weighted higher)
                                     exact_frequency = content.count(query_lower)
-                                    exact_score = (exact_frequency * 2.0) / content_length if content_length > 0 else 0
-                                    
+                                    exact_score = (
+                                        (exact_frequency * 2.0) / content_length
+                                        if content_length > 0
+                                        else 0
+                                    )
+
                                     # Factor 2: Individual term frequency
                                     term_score = 0.0
                                     for term in query_terms:
                                         if term in content:
-                                            term_score += content.count(term) / content_length if content_length > 0 else 0
-                                    
+                                            term_score += (
+                                                content.count(term) / content_length
+                                                if content_length > 0
+                                                else 0
+                                            )
+
                                     # Factor 3: Position bonus (early matches score higher)
                                     position_bonus = 0.0
                                     first_occurrence = content.find(query_lower)
                                     if first_occurrence != -1:
-                                        position_bonus = (content_length - first_occurrence) / content_length * 0.1
-                                    
+                                        position_bonus = (
+                                            (content_length - first_occurrence)
+                                            / content_length
+                                            * 0.1
+                                        )
+
                                     # Factor 4: Column importance (title/name columns get bonus)
                                     column_bonus = 0.0
-                                    if any(keyword in col.name.lower() for keyword in ['title', 'name', 'summary', 'description']):
+                                    if any(
+                                        keyword in col.name.lower()
+                                        for keyword in [
+                                            "title",
+                                            "name",
+                                            "summary",
+                                            "description",
+                                        ]
+                                    ):
                                         column_bonus = 0.2
-                                    
+
                                     # Combined relevance score
-                                    col_relevance = exact_score + term_score + position_bonus + column_bonus
+                                    col_relevance = (
+                                        exact_score
+                                        + term_score
+                                        + position_bonus
+                                        + column_bonus
+                                    )
                                     relevance_scores.append(col_relevance)
-                                    
+
                                     # Enhanced matched content with context
                                     snippet_start = max(0, first_occurrence - 50)
-                                    snippet_end = min(len(row_dict[col.name]), first_occurrence + len(query) + 50)
-                                    snippet = str(row_dict[col.name])[snippet_start:snippet_end]
+                                    snippet_end = min(
+                                        len(row_dict[col.name]),
+                                        first_occurrence + len(query) + 50,
+                                    )
+                                    snippet = str(row_dict[col.name])[
+                                        snippet_start:snippet_end
+                                    ]
                                     if snippet_start > 0:
                                         snippet = "..." + snippet
                                     if snippet_end < len(str(row_dict[col.name])):
                                         snippet = snippet + "..."
-                                    
+
                                     matched_content.append(f"{col.name}: {snippet}")
 
                         total_relevance = sum(relevance_scores)
@@ -460,8 +529,14 @@ class SQLiteMemoryDatabase:
                                     "row_data": row_dict,
                                     "matched_content": matched_content,
                                     "relevance": round(total_relevance, 4),
-                                    "match_quality": "high" if total_relevance > 0.5 else "medium" if total_relevance > 0.1 else "low",
-                                    "match_count": len(relevance_scores)
+                                    "match_quality": (
+                                        "high"
+                                        if total_relevance > 0.5
+                                        else (
+                                            "medium" if total_relevance > 0.1 else "low"
+                                        )
+                                    ),
+                                    "match_count": len(relevance_scores),
                                 }
                             )
 
@@ -471,6 +546,7 @@ class SQLiteMemoryDatabase:
                 if isinstance(rel, (int, float)):
                     return float(rel)
                 return 0.0
+
             results.sort(key=get_relevance, reverse=True)
             results = results[:limit]
 
@@ -486,16 +562,24 @@ class SQLiteMemoryDatabase:
                 raise e
             raise DatabaseError(f"Failed to search content: {str(e)}")
 
-    def explore_tables(self, pattern: Optional[str] = None, include_row_counts: bool = True) -> ToolResponse:
+    def explore_tables(
+        self, pattern: Optional[str] = None, include_row_counts: bool = True
+    ) -> ToolResponse:
         """Explore table structures and content."""
         try:
             self._refresh_metadata()
             table_names = list(self.metadata.tables.keys())
 
             if pattern:
-                table_names = [name for name in table_names if pattern.replace("%", "") in name]
-
-            exploration: Dict[str, Any] = {"tables": [], "total_tables": len(table_names), "total_rows": 0}
+                table_names = [
+                    name for name in table_names if pattern.replace("%", "") in name
+                ]
+
+            exploration: Dict[str, Any] = {
+                "tables": [],
+                "total_tables": len(table_names),
+                "total_rows": 0,
+            }
 
             with self.get_connection() as conn:
                 for table_name in table_names:
@@ -515,31 +599,48 @@ class SQLiteMemoryDatabase:
                         }
                         columns.append(col_data)
 
-                        if "TEXT" in str(col.type).upper() or "VARCHAR" in str(col.type).upper():
+                        if (
+                            "TEXT" in str(col.type).upper()
+                            or "VARCHAR" in str(col.type).upper()
+                        ):
                             text_columns.append(col.name)
 
-                    table_info: Dict[str, Any] = {"name": table_name, "columns": columns, "text_columns": text_columns}
+                    table_info: Dict[str, Any] = {
+                        "name": table_name,
+                        "columns": columns,
+                        "text_columns": text_columns,
+                    }
 
                     # Add row count if requested
                     if include_row_counts:
-                        count_result = conn.execute(select(text("COUNT(*)")).select_from(table))
+                        count_result = conn.execute(
+                            select(text("COUNT(*)")).select_from(table)
+                        )
                         row_count = count_result.scalar()
                         table_info["row_count"] = row_count
                         exploration["total_rows"] += row_count
 
                     # Add sample data
                     sample_result = conn.execute(select(table).limit(3))
-                    sample_rows = [dict(row._mapping) for row in sample_result.fetchall()]
+                    sample_rows = [
+                        dict(row._mapping) for row in sample_result.fetchall()
+                    ]
                     if sample_rows:
                         table_info["sample_data"] = sample_rows
 
                     # Add content preview for text columns
                     if text_columns:
                         content_preview: Dict[str, List[Any]] = {}
-                        for col_name in text_columns[:3]:  # Limit to first 3 text columns
+                        for col_name in text_columns[
+                            :3
+                        ]:  # Limit to first 3 text columns
                             col = table.c[col_name]
-                            preview_result = conn.execute(select(col).distinct().where(col.isnot(None)).limit(5))
-                            unique_values: List[Any] = [row[0] for row in preview_result.fetchall() if row[0]]
+                            preview_result = conn.execute(
+                                select(col).distinct().where(col.isnot(None)).limit(5)
+                            )
+                            unique_values: List[Any] = [
+                                row[0] for row in preview_result.fetchall() if row[0]
+                            ]
                             if unique_values:
                                 content_preview[col_name] = unique_values
 
@@ -554,22 +655,32 @@ class SQLiteMemoryDatabase:
 
     # --- Semantic Search Methods ---
 
-    def add_embedding_column(self, table_name: str, embedding_column: str = "embedding") -> EmbeddingColumnResponse:
+    def add_embedding_column(
+        self, table_name: str, embedding_column: str = "embedding"
+    ) -> EmbeddingColumnResponse:
         """Add an embedding column to a table for semantic search."""
         try:
             table = self._ensure_table_exists(table_name)
 
             # Check if embedding column already exists
             if embedding_column in [col.name for col in table.columns]:
-                return {"success": True, "message": f"Embedding column '{embedding_column}' already exists"}
+                return {
+                    "success": True,
+                    "message": f"Embedding column '{embedding_column}' already exists",
+                }
 
             # Add embedding column as TEXT (JSON storage)
             with self.get_connection() as conn:
-                conn.execute(text(f"ALTER TABLE {table_name} ADD COLUMN {embedding_column} TEXT"))
+                conn.execute(
+                    text(f"ALTER TABLE {table_name} ADD COLUMN {embedding_column} TEXT")
+                )
                 conn.commit()
 
             self._refresh_metadata()
-            return {"success": True, "message": f"Added embedding column '{embedding_column}' to table '{table_name}'"}
+            return {
+                "success": True,
+                "message": f"Added embedding column '{embedding_column}' to table '{table_name}'",
+            }
 
         except (ValidationError, SQLAlchemyError) as e:
             if isinstance(e, ValidationError):
@@ -586,7 +697,9 @@ class SQLiteMemoryDatabase:
     ) -> GenerateEmbeddingsResponse:
         """Generate embeddings for text content in a table."""
         if not is_semantic_search_available():
-            raise ValidationError("Semantic search is not available. Please install sentence-transformers.")
+            raise ValidationError(
+                "Semantic search is not available. Please install sentence-transformers."
+            )
 
         try:
             table = self._ensure_table_exists(table_name)
@@ -596,7 +709,9 @@ class SQLiteMemoryDatabase:
             table_columns = [col.name for col in table.columns]
             for col in text_columns:
                 if col not in table_columns:
-                    raise ValidationError(f"Column '{col}' not found in table '{table_name}'")
+                    raise ValidationError(
+                        f"Column '{col}' not found in table '{table_name}'"
+                    )
 
             # Add embedding column if it doesn't exist
             if embedding_column not in table_columns:
@@ -642,7 +757,9 @@ class SQLiteMemoryDatabase:
                             combined_text = " ".join(text_parts)
 
                             # Generate embedding
-                            embedding = semantic_engine.generate_embedding(combined_text)
+                            embedding = semantic_engine.generate_embedding(
+                                combined_text
+                            )
                             embedding_json = json.dumps(embedding)
 
                             # Update row with embedding
@@ -656,14 +773,17 @@ class SQLiteMemoryDatabase:
                             processed += 1
 
                     conn.commit()
-                    logging.info(f"Generated embeddings for batch {i//batch_size + 1}, processed {processed} rows")
+                    logging.info(
+                        f"Generated embeddings for batch {i//batch_size + 1}, processed {processed} rows"
+                    )
 
                 return {
                     "success": True,
                     "message": f"Generated embeddings for {processed} rows",
                     "processed": processed,
                     "model": model_name,
-                    "embedding_dimension": semantic_engine.get_embedding_dimensions() or 0,
+                    "embedding_dimension": semantic_engine.get_embedding_dimensions()
+                    or 0,
                 }
 
         except (ValidationError, SQLAlchemyError) as e:
@@ -683,7 +803,9 @@ class SQLiteMemoryDatabase:
     ) -> SemanticSearchResponse:
         """Perform semantic search across tables using vector embeddings."""
         if not is_semantic_search_available():
-            raise ValidationError("Semantic search is not available. Please install sentence-transformers.")
+            raise ValidationError(
+                "Semantic search is not available. Please install sentence-transformers."
+            )
 
         if not query or not query.strip():
             raise ValidationError("Search query cannot be empty")
@@ -704,7 +826,9 @@ class SQLiteMemoryDatabase:
 
                     # Check if table has embedding column
                     if embedding_column not in [col.name for col in table.columns]:
-                        logging.warning(f"Table '{table_name}' does not have embedding column '{embedding_column}'")
+                        logging.warning(
+                            f"Table '{table_name}' does not have embedding column '{embedding_column}'"
+                        )
                         continue
 
                     # Get all rows with embeddings
@@ -728,7 +852,8 @@ class SQLiteMemoryDatabase:
                         text_cols = [
                             col.name
                             for col in table.columns
-                            if "TEXT" in str(col.type).upper() or "VARCHAR" in str(col.type).upper()
+                            if "TEXT" in str(col.type).upper()
+                            or "VARCHAR" in str(col.type).upper()
                         ]
                     else:
                         text_cols = text_columns
@@ -753,6 +878,11 @@ class SQLiteMemoryDatabase:
             all_results.sort(key=lambda x: x.get("similarity_score", 0), reverse=True)
             final_results = all_results[:limit]
 
+            # Remove embedding data from results to keep LLM responses clean
+            for result in final_results:
+                if embedding_column in result:
+                    del result[embedding_column]
+
             return {
                 "success": True,
                 "results": final_results,
@@ -779,7 +909,9 @@ class SQLiteMemoryDatabase:
     ) -> RelatedContentResponse:
         """Find content related to a specific row by semantic similarity."""
         if not is_semantic_search_available():
-            raise ValidationError("Semantic search is not available. Please install sentence-transformers.")
+            raise ValidationError(
+                "Semantic search is not available. Please install sentence-transformers."
+            )
 
         try:
             table = self._ensure_table_exists(table_name)
@@ -791,7 +923,9 @@ class SQLiteMemoryDatabase:
                 target_row = conn.execute(target_stmt).fetchone()
 
                 if not target_row:
-                    raise ValidationError(f"Row with id {row_id} not found in table '{table_name}'")
+                    raise ValidationError(
+                        f"Row with id {row_id} not found in table '{table_name}'"
+                    )
 
                 target_dict = dict(target_row._mapping)
 
@@ -862,13 +996,23 @@ class SQLiteMemoryDatabase:
                 for candidate_idx, similarity_score in similar_indices:
                     original_idx = valid_indices[candidate_idx]
                     row_dict = content_data[original_idx].copy()
+
+                    # Remove embedding data to avoid polluting LLM responses
+                    if embedding_column in row_dict:
+                        del row_dict[embedding_column]
+
                     row_dict["similarity_score"] = round(similarity_score, 3)
                     results.append(row_dict)
 
+                # Remove embedding from target_row as well
+                target_dict_clean = target_dict.copy()
+                if embedding_column in target_dict_clean:
+                    del target_dict_clean[embedding_column]
+
                 return {
                     "success": True,
                     "results": results,
-                    "target_row": target_dict,
+                    "target_row": target_dict_clean,
                     "total_results": len(results),
                     "similarity_threshold": similarity_threshold,
                     "model": model_name,
@@ -950,13 +1094,23 @@ class SQLiteMemoryDatabase:
             # Enhance with text matching scores
             try:
                 semantic_engine = get_semantic_engine(model_name)
-                
+
                 # Verify the engine has the required method
-                if not hasattr(semantic_engine, 'hybrid_search') or not callable(getattr(semantic_engine, 'hybrid_search')):
-                    raise DatabaseError("Semantic engine hybrid_search method is not callable")
-                
+                if not hasattr(semantic_engine, "hybrid_search") or not callable(
+                    getattr(semantic_engine, "hybrid_search")
+                ):
+                    raise DatabaseError(
+                        "Semantic engine hybrid_search method is not callable"
+                    )
+
                 enhanced_results = semantic_engine.hybrid_search(
-                    query, semantic_results, text_columns or [], embedding_column, semantic_weight, text_weight, limit
+                    query,
+                    semantic_results,
+                    text_columns or [],
+                    embedding_column,
+                    semantic_weight,
+                    text_weight,
+                    limit,
                 )
             except Exception as e:
                 # If semantic enhancement fails, return semantic results without text enhancement
@@ -979,7 +1133,9 @@ class SQLiteMemoryDatabase:
                 raise e
             raise DatabaseError(f"Hybrid search failed: {str(e)}")
 
-    def get_embedding_stats(self, table_name: str, embedding_column: str = "embedding") -> ToolResponse:
+    def get_embedding_stats(
+        self, table_name: str, embedding_column: str = "embedding"
+    ) -> ToolResponse:
         """Get statistics about embeddings in a table."""
         try:
             table = self._ensure_table_exists(table_name)
@@ -989,8 +1145,13 @@ class SQLiteMemoryDatabase:
                 # Return 0% coverage when column doesn't exist (for compatibility with tests)
                 total_count = 0
                 with self.get_connection() as conn:
-                    total_count = conn.execute(select(text("COUNT(*)")).select_from(table)).scalar() or 0
-                
+                    total_count = (
+                        conn.execute(
+                            select(text("COUNT(*)")).select_from(table)
+                        ).scalar()
+                        or 0
+                    )
+
                 return {
                     "success": True,
                     "table_name": table_name,
@@ -1003,7 +1164,10 @@ class SQLiteMemoryDatabase:
 
             with self.get_connection() as conn:
                 # Count total rows
-                total_count = conn.execute(select(text("COUNT(*)")).select_from(table)).scalar() or 0
+                total_count = (
+                    conn.execute(select(text("COUNT(*)")).select_from(table)).scalar()
+                    or 0
+                )
 
                 # Count rows with embeddings
                 embedded_count = (
@@ -1043,7 +1207,9 @@ class SQLiteMemoryDatabase:
                     except json.JSONDecodeError:
                         pass
 
-                coverage_percent = (embedded_count / total_count * 100) if total_count > 0 else 0.0
+                coverage_percent = (
+                    (embedded_count / total_count * 100) if total_count > 0 else 0.0
+                )
 
                 return {
                     "success": True,
@@ -1072,7 +1238,7 @@ def get_database(db_path: Optional[str] = None) -> SQLiteMemoryDatabase:
     actual_path = db_path or os.environ.get("DB_PATH", "./test.db")
     if actual_path is None:
         actual_path = "./test.db"
-    
+
     if _db_instance is None or (db_path and db_path != _db_instance.db_path):
         # Close previous instance if it exists
         if _db_instance is not None: