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

pi_ai/api/azure_openai_config.py

@@ -0,0 +1,138 @@
+"""Azure OpenAI endpoint resolution (base URL, API version, deployment names)."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping
+from typing import Any, NamedTuple
+from urllib.parse import urlsplit, urlunsplit
+
+from ..types import StreamOptions
+from ..utils.misc import get_provider_env_value
+
+__all__ = [
+    "DEFAULT_AZURE_API_VERSION",
+    "AzureEndpointOptions",
+    "AzureConfig",
+    "resolve_deployment_name",
+    "resolve_azure_base_url",
+    "resolve_azure_config",
+]
+
+DEFAULT_AZURE_API_VERSION = "v1"
+
+_TRAILING_SLASHES = re.compile(r"/+$")
+_URL_SCHEME = re.compile(r"^[a-zA-Z][a-zA-Z0-9+.-]*$")
+_SPECIAL_SCHEMES = frozenset({"http", "https", "ws", "wss", "ftp", "file"})
+
+
+class AzureEndpointOptions(StreamOptions, total=False):
+    """Azure models ship without a base URL: one resource per user, resolved per request."""
+
+    azure_api_version: str
+    azure_resource_name: str
+    azure_base_url: str
+    azure_deployment_name: str
+
+
+class AzureConfig(NamedTuple):
+    base_url: str
+    api_version: str
+
+
+def _parse_deployment_name_map(value: str | None) -> dict[str, str]:
+    mapping: dict[str, str] = {}
+    if not value:
+        return mapping
+    for entry in value.split(","):
+        trimmed = entry.strip()
+        if not trimmed:
+            continue
+        model_id, _, deployment_name = trimmed.partition("=")
+        # JavaScript's split("=", 2) keeps only the text up to a second "=".
+        deployment_name = deployment_name.split("=", 1)[0]
+        if not model_id or not deployment_name:
+            continue
+        mapping[model_id.strip()] = deployment_name.strip()
+    return mapping
+
+
+def resolve_deployment_name(model: Any, options: Mapping[str, Any] | None = None) -> str:
+    """Deployment for ``model.id``: ``azure_deployment_name``, ``AZURE_OPENAI_DEPLOYMENT_NAME_MAP``, or the id."""
+    options = options or {}
+    explicit = options.get("azure_deployment_name")
+    if explicit:
+        return explicit
+    mapped = _parse_deployment_name_map(
+        get_provider_env_value("AZURE_OPENAI_DEPLOYMENT_NAME_MAP", options.get("env"))
+    ).get(model.id)
+    return mapped or model.id
+
+
+def _normalize_azure_base_url(base_url: str) -> str:
+    trimmed = _TRAILING_SLASHES.sub("", base_url.strip())
+    try:
+        parts = urlsplit(trimmed)
+    except ValueError as error:
+        raise ValueError(f"Invalid Azure OpenAI base URL: {base_url}") from error
+    scheme = parts.scheme.lower()
+    if not scheme or not _URL_SCHEME.match(scheme) or (scheme in _SPECIAL_SCHEMES and not parts.netloc):
+        raise ValueError(f"Invalid Azure OpenAI base URL: {base_url}")
+
+    hostname = parts.hostname or ""
+    is_azure_host = (
+        hostname.endswith(".openai.azure.com")
+        or hostname.endswith(".cognitiveservices.azure.com")
+        or hostname.endswith(".ai.azure.com")
+    )
+    normalized_path = _TRAILING_SLASHES.sub("", parts.path)
+    path = parts.path
+    query = parts.query
+
+    # Ensure Azure hosts have /openai/v1 as base path so /responses and ?api-version=v1
+    # are appended correctly.
+    if is_azure_host and normalized_path in ("", "/", "/openai", "/openai/v1/responses"):
+        path = "/openai/v1"
+        query = ""
+
+    netloc = parts.netloc if "@" in parts.netloc else parts.netloc.lower()
+    if not path and scheme in _SPECIAL_SCHEMES:
+        path = "/"
+    url = urlunsplit((scheme, netloc, path, query, parts.fragment))
+    return _TRAILING_SLASHES.sub("", url)
+
+
+def _build_default_base_url(resource_name: str) -> str:
+    return f"https://{resource_name}.openai.azure.com/openai/v1"
+
+
+def resolve_azure_base_url(model: Any, options: Mapping[str, Any] | None = None) -> str:
+    """Base URL from ``azure_base_url``, ``AZURE_OPENAI_BASE_URL``, the resource name, or ``model.base_url``."""
+    options = options or {}
+    env = options.get("env")
+    explicit = (options.get("azure_base_url") or "").strip()
+    base_url = explicit or (get_provider_env_value("AZURE_OPENAI_BASE_URL", env) or "").strip() or None
+    resource_name = options.get("azure_resource_name") or get_provider_env_value("AZURE_OPENAI_RESOURCE_NAME", env)
+
+    resolved_base_url = base_url
+    if not resolved_base_url and resource_name:
+        resolved_base_url = _build_default_base_url(resource_name)
+    if not resolved_base_url and getattr(model, "base_url", None):
+        resolved_base_url = model.base_url
+    if not resolved_base_url:
+        raise ValueError(
+            "Azure OpenAI base URL is required. Set AZURE_OPENAI_BASE_URL or AZURE_OPENAI_RESOURCE_NAME, "
+            "or pass azure_base_url, azure_resource_name, or model.base_url."
+        )
+    return _normalize_azure_base_url(resolved_base_url)
+
+
+def resolve_azure_config(model: Any, options: Mapping[str, Any] | None = None) -> AzureConfig:
+    """Base URL and API version (``azure_api_version``, ``AZURE_OPENAI_API_VERSION``, or ``"v1"``)."""
+    options = options or {}
+    return AzureConfig(
+        base_url=resolve_azure_base_url(model, options),
+        api_version=options.get("azure_api_version")
+        or get_provider_env_value("AZURE_OPENAI_API_VERSION", options.get("env"))
+        or DEFAULT_AZURE_API_VERSION,
+    )

pi_ai/api/azure_openai_responses.py

@@ -0,0 +1,245 @@
+# pyright: reportAttributeAccessIssue=false, reportArgumentType=false
+# Converters branch on ``role``/``type`` discriminators held in locals, which pyright does not narrow.
+"""Azure OpenAI Responses API (``azure-openai-responses``).
+
+Streams ``responses.create`` through the official ``openai`` SDK's ``AsyncAzureOpenAI``
+client (``POST {base_url}/responses?api-version=<version>`` with an ``api-key`` header).
+The deployment name goes in the body's ``model`` field.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, Literal
+
+from ..abort import abort_scope
+from ..event_stream import AssistantMessageEventStream
+from ..models import clamp_thinking_level
+from ..types import (
+    AssistantMessage,
+    DoneEvent,
+    ErrorEvent,
+    Model,
+    ProviderResponse,
+    StartEvent,
+    TranscriptContext,
+    Usage,
+)
+from ..utils.callbacks import maybe_await
+from ..utils.error_body import format_provider_error, normalize_provider_error
+from ..utils.misc import get_pi_user_agent, headers_to_record
+from ..utils.transcript import get_declared_tools, resolve_transcript, resolve_transcript_tools
+from .azure_openai_config import AzureEndpointOptions, resolve_azure_config, resolve_deployment_name
+from .constrained_sampling import create_grammar_tool_input_properties
+from .openai_prompt_cache import clamp_openai_prompt_cache_key
+from .openai_responses_shared import convert_responses_messages, convert_responses_tools, process_responses_stream
+from .openai_sdk import (
+    create_openai_client,
+    iter_sdk_stream,
+    merge_sdk_headers,
+    open_sdk_stream,
+    openai_request,
+    sdk_body_arguments,
+)
+from .simple_options import build_base_options, resolve_sampling_params
+
+__all__ = ["AzureOpenAIResponsesOptions", "AZURE_TOOL_CALL_PROVIDERS", "stream", "stream_simple"]
+
+AZURE_TOOL_CALL_PROVIDERS = frozenset({"openai", "openai-codex", "opencode", "azure"})
+# OpenAI Responses rejects max_output_tokens below 16: https://github.com/earendil-works/pi/issues/6265
+OPENAI_RESPONSES_MIN_OUTPUT_TOKENS = 16
+
+
+class AzureOpenAIResponsesOptions(AzureEndpointOptions, total=False):
+    reasoning_effort: Literal["minimal", "low", "medium", "high", "xhigh", "max"]
+    tool_choice: Any
+    reasoning_summary: Literal["auto", "detailed", "concise"] | None
+
+
+def _format_azure_openai_error(error: Any) -> str:
+    return format_provider_error(normalize_provider_error(error), "Azure OpenAI API error")
+
+
+def stream(
+    model: Model, context: TranscriptContext, options: Mapping[str, Any] | None = None
+) -> AssistantMessageEventStream:
+    """Stream a response from an Azure OpenAI Responses deployment."""
+    opts: dict[str, Any] = dict(options or {})
+    out = AssistantMessageEventStream()
+    normalized_context = resolve_transcript(context, (model.compat or {}).get("supportsMidConvoSystemMessages"))
+    out.start(_run(model, normalized_context, opts, out))
+    return out
+
+
+async def _run(
+    model: Model, context: TranscriptContext, options: dict[str, Any], out: AssistantMessageEventStream
+) -> None:
+    deployment_name = resolve_deployment_name(model, options)
+    output = AssistantMessage(
+        [],
+        api="azure-openai-responses",
+        provider=model.provider,
+        model=model.id,
+        usage=Usage(),
+        stop_reason="pending",
+    )
+    signal = options.get("signal")
+    try:
+        async with abort_scope(signal):
+            api_key = options.get("api_key")
+            if not api_key:
+                raise RuntimeError(f"No API key for provider: {model.provider}")
+            headers = merge_sdk_headers({"User-Agent": get_pi_user_agent()}, model.headers, options.get("headers"))
+            base_url, api_version = resolve_azure_config(model, options)
+            client = create_openai_client(
+                api_key=api_key,
+                # The SDK appends the path to the base URL and sends ``api-version`` as the query string.
+                base_url=base_url.partition("?")[0],
+                headers=headers,
+                options=options,
+                azure_api_version=api_version,
+            )
+            grammar_tool_input_properties = create_grammar_tool_input_properties(
+                get_declared_tools(context.messages), bool((model.compat or {}).get("supportsOpenAIGrammarTools"))
+            )
+            params: Any = _build_params(model, context, options, deployment_name, grammar_tool_input_properties)
+            on_payload = options.get("on_payload")
+            if on_payload is not None:
+                next_params = await maybe_await(on_payload(params, model))
+                if next_params is not None:
+                    params = next_params
+            create = client.responses.with_raw_response.create
+            body_arguments = sdk_body_arguments(params, headers=headers)
+            raw = await openai_request(lambda: create(stream=True, **body_arguments), options)
+            sdk_stream = open_sdk_stream(raw)
+            try:
+                on_response = options.get("on_response")
+                if on_response is not None:
+                    await maybe_await(
+                        on_response(ProviderResponse(raw.status_code, headers_to_record(raw.headers)), model)
+                    )
+                out.push(StartEvent(output))
+                await process_responses_stream(
+                    iter_sdk_stream(sdk_stream),
+                    output,
+                    out,
+                    model,
+                    {
+                        "on_provider_stream_event": options.get("on_provider_stream_event"),
+                        "grammar_tool_input_properties": grammar_tool_input_properties,
+                    },
+                )
+            finally:
+                await sdk_stream.close()
+
+        if signal is not None and signal.aborted:
+            raise RuntimeError("Request was aborted")
+        if output.stop_reason == "pending":
+            raise RuntimeError("Azure OpenAI Responses stream ended without a stop reason")
+        if output.stop_reason in ("aborted", "error"):
+            raise RuntimeError(output.error_message or "An unknown error occurred")
+
+        out.push(DoneEvent(output.stop_reason, output))  # type: ignore[arg-type]
+        out.end()
+    except Exception as error:
+        aborted = signal is not None and signal.aborted
+        output.stop_reason = "aborted" if aborted else "error"
+        output.error_message = "Request was aborted" if aborted else _format_azure_openai_error(error)
+        out.push(ErrorEvent(output.stop_reason, output))  # type: ignore[arg-type]
+        out.end()
+
+
+def stream_simple(
+    model: Model, context: TranscriptContext, options: Mapping[str, Any] | None = None
+) -> AssistantMessageEventStream:
+    """Stream with provider-neutral options (``reasoning``, ``tool_choice``, ...)."""
+    options = options or {}
+    api_key = options.get("api_key")
+    if not api_key:
+        raise RuntimeError(f"No API key for provider: {model.provider}")
+
+    base = build_base_options(model, context, options, api_key)
+    if options.get("tool_choice") is not None:
+        base["tool_choice"] = options["tool_choice"]
+    reasoning = options.get("reasoning")
+    clamped = clamp_thinking_level(model, reasoning) if reasoning else None
+    if clamped is not None and clamped != "off":
+        base["reasoning_effort"] = clamped
+    # Azure endpoint options are not part of the simple options; keep them when callers pass them.
+    for key in ("azure_api_version", "azure_resource_name", "azure_base_url", "azure_deployment_name"):
+        if options.get(key) is not None:
+            base[key] = options[key]
+    return stream(model, context, base)
+
+
+def _build_params(
+    model: Model,
+    context: TranscriptContext,
+    options: Mapping[str, Any],
+    deployment_name: str,
+    grammar_tool_input_properties: Mapping[str, str],
+) -> dict[str, Any]:
+    compat = model.compat or {}
+    supports_additional_tools = bool(compat.get("supportsAdditionalTools"))
+    supports_tool_search = bool(compat.get("supportsToolSearch"))
+    supports_strict_mode = compat.get("supportsStrictMode")
+    if supports_strict_mode is None:
+        supports_strict_mode = True
+    supports_openai_grammar_tools = bool(compat.get("supportsOpenAIGrammarTools"))
+    tool_options = {
+        "supports_strict_mode": supports_strict_mode,
+        "supports_openai_grammar_tools": supports_openai_grammar_tools,
+    }
+    transcript_tools = resolve_transcript_tools(context.messages, supports_additional_tools or supports_tool_search)
+    messages = convert_responses_messages(
+        model,
+        context,
+        AZURE_TOOL_CALL_PROVIDERS,
+        {
+            "grammar_tool_input_properties": grammar_tool_input_properties,
+            "supports_mid_convo_system_messages": bool(compat.get("supportsMidConvoSystemMessages")),
+            "supports_additional_tools": supports_additional_tools,
+            "supports_tool_search": supports_tool_search,
+            "tool_options": tool_options,  # type: ignore[typeddict-item]
+        },
+    )
+
+    params: dict[str, Any] = {"model": deployment_name, "input": messages, "stream": True}
+    prompt_cache_key = clamp_openai_prompt_cache_key(options.get("session_id"))
+    if prompt_cache_key is not None:
+        params["prompt_cache_key"] = prompt_cache_key
+    params["store"] = False
+
+    max_tokens = options.get("max_tokens")
+    if max_tokens:
+        params["max_output_tokens"] = max(max_tokens, OPENAI_RESPONSES_MIN_OUTPUT_TOKENS)
+
+    if options.get("temperature") is not None:
+        params["temperature"] = options["temperature"]
+
+    if transcript_tools.request_tools:
+        params["tools"] = convert_responses_tools(transcript_tools.request_tools, tool_options)  # type: ignore[arg-type]
+    if options.get("tool_choice") is not None:
+        params["tool_choice"] = options["tool_choice"]
+
+    requested = options.get("reasoning_effort")
+    reasoning_effort = requested or ("medium" if options.get("reasoning_summary") else None)
+    if model.reasoning:
+        level_map = model.thinking_level_map or {}
+        if reasoning_effort:
+            effort = reasoning_effort
+            if requested:
+                mapped = level_map.get(requested)
+                effort = mapped if mapped is not None else requested
+            params["reasoning"] = {"effort": effort, "summary": options.get("reasoning_summary") or "auto"}
+            params["include"] = ["reasoning.encrypted_content"]
+        elif not ("off" in level_map and level_map["off"] is None):
+            off = level_map.get("off")
+            params["reasoning"] = {"effort": off if off is not None else "none"}
+
+    # Last so model and request sampling parameters override named request fields.
+    sampling_params = resolve_sampling_params(model, reasoning_effort or "off", options.get("sampling_params"))
+    if sampling_params:
+        params.update(sampling_params)
+
+    return params