From 9cd880c9f89def00d4adaa89c4e4fb78e7d3ce0a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 12 Jun 2026 08:44:38 +0000 Subject: [PATCH] feat(translation): xai provider wrapping openai_compat (guard name knob, gate-free assemble_body, xai chunk dialect) The serializer is the openai_compat body over xai's own gates: v1's supported-list truth is RAISE-unless-drop_params (the mct rename arm is dead code), so every raise row is a typed fallback; reasoning_effort gates through deps over the xai/{model} map key; user rides verbatim; the function-level tool strict strip mirrors filter_value_from_dict with deeper strict keys failing closed in the raw guard. The guard also owns the R5 reroutes (web_search_options Responses bridge, use_xai_oauth PKCE) plus explicit stream:false, and runs the shared openai guard with a user-only message-name fallback (v1 strips non-user names, so the IR drop IS v1). The response parser is the openai parser plus the xai usage post-steps (reasoning fold, total normalize, num_sources_used -> web_search_requests); the finish_reason '' chain needs no arm because v1's own fix function is empirically dead and both sides share the live map_finish_reason. The stream parser pins the httpx dict path: per-chunk usage fold, no extras/system_fingerprint survival, reasoning rename, tool type 'function' default, and the choices:[] usage tail passes through folded for the seam contract. assemble_body is the gate-free entry point azure N4 deferred to this phase; openai_compat callers are unchanged. --- .../translation/inbound/openai_chat/stream.py | 12 +- .../providers/openai_compat/guard.py | 20 +- .../providers/openai_compat/response.py | 4 +- .../providers/openai_compat/serialize.py | 7 + litellm/translation/providers/xai/__init__.py | 12 + litellm/translation/providers/xai/guard.py | 97 ++++++++ litellm/translation/providers/xai/params.py | 57 +++++ litellm/translation/providers/xai/response.py | 113 +++++++++ .../translation/providers/xai/serialize.py | 58 +++++ litellm/translation/providers/xai/stream.py | 224 ++++++++++++++++++ 10 files changed, 595 insertions(+), 9 deletions(-) create mode 100644 litellm/translation/providers/xai/__init__.py create mode 100644 litellm/translation/providers/xai/guard.py create mode 100644 litellm/translation/providers/xai/params.py create mode 100644 litellm/translation/providers/xai/response.py create mode 100644 litellm/translation/providers/xai/serialize.py create mode 100644 litellm/translation/providers/xai/stream.py diff --git a/litellm/translation/inbound/openai_chat/stream.py b/litellm/translation/inbound/openai_chat/stream.py index 6e0e1a7a51..1aa721f26e 100644 --- a/litellm/translation/inbound/openai_chat/stream.py +++ b/litellm/translation/inbound/openai_chat/stream.py @@ -52,7 +52,7 @@ fold ``wire_chunk`` events and gemini folds composite ``chunk`` events; their parsers can never produce a per-block delta, and the step surfaces a loud error (never a fabricated placeholder body) if one ever arrives.""" -ChunkDialect = Literal[BlockDialect, "openai", "azure", "gemini"] +ChunkDialect = Literal[BlockDialect, "openai", "azure", "gemini", "xai"] @dataclass(frozen=True) @@ -131,7 +131,7 @@ def _as_block_dialect( match dialect: case "anthropic" | "bedrock_converse": return dialect - case "openai" | "azure" | "gemini": + case "openai" | "azure" | "gemini" | "xai": # wire/composite dialects fold whole chunks; a per-block delta # here is a wiring bug and must be loud, never a fabricated # anthropic-shaped placeholder (critic-google M5 / critic-azure M3). @@ -197,7 +197,7 @@ _OPENAI_WIRE_KEYS = frozenset({"id", "model", "system_fingerprint", "choices", " def _step_openai(state: StreamState, chunk: PlainJson) -> _StepResult: - if state.dialect not in ("openai", "azure"): + if state.dialect not in ("openai", "azure", "xai"): return _dialect_mismatch(state.dialect, "wire_chunk") if not isinstance(chunk, dict): return state, () # the parser only emits dict wire chunks @@ -307,6 +307,12 @@ def _openai_chunk_non_empty(state: StreamState, delta: dict[str, PlainJson]) -> tool_calls = delta.get("tool_calls") if isinstance(tool_calls, list) and len(tool_calls) > 0: return True + if state.dialect == "xai": + # Grok reasoning deltas carry only reasoning_content; v1's wrapper + # emits them (is_delta_empty consults reasoning_content too). + reasoning = delta.get("reasoning_content") + if isinstance(reasoning, str) and len(reasoning) > 0: + return True return not state.sent_role and delta.get("role") is not None diff --git a/litellm/translation/providers/openai_compat/guard.py b/litellm/translation/providers/openai_compat/guard.py index 8ffc17ab5f..ed4c14f3fe 100644 --- a/litellm/translation/providers/openai_compat/guard.py +++ b/litellm/translation/providers/openai_compat/guard.py @@ -21,8 +21,18 @@ from ...errors import TranslationError _Raw = Mapping[str, object] -def unsupported_request_shapes(raw: _Raw) -> TranslationError | None: - reason = _params_reason(raw) or _tools_reason(raw) or _messages_reason(raw) +def unsupported_request_shapes( + raw: _Raw, *, name_fallback_user_only: bool = False +) -> TranslationError | None: + """``name_fallback_user_only``: providers whose v1 transform STRIPS the + message ``name`` from non-user roles (xai ``strip_name_from_messages``) + only need the fallback for user messages, where v1 forwards it verbatim; + the IR's name-drop IS v1's behavior on the other roles.""" + reason = ( + _params_reason(raw) + or _tools_reason(raw) + or _messages_reason(raw, name_fallback_user_only) + ) if reason is None: return None return TranslationError.of_unsupported(f"{reason}; v1 forwards the original shape") @@ -71,7 +81,7 @@ def _tools_reason(raw: _Raw) -> str | None: return None -def _messages_reason(raw: _Raw) -> str | None: +def _messages_reason(raw: _Raw, name_fallback_user_only: bool = False) -> str | None: messages = _as_seq(raw.get("messages")) if messages is None: return None @@ -84,7 +94,9 @@ def _messages_reason(raw: _Raw) -> str | None: # keep scanning so the guard stays locally conservative too continue role = entry.get("role") - if entry.get("name") is not None: + if entry.get("name") is not None and ( + not name_fallback_user_only or role == "user" + ): return "message name field (not carried by the IR)" reason = _message_reason(entry, role, seen_non_system) if reason is not None: diff --git a/litellm/translation/providers/openai_compat/response.py b/litellm/translation/providers/openai_compat/response.py index a1e2060762..c197112a29 100644 --- a/litellm/translation/providers/openai_compat/response.py +++ b/litellm/translation/providers/openai_compat/response.py @@ -98,7 +98,7 @@ def parse_response(raw: PlainJson, request: ChatRequest) -> _ParseResult: model=model if isinstance(model, str) else request.model, content=Block.of_seq(blocks), finish=semantic_finish, - usage=_semantic_usage(raw.get("usage")), + usage=semantic_usage(raw.get("usage")), synthesized_json_content=False, wire=Some(JsonBlob(value=body)), ) @@ -327,7 +327,7 @@ def _tool_use_block(call: PlainJson) -> ContentBlock | TranslationError: ) -def _semantic_usage(raw: PlainJson) -> ResponseUsage: +def semantic_usage(raw: PlainJson) -> ResponseUsage: if not isinstance(raw, dict): return ResponseUsage( input_tokens=0, diff --git a/litellm/translation/providers/openai_compat/serialize.py b/litellm/translation/providers/openai_compat/serialize.py index ffec4aae40..90cb6c4af7 100644 --- a/litellm/translation/providers/openai_compat/serialize.py +++ b/litellm/translation/providers/openai_compat/serialize.py @@ -37,6 +37,13 @@ def serialize_request(request: ChatRequest, deps: TranslationDeps) -> _Serialize ) if reason is not None: return Error(TranslationError.of_unsupported(reason)) + return assemble_body(request) + + +def assemble_body(request: ChatRequest) -> _SerializeResult: + """The gate-free five-touch body assembly, for consumers (xai) whose v1 + config inherits OpenAIGPTConfig's transform_request but replaces the + openai param gates with their own.""" messages = serialize_messages(request) if isinstance(messages, TranslationError): return Error(messages) diff --git a/litellm/translation/providers/xai/__init__.py b/litellm/translation/providers/xai/__init__.py new file mode 100644 index 0000000000..489340941f --- /dev/null +++ b/litellm/translation/providers/xai/__init__.py @@ -0,0 +1,12 @@ +from .guard import unsupported_request_shapes +from .response import parse_response +from .serialize import serialize_request +from .stream import parse_event, parse_line + +__all__ = ( + "parse_event", + "parse_line", + "parse_response", + "serialize_request", + "unsupported_request_shapes", +) diff --git a/litellm/translation/providers/xai/guard.py b/litellm/translation/providers/xai/guard.py new file mode 100644 index 0000000000..bf0b92f248 --- /dev/null +++ b/litellm/translation/providers/xai/guard.py @@ -0,0 +1,97 @@ +"""Raw-shape fidelity guard for the xai (Grok) serializer. + +Three xai-only checks run before the shared openai guard: + +- ``web_search_options``: v1 reroutes xai chat to the Responses-API bridge + BEFORE any chat seam (``responses_api_bridge_check``, main.py:982-984), so + a chat-route v2 must never serve it. +- ``use_xai_oauth``: v1's ``validate_environment`` runs an interactive + browser PKCE flow with a localhost callback server (llms/xai/oauth.py); + envelope I/O the v2 surface cannot reproduce. +- tool definitions carrying a ``strict`` key anywhere BELOW the function + level: v1's ``filter_value_from_dict(tool, "strict")`` deletes EVERY key + named ``strict`` at any depth (including JSON-schema properties literally + named ``strict``); v2 reproduces only the standard function-level strip. + +The shared openai guard then runs with ``name_fallback_user_only``: v1's xai +transform strips message ``name`` from every non-user role +(``strip_name_from_messages``), so the IR's name-drop IS v1's behavior there +and only a user-message ``name`` (forwarded verbatim by v1) falls back. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import cast + +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH + +from ...errors import TranslationError +from ..openai_compat.guard import ( + unsupported_request_shapes as openai_unsupported_request_shapes, +) + +_Raw = Mapping[str, object] + + +def unsupported_request_shapes(raw: _Raw) -> TranslationError | None: + if raw.get("web_search_options") is not None: + return TranslationError.of_unsupported( + "web_search_options on xai: v1 reroutes the request to the " + "Responses-API bridge above the chat seam " + "(responses_api_bridge_check); v1 owns it" + ) + if raw.get("use_xai_oauth"): + return TranslationError.of_unsupported( + "use_xai_oauth: v1 runs an interactive browser PKCE flow inside " + "validate_environment (llms/xai/oauth.py); v1 owns it" + ) + if "stream" in raw and raw.get("stream") is False: + return TranslationError.of_unsupported( + "explicit stream: false (the xai httpx path keeps the key on the " + "wire; absent-vs-false is lost in the IR)" + ) + reason = _nested_tool_strict_reason(raw) + if reason is not None: + return TranslationError.of_unsupported(reason) + return openai_unsupported_request_shapes(raw, name_fallback_user_only=True) + + +def _nested_tool_strict_reason(raw: _Raw) -> str | None: + raw_tools = raw.get("tools") + if not isinstance(raw_tools, Sequence) or isinstance(raw_tools, str): + return None + for tool in cast(Sequence[object], raw_tools): + if not isinstance(tool, Mapping): + continue + entry = cast(_Raw, tool) + function = entry.get("function") + if not isinstance(function, Mapping): + continue + for key, value in cast(_Raw, function).items(): + if key == "strict": + continue # the standard slot; the serializer strips it like v1 + if _contains_strict_key(value, 0): + return ( + "tool definition carries a nested 'strict' key below the " + "function level; v1's filter_value_from_dict deletes it " + "at every depth" + ) + return None + + +def _contains_strict_key(value: object, depth: int) -> bool: + if depth > DEFAULT_MAX_RECURSE_DEPTH: + return False + if isinstance(value, Mapping): + mapping = cast(_Raw, value) + return any( + key == "strict" or _contains_strict_key(item, depth + 1) + for key, item in mapping.items() + ) + if isinstance(value, Sequence) and not isinstance(value, str): + return any( + _contains_strict_key(item, depth + 1) + for item in cast(Sequence[object], value) + ) + return False diff --git a/litellm/translation/providers/xai/params.py b/litellm/translation/providers/xai/params.py new file mode 100644 index 0000000000..1006a98d6a --- /dev/null +++ b/litellm/translation/providers/xai/params.py @@ -0,0 +1,57 @@ +"""Parameter gates for the xai (Grok) serializer. + +v1's gate is ``_check_valid_arg`` over ``XAIChatConfig. +get_supported_openai_params`` (utils.py:4162-4203): an unsupported param +RAISES ``UnsupportedParamsError`` unless ``drop_params``, in which case it is +popped BEFORE ``map_openai_params`` runs — so the ``max_completion_tokens -> +max_tokens`` rename arm (x_t:185-186) is dead code on the standard path. +v2 mirrors the SUPPORTED-LIST truth, never the raise-vs-drop interplay: +every shape v1 raises or drops on falls back typed so v1 serves its own +error (the v2-openai response_format precedent). + +Per-model gates mirror x_t:155-174 (substring checks) and the +``litellm.supports_reasoning(model, "xai")`` model-map read, reproduced +through ``deps.supports_capability`` over the ``xai/{model}`` map key (the +bare wire model has no model-map row; verified in-process at HEAD). +""" + +from __future__ import annotations + +from ...deps import TranslationDeps +from ...ir import ChatRequest + +_NO_STOP_FAMILIES = ("grok-3-mini", "grok-4", "grok-code-fast") + + +def supports_stop(model: str) -> bool: + return not any(family in model for family in _NO_STOP_FAMILIES) + + +def supports_reasoning(model: str, deps: TranslationDeps) -> bool: + return deps.supports_capability(f"xai/{model}", "supports_reasoning") + + +def unsupported_params(request: ChatRequest, deps: TranslationDeps) -> str | None: + if request.params.max_completion_tokens.is_some(): + return ( + "max_completion_tokens is outside xai's supported list; v1's " + "get_optional_params raises UnsupportedParamsError (or drops it " + "under drop_params) before the rename arm can run" + ) + if request.params.top_k.is_some(): + return "top_k is not an xai chat param; v1's get_optional_params raises or drops it" + if request.thinking.is_some(): + return "thinking is not an xai chat param; v1's get_optional_params raises or drops it" + if len(request.params.stop) > 0 and not supports_stop(request.model): + return ( + f"stop on {request.model}: outside xai's supported list " + "(grok-3-mini/grok-4/grok-code-fast); v1 raises or drops it" + ) + if request.reasoning_effort.is_some() and not supports_reasoning( + request.model, deps + ): + return ( + f"reasoning_effort on non-reasoning xai model {request.model}; " + "v1 raises or drops it" + ) + return None diff --git a/litellm/translation/providers/xai/response.py b/litellm/translation/providers/xai/response.py new file mode 100644 index 0000000000..395b7dd235 --- /dev/null +++ b/litellm/translation/providers/xai/response.py @@ -0,0 +1,113 @@ +"""xai (Grok) chat-completion response JSON -> IR ``ChatResponse``. + +The live v1 normalizer is ``XAIChatConfig.transform_response`` (xai is on +the dedicated httpx branch, main.py:2289, so transform_response RUNS — +the inverse of the openai SDK path): the shared +``convert_to_model_response_object`` conversion, then the xai post-steps. +v2 mirrors that chain over the openai_compat parser's normalized wire body: + +- finish_reason ``""`` (Grok's tool-call quirk) needs NO xai arm: v1's own + ``_fix_choice_finish_reason_for_tool_calls`` is empirically dead + (``Choices.__init__`` maps ``""`` -> ``"stop"`` before the check ever + runs), so v1-as-executed emits ``"stop"`` WITH tool_calls. The openai + parser already rides the native ``""`` on the wire body and the seam's + ``Choices`` runs the same live ``map_finish_reason`` — identical chain. +- ``_enhance_usage_with_xai_web_search_fields``: ``usage.num_sources_used`` + > 0 copies into ``prompt_tokens_details.web_search_requests`` (the + live-search billing hook). +- ``_fold_reasoning_tokens_into_completion``: fold reasoning into + completion_tokens when ``total == prompt + completion + reasoning``. +- ``_normalize_openai_compatible_usage_totals``: bump total_tokens up to + ``prompt + completion``. + +``citations`` and any future live-search top-level keys ride the parser's +unknown-key mirror exactly like v1's cdr:727-729 setattr. +""" + +from __future__ import annotations + +import dataclasses + +from expression import Result, Some + +from ...errors import TranslationError +from ...ir import ChatRequest, ChatResponse, JsonBlob, PlainJson +from ..openai_compat.response import parse_response as openai_parse_response +from ..openai_compat.response import semantic_usage as openai_semantic_usage + +_ParseResult = Result[ChatResponse, TranslationError] + + +def parse_response(raw: PlainJson, request: ChatRequest) -> _ParseResult: + return openai_parse_response(raw, request).map(_with_xai_usage_post_steps) + + +def _with_xai_usage_post_steps(response: ChatResponse) -> ChatResponse: + wire = response.wire.default_value(None) + if wire is None or not isinstance(wire.value, dict): + return response + usage = wire.value.get("usage") + if not isinstance(usage, dict): + return response + transformed = normalize_usage_totals( + fold_reasoning_tokens(_websearch_fields(usage)) + ) + if transformed == usage: + return response + body: dict[str, PlainJson] = {**wire.value, "usage": transformed} + return dataclasses.replace( + response, + usage=openai_semantic_usage(transformed), + wire=Some(JsonBlob(value=body)), + ) + + +def _websearch_fields(usage: dict[str, PlainJson]) -> dict[str, PlainJson]: + sources = usage.get("num_sources_used") + if not isinstance(sources, (int, float)) or isinstance(sources, bool): + return usage + if sources <= 0: + return usage + details = usage.get("prompt_tokens_details") + seeded: dict[str, PlainJson] = dict(details) if isinstance(details, dict) else {} + return { + **usage, + "num_sources_used": int(sources), + "prompt_tokens_details": {**seeded, "web_search_requests": int(sources)}, + } + + +def fold_reasoning_tokens(usage: dict[str, PlainJson]) -> dict[str, PlainJson]: + """Pure dict mirror of ``_fold_reasoning_tokens_into_completion`` (the + same arithmetic the stream chunk variant runs).""" + details = usage.get("completion_tokens_details") + reasoning = details.get("reasoning_tokens") if isinstance(details, dict) else 0 + reasoning_tokens = _int_of(reasoning) + if reasoning_tokens <= 0: + return usage + prompt_tokens = _int_of(usage.get("prompt_tokens")) + completion_tokens = _int_of(usage.get("completion_tokens")) + total_tokens = _int_of(usage.get("total_tokens")) + if total_tokens == prompt_tokens + completion_tokens: + return usage + if total_tokens != prompt_tokens + completion_tokens + reasoning_tokens: + return usage # v1's double-count guard + return {**usage, "completion_tokens": completion_tokens + reasoning_tokens} + + +def normalize_usage_totals(usage: dict[str, PlainJson]) -> dict[str, PlainJson]: + """Pure dict mirror of ``_normalize_openai_compatible_usage_totals``.""" + expected = _int_of(usage.get("prompt_tokens")) + _int_of( + usage.get("completion_tokens") + ) + if _int_of(usage.get("total_tokens")) >= expected: + return usage + return {**usage, "total_tokens": expected} + + +def _int_of(value: PlainJson) -> int: + if isinstance(value, bool): + return 0 + if isinstance(value, (int, float)): + return int(value) + return 0 diff --git a/litellm/translation/providers/xai/serialize.py b/litellm/translation/providers/xai/serialize.py new file mode 100644 index 0000000000..94c5b515b1 --- /dev/null +++ b/litellm/translation/providers/xai/serialize.py @@ -0,0 +1,58 @@ +"""Serialize the IR into an xai (Grok) ``/v1/chat/completions`` request body. + +v1's chain is ``XAIChatConfig.map_openai_params`` (own supported list, tools +``strict`` strip) then ``transform_request`` = ``strip_name_from_messages`` ++ the inherited OpenAIGPTConfig five-touch assembly. The body is therefore +the openai_compat assembly with three deltas: the function-level ``strict`` +key is stripped from every tool (v1 ``filter_value_from_dict(tool, +"strict")``; deeper ``strict`` keys fall back in the raw guard), and ``user`` +/ ``reasoning_effort`` — typed fallbacks on plain GPT — are emitted verbatim +because xai supports them (reasoning_effort gated per model in params.py). +The non-user message ``name`` strip needs no code: the IR never carries +``name`` and the xai raw guard only falls back on user-message names. +""" + +from __future__ import annotations + +from expression import Error, Result + +from ...deps import TranslationDeps +from ...errors import TranslationError +from ...ir import Body, ChatRequest, PlainJson +from ..openai_compat.serialize import assemble_body +from . import params as p + +_SerializeResult = Result[Body, TranslationError] + + +def serialize_request(request: ChatRequest, deps: TranslationDeps) -> _SerializeResult: + reason = p.unsupported_params(request, deps) + if reason is not None: + return Error(TranslationError.of_unsupported(reason)) + return assemble_body(request).map(lambda body: _with_xai_deltas(body, request)) + + +def _with_xai_deltas(body: Body, request: ChatRequest) -> Body: + extras: dict[str, PlainJson] = {} + user = request.user.default_value(None) + if user is not None: + extras = {**extras, "user": user} + effort = request.reasoning_effort.default_value(None) + if effort is not None: + extras = {**extras, "reasoning_effort": effort} + tools = body.get("tools") + if not isinstance(tools, list): + return {**body, **extras} + return {**body, "tools": [_without_strict(tool) for tool in tools], **extras} + + +def _without_strict(tool: PlainJson) -> PlainJson: + if not isinstance(tool, dict): + return tool + function = tool.get("function") + if not isinstance(function, dict): + return tool + return { + **tool, + "function": {key: value for key, value in function.items() if key != "strict"}, + } diff --git a/litellm/translation/providers/xai/stream.py b/litellm/translation/providers/xai/stream.py new file mode 100644 index 0000000000..f1c312f834 --- /dev/null +++ b/litellm/translation/providers/xai/stream.py @@ -0,0 +1,224 @@ +"""xai (Grok) SSE ``chat.completion.chunk`` payloads -> IR stream events. + +The v1 decode is the httpx dict path: ``BaseModelResponseIterator`` strips +``data:`` lines, ``XAIChatCompletionStreamingHandler.chunk_parser`` applies +the xai rewrites, and the rebuilt ``ModelResponseStream`` flows through +``CustomStreamWrapper``'s default openai branch. The xai deltas vs the +openai parser (all verified in-process at HEAD): + +- the chunk_parser rebuild keeps ONLY id/created/model/choices/usage: + ``system_fingerprint`` and every top-level extra (``citations``) are + DROPPED — the opposite of the SDK path's verbatim extras passthrough. +- a ``choices: []`` chunk carrying ``usage`` gets a dummy choice injected in + v1 purely so the wrapper machinery swallows it and re-synthesizes the + final usage chunk; v2 keeps the openai passthrough shape (``choices: []`` + + the FOLDED usage) so the seam's synthesized-final-usage contract from + the openai port applies unchanged. +- every usage-bearing chunk gets the reasoning fold + total normalize + (the dict variants of the response post-steps). +- ``delta.reasoning`` is renamed to ``reasoning_content`` (the base + handler's rename) and native ``reasoning_content`` deltas are admitted — + real Grok reasoning traffic, not an unreachable shape. +- a tool_call entry WITHOUT a ``type`` key gains ``type: "function"`` + (litellm's ``ChatCompletionDeltaToolCall`` default fires on the dict + path; the SDK path yields None there, so the openai parser must not). +""" + +from __future__ import annotations + +import json + +from expression import Error, Ok, Result + +from ...errors import BoundaryError, TranslationError +from ...ir import JsonBlob, PlainJson, StreamEvent +from .response import fold_reasoning_tokens, normalize_usage_totals + +_EventResult = Result[StreamEvent | None, TranslationError] + +_DELTA_KEYS = ( + "content", + "function_call", + "refusal", + "role", + "tool_calls", + "reasoning", + "reasoning_content", +) + + +def parse_line(line: str) -> _EventResult: + stripped = line.strip() + if not stripped.startswith("data:"): + return Ok(None) + payload = stripped[len("data:") :].strip() + if payload == "[DONE]": + return Ok(StreamEvent.of_stop()) + try: + event: PlainJson = json.loads(payload) + except ValueError: + return Error(_boundary(f"stream payload is not JSON: {payload[:120]!r}")) + return parse_event(event) + + +def parse_event(event: PlainJson) -> _EventResult: + if not isinstance(event, dict): + return Error(_boundary("stream chunk is not an object")) + if event.get("error") is not None: + return Error(_boundary(f"provider stream error: {event.get('error')!r}")) + choices = event.get("choices") + if not isinstance(choices, list): + return Error(_boundary("stream chunk 'choices' is missing")) + if len(choices) > 1: + return Error( + TranslationError.of_unsupported( + "multiple stream choices (n > 1); unreachable for v2-sent requests" + ) + ) + raw_usage = event.get("usage") + usage: PlainJson = ( + normalize_usage_totals(fold_reasoning_tokens(raw_usage)) + if isinstance(raw_usage, dict) + else None + ) + normalized_choices: list[PlainJson] = [] + if len(choices) == 1: + normalized = _normalize_choice(choices[0]) + if isinstance(normalized, TranslationError): + return Error(normalized) + normalized_choices = [normalized] + identifier = event.get("id") + chunk: dict[str, PlainJson] = { + # No extras passthrough and no system_fingerprint: the v1 chunk_parser + # rebuild keeps only id/created/model/choices/usage. + "id": identifier if isinstance(identifier, str) else None, + "system_fingerprint": None, + "choices": normalized_choices, + "usage": usage, + } + return Ok(StreamEvent.of_wire_chunk(JsonBlob(value=chunk))) + + +def _boundary(reason: str) -> TranslationError: + from expression.collections import Block + + return TranslationError.of_boundary(BoundaryError.of(Block.of_seq([reason]))) + + +def _string_or_none(value: PlainJson) -> PlainJson: + return value if isinstance(value, str) else None + + +def _normalize_choice(choice: PlainJson) -> PlainJson | TranslationError: + if not isinstance(choice, dict): + return _boundary("stream choice is not an object") + extra_keys = set(choice.keys()) - {"index", "delta", "logprobs", "finish_reason"} + if extra_keys: + return TranslationError.of_unsupported( + f"stream choice keys {sorted(extra_keys)!r}; unreachable for v2-sent requests" + ) + if choice.get("logprobs") is not None: + return TranslationError.of_unsupported( + "stream logprobs; unreachable for v2-sent requests" + ) + finish = choice.get("finish_reason") + if finish is not None and not isinstance(finish, str): + return _boundary("stream finish_reason is not a string") + if finish == "function_call": + return TranslationError.of_unsupported( + "legacy function_call stream finish; the v2 surface cannot send 'functions'" + ) + raw_delta = choice.get("delta") + delta = raw_delta if isinstance(raw_delta, dict) else {} + normalized_delta = _normalize_delta(delta) + if isinstance(normalized_delta, TranslationError): + return normalized_delta + if finish is not None and _delta_bears_content(normalized_delta): + return TranslationError.of_unsupported( + "finish chunk with a non-empty delta; v1's wrapper interleaves it" + ) + index = choice.get("index") + return { + "index": index if isinstance(index, int) else 0, + "delta": normalized_delta, + "logprobs": None, + "finish_reason": finish, + } + + +def _normalize_delta( + delta: dict[str, PlainJson], +) -> dict[str, PlainJson] | TranslationError: + extra_keys = set(delta.keys()) - set(_DELTA_KEYS) + if extra_keys: + return TranslationError.of_unsupported( + f"stream delta keys {sorted(extra_keys)!r}; unreachable for v2-sent requests" + ) + if delta.get("function_call") is not None: + return TranslationError.of_unsupported( + "legacy function_call stream delta; the v2 surface cannot send 'functions'" + ) + tool_calls = delta.get("tool_calls") + normalized_calls: PlainJson = None + if tool_calls is not None: + if not isinstance(tool_calls, list): + return _boundary("stream delta 'tool_calls' is not an array") + gathered: list[PlainJson] = [] + for call in tool_calls: + normalized = _normalize_tool_call(call) + if isinstance(normalized, TranslationError): + return normalized + gathered = [*gathered, normalized] + normalized_calls = gathered + # No "refusal" key and an explicit provider_specific_fields: None — the + # dict-path wrapper rebuild drops refusal and always stamps the null + # provider field on content-bearing deltas (verified in-process replay). + base: dict[str, PlainJson] = { + "content": _string_or_none(delta.get("content")), + "function_call": None, + "provider_specific_fields": None, + "role": _string_or_none(delta.get("role")), + "tool_calls": normalized_calls, + } + # the base handler renames delta.reasoning unconditionally + # (_map_reasoning_to_reasoning_content); the key is only present when the + # wire carried one, mirroring Delta's set-only serialization. + reasoning = ( + delta.get("reasoning") + if "reasoning" in delta + else (delta.get("reasoning_content") if "reasoning_content" in delta else None) + ) + if "reasoning" in delta or "reasoning_content" in delta: + return {**base, "reasoning_content": _string_or_none(reasoning)} + return base + + +def _normalize_tool_call(call: PlainJson) -> PlainJson | TranslationError: + if not isinstance(call, dict): + return _boundary("stream tool_call is not an object") + extra_keys = set(call.keys()) - {"index", "id", "function", "type"} + if extra_keys: + return TranslationError.of_unsupported( + f"stream tool_call keys {sorted(extra_keys)!r}; unreachable for v2-sent requests" + ) + raw_function = call.get("function") + function = raw_function if isinstance(raw_function, dict) else {} + index = call.get("index") + return { + "index": index if isinstance(index, int) else 0, + "id": _string_or_none(call.get("id")), + "function": { + "arguments": _string_or_none(function.get("arguments")), + "name": _string_or_none(function.get("name")), + }, + # dict-path default: ChatCompletionDeltaToolCall fills "function" when + # the wire omits the key (the SDK path keeps None there) + "type": _string_or_none(call.get("type")) if "type" in call else "function", + } + + +def _delta_bears_content(delta: dict[str, PlainJson]) -> bool: + content = delta.get("content") + return (isinstance(content, str) and len(content) > 0) or delta.get( + "tool_calls" + ) is not None