From 426b9094478bda9f97c1bb2ccca1f3db27735d46 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 10 Aug 2026 19:56:34 -0700 Subject: [PATCH 1/6] fix(proxy): inject streaming usage cost on openai passthrough streams --- litellm/proxy/common_request_processing.py | 148 +++++++++++------- .../streaming_handler.py | 24 +-- .../test_streaming_handler_interrupt.py | 105 ++++++++++++- .../proxy/test_common_request_processing.py | 130 +++++++++++++++ 4 files changed, 333 insertions(+), 74 deletions(-) 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 From 46fb1cd514cd2b21db62146fae06c704f8193e63 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:07:57 -0700 Subject: [PATCH 2/6] fix(proxy): reassemble fragmented SSE frames and inject logging dependency --- .../streaming_handler.py | 42 ++++++++++++++-- .../test_streaming_handler_interrupt.py | 49 +++++++++++++------ 2 files changed, 72 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 192600ba150..da4a8eceb89 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -1,5 +1,6 @@ +from collections.abc import Coroutine from datetime import datetime -from typing import Final +from typing import Final, Protocol import httpx @@ -24,6 +25,21 @@ from .llm_provider_handlers.vertex_passthrough_logging_handler import ( from .success_handler import PassThroughEndpointLogging +class RouteStreamingLogging(Protocol): + def __call__( + self, + *, + litellm_logging_obj: LiteLLMLoggingObj, + passthrough_success_handler_obj: PassThroughEndpointLogging, + url_route: str, + request_body: dict, + endpoint_type: EndpointType, + start_time: datetime, + raw_bytes: list[bytes], + end_time: datetime, + ) -> Coroutine[None, None, None]: ... + + class PassThroughStreamingHandler: @staticmethod def _stamp_first_chunk_if_needed(litellm_logging_obj: LiteLLMLoggingObj) -> None: @@ -39,7 +55,11 @@ class PassThroughStreamingHandler: start_time: datetime, passthrough_success_handler_obj: PassThroughEndpointLogging, url_route: str, + route_streaming_logging: RouteStreamingLogging | None = None, ): + resolved_route_streaming_logging: Final[RouteStreamingLogging] = ( + route_streaming_logging or PassThroughStreamingHandler._route_streaming_logging_to_handler + ) raw_bytes: Final[list[bytes]] = [] logging_scheduled = False model_name: Final = PassThroughStreamingHandler._extract_model_for_cost_injection( @@ -77,10 +97,19 @@ class PassThroughStreamingHandler: # -> ``str`` for the per-chunk call site. assert model_name is not None resolved_model_name: Final[str] = model_name + pending = b"" async for chunk in response.aiter_bytes(): raw_bytes.append(chunk) PassThroughStreamingHandler._stamp_first_chunk_if_needed(litellm_logging_obj) - yield ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, resolved_model_name) + complete_frames, pending = PassThroughStreamingHandler._split_complete_sse_frames( + pending + chunk + ) # rebind-ok: SSE frame reassembly buffer across transport chunks + if complete_frames: + yield ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + complete_frames, resolved_model_name + ) + if pending: + yield pending except Exception as e: verbose_proxy_logger.error("Error in chunk_processor: %s", e) raise @@ -96,7 +125,7 @@ class PassThroughStreamingHandler: logging_scheduled = True try: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=PassThroughStreamingHandler._route_streaming_logging_to_handler( + async_coroutine=resolved_route_streaming_logging( litellm_logging_obj=litellm_logging_obj, passthrough_success_handler_obj=passthrough_success_handler_obj, url_route=url_route, @@ -110,6 +139,13 @@ class PassThroughStreamingHandler: except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) + @staticmethod + def _split_complete_sse_frames(pending: bytes) -> tuple[bytes, bytes]: + frame_boundary: Final = pending.rfind(b"\n\n") + if frame_boundary == -1: + return b"", pending + return pending[: frame_boundary + 2], pending[frame_boundary + 2 :] + @staticmethod async def _route_streaming_logging_to_handler( litellm_logging_obj: LiteLLMLoggingObj, 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 6d47897f710..7b649db43ed 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 @@ -383,23 +383,19 @@ def _openai_passthrough_stream_chunks(): 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", + route_streaming_logging=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) + received.append(chunk) + await asyncio.sleep(0) return received @@ -426,6 +422,27 @@ async def test_chunk_processor_injects_cost_into_openai_passthrough_usage_frame( assert final_payload["usage"]["total_tokens"] == 15 +@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 + cost injected once the frame completes, instead of passing through untouched.""" + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + whole = _openai_passthrough_stream_chunks() + usage_frame = whole[2] + split_at = len(usage_frame) // 2 + chunks = [whole[0], whole[1], usage_frame[:split_at], usage_frame[split_at:], whole[3]] + + received = await _collect_openai_passthrough_chunks(chunks, EndpointType.OPENAI) + + reassembled = b"".join(received).decode("utf-8") + usage_lines = [ln for ln in reassembled.split("\n") if '"total_tokens"' in ln] + assert len(usage_lines) == 1 + final_payload = json.loads(usage_lines[0].split("data:", 1)[1].strip()) + assert final_payload["usage"]["cost"] > 0 + assert final_payload["usage"]["prompt_tokens"] == 11 + assert reassembled.endswith("data: [DONE]\n\n") + + @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) From 938396ef9099ebc5c0056169286325bf1d2bae0a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:37:07 -0700 Subject: [PATCH 3/6] fix(proxy): recognize crlf sse frame boundaries in passthrough reassembly --- litellm/proxy/common_request_processing.py | 2 +- .../streaming_handler.py | 8 +++++--- .../test_streaming_handler_interrupt.py | 18 ++++++++++++++++++ 3 files changed, 24 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 8dfa08ff19b..4e422ee49d7 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3066,7 +3066,7 @@ class ProxyBaseLLMRequestProcessing: elif isinstance(chunk, (bytes, bytearray)): try: s: Final = chunk.decode("utf-8") - if s.endswith("\n\n"): + if s.endswith(("\n\n", "\r\n\r\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") diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index da4a8eceb89..8428c7dcbe2 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -141,10 +141,12 @@ class PassThroughStreamingHandler: @staticmethod def _split_complete_sse_frames(pending: bytes) -> tuple[bytes, bytes]: - frame_boundary: Final = pending.rfind(b"\n\n") - if frame_boundary == -1: + lf_boundary_end: Final = pending.rfind(b"\n\n") + 2 + crlf_boundary_end: Final = pending.rfind(b"\r\n\r\n") + 4 + boundary_end: Final = max(lf_boundary_end if lf_boundary_end >= 2 else 0, crlf_boundary_end if crlf_boundary_end >= 4 else 0) + if boundary_end == 0: return b"", pending - return pending[: frame_boundary + 2], pending[frame_boundary + 2 :] + return pending[:boundary_end], pending[boundary_end:] @staticmethod async def _route_streaming_logging_to_handler( 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 7b649db43ed..d559faba1c2 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 @@ -443,6 +443,24 @@ async def test_chunk_processor_injects_cost_into_usage_frame_fragmented_across_c assert reassembled.endswith("data: [DONE]\n\n") +@pytest.mark.asyncio +async def test_chunk_processor_streams_crlf_delimited_frames_live_and_injects_cost(monkeypatch): + """Regression: CRLF-delimited SSE frames must flow as they complete instead of + buffering until EOF, and the usage frame must still get cost injected.""" + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + chunks = [chunk.replace(b"\n\n", b"\r\n\r\n") for chunk in _openai_passthrough_stream_chunks()] + + received = await _collect_openai_passthrough_chunks(chunks, EndpointType.OPENAI) + + assert len(received) == len(chunks) + assert received[0] == chunks[0] + reassembled = b"".join(received).decode("utf-8") + usage_lines = [ln for ln in reassembled.replace("\r\n", "\n").split("\n") if '"total_tokens"' in ln] + assert len(usage_lines) == 1 + final_payload = json.loads(usage_lines[0].split("data:", 1)[1].strip()) + assert final_payload["usage"]["cost"] > 0 + + @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) From 1819c97c271c58a527360d7bb967c1bdb3badd3e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:50:43 -0700 Subject: [PATCH 4/6] style: format streaming handler --- litellm/proxy/pass_through_endpoints/streaming_handler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 8428c7dcbe2..c7ccd2d0d0f 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -143,7 +143,9 @@ class PassThroughStreamingHandler: def _split_complete_sse_frames(pending: bytes) -> tuple[bytes, bytes]: lf_boundary_end: Final = pending.rfind(b"\n\n") + 2 crlf_boundary_end: Final = pending.rfind(b"\r\n\r\n") + 4 - boundary_end: Final = max(lf_boundary_end if lf_boundary_end >= 2 else 0, crlf_boundary_end if crlf_boundary_end >= 4 else 0) + boundary_end: Final = max( + lf_boundary_end if lf_boundary_end >= 2 else 0, crlf_boundary_end if crlf_boundary_end >= 4 else 0 + ) if boundary_end == 0: return b"", pending return pending[:boundary_end], pending[boundary_end:] From e7c8cff3b7c1a175b0e86638938ef7241631ca9a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 00:45:32 -0700 Subject: [PATCH 5/6] fix(proxy): preserve crlf line endings when injecting streamed usage cost --- litellm/proxy/common_request_processing.py | 3 +-- .../pass_through_endpoints/test_streaming_handler_interrupt.py | 3 +++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4e422ee49d7..10c54a25805 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3107,8 +3107,7 @@ class ProxyBaseLLMRequestProcessing: obj = json.loads(json_part) maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name) if maybe_modified is not None: - # Replace just this line with updated JSON using safe_dumps - lines[idx] = f"data: {safe_dumps(maybe_modified)}" + lines[idx] = "data: " + safe_dumps(maybe_modified) + ("\r" if ln.endswith("\r") else "") return "\n".join(lines) return None except Exception: 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 d559faba1c2..1d82a5dfc6e 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 @@ -454,6 +454,9 @@ async def test_chunk_processor_streams_crlf_delimited_frames_live_and_injects_co assert len(received) == len(chunks) assert received[0] == chunks[0] + injected_usage_frame = received[2] + assert injected_usage_frame.endswith(b"\r\n\r\n") + assert b"\n" not in injected_usage_frame.replace(b"\r\n", b"") reassembled = b"".join(received).decode("utf-8") usage_lines = [ln for ln in reassembled.replace("\r\n", "\n").split("\n") if '"total_tokens"' in ln] assert len(usage_lines) == 1 From dca7ba18d4a9f6e753285adcfc7454dde7a01087 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 11 Aug 2026 09:41:55 -0700 Subject: [PATCH 6/6] feat(ui): show models under each tier in routing benchmark chart (#36291) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(ui): show models under each tier in routing benchmark chart - Add TierTurnsChart: donut chart showing turns per complexity tier with tier-assigned models listed below each tier name in the legend - Only complexity routers show models; quality routers show tier name + % (quality tiers don't pin specific models) - Change 'Estimated spend at highest-cost model' wording to 'highest-tier' to clarify it's the most capable tier's estimated cost, not just the single-highest model Closes LIT-5302 * fix(ui): use categorical colors for tier donut, trim redundant turn count - Tier donut chart now uses a dedicated categorical palette instead of SEQUENTIAL_COLOR_RAMP, which is a blue monochrome gradient meant for magnitude series, not distinct categories. - Space out the tier legend rows (gap-3 -> gap-6) for readability. - Drop the turn count from "avg saved per session" since Routing by tier already shows the total turns. Co-Authored-By: Claude * test(ui): drop prohibited explanatory comments in TierTurnsChart test Per repo convention against source comments; the test name and assertions already communicate the scenario. Addresses Greptile review. Co-Authored-By: Claude * Remove dead modulo from color index in TierTurnsChart The colors array is built with length equal to slices.length, so idx % colors.length is always a no-op in the render loop. Simplify to idx for clarity. * fix(ui): wrap CostOptimizationView tests in QueryClientProvider The tests render CostOptimizationView which uses useCan() → useIsOrgAdmin() → useOrganizations() and useDailyActivityRange(), both of which call React Query's useQuery(). Without QueryClientProvider wrapping the render, React Query throws 'No QueryClient set' error. Also mock the required networking calls (organizationListCall, userDailyActivityCall) to prevent spurious network errors in test runs. All 7 tests now pass (CostOptimizationView + CostOptimizationView.activity). * style(ui): format test files and extract object literal to fix linting - Format CostOptimizationView.test.tsx with prettier - Extract getToolSpend mock response to named variable to satisfy eslint - Pass frontend-lint checks * fix(ui): hoist mockToolSpendResponse into vi.hoisted to fix test initialization Extracting the response object to a named variable violated hoisting rules: vi.mock() factories are evaluated at hoisting time before regular const declarations. Move mockToolSpendResponse into vi.hoisted() block. --------- Co-authored-by: Claude --- .../auto_router_endpoints.py | 4 +- .../AutoRouterBenchmarksTab.test.tsx | 31 +++- .../_components/AutoRouterBenchmarksTab.tsx | 45 +++-- .../CostOptimizationView.activity.test.tsx | 19 +- .../_components/CostOptimizationView.test.tsx | 16 +- .../_components/TierTurnsChart.test.tsx | 163 ++++++++++++++++++ .../_components/TierTurnsChart.tsx | 148 ++++++++++++++++ .../_components/autoRouterBenchmarks.ts | 5 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 9 files changed, 395 insertions(+), 38 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.tsx diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 6626dea6849..269d9b50414 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -126,8 +126,8 @@ class AutoRouterBenchmarkGroup(AutoRouterBenchmarkTotals): description="Turns per tier, keyed by the tier name the routing decision recorded at " "request time (never re-derived at read time, since the tier-to-model mapping is " "mutable config). Tier names are scoped to this group's router_type and are not " - "comparable across types: a complexity router reports 'simple'/'medium'/'complex'/" - "'reasoning', a quality router reports its numeric quality tier, and an adaptive router " + "comparable across types: a complexity router reports 'SIMPLE'/'MEDIUM'/'COMPLEX'/" + "'REASONING', a quality router reports its numeric quality tier, and an adaptive router " "records no tier at all. Turns no tier served (the classifier fell back to default_model) " "are absent rather than pooled under a sentinel key, so the values may sum to less than turns", ) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx index a5767383307..4a6ff77d99b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx @@ -1,10 +1,15 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { fireEvent, render, screen } from "@testing-library/react"; import React from "react"; -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels"; import { ApiError } from "@/lib/http/client"; vi.mock("./useAutoRouterBenchmarks", () => ({ useAutoRouterBenchmarks: vi.fn() })); +vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAutoRouters: vi.fn() })); + +import { useAutoRouters } from "@/app/(dashboard)/hooks/models/useModels"; import AutoRouterBenchmarksTab from "./AutoRouterBenchmarksTab"; import type { @@ -16,6 +21,10 @@ import { useAutoRouterBenchmarks } from "./useAutoRouterBenchmarks"; type HookResult = ReturnType; +const mockAutoRouters = (deployments: AutoRouterDeployment[] = []) => { + vi.mocked(useAutoRouters).mockReturnValue({ data: deployments } as unknown as ReturnType); +}; + const mockHook = (result: { data?: AutoRouterBenchmarksResponse; isPending?: boolean; error?: Error }) => { vi.mocked(useAutoRouterBenchmarks).mockReturnValue({ data: result.data, @@ -71,9 +80,20 @@ const response = (groups: AutoRouterBenchmarkGroup[], shared: Totals = totals()) groups, }); -const renderTab = () => render(); +const renderTab = () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return render( + + + , + ); +}; describe("AutoRouterBenchmarksTab", () => { + beforeEach(() => { + mockAutoRouters(); + }); + it("leads with total estimated savings, before the three session-shape metrics", () => { mockHook({ data: response([group(), group({ router_name: "gpt-auto" })]) }); renderTab(); @@ -97,7 +117,7 @@ describe("AutoRouterBenchmarksTab", () => { expect(screen.getByText("-86%")).toBeInTheDocument(); expect(screen.getByText("Actual auto-router spend")).toBeInTheDocument(); expect(screen.getByText("$359.86")).toBeInTheDocument(); - expect(screen.getByText("Estimated spend at highest-cost model")).toBeInTheDocument(); + expect(screen.getByText("Estimated spend at highest-tier model")).toBeInTheDocument(); expect(screen.getByText("$2,534.45")).toBeInTheDocument(); expect(screen.getByText("32.7")).toBeInTheDocument(); expect(screen.getByText("2.1h")).toBeInTheDocument(); @@ -108,12 +128,9 @@ describe("AutoRouterBenchmarksTab", () => { mockHook({ data: response([group(), group({ router_name: "gpt-auto" })]) }); renderTab(); - expect(screen.getByText("Total sessions")).toBeInTheDocument(); - expect(screen.getByText("94")).toBeInTheDocument(); - expect(screen.getByText("Total turns")).toBeInTheDocument(); - expect(screen.getByText("3,073")).toBeInTheDocument(); expect(screen.getByText("Avg saved per session")).toBeInTheDocument(); expect(screen.getByText("$23.13")).toBeInTheDocument(); + expect(screen.getByText("across 94 sessions")).toBeInTheDocument(); }); it("shows a cost increase as a positive delta rather than a saving", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx index ff0f52940b2..5d4fda765e7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx @@ -2,6 +2,8 @@ import React, { useState } from "react"; +import type { AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels"; +import { useAutoRouters } from "@/app/(dashboard)/hooks/models/useModels"; import { Badge } from "@/components/ui/badge"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; @@ -29,6 +31,7 @@ import { type BucketRow, } from "./autoRouterBenchmarks"; import { usd } from "./costOptimizationUtils"; +import TierTurnsChart from "./TierTurnsChart"; import { useAutoRouterBenchmarks } from "./useAutoRouterBenchmarks"; const Message: React.FC<{ children: React.ReactNode }> = ({ children }) => ( @@ -51,7 +54,7 @@ const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => { const cheaper = stats.saved_spend >= 0; return ( -
+

Total estimated savings

@@ -64,38 +67,22 @@ const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => { {Math.abs(stats.saved_pct).toFixed(0)}%
-
- -
Actual auto-router spend
{usd(stats.spend)}
-
Estimated spend at highest-cost model
+
Estimated spend at highest-tier model
{usd(stats.baseline_spend)}
-
-
-
-

Total sessions

-

{stats.sessions.toLocaleString()}

-
-
-

Total turns

-

{stats.turns.toLocaleString()}

-
-
-
-
-
Avg saved per session
-
{usd(stats.saved_per_session)}
-
-
+
+

Avg saved per session

+

{usd(stats.saved_per_session)}

+

across {stats.sessions.toLocaleString()} sessions

@@ -233,9 +220,10 @@ interface BenchmarksBodyProps { error: unknown; data: AutoRouterBenchmarksResponse | undefined; selectedKey: string; + autoRouters: readonly AutoRouterDeployment[]; } -const BenchmarksBody: React.FC = ({ isPending, error, data, selectedKey }) => { +const BenchmarksBody: React.FC = ({ isPending, error, data, selectedKey, autoRouters }) => { if (isPending) return Loading auto-router usage...; if (error instanceof ApiError && error.status === 403) { return Auto-router usage is visible to proxy admin roles only; @@ -249,6 +237,8 @@ const BenchmarksBody: React.FC = ({ isPending, error, data, <> + +
@@ -282,6 +272,7 @@ const AutoRouterBenchmarksTab: React.FC = ({ acces const [range, setRange] = useState("30d"); const { data, isPending, error } = useAutoRouterBenchmarks(accessToken, range); const [selectedKey, setSelectedKey] = useState(ALL_ROUTERS); + const { data: autoRouters } = useAutoRouters(); const groups = data?.groups ?? []; const selectedLabel = data ? viewFor(data, selectedKey).label : "All auto-routers"; @@ -319,7 +310,13 @@ const AutoRouterBenchmarksTab: React.FC = ({ acces
- +
); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx index ca7adf07941..96502cac953 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.activity.test.tsx @@ -1,12 +1,23 @@ +import React from "react"; import { fireEvent, render, waitFor } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; const mockUserDailyActivityCall = vi.fn(); +const { useAuthorizedMock, mockToolSpendResponse } = vi.hoisted(() => ({ + useAuthorizedMock: vi.fn(), + mockToolSpendResponse: { by_tool: [], daily: [], start_date: null, end_date: null }, +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: useAuthorizedMock, +})); vi.mock("@/components/networking", () => ({ userDailyActivityCall: (...args: unknown[]) => mockUserDailyActivityCall(...args), - getToolSpend: vi.fn().mockResolvedValue({ by_tool: [], daily: [], start_date: null, end_date: null }), + getToolSpend: vi.fn().mockResolvedValue(mockToolSpendResponse), getGeneralSettingsCall: vi.fn().mockResolvedValue([]), + organizationListCall: vi.fn().mockResolvedValue([]), })); vi.mock("@/components/shared/advanced_date_picker", () => ({ @@ -38,9 +49,13 @@ const singlePage = { describe("CostOptimizationView daily activity", () => { it("fetches daily activity once for the page and shares it with every tab that needs it", async () => { mockUserDailyActivityCall.mockResolvedValue(singlePage); + useAuthorizedMock.mockReturnValue({ accessToken: "test-token", userId: "u1", userRole: "proxy_admin" }); + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); const { getByRole, getByTestId } = render( - , + + + , ); await waitFor(() => expect(mockUserDailyActivityCall).toHaveBeenCalledTimes(1)); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx index c6d5a410418..60926f575bc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx @@ -1,5 +1,7 @@ +import React from "react"; import { fireEvent, render } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; const { useAuthorizedMock } = vi.hoisted(() => ({ useAuthorizedMock: vi.fn() })); @@ -7,6 +9,13 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: useAuthorizedMock, })); +vi.mock("@/components/networking", () => ({ + organizationListCall: vi.fn().mockResolvedValue([]), + userDailyActivityCall: vi + .fn() + .mockResolvedValue({ results: [], metadata: { total_pages: 1, has_more: false, page: 1 } }), +})); + vi.mock("./UsageTab", () => ({ __esModule: true, default: () =>
})); vi.mock("./PromptCompressionTab", () => ({ __esModule: true, default: () =>
})); vi.mock("./PromptCachingTab", () => ({ __esModule: true, default: () =>
})); @@ -19,7 +28,12 @@ import CostOptimizationView from "./CostOptimizationView"; const renderView = (userRole = "Admin") => { useAuthorizedMock.mockReturnValue({ accessToken: "test-token", userId: "u1", userRole }); - return render(); + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return render( + + + , + ); }; describe("CostOptimizationView", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.test.tsx new file mode 100644 index 00000000000..057eb54ee4e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.test.tsx @@ -0,0 +1,163 @@ +import { render, screen } from "@testing-library/react"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; + +import type { AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels"; + +vi.mock("@/components/shared/charts", () => ({ + DonutChart: ({ label }: { label: string }) =>
{label}
, + SEQUENTIAL_COLOR_RAMP: ["indigo", "blue"], + chartColorValue: (color: string) => color, +})); + +import TierTurnsChart, { tierDisplayLabel } from "./TierTurnsChart"; +import type { AutoRouterBenchmarkGroup, BenchmarkView } from "./autoRouterBenchmarks"; + +const totalsOnly = { + sessions: 3, + turns: 9, + avg_turns_per_session: 3, + avg_session_seconds: 60, + avg_tokens_per_session: 100, + spend: 1, + saved_spend: 1, + baseline_spend: 2, + saved_pct: 50, + saved_per_session: 0.33, + cache: { + coverage_pct: 0, + hit_rate_pct: 0, + same_model: { turns: 0, hits: 0, hit_rate_pct: 0 }, + first_visit: { turns: 0, hits: 0, hit_rate_pct: 0 }, + return_to_tier: { turns: 0, hits: 0, hit_rate_pct: 0 }, + unordered_turns: 0, + return_misses_expired: 0, + return_misses_within_ttl: 0, + return_misses_unknown: 0, + ttl_5m_turns: 0, + ttl_1h_turns: 0, + }, +}; + +const groupView = (overrides: Partial = {}): BenchmarkView => ({ + label: "claude-auto", + stats: { + ...totalsOnly, + router_name: "claude-auto", + router_type: "complexity", + tier_turns: { SIMPLE: 3, COMPLEX: 1 }, + ...overrides, + } as AutoRouterBenchmarkGroup, +}); + +const deployment = (config: unknown): AutoRouterDeployment => ({ + model_name: "claude-auto", + litellm_params: { model: "auto_router/claude-auto", complexity_router_config: config }, +}); + +describe("tierDisplayLabel", () => { + it("prefers the admin's custom label for a canonical complexity tier", () => { + expect(tierDisplayLabel("SIMPLE", { SIMPLE: "Cheap" })).toBe("Cheap"); + }); + + it("falls back to the canonical name when that tier has no custom label", () => { + expect(tierDisplayLabel("COMPLEX", { SIMPLE: "Cheap" })).toBe("Complex"); + expect(tierDisplayLabel("REASONING", undefined)).toBe("Reasoning"); + }); + + it("shows a non-complexity tier verbatim, since no label map covers a quality router's tier", () => { + expect(tierDisplayLabel("3", { SIMPLE: "Cheap" })).toBe("3"); + }); +}); + +describe("TierTurnsChart", () => { + it("labels each slice with its tier and share of the tiered turns", () => { + render(); + + expect(screen.getByText("Cheap 75%")).toBeInTheDocument(); + expect(screen.getByText("Complex 25%")).toBeInTheDocument(); + expect(screen.getByTestId("donut")).toHaveTextContent("4 total turns"); + }); + + it("reads tier_labels out of a config stored as a JSON string", () => { + const stored = JSON.stringify({ tier_labels: { SIMPLE: "Cheap" } }); + render(); + + expect(screen.getByText("Cheap 75%")).toBeInTheDocument(); + }); + + it("uses canonical names when the router is not in the deployment list", () => { + render(); + + expect(screen.getByText("Simple 75%")).toBeInTheDocument(); + expect(screen.getByText("Complex 25%")).toBeInTheDocument(); + }); + + it("lists each tier's assigned models below its name and share", () => { + render( + , + ); + + expect(screen.getByText("gpt-4o-mini")).toBeInTheDocument(); + expect(screen.getByText("gpt-4o, claude-3-opus")).toBeInTheDocument(); + }); + + it("widens a bare string tier (pinned single model) into its one-model list", () => { + render(); + + expect(screen.getByText("gpt-4o-mini")).toBeInTheDocument(); + }); + + it("omits the model line for a tier with no configured models", () => { + render(); + + expect(screen.getByText("Simple 75%")).toBeInTheDocument(); + }); + + it("shows no models for a quality router's numeric tier, which has no per-tier model list", () => { + render( + , + ); + + expect(screen.getByText("3 75%")).toBeInTheDocument(); + expect(screen.getByText("1 25%")).toBeInTheDocument(); + expect(screen.queryByText("gpt-4o")).not.toBeInTheDocument(); + }); + + it("ignores a same-named deployment of a different router type", () => { + const qualityDeployment = { + model_name: "claude-auto", + litellm_params: { model: "auto_router/claude-auto", quality_router_config: { available_models: ["gpt-4o"] } }, + }; + + render( + , + ); + + expect(screen.getByText("Simple 75%")).toBeInTheDocument(); + expect(screen.queryByText("gpt-4o")).not.toBeInTheDocument(); + }); + + it("renders nothing for the all-routers view, which carries no router identity", () => { + const { container } = render( + , + ); + + expect(container).toBeEmptyDOMElement(); + }); + + it("renders nothing when the router recorded no tiers", () => { + const { container } = render(); + + expect(container).toBeEmptyDOMElement(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.tsx new file mode 100644 index 00000000000..5b9b8563baa --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/TierTurnsChart.tsx @@ -0,0 +1,148 @@ +"use client"; + +import React from "react"; + +import type { AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels"; +import { hydrateTierLabels } from "@/components/add_model/build_complexity_router_config"; +import { + TIER_KEYS, + effectiveTierLabel, + type ComplexityTierLabels, + type ComplexityTiers, +} from "@/components/add_model/ComplexityRouterConfig"; +import { normalizeTierModels } from "@/components/add_model/complexity_router_tiers"; +import { chartColorValue, DonutChart, type ChartColor } from "@/components/shared/charts"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; + +import { viewGroup, type BenchmarkView } from "./autoRouterBenchmarks"; + +const safeParse = (value: string): unknown => { + try { + return JSON.parse(value); + } catch { + return null; + } +}; + +const asRecord = (value: unknown): Record => { + const parsed: unknown = typeof value === "string" ? safeParse(value) : value; + return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed) + ? (parsed as Record) + : {}; +}; + +const isComplexityTier = (tier: string): tier is keyof ComplexityTiers => + (TIER_KEYS as readonly string[]).includes(tier); + +export const tierDisplayLabel = (tier: string, tierLabels: ComplexityTierLabels | undefined): string => + isComplexityTier(tier) ? effectiveTierLabel(tier, tierLabels) : tier; + +const CONFIG_KEY_BY_ROUTER_TYPE: Record> = { + complexity: "complexity_router_config", + quality: "quality_router_config", + auto_router: "auto_router_config", + adaptive: "adaptive_router_config", +}; + +const deploymentFor = ( + routerName: string, + routerType: string, + autoRouters: readonly AutoRouterDeployment[], +): AutoRouterDeployment | undefined => { + const configKey = CONFIG_KEY_BY_ROUTER_TYPE[routerType]; + if (!configKey) return undefined; + return autoRouters.find((d) => d.model_name === routerName && d.litellm_params?.[configKey]); +}; + +const tierLabelsFor = ( + routerName: string, + routerType: string, + autoRouters: readonly AutoRouterDeployment[], +): ComplexityTierLabels | undefined => { + const deployment = deploymentFor(routerName, routerType, autoRouters); + if (!deployment) return undefined; + const config = asRecord(deployment.litellm_params?.complexity_router_config); + return hydrateTierLabels(config.tier_labels); +}; + +const tierModelsFor = ( + tier: string, + routerName: string, + routerType: string, + autoRouters: readonly AutoRouterDeployment[], +): string[] => { + if (!isComplexityTier(tier)) return []; + const deployment = deploymentFor(routerName, routerType, autoRouters); + if (!deployment) return []; + const config = asRecord(deployment.litellm_params?.complexity_router_config); + const tiers = asRecord(config.tiers); + return normalizeTierModels(tiers[tier]); +}; + +interface TierTurnsChartProps { + view: BenchmarkView; + autoRouters: readonly AutoRouterDeployment[]; +} + +const TIER_DONUT_COLORS: readonly ChartColor[] = ["#c7d2fe", "#1e293b", "#d4b483", "#87a878"]; + +const TierTurnsChart: React.FC = ({ view, autoRouters }) => { + const group = viewGroup(view); + const entries = Object.entries(group?.tier_turns ?? {}).filter(([, turns]) => turns > 0); + if (!group || entries.length === 0) return null; + + const tierLabels = tierLabelsFor(group.router_name, group.router_type, autoRouters); + const total = entries.reduce((sum, [, turns]) => sum + turns, 0); + const slices = entries.map(([tier, turns]) => ({ + tier: tierDisplayLabel(tier, tierLabels), + turns, + models: tierModelsFor(tier, group.router_name, group.router_type, autoRouters), + })); + const colors = slices.map((_, idx) => TIER_DONUT_COLORS[idx % TIER_DONUT_COLORS.length]); + + return ( + + + Routing by tier +

+ Turns each tier served. Turns the classifier sent to the default model belong to no tier and are not counted + here, so this can total less than the router's turns. +

+ + +
+ value.toLocaleString()} + showLabel + label={`${total.toLocaleString()} total turns`} + /> +
    + {slices.map((slice, idx) => ( +
  • + +
    +

    + {slice.tier} {Math.round((100 * slice.turns) / total).toLocaleString()}% +

    + {slice.models.length > 0 && ( +

    {slice.models.join(", ")}

    + )} +
    +
    • + ))} +
+
+ + + ); +}; + +export default TierTurnsChart; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.ts index 00793548278..2e1031ce701 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.ts @@ -24,9 +24,12 @@ export const windowFor = (range: BenchmarkWindow, now: Date): { start_date: stri export interface BenchmarkView { label: string; - stats: AutoRouterBenchmarkTotals; + stats: AutoRouterBenchmarkTotals | AutoRouterBenchmarkGroup; } +export const viewGroup = (view: BenchmarkView): AutoRouterBenchmarkGroup | null => + "router_name" in view.stats ? view.stats : null; + export const groupKey = (group: AutoRouterBenchmarkGroup): string => `${group.router_name} ${group.router_type}`; export const groupLabel = (group: AutoRouterBenchmarkGroup, groups: readonly AutoRouterBenchmarkGroup[]): string => { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 966d2162a62..8989e098f54 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21403,7 +21403,7 @@ export interface components { spend: number; /** * Tier Turns - * @description Turns per tier, keyed by the tier name the routing decision recorded at request time (never re-derived at read time, since the tier-to-model mapping is mutable config). Tier names are scoped to this group's router_type and are not comparable across types: a complexity router reports 'simple'/'medium'/'complex'/'reasoning', a quality router reports its numeric quality tier, and an adaptive router records no tier at all. Turns no tier served (the classifier fell back to default_model) are absent rather than pooled under a sentinel key, so the values may sum to less than turns + * @description Turns per tier, keyed by the tier name the routing decision recorded at request time (never re-derived at read time, since the tier-to-model mapping is mutable config). Tier names are scoped to this group's router_type and are not comparable across types: a complexity router reports 'SIMPLE'/'MEDIUM'/'COMPLEX'/'REASONING', a quality router reports its numeric quality tier, and an adaptive router records no tier at all. Turns no tier served (the classifier fell back to default_model) are absent rather than pooled under a sentinel key, so the values may sum to less than turns */ tier_turns?: { [key: string]: number;