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

pi_ai/providers/azure.py

@@ -0,0 +1,103 @@
+"""Azure OpenAI provider."""
+
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Mapping
+from typing import Any
+
+from ..api.lazy import lazy_stream
+from ..api.registry import azure_openai_responses_api, openai_completions_api
+from ..auth.helpers import env_api_key_auth
+from ..auth.types import ProviderAuth
+from ..event_stream import AssistantMessageEventStream
+from ..models import Provider, create_provider
+from ..types import Model, TranscriptContext
+from ..utils.callbacks import maybe_await
+from .catalog import catalog_models
+
+__all__ = ["azure_provider"]
+
+
+def _resolve_azure_model(model: Model, options: Mapping[str, Any] | None) -> Model:
+    from ..api.azure_openai_config import resolve_azure_base_url
+
+    return dataclasses.replace(model, base_url=resolve_azure_base_url(model, options))
+
+
+def _with_deployment_name(model: Model, options: Mapping[str, Any] | None) -> Mapping[str, Any] | None:
+    """Send the deployment name as the request's model, keeping ``model.id`` as the catalog id."""
+    from ..api.azure_openai_config import resolve_deployment_name
+
+    deployment_name = resolve_deployment_name(model, options)
+    if deployment_name == model.id:
+        return options
+    user_on_payload = options.get("on_payload") if options else None
+
+    async def on_payload(payload: Any, payload_model: Any) -> Any:
+        params = {**payload, "model": deployment_name}
+        if user_on_payload is not None:
+            replaced = await maybe_await(user_on_payload(params, payload_model))
+            if replaced is not None:
+                return replaced
+        return params
+
+    return {**(options or {}), "on_payload": on_payload}
+
+
+class _AzureStreams:
+    """Resolve the Azure endpoint and deployment before dispatch.
+
+    Resolution runs inside ``lazy_stream`` so an unconfigured endpoint errors on the
+    stream instead of raising out of ``stream()``.
+    """
+
+    __slots__ = ("_streams",)
+
+    def __init__(self, streams: Any) -> None:
+        self._streams = streams
+
+    def stream(
+        self, model: Model, context: TranscriptContext, options: Mapping[str, Any] | None = None
+    ) -> AssistantMessageEventStream:
+        async def setup() -> AssistantMessageEventStream:
+            return self._streams.stream(
+                _resolve_azure_model(model, options), context, _with_deployment_name(model, options)
+            )
+
+        return lazy_stream(model, setup)
+
+    def stream_simple(
+        self, model: Model, context: TranscriptContext, options: Mapping[str, Any] | None = None
+    ) -> AssistantMessageEventStream:
+        async def setup() -> AssistantMessageEventStream:
+            return self._streams.stream_simple(
+                _resolve_azure_model(model, options), context, _with_deployment_name(model, options)
+            )
+
+        return lazy_stream(model, setup)
+
+    def __repr__(self) -> str:
+        return f"azure_streams({self._streams!r})"
+
+
+def azure_streams(streams: Any) -> Any:
+    return _AzureStreams(streams)
+
+
+def azure_provider() -> Provider:
+    """Azure OpenAI (Responses and Chat Completions), authenticated with ``AZURE_OPENAI_API_KEY``.
+
+    Catalog models ship without a base URL: the endpoint comes from ``AZURE_OPENAI_BASE_URL``,
+    ``AZURE_OPENAI_RESOURCE_NAME`` or the matching request options.
+    """
+    return create_provider(
+        id="azure",
+        name="Azure",
+        auth=ProviderAuth(api_key=env_api_key_auth("Azure OpenAI API key", ["AZURE_OPENAI_API_KEY"])),
+        models=lambda: catalog_models("azure"),
+        api={
+            "azure-openai-responses": azure_openai_responses_api(),
+            "openai-completions": azure_streams(openai_completions_api()),
+        },
+    )

pi_ai/providers/baseten.py

