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 c3af6f534c2..8335d210671 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 @@ -12,6 +12,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import VERTEX_BATCH_PREDICTION_JOBS_ROUTE from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( InteractionsUsageObjectTransformation, ) @@ -668,7 +669,14 @@ class VertexPassthroughLoggingHandler: return None all_openai_chunks: Final = [parsed_chunk for parsed_chunk in parsed_chunks if parsed_chunk is not None] - 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 @@ -779,11 +787,16 @@ class VertexPassthroughLoggingHandler: """ + # A routed call carries the deployment it was routed to; price it at + # that deployment's custom pricing, as the Anthropic handler does, + # rather than at the cost map's shared entry for the backend model. response_cost: Final = litellm.completion_cost( completion_response=litellm_model_response, model=model, custom_llm_provider=custom_llm_provider, vertex_location=vertex_location, + custom_pricing=use_custom_pricing_for_model(litellm_params=getattr(logging_obj, "litellm_params", None)), + router_model_id=logging_obj.get_router_model_id(), ) kwargs["response_cost"] = response_cost 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 cc1246f4ede..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 @@ -127,3 +127,98 @@ def test_build_complete_streaming_response_assembles_stream_generate_content_sse assert response.choices[0].message.content == "Hello there!" assert response.choices[0].finish_reason == "stop" assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (10, 8, 18) + + +def test_generate_content_cost_uses_the_deployments_custom_pricing(monkeypatch: pytest.MonkeyPatch): + """A routed generateContent call (streamed calls are logged through this + path) is billed at the price set on the deployment it was routed to, not + the cost map's shared price for the backend model.""" + import litellm + from litellm.types.utils import Usage + + 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={"metadata": {"model_info": {"id": deployment_id, **custom_price}}}, + optional_params={}, + model="gemini-2.5-flash", + ) + response: Final = ModelResponse( + model="gemini-2.5-flash", + usage=Usage(prompt_tokens=1_000_000, completion_tokens=1_000_000, total_tokens=2_000_000), + ) + + kwargs: Final = VertexPassthroughLoggingHandler._create_vertex_response_logging_payload_for_generate_content( + litellm_model_response=response, + model="gemini-2.5-flash", + kwargs={}, + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), + logging_obj=logging_obj, + custom_llm_provider="vertex_ai", + vertex_location=None, + ) + + 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)