diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index fc29a6e3a08..ce7bf3da440 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -338,7 +338,7 @@ def resolve_litellm_call_id(client_call_id: str | None) -> str: return str(uuid.uuid4()) -def should_return_raw_model_name(request_data: dict[str, object]) -> bool: +def should_return_raw_model_name(request_data: Mapping[str, object]) -> bool: return any( isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True for metadata in (request_data.get("metadata"), request_data.get("litellm_metadata")) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4285b1f16bf..bbb62fcf490 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9916,13 +9916,15 @@ def _format_fallback_metadata_sse_event( fallback_model: str | None, fallback_errors: list[dict[str, object]], error_state: ResponsesStreamErrorState | None, + return_raw_model_name: bool, ) -> str | bytes: litellm_fallback: Final = {"fallback_model": fallback_model, "errors": fallback_errors} if error_state is not None: + stamped_model: Final = None if return_raw_model_name else fallback_model event: Final = { "type": _RESPONSES_FALLBACK_METADATA_EVENT_TYPE, "id": _FALLBACK_METADATA_ID, - **({"model": fallback_model} if fallback_model else {}), + **({"model": stamped_model} if stamped_model else {}), "litellm_fallback": litellm_fallback, } error_state.observe_chunk(event) @@ -9942,7 +9944,7 @@ def _restamp_streaming_chunk_model( *, chunk: object, requested_model_from_client: str, - request_data: dict, + request_data: Mapping[str, object], model_mismatch_logged: bool, fallback_was_attempted: bool = False, fallback_model_from_metadata: str | None = None, @@ -10491,6 +10493,7 @@ async def async_data_generator( model_mismatch_logged = False fallback_metadata_event_sent = False include_fallback_errors: Final = _should_include_fallback_errors(request_data) + return_raw_model_name: Final = should_return_raw_model_name(request_data) # Use a running string instead of list + join to avoid O(n^2) overhead. # Previously "".join(str_so_far_parts) was called every chunk, re-joining # the entire accumulated response. String += is O(n) amortized total. @@ -10586,6 +10589,7 @@ async def async_data_generator( fallback_model=fallback_model_from_metadata, fallback_errors=fallback_errors, error_state=error_state, + return_raw_model_name=return_raw_model_name, ) fallback_metadata_event_sent = True continue @@ -10640,6 +10644,7 @@ async def async_data_generator( fallback_model=fallback_model_from_metadata, fallback_errors=fallback_errors, error_state=error_state, + return_raw_model_name=return_raw_model_name, ) fallback_metadata_event_sent = True diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 596c86055e6..6c35d919f56 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -505,6 +505,7 @@ def test_format_fallback_metadata_sse_event(): fallback_model="fallback-model", fallback_errors=fallback_errors, error_state=None, + return_raw_model_name=False, ) assert isinstance(event, str) @@ -897,9 +898,11 @@ async def test_chat_stream_sends_fallback_metadata_in_place_of_a_stripped_usage_ @pytest.mark.asyncio @pytest.mark.parametrize("fallback_model", ["fallback-model", None]) +@pytest.mark.parametrize("return_raw_model_name", [False, True]) async def test_responses_stream_sends_fallback_metadata_as_a_typed_event_framed_like_its_neighbours( monkeypatch: pytest.MonkeyPatch, fallback_model: str | None, + return_raw_model_name: bool, ) -> None: _patch_logging_flags(monkeypatch) monkeypatch.setitem(ps.general_settings, "expose_fallback_errors_to_caller", True) @@ -936,7 +939,11 @@ async def test_responses_stream_sends_fallback_metadata_as_a_typed_event_framed_ async for frame in async_data_generator( response=upstream, user_api_key_dict=_user_auth(), - request_data={"model": "primary-model", "include_fallback_errors": True}, + request_data={ + "model": "primary-model", + "include_fallback_errors": True, + "metadata": {RETURN_RAW_MODEL_NAME_METADATA_KEY: return_raw_model_name}, + }, responses_stream_errors=True, ) ] @@ -952,16 +959,61 @@ async def test_responses_stream_sends_fallback_metadata_as_a_typed_event_framed_ "litellm.fallback_metadata", "response.completed", ) + stamped_model: Final = None if return_raw_model_name else fallback_model assert payloads[1] == { "type": "litellm.fallback_metadata", "id": "litellm-fallback-metadata", - **({"model": fallback_model} if fallback_model else {}), + **({"model": stamped_model} if stamped_model else {}), "litellm_fallback": {"fallback_model": fallback_model, "errors": fallback_errors}, } - assert {payload.get("model") for payload in payloads} == {fallback_model} + assert {payload.get("model") for payload in payloads} == {stamped_model} assert len({frame.startswith("event: ") for frame in event_frames}) == 1 +@pytest.mark.asyncio +async def test_responses_stream_names_the_fallback_event_after_a_raw_upstream_frame( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _patch_logging_flags(monkeypatch) + monkeypatch.setitem(ps.general_settings, "expose_fallback_errors_to_caller", True) + response: Final = ResponsesAPIResponse(id="resp_fallback", created_at=1, model="fallback-model", output=[]) + created_event: Final = {"type": "response.created", "sequence_number": 0, "response": response.model_dump(mode="json")} + upstream: Final = _FakeStream( + [ + f"event: response.created\ndata: {json.dumps(created_event)}\n\n".encode(), + ResponseCompletedEvent.model_validate( + {"type": "response.completed", "sequence_number": 1, "response": response} + ), + ], + hidden_params={ + "additional_headers": { + "x-litellm-attempted-fallbacks": 1, + "x-litellm-model-group": "fallback-model", + "x-litellm-fallback-errors": json.dumps([{"message": "upstream failed", "code": "500"}]), + } + }, + ) + + frames: Final = tuple( + [ + frame.decode() if isinstance(frame, bytes) else frame + async for frame in async_data_generator( + response=upstream, + user_api_key_dict=_user_auth(), + request_data={"model": "primary-model", "include_fallback_errors": True}, + responses_stream_errors=True, + ) + ] + ) + event_names: Final = tuple( + tuple(line.removeprefix("event: ") for line in frame.splitlines() if line.startswith("event: ")) + for frame in frames[1:3] + ) + + assert '"type": "litellm.fallback_metadata"' in frames[1] + assert event_names in (((), ()), (("litellm.fallback_metadata",), ("response.completed",))) + + @pytest.mark.asyncio async def test_async_data_generator_skips_fallback_error_event_without_opt_in( monkeypatch,