Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion tests/unit/vertexai/genai/test_evals.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand All @@ -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
Expand All @@ -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:
Expand Down
19 changes: 19 additions & 0 deletions vertexai/_genai/evals.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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.
Expand All @@ -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]]
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -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)

Expand Down Expand Up @@ -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.
Expand All @@ -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]]
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -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)

Expand Down
7 changes: 7 additions & 0 deletions vertexai/_genai/types/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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
Expand Down
Loading