diff --git a/tests/unit/vertexai/genai/test_genai_client.py b/tests/unit/vertexai/genai/test_genai_client.py index 16c7068b78..cd600df56d 100644 --- a/tests/unit/vertexai/genai/test_genai_client.py +++ b/tests/unit/vertexai/genai/test_genai_client.py @@ -52,6 +52,13 @@ def test_genai_client(self): assert test_client._api_client.project == _TEST_PROJECT assert test_client._api_client.location == _TEST_LOCATION + @pytest.mark.parametrize("location", ["us", "eu"]) + @pytest.mark.usefixtures("google_auth_mock") + def test_genai_client_mrep(self, location): + test_client = vertexai.Client(project=_TEST_PROJECT, location=location) + expected_url = f"https://aiplatform.{location}.rep.googleapis.com/" + assert test_client._api_client._http_options.base_url == expected_url + @pytest.mark.asyncio @pytest.mark.usefixtures("google_auth_mock") async def test_async_client(self): diff --git a/vertexai/_genai/client.py b/vertexai/_genai/client.py index 77d8a9aadc..f16c2d50ed 100644 --- a/vertexai/_genai/client.py +++ b/vertexai/_genai/client.py @@ -219,6 +219,10 @@ def __init__( if http_options.headers is None: http_options.headers = {} + # Set the base URL for MREP locations. + if location in ["us", "eu"] and not http_options.base_url: + http_options.base_url = f"https://aiplatform.{location}.rep.googleapis.com/" + tracking_label = f"{_GENAI_MODULES_TELEMETRY_HEADER}/{aip_version.__version__}" if "user-agent" in http_options.headers: