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

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"

sqlsaber/config/providers.py

@@ -0,0 +1,116 @@
+"""Central registry for supported AI providers.
+
+This module defines a single source of truth for providers used across the
+codebase (CLI, config, agents). Update this file to add or modify providers.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, Iterable, List, Optional
+
+
+@dataclass(frozen=True)
+class ProviderSpec:
+    """Specification for a provider."""
+
+    key: str
+    env_var: str
+    supports_oauth: bool = False
+    aliases: tuple[str, ...] = ()
+
+
+# Ordered definition -> used for CLI display order
+_PROVIDERS: List[ProviderSpec] = [
+    ProviderSpec(
+        key="anthropic",
+        env_var="ANTHROPIC_API_KEY",
+        supports_oauth=True,
+        aliases=(),
+    ),
+    ProviderSpec(
+        key="openai",
+        env_var="OPENAI_API_KEY",
+        aliases=(),
+    ),
+    ProviderSpec(
+        key="google",
+        env_var="GOOGLE_API_KEY",
+        # Historically some model IDs start with "google-gla"; treat as alias
+        aliases=("google-gla",),
+    ),
+    ProviderSpec(
+        key="groq",
+        env_var="GROQ_API_KEY",
+        aliases=(),
+    ),
+    ProviderSpec(
+        key="mistral",
+        env_var="MISTRAL_API_KEY",
+        aliases=(),
+    ),
+    ProviderSpec(
+        key="cohere",
+        env_var="COHERE_API_KEY",
+        aliases=(),
+    ),
+    ProviderSpec(
+        key="huggingface",
+        env_var="HUGGINGFACE_API_KEY",
+        aliases=(),
+    ),
+]
+
+
+# Fast lookup maps
+_BY_KEY: Dict[str, ProviderSpec] = {p.key: p for p in _PROVIDERS}
+_ALIAS_TO_KEY: Dict[str, str] = {
+    alias: p.key for p in _PROVIDERS for alias in p.aliases
+}
+
+
+def all_keys() -> List[str]:
+    """Return provider keys in display order."""
+    return [p.key for p in _PROVIDERS]
+
+
+def env_var_name(key: str) -> str:
+    """Return the expected environment variable for a provider.
+
+    Falls back to a generic name if the provider is unknown.
+    """
+    spec = _BY_KEY.get(key)
+    return spec.env_var if spec else "AI_API_KEY"
+
+
+def supports_oauth(key: str) -> bool:
+    """Return True if the provider supports OAuth in SQLsaber."""
+    spec = _BY_KEY.get(key)
+    return bool(spec and spec.supports_oauth)
+
+
+def canonical(key_or_alias: str) -> Optional[str]:
+    """Return the canonical provider key for a provider or alias.
+
+    Returns None if not recognized.
+    """
+    if key_or_alias in _BY_KEY:
+        return key_or_alias
+    return _ALIAS_TO_KEY.get(key_or_alias)
+
+
+def provider_from_model(model_name: str) -> Optional[str]:
+    """Infer the canonical provider key from a model identifier.
+
+    Accepts either "provider:model_id" or a bare provider string. Aliases are
+    normalized to their canonical provider key.
+    """
+    if not model_name:
+        return None
+    provider_raw = model_name.split(":", 1)[0]
+    return canonical(provider_raw)
+
+
+def specs() -> Iterable[ProviderSpec]:
+    """Iterate provider specifications (in display order)."""
+    return tuple(_PROVIDERS)

sqlsaber/config/settings.py

@@ -9,6 +9,7 @@ from typing import Any
 
 import platformdirs
 
+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
@@ -84,47 +85,66 @@ class Config:
         self.model_name = self.model_config_manager.get_model()
         self.api_key_manager = APIKeyManager()
         self.auth_config_manager = AuthConfigManager()
-        self.oauth_flow = AnthropicOAuthFlow()
 
-        # Get authentication credentials based on configured method
+        # Authentication method (API key or Anthropic OAuth)
         self.auth_method = self.auth_config_manager.get_auth_method()
-        self.api_key = None
-        self.oauth_token = None
-
-        if self.auth_method == AuthMethod.CLAUDE_PRO:
-            # Try to get OAuth token and refresh if needed
-            try:
-                token = self.oauth_flow.refresh_token_if_needed()
-                if token:
-                    self.oauth_token = token.access_token
-            except Exception:
-                # OAuth token unavailable, will need to re-authenticate
-                pass
+
+        # Optional Anthropic OAuth access token (only relevant for provider=='anthropic')
+        if self.auth_method == AuthMethod.CLAUDE_PRO and self.model_name.startswith(
+            "anthropic"
+        ):
+            self.oauth_token = self.get_oauth_access_token()
         else:
-            # Use API key authentication (default or explicitly configured)
             self.api_key = self._get_api_key()
+            # self.oauth_token = None
 
     def _get_api_key(self) -> str | None:
         """Get API key for the model provider using cascading logic."""
-        model = self.model_name
-        if model.startswith("anthropic:"):
-            return self.api_key_manager.get_api_key("anthropic")
+        model = self.model_name or ""
+        prov = providers.provider_from_model(model)
+        if prov in set(providers.all_keys()):
+            return self.api_key_manager.get_api_key(prov)  # type: ignore[arg-type]
+        return None
 
     def set_model(self, model: str) -> None:
         """Set the model and update configuration."""
         self.model_config_manager.set_model(model)
         self.model_name = model
 
+    def get_oauth_access_token(self) -> str | None:
+        """Return a valid Anthropic OAuth access token if configured, else None.
+
+        Uses the stored refresh token (if present) to refresh as needed.
+        Only relevant when provider is 'anthropic'.
+        """
+        if not self.model_name.startswith("anthropic"):
+            return None
+        try:
+            flow = AnthropicOAuthFlow()
+            token = flow.refresh_token_if_needed()
+            return token.access_token if token else None
+        except Exception:
+            return None
+
     def validate(self):
-        """Validate that necessary configuration is present."""
-        # 1. Claude-Pro flow → require OAuth token only
-        if self.auth_method == AuthMethod.CLAUDE_PRO:
-            if not self.oauth_token:
-                raise ValueError(
-                    "OAuth token not available. Run 'saber auth setup' to authenticate with Claude Pro."
-                )
-            return  # OAuth path satisfied – nothing more to check
-
-        # 2. Default / API-key flow → require API key
-        if not self.api_key:
-            raise ValueError("Anthropic API key not found.")
+        """Validate that necessary configuration is present.
+
+        Also ensure provider env var is set from keyring if needed for API-key flows.
+        """
+        model = self.model_name or ""
+        provider_key = providers.provider_from_model(model)
+        env_var = providers.env_var_name(provider_key or "") if provider_key else None
+        if env_var:
+            # Anthropic special-case: allow OAuth in lieu of API key only when explicitly configured
+            if (
+                provider_key == "anthropic"
+                and self.auth_method == AuthMethod.CLAUDE_PRO
+                and self.oauth_token
+            ):
+                return
+            # If we don't have a key resolved from env/keyring, raise
+            if not self.api_key:
+                raise ValueError(f"{provider_key.capitalize()} API key not found.")
+            # Hydrate env var for downstream SDKs if missing
+            if not os.getenv(env_var):
+                os.environ[env_var] = self.api_key

sqlsaber/database/connection.py

@@ -1,10 +1,10 @@
 """Database connection management."""
 
-from abc import ABC, abstractmethod
-from typing import Any
-from urllib.parse import urlparse, parse_qs
 import ssl
+from abc import ABC, abstractmethod
 from pathlib import Path
+from typing import Any
+from urllib.parse import parse_qs, urlparse
 
 import aiomysql
 import aiosqlite