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

pi_ai/api/system_one_shared.py

@@ -0,0 +1,148 @@
+"""TypeSafe System One protocol shared by the services that serve System One models."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Any, TypedDict
+
+from ..serialization import to_json_value
+from ..types import (
+    ClassifierAnswer,
+    ClassifierBoolAnswer,
+    ClassifierChoiceAnswer,
+    ClassifierContext,
+    ClassifierModel,
+    ClassifierQuestion,
+    ClassifierResult,
+    ClassifierScoreAnswer,
+    now_ms,
+)
+from ..utils.error_body import format_provider_error, normalize_provider_error
+from ..utils.model_operations import classifier_error_result
+from .classifier_shared import (
+    classifier_questions,
+    is_record,
+    parse_classifier_usage,
+    post_classifier_request,
+    required_number,
+)
+
+__all__ = ["SystemOneWireRequest", "SystemOneTransport", "wire_request", "classify_system_one"]
+
+
+class SystemOneWireRequest(TypedDict):
+    """TypeSafe System One request body without the transport-specific envelope."""
+
+    state: dict[str, Any]
+    questions: dict[str, Any]
+
+
+@dataclass(frozen=True, slots=True)
+class SystemOneTransport:
+    """Differences between services that serve System One models."""
+
+    #: Classifier API implemented by this transport.
+    api: str
+    #: Service name used in error messages.
+    label: str
+    #: Absolute request URL.
+    url: Callable[[ClassifierModel], str]
+    #: Wraps the System One request in the service's request envelope.
+    payload: Callable[[ClassifierModel, SystemOneWireRequest], Any]
+    #: Extracts the System One output (``{answers, usage}``) from the service's response envelope.
+    output: Callable[[Any], Mapping[str, Any]]
+
+
+def _probabilities(label: str, value: Any, question_id: str) -> dict[str, float]:
+    if not is_record(value):
+        raise ValueError(f"{label} returned invalid probabilities for {question_id}")
+    return {
+        key: required_number(label, probability, f"probability for {question_id}.{key}")
+        for key, probability in value.items()
+    }
+
+
+def _parse_answers(label: str, value: Any, questions: Mapping[str, ClassifierQuestion]) -> dict[str, ClassifierAnswer]:
+    if not is_record(value):
+        raise ValueError(f"{label} returned an unexpected response")
+    answers: dict[str, ClassifierAnswer] = {}
+    for question_id, question in questions.items():
+        answer = value.get(question_id)
+        if not is_record(answer):
+            raise ValueError(f"{label} did not return an answer for {question_id}")
+        if question.type == "choice":
+            choice = answer.get("choice")
+            if answer.get("type") != "choice" or not isinstance(choice, str):
+                raise ValueError(f"{label} did not return a choice answer for {question_id}")
+            answers[question_id] = ClassifierChoiceAnswer(
+                choice=choice,
+                probabilities=_probabilities(label, answer.get("probabilities"), question_id),
+                confidence=required_number(label, answer.get("confidence"), f"confidence for {question_id}"),
+            )
+        elif question.type == "score":
+            if answer.get("type") != "score":
+                raise ValueError(f"{label} did not return a score answer for {question_id}")
+            answers[question_id] = ClassifierScoreAnswer(
+                score=required_number(label, answer.get("score"), f"score for {question_id}"),
+                confidence=required_number(label, answer.get("confidence"), f"confidence for {question_id}"),
+            )
+        else:
+            if answer.get("type") != "noul":
+                raise ValueError(f"{label} did not return a bool answer for {question_id}")
+            answers[question_id] = ClassifierBoolAnswer(
+                probability=required_number(label, answer.get("noul"), f"probability for {question_id}")
+            )
+    return answers
+
+
+def wire_request(state: dict[str, Any], questions: Mapping[str, ClassifierQuestion]) -> SystemOneWireRequest:
+    """Map public ``bool`` questions to TypeSafe's wire-level ``noul`` type."""
+    wire_questions: dict[str, Any] = {}
+    for question_id, question in questions.items():
+        encoded = to_json_value(question)
+        wire_questions[question_id] = {**encoded, "type": "noul"} if question.type == "bool" else encoded
+    return {"state": state, "questions": wire_questions}
+
+
+async def classify_system_one(
+    transport: SystemOneTransport,
+    model: ClassifierModel,
+    context: ClassifierContext,
+    options: Mapping[str, Any] | None,
+) -> ClassifierResult:
+    """Run one System One classification over ``transport``. Never raises."""
+    options = options or {}
+    output = ClassifierResult(
+        api=model.api, provider=model.provider, model=model.id, answers={}, stop_reason="stop", timestamp=now_ms()
+    )
+    try:
+        if model.api != transport.api:
+            raise ValueError(f"Unsupported classifier API: {model.api}")
+        if context.images:
+            raise ValueError(f"{transport.label} does not support image input")
+        questions = classifier_questions(context)
+        body = await post_classifier_request(
+            transport.label,
+            transport.url(model),
+            model,
+            transport.payload(model, wire_request(context.state, questions)),
+            options,
+        )
+        result = transport.output(body)
+        # Set before parsing answers: a request with malformed answers was still billed.
+        usage = parse_classifier_usage(result.get("usage"), model)
+        if usage is not None:
+            output.usage = usage
+        output.answers = _parse_answers(transport.label, result.get("answers"), questions)
+        return output
+    except Exception as error:
+        signal = options.get("signal")
+        failed = classifier_error_result(
+            model,
+            format_provider_error(normalize_provider_error(error), f"{transport.label} error"),
+            aborted=signal is not None and signal.aborted,
+        )
+        failed.usage = output.usage
+        failed.timestamp = output.timestamp
+        return failed

