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

sqlsaber/application/db_setup.py

@@ -0,0 +1,223 @@
+"""Shared database setup logic for onboarding and CLI."""
+
+import getpass
+from dataclasses import dataclass
+from pathlib import Path
+
+from rich.console import Console
+
+from sqlsaber.application.prompts import Prompter
+from sqlsaber.config.database import DatabaseConfig, DatabaseConfigManager
+
+console = Console()
+
+
+@dataclass
+class DatabaseInput:
+    """Input data for database configuration."""
+
+    name: str
+    type: str
+    host: str
+    port: int
+    database: str
+    username: str
+    password: str | None
+    ssl_mode: str | None = None
+    ssl_ca: str | None = None
+    ssl_cert: str | None = None
+    ssl_key: str | None = None
+
+
+async def collect_db_input(
+    prompter: Prompter,
+    name: str,
+    db_type: str = "postgresql",
+    include_ssl: bool = True,
+) -> DatabaseInput | None:
+    """Collect database connection details interactively.
+
+    Args:
+        prompter: Prompter instance for interaction
+        name: Database connection name
+        db_type: Initial database type (can be changed via prompt)
+        include_ssl: Whether to prompt for SSL configuration
+
+    Returns:
+        DatabaseInput with collected values or None if cancelled
+    """
+    # Ask for database type
+    db_type = await prompter.select(
+        "Database type:",
+        choices=["postgresql", "mysql", "sqlite", "duckdb"],
+        default=db_type,
+    )
+
+    if db_type is None:
+        return None
+
+    # Handle file-based databases
+    if db_type in {"sqlite", "duckdb"}:
+        database_path = await prompter.path(
+            f"{db_type.upper()} file path:", only_directories=False
+        )
+
+        if database_path is None:
+            return None
+
+        database = str(Path(database_path).expanduser().resolve())
+        host = "localhost"
+        port = 0
+        username = db_type
+        password = ""
+        ssl_mode = None
+        ssl_ca = None
+        ssl_cert = None
+        ssl_key = None
+
+    else:
+        # PostgreSQL/MySQL need connection details
+        host = await prompter.text("Host:", default="localhost")
+        if host is None:
+            return None
+
+        default_port = 5432 if db_type == "postgresql" else 3306
+        port_str = await prompter.text("Port:", default=str(default_port))
+        if port_str is None:
+            return None
+
+        try:
+            port = int(port_str)
+        except ValueError:
+            console.print("[red]Invalid port number. Using default.[/red]")
+            port = default_port
+
+        database = await prompter.text("Database name:")
+        if database is None:
+            return None
+
+        username = await prompter.text("Username:")
+        if username is None:
+            return None
+
+        password = getpass.getpass("Password (stored in your OS keychain): ")
+
+        ssl_mode = None
+        ssl_ca = None
+        ssl_cert = None
+        ssl_key = None
+
+        # Ask for SSL configuration if enabled
+        if include_ssl:
+            configure_ssl = await prompter.confirm(
+                "Configure SSL/TLS settings?", default=False
+            )
+            if configure_ssl:
+                if db_type == "postgresql":
+                    ssl_mode = await prompter.select(
+                        "SSL mode for PostgreSQL:",
+                        choices=[
+                            "disable",
+                            "allow",
+                            "prefer",
+                            "require",
+                            "verify-ca",
+                            "verify-full",
+                        ],
+                        default="prefer",
+                    )
+                elif db_type == "mysql":
+                    ssl_mode = await prompter.select(
+                        "SSL mode for MySQL:",
+                        choices=[
+                            "DISABLED",
+                            "PREFERRED",
+                            "REQUIRED",
+                            "VERIFY_CA",
+                            "VERIFY_IDENTITY",
+                        ],
+                        default="PREFERRED",
+                    )
+
+                if ssl_mode and ssl_mode not in ["disable", "DISABLED"]:
+                    specify_certs = await prompter.confirm(
+                        "Specify SSL certificate files?", default=False
+                    )
+                    if specify_certs:
+                        ssl_ca = await prompter.path("SSL CA certificate file:")
+                        specify_client = await prompter.confirm(
+                            "Specify client certificate?", default=False
+                        )
+                        if specify_client:
+                            ssl_cert = await prompter.path(
+                                "SSL client certificate file:"
+                            )
+                            ssl_key = await prompter.path(
+                                "SSL client private key file:"
+                            )
+
+    return DatabaseInput(
+        name=name,
+        type=db_type,
+        host=host,
+        port=port,
+        database=database,
+        username=username,
+        password=password,
+        ssl_mode=ssl_mode,
+        ssl_ca=ssl_ca,
+        ssl_cert=ssl_cert,
+        ssl_key=ssl_key,
+    )
+
+
+def build_config(db_input: DatabaseInput) -> DatabaseConfig:
+    """Build DatabaseConfig from DatabaseInput."""
+    return DatabaseConfig(
+        name=db_input.name,
+        type=db_input.type,
+        host=db_input.host,
+        port=db_input.port,
+        database=db_input.database,
+        username=db_input.username,
+        ssl_mode=db_input.ssl_mode,
+        ssl_ca=db_input.ssl_ca,
+        ssl_cert=db_input.ssl_cert,
+        ssl_key=db_input.ssl_key,
+    )
+
+
+async def test_connection(config: DatabaseConfig, password: str | None) -> bool:
+    """Test database connection.
+
+    Args:
+        config: DatabaseConfig to test
+        password: Password for connection (not stored in config yet)
+
+    Returns:
+        True if connection successful, False otherwise
+    """
+    from sqlsaber.database import DatabaseConnection
+
+    try:
+        connection_string = config.to_connection_string()
+        db_conn = DatabaseConnection(connection_string)
+        await db_conn.execute_query("SELECT 1 as test")
+        await db_conn.close()
+        return True
+    except Exception as e:
+        console.print(f"[bold red]Connection failed:[/bold red] {e}", style="red")
+        return False
+
+
+def save_database(
+    config_manager: DatabaseConfigManager, config: DatabaseConfig, password: str | None
+) -> None:
+    """Save database configuration.
+
+    Args:
+        config_manager: DatabaseConfigManager instance
+        config: DatabaseConfig to save
+        password: Password to store in keyring (if provided)
+    """
+    config_manager.add_database(config, password if password else None)

