sqlsaber 0.29.0__py3-none-any.whl → 0.30.0__py3-none-any.whl

sqlsaber/cli/onboarding.py

@@ -32,16 +32,16 @@ def needs_onboarding(database_arg: str | None = None) -> bool:
 
 def welcome_screen() -> None:
     """Display welcome screen to new users."""
-    banner = """
+    banner = """[primary]
 ███████  ██████  ██      ███████  █████  ██████  ███████ ██████
 ██      ██    ██ ██      ██      ██   ██ ██   ██ ██      ██   ██
 ███████ ██    ██ ██      ███████ ███████ ██████  █████   ██████
      ██ ██ ▄▄ ██ ██           ██ ██   ██ ██   ██ ██      ██   ██
 ███████  ██████  ███████ ███████ ██   ██ ██████  ███████ ██   ██
            ▀▀
-    """
+    [/primary]"""
 
-    console.print(Panel.fit(banner, style="bold blue"))
+    console.print(Panel.fit(banner, style="primary"))
     console.print()
 
     welcome_message = """
@@ -52,7 +52,9 @@ SQLsaber is an agentic SQL assistant that lets you query your database using nat
 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(
+        Panel.fit(welcome_message.strip(), border_style="primary", padding=(1, 2))
+    )
     console.print()
 
 
@@ -69,9 +71,7 @@ async def setup_database_guided() -> str | None:
     )
     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("[heading]Step 1 of 2: Database Connection[/heading]")
     console.print()
 
     try:
@@ -92,7 +92,7 @@ async def setup_database_guided() -> str | None:
         db_manager = DatabaseConfigManager()
         if db_manager.get_database(name):
             console.print(
-                f"[yellow]Database connection '{name}' already exists.[/yellow]"
+                f"[warning]Database connection '{name}' already exists.[/warning]"
             )
             return name
 
@@ -108,7 +108,7 @@ async def setup_database_guided() -> str | None:
         db_config = build_config(db_input)
 
         # Test the connection
-        console.print(f"[dim]Testing connection to '{name}'...[/dim]")
+        console.print(f"[muted]Testing connection to '{name}'...[/muted]")
         connection_success = await test_connection(db_config, db_input.password)
 
         if not connection_success:
@@ -119,25 +119,25 @@ async def setup_database_guided() -> str | None:
                 return await setup_database_guided()
             else:
                 console.print(
-                    "[yellow]You can add a database later using 'saber db add'[/yellow]"
+                    "[warning]You can add a database later using 'saber db add'[/warning]"
                 )
                 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(f"[success]✓ Connection to '{name}' successful![/success]")
             console.print()
             return name
         except Exception as e:
-            console.print(f"[bold red]Error saving database:[/bold red] {e}")
+            console.print(f"[error]Error saving database:[/error] {e}")
             return None
 
     except KeyboardInterrupt:
-        console.print("\n[yellow]Setup cancelled.[/yellow]")
+        console.print("\n[warning]Setup cancelled.[/warning]")
         return None
     except Exception as e:
-        console.print(f"[bold red]Unexpected error:[/bold red] {e}")
+        console.print(f"[error]Unexpected error:[/error] {e}")
         return None
 
 
@@ -151,14 +151,14 @@ async def select_model_for_provider(provider: str) -> str | None:
 
     try:
         console.print()
-        console.print(f"[dim]Fetching available {provider.title()} models...[/dim]")
+        console.print(f"[muted]Fetching available {provider.title()} models...[/muted]")
 
         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]"
+                f"[warning]Could not fetch models for {provider}. Using default.[/warning]"
             )
             # Use provider-specific default or fallback to Anthropic
             default_model_id = ModelManager.RECOMMENDED_MODELS.get(
@@ -178,10 +178,10 @@ async def select_model_for_provider(provider: str) -> str | None:
         return selected_model
 
     except KeyboardInterrupt:
-        console.print("\n[yellow]Model selection cancelled.[/yellow]")
+        console.print("\n[warning]Model selection cancelled.[/warning]")
         return None
     except Exception as e:
-        console.print(f"[yellow]Error selecting model: {e}. Using default.[/yellow]")
+        console.print(f"[warning]Error selecting model: {e}. Using default.[/warning]")
         # Fallback to provider default
         if provider in ModelManager.RECOMMENDED_MODELS:
             return f"{provider}:{ModelManager.RECOMMENDED_MODELS[provider]}"
@@ -196,9 +196,7 @@ async def setup_auth_guided() -> tuple[bool, 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("[primary]Step 2 of 2: Authentication[/primary]")
     console.print()
 
     try:
@@ -218,7 +216,7 @@ async def setup_auth_guided() -> tuple[bool, str | None]:
 
         if not success:
             console.print(
-                "[yellow]You can set it up later using 'saber auth setup'[/yellow]"
+                "[warning]You can set it up later using 'saber auth setup'[/warning]"
             )
             console.print()
             return False, None
@@ -233,16 +231,16 @@ async def setup_auth_guided() -> tuple[bool, str | None]:
         if selected_model:
             model_manager = ModelManager()
             model_manager.set_model(selected_model)
-            console.print(f"[green]✓ Model set to: {selected_model}[/green]")
+            console.print(f"[success]✓ Model set to: {selected_model}[/success]")
         console.print()
         return True, selected_model
 
     except KeyboardInterrupt:
-        console.print("\n[yellow]Setup cancelled.[/yellow]")
+        console.print("\n[warning]Setup cancelled.[/warning]")
         console.print()
         return False, None
     except Exception as e:
-        console.print(f"[bold red]Unexpected error:[/bold red] {e}")
+        console.print(f"[error]Unexpected error:[/error] {e}")
         console.print()
         return False, None
 
@@ -251,35 +249,34 @@ 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("[success]You're all set! 🚀[/success]")
     console.print()
 
     if database_name and auth_configured:
         console.print(
-            f"[green]✓ Database '{database_name}' connected and ready to use[/green]"
+            f"[success]✓ Database '{database_name}' connected and ready to use[/success]"
         )
-        console.print("[green]✓ Authentication configured[/green]")
+        console.print("[success]✓ Authentication configured[/success]")
         if model_name:
-            console.print(f"[green]✓ Model: {model_name}[/green]")
+            console.print(f"[success]✓ Model: {model_name}[/success]")
     elif database_name:
         console.print(
-            f"[green]✓ Database '{database_name}' connected and ready to use[/green]"
+            f"[success]✓ Database '{database_name}' connected and ready to use[/success]"
         )
         console.print(
-            "[yellow]⚠ AI authentication not configured - you'll be prompted when needed[/yellow]"
+            "[warning]⚠ AI authentication not configured - you'll be prompted when needed[/warning]"
         )
     elif auth_configured:
-        console.print("[green]✓ AI authentication configured[/green]")
+        console.print("[success]✓ AI authentication configured[/success]")
         if model_name:
-            console.print(f"[green]✓ Model: {model_name}[/green]")
+            console.print(f"[success]✓ Model: {model_name}[/success]")
         console.print(
-            "[yellow]⚠ No database configured - you'll need to provide one via -d flag[/yellow]"
+            "[warning]⚠ No database configured - you'll need to provide one via -d flag[/warning]"
         )
 
     console.print()
-    console.print("[dim]Starting interactive session...[/dim]")
+    console.print("[muted]Starting interactive session...[/muted]")
     console.print()
 
 
@@ -298,9 +295,9 @@ async def run_onboarding() -> bool:
 
         # If user cancelled database setup, exit
         if database_name is None:
-            console.print("[yellow]Database setup is required to continue.[/yellow]")
+            console.print("[warning]Database setup is required to continue.[/warning]")
             console.print(
-                "[dim]You can also provide a connection string using: saber -d <connection-string>[/dim]"
+                "[muted]You can also provide a connection string using: saber -d <connection-string>[/muted]"
             )
             return False
 
@@ -313,13 +310,13 @@ async def run_onboarding() -> bool:
         return True
 
     except KeyboardInterrupt:
-        console.print("\n[yellow]Onboarding cancelled.[/yellow]")
+        console.print("\n[warning]Onboarding cancelled.[/warning]")
         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]"
+            "[muted]You can run setup commands manually:[/muted]\n"
+            "[muted]  - saber db add <name>  # Add database connection[/muted]\n"
+            "[muted]  - saber auth setup     # Configure authentication[/muted]"
         )
         sys.exit(0)
     except Exception as e:
-        console.print(f"[bold red]Onboarding failed:[/bold red] {e}")
+        console.print(f"[error]Onboarding failed:[/error] {e}")
         return False

sqlsaber/cli/streaming.py

@@ -144,17 +144,8 @@ class StreamingQueryHandler:
             prepared_prompt: str | list[str] = user_query
             no_history = not message_history
             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)
+                injected = sqlsaber_agent.system_prompt_text(include_memory=True)
+                if injected and str(injected).strip():
                     prepared_prompt = [injected, user_query]
 
             # Show a transient status until events start streaming

sqlsaber/cli/threads.py

@@ -170,7 +170,7 @@ def _render_transcript(
                                     Panel.fit(
                                         content_str,
                                         title=f"Tool result: {name}",
-                                        border_style="yellow",
+                                        border_style="warning",
                                     )
                                 )
                     except Exception:
@@ -183,7 +183,7 @@ def _render_transcript(
                                 Panel.fit(
                                     content_str,
                                     title=f"Tool result: {name}",
-                                    border_style="yellow",
+                                    border_style="warning",
                                 )
                             )
                 else:
@@ -194,7 +194,7 @@ def _render_transcript(
                             Panel.fit(
                                 content_str,
                                 title=f"Tool result: {name}",
-                                border_style="yellow",
+                                border_style="warning",
                             )
                         )
         # Thinking parts omitted

sqlsaber/config/api_keys.py

@@ -36,7 +36,7 @@ class APIKeyManager:
                 return api_key
         except Exception as e:
             # Keyring access failed, continue to prompt
-            console.print(f"Keyring access failed: {e}", style="dim yellow")
+            console.print(f"Keyring access failed: {e}", style="muted warning")
 
         # 3. Prompt user for API key
         return self._prompt_and_store_key(provider, env_var_name, service_name)
@@ -72,7 +72,7 @@ class APIKeyManager:
             if not api_key.strip():
                 console.print(
                     "No API key provided. Some functionality may not work.",
-                    style="yellow",
+                    style="warning",
                 )
                 return None
 
@@ -83,16 +83,16 @@ class APIKeyManager:
             except Exception as e:
                 console.print(
                     f"Warning: Could not store API key in your operating system's credentials store: {e}",
-                    style="yellow",
+                    style="warning",
                 )
                 console.print(
-                    "You may need to enter it again next time", style="yellow"
+                    "You may need to enter it again next time", style="warning"
                 )
 
             return api_key.strip()
 
         except KeyboardInterrupt:
-            console.print("\nOperation cancelled", style="yellow")
+            console.print("\nOperation cancelled", style="warning")
             return None
         except Exception as e:
             console.print(f"Error prompting for API key: {e}", style="red")

sqlsaber/config/oauth_flow.py

@@ -131,7 +131,7 @@ class AnthropicOAuthFlow:
         if not questionary.confirm(
             "Continue with browser-based authentication?", default=True
         ).ask():
-            console.print("[yellow]Authentication cancelled.[/yellow]")
+            console.print("[warning]Authentication cancelled.[/warning]")
             return False
 
         try:
@@ -168,7 +168,7 @@ class AnthropicOAuthFlow:
             ).ask()
 
             if not auth_code:
-                console.print("[yellow]Authentication cancelled.[/yellow]")
+                console.print("[warning]Authentication cancelled.[/warning]")
                 return False
 
             # Step 2: Exchange code for tokens
@@ -198,23 +198,23 @@ class AnthropicOAuthFlow:
                 )
 
                 if self.token_manager.store_oauth_token("anthropic", oauth_token):
-                    console.print(
-                        "\n[bold green]✓ Authentication successful![/bold green]"
-                    )
+                    console.print("\n[success]✓ Authentication successful![/success]")
                     console.print(
                         "Your Claude Pro/Max subscription is now configured for SQLSaber."
                     )
                     return True
                 else:
-                    console.print("[red]✗ Failed to store authentication tokens.[/red]")
+                    console.print(
+                        "[error]✗ Failed to store authentication tokens.[/error]"
+                    )
                     return False
 
         except KeyboardInterrupt:
-            console.print("\n[yellow]Authentication cancelled by user.[/yellow]")
+            console.print("\n[warning]Authentication cancelled by user.[/warning]")
             return False
         except Exception as e:
             logger.error(f"OAuth authentication failed: {e}")
-            console.print(f"[red]✗ Authentication failed: {str(e)}[/red]")
+            console.print(f"[error]✗ Authentication failed: {str(e)}[/error]")
             return False
 
     def refresh_token_if_needed(self) -> OAuthToken | None:
@@ -255,13 +255,14 @@ class AnthropicOAuthFlow:
                 console.print("OAuth token refreshed successfully", style="green")
                 return refreshed_token
             else:
-                console.print("Failed to store refreshed token", style="yellow")
+                console.print("Failed to store refreshed token", style="warning")
                 return current_token
 
         except Exception as e:
             logger.warning(f"Token refresh failed: {e}")
             console.print(
-                "Token refresh failed. You may need to re-authenticate.", style="yellow"
+                "Token refresh failed. You may need to re-authenticate.",
+                style="warning",
             )
             return current_token
 

sqlsaber/config/oauth_tokens.py

@@ -97,14 +97,14 @@ class OAuthTokenManager:
             if token.is_expired():
                 console.print(
                     f"OAuth token for {provider} has expired and needs refresh",
-                    style="dim yellow",
+                    style="muted warning",
                 )
                 return token  # Return anyway for refresh attempt
 
             if token.expires_soon():
                 console.print(
                     f"OAuth token for {provider} expires soon, consider refreshing",
-                    style="dim yellow",
+                    style="muted warning",
                 )
 
             return token
@@ -126,7 +126,7 @@ class OAuthTokenManager:
             logger.error(f"Failed to store OAuth token for {provider}: {e}")
             console.print(
                 f"Warning: Could not store OAuth token in keyring: {e}",
-                style="yellow",
+                style="warning",
             )
             return False
 
@@ -137,7 +137,7 @@ class OAuthTokenManager:
         existing_token = self.get_oauth_token(provider)
         if not existing_token:
             console.print(
-                f"No existing OAuth token found for {provider}", style="yellow"
+                f"No existing OAuth token found for {provider}", style="warning"
             )
             return False
 
@@ -161,7 +161,9 @@ class OAuthTokenManager:
             return True
         except Exception as e:
             logger.error(f"Failed to remove OAuth token for {provider}: {e}")
-            console.print(f"Warning: Could not remove OAuth token: {e}", style="yellow")
+            console.print(
+                f"Warning: Could not remove OAuth token: {e}", style="warning"
+            )
             return False
 
     def has_oauth_token(self, provider: str) -> bool:

sqlsaber/database/schema.py

@@ -158,7 +158,7 @@ class SchemaManager:
             table["schema"] = table["table_schema"]
             table["type"] = table["table_type"]  # Map table_type to type for display
 
-        return {"tables": tables_list}
+        return {"tables": tables_list, "total_tables": len(tables_list)}
 
     async def close(self):
         """Close database connection."""

sqlsaber/theme/manager.py

@@ -9,7 +9,7 @@ from typing import Dict
 from platformdirs import user_config_dir
 from prompt_toolkit.styles import Style as PTStyle
 from prompt_toolkit.styles.pygments import style_from_pygments_cls
-from pygments.styles import get_style_by_name
+from pygments.styles import get_all_styles, get_style_by_name
 from pygments.token import Token
 from pygments.util import ClassNotFound
 from rich.console import Console
@@ -18,14 +18,6 @@ from rich.theme import Theme
 DEFAULT_THEME_NAME = "nord"
 
 DEFAULT_ROLE_PALETTE = {
-    # base roles
-    "primary": "cyan",
-    "accent": "magenta",
-    "success": "green",
-    "warning": "yellow",
-    "error": "red",
-    "info": "cyan",
-    "muted": "dim",
     # components
     "table.header": "bold $primary",
     "panel.border.user": "$info",
@@ -208,6 +200,9 @@ def get_theme_manager() -> ThemeManager:
     user_cfg = _load_user_theme_config()
     env_name = os.getenv("SQLSABER_THEME")
 
+    if env_name and env_name.lower() not in get_all_styles():
+        env_name = None
+
     name = (
         env_name or user_cfg.get("theme", {}).get("name") or DEFAULT_THEME_NAME
     ).lower()

sqlsaber/tools/__init__.py

@@ -1,8 +1,6 @@
 """SQLSaber tools module."""
 
 from .base import Tool
-from .enums import ToolCategory, WorkflowPosition
-from .instructions import InstructionBuilder
 from .registry import ToolRegistry, register_tool, tool_registry
 
 # Import concrete tools to register them
@@ -10,12 +8,9 @@ from .sql_tools import ExecuteSQLTool, IntrospectSchemaTool, ListTablesTool, SQL
 
 __all__ = [
     "Tool",
-    "ToolCategory",
-    "WorkflowPosition",
     "ToolRegistry",
     "tool_registry",
     "register_tool",
-    "InstructionBuilder",
     "SQLTool",
     "ListTablesTool",
     "IntrospectSchemaTool",

sqlsaber/tools/base.py

@@ -3,8 +3,6 @@
 from abc import ABC, abstractmethod
 from typing import Any
 
-from .enums import ToolCategory, WorkflowPosition
-
 
 class Tool(ABC):
     """Abstract base class for all tools."""
@@ -42,32 +40,3 @@ class Tool(ABC):
             JSON string with the tool's output
         """
         pass