pi_ai/api/transform_messages.py

@@ -0,0 +1,211 @@
+"""Cross-provider message normalization shared by API implementations."""
+
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+from ..types import (
+    AssistantMessage,
+    ImageContent,
+    Message,
+    Model,
+    TextContent,
+    ToolCall,
+    ToolResultMessage,
+)
+
+__all__ = ["SplitErrorToolResult", "split_error_tool_result_images", "transform_messages"]
+
+_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)"
+_ERROR_TOOL_IMAGE_PLACEHOLDER = "(image attached after the tool results)"
+
+
+@dataclass(slots=True)
+class SplitErrorToolResult:
+    """An error tool result split into text-only content and images to send after the tool results."""
+
+    content: list[TextContent]
+    caption: str
+    images: list[ImageContent]
+
+
+def split_error_tool_result_images(message: ToolResultMessage) -> SplitErrorToolResult | None:
+    """Split the images out of an error tool result (for APIs that reject images in error results)."""
+    if not message.is_error:
+        return None
+    images = [block for block in message.content if block.type == "image"]
+    if not images:
+        return None
+    return SplitErrorToolResult(
+        content=_replace_images(message.content, _ERROR_TOOL_IMAGE_PLACEHOLDER),
+        caption=f"Images from the error result of tool call {message.tool_call_id} ({message.tool_name}):",
+        images=images,  # type: ignore[arg-type]
+    )
+
+
+def _replace_images(content: list[Any], placeholder: str) -> list[Any]:
+    result: list[Any] = []
+    previous_was_placeholder = False
+    for block in content:
+        if block.type == "image":
+            if not previous_was_placeholder:
+                result.append(TextContent(placeholder))
+            previous_was_placeholder = True
+            continue
+        result.append(block)
+        previous_was_placeholder = block.text == placeholder
+    return result
+
+
+def _has_images(content: Any) -> bool:
+    return not isinstance(content, str) and any(block.type == "image" for block in content)
+
+
+def _downgrade_unsupported_images(messages: list[Any], model: Model) -> list[Any]:
+    if "image" in model.input:
+        return messages
+    result: list[Any] = []
+    for message in messages:
+        role = message.role
+        if role == "user" and _has_images(message.content):
+            message = dataclasses.replace(
+                message, content=_replace_images(message.content, _NON_VISION_USER_IMAGE_PLACEHOLDER)
+            )
+        elif role == "toolResult" and _has_images(message.content):
+            message = dataclasses.replace(
+                message, content=_replace_images(message.content, _NON_VISION_TOOL_IMAGE_PLACEHOLDER)
+            )
+        result.append(message)
+    return result
+
+
+def transform_messages(
+    messages: list[Message],
+    model: Model,
+    normalize_tool_call_id: Callable[[str, Model, AssistantMessage], str] | None = None,
+) -> list[Message]:
+    """Prepare a transcript for ``model``.
+
+    * Images are replaced by placeholders for models without image input.
+    * Thinking from another model is dropped; same-model thinking is kept for replay.
+    * Tool call ids are normalized for the target API, and tool results follow.
+    * Errored/aborted assistant turns are skipped.
+    * Orphaned tool calls get synthetic error results so APIs accept the history.
+    """
+    tool_call_id_map: dict[str, str] = {}
+    normalized: list[Any] = [
+        dataclasses.replace(message, content=[]) if getattr(message, "content", []) is None else message
+        for message in messages
+    ]
+    image_aware = _downgrade_unsupported_images(normalized, model)
+
+    transformed: list[Any] = []
+    for message in image_aware:
+        role = message.role
+        if role in ("system", "user"):
+            transformed.append(message)
+            continue
+        if role == "toolResult":
+            mapped = tool_call_id_map.get(message.tool_call_id)
+            if mapped and mapped != message.tool_call_id:
+                message = dataclasses.replace(message, tool_call_id=mapped)
+            transformed.append(message)
+            continue
+        if role != "assistant":
+            transformed.append(message)
+            continue
+
+        is_same_model = message.provider == model.provider and message.api == model.api and message.model == model.id
+        content: list[Any] = []
+        changed = False
+        for block in message.content:
+            kind = block.type
+            if kind == "thinking":
+                if not is_same_model:
+                    changed = True
+                    continue
+                if block.redacted or block.thinking_signature:
+                    content.append(block)
+                    continue
+                if not block.thinking or not block.thinking.strip():
+                    changed = True
+                    continue
+                content.append(block)
+            elif kind == "text":
+                if is_same_model or block.text_signature is None:
+                    content.append(block)
+                else:
+                    content.append(TextContent(block.text))
+                    changed = True
+            elif kind == "toolCall":
+                tool_call: ToolCall = block
+                if not is_same_model and tool_call.thought_signature:
+                    tool_call = dataclasses.replace(tool_call, thought_signature=None)
+                    changed = True
+                if not is_same_model and normalize_tool_call_id is not None:
+                    normalized_id = normalize_tool_call_id(tool_call.id, model, message)
+                    if normalized_id != tool_call.id:
+                        tool_call_id_map[tool_call.id] = normalized_id
+                        tool_call = dataclasses.replace(tool_call, id=normalized_id)
+                        changed = True
+                content.append(tool_call)
+            else:
+                content.append(block)
+        if changed:
+            message = message.shallow_copy()
+            message.content = content
+        transformed.append(message)
+
+    result: list[Any] = []
+    pending_tool_calls: list[ToolCall] = []
+    existing_result_ids: set[str] = set()
+    held_system: list[Any] = []
+
+    def close_pending() -> None:
+        nonlocal pending_tool_calls, existing_result_ids
+        if pending_tool_calls:
+            for tool_call in pending_tool_calls:
+                if tool_call.id not in existing_result_ids:
+                    result.append(
+                        ToolResultMessage(
+                            tool_call.id,
+                            tool_call.name,
+                            [TextContent("No result provided")],
+                            True,
+                        )
+                    )
+            pending_tool_calls = []
+            existing_result_ids = set()
+        result.extend(held_system)
+        held_system.clear()
+
+    for message in transformed:
+        role = message.role
+        if role == "assistant":
+            close_pending()
+            if message.stop_reason in ("error", "aborted"):
+                continue
+            tool_calls = [block for block in message.content if block.type == "toolCall"]
+            if tool_calls:
+                pending_tool_calls = tool_calls
+                existing_result_ids = set()
+            result.append(message)
+        elif role == "toolResult":
+            existing_result_ids.add(message.tool_call_id)
+            result.append(message)
+        elif role == "system":
+            if pending_tool_calls:
+                held_system.append(message)
+            else:
+                result.append(message)
+        elif role == "user":
+            close_pending()
+            result.append(message)
+        else:
+            result.append(message)
+    close_pending()
+    return result

