Skip to content

Commit c0de42c

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Eval SDK: Route gemini-3 model traffic to global region in local agent scraping
PiperOrigin-RevId: 893194564
1 parent 09794ba commit c0de42c

3 files changed

Lines changed: 115 additions & 23 deletions

File tree

‎tests/unit/vertexai/genai/test_evals.py‎

Lines changed: 67 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@
4444
import pandas as pd
4545
import pytest
4646

47-
4847
_TEST_PROJECT = "test-project"
4948
_TEST_LOCATION = "us-central1"
5049

@@ -2079,6 +2078,73 @@ def test_has_tool_call_with_agent_event(self):
20792078

20802079

20812080
@pytest.mark.usefixtures("google_auth_mock")
2081+
class TestRunAgent:
2082+
"""Unit tests for the _run_agent function."""
2083+
2084+
@mock.patch.object(_evals_common, "_execute_inference_concurrently")
2085+
def test_run_agent_rewrites_gemini_3_model_name(
2086+
self, mock_execute_inference_concurrently, mock_api_client_fixture
2087+
):
2088+
mock_execute_inference_concurrently.return_value = []
2089+
user_simulator_config = vertexai_genai_types.evals.UserSimulatorConfig(
2090+
model_name="gemini-3-preview"
2091+
)
2092+
prompt_dataset = pd.DataFrame({"prompt": ["prompt1"]})
2093+
with mock.patch.dict(os.environ, clear=True):
2094+
os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
2095+
2096+
def mock_execute(*args, **kwargs):
2097+
assert os.environ["GOOGLE_CLOUD_LOCATION"] == "global"
2098+
return []
2099+
2100+
mock_execute_inference_concurrently.side_effect = mock_execute
2101+
2102+
_evals_common._run_agent(
2103+
api_client=mock_api_client_fixture,
2104+
agent_engine=mock.Mock(),
2105+
agent=None,
2106+
prompt_dataset=prompt_dataset,
2107+
user_simulator_config=user_simulator_config,
2108+
)
2109+
2110+
assert (
2111+
user_simulator_config.model_name
2112+
== f"projects/{mock_api_client_fixture.project}/locations/global/publishers/google/models/gemini-3-preview"
2113+
)
2114+
assert os.environ.get("GOOGLE_CLOUD_LOCATION") == "us-central1"
2115+
2116+
@mock.patch.object(_evals_common, "_execute_inference_concurrently")
2117+
def test_run_agent_rewrites_gemini_3_model_name_empty_env(
2118+
self, mock_execute_inference_concurrently, mock_api_client_fixture
2119+
):
2120+
mock_execute_inference_concurrently.return_value = []
2121+
user_simulator_config = vertexai_genai_types.evals.UserSimulatorConfig(
2122+
model_name="gemini-3-preview"
2123+
)
2124+
prompt_dataset = pd.DataFrame({"prompt": ["prompt1"]})
2125+
with mock.patch.dict(os.environ, clear=True):
2126+
2127+
def mock_execute(*args, **kwargs):
2128+
assert os.environ["GOOGLE_CLOUD_LOCATION"] == "global"
2129+
return []
2130+
2131+
mock_execute_inference_concurrently.side_effect = mock_execute
2132+
2133+
_evals_common._run_agent(
2134+
api_client=mock_api_client_fixture,
2135+
agent_engine=mock.Mock(),
2136+
agent=None,
2137+
prompt_dataset=prompt_dataset,
2138+
user_simulator_config=user_simulator_config,
2139+
)
2140+
2141+
assert (
2142+
user_simulator_config.model_name
2143+
== f"projects/{mock_api_client_fixture.project}/locations/global/publishers/google/models/gemini-3-preview"
2144+
)
2145+
assert "GOOGLE_CLOUD_LOCATION" not in os.environ
2146+
2147+
20822148
class TestRunAgentInternal:
20832149
"""Unit tests for the _run_agent_internal function."""
20842150

‎vertexai/_genai/_evals_common.py‎

