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

sqlsaber/agents/__init__.py

@@ -1,9 +1,7 @@
 """Agents module for SQLSaber."""
 
-from .anthropic import AnthropicSQLAgent
-from .base import BaseSQLAgent
+from .pydantic_ai_agent import build_sqlsaber_agent
 
 __all__ = [
-    "BaseSQLAgent",
-    "AnthropicSQLAgent",
+    "build_sqlsaber_agent",
 ]

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,
@@ -16,7 +14,7 @@ from sqlsaber.database.connection import (
     SQLiteConnection,
 )
 from sqlsaber.database.schema import SchemaManager
-from sqlsaber.models.events import StreamEvent
+from sqlsaber.tools import SQLTool, tool_registry
 
 
 class BaseSQLAgent(ABC):
@@ -32,13 +30,16 @@ 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,
         user_query: str,
         use_history: bool = True,
         cancellation_token: asyncio.Event | None = None,
-    ) -> AsyncIterator[StreamEvent]:
+    ) -> AsyncIterator:
         """Process a user query and stream responses.
 
         Args:
@@ -69,232 +70,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/agents/mcp.py

@@ -1,9 +1,9 @@
 """Generic SQL agent implementation for MCP tools."""
 
 from typing import AsyncIterator
+
 from sqlsaber.agents.base import BaseSQLAgent
 from sqlsaber.database.connection import BaseDatabaseConnection
-from sqlsaber.models.events import StreamEvent
 
 
 class MCPSQLAgent(BaseSQLAgent):
@@ -14,7 +14,7 @@ class MCPSQLAgent(BaseSQLAgent):
 
     async def query_stream(
         self, user_query: str, use_history: bool = True
-    ) -> AsyncIterator[StreamEvent]:
+    ) -> AsyncIterator:
         """Not implemented for generic agent as it's only used for tool operations."""
         raise NotImplementedError(
             "MCPSQLAgent does not support query streaming. Use specific agent implementations for conversation."

sqlsaber/agents/pydantic_ai_agent.py

@@ -0,0 +1,170 @@
+"""Pydantic-AI Agent for SQLSaber.
+
+This replaces the custom AnthropicSQLAgent and uses pydantic-ai's Agent,
+function tools, and streaming event types directly.
+"""
+
+import httpx
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.models.anthropic import AnthropicModel
+from pydantic_ai.models.google import GoogleModel
+from pydantic_ai.providers.anthropic import AnthropicProvider
+from pydantic_ai.providers.google import GoogleProvider
+
+from sqlsaber.config import providers
+from sqlsaber.config.settings import Config
+from sqlsaber.database.connection import (
+    BaseDatabaseConnection,
+    CSVConnection,
+    MySQLConnection,
+    PostgreSQLConnection,
+    SQLiteConnection,
+)
+from sqlsaber.memory.manager import MemoryManager
+from sqlsaber.tools.instructions import InstructionBuilder
+from sqlsaber.tools.registry import tool_registry
+from sqlsaber.tools.sql_tools import SQLTool
+
+
+def build_sqlsaber_agent(
+    db_connection: BaseDatabaseConnection,
+    database_name: str | None,
+) -> Agent:
+    """Create and configure a pydantic-ai Agent for SQLSaber.
+
+    - Registers function tools that delegate to the existing tool registry
+    - Attaches dynamic system prompt built from InstructionBuilder + MemoryManager
+    - Ensures SQL tools have the active DB connection
+    """
+    # Ensure SQL tools receive the active 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(db_connection)
+
+    cfg = Config()
+    # Ensure provider env var is hydrated from keyring for current provider (Config.validate handles it)
+    cfg.validate()
+
+    # Build model/agent. For some providers (e.g., google), construct provider model explicitly to
+    # allow arbitrary model IDs even if not in pydantic-ai's KnownModelName.
+    model_name_only = (
+        cfg.model_name.split(":", 1)[1] if ":" in cfg.model_name else cfg.model_name
+    )
+
+    provider = providers.provider_from_model(cfg.model_name) or ""
+    if provider == "google":
+        model_obj = GoogleModel(
+            model_name_only, provider=GoogleProvider(api_key=cfg.api_key)
+        )
+        agent = Agent(model_obj, name="sqlsaber")
+    elif provider == "anthropic" and bool(getattr(cfg, "oauth_token", None)):
+        # Build custom httpx client to inject OAuth headers for Anthropic
+        async def add_oauth_headers(request: httpx.Request) -> None:  # type: ignore[override]
+            # Remove API-key header if present and add OAuth headers
+            if "x-api-key" in request.headers:
+                del request.headers["x-api-key"]
+            request.headers.update(
+                {
+                    "Authorization": f"Bearer {cfg.oauth_token}",
+                    "anthropic-version": "2023-06-01",
+                    "anthropic-beta": "oauth-2025-04-20",
+                    "User-Agent": "ClaudeCode/1.0 (Anthropic Claude Code CLI)",
+                    "X-Client-Name": "claude-code",
+                    "X-Client-Version": "1.0.0",
+                }
+            )
+
+        http_client = httpx.AsyncClient(event_hooks={"request": [add_oauth_headers]})
+        provider_obj = AnthropicProvider(api_key="placeholder", http_client=http_client)
+        model_obj = AnthropicModel(model_name_only, provider=provider_obj)
+        agent = Agent(model_obj, name="sqlsaber")
+    else:
+        agent = Agent(cfg.model_name, name="sqlsaber")
+
+    # Memory + dynamic system prompt
+    memory_manager = MemoryManager()
+    instruction_builder = InstructionBuilder(tool_registry)
+
+    is_oauth = provider == "anthropic" and bool(getattr(cfg, "oauth_token", None))
+
+    if not is_oauth:
+
+        @agent.system_prompt(dynamic=True)
+        async def sqlsaber_system_prompt(ctx: RunContext) -> str:
+            db_type = _get_database_type_name(db_connection)
+            instructions = instruction_builder.build_instructions(db_type=db_type)
+
+            # Add memory context if available
+            if database_name:
+                mem = memory_manager.format_memories_for_prompt(database_name)
+            else:
+                mem = ""
+
+            parts = [p for p in (instructions, mem) if p and p.strip()]
+            return "\n\n".join(parts) if parts else ""
+    else:
+
+        @agent.system_prompt(dynamic=True)
+        async def sqlsaber_system_prompt(ctx: RunContext) -> str:
+            # Minimal system prompt in OAuth mode to match Claude Code identity
+            return "You are Claude Code, Anthropic's official CLI for Claude."
+
+    # Expose helpers and context on agent instance
+    agent._sqlsaber_memory_manager = memory_manager  # type: ignore[attr-defined]
+    agent._sqlsaber_database_name = database_name  # type: ignore[attr-defined]
+    agent._sqlsaber_instruction_builder = instruction_builder  # type: ignore[attr-defined]
+    agent._sqlsaber_db_type = _get_database_type_name(db_connection)  # type: ignore[attr-defined]
+    agent._sqlsaber_is_oauth = is_oauth  # type: ignore[attr-defined]
+
+    # Tool wrappers that invoke the registered tools
+    @agent.tool(name="list_tables")
+    async def list_tables(ctx: RunContext) -> str:
+        """
+        Get a list of all tables in the database with row counts.
+        Use this first to discover available tables.
+        """
+        tool = tool_registry.get_tool("list_tables")
+        return await tool.execute()
+
+    @agent.tool(name="introspect_schema")
+    async def introspect_schema(
+        ctx: RunContext, table_pattern: str | None = None
+    ) -> str:
+        """
+        Introspect database schema to understand table structures.
+
+        Args:
+            table_pattern: Optional pattern to filter tables (e.g., 'public.users', 'user%', '%order%')
+        """
+        tool = tool_registry.get_tool("introspect_schema")
+        return await tool.execute(table_pattern=table_pattern)
+
+    @agent.tool(name="execute_sql")
+    async def execute_sql(ctx: RunContext, query: str, limit: int | None = 100) -> str:
+        """
+        Execute a SQL query and return the results.
+
+        Args:
+            query: SQL query to execute
+            limit: Maximum number of rows to return (default: 100)
+        """
+        tool = tool_registry.get_tool("execute_sql")
+        return await tool.execute(query=query, limit=limit)
+
+    return agent
+
+
+def _get_database_type_name(db: BaseDatabaseConnection) -> str:
+    """Get the human-readable database type name (mirrors BaseSQLAgent)."""
+
+    if isinstance(db, PostgreSQLConnection):
+        return "PostgreSQL"
+    elif isinstance(db, MySQLConnection):
+        return "MySQL"
+    elif isinstance(db, SQLiteConnection):
+        return "SQLite"
+    elif isinstance(db, CSVConnection):
+        return "SQLite"
+    else:
+        return "database"