sqlsaber 0.25.0__py3-none-any.whl → 0.27.0__py3-none-any.whl

sqlsaber/agents/__init__.py

@@ -1,7 +1,7 @@
 """Agents module for SQLSaber."""
 
-from .pydantic_ai_agent import build_sqlsaber_agent
+from .pydantic_ai_agent import SQLSaberAgent
 
 __all__ = [
-    "build_sqlsaber_agent",
+    "SQLSaberAgent",
 ]

sqlsaber/agents/base.py

@@ -5,7 +5,7 @@ import json
 from abc import ABC, abstractmethod
 from typing import Any, AsyncIterator
 
-from sqlsaber.database.connection import (
+from sqlsaber.database import (
     BaseDatabaseConnection,
     CSVConnection,
     DuckDBConnection,

sqlsaber/agents/mcp.py

@@ -3,7 +3,7 @@
 from typing import AsyncIterator
 
 from sqlsaber.agents.base import BaseSQLAgent
-from sqlsaber.database.connection import BaseDatabaseConnection
+from sqlsaber.database import BaseDatabaseConnection
 
 
 class MCPSQLAgent(BaseSQLAgent):

sqlsaber/agents/pydantic_ai_agent.py

@@ -6,15 +6,16 @@ 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.models.openai import OpenAIResponsesModel
+from pydantic_ai.models.anthropic import AnthropicModel, AnthropicModelSettings
+from pydantic_ai.models.google import GoogleModel, GoogleModelSettings
+from pydantic_ai.models.groq import GroqModelSettings
+from pydantic_ai.models.openai import OpenAIResponsesModel, OpenAIResponsesModelSettings
 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 (
+from sqlsaber.database import (
     BaseDatabaseConnection,
     CSVConnection,
     DuckDBConnection,
@@ -28,47 +29,119 @@ 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)
+class SQLSaberAgent:
+    """Pydantic-AI Agent wrapper for SQLSaber with enhanced state management."""
+
+    def __init__(
+        self,
+        db_connection: BaseDatabaseConnection,
+        database_name: str | None = None,
+        memory_manager: MemoryManager | None = None,
+        thinking_enabled: bool | None = None,
+    ):
+        self.db_connection = db_connection
+        self.database_name = database_name
+        self.config = Config()
+        self.memory_manager = memory_manager or MemoryManager()
+        self.instruction_builder = InstructionBuilder(tool_registry)
+        self.db_type = self._get_database_type_name()
+
+        # Thinking configuration (CLI override or config default)
+        self.thinking_enabled = (
+            thinking_enabled
+            if thinking_enabled is not None
+            else self.config.thinking_enabled
         )
-        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
+
+        # Configure SQL tools with the database connection
+        self._configure_sql_tools()
+
+        # Create the pydantic-ai agent
+        self.agent = self._build_agent()
+
+    def _configure_sql_tools(self) -> None:
+        """Ensure SQL tools receive the active 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_connection)
+
+    def _build_agent(self) -> Agent:
+        """Create and configure the pydantic-ai Agent."""
+        self.config.validate()
+
+        model_name_only = (
+            self.config.model_name.split(":", 1)[1]
+            if ":" in self.config.model_name
+            else self.config.model_name
+        )
+
+        provider = providers.provider_from_model(self.config.model_name) or ""
+        self.is_oauth = provider == "anthropic" and bool(
+            getattr(self.config, "oauth_token", None)
+        )
+
+        agent = self._create_agent_for_provider(provider, model_name_only)
+        self._setup_system_prompt(agent)
+        self._register_tools(agent)
+
+        return agent
+
+    def _create_agent_for_provider(self, provider: str, model_name: str) -> Agent:
+        """Create the agent based on the provider type."""
+        if provider == "google":
+            model_obj = GoogleModel(
+                model_name, provider=GoogleProvider(api_key=self.config.api_key)
+            )
+            if self.thinking_enabled:
+                settings = GoogleModelSettings(
+                    google_thinking_config={"include_thoughts": True}
+                )
+                return Agent(model_obj, name="sqlsaber", model_settings=settings)
+            return Agent(model_obj, name="sqlsaber")
+        elif provider == "anthropic" and self.is_oauth:
+            return self._create_oauth_anthropic_agent(model_name)
+        elif provider == "anthropic":
+            if self.thinking_enabled:
+                settings = AnthropicModelSettings(
+                    anthropic_thinking={
+                        "type": "enabled",
+                        "budget_tokens": 2048,
+                    },
+                    max_tokens=8192,
+                )
+                return Agent(
+                    self.config.model_name, name="sqlsaber", model_settings=settings
+                )
+            return Agent(self.config.model_name, name="sqlsaber")
+        elif provider == "openai":
+            model_obj = OpenAIResponsesModel(model_name)
+            if self.thinking_enabled:
+                settings = OpenAIResponsesModelSettings(
+                    openai_reasoning_effort="medium",
+                    openai_reasoning_summary="auto",
+                )
+                return Agent(model_obj, name="sqlsaber", model_settings=settings)
+            return Agent(model_obj, name="sqlsaber")
+        elif provider == "groq":
+            if self.thinking_enabled:
+                settings = GroqModelSettings(groq_reasoning_format="parsed")
+                return Agent(
+                    self.config.model_name, name="sqlsaber", model_settings=settings
+                )
+            return Agent(self.config.model_name, name="sqlsaber")
+        else:
+            return Agent(self.config.model_name, name="sqlsaber")
+
+    def _create_oauth_anthropic_agent(self, model_name: str) -> Agent:
+        """Create an Anthropic agent with OAuth configuration."""
+
         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}",
+                    "Authorization": f"Bearer {self.config.oauth_token}",
                     "anthropic-version": "2023-06-01",
                     "anthropic-beta": "oauth-2025-04-20",
                     "User-Agent": "ClaudeCode/1.0 (Anthropic Claude Code CLI)",
@@ -79,100 +152,99 @@ def build_sqlsaber_agent(
 
         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")
-    elif provider == "openai":
-        # Use OpenAI Responses Model for structured output capabilities
-        model_obj = OpenAIResponsesModel(model_name_only)
-        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 = ""
+        model_obj = AnthropicModel(model_name, provider=provider_obj)
+        if self.thinking_enabled:
+            settings = AnthropicModelSettings(
+                anthropic_thinking={
+                    "type": "enabled",
+                    "budget_tokens": 2048,
+                },
+                max_tokens=8192,
+            )
+            return Agent(model_obj, name="sqlsaber", model_settings=settings)
+        return Agent(model_obj, name="sqlsaber")
+
+    def _setup_system_prompt(self, agent: Agent) -> None:
+        """Set up the dynamic system prompt for the agent."""
+        if not self.is_oauth:
 
-            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, DuckDBConnection):
-        return "DuckDB"
-    elif isinstance(db, CSVConnection):
-        return "DuckDB"
-    else:
-        return "database"
+            @agent.system_prompt(dynamic=True)
+            async def sqlsaber_system_prompt(ctx: RunContext) -> str:
+                instructions = self.instruction_builder.build_instructions(
+                    db_type=self.db_type
+                )
+
+                # Add memory context if available
+                mem = ""
+                if self.database_name:
+                    mem = self.memory_manager.format_memories_for_prompt(
+                        self.database_name
+                    )
+
+                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:
+                return "You are Claude Code, Anthropic's official CLI for Claude."
+
+    def _register_tools(self, agent: Agent) -> None:
+        """Register all the SQL tools with the agent."""
+
+        @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)
+
+    def set_thinking(self, enabled: bool) -> None:
+        """Update thinking settings and rebuild the agent."""
+        self.thinking_enabled = enabled
+        # Rebuild agent with new thinking settings
+        self.agent = self._build_agent()
+
+    def _get_database_type_name(self) -> str:
+        """Get the human-readable database type name."""
+        if isinstance(self.db_connection, PostgreSQLConnection):
+            return "PostgreSQL"
+        elif isinstance(self.db_connection, MySQLConnection):
+            return "MySQL"
+        elif isinstance(self.db_connection, SQLiteConnection):
+            return "SQLite"
+        elif isinstance(self.db_connection, DuckDBConnection):
+            return "DuckDB"
+        elif isinstance(self.db_connection, CSVConnection):
+            return "DuckDB"
+        else:
+            return "database"

sqlsaber/application/__init__.py

@@ -0,0 +1 @@
+"""Application layer for SQLsaber - shared business logic and interactive flows."""

sqlsaber/application/auth_setup.py

@@ -0,0 +1,164 @@
+"""Shared auth setup logic for onboarding and CLI."""
+
+import asyncio
+
+from questionary import Choice
+from rich.console import Console
+
+from sqlsaber.application.prompts import Prompter
+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
+
+console = Console()
+
+
+async def select_provider(prompter: Prompter, default: str = "anthropic") -> str | None:
+    """Interactive provider selection.
+
+    Args:
+        prompter: Prompter instance for interaction
+        default: Default provider to select
+
+    Returns:
+        Selected provider name or None if cancelled
+    """
+    provider = await prompter.select(
+        "Select AI provider:", choices=providers.all_keys(), default=default
+    )
+    return provider
+
+
+async def configure_oauth_anthropic(
+    auth_manager: AuthConfigManager, run_in_thread: bool = False
+) -> bool:
+    """Configure Anthropic OAuth.
+
+    Args:
+        auth_manager: AuthConfigManager instance
+        run_in_thread: Whether to run OAuth flow in a separate thread (for onboarding)
+
+    Returns:
+        True if OAuth configured successfully, False otherwise
+    """
+    flow = AnthropicOAuthFlow()
+
+    if run_in_thread:
+        # Run in thread to avoid event loop conflicts (onboarding)
+        oauth_success = await asyncio.to_thread(flow.authenticate)
+    else:
+        # Run directly (CLI)
+        oauth_success = flow.authenticate()
+
+    if oauth_success:
+        auth_manager.set_auth_method(AuthMethod.CLAUDE_PRO)
+        return True
+
+    return False
+
+
+async def configure_api_key(
+    provider: str, api_key_manager: APIKeyManager, auth_manager: AuthConfigManager
+) -> bool:
+    """Configure API key for a provider.
+
+    Args:
+        provider: Provider name
+        api_key_manager: APIKeyManager instance
+        auth_manager: AuthConfigManager instance
+
+    Returns:
+        True if API key configured successfully, False otherwise
+    """
+    # Get API key (cascades env -> keyring -> prompt)
+    api_key = api_key_manager.get_api_key(provider)
+
+    if api_key:
+        auth_manager.set_auth_method(AuthMethod.API_KEY)
+        return True
+
+    return False
+
+
+async def setup_auth(
+    prompter: Prompter,
+    auth_manager: AuthConfigManager,
+    api_key_manager: APIKeyManager,
+    allow_oauth: bool = True,
+    default_provider: str = "anthropic",
+    run_oauth_in_thread: bool = False,
+) -> tuple[bool, str | None]:
+    """Interactive authentication setup.
+
+    Args:
+        prompter: Prompter instance for interaction
+        auth_manager: AuthConfigManager instance
+        api_key_manager: APIKeyManager instance
+        allow_oauth: Whether to offer OAuth option for Anthropic
+        default_provider: Default provider to select
+        run_oauth_in_thread: Whether to run OAuth in thread (for onboarding)
+
+    Returns:
+        Tuple of (success: bool, provider: str | None)
+    """
+    # Check if auth is already configured
+    if auth_manager.has_auth_configured():
+        console.print("[green]✓ Authentication already configured![/green]")
+        return True, None
+
+    # Select provider
+    provider = await select_provider(prompter, default=default_provider)
+
+    if provider is None:
+        return False, None
+
+    # For Anthropic, offer OAuth or API key
+    if provider == "anthropic" and allow_oauth:
+        method_choice = await prompter.select(
+            "Authentication method:",
+            choices=[
+                Choice("API Key", value=AuthMethod.API_KEY),
+                Choice("Claude Pro/Max (OAuth)", value=AuthMethod.CLAUDE_PRO),
+            ],
+        )
+
+        if method_choice is None:
+            return False, None
+
+        if method_choice == AuthMethod.CLAUDE_PRO:
+            console.print()
+            oauth_success = await configure_oauth_anthropic(
+                auth_manager, run_in_thread=run_oauth_in_thread
+            )
+            if oauth_success:
+                console.print(
+                    "[green]✓ Anthropic OAuth configured successfully![/green]"
+                )
+                return True, provider
+            else:
+                console.print("[red]✗ Anthropic OAuth setup failed.[/red]")
+                return False, None
+
+    # API key flow
+    env_var = api_key_manager.get_env_var_name(provider)
+
+    console.print()
+    console.print(f"[dim]To use {provider.title()}, you need an API key.[/dim]")
+    console.print(f"[dim]You can set the {env_var} environment variable,[/dim]")
+    console.print("[dim]or enter it now to store securely in your OS keychain.[/dim]")
+    console.print()
+
+    # Configure API key
+    api_key_configured = await configure_api_key(
+        provider, api_key_manager, auth_manager
+    )
+
+    if api_key_configured:
+        console.print(
+            f"[green]✓ {provider.title()} API key configured successfully![/green]"
+        )
+        return True, provider
+    else:
+        console.print("[yellow]No API key provided.[/yellow]")
+        return False, None