sqlsaber/application/model_selection.py

@@ -0,0 +1,98 @@
+"""Shared model selection logic for onboarding and CLI."""
+
+from questionary import Choice
+from rich.console import Console
+
+from sqlsaber.application.prompts import Prompter
+from sqlsaber.cli.models import ModelManager
+
+console = Console()
+
+
+async def fetch_models(
+    model_manager: ModelManager, providers: list[str] | None = None
+) -> list[dict]:
+    """Fetch available models from models.dev API."""
+    return await model_manager.fetch_available_models(providers=providers)
+
+
+async def choose_model(
+    prompter: Prompter,
+    models: list[dict],
+    restrict_provider: str | None = None,
+    use_search_filter: bool = True,
+) -> str | None:
+    """Interactive model selection with recommended models prioritized.
+
+    Args:
+        prompter: Prompter instance for interaction
+        models: List of model dicts from fetch_models
+        restrict_provider: If set, only show models from this provider and use provider-specific recommendation
+        use_search_filter: Enable search filter for large lists
+
+    Returns:
+        Selected model ID (provider:model_id) or None if cancelled
+    """
+    if not models:
+        console.print("[yellow]No models available[/yellow]")
+        return None
+
+    # Filter by provider if restricted
+    if restrict_provider:
+        models = [m for m in models if m.get("provider") == restrict_provider]
+        if not models:
+            console.print(
+                f"[yellow]No models available for {restrict_provider}[/yellow]"
+            )
+            return None
+
+    # Get recommended model for the provider
+    recommended_id = None
+    if restrict_provider and restrict_provider in ModelManager.RECOMMENDED_MODELS:
+        recommended_id = ModelManager.RECOMMENDED_MODELS[restrict_provider]
+
+    # Build choices
+    choices = []
+    recommended_index = 0
+
+    for i, model in enumerate(models):
+        model_id_without_provider = model["id"].split(":", 1)[1]
+        is_recommended = recommended_id == model_id_without_provider
+
+        choice_text = model["name"]
+        if is_recommended:
+            choice_text += " (Recommended)"
+            recommended_index = i
+        elif model["description"]:
+            desc_short = model["description"][:40]
+            choice_text += (
+                f" ({desc_short}...)"
+                if len(model["description"]) > 40
+                else f" ({desc_short})"
+            )
+
+        choices.append(Choice(choice_text, value=model["id"]))
+
+    # Move recommended model to top if it exists
+    if recommended_index > 0:
+        choices.insert(0, choices.pop(recommended_index))
+
+    # Prompt user
+    selected_model = await prompter.select(
+        "Select a model:",
+        choices=choices,
+        use_search_filter=use_search_filter,
+    )
+
+    if selected_model:
+        return selected_model
+
+    # User cancelled, return recommended or first available
+    if recommended_id and restrict_provider:
+        return f"{restrict_provider}:{recommended_id}"
+    return models[0]["id"] if models else None
+
+
+def set_model(model_manager: ModelManager, model_id: str) -> bool:
+    """Set the current model."""
+    return model_manager.set_model(model_id)

