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

pi_durable/providers/retry.py

@@ -0,0 +1,127 @@
+"""Provider-level retries inside one request attempt, shared by the model adapters.
+
+Port of pi-ai 1.3.0 ``utils/provider-retry.ts``. Both SDKs are called with their own retries
+disabled (``max_retries=0``) and wrapped here instead, so a server-requested delay is interruptible
+and a delay above ``maxRetryDelayMs`` fails at once.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import random
+import time
+from collections.abc import Awaitable, Callable
+from email.utils import parsedate_to_datetime
+from typing import Any, TypeVar
+
+from ..context import AbortSignal
+from .stream import ABORT_MESSAGE, error_headers, error_status
+
+T = TypeVar("T")
+
+DEFAULT_MAX_RETRY_DELAY_MS = 60_000
+RETRYABLE_STATUSES = frozenset({408, 409, 429})
+
+
+def is_retryable_provider_error(error: BaseException) -> bool:
+    """The SDK retry policy: an explicit header first, then status, then a transport failure."""
+
+    should_retry = _header(error_headers(error), "x-should-retry")
+    if should_retry == "true":
+        return True
+    if should_retry == "false":
+        return False
+    status = error_status(error)
+    if status is None:
+        return True
+    return status in RETRYABLE_STATUSES or status >= 500
+
+
+def provider_retry_delay_ms(
+    error: BaseException, retry_index: int, max_retry_delay_ms: int | None
+) -> int:
+    """How long to wait before the next provider attempt."""
+
+    headers = error_headers(error)
+    retry_after_ms = _header(headers, "retry-after-ms")
+    if retry_after_ms is not None:
+        delay = float(retry_after_ms)
+        _validate_delay(delay, max_retry_delay_ms, str(error))
+        return max(0, int(delay))
+    retry_after = _header(headers, "retry-after")
+    if retry_after is not None:
+        delay = _retry_after_ms(retry_after)
+        _validate_delay(delay, max_retry_delay_ms, str(error))
+        return max(0, int(delay))
+    exponential = min(0.5 * 2**retry_index, 8) * 1000
+    return int(exponential * (1 - random.random() * 0.25))
+
+
+def _retry_after_ms(value: str) -> float:
+    """A ``Retry-After`` header, either delta-seconds or an HTTP date."""
+
+    try:
+        return float(value) * 1000
+    except ValueError:
+        pass
+    try:
+        retry_at = parsedate_to_datetime(value)
+    except (TypeError, ValueError):
+        return 0.0
+    if retry_at is None:
+        return 0.0
+    return (retry_at.timestamp() - time.time()) * 1000
+
+
+async def retry_provider_request(
+    request: Callable[[], Awaitable[T]],
+    *,
+    max_retries: int = 0,
+    max_retry_delay_ms: int | None = None,
+    signal: AbortSignal | None = None,
+) -> T:
+    """Run ``request``, retrying retryable provider failures while attempts remain."""
+
+    remaining = max(0, max_retries)
+    for attempt in range(remaining + 1):
+        if signal is not None and signal.aborted:
+            raise asyncio.CancelledError(ABORT_MESSAGE)
+        try:
+            return await request()
+        except Exception as error:
+            if signal is not None and signal.aborted:
+                raise
+            if attempt >= remaining or not is_retryable_provider_error(error):
+                raise
+            delay = provider_retry_delay_ms(error, attempt, max_retry_delay_ms)
+            await asyncio.sleep(delay / 1000.0)
+    raise RuntimeError("unreachable")  # pragma: no cover
+
+
+def _validate_delay(delay_ms: float, max_retry_delay_ms: int | None, message: str) -> None:
+    limit = DEFAULT_MAX_RETRY_DELAY_MS if max_retry_delay_ms is None else max_retry_delay_ms
+    if limit > 0 and delay_ms > limit:
+        raise RuntimeError(
+            f"Server requested {delay_ms / 1000:.0f}s retry delay (max: {limit / 1000:.0f}s). {message}"
+        )
+
+
+def _header(headers: Any, name: str) -> str | None:
+    if headers is None:
+        return None
+    get = getattr(headers, "get", None)
+    if get is None:
+        return None
+    value: Any = get(name)
+    if value is None:
+        value = get(name.title())
+    return None if value is None else str(value)
+
+
+__all__ = [
+    "DEFAULT_MAX_RETRY_DELAY_MS",
+    "RETRYABLE_STATUSES",
+    "is_retryable_provider_error",
+    "provider_retry_delay_ms",
+    "retry_provider_request",
+]

pi_durable/providers/stream.py

@@ -0,0 +1,183 @@
+"""Async plumbing shared by the model adapters.
+
+Each adapter turns one SDK request into a pi-ai style stream: partial events while the provider
+responds, then a terminal assistant message. :class:`ProviderStream` owns the queue, the terminal
+future, and the abort wiring so both adapters behave the same when a request fails or is cancelled.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence
+from typing import Any
+
+from ..context import AbortSignal
+
+ABORT_MESSAGE = "Request was aborted"
+
+
+def provider_field(value: Any, key: str, default: Any = None) -> Any:
+    """Read one field from an SDK object or a mapping."""
+
+    if isinstance(value, Mapping):
+        return value.get(key, default)
+    return getattr(value, key, default)
+
+
+def provider_mapping(value: Any) -> dict[str, Any]:
+    """Read a nested mapping from an SDK object, a mapping, or a value with ``model_dump()``."""
+
+    if isinstance(value, Mapping):
+        return dict(value)
+    dump = getattr(value, "model_dump", None)
+    if callable(dump):
+        dumped = dump()
+        return dict(dumped) if isinstance(dumped, Mapping) else {}
+    return dict(vars(value)) if hasattr(value, "__dict__") else {}
+
+
+def content_index(content: Sequence[Any], block: Any) -> int:
+    """The index of ``block`` in ``content``, matched by identity rather than by value."""
+
+    for index, candidate in enumerate(content):
+        if candidate is block:
+            return index
+    raise ValueError("block is not part of the message content")
+
+
+def error_status(error: BaseException) -> int | None:
+    """The HTTP status an SDK error carries, if any."""
+
+    status = provider_field(error, "status_code")
+    if status is None:
+        status = provider_field(error, "status")
+    return status if isinstance(status, int) else None
+
+
+def error_headers(error: BaseException) -> Any:
+    """The response headers an SDK error carries, if any."""
+
+    response = provider_field(error, "response")
+    headers = provider_field(response, "headers") if response is not None else None
+    return headers if headers is not None else provider_field(error, "headers")
+
+
+def format_provider_error(error: BaseException) -> str:
+    """The text a failed adapter attempt records, with the HTTP status when the SDK exposes it."""
+
+    message = str(error) or type(error).__name__
+    status = error_status(error)
+    if status is not None and str(status) not in message:
+        return f"{status}: {message}"
+    return message
+
+
+def now_ms() -> int:
+    """A Unix timestamp in milliseconds, as pi-ai's ``Date.now()``."""
+
+    return int(time.time() * 1000)
+
+
+async def complete_stream(stream: Any) -> dict[str, Any]:
+    """Consume a stream and return its terminal message."""
+
+    async for _event in stream:
+        pass
+    return await stream.result()
+
+
+class ProviderStream:
+    """A pi-ai style stream over one asynchronous producer.
+
+    Subclasses implement :meth:`_produce`, pushing partial events with :meth:`push` and returning
+    once the terminal message is complete. Every ``BaseException`` becomes an ``error`` event and a
+    terminal message with ``stopReason`` ``error`` or, when the invocation was signalled, ``aborted``.
+    """
+
+    def __init__(self, output: dict[str, Any], *, signal: AbortSignal | None = None) -> None:
+        self.output = output
+        self._signal = signal
+        self._queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue()
+        loop = asyncio.get_running_loop()
+        self._result: asyncio.Future[dict[str, Any]] = loop.create_future()
+        self._task: asyncio.Task[None] = asyncio.ensure_future(self._run())
+
+    # ── Stream surface ─────────────────────────────────────────────────────
+
+    def push(self, event: dict[str, Any]) -> None:
+        self._queue.put_nowait(event)
+
+    def emit(self, kind: str, *, content_index: int | None = None, **fields: Any) -> None:
+        """Push one assistant-message event carrying the live partial and any extra fields."""
+
+        event: dict[str, Any] = {"type": kind, "partial": self.output, **fields}
+        if content_index is not None:
+            event["contentIndex"] = content_index
+        self.push(event)
+
+    def result(self) -> Awaitable[dict[str, Any]]:
+        return self._result
+
+    def __aiter__(self) -> AsyncIterator[dict[str, Any]]:
+        return self._iterate()
+
+    async def _iterate(self) -> AsyncIterator[dict[str, Any]]:
+        try:
+            while True:
+                event = await self._queue.get()
+                if event is None:
+                    return
+                yield event
+        finally:
+            # A consumer that stops early cannot leave the provider request running.
+            if not self._task.done():
+                self._task.cancel()
+
+    # ── Producer plumbing ─────────────────────────────────────────────────
+
+    async def _produce(self) -> None:
+        raise NotImplementedError
+
+    async def _run(self) -> None:
+        remove: Any = None
+        try:
+            if self._signal is not None:
+                if self._signal.aborted:
+                    raise asyncio.CancelledError
+                remove = self._signal.add_listener(self._cancel)
+            await self._produce()
+        except asyncio.CancelledError:
+            self._fail("aborted", ABORT_MESSAGE)
+        except Exception as error:
+            aborted = self._signal is not None and self._signal.aborted
+            self._fail("aborted" if aborted else "error", format_provider_error(error))
+        finally:
+            if remove is not None:
+                remove()
+            if not self._result.done():
+                self._result.set_result(self.output)
+            self._queue.put_nowait(None)
+
+    def _cancel(self) -> None:
+        if not self._task.done():
+            self._task.cancel()
+
+    def _fail(self, stop_reason: str, message: str) -> None:
+        self.output["stopReason"] = stop_reason
+        self.output["errorMessage"] = message
+        self.emit("error", reason=stop_reason, error=self.output)
+
+
+__all__ = [
+    "ABORT_MESSAGE",
+    "ProviderStream",
+    "complete_stream",
+    "content_index",
+    "error_headers",
+    "error_status",
+    "format_provider_error",
+    "now_ms",
+    "provider_field",
+    "provider_mapping",
+]

