sqlsaber 0.14.0__tar.gz → 0.15.0__tar.gz

{sqlsaber-0.14.0 → sqlsaber-0.15.0}/CHANGELOG.md

@@ -4,6 +4,20 @@ All notable changes to SQLSaber will be documented in this file.
 
 ## [Unreleased]
 
+## [0.15.0] - 2025-08-18
+
+### Added
+
+- Tool abstraction system with centralized registry (new `Tool` base class, `ToolRegistry`, decorators)
+- Dynamic instruction generation system (`InstructionBuilder`)
+- Comprehensive test suite for the tools module
+
+### Changed
+
+- Refactored agents to use centralized tool registry instead of hardcoded tools
+- Enhanced MCP server with dynamic tool registration
+- Moved core SQL functionality to dedicated tool classes
+
 ## [0.14.0] - 2025-08-01
 
 ### Added

{sqlsaber-0.14.0 → sqlsaber-0.15.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlsaber
-Version: 0.14.0
+Version: 0.15.0
 Summary: SQLSaber - Agentic SQL assistant like Claude Code
 License-File: LICENSE
 Requires-Python: >=3.12

{sqlsaber-0.14.0 → sqlsaber-0.15.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sqlsaber"
-version = "0.14.0"
+version = "0.15.0"
 description = "SQLSaber - Agentic SQL assistant like Claude Code"
 readme = "README.md"
 requires-python = ">=3.12"

{sqlsaber-0.14.0 → sqlsaber-0.15.0}/src/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-0.15.0/src/sqlsaber/agents/base.py

@@ -0,0 +1,187 @@
+"""Abstract base class for SQL agents."""
+
+import asyncio
+import json
+from abc import ABC, abstractmethod
+from typing import Any, AsyncIterator
+
+from sqlsaber.conversation.manager import ConversationManager
+from sqlsaber.database.connection import (
+    BaseDatabaseConnection,
+    CSVConnection,
+    MySQLConnection,
+    PostgreSQLConnection,
+    SQLiteConnection,
+)
+from sqlsaber.database.schema import SchemaManager
+from sqlsaber.models.events import StreamEvent
+from sqlsaber.tools import SQLTool, tool_registry
+
+
+class BaseSQLAgent(ABC):
+    """Abstract base class for SQL agents."""
+
+    def __init__(self, db_connection: BaseDatabaseConnection):
+        self.db = db_connection
+        self.schema_manager = SchemaManager(db_connection)
+        self.conversation_history: list[dict[str, Any]] = []
+
+        # Conversation persistence
+        self._conv_manager = ConversationManager()
+        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,
+        user_query: str,
+        use_history: bool = True,
+        cancellation_token: asyncio.Event | None = None,
+    ) -> AsyncIterator[StreamEvent]:
+        """Process a user query and stream responses.
+
+        Args:
+            user_query: The user's query to process
+            use_history: Whether to include conversation history
+            cancellation_token: Optional event to signal cancellation
+        """
+        pass
+
+    async def clear_history(self):
+        """Clear conversation history."""
+        # End current conversation in storage
+        await self._end_conversation()
+
+        # Clear in-memory history
+        self.conversation_history = []
+
+    def _get_database_type_name(self) -> str:
+        """Get the human-readable database type name."""
+        if isinstance(self.db, PostgreSQLConnection):
+            return "PostgreSQL"
+        elif isinstance(self.db, MySQLConnection):
+            return "MySQL"
+        elif isinstance(self.db, SQLiteConnection):
+            return "SQLite"
+        elif isinstance(self.db, CSVConnection):
+            return "SQLite"  # we convert csv to in-memory sqlite
+        else:
+            return "database"  # Fallback
+
+    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."""
+        try:
+            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(
+                {"error": f"Error executing tool '{tool_name}': {str(e)}"}
+            )
+
+    # Conversation persistence helpers
+
+    async def _ensure_conversation(self) -> None:
+        """Ensure a conversation is active for storing messages."""
+        if self._conversation_id is None:
+            db_name = getattr(self, "database_name", "unknown")
+            self._conversation_id = await self._conv_manager.start_conversation(db_name)
+            self._msg_index = 0
+
+    async def _store_user_message(self, content: str | dict[str, Any]) -> None:
+        """Store a user message in conversation history."""
+        if self._conversation_id is None:
+            return
+
+        await self._conv_manager.add_user_message(
+            self._conversation_id, content, self._msg_index
+        )
+        self._msg_index += 1
+
+    async def _store_assistant_message(
+        self, content: list[dict[str, Any]] | dict[str, Any]
+    ) -> None:
+        """Store an assistant message in conversation history."""
+        if self._conversation_id is None:
+            return
+
+        await self._conv_manager.add_assistant_message(
+            self._conversation_id, content, self._msg_index
+        )
+        self._msg_index += 1
+
+    async def _store_tool_message(
+        self, content: list[dict[str, Any]] | dict[str, Any]
+    ) -> None:
+        """Store a tool/system message in conversation history."""
+        if self._conversation_id is None:
+            return
+
+        await self._conv_manager.add_tool_message(
+            self._conversation_id, content, self._msg_index
+        )
+        self._msg_index += 1
+
+    async def _end_conversation(self) -> None:
+        """End the current conversation."""
+        if self._conversation_id:
+            await self._conv_manager.end_conversation(self._conversation_id)
+        self._conversation_id = None
+        self._msg_index = 0
+
+    async def restore_conversation(self, conversation_id: str) -> bool:
+        """Restore a conversation from storage to in-memory history.
+
+        Args:
+            conversation_id: ID of the conversation to restore
+
+        Returns:
+            True if successfully restored, False otherwise
+        """
+        success = await self._conv_manager.restore_conversation_to_agent(
+            conversation_id, self.conversation_history
+        )
+
+        if success:
+            # Set up for continuing this conversation
+            self._conversation_id = conversation_id
+            self._msg_index = len(self.conversation_history)
+
+        return success
+
+    async def list_conversations(self, limit: int = 50) -> list:
+        """List conversations for this agent's database.
+
+        Args:
+            limit: Maximum number of conversations to return
+
+        Returns:
+            List of conversation data
+        """
+        db_name = getattr(self, "database_name", None)
+        conversations = await self._conv_manager.list_conversations(db_name, limit)
+
+        return [
+            {
+                "id": conv.id,
+                "database_name": conv.database_name,
+                "started_at": conv.formatted_start_time(),
+                "ended_at": conv.formatted_end_time(),
+                "duration": conv.duration_seconds(),
+            }
+            for conv in conversations
+        ]

sqlsaber-0.15.0/src/sqlsaber/mcp/mcp.py

@@ -0,0 +1,129 @@
+"""FastMCP server implementation for SQLSaber."""
+
+import json
+
+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
+
+# Initialize the instruction builder
+instruction_builder = InstructionBuilder(tool_registry)
+
+# Generate dynamic instructions
+DYNAMIC_INSTRUCTIONS = instruction_builder.build_mcp_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()
+
+
+async def _create_agent_for_database(database_name: str) -> MCPSQLAgent | None:
+    """Create a MCPSQLAgent for the specified database."""
+    try:
+        # Look up configured database connection
+        db_config = config_manager.get_database(database_name)
+        if not db_config:
+            return None
+        connection_string = db_config.to_connection_string()
+
+        # Create database connection
+        db_conn = DatabaseConnection(connection_string)
+
+        # Create and return the agent
+        agent = MCPSQLAgent(db_conn)
+        return agent
+
+    except Exception:
+        return None
+
+
+@mcp.tool
+def get_databases() -> dict:
+    """List all configured databases with their types."""
+    databases = []
+    for db_config in config_manager.list_databases():
+        databases.append(
+            {
+                "name": db_config.name,
+                "type": db_config.type,
+                "database": db_config.database,
+                "host": db_config.host,
+                "port": db_config.port,
+                "is_default": db_config.name == config_manager.get_default_name(),
+            }
+        )
+
+    return {"databases": databases, "count": len(databases)}
+
+
+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)
+        if not agent:
+            return json.dumps(
+                {"error": f"Database '{database}' not found or could not connect"}
+            )
+
+        # 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 in {tool_name}: {str(e)}"})
+
+
+# SQL Tool Wrappers with explicit signatures
+
+
+@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 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)
+
+
+@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():
+    """Entry point for the MCP server console script."""
+    mcp.run()
+
+
+if __name__ == "__main__":
+    main()

sqlsaber-0.15.0/src/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",
+]

sqlsaber-0.15.0/src/sqlsaber/tools/base.py

@@ -0,0 +1,83 @@
+"""Base class for SQLSaber tools."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from sqlsaber.clients.models import ToolDefinition
+
+from .enums import ToolCategory, WorkflowPosition
+
+
+class Tool(ABC):
+    """Abstract base class for all tools."""
+
+    def __init__(self):
+        """Initialize the tool."""
+        pass
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Return the tool name."""
+        pass
+
+    @property
+    @abstractmethod
+    def description(self) -> str:
+        """Return the tool description."""
+        pass
+
+    @property
+    @abstractmethod
+    def input_schema(self) -> dict[str, Any]:
+        """Return the tool's input schema."""
+        pass
+
+    @abstractmethod
+    async def execute(self, **kwargs) -> str:
+        """Execute the tool with given inputs.
+
+        Args:
+            **kwargs: Tool-specific keyword arguments
+
+        Returns:
+            JSON string with the tool's output
+        """
+        pass
+
+    def to_definition(self) -> ToolDefinition:
+        """Convert this tool to a ToolDefinition."""
+        return ToolDefinition(
+            name=self.name,
+            description=self.description,
+            input_schema=self.input_schema,
+        )
+
+    @property
+    def category(self) -> ToolCategory:
+        """Return the tool category. Override to customize."""
+        return ToolCategory.GENERAL
+
+    def get_usage_instructions(self) -> str | None:
+        """Return tool-specific usage instructions for LLM guidance.
+
+        Returns:
+            Usage instructions string, or None for no specific guidance
+        """
+        return None
+
+    def get_priority(self) -> int:
+        """Return priority for tool ordering in instructions.
+
+        Returns:
+            Priority number (lower = higher priority, default = 100)
+        """
+        return 100
+
+    def get_workflow_position(self) -> WorkflowPosition:
+        """Return the typical workflow position for this tool.
+
+        Returns:
+            WorkflowPosition enum value
+        """
+        return WorkflowPosition.OTHER

sqlsaber-0.15.0/src/sqlsaber/tools/enums.py

@@ -0,0 +1,21 @@
+"""Enums for tool categories and workflow positions."""
+
+from enum import Enum
+
+
+class ToolCategory(Enum):
+    """Tool categories for organizing and filtering tools."""
+
+    GENERAL = "general"
+    SQL = "sql"
+    VISUALIZATION = "visualization"
+
+
+class WorkflowPosition(Enum):
+    """Workflow positions for organizing tools by usage order."""
+
+    DISCOVERY = "discovery"
+    ANALYSIS = "analysis"
+    EXECUTION = "execution"
+    VISUALIZATION = "visualization"
+    OTHER = "other"