diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 97de2488b8d..22cd6e13298 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3868,6 +3868,7 @@ class ProxyBaseLLMRequestProcessing: request: Request | None = None, flush_tail: Callable[[], bytes] | None = None, seal_open_frame: Callable[[bytes], str] | None = None, + protocol_supports_stream_options: bool = True, ) -> AsyncGenerator[str, None]: """ Shared streaming data generator: runs proxy iterator hook, per-chunk hook, @@ -3883,6 +3884,10 @@ class ProxyBaseLLMRequestProcessing: relays raw upstream bytes, so an upstream that hangs up mid-frame leaves the client inside an open frame, where an error frame would be swallowed or misparsed instead of raised. + + ``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 @@ -3893,7 +3898,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 @@ -3935,7 +3943,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, @@ -4017,6 +4028,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 streaming data generator, which requires SSE events. @@ -4044,23 +4056,59 @@ class ProxyBaseLLMRequestProcessing: request=request, flush_tail=None if restamper is None else restamper.flush, seal_open_frame=seal_open_sse_frame, + 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. @@ -4069,11 +4117,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 19d8b063dd7..9d4a5a276ff 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.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks @@ -172,12 +171,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) @@ -220,7 +218,10 @@ class PassThroughStreamingHandler: complete_frames, pending = split_complete_sse_frames(pending + chunk) 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 00a3606bbf5..c347d438a4a 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 @@ -398,12 +398,16 @@ def _openai_passthrough_stream_chunks(): ] -async def _collect_openai_passthrough_chunks(chunks, endpoint_type): +def _openai_opted_in_body(): + return {"model": "gpt-4o-mini", "stream": True, "stream_options": {"include_usage": True}} + + +async def _collect_openai_passthrough_chunks(chunks, endpoint_type, request_body=None): 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=_openai_opted_in_body() if request_body is None else request_body, litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=endpoint_type, start_time=datetime.now(), @@ -481,6 +485,95 @@ 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): + 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): + 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): + 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): + 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 18b046cd83c..a1c84283cc7 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -78,11 +78,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 {}), } @@ -3115,14 +3116,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: @@ -3191,15 +3193,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 c17f41a8b8f..047fa297b5c 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -10226,3 +10226,109 @@ class TestStreamingContainerOwnershipRecordedBeforeDone: assert tuple(chunk for chunk, _ in observed) == self.CHUNKS assert tuple(count for _, count in observed) == (0, 0, 0, 0) recorder.assert_awaited_once() + + +class TestShouldInjectCostForRequest: + """Issue #38348: cost injection honours the caller's stream_options.include_usage.""" + + 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): + 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: + @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