mcp-sqlite-memory-bank 1.2.3__py3-none-any.whl → 1.3.0__py3-none-any.whl

mcp_sqlite_memory_bank/types.py

@@ -12,17 +12,19 @@ from enum import Enum, auto
 
 class SqliteType(str, Enum):
     """Valid SQLite column types."""
+
     TEXT = "TEXT"
     INTEGER = "INTEGER"
     REAL = "REAL"
     BLOB = "BLOB"
     NULL = "NULL"
     TIMESTAMP = "TIMESTAMP"  # Common extension
-    BOOLEAN = "BOOLEAN"      # Common extension
+    BOOLEAN = "BOOLEAN"  # Common extension
 
 
 class ErrorCategory(Enum):
     """Categories for SQLite Memory Bank errors."""
+
     VALIDATION = auto()
     DATABASE = auto()
     SCHEMA = auto()
@@ -33,18 +35,14 @@ class ErrorCategory(Enum):
 @dataclass
 class MemoryBankError(Exception):
     """Base class for SQLite Memory Bank errors."""
+
     message: str
     category: ErrorCategory
     details: Optional[Dict[str, Any]] = None
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert error to a dict format suitable for FastMCP responses."""
-        return {
-            "success": False,
-            "error": self.message,
-            "category": self.category.name,
-            "details": self.details or {}
-        }
+        return {"success": False, "error": self.message, "category": self.category.name, "details": self.details or {}}
 
 
 class ValidationError(MemoryBankError):
@@ -77,6 +75,7 @@ class DataError(MemoryBankError):
 
 class TableColumn(TypedDict):
     """Structure for table column definitions."""
+
     name: str
     type: str
     notnull: bool
@@ -86,11 +85,13 @@ class TableColumn(TypedDict):
 
 class SuccessResponse(TypedDict):
     """Base structure for successful responses."""
+
     success: Literal[True]
 
 
 class ErrorResponse(TypedDict):
     """Structure for error responses."""
+
     success: Literal[False]
     error: str
     category: str
@@ -99,59 +100,146 @@ class ErrorResponse(TypedDict):
 
 class CreateTableResponse(SuccessResponse):
     """Response for create_table tool."""
+
     pass
 
 
 class DropTableResponse(SuccessResponse):
     """Response for drop_table tool."""
+
     pass
 
 
 class RenameTableResponse(SuccessResponse):
     """Response for rename_table tool."""
+
     pass
 
 
 class ListTablesResponse(SuccessResponse):
     """Response for list_tables tool."""
+
     tables: List[str]
 
 
 class DescribeTableResponse(SuccessResponse):
     """Response for describe_table tool."""
-    columns: List[TableColumn]
+
+    columns: List[Dict[str, Any]]
 
 
 class ListAllColumnsResponse(SuccessResponse):
     """Response for list_all_columns tool."""
+
     schemas: Dict[str, List[str]]
 
 
 class CreateRowResponse(SuccessResponse):
     """Response for create_row tool."""
+
     id: int
 
 
 class ReadRowsResponse(SuccessResponse):
     """Response for read_rows tool."""
+
     rows: List[Dict[str, Any]]
 
 
 class UpdateRowsResponse(SuccessResponse):
     """Response for update_rows tool."""
+
     rows_affected: int
 
 
 class DeleteRowsResponse(SuccessResponse):
     """Response for delete_rows tool."""
+
     rows_affected: int
 
 
 class SelectQueryResponse(SuccessResponse):
     """Response for run_select_query tool."""
+
     rows: List[Dict[str, Any]]
 
 
+class SearchContentResponse(SuccessResponse):
+    """Response for search_content tool."""
+
+    results: List[Dict[str, Any]]
+    query: str
+    tables_searched: List[str]
+    total_results: int
+
+
+class ExploreTablesResponse(SuccessResponse):
+    """Response for explore_tables tool."""
+
+    exploration: Dict[str, Any]
+
+
+# Semantic Search Response Types
+class SemanticSearchResponse(TypedDict):
+    """Response type for semantic search operations."""
+    success: bool
+    results: List[Dict[str, Any]]
+    query: str
+    tables_searched: List[str]
+    total_results: int
+    model: str
+    similarity_threshold: float
+
+
+class RelatedContentResponse(TypedDict, total=False):
+    """Response type for find related content operations."""
+    success: bool
+    results: List[Dict[str, Any]]
+    target_row: Dict[str, Any]
+    total_results: int
+    similarity_threshold: float
+    model: str
+    message: str  # Optional field
+
+
+class HybridSearchResponse(TypedDict):
+    """Response type for hybrid search operations."""
+    success: bool
+    results: List[Dict[str, Any]]
+    query: str
+    search_type: str
+    semantic_weight: float
+    text_weight: float
+    total_results: int
+    model: str
+
+
+class EmbeddingStatsResponse(TypedDict):
+    """Response type for embedding statistics."""
+    success: bool
+    table_name: str
+    total_rows: int
+    embedded_rows: int
+    coverage_percent: float
+    embedding_dimensions: Optional[int]
+    embedding_column: str
+
+
+class GenerateEmbeddingsResponse(TypedDict):
+    """Response type for embedding generation operations."""
+    success: bool
+    message: str
+    processed: int
+    model: str
+    embedding_dimension: int
+
+
+class EmbeddingColumnResponse(TypedDict):
+    """Response type for adding embedding columns."""
+    success: bool
+    message: str
+
+
 # Type alias for all possible responses
 ToolResponse = Union[
     CreateTableResponse,
@@ -165,5 +253,13 @@ ToolResponse = Union[
     UpdateRowsResponse,
     DeleteRowsResponse,
     SelectQueryResponse,
-    ErrorResponse
+    SearchContentResponse,
+    ExploreTablesResponse,
+    ErrorResponse,
+    SemanticSearchResponse,
+    RelatedContentResponse,
+    HybridSearchResponse,
+    EmbeddingStatsResponse,
+    GenerateEmbeddingsResponse,
+    EmbeddingColumnResponse,
 ]

mcp_sqlite_memory_bank/utils.py

@@ -11,15 +11,9 @@ import sqlite3
 import logging
 from functools import wraps
 from typing import Any, Callable, Dict, List, TypeVar, cast, Union, Tuple
-from .types import (
-    ValidationError,
-    DatabaseError,
-    SchemaError,
-    MemoryBankError,
-    ToolResponse
-)
+from .types import ValidationError, DatabaseError, SchemaError, MemoryBankError, ToolResponse
 
-T = TypeVar('T', bound=Callable[..., ToolResponse])
+T = TypeVar("T", bound=Callable[..., ToolResponse])
 
 
 def catch_errors(f: T) -> T:
@@ -27,6 +21,7 @@ def catch_errors(f: T) -> T:
     Decorator to standardize error handling across tools.
     Catches exceptions and converts them to appropriate error responses.
     """
