diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9296e9635cd..c797b51b70b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -918,28 +918,9 @@ class ProxyBaseLLMRequestProcessing: ) ### 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 - # NOTE: Only apply to chat completions, NOT Responses API routes. - # Azure/OpenAI Responses API does not support stream_options (usage is - # included automatically in response.completed events). - _is_responses_api_route = route_type in { - "aresponses", - "_aresponses_websocket", - } - if ( - general_settings.get("always_include_stream_usage", False) is True - and self.data.get("stream", False) is True - and not _is_responses_api_route - ): - # 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 + ProxyBaseLLMRequestProcessing._apply_stream_options_for_usage( + self.data, general_settings, route_type + ) ### CALL HOOKS ### - modify/reject incoming data before calling the model ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call @@ -1569,6 +1550,33 @@ class ProxyBaseLLMRequestProcessing: return True return False + @staticmethod + def _apply_stream_options_for_usage( + data: dict, general_settings: dict, route_type: str + ) -> None: + """Inject stream_options={'include_usage': True} for streaming requests. + + Skipped for Responses API routes (aresponses, _aresponses_websocket) + because Azure/OpenAI Responses API does not support stream_options — + usage is included automatically in response.completed events. + """ + _is_responses_api_route = route_type in { + "aresponses", + "_aresponses_websocket", + } + if ( + general_settings.get("always_include_stream_usage", False) is True + and data.get("stream", False) is True + and not _is_responses_api_route + ): + if "stream_options" not in data: + data["stream_options"] = {"include_usage": True} + elif ( + isinstance(data["stream_options"], dict) + and "include_usage" not in data["stream_options"] + ): + data["stream_options"]["include_usage"] = True + @staticmethod def _has_post_call_guardrails() -> bool: """ diff --git a/tests/test_litellm/test_responses_api_stream_options.py b/tests/test_litellm/test_responses_api_stream_options.py index 5351fdfdf38..37ea5dba15e 100644 --- a/tests/test_litellm/test_responses_api_stream_options.py +++ b/tests/test_litellm/test_responses_api_stream_options.py @@ -9,25 +9,7 @@ when `always_include_stream_usage` is enabled. This must NOT happen for Response import pytest - -def _apply_stream_options_logic(data: dict, general_settings: dict, route_type: str): - """Reproduces the stream_options injection logic from common_processing_pre_call_logic.""" - _is_responses_api_route = route_type in { - "aresponses", - "_aresponses_websocket", - } - if ( - general_settings.get("always_include_stream_usage", False) is True - and data.get("stream", False) is True - and not _is_responses_api_route - ): - if "stream_options" not in data: - data["stream_options"] = {"include_usage": True} - elif ( - isinstance(data["stream_options"], dict) - and "include_usage" not in data["stream_options"] - ): - data["stream_options"]["include_usage"] = True +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing class TestStreamOptionsNotInjectedForResponsesAPI: @@ -37,7 +19,7 @@ class TestStreamOptionsNotInjectedForResponsesAPI: def test_stream_options_not_injected_for_responses_routes(self, route_type): """stream_options must NOT be added when route is a Responses API route.""" data = {"stream": True, "model": "gpt-4"} - _apply_stream_options_logic( + ProxyBaseLLMRequestProcessing._apply_stream_options_for_usage( data, {"always_include_stream_usage": True}, route_type ) assert "stream_options" not in data @@ -45,7 +27,7 @@ class TestStreamOptionsNotInjectedForResponsesAPI: def test_stream_options_injected_for_chat_completions(self): """stream_options SHOULD be added for acompletion route.""" data = {"stream": True, "model": "gpt-4"} - _apply_stream_options_logic( + ProxyBaseLLMRequestProcessing._apply_stream_options_for_usage( data, {"always_include_stream_usage": True}, "acompletion" ) assert data["stream_options"] == {"include_usage": True} @@ -53,15 +35,19 @@ class TestStreamOptionsNotInjectedForResponsesAPI: def test_stream_options_not_injected_when_disabled(self): """stream_options should NOT be added when always_include_stream_usage is False.""" data = {"stream": True, "model": "gpt-4"} - _apply_stream_options_logic( + ProxyBaseLLMRequestProcessing._apply_stream_options_for_usage( data, {"always_include_stream_usage": False}, "acompletion" ) assert "stream_options" not in data def test_existing_stream_options_not_overwritten(self): """If client already set stream_options with include_usage, don't overwrite.""" - data = {"stream": True, "model": "gpt-4", "stream_options": {"include_usage": False}} - _apply_stream_options_logic( + data = { + "stream": True, + "model": "gpt-4", + "stream_options": {"include_usage": False}, + } + ProxyBaseLLMRequestProcessing._apply_stream_options_for_usage( data, {"always_include_stream_usage": True}, "acompletion" ) assert data["stream_options"] == {"include_usage": False} @@ -69,7 +55,7 @@ class TestStreamOptionsNotInjectedForResponsesAPI: def test_non_streaming_request_skipped(self): """stream_options should NOT be added for non-streaming requests.""" data = {"stream": False, "model": "gpt-4"} - _apply_stream_options_logic( + ProxyBaseLLMRequestProcessing._apply_stream_options_for_usage( data, {"always_include_stream_usage": True}, "acompletion" ) assert "stream_options" not in data