From 366c5f38e72446006c1160ad2a151fc269a09de5 Mon Sep 17 00:00:00 2001 From: leilei3167 Date: Tue, 29 Sep 2026 12:36:38 +0000 Subject: [PATCH] fix(streaming): drop zero cost on Vertex stream chunks Vertex chunks reach chunk_creator with usage.cost=0 before cost resolution. Treat that non-positive cost as absent so token pricing runs for the stream. --- .../litellm_core_utils/streaming_handler.py | 22 ++++- .../test_streaming_handler.py | 88 +++++++++++++++++++ 2 files changed, 109 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 61d586251cd..22aaedc95af 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1658,7 +1658,27 @@ class CustomStreamWrapper: self.tool_call = True if hasattr(chunk, "usage") and chunk.usage is not None: - model_response.usage = chunk.usage + attached_usage = chunk.usage + if self.custom_llm_provider == "vertex_ai": + # The vertex_ai branch returns before cost propagation, so a + # leftover usage.cost=0 would otherwise be copied here and + # treated as a provider total. Resolve it the same way as + # stream assembly: non-positive is absent, token pricing runs. + raw_cost = ( + attached_usage.get("cost") + if isinstance(attached_usage, dict) + else getattr(attached_usage, "cost", None) + ) + if raw_cost is not None and CustomStreamWrapper._resolve_provider_reported_cost(raw_cost) is None: + if isinstance(attached_usage, dict): + attached_usage = {key: value for key, value in attached_usage.items() if key != "cost"} + elif hasattr(attached_usage, "model_copy"): + attached_usage = attached_usage.model_copy(update={"cost": None}) + try: + del attached_usage.cost + except (AttributeError, TypeError): + pass + model_response.usage = attached_usage ## RETURN ARG result: Final = self.return_processed_chunk_logic( diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 7ed784e6d89..071c97cd8bc 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -2079,6 +2079,94 @@ def test_stream_spend_prices_vertex_anthropic_without_cache_read(): assert stream_spend >= token_total - 1e-12 +def test_vertex_chunk_creator_drops_zero_cost_and_prices_tokens(): + """Vertex stream chunks can carry usage.cost=0 with real token counts. + + chunk_creator must not keep that 0 as a provider total. Token pricing on + the assembled stream still has to be a positive spend. + """ + logging_obj = Logging( + model="claude-opus-5", + messages=[{"role": "user", "content": "count to five"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="vertex-zero-cost-chunk", + function_id="1245", + ) + logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai" + logging_obj.optional_params = {} + wrapper = CustomStreamWrapper( + completion_stream=None, + model="claude-opus-5", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai", + ) + source_chunks = _vertex_anthropic_stream_chunks(cache_read=44616, prompt_tokens=70961, completion_tokens=4096) + zero_cost_usage = source_chunks[-1].usage.model_copy(update={"cost": 0}) + zero_cost_chunk = ModelResponseStream( + id="chatcmpl-vertex-zero-cost", + created=1745513207, + model="claude-opus-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="Hi", role="assistant"), + ) + ], + usage=zero_cost_usage, + ) + + processed = wrapper.chunk_creator(chunk=zero_cost_chunk) + + assert processed is not None + assert processed.usage.prompt_tokens == 70961 + assert processed.usage.completion_tokens == 4096 + assert processed.usage.cache_read_input_tokens == 44616 + assert getattr(processed.usage, "cost", None) in (None,) + + priced_chunk = ModelResponseStream( + id="chatcmpl-vertex-priced", + created=1745513208, + model="claude-opus-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="Hi", role="assistant"), + ) + ], + usage=zero_cost_usage.model_copy(update={"cost": 0.01}), + ) + priced = wrapper.chunk_creator(chunk=priced_chunk) + assert priced is not None + assert priced.usage.cost == 0.01 + + complete_response = litellm.stream_chunk_builder( + chunks=[processed], + messages=[{"role": "user", "content": "count to five"}], + logging_obj=logging_obj, + ) + assert complete_response is not None + token_prompt, token_completion = litellm.cost_per_token( + model="claude-opus-5", + custom_llm_provider="vertex_ai", + usage_object=complete_response.usage, + ) + token_total = token_prompt + token_completion + assert token_total > 0 + stream_spend = _stream_spend_via_logging( + complete_response, + model="claude-opus-5", + custom_llm_provider="vertex_ai", + ) + assert stream_spend > 0 + assert stream_spend >= token_total - 1e-12 + + def test_handle_special_delta_attributes( initialized_custom_stream_wrapper: CustomStreamWrapper, ):