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

pi_ai/utils/__init__.py

@@ -0,0 +1 @@
+"""Shared utilities."""

pi_ai/utils/callbacks.py

@@ -0,0 +1,15 @@
+"""Helpers for user callbacks that may be sync or async."""
+
+from __future__ import annotations
+
+import inspect
+from typing import Any
+
+__all__ = ["maybe_await"]
+
+
+async def maybe_await(value: Any) -> Any:
+    """Await ``value`` if it is awaitable, otherwise return it unchanged."""
+    if inspect.isawaitable(value):
+        return await value
+    return value

pi_ai/utils/diagnostics.py

@@ -0,0 +1,70 @@
+"""Redacted provider/runtime diagnostics attached to assistant messages."""
+
+from __future__ import annotations
+
+import traceback
+from dataclasses import dataclass, field
+from typing import Any
+
+from ..types import now_ms
+
+__all__ = [
+    "DiagnosticErrorInfo",
+    "AssistantMessageDiagnostic",
+    "format_thrown_value",
+    "extract_diagnostic_error",
+    "create_assistant_message_diagnostic",
+    "append_assistant_message_diagnostic",
+]
+
+
+@dataclass(slots=True)
+class DiagnosticErrorInfo:
+    message: str
+    name: str | None = None
+    stack: str | None = None
+    code: str | int | None = None
+
+
+@dataclass(slots=True)
+class AssistantMessageDiagnostic:
+    type: str
+    timestamp: int = field(default_factory=now_ms)
+    error: DiagnosticErrorInfo | None = None
+    details: dict[str, Any] | None = None
+
+
+def format_thrown_value(value: Any) -> str:
+    """Human-readable text for an exception or thrown value."""
+    if isinstance(value, BaseException):
+        return str(value) or type(value).__name__
+    if isinstance(value, str):
+        return value
+    return str(value)
+
+
+def extract_diagnostic_error(error: Any) -> DiagnosticErrorInfo:
+    if not isinstance(error, BaseException):
+        return DiagnosticErrorInfo(name="ThrownValue", message=format_thrown_value(error))
+    code = getattr(error, "code", None)
+    stack = (
+        "".join(traceback.format_exception(type(error), error, error.__traceback__)) if error.__traceback__ else None
+    )
+    return DiagnosticErrorInfo(
+        name=type(error).__name__,
+        message=str(error) or type(error).__name__,
+        stack=stack,
+        code=code if isinstance(code, (str, int)) and not isinstance(code, bool) else None,
+    )
+
+
+def create_assistant_message_diagnostic(
+    type: str, error: Any, details: dict[str, Any] | None = None
+) -> AssistantMessageDiagnostic:
+    return AssistantMessageDiagnostic(
+        type=type, timestamp=now_ms(), error=extract_diagnostic_error(error), details=details
+    )
+
+
+def append_assistant_message_diagnostic(message: Any, diagnostic: AssistantMessageDiagnostic) -> None:
+    message.diagnostics = [*(message.diagnostics or []), diagnostic]

pi_ai/utils/error_body.py

