Skip to content

feat: Delegate agent resolution to server via geminiAgentConfig #6997

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
+52 −75
Open

feat: Delegate agent resolution to server via geminiAgentConfig #6997

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
27 changes: 0 additions & 27 deletions agentplatform/_genai/_evals_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1016,33 +1016,6 @@ def _agent_data_response_text(agent_data: types.evals.AgentData) -> Optional[str
return "".join(text_parts) or None


def _agent_resource_to_agent_info(
agent: str, api_client: BaseApiClient
) -> "types.evals.AgentInfo":
"""Builds an `AgentInfo` from a Gemini Agents API agent resource name.

Fetches the agent through the SDK's `api_client` (so replay recording is
preserved) via `_fetch_agent_config_dict` and derives a single-agent
`AgentInfo`: the agent's short name is the agents-map key and
`root_agent_id`.

Args:
agent: The Gemini Agents API agent resource name
(`projects/{p}/locations/{l}/agents/{name}`).
api_client: The API client used to fetch the agent.

Returns:
An `AgentInfo` describing the fetched agent.
"""
agent_config = _fetch_agent_config_dict(api_client, agent)
short_name = agent_config.agent_id
return types.evals.AgentInfo( # pytype: disable=missing-parameter
name=short_name,
agents={short_name: agent_config},
root_agent_id=short_name,
)


_INTERACTION_TERMINAL_STATES = frozenset(
["completed", "failed", "cancelled", "incomplete", "budget_exceeded"]
)
Expand Down
61 changes: 37 additions & 24 deletions agentplatform/_genai/evals.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -725,6 +725,11 @@ def _GenerateUserScenariosParameters_to_vertex(
getv(from_object, ["allow_cross_region_model"]),
)

if getv(from_object, ["gemini_agent_config"]) is not None:
setv(
to_object, ["geminiAgentConfig"], getv(from_object, ["gemini_agent_config"])
)

return to_object


Expand Down Expand Up @@ -1518,6 +1523,7 @@ def _generate_user_scenarios(
] = None,
config: Optional[types.GenerateUserScenariosConfigOrDict] = None,
allow_cross_region_model: Optional[bool] = None,
gemini_agent_config: Optional[types.GeminiAgentConfigOrDict] = None,
) -> types.GenerateUserScenariosResponse:
"""
Generates user scenarios for agent evaluation.
Expand All @@ -1530,6 +1536,7 @@ def _generate_user_scenarios(
user_scenario_generation_config=user_scenario_generation_config,
config=config,
allow_cross_region_model=allow_cross_region_model,
gemini_agent_config=gemini_agent_config,
)

request_url_dict: Optional[dict[str, str]]
Expand Down Expand Up @@ -2938,16 +2945,16 @@ def generate_conversation_scenarios(
and the agent under test.

Exactly one of `agent_info` or `agent` must be provided. When `agent` is
a Gemini Agents API agent resource name, the agent is fetched and an
`AgentInfo` is derived from it.
a Gemini Agents API agent resource name, the agent config is resolved
server-side via ``gemini_agent_config``.

Args:
agent_info: The agent info to generate user scenarios for. Mutually
exclusive with `agent`.
agent: A Gemini Agents API agent resource name
(`projects/{p}/locations/{l}/agents/{name}`). When provided, the
agent is fetched and its configuration is used to build the agent
info. Mutually exclusive with `agent_info`.
server resolves the agent config directly. Mutually exclusive
with `agent_info`.
config: Configuration for generating user scenarios.
allow_cross_region_model: Opt-in flag to authorize cross-region
routing for model inference.
Expand All @@ -2967,21 +2974,23 @@ def generate_conversation_scenarios(
"`agent` must be a Gemini Agents API agent resource name of the"
" form projects/{project}/locations/{location}/agents/{agent}."
)
parsed_agent_info = _evals_common._agent_resource_to_agent_info(
agent, self._api_client
response = self._generate_user_scenarios(
gemini_agent_config=types.GeminiAgentConfig(gemini_agent=agent),
user_scenario_generation_config=config,
allow_cross_region_model=allow_cross_region_model,
)
else:
parsed_agent_info = (
evals_types.AgentInfo.model_validate(agent_info)
if isinstance(agent_info, dict)
else agent_info
)
response = self._generate_user_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,
)
response = self._generate_user_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)

def generate_loss_clusters(
Expand Down Expand Up @@ -3683,6 +3692,7 @@ async def _generate_user_scenarios(
] = None,
config: Optional[types.GenerateUserScenariosConfigOrDict] = None,
allow_cross_region_model: Optional[bool] = None,
gemini_agent_config: Optional[types.GeminiAgentConfigOrDict] = None,
) -> types.GenerateUserScenariosResponse:
"""
Generates user scenarios for agent evaluation.
Expand All @@ -3695,6 +3705,7 @@ async def _generate_user_scenarios(
user_scenario_generation_config=user_scenario_generation_config,
config=config,
allow_cross_region_model=allow_cross_region_model,
gemini_agent_config=gemini_agent_config,
)