sqlsaber/application/prompts.py

@@ -0,0 +1,115 @@
+"""Prompter abstraction for sync/async questionary interactions."""
+
+from abc import ABC, abstractmethod
+from typing import Any, Callable
+
+import questionary
+from questionary import Choice
+
+
+class Prompter(ABC):
+    """Abstract base class for interactive prompting."""
+
+    @abstractmethod
+    async def text(
+        self,
+        message: str,
+        default: str = "",
+        validate: Callable[[str], bool | str] | None = None,
+    ) -> str | None:
+        """Prompt for text input."""
+        pass
+
+    @abstractmethod
+    async def select(
+        self,
+        message: str,
+        choices: list[str] | list[Choice] | list[dict],
+        default: Any = None,
+        use_search_filter: bool = False,
+        use_jk_keys: bool = True,
+    ) -> Any:
+        """Prompt for selection from choices."""
+        pass
+
+    @abstractmethod
+    async def confirm(self, message: str, default: bool = False) -> bool | None:
+        """Prompt for yes/no confirmation."""
+        pass
+
+    @abstractmethod
+    async def path(self, message: str, only_directories: bool = False) -> str | None:
+        """Prompt for file/directory path."""
+        pass
+
+
+class AsyncPrompter(Prompter):
+    """Async prompter using questionary.ask_async() for onboarding."""
+
+    async def text(
+        self,
+        message: str,
+        default: str = "",
+        validate: Callable[[str], bool | str] | None = None,
+    ) -> str | None:
+        return await questionary.text(
+            message, default=default, validate=validate
+        ).ask_async()
+
+    async def select(
+        self,
+        message: str,
+        choices: list[str] | list[Choice] | list[dict],
+        default: Any = None,
+        use_search_filter: bool = True,
+        use_jk_keys: bool = False,
+    ) -> Any:
+        return await questionary.select(
+            message,
+            choices=choices,
+            default=default,
+            use_search_filter=use_search_filter,
+            use_jk_keys=use_jk_keys,
+        ).ask_async()
+
+    async def confirm(self, message: str, default: bool = False) -> bool | None:
+        return await questionary.confirm(message, default=default).ask_async()
+
+    async def path(self, message: str, only_directories: bool = False) -> str | None:
+        return await questionary.path(
+            message, only_directories=only_directories
+        ).ask_async()
+
+
+class SyncPrompter(Prompter):
+    """Sync prompter using questionary.ask() for CLI commands."""
+
+    async def text(
+        self,
+        message: str,
+        default: str = "",
+        validate: Callable[[str], bool | str] | None = None,
+    ) -> str | None:
+        return questionary.text(message, default=default, validate=validate).ask()
+
+    async def select(
+        self,
+        message: str,
+        choices: list[str] | list[Choice] | list[dict],
+        default: Any = None,
+        use_search_filter: bool = True,
+        use_jk_keys: bool = False,
+    ) -> Any:
+        return questionary.select(
+            message,
+            choices=choices,
+            default=default,
+            use_search_filter=use_search_filter,
+            use_jk_keys=use_jk_keys,
+        ).ask()
+
+    async def confirm(self, message: str, default: bool = False) -> bool | None:
+        return questionary.confirm(message, default=default).ask()
+
+    async def path(self, message: str, only_directories: bool = False) -> str | None:
+        return questionary.path(message, only_directories=only_directories).ask()

sqlsaber/cli/auth.py

@@ -10,7 +10,6 @@ from rich.console import Console
 from sqlsaber.config import providers
 from sqlsaber.config.api_keys import APIKeyManager
 from sqlsaber.config.auth import AuthConfigManager, AuthMethod
-from sqlsaber.config.oauth_flow import AnthropicOAuthFlow
 from sqlsaber.config.oauth_tokens import OAuthTokenManager
 
 # Global instances for CLI commands
