diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index f31d4205b82..800973573c9 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -113,6 +113,27 @@ _Key = TypeVar("_Key") # rebind-ok: a TypeVar cannot be Final _Value = TypeVar("_Value") # rebind-ok: a TypeVar cannot be Final +INCLUDE_COST_IN_STREAMING_USAGE_HEADER: Final = "x-litellm-include-cost-in-streaming-usage" + + +def proxy_request_headers(container: Mapping[str, object] | None) -> Mapping[str, object] | None: + proxy_server_request: Final = as_str_mapping(container.get("proxy_server_request")) if container else None + return as_str_mapping(proxy_server_request.get("headers")) if proxy_server_request else None + + +def stream_cost_injection_requested(request_headers: Mapping[str, object] | None) -> bool: + """The global include_cost_in_streaming_usage flag or the per-request header asks for usage.cost on streams.""" + if bool(getattr(litellm, "include_cost_in_streaming_usage", False)): + return True + if not request_headers: + return False + header_value: Final = next( + (value for name, value in request_headers.items() if name.lower() == INCLUDE_COST_IN_STREAMING_USAGE_HEADER), + None, + ) + return isinstance(header_value, str) and header_value.strip().lower() == "true" + + def without( mapping: Mapping[_Key, _Value], *keys: object ) -> dict[_Key, _Value]: # mutable-ok: a fresh copy the caller owns diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 55f2ee94d0e..887b8f98c2b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -70,8 +70,10 @@ from litellm.litellm_core_utils.core_helpers import ( has_attributes, independent_snapshot, is_expected_client_error, + proxy_request_headers, redact_nested_match_and_regex_keys, require_attributes, + stream_cost_injection_requested, ) from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer from litellm.litellm_core_utils.get_supported_openai_params import ( @@ -4815,7 +4817,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)) + request_headers: Final = proxy_request_headers( + request_data # pyright: ignore[reportUnknownArgumentType] # request_data is the untyped proxy request dict + ) + cost_injection_enabled: Final = stream_cost_injection_requested(request_headers) 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 @@ -4862,7 +4867,7 @@ 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"), request_headers=request_headers ) # Set before the yield: an async generator suspends at the yield, @@ -4991,18 +4996,30 @@ class ProxyBaseLLMRequestProcessing: @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, + *, + request_headers: Mapping[str, object] | 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, + *, + request_headers: Mapping[str, object] | 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, + *, + request_headers: Mapping[str, object] | None = None, ) -> object: """ Process a streaming chunk and inject cost information if enabled. @@ -5011,11 +5028,12 @@ 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 + request_headers: The proxy request headers, read for the per-request opt-in header Returns: The processed chunk with cost information injected if applicable """ - if not getattr(litellm, "include_cost_in_streaming_usage", False): + if not stream_cost_injection_requested(request_headers): return chunk try: diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index c94dc16a8f5..ef8ffdf640e 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -6,10 +6,14 @@ 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 +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + bind_budget_reservation_to_callbacks, + proxy_request_headers, + stream_cost_injection_requested, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception @@ -194,12 +198,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. + request_headers: Final = proxy_request_headers( + as_str_mapping(litellm_logging_obj.model_call_details.get("litellm_params")) + ) cost_injection_active: Final = ( - bool(getattr(litellm, "include_cost_in_streaming_usage", False)) + stream_cost_injection_requested(request_headers) and bool(model_name) and ( endpoint_type in (EndpointType.ANTHROPIC, EndpointType.OPENAI) @@ -242,7 +245,7 @@ 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, request_headers=request_headers ) if pending: yield pending diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 9adfcb60d74..78040ba7225 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -95,7 +95,7 @@ class Feature(str, Enum): CACHE_CONTROL_INJECTION = "cache_control_injection" """Prompt caching injection that Python applies to messages and system.""" STREAM_COST_INJECTION = "stream_cost_injection" - """litellm.include_cost_in_streaming_usage, where Python adds cost to streamed usage.""" + """litellm.include_cost_in_streaming_usage or the x-litellm-include-cost-in-streaming-usage header, where Python adds cost to streamed usage.""" THINKING_SIGNATURE_RETRY = "thinking_signature_retry" """Replayed thinking blocks, which Python strips and retries once if Anthropic rejects their signature.""" TOOL_CUSTOM_FIELD = "tool_custom_field" diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py index 6a0c867979b..99555bb4852 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -6,7 +6,11 @@ from typing import Final import litellm from litellm.anthropic_beta_headers_manager import is_foundry_api_base -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.core_helpers import ( + drops_unsupported_params, + proxy_request_headers, + stream_cost_injection_requested, +) from litellm.llms.opencode.common_utils import is_opencode_api_base from litellm.rust_bridge.detectors import api_base, overridden, present, streaming from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str @@ -64,8 +68,8 @@ def strict_tools_retry(request: NativeCall) -> bool: def stream_cost_injection(request: NativeCall) -> bool: - """Python adds response cost to streaming usage when the global setting is truthy.""" - return streaming(request) and bool(litellm.include_cost_in_streaming_usage) + """Python adds response cost to streaming usage when the global setting or the request header asks for it.""" + return streaming(request) and stream_cost_injection_requested(proxy_request_headers(request.resolved)) def pre_request_hook(_request: NativeCall) -> bool: diff --git a/tests/integration/spend/test_messages_stream_usage_cost.py b/tests/integration/spend/test_messages_stream_usage_cost.py index e28ab121de4..77ac95e2081 100644 --- a/tests/integration/spend/test_messages_stream_usage_cost.py +++ b/tests/integration/spend/test_messages_stream_usage_cost.py @@ -80,7 +80,9 @@ def _stream(message_id: str) -> bytes: ) -def _proxy_config(directory: Path, model: str, upstream_url: str) -> Path: +def _proxy_config( + directory: Path, model: str, upstream_url: str, litellm_settings: dict[str, JsonValue] | None = None +) -> Path: config: Final = directory / "streamed_usage_cost_config.yaml" config.write_text( json.dumps( @@ -109,7 +111,9 @@ def _proxy_config(directory: Path, model: str, upstream_url: str) -> Path: "proxy_batch_write_at": 1, "proxy_batch_polling_interval": 1, }, - "litellm_settings": {"include_cost_in_streaming_usage": True}, + "litellm_settings": {"include_cost_in_streaming_usage": True} + if litellm_settings is None + else litellm_settings, "router_settings": {"disable_cooldowns": True}, } ) @@ -176,3 +180,49 @@ def test_bedrock_messages_stream_usage_cost_matches_recorded_spend_with_custom_c "cost": pytest.approx(recorded_spend), }, (streamed_usage, rows, response.text) assert len(wire.drain()) == 1 + + +@pytest.mark.timeout(180) +def test_messages_stream_header_opts_into_usage_cost_while_global_flag_is_off(gateway: Gateway, tmp_path: Path) -> None: + model: Final = f"{BEDROCK_MODEL}-{uuid.uuid4().hex}" + message_ids: Final = {"plain": f"msg_{uuid.uuid4().hex}", "opted in": f"msg_{uuid.uuid4().hex}"} + + def respond(request: Request) -> Reply: + assert request.target == f"/model/{BEDROCK_MODEL}/invoke-with-response-stream", request.target + prompt: Final = json.loads(request.body)["messages"][0]["content"] + return Reply(content_type="application/vnd.amazon.eventstream", chunks=(_stream(message_ids[prompt]),)) + + def message_delta_usage(prompt: str, headers: dict[str, str]) -> dict[str, float]: + response: Final = candidate.request( + "POST", + "/v1/messages", + { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": OUTPUT_TOKENS, + "stream": True, + }, + headers=headers, + ) + assert response.status_code == 200, response.text + message_delta: Final = next(event for event in _data_events(response.text) if event["type"] == "message_delta") + return STREAMED_USAGE.validate_python(message_delta["usage"]) + + with wire_server(respond) as wire: + config: Final = _proxy_config(tmp_path, model, wire.url, litellm_settings={}) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate: + plain_usage: Final = message_delta_usage("plain", {}) + opted_in_usage: Final = message_delta_usage( + "opted in", {"x-litellm-include-cost-in-streaming-usage": "true"} + ) + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (message_ids["opted in"],) + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert "cost" not in plain_usage, plain_usage + assert opted_in_usage["cost"] == pytest.approx(float(str(rows[0]["spend"]))), (opted_in_usage, rows) + assert opted_in_usage["cost"] == pytest.approx(EXPECTED_SPEND), opted_in_usage + assert len(wire.drain()) == 2 diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 8127ffe558b..1817fd7d930 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -445,6 +445,47 @@ async def test_chunk_processor_injects_cost_into_openai_passthrough_usage_frame( assert final_payload["usage"]["total_tokens"] == 15 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("header_value", "cost_expected"), (("true", True), ("TRUE", True), ("1", False), ("false", False), (None, False)) +) +async def test_chunk_processor_honors_the_opt_in_header_while_global_flag_is_off( + monkeypatch: pytest.MonkeyPatch, header_value: str | None, cost_expected: bool +) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + chunks: Final = _openai_passthrough_stream_chunks() + logging_obj: Final = _unarmed_logging_obj(model="gpt-4o-mini") + headers: Final = {} if header_value is None else {"x-litellm-include-cost-in-streaming-usage": header_value} + logging_obj.model_call_details = {"litellm_params": {"proxy_server_request": {"headers": headers}}} + + received: Final = [ + chunk + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=_make_streaming_response(chunks), + request_body={"model": "gpt-4o-mini", "stream": True}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.OPENAI, + start_time=datetime(2026, 1, 1), + passthrough_success_handler_obj=MagicMock(), + url_route="/openai/v1/chat/completions", + route_streaming_logging=AsyncMock(), + ) + ] + await asyncio.sleep(0) + + usage_payloads: Final = [ + json.loads(chunk.decode("utf-8").split("data:", 1)[1].strip()) + for chunk in received + if b'"total_tokens"' in chunk + ] + assert len(usage_payloads) == 1 + assert ("cost" in usage_payloads[0]["usage"]) is cost_expected + if cost_expected: + assert usage_payloads[0]["usage"]["cost"] > 0 + else: + assert received == chunks + + @pytest.mark.asyncio async def test_chunk_processor_injects_cost_into_usage_frame_fragmented_across_chunks(monkeypatch): """Regression: an SSE usage frame split across transport chunks must still get @@ -826,7 +867,8 @@ async def test_dropped_attempt_logs_its_failure_without_claiming_the_shared_logg logging_obj.carry_failed_attempt_usage() retry: Final = _anthropic_messages_attempt( - logging_obj, _anthropic_stream_of((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())) + logging_obj, + _anthropic_stream_of((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())), ) async for _ in retry: pass @@ -858,6 +900,7 @@ async def test_every_dropped_attempt_logs_its_own_failure() -> None: stashed: Final = logging_obj.model_call_details.get("combined_usage_object") assert isinstance(stashed, litellm.Usage) and stashed.prompt_tokens == 9 + def _anthropic_message_start_frame(message_id: str = "msg_1", input_tokens: int = 52) -> bytes: return _anthropic_sse( "message_start", @@ -890,7 +933,11 @@ def _anthropic_completed_stream_tail() -> tuple[bytes, ...]: _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}), _anthropic_sse( "message_delta", - {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 3}}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, ), _anthropic_sse("message_stop", {"type": "message_stop"}), ) @@ -923,7 +970,9 @@ def _anthropic_stream_of(frames: tuple[bytes, ...]): "frames, billed", [ pytest.param( - (_anthropic_message_start_frame(), _anthropic_overloaded_error_frame()), True, id="error-after-message-start" + (_anthropic_message_start_frame(), _anthropic_overloaded_error_frame()), + True, + id="error-after-message-start", ), pytest.param((_anthropic_overloaded_error_frame(),), False, id="error-frame-only"), ], @@ -1102,7 +1151,8 @@ async def test_error_frame_then_exhausted_retry_bills_both_attempts() -> None: recorder: Final = _EventRecorder() logging_obj: Final = _anthropic_messages_logging_obj(recorder, "test-error-frame-terminal-billing") superseded: Final = _anthropic_messages_attempt( - logging_obj, _anthropic_stream_of((_anthropic_message_start_frame("msg_a1", 52), _anthropic_overloaded_error_frame())) + logging_obj, + _anthropic_stream_of((_anthropic_message_start_frame("msg_a1", 52), _anthropic_overloaded_error_frame())), ) async for chunk in superseded: if b"event: error" in chunk: @@ -1168,6 +1218,8 @@ async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_fr success_payload = recorder.success_kwargs[0]["standard_logging_object"] assert success_payload["status"] == "success" assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3) + + def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes: payload = { "candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}], diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index c59cde6d630..a519da41f61 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -5284,9 +5284,7 @@ class TestHandleLLMApiExceptionProviderOverload: assert proxy_exc.type == "internal_server_error" assert proxy_exc.to_dict()["code"] == "500" - async def test_openai_route_answers_wrapped_provider_overload_as_503( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: + async def test_openai_route_answers_wrapped_provider_overload_as_503(self, monkeypatch: pytest.MonkeyPatch) -> None: original: Final = InternalServerError( message="Overloaded", llm_provider="anthropic", @@ -5517,6 +5515,43 @@ class TestAsyncStreamingDataGeneratorFastPath: ProxyLogging._callback_capabilities_cache.clear() + async def test_opt_in_header_injects_cost_while_global_flag_is_off(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + monkeypatch.setattr(litellm, "callbacks", []) + ProxyLogging._callback_capabilities_cache.clear() + usage_frame: Final = ( + b'data: {"object":"chat.completion.chunk","choices":[],' + b'"usage":{"prompt_tokens":11,"completion_tokens":4,"total_tokens":15}}\n\n' + ) + + async def frames() -> AsyncGenerator[bytes, None]: + yield usage_frame + + async def stream(headers: dict[str, str]) -> list[object]: + return [ + c + async for c in ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=frames(), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + request_data={"model": "gpt-4o-mini", "proxy_server_request": {"headers": headers}}, + proxy_logging_obj=ProxyLogging(user_api_key_cache=MagicMock()), + serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk, + serialize_error=lambda e: "data: error\n\n", + ) + ] + + plain: Final = await stream({}) + opted_in: Final = await stream({"x-litellm-include-cost-in-streaming-usage": "true"}) + declined: Final = await stream({"x-litellm-include-cost-in-streaming-usage": "false"}) + ProxyLogging._callback_capabilities_cache.clear() + + assert plain == [usage_frame] + assert declined == [usage_frame] + assert isinstance(opted_in[0], bytes) + payload: Final = json.loads(opted_in[0].decode("utf-8").split("data:", 1)[1].strip()) + assert payload["usage"]["cost"] > 0 + assert payload["usage"]["total_tokens"] == 15 + class TestDisconnectGatherCleanup: def _disconnect_request(self) -> Request: @@ -10823,6 +10858,23 @@ class TestProcessChunkWithCostInjection: payload = json.loads(result.decode("utf-8").split("data:", 1)[1].strip()) assert payload["usage"]["cost"] > 0 + def test_request_header_opts_a_frame_into_cost_while_global_flag_is_off( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + chunk: Final = ( + b'data: {"object":"chat.completion.chunk","choices":[],' + b'"usage":{"prompt_tokens":11,"completion_tokens":4,"total_tokens":15}}\n\n' + ) + + assert ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection(chunk, "gpt-4o-mini") == chunk + injected: Final = ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection( + chunk, "gpt-4o-mini", request_headers={"X-LiteLLM-Include-Cost-In-Streaming-Usage": "True"} + ) + + payload: Final = json.loads(injected.decode("utf-8").split("data:", 1)[1].strip()) + assert payload["usage"]["cost"] > 0 + def test_chunk_ending_in_partial_frame_passes_through_byte_identical(self, monkeypatch): monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) chunk = ( diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index dfa29e24b2e..a6e6b2c0dca 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -352,6 +352,14 @@ def test_stream_cost_injection_needs_a_stream(monkeypatch: pytest.MonkeyPatch) - assert not messages_detectors.stream_cost_injection(call()) +def test_stream_cost_injection_reads_the_request_header_the_proxy_reads(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + opted_in: Final = {"proxy_server_request": {"headers": {"x-litellm-include-cost-in-streaming-usage": "true"}}} + assert messages_detectors.stream_cost_injection(call(stream=True, **opted_in)) + assert not messages_detectors.stream_cost_injection(call(**opted_in)) + assert not messages_detectors.stream_cost_injection(call(stream=True, proxy_server_request={"headers": {}})) + + def test_foundry_beta_row_reads_the_api_base_the_anthropic_handler_reads(monkeypatch: pytest.MonkeyPatch) -> None: for name in ("ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): monkeypatch.delenv(name, raising=False)