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

sqlsaber/clients/exceptions.py

@@ -0,0 +1,117 @@
+"""Exception classes for LLM client errors."""
+
+from typing import Any
+
+
+class LLMClientError(Exception):
+    """Base exception for LLM client errors."""
+
+    def __init__(
+        self,
+        message: str,
+        error_type: str | None = None,
+        status_code: int | None = None,
+        request_id: str | None = None,
+    ):
+        super().__init__(message)
+        self.error_type = error_type
+        self.status_code = status_code
+        self.request_id = request_id
+
+
+class AuthenticationError(LLMClientError):
+    """Authentication failed - invalid API key."""
+
+    def __init__(self, message: str = "Invalid API key", **kwargs):
+        super().__init__(message, "authentication_error", **kwargs)
+
+
+class PermissionError(LLMClientError):
+    """Permission denied for the requested resource."""
+
+    def __init__(self, message: str = "Permission denied", **kwargs):
+        super().__init__(message, "permission_error", **kwargs)
+
+
+class NotFoundError(LLMClientError):
+    """Requested resource not found."""
+
+    def __init__(self, message: str = "Resource not found", **kwargs):
+        super().__init__(message, "not_found_error", **kwargs)
+
+
+class InvalidRequestError(LLMClientError):
+    """Invalid request format or content."""
+
+    def __init__(self, message: str = "Invalid request", **kwargs):
+        super().__init__(message, "invalid_request_error", **kwargs)
+
+
+class RequestTooLargeError(LLMClientError):
+    """Request exceeds maximum allowed size."""
+
+    def __init__(self, message: str = "Request too large", **kwargs):
+        super().__init__(message, "request_too_large", **kwargs)
+
+
+class RateLimitError(LLMClientError):
+    """Rate limit exceeded."""
+
+    def __init__(self, message: str = "Rate limit exceeded", **kwargs):
+        super().__init__(message, "rate_limit_error", **kwargs)
+
+
+class APIError(LLMClientError):
+    """Internal API error."""
+
+    def __init__(self, message: str = "Internal API error", **kwargs):
+        super().__init__(message, "api_error", **kwargs)
+
+
+class OverloadedError(LLMClientError):
+    """API is temporarily overloaded."""
+
+    def __init__(self, message: str = "API temporarily overloaded", **kwargs):
+        super().__init__(message, "overloaded_error", **kwargs)
+
+
+# Mapping of HTTP status codes to exception classes
+STATUS_CODE_TO_EXCEPTION = {
+    400: InvalidRequestError,
+    401: AuthenticationError,
+    403: PermissionError,
+    404: NotFoundError,
+    413: RequestTooLargeError,
+    429: RateLimitError,
+    500: APIError,
+    529: OverloadedError,
+}
+
+
+def create_exception_from_response(
+    status_code: int,
+    response_data: dict[str, Any],
+    request_id: str | None = None,
+) -> LLMClientError:
+    """Create appropriate exception from HTTP response."""
+    error_data = response_data.get("error", {})
+    message = error_data.get("message", f"HTTP {status_code} error")
+    error_type = error_data.get("type")
+
+    exception_class = STATUS_CODE_TO_EXCEPTION.get(status_code, LLMClientError)
+
+    # Handle base vs subclass constructors
+    if exception_class == LLMClientError:
+        return exception_class(
+            message,
+            error_type,
+            status_code,
+            request_id,
+        )
+    else:
+        # Subclasses only take message and **kwargs
+        return exception_class(
+            message,
+            status_code=status_code,
+            request_id=request_id,
+        )

sqlsaber/clients/models.py

