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

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/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"

sqlsaber/tools/instructions.py

@@ -0,0 +1,251 @@
+"""Dynamic instruction builder for tools."""
+
+from .base import Tool
+from .enums import ToolCategory, WorkflowPosition
+from .registry import ToolRegistry
+
+
+class InstructionBuilder:
+    """Builds dynamic instructions based on available tools."""
+
+    def __init__(self, tool_registry: ToolRegistry):
+        """Initialize with a tool registry."""
+        self.registry = tool_registry
+
+    def build_instructions(
+        self,
+        db_type: str = "database",
+        category: str | ToolCategory | None = None,
+        include_base_instructions: bool = True,
+    ) -> str:
+        """Build dynamic instructions from available tools.
+
+        Args:
+            db_type: Type of database (PostgreSQL, MySQL, SQLite, etc.)
+            category: Optional category to filter tools by (string or ToolCategory enum)
+            include_base_instructions: Whether to include base SQL assistant instructions
+
+        Returns:
+            Complete instruction string for LLM
+        """
+        # Get available tools
+        tools = self.registry.get_all_tools(category)
+
+        if not tools:
+            return self._get_base_instructions(db_type)
+
+        # Sort tools by priority and workflow position
+        sorted_tools = self._sort_tools_by_workflow(tools)
+
+        # Build instruction components
+        instructions_parts = []
+
+        if include_base_instructions:
+            instructions_parts.append(self._get_base_instructions(db_type))
+
+        # Add tool-specific workflow guidance
+        workflow_instructions = self._build_workflow_instructions(sorted_tools)
+        if workflow_instructions:
+            instructions_parts.append(workflow_instructions)
+
+        # Add tool descriptions and guidelines
+        tool_guidelines = self._build_tool_guidelines(sorted_tools)
+        if tool_guidelines:
+            instructions_parts.append(tool_guidelines)
+
+        # Add general guidelines
+        general_guidelines = self._build_general_guidelines(sorted_tools)
+        if general_guidelines:
+            instructions_parts.append(general_guidelines)
+
+        return "\n\n".join(instructions_parts)
+
+    def _get_base_instructions(self, db_type: str) -> str:
+        """Get base SQL assistant instructions."""
+        return 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"""
+
+    def _sort_tools_by_workflow(self, tools: list[Tool]) -> list[Tool]:
+        """Sort tools by priority and workflow position."""
+        # Define workflow position ordering
+        position_order = {
+            WorkflowPosition.DISCOVERY: 1,
+            WorkflowPosition.ANALYSIS: 2,
+            WorkflowPosition.EXECUTION: 3,
+            WorkflowPosition.VISUALIZATION: 4,
+            WorkflowPosition.OTHER: 5,
+        }
+
+        return sorted(
+            tools,
+            key=lambda tool: (
+                position_order.get(tool.get_workflow_position(), 5),
+                tool.get_priority(),
+                tool.name,
+            ),
+        )
+
+    def _build_workflow_instructions(self, sorted_tools: list[Tool]) -> str:
+        """Build workflow-based instructions."""
+        # Group tools by workflow position
+        workflow_groups = {}
+        for tool in sorted_tools:
+            position = tool.get_workflow_position()
+            if position not in workflow_groups:
+                workflow_groups[position] = []
+            workflow_groups[position].append(tool)
+
+        # Build workflow instructions
+        instructions = ["IMPORTANT - Tool Usage Strategy:"]
+        step = 1
+
+        # Add discovery tools first
+        if WorkflowPosition.DISCOVERY in workflow_groups:
+            discovery_tools = workflow_groups[WorkflowPosition.DISCOVERY]
+            for tool in discovery_tools:
+                usage = tool.get_usage_instructions()
+                if usage:
+                    instructions.append(f"{step}. {usage}")
+                else:
+                    instructions.append(
+                        f"{step}. Use '{tool.name}' to {tool.description.lower()}"
+                    )
+                step += 1
+
+        # Add analysis tools
+        if WorkflowPosition.ANALYSIS in workflow_groups:
+            analysis_tools = workflow_groups[WorkflowPosition.ANALYSIS]
+            for tool in analysis_tools:
+                usage = tool.get_usage_instructions()
+                if usage:
+                    instructions.append(f"{step}. {usage}")
+                else:
+                    instructions.append(
+                        f"{step}. Use '{tool.name}' to {tool.description.lower()}"
+                    )
+                step += 1
+
+        # Add execution tools
+        if WorkflowPosition.EXECUTION in workflow_groups:
+            execution_tools = workflow_groups[WorkflowPosition.EXECUTION]
+            for tool in execution_tools:
+                usage = tool.get_usage_instructions()
+                if usage:
+                    instructions.append(f"{step}. {usage}")
+                else:
+                    instructions.append(
+                        f"{step}. Use '{tool.name}' to {tool.description.lower()}"
+                    )
+                step += 1
+
+        # Add visualization tools
+        if WorkflowPosition.VISUALIZATION in workflow_groups:
+            viz_tools = workflow_groups[WorkflowPosition.VISUALIZATION]
+            for tool in viz_tools:
+                usage = tool.get_usage_instructions()
+                if usage:
+                    instructions.append(f"{step}. {usage}")
+                else:
+                    instructions.append(
+                        f"{step}. Use '{tool.name}' when creating visualizations"
+                    )
+                step += 1
+
+        return "\n".join(instructions) if len(instructions) > 1 else ""
+
+    def _build_tool_guidelines(self, sorted_tools: list[Tool]) -> str:
+        """Build tool-specific guidelines."""
+        guidelines = []
+
+        for tool in sorted_tools:
+            usage = tool.get_usage_instructions()
+            if usage and not self._is_usage_in_workflow(usage):
+                guidelines.append(f"- {tool.name}: {usage}")
+
+        if guidelines:
+            return "Tool-Specific Guidelines:\n" + "\n".join(guidelines)
+        return ""
+
+    def _build_general_guidelines(self, sorted_tools: list[Tool]) -> str:
+        """Build general usage guidelines."""
+        guidelines = [
+            "Guidelines:",
+            "- 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",
+        ]
+
+        # Add category-specific guidelines
+        categories = {tool.category for tool in sorted_tools}
+
+        if ToolCategory.SQL in categories:
+            guidelines.extend(
+                [
+                    "- Timestamp columns must be converted to text when you write queries",
+                    "- Use table patterns like 'sample%' or '%experiment%' to filter related tables",
+                ]
+            )
+
+        if ToolCategory.VISUALIZATION in categories:
+            guidelines.append(
+                "- Create visualizations when they would enhance understanding of the data"
+            )
+
+        return "\n".join(guidelines)
+
+    def _is_usage_in_workflow(self, usage: str) -> bool:
+        """Check if usage instruction is already covered in workflow section."""
+        # Simple heuristic - if usage starts with workflow words, it's probably in workflow
+        workflow_words = ["always start", "first", "use this", "begin with", "start by"]
+        usage_lower = usage.lower()
+        return any(word in usage_lower for word in workflow_words)
+
+    def build_mcp_instructions(self) -> str:
+        """Build instructions specifically for MCP server."""
+        instructions = [
+            "This server provides helpful resources and tools that will help you address users queries on their database.",
+            "",
+        ]
+
+        # Add database discovery
+        instructions.append("- Get all databases using `get_databases()`")
+
+        # Add tool-specific instructions
+        sql_tools = self.registry.get_all_tools(category=ToolCategory.SQL)
+        sorted_tools = self._sort_tools_by_workflow(sql_tools)
+
+        for tool in sorted_tools:
+            instructions.append(f"- Call `{tool.name}()` to {tool.description.lower()}")
+
+        # Add workflow guidelines
+        instructions.extend(["", "Guidelines:"])
+
+        workflow_instructions = self._build_workflow_instructions(sorted_tools)
+        if workflow_instructions:
+            # Extract just the numbered steps without the "IMPORTANT" header
+            lines = workflow_instructions.split("\n")[1:]  # Skip header
+            for line in lines:
+                if line.strip():
+                    # Convert numbered steps to bullet points
+                    if line.strip()[0].isdigit():
+                        instructions.append(f"- {line.strip()[3:]}")  # Remove "X. "
+
+        # Add general guidelines
+        instructions.extend(
+            [
+                "- Use proper JOIN syntax and avoid cartesian products",
+                "- Include appropriate WHERE clauses to limit results",
+                "- Handle errors gracefully and suggest fixes",
+            ]
+        )
+
+        return "\n".join(instructions)

sqlsaber/tools/registry.py

@@ -0,0 +1,130 @@
+"""Tool registry for managing available tools."""
+
+from typing import Type
+
+from .base import Tool
+from .enums import ToolCategory
+
+
+class ToolRegistry:
+    """Registry for managing and discovering tools."""
+
+    def __init__(self):
+        """Initialize the registry."""
+        self._tools: dict[str, Type[Tool]] = {}
+        self._instances: dict[str, Tool] = {}
+
+    def register(self, tool_class: Type[Tool]) -> None:
+        """Register a tool class.
+
+        Args:
+            tool_class: The tool class to register
+        """
+        # Create a temporary instance to get the name
+        temp_instance = tool_class()
+        name = temp_instance.name
+
+        if name in self._tools:
+            raise ValueError(f"Tool '{name}' is already registered")
+
+        self._tools[name] = tool_class
+
+    def unregister(self, name: str) -> None:
+        """Unregister a tool.
+
+        Args:
+            name: Name of the tool to unregister
+        """
+        if name in self._tools:
+            del self._tools[name]
+        if name in self._instances:
+            del self._instances[name]
+
+    def get_tool(self, name: str) -> Tool:
+        """Get a tool instance by name.
+
+        Args:
+            name: Name of the tool
+
+        Returns:
+            Tool instance
+
+        Raises:
+            KeyError: If tool is not found
+        """
+        if name not in self._tools:
+            raise KeyError(f"Tool '{name}' not found in registry")
+
+        # Create instance if not already created (singleton pattern)
+        if name not in self._instances:
+            self._instances[name] = self._tools[name]()
+
+        return self._instances[name]
+
+    def list_tools(self, category: str | ToolCategory | None = None) -> list[str]:
+        """List all registered tool names.
+
+        Args:
+            category: Optional category to filter by (string or ToolCategory enum)
+
+        Returns:
+            List of tool names
+        """
+        if category is None:
+            return list(self._tools.keys())
+
+        # Convert string to enum
+        if isinstance(category, str):
+            try:
+                category = ToolCategory(category)
+            except ValueError:
+                # If string doesn't match any enum, return empty list
+                return []
+
+        # Filter by category
+        result = []
+        for name, tool_class in self._tools.items():
+            tool = self.get_tool(name)
+            if tool.category == category:
+                result.append(name)
+        return result
+
+    def get_all_tools(self, category: str | ToolCategory | None = None) -> list[Tool]:
+        """Get all tool instances.
+
+        Args:
+            category: Optional category to filter by (string or ToolCategory enum)
+
+        Returns:
+            List of tool instances
+        """
+        names = self.list_tools(category)
+        return [self.get_tool(name) for name in names]
+
+    def get_tool_definitions(self, category: str | ToolCategory | None = None) -> list:
+        """Get tool definitions for all tools.
+
+        Args:
+            category: Optional category to filter by (string or ToolCategory enum)
+
+        Returns:
+            List of ToolDefinition objects
+        """
+        tools = self.get_all_tools(category)
+        return [tool.to_definition() for tool in tools]
+
+
+# Global registry instance
+tool_registry = ToolRegistry()
+
+
+def register_tool(tool_class: Type[Tool]) -> Type[Tool]:
+    """Decorator to register a tool class.
+
+    Usage:
+        @register_tool
+        class MyTool(Tool):
+            ...
+    """
+    tool_registry.register(tool_class)
+    return tool_class