sqlsaber 0.25.0__py3-none-any.whl → 0.26.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/cli/commands.py

@@ -75,7 +75,7 @@ def query(
     query_text: Annotated[
         str | None,
         cyclopts.Parameter(
-            help="SQL query in natural language (if not provided, reads from stdin or starts interactive mode)",
+            help="Question in natural language (if not provided, reads from stdin or starts interactive mode)",
         ),
     ] = None,
     database: Annotated[
@@ -85,6 +85,7 @@ def query(
             help="Database connection name, file path (CSV/SQLite/DuckDB), or connection string (postgresql://, mysql://, duckdb://) (uses default if not specified)",
         ),
     ] = None,
+    thinking: bool = False,
 ):
     """Run a query against the database or start interactive mode.
 
@@ -109,16 +110,11 @@ def query(
     async def run_session():
         # Import heavy dependencies only when actually running a query
         # This is only done to speed up startup time
-        from sqlsaber.agents import build_sqlsaber_agent
+        from sqlsaber.agents import SQLSaberAgent
         from sqlsaber.cli.interactive import InteractiveSession
         from sqlsaber.cli.streaming import StreamingQueryHandler
-        from sqlsaber.database.connection import (
-            CSVConnection,
+        from sqlsaber.database import (
             DatabaseConnection,
-            DuckDBConnection,
-            MySQLConnection,
-            PostgreSQLConnection,
-            SQLiteConnection,
         )
         from sqlsaber.database.resolver import DatabaseResolutionError, resolve_database
         from sqlsaber.threads import ThreadStorage
@@ -147,45 +143,32 @@ def query(
             raise CLIError(f"Error creating database connection: {e}")
 
         # Create pydantic-ai agent instance with database name for memory context
-        agent = build_sqlsaber_agent(db_conn, db_name)
+        sqlsaber_agent = SQLSaberAgent(db_conn, db_name, thinking_enabled=thinking)
 
         try:
             if actual_query:
                 # Single query mode with streaming
                 streaming_handler = StreamingQueryHandler(console)
-                # Compute DB type for the greeting line
-                if isinstance(db_conn, PostgreSQLConnection):
-                    db_type = "PostgreSQL"
-                elif isinstance(db_conn, MySQLConnection):
-                    db_type = "MySQL"
-                elif isinstance(db_conn, DuckDBConnection):
-                    db_type = "DuckDB"
-                elif isinstance(db_conn, SQLiteConnection):
-                    db_type = "SQLite"
-                elif isinstance(db_conn, CSVConnection):
-                    db_type = "DuckDB"
-                else:
-                    db_type = "database"
+                db_type = sqlsaber_agent.db_type
                 console.print(
                     f"[bold blue]Connected to:[/bold blue] {db_name} ({db_type})\n"
                 )
                 run = await streaming_handler.execute_streaming_query(
-                    actual_query, agent
+                    actual_query, sqlsaber_agent
                 )
                 # Persist non-interactive run as a thread snapshot so it can be resumed later
                 try:
                     if run is not None:
                         threads = ThreadStorage()
-                        # Extract title and model name
-                        title = actual_query
-                        model_name: str | None = agent.model.model_name
 
                         thread_id = await threads.save_snapshot(
                             messages_json=run.all_messages_json(),
                             database_name=db_name,
                         )
                         await threads.save_metadata(
-                            thread_id=thread_id, title=title, model_name=model_name
+                            thread_id=thread_id,
+                            title=actual_query,
+                            model_name=sqlsaber_agent.agent.model.model_name,
                         )
                         await threads.end_thread(thread_id)
                         console.print(
@@ -198,7 +181,7 @@ def query(
                     await threads.prune_threads()
             else:
                 # Interactive mode
-                session = InteractiveSession(console, agent, db_conn, db_name)
+                session = InteractiveSession(console, sqlsaber_agent, db_conn, db_name)
                 await session.run()
 
         finally:

sqlsaber/cli/completers.py

@@ -19,6 +19,8 @@ class SlashCommandCompleter(Completer):
                 ("clear", "Clear conversation history"),
                 ("exit", "Exit the interactive session"),
                 ("quit", "Exit the interactive session"),
+                ("thinking on", "Enable extended thinking/reasoning"),
+                ("thinking off", "Disable extended thinking/reasoning"),
             ]
 
             # Yield completions that match the partial command

sqlsaber/cli/database.py

@@ -354,7 +354,7 @@ def test(
 
     async def test_connection():
         # Lazy import to keep CLI startup fast
-        from sqlsaber.database.connection import DatabaseConnection
+        from sqlsaber.database import DatabaseConnection
 
         if name:
             db_config = config_manager.get_database(name)

sqlsaber/cli/display.py

@@ -8,7 +8,7 @@ rendered with Live.
 import json
 from typing import Sequence, Type
 
-from pydantic_ai.messages import ModelResponsePart, TextPart
+from pydantic_ai.messages import ModelResponsePart, TextPart, ThinkingPart
 from rich.columns import Columns
 from rich.console import Console, ConsoleOptions, RenderResult
 from rich.live import Live
@@ -75,7 +75,7 @@ class LiveMarkdownRenderer:
             self.end()
             self.paragraph_break()
 
-        self._start()
+        self._start(kind)
         self._current_kind = kind
 
     def append(self, text: str | None) -> None:
@@ -87,7 +87,13 @@ class LiveMarkdownRenderer:
             self.ensure_segment(TextPart)
 
         self._buffer += text
-        self._live.update(Markdown(self._buffer))
+
+        # Apply dim styling for thinking segments
+        if self._current_kind == ThinkingPart:
+            content = Markdown(self._buffer, style="dim")
+            self._live.update(content)
+        else:
+            self._live.update(Markdown(self._buffer))
 
     def end(self) -> None:
         """Finalize and stop the current Live segment, if any."""
@@ -95,13 +101,17 @@ class LiveMarkdownRenderer:
             return
         # Persist the *final* render exactly once, then shut Live down.
         buf = self._buffer
+        kind = self._current_kind
         self._live.stop()
         self._live = None
         self._buffer = ""
         self._current_kind = None
         # Print the complete markdown to scroll-back for permanent reference
         if buf:
-            self.console.print(Markdown(buf))
+            if kind == ThinkingPart:
+                self.console.print(Text(buf, style="dim"))
+            else:
+                self.console.print(Markdown(buf))
 
     def end_if_active(self) -> None:
         self.end()
@@ -153,10 +163,20 @@ class LiveMarkdownRenderer:
         text = Text(f" {message}", style="yellow")
         return Columns([spinner, text], expand=False)
 
-    def _start(self, initial_markdown: str = "") -> None:
+    def _start(
+        self, kind: Type[ModelResponsePart] | None = None, initial_markdown: str = ""
+    ) -> None:
         if self._live is not None:
             self.end()
         self._buffer = initial_markdown or ""
+
+        # Add visual styling for thinking segments
+        if kind == ThinkingPart:
+            if self.console.is_terminal:
+                self.console.print("[dim]💭 Thinking...[/dim]")
+            else:
+                self.console.print("*Thinking...*\n")
+
         # NOTE: Use transient=True so the live widget disappears on exit,
         # giving a clean transition to the final printed result.
         live = Live(
@@ -219,7 +239,9 @@ class DisplayManager:
             if self.console.is_terminal:
                 self.console.print("[dim bold]:gear: Executing SQL:[/dim bold]")
                 self.show_newline()
-                syntax = Syntax(query, "sql", background_color="default", word_wrap=True)
+                syntax = Syntax(
+                    query, "sql", background_color="default", word_wrap=True
+                )
                 self.console.print(syntax)
             else:
                 self.console.print("**Executing SQL:**\n")
@@ -271,9 +293,7 @@ class DisplayManager:
                     f"[yellow]... and {len(results) - 20} more rows[/yellow]"
                 )
             else:
-                self.console.print(
-                    f"*... and {len(results) - 20} more rows*\n"
-                )
+                self.console.print(f"*... and {len(results) - 20} more rows*\n")
 
     def show_error(self, error_message: str):
         """Display error message."""