pi-python-core 0.8.1__py3-none-any.whl

pi_python/proxy.py

@@ -0,0 +1,304 @@
+"""Pi agent-core streamProxy wire protocol (v1.0.0, MIT)."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator, AsyncIterator, Callable
+from .messages import Message
+from .providers.transport import HTTPTransport
+from typing import Any
+from copy import deepcopy
+import json
+from .cancellation import CancelToken
+from .errors import ConfigurationError, ProviderProtocolError, UnsupportedCapabilityError
+from .messages import AssistantMessage, TextContent, ThinkingContent, ToolCall, message_to_dict
+from .provider import ModelEvent, ModelRequest
+from .tools import invoke
+from .stream import event_contract
+from .providers.common import RemoteProvider
+import asyncio
+from dataclasses import replace
+
+_KEYS = {
+    "tool_call": "toolCall",
+    "tool_result": "toolResult",
+    "tool_use": "toolUse",
+    "text_signature": "textSignature",
+    "thinking_signature": "thinkingSignature",
+    "thought_signature": "thoughtSignature",
+    "mime_type": "mimeType",
+    "call_id": "toolCallId",
+    "is_error": "isError",
+    "stop_reason": "stopReason",
+    "provider_thinking_level": "providerThinkingLevel",
+    "tools_added": "toolsAdded",
+    "tools_removed": "toolsRemoved",
+    "input_schema": "parameters",
+    "error": "errorMessage",
+    "response_id": "responseId",
+    "response_model": "responseModel",
+    "thinking_level": "thinkingLevel",
+    "raw_stop_reason": "rawStopReason",
+    "end_turn": "endTurn",
+    "nested_calls": "nestedCalls",
+    "duration_ms": "durationMs",
+    "arguments_bytes": "argumentsBytes",
+    "model_id": "modelId",
+    "expires_at": "expiresAt",
+    "poll_after_ms": "pollAfterMs",
+}
+_OPTIONS = {
+    "max_tokens": "maxTokens",
+    "sampling_params": "samplingParams",
+    "cache_retention": "cacheRetention",
+    "session_id": "sessionId",
+    "thinking_budgets": "thinkingBudgets",
+    "max_retry_delay_ms": "maxRetryDelayMs",
+}
+_ALLOWED = {
+    "temperature",
+    "samplingParams",
+    "maxTokens",
+    "reasoning",
+    "cacheRetention",
+    "sessionId",
+    "headers",
+    "metadata",
+    "transport",
+    "thinkingBudgets",
+    "maxRetryDelayMs",
+}
+
+
+def pi_usage(value: dict[str, Any], *, decode: bool = False) -> dict[str, Any]:
+    keys = {
+        "cache_read": "cacheRead",
+        "cache_write": "cacheWrite",
+        "cache_write_1h": "cacheWrite1h",
+        "total_tokens": "totalTokens",
+    }
+    if decode:
+        keys = {v: k for k, v in keys.items()}
+    return {
+        keys.get(k, k): pi_usage(v, decode=decode)
+        if k == "cost" and isinstance(v, dict)
+        else deepcopy(v)
+        for k, v in value.items()
+    }
+
+
+def pi_message(message: Message) -> dict[str, Any]:
+    def convert(value: Any) -> Any:
+        if isinstance(value, list):
+            return [convert(v) for v in value]
+        if isinstance(value, dict):
+            result = {
+                _KEYS.get(k, k): (
+                    deepcopy(v)
+                    if k
+                    in {
+                        "arguments",
+                        "input_schema",
+                        "data",
+                        "usage",
+                        "sections",
+                        "details",
+                        "structured_content",
+                    }
+                    else convert(v)
+                )
+                for k, v in value.items()
+            }
+            for key in ("type", "role", "stopReason"):
+                if key in result:
+                    result[key] = _KEYS.get(result[key], result[key])
+            if value.get("role") in {"assistant", "tool_result"} and isinstance(
+                value.get("usage"), dict
+            ):
+                result["usage"] = pi_usage(value["usage"])
+            if value.get("role") == "tool_result":
+                result["toolName"] = result.pop("name")
+            if value.get("role") == "system":
+                result["toolsRemoved"] = [{"name": name} for name in value.get("tools_removed", [])]
+            if "timestamp" in result:
+                result["timestamp"] *= 1000
+            return result
+        return value
+
+    return convert(message_to_dict(message))
+
+
+class ProxyProvider(RemoteProvider):
+    name = "proxy"
+
+    def __init__(
+        self,
+        *,
+        model: dict[str, Any],
+        proxy_url: str,
+        auth_token: str | Callable[[], Any],
+        transport: HTTPTransport | None = None,
+    ) -> None:
+        super().__init__(api_key=auth_token, transport=transport)
+        if not {"id", "provider", "api"} <= model.keys():
+            raise ConfigurationError(
+                "Proxy model needs id, provider, api; use the server model descriptor"
+            )
+        self.model = deepcopy(model)
+        self.proxy_url = proxy_url.rstrip("/")
+
+    @event_contract
+    async def stream(
+        self, request: ModelRequest, cancel: CancelToken
+    ) -> AsyncGenerator[ModelEvent, None]:
+        token = await self.credential(request, cancel)
+        options = {
+            _OPTIONS.get(k, k): v
+            for k, v in request.options.items()
+            if _OPTIONS.get(k, k) in _ALLOWED
+        }
+        # Transcript tool declarations are carried by system messages, as upstream does.
+        messages = [pi_message(m) for m in request.messages]
+        if request.tools:
+            messages.append(
+                {
+                    "role": "system",
+                    "content": "",
+                    "sections": {},
+                    "toolsAdded": [
+                        {"name": t.name, "description": t.description, "parameters": t.input_schema}
+                        for t in request.tools
+                    ],
+                    "toolsRemoved": [],
+                    "timestamp": 0,
+                }
+            )
+        body = await self.payload(
+            request, {"model": self.model, "context": {"messages": messages}, "options": options}
+        )
+        events = self.transport.stream(
+            self.proxy_url + "/api/stream",
+            body,
+            {"authorization": f"Bearer {token}", "content-type": "application/json"},
+            cancel,
+            on_response=request.on_response,
+        )
+        blocks: dict[int, Any] = {}
+        arguments = {}
+        final = None
+        closed_blocks = set()
+        yield ModelEvent("start")
+        try:
+            async for event in events:
+                await invoke(request.on_provider_stream_event, deepcopy(event))
+                kind = event.get("type")
+                raw_index = event.get("contentIndex")
+                index = raw_index if type(raw_index) is int else -1
+                if str(kind).endswith(("_start", "_delta", "_end")) and index < 0:
+                    raise ProviderProtocolError(f"Proxy {kind} without contentIndex")
+                if final is not None:
+                    raise ProviderProtocolError("Proxy event after completion")
+                if kind in {"text_start", "thinking_start", "toolcall_start"}:
+                    if type(index) is not int or index != len(blocks):
+                        raise ProviderProtocolError("Invalid proxy block index")
+                    if kind == "text_start":
+                        blocks[index] = TextContent("")
+                    elif kind == "thinking_start":
+                        blocks[index] = ThinkingContent("")
+                    else:
+                        blocks[index] = ToolCall(event["id"], event["toolName"], {})
+                        arguments[index] = ""
+                    yield ModelEvent.boundary("start", index, blocks[index])
+                elif (
+                    kind in {"text_delta", "thinking_delta", "toolcall_delta"}
+                    and index in closed_blocks
+                ):
+                    raise ProviderProtocolError("Proxy delta after block end")
+                elif kind == "text_delta":
+                    blocks[index].text += event["delta"]
+                    yield ModelEvent.text(event["delta"], index)
+                elif kind == "thinking_delta":
+                    blocks[index].thinking += event["delta"]
+                    yield ModelEvent.thinking(event["delta"], index)
+                elif kind == "toolcall_delta":
+                    arguments[index] += event["delta"]
+                    b = blocks[index]
+                    yield ModelEvent.toolcall(event["delta"], index)
+                elif kind in {"text_end", "thinking_end"}:
+                    if index in closed_blocks:
+                        raise ProviderProtocolError("Duplicate proxy block end")
+                    closed_blocks.add(index)
+                    b = blocks[index]
+                    if kind == "text_end":
+                        b.text_signature = event.get("contentSignature")
+                    else:
+                        b.thinking_signature = event.get("contentSignature")
+                    yield ModelEvent.boundary("end", index, b)
+                elif kind == "toolcall_end":
+                    if index in closed_blocks:
+                        raise ProviderProtocolError("Duplicate proxy block end")
+                    closed_blocks.add(index)
+                    value = event["toolCall"]
+                    b = blocks[index]
+                    if value["id"] != b.id or value["name"] != b.name:
+                        raise ProviderProtocolError("Proxy tool identity changed")
+                    b.arguments = value["arguments"]
+                    b.thought_signature = value.get("thoughtSignature")
+                    b.namespace = value.get("namespace")
+                    if arguments[index] and json.loads(arguments[index]) != b.arguments:
+                        raise ProviderProtocolError("Proxy tool arguments mismatch")
+                    yield ModelEvent.boundary("end", index, b)
+                elif kind == "done":
+                    if closed_blocks != set(blocks):
+                        raise ProviderProtocolError("Proxy completed with unfinished blocks")
+                    reason = {"stop": "stop", "length": "length", "toolUse": "tool_use"}.get(
+                        event["reason"]
+                    )
+                    if reason is None:
+                        raise ProviderProtocolError("Invalid proxy stop reason")
+                    final = AssistantMessage(
+                        list(blocks.values()),
+                        reason,
+                        self.model["provider"],
+                        self.model["id"],
+                        pi_usage(event.get("usage", {}), decode=True),
+                        api=self.model["api"],
+                        provider_thinking_level=event.get("providerThinkingLevel"),
+                    )
+                elif kind == "error":
+                    if event.get("reason") == "aborted":
+                        raise asyncio.CancelledError("Proxy request aborted")
+                    raise ProviderProtocolError("Proxy stream reported an error")
+                elif kind != "start":
+                    raise UnsupportedCapabilityError(f"Unsupported proxy event: {kind}")
+        finally:
+            await events.aclose()
+        if final is None:
+            raise ProviderProtocolError("Proxy stream ended without completion")
+        message_to_dict(final)
+        yield ModelEvent.done(final)
+
+
+def stream_proxy(
+    model: dict[str, Any],
+    context: ModelRequest | list[Message],
+    options: dict[str, Any],
+    cancel: CancelToken | None = None,
+    *,
+    transport: HTTPTransport | None = None,
+) -> AsyncIterator[ModelEvent]:
+    """Return an async iterator of ModelEvent; context is ModelRequest or message list."""
+    options = dict(options)
+    provider = ProxyProvider(
+        model=model,
+        proxy_url=options.pop("proxy_url"),
+        auth_token=options.pop("auth_token"),
+        transport=transport,
+    )
+    request = (
+        context
+        if isinstance(context, ModelRequest)
+        else ModelRequest(context, model=model["id"], options=options)
+    )
+    if isinstance(context, ModelRequest):
+        request = replace(context, options={**context.options, **options})
+    return provider.stream(request, cancel or CancelToken())

pi_python/queues.py

@@ -0,0 +1,76 @@
+"""Steering and follow-up queues, consumed at Pi's safe boundaries."""
+
+from __future__ import annotations
+from copy import deepcopy
+
+from .errors import ConfigurationError
+from .messages import Message
+
+MODES = {"all", "one_at_a_time"}
+
+
+class MessageQueues:
+    """Two input queues. A taken message is reserved until history commits it.
+
+    If a run ends before committing a reserved message, `restore` puts it back at the
+    front of its queue, so an interrupted run never loses queued input.
+    """
+
+    def __init__(self, steering_mode: str = "one_at_a_time", follow_up_mode: str = "one_at_a_time"):
+        if steering_mode not in MODES or follow_up_mode not in MODES:
+            raise ConfigurationError("Invalid queue mode")
+        self.steering_mode, self.follow_up_mode = steering_mode, follow_up_mode
+        self.steering: list[Message] = []
+        self.follow_up: list[Message] = []
+        self._reserved: list[tuple[list[Message], Message]] = []
+
+    def take(self, steering: bool) -> list[Message]:
+        queue = self.steering if steering else self.follow_up
+        mode = self.steering_mode if steering else self.follow_up_mode
+        count = len(queue) if mode == "all" else min(1, len(queue))
+        taken = queue[:count]
+        del queue[:count]
+        self._reserved.extend((queue, message) for message in taken)
+        return taken
+
+    def committed(self, message: Message, pending: list[Message]) -> None:
+        """Release the reservation for a committed message from `pending`.
+
+        Declaration reconciliation can change a system message's tool fields, so the
+        match uses identity within `pending` plus timestamp and role.
+        """
+        for i, (_, reserved) in enumerate(self._reserved):
+            if (
+                any(reserved is item for item in pending)
+                and reserved.timestamp == message.timestamp
+                and reserved.role == message.role
+            ):
+                del self._reserved[i]
+                return
+
+    def restore(self) -> None:
+        for queue, message in reversed(self._reserved):
+            queue.insert(0, message)
+        self._reserved.clear()
+
+    def clear(self, *, steering: bool, follow_up: bool) -> None:
+        if steering:
+            self.steering.clear()
+        if follow_up:
+            self.follow_up.clear()
+        self._reserved = [
+            (queue, message)
+            for queue, message in self._reserved
+            if not (
+                (steering and queue is self.steering) or (follow_up and queue is self.follow_up)
+            )
+        ]
+
+    def peek(self) -> list[Message]:
+        """What the next boundary would take: steering first, else follow-up."""
+        steering = self.steering if self.steering_mode == "all" else self.steering[:1]
+        follow_up = self.follow_up if self.follow_up_mode == "all" else self.follow_up[:1]
+        return deepcopy(steering if steering else follow_up)
+
+    def __bool__(self) -> bool:
+        return bool(self.steering or self.follow_up)

