fix(vertex): bill routed generateContent at the deployment's custom pricing

The Vertex pass-through logging handler priced generateContent calls,
including every streamed /v1beta/models/{model}:streamGenerateContent
call, with completion_cost(model, custom_llm_provider) alone, so it read
the cost map's shared entry for the backend model, from which the router
strips custom pricing on purpose. A deployment with its own
input_cost_per_token and output_cost_per_token was billed at the bundled
price. Pass custom_pricing and router_model_id, as the Anthropic handler
already does, so the deployment's registered price wins.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Andrew 2026-10-07 10:16:01 -04:00
parent 736ff14f11
commit 8319c2d035
2 changed files with 53 additions and 0 deletions

View file

@ -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,
)
@ -769,11 +770,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

View file

@ -127,3 +127,50 @@ 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)