diff --git a/litellm/llms/vertex_httpx.py b/litellm/llms/vertex_httpx.py index 01fef7e98bc..05ea17c7c86 100644 --- a/litellm/llms/vertex_httpx.py +++ b/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 ## diff --git a/litellm/main.py b/litellm/main.py index a1b8a42abb9..bd26528178c 100644 --- a/litellm/main.py +++ b/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": diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 6433b064108..f764324479a 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/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"])