pi_ai/api/typesafe_system_one.py

@@ -0,0 +1,43 @@
+"""TypeSafe's native System One protocol.
+
+OpenRouter (and other gateways) serve the same protocol, so those providers use
+this API with different base URLs.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+from urllib.parse import urljoin
+
+from ..types import ClassifierContext, ClassifierModel, ClassifierResult
+from .classifier_shared import is_record
+from .system_one_shared import SystemOneTransport, SystemOneWireRequest, classify_system_one
+
+__all__ = ["classify"]
+
+
+def _url(model: ClassifierModel) -> str:
+    return urljoin(f"{model.base_url.rstrip('/')}/", "systemone")
+
+
+def _payload(model: ClassifierModel, request: SystemOneWireRequest) -> Any:
+    return {"model": model.id, **request}
+
+
+def _output(body: Any) -> Mapping[str, Any]:
+    if not is_record(body):
+        raise ValueError("System One API returned an unexpected response")
+    return body
+
+
+_TRANSPORT = SystemOneTransport(
+    api="typesafe-system-one", label="System One API", url=_url, payload=_payload, output=_output
+)
+
+
+async def classify(
+    model: ClassifierModel, context: ClassifierContext, options: Mapping[str, Any] | None = None
+) -> ClassifierResult:
+    """TypeSafe System One classification with public ``bool`` values mapped to wire-level ``noul``."""
+    return await classify_system_one(_TRANSPORT, model, context, options)

pi_ai/auth/__init__.py

@@ -0,0 +1 @@
+"""Authentication: credentials, stores, resolution and OAuth flows."""

pi_ai/auth/context.py

@@ -0,0 +1,47 @@
+"""Default auth context backed by ``os.environ`` and the local filesystem."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Mapping
+
+from .types import AuthContext
+
+__all__ = ["DefaultAuthContext", "default_provider_auth_context", "overlay_env_auth_context"]
+
+
+class DefaultAuthContext:
+    """Env vars from ``os.environ`` (blank values count as unset); file checks via ``os.path``."""
+
+    async def env(self, name: str) -> str | None:
+        value = os.environ.get(name)
+        return value if value and value.strip() else None
+
+    async def file_exists(self, path: str) -> bool:
+        return os.path.exists(os.path.expanduser(path))
+
+
+def default_provider_auth_context() -> AuthContext:
+    return DefaultAuthContext()
+
+
+class _OverlayAuthContext:
+    __slots__ = ("_base", "_env")
+
+    def __init__(self, base: AuthContext, env: Mapping[str, str]) -> None:
+        self._base = base
+        self._env = env
+
+    async def env(self, name: str) -> str | None:
+        value = self._env.get(name)
+        if value:
+            return value
+        return await self._base.env(name)
+
+    async def file_exists(self, path: str) -> bool:
+        return await self._base.file_exists(path)
+
+
+def overlay_env_auth_context(base: AuthContext, env: Mapping[str, str]) -> AuthContext:
+    """An auth context where ``env`` values take precedence over ``base``."""
+    return _OverlayAuthContext(base, env)

pi_ai/auth/credential_store.py

@@ -0,0 +1,208 @@
+"""Credential stores: in-memory (default) and ``auth.json`` file-backed."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import copy
+import json
+import os
+import tempfile
+from collections.abc import Awaitable, Callable, Iterator, Sequence
+from pathlib import Path
+from typing import Any
+
+from ..abort import AbortSignal, race_with_abort_signal
+from .types import Credential, CredentialInfo
+
+__all__ = ["InMemoryCredentialStore", "FileCredentialStore"]
+
+try:  # POSIX advisory locks for cross-process serialization
+    import fcntl as _fcntl
+except ImportError:  # pragma: no cover - Windows
+    _fcntl = None  # type: ignore[assignment]
+
+
+class _ProviderLocks:
+    def __init__(self) -> None:
+        self._locks: dict[str, asyncio.Lock] = {}
+
+    def get(self, provider_id: str) -> asyncio.Lock:
+        lock = self._locks.get(provider_id)
+        if lock is None:
+            lock = asyncio.Lock()
+            self._locks[provider_id] = lock
+        return lock
+
+
+async def _enqueue(
+    locks: _ProviderLocks,
+    provider_id: str,
+    task: Callable[[], Awaitable[Any]],
+    signal: AbortSignal | None,
+) -> Any:
+    """Serialize ``task`` per provider. ``signal`` stops the caller waiting; started work completes."""
+    lock = locks.get(provider_id)
+
+    async def run() -> Any:
+        async with lock:
+            if signal is not None:
+                signal.throw_if_aborted()
+            return await task()
+
+    return await race_with_abort_signal(run(), signal)
+
+
+class InMemoryCredentialStore:
+    """Default in-memory credential store; writes are serialized per provider id."""
+
+    def __init__(self, initial: dict[str, Credential] | None = None) -> None:
+        self._credentials: dict[str, Credential] = copy.deepcopy(initial) if initial else {}
+        self._locks = _ProviderLocks()
+
+    async def read(self, provider_id: str, *, signal: AbortSignal | None = None) -> Credential | None:
+        if signal is not None:
+            signal.throw_if_aborted()
+        credential = self._credentials.get(provider_id)
+        return copy.deepcopy(credential) if credential is not None else None
+
+    async def list(self, *, signal: AbortSignal | None = None) -> Sequence[CredentialInfo]:
+        if signal is not None:
+            signal.throw_if_aborted()
+        return [
+            CredentialInfo(provider_id, credential.get("type", "api_key"))
+            for provider_id, credential in self._credentials.items()
+        ]
+
+    async def modify(
+        self,
+        provider_id: str,
+        fn: Callable[[Credential | None], Awaitable[Credential | None]],
+        *,
+        signal: AbortSignal | None = None,
+    ) -> Credential | None:
+        async def task() -> Credential | None:
+            current = self._credentials.get(provider_id)
+            following = await fn(copy.deepcopy(current) if current is not None else None)
+            if signal is not None:
+                signal.throw_if_aborted()
+            if following is not None:
+                self._credentials[provider_id] = copy.deepcopy(following)
+                return following
+            return current
+
+        return await _enqueue(self._locks, provider_id, task, signal)
+
+    async def delete(self, provider_id: str, *, signal: AbortSignal | None = None) -> None:
+        async def task() -> None:
+            self._credentials.pop(provider_id, None)
+
+        await _enqueue(self._locks, provider_id, task, signal)
+
+
+class FileCredentialStore:
+    """Credentials in an ``auth.json`` file shared with TypeScript pi.
+
+    The file maps provider ids to credentials (``{"anthropic": {"type": "oauth", ...}}``)
+    and is written atomically with ``0600`` permissions. On POSIX, ``modify`` also takes
+    an advisory file lock so concurrent processes cannot double-refresh OAuth tokens.
+    """
+
+    def __init__(self, path: str | os.PathLike[str] = "auth.json") -> None:
+        self.path = Path(path).expanduser()
+        self._locks = _ProviderLocks()
+
+    def _load(self) -> dict[str, Credential]:
+        try:
+            data = json.loads(self.path.read_text("utf-8"))
+        except FileNotFoundError:
+            return {}
+        if not isinstance(data, dict):
+            raise ValueError(f"{self.path} does not contain a JSON object")
+        return data
+
+    def _save(self, data: dict[str, Credential]) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        fd, tmp = tempfile.mkstemp(dir=self.path.parent, prefix=".auth-", suffix=".json")
+        try:
+            os.chmod(tmp, 0o600)
+            with os.fdopen(fd, "w", encoding="utf-8") as handle:
+                json.dump(data, handle, ensure_ascii=False, indent=2)
+                handle.write("\n")
+            os.replace(tmp, self.path)
+        except BaseException:
+            with contextlib.suppress(FileNotFoundError):
+                os.unlink(tmp)
+            raise
+
+    @contextlib.contextmanager
+    def _file_lock(self) -> Iterator[None]:
+        if _fcntl is None:
+            yield
+            return
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        lock_path = self.path.with_name(self.path.name + ".lock")
+        with open(lock_path, "a+") as handle:
+            _fcntl.flock(handle.fileno(), _fcntl.LOCK_EX)
+            try:
+                yield
+            finally:
+                _fcntl.flock(handle.fileno(), _fcntl.LOCK_UN)
+
+    async def read(self, provider_id: str, *, signal: AbortSignal | None = None) -> Credential | None:
+        if signal is not None:
+            signal.throw_if_aborted()
+        data = await asyncio.to_thread(self._load)
+        credential = data.get(provider_id)
+        return credential if isinstance(credential, dict) else None
+
+    async def list(self, *, signal: AbortSignal | None = None) -> Sequence[CredentialInfo]:
+        if signal is not None:
+            signal.throw_if_aborted()
+        data = await asyncio.to_thread(self._load)
+        return [
+            CredentialInfo(provider_id, credential.get("type", "api_key"))
+            for provider_id, credential in data.items()
+            if isinstance(credential, dict)
+        ]
+
+    async def modify(
+        self,
+        provider_id: str,
+        fn: Callable[[Credential | None], Awaitable[Credential | None]],
+        *,
+        signal: AbortSignal | None = None,
+    ) -> Credential | None:
+        loop = asyncio.get_running_loop()
+
+        async def task() -> Credential | None:
+            lock = self._file_lock()
+            await asyncio.to_thread(lock.__enter__)
+            try:
+                data = await asyncio.to_thread(self._load)
+                current = data.get(provider_id)
+                following = await fn(copy.deepcopy(current) if isinstance(current, dict) else None)
+                if signal is not None:
+                    signal.throw_if_aborted()
+                if following is None:
+                    return current if isinstance(current, dict) else None
+                data[provider_id] = following
+                await asyncio.to_thread(self._save, data)
+                return following
+            finally:
+                await loop.run_in_executor(None, lock.__exit__, None, None, None)
+
+        return await _enqueue(self._locks, provider_id, task, signal)
+
+    async def delete(self, provider_id: str, *, signal: AbortSignal | None = None) -> None:
+        async def task() -> None:
+            lock = self._file_lock()
+            await asyncio.to_thread(lock.__enter__)
+            try:
+                data = await asyncio.to_thread(self._load)
+                if data.pop(provider_id, None) is not None:
+                    await asyncio.to_thread(self._save, data)
+            finally:
+                await asyncio.to_thread(lock.__exit__, None, None, None)
+
+        await _enqueue(self._locks, provider_id, task, signal)