+
     @wraps(f)
     def wrapper(*args: Any, **kwargs: Any) -> ToolResponse:
         try:
@@ -36,15 +31,11 @@ def catch_errors(f: T) -> T:
             return cast(ToolResponse, e.to_dict())
         except sqlite3.Error as e:
             logging.error(f"{f.__name__} database error: {e}")
-            return cast(ToolResponse, DatabaseError(
-                f"Database error in {f.__name__}: {e}",
-                {"sqlite_error": str(e)}
-            ).to_dict())
+            return cast(ToolResponse, DatabaseError(f"Database error in {f.__name__}: {e}", {"sqlite_error": str(e)}).to_dict())
         except Exception as e:
             logging.error(f"Unexpected error in {f.__name__}: {e}")
-            return cast(ToolResponse, DatabaseError(
-                f"Unexpected error in {f.__name__}: {e}"
-            ).to_dict())
+            return cast(ToolResponse, DatabaseError(f"Unexpected error in {f.__name__}: {e}").to_dict())
+
     return cast(T, wrapper)
 
 
@@ -57,11 +48,10 @@ def validate_identifier(name: str, context: str = "identifier") -> None:
         name: The identifier to validate
         context: Description of what's being validated (for error messages)
     """
-    if not bool(re.match(r'^[A-Za-z_][A-Za-z0-9_]*$', name)):
+    if not bool(re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", name)):
         raise ValidationError(
-            f"Invalid {context}: {name}. Must start with letter/underscore and "
-            f"contain only letters, numbers, underscores.",
-            {"invalid_name": name}
+            f"Invalid {context}: {name}. Must start with letter/underscore and " f"contain only letters, numbers, underscores.",
+            {"invalid_name": name},
         )
 
 
@@ -74,14 +64,11 @@ def validate_column_definition(column: Dict[str, Any]) -> None:
         column: Dictionary with column definition (must have 'name' and 'type' keys)
     """
     if not isinstance(column, dict):