pi_durable/providers/transcript.py

@@ -0,0 +1,224 @@
+"""Provider-facing transcript replay shared by the model adapters.
+
+Port of pi-ai 1.3.0 ``utils/transcript.ts`` and the cross-model parts of
+``api/transform-messages.ts``. Durable hands every provider a pi-ai-shaped message list whose system
+messages carry the prompt and tool changes. Adapters collapse those deltas into one leading system
+message (pi-ai's default for providers without mid-conversation system support), replay tools, drop
+thinking signed by another model, and replace images a model cannot take.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any, cast
+
+from . import Model
+
+NON_VISION_USER_IMAGE_PLACEHOLDER = "(image omitted: model does not support images)"
+NON_VISION_TOOL_IMAGE_PLACEHOLDER = "(tool image omitted: model does not support images)"
+
+
+def content_text(content: Any, separator: str = "\n") -> str:
+    """Extract and join the text of a message content."""
+
+    if isinstance(content, str):
+        return content
+    return separator.join(
+        cast(str, block.get("text") or "")
+        for block in cast(Sequence[Mapping[str, Any]], content or ())
+        if block.get("type") == "text"
+    )
+
+
+def system_message_text(message: Mapping[str, Any]) -> str:
+    """Render a system message as a complete prompt: its content followed by its sections."""
+
+    parts = [content_text(message.get("content"))]
+    for value in cast(Mapping[str, str | None], message.get("sections") or {}).values():
+        if value is not None:
+            parts.append(value)
+    return "\n\n".join(part for part in parts if part)
+
+
+def current_tools(messages: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]:
+    """Resolve the tools available after applying every transcript delta in order."""
+
+    tools: dict[str, dict[str, Any]] = {}
+    for message in messages:
+        if message.get("role") != "system":
+            continue
+        for reference in cast(Sequence[Mapping[str, Any]], message.get("toolsRemoved") or ()):
+            tools.pop(cast(str, reference["name"]), None)
+        for tool in cast(Sequence[Mapping[str, Any]], message.get("toolsAdded") or ()):
+            tools[cast(str, tool["name"])] = dict(tool)
+    return list(tools.values())
+
+
+def current_system_message(messages: Sequence[Mapping[str, Any]]) -> dict[str, Any] | None:
+    """Replay every system message into one leading message holding the prompt and tools."""
+
+    contents: list[str] = []
+    sections: dict[str, str] = {}
+    timestamp: int | None = None
+    found = False
+    for message in messages:
+        if message.get("role") != "system":
+            continue
+        found = True
+        if timestamp is None:
+            timestamp = cast(int, message.get("timestamp") or 0)
+        text = content_text(message.get("content"))
+        if text:
+            contents.append(text)
+        for name, value in cast(Mapping[str, str | None], message.get("sections") or {}).items():
+            if value is None:
+                sections.pop(name, None)
+            else:
+                sections[name] = value
+    tools = current_tools(messages)
+    if not found:
+        return None
+    head: dict[str, Any] = {
+        "role": "system",
+        "content": "\n\n".join(contents),
+        "timestamp": 0 if timestamp is None else timestamp,
+    }
+    if sections:
+        head["sections"] = sections
+    if tools:
+        head["toolsAdded"] = tools
+    return head
+
+
+def collapse_system_messages(
+    messages: Sequence[Mapping[str, Any]],
+) -> tuple[dict[str, Any] | None, list[dict[str, Any]]]:
+    """Return the replayed leading system message and the conversation messages without system entries."""
+
+    head = current_system_message(messages)
+    rest = [dict(message) for message in messages if message.get("role") != "system"]
+    return head, rest
+
+
+def transform_messages(
+    messages: Sequence[Mapping[str, Any]],
+    model: Model,
+    *,
+    api: str,
+    normalize_tool_call_id: Callable[[str], str] | None = None,
+) -> list[dict[str, Any]]:
+    """Adapt a transcript to ``model``: image downgrade, cross-model thinking, and tool-call IDs.
+
+    Thinking signed by another model is dropped because signatures only replay for their producer.
+    A tool-call ID another provider minted is normalized when the adapter supplies a rule, and the
+    matching tool result is rewritten too. Failed and aborted assistant messages are skipped.
+    """
+
+    tool_call_ids: dict[str, str] = {}
+    transformed: list[dict[str, Any]] = []
+    for message in messages:
+        role = message.get("role")
+        if role == "system":
+            transformed.append(dict(message))
+            continue
+        if role == "user":
+            transformed.append(_downgrade_images(dict(message), model, NON_VISION_USER_IMAGE_PLACEHOLDER))
+            continue
+        if role == "toolResult":
+            result = _downgrade_images(dict(message), model, NON_VISION_TOOL_IMAGE_PLACEHOLDER)
+            normalized = tool_call_ids.get(cast(str, message.get("toolCallId")))
+            if normalized is not None:
+                result["toolCallId"] = normalized
+            transformed.append(result)
+            continue
+        if role != "assistant":
+            transformed.append(dict(message))
+            continue
+        if message.get("stopReason") in ("error", "aborted"):
+            continue
+        same_model = (
+            message.get("provider") == model.provider
+            and message.get("api") == api
+            and message.get("model") == model.model_id
+        )
+        content: list[dict[str, Any]] = []
+        for block in cast(Sequence[Mapping[str, Any]], message.get("content") or ()):
+            kind = block.get("type")
+            if kind == "thinking":
+                if not same_model:
+                    continue
+                if (
+                    not block.get("redacted")
+                    and not block.get("thinkingSignature")
+                    and not block.get("thinking")
+                ):
+                    continue
+                content.append(dict(block))
+            elif kind == "text":
+                content.append(dict(block))
+            elif kind == "toolCall":
+                call = dict(block)
+                if not same_model:
+                    call.pop("thoughtSignature", None)
+                    if normalize_tool_call_id is not None:
+                        original = cast(str, call.get("id") or "")
+                        normalized = normalize_tool_call_id(original)
+                        if normalized != original:
+                            tool_call_ids[original] = normalized
+                            call["id"] = normalized
+                content.append(call)
+            else:
+                content.append(dict(block))
+        assistant = dict(message)
+        assistant["content"] = content
+        transformed.append(assistant)
+    return transformed
+
+
+def _downgrade_images(
+    message: dict[str, Any], model: Model, placeholder: str
+) -> dict[str, Any]:
+    if "image" in model.input:
+        return message
+    content = message.get("content")
+    if not isinstance(content, (list, tuple)):
+        return message
+    replaced: list[dict[str, Any]] = []
+    previous_was_placeholder = False
+    for block in cast(Sequence[Mapping[str, Any]], content):
+        if block.get("type") != "image":
+            replaced.append(dict(block))
+            previous_was_placeholder = block.get("type") == "text" and block.get("text") == placeholder
+            continue
+        if not previous_was_placeholder:
+            replaced.append({"type": "text", "text": placeholder})
+        previous_was_placeholder = True
+    message["content"] = replaced
+    return message
+
+
+def has_tool_history(messages: Sequence[Mapping[str, Any]]) -> bool:
+    """Whether the transcript carries any tool call or tool result."""
+
+    for message in messages:
+        if message.get("role") == "toolResult":
+            return True
+        if message.get("role") == "assistant" and any(
+            block.get("type") == "toolCall"
+            for block in cast(Sequence[Mapping[str, Any]], message.get("content") or ())
+        ):
+            return True
+    return False
+
+
+__all__ = [
+    "NON_VISION_TOOL_IMAGE_PLACEHOLDER",
+    "NON_VISION_USER_IMAGE_PLACEHOLDER",
+    "collapse_system_messages",
+    "content_text",
+    "current_system_message",
+    "current_tools",
+    "has_tool_history",
+    "system_message_text",
+    "transform_messages",
+]

pi_durable/providers/usage.py

@@ -0,0 +1,58 @@
+"""Provider usage accounting shared by the model adapters.
+
+Port of the pi-ai 1.3.0 usage shape and ``calculateCost()``. Every adapter reports tokens in the
+same four buckets and prices them at the model's catalog rates so the durable usage ledger sums
+providers without knowing which one produced a response.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from . import Model
+
+
+def zero_cost() -> dict[str, float]:
+    """A complete pi-ai ``Usage.cost`` value with every rate at zero."""
+
+    return {"input": 0.0, "output": 0.0, "cacheRead": 0.0, "cacheWrite": 0.0, "total": 0.0}
+
+
+def zero_usage() -> dict[str, Any]:
+    """A complete pi-ai ``Usage`` value with every counter at zero."""
+
+    return {
+        "input": 0,
+        "output": 0,
+        "cacheRead": 0,
+        "cacheWrite": 0,
+        "totalTokens": 0,
+        "cost": zero_cost(),
+    }
+
+
+def calculate_cost(model: Model, usage: dict[str, Any]) -> dict[str, float]:
+    """Fill ``usage["cost"]`` from the model's USD-per-million-token rates."""
+
+    rates = model.cost or {}
+    input_rate = float(rates.get("input", 0.0))
+    output_rate = float(rates.get("output", 0.0))
+    cache_read_rate = float(rates.get("cacheRead", 0.0))
+    cache_write_rate = float(rates.get("cacheWrite", 0.0))
+    long_write = int(usage.get("cacheWrite1h") or 0)
+    short_write = max(0, int(usage.get("cacheWrite") or 0) - long_write)
+    cost = {
+        "input": input_rate / 1_000_000 * int(usage.get("input") or 0),
+        "output": output_rate / 1_000_000 * int(usage.get("output") or 0),
+        "cacheRead": cache_read_rate / 1_000_000 * int(usage.get("cacheRead") or 0),
+        "cacheWrite": (
+            cache_write_rate * short_write + input_rate * 2 * long_write
+        )
+        / 1_000_000,
+    }
+    cost["total"] = cost["input"] + cost["output"] + cost["cacheRead"] + cost["cacheWrite"]
+    usage["cost"] = cost
+    return cost
+
+
+__all__ = ["calculate_cost", "zero_cost", "zero_usage"]