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

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(

sqlsaber/cli/models.py

@@ -3,12 +3,13 @@
 import asyncio
 import sys
 
+import cyclopts
 import httpx
 import questionary
-import cyclopts
 from rich.console import Console
 from rich.table import Table
 
+from sqlsaber.config import providers
 from sqlsaber.config.settings import Config
 
 # Global instances for CLI commands
@@ -26,49 +27,75 @@ class ModelManager:
 
     DEFAULT_MODEL = "anthropic:claude-sonnet-4-20250514"
     MODELS_API_URL = "https://models.dev/api.json"
+    # Providers come from central registry
+    SUPPORTED_PROVIDERS = providers.all_keys()
+
+    async def fetch_available_models(
+        self, providers: list[str] | None = None
+    ) -> list[dict]:
+        """Fetch available models across providers from models.dev API.
 
-    async def fetch_available_models(self) -> list[dict]:
-        """Fetch available models from models.dev API."""
+        Returns list of dicts with keys: id (provider:model_id), provider, name, description, context_length, knowledge.
+        """
         try:
             async with httpx.AsyncClient(timeout=10.0) as client:
                 response = await client.get(self.MODELS_API_URL)
                 response.raise_for_status()
                 data = response.json()
 
-                # Filter for Anthropic models only
-                anthropic_models = []
-                anthropic_data = data.get("anthropic", {})
-
-                if "models" in anthropic_data:
-                    for model_id, model_info in anthropic_data["models"].items():
-                        # Convert to our format (anthropic:model-name)
-                        formatted_id = f"anthropic:{model_id}"
-
-                        # Extract cost information for display
-                        cost_info = model_info.get("cost", {})
+                providers = providers or self.SUPPORTED_PROVIDERS
+                results: list[dict] = []
+
+                for provider in providers:
+                    prov_data = data.get(provider, {})
+                    models_obj = (
+                        prov_data.get("models") or prov_data.get("Models") or {}
+                    )
+                    if not isinstance(models_obj, dict):
+                        continue
+                    for model_id, model_info in models_obj.items():
+                        formatted_id = f"{provider}:{model_id}"
+                        # cost
+                        cost_info = (
+                            model_info.get("cost", {})
+                            if isinstance(model_info, dict)
+                            else {}
+                        )
                         cost_display = ""
-                        if cost_info:
+                        if isinstance(cost_info, dict) and cost_info:
                             input_cost = cost_info.get("input", 0)
                             output_cost = cost_info.get("output", 0)
                             cost_display = f"${input_cost}/{output_cost} per 1M tokens"
+                        # context
+                        limit_info = (
+                            model_info.get("limit", {})
+                            if isinstance(model_info, dict)
+                            else {}
+                        )
+                        context_length = (
+                            limit_info.get("context", 0)
+                            if isinstance(limit_info, dict)
+                            else 0
+                        )
 
-                        # Extract context length
-                        limit_info = model_info.get("limit", {})
-                        context_length = limit_info.get("context", 0)
-
-                        anthropic_models.append(
+                        results.append(
                             {
                                 "id": formatted_id,
-                                "name": model_info.get("name", model_id),
+                                "provider": provider,
+                                "name": model_info.get("name", model_id)
+                                if isinstance(model_info, dict)
+                                else model_id,
                                 "description": cost_display,
                                 "context_length": context_length,
-                                "knowledge": model_info.get("knowledge", ""),
+                                "knowledge": model_info.get("knowledge", "")
+                                if isinstance(model_info, dict)
+                                else "",
                             }
                         )
 
-                # Sort by name for better display
-                anthropic_models.sort(key=lambda x: x["name"])
-                return anthropic_models
+                # Sort by provider then by name
+                results.sort(key=lambda x: (x["provider"], x["name"]))
+                return results
         except Exception as e:
             console.print(f"[red]Error fetching models: {e}[/red]")
             return []
@@ -110,7 +137,8 @@ def list():
             )
             return
 
-        table = Table(title="Available Anthropic Models")
+        table = Table(title="Available Models")
+        table.add_column("Provider", style="magenta")
         table.add_column("ID", style="cyan")
         table.add_column("Name", style="green")
         table.add_column("Description", style="white")
@@ -133,6 +161,7 @@ def list():
             )
 
             table.add_row(
+                model.get("provider", "-"),
                 model["id"],
                 model["name"],
                 description,
@@ -161,8 +190,9 @@ def set():
         # Create choices for questionary
         choices = []
         for model in models:
-            # Format: "ID - Name (Description)"
-            choice_text = f"{model['id']} - {model['name']}"
+            # Format: "[provider] ID - Name (Description)"
+            prov = model.get("provider", "?")
+            choice_text = f"[{prov}] {model['id']} - {model['name']}"
             if model["description"]:
                 choice_text += f" ({model['description'][:50]}{'...' if len(model['description']) > 50 else ''})"
 
@@ -179,7 +209,6 @@ def set():
         selected_model = await questionary.select(
             "Select a model:",
             choices=choices,
-            use_shortcuts=True,
             use_search_filter=True,
             use_jk_keys=False,  # Disable j/k keys when using search filter
             default=choices[default_index] if choices else None,

sqlsaber/cli/streaming.py

@@ -1,100 +1,137 @@
-"""Streaming query handling for the CLI."""
+"""Streaming query handling for the CLI (pydantic-ai based)."""
 
 import asyncio
-
+import json
+from typing import AsyncIterable
+
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.messages import (
+    AgentStreamEvent,
+    FunctionToolCallEvent,
+    FunctionToolResultEvent,
+    PartDeltaEvent,
+    PartStartEvent,
+    TextPart,
+    TextPartDelta,
+    ThinkingPart,
+    ThinkingPartDelta,
+)
 from rich.console import Console
 
-from sqlsaber.agents.base import BaseSQLAgent
 from sqlsaber.cli.display import DisplayManager
 
 
 class StreamingQueryHandler:
-    """Handles streaming query execution and display."""
+    """Handles streaming query execution and display using pydantic-ai events."""
 
     def __init__(self, console: Console):
         self.console = console
         self.display = DisplayManager(console)
 
+        self.status = self.console.status(
+            "[yellow]Crunching data...[/yellow]", spinner="bouncingBall"
+        )
+
+    async def _event_stream_handler(
+        self, ctx: RunContext, event_stream: AsyncIterable[AgentStreamEvent]
+    ) -> None:
+        async for event in event_stream:
+            if isinstance(event, PartStartEvent):
+                if isinstance(event.part, (TextPart, ThinkingPart)):
+                    self.status.stop()
+                    self.display.show_text_stream(event.part.content)
+
+            elif isinstance(event, PartDeltaEvent):
+                if isinstance(event.delta, (TextPartDelta, ThinkingPartDelta)):
+                    delta = event.delta.content_delta or ""
+                    if delta:
+                        self.status.stop()
+                        self.display.show_text_stream(delta)
+
+            elif isinstance(event, FunctionToolCallEvent):
+                # Show tool execution start
+                self.status.stop()
+                args = event.part.args_as_dict()
+                self.display.show_newline()
+                self.display.show_tool_executing(event.part.tool_name, args)
+
+            elif isinstance(event, FunctionToolResultEvent):
+                self.status.stop()
+                # Route tool result to appropriate display
+                tool_name = event.result.tool_name
+                content = event.result.content
+                if tool_name == "list_tables":
+                    self.display.show_table_list(content)
+                elif tool_name == "introspect_schema":
+                    self.display.show_schema_info(content)
+                elif tool_name == "execute_sql":
+                    try:
+                        data = json.loads(content)
+                        if data.get("success") and data.get("results"):
+                            self.display.show_query_results(data["results"])  # type: ignore[arg-type]
+                    except json.JSONDecodeError:
+                        # If not JSON, ignore here
+                        pass
+                elif tool_name == "plot_data":
+                    self.display.show_plot(
+                        {"tool_name": tool_name, "result": content, "input": {}}
+                    )
+
     async def execute_streaming_query(
         self,
         user_query: str,
-        agent: BaseSQLAgent,
+        agent: Agent,
         cancellation_token: asyncio.Event | None = None,
+        message_history: list | None = None,
     ):
-        """Execute a query with streaming display."""
-
-        status = self.console.status(
-            "[yellow]Crunching data...[/yellow]", spinner="bouncingBall"
-        )
-        status.start()
-
+        self.status.start()
         try:
-            async for event in agent.query_stream(
-                user_query, cancellation_token=cancellation_token
-            ):
-                if cancellation_token is not None and cancellation_token.is_set():
-                    break
-
-                if event.type == "tool_use":
-                    self._stop_status(status)
-
-                    if event.data["status"] == "executing":
-                        self.display.show_newline()
-                        self.display.show_tool_executing(
-                            event.data["name"], event.data["input"]
-                        )
-
-                elif event.type == "text":
-                    # Always stop status when text streaming starts
-                    self._stop_status(status)
-                    self.display.show_text_stream(event.data)
-
-                elif event.type == "query_result":
-                    if event.data["results"]:
-                        self.display.show_query_results(event.data["results"])
-
-                elif event.type == "tool_result":
-                    # 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"])
-                    elif event.data.get("tool_name") == "introspect_schema":
-                        self.display.show_schema_info(event.data["result"])
-
-                elif event.type == "plot_result":
-                    # Handle plot results
-                    self.display.show_plot(event.data)
-
-                elif event.type == "processing":
-                    self.display.show_newline()  # Add newline after explanation text
-                    self._stop_status(status)
-                    status = self.display.show_processing(event.data)
-                    status.start()
-
-                elif event.type == "error":
-                    self._stop_status(status)
-                    self.display.show_error(event.data)
-
+            # If Anthropic OAuth, inject SQLsaber instructions before the first user prompt
+            prepared_prompt: str | list[str] = user_query
+            is_oauth = bool(getattr(agent, "_sqlsaber_is_oauth", False))
+            no_history = not message_history
+            if is_oauth and no_history:
+                ib = getattr(agent, "_sqlsaber_instruction_builder", None)
+                mm = getattr(agent, "_sqlsaber_memory_manager", None)
+                db_type = getattr(agent, "_sqlsaber_db_type", "database")
+                db_name = getattr(agent, "_sqlsaber_database_name", None)
+                instructions = (
+                    ib.build_instructions(db_type=db_type) if ib is not None else ""
+                )
+                mem = (
+                    mm.format_memories_for_prompt(db_name)
+                    if (mm is not None and db_name)
+                    else ""
+                )
+                parts = [p for p in (instructions, mem) if p and str(p).strip()]
+                if parts:
+                    injected = "\n\n".join(parts)
+                    prepared_prompt = [injected, user_query]
+
+            # Run the agent with our event stream handler
+            run = await agent.run(
+                prepared_prompt,
+                message_history=message_history,
+                event_stream_handler=self._event_stream_handler,
+            )
+            # After the run completes, show the assistant's final text as markdown if available
+            try:
+                output = run.output
+                if isinstance(output, str) and output.strip():
+                    self.display.show_newline()
+                    self.display.show_markdown_response(
+                        [{"type": "text", "text": output}]
+                    )
+            except Exception as e:
+                self.display.show_error(str(e))
+                self.display.show_newline()
+            return run
         except asyncio.CancelledError:
-            # Handle cancellation gracefully
-            self._stop_status(status)
             self.display.show_newline()
             self.console.print("[yellow]Query interrupted[/yellow]")
-            return
+            return None
         finally:
-            # Make sure status is stopped
-            self._stop_status(status)
-
-            # Display the last assistant response as markdown
-            if hasattr(agent, "conversation_history") and agent.conversation_history:
-                last_message = agent.conversation_history[-1]
-                if last_message.get("role") == "assistant" and last_message.get(
-                    "content"
-                ):
-                    self.display.show_markdown_response(last_message["content"])
-
-    def _stop_status(self, status):
-        """Safely stop a status spinner."""
-        try:
-            status.stop()
-        except Exception:
-            pass  # Status might already be stopped
+            try:
+                self.status.stop()
+            except Exception:
+                pass

sqlsaber/config/api_keys.py

@@ -6,6 +6,8 @@ import os
 import keyring
 from rich.console import Console
 
+from sqlsaber.config import providers
+
 console = Console()
 
 
@@ -30,9 +32,7 @@ class APIKeyManager:
         try:
             api_key = keyring.get_password(service_name, provider)
             if api_key:
-                console.print(
-                    f"Using stored {provider} API key from keyring", style="dim"
-                )
+                console.print(f"Using stored {provider} API key", style="dim")
                 return api_key
         except Exception as e:
             # Keyring access failed, continue to prompt
@@ -43,12 +43,9 @@ class APIKeyManager:
 
     def _get_env_var_name(self, provider: str) -> str:
         """Get the expected environment variable name for a provider."""
-        if provider == "openai":
-            return "OPENAI_API_KEY"
-        elif provider == "anthropic":
-            return "ANTHROPIC_API_KEY"
-        else:
-            return "AI_API_KEY"
+        # Normalize aliases to canonical provider keys
+        key = providers.canonical(provider) or provider
+        return providers.env_var_name(key)
 
     def _get_service_name(self, provider: str) -> str:
         """Get the keyring service name for a provider."""
@@ -60,7 +57,7 @@ class APIKeyManager:
         """Prompt user for API key and store it in keyring."""
         try:
             console.print(
-                f"\n{provider.title()} API key not found in environment or keyring."
+                f"\n{provider.title()} API key not found in environment or your OS's credentials store."
             )
             console.print("You can either:")
             console.print(f"  1. Set the {env_var_name} environment variable")
@@ -85,7 +82,8 @@ class APIKeyManager:
                 console.print("API key stored securely for future use", style="green")
             except Exception as e:
                 console.print(
-                    f"Warning: Could not store API key in keyring: {e}", style="yellow"
+                    f"Warning: Could not store API key in your operating system's credentials store: {e}",
+                    style="yellow",
                 )
                 console.print(
                     "You may need to enter it again next time", style="yellow"