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

sqlsaber/tools/sql_tools.py

@@ -0,0 +1,275 @@
+"""SQL-related tools for database operations."""
+
+import json
+from typing import Any
+
+from sqlsaber.database.connection import BaseDatabaseConnection
+from sqlsaber.database.schema import SchemaManager
+
+from .base import Tool
+from .enums import ToolCategory, WorkflowPosition
+from .registry import register_tool
+
+
+class SQLTool(Tool):
+    """Base class for SQL tools that need database access."""
+
+    def __init__(self, db_connection: BaseDatabaseConnection | None = None):
+        """Initialize with optional database connection."""
+        super().__init__()
+        self.db = db_connection
+        self.schema_manager = SchemaManager(db_connection) if db_connection else None
+
+    def set_connection(self, db_connection: BaseDatabaseConnection) -> None:
+        """Set the database connection after initialization."""
+        self.db = db_connection
+        self.schema_manager = SchemaManager(db_connection)
+
+    @property
+    def category(self) -> ToolCategory:
+        """SQL tools belong to the 'sql' category."""
+        return ToolCategory.SQL
+
+
+@register_tool
+class ListTablesTool(SQLTool):
+    """Tool for listing database tables."""
+
+    @property
+    def name(self) -> str:
+        return "list_tables"
+
+    @property
+    def description(self) -> str:
+        return "Get a list of all tables in the database with row counts. Use this first to discover available tables."
+
+    @property
+    def input_schema(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {},
+            "required": [],
+        }
+
+    def get_usage_instructions(self) -> str | None:
+        """Return usage instructions for this tool."""
+        return "ALWAYS start with 'list_tables' to see available tables and row counts. Use this first to discover available tables."
+
+    def get_priority(self) -> int:
+        """Return priority for tool ordering."""
+        return 10  # High priority - should be used first
+
+    def get_workflow_position(self) -> WorkflowPosition:
+        """Return workflow position."""
+        return WorkflowPosition.DISCOVERY
+
+    async def execute(self, **kwargs) -> str:
+        """List all tables in the database."""
+        if not self.db or not self.schema_manager:
+            return json.dumps({"error": "No database connection available"})
+
+        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)}"})
+
+
+@register_tool
+class IntrospectSchemaTool(SQLTool):
+    """Tool for introspecting database schema."""
+
+    @property
+    def name(self) -> str:
+        return "introspect_schema"
+
+    @property
+    def description(self) -> str:
+        return "Introspect database schema to understand table structures."
+
+    @property
+    def input_schema(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "table_pattern": {
+                    "type": "string",
+                    "description": "Optional pattern to filter tables (e.g., 'public.users', 'user%', '%order%')",
+                }
+            },
+            "required": [],
+        }
+
+    def get_usage_instructions(self) -> str | None:
+        """Return usage instructions for this tool."""
+        return "Use 'introspect_schema' with a table_pattern to get details ONLY for relevant tables. Use table patterns like 'sample%' or '%experiment%' to filter related tables."
+
+    def get_priority(self) -> int:
+        """Return priority for tool ordering."""
+        return 20  # Should come after list_tables
+
+    def get_workflow_position(self) -> WorkflowPosition:
+        """Return workflow position."""
+        return WorkflowPosition.ANALYSIS
+
+    async def execute(self, **kwargs) -> str:
+        """Introspect database schema."""
+        if not self.db or not self.schema_manager:
+            return json.dumps({"error": "No database connection available"})
+
+        try:
+            table_pattern = kwargs.get("table_pattern")
+            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)}"})
+
+
+@register_tool
+class ExecuteSQLTool(SQLTool):
+    """Tool for executing SQL queries."""
+
+    DEFAULT_LIMIT = 100
+
+    @property
+    def name(self) -> str:
+        return "execute_sql"
+
+    @property
+    def description(self) -> str:
+        return "Execute a SQL query against the database."
+
+    @property
+    def input_schema(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "SQL query to execute",
+                },
+                "limit": {
+                    "type": "integer",
+                    "description": f"Maximum number of rows to return (default: {self.DEFAULT_LIMIT})",
+                    "default": self.DEFAULT_LIMIT,
+                },
+            },
+            "required": ["query"],
+        }
+
+    def get_usage_instructions(self) -> str | None:
+        """Return usage instructions for this tool."""
+        return "Execute SQL queries safely with automatic LIMIT clauses for SELECT statements. Only SELECT queries are permitted for security."
+
+    def get_priority(self) -> int:
+        """Return priority for tool ordering."""
+        return 30  # Should come after schema tools
+
+    def get_workflow_position(self) -> WorkflowPosition:
+        """Return workflow position."""
+        return WorkflowPosition.EXECUTION
+
+    async def execute(self, **kwargs) -> str:
+        """Execute a SQL query."""
+        if not self.db:
+            return json.dumps({"error": "No database connection available"})
+
+        query = kwargs.get("query")
+        if not query:
+            return json.dumps({"error": "No query provided"})
+
+        limit = kwargs.get("limit", self.DEFAULT_LIMIT)
+
+        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
+            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],
+                    "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 _validate_write_operation(self, query: str) -> str | None:
+        """Validate if a write operation is 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

sqlsaber/tools/visualization_tools.py

@@ -0,0 +1,144 @@
+"""Visualization tools for data plotting."""
+
+import json
+from typing import Any
+
+from uniplot import histogram, plot
+
+from .base import Tool
+from .enums import ToolCategory, WorkflowPosition
+from .registry import register_tool
+
+
+@register_tool
+class PlotDataTool(Tool):
+    """Tool for creating terminal plots using uniplot."""
+
+    @property
+    def name(self) -> str:
+        return "plot_data"
+
+    @property
+    def description(self) -> str:
+        return "Create a plot of query results."
+
+    @property
+    def input_schema(self) -> dict[str, Any]:
+        return {
+            "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"],
+        }
+
+    @property
+    def category(self) -> ToolCategory:
+        return ToolCategory.VISUALIZATION
+
+    def get_usage_instructions(self) -> str | None:
+        """Return usage instructions for this tool."""
+        return "Create terminal plots from query results when visualization would enhance understanding of the data."
+
+    def get_priority(self) -> int:
+        """Return priority for tool ordering."""
+        return 40  # Should come after SQL execution
+
+    def get_workflow_position(self) -> WorkflowPosition:
+        """Return workflow position."""
+        return WorkflowPosition.VISUALIZATION
+
+    async def execute(self, **kwargs) -> str:
+        """Create a terminal plot."""
+        y_values = kwargs.get("y_values", [])
+        x_values = kwargs.get("x_values")
+        plot_type = kwargs.get("plot_type", "line")
+        title = kwargs.get("title")
+        x_label = kwargs.get("x_label")
+        y_label = kwargs.get("y_label")
+
+        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}"})
+
+            return json.dumps(
+                {"success": True, "plot_rendered": True, "plot_info": plot_info}
+            )
+
+        except Exception as e:
+            return json.dumps({"error": f"Error creating plot: {str(e)}"})

{sqlsaber-0.14.0.dist-info → sqlsaber-0.15.0.dist-info}/METADATA

@@ -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.dist-info → sqlsaber-0.15.0.dist-info}/RECORD

@@ -1,8 +1,8 @@
 sqlsaber/__init__.py,sha256=HjS8ULtP4MGpnTL7njVY45NKV9Fi4e_yeYuY-hyXWQc,73
 sqlsaber/__main__.py,sha256=RIHxWeWh2QvLfah-2OkhI5IJxojWfy4fXpMnVEJYvxw,78
 sqlsaber/agents/__init__.py,sha256=LWeSeEUE4BhkyAYFF3TE-fx8TtLud3oyEtyB8ojFJgo,167
-sqlsaber/agents/anthropic.py,sha256=xH4xY1lq5idDgY8Vklgu5hGLoKLHWQMMfPgjvQjfG7I,23615
-sqlsaber/agents/base.py,sha256=9A-iceb93eHrQHTdLoeSicP3ZqDHOFTx3Fe5OvHvUkg,14093
+sqlsaber/agents/anthropic.py,sha256=OSrf_2fancbGH04ckouxChNyKj1H54wt_fWGP0hECII,19564
+sqlsaber/agents/base.py,sha256=jOchhaEQWiW__koy2WF4e_YvvsF5l_WYru_q0PkBF7g,6413
 sqlsaber/agents/mcp.py,sha256=FKtXgDrPZ2-xqUYCw2baI5JzrWekXaC5fjkYW1_Mg50,827
 sqlsaber/agents/streaming.py,sha256=LaSeMTlxuJFRArJVqDly5-_KgcePiCCKPKfMxfB4oGs,521
 sqlsaber/cli/__init__.py,sha256=qVSLVJLLJYzoC6aj6y9MFrzZvAwc4_OgxU9DlkQnZ4M,86
@@ -37,15 +37,22 @@ sqlsaber/database/connection.py,sha256=sZVGNMzMwiM11GrsLLPwR8A5ugzJ5O0TCdkrt0KVR
 sqlsaber/database/resolver.py,sha256=RPXF5EoKzvQDDLmPGNHYd2uG_oNICH8qvUjBp6iXmNY,3348
 sqlsaber/database/schema.py,sha256=OC93dnZkijCoVNqb6itSpQ2XsiZ85PjVUW-VZDwrPrk,25989
 sqlsaber/mcp/__init__.py,sha256=COdWq7wauPBp5Ew8tfZItFzbcLDSEkHBJSMhxzy8C9c,112
-sqlsaber/mcp/mcp.py,sha256=YH4crygqb5_Y94nsns6d-26FZCTlDPOh3tf-ghihzDM,4440
+sqlsaber/mcp/mcp.py,sha256=X12oCMZYAtgJ7MNuh5cqz8y3lALrOzkXWcfpuY0Ijxk,3950
 sqlsaber/memory/__init__.py,sha256=GiWkU6f6YYVV0EvvXDmFWe_CxarmDCql05t70MkTEWs,63
 sqlsaber/memory/manager.py,sha256=p3fybMVfH-E4ApT1ZRZUnQIWSk9dkfUPCyfkmA0HALs,2739
 sqlsaber/memory/storage.py,sha256=ne8szLlGj5NELheqLnI7zu21V8YS4rtpYGGC7tOmi-s,5745
 sqlsaber/models/__init__.py,sha256=RJ7p3WtuSwwpFQ1Iw4_DHV2zzCtHqIzsjJzxv8kUjUE,287
 sqlsaber/models/events.py,sha256=89SXKb5GGpH01yTr2kPEBhzp9xv35RFIYuFdAZSIPoE,721
 sqlsaber/models/types.py,sha256=w-zk81V2dtveuteej36_o1fDK3So428j3P2rAejU62U,862
-sqlsaber-0.14.0.dist-info/METADATA,sha256=xAj3wcH-3OJWud9grNykoiHhuVp2LIdMV5bbxNUQOEk,6877
-sqlsaber-0.14.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-sqlsaber-0.14.0.dist-info/entry_points.txt,sha256=qEbOB7OffXPFgyJc7qEIJlMEX5RN9xdzLmWZa91zCQQ,162
-sqlsaber-0.14.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sqlsaber-0.14.0.dist-info/RECORD,,
+sqlsaber/tools/__init__.py,sha256=a-JNOhHsC7WEVhTsQY_IHckaPOmswoM3Q85YOd2iC_E,652
+sqlsaber/tools/base.py,sha256=0pCC2OMZc4uyxv-ErRHuXbu_9MeVXPzSyGIyix6Hvw4,2085
+sqlsaber/tools/enums.py,sha256=TnlvEOpkGtuzEzg_JBdSb_S38jg4INhtxw4qZ20kU_E,483
+sqlsaber/tools/instructions.py,sha256=nnItVvJBtN-FLB3-PSR6Y23ix6AiOB5hNX3r4TtYFKw,9869
+sqlsaber/tools/registry.py,sha256=decVRNf50JPNT4i-OHZIL2B8cmoOMDaE4y0Av6q6v0I,3619
+sqlsaber/tools/sql_tools.py,sha256=hM6tKqW5MDhFUt6MesoqhTUqIpq_5baIIDoN1MjDCXY,9647
+sqlsaber/tools/visualization_tools.py,sha256=059Pe3aOZvgpqT9487Ydv2PhY7T1pVmfALPTvfqPisI,4973
+sqlsaber-0.15.0.dist-info/METADATA,sha256=XEF-xwN39iDs9Gd_uSdeHxyE50i429ol5iU2HSnkdZc,6877
+sqlsaber-0.15.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+sqlsaber-0.15.0.dist-info/entry_points.txt,sha256=qEbOB7OffXPFgyJc7qEIJlMEX5RN9xdzLmWZa91zCQQ,162
+sqlsaber-0.15.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sqlsaber-0.15.0.dist-info/RECORD,,