scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

fix(proxy): gate streaming cost injection on stream_options.include_usage

Browse source 
include_cost_in_streaming_usage is a process-wide flag, so every implementing
path injected usage.cost for every caller on every route, whether or not the
caller asked for usage. Turning it on also disabled the streaming fast path for
all traffic, including requests that never carry a usage dict.

Injection now consults the caller's stream_options.include_usage alongside the
global flag. An explicit include_usage: false opts out on any protocol. Anthropic
Messages, Vertex rawPredict and Gemini generateContent have no such field for a
caller to set, so injection stays always-on there, now documented rather than
accidental. The fast path is resolved per request, so streams that will never be
injected into keep it.

Fixes #38348
This commit is contained in:
Priyansh Nandwana 2026-08-27 16:29:18 +05:30
parent aea5358c48
commit 3a17303fdf
5 changed files with 291 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

64
litellm/proxy/common_request_processing.py
View file

@ -3585,6 +3585,7 @@ class ProxyBaseLLMRequestProcessing:
serialize_error: StreamErrorSerializer,
request: Request | None = None,
flush_tail: Callable[[], bytes] | None = None,
protocol_supports_stream_options: bool = True,
) -> AsyncGenerator[str, None]:
"""
Shared streaming data generator: runs proxy iterator hook, per-chunk hook,
@ -3594,6 +3595,9 @@ class ProxyBaseLLMRequestProcessing:
``flush_tail`` runs once after the upstream iterator completes cleanly and
its non-empty result is yielded, so a serializer that buffers bytes across
chunks can emit anything still held at end of stream.
``protocol_supports_stream_options`` says whether this route's protocol gives
callers a ``stream_options.include_usage`` to opt in with, which gates cost
injection; see ``_should_inject_cost_for_request``.
"""
verbose_proxy_logger.debug("inside generator")
# Resolve per-stream (not per-chunk) whether the heavy per-chunk path
@ -3604,7 +3608,10 @@ class ProxyBaseLLMRequestProcessing:
# await, response-string materialization, and cost-injection call are
# pure overhead on the streaming hot path (the default config).
caps: Final = ProxyLogging._callback_capabilities()
cost_injection_enabled: Final = bool(getattr(litellm, "include_cost_in_streaming_usage", False))
cost_injection_enabled: Final = ProxyBaseLLMRequestProcessing._should_inject_cost_for_request(
request_data,
protocol_supports_stream_options=protocol_supports_stream_options,
)
fast_path = not caps.has_streaming_chunk_override and not caps.has_guardrail and not cost_injection_enabled
debug_enabled: Final = verbose_proxy_logger.isEnabledFor(logging.DEBUG)
stream_completed = False
@ -3645,7 +3652,10 @@ class ProxyBaseLLMRequestProcessing:
model_name = request_data.get("model", "")
chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
chunk, model_name, request_data.get("litellm_logging_obj")
chunk,
model_name,
request_data.get("litellm_logging_obj"),
enabled=cost_injection_enabled,
)
# Set before the yield: an async generator suspends at the yield,
@ -3723,6 +3733,7 @@ class ProxyBaseLLMRequestProcessing:
proxy_logging_obj: ProxyLogging,
request: Request | None = None,
restamp_model: str | None = None,
protocol_supports_stream_options: bool = False,
) -> AsyncGenerator[str, None]:
"""
Anthropic /messages and Google /generateContent streaming data generator require SSE events.
@ -3748,23 +3759,59 @@ class ProxyBaseLLMRequestProcessing:
),
request=request,
flush_tail=None if restamper is None else restamper.flush,
protocol_supports_stream_options=protocol_supports_stream_options,
)
@staticmethod
def _should_inject_cost_for_request(
request_data: Mapping[str, Any] | None,
*,
protocol_supports_stream_options: bool = True,
) -> bool:
"""
Whether this request's streamed usage events should carry ``usage.cost``.
``litellm.include_cost_in_streaming_usage`` is process-wide, so on its own it
injects for every caller on every route. OpenAI-protocol callers opt into usage
reporting per request via ``stream_options.include_usage``, so that opt-in gates
injection too. An explicit ``include_usage: false`` opts out on any protocol.
Anthropic Messages, Vertex ``rawPredict`` and Gemini ``generateContent`` have no
such field for a caller to set, so injection stays always-on there.
"""
if not getattr(litellm, "include_cost_in_streaming_usage", False):
return False
stream_options: Final = request_data.get("stream_options") if isinstance(request_data, Mapping) else None
if isinstance(stream_options, Mapping):
return bool(stream_options.get("include_usage", False))
return not protocol_supports_stream_options
@overload
@staticmethod
def _process_chunk_with_cost_injection(
chunk: bytes, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None
chunk: bytes,
model_name: str,
litellm_logging_obj: LiteLLMLoggingObj | None = None,
*,
enabled: bool | None = None,
) -> bytes: ...
@overload
@staticmethod
def _process_chunk_with_cost_injection(
chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None
chunk: object,
model_name: str,
litellm_logging_obj: LiteLLMLoggingObj | None = None,
*,
enabled: bool | None = None,
) -> object: ...
@staticmethod
def _process_chunk_with_cost_injection(
chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None
chunk: object,
model_name: str,
litellm_logging_obj: LiteLLMLoggingObj | None = None,
*,
enabled: bool | None = None,
) -> object:
"""
Process a streaming chunk and inject cost information if enabled.
@ -3773,11 +3820,16 @@ class ProxyBaseLLMRequestProcessing:
chunk: The streaming chunk (dict, str, bytes, or bytearray)
model_name: Model name for cost calculation
litellm_logging_obj: The call's logging object, used for pricing
enabled: Per-stream decision from ``_should_inject_cost_for_request``.
Falls back to the global flag alone when not passed.
Returns:
The processed chunk with cost information injected if applicable
"""
if not getattr(litellm, "include_cost_in_streaming_usage", False):
injection_enabled: Final = (
enabled if enabled is not None else bool(getattr(litellm, "include_cost_in_streaming_usage", False))
)
if not injection_enabled:
return chunk
try:

2
litellm/proxy/pass_through_endpoints/architecture.md
View file

@ -62,7 +62,7 @@ sequenceDiagram
| Streaming chunk collection | Collect chunks async for logging after stream completes |
| Multipart form handling | Reconstruct multipart/form-data requests for file uploads |
| Guardrails (opt-in) | Run content filtering when explicitly configured |
| Cost injection | Inject cost into streaming chunks when `include_cost_in_streaming_usage` enabled |
| Cost injection | Inject cost into streaming chunks when `include_cost_in_streaming_usage` enabled. OpenAI-protocol callers must also opt in via `stream_options.include_usage`; Anthropic and Vertex `rawPredict` have no such field, so injection is always-on there unless the caller sends an explicit `stream_options.include_usage: false` |
## What Does NOT Change

15
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -6,7 +6,6 @@ from typing import Final, Protocol
import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
@ -162,12 +161,11 @@ class PassThroughStreamingHandler:
litellm_logging_obj=litellm_logging_obj,
)
# Resolve once per stream rather than re-reading the global +
# re-branching on every chunk. ``include_cost_in_streaming_usage`` is
# set at config load and stable for the process, matching how the
# proxy-level streaming fast path resolves it.
cost_injection_active: Final = (
bool(getattr(litellm, "include_cost_in_streaming_usage", False))
ProxyBaseLLMRequestProcessing._should_inject_cost_for_request(
request_body,
protocol_supports_stream_options=endpoint_type == EndpointType.OPENAI,
)
and bool(model_name)
and (
endpoint_type in (EndpointType.ANTHROPIC, EndpointType.OPENAI)
@ -199,7 +197,10 @@ class PassThroughStreamingHandler:
) # rebind-ok: SSE frame reassembly buffer across transport chunks
if complete_frames:
yield ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
complete_frames, resolved_model_name, litellm_logging_obj
complete_frames,
resolved_model_name,
litellm_logging_obj,
enabled=True,
)
if pending:
yield pending

112
tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py
View file

@ -392,12 +392,20 @@ def _openai_passthrough_stream_chunks():
]
async def _collect_openai_passthrough_chunks(chunks, endpoint_type):
async def _collect_openai_passthrough_chunks(chunks, endpoint_type, request_body=None):
# Default to the opt-in body a real caller must send for the OpenAI protocol
# to emit a usage frame at all -- cost injection is gated on that opt-in.
if request_body is None:
request_body = {
"model": "gpt-4o-mini",
"stream": True,
"stream_options": {"include_usage": True},
}
response = _make_streaming_response(chunks)
received = []
async for chunk in PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={"model": "gpt-4o-mini", "stream": True},
request_body=request_body,
litellm_logging_obj=_unarmed_logging_obj(),
endpoint_type=endpoint_type,
start_time=datetime.now(),
@ -475,6 +483,106 @@ async def test_chunk_processor_streams_crlf_delimited_frames_live_and_injects_co
assert final_payload["usage"]["cost"] > 0
@pytest.mark.asyncio
async def test_chunk_processor_skips_injection_when_openai_caller_did_not_opt_in(monkeypatch):
"""Regression: issue #38348 -- ``include_cost_in_streaming_usage`` is a process-wide
flag, but OpenAI-protocol callers opt into usage reporting per request via
``stream_options.include_usage``. A caller that never asked for usage must not have
``usage.cost`` injected into its stream just because the flag is on proxy-wide."""
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
chunks = _openai_passthrough_stream_chunks()
received = await _collect_openai_passthrough_chunks(
chunks,
EndpointType.OPENAI,
request_body={"model": "gpt-4o-mini", "stream": True},
)
assert received == chunks
@pytest.mark.asyncio
async def test_chunk_processor_respects_explicit_include_usage_false(monkeypatch):
"""Regression: issue #38348 -- an explicit ``include_usage: false`` is a caller
opting out, and must be honoured even with the global flag on."""
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
chunks = _openai_passthrough_stream_chunks()
received = await _collect_openai_passthrough_chunks(
chunks,
EndpointType.OPENAI,
request_body={
"model": "gpt-4o-mini",
"stream": True,
"stream_options": {"include_usage": False},
},
)
assert received == chunks
@pytest.mark.asyncio
async def test_chunk_processor_anthropic_injects_without_stream_options(monkeypatch):
"""The Anthropic Messages protocol has no ``stream_options`` for a caller to opt in
with, so injection stays always-on there while the flag is set -- issue #38348 asks
for that behaviour to be explicit rather than accidental."""
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
frame = (
b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},'
b'"usage":{"input_tokens":11,"output_tokens":4}}\n\n'
)
response = _make_streaming_response([frame])
received = []
async for chunk in PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={"model": "claude-haiku-4-5", "stream": True},
litellm_logging_obj=MagicMock(),
endpoint_type=EndpointType.ANTHROPIC,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
url_route="/v1/messages",
route_streaming_logging=AsyncMock(),
):
received.append(chunk)
await asyncio.sleep(0)
payload = json.loads(b"".join(received).decode("utf-8").split("data:", 1)[1].strip())
assert payload["usage"]["cost"] > 0
@pytest.mark.asyncio
async def test_chunk_processor_anthropic_respects_explicit_opt_out(monkeypatch):
"""Even on Anthropic, a caller that explicitly sends ``include_usage: false`` opts
out of cost injection."""
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
frame = (
b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},'
b'"usage":{"input_tokens":11,"output_tokens":4}}\n\n'
)
response = _make_streaming_response([frame])
received = []
async for chunk in PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={
"model": "claude-haiku-4-5",
"stream": True,
"stream_options": {"include_usage": False},
},
litellm_logging_obj=MagicMock(),
endpoint_type=EndpointType.ANTHROPIC,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
url_route="/v1/messages",
route_streaming_logging=AsyncMock(),
):
received.append(chunk)
await asyncio.sleep(0)
assert received == [frame]
@pytest.mark.asyncio
async def test_chunk_processor_flag_off_leaves_openai_passthrough_stream_byte_identical(monkeypatch):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)

114
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -7945,3 +7945,117 @@ class TestDetachedStreamFailureHook:
await logging_obj._on_detached_stream_failure(failure)
assert [call["original_exception"] for call in recorder.calls] == [failure]
class TestShouldInjectCostForRequest:
"""Issue #38348: ``include_cost_in_streaming_usage`` is a process-wide flag, so on its
own it injects ``usage.cost`` for every caller on every route. Injection must also
consult the caller's per-request ``stream_options.include_usage`` opt-in."""
def test_global_flag_off_never_injects(self, monkeypatch):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
assert (
ProxyBaseLLMRequestProcessing._should_inject_cost_for_request(
{"stream_options": {"include_usage": True}}
)
is False
)
def test_openai_protocol_requires_caller_opt_in(self, monkeypatch):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
assert (
ProxyBaseLLMRequestProcessing._should_inject_cost_for_request(
{"model": "gpt-4o-mini", "stream": True},
protocol_supports_stream_options=True,
)
is False
)
def test_openai_protocol_opted_in_injects(self, monkeypatch):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
assert (
ProxyBaseLLMRequestProcessing._should_inject_cost_for_request(
{"stream_options": {"include_usage": True}},
protocol_supports_stream_options=True,
)
is True
)
@pytest.mark.parametrize("protocol_supports_stream_options", [True, False])
def test_explicit_opt_out_is_honoured_on_every_protocol(self, monkeypatch, protocol_supports_stream_options):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
assert (
ProxyBaseLLMRequestProcessing._should_inject_cost_for_request(
{"stream_options": {"include_usage": False}},
protocol_supports_stream_options=protocol_supports_stream_options,
)
is False
)
def test_protocol_without_stream_options_stays_always_on(self, monkeypatch):
"""Anthropic Messages / Vertex rawPredict / Gemini give a caller no way to opt
in, so the flag remains always-on there."""
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
assert (
ProxyBaseLLMRequestProcessing._should_inject_cost_for_request(
{"model": "claude-haiku-4-5", "stream": True},
protocol_supports_stream_options=False,
)
is True
)
def test_missing_request_data_falls_back_to_protocol_default(self, monkeypatch):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
assert ProxyBaseLLMRequestProcessing._should_inject_cost_for_request(None) is False
assert (
ProxyBaseLLMRequestProcessing._should_inject_cost_for_request(
None, protocol_supports_stream_options=False
)
is True
)
def test_malformed_stream_options_falls_back_to_protocol_default(self, monkeypatch):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
assert (
ProxyBaseLLMRequestProcessing._should_inject_cost_for_request(
{"stream_options": "include_usage"},
protocol_supports_stream_options=False,
)
is True
)
class TestProcessChunkCostInjectionGate:
"""``_process_chunk_with_cost_injection`` takes the per-stream decision from
``_should_inject_cost_for_request`` and falls back to the global flag when the
caller does not pass one."""
@staticmethod
def _usage_chunk():
return {
"id": "chatcmpl-1",
"object": "chat.completion.chunk",
"choices": [],
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
}
def test_enabled_false_leaves_chunk_untouched(self, monkeypatch):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
chunk = self._usage_chunk()
assert ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
chunk, "gpt-4o-mini", None, enabled=False
) is chunk
def test_enabled_true_injects_even_with_global_flag_off(self, monkeypatch):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
result = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
self._usage_chunk(), "gpt-4o-mini", None, enabled=True
)
assert result["usage"]["cost"] > 0
def test_omitted_enabled_falls_back_to_global_flag(self, monkeypatch):
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
result = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
self._usage_chunk(), "gpt-4o-mini"
)
assert result["usage"]["cost"] > 0