@@ -0,0 +1,23 @@
+"""Baseten 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__ = ["baseten_provider"]
+
+
+def baseten_provider() -> Provider:
+    """Baseten (OpenAI-compatible Chat Completions), authenticated with ``BASETEN_API_KEY``."""
+    return create_provider(
+        id="baseten",
+        name="Baseten",
+        base_url="https://inference.baseten.co/v1",
+        auth=ProviderAuth(api_key=env_api_key_auth("Baseten API key", ["BASETEN_API_KEY"])),
+        models=lambda: catalog_models("baseten"),
+        api=openai_completions_api(),
+    )

pi_ai/providers/catalog.py

@@ -0,0 +1,167 @@
+"""Built-in model catalog data.
+
+Catalog files (``data/<provider>.json``) are the generated model data shipped by
+``@earendil-works/pi-ai``. Each file groups entries by API id and keys them as
+``"<type>:<model id>"``. Files load on first access and are cached per process.
+
+Returned models are shared instances, as in pi. Derive variants with
+``dataclasses.replace(model, ...)`` instead of mutating them.
+"""
+
+from __future__ import annotations
+
+import json
+import threading
+from datetime import datetime
+from functools import lru_cache
+from importlib import resources
+from typing import Any
+
+from ..serialization import model_from_dict
+from ..types import AnyModel, ClassifierModel, ImageModel, Model
+
+__all__ = [
+    "BUILTIN_PROVIDERS",
+    "load_catalog",
+    "catalog_models",
+    "get_builtin_providers",
+    "get_builtin_model",
+    "get_builtin_image_model",
+    "get_builtin_classifier_model",
+    "get_builtin_models",
+    "get_builtin_image_models",
+    "get_builtin_classifier_models",
+    "get_all_builtin_models",
+    "get_builtin_model_data_generated_at",
+]
+
+#: Providers present in the generated catalog.
+BUILTIN_PROVIDERS: tuple[str, ...] = (
+    "amazon-bedrock",
+    "ant-ling",
+    "anthropic",
+    "azure",
+    "baseten",
+    "cerebras",
+    "cloudflare-ai-gateway",
+    "cloudflare-workers-ai",
+    "deepseek",
+    "fireworks",
+    "github-copilot",
+    "google",
+    "google-vertex",
+    "groq",
+    "huggingface",
+    "kimi-coding",
+    "meta",
+    "minimax",
+    "minimax-cn",
+    "mistral",
+    "moonshotai",
+    "moonshotai-cn",
+    "nvidia",
+    "openai",
+    "openai-codex",
+    "opencode",
+    "opencode-go",
+    "openrouter",
+    "qwen-token-plan",
+    "qwen-token-plan-cn",
+    "qwen-token-plan-individual",
+    "radius",
+    "together",
+    "typesafe",
+    "vercel-ai-gateway",
+    "xai",
+    "xiaomi",
+    "xiaomi-token-plan-ams",
+    "xiaomi-token-plan-cn",
+    "xiaomi-token-plan-sgp",
+    "zai",
+    "zai-coding-cn",
+)
+
+_lock = threading.Lock()
+_cache: dict[str, dict[str, dict[str, AnyModel]]] = {}
+
+
+def _read_json(name: str) -> Any:
+    data_dir = resources.files("pi_ai.providers").joinpath("data")
+    with data_dir.joinpath(name).open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def load_catalog(provider: str) -> dict[str, dict[str, AnyModel]]:
+    """Catalog of ``provider`` as ``{"chat": {id: model}, "image": {...}, "classifier": {...}}``."""
+    cached = _cache.get(provider)
+    if cached is not None:
+        return cached
+    with _lock:
+        cached = _cache.get(provider)
+        if cached is not None:
+            return cached
+        groups: dict[str, dict[str, AnyModel]] = {"chat": {}, "image": {}, "classifier": {}}
+        try:
+            raw = _read_json(f"{provider}.json")
+        except FileNotFoundError:
+            raw = {}
+        for entries in raw.values():
+            for key, value in entries.items():
+                model_type, _, _ = key.partition(":")
+                group = groups.get(model_type or "chat")
+                if group is None:
+                    continue
+                model = model_from_dict(value)
+                group[model.id] = model
+        _cache[provider] = groups
+        return groups
+
+
+def catalog_models(provider: str) -> list[AnyModel]:
+    """Every built-in model of ``provider`` (chat, then image, then classifier)."""
+    groups = load_catalog(provider)
+    return [*groups["chat"].values(), *groups["image"].values(), *groups["classifier"].values()]
+
+
+def get_builtin_providers() -> list[str]:
+    return list(BUILTIN_PROVIDERS)
+
+
+def get_builtin_model(provider: str, model_id: str) -> Model | None:
+    """One generated built-in chat model, independent of any ``Models`` collection."""
+    return load_catalog(provider)["chat"].get(model_id)  # type: ignore[return-value]
+
+
+def get_builtin_image_model(provider: str, model_id: str) -> ImageModel | None:
+    return load_catalog(provider)["image"].get(model_id)  # type: ignore[return-value]
+
+
+def get_builtin_classifier_model(provider: str, model_id: str) -> ClassifierModel | None:
+    return load_catalog(provider)["classifier"].get(model_id)  # type: ignore[return-value]
+
+
+def get_builtin_models(provider: str) -> list[Model]:
+    return list(load_catalog(provider)["chat"].values())  # type: ignore[arg-type]
+
+
+def get_builtin_image_models(provider: str) -> list[ImageModel]:
+    return list(load_catalog(provider)["image"].values())  # type: ignore[arg-type]
+
+
+def get_builtin_classifier_models(provider: str) -> list[ClassifierModel]:
+    return list(load_catalog(provider)["classifier"].values())  # type: ignore[arg-type]
+
+
+def get_all_builtin_models(provider: str) -> list[AnyModel]:
+    return catalog_models(provider)
+
+
+@lru_cache(maxsize=1)
+def get_builtin_model_data_generated_at() -> int | None:
+    """Generation timestamp (ms) shared by all built-in provider catalogs."""
+    try:
+        manifest = _read_json("manifest.json")
+        stamp = datetime.fromisoformat(manifest["generatedAt"].replace("Z", "+00:00"))
+    except (FileNotFoundError, KeyError, ValueError):
+        return None
+    return int(stamp.timestamp() * 1000)

pi_ai/providers/cerebras.py

@@ -0,0 +1,23 @@
+"""Cerebras 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__ = ["cerebras_provider"]
+
+
+def cerebras_provider() -> Provider:
+    """Cerebras (OpenAI-compatible Chat Completions), authenticated with ``CEREBRAS_API_KEY``."""
+    return create_provider(
+        id="cerebras",
+        name="Cerebras",
+        base_url="https://api.cerebras.ai/v1",
+        auth=ProviderAuth(api_key=env_api_key_auth("Cerebras API key", ["CEREBRAS_API_KEY"])),
+        models=lambda: catalog_models("cerebras"),
+        api=openai_completions_api(),
+    )

pi_ai/providers/cloudflare_ai_gateway.py

@@ -0,0 +1,30 @@
+"""Cloudflare AI Gateway provider."""
+
+from __future__ import annotations
+
+from ..api.registry import anthropic_messages_api, openai_completions_api, openai_responses_api
+from ..auth.types import ProviderAuth
+from ..models import Provider, create_provider
+from .catalog import catalog_models
+from .cloudflare_auth import cloudflare_ai_gateway_auth
+from .cloudflare_stream import cloudflare_streams
+
+__all__ = ["cloudflare_ai_gateway_provider"]
+
+
+def cloudflare_ai_gateway_provider() -> Provider:
+    """Cloudflare AI Gateway: ``CLOUDFLARE_API_KEY``, ``CLOUDFLARE_ACCOUNT_ID`` and ``CLOUDFLARE_GATEWAY_ID``."""
+    # The api map is pinned to all three APIs: models.dev's gateway catalog drops and
+    # restores `workers-ai/*` (openai-completions) entries over time, so the map must
+    # not depend on which APIs the generated catalog happens to contain.
+    return create_provider(
+        id="cloudflare-ai-gateway",
+        name="Cloudflare AI Gateway",
+        auth=ProviderAuth(api_key=cloudflare_ai_gateway_auth()),
+        models=lambda: catalog_models("cloudflare-ai-gateway"),
+        api={
+            "anthropic-messages": cloudflare_streams(anthropic_messages_api()),
+            "openai-completions": cloudflare_streams(openai_completions_api()),
+            "openai-responses": cloudflare_streams(openai_responses_api()),
+        },
+    )

pi_ai/providers/cloudflare_auth.py