Lines changed: 44 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1877,28 +1877,50 @@ def _run_agent(
18771877
]
18781878
]:
18791879
"""Internal helper to run inference using Gemini model with concurrency."""
1880-
if agent_engine:
1881-
return _execute_inference_concurrently(
1882-
api_client=api_client,
1883-
agent_engine=agent_engine,
1884-
prompt_dataset=prompt_dataset,
1885-
progress_desc="Agent Run",
1886-
gemini_config=None,
1887-
user_simulator_config=None,
1888-
inference_fn=_execute_agent_run_with_retry,
1889-
)
1890-
elif agent:
1891-
return _execute_inference_concurrently(
1892-
api_client=api_client,
1893-
agent=agent,
1894-
prompt_dataset=prompt_dataset,
1895-
progress_desc="Local Agent Run",
1896-
gemini_config=None,
1897-
user_simulator_config=user_simulator_config,
1898-
inference_fn=_execute_local_agent_run_with_retry,
1899-
)
1900-
else:
1901-
raise ValueError("Neither agent_engine nor agent is provided.")
1880+
original_location = os.environ.get("GOOGLE_CLOUD_LOCATION")
1881+
location_overridden = False
1882+
1883+
if user_simulator_config and user_simulator_config.model_name:
1884+
model_name = user_simulator_config.model_name
1885+
if model_name.startswith("gemini-3") and "/" not in model_name:
1886+
logger.warning(
1887+
"Model %s is only available in the global region. Routing to global.",
1888+
model_name,
1889+
)
1890+
user_simulator_config.model_name = f"projects/{api_client.project}/locations/global/publishers/google/models/{model_name}"
1891+
if original_location != "global":
1892+
os.environ["GOOGLE_CLOUD_LOCATION"] = "global"
1893+
location_overridden = True
1894+
1895+
try:
1896+
if agent_engine:
1897+
return _execute_inference_concurrently(
1898+
api_client=api_client,
1899+
agent_engine=agent_engine,
1900+
prompt_dataset=prompt_dataset,
1901+
progress_desc="Agent Run",
1902+
gemini_config=None,
1903+
user_simulator_config=None,
1904+
inference_fn=_execute_agent_run_with_retry,
1905+
)
1906+
elif agent:
1907+
return _execute_inference_concurrently(
1908+
api_client=api_client,
1909+
agent=agent,
1910+
prompt_dataset=prompt_dataset,
1911+
progress_desc="Local Agent Run",
1912+
gemini_config=None,
1913+
user_simulator_config=user_simulator_config,
1914+
inference_fn=_execute_local_agent_run_with_retry,
1915+
)
1916+
else:
1917+
raise ValueError("Neither agent_engine nor agent is provided.")
1918+
finally:
1919+
if location_overridden:
1920+
if original_location is None:
1921+
del os.environ["GOOGLE_CLOUD_LOCATION"]
1922+
else:
1923+
os.environ["GOOGLE_CLOUD_LOCATION"] = original_location
19021924

19031925

19041926
def _execute_agent_run_with_retry(

‎vertexai/_genai/evals.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2216,6 +2216,8 @@ def create_evaluation_run(
22162216
If `agent_info` is provided without `inference_configs`, this config is used
22172217
to automatically construct the inference configuration. If not specified,
22182218
or if `max_turn` is not set, `max_turn` defaults to 5.
2219+
The `model_name` inside this config can be either a full model path or a
2220+
short model name, e.g. `gemini-3-preview-flash`.
22192221
inference_configs: The candidate to inference config map for the evaluation run.
22202222
The key is the candidate name, and the value is the inference config.
22212223
If provided, `agent_info` must be None. If omitted and `agent_info` is provided,
@@ -3486,6 +3488,8 @@ async def create_evaluation_run(
34863488
If `agent_info` is provided without `inference_configs`, this config is used
34873489
to automatically construct the inference configuration. If not specified,
34883490
or if `max_turn` is not set, `max_turn` defaults to 5.
3491+
The `model_name` inside this config can be either a full model path or a
3492+
short model name, e.g. `gemini-3-preview-flash`.
34893493
inference_configs: The candidate to inference config map for the evaluation run.
34903494
The key is the candidate name, and the value is the inference config.
34913495
If provided, `agent_info` must be None. If omitted and `agent_info` is provided,

0 commit comments

Comments
 (0)