sqlsaber 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

sqlsaber/cli/interactive.py

@@ -1,46 +1,21 @@
 """Interactive mode handling for the CLI."""
 
 import asyncio
-from typing import Optional
 
 import questionary
-from prompt_toolkit.completion import Completer, Completion
 from rich.console import Console
 from rich.panel import Panel
 
 from sqlsaber.agents.base import BaseSQLAgent
+from sqlsaber.cli.completers import (
+    CompositeCompleter,
+    SlashCommandCompleter,
+    TableNameCompleter,
+)
 from sqlsaber.cli.display import DisplayManager
 from sqlsaber.cli.streaming import StreamingQueryHandler
 
 
-class SlashCommandCompleter(Completer):
-    """Custom completer for slash commands."""
-
-    def get_completions(self, document, complete_event):
-        """Get completions for slash commands."""
-        # Only provide completions if the line starts with "/"
-        text = document.text
-        if text.startswith("/"):
-            # Get the partial command after the slash
-            partial_cmd = text[1:]
-
-            # Define available commands with descriptions
-            commands = [
-                ("clear", "Clear conversation history"),
-                ("exit", "Exit the interactive session"),
-                ("quit", "Exit the interactive session"),
-            ]
-
-            # Yield completions that match the partial command
-            for cmd, description in commands:
-                if cmd.startswith(partial_cmd):
-                    yield Completion(
-                        cmd,
-                        start_position=-len(partial_cmd),
-                        display_meta=description,
-                    )
-
-
 class InteractiveSession:
     """Manages interactive CLI sessions."""
 
@@ -49,8 +24,9 @@ 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):
         """Display welcome message for interactive mode."""
@@ -63,7 +39,8 @@ class InteractiveSession:
                 "[bold green]SQLSaber - Use the agent Luke![/bold green]\n\n"
                 "[bold]Your agentic SQL assistant.[/bold]\n\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]",
+                "[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",
             )
         )
@@ -75,6 +52,39 @@ class InteractiveSession:
             "[dim]Press Ctrl+C during query execution to interrupt and return to prompt.[/dim]\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()
+
+            # Parse the table information
+            table_list = []
+            if isinstance(tables_data, dict) and "tables" in tables_data:
+                for table in tables_data["tables"]:
+                    if isinstance(table, dict):
+                        name = table.get("name", "")
+                        schema = table.get("schema", "")
+                        full_name = table.get("full_name", "")
+
+                        # Use full_name if available, otherwise construct it
+                        if full_name:
+                            table_name = full_name
+                        elif schema and schema != "main":
+                            table_name = f"{schema}.{name}"
+                        else:
+                            table_name = name
+
+                        # No description needed - cleaner completions
+                        table_list.append((table_name, ""))
+
+            # Update the completer cache
+            self.table_completer.update_cache(table_list)
+
+        except Exception:
+            # If there's an error, just use empty cache
+            self.table_completer.update_cache([])
+
     async def _execute_query_with_cancellation(self, user_query: str):
         """Execute a query with cancellation support."""
         # Create cancellation token
@@ -101,6 +111,9 @@ class InteractiveSession:
         """Run the interactive session loop."""
         self.show_welcome_message()
 
