mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
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:
parent
2b3070890a
commit
4282c437be
2 changed files with 63 additions and 2 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue