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

pi_ai/models_store.py

@@ -0,0 +1,121 @@
+"""Persistent storage for dynamic model catalogs."""
+
+from __future__ import annotations
+
+import asyncio
+import copy
+import json
+import os
+import tempfile
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Protocol
+
+from .abort import AbortSignal
+
+__all__ = ["ModelsStoreEntry", "ModelsStore", "InMemoryModelsStore", "FileModelsStore"]
+
+
+@dataclass(slots=True)
+class ModelsStoreEntry:
+    #: Persisted models of every type.
+    models: list[Any] = field(default_factory=list)
+    #: Unix timestamp (ms) from the remote catalog's Last-Modified header.
+    last_modified: int | None = None
+    #: Unix timestamp (ms) of the last completed remote check.
+    checked_at: int | None = None
+    #: Opaque validator from the remote catalog's ETag header, echoed back as If-None-Match.
+    etag: str | None = None
+
+
+class ModelsStore(Protocol):
+    """Persistent model catalogs keyed by provider id."""
+
+    async def read(self, provider_id: str, *, signal: AbortSignal | None = None) -> ModelsStoreEntry | None: ...
+
+    async def write(self, provider_id: str, entry: ModelsStoreEntry, *, signal: AbortSignal | None = None) -> None: ...
+
+    async def delete(self, provider_id: str, *, signal: AbortSignal | None = None) -> None: ...
+
+
+class InMemoryModelsStore:
+    """Default in-process store. Entries are deep-copied on read and write."""
+
+    def __init__(self) -> None:
+        self._entries: dict[str, ModelsStoreEntry] = {}
+
+    async def read(self, provider_id: str, *, signal: AbortSignal | None = None) -> ModelsStoreEntry | None:
+        if signal is not None:
+            signal.throw_if_aborted()
+        entry = self._entries.get(provider_id)
+        return copy.deepcopy(entry) if entry is not None else None
+
+    async def write(self, provider_id: str, entry: ModelsStoreEntry, *, signal: AbortSignal | None = None) -> None:
+        if signal is not None:
+            signal.throw_if_aborted()
+        self._entries[provider_id] = copy.deepcopy(entry)
+
+    async def delete(self, provider_id: str, *, signal: AbortSignal | None = None) -> None:
+        if signal is not None:
+            signal.throw_if_aborted()
+        self._entries.pop(provider_id, None)
+
+
+class FileModelsStore:
+    """JSON-file store (``{"<provider>": {"models": [...], "checkedAt": ...}}``) in pi's format."""
+
+    def __init__(self, path: str | os.PathLike[str]) -> None:
+        self.path = Path(path).expanduser()
+        self._lock = asyncio.Lock()
+
+    def _load(self) -> dict[str, Any]:
+        try:
+            return json.loads(self.path.read_text("utf-8"))
+        except FileNotFoundError:
+            return {}
+
+    def _save(self, data: dict[str, Any]) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        fd, tmp = tempfile.mkstemp(dir=self.path.parent, prefix=".models-", suffix=".json")
+        with os.fdopen(fd, "w", encoding="utf-8") as handle:
+            json.dump(data, handle, ensure_ascii=False, indent=2)
+        os.replace(tmp, self.path)
+
+    async def read(self, provider_id: str, *, signal: AbortSignal | None = None) -> ModelsStoreEntry | None:
+        from .serialization import model_from_dict
+
+        if signal is not None:
+            signal.throw_if_aborted()
+        raw = (await asyncio.to_thread(self._load)).get(provider_id)
+        if not isinstance(raw, dict):
+            return None
+        models = []
+        for item in raw.get("models", []):
+            try:
+                models.append(model_from_dict(item))
+            except ValueError:
+                continue  # unknown model types from newer versions
+        return ModelsStoreEntry(
+            models=models,
+            last_modified=raw.get("lastModified"),
+            checked_at=raw.get("checkedAt"),
+            etag=raw.get("etag"),
+        )
+
+    async def write(self, provider_id: str, entry: ModelsStoreEntry, *, signal: AbortSignal | None = None) -> None:
+        from .serialization import to_json_value
+
+        if signal is not None:
+            signal.throw_if_aborted()
+        async with self._lock:
+            data = await asyncio.to_thread(self._load)
+            data[provider_id] = to_json_value(entry)
+            await asyncio.to_thread(self._save, data)
+
+    async def delete(self, provider_id: str, *, signal: AbortSignal | None = None) -> None:
+        if signal is not None:
+            signal.throw_if_aborted()
+        async with self._lock:
+            data = await asyncio.to_thread(self._load)
+            if data.pop(provider_id, None) is not None:
+                await asyncio.to_thread(self._save, data)

