pi-ai-python 1.3.0__py3-none-any.whl

pi_ai/utils/models_error.py

@@ -0,0 +1,35 @@
+"""The error type raised by the model runtime (auth, provider lookup, catalogs)."""
+
+from __future__ import annotations
+
+from typing import Any, Literal, TypeAlias
+
+from .diagnostics import format_thrown_value
+
+__all__ = ["ModelsError", "ModelsErrorCode"]
+
+ModelsErrorCode: TypeAlias = Literal["model_source", "model_validation", "provider", "stream", "auth", "oauth"]
+
+
+class ModelsError(Exception):
+    """Runtime error with a machine-readable ``code``.
+
+    The message includes the cause's detail, because callers usually surface only
+    ``str(error)``.
+    """
+
+    def __init__(self, code: ModelsErrorCode, message: str, *, cause: Any = None) -> None:
+        super().__init__(_with_cause_detail(message, cause))
+        self.code: ModelsErrorCode = code
+        if isinstance(cause, BaseException):
+            self.__cause__ = cause
+        self.cause = cause
+
+
+def _with_cause_detail(message: str, cause: Any) -> str:
+    if cause is None:
+        return message
+    detail = format_thrown_value(cause).strip()
+    if not detail or detail in message:
+        return message
+    return f"{message}: {detail}"

pi_ai/utils/oauth_page.py

@@ -0,0 +1,121 @@
+"""HTML pages shown by the local OAuth callback servers."""
+
+from __future__ import annotations
+
+from string import Template
+
+__all__ = ["oauth_success_html", "oauth_error_html"]
+
+_LOGO_SVG = (
+    '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 800 800" aria-hidden="true">'
+    '<path fill="#F09082" d="M165.29 165.29H517.36V400H400V282.65H165.29Z"/>'
+    '<path fill="#4D9ABF" d="M165.29 282.65H282.65V400H400V517.36H282.65V634.72H165.29Z"/>'
+    '<path fill="#F1BE58" d="M517.36 400H634.72V634.72H517.36Z"/></svg>'
+)
+
+_PAGE = Template(
+    """<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="utf-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1" />
+  <title>$title</title>
+  <style>
+    :root {
+      --text: #fafafa;
+      --text-dim: #a1a1aa;
+      --page-bg: #09090b;
+      --font-sans: ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif, "Apple Color Emoji", "Segoe UI Emoji", "Segoe UI Symbol", "Noto Color Emoji";
+      --font-mono: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace;
+    }
+    * { box-sizing: border-box; }
+    html { color-scheme: dark; }
+    body {
+      margin: 0;
+      min-height: 100vh;
+      display: flex;
+      align-items: center;
+      justify-content: center;
+      padding: 24px;
+      background: var(--page-bg);
+      color: var(--text);
+      font-family: var(--font-sans);
+      text-align: center;
+    }
+    main {
+      width: 100%;
+      max-width: 560px;
+      display: flex;
+      flex-direction: column;
+      align-items: center;
+      justify-content: center;
+    }
+    .logo {
+      width: 72px;
+      height: 72px;
+      display: block;
+      margin-bottom: 24px;
+    }
+    h1 {
+      margin: 0 0 10px;
+      font-size: 28px;
+      line-height: 1.15;
+      font-weight: 650;
+      color: var(--text);
+    }
+    p {
+      margin: 0;
+      line-height: 1.7;
+      color: var(--text-dim);
+      font-size: 15px;
+    }
+    .details {
+      margin-top: 16px;
+      font-family: var(--font-mono);
+      font-size: 13px;
+      color: var(--text-dim);
+      white-space: pre-wrap;
+      word-break: break-word;
+    }
+  </style>
+</head>
+<body>
+  <main>
+    <div class="logo">$logo</div>
+    <h1>$heading</h1>
+    <p>$message</p>
+    $details
+  </main>
+</body>
+</html>"""
+)
+
+
+def _escape_html(value: str) -> str:
+    return (
+        value.replace("&", "&amp;")
+        .replace("<", "&lt;")
+        .replace(">", "&gt;")
+        .replace('"', "&quot;")
+        .replace("'", "&#39;")
+    )
+
+
+def _render_page(*, title: str, heading: str, message: str, details: str | None = None) -> str:
+    return _PAGE.substitute(
+        title=_escape_html(title),
+        heading=_escape_html(heading),
+        message=_escape_html(message),
+        logo=_LOGO_SVG,
+        details=f'<div class="details">{_escape_html(details)}</div>' if details else "",
+    )
+
+
+def oauth_success_html(message: str) -> str:
+    return _render_page(title="Authentication successful", heading="Authentication successful", message=message)
+
+
+def oauth_error_html(message: str, details: str | None = None) -> str:
+    return _render_page(
+        title="Authentication failed", heading="Authentication failed", message=message, details=details
+    )

