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

pi_ai/utils/json_parse.py

@@ -0,0 +1,424 @@
+"""Tolerant JSON parsing for streamed tool-call arguments.
+
+``parse_streaming_json`` mirrors pi's helper (best-effort parse of a JSON prefix).
+:class:`StreamingJsonParser` is an incremental variant used by the provider
+adapters: it keeps the lexer state between deltas, so each delta costs time
+proportional to the delta (plus one C-level ``json.loads``) instead of rescanning
+the whole argument text in Python.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import Any
+
+__all__ = [
+    "repair_json",
+    "parse_json_with_repair",
+    "parse_tool_call_arguments",
+    "parse_streaming_json",
+    "complete_partial_json",
+    "StreamingJsonParser",
+]
+
+_VALID_ESCAPES = frozenset('"\\/bfnrt')
+_HEX4 = re.compile(r"[0-9a-fA-F]{4}")
+_lenient_decode = json.JSONDecoder(strict=False).decode
+_strict_decode = json.JSONDecoder(strict=True).decode
+
+
+def _escape_control(char: str) -> str:
+    match char:
+        case "\b":
+            return "\\b"
+        case "\f":
+            return "\\f"
+        case "\n":
+            return "\\n"
+        case "\r":
+            return "\\r"
+        case "\t":
+            return "\\t"
+        case _:
+            return f"\\u{ord(char):04x}"
+
+
+def repair_json(text: str) -> str:
+    """Repair malformed JSON string literals.
+
+    Escapes raw control characters inside strings and doubles backslashes before
+    invalid escape characters.
+    """
+    out: list[str] = []
+    in_string = False
+    index = 0
+    length = len(text)
+    while index < length:
+        char = text[index]
+        if not in_string:
+            out.append(char)
+            if char == '"':
+                in_string = True
+            index += 1
+            continue
+        if char == '"':
+            out.append(char)
+            in_string = False
+            index += 1
+            continue
+        if char == "\\":
+            if index + 1 >= length:
+                out.append("\\\\")
+                index += 1
+                continue
+            next_char = text[index + 1]
+            if next_char == "u" and _HEX4.fullmatch(text, index + 2, index + 6):
+                out.append(text[index : index + 6])
+                index += 6
+                continue
+            if next_char in _VALID_ESCAPES:
+                out.append("\\" + next_char)
+                index += 2
+                continue
+            out.append("\\\\")
+            index += 1
+            continue
+        out.append(_escape_control(char) if ord(char) < 0x20 else char)
+        index += 1
+    return "".join(out)
+
+
+def parse_json_with_repair(text: str) -> Any:
+    """``json.loads`` that retries once with :func:`repair_json` applied."""
+    try:
+        return _strict_decode(text)
+    except ValueError:
+        repaired = repair_json(text)
+        if repaired != text:
+            return _strict_decode(repaired)
+        raise
+
+
+def parse_tool_call_arguments(text: str | None) -> tuple[dict[str, Any], str | None]:
+    """Parse the complete arguments text of a finished tool call.
+
+    Returns ``(arguments, arguments_error)``. Unlike :func:`parse_streaming_json`,
+    this never keeps a partially parsed prefix: text that is not a valid JSON object
+    yields ``({}, error)`` where the error includes the text the model sent, so tool
+    validation can report it back to the model.
+    """
+    if not text or not text.strip():
+        return {}, None
+    try:
+        parsed = parse_json_with_repair(text)
+        if not isinstance(parsed, dict):
+            raise ValueError("Tool arguments must be a JSON object")
+        return parsed, None
+    except ValueError as error:
+        return {}, f"{error}\n\nReceived arguments:\n{text}"
+
+
+def parse_streaming_json(partial_json: str | None) -> Any:
+    """Best-effort parse of possibly incomplete JSON. Always returns a value (``{}`` on failure)."""
+    if not partial_json or not partial_json.strip():
+        return {}
+    try:
+        return parse_json_with_repair(partial_json)
+    except ValueError:
+        pass
+    parser = StreamingJsonParser()
+    parser.feed(partial_json)
+    value = parser.value(default=None)
+    return {} if value is None else value
+
+
+def complete_partial_json(partial_json: str) -> str | None:
+    """Close an incomplete JSON prefix into parseable JSON text, or ``None`` if nothing is usable."""
+    scanner = _Scanner()
+    scanner.feed(partial_json)
+    return scanner.complete(partial_json)
+
+
+# ---------------------------------------------------------------------------
+# Incremental scanner
+# ---------------------------------------------------------------------------
+
+# Frame states
+_OBJ_KEY = 0  # after "{" or ",": expecting a key or "}"
+_OBJ_IN_KEY = 1  # inside the key string
+_OBJ_COLON = 2  # key done, expecting ":"
+_OBJ_VALUE = 3  # after ":", expecting a value
+_AFTER_VALUE = 4  # value done, expecting "," or the closer
+_ARR_VALUE = 5  # after "[" or ",": expecting a value or "]"
+
+_LITERALS = ("true", "false", "null")
+_NUMBER_RE = re.compile(r"-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?")
+_STRING_SPECIAL = re.compile(r'["\\]')
+_SCALAR_END = re.compile(r"[\s,\]}]")
+
+
+class _Frame:
+    __slots__ = ("closer", "safe", "state")
+
+    def __init__(self, closer: str, state: int, safe: int) -> None:
+        self.closer = closer
+        self.state = state
+        #: Prefix length that is structurally valid for this container once closed.
+        self.safe = safe
+
+
+class _Scanner:
+    """Resumable JSON lexer that tracks where a prefix can be safely cut and closed."""
+
+    __slots__ = ("in_string", "invalid", "pos", "scalar_start", "stack", "string_is_key", "top_done")
+
+    def __init__(self) -> None:
+        self.pos = 0
+        self.stack: list[_Frame] = []
+        self.in_string = False
+        self.string_is_key = False
+        self.scalar_start = -1  # start of a number/literal token in progress
+        self.top_done = False
+        self.invalid = False
+
+    def feed(self, text: str) -> None:
+        if self.invalid:
+            return
+        pos = self.pos
+        length = len(text)
+        stack = self.stack
+        while pos < length:
+            if self.in_string:
+                match = _STRING_SPECIAL.search(text, pos)
+                if match is None:
+                    pos = length
+                    break
+                index = match.start()
+                if text[index] == "\\":
+                    if index + 1 >= length:
+                        pos = index  # wait for the escaped character
+                        break
+                    pos = index + 2
+                    continue
+                pos = index + 1
+                self.in_string = False
+                self._value_done(pos, is_key=self.string_is_key)
+                continue
+
+            if self.scalar_start >= 0:
+                match = _SCALAR_END.search(text, pos)
+                if match is None:
+                    pos = length
+                    break
+                pos = match.start()
+                self.scalar_start = -1
+                self._value_done(pos, is_key=False)
+                continue
+
+            char = text[pos]
+            if char in " \t\r\n":
+                pos += 1
+                continue
+            frame = stack[-1] if stack else None
+            if char == '"':
+                if frame is not None and frame.state == _OBJ_KEY:
+                    frame.state = _OBJ_IN_KEY
+                    self.string_is_key = True
+                else:
+                    if not self._expect_value(frame):
+                        break
+                    self.string_is_key = False
+                self.in_string = True
+                pos += 1
+                continue
+            if char == "{" or char == "[":
+                if not self._expect_value(frame):
+                    break
+                pos += 1
+                stack.append(_Frame("}" if char == "{" else "]", _OBJ_KEY if char == "{" else _ARR_VALUE, pos))
+                continue
+            if char == "}" or char == "]":
+                if frame is None or frame.closer != char or frame.state not in (_OBJ_KEY, _AFTER_VALUE, _ARR_VALUE):
+                    self.invalid = True
+                    break
+                stack.pop()
+                pos += 1
+                self._value_done(pos, is_key=False)
+                continue
+            if char == ",":
+                if frame is None or frame.state != _AFTER_VALUE:
+                    self.invalid = True
+                    break
+                frame.state = _OBJ_KEY if frame.closer == "}" else _ARR_VALUE
+                pos += 1
+                continue
+            if char == ":":
+                if frame is None or frame.state != _OBJ_COLON:
+                    self.invalid = True
+                    break
+                frame.state = _OBJ_VALUE
+                pos += 1
+                continue
+            # number or literal
+            if not self._expect_value(frame):
+                break
+            self.scalar_start = pos
+            pos += 1
+        self.pos = pos
+
+    def _expect_value(self, frame: _Frame | None) -> bool:
+        if frame is None:
+            if self.top_done:
+                self.invalid = True
+                return False
+            return True
+        if frame.state in (_OBJ_VALUE, _ARR_VALUE):
+            return True
+        self.invalid = True
+        return False
+
+    def _value_done(self, end: int, *, is_key: bool) -> None:
+        if not self.stack:
+            self.top_done = True
+            return
+        frame = self.stack[-1]
+        if is_key:
+            frame.state = _OBJ_COLON
+        else:
+            frame.state = _AFTER_VALUE
+            frame.safe = end
+
+    def complete(self, text: str) -> str | None:
+        """Return ``text`` cut and closed into parseable JSON (``None`` if nothing usable)."""
+        if self.invalid:
+            return None
+        stack = self.stack
+        if not stack:
+            if self.top_done:
+                return text
+            if self.in_string:
+                return _close_string(text)
+            if self.scalar_start >= 0:
+                return _complete_scalar(text[self.scalar_start :])
+            return None
+
+        frame = stack[-1]
+        suffix = "".join(f.closer for f in reversed(stack))
+        if self.in_string:
+            if self.string_is_key:
+                return text[: frame.safe] + suffix
+            return _close_string(text) + suffix
+        if self.scalar_start >= 0:
+            scalar = _complete_scalar(text[self.scalar_start :])
+            if scalar is None:
+                return text[: frame.safe] + suffix
+            return text[: self.scalar_start] + scalar + suffix
+        if frame.state == _AFTER_VALUE:
+            return text + suffix
+        return text[: frame.safe] + suffix
+
+
+def _close_string(text: str) -> str:
+    # Drop a dangling escape so the closing quote is not escaped.
+    end = len(text)
+    backslashes = 0
+    while end - backslashes - 1 >= 0 and text[end - backslashes - 1] == "\\":
+        backslashes += 1
+    if backslashes % 2 == 1:
+        text = text[:-1]
+    else:
+        unicode_start = text.rfind("\\u", max(0, len(text) - 5))
+        if unicode_start != -1 and len(text) - unicode_start < 6:
+            preceding = 0
+            while unicode_start - preceding - 1 >= 0 and text[unicode_start - preceding - 1] == "\\":
+                preceding += 1
+            if preceding % 2 == 0:
+                text = text[:unicode_start]
+    return text + '"'
+
+
+def _complete_scalar(token: str) -> str | None:
+    if not token:
+        return None
+    first = token[0]
+    if first in "tfn":
+        for literal in _LITERALS:
+            if literal.startswith(token):
+                return literal
+        return None
+    match = _NUMBER_RE.match(token)
+    if match is None:
+        return None
+    return match.group(0)
+
+
+class StreamingJsonParser:
+    """Incrementally parse a growing JSON document.
+
+    ::
+
+        parser = StreamingJsonParser()
+        for delta in deltas:
+            parser.feed(delta)
+            partial_arguments = parser.value()
+
+    :meth:`value` returns the best parse of the text so far. When the current prefix
+    cannot be parsed (for example invalid escapes the repair pass cannot fix), the
+    last successful value is returned instead of an empty object, so partial
+    arguments never regress during streaming.
+    """
+
+    __slots__ = ("_dirty", "_last", "_scanner", "_text")
+
+    def __init__(self) -> None:
+        self._text = ""
+        self._scanner = _Scanner()
+        self._last: Any = None
+        self._dirty = False
+
+    @property
+    def text(self) -> str:
+        return self._text
+
+    def feed(self, delta: str) -> None:
+        if not delta:
+            return
+        self._text += delta
+        self._scanner.feed(self._text)
+        self._dirty = True
+
+    def value(self, default: Any = ...) -> Any:
+        if not self._dirty:
+            if self._last is not None:
+                return self._last
+            return {} if default is ... else default
+        self._dirty = False
+        text = self._text
+        scanner = self._scanner
+        # While the scanner knows the document is still open, a full parse can only fail.
+        if not scanner.invalid and (scanner.stack or scanner.in_string or not scanner.top_done):
+            parsed = _FAIL
+        else:
+            parsed = _try_parse(text)
+        if parsed is _FAIL:
+            completed = self._scanner.complete(text)
+            if completed is not None:
+                parsed = _try_parse(completed)
+                if parsed is _FAIL:
+                    parsed = _try_parse(repair_json(completed))
+        if parsed is not _FAIL:
+            self._last = parsed
+        if self._last is None:
+            return {} if default is ... else default
+        return self._last
+
+
+_FAIL: Any = object()
+
+
+def _try_parse(text: str) -> Any:
+    try:
+        return _lenient_decode(text)
+    except ValueError:
+        return _FAIL

pi_ai/utils/misc.py

@@ -0,0 +1,193 @@
+"""Small shared helpers: env lookup, headers, hashing, ids, unicode, user agent."""
+
+from __future__ import annotations
+
+import os
+import platform
+import re
+import secrets
+import sys
+import threading
+import time
+from collections.abc import Iterable, Mapping
+from typing import Any
+
+__all__ = [
+    "get_provider_env_value",
+    "headers_to_record",
+    "provider_headers_to_record",
+    "merge_headers",
+    "has_header",
+    "short_hash",
+    "uuidv7",
+    "sanitize_surrogates",
+    "get_pi_user_agent",
+]
+
+
+def get_provider_env_value(name: str, env: Mapping[str, str] | None = None) -> str | None:
+    """Resolve a provider env value from scoped overrides, then ``os.environ``.
+
+    Empty strings count as unset, like JavaScript's ``||``.
+    """
+    if env:
+        value = env.get(name)
+        if value:
+            return value
+    value = os.environ.get(name)
+    return value or None
+
+
+def headers_to_record(headers: Any) -> dict[str, str]:
+    """Convert response headers (``httpx.Headers`` or a mapping) to a lowercase-key dict."""
+    items = headers.multi_items() if hasattr(headers, "multi_items") else headers.items()
+    result: dict[str, str] = {}
+    for key, value in items:
+        lower = key.lower()
+        result[lower] = f"{result[lower]}, {value}" if lower in result else value
+    return result
+
+
+def provider_headers_to_record(*sources: Mapping[str, str | None] | None) -> dict[str, str] | None:
+    """Merge header sources case-insensitively; later sources win and ``None`` deletes."""
+    merged: dict[str, tuple[str, str]] = {}
+    for source in sources:
+        if not source:
+            continue
+        for name, value in source.items():
+            lower = name.lower()
+            merged.pop(lower, None)
+            if value is not None:
+                merged[lower] = (name, value)
+    return {name: value for name, value in merged.values()} if merged else None
+
+
+def merge_headers(
+    base: Mapping[str, str | None] | None, override: Mapping[str, str | None] | None
+) -> dict[str, str | None] | None:
+    """Case-insensitive header merge that keeps ``None`` suppression markers."""
+    if not base and not override:
+        return None
+    merged: dict[str, str | None] = dict(base or {})
+    for name, value in (override or {}).items():
+        lower = name.lower()
+        for existing in [key for key in merged if key.lower() == lower]:
+            del merged[existing]
+        merged[name] = value
+    return merged
+
+
+def has_header(headers: Mapping[str, Any] | None, name: str) -> bool:
+    if not headers:
+        return False
+    lower = name.lower()
+    return any(key.lower() == lower and value is not None for key, value in headers.items())
+
+
+_MASK = 0xFFFFFFFF
+_BASE36 = "0123456789abcdefghijklmnopqrstuvwxyz"
+
+
+def _imul(a: int, b: int) -> int:
+    return ((a & _MASK) * (b & _MASK)) & _MASK
+
+
+def _to_base36(value: int) -> str:
+    if value == 0:
+        return "0"
+    digits: list[str] = []
+    while value:
+        value, remainder = divmod(value, 36)
+        digits.append(_BASE36[remainder])
+    return "".join(reversed(digits))
+
+
+def _utf16_units(text: str) -> Iterable[int]:
+    for char in text:
+        code = ord(char)
+        if code > 0xFFFF:
+            code -= 0x10000
+            yield 0xD800 + (code >> 10)
+            yield 0xDC00 + (code & 0x3FF)
+        else:
+            yield code
+
+
+def short_hash(text: str) -> str:
+    """Fast deterministic hash used to shorten long ids (bit-identical to pi's ``shortHash``)."""
+    h1 = 0xDEADBEEF
+    h2 = 0x41C6CE57
+    for unit in _utf16_units(text):
+        h1 = _imul(h1 ^ unit, 2654435761)
+        h2 = _imul(h2 ^ unit, 1597334677)
+    h1 = _imul(h1 ^ (h1 >> 16), 2246822507) ^ _imul(h2 ^ (h2 >> 13), 3266489909)
+    h2 = _imul(h2 ^ (h2 >> 16), 2246822507) ^ _imul(h1 ^ (h1 >> 13), 3266489909)
+    return _to_base36(h2 & _MASK) + _to_base36(h1 & _MASK)
+
+
+_MAX_UUID_V7_TIMESTAMP = 0xFFFFFFFFFFFF
+_MAX_SEQUENCE = (1 << 41) - 1
+_uuid_lock = threading.Lock()
+_last_ordinary_timestamp = -1
+_sequence: int | None = None
+
+
+def uuidv7(timestamp_ms: int | None = None) -> str:
+    """Generate a time-ordered UUIDv7. A supplied timestamp is preserved for follower ids."""
+    global _last_ordinary_timestamp, _sequence
+    requested = timestamp_ms if timestamp_ms is not None else time.time_ns() // 1_000_000
+    if not isinstance(requested, int) or requested < 0 or requested > _MAX_UUID_V7_TIMESTAMP:
+        raise ValueError(f"UUIDv7 timestamp must be an integer between 0 and {_MAX_UUID_V7_TIMESTAMP}")
+    random_bytes = bytearray(secrets.token_bytes(16))
+    with _uuid_lock:
+        if timestamp_ms is None:
+            effective = max(requested, _last_ordinary_timestamp)
+            _last_ordinary_timestamp = effective
+        else:
+            effective = timestamp_ms
+        if _sequence is None:
+            _sequence = int.from_bytes(random_bytes[1:6], "big")
+        else:
+            if _sequence == _MAX_SEQUENCE:
+                raise OverflowError("UUIDv7 generator sequence exhausted")
+            _sequence += 1
+        sequence = _sequence
+    data = bytearray(16)
+    data[0:6] = effective.to_bytes(6, "big")
+    data[6] = 0x70 | ((sequence >> 37) & 0x0F)
+    data[7] = (sequence >> 29) & 0xFF
+    data[8] = 0x80 | ((sequence >> 23) & 0x3F)
+    data[9] = (sequence >> 15) & 0xFF
+    data[10] = (sequence >> 7) & 0xFF
+    data[11] = ((sequence & 0x7F) << 1) | (random_bytes[11] & 0x01)
+    data[12:16] = random_bytes[12:16]
+    hex_text = data.hex()
+    return f"{hex_text[0:8]}-{hex_text[8:12]}-{hex_text[12:16]}-{hex_text[16:20]}-{hex_text[20:32]}"
+
+
+_SURROGATE_RE = re.compile("[\ud800-\udfff]")
+
+
+def sanitize_surrogates(text: str) -> str:
+    """Remove unpaired UTF-16 surrogates (they break JSON/UTF-8 encoding).
+
+    Properly paired surrogates are joined into their code point, so emoji survive.
+    """
+    if _SURROGATE_RE.search(text) is None:
+        return text
+    joined = text.encode("utf-16-le", "surrogatepass").decode("utf-16-le", "ignore")
+    return _SURROGATE_RE.sub("", joined)
+
+
+_user_agent: str | None = None
+
+
+def get_pi_user_agent() -> str:
+    global _user_agent
+    if _user_agent is None:
+        machine = platform.machine().lower()
+        arch = {"x86_64": "x64", "amd64": "x64", "aarch64": "arm64", "i386": "ia32", "i686": "ia32"}.get(
+            machine, machine
+        )
+        _user_agent = f"pi ({sys.platform} {platform.release()}; {arch})"
+    return _user_agent

pi_ai/utils/model_operations.py

@@ -0,0 +1,73 @@
+"""Model-type narrowing and operation result helpers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from ..types import AssistantImages, ClassifierContext, ClassifierModel, ClassifierResult, ImageModel, now_ms
+from .diagnostics import format_thrown_value
+from .models_error import ModelsError
+
+__all__ = [
+    "get_model_type",
+    "is_model_type",
+    "assert_chat_model",
+    "assert_image_model",
+    "assert_classifier_model",
+    "assert_classifier_input_supported",
+    "image_error_result",
+    "classifier_error_result",
+]
+
+
+def get_model_type(model: Any) -> str:
+    """The type of a model. Models without ``type`` are chat models."""
+    return getattr(model, "type", None) or "chat"
+
+
+def is_model_type(model: Any, type: str) -> bool:
+    return get_model_type(model) == type
+
+
+def assert_chat_model(model: Any) -> None:
+    if not is_model_type(model, "chat"):
+        raise ModelsError("provider", f"Model {model.provider}/{model.id} is not a chat model")
+
+
+def assert_image_model(model: Any) -> None:
+    if not is_model_type(model, "image"):
+        raise ModelsError("provider", f"Model {model.provider}/{model.id} is not an image model")
+
+
+def assert_classifier_model(model: Any) -> None:
+    if not is_model_type(model, "classifier"):
+        raise ModelsError("provider", f"Model {model.provider}/{model.id} is not a classifier model")
+
+
+def assert_classifier_input_supported(model: ClassifierModel, context: ClassifierContext) -> None:
+    if context.images and "image" not in model.input:
+        raise ModelsError("provider", f"Model {model.provider}/{model.id} does not accept image input")
+
+
+def image_error_result(model: ImageModel, error: Any, aborted: bool = False) -> AssistantImages:
+    return AssistantImages(
+        api=model.api,
+        provider=model.provider,
+        model=model.id,
+        output=[],
+        stop_reason="aborted" if aborted else "error",
+        error_message=format_thrown_value(error),
+        timestamp=now_ms(),
+    )
+
+
+def classifier_error_result(model: ClassifierModel, error: Any, aborted: bool = False) -> ClassifierResult:
+    return ClassifierResult(
+        api=model.api,
+        provider=model.provider,
+        model=model.id,
+        answers={},
+        stop_reason="aborted" if aborted else "error",
+        error_message=format_thrown_value(error),
+        timestamp=now_ms(),
+    )