pi_ai/oauth.py

@@ -0,0 +1,90 @@
+"""Type-only compatibility entry point for coding-agent extension OAuth declarations.
+
+These legacy callback shapes are kept so extensions written against the older
+``OAuthLoginCallbacks`` surface keep type-checking. Built-in flows use
+:class:`pi_ai.OAuthAuth` with :class:`pi_ai.ProviderAuthInteraction` instead, and
+this module never imports them.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Protocol
+
+from .auth.oauth._common import OAuthError
+from .auth.types import OAuthCredentials
+
+if TYPE_CHECKING:
+    from .abort import AbortSignal
+
+__all__ = [
+    "OAuthError",
+    "OAuthAuthInfo",
+    "OAuthCredentials",
+    "OAuthDeviceCodeInfo",
+    "OAuthLoginCallbacks",
+    "OAuthPrompt",
+    "OAuthSelectOption",
+    "OAuthSelectPrompt",
+]
+
+
+@dataclass(slots=True)
+class OAuthPrompt:
+    """Legacy extension OAuth prompt."""
+
+    message: str
+    placeholder: str | None = None
+    allow_empty: bool | None = None
+
+
+@dataclass(slots=True)
+class OAuthAuthInfo:
+    """Legacy extension OAuth authorization link."""
+
+    url: str
+    instructions: str | None = None
+
+
+@dataclass(slots=True)
+class OAuthDeviceCodeInfo:
+    """Legacy extension OAuth device-code notification."""
+
+    user_code: str
+    verification_uri: str
+    interval_seconds: float | None = None
+    expires_in_seconds: float | None = None
+
+
+@dataclass(slots=True)
+class OAuthSelectOption:
+    id: str
+    label: str
+
+
+@dataclass(slots=True)
+class OAuthSelectPrompt:
+    message: str
+    options: list[OAuthSelectOption] = field(default_factory=list)
+
+
+class OAuthLoginCallbacks(Protocol):
+    """Callback surface retained only for coding-agent extension compatibility.
+
+    ``on_progress`` and ``on_manual_code_input`` are optional (``None`` when unsupported).
+    """
+
+    #: Optional progress message sink.
+    on_progress: Callable[[str], None] | None
+    #: Optional manual code entry that races the browser callback.
+    on_manual_code_input: Callable[[], Awaitable[str]] | None
+    signal: AbortSignal | None
+
+    def on_auth(self, info: OAuthAuthInfo) -> None: ...
+
+    def on_device_code(self, info: OAuthDeviceCodeInfo) -> None: ...
+
+    async def on_prompt(self, prompt: OAuthPrompt) -> str: ...
+
+    async def on_select(self, prompt: OAuthSelectPrompt) -> str | None: ...

pi_ai/providers/__init__.py

@@ -0,0 +1 @@
+"""Provider factories. Import ``pi_ai.providers.all`` for every built-in provider."""

pi_ai/providers/all.py

@@ -0,0 +1,154 @@
+"""Every built-in provider, plus static reads of the generated catalog.
+
+::
+
+    from pi_ai.providers.all import builtin_models
+
+    models = builtin_models()
+    model = models.get_model("anthropic", "claude-sonnet-4-5")
+
+Building providers is cheap: catalogs load on first access to a provider's models
+and API implementations on the first request.
+"""
+
+from __future__ import annotations
+
+from typing import TypeAlias
+
+from ..auth.types import AuthContext, CredentialStore
+from ..models import Models, Provider, create_models
+from ..models_store import ModelsStore
+from .amazon_bedrock import amazon_bedrock_provider
+from .ant_ling import ant_ling_provider
+from .anthropic import anthropic_provider
+from .azure import azure_provider
+from .baseten import baseten_provider
+from .catalog import (
+    BUILTIN_PROVIDERS,
+    get_all_builtin_models,
+    get_builtin_classifier_model,
+    get_builtin_classifier_models,
+    get_builtin_image_model,
+    get_builtin_image_models,
+    get_builtin_model,
+    get_builtin_model_data_generated_at,
+    get_builtin_models,
+    get_builtin_providers,
+)
+from .cerebras import cerebras_provider
+from .cloudflare_ai_gateway import cloudflare_ai_gateway_provider
+from .cloudflare_workers_ai import cloudflare_workers_ai_provider
+from .deepseek import deepseek_provider
+from .fireworks import fireworks_provider
+from .github_copilot import github_copilot_provider
+from .google import google_provider
+from .google_vertex import google_vertex_provider
+from .groq import groq_provider
+from .huggingface import huggingface_provider
+from .kimi_coding import kimi_coding_provider
+from .meta import meta_provider
+from .minimax import minimax_provider
+from .minimax_cn import minimax_cn_provider
+from .mistral import mistral_provider
+from .moonshotai import moonshotai_provider
+from .moonshotai_cn import moonshotai_cn_provider
+from .nvidia import nvidia_provider
+from .openai import openai_provider
+from .openai_codex import openai_codex_provider
+from .opencode import opencode_provider
+from .opencode_go import opencode_go_provider
+from .openrouter import openrouter_provider
+from .qwen_token_plan import qwen_token_plan_provider
+from .qwen_token_plan_cn import qwen_token_plan_cn_provider
+from .qwen_token_plan_individual import qwen_token_plan_individual_provider
+from .radius import radius_provider
+from .together import together_provider
+from .typesafe import typesafe_provider
+from .vercel_ai_gateway import vercel_ai_gateway_provider
+from .xai import xai_provider
+from .xiaomi import xiaomi_provider
+from .xiaomi_token_plan_ams import xiaomi_token_plan_ams_provider
+from .xiaomi_token_plan_cn import xiaomi_token_plan_cn_provider
+from .xiaomi_token_plan_sgp import xiaomi_token_plan_sgp_provider
+from .zai import zai_provider
+from .zai_coding_cn import zai_coding_cn_provider
+
+__all__ = [
+    "BuiltinProvider",
+    "BUILTIN_PROVIDERS",
+    "builtin_providers",
+    "builtin_models",
+    "radius_provider",
+    "get_builtin_model",
+    "get_builtin_image_model",
+    "get_builtin_classifier_model",
+    "get_builtin_providers",
+    "get_builtin_model_data_generated_at",
+    "get_builtin_models",
+    "get_builtin_image_models",
+    "get_builtin_classifier_models",
+    "get_all_builtin_models",
+]
+
+#: Providers present in the generated catalog (see :data:`BUILTIN_PROVIDERS`).
+BuiltinProvider: TypeAlias = str
+
+
+def builtin_providers() -> list[Provider]:
+    """All built-in providers, freshly constructed."""
+    return [
+        amazon_bedrock_provider(),
+        ant_ling_provider(),
+        anthropic_provider(),
+        azure_provider(),
+        baseten_provider(),
+        cerebras_provider(),
+        cloudflare_ai_gateway_provider(),
+        cloudflare_workers_ai_provider(),
+        deepseek_provider(),
+        fireworks_provider(),
+        github_copilot_provider(),
+        google_provider(),
+        google_vertex_provider(),
+        groq_provider(),
+        huggingface_provider(),
+        kimi_coding_provider(),
+        meta_provider(),
+        minimax_provider(),
+        minimax_cn_provider(),
+        mistral_provider(),
+        moonshotai_provider(),
+        moonshotai_cn_provider(),
+        nvidia_provider(),
+        openai_provider(),
+        openai_codex_provider(),
+        opencode_provider(),
+        opencode_go_provider(),
+        openrouter_provider(),
+        qwen_token_plan_provider(),
+        qwen_token_plan_cn_provider(),
+        qwen_token_plan_individual_provider(),
+        radius_provider(),
+        together_provider(),
+        typesafe_provider(),
+        vercel_ai_gateway_provider(),
+        xai_provider(),
+        xiaomi_provider(),
+        xiaomi_token_plan_ams_provider(),
+        xiaomi_token_plan_cn_provider(),
+        xiaomi_token_plan_sgp_provider(),
+        zai_provider(),
+        zai_coding_cn_provider(),
+    ]
+
+
+def builtin_models(
+    credentials: CredentialStore | None = None,
+    models_store: ModelsStore | None = None,
+    auth_context: AuthContext | None = None,
+) -> Models:
+    """A :class:`~pi_ai.Models` collection with every built-in provider registered."""
+    models = create_models(credentials=credentials, models_store=models_store, auth_context=auth_context)
+    for provider in builtin_providers():
+        models.set_provider(provider)
+    return models

pi_ai/providers/amazon_bedrock.py

@@ -0,0 +1,116 @@
+"""Amazon Bedrock provider."""
+
+from __future__ import annotations
+
+from ..abort import AbortSignal
+from ..api.registry import bedrock_converse_stream_api
+from ..auth.types import (
+    ApiKeyAuth,
+    AuthContext,
+    AuthEvent,
+    AuthInfoLink,
+    AuthPrompt,
+    AuthResult,
+    AuthSelectOption,
+    Credential,
+    ModelAuth,
+    ProviderAuth,
+    ProviderAuthInteraction,
+)
+from ..models import Provider, create_provider
+from .catalog import catalog_models
+
+__all__ = ["amazon_bedrock_provider", "bedrock_auth"]
+
+
+async def _login(interaction: ProviderAuthInteraction) -> Credential:
+    interaction.signal.throw_if_aborted()
+    method = await interaction.prompt(
+        AuthPrompt(
+            type="select",
+            message="Select Amazon Bedrock authentication method:",
+            options=[
+                AuthSelectOption(id="bearer-token", label="Bearer token"),
+                AuthSelectOption(id="aws-profile", label="AWS profile"),
+                AuthSelectOption(id="credential-chain", label="Existing AWS credential chain"),
+            ],
+        )
+    )
+    interaction.signal.throw_if_aborted()
+    if method == "bearer-token":
+        return {
+            "type": "api_key",
+            "key": await interaction.prompt(AuthPrompt(type="secret", message="Enter Amazon Bedrock bearer token")),
+        }
+    interaction.notify(
+        AuthEvent(
+            type="info",
+            message="Amazon Bedrock supports AWS profiles, IAM credentials, and role-based credentials.",
+            links=[
+                AuthInfoLink(
+                    label="AWS credential provider chain",
+                    url="https://docs.aws.amazon.com/sdkref/latest/guide/standardized-credentials.html",
+                )
+            ],
+        )
+    )
+    if method == "aws-profile":
+        return {
+            "type": "api_key",
+            "env": {"AWS_PROFILE": await interaction.prompt(AuthPrompt(type="text", message="Enter AWS profile name"))},
+        }
+    if method != "credential-chain":
+        raise ValueError(f"Unknown Amazon Bedrock auth method: {method}")
+    await interaction.prompt(AuthPrompt(type="text", message="Configure AWS credentials, then press Enter to continue"))
+    return {"type": "api_key"}
+
+
+async def _resolve(*, ctx: AuthContext, credential: Credential | None, signal: AbortSignal) -> AuthResult | None:
+    async def env(name: str) -> str | None:
+        signal.throw_if_aborted()
+        value = await ctx.env(name)
+        signal.throw_if_aborted()
+        return value
+
+    if credential and credential.get("key"):
+        return AuthResult(ModelAuth(api_key=credential["key"]), env=credential.get("env"), source="stored credential")
+    if await env("AWS_BEARER_TOKEN_BEDROCK"):
+        return AuthResult(ModelAuth(), source="AWS_BEARER_TOKEN_BEDROCK")
+    credential_env = credential.get("env") if credential else None
+    stored_profile = credential_env.get("AWS_PROFILE") if credential_env else None
+    profile = stored_profile if stored_profile is not None else await env("AWS_PROFILE")
+    if profile:
+        return AuthResult(
+            ModelAuth(),
+            env=credential_env,
+            source="stored credential" if stored_profile else "AWS_PROFILE",
+        )
+    if (await env("AWS_ACCESS_KEY_ID")) and (await env("AWS_SECRET_ACCESS_KEY")):
+        return AuthResult(ModelAuth(), source="AWS access keys")
+    if await env("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"):
+        return AuthResult(ModelAuth(), source="ECS task role")
+    if await env("AWS_CONTAINER_CREDENTIALS_FULL_URI"):
+        return AuthResult(ModelAuth(), source="ECS task role")
+    if await env("AWS_WEB_IDENTITY_TOKEN_FILE"):
+        return AuthResult(ModelAuth(), source="web identity token")
+    return None
+
+
+def bedrock_auth() -> ApiKeyAuth:
+    """Bedrock accepts a bearer token or the AWS default credential chain.
+
+    The login flow can store a token/profile choice; ``resolve`` also detects ambient
+    AWS credentials without copying them into pi's credential store.
+    """
+    return ApiKeyAuth("AWS credentials or bearer token", _resolve, login=_login)
+
+
+def amazon_bedrock_provider() -> Provider:
+    """Amazon Bedrock (Converse Stream API) with bearer-token or AWS credential-chain auth."""
+    return create_provider(
+        id="amazon-bedrock",
+        name="Amazon Bedrock",
+        auth=ProviderAuth(api_key=bedrock_auth()),
+        models=lambda: catalog_models("amazon-bedrock"),
+        api=bedrock_converse_stream_api(),
+    )

pi_ai/providers/ant_ling.py

@@ -0,0 +1,23 @@
+"""Ant Ling provider."""
+
+from __future__ import annotations
+
+from ..api.registry import openai_completions_api
+from ..auth.helpers import env_api_key_auth
+from ..auth.types import ProviderAuth
+from ..models import Provider, create_provider
+from .catalog import catalog_models
+
+__all__ = ["ant_ling_provider"]
+
+
+def ant_ling_provider() -> Provider:
+    """Ant Ling (OpenAI-compatible Chat Completions), authenticated with ``ANT_LING_API_KEY``."""
+    return create_provider(
+        id="ant-ling",
+        name="Ant Ling",
+        base_url="https://api.ant-ling.com/v1",
+        auth=ProviderAuth(api_key=env_api_key_auth("Ant Ling API key", ["ANT_LING_API_KEY"])),
+        models=lambda: catalog_models("ant-ling"),
+        api=openai_completions_api(),
+    )

pi_ai/providers/anthropic.py

@@ -0,0 +1,105 @@
+"""Anthropic provider."""
+
+from __future__ import annotations
+
+from ..abort import AbortSignal
+from ..api.registry import anthropic_messages_api
+from ..auth.helpers import lazy_oauth
+from ..auth.oauth.load import load_anthropic_oauth
+from ..auth.types import (
+    ApiKeyAuth,
+    AuthContext,
+    AuthPrompt,
+    AuthResult,
+    Credential,
+    ModelAuth,
+    ProviderAuth,
+    ProviderAuthInteraction,
+)
+from ..env_api_keys import (
+    ANTHROPIC_API_KEY_ENV,
+    ANTHROPIC_AUTH_TOKEN_ENV,
+    ANTHROPIC_FEDERATION_RULE_ID_ENV,
+    ANTHROPIC_IDENTITY_TOKEN_FILE_ENV,
+    ANTHROPIC_OAUTH_TOKEN_ENV,
+    ANTHROPIC_ORGANIZATION_ID_ENV,
+    ANTHROPIC_SERVICE_ACCOUNT_ID_ENV,
+    ANTHROPIC_WORKSPACE_ID_ENV,
+)
+from ..models import Provider, create_provider
+from .catalog import catalog_models
+
+__all__ = ["anthropic_provider", "anthropic_api_key_auth"]
+
+
+def anthropic_api_key_auth() -> ApiKeyAuth:
+    """Anthropic API-key auth.
+
+    Order: stored key, ``ANTHROPIC_AUTH_TOKEN`` (sent as ``Authorization: Bearer``),
+    ``ANTHROPIC_OAUTH_TOKEN``, ``ANTHROPIC_API_KEY``, then workload identity federation.
+    """
+
+    async def login(interaction: ProviderAuthInteraction) -> Credential:
+        interaction.signal.throw_if_aborted()
+        key = await interaction.prompt(AuthPrompt(type="secret", message="Enter Anthropic API key"))
+        interaction.signal.throw_if_aborted()
+        return {"type": "api_key", "key": key}
+
+    async def resolve(*, ctx: AuthContext, credential: Credential | None, signal: AbortSignal) -> AuthResult | None:
+        signal.throw_if_aborted()
+        if credential and credential.get("key"):
+            return AuthResult(
+                ModelAuth(api_key=credential["key"]), env=credential.get("env"), source="stored credential"
+            )
+
+        auth_token = await ctx.env(ANTHROPIC_AUTH_TOKEN_ENV)
+        signal.throw_if_aborted()
+        if auth_token:
+            return AuthResult(
+                ModelAuth(headers={"Authorization": f"Bearer {auth_token}"}), source=ANTHROPIC_AUTH_TOKEN_ENV
+            )
+
+        for env_var in (ANTHROPIC_OAUTH_TOKEN_ENV, ANTHROPIC_API_KEY_ENV):
+            api_key = await ctx.env(env_var)
+            signal.throw_if_aborted()
+            if api_key:
+                return AuthResult(ModelAuth(api_key=api_key), source=env_var)
+
+        # Workload identity federation: the API implementation exchanges the identity
+        # token for a short-lived access token and refreshes it itself. Last in
+        # line so keys and ANTHROPIC_AUTH_TOKEN keep winning, as in the SDK. The
+        # ids are provider config rather than auth, so they travel in `env`.
+        federation: dict[str, str] = {}
+        for env_var in (
+            ANTHROPIC_FEDERATION_RULE_ID_ENV,
+            ANTHROPIC_ORGANIZATION_ID_ENV,
+            ANTHROPIC_IDENTITY_TOKEN_FILE_ENV,
+        ):
+            value = await ctx.env(env_var)
+            signal.throw_if_aborted()
+            if not value:
+                return None
+            federation[env_var] = value
+        for env_var in (ANTHROPIC_SERVICE_ACCOUNT_ID_ENV, ANTHROPIC_WORKSPACE_ID_ENV):
+            value = await ctx.env(env_var)
+            signal.throw_if_aborted()
+            if value:
+                federation[env_var] = value
+        return AuthResult(ModelAuth(), env=federation, source="workload identity federation")
+
+    return ApiKeyAuth("Anthropic API key", resolve, login=login)
+
+
+def anthropic_provider() -> Provider:
+    """Anthropic (Messages API): API key, ``ANTHROPIC_AUTH_TOKEN``, federation, or Claude Pro/Max OAuth."""
+    return create_provider(
+        id="anthropic",
+        name="Anthropic",
+        base_url="https://api.anthropic.com",
+        auth=ProviderAuth(
+            api_key=anthropic_api_key_auth(),
+            oauth=lazy_oauth("Anthropic (Claude Pro/Max)", load_anthropic_oauth, is_subscription=True),
+        ),
+        models=lambda: catalog_models("anthropic"),
+        api=anthropic_messages_api(),
+    )