mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-03 06:23:06 +00:00
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.
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"},
|
||||
}
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user