@@ -0,0 +1,89 @@
+"""Normalization of provider HTTP error objects into display strings."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+__all__ = [
+    "MAX_PROVIDER_ERROR_BODY_CHARS",
+    "NormalizedProviderError",
+    "normalize_provider_error",
+    "format_provider_error",
+    "truncate_error_text",
+    "safe_json_stringify",
+]
+
+MAX_PROVIDER_ERROR_BODY_CHARS = 4000
+
+
+@dataclass(slots=True)
+class NormalizedProviderError:
+    message: str
+    #: True when ``message`` already contains the body.
+    message_carries_body: bool
+    status: int | None = None
+    body: str | None = None
+
+
+def safe_json_stringify(value: Any) -> str:
+    try:
+        return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
+    except (TypeError, ValueError):
+        return str(value)
+
+
+def truncate_error_text(text: str, max_chars: int) -> str:
+    if len(text) <= max_chars:
+        return text
+    return f"{text[:max_chars]}... [truncated {len(text) - max_chars} chars]"
+
+
+def _extract_status(error: Any) -> int | None:
+    for name in ("status_code", "status"):
+        value = getattr(error, name, None)
+        if isinstance(value, int) and not isinstance(value, bool):
+            return value
+    return None
+
+
+def _pick_body_text(error: Any) -> str | None:
+    body = getattr(error, "body", None)
+    if isinstance(body, str):
+        return body
+    parsed = getattr(error, "error", None)
+    if isinstance(parsed, dict) and parsed:
+        return safe_json_stringify(parsed)
+    return None
+
+
+def normalize_provider_error(error: Any) -> NormalizedProviderError:
+    """Probe an exception for HTTP status and body so the body reason is not lost."""
+    if not isinstance(error, BaseException):
+        return NormalizedProviderError(message=safe_json_stringify(error), message_carries_body=False)
+    message = str(error) or type(error).__name__
+    status = _extract_status(error)
+    body_text = _pick_body_text(error)
+    body: str | None = None
+    if body_text is not None:
+        trimmed = body_text.strip()
+        if trimmed:
+            body = truncate_error_text(trimmed, MAX_PROVIDER_ERROR_BODY_CHARS)
+    return NormalizedProviderError(
+        message=message,
+        message_carries_body=body is None or body in message,
+        status=status,
+        body=body,
+    )
+
+
+def format_provider_error(norm: NormalizedProviderError, prefix: str | None = None) -> str:
+    """Compose a display string: ``"<status>: <body>"`` or ``"<prefix> (<status>): <body>"``."""
+    if norm.message_carries_body or norm.status is None or norm.body is None:
+        if prefix is not None and norm.status is not None:
+            return f"{prefix} ({norm.status}): {norm.message}"
+        return norm.message
+    if prefix is not None:
+        return f"{prefix} ({norm.status}): {norm.body}"
+    return f"{norm.status}: {norm.body}"

pi_ai/utils/estimate.py

@@ -0,0 +1,125 @@
+"""Rough token estimates for context accounting."""
+
+from __future__ import annotations
+
+import json
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from ..types import TranscriptContext, Usage
+from .text import get_system_message_text
+
+__all__ = [
+    "ContextUsageEstimate",
+    "calculate_context_tokens",
+    "estimate_text_tokens",
+    "estimate_text_and_image_content_tokens",
+    "estimate_message_tokens",
+    "estimate_context_tokens",
+]
+
+_CHARS_PER_TOKEN = 3.5
+_ESTIMATED_IMAGE_CHARS = 4800
+
+
+@dataclass(slots=True)
+class ContextUsageEstimate:
+    #: Estimated total context tokens.
+    tokens: int
+    #: Tokens reported by the most recent applicable assistant usage block.
+    usage_tokens: int
+    #: Estimated tokens after that usage block.
+    trailing_tokens: int
+    #: Index of the message that provided usage, or ``None``.
+    last_usage_index: int | None
+
+
+def calculate_context_tokens(usage: Usage) -> int:
+    return usage.total_tokens or usage.input + usage.output + usage.cache_read + usage.cache_write
+
+
+def _safe_json(value: Any) -> str:
+    try:
+        from ..serialization import to_json_value
+
+        return json.dumps(to_json_value(value), separators=(",", ":"))
+    except Exception:
+        return "[unserializable]"
+
+
+def _content_chars(content: Any) -> int:
+    if isinstance(content, str):
+        return len(content)
+    chars = 0
+    for block in content:
+        chars += len(block.text) if block.type == "text" else _ESTIMATED_IMAGE_CHARS
+    return chars
+
+
+def estimate_text_tokens(text: str) -> int:
+    return math.ceil(len(text) / _CHARS_PER_TOKEN)
+
+
+def estimate_text_and_image_content_tokens(content: Any) -> int:
+    return math.ceil(_content_chars(content) / _CHARS_PER_TOKEN)
+
+
+def _estimate_tools_tokens(tools: Sequence[Any] | None) -> int:
+    if not tools:
+        return 0
+    return estimate_text_tokens(_safe_json(list(tools)))
+
+
+def estimate_message_tokens(message: Any) -> int:
+    role = message.role
+    if role == "system":
+        return (
+            estimate_text_tokens(get_system_message_text(message))
+            + _estimate_tools_tokens(message.tools_added)
+            + _estimate_tools_tokens(message.tools_removed)
+        )
+    if role in ("user", "toolResult"):
+        return estimate_text_and_image_content_tokens(message.content)
+    if role != "assistant":
+        return 0
+    chars = 0
+    for block in message.content:
+        kind = block.type
+        if kind == "text":
+            chars += len(block.text)
+        elif kind == "thinking":
+            chars += len(block.thinking)
+        else:
+            chars += len(block.name) + len(_safe_json(block.arguments))
+    return math.ceil(chars / _CHARS_PER_TOKEN)
+
+
+def _last_assistant_usage(messages: Sequence[Any]) -> tuple[Usage, int] | None:
+    latest_prefix_timestamp = -math.inf
+    info: tuple[Usage, int] | None = None
+    for index, message in enumerate(messages):
+        if message.role == "assistant":
+            applies = message.timestamp >= latest_prefix_timestamp
+            if (
+                applies
+                and message.stop_reason not in ("aborted", "error")
+                and calculate_context_tokens(message.usage) > 0
+            ):
+                info = (message.usage, index)
+        latest_prefix_timestamp = max(latest_prefix_timestamp, getattr(message, "timestamp", 0))
+    return info
+
+
+def estimate_context_tokens(context: TranscriptContext | Sequence[Any]) -> ContextUsageEstimate:
+    """Estimate context tokens, anchored on the most recent assistant usage report."""
+    messages = context.messages if isinstance(context, TranscriptContext) else context
+    usage_info = _last_assistant_usage(messages)
+    if usage_info is not None:
+        usage, index = usage_info
+        usage_tokens = calculate_context_tokens(usage)
+        trailing = sum(estimate_message_tokens(message) for message in messages[index + 1 :])
+        return ContextUsageEstimate(usage_tokens + trailing, usage_tokens, trailing, index)
+    tokens = sum(estimate_message_tokens(message) for message in messages)
+    return ContextUsageEstimate(tokens, 0, tokens, None)

pi_ai/utils/fastjson.py

@@ -0,0 +1,63 @@
+"""JSON encode/decode with an optional ``orjson`` fast path.
+
+Install ``pi-ai-python[fast]`` to use orjson for SSE chunk decoding and request bodies.
+Output is compact and UTF-8 in both cases.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+__all__ = ["loads", "dumps", "dumps_bytes", "HAS_ORJSON"]
+
+try:  # pragma: no cover - exercised when orjson is installed
+    import orjson as _orjson
+
+    HAS_ORJSON = True
+except ImportError:  # pragma: no cover
+    _orjson = None
+    HAS_ORJSON = False
+
+_decode = json.JSONDecoder(strict=True).decode
+
+
+def _default(value: Any) -> Any:
+    from ..serialization import to_json_value
+
+    return to_json_value(value)
+
+
+if _orjson is not None:
+    # Dataclasses go through ``_default`` so they serialize in pi's camelCase format.
+    _ORJSON_OPTS = _orjson.OPT_NON_STR_KEYS | _orjson.OPT_PASSTHROUGH_DATACLASS
+
+    loads = _orjson.loads  # pyright: ignore[reportAssignmentType]
+    _orjson_dumps = _orjson.dumps
+
+    def dumps_bytes(value: Any) -> bytes:
+        try:
+            return _orjson_dumps(value, default=_default, option=_ORJSON_OPTS)
+        except TypeError:
+            # orjson rejects lone surrogates and integers above 64 bits; the stdlib escapes them.
+            return json.dumps(value, ensure_ascii=True, separators=(",", ":"), default=_default).encode()
+
+    def dumps(value: Any) -> str:
+        return dumps_bytes(value).decode()
+
+else:
+
+    def loads(data: str | bytes | bytearray | memoryview) -> Any:
+        if not isinstance(data, str):
+            data = bytes(data).decode()
+        return _decode(data)
+
+    def dumps(value: Any) -> str:
+        return json.dumps(value, ensure_ascii=False, separators=(",", ":"), default=_default)
+
+    def dumps_bytes(value: Any) -> bytes:
+        text = dumps(value)
+        try:
+            return text.encode()
+        except UnicodeEncodeError:  # lone surrogates: emit them as \u escapes
+            return json.dumps(value, ensure_ascii=True, separators=(",", ":"), default=_default).encode()