This commit is contained in:
Priyansh Nandwana 2026-09-05 06:55:27 +00:00 committed by GitHub
commit 6ab2804f93
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 300 additions and 22 deletions

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:

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

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

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)

View file

@ -73,11 +73,12 @@ def spend_counter_state():
ps.prisma_client = original_prisma_client
def _request_body() -> dict:
def _request_body(*, include_usage: bool = False) -> dict:
return {
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 10,
**({"stream_options": {"include_usage": True}} if include_usage else {}),
}
@ -2717,14 +2718,15 @@ async def test_streaming_cancel_in_slow_path_before_yield_refunds(spend_counter_
generator = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
response=MagicMock(),
user_api_key_dict=valid_token,
request_data=_request_body(),
request_data=_request_body(include_usage=True),
proxy_logging_obj=streaming_logging_obj,
serialize_chunk=lambda chunk: chunk,
serialize_error=lambda exc: str(exc),
)
received = []
# include_cost_in_streaming_usage forces fast_path off, so the hook above runs
# include_cost_in_streaming_usage plus the caller's stream_options.include_usage
# opt-in forces fast_path off, so the hook above runs
with patch.object(litellm, "include_cost_in_streaming_usage", True, create=True):
async def _drain():
async for chunk in generator:
@ -2793,15 +2795,16 @@ async def test_streaming_slow_path_processes_and_yields_chunk(spend_counter_stat
generator = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
response=MagicMock(),
user_api_key_dict=valid_token,
request_data=_request_body(),
request_data=_request_body(include_usage=True),
proxy_logging_obj=streaming_logging_obj,
serialize_chunk=lambda chunk: chunk,
serialize_error=lambda exc: str(exc),
)
received = []
# include_cost_in_streaming_usage forces the slow path so the per-chunk hook,
# content accumulation, and cost-injection branch all run to a successful yield
# include_cost_in_streaming_usage plus the caller's stream_options.include_usage
# opt-in forces the slow path, so the per-chunk hook, content accumulation, and
# cost-injection branch all run to a successful yield
with patch.object(litellm, "include_cost_in_streaming_usage", True, create=True):
async for chunk in generator:
received.append(chunk)

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