@@ -27,60 +26,33 @@ auth_app = cyclopts.App(
 @auth_app.command
 def setup():
     """Configure authentication for SQLsaber (API keys and Anthropic OAuth)."""
-    console.print("\n[bold]SQLsaber Authentication Setup[/bold]\n")
+    import asyncio
 
-    provider = questionary.select(
-        "Select provider to configure:",
-        choices=providers.all_keys(),
-    ).ask()
+    from sqlsaber.application.auth_setup import setup_auth
+    from sqlsaber.application.prompts import AsyncPrompter
 
-    if provider is None:
-        console.print("[yellow]Setup cancelled.[/yellow]")
-        return
+    console.print("\n[bold]SQLsaber Authentication Setup[/bold]\n")
 
-    if provider == "anthropic":
-        # Let user choose API key or OAuth
-        method_choice = questionary.select(
-            "Select Anthropic authentication method:",
-            choices=[
-                {"name": "API key", "value": AuthMethod.API_KEY},
-                {"name": "Claude Pro/Max (OAuth)", "value": AuthMethod.CLAUDE_PRO},
-            ],
-        ).ask()
-
-        if method_choice == AuthMethod.CLAUDE_PRO:
-            flow = AnthropicOAuthFlow()
-            if flow.authenticate():
-                config_manager.set_auth_method(AuthMethod.CLAUDE_PRO)
-                console.print(
-                    "\n[bold green]✓ Anthropic OAuth configured successfully![/bold green]"
-                )
-            else:
-                console.print("\n[red]✗ Anthropic OAuth setup failed.[/red]")
-            console.print(
-                "You can change this anytime by running [cyan]saber auth setup[/cyan] again."
-            )
-            return
-
-    # API key flow (all providers + Anthropic when selected above)
-    api_key_manager = APIKeyManager()
-    env_var = api_key_manager._get_env_var_name(provider)
-    console.print("\nTo configure your API key, you can either:")
-    console.print(f"• Set the {env_var} environment variable")
-    console.print("• Let SQLsaber prompt you for the key when needed (stored securely)")
-
-    # Fetch/store key (cascades env -> keyring -> prompt)
-    api_key = api_key_manager.get_api_key(provider)
-    if api_key:
-        config_manager.set_auth_method(AuthMethod.API_KEY)
-        console.print(
-            f"\n[bold green]✓ {provider.title()} API key configured successfully![/bold green]"
+    async def run_setup():
+        prompter = AsyncPrompter()
+        api_key_manager = APIKeyManager()
+        success, provider = await setup_auth(
+            prompter=prompter,
+            auth_manager=config_manager,
+            api_key_manager=api_key_manager,
+            allow_oauth=True,
+            default_provider="anthropic",
+            run_oauth_in_thread=False,
         )
-    else:
-        console.print("\n[yellow]No API key configured.[/yellow]")
+        return success, provider
+
+    success, _ = asyncio.run(run_setup())
+
+    if not success:
+        console.print("\n[yellow]No authentication configured.[/yellow]")
 
     console.print(
-        "You can change this anytime by running [cyan]saber auth setup[/cyan] again."
+        "\nYou can change this anytime by running [cyan]saber auth setup[/cyan] again."
     )
 
 
@@ -109,7 +81,7 @@ def status():
             # Include OAuth status
             if OAuthTokenManager().has_oauth_token("anthropic"):
                 console.print("> anthropic (oauth): [green]configured[/green]")
-        env_var = api_key_manager._get_env_var_name(provider)
+        env_var = api_key_manager.get_env_var_name(provider)
         service = api_key_manager._get_service_name(provider)
         from_env = bool(os.getenv(env_var))
         from_keyring = bool(keyring.get_password(service, provider))

sqlsaber/cli/commands.py

@@ -11,6 +11,7 @@ from sqlsaber.cli.auth import create_auth_app
 from sqlsaber.cli.database import create_db_app
 from sqlsaber.cli.memory import create_memory_app
 from sqlsaber.cli.models import create_models_app
+from sqlsaber.cli.onboarding import needs_onboarding, run_onboarding
 from sqlsaber.cli.threads import create_threads_app
 
 # Lazy imports - only import what's needed for CLI parsing
@@ -75,7 +76,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 +86,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 +111,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
@@ -132,6 +129,16 @@ def query(
                 # If stdin was empty, fall back to interactive mode
                 actual_query = None
 
+        # Check if onboarding is needed (only for interactive mode or when no database is configured)
+        if needs_onboarding(database):
+            # Run onboarding flow
+            onboarding_success = await run_onboarding()
+            if not onboarding_success:
+                # User cancelled or onboarding failed
+                raise CLIError(
+                    "Setup incomplete. Please configure your database and try again."
+                )
+
         # Resolve database from CLI input
         try:
             resolved = resolve_database(database, config_manager)
@@ -147,45 +154,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 +192,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