sqlsaber 0.24.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,9 +5,10 @@ import json
 from abc import ABC, abstractmethod
 from typing import Any, AsyncIterator
 
-from sqlsaber.database.connection import (
+from sqlsaber.database import (
     BaseDatabaseConnection,
     CSVConnection,
+    DuckDBConnection,
     MySQLConnection,
     PostgreSQLConnection,
     SQLiteConnection,
@@ -51,7 +52,9 @@ class BaseSQLAgent(ABC):
         elif isinstance(self.db, SQLiteConnection):
             return "SQLite"
         elif isinstance(self.db, CSVConnection):
-            return "SQLite"  # we convert csv to in-memory sqlite
+            return "DuckDB"
+        elif isinstance(self.db, DuckDBConnection):
+            return "DuckDB"
         else:
             return "database"  # Fallback
 

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,17 +6,19 @@ 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,
     MySQLConnection,
     PostgreSQLConnection,
     SQLiteConnection,
@@ -27,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)",
@@ -78,98 +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, CSVConnection):
-        return "SQLite"
-    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

@@ -46,7 +46,7 @@ def meta_handler(
         str | None,
         cyclopts.Parameter(
             ["--database", "-d"],
-            help="Database connection name, file path (CSV/SQLite), or connection string (postgresql://, mysql://) (uses default if not specified)",
+            help="Database connection name, file path (CSV/SQLite/DuckDB), or connection string (postgresql://, mysql://, duckdb://) (uses default if not specified)",
         ),
     ] = None,
 ):
@@ -59,8 +59,10 @@ def meta_handler(
         saber -d mydb "show me users"          # Run a query with specific database
         saber -d data.csv "show me users"      # Run a query with ad-hoc CSV file
         saber -d data.db "show me users"       # Run a query with ad-hoc SQLite file
+        saber -d data.duckdb "show me users"   # Run a query with ad-hoc DuckDB file
         saber -d "postgresql://user:pass@host:5432/db" "show users"  # PostgreSQL connection string
         saber -d "mysql://user:pass@host:3306/db" "show users"       # MySQL connection string
+        saber -d "duckdb:///data.duckdb" "show users"                 # DuckDB connection string
         echo "show me all users" | saber       # Read query from stdin
         cat query.txt | saber                  # Read query from file via stdin
     """
@@ -73,16 +75,17 @@ 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[
         str | None,
         cyclopts.Parameter(
             ["--database", "-d"],
-            help="Database connection name, file path (CSV/SQLite), or connection string (postgresql://, mysql://) (uses default if not specified)",
+            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.
 
@@ -97,23 +100,21 @@ def query(
         saber "show me all users"         # Run a single query
         saber -d data.csv "show users"    # Run a query with ad-hoc CSV file
         saber -d data.db "show users"     # Run a query with ad-hoc SQLite file
+        saber -d data.duckdb "show users" # Run a query with ad-hoc DuckDB file
         saber -d "postgresql://user:pass@host:5432/db" "show users"  # PostgreSQL connection string
         saber -d "mysql://user:pass@host:3306/db" "show users"       # MySQL connection string
+        saber -d "duckdb:///data.duckdb" "show users"                 # DuckDB connection string
         echo "show me all users" | saber  # Read query from stdin
     """
 
     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,
-            MySQLConnection,
-            PostgreSQLConnection,
-            SQLiteConnection,
         )
         from sqlsaber.database.resolver import DatabaseResolutionError, resolve_database
         from sqlsaber.threads import ThreadStorage
@@ -142,42 +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
-                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"
-                )
+                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(
@@ -190,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

@@ -31,7 +31,7 @@ def add(
         str,
         cyclopts.Parameter(
             ["--type", "-t"],
-            help="Database type (postgresql, mysql, sqlite)",
+            help="Database type (postgresql, mysql, sqlite, duckdb)",
         ),
     ] = "postgresql",
     host: Annotated[
@@ -87,17 +87,17 @@ def add(
         if not type or type == "postgresql":
             type = questionary.select(
                 "Database type:",
-                choices=["postgresql", "mysql", "sqlite"],
+                choices=["postgresql", "mysql", "sqlite", "duckdb"],
                 default="postgresql",
             ).ask()
 
-        if type == "sqlite":
-            # SQLite only needs database path
+        if type in {"sqlite", "duckdb"}:
+            # SQLite/DuckDB only need database file path
             database = database or questionary.path("Database file path:").ask()
             database = str(Path(database).expanduser().resolve())
             host = "localhost"
             port = 0
-            username = "sqlite"
+            username = type
             password = ""
         else:
             # PostgreSQL/MySQL need connection details
@@ -182,6 +182,17 @@ def add(
             port = 0
             username = "sqlite"
             password = ""
+        elif type == "duckdb":
+            if not database:
+                console.print(
+                    "[bold red]Error:[/bold red] Database file path is required for DuckDB"
+                )
+                sys.exit(1)
+            database = str(Path(database).expanduser().resolve())
+            host = "localhost"
+            port = 0
+            username = "duckdb"
+            password = ""
         else:
             if not all([host, database, username]):
                 console.print(
@@ -264,7 +275,7 @@ def list():
             if db.ssl_ca or db.ssl_cert:
                 ssl_status += " (certs)"
         else:
-            ssl_status = "disabled" if db.type != "sqlite" else "N/A"
+            ssl_status = "disabled" if db.type not in {"sqlite", "duckdb"} else "N/A"
 
         table.add_row(
             db.name,
@@ -343,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)