+        # Initialize table cache
+        await self._update_table_cache()
+
         while True:
             try:
                 user_query = await questionary.text(
@@ -108,7 +121,9 @@ class InteractiveSession:
                     qmark="",
                     multiline=True,
                     instruction="",
-                    completer=SlashCommandCompleter(),
+                    completer=CompositeCompleter(
+                        SlashCommandCompleter(), self.table_completer
+                    ),
                 ).ask_async()
 
                 if not user_query:

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,45 +57,40 @@ 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]
+                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."""

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"]

sqlsaber/clients/anthropic.py

@@ -0,0 +1,285 @@
+"""Anthropic API client implementation."""
+
+import asyncio
+import json
+import logging
+from typing import Any, AsyncIterator
+
+import httpx
+
+from .base import BaseLLMClient
+from .exceptions import LLMClientError, create_exception_from_response
+from .models import CreateMessageRequest
+from .streaming import AnthropicStreamAdapter, StreamingResponse
+
+logger = logging.getLogger(__name__)
+
+
+class AnthropicClient(BaseLLMClient):
+    """Client for Anthropic's Claude API."""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        oauth_token: str | None = None,
+        base_url: str | None = None,
+    ):
+        """Initialize the Anthropic client.
+
+        Args:
+            api_key: Anthropic API key
+            base_url: Base URL for the API (defaults to Anthropic's API)
+        """
+        super().__init__(api_key or "", base_url)
+
+        if not api_key and not oauth_token:
+            raise ValueError("Either api_key or oauth_token must be provided")
+
+        self.oauth_token = oauth_token
+        self.use_oauth = oauth_token is not None
+        self.base_url = base_url or "https://api.anthropic.com"
+        self.client: httpx.AsyncClient | None = None
+
+    def _get_client(self) -> httpx.AsyncClient:
+        """Get or create the HTTP client."""
+        if self.client is None or self.client.is_closed:
+            # Configure timeouts and connection limits for reliability
+            timeout = httpx.Timeout(
+                connect=10.0,  # Connection timeout
+                read=60.0,  # Read timeout for streaming
+                write=10.0,  # Write timeout
+                pool=10.0,  # Pool timeout
+            )
+            limits = httpx.Limits(
+                max_keepalive_connections=20, max_connections=100, keepalive_expiry=30.0
+            )
+            self.client = httpx.AsyncClient(
+                timeout=timeout, limits=limits, follow_redirects=True
+            )
+        return self.client
+
+    def _get_headers(self) -> dict[str, str]:
+        """Get the standard headers for API requests."""
+        if self.use_oauth:
+            # OAuth headers for Claude Pro authentication (matching Claude Code CLI)
+            return {
+                "Authorization": f"Bearer {self.oauth_token}",
+                "Content-Type": "application/json",
+                "anthropic-version": "2023-06-01",
+                "anthropic-beta": "oauth-2025-04-20",
+                "User-Agent": "ClaudeCode/1.0 (Anthropic Claude Code CLI)",
+                "Accept": "application/json",
+                "X-Client-Name": "claude-code",
+                "X-Client-Version": "1.0.0",
+            }
+        else:
+            # API key headers for standard authentication
+            return {
+                "x-api-key": self.api_key,
+                "anthropic-version": "2023-06-01",
+                "content-type": "application/json",
+            }
+
+    async def create_message_with_tools(
+        self,
+        request: CreateMessageRequest,
+        cancellation_token: asyncio.Event | None = None,
+    ) -> AsyncIterator[Any]:
+        """Create a message with tool support and stream the response.
+
+        This method handles the full message creation flow including tool use,
+        similar to what the current AnthropicSQLAgent expects.
+
+        Args:
+            request: The message creation request
+            cancellation_token: Optional event to signal cancellation
+
+        Yields:
+            Stream events and final StreamingResponse
+        """
+        request.stream = True
+
+        client = self._get_client()
+        url = f"{self.base_url}/v1/messages"
+        headers = self._get_headers()
+        data = request.to_dict()
+
+        try:
+            async with client.stream(
+                "POST", url, headers=headers, json=data
+            ) as response:
+                request_id = response.headers.get("request-id")
+
+                if response.status_code != 200:
+                    response_content = await response.aread()
+                    response_data = json.loads(response_content.decode())
+                    raise create_exception_from_response(
+                        response.status_code, response_data, request_id
+                    )
+
+                # Use stream adapter to convert raw events and track state
+                adapter = AnthropicStreamAdapter()
+                raw_stream = self._process_sse_stream(response, cancellation_token)
+
+                async for event in adapter.process_stream(
+                    raw_stream, cancellation_token
+                ):
+                    yield event
+
+                # Create final response object with proper state
+                response_obj = StreamingResponse(
+                    content=adapter.get_content_blocks(),
+                    stop_reason=adapter.get_stop_reason(),
+                )
+
+                # Yield special event with response
+                yield {"type": "response_ready", "data": response_obj}
+
+        except asyncio.CancelledError:
+            # Handle cancellation gracefully
+            logger.debug("Stream cancelled")
+            return
+        except Exception as e:
+            if not isinstance(e, LLMClientError):
+                raise LLMClientError(f"Stream processing error: {str(e)}")
+            raise
+
+    def _handle_ping_event(self, event_data: str) -> dict[str, Any]:
+        """Handle ping event data.
+
+        Args:
+            event_data: Raw event data string
+
+        Returns:
+            Parsed ping event
+        """
+        try:
+            return {"type": "ping", "data": json.loads(event_data)}
+        except json.JSONDecodeError:
+            return {"type": "ping", "data": {}}
+
+    def _handle_error_event(self, event_data: str) -> None:
+        """Handle error event data.
+
+        Args:
+            event_data: Raw event data string
+
+        Raises:
+            LLMClientError: Always raises with error details
+        """
+        try:
+            error_data = json.loads(event_data)
+            raise LLMClientError(
+                error_data.get("message", "Stream error"),
+                error_data.get("type", "stream_error"),
+            )
+        except json.JSONDecodeError:
+            raise LLMClientError("Stream error with invalid JSON")
+
+    def _parse_event_data(
+        self, event_type: str | None, event_data: str
+    ) -> dict[str, Any] | None:
+        """Parse event data based on event type.
+
+        Args:
+            event_type: Type of the event
+            event_data: Raw event data string
+
+        Returns:
+            Parsed event or None if parsing failed
+        """
+        try:
+            parsed_data = json.loads(event_data)
+            return {"type": event_type, "data": parsed_data}
+        except json.JSONDecodeError as e:
+            logger.warning(f"Failed to parse stream data for event {event_type}: {e}")
+            return None
+
+    def _process_sse_line(
+        self, line: str, event_type: str | None
+    ) -> tuple[str | None, dict[str, Any] | None]:
+        """Process a single SSE line.
+
+        Args:
+            line: Line to process
+            event_type: Current event type
+
+        Returns:
+            Tuple of (new_event_type, event_to_yield)
+        """
+        if line.startswith("event: "):
+            return line[7:], None
+        elif line.startswith("data: "):
+            event_data = line[6:]
+
+            if event_type == "ping":
+                return event_type, self._handle_ping_event(event_data)
+            elif event_type == "error":
+                self._handle_error_event(event_data)
+                return event_type, None  # Never reached due to exception
+            else:
+                parsed_event = self._parse_event_data(event_type, event_data)
+                return event_type, parsed_event
+
+        return event_type, None
+
+    async def _process_sse_stream(
+        self,
+        response: httpx.Response,
+        cancellation_token: asyncio.Event | None = None,
+    ) -> AsyncIterator[dict[str, Any]]:
+        """Process server-sent events from the response stream.
+
+        Args:
+            response: The HTTP response object
+            cancellation_token: Optional event to signal cancellation
+
+        Yields:
+            Parsed stream events
+
+        Raises:
+            LLMClientError: If stream processing fails
+        """
+        buffer = ""
+        event_type = None
+
+        try:
+            async for chunk in response.aiter_bytes():
+                if cancellation_token is not None and cancellation_token.is_set():
+                    return
+
+                try:
+                    buffer += chunk.decode("utf-8")
+                except UnicodeDecodeError as e:
+                    logger.warning(f"Failed to decode chunk: {e}")
+                    continue
+
+                while "\n" in buffer:
+                    line, buffer = buffer.split("\n", 1)
+                    line = line.strip()
+
+                    if not line:
+                        continue
+
+                    event_type, event_to_yield = self._process_sse_line(
+                        line, event_type
+                    )
+                    if event_to_yield is not None:
+                        yield event_to_yield
+
+        except httpx.TimeoutException as e:
+            raise LLMClientError(f"Stream timeout error: {str(e)}")
+        except httpx.NetworkError as e:
+            raise LLMClientError(f"Network error during streaming: {str(e)}")
+        except httpx.HTTPError as e:
+            raise LLMClientError(f"HTTP error during streaming: {str(e)}")
+        except asyncio.TimeoutError:
+            raise LLMClientError("Stream timeout")
+        except Exception as e:
+            raise LLMClientError(f"Unexpected error during streaming: {str(e)}")
+
+    async def close(self):
+        """Close the HTTP client."""
+        if self.client and not self.client.is_closed:
+            await self.client.aclose()
+            self.client = None

sqlsaber/clients/base.py

@@ -0,0 +1,31 @@
+"""Abstract base class for LLM clients."""
+
+from abc import ABC
+
+
+class BaseLLMClient(ABC):
+    """Abstract base class for LLM API clients."""
+
+    def __init__(self, api_key: str, base_url: str | None = None):
+        """Initialize the client with API key and optional base URL.
+
+        Args:
+            api_key: API key for authentication
+            base_url: Base URL for the API (optional, uses default if not provided)
+        """
+        self.api_key = api_key
+        self.base_url = base_url
+
+    async def close(self):
+        """Close the client and clean up resources."""
+        # Default implementation does nothing
+        # Subclasses can override to clean up HTTP sessions, etc.
+        pass
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.close()