fix(cost): propagate vertex_location as region_name for vertex_ai region-based pricing

Region-based cost tracking worked for bedrock and azure but never for vertex_ai.
cost_calculator already looks up a region-scoped "{provider}/{region}/{model}"
price key when _hidden_params["region_name"] is set, but completion() only ever
populated region_name from aws_region_name (bedrock). For vertex_ai the region is
supplied via vertex_location and was dropped, so every vertex region billed at the
same model-name price even though Google charges differently for the global vs
regional endpoints.

Populate region_name from vertex_location as well as aws_region_name. This unblocks
region-scoped vertex_ai pricing the same way azure/eu/... and bedrock region entries
already work; region-scoped vertex_ai/{region}/... entries can then be added to the
price map as their rates are sourced.

Fixes #34393
This commit is contained in:
Anuj7411 2026-07-31 17:25:54 +05:30
parent 2b3070890a
commit 4282c437be
2 changed files with 63 additions and 2 deletions

View file

@ -5172,8 +5172,8 @@ def completion( # type: ignore
if model_response is not None and hasattr(model_response, "_hidden_params"):
model_response._hidden_params["custom_llm_provider"] = custom_llm_provider
model_response._hidden_params["region_name"] = kwargs.get(
"aws_region_name", None
) # support region-based pricing for bedrock
"aws_region_name"
) or kwargs.get("vertex_location")
### TIMEOUT LOGIC ###
timeout = CompletionTimeout.resolve(

View file

@ -5553,3 +5553,64 @@ def test_accumulated_json_skips_non_dict_leading_value():
assert len(out) == 1
assert out[0].choices[0].delta.content == "a"
@pytest.mark.parametrize(
"vertex_location, expected_region_name",
[
("europe-west1", "europe-west1"),
("global", "global"),
(None, None),
],
)
def test_vertex_region_name_propagated_from_vertex_location(
vertex_location, expected_region_name
):
"""
Regression for #34393: region-based cost tracking worked for bedrock (aws_region_name)
and azure but never for vertex_ai, because completion() only copied aws_region_name into
_hidden_params["region_name"]. vertex_location was dropped, so cost_calculator could never
build the region-scoped "vertex_ai/{region}/{model}" price key and every vertex region was
billed at the same model-name price.
completion() must now propagate vertex_location as region_name (and leave it None when no
location is supplied, so a request that opts out does not accidentally match a region key).
"""
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
response_body = {
"candidates": [
{
"content": {"parts": [{"text": "hi"}], "role": "model"},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 5,
"totalTokenCount": 15,
},
"modelVersion": "gemini-2.5-flash",
}
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = response_body
completion_kwargs = dict(
model="vertex_ai/gemini-2.5-flash",
messages=[{"role": "user", "content": "hi"}],
vertex_project="fake-project",
vertex_credentials='{"type": "service_account"}',
client=client,
)
if vertex_location is not None:
completion_kwargs["vertex_location"] = vertex_location
with patch(
"litellm.llms.vertex_ai.vertex_llm_base.VertexBase._ensure_access_token",
return_value=("fake-token", "fake-project"),
):
with patch.object(client, "post", return_value=mock_response):
response = completion(**completion_kwargs)
assert response._hidden_params.get("custom_llm_provider") == "vertex_ai"
assert response._hidden_params.get("region_name") == expected_region_name