From 4282c437beda6129eddeac69c96b7ff429b5b35d Mon Sep 17 00:00:00 2001
From: Anuj7411 <ojhaanuj224@gmail.com>
Date: Fri, 31 Jul 2026 17:25:54 +0530
Subject: [PATCH] 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
---
 litellm/main.py                               |  4 +-
 ...test_vertex_and_google_ai_studio_gemini.py | 61 +++++++++++++++++++
 2 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/litellm/main.py b/litellm/main.py
index acdec7385da..4ae70f5a8f4 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -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(
diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
index 51cc2857252..6eeca04be4a 100644
--- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
+++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
@@ -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