fix: record vertex_location on the generate_content path so cost uses the configured region

The Vertex location decides the regional pricing uplift. The completion path
sets it on litellm_params and async_anthropic_messages_handler records it on
the logging object, but generate_content passed only litellm_call_id, so the
cost calculator never saw it.

The result is that a model configured `vertex_location: global` is priced as
us-central1 on POST /v1beta/models/{model}:generateContent and costs 10% more
than the same model over /v1/chat/completions. Two requests differing only in
the URL are billed differently, and spend cannot be reconciled against the
cloud bill.

Uses the same VertexBase.explicit_vertex_ai_location helper the anthropic
messages path already uses, so an unset location records nothing and the cost
calculator keeps its own fallback.

Fixes #40692

Signed-off-by: Ankit Jha <jhaankit373@gmail.com>
This commit is contained in:
Ankit Jha 2026-09-11 15:49:20 +05:30
parent 362033cb2a
commit b00c00e5f8
No known key found for this signature in database
2 changed files with 85 additions and 0 deletions

View file

@ -2,6 +2,7 @@ import asyncio
import contextvars
from collections.abc import Iterator
from functools import partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final
import httpx
@ -192,12 +193,26 @@ class GenerateContentHelper:
if litellm_logging_obj is None:
raise ValueError("litellm_logging_obj is required, but got None")
# The configured Vertex location decides the regional pricing uplift. The
# completion and anthropic_messages paths both record it here; without it
# the cost calculator falls back to a default region and prices a `global`
# model as us-central1, so the same model costs 10% more on this route.
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
explicit_vertex_location: Final = VertexBase.explicit_vertex_ai_location(
MappingProxyType(litellm_params.model_dump(exclude_none=True))
)
vertex_location_params: Final = (
{"vertex_location": explicit_vertex_location} if explicit_vertex_location else {}
)
litellm_logging_obj.update_from_kwargs(
kwargs=kwargs,
model=model,
optional_params=dict(generate_content_config_dict),
litellm_params={
"litellm_call_id": litellm_call_id,
**vertex_location_params,
},
custom_llm_provider=custom_llm_provider,
)

View file

@ -259,3 +259,73 @@ async def test_native_fields_forwarded_on_async_stream():
body = mock_post.call_args.kwargs["json"]
assert body["safetySettings"] == safety_settings
assert "safetySettings" not in body.get("generationConfig", {})
def _generate_content_logging_obj(call_id: str):
from datetime import datetime
from litellm.litellm_core_utils.litellm_logging import Logging
return Logging(
model="gemini-2.0-flash",
messages=[],
stream=False,
call_type="generate_content",
start_time=datetime.now(),
litellm_call_id=call_id,
function_id=call_id,
)
@pytest.mark.parametrize("configured_location", ["global", "us-east5"])
def test_vertex_location_recorded_for_cost_calculation(configured_location):
"""
Regression for https://github.com/BerriAI/litellm/issues/40692
The Vertex location decides the regional pricing uplift. The completion and
anthropic_messages paths record it on the logging object so the cost
calculator can price against the configured region. generate_content did
not, so a model configured `vertex_location: global` was priced as
us-central1 and cost 10% more than the same model over /v1/chat/completions.
"""
from litellm.google_genai.main import GenerateContentHelper
logging_obj = _generate_content_logging_obj("vertex-location-test")
GenerateContentHelper.setup_generate_content_call(
model="vertex_ai/gemini-2.0-flash",
contents=[{"role": "user", "parts": [{"text": "say ok"}]}],
custom_llm_provider="vertex_ai",
litellm_logging_obj=logging_obj,
litellm_call_id="vertex-location-test",
vertex_project="test-project",
vertex_location=configured_location,
)
recorded = logging_obj.model_call_details["litellm_params"]
assert recorded.get("vertex_location") == configured_location, (
"the configured vertex_location must reach the cost calculator; "
f"got {recorded.get('vertex_location')!r}"
)
def test_vertex_location_absent_when_not_configured():
"""
Nothing configured means nothing recorded, so the cost calculator keeps its
own fallback rather than being handed an empty value here.
"""
from litellm.google_genai.main import GenerateContentHelper
logging_obj = _generate_content_logging_obj("vertex-location-unset")
GenerateContentHelper.setup_generate_content_call(
model="gemini/gemini-2.0-flash",
contents=[{"role": "user", "parts": [{"text": "say ok"}]}],
custom_llm_provider="gemini",
litellm_logging_obj=logging_obj,
litellm_call_id="vertex-location-unset",
api_key="test-key",
)
recorded = logging_obj.model_call_details["litellm_params"]
assert "vertex_location" not in recorded