scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

refactor: extract stream_options logic to static method for testability

Browse source 
Addresses greptile review concern: tests were exercising a hand-rolled
copy of the guard logic instead of the real production code. Extracts
_apply_stream_options_for_usage as a @staticmethod on
ProxyBaseLLMRequestProcessing so tests import and call the real method.
This commit is contained in:
Terrajlz 2026-05-22 13:56:19 -04:00
parent 5fd2e8941d
commit 0b8333f070
2 changed files with 41 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

52
litellm/proxy/common_request_processing.py
View file

@ -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:
"""

36
tests/test_litellm/test_responses_api_stream_options.py
View file

@ -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