sqlsaber 0.14.0__py3-none-any.whl → 0.15.0__py3-none-any.whl

sqlsaber/agents/anthropic.py

@@ -21,6 +21,8 @@ from sqlsaber.config.settings import Config
 from sqlsaber.database.connection import BaseDatabaseConnection
 from sqlsaber.memory.manager import MemoryManager
 from sqlsaber.models.events import StreamEvent
+from sqlsaber.tools import tool_registry
+from sqlsaber.tools.instructions import InstructionBuilder
 
 
 class AnthropicSQLAgent(BaseSQLAgent):
@@ -51,89 +53,11 @@ class AnthropicSQLAgent(BaseSQLAgent):
         self._last_results = None
         self._last_query = None
 
-        # Define tools in the new format
-        self.tools: list[ToolDefinition] = [
-            ToolDefinition(
-                name="list_tables",
-                description="Get a list of all tables in the database with row counts. Use this first to discover available tables.",
-                input_schema={
-                    "type": "object",
-                    "properties": {},
-                    "required": [],
-                },
-            ),
-            ToolDefinition(
-                name="introspect_schema",
-                description="Introspect database schema to understand table structures.",
-                input_schema={
-                    "type": "object",
-                    "properties": {
-                        "table_pattern": {
-                            "type": "string",
-                            "description": "Optional pattern to filter tables (e.g., 'public.users', 'user%', '%order%')",
-                        }
-                    },
-                    "required": [],
-                },
-            ),
-            ToolDefinition(
-                name="execute_sql",
-                description="Execute a SQL query against the database.",
-                input_schema={
-                    "type": "object",
-                    "properties": {
-                        "query": {
-                            "type": "string",
-                            "description": "SQL query to execute",
-                        },
-                        "limit": {
-                            "type": "integer",
-                            "description": f"Maximum number of rows to return (default: {AnthropicSQLAgent.DEFAULT_SQL_LIMIT})",
-                            "default": AnthropicSQLAgent.DEFAULT_SQL_LIMIT,
-                        },
-                    },
-                    "required": ["query"],
-                },
-            ),
-            ToolDefinition(
-                name="plot_data",
-                description="Create a plot of query results.",
-                input_schema={
-                    "type": "object",
-                    "properties": {
-                        "y_values": {
-                            "type": "array",
-                            "items": {"type": ["number", "null"]},
-                            "description": "Y-axis data points (required)",
-                        },
-                        "x_values": {
-                            "type": "array",
-                            "items": {"type": ["number", "null"]},
-                            "description": "X-axis data points (optional, will use indices if not provided)",
-                        },
-                        "plot_type": {
-                            "type": "string",
-                            "enum": ["line", "scatter", "histogram"],
-                            "description": "Type of plot to create (default: line)",
-                            "default": "line",
-                        },
-                        "title": {
-                            "type": "string",
-                            "description": "Title for the plot",
-                        },
-                        "x_label": {
-                            "type": "string",
-                            "description": "Label for X-axis",
-                        },
-                        "y_label": {
-                            "type": "string",
-                            "description": "Label for Y-axis",
-                        },
-                    },
-                    "required": ["y_values"],
-                },
-            ),
-        ]
+        # Get tool definitions from registry
+        self.tools: list[ToolDefinition] = tool_registry.get_tool_definitions()
+
+        # Initialize instruction builder
+        self.instruction_builder = InstructionBuilder(tool_registry)
 
         # Build system prompt with memories if available
         self.system_prompt = self._build_system_prompt()
@@ -157,31 +81,9 @@ class AnthropicSQLAgent(BaseSQLAgent):
     def _get_sql_assistant_instructions(self) -> str:
         """Get the detailed SQL assistant instructions."""
         db_type = self._get_database_type_name()
-        instructions = f"""You are also a helpful SQL assistant that helps users query their {db_type} database.
-
-Your responsibilities:
-1. Understand user's natural language requests, think and convert them to SQL
-2. Use the provided tools efficiently to explore database schema
-3. Generate appropriate SQL queries
-4. Execute queries safely - queries that modify the database are not allowed
-5. Format and explain results clearly
-6. Create visualizations when requested or when they would be helpful
-
-IMPORTANT - Schema Discovery Strategy:
-1. ALWAYS start with 'list_tables' to see available tables and row counts
-2. Based on the user's query, identify which specific tables are relevant
-3. Use 'introspect_schema' with a table_pattern to get details ONLY for relevant tables
-4. Timestamp columns must be converted to text when you write queries
-
-Guidelines:
-- Use list_tables first, then introspect_schema for specific tables only
-- Use table patterns like 'sample%' or '%experiment%' to filter related tables
-- Use proper JOIN syntax and avoid cartesian products
-- Include appropriate WHERE clauses to limit results
-- Explain what the query does in simple terms
-- Handle errors gracefully and suggest fixes
-- Be security conscious - use parameterized queries when needed
-"""
+
+        # Build dynamic instructions from available tools
+        instructions = self.instruction_builder.build_instructions(db_type=db_type)
 
         # Add memory context if database name is available
         if self.database_name:
