diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f25fa46197e..464d16b752d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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: diff --git a/litellm/proxy/pass_through_endpoints/architecture.md b/litellm/proxy/pass_through_endpoints/architecture.md index f7dd8077ab5..c47f6f9054e 100644 --- a/litellm/proxy/pass_through_endpoints/architecture.md +++ b/litellm/proxy/pass_through_endpoints/architecture.md @@ -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 diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 4be0235adbb..0b8febde728 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -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 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index ea6adc35b9a..9fe3b97e9eb 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -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) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index b8fb6170d34..fb379fd541c 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -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) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index f7fe6ad9d39..ee8749a7c9b 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -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