@@ -0,0 +1,112 @@
+"""Cloudflare Workers AI / AI Gateway auth: API key plus account (and gateway) ids."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from ..abort import AbortSignal
+from ..auth.types import (
+    ApiKeyAuth,
+    AuthContext,
+    AuthPrompt,
+    AuthResult,
+    Credential,
+    ModelAuth,
+    ProviderAuthInteraction,
+)
+
+__all__ = ["cloudflare_workers_ai_auth", "cloudflare_ai_gateway_auth"]
+
+CLOUDFLARE_API_KEY = "CLOUDFLARE_API_KEY"
+CLOUDFLARE_ACCOUNT_ID = "CLOUDFLARE_ACCOUNT_ID"
+CLOUDFLARE_GATEWAY_ID = "CLOUDFLARE_GATEWAY_ID"
+
+
+@dataclass(slots=True)
+class _ResolvedCloudflare:
+    api_key: str
+    env: dict[str, str]
+    source: str
+
+
+async def _resolve_value(name: str, ctx: AuthContext, credential: Credential | None, signal: AbortSignal) -> str | None:
+    # Per-field merge: prefer the credential value, fall back to ambient env.
+    # A credential carrying only the API key must still pick up the account /
+    # gateway id from the environment.
+    from_credential: str | None = None
+    if credential is not None:
+        if name == CLOUDFLARE_API_KEY:
+            from_credential = credential.get("key")
+        else:
+            from_credential = (credential.get("env") or {}).get(name)
+    if from_credential is not None:
+        return from_credential
+    signal.throw_if_aborted()
+    value = await ctx.env(name)
+    signal.throw_if_aborted()
+    return value
+
+
+async def _resolve_cloudflare_env(
+    kind: str, ctx: AuthContext, credential: Credential | None, signal: AbortSignal
+) -> _ResolvedCloudflare | None:
+    api_key = await _resolve_value(CLOUDFLARE_API_KEY, ctx, credential, signal)
+    account_id = await _resolve_value(CLOUDFLARE_ACCOUNT_ID, ctx, credential, signal)
+    gateway_id = await _resolve_value(CLOUDFLARE_GATEWAY_ID, ctx, credential, signal) if kind == "ai-gateway" else None
+
+    if not api_key or not account_id or (kind == "ai-gateway" and not gateway_id):
+        return None
+
+    env = {CLOUDFLARE_ACCOUNT_ID: account_id}
+    if gateway_id:
+        env[CLOUDFLARE_GATEWAY_ID] = gateway_id
+    return _ResolvedCloudflare(api_key, env, "stored credential" if credential is not None else CLOUDFLARE_API_KEY)
+
+
+def cloudflare_workers_ai_auth() -> ApiKeyAuth:
+    """API key auth that also requires ``CLOUDFLARE_ACCOUNT_ID`` (returned as provider env)."""
+
+    async def login(interaction: ProviderAuthInteraction) -> Credential:
+        key = await interaction.prompt(AuthPrompt(type="secret", message="Enter Cloudflare API key"))
+        account_id = await interaction.prompt(AuthPrompt(type="text", message="Enter Cloudflare account ID"))
+        return {"type": "api_key", "key": key, "env": {CLOUDFLARE_ACCOUNT_ID: account_id}}
+
+    async def resolve(*, ctx: AuthContext, credential: Credential | None, signal: AbortSignal) -> AuthResult | None:
+        resolved = await _resolve_cloudflare_env("workers-ai", ctx, credential, signal)
+        if resolved is None:
+            return None
+        return AuthResult(ModelAuth(api_key=resolved.api_key), env=resolved.env, source=resolved.source)
+
+    return ApiKeyAuth("Cloudflare API key", resolve, login=login)
+
+
+def cloudflare_ai_gateway_auth() -> ApiKeyAuth:
+    """AI Gateway auth: the key travels as ``cf-aig-authorization``; account and gateway ids as provider env."""
+
+    async def login(interaction: ProviderAuthInteraction) -> Credential:
+        key = await interaction.prompt(AuthPrompt(type="secret", message="Enter Cloudflare API key"))
+        account_id = await interaction.prompt(AuthPrompt(type="text", message="Enter Cloudflare account ID"))
+        gateway_id = await interaction.prompt(AuthPrompt(type="text", message="Enter Cloudflare AI Gateway ID"))
+        return {
+            "type": "api_key",
+            "key": key,
+            "env": {CLOUDFLARE_ACCOUNT_ID: account_id, CLOUDFLARE_GATEWAY_ID: gateway_id},
+        }
+
+    async def resolve(*, ctx: AuthContext, credential: Credential | None, signal: AbortSignal) -> AuthResult | None:
+        resolved = await _resolve_cloudflare_env("ai-gateway", ctx, credential, signal)
+        if resolved is None:
+            return None
+        return AuthResult(
+            ModelAuth(
+                headers={
+                    "cf-aig-authorization": f"Bearer {resolved.api_key}",
+                    "Authorization": None,
+                    "x-api-key": None,
+                }
+            ),
+            env=resolved.env,
+            source=resolved.source,
+        )
+
+    return ApiKeyAuth("Cloudflare API key", resolve, login=login)

pi_ai/providers/cloudflare_stream.py

@@ -0,0 +1,84 @@
+"""Materialize Cloudflare account/gateway placeholders in model base URLs before dispatch."""
+
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Mapping
+from typing import Any, TypeVar
+
+from ..event_stream import AssistantMessageEventStream
+from ..types import ClassifierContext, ClassifierModel, ClassifierResult, Model, TranscriptContext
+
+__all__ = ["resolve_cloudflare_model", "cloudflare_streams", "cloudflare_classifier"]
+
+CLOUDFLARE_ACCOUNT_ID = "CLOUDFLARE_ACCOUNT_ID"
+CLOUDFLARE_GATEWAY_ID = "CLOUDFLARE_GATEWAY_ID"
+
+_ModelT = TypeVar("_ModelT")
+
+
+def resolve_cloudflare_model(model: _ModelT, env: Mapping[str, str] | None) -> _ModelT:
+    """Replace ``{CLOUDFLARE_ACCOUNT_ID}``/``{CLOUDFLARE_GATEWAY_ID}`` in ``model.base_url`` from ``env``.
+
+    Unresolved placeholders are kept; the model is returned unchanged when nothing changes.
+    """
+    if not env:
+        return model
+    base_url: str = model.base_url  # type: ignore[attr-defined]
+    resolved = base_url
+    for name in (CLOUDFLARE_ACCOUNT_ID, CLOUDFLARE_GATEWAY_ID):
+        value = env.get(name)
+        placeholder = "{" + name + "}"
+        resolved = resolved.replace(placeholder, value if value is not None else placeholder)
+    if resolved == base_url:
+        return model
+    return dataclasses.replace(model, base_url=resolved)  # type: ignore[type-var]
+
+
+def _env(options: Mapping[str, Any] | None) -> Mapping[str, str] | None:
+    return options.get("env") if options else None
+
+
+class _CloudflareStreams:
+    __slots__ = ("_streams",)
+
+    def __init__(self, streams: Any) -> None:
+        self._streams = streams
+
+    def stream(
+        self, model: Model, context: TranscriptContext, options: Mapping[str, Any] | None = None
+    ) -> AssistantMessageEventStream:
+        return self._streams.stream(resolve_cloudflare_model(model, _env(options)), context, options)
+
+    def stream_simple(
+        self, model: Model, context: TranscriptContext, options: Mapping[str, Any] | None = None
+    ) -> AssistantMessageEventStream:
+        return self._streams.stream_simple(resolve_cloudflare_model(model, _env(options)), context, options)
+
+    def __repr__(self) -> str:
+        return f"cloudflare_streams({self._streams!r})"
+
+
+class _CloudflareClassifier:
+    __slots__ = ("_classifier",)
+
+    def __init__(self, classifier: Any) -> None:
+        self._classifier = classifier
+
+    async def classify(
+        self, model: ClassifierModel, context: ClassifierContext, options: Mapping[str, Any] | None = None
+    ) -> ClassifierResult:
+        return await self._classifier.classify(resolve_cloudflare_model(model, _env(options)), context, options)
+
+    def __repr__(self) -> str:
+        return f"cloudflare_classifier({self._classifier!r})"
+
+
+def cloudflare_streams(streams: Any) -> Any:
+    """Wrap an API implementation so Cloudflare endpoint placeholders materialize from the resolved provider env."""
+    return _CloudflareStreams(streams)
+
+
+def cloudflare_classifier(classifier: Any) -> Any:
+    """Classifier counterpart of :func:`cloudflare_streams`."""
+    return _CloudflareClassifier(classifier)

pi_ai/providers/cloudflare_workers_ai.py

@@ -0,0 +1,26 @@
+"""Cloudflare Workers AI provider."""
+
+from __future__ import annotations
+
+from ..api.registry import cloudflare_workers_ai_system_one_api, openai_completions_api
+from ..auth.types import ProviderAuth
+from ..models import Provider, create_provider
+from .catalog import catalog_models
+from .cloudflare_auth import cloudflare_workers_ai_auth
+from .cloudflare_stream import cloudflare_classifier, cloudflare_streams
+
+__all__ = ["cloudflare_workers_ai_provider"]
+
+
+def cloudflare_workers_ai_provider() -> Provider:
+    """Cloudflare Workers AI: ``CLOUDFLARE_API_KEY`` plus ``CLOUDFLARE_ACCOUNT_ID``."""
+    return create_provider(
+        id="cloudflare-workers-ai",
+        name="Cloudflare Workers AI",
+        auth=ProviderAuth(api_key=cloudflare_workers_ai_auth()),
+        models=lambda: catalog_models("cloudflare-workers-ai"),
+        api=cloudflare_streams(openai_completions_api()),
+        classifiers={
+            "cloudflare-workers-ai-system-one": cloudflare_classifier(cloudflare_workers_ai_system_one_api()),
+        },
+    )