pi_ai/utils/overflow.py

@@ -0,0 +1,87 @@
+"""Context overflow detection across providers.
+
+See the TypeScript source for the catalogue of provider error messages each
+pattern matches; the patterns are kept identical.
+"""
+
+from __future__ import annotations
+
+import re
+
+from ..types import AssistantMessage
+
+__all__ = ["is_context_overflow", "is_recoverable_length", "get_overflow_patterns"]
+
+_OVERFLOW_PATTERNS: list[re.Pattern[str]] = [
+    re.compile(pattern, re.IGNORECASE)
+    for pattern in (
+        r"prompt (?:is )?too long",  # Anthropic and z.ai token overflow
+        r"prompt exceeds max length",  # z.ai CN endpoint token overflow
+        r"request_too_large",  # Anthropic request byte-size overflow (HTTP 413)
+        r"input is too long for requested model",  # Amazon Bedrock
+        r"exceeds the context window",  # OpenAI (Completions & Responses API)
+        r"exceeds (?:the )?(?:model'?s )?maximum context length(?: of [\d,]+ tokens?|\s*\([\d,]+\))",  # LiteLLM
+        r"input token count.*exceeds the maximum",  # Google (Gemini)
+        r"maximum prompt length is \d+",  # xAI (Grok)
+        r"reduce the length of the messages",  # Groq
+        r"maximum context length is \d+ tokens",  # OpenRouter (most backends)
+        r"exceeds (?:the )?maximum allowed input length of [\d,]+ tokens?",  # OpenRouter/Poolside
+        r"input \(\d+ tokens\) is longer than the model'?s context length \(\d+ tokens\)",  # Together AI
+        r"exceeds the limit of \d+",  # GitHub Copilot
+        r"exceeds the available context size",  # llama.cpp server
+        r"greater than the context length",  # LM Studio
+        r"context window exceeds limit",  # MiniMax
+        r"exceeded model token limit",  # Kimi For Coding
+        r"too large for model with \d+ maximum context length",  # Mistral
+        r"prompt has [\d,]+ tokens?, but the configured context size is [\d,]+ tokens?",  # DS4 server
+        r"model_context_window_exceeded",  # z.ai non-standard finish_reason
+        r"prompt too long; exceeded (?:max )?context length",  # Ollama
+        r"range of input length should be",  # DashScope / Qwen Token Plan
+        r"context[_ ]length[_ ]exceeded",  # Generic fallback
+        r"too many tokens",  # Generic fallback
+        r"token limit exceeded",  # Generic fallback
+    )
+]
+
+_CEREBRAS_BODYLESS_OVERFLOW = re.compile(r"^4(?:00|13)\s*(?:status code)?\s*\(no body\)", re.IGNORECASE)
+
+_NON_OVERFLOW_PATTERNS: list[re.Pattern[str]] = [
+    re.compile(r"^(Throttling error|Service unavailable):", re.IGNORECASE),  # AWS Bedrock non-overflow errors
+    re.compile(r"rate limit", re.IGNORECASE),
+    re.compile(r"too many requests", re.IGNORECASE),
+]
+
+
+def is_context_overflow(message: AssistantMessage, context_window: int | None = None) -> bool:
+    """Whether an assistant message represents a context overflow.
+
+    Handles error-message overflow (most providers), silent overflow where
+    ``usage.input`` exceeds the context window (z.ai), and length-stop overflow with
+    zero output (Xiaomi MiMo). Pass ``context_window`` for the latter two.
+    """
+    error_message = message.error_message
+    if message.stop_reason == "error" and error_message:
+        if not any(pattern.search(error_message) for pattern in _NON_OVERFLOW_PATTERNS):
+            if any(pattern.search(error_message) for pattern in _OVERFLOW_PATTERNS):
+                return True
+            if message.provider == "cerebras" and _CEREBRAS_BODYLESS_OVERFLOW.search(error_message):
+                return True
+
+    if context_window and message.stop_reason == "stop":
+        if message.usage.input + message.usage.cache_read > context_window:
+            return True
+
+    if context_window and message.stop_reason == "length" and message.usage.output == 0:
+        if message.usage.input + message.usage.cache_read >= context_window * 0.99:
+            return True
+
+    return False
+
+
+def is_recoverable_length(message: AssistantMessage, desired_max_output: int) -> bool:
+    """Whether a length stop ended below the intended output limit (context pressure)."""
+    return message.stop_reason == "length" and desired_max_output > 0 and message.usage.output < desired_max_output
+
+
+def get_overflow_patterns() -> list[re.Pattern[str]]:
+    return list(_OVERFLOW_PATTERNS)

