support vertex_ai global endpoints for chat

Author: sorenmat

litellm/llms/vertex_ai/common_utils.py

@@ -84,9 +84,15 @@ def _get_vertex_url(
         endpoint = "generateContent"
         if stream is True:
             endpoint = "streamGenerateContent"
-            url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}?alt=sse"
+            if vertex_location== "global":
+                url = f"https://aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/global/publishers/google/models/{model}:{endpoint}?alt=sse"
+            else:
+                url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}?alt=sse"
         else:
-            url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}"
+            if vertex_location == "global":
+                url = f"https://aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/global/publishers/google/models/{model}:{endpoint}"
+            else:
+                url = f"https://{vertex_location}-aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}"
 
         # if model is only numeric chars then it's a fine tuned gemini model
         # model = 4965075652664360960

tests/litellm/llms/vertex_ai/test_vertex_ai_common_utils.py

@@ -17,6 +17,7 @@
     get_vertex_location_from_url,
     get_vertex_project_id_from_url,
     set_schema_property_ordering,
+    _get_vertex_url
 )
 
 
@@ -292,3 +293,38 @@ def test_process_items_basic():
     }
     process_items(schema)
     assert schema["properties"]["nested"]["items"] == {"type": "object"}
+
+def test_get_vertex_url_global_region(stream):
+    """
+    Test _get_vertex_url when vertex_location is 'global' for chat mode.
+    """
+    mode = "chat"
+    model = "gemini-1.5-pro-preview-0409"
+    vertex_project = "test-g-project"
+    vertex_location = "global"
+    vertex_api_version = "v1"
+
+    # Mock litellm.VertexGeminiConfig.get_model_for_vertex_ai_url to return model as is
+    # as we are not testing that part here, just the URL construction
+    with patch("litellm.VertexGeminiConfig.get_model_for_vertex_ai_url", side_effect=lambda model: model):
+        url, endpoint = _get_vertex_url(
+            mode=mode,
+            model=model,
+            stream=stream,
+            vertex_project=vertex_project,
+            vertex_location=vertex_location,
+            vertex_api_version=vertex_api_version,
+        )
+
+    expected_url_base = f"https://aiplatform.googleapis.com/{vertex_api_version}/projects/{vertex_project}/locations/global/publishers/google/models/{model}"
+    
+    if stream:
+        expected_endpoint = "streamGenerateContent"
+        expected_url = f"{expected_url_base}:{expected_endpoint}?alt=sse"
+    else:
+        expected_endpoint = "generateContent"
+        expected_url = f"{expected_url_base}:{expected_endpoint}"
+
+
+    assert endpoint == expected_endpoint
+    assert url == expected_url