sqlsaber 0.15.0__py3-none-any.whl → 0.16.1__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

@@ -14,7 +14,6 @@ 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
 
 
@@ -40,7 +39,7 @@ class BaseSQLAgent(ABC):
         user_query: str,
         use_history: bool = True,
         cancellation_token: asyncio.Event | None = None,
-    ) -> AsyncIterator[StreamEvent]:
+    ) -> AsyncIterator:
         """Process a user query and stream responses.
 
         Args:

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"

sqlsaber/cli/auth.py

@@ -1,11 +1,17 @@
 """Authentication CLI commands."""
 
-import questionary
+import os
+
 import cyclopts
+import keyring
+import questionary
 from rich.console import Console
 
+from sqlsaber.config import providers
+from sqlsaber.config.api_keys import APIKeyManager
 from sqlsaber.config.auth import AuthConfigManager, AuthMethod
 from sqlsaber.config.oauth_flow import AnthropicOAuthFlow
+from sqlsaber.config.oauth_tokens import OAuthTokenManager
 
 # Global instances for CLI commands
 console = Console()
@@ -20,59 +26,59 @@ auth_app = cyclopts.App(
 
 @auth_app.command
 def setup():
-    """Configure authentication method for SQLSaber."""
-    console.print("\n[bold]SQLSaber Authentication Setup[/bold]\n")
-
-    # Use questionary for selection
-    auth_choice = questionary.select(
-        "Choose your authentication method:",
-        choices=[
-            questionary.Choice(
-                title="Anthropic API Key",
-                value=AuthMethod.API_KEY,
-                description="You can create one by visiting https://console.anthropic.com",
-            ),
-            questionary.Choice(
-                title="Claude Pro or Max Subscription",
-                value=AuthMethod.CLAUDE_PRO,
-                description="This does not require creating an API Key, but requires a subscription at https://claude.ai",
-            ),
-        ],
+    """Configure authentication for SQLsaber (API keys and Anthropic OAuth)."""
+    console.print("\n[bold]SQLsaber Authentication Setup[/bold]\n")
+
+    provider = questionary.select(
+        "Select provider to configure:",
+        choices=providers.all_keys(),
     ).ask()
 
-    if auth_choice is None:
+    if provider is None:
         console.print("[yellow]Setup cancelled.[/yellow]")
         return
 
-    # Handle auth method setup
-    if auth_choice == AuthMethod.API_KEY:
-        console.print("\nTo configure your API key, you can either:")
-        console.print("• Set the ANTHROPIC_API_KEY environment variable")
-        console.print(
-            "• Let SQLsaber prompt you for the key when needed (stored securely)"
-        )
-
-        config_manager.set_auth_method(auth_choice)
-        console.print("\n[bold green]Authentication method saved![/bold green]")
-
-    elif auth_choice == AuthMethod.CLAUDE_PRO:
-        oauth_flow = AnthropicOAuthFlow()
-        try:
-            success = oauth_flow.authenticate()
-            if success:
-                config_manager.set_auth_method(auth_choice)
+    if provider == "anthropic":
+        # Let user choose API key or OAuth
+        method_choice = questionary.select(
+            "Select Anthropic authentication method:",
+            choices=[
+                {"name": "API key", "value": AuthMethod.API_KEY},
+                {"name": "Claude Pro/Max (OAuth)", "value": AuthMethod.CLAUDE_PRO},
+            ],
+        ).ask()
+
+        if method_choice == AuthMethod.CLAUDE_PRO:
+            flow = AnthropicOAuthFlow()
+            if flow.authenticate():
+                config_manager.set_auth_method(AuthMethod.CLAUDE_PRO)
                 console.print(
-                    "\n[bold green]Authentication setup complete![/bold green]"
+                    "\n[bold green]✓ Anthropic OAuth configured successfully![/bold green]"
                 )
             else:
-                console.print(
-                    "\n[yellow]OAuth authentication failed. Please try again.[/yellow]"
-                )
-                return
-        except Exception as e:
-            console.print(f"\n[red]Authentication setup failed: {str(e)}[/red]")
+                console.print("\n[red]✗ Anthropic OAuth setup failed.[/red]")
+            console.print(
+                "You can change this anytime by running [cyan]saber auth setup[/cyan] again."
+            )
             return
 
+    # API key flow (all providers + Anthropic when selected above)
+    api_key_manager = APIKeyManager()
+    env_var = api_key_manager._get_env_var_name(provider)
+    console.print("\nTo configure your API key, you can either:")
+    console.print(f"• Set the {env_var} environment variable")
+    console.print("• Let SQLsaber prompt you for the key when needed (stored securely)")
+
+    # Fetch/store key (cascades env -> keyring -> prompt)
+    api_key = api_key_manager.get_api_key(provider)
+    if api_key:
+        config_manager.set_auth_method(AuthMethod.API_KEY)
+        console.print(
+            f"\n[bold green]✓ {provider.title()} API key configured successfully![/bold green]"
+        )
+    else:
+        console.print("\n[yellow]No API key configured.[/yellow]")
+
     console.print(
         "You can change this anytime by running [cyan]saber auth setup[/cyan] again."
     )
@@ -80,7 +86,7 @@ def setup():
 
 @auth_app.command
 def status():
-    """Show current authentication configuration."""
+    """Show current authentication configuration and provider key status."""
     auth_method = config_manager.get_auth_method()
 
     console.print("\n[bold blue]Authentication Status[/bold blue]")
@@ -88,52 +94,113 @@ def status():
     if auth_method is None:
         console.print("[yellow]No authentication method configured[/yellow]")
         console.print("Run [cyan]saber auth setup[/cyan] to configure authentication.")
+        return
+
+    # Show configured method summary
+    if auth_method == AuthMethod.CLAUDE_PRO:
+        console.print("[green]✓ Anthropic Claude Pro/Max (OAuth) configured[/green]\n")
     else:
-        if auth_method == AuthMethod.API_KEY:
-            console.print("[green]✓ API Key authentication configured[/green]")
-            console.print("Using Anthropic API key for authentication")
-        elif auth_method == AuthMethod.CLAUDE_PRO:
-            console.print("[green]✓ Claude Pro/Max subscription configured[/green]")
-
-            # Check OAuth token status
-            oauth_flow = AnthropicOAuthFlow()
-            if oauth_flow.has_valid_authentication():
-                console.print("OAuth token is valid and ready to use")
-            else:
-                console.print("[yellow]OAuth token missing or expired[/yellow]")
+        console.print("[green]✓ API Key authentication configured[/green]\n")
+
+    # Show per-provider status without prompting
+    api_key_manager = APIKeyManager()
+    for provider in providers.all_keys():
+        if provider == "anthropic":
+            # Include OAuth status
+            if OAuthTokenManager().has_oauth_token("anthropic"):
+                console.print("> anthropic (oauth): [green]configured[/green]")
+        env_var = api_key_manager._get_env_var_name(provider)
+        service = api_key_manager._get_service_name(provider)
+        from_env = bool(os.getenv(env_var))
+        from_keyring = bool(keyring.get_password(service, provider))
+        if from_env:
+            console.print(f"> {provider}: configured via {env_var}")
+        elif from_keyring:
+            console.print(f"> {provider}: [green]configured[/green]")
+        else:
+            console.print(f"> {provider}: [yellow]not configured[/yellow]")
 
 
 @auth_app.command
 def reset():
-    """Reset authentication configuration."""
-    if not config_manager.has_auth_configured():
-        console.print("[yellow]No authentication configuration to reset.[/yellow]")
+    """Reset credentials for a selected provider (API key and/or OAuth)."""
+    console.print("\n[bold]SQLsaber Authentication Reset[/bold]\n")
+
+    # Choose provider to reset (mirrors setup)
+    provider = questionary.select(
+        "Select provider to reset:",
+        choices=providers.all_keys(),
+    ).ask()
+
+    if provider is None:
+        console.print("[yellow]Reset cancelled.[/yellow]")
         return
 
-    current_method = config_manager.get_auth_method()
-    method_name = (
-        "API Key" if current_method == AuthMethod.API_KEY else "Claude Pro/Max"
+    api_key_manager = APIKeyManager()
+    service = api_key_manager._get_service_name(provider)
+
+    # Determine what exists in keyring
+    api_key_present = bool(keyring.get_password(service, provider))
+    oauth_present = (
+        OAuthTokenManager().has_oauth_token("anthropic")
+        if provider == "anthropic"
+        else False
     )
 
-    if questionary.confirm(
-        f"Are you sure you want to reset the current authentication method ({method_name})?",
-        default=False,
-    ).ask():
-        # If Claude Pro, also remove OAuth tokens
-        if current_method == AuthMethod.CLAUDE_PRO:
-            oauth_flow = AnthropicOAuthFlow()
-            oauth_flow.remove_authentication()
-
-        # Clear the auth config by setting it to None
-        config = config_manager._load_config()
-        config["auth_method"] = None
-        config_manager._save_config(config)
-        console.print("[green]Authentication configuration reset.[/green]")
+    if not api_key_present and not oauth_present:
         console.print(
-            "Run [cyan]saber auth setup[/cyan] to configure authentication again."
+            f"[yellow]No stored credentials found for {provider}. Nothing to reset.[/yellow]"
         )
-    else:
+        return
+
+    # Build confirmation message
+    to_remove: list[str] = []
+    if oauth_present:
+        to_remove.append("Anthropic OAuth token")
+    if api_key_present:
+        to_remove.append(f"{provider.title()} API key")
+
+    summary = ", ".join(to_remove)
+    confirmed = questionary.confirm(
+        f"Remove the following for {provider}: {summary}?",
+        default=False,
+    ).ask()
+
+    if not confirmed:
         console.print("Reset cancelled.")
+        return
+
+    # Perform deletions
+    if oauth_present:
+        OAuthTokenManager().remove_oauth_token("anthropic")
+    if api_key_present:
+        try:
+            keyring.delete_password(service, provider)
+            console.print(f"Removed {provider} API key from keyring", style="green")
+        except keyring.errors.PasswordDeleteError:
+            # Already absent; treat as success
+            pass
+        except Exception as e:
+            console.print(f"Warning: Could not remove API key: {e}", style="yellow")
+
+    # Optionally clear global auth method if removing Anthropic OAuth configuration
+    if provider == "anthropic" and oauth_present:
+        current_method = config_manager.get_auth_method()
+        if current_method == AuthMethod.CLAUDE_PRO:
+            also_clear = questionary.confirm(
+                "Anthropic OAuth was removed. Also unset the global auth method?",
+                default=False,
+            ).ask()
+            if also_clear:
+                config = config_manager._load_config()
+                config["auth_method"] = None
+                config_manager._save_config(config)
+                console.print("Global auth method unset.", style="green")
+
+    console.print("\n[bold green]✓ Reset complete.[/bold green]")
+    console.print(
+        "Environment variables are not modified by this command.", style="dim"
+    )
 
 
 def create_auth_app() -> cyclopts.App:

sqlsaber/cli/commands.py

@@ -7,7 +7,7 @@ from typing import Annotated
 import cyclopts
 from rich.console import Console
 
-from sqlsaber.agents.anthropic import AnthropicSQLAgent
+from sqlsaber.agents import build_sqlsaber_agent
 from sqlsaber.cli.auth import create_auth_app
 from sqlsaber.cli.database import create_db_app
 from sqlsaber.cli.interactive import InteractiveSession
@@ -15,7 +15,13 @@ from sqlsaber.cli.memory import create_memory_app
 from sqlsaber.cli.models import create_models_app
 from sqlsaber.cli.streaming import StreamingQueryHandler
 from sqlsaber.config.database import DatabaseConfigManager
-from sqlsaber.database.connection import DatabaseConnection
+from sqlsaber.database.connection import (
+    CSVConnection,
+    DatabaseConnection,
+    MySQLConnection,
+    PostgreSQLConnection,
+    SQLiteConnection,
+)
 from sqlsaber.database.resolver import DatabaseResolutionError, resolve_database
 
 
@@ -29,7 +35,7 @@ class CLIError(Exception):
 
 app = cyclopts.App(
     name="sqlsaber",
-    help="SQLSaber - Use the agent Luke!\n\nSQL assistant for your database",
+    help="SQLsaber - Open-source agentic SQL assistant for your database",
 )
 
 
@@ -123,25 +129,34 @@ def query(
         except Exception as e:
             raise CLIError(f"Error creating database connection: {e}")
 
-        # Create agent instance with database name for memory context
-        agent = AnthropicSQLAgent(db_conn, db_name)
+        # Create pydantic-ai agent instance with database name for memory context
+        agent = build_sqlsaber_agent(db_conn, db_name)
 
         try:
             if actual_query:
                 # Single query mode with streaming
                 streaming_handler = StreamingQueryHandler(console)
+                # Compute DB type for the greeting line
+                db_type = (
+                    "PostgreSQL"
+                    if isinstance(db_conn, PostgreSQLConnection)
+                    else "MySQL"
+                    if isinstance(db_conn, MySQLConnection)
+                    else "SQLite"
+                    if isinstance(db_conn, (SQLiteConnection, CSVConnection))
+                    else "database"
+                )
                 console.print(
-                    f"[bold blue]Connected to:[/bold blue] {db_name} {agent._get_database_type_name()}\n"
+                    f"[bold blue]Connected to:[/bold blue] {db_name} ({db_type})\n"
                 )
                 await streaming_handler.execute_streaming_query(actual_query, agent)
             else:
                 # Interactive mode
-                session = InteractiveSession(console, agent)
+                session = InteractiveSession(console, agent, db_conn, db_name)
                 await session.run()
 
         finally:
             # Clean up
-            await agent.close()  # Close the agent's HTTP client
             await db_conn.close()
             console.print("\n[green]Goodbye![/green]")
 

sqlsaber/cli/database.py

@@ -6,8 +6,8 @@ import sys
 from pathlib import Path
 from typing import Annotated
 
-import questionary
 import cyclopts
+import questionary
 from rich.console import Console
 from rich.table import Table