fix(vertex_httpx.py): support setting custom api base for vertex ai calls

Closes https://github.com/BerriAI/litellm/issues/4317

litellm/llms/vertex_httpx.py

@@ -551,6 +551,7 @@ class VertexLLM(BaseLLM):
         vertex_credentials: Optional[str],
         stream: Optional[bool],
         custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
+        api_base: Optional[str],
     ) -> Tuple[Optional[str], str]:
         """
         Internal function. Returns the token and url for the call.
@@ -584,6 +585,17 @@ class VertexLLM(BaseLLM):
                 endpoint = "streamGenerateContent"
             url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}"
 
+        if (
+            api_base is not None
+        ):  # for cloudflare ai gateway - https://github.com/BerriAI/litellm/issues/4317
+            if custom_llm_provider == "gemini":
+                url = "{}/{}".format(api_base, endpoint)
+                auth_header = (
+                    gemini_api_key  # cloudflare expects api key as bearer token
+                )
+            else:
+                url = "{}:{}".format(api_base, endpoint)
+
         return auth_header, url
 
     async def async_streaming(
@@ -694,6 +706,7 @@ class VertexLLM(BaseLLM):
         logger_fn=None,
         extra_headers: Optional[dict] = None,
         client: Optional[Union[AsyncHTTPHandler, HTTPHandler]] = None,
+        api_base: Optional[str] = None,
     ) -> Union[ModelResponse, CustomStreamWrapper]:
         stream: Optional[bool] = optional_params.pop("stream", None)  # type: ignore
 
@@ -705,6 +718,7 @@ class VertexLLM(BaseLLM):
             vertex_credentials=vertex_credentials,
             stream=stream,
             custom_llm_provider=custom_llm_provider,
+            api_base=api_base,
         )
 
         ## TRANSFORMATION ##

litellm/main.py

@@ -1930,6 +1930,7 @@ def completion(
                 timeout=timeout,
                 custom_llm_provider=custom_llm_provider,
                 client=client,
+                api_base=api_base,
             )
 
         elif custom_llm_provider == "vertex_ai":

litellm/tests/test_amazing_vertex_completion.py

@@ -854,6 +854,38 @@ Using this JSON schema:
         mock_call.assert_called_once()
 
 
+@pytest.mark.parametrize("provider", ["vertex_ai_beta"])  # "vertex_ai",
+@pytest.mark.asyncio
+async def test_gemini_pro_httpx_custom_api_base(provider):
+    load_vertex_ai_credentials()
+    litellm.set_verbose = True
+    messages = [
+        {
+            "role": "user",
+            "content": "Hello world",
+        }
+    ]
+    from litellm.llms.custom_httpx.http_handler import HTTPHandler
+
+    client = HTTPHandler()
+
+    with patch.object(client, "post", new=MagicMock()) as mock_call:
+        try:
+            response = completion(
+                model="vertex_ai_beta/gemini-1.5-flash",
+                messages=messages,
+                response_format={"type": "json_object"},
+                client=client,
+                api_base="my-custom-api-base",
+            )
+        except Exception as e:
+            pass
+
+        mock_call.assert_called_once()
+
+        assert "my-custom-api-base:generateContent" == mock_call.call_args.kwargs["url"]
+
+
 @pytest.mark.skip(reason="exhausted vertex quota. need to refactor to mock the call")
 @pytest.mark.parametrize("sync_mode", [True])
 @pytest.mark.parametrize("provider", ["vertex_ai"])