axion-code 1.0.0__py3-none-any.whl

axion/api/client.py

@@ -0,0 +1,259 @@
+"""Provider client factory - dispatches to Anthropic or OpenAI-compatible.
+
+Maps to: rust/crates/api/src/client.rs
+"""
+
+from __future__ import annotations
+
+import enum
+from typing import Any, AsyncIterator
+
+from axion.api.anthropic import AnthropicClient, AuthCredentials
+from axion.api.error import ApiError
+from axion.api.ollama import OllamaClient, is_ollama_model
+from axion.api.openai_compat import OpenAiCompatClient, OpenAiCompatConfig
+from axion.api.types import MessageRequest, MessageResponse, StreamEvent
+
+# Model alias resolution
+MODEL_ALIASES: dict[str, str] = {
+    # Anthropic Claude
+    "opus": "claude-opus-4-6",
+    "sonnet": "claude-sonnet-4-6",
+    "haiku": "claude-haiku-4-5",
+    "opus[1m]": "claude-opus-4-6",
+    "sonnet[1m]": "claude-sonnet-4-6",
+    "haiku[1m]": "claude-haiku-4-5",
+
+    # OpenAI — GPT-4 series
+    "gpt4": "gpt-4o",
+    "gpt4o": "gpt-4o",
+    "gpt-4": "gpt-4o",
+    "4o": "gpt-4o",
+    "gpt-4o-mini": "gpt-4o-mini",
+    "4o-mini": "gpt-4o-mini",
+    "gpt-4.1": "gpt-4.1",
+    "gpt-4.1-mini": "gpt-4.1-mini",
+    "gpt-4.1-nano": "gpt-4.1-nano",
+
+    # OpenAI — GPT-5 series
+    "gpt5": "gpt-5",
+    "gpt-5": "gpt-5",
+    "5": "gpt-5",
+    "gpt-5-mini": "gpt-5-mini",
+    "5-mini": "gpt-5-mini",
+    "gpt-5-nano": "gpt-5-nano",
+    "5-nano": "gpt-5-nano",
+    "gpt-5-pro": "gpt-5-pro",
+    "5-pro": "gpt-5-pro",
+    "gpt-5.4": "gpt-5.4",
+    "gpt-5.4-mini": "gpt-5.4-mini",
+    "gpt-5.4-nano": "gpt-5.4-nano",
+    "gpt-5.4-pro": "gpt-5.4-pro",
+
+    # OpenAI — Codex (real Codex models via the /v1/responses endpoint)
+    "codex": "gpt-5-codex",
+    "codex-mini": "gpt-5-codex-mini",
+    "gpt-5-codex": "gpt-5-codex",
+    "gpt-5-codex-mini": "gpt-5-codex-mini",
+
+    # OpenAI — o-series (reasoning)
+    "o1": "o1",
+    "o1-pro": "o1-pro",
+    "o3": "o3",
+    "o3-mini": "o3-mini",
+    "o3-pro": "o3-pro",
+    "o4-mini": "o4-mini",
+
+    # xAI
+    "grok": "grok-2",
+    "grok2": "grok-2",
+    "grok-3": "grok-3",
+
+    # Ollama / local
+    "local": "llama3.1",
+    "llama": "llama3.1",
+    "llama4": "llama4-scout",
+    "mistral": "mistral",
+    "codellama": "codellama",
+    "deepseek": "deepseek-coder-v2",
+    "phi": "phi3",
+    "gemma": "gemma2",
+    "qwen": "qwen2.5-coder",
+}
+
+
+class ProviderKind(enum.Enum):
+    ANTHROPIC = "anthropic"
+    OPENAI = "openai"
+    OPENAI_CODEX = "openai_codex"  # /v1/responses endpoint
+    XAI = "xai"
+    OLLAMA = "ollama"
+
+
+def _is_codex_model(resolved: str) -> bool:
+    """Codex models require the /v1/responses endpoint."""
+    name = resolved.lower()
+    return "codex" in name
+
+
+def resolve_model_alias(model: str | None) -> str:
+    """Resolve short model aliases to full model IDs.
+
+    Handles Claude Code format like "opus[1m]", "sonnet[1m]".
+    Returns default model if None.
+    """
+    if not model:
+        return "claude-sonnet-4-6"
+    lower = model.lower().strip()
+
+    # Direct match
+    if lower in MODEL_ALIASES:
+        return MODEL_ALIASES[lower]
+
+    # Strip [context] suffix (e.g. "opus[1m]" -> "opus")
+    import re
+    stripped = re.sub(r"\[.*?\]$", "", lower).strip()
+    if stripped in MODEL_ALIASES:
+        return MODEL_ALIASES[stripped]
+
+    return model
+
+
+def detect_provider_kind(model: str) -> ProviderKind:
+    """Detect the provider from a model name."""
+    resolved = resolve_model_alias(model)
+    if resolved.startswith("claude-"):
+        return ProviderKind.ANTHROPIC
+    if resolved.startswith("grok-"):
+        return ProviderKind.XAI
+    # Codex models route to the Responses API, not Chat Completions
+    if _is_codex_model(resolved):
+        return ProviderKind.OPENAI_CODEX
+    if any(resolved.startswith(p) for p in ("gpt-", "o1", "o3", "o4", "codex", "gpt-5")):
+        return ProviderKind.OPENAI
+    if is_ollama_model(resolved):
+        return ProviderKind.OLLAMA
+    # Default to Anthropic
+    return ProviderKind.ANTHROPIC
+
+
+# Max tokens per model
+MAX_TOKENS_FOR_MODEL: dict[str, int] = {
+    "claude-opus-4-6": 32_000,
+    "claude-sonnet-4-6": 64_000,
+    "claude-haiku-4-5": 64_000,
+}
+
+DEFAULT_MAX_TOKENS = 16_000
+
+
+def max_tokens_for_model(model: str) -> int:
+    """Get the max output tokens for a model."""
+    resolved = resolve_model_alias(model)
+    return MAX_TOKENS_FOR_MODEL.get(resolved, DEFAULT_MAX_TOKENS)
+
+
+class ProviderClient:
+    """Unified provider client that dispatches to the correct backend.
+
+    Maps to: rust/crates/api/src/client.rs::ProviderClient
+    """
+
+    def __init__(
+        self,
+        kind: ProviderKind,
+        anthropic: AnthropicClient | None = None,
+        openai_compat: OpenAiCompatClient | None = None,
+        openai_responses: Any = None,  # OpenAiResponsesClient — Any to avoid import cycle
+        ollama: OllamaClient | None = None,
+    ) -> None:
+        self._kind = kind
+        self._anthropic = anthropic
+        self._openai_compat = openai_compat
+        self._openai_responses = openai_responses
+        self._ollama = ollama
+
+    @classmethod
+    def from_model(
+        cls,
+        model: str,
+        auth: AuthCredentials | None = None,
+    ) -> ProviderClient:
+        """Create a provider client based on the model name."""
+        resolved = resolve_model_alias(model)
+        kind = detect_provider_kind(resolved)
+
+        if kind == ProviderKind.ANTHROPIC:
+            if auth is not None:
+                client = AnthropicClient(auth=auth)
+            else:
+                client = AnthropicClient.from_env()
+            return cls(kind=kind, anthropic=client)
+
+        if kind == ProviderKind.XAI:
+            client = OpenAiCompatClient.from_env(OpenAiCompatConfig.xai())
+            return cls(kind=kind, openai_compat=client)
+
+        if kind == ProviderKind.OPENAI:
+            client = OpenAiCompatClient.from_env(OpenAiCompatConfig.openai())
+            return cls(kind=kind, openai_compat=client)
+
+        if kind == ProviderKind.OPENAI_CODEX:
+            from axion.api.openai_responses import OpenAiResponsesClient
+            responses_client = OpenAiResponsesClient.from_env()
+            return cls(kind=kind, openai_responses=responses_client)
+
+        if kind == ProviderKind.OLLAMA:
+            ollama_client = OllamaClient.from_env(model=resolved)
+            return cls(kind=kind, ollama=ollama_client)
+
+        raise ApiError(f"Provider {kind.value} not yet implemented")
+
+    @property
+    def provider_kind(self) -> ProviderKind:
+        return self._kind
+
+    async def send_message(self, request: MessageRequest) -> MessageResponse:
+        """Send a non-streaming message request."""
+        if self._anthropic is not None:
+            return await self._anthropic.send_message(request)
+        if self._openai_responses is not None:
+            return await self._openai_responses.send_message(request)
+        if self._openai_compat is not None:
+            return await self._openai_compat.send_message(request)
+        if self._ollama is not None:
+            return await self._ollama.send_message(request)
+        raise ApiError("No provider client configured")
+
+    async def stream_message(
+        self, request: MessageRequest
+    ) -> AsyncIterator[StreamEvent]:
+        """Send a streaming request and yield events."""
+        if self._anthropic is not None:
+            async for event in self._anthropic.stream_message(request):
+                yield event
+            return
+        if self._openai_responses is not None:
+            async for event in self._openai_responses.stream_message(request):
+                yield event
+            return
+        if self._openai_compat is not None:
+            async for event in self._openai_compat.stream_message(request):
+                yield event
+            return
+        if self._ollama is not None:
+            async for event in self._ollama.stream_message(request):
+                yield event
+            return
+        raise ApiError("No provider client configured")
+
+    async def close(self) -> None:
+        """Close underlying HTTP clients."""
+        if self._anthropic is not None:
+            await self._anthropic.close()
+        if self._openai_responses is not None:
+            await self._openai_responses.close()
+        if self._openai_compat is not None:
+            await self._openai_compat.close()
+        if self._ollama is not None:
+            await self._ollama.close()