-        raise ValidationError(
-            "Column definition must be a dictionary",
-            {"received": str(type(column))}
-        )
+        raise ValidationError("Column definition must be a dictionary", {"received": str(type(column))})
     if "name" not in column or "type" not in column:
         raise ValidationError(
             "Column definition must have 'name' and 'type' keys",
-            {"missing_keys": [k for k in ["name", "type"] if k not in column]}
+            {"missing_keys": [k for k in ["name", "type"] if k not in column]},
         )
     validate_identifier(column["name"], "column name")
 
@@ -103,17 +90,14 @@ def get_table_columns(conn: sqlite3.Connection, table_name: str) -> List[str]:
     cur.execute(f"PRAGMA table_info({table_name})")
     columns = [row[1] for row in cur.fetchall()]
     if not columns:
-        raise SchemaError(
-            f"Table does not exist: {table_name}",
-            {"table_name": table_name}
-        )
+        raise SchemaError(f"Table does not exist: {table_name}", {"table_name": table_name})
     return columns
 
+
 # Compatibility function for direct table_name usage
 
 
-def get_table_columns_by_name(
-        table_name: str) -> Union[List[str], Dict[str, Any]]:
+def get_table_columns_by_name(table_name: str) -> Union[List[str], Dict[str, Any]]:
     """
     Get list of column names for a table by name.
     Compatibility function for the old implementation.
@@ -128,14 +112,9 @@ def get_table_columns_by_name(
         validate_identifier(table_name, "table name")
         with sqlite3.connect(os.environ.get("DB_PATH", "./test.db")) as conn:
             cur = conn.cursor()
-            cur.execute(
-                "SELECT name FROM sqlite_master WHERE type='table' AND name=?",
-                (table_name,)
-            )
+            cur.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
             if not cur.fetchone():
-                return {
-                    "success": False,
-                    "error": f"Table '{table_name}' does not exist"}
+                return {"success": False, "error": f"Table '{table_name}' does not exist"}
 
             # Get column information
             cur.execute(f"PRAGMA table_info({table_name})")
@@ -144,9 +123,7 @@ def get_table_columns_by_name(
     except MemoryBankError as e:
         return e.to_dict()
     except Exception as e:
-        return {
-            "success": False,
-            "error": f"Exception in get_table_columns: {e}"}
+        return {"success": False, "error": f"Exception in get_table_columns: {e}"}
 
 
 def validate_table_exists(conn: sqlite3.Connection, table_name: str) -> None:
@@ -159,19 +136,12 @@ def validate_table_exists(conn: sqlite3.Connection, table_name: str) -> None:
         table_name: Name of table to check
     """
     cur = conn.cursor()
-    cur.execute(
-        "SELECT name FROM sqlite_master WHERE type='table' AND name=?",
-        (table_name,)
-    )
+    cur.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
     if not cur.fetchone():
-        raise SchemaError(
-            f"Table does not exist: {table_name}",
-            {"table_name": table_name}
-        )
+        raise SchemaError(f"Table does not exist: {table_name}", {"table_name": table_name})
 
 
-def build_where_clause(
-        where: Dict[str, Any], valid_columns: List[str]) -> Union[Tuple[str, list], Dict[str, Any]]:
+def build_where_clause(where: Dict[str, Any], valid_columns: List[str]) -> Union[Tuple[str, list], Dict[str, Any]]:
     """
     Build a WHERE clause from a dictionary of column-value pairs.
 
@@ -192,13 +162,11 @@ def build_where_clause(
 
         for col, val in where.items():
             if col not in valid_columns:
-                return {
-                    "success": False,
-                    "error": f"Invalid column in where clause: {col}"}
+                return {"success": False, "error": f"Invalid column in where clause: {col}"}
             conditions.append(f"{col}=?")
             values.append(val)
 
-        clause = ' AND '.join(conditions)
+        clause = " AND ".join(conditions)
         return clause, values
     except Exception as e:
         logging.error(f"Error in build_where_clause: {e}")