Merge 6ce2248bf3 into 82ca10cc2d

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

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)