sqlsaber 0.14.0__py3-none-any.whl → 0.16.0__py3-none-any.whl

sqlsaber/cli/auth.py

@@ -1,11 +1,17 @@
 """Authentication CLI commands."""
 
-import questionary
+import os
+
 import cyclopts
+import keyring
+import questionary
 from rich.console import Console
 
+from sqlsaber.config import providers
+from sqlsaber.config.api_keys import APIKeyManager
 from sqlsaber.config.auth import AuthConfigManager, AuthMethod
 from sqlsaber.config.oauth_flow import AnthropicOAuthFlow
+from sqlsaber.config.oauth_tokens import OAuthTokenManager
 
 # Global instances for CLI commands
 console = Console()
@@ -20,59 +26,59 @@ auth_app = cyclopts.App(
 
 @auth_app.command
 def setup():
-    """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",
-            ),
-        ],
+    """Configure authentication for SQLsaber (API keys and Anthropic OAuth)."""
+    console.print("\n[bold]SQLsaber Authentication Setup[/bold]\n")
+
+    provider = questionary.select(
+        "Select provider to configure:",
+        choices=providers.all_keys(),
     ).ask()
 
-    if auth_choice is None:
+    if provider 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)
+    if provider == "anthropic":
+        # Let user choose API key or OAuth
+        method_choice = questionary.select(
+            "Select Anthropic authentication method:",
+            choices=[
+                {"name": "API key", "value": AuthMethod.API_KEY},
+                {"name": "Claude Pro/Max (OAuth)", "value": AuthMethod.CLAUDE_PRO},
+            ],
+        ).ask()
+
+        if method_choice == AuthMethod.CLAUDE_PRO:
+            flow = AnthropicOAuthFlow()
+            if flow.authenticate():
+                config_manager.set_auth_method(AuthMethod.CLAUDE_PRO)
                 console.print(
-                    "\n[bold green]Authentication setup complete![/bold green]"
+                    "\n[bold green]✓ Anthropic OAuth configured successfully![/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]")
+                console.print("\n[red]✗ Anthropic OAuth setup failed.[/red]")
+            console.print(
+                "You can change this anytime by running [cyan]saber auth setup[/cyan] again."
+            )
             return
 
+    # API key flow (all providers + Anthropic when selected above)
+    api_key_manager = APIKeyManager()
+    env_var = api_key_manager._get_env_var_name(provider)
+    console.print("\nTo configure your API key, you can either:")
+    console.print(f"• Set the {env_var} environment variable")
+    console.print("• Let SQLsaber prompt you for the key when needed (stored securely)")
+
+    # Fetch/store key (cascades env -> keyring -> prompt)
+    api_key = api_key_manager.get_api_key(provider)
+    if api_key:
+        config_manager.set_auth_method(AuthMethod.API_KEY)
+        console.print(
+            f"\n[bold green]✓ {provider.title()} API key configured successfully![/bold green]"
+        )
+    else:
+        console.print("\n[yellow]No API key configured.[/yellow]")
+
     console.print(
         "You can change this anytime by running [cyan]saber auth setup[/cyan] again."
     )
@@ -80,7 +86,7 @@ def setup():
 
 @auth_app.command
 def status():
-    """Show current authentication configuration."""
+    """Show current authentication configuration and provider key status."""
     auth_method = config_manager.get_auth_method()
 
     console.print("\n[bold blue]Authentication Status[/bold blue]")
@@ -88,52 +94,113 @@ def status():
     if auth_method is None:
         console.print("[yellow]No authentication method configured[/yellow]")
         console.print("Run [cyan]saber auth setup[/cyan] to configure authentication.")
+        return
+
+    # Show configured method summary
+    if auth_method == AuthMethod.CLAUDE_PRO:
+        console.print("[green]✓ Anthropic Claude Pro/Max (OAuth) configured[/green]\n")
     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]")
+        console.print("[green]✓ API Key authentication configured[/green]\n")
+
+    # Show per-provider status without prompting
+    api_key_manager = APIKeyManager()
+    for provider in providers.all_keys():
+        if provider == "anthropic":
+            # Include OAuth status
+            if OAuthTokenManager().has_oauth_token("anthropic"):
+                console.print("> anthropic (oauth): [green]configured[/green]")
+        env_var = api_key_manager._get_env_var_name(provider)
+        service = api_key_manager._get_service_name(provider)
+        from_env = bool(os.getenv(env_var))
+        from_keyring = bool(keyring.get_password(service, provider))
+        if from_env:
+            console.print(f"> {provider}: configured via {env_var}")
+        elif from_keyring:
+            console.print(f"> {provider}: [green]configured[/green]")
+        else:
+            console.print(f"> {provider}: [yellow]not configured[/yellow]")
 
 
 @auth_app.command
 def reset():
-    """Reset authentication configuration."""
-    if not config_manager.has_auth_configured():
-        console.print("[yellow]No authentication configuration to reset.[/yellow]")
+    """Reset credentials for a selected provider (API key and/or OAuth)."""
+    console.print("\n[bold]SQLsaber Authentication Reset[/bold]\n")
+
+    # Choose provider to reset (mirrors setup)
+    provider = questionary.select(
+        "Select provider to reset:",
+        choices=providers.all_keys(),
+    ).ask()
+
+    if provider is None:
+        console.print("[yellow]Reset cancelled.[/yellow]")
         return
 
-    current_method = config_manager.get_auth_method()
-    method_name = (
-        "API Key" if current_method == AuthMethod.API_KEY else "Claude Pro/Max"
+    api_key_manager = APIKeyManager()
+    service = api_key_manager._get_service_name(provider)
+
+    # Determine what exists in keyring
+    api_key_present = bool(keyring.get_password(service, provider))
+    oauth_present = (
+        OAuthTokenManager().has_oauth_token("anthropic")
+        if provider == "anthropic"
+        else False
     )
 
-    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]")
+    if not api_key_present and not oauth_present:
         console.print(
-            "Run [cyan]saber auth setup[/cyan] to configure authentication again."
+            f"[yellow]No stored credentials found for {provider}. Nothing to reset.[/yellow]"
         )
-    else:
+        return
+
+    # Build confirmation message
+    to_remove: list[str] = []
+    if oauth_present:
+        to_remove.append("Anthropic OAuth token")
+    if api_key_present:
+        to_remove.append(f"{provider.title()} API key")
+
+    summary = ", ".join(to_remove)
+    confirmed = questionary.confirm(
+        f"Remove the following for {provider}: {summary}?",
+        default=False,
+    ).ask()
+
+    if not confirmed:
         console.print("Reset cancelled.")
+        return
+
+    # Perform deletions
+    if oauth_present:
+        OAuthTokenManager().remove_oauth_token("anthropic")
+    if api_key_present:
+        try:
+            keyring.delete_password(service, provider)
+            console.print(f"Removed {provider} API key from keyring", style="green")
+        except keyring.errors.PasswordDeleteError:
+            # Already absent; treat as success
+            pass
+        except Exception as e:
+            console.print(f"Warning: Could not remove API key: {e}", style="yellow")
+
+    # Optionally clear global auth method if removing Anthropic OAuth configuration
+    if provider == "anthropic" and oauth_present:
+        current_method = config_manager.get_auth_method()
+        if current_method == AuthMethod.CLAUDE_PRO:
+            also_clear = questionary.confirm(
+                "Anthropic OAuth was removed. Also unset the global auth method?",
+                default=False,
+            ).ask()
+            if also_clear:
+                config = config_manager._load_config()
+                config["auth_method"] = None
+                config_manager._save_config(config)
+                console.print("Global auth method unset.", style="green")
+
+    console.print("\n[bold green]✓ Reset complete.[/bold green]")
+    console.print(
+        "Environment variables are not modified by this command.", style="dim"
+    )
 
 
 def create_auth_app() -> cyclopts.App:

sqlsaber/cli/commands.py

@@ -7,7 +7,7 @@ from typing import Annotated
 import cyclopts
 from rich.console import Console
 
-from sqlsaber.agents.anthropic import AnthropicSQLAgent
+from sqlsaber.agents import build_sqlsaber_agent
 from sqlsaber.cli.auth import create_auth_app
 from sqlsaber.cli.database import create_db_app
 from sqlsaber.cli.interactive import InteractiveSession
@@ -15,7 +15,13 @@ from sqlsaber.cli.memory import create_memory_app
 from sqlsaber.cli.models import create_models_app
 from sqlsaber.cli.streaming import StreamingQueryHandler
 from sqlsaber.config.database import DatabaseConfigManager
-from sqlsaber.database.connection import DatabaseConnection
+from sqlsaber.database.connection import (
+    CSVConnection,
+    DatabaseConnection,
+    MySQLConnection,
+    PostgreSQLConnection,
+    SQLiteConnection,
+)
 from sqlsaber.database.resolver import DatabaseResolutionError, resolve_database
 
 
@@ -123,25 +129,34 @@ def query(
         except Exception as e:
             raise CLIError(f"Error creating database connection: {e}")
 
-        # Create agent instance with database name for memory context
-        agent = AnthropicSQLAgent(db_conn, db_name)
+        # Create pydantic-ai agent instance with database name for memory context
+        agent = build_sqlsaber_agent(db_conn, db_name)
 
         try:
             if actual_query:
                 # Single query mode with streaming
                 streaming_handler = StreamingQueryHandler(console)
+                # Compute DB type for the greeting line
+                db_type = (
+                    "PostgreSQL"
+                    if isinstance(db_conn, PostgreSQLConnection)
+                    else "MySQL"
+                    if isinstance(db_conn, MySQLConnection)
+                    else "SQLite"
+                    if isinstance(db_conn, (SQLiteConnection, CSVConnection))
+                    else "database"
+                )
                 console.print(
-                    f"[bold blue]Connected to:[/bold blue] {db_name} {agent._get_database_type_name()}\n"
+                    f"[bold blue]Connected to:[/bold blue] {db_name} ({db_type})\n"
                 )
                 await streaming_handler.execute_streaming_query(actual_query, agent)
             else:
                 # Interactive mode
-                session = InteractiveSession(console, agent)
+                session = InteractiveSession(console, agent, db_conn, db_name)
                 await session.run()
 
         finally:
             # Clean up
-            await agent.close()  # Close the agent's HTTP client
             await db_conn.close()
             console.print("\n[green]Goodbye![/green]")
 

sqlsaber/cli/database.py

@@ -6,8 +6,8 @@ import sys
 from pathlib import Path
 from typing import Annotated
 
-import questionary
 import cyclopts
+import questionary
 from rich.console import Console
 from rich.table import Table
 

sqlsaber/cli/interactive.py

@@ -3,10 +3,10 @@
 import asyncio
 
 import questionary
+from pydantic_ai import Agent
 from rich.console import Console
 from rich.panel import Panel
 
-from sqlsaber.agents.base import BaseSQLAgent
 from sqlsaber.cli.completers import (
     CompositeCompleter,
     SlashCommandCompleter,
@@ -14,25 +14,44 @@ from sqlsaber.cli.completers import (
 )
 from sqlsaber.cli.display import DisplayManager
 from sqlsaber.cli.streaming import StreamingQueryHandler
+from sqlsaber.database.schema import SchemaManager
 
 
 class InteractiveSession:
     """Manages interactive CLI sessions."""
 
-    def __init__(self, console: Console, agent: BaseSQLAgent):
+    def __init__(self, console: Console, agent: Agent, db_conn, database_name: str):
         self.console = console
         self.agent = agent
+        self.db_conn = db_conn
+        self.database_name = database_name
         self.display = DisplayManager(console)
         self.streaming_handler = StreamingQueryHandler(console)
         self.current_task: asyncio.Task | None = None
         self.cancellation_token: asyncio.Event | None = None
         self.table_completer = TableNameCompleter()
+        self.message_history: list | None = []
 
     def show_welcome_message(self):
         """Display welcome message for interactive mode."""
         # Show database information
-        db_name = getattr(self.agent, "database_name", None) or "Unknown"
-        db_type = self.agent._get_database_type_name()
+        db_name = self.database_name or "Unknown"
+        from sqlsaber.database.connection import (
+            CSVConnection,
+            MySQLConnection,
+            PostgreSQLConnection,
+            SQLiteConnection,
+        )
+
+        db_type = (
+            "PostgreSQL"
+            if isinstance(self.db_conn, PostgreSQLConnection)
+            else "MySQL"
+            if isinstance(self.db_conn, MySQLConnection)
+            else "SQLite"
+            if isinstance(self.db_conn, (SQLiteConnection, CSVConnection))
+            else "database"
+        )
 
         self.console.print(
             Panel.fit(
@@ -44,26 +63,27 @@ class InteractiveSession:
 ███████  ██████  ███████ ███████ ██   ██ ██████  ███████ ██   ██
             ▀▀
 """
-                "\n\n"
-                "[dim]Use '/clear' to reset conversation, '/exit' or '/quit' to leave.[/dim]\n\n"
-                "[dim]Start a message with '#' to add something to agent's memory for this database.[/dim]\n\n"
-                "[dim]Type '@' to get table name completions.[/dim]",
-                border_style="green",
             )
         )
         self.console.print(
-            f"[bold blue]Connected to:[/bold blue] {db_name} ({db_type})\n"
+            "\n",
+            "[dim] ≥ Use '/clear' to reset conversation",
+            "[dim] ≥ Use '/exit' or '/quit' to leave[/dim]",
+            "[dim] ≥ Use 'Ctrl+C' to interrupt and return to prompt\n\n",
+            "[dim] ≥ Start message with '#' to add something to agent's memory for this database",
+            "[dim] ≥ Type '@' to get table name completions",
+            "[dim] ≥ Press 'Esc-Enter' or 'Meta-Enter' to submit your question",
+            sep="\n",
         )
