Skip to content

Commit 181e41d

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
BREAKING_CHANGE(agentplatform): agent_engines module renamed to runtimes.
BREAKING_CHANGE(agentplatform): a2a tasks module is removed. BREAKING_CHANGE(agentplatform): sessions and sandboxes as top level modules. BREAKING_CHANGE(agentplatform): removed global initializer dependency - project and location from environment variables in agent frameworks. PiperOrigin-RevId: 957042628
1 parent 2a18bff commit 181e41d

105 files changed

Lines changed: 14090 additions & 18893 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

agentplatform/_genai/_evals_common.py

Lines changed: 64 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -139,21 +139,21 @@ def _get_api_client_with_location(
139139
)._api_client
140140

141141

142-
def _get_agent_engine_instance(
142+
def _get_runtime_instance(
143143
agent_name: str, api_client: BaseApiClient
144-
) -> Union[types.AgentEngine, Any]:
144+
) -> Union[types.Runtime, Any]:
145145
"""Gets or creates an agent engine instance for the current thread."""
146-
if not hasattr(_thread_local_data, "agent_engine_instances"):
147-
_thread_local_data.agent_engine_instances = {}
148-
if agent_name not in _thread_local_data.agent_engine_instances:
146+
if not hasattr(_thread_local_data, "runtime_instances"):
147+
_thread_local_data.runtime_instances = {}
148+
if agent_name not in _thread_local_data.runtime_instances:
149149
client = agentplatform.Client(
150150
project=api_client.project,
151151
location=api_client.location,
152152
)
153-
_thread_local_data.agent_engine_instances[agent_name] = (
154-
client.agent_engines.get(name=agent_name)
153+
_thread_local_data.runtime_instances[agent_name] = client.runtimes.get(
154+
name=agent_name
155155
)
156-
return _thread_local_data.agent_engine_instances[agent_name]
156+
return _thread_local_data.runtime_instances[agent_name]
157157

158158

159159
def _generate_content_with_retry(
@@ -1754,7 +1754,7 @@ def _execute_inference_concurrently(
17541754
model_or_fn: Optional[Union[str, Callable[[Any], Any]]] = None,
17551755
gemini_config: Optional[genai_types.GenerateContentConfig] = None,
17561756
inference_fn: Optional[Callable[..., Any]] = None,
1757-
agent_engine: Optional[Union[str, types.AgentEngine]] = None,
1757+
runtime: Optional[Union[str, types.Runtime]] = None,
17581758
agent: Optional["LlmAgent"] = None, # type: ignore # noqa: F821
17591759
user_simulator_config: Optional[types.evals.UserSimulatorConfig] = None,
17601760
) -> list[
@@ -1784,7 +1784,7 @@ def _execute_inference_concurrently(
17841784
# prompt from the structured agent_data rather than requiring a flat
17851785
# prompt/request column.
17861786
has_agent_data = (
1787-
agent is not None or agent_engine is not None
1787+
agent is not None or runtime is not None
17881788
) and AGENT_DATA in prompt_dataset.columns
17891789

17901790
primary_prompt_column: Optional[str] = None
@@ -1801,7 +1801,7 @@ def _execute_inference_concurrently(
18011801
f" Found: {prompt_dataset.columns.tolist()}"
18021802
)
18031803

1804-
max_workers = AGENT_MAX_WORKERS if agent_engine or agent else MAX_WORKERS
1804+
max_workers = AGENT_MAX_WORKERS if runtime or agent else MAX_WORKERS
18051805
with tqdm(total=len(prompt_dataset), desc=progress_desc) as pbar:
18061806
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
18071807
for index, row in prompt_dataset.iterrows():
@@ -1857,29 +1857,29 @@ def _execute_inference_concurrently(
18571857
pbar.update(1)
18581858
continue
18591859

1860-
if agent_engine or agent:
1860+
if runtime or agent:
18611861

18621862
def agent_run_wrapper( # type: ignore[no-untyped-def]
18631863
row_arg,
18641864
contents_arg,
1865-
agent_engine_arg,
1865+
runtime_arg,
18661866
agent_arg,
18671867
inference_fn_arg,
18681868
api_client_arg,
18691869
user_simulator_config_arg,
18701870
) -> Any:
1871-
if agent_engine_arg:
1872-
if isinstance(agent_engine_arg, str):
1873-
agent_engine_instance = _get_agent_engine_instance(
1874-
agent_engine_arg, api_client_arg
1871+
if runtime_arg:
1872+
if isinstance(runtime_arg, str):
1873+
runtime_instance = _get_runtime_instance(
1874+
runtime_arg, api_client_arg
18751875
)
18761876
else:
1877-
agent_engine_instance = agent_engine_arg
1877+
runtime_instance = runtime_arg
18781878

18791879
return inference_fn_arg(
18801880
row=row_arg,
18811881
contents=contents_arg,
1882-
agent_engine=agent_engine_instance,
1882+
runtime=runtime_instance,
18831883
)
18841884
elif agent_arg:
18851885
return inference_fn_arg(
@@ -1894,7 +1894,7 @@ def agent_run_wrapper( # type: ignore[no-untyped-def]
18941894
agent_run_wrapper,
18951895
row,
18961896
contents,
1897-
agent_engine,
1897+
runtime,
18981898
agent,
18991899
inference_fn,
19001900
api_client,
@@ -2509,7 +2509,7 @@ def _execute_inference(
25092509
api_client: BaseApiClient,
25102510
src: Union[str, pd.DataFrame],
25112511
model: Optional[Union[Callable[[Any], Any], str]] = None,
2512-
agent_engine: Optional[Union[str, types.AgentEngine]] = None,
2512+
runtime: Optional[Union[str, types.Runtime]] = None,
25132513
agent: Optional["LlmAgent"] = None, # type: ignore # noqa: F821
25142514
gemini_agent: Optional[str] = None,
25152515
dest: Optional[str] = None,
@@ -2527,8 +2527,8 @@ def _execute_inference(
25272527
GCS path, or a BigQuery table) or a Pandas DataFrame.
25282528
model: The model to use for inference. Can be a callable function or a
25292529
string representing a model.
2530-
agent_engine: The agent engine to use for inference. Can be a resource
2531-
name string or an `AgentEngine` instance.
2530+
runtime: The agent engine to use for inference. Can be a resource
2531+
name string or an `Runtime` instance.
25322532
agent: The local agent to use for inference. Can be an ADK agent instance.
25332533
gemini_agent: The Gemini Agents API agent resource name to run inference
25342534
against via the Interactions API.
@@ -2549,10 +2549,9 @@ def _execute_inference(
25492549
if location:
25502550
api_client = _get_api_client_with_location(api_client, location)
25512551

2552-
if sum(x is not None for x in [model, agent_engine, agent, gemini_agent]) != 1:
2552+
if sum(x is not None for x in [model, runtime, agent, gemini_agent]) != 1:
25532553
raise ValueError(
2554-
"Exactly one of model, agent_engine, agent, or gemini_agent must be"
2555-
" provided."
2554+
"Exactly one of model, runtime, agent, or gemini_agent must be" " provided."
25562555
)
25572556

25582557
prompt_dataset = _load_dataframe(api_client, src)
@@ -2615,27 +2614,26 @@ def _execute_inference(
26152614
eval_dataset_df=results_df,
26162615
candidate_name=candidate_name,
26172616
)
2618-
elif agent_engine or agent:
2617+
elif runtime or agent:
26192618
candidate_name = None
2620-
if agent_engine:
2621-
candidate_name = "agent_engine_0"
2619+
if runtime:
2620+
candidate_name = "runtime_0"
26222621
elif agent:
26232622
agent_config = types.evals.AgentConfig.from_agent(agent)
26242623
candidate_name = agent_config.agent_id or "agent_0"
26252624

26262625
if (
2627-
agent_engine
2628-
and not isinstance(agent_engine, str)
2626+
runtime
2627+
and not isinstance(runtime, str)
26292628
and not (
2630-
hasattr(agent_engine, "api_client")
2631-
and type(agent_engine).__name__ == "AgentEngine"
2629+
hasattr(runtime, "api_client") and type(runtime).__name__ == "Runtime"
26322630
)
26332631
):
26342632
raise TypeError(
2635-
f"Unsupported agent_engine type: {type(agent_engine)}. Expecting a"
2633+
f"Unsupported runtime type: {type(runtime)}. Expecting a"
26362634
" string (agent engine resource name in"
26372635
" 'projects/{project_id}/locations/{location_id}/reasoningEngines/{reasoning_engine_id}'"
2638-
" format) or a types.AgentEngine instance."
2636+
" format) or a types.Runtime instance."
26392637
)
26402638
if (
26412639
_evals_constant.INTERMEDIATE_EVENTS in prompt_dataset.columns
@@ -2651,7 +2649,7 @@ def _execute_inference(
26512649
logger.debug("Starting Agent Run process ...")
26522650
results_df = _run_agent_internal(
26532651
api_client=api_client,
2654-
agent_engine=agent_engine,
2652+
runtime=runtime,
26552653
agent=agent,
26562654
prompt_dataset=prompt_dataset,
26572655
user_simulator_config=user_simulator_config,
@@ -2666,7 +2664,7 @@ def _execute_inference(
26662664
candidate_name=candidate_name,
26672665
)
26682666
else:
2669-
raise ValueError("Either model, agent_engine or agent must be provided.")
2667+
raise ValueError("Either model, runtime or agent must be provided.")
26702668

26712669
if dest:
26722670
file_name = "inference_results.jsonl" if model else "agent_run_results.jsonl"
@@ -3263,7 +3261,7 @@ def _create_agent_results_dataframe(
32633261

32643262
def _run_agent_internal(
32653263
api_client: BaseApiClient,
3266-
agent_engine: Optional[Union[str, types.AgentEngine]],
3264+
runtime: Optional[Union[str, types.Runtime]],
32673265
agent: Optional["LlmAgent"], # type: ignore # noqa: F821
32683266
prompt_dataset: pd.DataFrame,
32693267
user_simulator_config: Optional[types.evals.UserSimulatorConfig] = None,
@@ -3272,7 +3270,7 @@ def _run_agent_internal(
32723270
"""Runs an agent."""
32733271
raw_responses = _run_agent(
32743272
api_client=api_client,
3275-
agent_engine=agent_engine,
3273+
runtime=runtime,
32763274
agent=agent,
32773275
prompt_dataset=prompt_dataset,
32783276
user_simulator_config=user_simulator_config,
@@ -3314,7 +3312,7 @@ def _run_agent_internal(
33143312

33153313
def _run_agent(
33163314
api_client: BaseApiClient,
3317-
agent_engine: Optional[Union[str, types.AgentEngine]],
3315+
runtime: Optional[Union[str, types.Runtime]],
33183316
agent: Optional["LlmAgent"], # type: ignore # noqa: F821
33193317
prompt_dataset: pd.DataFrame,
33203318
user_simulator_config: Optional[types.evals.UserSimulatorConfig] = None,
@@ -3333,10 +3331,10 @@ def _run_agent(
33333331
simulator is never routed to a different region.
33343332
"""
33353333
del allow_cross_region_model # Simulator always runs in the client region.
3336-
if agent_engine:
3334+
if runtime:
33373335
return _execute_inference_concurrently(
33383336
api_client=api_client,
3339-
agent_engine=agent_engine,
3337+
runtime=runtime,
33403338
prompt_dataset=prompt_dataset,
33413339
progress_desc="Agent Run",
33423340
gemini_config=None,
@@ -3354,12 +3352,12 @@ def _run_agent(
33543352
inference_fn=_execute_local_agent_run_with_retry,
33553353
)
33563354
else:
3357-
raise ValueError("Neither agent_engine nor agent is provided.")
3355+
raise ValueError("Neither runtime nor agent is provided.")
33583356

33593357

3360-
def _create_agent_engine_session(
3358+
def _create_runtime_session(
33613359
*,
3362-
agent_engine: types.AgentEngine,
3360+
runtime: types.Runtime,
33633361
user_id: str,
33643362
session_state: Optional[dict[str, Any]] = None,
33653363
) -> Any:
@@ -3371,7 +3369,7 @@ def _create_agent_engine_session(
33713369
Sessions API.
33723370
33733371
Args:
3374-
agent_engine: The AgentEngine instance.
3372+
runtime: The Runtime instance.
33753373
user_id: The user ID for the session.
33763374
session_state: Optional initial state for the session.
33773375
@@ -3382,7 +3380,7 @@ def _create_agent_engine_session(
33823380
RuntimeError: If the session could not be created via either path.
33833381
"""
33843382
try:
3385-
session = agent_engine.create_session( # type: ignore[attr-defined]
3383+
session = runtime.create_session( # type: ignore[attr-defined]
33863384
user_id=user_id,
33873385
state=session_state,
33883386
)
@@ -3395,18 +3393,18 @@ def _create_agent_engine_session(
33953393
"Agent engine does not have 'create_session' operation registered."
33963394
" Falling back to managed Sessions API."
33973395
)
3398-
if agent_engine.api_resource is None:
3396+
if runtime.api_resource is None:
33993397
raise RuntimeError(
3400-
"Failed to create session: agent_engine.api_resource is None."
3398+
"Failed to create session: runtime.api_resource is None."
34013399
) from exc
3402-
if agent_engine.api_client is None:
3400+
if runtime.api_client is None:
34033401
raise RuntimeError(
3404-
"Failed to create session: agent_engine.api_client is None."
3402+
"Failed to create session: runtime.api_client is None."
34053403
) from exc
3406-
operation = agent_engine.api_client.sessions.create(
3407-
name=agent_engine.api_resource.name,
3404+
operation = runtime.api_client.sessions.create(
3405+
name=runtime.api_resource.name,
34083406
user_id=user_id,
3409-
config=types.CreateAgentEngineSessionConfig(
3407+
config=types.CreateRuntimeSessionConfig(
34103408
session_state=session_state,
34113409
),
34123410
)
@@ -3428,7 +3426,7 @@ def _create_agent_engine_session(
34283426
def _execute_agent_run_with_retry(
34293427
row: pd.Series,
34303428
contents: Union[genai_types.ContentListUnion, genai_types.ContentListUnionDict],
3431-
agent_engine: types.AgentEngine,
3429+
runtime: types.Runtime,
34323430
max_retries: int = 3,
34333431
) -> Union[list[dict[str, Any]], dict[str, Any]]:
34343432
"""Executes agent run over agent engine for a single prompt."""
@@ -3444,8 +3442,8 @@ def _execute_agent_run_with_retry(
34443442
return {"error": f"Failed to get all required agent engine inputs: {e}"}
34453443

34463444
try:
3447-
session_id = _create_agent_engine_session(
3448-
agent_engine=agent_engine,
3445+
session_id = _create_runtime_session(
3446+
runtime=runtime,
34493447
user_id=user_id,
34503448
session_state=session_state,
34513449
)
@@ -3463,19 +3461,19 @@ def _execute_agent_run_with_retry(
34633461
agent_data_obj = types.evals.AgentData.model_validate(agent_data_obj)
34643462
_, history_events = _extract_prompt_from_agent_data(agent_data_obj)
34653463

3466-
if agent_engine.api_resource is None:
3467-
return {"error": "agent_engine.api_resource is None."}
3468-
if agent_engine.api_client is None:
3469-
return {"error": "agent_engine.api_client is None."}
3470-
session_name = f"{agent_engine.api_resource.name}/sessions/{session_id}"
3464+
if runtime.api_resource is None:
3465+
return {"error": "runtime.api_resource is None."}
3466+
if runtime.api_client is None:
3467+
return {"error": "runtime.api_client is None."}
3468+
session_name = f"{runtime.api_resource.name}/sessions/{session_id}"
34713469
base_ts = datetime.datetime(2000, 1, 1, tzinfo=datetime.timezone.utc)
34723470
for i, ag_event in enumerate(history_events):
3473-
agent_engine.api_client.sessions.events.append(
3471+
runtime.api_client.sessions.events.append(
34743472
name=session_name,
34753473
author=ag_event.author or "user",
34763474
invocation_id="history",
34773475
timestamp=base_ts + datetime.timedelta(seconds=i),
3478-
config=types.AppendAgentEngineSessionEventConfig(
3476+
config=types.AppendRuntimeSessionEventConfig(
34793477
content=ag_event.content,
34803478
),
34813479
)
@@ -3484,7 +3482,7 @@ def _execute_agent_run_with_retry(
34843482
for attempt in range(max_retries):
34853483
try:
34863484
responses = []
3487-
for event in agent_engine.stream_query( # type: ignore[attr-defined]
3485+
for event in runtime.stream_query( # type: ignore[attr-defined]
34883486
user_id=user_id,
34893487
session_id=session_id,
34903488
message=contents,
@@ -4106,7 +4104,7 @@ def _create_evaluation_set_from_dataframe(
41064104
agent_data_obj = agent_data_val
41074105

41084106
# When agent_data exists but has no agents map (e.g. from remote
4109-
# agent_engine inference), inject the agents map from agent_info so
4107+
# runtime inference), inject the agents map from agent_info so
41104108
# the server-side autorater can access tool definitions and
41114109
# instructions.
41124110
if (

0 commit comments

Comments
 (0)