sqlsaber 0.7.0__py3-none-any.whl → 0.8.1__py3-none-any.whl

sqlsaber/agents/base.py

@@ -3,7 +3,7 @@
 import asyncio
 import json
 from abc import ABC, abstractmethod
-from typing import Any, AsyncIterator, Dict, List, Optional
+from typing import Any, AsyncIterator
 
 from uniplot import histogram, plot
 
@@ -24,7 +24,7 @@ class BaseSQLAgent(ABC):
     def __init__(self, db_connection: BaseDatabaseConnection):
         self.db = db_connection
         self.schema_manager = SchemaManager(db_connection)
-        self.conversation_history: List[Dict[str, Any]] = []
+        self.conversation_history: list[dict[str, Any]] = []
 
     @abstractmethod
     async def query_stream(
@@ -59,7 +59,7 @@ class BaseSQLAgent(ABC):
         else:
             return "database"  # Fallback
 
-    async def introspect_schema(self, table_pattern: Optional[str] = None) -> str:
+    async def introspect_schema(self, table_pattern: str | None = None) -> str:
         """Introspect database schema to understand table structures."""
         try:
             # Pass table_pattern to get_schema_info for efficient filtering at DB level
@@ -96,7 +96,7 @@ class BaseSQLAgent(ABC):
         except Exception as e:
             return json.dumps({"error": f"Error listing tables: {str(e)}"})
 
-    async def execute_sql(self, query: str, limit: Optional[int] = None) -> str:
+    async def execute_sql(self, query: str, limit: int | None = None) -> str:
         """Execute a SQL query against the database."""
         try:
             # Security check - only allow SELECT queries unless write is enabled
@@ -147,7 +147,7 @@ class BaseSQLAgent(ABC):
             return json.dumps({"error": error_msg, "suggestions": suggestions})
 
     async def process_tool_call(
-        self, tool_name: str, tool_input: Dict[str, Any]
+        self, tool_name: str, tool_input: dict[str, Any]
     ) -> str:
         """Process a tool call and return the result."""
         if tool_name == "list_tables":
@@ -170,7 +170,7 @@ class BaseSQLAgent(ABC):
         else:
             return json.dumps({"error": f"Unknown tool: {tool_name}"})
 
-    def _validate_write_operation(self, query: str) -> Optional[str]:
+    def _validate_write_operation(self, query: str) -> str | None:
         """Validate if a write operation is allowed.
 
         Returns:
@@ -206,12 +206,12 @@ class BaseSQLAgent(ABC):
 
     async def plot_data(
         self,
-        y_values: List[float],
-        x_values: Optional[List[float]] = None,
+        y_values: list[float],
+        x_values: list[float] | None = None,
         plot_type: str = "line",
-        title: Optional[str] = None,
-        x_label: Optional[str] = None,
-        y_label: Optional[str] = None,
+        title: str | None = None,
+        x_label: str | None = None,
+        y_label: str | None = None,
     ) -> str:
         """Create a terminal plot using uniplot.
 

sqlsaber/agents/streaming.py

@@ -1,16 +1,16 @@
 """Streaming utilities for agents."""
 
-from typing import Any, Dict, List
+from typing import Any
 
 
 class StreamingResponse:
     """Helper class to manage streaming response construction."""
 
-    def __init__(self, content: List[Dict[str, Any]], stop_reason: str):
+    def __init__(self, content: list[dict[str, Any]], stop_reason: str):
         self.content = content
         self.stop_reason = stop_reason
 
 
-def build_tool_result_block(tool_use_id: str, content: str) -> Dict[str, Any]:
+def build_tool_result_block(tool_use_id: str, content: str) -> dict[str, Any]:
     """Build a tool result block for the conversation."""
     return {"type": "tool_result", "tool_use_id": tool_use_id, "content": content}

sqlsaber/cli/auth.py

@@ -0,0 +1,142 @@
+"""Authentication CLI commands."""
+
+import questionary
+import typer
+from rich.console import Console
+
+from sqlsaber.config.auth import AuthConfigManager, AuthMethod
+from sqlsaber.config.oauth_flow import AnthropicOAuthFlow
+
+# Global instances for CLI commands
+console = Console()
+config_manager = AuthConfigManager()
+
+# Create the authentication management CLI app
+auth_app = typer.Typer(
+    name="auth",
+    help="Manage authentication configuration",
+    add_completion=True,
+)
+
+
+@auth_app.command("setup")
+def setup_auth():
+    """Configure authentication method for SQLSaber."""
+    console.print("\n[bold]SQLSaber Authentication Setup[/bold]\n")
+
+    # Use questionary for selection
+    auth_choice = questionary.select(
+        "Choose your authentication method:",
+        choices=[
+            questionary.Choice(
+                title="Anthropic API Key",
+                value=AuthMethod.API_KEY,
+                description="You can create one by visiting https://console.anthropic.com",
+            ),
+            questionary.Choice(
+                title="Claude Pro or Max Subscription",
+                value=AuthMethod.CLAUDE_PRO,
+                description="This does not require creating an API Key, but requires a subscription at https://claude.ai",
+            ),
+        ],
+    ).ask()
+
+    if auth_choice is None:
+        console.print("[yellow]Setup cancelled.[/yellow]")
+        return
+
+    # Handle auth method setup
+    if auth_choice == AuthMethod.API_KEY:
+        console.print("\nTo configure your API key, you can either:")
+        console.print("• Set the ANTHROPIC_API_KEY environment variable")
+        console.print(
+            "• Let SQLsaber prompt you for the key when needed (stored securely)"
+        )
+
+        config_manager.set_auth_method(auth_choice)
+        console.print("\n[bold green]Authentication method saved![/bold green]")
+
+    elif auth_choice == AuthMethod.CLAUDE_PRO:
+        oauth_flow = AnthropicOAuthFlow()
+        try:
+            success = oauth_flow.authenticate()
+            if success:
+                config_manager.set_auth_method(auth_choice)
+                console.print(
+                    "\n[bold green]Authentication setup complete![/bold green]"
+                )
+            else:
+                console.print(
+                    "\n[yellow]OAuth authentication failed. Please try again.[/yellow]"
+                )
+                return
+        except Exception as e:
+            console.print(f"\n[red]Authentication setup failed: {str(e)}[/red]")
+            return
+
+    console.print(
+        "You can change this anytime by running [cyan]saber auth setup[/cyan] again."
+    )
+
+
+@auth_app.command("status")
+def show_auth_status():
+    """Show current authentication configuration."""
+    auth_method = config_manager.get_auth_method()
+
+    console.print("\n[bold blue]Authentication Status[/bold blue]")
+
+    if auth_method is None:
+        console.print("[yellow]No authentication method configured[/yellow]")
+        console.print("Run [cyan]saber auth setup[/cyan] to configure authentication.")
+    else:
+        if auth_method == AuthMethod.API_KEY:
+            console.print("[green]✓ API Key authentication configured[/green]")
+            console.print("Using Anthropic API key for authentication")
+        elif auth_method == AuthMethod.CLAUDE_PRO:
+            console.print("[green]✓ Claude Pro/Max subscription configured[/green]")
+
+            # Check OAuth token status
+            oauth_flow = AnthropicOAuthFlow()
+            if oauth_flow.has_valid_authentication():
+                console.print("OAuth token is valid and ready to use")
+            else:
+                console.print("[yellow]OAuth token missing or expired[/yellow]")
+
+
+@auth_app.command("reset")
+def reset_auth():
+    """Reset authentication configuration."""
+    if not config_manager.has_auth_configured():
+        console.print("[yellow]No authentication configuration to reset.[/yellow]")
+        return
+
+    current_method = config_manager.get_auth_method()
+    method_name = (
+        "API Key" if current_method == AuthMethod.API_KEY else "Claude Pro/Max"
+    )
+
+    if questionary.confirm(
+        f"Are you sure you want to reset the current authentication method ({method_name})?",
+        default=False,
+    ).ask():
+        # If Claude Pro, also remove OAuth tokens
+        if current_method == AuthMethod.CLAUDE_PRO:
+            oauth_flow = AnthropicOAuthFlow()
+            oauth_flow.remove_authentication()
+
+        # Clear the auth config by setting it to None
+        config = config_manager._load_config()
+        config["auth_method"] = None
+        config_manager._save_config(config)
+        console.print("[green]Authentication configuration reset.[/green]")
+        console.print(
+            "Run [cyan]saber auth setup[/cyan] to configure authentication again."
+        )
+    else:
+        console.print("Reset cancelled.")
+
+
+def create_auth_app() -> typer.Typer:
+    """Return the authentication management CLI app."""
+    return auth_app

sqlsaber/cli/commands.py

@@ -2,12 +2,12 @@
 
 import asyncio
 from pathlib import Path
-from typing import Optional
 
 import typer
 from rich.console import Console
 
 from sqlsaber.agents.anthropic import AnthropicSQLAgent
+from sqlsaber.cli.auth import create_auth_app
 from sqlsaber.cli.database import create_db_app
 from sqlsaber.cli.interactive import InteractiveSession
 from sqlsaber.cli.memory import create_memory_app
@@ -29,7 +29,7 @@ config_manager = DatabaseConfigManager()
 
 @app.callback()
 def main_callback(
-    database: Optional[str] = typer.Option(
+    database: str | None = typer.Option(
         None,
         "--database",
         "-d",
@@ -49,11 +49,11 @@ def main_callback(
 
 @app.command()
 def query(
-    query_text: Optional[str] = typer.Argument(
+    query_text: str | None = typer.Argument(
         None,
         help="SQL query in natural language (if not provided, starts interactive mode)",
     ),
-    database: Optional[str] = typer.Option(
+    database: str | None = typer.Option(
         None,
         "--database",
         "-d",
@@ -128,6 +128,7 @@ def query(
 
         finally:
             # Clean up
+            await agent.close()  # Close the agent's HTTP client
             await db_conn.close()
             console.print("\n[green]Goodbye![/green]")
 
@@ -135,6 +136,10 @@ def query(
     asyncio.run(run_session())
 
 
+# Add authentication management commands
+auth_app = create_auth_app()
+app.add_typer(auth_app, name="auth")
+
 # Add database management commands after main callback is defined
 db_app = create_db_app()
 app.add_typer(db_app, name="db")

sqlsaber/cli/completers.py

@@ -1,7 +1,5 @@
 """Command line completers for the CLI interface."""
 
-from typing import List, Tuple
-
 from prompt_toolkit.completion import Completer, Completion
 
 
@@ -37,13 +35,13 @@ class TableNameCompleter(Completer):
     """Custom completer for table names."""
 
     def __init__(self):
-        self._table_cache: List[Tuple[str, str]] = []
+        self._table_cache: list[tuple[str, str]] = []
 
-    def update_cache(self, tables_data: List[Tuple[str, str]]):
+    def update_cache(self, tables_data: list[tuple[str, str]]):
         """Update the cache with fresh table data."""
         self._table_cache = tables_data
 
-    def _get_table_names(self) -> List[Tuple[str, str]]:
+    def _get_table_names(self) -> list[tuple[str, str]]:
         """Get table names from cache."""
         return self._table_cache
 

sqlsaber/cli/database.py

@@ -3,7 +3,6 @@
 import asyncio
 import getpass
 from pathlib import Path
-from typing import Optional
 
 import questionary
 import typer
@@ -34,24 +33,24 @@ def add_database(
         "-t",
         help="Database type (postgresql, mysql, sqlite)",
     ),
-    host: Optional[str] = typer.Option(None, "--host", "-h", help="Database host"),
-    port: Optional[int] = typer.Option(None, "--port", "-p", help="Database port"),
-    database: Optional[str] = typer.Option(
+    host: str | None = typer.Option(None, "--host", "-h", help="Database host"),
+    port: int | None = typer.Option(None, "--port", "-p", help="Database port"),
+    database: str | None = typer.Option(
         None, "--database", "--db", help="Database name"
     ),
-    username: Optional[str] = typer.Option(None, "--username", "-u", help="Username"),
-    ssl_mode: Optional[str] = typer.Option(
+    username: str | None = typer.Option(None, "--username", "-u", help="Username"),
+    ssl_mode: str | None = typer.Option(
         None,
         "--ssl-mode",
         help="SSL mode (disable, allow, prefer, require, verify-ca, verify-full for PostgreSQL; DISABLED, PREFERRED, REQUIRED, VERIFY_CA, VERIFY_IDENTITY for MySQL)",
     ),
-    ssl_ca: Optional[str] = typer.Option(
+    ssl_ca: str | None = typer.Option(
         None, "--ssl-ca", help="SSL CA certificate file path"
     ),
-    ssl_cert: Optional[str] = typer.Option(
+    ssl_cert: str | None = typer.Option(
         None, "--ssl-cert", help="SSL client certificate file path"
     ),
-    ssl_key: Optional[str] = typer.Option(
+    ssl_key: str | None = typer.Option(
         None, "--ssl-key", help="SSL client private key file path"
     ),
     interactive: bool = typer.Option(
@@ -310,7 +309,7 @@ def set_default_database(
 
 @db_app.command("test")
 def test_database(
-    name: Optional[str] = typer.Argument(
+    name: str | None = typer.Argument(
         None,
         help="Name of the database connection to test (uses default if not specified)",
     ),

sqlsaber/cli/display.py

@@ -1,7 +1,6 @@
 """Display utilities for the CLI interface."""
 
 import json
-from typing import Optional
 
 from rich.console import Console
 from rich.markdown import Markdown
@@ -19,7 +18,7 @@ class DisplayManager:
         self,
         columns: list,
         header_style: str = "bold blue",
-        title: Optional[str] = None,
+        title: str | None = None,
     ) -> Table:
         """Create a Rich table with specified columns."""
         table = Table(show_header=True, header_style=header_style, title=title)
@@ -32,12 +31,9 @@ class DisplayManager:
                 table.add_column(col)
         return table
 
-    def show_tool_started(self, tool_name: str):
-        """Display tool started message."""
-        self.console.print(f"\n[yellow]🔧 Using tool: {tool_name}[/yellow]")
-
     def show_tool_executing(self, tool_name: str, tool_input: dict):
         """Display tool execution details."""
+        self.console.print(f"\n[yellow]🔧 Using tool: {tool_name}[/yellow]")
         if tool_name == "list_tables":
             self.console.print("[dim]  → Discovering available tables[/dim]")
         elif tool_name == "introspect_schema":
@@ -46,12 +42,14 @@ class DisplayManager:
         elif tool_name == "execute_sql":
             query = tool_input.get("query", "")
             self.console.print("\n[bold green]Executing SQL:[/bold green]")
+            self.show_newline()
             syntax = Syntax(query, "sql")
             self.console.print(syntax)
 
     def show_text_stream(self, text: str):
         """Display streaming text."""
-        self.console.print(text, end="", markup=False)
+        if text is not None:  # Extra safety check
+            self.console.print(text, end="", markup=False)
 
     def show_query_results(self, results: list):
         """Display query results in a formatted table."""

sqlsaber/cli/interactive.py

@@ -1,7 +1,6 @@
 """Interactive mode handling for the CLI."""
 
 import asyncio
-from typing import Optional
 
 import questionary
 from rich.console import Console
@@ -25,8 +24,8 @@ class InteractiveSession:
         self.agent = agent
         self.display = DisplayManager(console)
         self.streaming_handler = StreamingQueryHandler(console)
-        self.current_task: Optional[asyncio.Task] = None
-        self.cancellation_token: Optional[asyncio.Event] = None
+        self.current_task: asyncio.Task | None = None
+        self.cancellation_token: asyncio.Event | None = None
         self.table_completer = TableNameCompleter()
 
     def show_welcome_message(self):

sqlsaber/cli/memory.py

@@ -1,7 +1,5 @@
 """Memory management CLI commands."""
 
-from typing import Optional
-
 import typer
 from rich.console import Console
 from rich.table import Table
@@ -22,7 +20,7 @@ memory_app = typer.Typer(
 )
 
 
-def _get_database_name(database: Optional[str] = None) -> str:
+def _get_database_name(database: str | None = None) -> str:
     """Get the database name to use, either specified or default."""
     if database:
         db_config = config_manager.get_database(database)
@@ -46,7 +44,7 @@ def _get_database_name(database: Optional[str] = None) -> str:
 @memory_app.command("add")
 def add_memory(
     content: str = typer.Argument(..., help="Memory content to add"),
-    database: Optional[str] = typer.Option(
+    database: str | None = typer.Option(
         None,
         "--database",
         "-d",
@@ -68,7 +66,7 @@ def add_memory(
 
 @memory_app.command("list")
 def list_memories(
-    database: Optional[str] = typer.Option(
+    database: str | None = typer.Option(
         None,
         "--database",
         "-d",
@@ -107,7 +105,7 @@ def list_memories(
 @memory_app.command("show")
 def show_memory(
     memory_id: str = typer.Argument(..., help="Memory ID to show"),
-    database: Optional[str] = typer.Option(
+    database: str | None = typer.Option(
         None,
         "--database",
         "-d",
@@ -135,7 +133,7 @@ def show_memory(
 @memory_app.command("remove")
 def remove_memory(
     memory_id: str = typer.Argument(..., help="Memory ID to remove"),
-    database: Optional[str] = typer.Option(
+    database: str | None = typer.Option(
         None,
         "--database",
         "-d",
@@ -170,7 +168,7 @@ def remove_memory(
 
 @memory_app.command("clear")
 def clear_memories(
-    database: Optional[str] = typer.Option(
+    database: str | None = typer.Option(
         None,
         "--database",
         "-d",
@@ -213,7 +211,7 @@ def clear_memories(
 
 @memory_app.command("summary")
 def memory_summary(
-    database: Optional[str] = typer.Option(
+    database: str | None = typer.Option(
         None,
         "--database",
         "-d",

sqlsaber/cli/models.py

@@ -1,7 +1,6 @@
 """Model management CLI commands."""
 
 import asyncio
-from typing import Dict, List
 
 import httpx
 import questionary
@@ -28,7 +27,7 @@ class ModelManager:
     DEFAULT_MODEL = "anthropic:claude-sonnet-4-20250514"
     MODELS_API_URL = "https://models.dev/api.json"
 
-    async def fetch_available_models(self) -> List[Dict]:
+    async def fetch_available_models(self) -> list[dict]:
         """Fetch available models from models.dev API."""
         try:
             async with httpx.AsyncClient(timeout=10.0) as client:

sqlsaber/cli/streaming.py

@@ -23,8 +23,6 @@ class StreamingQueryHandler:
     ):
         """Execute a query with streaming display."""
 
-        has_content = False
-        explanation_started = False
         status = self.console.status(
             "[yellow]Crunching data...[/yellow]", spinner="bouncingBall"
         )
@@ -38,15 +36,10 @@ class StreamingQueryHandler:
                     break
 
                 if event.type == "tool_use":
-                    # Stop any ongoing status, but don't mark has_content yet
                     self._stop_status(status)
 
-                    if event.data["status"] == "started":
-                        # If explanation was streaming, add newline before tool use
-                        if explanation_started:
-                            self.display.show_newline()
-                        self.display.show_tool_started(event.data["name"])
-                    elif event.data["status"] == "executing":
+                    if event.data["status"] == "executing":
+                        self.display.show_newline()
                         self.display.show_tool_executing(
                             event.data["name"], event.data["input"]
                         )
@@ -54,12 +47,6 @@ class StreamingQueryHandler:
                 elif event.type == "text":
                     # Always stop status when text streaming starts
                     self._stop_status(status)
-
-                    if not explanation_started:
-                        explanation_started = True
-                        has_content = True
-
-                    # Print text as it streams
                     self.display.show_text_stream(event.data)
 
                 elif event.type == "query_result":
@@ -70,46 +57,33 @@ class StreamingQueryHandler:
                     # Handle tool results - particularly list_tables and introspect_schema
                     if event.data.get("tool_name") == "list_tables":
                         self.display.show_table_list(event.data["result"])
-                        has_content = True
                     elif event.data.get("tool_name") == "introspect_schema":
                         self.display.show_schema_info(event.data["result"])
-                        has_content = True
 
                 elif event.type == "plot_result":
                     # Handle plot results
                     self.display.show_plot(event.data)
-                    has_content = True
 
                 elif event.type == "processing":
-                    # Show status when processing tool results
-                    if explanation_started:
-                        self.display.show_newline()  # Add newline after explanation text
+                    self.display.show_newline()  # Add newline after explanation text
                     self._stop_status(status)
                     status = self.display.show_processing(event.data)
                     status.start()
-                    has_content = True
 
                 elif event.type == "error":
-                    if not has_content:
-                        self._stop_status(status)
-                        has_content = True
+                    self._stop_status(status)
                     self.display.show_error(event.data)
 
         except asyncio.CancelledError:
             # Handle cancellation gracefully
             self._stop_status(status)
-            if explanation_started:
-                self.display.show_newline()
+            self.display.show_newline()
             self.console.print("[yellow]Query interrupted[/yellow]")
             return
         finally:
             # Make sure status is stopped
             self._stop_status(status)
 
-            # Add a newline after streaming completes if explanation was shown
-            if explanation_started:
-                self.display.show_newline()  # Empty line for better readability
-
             # Display the last assistant response as markdown
             if hasattr(agent, "conversation_history") and agent.conversation_history:
                 last_message = agent.conversation_history[-1]

sqlsaber/clients/__init__.py

@@ -0,0 +1,6 @@
+"""Client implementations for various LLM APIs."""
+
+from .base import BaseLLMClient
+from .anthropic import AnthropicClient
+
+__all__ = ["BaseLLMClient", "AnthropicClient"]