mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-24 05:45:11 +00:00
fix(proxy): report requested model in streamed Responses response.model (#46662)
* chore: Add failing test case Signed-off-by: Alberto Perdomo <aperdomo@redhat.com> * fix: Restamp nested response.model in streamed Responses events Signed-off-by: Alberto Perdomo <aperdomo@redhat.com> * chore: Add tests Signed-off-by: Alberto Perdomo <aperdomo@redhat.com> * chore: Run format and lint, and fix Signed-off-by: Alberto Perdomo <aperdomo@redhat.com> * fix: Type nested Responses restamp to satisfy basedpyright gate Signed-off-by: Alberto Perdomo <aperdomo@redhat.com> * fix(proxy): restamp the nested model on streamed response.queued events Background streams send response.queued as a GenericEvent, so its nested response.model still carried the backend name. Also drop the new comments * fix(proxy): restamp a nested response kept as a dict on streamed lifecycle events * test(proxy): drive typed events through the nested restamp no-op test * fix(proxy): restamp the nested model on guardrail block frames in streamed responses * test(proxy): assert a cache-polled background response reports the requested model * fix(proxy): restamp nested model on dict events and on every SSE data line --------- Signed-off-by: Alberto Perdomo <aperdomo@redhat.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
alberto
GitHub
parent
4846cfabc7
commit
bbd446b476
3 changed files with 500 additions and 24 deletions
|
|
@ -954,7 +954,14 @@ from litellm.types.llms.openai import (
|
|||
AllMessageValues,
|
||||
ChatCompletionSystemMessage,
|
||||
ChatCompletionToolParam,
|
||||
GenericEvent,
|
||||
HttpxBinaryResponseContent,
|
||||
ResponseCompletedEvent,
|
||||
ResponseCreatedEvent,
|
||||
ResponseFailedEvent,
|
||||
ResponseIncompleteEvent,
|
||||
ResponseInProgressEvent,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
|
|
@ -9922,7 +9929,7 @@ def _restamp_streaming_chunk_model(
|
|||
# Note: This warning is intentionally verbose. A mismatch is a useful signal that an
|
||||
# internal provider/deployment identifier is leaking into the public API, and helps
|
||||
# maintainers/operators catch regressions while preserving OpenAI-compatible output.
|
||||
if not target_model or not isinstance(chunk, (BaseModel, dict)):
|
||||
if not target_model or not isinstance(chunk, (BaseModel, dict, bytes)):
|
||||
return chunk, model_mismatch_logged
|
||||
|
||||
# For Azure Model Router, preserve the actual model used in each chunk
|
||||
|
|
@ -9934,9 +9941,21 @@ def _restamp_streaming_chunk_model(
|
|||
if not fallback_was_attempted and request_data.get("fastest_response", False):
|
||||
return chunk, model_mismatch_logged
|
||||
|
||||
downstream_model: Final = chunk.get("model") if isinstance(chunk, dict) else getattr(chunk, "model", None)
|
||||
if isinstance(chunk, bytes):
|
||||
if request_data.get("_litellm_raw_sse_stream"):
|
||||
return chunk, model_mismatch_logged
|
||||
return _restamp_nested_response_model_in_sse_frame(chunk, target_model), model_mismatch_logged
|
||||
|
||||
typed_chunk: Final[BaseModel | Mapping[str, object]] = chunk
|
||||
restamped_chunk: Final = _restamp_nested_response_model(typed_chunk, target_model)
|
||||
|
||||
downstream_model: Final = (
|
||||
restamped_chunk.get("model")
|
||||
if isinstance(restamped_chunk, Mapping)
|
||||
else getattr(restamped_chunk, "model", None)
|
||||
)
|
||||
if downstream_model == target_model:
|
||||
return chunk, model_mismatch_logged
|
||||
return restamped_chunk, model_mismatch_logged
|
||||
|
||||
if not model_mismatch_logged and downstream_model != target_model:
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -9954,22 +9973,92 @@ def _restamp_streaming_chunk_model(
|
|||
logging_obj: Final = request_data.get("litellm_logging_obj")
|
||||
if isinstance(logging_obj, LiteLLMLoggingObj):
|
||||
logging_obj.client_facing_stream_model = target_model
|
||||
if isinstance(chunk, dict):
|
||||
return {**chunk, "model": target_model}, model_mismatch_logged
|
||||
if isinstance(restamped_chunk, Mapping):
|
||||
return {**restamped_chunk, "model": target_model}, model_mismatch_logged
|
||||
|
||||
try:
|
||||
return chunk.model_copy(update={"model": target_model}), model_mismatch_logged
|
||||
return restamped_chunk.model_copy(update={"model": target_model}), model_mismatch_logged
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
"litellm_call_id=%s: failed to override chunk.model=%r on chunk_type=%s. error=%s",
|
||||
request_data.get("litellm_call_id"),
|
||||
target_model,
|
||||
type(chunk),
|
||||
type(restamped_chunk),
|
||||
str(e),
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
return chunk, model_mismatch_logged
|
||||
return restamped_chunk, model_mismatch_logged
|
||||
|
||||
|
||||
def _nested_response_of(chunk: BaseModel) -> object:
|
||||
if isinstance(
|
||||
chunk,
|
||||
(
|
||||
ResponseCreatedEvent,
|
||||
ResponseInProgressEvent,
|
||||
ResponseCompletedEvent,
|
||||
ResponseFailedEvent,
|
||||
ResponseIncompleteEvent,
|
||||
),
|
||||
):
|
||||
return chunk.response
|
||||
if isinstance(chunk, GenericEvent):
|
||||
extra_fields: Final[Mapping[str, object]] = chunk.model_extra or {}
|
||||
return extra_fields.get("response")
|
||||
return None
|
||||
|
||||
|
||||
def _restamp_nested_response_model(
|
||||
chunk: BaseModel | Mapping[str, object], target_model: str
|
||||
) -> BaseModel | Mapping[str, object]:
|
||||
if not isinstance(chunk, BaseModel):
|
||||
restamped_event: Final = _event_with_restamped_nested_model(chunk, target_model)
|
||||
return chunk if restamped_event is None else restamped_event
|
||||
nested_response: Final = _nested_response_of(chunk)
|
||||
if isinstance(nested_response, ResponsesAPIResponse):
|
||||
if nested_response.model in (None, target_model):
|
||||
return chunk
|
||||
return chunk.model_copy(update={"response": nested_response.model_copy(update={"model": target_model})})
|
||||
untyped_response: Final = as_str_mapping(nested_response)
|
||||
if untyped_response is None or untyped_response.get("model") in (None, target_model):
|
||||
return chunk
|
||||
return chunk.model_copy(update={"response": {**untyped_response, "model": target_model}})
|
||||
|
||||
|
||||
def _event_with_restamped_nested_model(event: Mapping[str, object], target_model: str) -> Mapping[str, object] | None:
|
||||
nested_response: Final = as_str_mapping(event.get("response"))
|
||||
if nested_response is None or nested_response.get("model") in (None, target_model):
|
||||
return None
|
||||
return {**event, "response": {**nested_response, "model": target_model}}
|
||||
|
||||
|
||||
_SSE_DATA_FIELD: Final = b"data:"
|
||||
_SSE_NESTED_RESPONSE_MARKER: Final = b'"response"'
|
||||
_SSE_EVENT_OBJECT: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _restamp_nested_response_model_in_sse_frame(frame: bytes, target_model: str) -> bytes:
|
||||
if _SSE_NESTED_RESPONSE_MARKER not in frame:
|
||||
return frame
|
||||
return b"".join(
|
||||
_restamp_nested_response_model_in_sse_line(line, target_model) for line in frame.splitlines(keepends=True)
|
||||
)
|
||||
|
||||
|
||||
def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) -> bytes:
|
||||
if not line.startswith(_SSE_DATA_FIELD):
|
||||
return line
|
||||
try:
|
||||
event: Final = _SSE_EVENT_OBJECT.validate_json(line.removeprefix(_SSE_DATA_FIELD).strip())
|
||||
except ValidationError:
|
||||
return line
|
||||
restamped_event: Final = _event_with_restamped_nested_model(event, target_model)
|
||||
if restamped_event is None:
|
||||
return line
|
||||
serialized_event: Final = json.dumps(restamped_event, separators=(",", ":"), ensure_ascii=False)
|
||||
line_ending: Final = line[len(line.rstrip(b"\r\n")) :]
|
||||
return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending
|
||||
|
||||
|
||||
def _fast_serialize_simple_model_response_stream(
|
||||
|
|
|
|||
|
|
@ -549,7 +549,7 @@ def _cache_polled_create(proxy: Gateway, key: str, model: str, create: Mapping[s
|
|||
|
||||
@pytest.mark.timeout(OWNED_PROXY_SECONDS + 120)
|
||||
@pytest.mark.parametrize("create", CACHE_POLLED_CREATES.keys())
|
||||
def test_a_background_response_the_proxy_polls_through_its_cache_bills_its_creator_once(
|
||||
def test_a_background_response_polled_through_the_cache_reports_its_alias_and_bills_its_creator_once(
|
||||
gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, create: str
|
||||
) -> None:
|
||||
body, provider_sees_background = CACHE_POLLED_CREATES[create]
|
||||
|
|
@ -564,6 +564,8 @@ def test_a_background_response_the_proxy_polls_through_its_cache_bills_its_creat
|
|||
seconds=SPEND_RETRY_SECONDS,
|
||||
)
|
||||
assert finished.text == "pong"
|
||||
polled: Final = object_value(proxy.request("GET", f"/v1/responses/{polling_id}", key=key).json())
|
||||
assert polled["model"] == handle.scenario_id
|
||||
assert_charged_exactly(lambda: key_spend(key), PRICE)
|
||||
assert [read_response(proxy, key, polling_id).text for _ in range(READS)] == ["pong"] * READS
|
||||
assert_charged_exactly(lambda: key_spend(key), PRICE)
|
||||
|
|
|
|||
|
|
@ -68,6 +68,19 @@ def _decode_sse_chunk(chunk) -> str:
|
|||
return chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk
|
||||
|
||||
|
||||
def _make_response_created_event(model: str | None):
|
||||
from litellm.types.llms.openai import (
|
||||
ResponseCreatedEvent,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
)
|
||||
|
||||
return ResponseCreatedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_CREATED,
|
||||
response=ResponsesAPIResponse(id="resp_test", created_at=0, output=[], model=model),
|
||||
)
|
||||
|
||||
|
||||
def test_restamp_streaming_chunk_skips_matching_model():
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
|
|
@ -117,19 +130,13 @@ def test_fast_serialize_returns_none_when_model_field_is_missing():
|
|||
# produce a serialized result via the slow-path fallback, and it must
|
||||
# not contain ``"model": null``.
|
||||
serialized = _serialize_streaming_chunk(chunk)
|
||||
payload_str = (
|
||||
serialized.decode("utf-8") if isinstance(serialized, bytes) else serialized
|
||||
)
|
||||
payload_str = serialized.decode("utf-8") if isinstance(serialized, bytes) else serialized
|
||||
assert '"model": null' not in payload_str
|
||||
assert '"model":null' not in payload_str
|
||||
assert json.loads(payload_str) == json.loads(
|
||||
chunk.model_dump_json(exclude_none=True, exclude_unset=True)
|
||||
)
|
||||
assert json.loads(payload_str) == json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True))
|
||||
|
||||
|
||||
def test_proxy_chat_completion_does_not_return_provider_prefixed_model(
|
||||
tmp_path, monkeypatch
|
||||
):
|
||||
def test_proxy_chat_completion_does_not_return_provider_prefixed_model(tmp_path, monkeypatch):
|
||||
"""
|
||||
Regression test:
|
||||
|
||||
|
|
@ -159,15 +166,11 @@ def test_proxy_chat_completion_does_not_return_provider_prefixed_model(
|
|||
monkeypatch.setattr(
|
||||
proxy_server.llm_router, # type: ignore[arg-type]
|
||||
"acompletion",
|
||||
AsyncMock(
|
||||
return_value=_make_minimal_chat_completion_response(model=internal_model)
|
||||
),
|
||||
AsyncMock(return_value=_make_minimal_chat_completion_response(model=internal_model)),
|
||||
)
|
||||
|
||||
# Also no-op proxy logging hooks to keep this test focused and deterministic.
|
||||
monkeypatch.setattr(
|
||||
proxy_server.proxy_logging_obj, "during_call_hook", AsyncMock(return_value=None)
|
||||
)
|
||||
monkeypatch.setattr(proxy_server.proxy_logging_obj, "during_call_hook", AsyncMock(return_value=None))
|
||||
monkeypatch.setattr(
|
||||
proxy_server.proxy_logging_obj,
|
||||
"update_request_status",
|
||||
|
|
@ -482,3 +485,385 @@ async def test_proxy_streaming_fastest_response_preserves_winning_model(monkeypa
|
|||
payload = json.loads(first[len("data: ") :].strip())
|
||||
assert payload["model"] == winning_model
|
||||
assert payload["model"] != comma_separated_models
|
||||
|
||||
|
||||
def test_restamp_streaming_chunk_rewrites_nested_response_model():
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
chunk = _make_response_created_event("hosted_vllm/some-model")
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=chunk,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"litellm_call_id": "test-call-id"},
|
||||
model_mismatch_logged=False,
|
||||
)
|
||||
|
||||
assert result.response.model == "client-model"
|
||||
assert chunk.response.model == "hosted_vllm/some-model"
|
||||
|
||||
|
||||
def _make_response_completed_event(model: str):
|
||||
from litellm.types.llms.openai import (
|
||||
ResponseCompletedEvent,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
)
|
||||
|
||||
return ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=ResponsesAPIResponse(id="resp_test", created_at=0, output=[], model=model),
|
||||
)
|
||||
|
||||
|
||||
def test_restamp_streaming_chunk_rewrites_nested_model_on_completed_event():
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
chunk = _make_response_completed_event("hosted_vllm/some-model")
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=chunk,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"litellm_call_id": "test-call-id"},
|
||||
model_mismatch_logged=False,
|
||||
)
|
||||
|
||||
assert result.response.model == "client-model"
|
||||
assert chunk.response.model == "hosted_vllm/some-model", "the priced original must not change"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("event_type", ["response.in_progress", "response.failed", "response.incomplete"])
|
||||
def test_restamp_streaming_chunk_rewrites_nested_model_on_other_lifecycle_events(event_type):
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
payload = {
|
||||
"type": event_type,
|
||||
"sequence_number": 1,
|
||||
"response": {"id": "resp_test", "object": "response", "created_at": 0, "output": [], "model": "gpt-backend"},
|
||||
}
|
||||
chunk = OpenAIResponsesAPIConfig.get_event_model_class(event_type).model_validate(payload)
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=chunk,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"litellm_call_id": "test-call-id"},
|
||||
model_mismatch_logged=False,
|
||||
)
|
||||
|
||||
assert result.response.model == "client-model"
|
||||
assert chunk.response.model == "gpt-backend", "the priced original must not change"
|
||||
|
||||
|
||||
def test_restamp_streaming_chunk_leaves_nested_none_model_untouched():
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
chunk = _make_response_created_event(None)
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=chunk,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"litellm_call_id": "test-call-id"},
|
||||
model_mismatch_logged=False,
|
||||
)
|
||||
|
||||
assert result.response.model is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("event_type", ["response.output_text.delta", "response.custom_event"])
|
||||
def test_restamp_streaming_chunk_does_not_add_nested_response_to_other_events(event_type):
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
chunk = OpenAIResponsesAPIConfig.get_event_model_class(event_type).model_validate(
|
||||
{
|
||||
"type": event_type,
|
||||
"sequence_number": 3,
|
||||
"item_id": "msg_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "hi",
|
||||
}
|
||||
)
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=chunk,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"litellm_call_id": "test-call-id"},
|
||||
model_mismatch_logged=False,
|
||||
)
|
||||
|
||||
assert "response" not in result.model_dump()
|
||||
assert result.delta == "hi"
|
||||
|
||||
|
||||
def test_restamp_streaming_chunk_keeps_nested_model_when_raw_name_requested():
|
||||
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
chunk = _make_response_created_event("hosted_vllm/some-model")
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=chunk,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"metadata": {RETURN_RAW_MODEL_NAME_METADATA_KEY: True}},
|
||||
model_mismatch_logged=False,
|
||||
)
|
||||
|
||||
assert result.response.model == "hosted_vllm/some-model"
|
||||
|
||||
|
||||
def test_restamp_streaming_chunk_uses_fallback_model_for_nested_response():
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
chunk = _make_response_created_event("hosted_vllm/some-model")
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=chunk,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"litellm_call_id": "test-call-id"},
|
||||
model_mismatch_logged=False,
|
||||
fallback_was_attempted=True,
|
||||
fallback_model_from_metadata="fallback-public",
|
||||
)
|
||||
|
||||
assert result.response.model == "fallback-public"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_responses_stream_reports_requested_model_in_nested_response(monkeypatch):
|
||||
"""
|
||||
End-to-end regression for /v1/responses streaming:
|
||||
|
||||
The nested `response.model` on lifecycle events reaching the wire must report the
|
||||
client-requested model, not the provider-prefixed deployment name.
|
||||
"""
|
||||
client_model = "vllm-model"
|
||||
internal_model = f"hosted_vllm/{client_model}"
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
async def _iterator_hook(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: AsyncGenerator,
|
||||
request_data: dict,
|
||||
):
|
||||
yield _make_response_created_event(model=internal_model)
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_server.proxy_logging_obj,
|
||||
"async_post_call_streaming_iterator_hook",
|
||||
_iterator_hook,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
proxy_server.proxy_logging_obj,
|
||||
"async_post_call_streaming_hook",
|
||||
AsyncMock(side_effect=lambda **kwargs: kwargs["response"]),
|
||||
)
|
||||
for name in ("has_streaming_callbacks", "needs_iterator_wrap", "needs_per_chunk_streaming_hook"):
|
||||
monkeypatch.setattr(proxy_server.proxy_logging_obj, name, MagicMock(return_value=True))
|
||||
|
||||
gen = proxy_server.async_data_generator(
|
||||
response=MagicMock(),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key=MASTER_KEY),
|
||||
request_data={"model": client_model},
|
||||
responses_stream_errors=True,
|
||||
)
|
||||
|
||||
payloads = []
|
||||
async for item in gen:
|
||||
text = _decode_sse_chunk(item)
|
||||
for line in text.splitlines():
|
||||
if line.startswith("data: ") and line.strip() != "data: [DONE]":
|
||||
payloads.append(json.loads(line[len("data: ") :].strip()))
|
||||
|
||||
created = next(p for p in payloads if p.get("type") == "response.created")
|
||||
assert created["response"]["model"] == client_model
|
||||
assert not created["response"]["model"].startswith("hosted_vllm/")
|
||||
|
||||
|
||||
def test_restamp_streaming_chunk_rewrites_nested_model_on_queued_event():
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
queued_payload = {
|
||||
"type": "response.queued",
|
||||
"sequence_number": 0,
|
||||
"response": {"id": "resp_test", "object": "response", "status": "queued", "model": "gpt-backend-2026-01-01"},
|
||||
}
|
||||
chunk = OpenAIResponsesAPIConfig.get_event_model_class("response.queued").model_validate(queued_payload)
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=chunk,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"litellm_call_id": "test-call-id"},
|
||||
model_mismatch_logged=False,
|
||||
)
|
||||
|
||||
assert result.model_dump()["response"] == {**queued_payload["response"], "model": "client-model"}
|
||||
assert chunk.model_dump()["response"]["model"] == "gpt-backend-2026-01-01", "the priced original must not change"
|
||||
|
||||
|
||||
def test_restamp_streaming_chunk_rewrites_nested_model_when_lifecycle_event_kept_raw_response():
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
partial_response = {"id": "resp_test", "object": "response", "status": "in_progress", "model": "gpt-backend-2026-01-01"}
|
||||
chunk = OpenAIResponsesAPIConfig().transform_streaming_response(
|
||||
model="gpt-backend",
|
||||
parsed_chunk={"type": "response.created", "sequence_number": 0, "response": partial_response},
|
||||
logging_obj=MagicMock(),
|
||||
)
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=chunk,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"litellm_call_id": "test-call-id"},
|
||||
model_mismatch_logged=False,
|
||||
)
|
||||
|
||||
assert result.response == {**partial_response, "model": "client-model"}
|
||||
assert chunk.response["model"] == "gpt-backend-2026-01-01", "the priced original must not change"
|
||||
|
||||
|
||||
def test_restamp_streaming_chunk_leaves_generic_event_without_nested_response_untouched():
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
shell_payload = {"type": "response.shell_call.in_progress", "sequence_number": 3, "item_id": "sh_1"}
|
||||
chunk = OpenAIResponsesAPIConfig.get_event_model_class(shell_payload["type"]).model_validate(shell_payload)
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=chunk,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"litellm_call_id": "test-call-id"},
|
||||
model_mismatch_logged=False,
|
||||
)
|
||||
|
||||
assert "response" not in result.model_dump()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream_started", [False, True])
|
||||
def test_restamp_streaming_chunk_rewrites_nested_model_in_guardrail_block_frames(stream_started):
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
sent_before_block = (
|
||||
[{"type": "response.created", "response": {"id": "resp_live", "model": "gpt-backend-2026-01-01"}}]
|
||||
if stream_started
|
||||
else []
|
||||
)
|
||||
block = ModifyResponseException(
|
||||
message="Blocked by policy.",
|
||||
model="gpt-backend-2026-01-01",
|
||||
request_data={},
|
||||
original_response=sent_before_block,
|
||||
)
|
||||
frames = OpenAIResponsesHandler().build_block_sse_chunks(
|
||||
block, stream_started=stream_started, responses_so_far=sent_before_block
|
||||
)
|
||||
client_frames = [
|
||||
_restamp_streaming_chunk_model(
|
||||
chunk=frame,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"litellm_call_id": "test-call-id"},
|
||||
model_mismatch_logged=False,
|
||||
)[0]
|
||||
for frame in frames
|
||||
]
|
||||
events = [json.loads(frame.decode().removeprefix("data: ")) for frame in client_frames]
|
||||
|
||||
assert all(frame.endswith(b"\n\n") for frame in client_frames)
|
||||
assert events[-1]["type"] == "response.completed"
|
||||
assert events[-1]["response"]["output"][0]["content"][0]["text"] == "Blocked by policy."
|
||||
assert {event["response"]["model"] for event in events if "response" in event} == {"client-model"}
|
||||
assert [frame for frame, event in zip(client_frames, events) if "response" not in event] == [
|
||||
frame for frame in frames if b'"response"' not in frame
|
||||
]
|
||||
|
||||
|
||||
def test_restamp_streaming_chunk_rewrites_nested_model_on_dict_lifecycle_events():
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
chunk = {"type": "response.completed", "response": {"id": "resp_test", "model": "gpt-backend-2026-01-01"}}
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=chunk,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"litellm_call_id": "test-call-id"},
|
||||
model_mismatch_logged=False,
|
||||
)
|
||||
|
||||
assert result["response"] == {"id": "resp_test", "model": "client-model"}
|
||||
assert chunk["response"]["model"] == "gpt-backend-2026-01-01", "the priced original must not change"
|
||||
|
||||
|
||||
_BACKEND_COMPLETED_EVENT = b'{"type":"response.completed","response":{"id":"resp_1","model":"gpt-backend-2026-01-01"}}'
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"frame",
|
||||
[
|
||||
pytest.param(b"event: response.completed\ndata: " + _BACKEND_COMPLETED_EVENT + b"\n\n", id="event_line_first"),
|
||||
pytest.param(b"data:" + _BACKEND_COMPLETED_EVENT + b"\n\n", id="no_space_after_data"),
|
||||
pytest.param(
|
||||
b"data: " + _BACKEND_COMPLETED_EVENT + b"\n\ndata: " + _BACKEND_COMPLETED_EVENT + b"\r\n\r\n",
|
||||
id="coalesced_frames",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_restamp_streaming_chunk_rewrites_nested_model_on_every_sse_data_line(frame):
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=frame,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"litellm_call_id": "test-call-id"},
|
||||
model_mismatch_logged=False,
|
||||
)
|
||||
|
||||
def is_data_line(line: bytes) -> bool:
|
||||
return line.startswith(b"data:")
|
||||
|
||||
result_lines = result.splitlines(keepends=True)
|
||||
frame_lines = frame.splitlines(keepends=True)
|
||||
data_events = [json.loads(line.removeprefix(b"data:")) for line in result_lines if is_data_line(line)]
|
||||
assert len(data_events) == len([line for line in frame_lines if is_data_line(line)])
|
||||
assert {event["response"]["model"] for event in data_events} == {"client-model"}
|
||||
assert [line for line in result_lines if not is_data_line(line)] == [
|
||||
line for line in frame_lines if not is_data_line(line)
|
||||
]
|
||||
assert [line[len(line.rstrip(b"\r\n")) :] for line in result_lines] == [
|
||||
line[len(line.rstrip(b"\r\n")) :] for line in frame_lines
|
||||
]
|
||||
|
||||
|
||||
_SKIPPED_RESTAMPS = {
|
||||
"azure_model_router": ("azure_ai/model_router", {}),
|
||||
"fastest_response": ("openai/gpt-4o,gemini/gemini-2.5-flash", {"fastest_response": True}),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("skip", _SKIPPED_RESTAMPS.keys())
|
||||
def test_restamp_streaming_chunk_keeps_nested_model_where_the_top_level_model_is_kept(skip):
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
requested_model, request_data = _SKIPPED_RESTAMPS[skip]
|
||||
frame = b"data: " + _BACKEND_COMPLETED_EVENT + b"\n\n"
|
||||
results = [
|
||||
_restamp_streaming_chunk_model(
|
||||
chunk=chunk,
|
||||
requested_model_from_client=requested_model,
|
||||
request_data=request_data,
|
||||
model_mismatch_logged=False,
|
||||
)[0]
|
||||
for chunk in (_make_response_created_event("gpt-backend-2026-01-01"), frame)
|
||||
]
|
||||
|
||||
assert results[0].response.model == "gpt-backend-2026-01-01"
|
||||
assert results[1] == frame
|
||||
|
||||
|
||||
def test_restamp_streaming_chunk_leaves_raw_sse_passthrough_frames_untouched():
|
||||
from litellm.proxy.proxy_server import _restamp_streaming_chunk_model
|
||||
|
||||
frame = b"data: " + _BACKEND_COMPLETED_EVENT + b"\n\n"
|
||||
result, _ = _restamp_streaming_chunk_model(
|
||||
chunk=frame,
|
||||
requested_model_from_client="client-model",
|
||||
request_data={"_litellm_raw_sse_stream": True},
|
||||
model_mismatch_logged=False,
|
||||
)
|
||||
|
||||
assert result == frame
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue