diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f5a50d1697a4..1ca14ea36840 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -5,6 +5,7 @@ import time import traceback from datetime import datetime +from functools import lru_cache from typing import ( TYPE_CHECKING, Any, @@ -12,6 +13,7 @@ Callable, Dict, Literal, + Mapping, Optional, Tuple, Union, @@ -38,6 +40,9 @@ ) from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer +from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, @@ -244,6 +249,71 @@ async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[Any]"]) -> None pass +@lru_cache(maxsize=512) +def _litellm_model_supports_stream_options(litellm_model: str) -> bool: + try: + supported_params = get_supported_openai_params(model=litellm_model) + except Exception: # noqa: BLE001 # unmapped or malformed model strings must disable injection, not fail the request + return False + return supported_params is not None and "stream_options" in supported_params + + +def _deployment_litellm_model(deployment: Mapping[str, object]) -> str | None: + litellm_params = deployment.get("litellm_params") + if isinstance(litellm_params, Mapping): + litellm_model = litellm_params.get("model") + else: + litellm_model = getattr(litellm_params, "model", None) + return litellm_model if isinstance(litellm_model, str) else None + + +def _model_deployments_support_stream_options( + model: object, + llm_router: Router | None, + team_id: str | None, +) -> bool: + if not isinstance(model, str): + return False + deployments = llm_router.get_model_list(model_name=model, team_id=team_id) if llm_router is not None else None + deployment_models = tuple( + litellm_model + for deployment in deployments or () + if (litellm_model := _deployment_litellm_model(deployment)) is not None + ) + candidate_models = deployment_models if deployment_models else (model,) + return all(_litellm_model_supports_stream_options(m) for m in candidate_models) + + +def _stream_usage_tracking_updates( + data: Mapping[str, object], + general_settings: Mapping[str, object], + route_type: str, + supports_stream_options: Callable[[], bool], +) -> Mapping[str, object]: + scrub = {"_litellm_strip_stream_usage": False} if "_litellm_strip_stream_usage" in data else {} + if data.get("stream", False) is not True: + return scrub + always_include = general_settings.get("always_include_stream_usage") + stream_options = data.get("stream_options") + if always_include is True: + if "stream_options" not in data: + return {**scrub, "stream_options": {"include_usage": True}} + if isinstance(stream_options, dict) and "include_usage" not in stream_options: + return {**scrub, "stream_options": {**stream_options, "include_usage": True}} + return scrub + if always_include is False or route_type != "acompletion": + return scrub + if isinstance(stream_options, dict) and stream_options.get("include_usage") is True: + return scrub + if not supports_stream_options(): + return scrub + merged_stream_options = {**stream_options} if isinstance(stream_options, dict) else {} + return { + "stream_options": {**merged_stream_options, "include_usage": True}, + "_litellm_strip_stream_usage": True, + } + + def _serialize_http_exception_detail( detail: Any, ) -> Tuple[str, Optional[dict]]: @@ -1232,17 +1302,18 @@ async def common_processing_pre_call_logic( ) ### AUTO STREAM USAGE TRACKING ### - # If always_include_stream_usage is enabled and this is a streaming request - # automatically add stream_options={'include_usage': True} if not already set - if ( - general_settings.get("always_include_stream_usage", False) is True - and self.data.get("stream", False) is True - ): - # Only set if stream_options is not already provided by the client - if "stream_options" not in self.data: - self.data["stream_options"] = {"include_usage": True} - elif isinstance(self.data["stream_options"], dict) and "include_usage" not in self.data["stream_options"]: - self.data["stream_options"]["include_usage"] = True + self.data.update( + _stream_usage_tracking_updates( + data=self.data, + general_settings=general_settings, + route_type=route_type, + supports_stream_options=lambda: _model_deployments_support_stream_options( + model=self.data.get("model"), + llm_router=llm_router, + team_id=user_api_key_dict.team_id, + ), + ) + ) ### CALL HOOKS ### - modify/reject incoming data before calling the model ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c72e3d4ee5b1..a60ea2da0192 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -119,6 +119,7 @@ from litellm.types.utils import ( ModelResponse, ModelResponseStream, + StreamingChoices, TextCompletionResponse, TokenCountResponse, ) @@ -7368,6 +7369,25 @@ def _serialize_streaming_chunk(chunk: BaseModel) -> Union[str, bytes]: return chunk.model_dump_json(exclude_none=True, exclude_unset=True) +def _is_injected_stream_usage_artifact(chunk: object) -> bool: + if not isinstance(chunk, ModelResponseStream): + return False + if chunk.provider_specific_fields is not None: + return False + return all(_is_empty_streaming_choice(choice) for choice in chunk.choices or []) + + +def _is_empty_streaming_choice(choice: StreamingChoices) -> bool: + if choice.finish_reason is not None: + return False + if getattr(choice, "logprobs", None) is not None: + return False + delta = getattr(choice, "delta", None) + if delta is None: + return True + return all(value is None for value in delta.model_dump().values()) + + async def _apply_streaming_chunk_hooks( *, chunk: Any, @@ -7447,6 +7467,7 @@ async def async_data_generator( needs_iterator_wrap = proxy_logging_obj.needs_iterator_wrap() needs_per_chunk_hook = proxy_logging_obj.needs_per_chunk_streaming_hook() is_raw_sse_stream = bool(request_data.get("_litellm_raw_sse_stream")) + strip_stream_usage = bool(request_data.get("_litellm_strip_stream_usage")) raw_sse_buffer = "" if needs_iterator_wrap: @@ -7498,6 +7519,15 @@ async def async_data_generator( fallback_model_from_metadata=fallback_model_from_metadata, ) + if strip_stream_usage and _is_injected_stream_usage_artifact(chunk): + if pending_fallback_event: + yield _format_fallback_metadata_sse_event( + fallback_model=fallback_model_from_metadata, + fallback_errors=fallback_errors, + ) + fallback_metadata_event_sent = True + continue + raw_passthrough = False if isinstance(chunk, BaseModel): chunk = _serialize_streaming_chunk(chunk) @@ -13470,6 +13500,7 @@ async def async_queue_request( data = {} try: data = await request.json() # type: ignore + data.pop("_litellm_strip_stream_usage", None) # Include original request and headers in the data data["proxy_server_request"] = { diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 9df44c6202ce..5ba994b921c9 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3296,6 +3296,7 @@ def shared_backend_model_info(model_info: Dict[str, Any]) -> Dict[str, Any]: "model_file_id_mapping", "litellm_logging_obj", "litellm_call_id", + "_litellm_strip_stream_usage", "use_client", "id", "fallbacks", diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 58f81cdad35f..3bb84e095a04 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1,7 +1,7 @@ import asyncio import copy import datetime -from typing import AsyncGenerator, Optional +from typing import AsyncGenerator, Callable, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -5111,3 +5111,246 @@ async def test_disconnect_without_billable_chunks_releases_slot(self): ) proxy_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + + +def _apply_stream_usage_tracking( + data: dict, + general_settings: dict, + route_type: str, + supports_stream_options: Callable[[], bool] = lambda: True, +) -> None: + from litellm.proxy.common_request_processing import _stream_usage_tracking_updates + + data.update( + _stream_usage_tracking_updates( + data=data, + general_settings=general_settings, + route_type=route_type, + supports_stream_options=supports_stream_options, + ) + ) + + +class TestApplyStreamUsageTracking: + def test_default_injects_usage_and_marks_strip_for_chat_completions(self): + data = {"stream": True, "model": "gpt-5.4-nano"} + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion") + + assert data["stream_options"] == {"include_usage": True} + assert data["_litellm_strip_stream_usage"] is True + + def test_default_preserves_other_client_stream_options_keys(self): + data = {"stream": True, "stream_options": {"include_obfuscation": True}} + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion") + + assert data["stream_options"] == {"include_obfuscation": True, "include_usage": True} + assert data["_litellm_strip_stream_usage"] is True + + def test_client_requested_usage_is_left_untouched_and_not_stripped(self): + data = {"stream": True, "stream_options": {"include_usage": True}} + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion") + + assert data["stream_options"] == {"include_usage": True} + assert "_litellm_strip_stream_usage" not in data + + def test_client_include_usage_false_is_overridden_and_stripped(self): + data = {"stream": True, "stream_options": {"include_usage": False}} + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion") + + assert data["stream_options"]["include_usage"] is True + assert data["_litellm_strip_stream_usage"] is True + + def test_explicit_false_flag_disables_injection_entirely(self): + data = {"stream": True} + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": False}, + route_type="acompletion", + ) + + assert "stream_options" not in data + assert "_litellm_strip_stream_usage" not in data + + def test_flag_true_injects_without_strip_marker(self): + data = {"stream": True} + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["stream_options"] == {"include_usage": True} + assert "_litellm_strip_stream_usage" not in data + + def test_flag_true_respects_client_explicit_include_usage_false(self): + data = {"stream": True, "stream_options": {"include_usage": False}} + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["stream_options"] == {"include_usage": False} + assert "_litellm_strip_stream_usage" not in data + + def test_default_does_not_touch_non_chat_completion_routes(self): + data = {"stream": True} + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="anthropic_messages") + + assert "stream_options" not in data + assert "_litellm_strip_stream_usage" not in data + + def test_non_streaming_request_is_untouched(self): + data = {"model": "gpt-5.4-nano"} + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion") + + assert "stream_options" not in data + assert "_litellm_strip_stream_usage" not in data + + def test_default_skips_injection_when_provider_lacks_stream_options_support(self): + data = {"stream": True, "model": "bytez-model"} + + _apply_stream_usage_tracking( + data=data, + general_settings={}, + route_type="acompletion", + supports_stream_options=lambda: False, + ) + + assert "stream_options" not in data + assert "_litellm_strip_stream_usage" not in data + + def test_client_supplied_strip_marker_is_neutralized(self): + data = { + "stream": True, + "stream_options": {"include_usage": True}, + "_litellm_strip_stream_usage": True, + } + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion") + + assert data["_litellm_strip_stream_usage"] is False + assert data["stream_options"] == {"include_usage": True} + + def test_client_supplied_strip_marker_is_neutralized_with_flag_true(self): + data = { + "stream": True, + "stream_options": {"include_usage": True}, + "_litellm_strip_stream_usage": True, + } + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["_litellm_strip_stream_usage"] is False + + def test_client_supplied_strip_marker_is_neutralized_on_non_streaming_request(self): + data = {"_litellm_strip_stream_usage": True} + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion") + + assert data["_litellm_strip_stream_usage"] is False + + +class TestModelDeploymentsSupportStreamOptions: + def _support(self, model, llm_router=None, team_id=None) -> bool: + from litellm.proxy.common_request_processing import ( + _model_deployments_support_stream_options, + ) + + return _model_deployments_support_stream_options(model=model, llm_router=llm_router, team_id=team_id) + + def test_openai_compatible_deployment_supports_stream_options(self): + router = litellm.Router( + model_list=[ + { + "model_name": "azure-nano", + "litellm_params": { + "model": "azure/gpt-5.4-nano", + "api_key": "fake", + "api_base": "https://example.openai.azure.com", + }, + } + ] + ) + + assert self._support("azure-nano", router) is True + + def test_deployment_on_provider_rejecting_stream_options_is_not_injected(self): + router = litellm.Router( + model_list=[ + { + "model_name": "tiny", + "litellm_params": {"model": "bytez/openai-community/gpt2", "api_key": "fake"}, + } + ] + ) + + assert self._support("tiny", router) is False + + def test_mixed_provider_model_group_is_not_injected(self): + router = litellm.Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "oci/cohere.command-r-plus", "api_key": "fake"}, + }, + ] + ) + + assert self._support("mixed", router) is False + + def test_wildcard_route_resolves_provider_support(self): + router = litellm.Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "fake"}, + } + ] + ) + + assert self._support("openai/gpt-4o", router) is True + + def test_provider_prefixed_model_without_router_is_resolved_directly(self): + assert self._support("openai/gpt-4o", None) is True + assert self._support("bytez/openai-community/gpt2", None) is False + + def test_unmapped_model_name_is_not_injected(self): + assert self._support("some-unmapped-public-alias", None) is False + + def test_team_alias_model_resolves_with_team_id(self): + router = litellm.Router( + model_list=[ + { + "model_name": "model_name_team-1_8b6a0b3f", + "litellm_params": {"model": "azure/gpt-5.4-nano", "api_key": "fake"}, + "model_info": { + "team_id": "team-1", + "team_public_model_name": "team-gpt", + }, + } + ] + ) + + assert self._support("team-gpt", router, team_id="team-1") is True + assert self._support("team-gpt", router, team_id=None) is False + + def test_non_string_model_is_not_injected(self): + assert self._support(None, None) is False diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 5646d202e318..b9a33bd2cefa 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -10572,3 +10572,109 @@ async def test_startup_survives_database_read_failure_for_coordination_redis(): ) assert result is None + + +def _stream_usage_test_chunks(): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage + + content_chunk = ModelResponseStream( + model="gpt-5.4-nano", + choices=[StreamingChoices(delta=Delta(content="pong"))], + ) + finish_chunk = ModelResponseStream( + model="gpt-5.4-nano", + choices=[StreamingChoices(finish_reason="stop")], + ) + usage_chunk = ModelResponseStream(model="gpt-5.4-nano", choices=[]) + usage_chunk.usage = Usage(prompt_tokens=50, completion_tokens=188, total_tokens=238) + return content_chunk, finish_chunk, usage_chunk + + +def _stream_usage_generator_chunks(): + from litellm.types.utils import ModelResponseStream + + content_chunk, finish_chunk, usage_chunk = _stream_usage_test_chunks() + prompt_filter_chunk = ModelResponseStream(model="gpt-5.4-nano", choices=[]) + return prompt_filter_chunk, content_chunk, finish_chunk, usage_chunk + + +def test_is_injected_stream_usage_artifact(): + from litellm.proxy.proxy_server import _is_injected_stream_usage_artifact + from litellm.types.utils import ModelResponseStream, Usage + + content_chunk, finish_chunk, empty_choices_usage_chunk = _stream_usage_test_chunks() + assert _is_injected_stream_usage_artifact(empty_choices_usage_chunk) is True + + synthetic_final_chunk = ModelResponseStream(model="gpt-5.4-nano") + synthetic_final_chunk.usage = Usage(prompt_tokens=50, completion_tokens=188, total_tokens=238) + assert _is_injected_stream_usage_artifact(synthetic_final_chunk) is True + + azure_prompt_filter_chunk = ModelResponseStream(model="gpt-5.4-nano", choices=[]) + assert _is_injected_stream_usage_artifact(azure_prompt_filter_chunk) is True + + assert _is_injected_stream_usage_artifact(content_chunk) is False + assert _is_injected_stream_usage_artifact(finish_chunk) is False + + content_chunk_with_usage, finish_chunk_with_usage, _ = _stream_usage_test_chunks() + content_chunk_with_usage.usage = Usage(prompt_tokens=50, completion_tokens=188, total_tokens=238) + finish_chunk_with_usage.usage = Usage(prompt_tokens=50, completion_tokens=188, total_tokens=238) + assert _is_injected_stream_usage_artifact(content_chunk_with_usage) is False + assert _is_injected_stream_usage_artifact(finish_chunk_with_usage) is False + + assert _is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False + + +async def _collect_async_data_generator_frames(request_data: dict) -> list: + from litellm.proxy.proxy_server import async_data_generator + from litellm.proxy.utils import ProxyLogging + + chunks = _stream_usage_generator_chunks() + + class MockStream: + def __aiter__(self): + return self._stream() + + async def _stream(self): + for chunk in chunks: + yield chunk + + async def aclose(self): + pass + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.needs_iterator_wrap.return_value = False + mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock() + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): + with patch.object(proxy_server_module.ProxyLogging, "_fire_deferred_stream_logging"): + return [ + frame.decode("utf-8") if isinstance(frame, bytes) else frame + async for frame in async_data_generator( + MockStream(), MagicMock(spec=UserAPIKeyAuth), request_data + ) + ] + + +@pytest.mark.asyncio +async def test_async_data_generator_strips_injected_usage_chunk(): + frames = await _collect_async_data_generator_frames( + {"model": "gpt-5.4-nano", "_litellm_strip_stream_usage": True} + ) + + data_frames = [frame for frame in frames if frame.startswith("data: {")] + assert len(data_frames) == 2 + assert any("pong" in frame for frame in data_frames) + assert any("finish_reason" in frame for frame in data_frames) + assert not any('"usage"' in frame for frame in data_frames) + assert frames[-1] == "data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_async_data_generator_forwards_usage_chunk_without_strip_marker(): + frames = await _collect_async_data_generator_frames({"model": "gpt-5.4-nano"}) + + data_frames = [frame for frame in frames if frame.startswith("data: {")] + assert len(data_frames) == 4 + assert any('"usage"' in frame and '"completion_tokens":188' in frame.replace(" ", "") for frame in data_frames) + assert frames[-1] == "data: [DONE]\n\n"