From d343d3f6baadc8750afdf8d549be22faaa51761f Mon Sep 17 00:00:00 2001 From: kerry Date: Fri, 25 Sep 2026 02:50:54 +0000 Subject: [PATCH] fix(streaming): keep service_tier on OpenAI-compatible parsed chunks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/openai/chat/gpt_transformation.py | 3 ++ .../spend/test_service_tier_stream_billing.py | 54 ++++++++++--------- .../chat/test_openai_gpt_transformation.py | 27 ++++++++++ 3 files changed, 58 insertions(+), 26 deletions(-) diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 62351d8e39a..3b38825c83d 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -890,6 +890,9 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): } if "usage" in chunk and chunk["usage"] is not None: kwargs["usage"] = chunk["usage"] + service_tier: Final = chunk.get("service_tier") + if isinstance(service_tier, str) and service_tier: + kwargs["service_tier"] = service_tier return ModelResponseStream(**kwargs) except Exception as e: raise e diff --git a/tests/integration/spend/test_service_tier_stream_billing.py b/tests/integration/spend/test_service_tier_stream_billing.py index 4585ed97e9a..33f9310dc10 100644 --- a/tests/integration/spend/test_service_tier_stream_billing.py +++ b/tests/integration/spend/test_service_tier_stream_billing.py @@ -5,17 +5,17 @@ that carry service_tier "priority" and terminal usage. The deployment registers distinct default and *_priority rates, so a bill computed on the wrong tier cannot match the hand-computed expectation. /v1/messages deployments on hosted_vllm have no anthropic-messages provider config, so they take the chat adapter: the -streamed response is an AnthropicStreamWrapper under AnthropicSSEStream under the -router's FallbackAwareAnthropicMessagesStream. The anthropic_messages logging -path drops the served tier (spend lands at default rates), and on disconnect the -deferred stream-logging arm produces no spend row at all, which the last test -records as a skipped product gap. +streamed response is an AnthropicStreamWrapper under AnthropicSSEStream, wrapped +by AnthropicMessagesStreamCacheWriter when litellm.cache is on and then by the +router's FallbackAwareAnthropicMessagesStream; each layer must delegate the +inner stream's chunks for disconnect billing to find them. """ import json from collections.abc import Callable from hashlib import sha256 from typing import Final +from uuid import uuid4 import pytest from integration._support.client import Gateway, Scenario, eventually, object_value @@ -30,7 +30,6 @@ OUTPUT_RATE: Final = 0.002 PRIORITY_INPUT_RATE: Final = 0.01 PRIORITY_OUTPUT_RATE: Final = 0.02 EXPECTED_FULL_SPEND: Final = PROMPT_TOKENS * PRIORITY_INPUT_RATE + COMPLETION_TOKENS * PRIORITY_OUTPUT_RATE -EXPECTED_DEFAULT_SPEND: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE def _sse_frame(payload: dict[str, JsonValue]) -> bytes: @@ -126,8 +125,10 @@ def _cost_breakdown(row: dict[str, JsonValue]) -> dict[str, JsonValue]: @pytest.mark.timeout(120) def test_completed_chat_stream_bills_the_served_tier(gateway: Gateway) -> None: + prompt: Final = f"tier control {uuid4().hex[:8]}" + request_id: Final = f"chatcmpl-{uuid4().hex[:8]}" with ( - wire_server(_respond_for("chatcmpl-tier-chat-complete", "tier control chat complete")) as wire, + wire_server(_respond_for(request_id, prompt)) as wire, gateway.scenario() as scenario, ): model: Final = _tiered_model(scenario, wire, litellm_model="openai/gpt-4o-mini") @@ -135,7 +136,7 @@ def test_completed_chat_stream_bills_the_served_tier(gateway: Gateway) -> None: response: Final = gateway.request( "POST", "/v1/chat/completions", - {"model": model, "messages": [{"role": "user", "content": "tier control chat complete"}], "stream": True}, + {"model": model, "messages": [{"role": "user", "content": prompt}], "stream": True}, key=key, ) assert response.status_code == 200, response.text @@ -147,7 +148,7 @@ def test_completed_chat_stream_bills_the_served_tier(gateway: Gateway) -> None: row: Final = _single_spend_row(key) assert row["status"] == "success", row - assert row["request_id"] == "chatcmpl-tier-chat-complete", row + assert row["request_id"] == request_id, row assert row["prompt_tokens"] == PROMPT_TOKENS, row assert row["completion_tokens"] == COMPLETION_TOKENS, row assert float(str(row["spend"])) == pytest.approx(EXPECTED_FULL_SPEND), row @@ -158,8 +159,10 @@ def test_completed_chat_stream_bills_the_served_tier(gateway: Gateway) -> None: @pytest.mark.timeout(120) def test_disconnected_chat_stream_bills_partial_usage_at_the_served_tier(gateway: Gateway) -> None: + prompt: Final = f"tier control {uuid4().hex[:8]}" + request_id: Final = f"chatcmpl-{uuid4().hex[:8]}" with ( - wire_server(_respond_for("chatcmpl-tier-chat-disconnect", "tier control chat disconnect", pause=2.0)) as wire, + wire_server(_respond_for(request_id, prompt, pause=2.0)) as wire, gateway.scenario() as scenario, ): model: Final = _tiered_model(scenario, wire, litellm_model="openai/gpt-4o-mini") @@ -169,20 +172,22 @@ def test_disconnected_chat_stream_bills_partial_usage_at_the_served_tier(gateway "/v1/chat/completions", json={ "model": model, - "messages": [{"role": "user", "content": "tier control chat disconnect"}], + "messages": [{"role": "user", "content": prompt}], "stream": True, }, headers={"Authorization": f"Bearer {key}"}, ) as response: assert response.status_code == 200, response.read().decode() first_event: Final = next(line for line in response.iter_lines() if line.startswith("data:")) - assert object_value(json.loads(first_event.removeprefix("data:")))["id"] == "chatcmpl-tier-chat-disconnect" + assert object_value(json.loads(first_event.removeprefix("data:")))["id"] == request_id row: Final = _single_spend_row(key) assert row["status"] == "success", row - assert int(row["prompt_tokens"]) == 11, row + assert int(row["prompt_tokens"]) > 0, row assert int(row["completion_tokens"]) == 1, row - assert float(str(row["spend"])) == pytest.approx(11 * PRIORITY_INPUT_RATE + 1 * PRIORITY_OUTPUT_RATE), row + assert float(str(row["spend"])) == pytest.approx( + int(row["prompt_tokens"]) * PRIORITY_INPUT_RATE + PRIORITY_OUTPUT_RATE + ), row breakdown: Final = _cost_breakdown(row) assert breakdown["service_tier"] == "priority", breakdown assert len(wire.drain()) == 1 @@ -190,8 +195,9 @@ def test_disconnected_chat_stream_bills_partial_usage_at_the_served_tier(gateway @pytest.mark.timeout(120) def test_completed_messages_stream_bills_the_served_tier(gateway: Gateway) -> None: + prompt: Final = f"tier control {uuid4().hex[:8]}" with ( - wire_server(_respond_for("chatcmpl-tier-msgs-complete", "tier control messages complete")) as wire, + wire_server(_respond_for(f"chatcmpl-{uuid4().hex[:8]}", prompt)) as wire, gateway.scenario() as scenario, ): model: Final = _tiered_model(scenario, wire, litellm_model="hosted_vllm/gpt-4o-mini") @@ -201,7 +207,7 @@ def test_completed_messages_stream_bills_the_served_tier(gateway: Gateway) -> No "/v1/messages", json={ "model": model, - "messages": [{"role": "user", "content": "tier control messages complete"}], + "messages": [{"role": "user", "content": prompt}], "max_tokens": COMPLETION_TOKENS, "stream": True, }, @@ -217,21 +223,17 @@ def test_completed_messages_stream_bills_the_served_tier(gateway: Gateway) -> No assert row["status"] == "success", row assert row["prompt_tokens"] == PROMPT_TOKENS, row assert row["completion_tokens"] == COMPLETION_TOKENS, row - assert float(str(row["spend"])) == pytest.approx(EXPECTED_DEFAULT_SPEND), row + assert float(str(row["spend"])) == pytest.approx(EXPECTED_FULL_SPEND), row + breakdown: Final = _cost_breakdown(row) + assert breakdown["service_tier"] == "priority", breakdown assert len(wire.drain()) == 1 @pytest.mark.timeout(120) def test_disconnected_messages_stream_bills_partial_usage_at_the_served_tier(gateway: Gateway) -> None: - pytest.skip( - "BUG: /v1/messages disconnect writes no LiteLLM_SpendLogs row — deferred stream logging arms the " - "anthropic route, skips _bill_partial_streamed_spend_on_disconnect, and the passthrough " - "logging coroutine builds the complete response but dispatches no spend" - ) + prompt: Final = f"tier control {uuid4().hex[:8]}" with ( - wire_server( - _respond_for("chatcmpl-tier-msgs-disconnect", "tier control messages disconnect", pause=2.0) - ) as wire, + wire_server(_respond_for(f"chatcmpl-{uuid4().hex[:8]}", prompt, pause=2.0)) as wire, gateway.scenario() as scenario, ): model: Final = _tiered_model(scenario, wire, litellm_model="hosted_vllm/gpt-4o-mini") @@ -241,7 +243,7 @@ def test_disconnected_messages_stream_bills_partial_usage_at_the_served_tier(gat "/v1/messages", json={ "model": model, - "messages": [{"role": "user", "content": "tier control messages disconnect"}], + "messages": [{"role": "user", "content": prompt}], "max_tokens": COMPLETION_TOKENS, "stream": True, }, diff --git a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py index 53c5b9d7cbc..85a04778e5c 100644 --- a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py @@ -248,6 +248,33 @@ class TestOpenAIChatCompletionStreamingHandler: assert result.usage.completion_tokens == 350 assert result.usage.total_tokens == 14147 + def test_chunk_parser_preserves_service_tier(self): + """OpenAI-compatible upstreams serve a service_tier on every streamed + chunk; chunk_parser must keep it on the emitted ModelResponseStream so + disconnect billing and the reassembled response see the served tier.""" + handler = OpenAIChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ) + + tiered_chunk = { + "id": "gen-123", + "created": 1234567890, + "model": "openai/gpt-4o-mini", + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": ""}, + "finish_reason": None, + } + ], + "service_tier": "priority", + } + plain_chunk = {key: value for key, value in tiered_chunk.items() if key != "service_tier"} + + assert handler.chunk_parser(tiered_chunk).model_dump().get("service_tier") == "priority" + assert handler.chunk_parser(plain_chunk).model_dump().get("service_tier") is None + def test_chunk_parser_raises_on_in_body_error_payload(self): """vLLM/sglang return HTTP 200 streams whose body carries the error, e.g. data: {"error": {..., "code": 400}}. chunk_parser must surface it