diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 159d7508f4e..8dfa08ff19b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -7,6 +7,7 @@ import traceback from collections.abc import AsyncGenerator, Callable, Mapping from datetime import datetime from functools import lru_cache +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal import anyio @@ -3063,12 +3064,12 @@ class ProxyBaseLLMRequestProcessing: if maybe_modified is not None: return maybe_modified elif isinstance(chunk, (bytes, bytearray)): - # Decode to str, inject, and rebuild as bytes try: - s: Final = chunk.decode("utf-8", errors="ignore") - maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(s, model_name) - if maybe_mod is not None: - return (maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n")).encode("utf-8") + s: Final = chunk.decode("utf-8") + if s.endswith("\n\n"): + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(s, model_name) + if maybe_mod is not None: + return maybe_mod.encode("utf-8") except Exception: pass elif isinstance(chunk, str): @@ -3113,10 +3114,79 @@ class ProxyBaseLLMRequestProcessing: except Exception: return None + @staticmethod + def _anthropic_stream_usage_kwargs(usage: Mapping[str, Any]) -> Mapping[str, Any]: + prompt_tokens: Final = int(usage.get("input_tokens", 0) or 0) + completion_tokens: Final = int(usage.get("output_tokens", 0) or 0) + total_tokens: Final = int( + usage.get("total_tokens", prompt_tokens + completion_tokens) or (prompt_tokens + completion_tokens) + ) + web_search_requests: Final = usage.get("web_search_requests") + server_tool_use: Final = ( + ServerToolUse(web_search_requests=web_search_requests) if web_search_requests is not None else None + ) + return MappingProxyType( + { + key: value + for key, value in ( + ("prompt_tokens", prompt_tokens), + ("completion_tokens", completion_tokens), + ("total_tokens", total_tokens), + ("completion_tokens_details", usage.get("completion_tokens_details")), + ("prompt_tokens_details", usage.get("prompt_tokens_details")), + ("cache_creation_input_tokens", usage.get("cache_creation_input_tokens")), + ("cache_read_input_tokens", usage.get("cache_read_input_tokens")), + ("server_tool_use", server_tool_use), + ) + if value is not None + } + ) + + @staticmethod + def _openai_stream_usage_kwargs(usage: Mapping[str, Any]) -> Mapping[str, Any]: + prompt_tokens: Final = int(usage.get("prompt_tokens", 0) or 0) + completion_tokens: Final = int(usage.get("completion_tokens", 0) or 0) + total_tokens: Final = int( + usage.get("total_tokens", prompt_tokens + completion_tokens) or (prompt_tokens + completion_tokens) + ) + return MappingProxyType( + { + key: value + for key, value in ( + ("prompt_tokens", prompt_tokens), + ("completion_tokens", completion_tokens), + ("total_tokens", total_tokens), + ("completion_tokens_details", usage.get("completion_tokens_details")), + ("prompt_tokens_details", usage.get("prompt_tokens_details")), + ) + if value is not None + } + ) + + @staticmethod + def _stream_usage_kwargs_for_event(obj: Mapping[str, object], usage: Mapping[str, Any]) -> Mapping[str, Any] | None: + if obj.get("type") == "message_delta": + return ProxyBaseLLMRequestProcessing._anthropic_stream_usage_kwargs(usage) + if obj.get("object") == "chat.completion.chunk": + return ProxyBaseLLMRequestProcessing._openai_stream_usage_kwargs(usage) + return None + + @staticmethod + def _completion_cost_or_none( + model_response: ModelResponse, model_name: str, service_tier: str | None + ) -> float | None: + try: + return litellm.completion_cost( + completion_response=model_response, model=model_name, service_tier=service_tier + ) + except Exception: + return None + @staticmethod def _inject_cost_into_usage_dict(obj: dict, model_name: str) -> dict | None: """ - Inject cost information into a usage dictionary for message_delta events. + Inject cost information into the usage object of a streamed usage event + (Anthropic ``message_delta`` or OpenAI ``chat.completion.chunk``). Args: obj: Dictionary containing the SSE event data @@ -3125,57 +3195,21 @@ class ProxyBaseLLMRequestProcessing: Returns: Modified dictionary with cost injected, or None if no modification needed """ - if obj.get("type") == "message_delta" and isinstance(obj.get("usage"), dict): - _usage: Final = obj["usage"] - prompt_tokens: Final = int(_usage.get("input_tokens", 0) or 0) - completion_tokens: Final = int(_usage.get("output_tokens", 0) or 0) - total_tokens: Final = int( - _usage.get("total_tokens", prompt_tokens + completion_tokens) or (prompt_tokens + completion_tokens) - ) - - # Extract additional usage fields - cache_creation_input_tokens: Final = _usage.get("cache_creation_input_tokens") - cache_read_input_tokens: Final = _usage.get("cache_read_input_tokens") - web_search_requests: Final = _usage.get("web_search_requests") - completion_tokens_details: Final = _usage.get("completion_tokens_details") - prompt_tokens_details: Final = _usage.get("prompt_tokens_details") - - usage_kwargs: Final[dict[str, Any]] = { - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "total_tokens": total_tokens, - } - - # Add optional named parameters - if completion_tokens_details is not None: - usage_kwargs["completion_tokens_details"] = completion_tokens_details - if prompt_tokens_details is not None: - usage_kwargs["prompt_tokens_details"] = prompt_tokens_details - - # Handle web_search_requests by wrapping in ServerToolUse - if web_search_requests is not None: - usage_kwargs["server_tool_use"] = ServerToolUse(web_search_requests=web_search_requests) - - # Add cache-related fields to **params (handled by Usage.__init__) - if cache_creation_input_tokens is not None: - usage_kwargs["cache_creation_input_tokens"] = cache_creation_input_tokens - if cache_read_input_tokens is not None: - usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens - - _mr: Final = ModelResponse(usage=Usage(**usage_kwargs)) - - try: - cost_val = litellm.completion_cost( - completion_response=_mr, - model=model_name, - ) - except Exception: - cost_val = None - - if cost_val is not None: - obj.setdefault("usage", {})["cost"] = cost_val - return obj - return None + usage: Final = obj.get("usage") + if not isinstance(usage, dict): + return None + usage_kwargs: Final = ProxyBaseLLMRequestProcessing._stream_usage_kwargs_for_event(obj, usage) + if usage_kwargs is None: + return None + service_tier: Final = obj.get("service_tier") + cost_val: Final = ProxyBaseLLMRequestProcessing._completion_cost_or_none( + ModelResponse(usage=Usage(**usage_kwargs)), + model_name, + service_tier if isinstance(service_tier, str) else None, + ) + if cost_val is None: + return None + return {**obj, "usage": {**usage, "cost": cost_val}} def maybe_get_model_id(self, _logging_obj: LiteLLMLoggingObj | None) -> str | None: """ diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index ff1c12d08d7..192600ba150 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -56,7 +56,13 @@ class PassThroughStreamingHandler: cost_injection_active: Final = ( bool(getattr(litellm, "include_cost_in_streaming_usage", False)) and bool(model_name) - and endpoint_type in (EndpointType.VERTEX_AI, EndpointType.ANTHROPIC) + and ( + endpoint_type in (EndpointType.ANTHROPIC, EndpointType.OPENAI) + or ( + endpoint_type == EndpointType.VERTEX_AI + and ("streamRawPredict" in url_route or "rawPredict" in url_route) + ) + ) ) try: if not cost_injection_active: @@ -74,21 +80,7 @@ class PassThroughStreamingHandler: async for chunk in response.aiter_bytes(): raw_bytes.append(chunk) PassThroughStreamingHandler._stamp_first_chunk_if_needed(litellm_logging_obj) - if endpoint_type == EndpointType.VERTEX_AI: - if "streamRawPredict" in url_route or "rawPredict" in url_route: - modified_chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( - chunk, resolved_model_name - ) - if modified_chunk is not None: - chunk = modified_chunk - else: # EndpointType.ANTHROPIC - modified_chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( - chunk, resolved_model_name - ) - if modified_chunk is not None: - chunk = modified_chunk - - yield chunk + yield ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, resolved_model_name) except Exception as e: verbose_proxy_logger.error("Error in chunk_processor: %s", e) raise 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 163a0cbff3c..6d47897f710 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 @@ -1,12 +1,14 @@ -"""Regression tests for LIT-2642 — interrupted pass-through streams must still log usage.""" +"""Regression tests for PassThroughStreamingHandler.chunk_processor.""" import asyncio +import json from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import litellm from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.proxy.pass_through_endpoints.streaming_handler import ( PassThroughStreamingHandler, @@ -361,6 +363,107 @@ async def test_chunk_processor_stamps_completion_start_time_on_cost_injection_pa mock_logging_obj._update_completion_start_time.assert_called_once() +def _openai_passthrough_stream_chunks(): + return [ + ( + b'data: {"id":"chatcmpl-1","object":"chat.completion.chunk",' + b'"choices":[{"index":0,"delta":{"content":"Hi"}}],"usage":null}\n\n' + ), + b": keepalive\n\n", + ( + b'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[],' + b'"usage":{"prompt_tokens":11,"completion_tokens":4,"total_tokens":15,' + b'"prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},' + b'"completion_tokens_details":{"reasoning_tokens":0,"audio_tokens":0,' + b'"accepted_prediction_tokens":0,"rejected_prediction_tokens":0}}}\n\n' + ), + b"data: [DONE]\n\n", + ] + + +async def _collect_openai_passthrough_chunks(chunks, endpoint_type): + response = _make_streaming_response(chunks) + with patch.object( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + new=AsyncMock(), + ): + received = [] + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "gpt-4o-mini", "stream": True}, + litellm_logging_obj=MagicMock(), + endpoint_type=endpoint_type, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/openai/v1/chat/completions", + ): + received.append(chunk) + await asyncio.sleep(0) + return received + + +@pytest.mark.asyncio +async def test_chunk_processor_injects_cost_into_openai_passthrough_usage_frame(monkeypatch): + """Regression: issue #36492 — with include_cost_in_streaming_usage on, the final + OpenAI passthrough chat.completion.chunk usage frame must carry usage.cost, like + every other streaming surface already does.""" + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + chunks = _openai_passthrough_stream_chunks() + + received = await _collect_openai_passthrough_chunks(chunks, EndpointType.OPENAI) + + assert received[0] == chunks[0] + assert received[1] == chunks[1] + assert received[3] == chunks[3] + final_payload = json.loads(received[2].decode("utf-8").split("data:", 1)[1].strip()) + pricing = litellm.model_cost["gpt-4o-mini"] + expected_cost = 11 * pricing["input_cost_per_token"] + 4 * pricing["output_cost_per_token"] + assert final_payload["usage"]["cost"] == pytest.approx(expected_cost) + assert final_payload["usage"]["cost"] > 0 + assert final_payload["usage"]["prompt_tokens"] == 11 + assert final_payload["usage"]["completion_tokens"] == 4 + assert final_payload["usage"]["total_tokens"] == 15 + + +@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) + chunks = _openai_passthrough_stream_chunks() + + received = await _collect_openai_passthrough_chunks(chunks, EndpointType.OPENAI) + + assert received == chunks + + +@pytest.mark.asyncio +async def test_chunk_processor_flag_on_leaves_openai_frames_without_usage_untouched(monkeypatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + chunks = [ + ( + b'data: {"id":"chatcmpl-1","object":"chat.completion.chunk",' + b'"choices":[{"index":0,"delta":{"content":"Hi"}}],"usage":null}\n\n' + ), + b": keepalive\n\n", + b"not json at all\n\n", + b"data: [DONE]\n\n", + ] + + received = await _collect_openai_passthrough_chunks(chunks, EndpointType.OPENAI) + + assert received == chunks + + +@pytest.mark.asyncio +async def test_chunk_processor_flag_on_leaves_generic_passthrough_untouched(monkeypatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + chunks = _openai_passthrough_stream_chunks() + + received = await _collect_openai_passthrough_chunks(chunks, EndpointType.GENERIC) + + assert received == chunks + + def test_convert_raw_bytes_survives_truncated_multibyte_sequence(): """A stream cut mid-multibyte-sequence (client disconnect) must still decode via errors="replace" so the usage events already received are logged, instead diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index aacc7498ccb..a3c0f0089fe 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1,6 +1,7 @@ import asyncio import copy import datetime +import json from types import SimpleNamespace from typing import AsyncGenerator, Callable, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -5746,3 +5747,132 @@ class TestPerRequestModelGroupAlias: ) assert merged_for == ["group-b"] + + +class TestInjectCostIntoUsageDict: + @staticmethod + def _expected_cost(model, prompt_tokens, completion_tokens): + pricing = litellm.model_cost[model] + return prompt_tokens * pricing["input_cost_per_token"] + completion_tokens * pricing["output_cost_per_token"] + + def test_openai_chat_completion_chunk_usage_gets_cost(self): + event = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "choices": [], + "usage": { + "prompt_tokens": 11, + "completion_tokens": 4, + "total_tokens": 15, + "prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 0}, + "completion_tokens_details": { + "reasoning_tokens": 0, + "audio_tokens": 0, + "accepted_prediction_tokens": 0, + "rejected_prediction_tokens": 0, + }, + }, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, "gpt-4o-mini") + + assert result is not None + assert result["usage"]["cost"] == pytest.approx(self._expected_cost("gpt-4o-mini", 11, 4)) + assert result["usage"]["cost"] > 0 + assert result["usage"]["prompt_tokens"] == 11 + assert result["id"] == "chatcmpl-1" + assert "cost" not in event["usage"] + + def test_anthropic_message_delta_usage_still_gets_cost(self): + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"input_tokens": 11, "output_tokens": 4}, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, "claude-haiku-4-5") + + assert result is not None + assert result["usage"]["cost"] == pytest.approx(self._expected_cost("claude-haiku-4-5", 11, 4)) + assert result["usage"]["cost"] > 0 + assert result["usage"]["output_tokens"] == 4 + + def test_openai_chunk_with_flex_service_tier_uses_flex_pricing(self): + event = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "service_tier": "flex", + "choices": [], + "usage": {"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100}, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, "gpt-5-mini") + + assert result is not None + pricing = litellm.model_cost["gpt-5-mini"] + expected_flex_cost = 1000 * pricing["input_cost_per_token_flex"] + 100 * pricing["output_cost_per_token_flex"] + assert result["usage"]["cost"] == pytest.approx(expected_flex_cost) + assert result["usage"]["cost"] < self._expected_cost("gpt-5-mini", 1000, 100) + + def test_openai_chunk_with_null_usage_is_not_modified(self): + event = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {"content": "Hi"}}], + "usage": None, + } + + assert ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, "gpt-4o-mini") is None + + def test_unrecognized_event_shape_with_usage_is_not_modified(self): + event = {"kind": "custom", "usage": {"prompt_tokens": 11, "completion_tokens": 4}} + + assert ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, "gpt-4o-mini") is None + + def test_sse_frame_with_coalesced_done_line_injects_into_usage_frame(self): + frame = ( + 'data: {"object":"chat.completion.chunk","choices":[],' + '"usage":{"prompt_tokens":11,"completion_tokens":4,"total_tokens":15}}\n\n' + "data: [DONE]\n\n" + ) + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(frame, "gpt-4o-mini") + + assert result is not None + assert "data: [DONE]" in result + injected = json.loads(result.split("\n")[0].split("data:", 1)[1].strip()) + assert injected["usage"]["cost"] == pytest.approx(self._expected_cost("gpt-4o-mini", 11, 4)) + + +class TestProcessChunkWithCostInjection: + def test_complete_usage_frame_chunk_is_injected(self, monkeypatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + chunk = ( + b'data: {"object":"chat.completion.chunk","choices":[],' + b'"usage":{"prompt_tokens":11,"completion_tokens":4,"total_tokens":15}}\n\n' + ) + + result = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, "gpt-4o-mini") + + assert result != chunk + assert result.endswith(b"\n\n") + payload = json.loads(result.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 = ( + b'data: {"object":"chat.completion.chunk","choices":[],' + b'"usage":{"prompt_tokens":11,"completion_tokens":4,"total_tokens":15}}\n\ndata: [DO' + ) + + assert ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, "gpt-4o-mini") == chunk + + def test_chunk_with_invalid_utf8_passes_through_byte_identical(self, monkeypatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + chunk = ( + b'\xa8data: {"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