From 6ce2248bf3382b35efbd37b6fdd30361836cc0f8 Mon Sep 17 00:00:00 2001 From: Andrew Date: Wed, 7 Oct 2026 15:02:29 -0400 Subject: [PATCH] fix(vertex): leave a streamed generateContent response unpriced for the logger The Vertex pass-through handler assembled a streamed generateContent response with stream_chunk_builder but no logging object, so the builder priced the assembled response by its model name alone and stamped that cost into its hidden params; the logger then reused the stamped cost as already calculated, overriding the deployment's custom pricing. Pass the logging object, as the Anthropic handler does, so the builder leaves the pricing to the logger. Co-Authored-By: Claude Opus 5.5 --- .../vertex_passthrough_logging_handler.py | 9 +++- ...test_vertex_passthrough_logging_handler.py | 48 +++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index cc8145bd162..3b27414e2a8 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -659,7 +659,14 @@ class VertexPassthroughLoggingHandler: continue all_openai_chunks.append(parsed_chunk) - complete_streaming_response: Final = litellm.stream_chunk_builder(chunks=all_openai_chunks) + # With the logging object the builder leaves pricing to the logger, + # which honours the deployment's custom pricing; without it the + # builder stamps a cost priced by model name alone, which the logger + # then reuses as already calculated. + complete_streaming_response: Final = litellm.stream_chunk_builder( + chunks=all_openai_chunks, + logging_obj=litellm_logging_obj, + ) return complete_streaming_response diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py index dd271aab380..1f98549071d 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py @@ -174,3 +174,51 @@ def test_generate_content_cost_uses_the_deployments_custom_pricing(monkeypatch: ) assert kwargs["response_cost"] == pytest.approx(1.5 + 7.5) + + +def test_assembled_stream_is_priced_at_the_deployments_custom_pricing(monkeypatch: pytest.MonkeyPatch): + """The logger prices a streamed generateContent call from the response the + handler assembles; that response must not arrive carrying a cost priced + by model name alone, which the logger would reuse as already calculated.""" + import litellm + + deployment_id: Final = "vertex-gemini-with-custom-price" + custom_price: Final = {"input_cost_per_token": 1.5e-6, "output_cost_per_token": 7.5e-6} + monkeypatch.setitem( + litellm.model_cost, + deployment_id, + {**custom_price, "litellm_provider": "vertex_ai-language-models", "mode": "chat"}, + ) + logging_obj: Final = Logging( + model="gemini-2.5-flash", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime(2026, 1, 1), + litellm_call_id="call-1", + function_id="fn-1", + ) + logging_obj.update_environment_variables( + litellm_params={"litellm_metadata": {"model_info": {"id": deployment_id, **custom_price}}}, + optional_params={}, + model="gemini-2.5-flash", + custom_llm_provider="vertex_ai", + ) + last_chunk: Final = { + "candidates": [{"content": {"parts": [{"text": "Hi."}], "role": "model"}, "finishReason": "STOP", "index": 0}], + "usageMetadata": { + "promptTokenCount": 1_000_000, + "candidatesTokenCount": 1_000_000, + "totalTokenCount": 2_000_000, + }, + "modelVersion": "gemini-2.5-flash", + } + + response: Final = VertexPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=[f"data: {json.dumps(last_chunk)}"], + litellm_logging_obj=logging_obj, + model="gemini-2.5-flash", + url_route="/v1/projects/p/locations/global/publishers/google/models/gemini-2.5-flash:streamGenerateContent", + ) + + assert logging_obj.response_cost_calculator(result=response) == pytest.approx(1.5 + 7.5)