@@ -189,7 +91,7 @@ Guidelines:
                 self.database_name
             )
             if memory_context.strip():
-                instructions += memory_context
+                instructions += "\n\n" + memory_context
 
         return instructions
 
@@ -199,16 +101,19 @@ Guidelines:
             return None
 
         memory = self.memory_manager.add_memory(self.database_name, content)
-        # Rebuild system prompt with new memory
+        # Rebuild system prompt with new memory (includes dynamic instructions)
         self.system_prompt = self._build_system_prompt()
         return memory.id
 
-    async def execute_sql(self, query: str, limit: int | None = None) -> str:
-        """Execute a SQL query against the database with streaming support."""
-        # Call parent implementation for core functionality
-        result = await super().execute_sql(query, limit)
+    async def _execute_sql_with_tracking(
+        self, query: str, limit: int | None = None
+    ) -> str:
+        """Execute SQL and track results for streaming."""
+        # Get the execute_sql tool and run it
+        tool = tool_registry.get_tool("execute_sql")
+        result = await tool.execute(query=query, limit=limit)
 
-        # Parse result to extract data for streaming (AnthropicSQLAgent specific)
+        # Parse result to extract data for streaming
         try:
             result_data = json.loads(result)
             if result_data.get("success") and "results" in result_data:
@@ -228,7 +133,14 @@ Guidelines:
         self, tool_name: str, tool_input: dict[str, Any]
     ) -> str:
         """Process a tool call and return the result."""