+
         self.console.print(
-            "[dim]Press Esc-Enter or Meta-Enter to submit your query.[/dim]\n"
-            "[dim]Press Ctrl+C during query execution to interrupt and return to prompt.[/dim]\n"
+            f"[bold blue]\n\nConnected to:[/bold blue] {db_name} ({db_type})\n"
         )
 
     async def _update_table_cache(self):
         """Update the table completer cache with fresh data."""
         try:
-            # Use the schema manager directly which has built-in caching
-            tables_data = await self.agent.schema_manager.list_tables()
+            tables_data = await SchemaManager(self.db_conn).list_tables()
 
             # Parse the table information
             table_list = []
@@ -100,16 +120,20 @@ class InteractiveSession:
         # Create the query task
         query_task = asyncio.create_task(
             self.streaming_handler.execute_streaming_query(
-                user_query, self.agent, self.cancellation_token
+                user_query, self.agent, self.cancellation_token, self.message_history
             )
         )
         self.current_task = query_task
 
         try:
-            # Simply await the query task
-            # Ctrl+C will be handled by the KeyboardInterrupt exception in run()
-            await query_task
-
+            run_result = await query_task
+            # Persist message history from this run using pydantic-ai API
+            if run_result is not None:
+                try:
+                    # Use all_messages() so the system prompt and all prior turns are preserved
+                    self.message_history = run_result.all_messages()
+                except Exception:
+                    pass
         finally:
             self.current_task = None
             self.cancellation_token = None
@@ -144,7 +168,8 @@ class InteractiveSession:
                     break
 
                 if user_query == "/clear":
-                    await self.agent.clear_history()
+                    # Reset local history (pydantic-ai call will receive empty history on next run)
+                    self.message_history = []
                     self.console.print("[green]Conversation history cleared.[/green]\n")
                     continue
 
@@ -153,18 +178,28 @@ class InteractiveSession:
                     if memory_text.startswith("#"):
                         memory_content = memory_text[1:].strip()  # Remove # and trim
                         if memory_content:
-                            # Add memory
-                            memory_id = self.agent.add_memory(memory_content)
-                            if memory_id:
-                                self.console.print(
-                                    f"[green]✓ Memory added:[/green] {memory_content}"
-                                )
-                                self.console.print(
-                                    f"[dim]Memory ID: {memory_id}[/dim]\n"
+                            # Add memory via the agent's memory manager
+                            try:
+                                mm = getattr(
+                                    self.agent, "_sqlsaber_memory_manager", None
                                 )
-                            else:
+                                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 (no database context)[/yellow]\n"
+                                    "[yellow]Could not add memory[/yellow]\n"
                                 )
                         else:
                             self.console.print(