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

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
@@ -128,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)

sqlsaber/cli/database.py

@@ -81,95 +81,34 @@ def add(
 
     if interactive:
         # Interactive mode - prompt for all required fields
-        console.print(f"[bold]Adding database connection: {name}[/bold]")
+        from sqlsaber.application.db_setup import collect_db_input
+        from sqlsaber.application.prompts import AsyncPrompter
 
-        # Database type
-        if not type or type == "postgresql":
-            type = questionary.select(
-                "Database type:",
-                choices=["postgresql", "mysql", "sqlite", "duckdb"],
-                default="postgresql",
-            ).ask()
-
-        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 = type
-            password = ""
-        else:
-            # PostgreSQL/MySQL need connection details
-            host = host or questionary.text("Host:", default="localhost").ask()
+        console.print(f"[bold]Adding database connection: {name}[/bold]")
 
-            default_port = 5432 if type == "postgresql" else 3306
-            port = port or int(
-                questionary.text("Port:", default=str(default_port)).ask()
+        async def collect_input():
+            prompter = AsyncPrompter()
+            return await collect_db_input(
+                prompter=prompter, name=name, db_type=type, include_ssl=True
             )
 
-            database = database or questionary.text("Database name:").ask()
-            username = username or questionary.text("Username:").ask()
-
-            # Ask for password
-            password = getpass.getpass("Password (stored in your OS keychain): ")
-
-            # Ask for SSL configuration
-            if questionary.confirm("Configure SSL/TLS settings?", default=False).ask():
-                if type == "postgresql":
-                    ssl_mode = (
-                        ssl_mode
-                        or questionary.select(
-                            "SSL mode for PostgreSQL:",
-                            choices=[
-                                "disable",
-                                "allow",
-                                "prefer",
-                                "require",
-                                "verify-ca",
-                                "verify-full",
-                            ],
-                            default="prefer",
-                        ).ask()
-                    )
-                elif type == "mysql":
-                    ssl_mode = (
-                        ssl_mode
-                        or questionary.select(
-                            "SSL mode for MySQL:",
-                            choices=[
-                                "DISABLED",
-                                "PREFERRED",
-                                "REQUIRED",
-                                "VERIFY_CA",
-                                "VERIFY_IDENTITY",
-                            ],
-                            default="PREFERRED",
-                        ).ask()
-                    )
-
-                if ssl_mode and ssl_mode not in ["disable", "DISABLED"]:
-                    if questionary.confirm(
-                        "Specify SSL certificate files?", default=False
-                    ).ask():
-                        ssl_ca = (
-                            ssl_ca or questionary.path("SSL CA certificate file:").ask()
-                        )
-                        if questionary.confirm(
-                            "Specify client certificate?", default=False
-                        ).ask():
-                            ssl_cert = (
-                                ssl_cert
-                                or questionary.path(
-                                    "SSL client certificate file:"
-                                ).ask()
-                            )
-                            ssl_key = (
-                                ssl_key
-                                or questionary.path(
-                                    "SSL client private key file:"
-                                ).ask()
-                            )
+        db_input = asyncio.run(collect_input())
+
+        if db_input is None:
+            console.print("[yellow]Operation cancelled[/yellow]")
+            return
+
+        # Extract values from db_input
+        type = db_input.type
+        host = db_input.host
+        port = db_input.port
+        database = db_input.database
+        username = db_input.username
+        password = db_input.password
+        ssl_mode = db_input.ssl_mode
+        ssl_ca = db_input.ssl_ca
+        ssl_cert = db_input.ssl_cert
+        ssl_key = db_input.ssl_key
     else:
         # Non-interactive mode - use provided values or defaults
         if type == "sqlite":

sqlsaber/cli/display.py

@@ -109,7 +109,7 @@ class LiveMarkdownRenderer:
         # Print the complete markdown to scroll-back for permanent reference
         if buf:
             if kind == ThinkingPart:
-                self.console.print(Text(buf, style="dim"))
+                self.console.print(Markdown(buf, style="dim"))
             else:
                 self.console.print(Markdown(buf))
 

sqlsaber/cli/interactive.py

@@ -35,26 +35,12 @@ if TYPE_CHECKING:
     from sqlsaber.agents.pydantic_ai_agent import SQLSaberAgent
 
 
-def bottom_toolbar():
-    return [
-        (
-            "class:bottom-toolbar",
-            " Use 'Esc-Enter' or 'Meta-Enter' to submit.",
-        )
-    ]
-
-
-style = Style.from_dict(
-    {
-        "frame.border": "#ebbcba",
-        "bottom-toolbar": "#ebbcba bg:#21202e",
-    }
-)
-
-
 class InteractiveSession:
     """Manages interactive CLI sessions."""
 
+    exit_commands = {"/exit", "/quit", "exit", "quit"}
+    resume_command_template = "saber threads resume {thread_id}"
+
     def __init__(
         self,
         console: Console,
@@ -80,28 +66,33 @@ class InteractiveSession:
         self._thread_id: str | None = initial_thread_id
         self.first_message = not self._thread_id
 
-    def show_welcome_message(self):
-        """Display welcome message for interactive mode."""
-        # Show database information
-        db_name = self.database_name or "Unknown"
-        db_type = (
-            "PostgreSQL"
-            if isinstance(self.db_conn, PostgreSQLConnection)
-            else "MySQL"
-            if isinstance(self.db_conn, MySQLConnection)
-            else "DuckDB"
-            if isinstance(self.db_conn, DuckDBConnection)
-            else "DuckDB"
-            if isinstance(self.db_conn, CSVConnection)
-            else "SQLite"
-            if isinstance(self.db_conn, SQLiteConnection)
-            else "database"
+    def _history_path(self) -> Path:
+        """Get the history file path, ensuring directory exists."""
+        history_dir = Path(platformdirs.user_config_dir("sqlsaber"))
+        history_dir.mkdir(parents=True, exist_ok=True)
+        return history_dir / "history"
+
+    def _prompt_style(self) -> Style:
+        """Get the prompt style configuration."""
+        return Style.from_dict(
+            {
+                "frame.border": "gray",
+                "bottom-toolbar": "white bg:#21202e",
+            }
         )
 
-        if self.first_message:
-            self.console.print(
-                Panel.fit(
-                    """
+    def _bottom_toolbar(self):
+        """Get the bottom toolbar text."""
+        return [
+            (
+                "class:bottom-toolbar",
+                " Use 'Esc-Enter' or 'Meta-Enter' to submit.",
+            )
+        ]
+
+    def _banner(self) -> str:
+        """Get the ASCII banner."""
+        return """
 ███████  ██████  ██      ███████  █████  ██████  ███████ ██████
 ██      ██    ██ ██      ██      ██   ██ ██   ██ ██      ██   ██
 ███████ ██    ██ ██      ███████ ███████ ██████  █████   ██████
@@ -109,35 +100,110 @@ class InteractiveSession:
 ███████  ██████  ███████ ███████ ██   ██ ██████  ███████ ██   ██
             ▀▀
     """
-                )
-            )
-            self.console.print(
-                Markdown(
-                    dedent("""
+
+    def _instructions(self) -> str:
+        """Get the instruction text."""
+        return dedent("""
                     - Use `/` for slash commands
                     - Type `@` to get table name completions
                     - Start message with `#` to add something to agent's memory
                     - Use `Ctrl+C` to interrupt and `Ctrl+D` to exit
                     """)
-                )
-            )
 
+    def _db_type_name(self) -> str:
+        """Get human-readable database type name."""
+        mapping = {
+            PostgreSQLConnection: "PostgreSQL",
+            MySQLConnection: "MySQL",
+            DuckDBConnection: "DuckDB",
+            CSVConnection: "DuckDB",
+            SQLiteConnection: "SQLite",
+        }
+        for cls, name in mapping.items():
+            if isinstance(self.db_conn, cls):
+                return name
+        return "database"
+
+    def _resume_hint(self, thread_id: str) -> str:
+        """Build resume command hint."""
+        return self.resume_command_template.format(thread_id=thread_id)
+
+    def show_welcome_message(self):
+        """Display welcome message for interactive mode."""
+        if self.first_message:
+            self.console.print(Panel.fit(self._banner()))
+            self.console.print(Markdown(self._instructions()))
+
+        db_name = self.database_name or "Unknown"
         self.console.print(
-            f"[bold blue]\n\nConnected to:[/bold blue] {db_name} ({db_type})\n"
+            f"[bold blue]\n\nConnected to:[/bold blue] {db_name} ({self._db_type_name()})\n"
         )
-        # If resuming a thread, show a notice
+
         if self._thread_id:
             self.console.print(f"[dim]Resuming thread:[/dim] {self._thread_id}\n")
 
-    async def _end_thread_and_display_resume_hint(self):
-        """End thread and display command to resume thread"""
-        # Print resume hint if there is an active thread
+    async def _end_thread(self):
+        """End thread and display resume hint."""
         if self._thread_id:
             await self._threads.end_thread(self._thread_id)
             self.console.print(
-                f"[dim]You can continue this thread using:[/dim] saber threads resume {self._thread_id}"
+                f"[dim]You can continue this thread using:[/dim] {self._resume_hint(self._thread_id)}"
             )
 
+    async def _handle_memory(self, content: str):
+        """Handle memory addition command."""
+        if not content:
+            self.console.print("[yellow]Empty memory content after '#'[/yellow]\n")
+            return
+
+        try:
+            mm = self.sqlsaber_agent.memory_manager
+            if mm and self.database_name:
+                memory = mm.add_memory(self.database_name, content)
+                self.console.print(f"[green]✓ Memory added:[/green] {content}")
+                self.console.print(f"[dim]Memory ID: {memory.id}[/dim]\n")
+            else:
+                self.console.print(
+                    "[yellow]Could not add memory (no database context)[/yellow]\n"
+                )
+        except Exception as exc:
+            self.console.print(f"[yellow]Could not add memory:[/yellow] {exc}\n")
+
+    async def _cmd_clear(self):
+        """Clear conversation history."""
+        self.message_history = []
+        try:
+            if self._thread_id:
+                await self._threads.end_thread(self._thread_id)
+        except Exception:
+            pass
+        self.console.print("[green]Conversation history cleared.[/green]\n")
+        self._thread_id = None
+        self.first_message = True
+
+    async def _cmd_thinking_on(self):
+        """Enable thinking mode."""
+        self.sqlsaber_agent.set_thinking(enabled=True)
+        self.console.print("[green]✓ Thinking enabled[/green]\n")
+
+    async def _cmd_thinking_off(self):
+        """Disable thinking mode."""
+        self.sqlsaber_agent.set_thinking(enabled=False)
+        self.console.print("[green]✓ Thinking disabled[/green]\n")
+
+    async def _handle_command(self, user_query: str) -> bool:
+        """Handle slash commands. Returns True if command was handled."""
+        if user_query == "/clear":
+            await self._cmd_clear()
+            return True
+        if user_query == "/thinking on":
+            await self._cmd_thinking_on()
+            return True
+        if user_query == "/thinking off":
+            await self._cmd_thinking_off()
+            return True
+        return False
+
     async def _update_table_cache(self):
         """Update the table completer cache with fresh data."""
         try:
@@ -170,6 +236,10 @@ class InteractiveSession:
             # If there's an error, just use empty cache
             self.table_completer.update_cache([])
 
+    async def before_prompt_loop(self):
+        """Hook to refresh context before prompt loop."""
+        await self._update_table_cache()
+
     async def _execute_query_with_cancellation(self, user_query: str):
         """Execute a query with cancellation support."""
         # Create cancellation token
@@ -207,6 +277,7 @@ class InteractiveSession:
                             title=user_query,
                             model_name=self.sqlsaber_agent.agent.model.model_name,
                         )
+                        self.first_message = False
                 except Exception:
                     pass
                 finally:
@@ -218,15 +289,9 @@ class InteractiveSession:
     async def run(self):
         """Run the interactive session loop."""
         self.show_welcome_message()
+        await self.before_prompt_loop()
 
-        # Initialize table cache
-        await self._update_table_cache()
-
-        session = PromptSession(
-            history=FileHistory(
-                Path(platformdirs.user_config_dir("sqlsaber")) / "history"
-            )
-        )
+        session = PromptSession(history=FileHistory(self._history_path()))
 
         while True:
             try:
@@ -238,81 +303,32 @@ class InteractiveSession:
                             SlashCommandCompleter(), self.table_completer
                         ),
                         show_frame=True,
-                        bottom_toolbar=bottom_toolbar,
-                        style=style,
+                        bottom_toolbar=self._bottom_toolbar,
+                        style=self._prompt_style(),
                     )
 
                 if not user_query:
                     continue
 
-                if (
-                    user_query in ["/exit", "/quit", "exit", "quit"]
-                    or user_query.startswith("/exit")
-                    or user_query.startswith("/quit")
+                # Handle exit commands
+                if user_query in self.exit_commands or any(
+                    user_query.startswith(cmd) for cmd in self.exit_commands
                 ):
-                    await self._end_thread_and_display_resume_hint()
+                    await self._end_thread()
                     break
 
-                if user_query == "/clear":
-                    # Reset local history (pydantic-ai call will receive empty history on next run)
-                    self.message_history = []
-                    # End current thread (if any) so the next turn creates a fresh one
-                    try:
-                        if self._thread_id:
-                            await self._threads.end_thread(self._thread_id)
-                    except Exception:
-                        pass
-                    self.console.print("[green]Conversation history cleared.[/green]\n")
-                    # Do not print resume hint when clearing; a new thread will be created on next turn
-                    self._thread_id = None
-                    continue
-
-                # Thinking commands
-                if user_query == "/thinking on":
-                    self.sqlsaber_agent.set_thinking(enabled=True)
-                    self.console.print("[green]✓ Thinking enabled[/green]\n")
+                # Handle slash commands
+                if await self._handle_command(user_query):
                     continue
 
-                if user_query == "/thinking off":
-                    self.sqlsaber_agent.set_thinking(enabled=False)
-                    self.console.print("[green]✓ Thinking disabled[/green]\n")
+                # Handle memory addition
+                if user_query.strip().startswith("#"):
+                    await self._handle_memory(user_query.strip()[1:].strip())
                     continue
 
-                if memory_text := user_query.strip():
-                    # Check if query starts with # for memory addition
-                    if memory_text.startswith("#"):
-                        memory_content = memory_text[1:].strip()  # Remove # and trim
-                        if memory_content:
-                            # Add memory via the agent's memory manager
-                            try:
-                                mm = self.sqlsaber_agent.memory_manager
-                                if mm and self.database_name:
-                                    memory = mm.add_memory(
-                                        self.database_name, memory_content
-                                    )
-                                    self.console.print(
-                                        f"[green]✓ Memory added:[/green] {memory_content}"
-                                    )
-                                    self.console.print(
-                                        f"[dim]Memory ID: {memory.id}[/dim]\n"
-                                    )
-                                else:
-                                    self.console.print(
-                                        "[yellow]Could not add memory (no database context)[/yellow]\n"
-                                    )
-                            except Exception:
-                                self.console.print(
-                                    "[yellow]Could not add memory[/yellow]\n"
-                                )
-                        else:
-                            self.console.print(
-                                "[yellow]Empty memory content after '#'[/yellow]\n"
-                            )
-                        continue
-
-                    # Execute query with cancellation support
-                    await self._execute_query_with_cancellation(user_query)
-                    self.display.show_newline()  # Empty line for readability
+                # Execute query with cancellation support
+                await self._execute_query_with_cancellation(user_query)
+                self.display.show_newline()
 
             except KeyboardInterrupt:
                 # Handle Ctrl+C - cancel current task if running
@@ -331,7 +347,7 @@ class InteractiveSession:
                     )
             except EOFError:
                 # Exit when Ctrl+D is pressed
-                await self._end_thread_and_display_resume_hint()
+                await self._end_thread()
                 break
-            except Exception as e:
-                self.console.print(f"[bold red]Error:[/bold red] {str(e)}")
+            except Exception as exc:
+                self.console.print(f"[bold red]Error:[/bold red] {exc}")

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):