Skip to content

Commit fb95b6d

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: Add detailed agent tool definitions
FUTURE_COPYBARA_INTEGRATE_REVIEW=#6967 from googleapis:release-please--branches--main 4a0d1d2 PiperOrigin-RevId: 948478465
1 parent 3335750 commit fb95b6d

5 files changed

Lines changed: 423 additions & 122 deletions

File tree

agentplatform/_genai/_evals_common.py

Lines changed: 199 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -837,47 +837,212 @@ def _merge_text_parts_in_agent_data(
837837
content.parts = merged_parts
838838

839839

840+
# ---------------------------------------------------------------------------
841+
# Built-in tool catalog (display-only)
842+
#
843+
# The Gemini Agents API (``GET agents/{id}``) returns each tool as a bare type
844+
# discriminator (e.g. ``{"type": "code_execution"}``) with no parameter schema
845+
# or description. The authoritative, full-fidelity expansion lives server-side
846+
# in ``cloud/ai/platform/evaluation/utils/interaction_converter.py``. This
847+
# catalog is a **display-only duplicate** of that server catalog, kept here so
848+
# ``show()`` can render tools with full names, descriptions, and parameter
849+
# schemas without a server round-trip.
850+
#
851+
# If the server catalog changes, this SDK-side copy must be updated to match.
852+
# ---------------------------------------------------------------------------
853+
854+
855+
def _str_schema(description: str) -> genai_types.Schema:
856+
return genai_types.Schema(type="STRING", description=description)
857+
858+
859+
_BUILTIN_TOOL_DECLARATIONS: dict[str, list[genai_types.FunctionDeclaration]] = {
860+
"code_execution": [
861+
genai_types.FunctionDeclaration(
862+
name="run_command",
863+
description=(
864+
"Runs a shell command on the sandbox VM. If the command does"
865+
" not complete within WaitMsBeforeAsync, it is sent to the"
866+
" background and a CommandId is returned for use with"
867+
" command_status."
868+
),
869+
parameters=genai_types.Schema(
870+
type="OBJECT",
871+
properties={
872+
"CommandLine": _str_schema("The shell command to run."),
873+
"Cwd": _str_schema(
874+
"The current working directory for the command."
875+
),
876+
"WaitMsBeforeAsync": genai_types.Schema(
877+
type="INTEGER",
878+
description=(
879+
"Milliseconds to wait for the command to complete"
880+
" before sending it to the background. Default:"
881+
" 10000."
882+
),
883+
),
884+
"SafeToAutoRun": genai_types.Schema(
885+
type="BOOLEAN",
886+
description=(
887+
"Whether the command is safe to auto-run without"
888+
" user approval."
889+
),
890+
),
891+
},
892+
required=["CommandLine", "Cwd"],
893+
),
894+
),
895+
],
896+
"filesystem": [
897+
genai_types.FunctionDeclaration(
898+
name="view_file",
899+
description="Reads the content of a workspace file.",
900+
),
901+
genai_types.FunctionDeclaration(
902+
name="create_file",
903+
description="Writes content to a new or existing file.",
904+
),
905+
genai_types.FunctionDeclaration(
906+
name="edit_file",
907+
description="Replaces a specific block of text in a file.",
908+
),
909+
genai_types.FunctionDeclaration(
910+
name="list_dir",
911+
description="Lists the files in a directory.",
912+
),
913+
genai_types.FunctionDeclaration(
914+
name="delete_file",
915+
description="Removes a file from the workspace.",
916+
),
917+
genai_types.FunctionDeclaration(
918+
name="move_file",
919+
description="Renames or moves a file.",
920+
),
921+
],
922+
}
923+
924+
925+
_SANDBOX_DECLARATIONS: list[genai_types.FunctionDeclaration] = [
926+
genai_types.FunctionDeclaration(
927+
name="provision_sandbox",
928+
description="Provisions a sandbox environment.",
929+
parameters=genai_types.Schema(
930+
type="OBJECT",
931+
properties={
932+
"display_name": _str_schema(
933+
"The display name of the sandbox environment."
934+
),
935+
"poll_creation_lro": genai_types.Schema(
936+
type="BOOLEAN",
937+
description=(
938+
"Whether to poll the creation long-running operation."
939+
),
940+
),
941+
},
942+
),
943+
),
944+
genai_types.FunctionDeclaration(
945+
name="load_sandbox",
946+
description="Loads a previously provisioned sandbox environment.",
947+
parameters=genai_types.Schema(
948+
type="OBJECT",
949+
properties={
950+
"reasoning_engine_resource_name": _str_schema(
951+
"The resource name of the reasoning engine. Format:"
952+
" projects/{project}/locations/{location}/reasoningEngines/{id}"
953+
),
954+
"display_name": _str_schema(
955+
"The display name of the sandbox environment. Format: any"
956+
" string."
957+
),
958+
},
959+
),
960+
),
961+
]
962+
963+
840964
def _agent_tools_to_config_tools(
841965
agent_tools: Optional[list[Any]],
966+
has_environment: bool = False,
842967
) -> Optional[list[genai_types.Tool]]:
843-
"""Maps Gemini Agents API tools to ``genai_types.Tool`` for an AgentConfig.
844-
845-
The Gemini Agents API returns built-in tool variants (``google_search``,
846-
``code_execution``, ``url_context``) whose schema differs from
847-
``genai_types.Tool``. Each recognised built-in variant is mapped to the
848-
matching ``genai_types.Tool`` field. Tools with a non-empty body after
849-
stripping the ``type`` key (e.g. ``function_declarations``) are passed
850-
through ``model_validate``. Variants without a ``genai_types.Tool``
851-
equivalent (e.g. ``filesystem``, ``mcp_server``) are skipped.
968+
"""Maps Gemini Agents API tools to ``genai_types.Tool`` for display.
969+
970+
Expands built-in agent tool types into their concrete function declarations
971+
using ``_BUILTIN_TOOL_DECLARATIONS`` (a display-only duplicate of the
972+
server-side catalog in ``interaction_converter.py``).
973+
974+
Mapping rules:
975+
* ``code_execution`` is expanded to ``run_command`` with full parameter
976+
schema.
977+
* ``filesystem`` is expanded to ``view_file``, ``create_file``,
978+
``edit_file``, ``list_dir``, ``delete_file``, ``move_file``.
979+
* ``google_search`` and ``url_context`` are mapped to their typed
980+
``genai_types.Tool`` variant.
981+
* ``mcp_server`` is represented as a named declaration with a
982+
human-readable label.
983+
* Tools carrying explicit ``function_declarations`` are passed through.
984+
* When ``has_environment`` is True, sandbox orchestration tools
985+
(``provision_sandbox``, ``load_sandbox``) are appended.
852986
853987
Args:
854988
agent_tools: The ``tools`` list from a fetched Gemini agent dict.
989+
has_environment: Whether the agent has a sandbox environment configured.
855990
856991
Returns:
857992
A list of ``genai_types.Tool``, or ``None`` if there are no mappable
858993
tools.
859994
"""
860-
if not agent_tools:
995+
if not agent_tools and not has_environment:
861996
return None
862997
tools: list[genai_types.Tool] = []
863-
for tool in agent_tools:
998+
for tool in agent_tools or []:
864999
if not isinstance(tool, dict):
8651000
continue
8661001
tool_type = tool.get("type")
867-
if tool_type == "google_search":
868-
tools.append(genai_types.Tool(google_search=genai_types.GoogleSearch()))
869-
elif tool_type == "code_execution":
1002+
remainder = {k: v for k, v in tool.items() if k != "type"}
1003+
1004+
# Check the built-in catalog first (code_execution, filesystem).
1005+
catalog_decls = _BUILTIN_TOOL_DECLARATIONS.get(tool_type or "")
1006+
if catalog_decls:
8701007
tools.append(
871-
genai_types.Tool(code_execution=genai_types.ToolCodeExecution())
1008+
genai_types.Tool(function_declarations=list(catalog_decls))
8721009
)
1010+
elif tool_type == "google_search":
1011+
tools.append(genai_types.Tool(google_search=genai_types.GoogleSearch()))
8731012
elif tool_type == "url_context":
8741013
tools.append(genai_types.Tool(url_context=genai_types.UrlContext()))
875-
else:
876-
# For non-built-in tools (e.g. function_declarations), strip the
877-
# type key and validate through genai_types.Tool.
878-
remainder = {k: v for k, v in tool.items() if k != "type"}
879-
if remainder:
880-
tools.append(genai_types.Tool.model_validate(remainder))
1014+
elif "function_declarations" in remainder:
1015+
# Real function tool with explicit declarations.
1016+
tools.append(genai_types.Tool.model_validate(remainder))
1017+
elif tool_type == "mcp_server":
1018+
label = remainder.get("name") or remainder.get("url")
1019+
description = f"MCP server: {label}" if label else "MCP server."
1020+
tools.append(
1021+
genai_types.Tool(
1022+
function_declarations=[
1023+
genai_types.FunctionDeclaration(
1024+
name="mcp_server", description=description
1025+
)
1026+
]
1027+
)
1028+
)
1029+
elif tool_type:
1030+
# Unknown built-in: show by name so it isn't silently dropped.
1031+
tools.append(
1032+
genai_types.Tool(
1033+
function_declarations=[
1034+
genai_types.FunctionDeclaration(name=tool_type)
1035+
]
1036+
)
1037+
)
1038+
elif remainder:
1039+
tools.append(genai_types.Tool.model_validate(remainder))
1040+
1041+
if has_environment:
1042+
tools.append(
1043+
genai_types.Tool(function_declarations=list(_SANDBOX_DECLARATIONS))
1044+
)
1045+
8811046
return tools or None
8821047

8831048

@@ -916,7 +1081,13 @@ def _fetch_agent_config_dict(
9161081
instruction = agent_dict.get("system_instruction") or None
9171082
description = agent_dict.get("description") or None
9181083
agent_type = agent_dict.get("base_agent") or None
919-
tools = _agent_tools_to_config_tools(agent_dict.get("tools"))
1084+
has_environment = bool(
1085+
agent_dict.get("environment_config")
1086+
or agent_dict.get("base_environment")
1087+
)
1088+
tools = _agent_tools_to_config_tools(
1089+
agent_dict.get("tools"), has_environment=has_environment
1090+
)
9201091
except Exception as e: # pylint: disable=broad-exception-caught
9211092
logger.warning(
9221093
"Failed to fetch agent config for '%s' (continuing without it): %s",
@@ -988,33 +1159,6 @@ def _agent_data_response_text(agent_data: types.evals.AgentData) -> Optional[str
9881159
return "".join(text_parts) or None
9891160

9901161

991-
def _agent_resource_to_agent_info(
992-
agent: str, api_client: BaseApiClient
993-
) -> "types.evals.AgentInfo":
994-
"""Builds an `AgentInfo` from a Gemini Agents API agent resource name.
995-
996-
Fetches the agent through the SDK's `api_client` (so replay recording is
997-
preserved) via `_fetch_agent_config_dict` and derives a single-agent
998-
`AgentInfo`: the agent's short name is the agents-map key and
999-
`root_agent_id`.
1000-
1001-
Args:
1002-
agent: The Gemini Agents API agent resource name
1003-
(`projects/{p}/locations/{l}/agents/{name}`).
1004-
api_client: The API client used to fetch the agent.
1005-
1006-
Returns:
1007-
An `AgentInfo` describing the fetched agent.
1008-
"""
1009-
agent_config = _fetch_agent_config_dict(api_client, agent)
1010-
short_name = agent_config.agent_id
1011-
return types.evals.AgentInfo( # pytype: disable=missing-parameter
1012-
name=short_name,
1013-
agents={short_name: agent_config},
1014-
root_agent_id=short_name,
1015-
)
1016-
1017-
10181162
_INTERACTION_TERMINAL_STATES = frozenset(
10191163
["completed", "failed", "cancelled", "incomplete", "budget_exceeded"]
10201164
)
@@ -1108,6 +1252,11 @@ def _run_gemini_agent_inference(
11081252

11091253
interactions_client = _get_interactions_client(api_client)
11101254

1255+
# Best-effort: fetch the agent config (instruction, tools, description)
1256+
# once, so every row's agent_data carries the agents map and the display
1257+
# can render the System Topology section.
1258+
agent_config = _fetch_agent_config_dict(api_client, gemini_agent)
1259+
11111260
agent_short_id = gemini_agent.split("/")[-1]
11121261
prompts: list[str] = []
11131262
responses: list[Optional[str]] = []
@@ -1128,6 +1277,7 @@ def _run_gemini_agent_inference(
11281277
)
11291278
interaction = _await_interaction(interactions_client, interaction)
11301279
agent_data_obj = _interaction_dict_to_agent_data(interaction)
1280+
agent_data_obj.agents = {agent_config.agent_id: agent_config}
11311281
_merge_text_parts_in_agent_data(agent_data_obj)
11321282
responses.append(_agent_data_response_text(agent_data_obj))
11331283
interaction_ids.append(interaction.get("id"))

0 commit comments

Comments
 (0)