mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-23 05:34:41 +00:00
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:
parent
736ff14f11
commit
8319c2d035
2 changed files with 53 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue