diff --git a/tests/unit/vertexai/genai/test_evals.py b/tests/unit/vertexai/genai/test_evals.py index 1d400584f7..1f15ac95a3 100644 --- a/tests/unit/vertexai/genai/test_evals.py +++ b/tests/unit/vertexai/genai/test_evals.py @@ -3674,12 +3674,16 @@ def test_run_inference_with_litellm_openai_request_format( mock_api_client_fixture, ): """Tests inference with LiteLLM where the row contains a chat completion request body.""" + # fmt: off with ( - mock.patch("vertexai._genai._evals_common.litellm") as mock_litellm, + mock.patch( + "vertexai._genai._evals_common.litellm" + ) as mock_litellm, mock.patch( "vertexai._genai._evals_common._call_litellm_completion" ) as mock_call_litellm_completion, ): + # fmt: on mock_litellm.get_llm_provider.return_value = ( "gpt-4o", "openai", @@ -8452,6 +8456,7 @@ def test_generate_conversation_scenarios(self): root_agent_id="agent_1", ), config={"count": 2}, + allow_cross_region_model=True, ) assert isinstance(eval_dataset, vertexai_genai_types.EvaluationDataset) assert len(eval_dataset.eval_cases) == 2 @@ -8468,6 +8473,9 @@ def test_generate_conversation_scenarios(self): assert eval_dataset.eval_dataset_df.iloc[1]["conversation_plan"] == "Plan 2" self.mock_api_client.request.assert_called_once() + call_args = self.mock_api_client.request.call_args + request_body = call_args[0][2] # Third positional arg is the request dict + assert request_body.get("allowCrossRegionModel") is True @pytest.mark.asyncio async def test_async_generate_conversation_scenarios(self): @@ -8484,6 +8492,7 @@ async def test_async_generate_conversation_scenarios(self): root_agent_id="agent_1", ), config={"count": 2}, + allow_cross_region_model=True, ) assert isinstance(eval_dataset, vertexai_genai_types.EvaluationDataset) assert len(eval_dataset.eval_cases) == 2 @@ -8497,6 +8506,9 @@ async def test_async_generate_conversation_scenarios(self): assert eval_dataset.eval_dataset_df.iloc[1]["conversation_plan"] == "Plan 2" self.mock_api_client.async_request.assert_called_once() + call_args = self.mock_api_client.async_request.call_args + request_body = call_args[0][2] # Third positional arg is the request dict + assert request_body.get("allowCrossRegionModel") is True class TestTransformDataframe: diff --git a/vertexai/_genai/evals.py b/vertexai/_genai/evals.py index ff7af8e7f5..5567036989 100644 --- a/vertexai/_genai/evals.py +++ b/vertexai/_genai/evals.py @@ -683,6 +683,13 @@ def _GenerateUserScenariosParameters_to_vertex( if getv(from_object, ["config"]) is not None: setv(to_object, ["config"], getv(from_object, ["config"])) + if getv(from_object, ["allow_cross_region_model"]) is not None: + setv( + to_object, + ["allowCrossRegionModel"], + getv(from_object, ["allow_cross_region_model"]), + ) + return to_object @@ -1473,6 +1480,7 @@ def _generate_user_scenarios( evals_types.UserScenarioGenerationConfigOrDict ] = None, config: Optional[types.GenerateUserScenariosConfigOrDict] = None, + allow_cross_region_model: Optional[bool] = None, ) -> types.GenerateUserScenariosResponse: """ Generates user scenarios for agent evaluation. @@ -1484,6 +1492,7 @@ def _generate_user_scenarios( root_agent_id=root_agent_id, user_scenario_generation_config=user_scenario_generation_config, config=config, + allow_cross_region_model=allow_cross_region_model, ) request_url_dict: Optional[dict[str, str]] @@ -2852,6 +2861,7 @@ def generate_conversation_scenarios( *, agent_info: evals_types.AgentInfoOrDict, config: evals_types.UserScenarioGenerationConfigOrDict, + allow_cross_region_model: Optional[bool] = None, ) -> types.EvaluationDataset: """Generates an evaluation dataset with user scenarios, which helps to generate conversations between a simulated user @@ -2860,6 +2870,8 @@ def generate_conversation_scenarios( Args: agent_info: The agent info to generate user scenarios for. config: Configuration for generating user scenarios. + allow_cross_region_model: Opt-in flag to authorize cross-region + routing for model inference. Returns: An EvaluationDataset containing the generated user scenarios. @@ -2873,6 +2885,7 @@ def generate_conversation_scenarios( agents=parsed_agent_info.agents, root_agent_id=parsed_agent_info.root_agent_id, user_scenario_generation_config=config, + allow_cross_region_model=allow_cross_region_model, ) return _evals_utils._postprocess_user_scenarios_response(response) @@ -3592,6 +3605,7 @@ async def _generate_user_scenarios( evals_types.UserScenarioGenerationConfigOrDict ] = None, config: Optional[types.GenerateUserScenariosConfigOrDict] = None, + allow_cross_region_model: Optional[bool] = None, ) -> types.GenerateUserScenariosResponse: """ Generates user scenarios for agent evaluation. @@ -3603,6 +3617,7 @@ async def _generate_user_scenarios( root_agent_id=root_agent_id, user_scenario_generation_config=user_scenario_generation_config, config=config, + allow_cross_region_model=allow_cross_region_model, ) request_url_dict: Optional[dict[str, str]] @@ -4622,6 +4637,7 @@ async def generate_conversation_scenarios( *, agent_info: evals_types.AgentInfoOrDict, config: evals_types.UserScenarioGenerationConfigOrDict, + allow_cross_region_model: Optional[bool] = None, ) -> types.EvaluationDataset: """Generates an evaluation dataset with user scenarios, which helps to generate conversations between a simulated user @@ -4630,6 +4646,8 @@ async def generate_conversation_scenarios( Args: agent_info: The agent info to generate user scenarios for. config: Configuration for generating user scenarios. + allow_cross_region_model: Opt-in flag to authorize cross-region + routing for model inference. Returns: An EvaluationDataset containing the generated user scenarios. @@ -4643,6 +4661,7 @@ async def generate_conversation_scenarios( agents=parsed_agent_info.agents, root_agent_id=parsed_agent_info.root_agent_id, user_scenario_generation_config=config, + allow_cross_region_model=allow_cross_region_model, ) return _evals_utils._postprocess_user_scenarios_response(response) diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index da753f6f78..81c505f81a 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -5012,6 +5012,10 @@ class _GenerateUserScenariosParameters(_common.BaseModel): config: Optional[GenerateUserScenariosConfig] = Field( default=None, description="""""" ) + allow_cross_region_model: Optional[bool] = Field( + default=None, + description="""Opt-in flag to authorize cross-region routing for LLM models.""", + ) class _GenerateUserScenariosParametersDict(TypedDict, total=False): @@ -5032,6 +5036,9 @@ class _GenerateUserScenariosParametersDict(TypedDict, total=False): config: Optional[GenerateUserScenariosConfigDict] """""" + allow_cross_region_model: Optional[bool] + """Opt-in flag to authorize cross-region routing for LLM models.""" + _GenerateUserScenariosParametersOrDict = Union[ _GenerateUserScenariosParameters, _GenerateUserScenariosParametersDict