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

sqlsaber/cli/models.py

@@ -25,11 +25,20 @@ models_app = cyclopts.App(
 class ModelManager:
     """Manages AI model configuration and fetching."""
 
-    DEFAULT_MODEL = "anthropic:claude-sonnet-4-20250514"
+    DEFAULT_MODEL = "anthropic:claude-sonnet-4-5-20250929"
     MODELS_API_URL = "https://models.dev/api.json"
     # Providers come from central registry
     SUPPORTED_PROVIDERS = providers.all_keys()
 
+    RECOMMENDED_MODELS = {
+        "anthropic": "claude-sonnet-4-5-20250929",
+        "openai": "gpt-5",
+        "google": "gemini-2.5-pro",
+        "groq": "llama-3-3-70b-versatile",
+        "mistral": "mistral-large-latest",
+        "cohere": "command-r-plus",
+    }
+
     async def fetch_available_models(
         self, providers: list[str] | None = None
     ) -> list[dict]:
@@ -180,39 +189,20 @@ def set():
     """Set the AI model to use."""
 
     async def interactive_set():
+        from sqlsaber.application.model_selection import choose_model, fetch_models
+        from sqlsaber.application.prompts import AsyncPrompter
+
         console.print("[blue]Fetching available models...[/blue]")
-        models = await model_manager.fetch_available_models()
+        models = await fetch_models(model_manager)
 
         if not models:
             console.print("[red]Failed to fetch models. Cannot set model.[/red]")
             sys.exit(1)
 
-        # Create choices for questionary
-        choices = []
-        for model in models:
-            # Format: "[provider] ID - Name (Description)"
-            prov = model.get("provider", "?")
-            choice_text = f"[{prov}] {model['id']} - {model['name']}"
-            if model["description"]:
-                choice_text += f" ({model['description'][:50]}{'...' if len(model['description']) > 50 else ''})"
-
-            choices.append({"name": choice_text, "value": model["id"]})
-
-        # Get current model to set as default
-        current_model = model_manager.get_current_model()
-        default_index = 0
-        for i, choice in enumerate(choices):
-            if choice["value"] == current_model:
-                default_index = i
-                break
-
-        selected_model = await questionary.select(
-            "Select a model:",
-            choices=choices,
-            use_search_filter=True,
-            use_jk_keys=False,  # Disable j/k keys when using search filter
-            default=choices[default_index] if choices else None,
-        ).ask_async()
+        prompter = AsyncPrompter()
+        selected_model = await choose_model(
+            prompter, models, restrict_provider=None, use_search_filter=True
+        )
 
         if selected_model:
             if model_manager.set_model(selected_model):

sqlsaber/cli/onboarding.py

@@ -0,0 +1,325 @@
+"""Interactive onboarding flow for first-time SQLSaber users."""
+
+import sys
+
+from rich.console import Console
+from rich.panel import Panel
+
+from sqlsaber.cli.models import ModelManager
+from sqlsaber.config.api_keys import APIKeyManager
+from sqlsaber.config.auth import AuthConfigManager
+from sqlsaber.config.database import DatabaseConfigManager
+
+console = Console()
+
+
+def needs_onboarding(database_arg: str | None = None) -> bool:
+    """Check if user needs onboarding.
+
+    Onboarding is needed if:
+    - No database is configured AND no database connection string provided via CLI
+    """
+    # If user provided a database argument, skip onboarding
+    if database_arg:
+        return False
+
+    # Check if databases are configured
+    db_manager = DatabaseConfigManager()
+    has_db = db_manager.has_databases()
+
+    return not has_db
+
+
+def welcome_screen() -> None:
+    """Display welcome screen to new users."""
+    banner = """
+███████  ██████  ██      ███████  █████  ██████  ███████ ██████
+██      ██    ██ ██      ██      ██   ██ ██   ██ ██      ██   ██
+███████ ██    ██ ██      ███████ ███████ ██████  █████   ██████
+     ██ ██ ▄▄ ██ ██           ██ ██   ██ ██   ██ ██      ██   ██
+███████  ██████  ███████ ███████ ██   ██ ██████  ███████ ██   ██
+           ▀▀
+    """
+
+    console.print(Panel.fit(banner, style="bold blue"))
+    console.print()
+
+    welcome_message = """
+[bold]Welcome to SQLsaber! 🎉[/bold]
+
+SQLsaber is an agentic SQL assistant that lets you query your database using natural language.
+
+Let's get you set up in just a few steps.
+    """
+
+    console.print(Panel(welcome_message.strip(), border_style="blue", padding=(1, 2)))
+    console.print()
+
+
+async def setup_database_guided() -> str | None:
+    """Guide user through database setup.
+
+    Returns the name of the configured database or None if cancelled.
+    """
+    from sqlsaber.application.db_setup import (
+        build_config,
+        collect_db_input,
+        save_database,
+        test_connection,
+    )
+    from sqlsaber.application.prompts import AsyncPrompter
+
+    console.print("━" * 80, style="dim")
+    console.print("[bold cyan]Step 1 of 2: Database Connection[/bold cyan]")
+    console.print("━" * 80, style="dim")
+    console.print()
+
+    try:
+        # Ask for connection name
+        prompter = AsyncPrompter()
+        name = await prompter.text(
+            "What would you like to name this connection?",
+            default="mydb",
+            validate=lambda x: bool(x.strip()) or "Name cannot be empty",
+        )
+
+        if name is None:
+            return None
+
+        name = name.strip()
+
+        # Check if name already exists
+        db_manager = DatabaseConfigManager()
+        if db_manager.get_database(name):
+            console.print(
+                f"[yellow]Database connection '{name}' already exists.[/yellow]"
+            )
+            return name
+
+        # Collect database input (simplified - no SSL in onboarding)
+        db_input = await collect_db_input(
+            prompter=prompter, name=name, db_type="postgresql", include_ssl=False
+        )
+
+        if db_input is None:
+            return None
+
+        # Build config
+        db_config = build_config(db_input)
+
+        # Test the connection
+        console.print(f"[dim]Testing connection to '{name}'...[/dim]")
+        connection_success = await test_connection(db_config, db_input.password)
+
+        if not connection_success:
+            retry = await prompter.confirm(
+                "Would you like to try again with different settings?", default=True
+            )
+            if retry:
+                return await setup_database_guided()
+            else:
+                console.print(
+                    "[yellow]You can add a database later using 'saber db add'[/yellow]"
+                )
+                return None
+
+        # Save the configuration
+        try:
+            save_database(db_manager, db_config, db_input.password)
+            console.print(f"[green]✓ Connection to '{name}' successful![/green]")
+            console.print()
+            return name
+        except Exception as e:
+            console.print(f"[bold red]Error saving database:[/bold red] {e}")
+            return None
+
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Setup cancelled.[/yellow]")
+        return None
+    except Exception as e:
+        console.print(f"[bold red]Unexpected error:[/bold red] {e}")
+        return None
+
+
+async def select_model_for_provider(provider: str) -> str | None:
+    """Fetch and let user select a model for the given provider.
+
+    Returns the selected model ID or None if cancelled/failed.
+    """
+    from sqlsaber.application.model_selection import choose_model, fetch_models
+    from sqlsaber.application.prompts import AsyncPrompter
+
+    try:
+        console.print()
+        console.print(f"[dim]Fetching available {provider.title()} models...[/dim]")
+
+        model_manager = ModelManager()
+        models = await fetch_models(model_manager, providers=[provider])
+
+        if not models:
+            console.print(
+                f"[yellow]Could not fetch models for {provider}. Using default.[/yellow]"
+            )
+            # Use provider-specific default or fallback to Anthropic
+            default_model_id = ModelManager.RECOMMENDED_MODELS.get(
+                provider, ModelManager.DEFAULT_MODEL
+            )
+            # Format it properly if we have a recommended model for this provider
+            if provider in ModelManager.RECOMMENDED_MODELS:
+                return f"{provider}:{ModelManager.RECOMMENDED_MODELS[provider]}"
+            return default_model_id
+
+        prompter = AsyncPrompter()
+        console.print()
+        selected_model = await choose_model(
+            prompter, models, restrict_provider=provider, use_search_filter=True
+        )
+
+        return selected_model
+
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Model selection cancelled.[/yellow]")
+        return None
+    except Exception as e:
+        console.print(f"[yellow]Error selecting model: {e}. Using default.[/yellow]")
+        # Fallback to provider default
+        if provider in ModelManager.RECOMMENDED_MODELS:
+            return f"{provider}:{ModelManager.RECOMMENDED_MODELS[provider]}"
+        return ModelManager.DEFAULT_MODEL
+
+
+async def setup_auth_guided() -> tuple[bool, str | None]:
+    """Guide user through auth setup.
+
+    Returns tuple of (success: bool, selected_model: str | None).
+    """
+    from sqlsaber.application.auth_setup import setup_auth
+    from sqlsaber.application.prompts import AsyncPrompter
+
+    console.print("━" * 80, style="dim")
+    console.print("[bold cyan]Step 2 of 2: Authentication[/bold cyan]")
+    console.print("━" * 80, style="dim")
+    console.print()
+
+    try:
+        # Run auth setup
+        prompter = AsyncPrompter()
+        auth_manager = AuthConfigManager()
+        api_key_manager = APIKeyManager()
+
+        success, provider = await setup_auth(
+            prompter=prompter,
+            auth_manager=auth_manager,
+            api_key_manager=api_key_manager,
+            allow_oauth=True,
+            default_provider="anthropic",
+            run_oauth_in_thread=True,
+        )
+
+        if not success:
+            console.print(
+                "[yellow]You can set it up later using 'saber auth setup'[/yellow]"
+            )
+            console.print()
+            return False, None
+
+        # If auth configured but we don't know the provider (already configured case)
+        if provider is None:
+            console.print()
+            return True, None
+
+        # Select model for this provider
+        selected_model = await select_model_for_provider(provider)
+        if selected_model:
+            model_manager = ModelManager()
+            model_manager.set_model(selected_model)
+            console.print(f"[green]✓ Model set to: {selected_model}[/green]")
+        console.print()
+        return True, selected_model
+
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Setup cancelled.[/yellow]")
+        console.print()
+        return False, None
+    except Exception as e:
+        console.print(f"[bold red]Unexpected error:[/bold red] {e}")
+        console.print()
+        return False, None
+
+
+def success_screen(
+    database_name: str | None, auth_configured: bool, model_name: str | None = None
+) -> None:
+    """Display success screen after onboarding."""
+    console.print("━" * 80, style="dim")
+    console.print("[bold green]You're all set! 🚀[/bold green]")
+    console.print("━" * 80, style="dim")
+    console.print()
+
+    if database_name and auth_configured:
+        console.print(
+            f"[green]✓ Database '{database_name}' connected and ready to use[/green]"
+        )
+        console.print("[green]✓ Authentication configured[/green]")
+        if model_name:
+            console.print(f"[green]✓ Model: {model_name}[/green]")
+    elif database_name:
+        console.print(
+            f"[green]✓ Database '{database_name}' connected and ready to use[/green]"
+        )
+        console.print(
+            "[yellow]⚠ AI authentication not configured - you'll be prompted when needed[/yellow]"
+        )
+    elif auth_configured:
+        console.print("[green]✓ AI authentication configured[/green]")
+        if model_name:
+            console.print(f"[green]✓ Model: {model_name}[/green]")
+        console.print(
+            "[yellow]⚠ No database configured - you'll need to provide one via -d flag[/yellow]"
+        )
+
+    console.print()
+    console.print("[dim]Starting interactive session...[/dim]")
+    console.print()
+
+
+async def run_onboarding() -> bool:
+    """Run the complete onboarding flow.
+
+    Returns True if onboarding completed successfully (at least database configured),
+    False if user cancelled or onboarding failed.
+    """
+    try:
+        # Welcome screen
+        welcome_screen()
+
+        # Database setup
+        database_name = await setup_database_guided()
+
+        # If user cancelled database setup, exit
+        if database_name is None:
+            console.print("[yellow]Database setup is required to continue.[/yellow]")
+            console.print(
+                "[dim]You can also provide a connection string using: saber -d <connection-string>[/dim]"
+            )
+            return False
+
+        # Auth setup
+        auth_configured, model_name = await setup_auth_guided()
+
+        # Show success screen
+        success_screen(database_name, auth_configured, model_name)
+
+        return True
+
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Onboarding cancelled.[/yellow]")
+        console.print(
+            "[dim]You can run setup commands manually:[/dim]\n"
+            "[dim]  - saber db add <name>  # Add database connection[/dim]\n"
+            "[dim]  - saber auth setup     # Configure authentication[/dim]"
+        )
+        sys.exit(0)
+    except Exception as e:
+        console.print(f"[bold red]Onboarding failed:[/bold red] {e}")
+        return False

sqlsaber/cli/streaming.py

@@ -8,9 +8,9 @@ rendered via DisplayManager helpers.
 import asyncio
 import json
 from functools import singledispatchmethod
-from typing import AsyncIterable
+from typing import TYPE_CHECKING, AsyncIterable
 
-from pydantic_ai import Agent, RunContext
+from pydantic_ai import RunContext
 from pydantic_ai.messages import (
     AgentStreamEvent,
     FunctionToolCallEvent,
@@ -26,6 +26,9 @@ from rich.console import Console
 
 from sqlsaber.cli.display import DisplayManager
 
+if TYPE_CHECKING:
+    from sqlsaber.agents.pydantic_ai_agent import SQLSaberAgent
+
 
 class StreamingQueryHandler:
     """
@@ -130,7 +133,7 @@ class StreamingQueryHandler:
     async def execute_streaming_query(
         self,
         user_query: str,
-        agent: Agent,
+        sqlsaber_agent: "SQLSaberAgent",
         cancellation_token: asyncio.Event | None = None,
         message_history: list | None = None,
     ):
@@ -139,21 +142,16 @@ class StreamingQueryHandler:
         try:
             # If Anthropic OAuth, inject SQLsaber instructions before the first user prompt
             prepared_prompt: str | list[str] = user_query
-            is_oauth = bool(getattr(agent, "_sqlsaber_is_oauth", False))
             no_history = not message_history
-            if is_oauth and no_history:
-                ib = getattr(agent, "_sqlsaber_instruction_builder", None)
-                mm = getattr(agent, "_sqlsaber_memory_manager", None)
-                db_type = getattr(agent, "_sqlsaber_db_type", "database")
-                db_name = getattr(agent, "_sqlsaber_database_name", None)
-                instructions = (
-                    ib.build_instructions(db_type=db_type) if ib is not None else ""
-                )
-                mem = (
-                    mm.format_memories_for_prompt(db_name)
-                    if (mm is not None and db_name)
-                    else ""
+            if sqlsaber_agent.is_oauth and no_history:
+                instructions = sqlsaber_agent.instruction_builder.build_instructions(
+                    db_type=sqlsaber_agent.db_type
                 )
+                mem = ""
+                if sqlsaber_agent.database_name:
+                    mem = sqlsaber_agent.memory_manager.format_memories_for_prompt(
+                        sqlsaber_agent.database_name
+                    )
                 parts = [p for p in (instructions, mem) if p and str(p).strip()]
                 if parts:
                     injected = "\n\n".join(parts)
@@ -163,7 +161,7 @@ class StreamingQueryHandler:
             self.display.live.start_status("Crunching data...")
 
             # Run the agent with our event stream handler
-            run = await agent.run(
+            run = await sqlsaber_agent.agent.run(
                 prepared_prompt,
                 message_history=message_history,
                 event_stream_handler=self._event_stream_handler,

sqlsaber/cli/threads.py

@@ -148,7 +148,9 @@ def _render_transcript(
                             )
                         else:
                             if is_redirected:
-                                console.print(f"**Tool result ({name}):**\n\n{content_str}\n")
+                                console.print(
+                                    f"**Tool result ({name}):**\n\n{content_str}\n"
+                                )
                             else:
                                 console.print(
                                     Panel.fit(
@@ -159,7 +161,9 @@ def _render_transcript(
                                 )
                     except Exception:
                         if is_redirected:
-                            console.print(f"**Tool result ({name}):**\n\n{content_str}\n")
+                            console.print(
+                                f"**Tool result ({name}):**\n\n{content_str}\n"
+                            )
                         else:
                             console.print(
                                 Panel.fit(
@@ -258,10 +262,10 @@ def resume(
 
     async def _run() -> None:
         # Lazy imports to avoid heavy modules at CLI startup
-        from sqlsaber.agents import build_sqlsaber_agent
+        from sqlsaber.agents import SQLSaberAgent
         from sqlsaber.cli.interactive import InteractiveSession
         from sqlsaber.config.database import DatabaseConfigManager
-        from sqlsaber.database.connection import DatabaseConnection
+        from sqlsaber.database import DatabaseConnection
         from sqlsaber.database.resolver import (
             DatabaseResolutionError,
             resolve_database,
@@ -288,7 +292,7 @@ def resume(
 
         db_conn = DatabaseConnection(connection_string)
         try:
-            agent = build_sqlsaber_agent(db_conn, db_name)
+            sqlsaber_agent = SQLSaberAgent(db_conn, db_name)
             history = await store.get_thread_messages(thread_id)
             if console.is_terminal:
                 console.print(Panel.fit(f"Thread: {thread.id}", border_style="blue"))
@@ -297,7 +301,7 @@ def resume(
             _render_transcript(console, history, None)
             session = InteractiveSession(
                 console=console,
-                agent=agent,
+                sqlsaber_agent=sqlsaber_agent,
                 db_conn=db_conn,
                 database_name=db_name,
                 initial_thread_id=thread_id,

sqlsaber/config/api_keys.py

@@ -19,7 +19,7 @@ class APIKeyManager:
 
     def get_api_key(self, provider: str) -> str | None:
         """Get API key for the specified provider using cascading logic."""
-        env_var_name = self._get_env_var_name(provider)
+        env_var_name = self.get_env_var_name(provider)
         service_name = self._get_service_name(provider)
 
         # 1. Check environment variable first
@@ -41,7 +41,7 @@ class APIKeyManager:
         # 3. Prompt user for API key
         return self._prompt_and_store_key(provider, env_var_name, service_name)
 
-    def _get_env_var_name(self, provider: str) -> str:
+    def get_env_var_name(self, provider: str) -> str:
         """Get the expected environment variable name for a provider."""
         # Normalize aliases to canonical provider keys
         key = providers.canonical(provider) or provider

sqlsaber/config/settings.py

@@ -46,7 +46,10 @@ class ModelConfigManager:
     def _load_config(self) -> dict[str, Any]:
         """Load configuration from file."""
         if not self.config_file.exists():
-            return {"model": self.DEFAULT_MODEL}
+            return {
+                "model": self.DEFAULT_MODEL,
+                "thinking_enabled": False,
+            }
 
         try:
             with open(self.config_file, "r") as f:
@@ -54,9 +57,15 @@ class ModelConfigManager:
                 # Ensure we have a model set
                 if "model" not in config:
                     config["model"] = self.DEFAULT_MODEL
+                # Set defaults for thinking if not present
+                if "thinking_enabled" not in config:
+                    config["thinking_enabled"] = False
                 return config
         except (json.JSONDecodeError, IOError):
-            return {"model": self.DEFAULT_MODEL}
+            return {
+                "model": self.DEFAULT_MODEL,
+                "thinking_enabled": False,
+            }
 
     def _save_config(self, config: dict[str, Any]) -> None:
         """Save configuration to file."""
@@ -76,6 +85,17 @@ class ModelConfigManager:
         config["model"] = model
         self._save_config(config)
 
+    def get_thinking_enabled(self) -> bool:
+        """Get whether thinking is enabled."""
+        config = self._load_config()
+        return config.get("thinking_enabled", False)
+
+    def set_thinking_enabled(self, enabled: bool) -> None:
+        """Set whether thinking is enabled."""
+        config = self._load_config()
+        config["thinking_enabled"] = enabled
+        self._save_config(config)
+
 
 class Config:
     """Configuration class for SQLSaber."""
@@ -86,6 +106,9 @@ class Config:
         self.api_key_manager = APIKeyManager()
         self.auth_config_manager = AuthConfigManager()
 
+        # Thinking configuration
+        self.thinking_enabled = self.model_config_manager.get_thinking_enabled()
+
         # Authentication method (API key or Anthropic OAuth)
         self.auth_method = self.auth_config_manager.get_auth_method()
 

sqlsaber/database/__init__.py

@@ -1,9 +1,63 @@
 """Database module for SQLSaber."""
 
-from .connection import DatabaseConnection
+from .base import (
+    DEFAULT_QUERY_TIMEOUT,
+    BaseDatabaseConnection,
+    BaseSchemaIntrospector,
+    ColumnInfo,
+    ForeignKeyInfo,
+    IndexInfo,
+    QueryTimeoutError,
+    SchemaInfo,
+)
+from .csv import CSVConnection, CSVSchemaIntrospector
+from .duckdb import DuckDBConnection, DuckDBSchemaIntrospector
+from .mysql import MySQLConnection, MySQLSchemaIntrospector
+from .postgresql import PostgreSQLConnection, PostgreSQLSchemaIntrospector
 from .schema import SchemaManager
+from .sqlite import SQLiteConnection, SQLiteSchemaIntrospector
+
+
+def DatabaseConnection(connection_string: str) -> BaseDatabaseConnection:
+    """Factory function to create appropriate database connection based on connection string."""
+    if connection_string.startswith("postgresql://"):
+        return PostgreSQLConnection(connection_string)
+    elif connection_string.startswith("mysql://"):
+        return MySQLConnection(connection_string)
+    elif connection_string.startswith("sqlite:///"):
+        return SQLiteConnection(connection_string)
+    elif connection_string.startswith("duckdb://"):
+        return DuckDBConnection(connection_string)
+    elif connection_string.startswith("csv:///"):
+        return CSVConnection(connection_string)
+    else:
+        raise ValueError(
+            f"Unsupported database type in connection string: {connection_string}"
+        )
+
 
 __all__ = [
+    # Base classes and types
+    "BaseDatabaseConnection",
+    "BaseSchemaIntrospector",
+    "ColumnInfo",
+    "DEFAULT_QUERY_TIMEOUT",
+    "ForeignKeyInfo",
+    "IndexInfo",
+    "QueryTimeoutError",
+    "SchemaInfo",
+    # Concrete implementations
+    "PostgreSQLConnection",
+    "MySQLConnection",
+    "SQLiteConnection",
+    "DuckDBConnection",
+    "CSVConnection",
+    "PostgreSQLSchemaIntrospector",
+    "MySQLSchemaIntrospector",
+    "SQLiteSchemaIntrospector",
+    "DuckDBSchemaIntrospector",
+    "CSVSchemaIntrospector",
+    # Factory function and manager
     "DatabaseConnection",
     "SchemaManager",
 ]