axion/api/error.py

@@ -0,0 +1,161 @@
+"""API error types.
+
+Maps to: rust/crates/api/src/error.rs
+"""
+
+from __future__ import annotations
+
+CONTEXT_WINDOW_ERROR_MARKERS = [
+    "maximum context length",
+    "context window",
+    "context length",
+    "too many tokens",
+    "prompt is too long",
+    "input is too long",
+    "request is too large",
+]
+
+GENERIC_FATAL_WRAPPER_MARKERS = [
+    "something went wrong while processing your request",
+    "please try again, or use /new to start a fresh session",
+]
+
+
+class ApiError(Exception):
+    """Base class for all API errors."""
+
+    def is_retryable(self) -> bool:
+        return False
+
+    def request_id(self) -> str | None:
+        return None
+
+
+class MissingCredentialsError(ApiError):
+    """No API key or OAuth token available."""
+
+    def __init__(self, provider: str, env_vars: list[str]) -> None:
+        self.provider = provider
+        self.env_vars = env_vars
+        super().__init__(
+            f"Missing credentials for {provider}. "
+            f"Set one of: {', '.join(env_vars)}"
+        )
+
+
+class ContextWindowExceededError(ApiError):
+    """Request exceeds the model's context window."""
+
+    def __init__(
+        self,
+        model: str,
+        estimated_input_tokens: int,
+        requested_output_tokens: int,
+        estimated_total_tokens: int,
+        context_window_tokens: int,
+    ) -> None:
+        self.model = model
+        self.estimated_input_tokens = estimated_input_tokens
+        self.requested_output_tokens = requested_output_tokens
+        self.estimated_total_tokens = estimated_total_tokens
+        self.context_window_tokens = context_window_tokens
+        super().__init__(
+            f"Context window exceeded for {model}: "
+            f"{estimated_total_tokens} tokens > {context_window_tokens} limit"
+        )
+
+
+class ExpiredOAuthTokenError(ApiError):
+    """OAuth token has expired."""
+
+    def __init__(self) -> None:
+        super().__init__("OAuth token has expired")
+
+
+class AuthError(ApiError):
+    """Authentication failed."""
+
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class HttpError(ApiError):
+    """Low-level HTTP transport error."""
+
+    def __init__(self, message: str, cause: Exception | None = None) -> None:
+        self.cause = cause
+        super().__init__(message)
+
+    def is_retryable(self) -> bool:
+        return True
+
+
+class ApiResponseError(ApiError):
+    """API returned an error response."""
+
+    def __init__(
+        self,
+        status: int,
+        error_type: str | None = None,
+        message: str | None = None,
+        request_id_val: str | None = None,
+        body: str = "",
+        retryable: bool = False,
+    ) -> None:
+        self.status = status
+        self.error_type = error_type
+        self._message = message
+        self._request_id = request_id_val
+        self.body = body
+        self.retryable = retryable
+        detail = message or body[:200]
+        super().__init__(f"API error {status}: {detail}")
+
+    def is_retryable(self) -> bool:
+        return self.retryable
+
+    def request_id(self) -> str | None:
+        return self._request_id
+
+
+class RetriesExhaustedError(ApiError):
+    """All retry attempts failed."""
+
+    def __init__(self, attempts: int, last_error: ApiError) -> None:
+        self.attempts = attempts
+        self.last_error = last_error
+        super().__init__(f"Exhausted {attempts} retries. Last error: {last_error}")
+
+    def is_retryable(self) -> bool:
+        return self.last_error.is_retryable()
+
+    def request_id(self) -> str | None:
+        return self.last_error.request_id()
+
+
+class InvalidSseFrameError(ApiError):
+    """SSE frame could not be parsed."""
+
+    def __init__(self, reason: str) -> None:
+        super().__init__(f"Invalid SSE frame: {reason}")
+
+
+class BackoffOverflowError(ApiError):
+    """Backoff delay calculation overflowed."""
+
+    def __init__(self, attempt: int, base_delay_ms: int) -> None:
+        self.attempt = attempt
+        self.base_delay_ms = base_delay_ms
+        super().__init__(f"Backoff overflow at attempt {attempt}")
+
+
+def looks_like_context_window_error(message: str) -> bool:
+    """Check if an error message indicates a context window exceeded error."""
+    lower = message.lower()
+    return any(marker in lower for marker in CONTEXT_WINDOW_ERROR_MARKERS)
+
+
+def looks_like_generic_fatal_error(message: str) -> bool:
+    """Check if an error message indicates a generic fatal wrapper error."""
+    lower = message.lower()
+    return any(marker in lower for marker in GENERIC_FATAL_WRAPPER_MARKERS)