-
-    @property
-    def category(self) -> ToolCategory:
-        """Return the tool category. Override to customize."""
-        return ToolCategory.GENERAL
-
-    def get_usage_instructions(self) -> str | None:
-        """Return tool-specific usage instructions for LLM guidance.
-
-        Returns:
-            Usage instructions string, or None for no specific guidance
-        """
-        return None
-
-    def get_priority(self) -> int:
-        """Return priority for tool ordering in instructions.
-
-        Returns:
-            Priority number (lower = higher priority, default = 100)
-        """
-        return 100
-
-    def get_workflow_position(self) -> WorkflowPosition:
-        """Return the typical workflow position for this tool.
-
-        Returns:
-            WorkflowPosition enum value
-        """
-        return WorkflowPosition.OTHER

sqlsaber/tools/registry.py

@@ -3,7 +3,6 @@
 from typing import Type
 
 from .base import Tool
-from .enums import ToolCategory
 
 
 class ToolRegistry:
@@ -61,45 +60,13 @@ class ToolRegistry:
 
         return self._instances[name]
 
-    def list_tools(self, category: str | ToolCategory | None = None) -> list[str]:
-        """List all registered tool names.
+    def list_tools(self) -> list[str]:
+        """List all registered tool names."""
+        return list(self._tools.keys())
 
-        Args:
-            category: Optional category to filter by (string or ToolCategory enum)
-
-        Returns:
-            List of tool names
-        """
-        if category is None:
-            return list(self._tools.keys())
-
-        # Convert string to enum
-        if isinstance(category, str):
-            try:
-                category = ToolCategory(category)
-            except ValueError:
-                # If string doesn't match any enum, return empty list
-                return []
-
-        # Filter by category
-        result = []
-        for name, tool_class in self._tools.items():
-            tool = self.get_tool(name)
-            if tool.category == category:
-                result.append(name)
-        return result
-
-    def get_all_tools(self, category: str | ToolCategory | None = None) -> list[Tool]:
-        """Get all tool instances.
-
-        Args:
-            category: Optional category to filter by (string or ToolCategory enum)
-
-        Returns:
-            List of tool instances
-        """
-        names = self.list_tools(category)
-        return [self.get_tool(name) for name in names]
+    def get_all_tools(self) -> list[Tool]:
+        """Get all tool instances."""
+        return [self.get_tool(name) for name in self.list_tools()]
 
 
 # Global registry instance