pi_ai/utils/retry.py

@@ -0,0 +1,193 @@
+"""Retry classification and a policy-driven retry loop for assistant calls."""
+
+from __future__ import annotations
+
+import dataclasses
+import re
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from typing import Any
+
+from ..abort import AbortError, AbortSignal, sleep
+from ..types import AssistantMessage
+from .callbacks import maybe_await
+
+__all__ = [
+    "RetryPolicy",
+    "RetryCallbacks",
+    "DEFAULT_MAX_AGENT_RETRY_DELAY_MS",
+    "retry_delay_ms",
+    "retry_assistant_call",
+    "is_retryable_assistant_error",
+]
+
+
+def _pattern(parts: list[str]) -> re.Pattern[str]:
+    return re.compile("|".join(parts), re.IGNORECASE)
+
+
+_NON_RETRYABLE = _pattern(
+    [
+        "GoUsageLimitError",
+        "FreeUsageLimitError",
+        "Monthly usage limit reached",
+        "available balance",
+        "insufficient_quota",
+        "out of budget",
+        "quota exceeded",
+        "billing",
+        "subscription_sharing_usage_limit_exceeded",
+    ]
+)
+
+_RETRYABLE = _pattern(
+    [
+        "overloaded",
+        "server_busy",
+        "servers are currently busy",
+        "currently experiencing high demand",
+        "model is at capacity",
+        "rate.?limit",
+        "too many requests",
+        "429",
+        "500",
+        "502",
+        "503",
+        "504",
+        "520",
+        "524",
+        r"service.?(?:is.?)?(?:temporarily.?)?unavailable",
+        "server.?error",
+        "internal.?error",
+        r"provider.?returned.?(?:an.?)?error",
+        "exceeded request buffer limit while retrying upstream",
+        "unable to forward upstream",
+        "upstream.?request.?failed",
+        "upstream_stream_read_error",
+        "stream was interrupted",
+        "http/?2 stream failed",
+        "network.?error",
+        "connection.?error",
+        "connection.?refused",
+        "connection.?lost",
+        "other side closed",
+        "fetch failed",
+        r"request to \S+ failed, reason:",
+        "getaddrinfo",
+        "ENOTFOUND",
+        "EAI_AGAIN",
+        "upstream.?connect",
+        "reset before headers",
+        "socket hang up",
+        "socket connection was closed",
+        "timed? out",
+        "timeout",
+        "terminated",
+        "websocket.?closed",
+        "websocket.?error",
+        "ended without",
+        "stream ended before message_stop",
+        "stream ended before a terminal response event",
+        "http2 request did not get a response",
+        "pending stream has been canceled",
+        "retry delay",
+        "you can retry your request",
+        "try your request again",
+        "please retry your request",
+        "ResourceExhausted",
+        "subscription_sharing_usage_unavailable",
+        "subscription_sharing_user_unavailable",
+        # Python transport errors (httpx/httpcore wording)
+        "remote ?protocol ?error",
+        "server disconnected",
+        "connection reset",
+        "read ?error",
+        "connect ?error",
+    ]
+)
+
+DEFAULT_MAX_AGENT_RETRY_DELAY_MS = 60_000
+
+
+@dataclass(slots=True)
+class RetryPolicy:
+    """Bounded attempts with exponential backoff (``base_delay_ms * 2^(attempt-1)``)."""
+
+    enabled: bool = True
+    #: Max retry attempts (0 = no retries). The initial call never counts as a retry.
+    max_retries: int = 3
+    base_delay_ms: float = 2000
+    #: Cap for each computed delay. Defaults to 60 seconds.
+    max_agent_delay_ms: float | None = None
+
+
+def retry_delay_ms(policy: RetryPolicy, attempt: int) -> float:
+    delay = policy.base_delay_ms * 2 ** max(0, attempt - 1)
+    cap = policy.max_agent_delay_ms if policy.max_agent_delay_ms is not None else DEFAULT_MAX_AGENT_RETRY_DELAY_MS
+    return min(delay, cap)
+
+
+@dataclass(slots=True)
+class RetryCallbacks:
+    #: Before the backoff sleep of each retry attempt (1-indexed): (attempt, max_attempts, delay_ms, error_message).
+    on_retry_scheduled: Callable[[int, int, float, str], Any] | None = None
+    #: After the backoff sleep, immediately before the retried call starts.
+    on_retry_attempt_start: Callable[[], Any] | None = None
+    #: Once when the loop ends: (success, attempt, final_error).
+    on_retry_finished: Callable[[bool, int, str | None], Any] | None = None
+
+
+async def retry_assistant_call(
+    produce: Callable[[], Awaitable[AssistantMessage]],
+    policy: RetryPolicy | None,
+    signal: AbortSignal | None = None,
+    callbacks: RetryCallbacks | None = None,
+) -> AssistantMessage:
+    """Run an assistant-producing call with bounded retry on transient errors.
+
+    Aborts are terminal and never retried; aborts during the backoff sleep are
+    normalized to an aborted message. Non-retryable errors (quota, billing) return
+    immediately.
+    """
+    max_attempts = policy.max_retries if policy is not None and policy.enabled else 0
+    callbacks = callbacks or RetryCallbacks()
+    attempt = 0
+    last_retry: tuple[int, str] | None = None
+    while True:
+        response = await produce()
+        if response.stop_reason == "aborted":
+            if last_retry and callbacks.on_retry_finished:
+                await maybe_await(callbacks.on_retry_finished(False, last_retry[0], None))
+            return response
+        if response.stop_reason != "error":
+            if last_retry and callbacks.on_retry_finished:
+                await maybe_await(callbacks.on_retry_finished(True, last_retry[0], None))
+            return response
+        if attempt >= max_attempts or not is_retryable_assistant_error(response):
+            if last_retry and callbacks.on_retry_finished:
+                await maybe_await(callbacks.on_retry_finished(False, last_retry[0], response.error_message))
+            return response
+
+        attempt += 1
+        last_retry = (attempt, response.error_message or "Unknown error")
+        assert policy is not None
+        delay = retry_delay_ms(policy, attempt)
+        if callbacks.on_retry_scheduled:
+            await maybe_await(callbacks.on_retry_scheduled(attempt, max_attempts, delay, last_retry[1]))
+        try:
+            await sleep(delay / 1000, signal)
+        except AbortError:
+            if callbacks.on_retry_finished:
+                await maybe_await(callbacks.on_retry_finished(False, attempt, last_retry[1]))
+            return dataclasses.replace(response, stop_reason="aborted", error_message=None)
+        if callbacks.on_retry_attempt_start:
+            await maybe_await(callbacks.on_retry_attempt_start())
+
+
+def is_retryable_assistant_error(message: AssistantMessage) -> bool:
+    """Whether a failed assistant message looks like a transient provider/transport error."""
+    if message.stop_reason != "error" or not message.error_message:
+        return False
+    if _NON_RETRYABLE.search(message.error_message):
+        return False
+    return _RETRYABLE.search(message.error_message) is not None