@@ -0,0 +1,282 @@
+"""Data models for LLM client requests and responses."""
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+
+class MessageRole(str, Enum):
+    """Message roles in a conversation."""
+
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+
+
+class ContentType(str, Enum):
+    """Content block types."""
+
+    TEXT = "text"
+    IMAGE = "image"
+    TOOL_USE = "tool_use"
+    TOOL_RESULT = "tool_result"
+
+
+class ToolChoiceType(str, Enum):
+    """Tool choice types."""
+
+    AUTO = "auto"
+    ANY = "any"
+    TOOL = "tool"
+    NONE = "none"
+
+
+class StopReason(str, Enum):
+    """Stop reasons for message completion."""
+
+    END_TURN = "end_turn"
+    MAX_TOKENS = "max_tokens"
+    STOP_SEQUENCE = "stop_sequence"
+    TOOL_USE = "tool_use"
+
+
+@dataclass
+class ContentBlock:
+    """A content block in a message."""
+
+    type: ContentType
+    content: str | dict[str, Any]
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary format."""
+        if self.type == ContentType.TEXT:
+            return {"type": "text", "text": self.content}
+        elif self.type == ContentType.TOOL_USE:
+            return {
+                "type": "tool_use",
+                "id": self.content["id"],
+                "name": self.content["name"],
+                "input": self.content["input"],
+            }
+        elif self.type == ContentType.TOOL_RESULT:
+            return {
+                "type": "tool_result",
+                "tool_use_id": self.content["tool_use_id"],
+                "content": self.content["content"],
+            }
+        else:
+            return {"type": self.type.value, **self.content}
+
+
+@dataclass
+class Message:
+    """A message in a conversation."""
+
+    role: MessageRole
+    content: str | list[ContentBlock]
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary format for API requests."""
+        if isinstance(self.content, str):
+            return {"role": self.role.value, "content": self.content}
+        else:
+            return {
+                "role": self.role.value,
+                "content": [block.to_dict() for block in self.content],
+            }
+
+
+@dataclass
+class ToolDefinition:
+    """Definition of a tool that can be called."""
+
+    name: str
+    description: str
+    input_schema: dict[str, Any]
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary format for API requests."""
+        return {
+            "name": self.name,
+            "description": self.description,
+            "input_schema": self.input_schema,
+        }
+
+
+@dataclass
+class ToolChoice:
+    """Tool choice configuration."""
+
+    type: ToolChoiceType
+    name: str | None = None
+    disable_parallel_tool_use: bool = False
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary format for API requests."""
+        result = {"type": self.type.value}
+        if self.name:
+            result["name"] = self.name
+        if self.disable_parallel_tool_use:
+            result["disable_parallel_tool_use"] = True
+        return result
+
+
+@dataclass
+class CreateMessageRequest:
+    """Request to create a message."""
+
+    model: str
+    messages: list[Message]
+    max_tokens: int
+    system: str | None = None
+    tools: list[ToolDefinition] | None = None
+    tool_choice: ToolChoice | None = None
+    temperature: float | None = None
+    stream: bool = False
+    stop_sequences: list[str] | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary format for API requests."""
+        data = {
+            "model": self.model,
+            "messages": [msg.to_dict() for msg in self.messages],
+            "max_tokens": self.max_tokens,
+        }
+
+        if self.system:
+            data["system"] = self.system
+        if self.tools:
+            data["tools"] = [tool.to_dict() for tool in self.tools]
+        if self.tool_choice:
+            data["tool_choice"] = self.tool_choice.to_dict()
+        if self.temperature is not None:
+            data["temperature"] = self.temperature
+        if self.stream:
+            data["stream"] = True
+        if self.stop_sequences:
+            data["stop_sequences"] = self.stop_sequences
+
+        return data
+
+
+@dataclass
+class Usage:
+    """Token usage information."""
+
+    input_tokens: int
+    output_tokens: int
+
+
+@dataclass
+class MessageResponse:
+    """Response from message creation."""
+
+    id: str
+    model: str
+    role: MessageRole
+    content: list[ContentBlock]
+    stop_reason: StopReason
+    stop_sequence: str | None
+    usage: Usage
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "MessageResponse":
+        """Create from API response dictionary."""
+        content_blocks = []
+        for block_data in data["content"]:
+            if block_data["type"] == "text":
+                content_blocks.append(
+                    ContentBlock(ContentType.TEXT, block_data["text"])
+                )
+            elif block_data["type"] == "tool_use":
+                content_blocks.append(
+                    ContentBlock(
+                        ContentType.TOOL_USE,
+                        {
+                            "id": block_data["id"],
+                            "name": block_data["name"],
+                            "input": block_data["input"],
+                        },
+                    )
+                )
+
+        return cls(
+            id=data["id"],
+            model=data["model"],
+            role=MessageRole(data["role"]),
+            content=content_blocks,
+            stop_reason=StopReason(data["stop_reason"]),
+            stop_sequence=data.get("stop_sequence"),
+            usage=Usage(
+                input_tokens=data["usage"]["input_tokens"],
+                output_tokens=data["usage"]["output_tokens"],
+            ),
+        )
+
+
+# Stream event types and data models for streaming
+@dataclass
+class StreamEventData:
+    """Base class for stream event data."""
+
+    pass
+
+
+@dataclass
+class TextDeltaData(StreamEventData):
+    """Text delta stream event data."""
+
+    text: str
+
+
+@dataclass
+class ToolUseStartData(StreamEventData):
+    """Tool use start stream event data."""
+
+    id: str
+    name: str
+
+
+@dataclass
+class ToolUseInputData(StreamEventData):
+    """Tool use input stream event data."""
+
+    partial_json: str
+
+
+@dataclass
+class MessageStartData(StreamEventData):
+    """Message start stream event data."""
+
+    message: dict[str, Any]
+
+
+@dataclass
+class MessageDeltaData(StreamEventData):
+    """Message delta stream event data."""
+
+    delta: dict[str, Any]
+    usage: dict[str, Any] | None = None
+
+
+@dataclass
+class ContentBlockStartData(StreamEventData):
+    """Content block start stream event data."""
+
+    index: int
+    content_block: dict[str, Any]
+
+
+@dataclass
+class ContentBlockDeltaData(StreamEventData):
+    """Content block delta stream event data."""
+
+    index: int
+    delta: dict[str, Any]
+
+
+@dataclass
+class ContentBlockStopData(StreamEventData):
+    """Content block stop stream event data."""
+
+    index: int

sqlsaber/clients/streaming.py

@@ -0,0 +1,257 @@
+"""Streaming adapters and utilities for LLM clients."""
+
+import asyncio
+import json
+import logging
+from typing import Any, AsyncIterator
+
+logger = logging.getLogger(__name__)
+
+
+class AnthropicStreamAdapter:
+    """Adapter to convert raw Anthropic stream events to standardized format.
+
+    This adapter converts the raw SSE events from Anthropic API into objects
+    that match the structure expected by the current AnthropicSQLAgent.
+    """
+
+    def __init__(self):
+        self.content_blocks: list[dict[str, Any]] = []
+        self.tool_use_blocks: list[dict[str, Any]] = []
+
+    async def process_stream(
+        self,
+        raw_stream: AsyncIterator[dict[str, Any]],
+        cancellation_token: asyncio.Event | None = None,
+    ) -> AsyncIterator[Any]:
+        """Process raw stream events and yield adapted events.
+
+        Args:
+            raw_stream: Raw stream events from the API
+            cancellation_token: Optional cancellation token
+
+        Yields:
+            Adapted stream events that match the SDK format
+        """
+        async for raw_event in raw_stream:
+            # Check for cancellation
+            if cancellation_token is not None and cancellation_token.is_set():
+                return
+
+            # Convert raw event to SDK-like event
+            adapted_event = self._adapt_event(raw_event)
+            if adapted_event:
+                yield adapted_event
+
+    def _adapt_event(self, raw_event: dict[str, Any]) -> Any | None:
+        """Adapt a raw stream event to match SDK format.
+
+        Args:
+            raw_event: Raw event from the API
+
+        Returns:
+            Adapted event object or None if event should be filtered
+        """
+        event_type = raw_event.get("type")
+        event_data = raw_event.get("data", {})
+
+        if event_type == "ping":
+            # Create a ping event object
+            return PingEvent()
+
+        elif event_type == "message_start":
+            # Create message start event
+            return MessageStartEvent(event_data.get("message", {}))
+
+        elif event_type == "content_block_start":
+            # Create content block start event
+            index = event_data.get("index", 0)
+            content_block = event_data.get("content_block", {})
+
+            # Initialize content blocks list if needed
+            while len(self.content_blocks) <= index:
+                self.content_blocks.append({"type": "text", "text": ""})
+
+            if content_block.get("type") == "tool_use":
+                # Add to tool use blocks tracking
+                tool_block = {
+                    "id": content_block.get("id"),
+                    "name": content_block.get("name"),
+                    "input": {},
+                    "_partial": "",
+                }
+                self.tool_use_blocks.append(tool_block)
+
+            return ContentBlockStartEvent(index, content_block)
+
+        elif event_type == "content_block_delta":
+            # Create content block delta event
+            index = event_data.get("index", 0)
+            delta = event_data.get("delta", {})
+
+            # Update content blocks tracking
+            if index < len(self.content_blocks):
+                if delta.get("type") == "text_delta":
+                    self.content_blocks[index]["text"] += delta.get("text", "")
+                elif delta.get("type") == "input_json_delta":
+                    # Update tool use input tracking
+                    if self.tool_use_blocks:
+                        current_tool = self.tool_use_blocks[-1]
+                        current_tool["_partial"] += delta.get("partial_json", "")
+                        try:
+                            current_tool["input"] = json.loads(current_tool["_partial"])
+                        except json.JSONDecodeError:
+                            pass  # Partial JSON, continue accumulating
+
+            return ContentBlockDeltaEvent(index, delta)
+
+        elif event_type == "content_block_stop":
+            # Create content block stop event
+            index = event_data.get("index", 0)
+            return ContentBlockStopEvent(index)
+
+        elif event_type == "message_delta":
+            # Create message delta event
+            delta = event_data.get("delta", {})
+            usage = event_data.get("usage", {})
+            return MessageDeltaEvent(delta, usage)
+
+        elif event_type == "message_stop":
+            # Finalize tool blocks
+            self._finalize_tool_blocks()
+            return MessageStopEvent()
+
+        elif event_type == "error":
+            # Create error event
+            return ErrorEvent(event_data)
+
+        else:
+            # Unknown event type, log and ignore
+            logger.debug(f"Unknown event type: {event_type}")
+            return None
+
+    def _finalize_tool_blocks(self):
+        """Finalize tool use blocks by cleaning up and adding to content blocks."""
+        for block in self.tool_use_blocks:
+            block["type"] = "tool_use"
+            if "_partial" in block:
+                del block["_partial"]
+            self.content_blocks.append(block)
+
+    def get_stop_reason(self) -> str:
+        """Get the stop reason based on current state."""
+        if self.tool_use_blocks:
+            return "tool_use"
+        return "stop"
+
+    def get_content_blocks(self) -> list[dict[str, Any]]:
+        """Get the current content blocks."""
+        return self.content_blocks.copy()
+
+
+# Event classes that match the SDK structure
+class BaseStreamEvent:
+    """Base class for stream events."""
+
+    def __init__(self, event_type: str):
+        self.type = event_type
+
+
+class PingEvent(BaseStreamEvent):
+    """Ping event."""
+
+    def __init__(self):
+        super().__init__("ping")
+
+
+class MessageStartEvent(BaseStreamEvent):
+    """Message start event."""
+
+    def __init__(self, message: dict[str, Any]):
+        super().__init__("message_start")
+        self.message = message
+
+
+class ContentBlockStartEvent(BaseStreamEvent):
+    """Content block start event."""
+
+    def __init__(self, index: int, content_block: dict[str, Any]):
+        super().__init__("content_block_start")
+        self.index = index
+        self.content_block = MockContentBlock(content_block)
+
+
+class ContentBlockDeltaEvent(BaseStreamEvent):
+    """Content block delta event."""
+
+    def __init__(self, index: int, delta: dict[str, Any]):
+        super().__init__("content_block_delta")
+        self.index = index
+        self.delta = MockDelta(delta)
+
+
+class ContentBlockStopEvent(BaseStreamEvent):
+    """Content block stop event."""
+
+    def __init__(self, index: int):
+        super().__init__("content_block_stop")
+        self.index = index
+
+
+class MessageDeltaEvent(BaseStreamEvent):
+    """Message delta event."""
+
+    def __init__(self, delta: dict[str, Any], usage: dict[str, Any]):
+        super().__init__("message_delta")
+        self.delta = delta
+        self.usage = usage
+
+
+class MessageStopEvent(BaseStreamEvent):
+    """Message stop event."""
+
+    def __init__(self):
+        super().__init__("message_stop")
+
+
+class ErrorEvent(BaseStreamEvent):
+    """Error event."""
+
+    def __init__(self, error_data: dict[str, Any]):
+        super().__init__("error")
+        self.error = error_data
+
+
+# Mock classes to match SDK object structure
+class MockContentBlock:
+    """Mock content block object that matches SDK structure."""
+
+    def __init__(self, data: dict[str, Any]):
+        self.type = data.get("type")
+        self.id = data.get("id")
+        self.name = data.get("name")
+        self.input = data.get("input", {})
+
+
+class MockDelta:
+    """Mock delta object that matches SDK structure."""
+
+    def __init__(self, data: dict[str, Any]):
+        self.data = data
+        # Set attributes based on delta type
+        if data.get("type") == "text_delta":
+            self.text = data.get("text", "")
+        elif data.get("type") == "input_json_delta":
+            self.partial_json = data.get("partial_json", "")
+
+    def __getattr__(self, name):
+        """Allow access to any attribute from the data."""
+        return self.data.get(name)
+
+
+class StreamingResponse:
+    """Response object for streaming that matches the current agent's expectations."""
+
+    def __init__(self, content: list[dict[str, Any]], stop_reason: str):
+        self.content = content
+        self.stop_reason = stop_reason

sqlsaber/config/api_keys.py

@@ -2,7 +2,6 @@
 
 import getpass
 import os
-from typing import Optional
 
 import keyring
 from rich.console import Console
@@ -16,7 +15,7 @@ class APIKeyManager:
     def __init__(self):
         self.service_prefix = "sqlsaber"
 
-    def get_api_key(self, provider: str) -> Optional[str]:
+    def get_api_key(self, provider: str) -> str | None:
         """Get API key for the specified provider using cascading logic."""
         env_var_name = self._get_env_var_name(provider)
         service_name = self._get_service_name(provider)
@@ -57,7 +56,7 @@ class APIKeyManager:
 
     def _prompt_and_store_key(
         self, provider: str, env_var_name: str, service_name: str
-    ) -> Optional[str]:
+    ) -> str | None:
         """Prompt user for API key and store it in keyring."""
         try:
             console.print(