request_url_dict: Optional[dict[str, str]]
Expand Down Expand Up @@ -4729,16 +4740,16 @@ async def generate_conversation_scenarios(
and the agent under test.

Exactly one of `agent_info` or `agent` must be provided. When `agent` is
a Gemini Agents API agent resource name, the agent is fetched and an
`AgentInfo` is derived from it.
a Gemini Agents API agent resource name, the agent config is resolved
server-side via ``gemini_agent_config``.

Args:
agent_info: The agent info to generate user scenarios for. Mutually
exclusive with `agent`.
agent: A Gemini Agents API agent resource name
(`projects/{p}/locations/{l}/agents/{name}`). When provided, the
agent is fetched and its configuration is used to build the agent
info. Mutually exclusive with `agent_info`.
server resolves the agent config directly. Mutually exclusive
with `agent_info`.
config: Configuration for generating user scenarios.
allow_cross_region_model: Opt-in flag to authorize cross-region
routing for model inference.
Expand All @@ -4758,21 +4769,23 @@ async def generate_conversation_scenarios(
"`agent` must be a Gemini Agents API agent resource name of the"
" form projects/{project}/locations/{location}/agents/{agent}."
)
parsed_agent_info = _evals_common._agent_resource_to_agent_info(
agent, self._api_client
response = await self._generate_user_scenarios(
gemini_agent_config=types.GeminiAgentConfig(gemini_agent=agent),
user_scenario_generation_config=config,
allow_cross_region_model=allow_cross_region_model,
)
else:
parsed_agent_info = (
evals_types.AgentInfo.model_validate(agent_info)
if isinstance(agent_info, dict)
else agent_info
)
response = await self._generate_user_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,
)
response = await self._generate_user_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)

async def generate_loss_clusters(
Expand Down
9 changes: 9 additions & 0 deletions agentplatform/_genai/types/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5501,6 +5501,11 @@ class _GenerateUserScenariosParameters(_common.BaseModel):
default=None,
description="""Opt-in flag to authorize cross-region routing for LLM models.""",
)
gemini_agent_config: Optional[GeminiAgentConfig] = Field(
default=None,
description="""If set, the server derives the agents map and root_agent_id
from the referenced Gemini Agent server-side.""",
)


class _GenerateUserScenariosParametersDict(TypedDict, total=False):
Expand All @@ -5524,6 +5529,10 @@ class _GenerateUserScenariosParametersDict(TypedDict, total=False):
allow_cross_region_model: Optional[bool]
"""Opt-in flag to authorize cross-region routing for LLM models."""

gemini_agent_config: Optional[GeminiAgentConfigDict]
"""If set, the server derives the agents map and root_agent_id
from the referenced Gemini Agent server-side."""


_GenerateUserScenariosParametersOrDict = Union[
_GenerateUserScenariosParameters, _GenerateUserScenariosParametersDict
Expand Down
30 changes: 6 additions & 24 deletions tests/unit/agentplatform/genai/test_evals.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9409,19 +9409,9 @@ async def test_async_generate_conversation_scenarios(self):
request_body = call_args[0][2] # Third positional arg is the request dict
assert request_body.get("allowCrossRegionModel") is True

@mock.patch.object(_evals_common, "_fetch_agent_config_dict")
def test_generate_conversation_scenarios_from_gemini_agent(
self, mock_fetch_agent_config
):
mock_fetch_agent_config.return_value = (
agentplatform_genai_types.evals.AgentConfig(
agent_id="test-agent",
instruction="You are a helpful travel assistant.",
description="An agent that books flights.",
tools=[genai_types.Tool(google_search=genai_types.GoogleSearch())],
)
)

def test_generate_conversation_scenarios_from_gemini_agent(self):
"""When `agent` is a Gemini agent resource, gemini_agent_config is
forwarded to the server (no client-side synthesis)."""
evals_module = evals.Evals(api_client_=self.mock_api_client)

with mock.patch.object(
Expand All @@ -9433,18 +9423,10 @@ def test_generate_conversation_scenarios_from_gemini_agent(
config={"count": 2},
)

mock_fetch_agent_config.assert_called_once_with(
self.mock_api_client, _TEST_GEMINI_AGENT
)
call_kwargs = mock_generate_user_scenarios.call_args.kwargs
assert call_kwargs["root_agent_id"] == "test-agent"
agents = call_kwargs["agents"]
assert "test-agent" in agents
derived_config = agents["test-agent"]
assert derived_config.instruction == "You are a helpful travel assistant."
assert derived_config.description == "An agent that books flights."
assert derived_config.tools is not None
assert derived_config.tools[0].google_search is not None
assert call_kwargs["gemini_agent_config"].gemini_agent == _TEST_GEMINI_AGENT
assert call_kwargs.get("agents") is None
assert call_kwargs.get("root_agent_id") is None

def test_generate_conversation_scenarios_agent_and_agent_info_raises(self):
evals_module = evals.Evals(api_client_=self.mock_api_client)
Expand Down
Loading