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

pi_ai/api/openai_sdk.py

@@ -0,0 +1,204 @@
+"""Shared OpenAI SDK plumbing for the adapters built on the official ``openai`` package.
+
+Upstream pi creates an ``OpenAI`` (or ``AzureOpenAI``) client per request, calls
+``create(params, { maxRetries: 0 }).withResponse()`` and wraps that call in
+``retryProviderRequest`` so the backoff sleep stays abortable. The helpers here do
+the same with ``AsyncOpenAI`` / ``AsyncAzureOpenAI`` and translate SDK exceptions
+into :class:`~pi_ai.utils.http.ProviderHttpError` with the JS SDK's message wording,
+which pi's retry and context-overflow detection match against.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import AsyncIterator, Awaitable, Callable, Mapping
+from typing import Any, TypeVar, cast
+
+import httpx
+import openai
+from openai import AsyncAzureOpenAI, AsyncOpenAI, AsyncStream, omit
+
+from ..utils.http import (
+    DEFAULT_TIMEOUT_MS,
+    ProviderHttpError,
+    get_http_client,
+    make_status_error,
+    retry_provider_request,
+)
+from ..utils.misc import get_provider_env_value
+
+__all__ = [
+    "create_openai_client",
+    "merge_sdk_headers",
+    "sdk_body_arguments",
+    "openai_request",
+    "open_sdk_stream",
+    "iter_sdk_stream",
+    "status_error_from_sdk",
+    "stream_error_from_sdk",
+]
+
+_T = TypeVar("_T")
+
+#: ``AsyncStream`` yielding the decoded JSON of each event. Like the JS SDK, which hands
+#: back the parsed objects as-is; building pydantic models for every event would cost more
+#: than pi's own event processing.
+_RawStream = AsyncStream[object]
+
+
+def merge_sdk_headers(*sources: Mapping[str, str | None] | None) -> dict[str, str | None]:
+    """Merge header sources case-insensitively; later sources win.
+
+    ``None`` values are kept: they remove the header, including the SDK's own defaults
+    (``Authorization``, ``User-Agent``, ...), like ``null`` in the JS SDK's ``defaultHeaders``.
+    """
+    merged: dict[str, tuple[str, str | None]] = {}
+    for source in sources:
+        if not source:
+            continue
+        for name, value in source.items():
+            lower = name.lower()
+            merged.pop(lower, None)
+            merged[lower] = (name, value)
+    return {name: value for name, value in merged.values()}
+
+
+def create_openai_client(
+    *,
+    api_key: str,
+    base_url: str | None,
+    headers: Mapping[str, str | None],
+    options: Mapping[str, Any],
+    azure_api_version: str | None = None,
+) -> AsyncOpenAI:
+    """An ``AsyncOpenAI`` client (``AsyncAzureOpenAI`` with ``azure_api_version``) for one request.
+
+    ``None`` header values are left out here and removed per request (see
+    :func:`sdk_body_arguments`). Retries are left to :func:`openai_request`. The transport is ``options["http_client"]`` or
+    pi's pooled client for the running loop, so connections are reused across requests.
+    Construction is cheap (tens of microseconds), so clients are not cached.
+    """
+    timeout_ms = options.get("timeout_ms")
+    env = options.get("env")
+    default_headers = {name: value for name, value in headers.items() if value is not None}
+    http_client = cast(Any, get_http_client(base_url or "", options))
+    common: dict[str, Any] = {
+        "api_key": api_key,
+        "organization": get_provider_env_value("OPENAI_ORG_ID", env),
+        "project": get_provider_env_value("OPENAI_PROJECT_ID", env),
+        # Same overall budget as the JS SDK default (10 minutes); ``read`` is the idle gap
+        # between chunks, so an active stream is never cut off.
+        "timeout": (timeout_ms if timeout_ms is not None else DEFAULT_TIMEOUT_MS) / 1000,
+        "max_retries": 0,
+        "default_headers": default_headers,
+        "http_client": http_client,
+    }
+    if azure_api_version is not None:
+        return AsyncAzureOpenAI(api_version=azure_api_version, base_url=base_url or "", **common)
+    return AsyncOpenAI(base_url=base_url or None, **common)
+
+
+def sdk_body_arguments(
+    params: Mapping[str, Any], *required: str, headers: Mapping[str, str | None] | None = None
+) -> dict[str, Any]:
+    """``create()`` keyword arguments that send ``params`` as the request body unchanged.
+
+    pi builds (and ``on_payload`` may replace) the body in wire format, including fields the
+    SDK does not declare, so it goes through ``extra_body``. ``required`` parameters are
+    passed as ``omit`` so the body keys come only from ``params``. This also skips the SDK's
+    parameter transform, which walks the whole transcript on every request.
+
+    Headers that ``headers`` (the client's header record) maps to ``None`` are omitted per
+    request: the Python SDK only accepts a removed ``Authorization`` header there.
+    """
+    arguments: dict[str, Any] = dict.fromkeys(required, omit)
+    arguments["extra_body"] = params
+    if headers:
+        removed = {name: omit for name, value in headers.items() if value is None}
+        if removed:
+            arguments["extra_headers"] = removed
+    return arguments
+
+
+def _legacy_headers(headers: Any) -> httpx.Headers:
+    return headers if isinstance(headers, httpx.Headers) else httpx.Headers(list(headers.multi_items()))
+
+
+def status_error_from_sdk(error: openai.APIStatusError) -> ProviderHttpError:
+    """``"<status> <error.message>"`` from the raw body, as the JS SDK formats ``APIError``."""
+    response = error.response
+    try:
+        text = response.text
+    except Exception:  # body could not be read
+        text = ""
+    return make_status_error(error.status_code, text, _legacy_headers(response.headers), "openai")
+
+
+def stream_error_from_sdk(error: openai.APIError, headers: Any) -> ProviderHttpError:
+    """The JS SDK's ``APIError`` for a stream event carrying an ``error`` field."""
+    body = error.body
+    message = body.get("message") if isinstance(body, dict) else None
+    if message:
+        text = message if isinstance(message, str) else _json_compact(message)
+    elif body:
+        text = _json_compact(body)
+    else:
+        text = "(no status code or body)"
+    return ProviderHttpError(text, headers=_legacy_headers(headers), error=body)
+
+
+def _json_compact(value: Any) -> str:
+    try:
+        return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
+    except (TypeError, ValueError):
+        return str(value)
+
+
+async def openai_request(call: Callable[[], Awaitable[_T]], options: Mapping[str, Any]) -> _T:
+    """Run an SDK request with pi's abortable, SDK-compatible retry policy.
+
+    HTTP errors become :class:`ProviderHttpError` (status, headers, parsed ``error`` body) and
+    connection failures a status-less one (``"Connection error."``, ``"Request timed out."``).
+    """
+
+    async def attempt() -> _T:
+        try:
+            return await call()
+        except openai.APIStatusError as error:
+            raise status_error_from_sdk(error) from error
+        except openai.APIConnectionError as error:
+            raise ProviderHttpError(error.message) from error
+
+    return await retry_provider_request(
+        attempt,
+        max_retries=options.get("max_retries") or 0,
+        max_retry_delay_ms=options.get("max_retry_delay_ms"),
+        signal=options.get("signal"),
+    )
+
+
+# Streams are decoded by the SDK, as upstream iterates the JS SDK's ``Stream``. The SDK's SSE
+# decoder is the throughput limit of these adapters: about 70k events/s in
+# ``benchmarks/bench_stream.py`` versus about 120k with pi's ``iter_sse`` + ``fastjson``.
+# Two ways to recover that speed, neither taken:
+# - overriding ``AsyncOpenAI._make_sse_decoder`` relies on a private SDK method that can change
+#   in any release;
+# - parsing ``raw.http_response`` with pi's own SSE parser uses only public API (the Anthropic
+#   adapter does this), but stops using the SDK's stream decoding.
+# Real model output arrives far slower than either rate, so per-event cost rarely matters.
+def open_sdk_stream(raw: Any) -> AsyncStream[Any]:
+    """The event stream of a ``with_raw_response`` streaming call, yielding plain JSON."""
+    return raw.parse(to=_RawStream)
+
+
+async def iter_sdk_stream(stream: AsyncStream[Any]) -> AsyncIterator[Any]:
+    """Iterate SDK stream events; ``error`` events raise like the JS SDK's ``Stream``."""
+    try:
+        async for event in stream:
+            yield event
+    except (openai.APIStatusError, openai.APIConnectionError):
+        raise
+    except openai.APIError as error:
+        raise stream_error_from_sdk(error, stream.response.headers) from error
+    finally:
+        await stream.close()

pi_ai/api/openrouter_images.py

@@ -0,0 +1,153 @@
+"""Image generation over OpenRouter's chat completions endpoint."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping
+from typing import Any
+
+from ..abort import AbortError, abort_scope
+from ..types import (
+    AssistantImages,
+    ImageContent,
+    ImageModel,
+    ImagesContext,
+    ProviderResponse,
+    TextContent,
+    Usage,
+    UsageCost,
+    now_ms,
+)
+from ..utils import fastjson
+from ..utils.callbacks import maybe_await
+from ..utils.error_body import format_provider_error, normalize_provider_error
+from ..utils.misc import headers_to_record, provider_headers_to_record, sanitize_surrogates
+from ..utils.model_operations import image_error_result
+from .openai_sdk import create_openai_client, openai_request, sdk_body_arguments
+
+__all__ = ["generate_images"]
+
+_DATA_URL = re.compile(r"data:([^;]+);base64,(.+)")
+
+
+async def generate_images(
+    model: ImageModel, context: ImagesContext, options: Mapping[str, Any] | None = None
+) -> AssistantImages:
+    """Generate images with an OpenRouter image model. Never raises."""
+    options = options or {}
+    output = AssistantImages(
+        api=model.api, provider=model.provider, model=model.id, output=[], stop_reason="stop", timestamp=now_ms()
+    )
+    signal = options.get("signal")
+    try:
+        api_key = options.get("api_key")
+        if not api_key:
+            raise ValueError(f"No API key for provider: {model.provider}")
+        client = create_openai_client(
+            api_key=api_key,
+            base_url=model.base_url,
+            headers=provider_headers_to_record(model.headers, options.get("headers")) or {},
+            options=options,
+        )
+        params: Any = _build_params(model, context)
+        on_payload = options.get("on_payload")
+        if on_payload is not None:
+            next_params = await maybe_await(on_payload(params, model))
+            if next_params is not None:
+                params = next_params
+        create = client.chat.completions.with_raw_response.create
+        body_arguments = sdk_body_arguments(params, "model", "messages")
+        try:
+            async with abort_scope(signal):
+                response = await openai_request(lambda: create(**body_arguments), options)
+        except AbortError as error:
+            raise AbortError("Request aborted") from error
+        on_response = options.get("on_response")
+        if on_response is not None:
+            await maybe_await(
+                on_response(ProviderResponse(response.status_code, headers_to_record(response.headers)), model)
+            )
+
+        body = fastjson.loads(response.content)
+        if not isinstance(body, Mapping):
+            raise ValueError("OpenRouter returned an unexpected image generation response")
+        output.response_id = body.get("id")
+        raw_usage = body.get("usage")
+        if isinstance(raw_usage, Mapping):
+            output.usage = _parse_usage(raw_usage, model)
+
+        choices = body.get("choices")
+        if not isinstance(choices, list):
+            raise ValueError("OpenRouter returned an image generation response without choices")
+        choice = choices[0] if choices else None
+        if isinstance(choice, Mapping):
+            message = choice.get("message") or {}
+            content = message.get("content")
+            if isinstance(content, str) and content:
+                output.output.append(TextContent(content))
+            for image in message.get("images") or ():
+                image_url = image.get("image_url") if isinstance(image, Mapping) else None
+                url = image_url.get("url") if isinstance(image_url, Mapping) else image_url
+                if not isinstance(url, str) or not url.startswith("data:"):
+                    continue
+                matches = _DATA_URL.fullmatch(url)
+                if matches is None:
+                    continue
+                output.output.append(ImageContent(data=matches.group(2), mime_type=matches.group(1)))
+        return output
+    except Exception as error:
+        failed = image_error_result(
+            model,
+            format_provider_error(normalize_provider_error(error)),
+            aborted=signal is not None and signal.aborted,
+        )
+        failed.response_id = output.response_id
+        failed.usage = output.usage
+        failed.timestamp = output.timestamp
+        return failed
+
+
+def _build_params(model: ImageModel, context: ImagesContext) -> dict[str, Any]:
+    """The chat completions request body for an image generation."""
+    content: list[dict[str, Any]] = []
+    for item in context.input:
+        if item.type == "text":
+            content.append({"type": "text", "text": sanitize_surrogates(item.text)})  # type: ignore[union-attr]
+        else:
+            content.append(
+                {"type": "image_url", "image_url": {"url": f"data:{item.mime_type};base64,{item.data}"}}  # type: ignore[union-attr]
+            )
+    return {
+        "model": model.id,
+        "messages": [{"role": "user", "content": content}],
+        "stream": False,
+        "modalities": ["image", "text"] if "text" in model.output else ["image"],
+    }
+
+
+def _parse_usage(raw_usage: Mapping[str, Any], model: ImageModel) -> Usage:
+    """Usage from an OpenAI-style ``usage`` object, priced from the model's base rates."""
+    details = raw_usage.get("prompt_tokens_details") or {}
+    prompt_tokens = raw_usage.get("prompt_tokens") or 0
+    reported_cached_tokens = details.get("cached_tokens") or 0
+    cache_write_tokens = details.get("cache_write_tokens") or 0
+    cache_read_tokens = (
+        max(0, reported_cached_tokens - cache_write_tokens) if cache_write_tokens > 0 else reported_cached_tokens
+    )
+    input_tokens = max(0, prompt_tokens - cache_read_tokens - cache_write_tokens)
+    output_tokens = raw_usage.get("completion_tokens") or 0
+    cost = UsageCost(
+        input=(model.cost.input / 1000000) * input_tokens,
+        output=(model.cost.output / 1000000) * output_tokens,
+        cache_read=(model.cost.cache_read / 1000000) * cache_read_tokens,
+        cache_write=(model.cost.cache_write / 1000000) * cache_write_tokens,
+    )
+    cost.total = cost.input + cost.output + cost.cache_read + cost.cache_write
+    return Usage(
+        input_tokens,
+        output_tokens,
+        cache_read_tokens,
+        cache_write_tokens,
+        input_tokens + output_tokens + cache_read_tokens + cache_write_tokens,
+        cost,
+    )