-        # Use parent implementation for core tools
+        # Special handling for execute_sql to track results
+        if tool_name == "execute_sql":
+            return await self._execute_sql_with_tracking(
+                tool_input.get("query", ""),
+                tool_input.get("limit", self.DEFAULT_SQL_LIMIT),
+            )
+
+        # Use parent implementation for all other tools
         return await super().process_tool_call(tool_name, tool_input)
 
     def _convert_user_message_to_message(

sqlsaber/agents/base.py

@@ -5,8 +5,6 @@ import json
 from abc import ABC, abstractmethod
 from typing import Any, AsyncIterator
 
-from uniplot import histogram, plot
-
 from sqlsaber.conversation.manager import ConversationManager
 from sqlsaber.database.connection import (
     BaseDatabaseConnection,
@@ -17,6 +15,7 @@ from sqlsaber.database.connection import (
 )
 from sqlsaber.database.schema import SchemaManager
 from sqlsaber.models.events import StreamEvent
+from sqlsaber.tools import SQLTool, tool_registry
 
 
 class BaseSQLAgent(ABC):
@@ -32,6 +31,9 @@ class BaseSQLAgent(ABC):
         self._conversation_id: str | None = None
         self._msg_index: int = 0
 
+        # Initialize SQL tools with database connection
+        self._init_tools()
+
     @abstractmethod
     async def query_stream(
         self,
@@ -69,232 +71,28 @@ class BaseSQLAgent(ABC):
         else:
             return "database"  # Fallback
 
-    async def introspect_schema(self, table_pattern: str | None = None) -> str:
-        """Introspect database schema to understand table structures."""
-        try:
-            # Pass table_pattern to get_schema_info for efficient filtering at DB level
-            schema_info = await self.schema_manager.get_schema_info(table_pattern)
-
-            # Format the schema information
-            formatted_info = {}
-            for table_name, table_info in schema_info.items():
-                formatted_info[table_name] = {
-                    "columns": {
-                        col_name: {
-                            "type": col_info["data_type"],
-                            "nullable": col_info["nullable"],
-                            "default": col_info["default"],
-                        }
-                        for col_name, col_info in table_info["columns"].items()
-                    },
-                    "primary_keys": table_info["primary_keys"],
-                    "foreign_keys": [
-                        f"{fk['column']} -> {fk['references']['table']}.{fk['references']['column']}"
-                        for fk in table_info["foreign_keys"]
-                    ],
-                }
-
-            return json.dumps(formatted_info)
-        except Exception as e:
-            return json.dumps({"error": f"Error introspecting schema: {str(e)}"})
-
-    async def list_tables(self) -> str:
-        """List all tables in the database with basic information."""
-        try:
-            tables_info = await self.schema_manager.list_tables()
-            return json.dumps(tables_info)
-        except Exception as e:
-            return json.dumps({"error": f"Error listing tables: {str(e)}"})
-
-    async def execute_sql(self, query: str, limit: int | None = None) -> str:
-        """Execute a SQL query against the database."""
-        try:
-            # Security check - only allow SELECT queries unless write is enabled
-            write_error = self._validate_write_operation(query)
-            if write_error:
-                return json.dumps(
-                    {
-                        "error": write_error,
-                    }
-                )
-
-            # Add LIMIT if not present and it's a SELECT query
-            query = self._add_limit_to_query(query, limit)
-
-            # Execute the query (wrapped in a transaction for safety)
-            results = await self.db.execute_query(query)
-
-            # Format results
-            actual_limit = limit if limit is not None else len(results)
-
-            return json.dumps(
-                {
-                    "success": True,
-                    "row_count": len(results),
-                    "results": results[:actual_limit],  # Extra safety for limit
-                    "truncated": len(results) > actual_limit,
-                }
-            )
-
-        except Exception as e:
-            error_msg = str(e)
-
-            # Provide helpful error messages
-            suggestions = []
-            if "column" in error_msg.lower() and "does not exist" in error_msg.lower():
-                suggestions.append(
-                    "Check column names using the schema introspection tool"
-                )
-            elif "table" in error_msg.lower() and "does not exist" in error_msg.lower():
-                suggestions.append(
-                    "Check table names using the schema introspection tool"
-                )
-            elif "syntax error" in error_msg.lower():
-                suggestions.append(
-                    "Review SQL syntax, especially JOIN conditions and WHERE clauses"
-                )
-
-            return json.dumps({"error": error_msg, "suggestions": suggestions})
+    def _init_tools(self) -> None:
+        """Initialize SQL tools with database connection."""
+        # Get all SQL tools and set their database connection
+        for tool_name in tool_registry.list_tools(category="sql"):
+            tool = tool_registry.get_tool(tool_name)
+            if isinstance(tool, SQLTool):
+                tool.set_connection(self.db)
 
     async def process_tool_call(
         self, tool_name: str, tool_input: dict[str, Any]
     ) -> str:
         """Process a tool call and return the result."""
-        if tool_name == "list_tables":
-            return await self.list_tables()
-        elif tool_name == "introspect_schema":
-            return await self.introspect_schema(tool_input.get("table_pattern"))
-        elif tool_name == "execute_sql":
-            return await self.execute_sql(
-                tool_input["query"], tool_input.get("limit", 100)
-            )
-        elif tool_name == "plot_data":
-            return await self.plot_data(
-                y_values=tool_input["y_values"],
-                x_values=tool_input.get("x_values"),
-                plot_type=tool_input.get("plot_type", "line"),
-                title=tool_input.get("title"),
-                x_label=tool_input.get("x_label"),
-                y_label=tool_input.get("y_label"),
-            )
-        else:
-            return json.dumps({"error": f"Unknown tool: {tool_name}"})
-
-    def _validate_write_operation(self, query: str) -> str | None:
-        """Validate if a write operation is allowed.
-
-        Returns:
-            None if operation is allowed, error message if not allowed.
-        """
-        query_upper = query.strip().upper()
-
-        # Check for write operations
-        write_keywords = [
-            "INSERT",
-            "UPDATE",
-            "DELETE",
-            "DROP",
-            "CREATE",
-            "ALTER",
-            "TRUNCATE",
-        ]
-        is_write_query = any(query_upper.startswith(kw) for kw in write_keywords)
-
-        if is_write_query:
-            return (
-                "Write operations are not allowed. Only SELECT queries are permitted."
-            )
-
-        return None
-
-    def _add_limit_to_query(self, query: str, limit: int = 100) -> str:
-        """Add LIMIT clause to SELECT queries if not present."""
-        query_upper = query.strip().upper()
-        if query_upper.startswith("SELECT") and "LIMIT" not in query_upper:
-            return f"{query.rstrip(';')} LIMIT {limit};"
-        return query
-
-    async def plot_data(
-        self,
-        y_values: list[float],
-        x_values: list[float] | None = None,
-        plot_type: str = "line",
-        title: str | None = None,
-        x_label: str | None = None,
-        y_label: str | None = None,
-    ) -> str:
-        """Create a terminal plot using uniplot.
-
-        Args:
-            y_values: Y-axis data points
-            x_values: X-axis data points (optional)
-            plot_type: Type of plot - "line", "scatter", or "histogram"
-            title: Plot title
-            x_label: X-axis label
-            y_label: Y-axis label
-
-        Returns:
-            JSON string with success status and plot details
-        """
         try:
-            # Validate inputs
-            if not y_values:
-                return json.dumps({"error": "No data provided for plotting"})
-
-            # Convert to floats if needed
-            try:
-                y_values = [float(v) if v is not None else None for v in y_values]
-                if x_values:
-                    x_values = [float(v) if v is not None else None for v in x_values]
-            except (ValueError, TypeError) as e:
-                return json.dumps({"error": f"Invalid data format: {str(e)}"})
-
-            # Create the plot
-            if plot_type == "histogram":
-                # For histogram, we only need y_values
-                histogram(
-                    y_values,
-                    title=title,
-                    bins=min(20, len(set(y_values))),  # Adaptive bin count
-                )
-                plot_info = {
-                    "type": "histogram",
-                    "data_points": len(y_values),
-                    "title": title or "Histogram",
-                }
-            elif plot_type in ["line", "scatter"]:
-                # For line/scatter plots
-                plot_kwargs = {
-                    "ys": y_values,
-                    "title": title,
-                    "lines": plot_type == "line",
-                }
-
-                if x_values:
-                    plot_kwargs["xs"] = x_values
-                if x_label:
-                    plot_kwargs["x_unit"] = x_label
-                if y_label:
-                    plot_kwargs["y_unit"] = y_label
-
-                plot(**plot_kwargs)
-
-                plot_info = {
-                    "type": plot_type,
-                    "data_points": len(y_values),
-                    "title": title or f"{plot_type.capitalize()} Plot",
-                    "has_x_values": x_values is not None,
-                }
-            else:
-                return json.dumps({"error": f"Unsupported plot type: {plot_type}"})
-
+            tool = tool_registry.get_tool(tool_name)
+            return await tool.execute(**tool_input)
+        except KeyError:
+            return json.dumps({"error": f"Unknown tool: {tool_name}"})
+        except Exception as e:
             return json.dumps(
-                {"success": True, "plot_rendered": True, "plot_info": plot_info}
+                {"error": f"Error executing tool '{tool_name}': {str(e)}"}
             )
 
-        except Exception as e:
-            return json.dumps({"error": f"Error creating plot: {str(e)}"})
-
     # Conversation persistence helpers
 
     async def _ensure_conversation(self) -> None:

sqlsaber/mcp/mcp.py

@@ -7,25 +7,17 @@ from fastmcp import FastMCP
 from sqlsaber.agents.mcp import MCPSQLAgent
 from sqlsaber.config.database import DatabaseConfigManager
 from sqlsaber.database.connection import DatabaseConnection
+from sqlsaber.tools import SQLTool, tool_registry
+from sqlsaber.tools.instructions import InstructionBuilder
 
-INSTRUCTIONS = """
-This server provides helpful resources and tools that will help you address users queries on their database.
+# Initialize the instruction builder
+instruction_builder = InstructionBuilder(tool_registry)
 
-- Get all databases using `get_databases()`
-- Call `list_tables()` to get a list of all tables in the database with row counts. Use this first to discover available tables.
-- Call `introspect_schema()` to introspect database schema to understand table structures.
-- Call `execute_sql()` to execute SQL queries against the database and retrieve results.
+# Generate dynamic instructions
+DYNAMIC_INSTRUCTIONS = instruction_builder.build_mcp_instructions()
 
-Guidelines:
-- Use list_tables first, then introspect_schema for specific tables only
-- Use table patterns like 'sample%' or '%experiment%' to filter related tables
-- Use proper JOIN syntax and avoid cartesian products
-- Include appropriate WHERE clauses to limit results
-- Handle errors gracefully and suggest fixes
-"""
-
-# Create the FastMCP server instance
-mcp = FastMCP(name="SQL Assistant", instructions=INSTRUCTIONS)
+# Create the FastMCP server instance with dynamic instructions
+mcp = FastMCP(name="SQL Assistant", instructions=DYNAMIC_INSTRUCTIONS)
 
 # Initialize the database config manager
 config_manager = DatabaseConfigManager()
@@ -70,10 +62,16 @@ def get_databases() -> dict:
     return {"databases": databases, "count": len(databases)}
 
 
-@mcp.tool
-async def list_tables(database: str) -> str:
-    """
-    Get a list of all tables in the database with row counts. Use this first to discover available tables.
+async def _execute_with_connection(tool_name: str, database: str, **kwargs) -> str:
+    """Execute a SQL tool with database connection management.
+
+    Args:
+        tool_name: Name of the tool to execute
+        database: Database name to connect to
+        **kwargs: Tool-specific parameters
+
+    Returns:
+        JSON string with the tool's output
     """
     try:
         agent = await _create_agent_for_database(database)
@@ -82,50 +80,44 @@ async def list_tables(database: str) -> str:
                 {"error": f"Database '{database}' not found or could not connect"}
             )
 
-        result = await agent.list_tables()
+        # Get the tool and set up connection
+        tool = tool_registry.get_tool(tool_name)
+        if isinstance(tool, SQLTool):
+            tool.set_connection(agent.db)
+
+        # Execute the tool
+        result = await tool.execute(**kwargs)
         await agent.db.close()
         return result
 
     except Exception as e:
-        return json.dumps({"error": f"Error listing tables: {str(e)}"})
+        return json.dumps({"error": f"Error in {tool_name}: {str(e)}"})
 
 
-@mcp.tool
-async def introspect_schema(database: str, table_pattern: str | None = None) -> str:
-    """
-    Introspect database schema to understand table structures. Use optional pattern to filter tables (e.g., 'public.users', 'user%', '%order%').
-    """
-    try:
-        agent = await _create_agent_for_database(database)
-        if not agent:
-            return json.dumps(
-                {"error": f"Database '{database}' not found or could not connect"}
-            )
+# SQL Tool Wrappers with explicit signatures
 
-        result = await agent.introspect_schema(table_pattern)
-        await agent.db.close()
-        return result
 
-    except Exception as e:
-        return json.dumps({"error": f"Error introspecting schema: {str(e)}"})
+@mcp.tool
+async def list_tables(database: str) -> str:
+    """Get a list of all tables in the database with row counts. Use this first to discover available tables."""
+    return await _execute_with_connection("list_tables", database)
 
 
 @mcp.tool
-async def execute_sql(database: str, query: str, limit: int | None = 100) -> str:
-    """Execute a SQL query against the specified database."""
-    try:
-        agent = await _create_agent_for_database(database)
-        if not agent:
-            return json.dumps(
-                {"error": f"Database '{database}' not found or could not connect"}
-            )
+async def introspect_schema(database: str, table_pattern: str = None) -> str:
+    """Introspect database schema to understand table structures."""
+    kwargs = {}
+    if table_pattern is not None:
+        kwargs["table_pattern"] = table_pattern
+    return await _execute_with_connection("introspect_schema", database, **kwargs)
 
-        result = await agent.execute_sql(query, limit)
-        await agent.db.close()
-        return result
 
-    except Exception as e:
-        return json.dumps({"error": f"Error executing SQL: {str(e)}"})
+@mcp.tool
+async def execute_sql(database: str, query: str, limit: int = 100) -> str:
+    """Execute a SQL query against the database."""
+    return await _execute_with_connection(
+        "execute_sql", database, query=query, limit=limit
+    )
 
 
 def main():

sqlsaber/tools/__init__.py

@@ -0,0 +1,25 @@
+"""SQLSaber tools module."""
+
+from .base import Tool
+from .enums import ToolCategory, WorkflowPosition
+from .instructions import InstructionBuilder
+from .registry import ToolRegistry, register_tool, tool_registry
+
+# Import concrete tools to register them
+from .sql_tools import ExecuteSQLTool, IntrospectSchemaTool, ListTablesTool, SQLTool
+from .visualization_tools import PlotDataTool
+
+__all__ = [
+    "Tool",
+    "ToolCategory",
+    "WorkflowPosition",
+    "ToolRegistry",
+    "tool_registry",
+    "register_tool",
+    "InstructionBuilder",
+    "SQLTool",
+    "ListTablesTool",
+    "IntrospectSchemaTool",
+    "ExecuteSQLTool",
+    "PlotDataTool",
+]