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

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

sqlsaber/cli/display.py

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

sqlsaber/cli/interactive.py

@@ -3,13 +3,13 @@
 import asyncio
 from pathlib import Path
 from textwrap import dedent
+from typing import TYPE_CHECKING
 
 import platformdirs
 from prompt_toolkit import PromptSession
 from prompt_toolkit.history import FileHistory
 from prompt_toolkit.patch_stdout import patch_stdout
 from prompt_toolkit.styles import Style
-from pydantic_ai import Agent
 from rich.console import Console
 from rich.markdown import Markdown
 from rich.panel import Panel
@@ -21,7 +21,7 @@ from sqlsaber.cli.completers import (
 )
 from sqlsaber.cli.display import DisplayManager
 from sqlsaber.cli.streaming import StreamingQueryHandler
-from sqlsaber.database.connection import (
+from sqlsaber.database import (
     CSVConnection,
     DuckDBConnection,
     MySQLConnection,
@@ -31,31 +31,20 @@ from sqlsaber.database.connection import (
 from sqlsaber.database.schema import SchemaManager
 from sqlsaber.threads import ThreadStorage
 
-
-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",
-    }
-)
+if TYPE_CHECKING:
+    from sqlsaber.agents.pydantic_ai_agent import SQLSaberAgent
 
 
 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,
-        agent: Agent,
+        sqlsaber_agent: "SQLSaberAgent",
         db_conn,
         database_name: str,
         *,
@@ -63,7 +52,7 @@ class InteractiveSession:
         initial_history: list | None = None,
     ):
         self.console = console
-        self.agent = agent
+        self.sqlsaber_agent = sqlsaber_agent
         self.db_conn = db_conn
         self.database_name = database_name
         self.display = DisplayManager(console)
@@ -77,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 """
 ███████  ██████  ██      ███████  █████  ██████  ███████ ██████
 ██      ██    ██ ██      ██      ██   ██ ██   ██ ██      ██   ██
 ███████ ██    ██ ██      ███████ ███████ ██████  █████   ██████
@@ -106,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:
@@ -167,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
@@ -176,7 +249,7 @@ class InteractiveSession:
         query_task = asyncio.create_task(
             self.streaming_handler.execute_streaming_query(
                 user_query,
-                self.agent,
+                self.sqlsaber_agent,
                 self.cancellation_token,
                 self.message_history,
             )
@@ -191,11 +264,6 @@ class InteractiveSession:
                     # Use all_messages() so the system prompt and all prior turns are preserved
                     self.message_history = run_result.all_messages()
 
-                    # Extract title (first user prompt) and model name
-                    if not self._thread_id:
-                        title = user_query
-                        model_name = self.agent.model.model_name
-
                     # Persist snapshot to thread storage (create or overwrite)
                     self._thread_id = await self._threads.save_snapshot(
                         messages_json=run_result.all_messages_json(),
@@ -206,9 +274,10 @@ class InteractiveSession:
                     if self.first_message:
                         await self._threads.save_metadata(
                             thread_id=self._thread_id,
-                            title=title,
-                            model_name=model_name,
+                            title=user_query,
+                            model_name=self.sqlsaber_agent.agent.model.model_name,
                         )
+                        self.first_message = False
                 except Exception:
                     pass
                 finally:
@@ -220,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:
@@ -240,72 +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
+                # Handle slash commands
+                if await self._handle_command(user_query):
                     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 = getattr(
-                                    self.agent, "_sqlsaber_memory_manager", None
-                                )
-                                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
+                # Handle memory addition
+                if user_query.strip().startswith("#"):
+                    await self._handle_memory(user_query.strip()[1:].strip())
+                    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
@@ -324,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}")