pi_python/recovery.py

@@ -0,0 +1,209 @@
+"""Recognize failed responses an application can recover from.
+
+Ported from pi-ai (utils/overflow.ts, utils/retry.ts at the pinned Pi revision). The
+functions read a committed assistant message; the application decides what to do, for
+example compact the history and continue after an overflow, or back off and continue
+after a transient error. `Agent.continue_run()` retries a failed or aborted last response.
+"""
+
+from __future__ import annotations
+
+import re
+
+from .messages import AssistantMessage
+
+# Context overflow messages, per provider (see the pinned overflow.ts for examples).
+_OVERFLOW = [
+    re.compile(pattern, re.I)
+    for pattern in (
+        r"prompt (?:is )?too long",  # Anthropic, z.ai
+        r"prompt exceeds max length",  # z.ai CN
+        r"request_too_large",  # Anthropic byte-size overflow (HTTP 413)
+        r"input is too long for requested model",  # Amazon Bedrock
+        r"exceeds the context window",  # OpenAI Completions and Responses
+        r"exceeds (?:the )?(?:model'?s )?maximum context length(?: of [\d,]+ tokens?|\s*\([\d,]+\))",
+        r"input token count.*exceeds the maximum",  # Google Gemini
+        r"maximum prompt length is \d+",  # xAI
+        r"reduce the length of the messages",  # Groq
+        r"maximum context length is \d+ tokens",  # OpenRouter
+        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
+        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
+        r"model_context_window_exceeded",  # z.ai finish reason surfaced as an error
+        r"prompt too long; exceeded (?:max )?context length",  # Ollama
+        r"range of input length should be",  # DashScope / Qwen
+        r"context[_ ]length[_ ]exceeded",
+        r"too many tokens",
+        r"token limit exceeded",
+    )
+]
+_CEREBRAS_OVERFLOW = re.compile(r"^4(?:00|13)\s*(?:status code)?\s*\(no body\)", re.I)
+# Throttling text that would otherwise match an overflow pattern ("Too many tokens, please wait").
+_NOT_OVERFLOW = [
+    re.compile(pattern, re.I)
+    for pattern in (
+        r"^(Throttling error|Service unavailable):",
+        r"rate limit",
+        r"too many requests",
+    )
+]
+
+# Account limits look like throttling but do not clear in seconds.
+_NOT_RETRYABLE = re.compile(
+    "|".join(
+        (
+            "GoUsageLimitError",
+            "FreeUsageLimitError",
+            "Monthly usage limit reached",
+            "available balance",
+            "insufficient_quota",
+            "out of budget",
+            "quota exceeded",
+            "billing",
+            "subscription_sharing_usage_limit_exceeded",
+        )
+    ),
+    re.I,
+)
+_RETRYABLE = re.compile(
+    "|".join(
+        (
+            "overloaded",
+            "currently experiencing high demand",
+            "rate.?limit",
+            "too many requests",
+            "429",
+            "500",
+            "502",
+            "503",
+            "504",
+            "520",
+            "524",
+            "service.?unavailable",
+            "server.?error",
+            "internal.?error",
+            "provider.?returned.?error",
+            "exceeded request buffer limit while retrying upstream",
+            "network.?error",
+            "connection.?error",
+            "connection.?refused",
+            "connection.?lost",
+            "other side closed",
+            "fetch failed",
+            "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",
+            "retry delay",
+            "you can retry your request",
+            "try your request again",
+            "please retry your request",
+            "ResourceExhausted",
+            "subscription_sharing_usage_unavailable",
+            "subscription_sharing_user_unavailable",
+        )
+    ),
+    re.I,
+)
+
+
+# Python's own transport failures (httpx, websockets, the standard library), which the
+# upstream patterns, written for JavaScript runtimes, do not name.
+_PYTHON_TRANSIENT = re.compile(
+    "|".join(
+        (
+            r"\b(?:Connect|Read|Write|Pool)(?:Error|Timeout)\b",
+            r"RemoteProtocolError",
+            r"ConnectionResetError|ConnectionAbortedError|BrokenPipeError|IncompleteRead",
+            r"connection reset",
+            r"ConnectionClosed",
+            r"server disconnected",
+            r"peer closed",
+            r"name resolution",
+            r"nodename nor servname",
+        )
+    ),
+    re.I,
+)
+# This library's HTTP failures read "... (HTTP 503; server) ...": decide by the status,
+# not by digits that happen to appear in the response body.
+_HTTP_STATUS = re.compile(r"\(HTTP (\d{3});")
+_RETRYABLE_STATUS = {408, 429}
+
+
+def _input_tokens(message: AssistantMessage) -> int:
+    usage = message.usage or {}
+    return int(usage.get("input") or 0) + int(usage.get("cache_read") or 0)
+
+
+def is_context_overflow(message: AssistantMessage, context_window: int | None = None) -> bool:
+    """Whether a response failed because the input exceeded the model's context window.
+
+    Most providers report it as an error message. Pass `context_window` to also catch
+    providers that accept an oversized input silently (input usage above the window) or
+    truncate it and stop for length with no output.
+    """
+    error = message.error or ""
+    if message.stop_reason == "error" and error:
+        if not any(p.search(error) for p in _NOT_OVERFLOW):
+            if any(p.search(error) for p in _OVERFLOW):
+                return True
+            if message.provider == "cerebras" and _CEREBRAS_OVERFLOW.search(error):
+                return True
+    if context_window and message.stop_reason == "stop":
+        if _input_tokens(message) > context_window:
+            return True
+    if context_window and message.stop_reason == "length":
+        if int((message.usage or {}).get("output") or 0) == 0:
+            if _input_tokens(message) >= context_window * 0.99:
+                return True
+    return False
+
+
+def is_recoverable_length(message: AssistantMessage, desired_max_output: int) -> bool:
+    """A length stop below the intended output limit, possibly from context pressure."""
+    output = int((message.usage or {}).get("output") or 0)
+    return (
+        message.stop_reason == "length" and desired_max_output > 0 and output < desired_max_output
+    )
+
+
+def is_retryable_error(message: AssistantMessage) -> bool:
+    """Whether a failed response looks transient (overload, rate limit, network, server error).
+
+    Check `is_context_overflow` first: an overflow needs a smaller context, not a retry.
+    Quota and billing limits are never retryable.
+    """
+    error = message.error or ""
+    if message.stop_reason != "error" or not error:
+        return False
+    if _NOT_RETRYABLE.search(error):
+        return False
+    status = _HTTP_STATUS.search(error)
+    if status:
+        code = int(status.group(1))
+        return code >= 500 or code in _RETRYABLE_STATUS
+    return bool(_RETRYABLE.search(error) or _PYTHON_TRANSIENT.search(error))
+
+
+def retry_delay(attempt: int, base: float = 2.0, max_delay: float = 60.0) -> float:
+    """Exponential backoff in seconds for the 1-based retry `attempt`, as in Pi."""
+    return min(base * 2 ** max(0, attempt - 1), max_delay)