From fecccc68bf73d38bdab3262c7bfc9221244c4ddb Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Fri, 31 Jul 2026 03:46:37 -0700 Subject: [PATCH] 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 --- agentplatform/_genai/_evals_common.py | 130 +- ...nt_engines_utils.py => _runtimes_utils.py} | 195 +- agentplatform/_genai/agent_engines.py | 4115 ---- agentplatform/_genai/client.py | 115 +- agentplatform/_genai/evals.py | 22 +- agentplatform/_genai/feedback_contexts.py | 6 +- agentplatform/_genai/feedback_entries.py | 14 +- agentplatform/_genai/live.py | 20 +- ...live_agent_engines.py => live_runtimes.py} | 47 +- agentplatform/_genai/memories.py | 608 +- agentplatform/_genai/memory_revisions.py | 92 +- agentplatform/_genai/runtime_revisions.py | 328 +- agentplatform/_genai/runtimes.py | 3780 ++- agentplatform/_genai/sandbox_snapshots.py | 68 +- agentplatform/_genai/sandbox_templates.py | 46 +- agentplatform/_genai/sandboxes.py | 230 +- agentplatform/_genai/session_events.py | 120 +- agentplatform/_genai/sessions.py | 298 +- agentplatform/_genai/types/__init__.py | 1058 +- agentplatform/_genai/types/common.py | 19039 ++++++++-------- agentplatform/agent_engines/__init__.py | 428 - agentplatform/agent_engines/_agent_engines.py | 2028 -- agentplatform/frameworks/__init__.py | 39 + .../templates => frameworks}/a2a.py | 22 +- .../templates => frameworks}/adk.py | 79 +- .../templates => frameworks}/ag2.py | 91 +- .../templates => frameworks}/langchain.py | 44 +- .../templates => frameworks}/langgraph.py | 44 +- .../templates => frameworks}/llama_index.py | 50 +- .../frameworks/test_frameworks_a2a.py | 325 +- .../frameworks/test_frameworks_adk.py | 146 +- .../frameworks/test_frameworks_ag2.py | 72 +- .../frameworks/test_frameworks_langchain.py | 48 +- .../frameworks/test_frameworks_langgraph.py | 64 +- .../frameworks/test_frameworks_llama_index.py | 72 +- .../agentplatform/genai/replays/conftest.py | 18 +- .../genai/replays/test_ae_memories_delete.py | 6 +- .../genai/replays/test_ae_memories_get.py | 4 +- .../test_ae_memories_private_create.py | 6 +- .../test_ae_memories_private_generate.py | 6 +- ...private_get_generate_memories_operation.py | 6 +- ...e_memories_private_get_memory_operation.py | 6 +- .../replays/test_ae_memories_private_list.py | 4 +- .../replays/test_ae_memories_private_purge.py | 6 +- .../test_ae_memories_private_retrieve.py | 4 +- .../test_ae_memories_private_rollback.py | 6 +- .../test_ae_memories_private_update.py | 6 +- .../replays/test_ae_memory_revisions_get.py | 4 +- .../test_ae_memory_revisions_private_list.py | 6 +- .../test_ae_sandbox_snapshots_create.py | 6 +- .../test_ae_sandbox_snapshots_delete.py | 4 +- .../replays/test_ae_sandbox_snapshots_get.py | 4 +- .../replays/test_ae_sandbox_snapshots_list.py | 4 +- .../test_ae_sandbox_templates_byoc_create.py | 4 +- ...est_ae_sandbox_templates_default_create.py | 4 +- .../test_ae_sandbox_templates_delete.py | 4 +- .../replays/test_ae_sandbox_templates_get.py | 4 +- ...emplates_get_sandbox_template_operation.py | 4 +- .../replays/test_ae_sandbox_templates_list.py | 4 +- .../test_ae_sandboxes_private_create.py | 6 +- .../test_ae_sandboxes_private_delete.py | 6 +- .../test_ae_sandboxes_private_execute_code.py | 4 +- .../replays/test_ae_sandboxes_private_get.py | 8 +- ...sandboxes_private_get_sandbox_operation.py | 8 +- .../replays/test_ae_sandboxes_private_list.py | 8 +- .../genai/replays/test_ae_session_delete.py | 6 +- .../replays/test_ae_session_events_append.py | 6 +- .../test_ae_session_events_private_list.py | 6 +- .../replays/test_ae_session_private_create.py | 6 +- .../replays/test_ae_session_private_get.py | 6 +- .../replays/test_ae_session_private_list.py | 4 +- .../replays/test_ae_session_private_update.py | 8 +- .../replays/test_agent_engine_a2a_methods.py | 10 +- .../test_agent_engine_a2a_v1_methods.py | 10 +- .../test_agent_engine_private_create.py | 6 +- .../test_agent_engine_private_delete.py | 6 +- .../replays/test_agent_engine_private_get.py | 8 +- .../test_agent_engine_private_update.py | 8 +- ...est_append_agent_engine_a2a_task_events.py | 98 - .../test_append_agent_engine_session_event.py | 55 - .../genai/replays/test_create_agent_engine.py | 40 +- .../replays/test_create_agent_engine_a2a.py | 14 +- .../test_create_agent_engine_a2a_task.py | 173 - ...t_create_agent_engine_developer_connect.py | 14 +- .../test_create_agent_engine_docker.py | 16 +- .../test_create_agent_engine_memory.py | 56 +- .../test_create_agent_engine_sandbox.py | 18 +- .../test_create_agent_engine_session.py | 56 +- .../replays/test_create_feedback_entry.py | 24 +- .../test_delete_ae_runtime_revision.py | 64 +- .../genai/replays/test_delete_agent_engine.py | 24 +- .../test_delete_agent_engine_a2a_task.py | 96 - .../test_delete_agent_engine_memory.py | 28 +- .../test_delete_agent_engine_sandbox.py | 20 +- .../test_delete_agent_engine_session.py | 37 - .../replays/test_delete_feedback_entry.py | 24 +- .../genai/replays/test_evaluate_instances.py | 6 +- .../test_execute_code_agent_engine_sandbox.py | 18 +- .../test_generate_agent_engine_memories.py | 132 +- .../replays/test_get_ae_runtime_revision.py | 52 +- .../replays/test_get_agent_engine_a2a_task.py | 89 - .../replays/test_get_agent_engine_memory.py | 28 +- .../replays/test_get_agent_engine_sandbox.py | 18 +- .../replays/test_get_agent_engine_session.py | 39 - .../genai/replays/test_get_feedback_entry.py | 22 +- .../replays/test_ingest_events_memory_bank.py | 28 +- .../replays/test_list_ae_runtime_revisions.py | 44 +- .../test_list_agent_engine_a2a_task_events.py | 160 - .../test_list_agent_engine_a2a_tasks.py | 117 - .../test_list_agent_engine_memories.py | 42 +- .../test_list_agent_engine_sandboxes.py | 20 +- .../test_list_agent_engine_session_events.py | 28 +- .../test_list_agent_engine_sessions.py | 34 +- .../replays/test_list_feedback_entries.py | 32 +- .../test_purge_agent_engine_memories.py | 78 +- .../test_retrieve_agent_engine_memories.py | 86 +- .../genai/replays/test_run_inference.py | 49 +- .../genai/replays/test_structured_memories.py | 22 +- .../genai/replays/test_update_agent_engine.py | 20 +- .../replays/test_update_feedback_context.py | 16 +- .../replays/test_update_feedback_entry.py | 24 +- .../test_update_traffic_agent_engine.py | 74 +- .../test_agent_engine_runtime_revisions.py | 136 +- .../agentplatform/genai/test_agent_engines.py | 778 +- tests/unit/agentplatform/genai/test_evals.py | 116 +- .../genai/test_live_agent_engines.py | 14 +- .../unit/agentplatform/genai/test_sandbox.py | 4 +- 127 files changed, 16192 insertions(+), 21134 deletions(-) rename agentplatform/_genai/{_agent_engines_utils.py => _runtimes_utils.py} (93%) delete mode 100644 agentplatform/_genai/agent_engines.py rename agentplatform/_genai/{live_agent_engines.py => live_runtimes.py} (78%) delete mode 100644 agentplatform/agent_engines/__init__.py delete mode 100644 agentplatform/agent_engines/_agent_engines.py create mode 100644 agentplatform/frameworks/__init__.py rename agentplatform/{agent_engines/templates => frameworks}/a2a.py (97%) rename agentplatform/{agent_engines/templates => frameworks}/adk.py (96%) rename agentplatform/{agent_engines/templates => frameworks}/ag2.py (89%) rename agentplatform/{agent_engines/templates => frameworks}/langchain.py (95%) rename agentplatform/{agent_engines/templates => frameworks}/langgraph.py (95%) rename agentplatform/{agent_engines/templates => frameworks}/llama_index.py (93%) delete mode 100644 tests/unit/agentplatform/genai/replays/test_append_agent_engine_a2a_task_events.py delete mode 100644 tests/unit/agentplatform/genai/replays/test_append_agent_engine_session_event.py delete mode 100644 tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a_task.py delete mode 100644 tests/unit/agentplatform/genai/replays/test_delete_agent_engine_a2a_task.py delete mode 100644 tests/unit/agentplatform/genai/replays/test_delete_agent_engine_session.py delete mode 100644 tests/unit/agentplatform/genai/replays/test_get_agent_engine_a2a_task.py delete mode 100644 tests/unit/agentplatform/genai/replays/test_get_agent_engine_session.py delete mode 100644 tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_task_events.py delete mode 100644 tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_tasks.py diff --git a/agentplatform/_genai/_evals_common.py b/agentplatform/_genai/_evals_common.py index c6eb338c04..f8daedd2a7 100644 --- a/agentplatform/_genai/_evals_common.py +++ b/agentplatform/_genai/_evals_common.py @@ -139,21 +139,21 @@ def _get_api_client_with_location( )._api_client -def _get_agent_engine_instance( +def _get_runtime_instance( agent_name: str, api_client: BaseApiClient -) -> Union[types.AgentEngine, Any]: +) -> Union[types.Runtime, Any]: """Gets or creates an agent engine instance for the current thread.""" - if not hasattr(_thread_local_data, "agent_engine_instances"): - _thread_local_data.agent_engine_instances = {} - if agent_name not in _thread_local_data.agent_engine_instances: + if not hasattr(_thread_local_data, "runtime_instances"): + _thread_local_data.runtime_instances = {} + if agent_name not in _thread_local_data.runtime_instances: client = agentplatform.Client( project=api_client.project, location=api_client.location, ) - _thread_local_data.agent_engine_instances[agent_name] = ( - client.agent_engines.get(name=agent_name) + _thread_local_data.runtime_instances[agent_name] = client.runtimes.get( + name=agent_name ) - return _thread_local_data.agent_engine_instances[agent_name] + return _thread_local_data.runtime_instances[agent_name] def _generate_content_with_retry( @@ -1754,7 +1754,7 @@ def _execute_inference_concurrently( model_or_fn: Optional[Union[str, Callable[[Any], Any]]] = None, gemini_config: Optional[genai_types.GenerateContentConfig] = None, inference_fn: Optional[Callable[..., Any]] = None, - agent_engine: Optional[Union[str, types.AgentEngine]] = None, + runtime: Optional[Union[str, types.Runtime]] = None, agent: Optional["LlmAgent"] = None, # type: ignore # noqa: F821 user_simulator_config: Optional[types.evals.UserSimulatorConfig] = None, ) -> list[ @@ -1784,7 +1784,7 @@ def _execute_inference_concurrently( # prompt from the structured agent_data rather than requiring a flat # prompt/request column. has_agent_data = ( - agent is not None or agent_engine is not None + agent is not None or runtime is not None ) and AGENT_DATA in prompt_dataset.columns primary_prompt_column: Optional[str] = None @@ -1801,7 +1801,7 @@ def _execute_inference_concurrently( f" Found: {prompt_dataset.columns.tolist()}" ) - max_workers = AGENT_MAX_WORKERS if agent_engine or agent else MAX_WORKERS + max_workers = AGENT_MAX_WORKERS if runtime or agent else MAX_WORKERS with tqdm(total=len(prompt_dataset), desc=progress_desc) as pbar: with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: for index, row in prompt_dataset.iterrows(): @@ -1857,29 +1857,29 @@ def _execute_inference_concurrently( pbar.update(1) continue - if agent_engine or agent: + if runtime or agent: def agent_run_wrapper( # type: ignore[no-untyped-def] row_arg, contents_arg, - agent_engine_arg, + runtime_arg, agent_arg, inference_fn_arg, api_client_arg, user_simulator_config_arg, ) -> Any: - if agent_engine_arg: - if isinstance(agent_engine_arg, str): - agent_engine_instance = _get_agent_engine_instance( - agent_engine_arg, api_client_arg + if runtime_arg: + if isinstance(runtime_arg, str): + runtime_instance = _get_runtime_instance( + runtime_arg, api_client_arg ) else: - agent_engine_instance = agent_engine_arg + runtime_instance = runtime_arg return inference_fn_arg( row=row_arg, contents=contents_arg, - agent_engine=agent_engine_instance, + runtime=runtime_instance, ) elif agent_arg: return inference_fn_arg( @@ -1894,7 +1894,7 @@ def agent_run_wrapper( # type: ignore[no-untyped-def] agent_run_wrapper, row, contents, - agent_engine, + runtime, agent, inference_fn, api_client, @@ -2509,7 +2509,7 @@ def _execute_inference( api_client: BaseApiClient, src: Union[str, pd.DataFrame], model: Optional[Union[Callable[[Any], Any], str]] = None, - agent_engine: Optional[Union[str, types.AgentEngine]] = None, + runtime: Optional[Union[str, types.Runtime]] = None, agent: Optional["LlmAgent"] = None, # type: ignore # noqa: F821 gemini_agent: Optional[str] = None, dest: Optional[str] = None, @@ -2527,8 +2527,8 @@ def _execute_inference( GCS path, or a BigQuery table) or a Pandas DataFrame. model: The model to use for inference. Can be a callable function or a string representing a model. - agent_engine: The agent engine to use for inference. Can be a resource - name string or an `AgentEngine` instance. + runtime: The agent engine to use for inference. Can be a resource + name string or an `Runtime` instance. agent: The local agent to use for inference. Can be an ADK agent instance. gemini_agent: The Gemini Agents API agent resource name to run inference against via the Interactions API. @@ -2549,10 +2549,9 @@ def _execute_inference( if location: api_client = _get_api_client_with_location(api_client, location) - if sum(x is not None for x in [model, agent_engine, agent, gemini_agent]) != 1: + if sum(x is not None for x in [model, runtime, agent, gemini_agent]) != 1: raise ValueError( - "Exactly one of model, agent_engine, agent, or gemini_agent must be" - " provided." + "Exactly one of model, runtime, agent, or gemini_agent must be" " provided." ) prompt_dataset = _load_dataframe(api_client, src) @@ -2615,27 +2614,26 @@ def _execute_inference( eval_dataset_df=results_df, candidate_name=candidate_name, ) - elif agent_engine or agent: + elif runtime or agent: candidate_name = None - if agent_engine: - candidate_name = "agent_engine_0" + if runtime: + candidate_name = "runtime_0" elif agent: agent_config = types.evals.AgentConfig.from_agent(agent) candidate_name = agent_config.agent_id or "agent_0" if ( - agent_engine - and not isinstance(agent_engine, str) + runtime + and not isinstance(runtime, str) and not ( - hasattr(agent_engine, "api_client") - and type(agent_engine).__name__ == "AgentEngine" + hasattr(runtime, "api_client") and type(runtime).__name__ == "Runtime" ) ): raise TypeError( - f"Unsupported agent_engine type: {type(agent_engine)}. Expecting a" + f"Unsupported runtime type: {type(runtime)}. Expecting a" " string (agent engine resource name in" " 'projects/{project_id}/locations/{location_id}/reasoningEngines/{reasoning_engine_id}'" - " format) or a types.AgentEngine instance." + " format) or a types.Runtime instance." ) if ( _evals_constant.INTERMEDIATE_EVENTS in prompt_dataset.columns @@ -2651,7 +2649,7 @@ def _execute_inference( logger.debug("Starting Agent Run process ...") results_df = _run_agent_internal( api_client=api_client, - agent_engine=agent_engine, + runtime=runtime, agent=agent, prompt_dataset=prompt_dataset, user_simulator_config=user_simulator_config, @@ -2666,7 +2664,7 @@ def _execute_inference( candidate_name=candidate_name, ) else: - raise ValueError("Either model, agent_engine or agent must be provided.") + raise ValueError("Either model, runtime or agent must be provided.") if dest: file_name = "inference_results.jsonl" if model else "agent_run_results.jsonl" @@ -3263,7 +3261,7 @@ def _create_agent_results_dataframe( def _run_agent_internal( api_client: BaseApiClient, - agent_engine: Optional[Union[str, types.AgentEngine]], + runtime: Optional[Union[str, types.Runtime]], agent: Optional["LlmAgent"], # type: ignore # noqa: F821 prompt_dataset: pd.DataFrame, user_simulator_config: Optional[types.evals.UserSimulatorConfig] = None, @@ -3272,7 +3270,7 @@ def _run_agent_internal( """Runs an agent.""" raw_responses = _run_agent( api_client=api_client, - agent_engine=agent_engine, + runtime=runtime, agent=agent, prompt_dataset=prompt_dataset, user_simulator_config=user_simulator_config, @@ -3314,7 +3312,7 @@ def _run_agent_internal( def _run_agent( api_client: BaseApiClient, - agent_engine: Optional[Union[str, types.AgentEngine]], + runtime: Optional[Union[str, types.Runtime]], agent: Optional["LlmAgent"], # type: ignore # noqa: F821 prompt_dataset: pd.DataFrame, user_simulator_config: Optional[types.evals.UserSimulatorConfig] = None, @@ -3333,10 +3331,10 @@ def _run_agent( simulator is never routed to a different region. """ del allow_cross_region_model # Simulator always runs in the client region. - if agent_engine: + if runtime: return _execute_inference_concurrently( api_client=api_client, - agent_engine=agent_engine, + runtime=runtime, prompt_dataset=prompt_dataset, progress_desc="Agent Run", gemini_config=None, @@ -3354,12 +3352,12 @@ def _run_agent( inference_fn=_execute_local_agent_run_with_retry, ) else: - raise ValueError("Neither agent_engine nor agent is provided.") + raise ValueError("Neither runtime nor agent is provided.") -def _create_agent_engine_session( +def _create_runtime_session( *, - agent_engine: types.AgentEngine, + runtime: types.Runtime, user_id: str, session_state: Optional[dict[str, Any]] = None, ) -> Any: @@ -3371,7 +3369,7 @@ def _create_agent_engine_session( Sessions API. Args: - agent_engine: The AgentEngine instance. + runtime: The Runtime instance. user_id: The user ID for the session. session_state: Optional initial state for the session. @@ -3382,7 +3380,7 @@ def _create_agent_engine_session( RuntimeError: If the session could not be created via either path. """ try: - session = agent_engine.create_session( # type: ignore[attr-defined] + session = runtime.create_session( # type: ignore[attr-defined] user_id=user_id, state=session_state, ) @@ -3395,18 +3393,18 @@ def _create_agent_engine_session( "Agent engine does not have 'create_session' operation registered." " Falling back to managed Sessions API." ) - if agent_engine.api_resource is None: + if runtime.api_resource is None: raise RuntimeError( - "Failed to create session: agent_engine.api_resource is None." + "Failed to create session: runtime.api_resource is None." ) from exc - if agent_engine.api_client is None: + if runtime.api_client is None: raise RuntimeError( - "Failed to create session: agent_engine.api_client is None." + "Failed to create session: runtime.api_client is None." ) from exc - operation = agent_engine.api_client.sessions.create( - name=agent_engine.api_resource.name, + operation = runtime.api_client.sessions.create( + name=runtime.api_resource.name, user_id=user_id, - config=types.CreateAgentEngineSessionConfig( + config=types.CreateRuntimeSessionConfig( session_state=session_state, ), ) @@ -3428,7 +3426,7 @@ def _create_agent_engine_session( def _execute_agent_run_with_retry( row: pd.Series, contents: Union[genai_types.ContentListUnion, genai_types.ContentListUnionDict], - agent_engine: types.AgentEngine, + runtime: types.Runtime, max_retries: int = 3, ) -> Union[list[dict[str, Any]], dict[str, Any]]: """Executes agent run over agent engine for a single prompt.""" @@ -3444,8 +3442,8 @@ def _execute_agent_run_with_retry( return {"error": f"Failed to get all required agent engine inputs: {e}"} try: - session_id = _create_agent_engine_session( - agent_engine=agent_engine, + session_id = _create_runtime_session( + runtime=runtime, user_id=user_id, session_state=session_state, ) @@ -3463,19 +3461,19 @@ def _execute_agent_run_with_retry( agent_data_obj = types.evals.AgentData.model_validate(agent_data_obj) _, history_events = _extract_prompt_from_agent_data(agent_data_obj) - if agent_engine.api_resource is None: - return {"error": "agent_engine.api_resource is None."} - if agent_engine.api_client is None: - return {"error": "agent_engine.api_client is None."} - session_name = f"{agent_engine.api_resource.name}/sessions/{session_id}" + if runtime.api_resource is None: + return {"error": "runtime.api_resource is None."} + if runtime.api_client is None: + return {"error": "runtime.api_client is None."} + session_name = f"{runtime.api_resource.name}/sessions/{session_id}" base_ts = datetime.datetime(2000, 1, 1, tzinfo=datetime.timezone.utc) for i, ag_event in enumerate(history_events): - agent_engine.api_client.sessions.events.append( + runtime.api_client.sessions.events.append( name=session_name, author=ag_event.author or "user", invocation_id="history", timestamp=base_ts + datetime.timedelta(seconds=i), - config=types.AppendAgentEngineSessionEventConfig( + config=types.AppendRuntimeSessionEventConfig( content=ag_event.content, ), ) @@ -3484,7 +3482,7 @@ def _execute_agent_run_with_retry( for attempt in range(max_retries): try: responses = [] - for event in agent_engine.stream_query( # type: ignore[attr-defined] + for event in runtime.stream_query( # type: ignore[attr-defined] user_id=user_id, session_id=session_id, message=contents, @@ -4106,7 +4104,7 @@ def _create_evaluation_set_from_dataframe( agent_data_obj = agent_data_val # When agent_data exists but has no agents map (e.g. from remote - # agent_engine inference), inject the agents map from agent_info so + # runtime inference), inject the agents map from agent_info so # the server-side autorater can access tool definitions and # instructions. if ( diff --git a/agentplatform/_genai/_agent_engines_utils.py b/agentplatform/_genai/_runtimes_utils.py similarity index 93% rename from agentplatform/_genai/_agent_engines_utils.py rename to agentplatform/_genai/_runtimes_utils.py index d6b1820371..ca1cbce09a 100644 --- a/agentplatform/_genai/_agent_engines_utils.py +++ b/agentplatform/_genai/_runtimes_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""Utility functions for agent engines.""" +"""Utility functions for runtimes.""" import abc import asyncio @@ -182,9 +182,9 @@ _DEFAULT_ASYNC_STREAM_METHOD_RETURN_TYPE = "AsyncIterable[Any]" _DEFAULT_GCS_DIR_NAME = "agent_engine" _DEFAULT_METHOD_DOCSTRING_TEMPLATE = """ - Runs the Agent Engine to serve the user request. + Runs the Agent Runtime to serve the user request. This will be based on the `.{method_name}(...)` of the python object that - was passed in when creating the Agent Engine. The method will invoke the + was passed in when creating the Agent Runtime. The method will invoke the `{default_method_name}` API client of the python object. Args: **kwargs: @@ -236,25 +236,25 @@ @typing.runtime_checkable class Queryable(Protocol): - """Protocol for Agent Engines that can be queried.""" + """Protocol for Agent Runtimes that can be queried.""" @abc.abstractmethod def query(self, **kwargs): # type: ignore[no-untyped-def] - """Runs the Agent Engine to serve the user query.""" + """Runs the Agent Runtime to serve the user query.""" @typing.runtime_checkable class AsyncQueryable(Protocol): - """Protocol for Agent Engines that can be queried asynchronously.""" + """Protocol for Agent Runtimes that can be queried asynchronously.""" @abc.abstractmethod def async_query(self, **kwargs): # type: ignore[no-untyped-def] - """Runs the Agent Engine to serve the user query asynchronously.""" + """Runs the Agent Runtime to serve the user query asynchronously.""" @typing.runtime_checkable class AsyncStreamQueryable(Protocol): - """Protocol for Agent Engines that can stream responses asynchronously.""" + """Protocol for Agent Runtimes that can stream responses asynchronously.""" @abc.abstractmethod async def async_stream_query(self, **kwargs) -> AsyncIterator[Any]: # type: ignore[no-untyped-def] @@ -263,7 +263,7 @@ async def async_stream_query(self, **kwargs) -> AsyncIterator[Any]: # type: ign @typing.runtime_checkable class StreamQueryable(Protocol): - """Protocol for Agent Engines that can stream responses.""" + """Protocol for Agent Runtimes that can stream responses.""" @abc.abstractmethod def stream_query(self, **kwargs) -> Iterator[Any]: # type: ignore[no-untyped-def] @@ -272,7 +272,7 @@ def stream_query(self, **kwargs) -> Iterator[Any]: # type: ignore[no-untyped-de @typing.runtime_checkable class BidiStreamQueryable(Protocol): - """Protocol for Agent Engines that can stream requests and responses.""" + """Protocol for Agent Runtimes that can stream requests and responses.""" @abc.abstractmethod async def bidi_stream_query( @@ -283,7 +283,7 @@ async def bidi_stream_query( @typing.runtime_checkable class Cloneable(Protocol): - """Protocol for Agent Engines that can be cloned.""" + """Protocol for Agent Runtimes that can be cloned.""" @abc.abstractmethod def clone(self) -> Any: @@ -312,7 +312,7 @@ def register_operations(self, **kwargs: Any) -> dict[str, list[str]]: except (ImportError, AttributeError): ADKAgent = None # type: ignore[no-redef] -_AgentEngineInterface = Union[ +_RuntimeInterface = Union[ ADKAgent, AsyncQueryable, AsyncStreamQueryable, @@ -328,7 +328,7 @@ class _ModuleAgentAttributes(TypedDict, total=False): agent_name: str register_operations: Dict[str, list[str]] sys_paths: Optional[Sequence[str]] - agent: _AgentEngineInterface + agent: _RuntimeInterface class ModuleAgent(Cloneable, OperationRegistrable): @@ -441,24 +441,22 @@ class _RequirementsValidationResult(TypedDict): actions: _RequirementsValidationActions -AgentEngineOperationUnion = Union[ - genai_types.AgentEngineOperation, - genai_types.AgentEngineMemoryOperation, - genai_types.AgentEngineGenerateMemoriesOperation, +RuntimeOperationUnion = Union[ + genai_types.RuntimeOperation, + genai_types.RuntimeMemoryOperation, + genai_types.RuntimeGenerateMemoriesOperation, ] class GetOperationFunction(Protocol): - def __call__( - self, *, operation_name: str, **kwargs: Any - ) -> AgentEngineOperationUnion: + def __call__(self, *, operation_name: str, **kwargs: Any) -> RuntimeOperationUnion: pass class GetAsyncOperationFunction(Protocol): async def __call__( self, *, operation_name: str, **kwargs: Any - ) -> AgentEngineOperationUnion: + ) -> RuntimeOperationUnion: pass @@ -490,8 +488,7 @@ def _get_reasoning_engine_id(operation_name: str = "", resource_name: str = "") if match: return match.group(1) raise ValueError( - "Failed to parse reasoning engine ID from operation name: " - f"`{operation_name}`" + f"Failed to parse reasoning engine ID from operation name: `{operation_name}`" ) @@ -501,11 +498,11 @@ async def _await_async_operation( get_operation_fn: GetAsyncOperationFunction, poll_interval_seconds: float = 10, ) -> Any: - """Waits for the operation for creating an agent engine to complete. + """Waits for the operation for creating an agent runtime to complete. Args: operation_name (str): - Required. The name of the operation for creating the Agent Engine. + Required. The name of the operation for creating the Agent Runtime. poll_interval_seconds (float): The number of seconds to wait between each poll. get_operation_fn (Callable[[str], Awaitable[Any]]): @@ -529,11 +526,11 @@ def _await_operation( get_operation_fn: GetOperationFunction, poll_interval_seconds: float = 10, ) -> Any: - """Waits for the operation for creating an agent engine to complete. + """Waits for the operation for creating an agent runtime to complete. Args: operation_name (str): - Required. The name of the operation for creating the Agent Engine. + Required. The name of the operation for creating the Agent Runtime. poll_interval_seconds (float): The number of seconds to wait between each poll. get_operation_fn (Callable[[str], Any]): @@ -569,7 +566,7 @@ def _compare_requirements( required_packages (Iterator[str]): Optional. The set of packages that are required to be in the constraints. It defaults to the set of packages that are required - for deployment on Agent Engine. + for deployment on Agent Runtime. Returns: dict[str, dict[str, Any]]: The comparison result as a dictionary containing: @@ -611,7 +608,7 @@ def _compare_requirements( def _generate_class_methods_spec_or_raise( *, - agent: _AgentEngineInterface, + agent: _RuntimeInterface, operations: Dict[str, List[str]], ) -> List[proto.Message]: """Generates a ReasoningEngineSpec based on the registered operations. @@ -628,7 +625,7 @@ def _generate_class_methods_spec_or_raise( the AgentEngine. """ if isinstance(agent, ModuleAgent): - # We do a dry-run of setting up the agent engine to have the operations + # We do a dry-run of setting up the agent runtime to have the operations # needed for registration. agent: ModuleAgent = agent.clone() # type: ignore[no-redef] try: @@ -811,7 +808,7 @@ def _generate_schema( def _get_agent_framework( *, agent_framework: Optional[str], - agent: _AgentEngineInterface, + agent: _RuntimeInterface, ) -> Union[str, Any]: """Gets the agent framework to use. @@ -823,8 +820,8 @@ def _get_agent_framework( Args: agent_framework (str): The agent framework provided by the user. - agent (_AgentEngineInterface): - The agent engine instance. + agent (_RuntimeInterface): + The agent runtime instance. Returns: str: The name of the agent framework to use. @@ -871,7 +868,7 @@ def _get_gcs_bucket( def _get_registered_operations( *, - agent: _AgentEngineInterface, + agent: _RuntimeInterface, ) -> dict[str, list[str]]: """Retrieves registered operations for a AgentEngine.""" if isinstance(agent, OperationRegistrable): @@ -985,7 +982,7 @@ def _parse_constraints( def _prepare( *, - agent: Optional[_AgentEngineInterface], + agent: Optional[_RuntimeInterface], requirements: Optional[Sequence[str]], extra_packages: Optional[Sequence[str]], project: str, @@ -994,16 +991,16 @@ def _prepare( gcs_dir_name: str, credentials: Optional[Any] = None, ) -> None: - """Prepares the agent engine for creation or updates in Vertex AI. + """Prepares the agent runtime for creation or updates in Vertex AI. This involves packaging and uploading artifacts to Cloud Storage. Note that - 1. This does not actually update the Agent Engine in Vertex AI. + 1. This does not actually update the Agent Runtime in Vertex AI. 2. This will only generate and upload a pickled object if specified. 3. This will only generate and upload the dependencies.tar.gz file if extra_packages is non-empty. Args: - agent: The agent engine to be prepared. + agent: The agent runtime to be prepared. requirements (Sequence[str]): The set of PyPI dependencies needed. extra_packages (Sequence[str]): The set of extra user-provided packages. project (str): The project for the staging bucket. @@ -1021,7 +1018,7 @@ def _prepare( staging_bucket=staging_bucket, credentials=credentials, ) - _upload_agent_engine( + _upload_runtime( agent=agent, gcs_bucket=gcs_bucket, gcs_dir_name=gcs_dir_name, @@ -1042,21 +1039,21 @@ def _prepare( def _register_api_methods_or_raise( *, - agent_engine: genai_types.AgentEngine | genai_types.AgentEngineRuntimeRevision, + runtime: genai_types.Runtime | genai_types.RuntimeRevision, wrap_operation_fn: Optional[ dict[str, Callable[[str, str], Callable[..., Any]]] ] = None, ) -> None: - """Registers Agent Engine API methods based on operation schemas. + """Registers Agent Runtime API methods based on operation schemas. This function iterates through operation schemas provided by the - `agent_engine`. Each schema defines an API mode and method name. - It dynamically creates and registers methods on the `agent_engine` + `runtime`. Each schema defines an API mode and method name. + It dynamically creates and registers methods on the `runtime` to handle API calls based on the specified API mode. Currently, only standard API mode `` is supported. Args: - agent_engine: The AgentEngine to augment with API methods. + runtime: The AgentEngine to augment with API methods. wrap_operation_fn: A dictionary of API modes and method wrapping functions. @@ -1064,7 +1061,7 @@ def _register_api_methods_or_raise( ValueError: If the API mode is not supported or if the operation schema is missing any required fields (e.g. `api_mode` or `name`). """ - operation_schemas = agent_engine.operation_schemas() + operation_schemas = runtime.operation_schemas() if not operation_schemas: return for operation_schema in operation_schemas: @@ -1126,15 +1123,13 @@ def _register_api_methods_or_raise( # Bind the method to the object. if api_mode == _A2A_EXTENSION_MODE: agent_card = operation_schema.get(_A2A_AGENT_CARD) - method = _wrap_operation( - method_name=method_name, agent_card=agent_card - ) # type: ignore[call-arg] + method = _wrap_operation(method_name=method_name, agent_card=agent_card) # type: ignore[call-arg] else: method = _wrap_operation(method_name=method_name) # type: ignore[call-arg] method.__name__ = method_name if method_description and isinstance(method_description, str): method.__doc__ = method_description - setattr(agent_engine, method_name, types.MethodType(method, agent_engine)) + setattr(runtime, method_name, types.MethodType(method, runtime)) def _scan_requirements( @@ -1241,13 +1236,13 @@ def _to_proto( return message -def _upload_agent_engine( +def _upload_runtime( *, - agent: _AgentEngineInterface, + agent: _RuntimeInterface, gcs_bucket: _StorageBucket, gcs_dir_name: str, ) -> None: - """Uploads the agent engine to GCS.""" + """Uploads the agent runtime to GCS.""" cloudpickle = _import_cloudpickle_or_raise() blob = gcs_bucket.blob(f"{gcs_dir_name}/{_BLOB_FILENAME}") with blob.open("wb") as f: @@ -1255,7 +1250,7 @@ def _upload_agent_engine( cloudpickle.dump(agent, f) except Exception as e: url = "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#deployment-considerations" - error_msg = f"Failed to serialize agent engine. Visit {url} for details." + error_msg = f"Failed to serialize agent runtime. Visit {url} for details." if "google._upb._message" in str(e) or "Descriptor" in str(e): error_msg += ( " This is often caused by protobuf objects (like Part, AgentCard) " @@ -1269,7 +1264,7 @@ def _upload_agent_engine( try: _ = cloudpickle.load(f) except Exception as e: - raise TypeError("Agent engine serialized to an invalid format") from e + raise TypeError("agent runtime serialized to an invalid format") from e dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" logger.info(f"Wrote to {dir_name}/{_BLOB_FILENAME}") @@ -1412,7 +1407,7 @@ def _validate_staging_bucket_or_raise(*, staging_bucket: str) -> str: """Tries to validate the staging bucket.""" if not staging_bucket: raise ValueError( - "Please provide a `staging_bucket` in `client.agent_engines.create(...)`." + "Please provide a `staging_bucket` in `client.runtimes.create(...)`." ) if not staging_bucket.startswith("gs://"): raise ValueError(f"{staging_bucket=} must start with `gs://`") @@ -1475,11 +1470,11 @@ def _validate_requirements_or_raise( def _validate_agent_or_raise( *, - agent: _AgentEngineInterface, -) -> _AgentEngineInterface: - """Tries to validate the agent engine. + agent: _RuntimeInterface, +) -> _RuntimeInterface: + """Tries to validate the agent runtime. - The agent engine must have one of the following: + The agent runtime must have one of the following: * a callable method named `query` * a callable method named `stream_query` * a callable method named `async_stream_query` @@ -1490,12 +1485,12 @@ def _validate_agent_or_raise( agent: The agent to be validated. Returns: - The validated agent engine. + The validated agent runtime. Raises: - TypeError: If `agent_engine` has no callable method named `query`, + TypeError: If `runtime` has no callable method named `query`, `stream_query` or `register_operations`. - ValueError: If `agent_engine` has an invalid `query`, `stream_query` or + ValueError: If `runtime` has an invalid `query`, `stream_query` or `register_operations` signature. """ try: @@ -1503,9 +1498,9 @@ def _validate_agent_or_raise( if isinstance(agent, BaseAgent): logger.info("Deploying google.adk.agents.Agent as an application.") - from agentplatform import agent_engines + from agentplatform import frameworks - agent = agent_engines.AdkApp(agent=agent) + agent = frameworks.AdkApp(agent=agent) except Exception: pass is_queryable = isinstance(agent, Queryable) and callable(agent.query) @@ -1534,7 +1529,7 @@ def _validate_agent_or_raise( or is_async_stream_queryable ): raise TypeError( - "agent_engine has none of the following callable methods: " + "runtime has none of the following callable methods: " "`query`, `async_query`, `stream_query`, `async_stream_query`, " "`bidi_stream_query`, or `register_operations`." ) @@ -1613,10 +1608,10 @@ def _method(self, **kwargs) -> Any: # type: ignore[no-untyped-def] def _wrap_query_operation(*, method_name: str) -> Callable[..., Any]: - """Wraps an Agent Engine method, creating a callable for `query` API. + """Wraps an Agent Runtime method, creating a callable for `query` API. This function creates a callable object that executes the specified - Agent Engine method using the `query` API. It handles the creation of + Agent Runtime method using the `query` API. It handles the creation of the API request and the processing of the API response. The reserved keyword argument `http_options` is consumed by this @@ -1626,21 +1621,21 @@ def _wrap_query_operation(*, method_name: str) -> Callable[..., Any]: from google.genai.types import HttpOptions - agent_engine.query( + runtime.query( input="hello", http_options=HttpOptions(headers={"x-my-header": "value"}), ) Args: - method_name: The name of the Agent Engine method to call. + method_name: The name of the Agent Runtime method to call. doc: Documentation string for the method. Returns: - A callable object that executes the method on the Agent Engine via + A callable object that executes the method on the Agent Runtime via the `query` API. """ - def _method(self: genai_types.AgentEngine, **kwargs) -> Any: # type: ignore[no-untyped-def] + def _method(self: genai_types.Runtime, **kwargs) -> Any: # type: ignore[no-untyped-def] if not self.api_client: raise ValueError("api_client is not initialized.") if not self.api_resource: @@ -1663,10 +1658,10 @@ def _method(self: genai_types.AgentEngine, **kwargs) -> Any: # type: ignore[no- def _wrap_async_query_operation( *, method_name: str ) -> Callable[..., Coroutine[Any, Any, Any]]: - """Wraps an Agent Engine method, creating an async callable for `query` API. + """Wraps an Agent Runtime method, creating an async callable for `query` API. This function creates a callable object that executes the specified - Agent Engine method asynchronously using the `query` API. It handles the + Agent Runtime method asynchronously using the `query` API. It handles the creation of the API request and the processing of the API response. The reserved keyword argument `http_options` is consumed by this @@ -1674,16 +1669,16 @@ def _wrap_async_query_operation( `input`) and is propagated to the underlying HTTP call. Args: - method_name: The name of the Agent Engine method to call. + method_name: The name of the Agent Runtime method to call. doc: Documentation string for the method. Returns: - A callable object that executes the method on the Agent Engine via + A callable object that executes the method on the Agent Runtime via the `query` API. """ async def _method( - self: genai_types.AgentEngine, **kwargs: Any + self: genai_types.Runtime, **kwargs: Any ) -> Union[Coroutine[Any, Any, Any], Any]: if not self.api_async_client: raise ValueError("api_async_client is not initialized.") @@ -1705,10 +1700,10 @@ async def _method( def _wrap_stream_query_operation(*, method_name: str) -> Callable[..., Iterator[Any]]: - """Wraps an Agent Engine method, creating a callable for `stream_query` API. + """Wraps an Agent Runtime method, creating a callable for `stream_query` API. This function creates a callable object that executes the specified - Agent Engine method using the `stream_query` API. It handles the + Agent Runtime method using the `stream_query` API. It handles the creation of the API request and the processing of the API response. The reserved keyword argument `http_options` is consumed by this @@ -1716,15 +1711,15 @@ def _wrap_stream_query_operation(*, method_name: str) -> Callable[..., Iterator[ `input`) and is propagated to the underlying HTTP call. Args: - method_name: The name of the Agent Engine method to call. + method_name: The name of the Agent Runtime method to call. doc: Documentation string for the method. Returns: - A callable object that executes the method on the Agent Engine via + A callable object that executes the method on the Agent Runtime via the `stream_query` API. """ - def _method(self: genai_types.AgentEngine, **kwargs) -> Iterator[Any]: # type: ignore[no-untyped-def] + def _method(self: genai_types.Runtime, **kwargs) -> Iterator[Any]: # type: ignore[no-untyped-def] if not self.api_client: raise ValueError("api_client is not initialized.") if not self.api_resource: @@ -1749,10 +1744,10 @@ def _method(self: genai_types.AgentEngine, **kwargs) -> Iterator[Any]: # type: def _wrap_async_stream_query_operation( *, method_name: str ) -> Callable[..., AsyncIterator[Any]]: - """Wraps an Agent Engine method, creating an async callable for `stream_query` API. + """Wraps an Agent Runtime method, creating an async callable for `stream_query` API. This function creates a callable object that executes the specified - Agent Engine method using the `stream_query` API. It handles the + Agent Runtime method using the `stream_query` API. It handles the creation of the API request and the processing of the API response. The reserved keyword argument `http_options` is consumed by this @@ -1760,15 +1755,15 @@ def _wrap_async_stream_query_operation( `input`) and is propagated to the underlying HTTP call. Args: - method_name: The name of the Agent Engine method to call. + method_name: The name of the Agent Runtime method to call. doc: Documentation string for the method. Returns: - A callable object that executes the method on the Agent Engine via + A callable object that executes the method on the Agent Runtime via the `stream_query` API. """ - async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any]: # type: ignore[no-untyped-def] + async def _method(self: genai_types.Runtime, **kwargs) -> AsyncIterator[Any]: # type: ignore[no-untyped-def] if not self.api_client: raise ValueError("api_client is not initialized.") if not self.api_resource: @@ -1791,10 +1786,10 @@ async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any] def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list[Any]]: - """Wraps an Agent Engine method, creating a callable for A2A API. + """Wraps an Agent Runtime method, creating a callable for A2A API. Args: - method_name: The name of the Agent Engine method to call. + method_name: The name of the Agent Runtime method to call. agent_card: The agent card to use for the A2A API call. Example: { 'name': 'Sample Agent', 'description': ( 'A helpful assistant agent that can answer questions.' ), @@ -1810,7 +1805,7 @@ def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list ], 'inputModes': ['text'], 'outputModes': ['text'], }], } Returns: - A callable object that executes the method on the Agent Engine via + A callable object that executes the method on the Agent Runtime via the A2A API. """ @@ -1991,7 +1986,7 @@ def _validate_resource_limits_or_raise(resource_limits: dict[str, str]) -> None: if cpu not in [1, 2, 4, 6, 8]: raise ValueError( - "resource_limits['cpu'] must be one of 1, 2, 4, 6, 8. Got" f" {cpu}" + f"resource_limits['cpu'] must be one of 1, 2, 4, 6, 8. Got {cpu}" ) if not isinstance(memory_str, str) or not memory_str.endswith("Gi"): @@ -2004,8 +1999,7 @@ def _validate_resource_limits_or_raise(resource_limits: dict[str, str]) -> None: memory_gb = int(memory_str[:-2]) except ValueError: raise ValueError( - f"Invalid memory value: {memory_str}. Must be an integer" - " followed by 'Gi'." + f"Invalid memory value: {memory_str}. Must be an integer followed by 'Gi'." ) # https://cloud.google.com/run/docs/configuring/memory-limits @@ -2026,28 +2020,27 @@ def _validate_resource_limits_or_raise(resource_limits: dict[str, str]) -> None: if cpu < min_cpu: raise ValueError( - f"Memory size of {memory_str} requires at least {min_cpu} CPUs." - f" Got {cpu}" + f"Memory size of {memory_str} requires at least {min_cpu} CPUs. Got {cpu}" ) -def _is_adk_agent(agent_engine: _AgentEngineInterface) -> bool: - """Checks if the agent engine is an ADK agent. +def _is_adk_agent(runtime: _RuntimeInterface) -> bool: + """Checks if the agent runtime is an ADK agent. Args: - agent_engine: The agent engine to check. + runtime: The agent runtime to check. Returns: - True if the agent engine is an ADK agent, False otherwise. + True if the agent runtime is an ADK agent, False otherwise. """ - from agentplatform.agent_engines.templates import adk + from agentplatform.frameworks import AdkApp - return isinstance(agent_engine, adk.AdkApp) + return isinstance(runtime, AdkApp) def _add_telemetry_enablement_env( - env_vars: Optional[Dict[str, Union[str, Any]]] + env_vars: Optional[Dict[str, Union[str, Any]]], ) -> Optional[Dict[str, Union[str, Any]]]: """Adds telemetry enablement env var to the env vars. diff --git a/agentplatform/_genai/agent_engines.py b/agentplatform/_genai/agent_engines.py deleted file mode 100644 index ef4d6f2a09..0000000000 --- a/agentplatform/_genai/agent_engines.py +++ /dev/null @@ -1,4115 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -# Code generated by the Google Gen AI SDK generator DO NOT EDIT. - -import builtins -import datetime -import importlib -import json -import logging -import typing -from typing import Any, AsyncIterator, Iterator, Optional, Sequence, Tuple, Union -from urllib.parse import urlencode -import warnings - -from google.genai import _api_module -from google.genai import _common -from google.genai import types as genai_types -from google.genai._common import get_value_by_path as getv -from google.genai._common import set_value_by_path as setv -from google.genai.pagers import Pager - -from . import _agent_engines_utils -from . import types - -if typing.TYPE_CHECKING: - from . import sessions as sessions_module - from . import memories as memories_module - from . import a2a_tasks as a2a_tasks_module - from . import runtimes as runtimes_module - - _ = sessions_module - __ = memories_module - ___ = a2a_tasks_module - ____ = runtimes_module - - -logger = logging.getLogger("agentplatform_genai.agentengines") - -logger.setLevel(logging.INFO) - - -def _AgentEngineOperation_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["name"], getv(from_object, ["name"])) - - if getv(from_object, ["metadata"]) is not None: - setv(to_object, ["metadata"], getv(from_object, ["metadata"])) - - if getv(from_object, ["done"]) is not None: - setv(to_object, ["done"], getv(from_object, ["done"])) - - if getv(from_object, ["error"]) is not None: - setv(to_object, ["error"], getv(from_object, ["error"])) - - if getv(from_object, ["response"]) is not None: - setv( - to_object, - ["response"], - _ReasoningEngine_from_vertex(getv(from_object, ["response"]), to_object), - ) - - return to_object - - -def _CancelQueryJobAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["operation_name"]) is not None: - setv(parent_object, ["operationName"], getv(from_object, ["operation_name"])) - - return to_object - - -def _CancelQueryJobAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - if getv(from_object, ["config"]) is not None: - setv( - to_object, - ["config"], - _CancelQueryJobAgentEngineConfig_to_vertex( - getv(from_object, ["config"]), to_object - ), - ) - - return to_object - - -def _CheckQueryJobAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["retrieve_result"]) is not None: - setv(parent_object, ["retrieveResult"], getv(from_object, ["retrieve_result"])) - - return to_object - - -def _CheckQueryJobAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - if getv(from_object, ["config"]) is not None: - setv( - to_object, - ["config"], - _CheckQueryJobAgentEngineConfig_to_vertex( - getv(from_object, ["config"]), to_object - ), - ) - - return to_object - - -def _CheckQueryJobResult_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(parent_object, ["operationName"]) is not None: - setv(to_object, ["operation_name"], getv(parent_object, ["operationName"])) - - if getv(parent_object, ["outputGcsUri"]) is not None: - setv(to_object, ["output_gcs_uri"], getv(parent_object, ["outputGcsUri"])) - - if getv(parent_object, ["status"]) is not None: - setv(to_object, ["status"], getv(parent_object, ["status"])) - - if getv(parent_object, ["result"]) is not None: - setv(to_object, ["result"], getv(parent_object, ["result"])) - - return to_object - - -def _CreateAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["display_name"]) is not None: - setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) - - if getv(from_object, ["description"]) is not None: - setv(parent_object, ["description"], getv(from_object, ["description"])) - - if getv(from_object, ["spec"]) is not None: - setv(parent_object, ["spec"], getv(from_object, ["spec"])) - - if getv(from_object, ["context_spec"]) is not None: - setv( - parent_object, - ["contextSpec"], - _ReasoningEngineContextSpec_to_vertex( - getv(from_object, ["context_spec"]), to_object - ), - ) - - if getv(from_object, ["psc_interface_config"]) is not None: - setv( - parent_object, - ["pscInterfaceConfig"], - getv(from_object, ["psc_interface_config"]), - ) - - if getv(from_object, ["encryption_spec"]) is not None: - setv(parent_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"])) - - if getv(from_object, ["labels"]) is not None: - setv(parent_object, ["labels"], getv(from_object, ["labels"])) - - if getv(from_object, ["source_packages"]) is not None: - setv(parent_object, ["sourcePackages"], getv(from_object, ["source_packages"])) - - if getv(from_object, ["entrypoint_module"]) is not None: - setv( - parent_object, - ["entrypointModule"], - getv(from_object, ["entrypoint_module"]), - ) - - if getv(from_object, ["entrypoint_object"]) is not None: - setv( - parent_object, - ["entrypointObject"], - getv(from_object, ["entrypoint_object"]), - ) - - if getv(from_object, ["requirements_file"]) is not None: - setv( - parent_object, - ["requirementsFile"], - getv(from_object, ["requirements_file"]), - ) - - if getv(from_object, ["agent_framework"]) is not None: - setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) - - if getv(from_object, ["python_version"]) is not None: - setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) - - if getv(from_object, ["agent_gateway_config"]) is not None: - setv( - parent_object, - ["agentGatewayConfig"], - getv(from_object, ["agent_gateway_config"]), - ) - - return to_object - - -def _CreateAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["config"]) is not None: - _CreateAgentEngineConfig_to_vertex(getv(from_object, ["config"]), to_object) - - return to_object - - -def _DeleteAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - if getv(from_object, ["force"]) is not None: - setv(to_object, ["force"], getv(from_object, ["force"])) - - return to_object - - -def _GetAgentEngineOperationParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["operation_name"]) is not None: - setv( - to_object, ["_url", "operationName"], getv(from_object, ["operation_name"]) - ) - - return to_object - - -def _GetAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - return to_object - - -def _ListAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["page_size"]) is not None: - setv(parent_object, ["_query", "pageSize"], getv(from_object, ["page_size"])) - - if getv(from_object, ["page_token"]) is not None: - setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"])) - - if getv(from_object, ["filter"]) is not None: - setv(parent_object, ["_query", "filter"], getv(from_object, ["filter"])) - - return to_object - - -def _ListAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["config"]) is not None: - _ListAgentEngineConfig_to_vertex(getv(from_object, ["config"]), to_object) - - return to_object - - -def _ListReasoningEnginesResponse_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["sdkHttpResponse"]) is not None: - setv(to_object, ["sdk_http_response"], getv(from_object, ["sdkHttpResponse"])) - - if getv(from_object, ["nextPageToken"]) is not None: - setv(to_object, ["next_page_token"], getv(from_object, ["nextPageToken"])) - - if getv(from_object, ["reasoningEngines"]) is not None: - setv( - to_object, - ["reasoning_engines"], - [ - _ReasoningEngine_from_vertex(item, to_object) - for item in getv(from_object, ["reasoningEngines"]) - ], - ) - - return to_object - - -def _QueryAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["class_method"]) is not None: - setv(parent_object, ["classMethod"], getv(from_object, ["class_method"])) - - if getv(from_object, ["input"]) is not None: - setv(parent_object, ["input"], getv(from_object, ["input"])) - - if getv(from_object, ["include_all_fields"]) is not None: - setv(to_object, ["includeAllFields"], getv(from_object, ["include_all_fields"])) - - return to_object - - -def _QueryAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - if getv(from_object, ["config"]) is not None: - _QueryAgentEngineConfig_to_vertex(getv(from_object, ["config"]), to_object) - - return to_object - - -def _ReasoningEngineContextSpecMemoryBankConfig_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["customizationConfigs"]) is not None: - setv( - to_object, - ["customization_configs"], - [item for item in getv(from_object, ["customizationConfigs"])], - ) - - if getv(from_object, ["disableMemoryRevisions"]) is not None: - setv( - to_object, - ["disable_memory_revisions"], - getv(from_object, ["disableMemoryRevisions"]), - ) - - if getv(from_object, ["generationConfig"]) is not None: - setv(to_object, ["generation_config"], getv(from_object, ["generationConfig"])) - - if getv(from_object, ["similaritySearchConfig"]) is not None: - setv( - to_object, - ["similarity_search_config"], - getv(from_object, ["similaritySearchConfig"]), - ) - - if getv(from_object, ["ttlConfig"]) is not None: - setv(to_object, ["ttl_config"], getv(from_object, ["ttlConfig"])) - - if getv(from_object, ["structuredMemoryConfigs"]) is not None: - setv( - to_object, - ["structured_memory_configs"], - [ - _StructuredMemoryConfig_from_vertex(item, to_object) - for item in getv(from_object, ["structuredMemoryConfigs"]) - ], - ) - - return to_object - - -def _ReasoningEngineContextSpecMemoryBankConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["customization_configs"]) is not None: - setv( - to_object, - ["customizationConfigs"], - [item for item in getv(from_object, ["customization_configs"])], - ) - - if getv(from_object, ["disable_memory_revisions"]) is not None: - setv( - to_object, - ["disableMemoryRevisions"], - getv(from_object, ["disable_memory_revisions"]), - ) - - if getv(from_object, ["generation_config"]) is not None: - setv(to_object, ["generationConfig"], getv(from_object, ["generation_config"])) - - if getv(from_object, ["similarity_search_config"]) is not None: - setv( - to_object, - ["similaritySearchConfig"], - getv(from_object, ["similarity_search_config"]), - ) - - if getv(from_object, ["ttl_config"]) is not None: - setv(to_object, ["ttlConfig"], getv(from_object, ["ttl_config"])) - - if getv(from_object, ["structured_memory_configs"]) is not None: - setv( - to_object, - ["structuredMemoryConfigs"], - [ - _StructuredMemoryConfig_to_vertex(item, to_object) - for item in getv(from_object, ["structured_memory_configs"]) - ], - ) - - return to_object - - -def _ReasoningEngineContextSpec_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["memoryBankConfig"]) is not None: - setv( - to_object, - ["memory_bank_config"], - _ReasoningEngineContextSpecMemoryBankConfig_from_vertex( - getv(from_object, ["memoryBankConfig"]), to_object - ), - ) - - return to_object - - -def _ReasoningEngineContextSpec_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["memory_bank_config"]) is not None: - setv( - to_object, - ["memoryBankConfig"], - _ReasoningEngineContextSpecMemoryBankConfig_to_vertex( - getv(from_object, ["memory_bank_config"]), to_object - ), - ) - - return to_object - - -def _ReasoningEngine_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["encryptionSpec"]) is not None: - setv(to_object, ["encryption_spec"], getv(from_object, ["encryptionSpec"])) - - if getv(from_object, ["contextSpec"]) is not None: - setv( - to_object, - ["context_spec"], - _ReasoningEngineContextSpec_from_vertex( - getv(from_object, ["contextSpec"]), to_object - ), - ) - - if getv(from_object, ["createTime"]) is not None: - setv(to_object, ["create_time"], getv(from_object, ["createTime"])) - - if getv(from_object, ["description"]) is not None: - setv(to_object, ["description"], getv(from_object, ["description"])) - - if getv(from_object, ["displayName"]) is not None: - setv(to_object, ["display_name"], getv(from_object, ["displayName"])) - - if getv(from_object, ["etag"]) is not None: - setv(to_object, ["etag"], getv(from_object, ["etag"])) - - if getv(from_object, ["labels"]) is not None: - setv(to_object, ["labels"], getv(from_object, ["labels"])) - - if getv(from_object, ["name"]) is not None: - setv(to_object, ["name"], getv(from_object, ["name"])) - - if getv(from_object, ["spec"]) is not None: - setv(to_object, ["spec"], getv(from_object, ["spec"])) - - if getv(from_object, ["updateTime"]) is not None: - setv(to_object, ["update_time"], getv(from_object, ["updateTime"])) - - if getv(from_object, ["trafficConfig"]) is not None: - setv(to_object, ["traffic_config"], getv(from_object, ["trafficConfig"])) - - return to_object - - -def _RunQueryJobAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["input_gcs_uri"]) is not None: - setv(parent_object, ["inputGcsUri"], getv(from_object, ["input_gcs_uri"])) - - if getv(from_object, ["output_gcs_uri"]) is not None: - setv(parent_object, ["outputGcsUri"], getv(from_object, ["output_gcs_uri"])) - - return to_object - - -def _RunQueryJobAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - if getv(from_object, ["config"]) is not None: - setv( - to_object, - ["config"], - _RunQueryJobAgentEngineConfig_to_vertex( - getv(from_object, ["config"]), to_object - ), - ) - - return to_object - - -def _StructuredMemoryConfig_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["schemaConfigs"]) is not None: - setv( - to_object, - ["schema_configs"], - [ - _StructuredMemorySchemaConfig_from_vertex(item, to_object) - for item in getv(from_object, ["schemaConfigs"]) - ], - ) - - if getv(from_object, ["scopeKeys"]) is not None: - setv(to_object, ["scope_keys"], getv(from_object, ["scopeKeys"])) - - return to_object - - -def _StructuredMemoryConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["schema_configs"]) is not None: - setv( - to_object, - ["schemaConfigs"], - [ - _StructuredMemorySchemaConfig_to_vertex(item, to_object) - for item in getv(from_object, ["schema_configs"]) - ], - ) - - if getv(from_object, ["scope_keys"]) is not None: - setv(to_object, ["scopeKeys"], getv(from_object, ["scope_keys"])) - - return to_object - - -def _StructuredMemorySchemaConfig_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["schema"]) is not None: - setv(to_object, ["memory_schema"], getv(from_object, ["schema"])) - - if getv(from_object, ["id"]) is not None: - setv(to_object, ["id"], getv(from_object, ["id"])) - - if getv(from_object, ["memoryType"]) is not None: - setv(to_object, ["memory_type"], getv(from_object, ["memoryType"])) - - return to_object - - -def _StructuredMemorySchemaConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["memory_schema"]) is not None: - setv(to_object, ["schema"], getv(from_object, ["memory_schema"])) - - if getv(from_object, ["id"]) is not None: - setv(to_object, ["id"], getv(from_object, ["id"])) - - if getv(from_object, ["memory_type"]) is not None: - setv(to_object, ["memoryType"], getv(from_object, ["memory_type"])) - - return to_object - - -def _UpdateAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["display_name"]) is not None: - setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) - - if getv(from_object, ["description"]) is not None: - setv(parent_object, ["description"], getv(from_object, ["description"])) - - if getv(from_object, ["spec"]) is not None: - setv(parent_object, ["spec"], getv(from_object, ["spec"])) - - if getv(from_object, ["context_spec"]) is not None: - setv( - parent_object, - ["contextSpec"], - _ReasoningEngineContextSpec_to_vertex( - getv(from_object, ["context_spec"]), to_object - ), - ) - - if getv(from_object, ["psc_interface_config"]) is not None: - setv( - parent_object, - ["pscInterfaceConfig"], - getv(from_object, ["psc_interface_config"]), - ) - - if getv(from_object, ["encryption_spec"]) is not None: - setv(parent_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"])) - - if getv(from_object, ["labels"]) is not None: - setv(parent_object, ["labels"], getv(from_object, ["labels"])) - - if getv(from_object, ["source_packages"]) is not None: - setv(parent_object, ["sourcePackages"], getv(from_object, ["source_packages"])) - - if getv(from_object, ["entrypoint_module"]) is not None: - setv( - parent_object, - ["entrypointModule"], - getv(from_object, ["entrypoint_module"]), - ) - - if getv(from_object, ["entrypoint_object"]) is not None: - setv( - parent_object, - ["entrypointObject"], - getv(from_object, ["entrypoint_object"]), - ) - - if getv(from_object, ["requirements_file"]) is not None: - setv( - parent_object, - ["requirementsFile"], - getv(from_object, ["requirements_file"]), - ) - - if getv(from_object, ["agent_framework"]) is not None: - setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) - - if getv(from_object, ["python_version"]) is not None: - setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) - - if getv(from_object, ["agent_gateway_config"]) is not None: - setv( - parent_object, - ["agentGatewayConfig"], - getv(from_object, ["agent_gateway_config"]), - ) - - if getv(from_object, ["update_mask"]) is not None: - setv( - parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"]) - ) - - if getv(from_object, ["traffic_config"]) is not None: - setv(parent_object, ["trafficConfig"], getv(from_object, ["traffic_config"])) - - return to_object - - -def _UpdateAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - if getv(from_object, ["config"]) is not None: - _UpdateAgentEngineConfig_to_vertex(getv(from_object, ["config"]), to_object) - - return to_object - - -class AgentEngines(_api_module.BaseModule): - - def cancel_query_job( - self, - *, - name: str, - config: Optional[types.CancelQueryJobAgentEngineConfigOrDict] = None, - ) -> types.CancelQueryJobResult: - """ - Cancels a long-running query job on an Agent Engine. - - Args: - name (str): - Required. The reasoning engine resource name. - config (CancelQueryJobAgentEngineConfigOrDict): - Optional. The configuration for the cancel_query_job. - - """ - - parameter_model = types._CancelQueryJobAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _CancelQueryJobAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:cancelAsyncQuery".format_map(request_url_dict) - else: - path = "{name}:cancelAsyncQuery" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("post", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.CancelQueryJobResult._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _check_query_job( - self, - *, - name: str, - config: Optional[types.CheckQueryJobAgentEngineConfigOrDict] = None, - ) -> types.CheckQueryJobResult: - """ - Query an Agent Engine asynchronously. - """ - - parameter_model = types._CheckQueryJobAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _CheckQueryJobAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:checkQueryJob".format_map(request_url_dict) - else: - path = "{name}:checkQueryJob" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("post", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _CheckQueryJobResult_from_vertex(response_dict) - - return_value = types.CheckQueryJobResult._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _run_query_job( - self, - *, - name: str, - config: Optional[types._RunQueryJobAgentEngineConfigOrDict] = None, - ) -> types.AgentEngineOperation: - """ - Run a query job on an agent engine. - """ - - parameter_model = types._RunQueryJobAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _RunQueryJobAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:asyncQuery".format_map(request_url_dict) - else: - path = "{name}:asyncQuery" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("post", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _create( - self, *, config: Optional[types.CreateAgentEngineConfigOrDict] = None - ) -> types.AgentEngineOperation: - """ - Creates a new Agent Engine. - """ - - parameter_model = types._CreateAgentEngineRequestParameters( - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _CreateAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "reasoningEngines".format_map(request_url_dict) - else: - path = "reasoningEngines" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("post", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _delete( - self, - *, - name: str, - force: Optional[bool] = None, - config: Optional[types.DeleteAgentEngineConfigOrDict] = None, - ) -> types.DeleteAgentEngineOperation: - """ - Delete an Agent Engine resource. - - Args: - name (str): - Required. The name of the Agent Engine to be deleted. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` - or `reasoningEngines/{resource_id}`. - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error when - the Agent Engine has undeleted child resources. Defaults to False. - config (DeleteAgentEngineConfig): - Optional. Additional configurations for deleting the Agent Engine. - - """ - - parameter_model = types._DeleteAgentEngineRequestParameters( - name=name, - force=force, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _DeleteAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}".format_map(request_url_dict) - else: - path = "{name}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("delete", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.DeleteAgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _get( - self, *, name: str, config: Optional[types.GetAgentEngineConfigOrDict] = None - ) -> types.ReasoningEngine: - """ - Get an Agent Engine instance. - """ - - parameter_model = types._GetAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _GetAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}".format_map(request_url_dict) - else: - path = "{name}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("get", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _ReasoningEngine_from_vertex(response_dict) - - return_value = types.ReasoningEngine._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _list( - self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None - ) -> types.ListReasoningEnginesResponse: - """ - Lists Agent Engines. - """ - - parameter_model = types._ListAgentEngineRequestParameters( - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _ListAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "reasoningEngines".format_map(request_url_dict) - else: - path = "reasoningEngines" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("get", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) - - return_value = types.ListReasoningEnginesResponse._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _get_agent_operation( - self, - *, - operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineOperation: - parameter_model = types._GetAgentEngineOperationParameters( - operation_name=operation_name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _GetAgentEngineOperationParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{operationName}".format_map(request_url_dict) - else: - path = "{operationName}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("get", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _query( - self, *, name: str, config: Optional[types.QueryAgentEngineConfigOrDict] = None - ) -> types.QueryReasoningEngineResponse: - """ - Query an Agent Engine. - """ - - parameter_model = types._QueryAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _QueryAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:query".format_map(request_url_dict) - else: - path = "{name}:query" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("post", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.QueryReasoningEngineResponse._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _update( - self, *, name: str, config: Optional[types.UpdateAgentEngineConfigOrDict] = None - ) -> types.AgentEngineOperation: - """ - Updates an Agent Engine. - """ - - parameter_model = types._UpdateAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _UpdateAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}".format_map(request_url_dict) - else: - path = "{name}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("patch", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - _a2a_tasks = None - _memories = None - _sandboxes = None - _sessions = None - _runtimes = None - - @property - def runtimes(self) -> "runtimes_module.Runtimes": - if self._runtimes is None: - try: - # We need to lazy load the runtimes module to handle the - # possibility of ImportError when dependencies are not installed. - self._runtimes = importlib.import_module(".runtimes", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.runtimes' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._runtimes.Runtimes(self._api_client) # type: ignore[no-any-return] - - @property - def a2a_tasks(self) -> "a2a_tasks_module.A2aTasks": - if self._a2a_tasks is None: - try: - # We need to lazy load the a2a_tasks module to handle the - # possibility of ImportError when dependencies are not installed. - self._a2a_tasks = importlib.import_module(".a2a_tasks", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.a2a_tasks' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._a2a_tasks.A2aTasks(self._api_client) # type: ignore[no-any-return] - - @property - def memories(self) -> "memories_module.Memories": - if self._memories is None: - try: - # We need to lazy load the memories module to handle the - # possibility of ImportError when dependencies are not installed. - self._memories = importlib.import_module(".memories", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.memories' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._memories.Memories(self._api_client) # type: ignore[no-any-return] - - @property - def sandboxes(self) -> Any: - if self._sandboxes is None: - try: - # We need to lazy load the sandboxes module to handle the - # possibility of ImportError when dependencies are not installed. - self._sandboxes = importlib.import_module(".sandboxes", __package__) - except ImportError as e: - raise ImportError( - "The agent_engines.sandboxes module requires additional packages. " - "Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._sandboxes.Sandboxes(self._api_client) - - @property - def sessions(self) -> "sessions_module.Sessions": - if self._sessions is None: - try: - # We need to lazy load the sessions module to handle the - # possibility of ImportError when dependencies are not installed. - self._sessions = importlib.import_module(".sessions", __package__) - except ImportError as e: - raise ImportError( - "The agent_engines.sessions module requires additional packages. " - "Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._sessions.Sessions(self._api_client) # type: ignore[no-any-return] - - def _list_pager( - self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None - ) -> Pager[types.ReasoningEngine]: - return Pager( - "reasoning_engines", - self._list, - self._list(config=config), - config, - ) - - def check_query_job( - self, - *, - name: str, - config: Optional[types.CheckQueryJobAgentEngineConfigOrDict] = None, - ) -> types.CheckQueryJobResult: - """Checks a query job on an agent engine and optionally returns the results. - - Args: - name (str): - Required. A fully-qualified resource name or ID. - config (CheckQueryJobAgentEngineConfigOrDict): - Optional. The configuration for the check_query_job. If not provided, - the default configuration will be used. This can be used to specify - the following fields: - - retrieve_result: Whether to retrieve the results of the query job. - """ - from google.cloud import storage # type: ignore[attr-defined] - import json - - if config is None: - config = types.CheckQueryJobAgentEngineConfig() - elif isinstance(config, dict): - config = types.CheckQueryJobAgentEngineConfig(**config) - - raw_response = self._api_client.request("get", name, {}) - if hasattr(raw_response, "body"): - operation = ( - json.loads(raw_response.body) - if isinstance(raw_response.body, str) - else raw_response.body - ) - else: - operation = raw_response - - status = "RUNNING" - if isinstance(operation, dict): - if operation.get("done"): - status = "FAILED" if operation.get("error") else "SUCCESS" - - response_dict = operation.get("response", {}) - output_gcs_uri = response_dict.get("outputGcsUri") or response_dict.get( - "output_gcs_uri" - ) - error = operation.get("error") - else: - if getattr(operation, "done", False): - status = "FAILED" if getattr(operation, "error", None) else "SUCCESS" - - response_obj = getattr(operation, "response", None) - if isinstance(response_obj, dict): - output_gcs_uri = response_obj.get("outputGcsUri") or response_obj.get( - "output_gcs_uri" - ) - else: - output_gcs_uri = ( - getattr( - response_obj, - "output_gcs_uri", - getattr(response_obj, "outputGcsUri", None), - ) - if response_obj - else None - ) - error = getattr(operation, "error", None) - - result_str = None - if status == "SUCCESS" and config.retrieve_result and output_gcs_uri: - storage_client = storage.Client( - project=self._api_client.project, - credentials=self._api_client._credentials, - ) - bucket_name = output_gcs_uri.replace("gs://", "").split("/")[0] - blob_name = output_gcs_uri.replace(f"gs://{bucket_name}/", "") - bucket = storage_client.bucket(bucket_name) - blob = bucket.blob(blob_name) - if blob.exists(): - result_str = blob.download_as_string().decode("utf-8") - else: - raise ValueError( - f"Failed to retrieve blob results for {output_gcs_uri}" - ) - - elif status == "FAILED" and error: - result_str = str(error) - - return types.CheckQueryJobResult( - operation_name=name, - output_gcs_uri=output_gcs_uri, - status=status, - result=result_str, - ) - - def _is_lightweight_creation( - self, agent: Any, config: types.AgentEngineConfig - ) -> bool: - if ( - agent - or config.source_packages - or config.developer_connect_source - or config.agent_config_source - or config.container_spec - ): - return False - return True - - def run_query_job( - self, - *, - name: str, - config: Optional[types.RunQueryJobAgentEngineConfigOrDict] = None, - ) -> types.RunQueryJobResult: - """Launches a long-running query job on an Agent Engine - - Args: - name (str): - Required. A fully-qualified resource name or ID. - config (RunQueryJobAgentEngineConfigOrDict): - Optional. The configuration for the async query. If not provided, - the default configuration will be used. This can be used to specify - the following fields: - - query: The query to send to the agent engine. - - output_gcs_uri: The GCS URI to use for the output. - """ - from google.cloud import storage # type: ignore[attr-defined] - from google.api_core import exceptions - import uuid - - if config is None: - config = types.RunQueryJobAgentEngineConfig() - elif isinstance(config, dict): - config = types.RunQueryJobAgentEngineConfig(**config) - - if not config.query: - raise ValueError("`query` is required in the config object.") - if not config.output_gcs_uri: - raise ValueError("`output_gcs_uri` is required in the config object.") - - output_gcs_uri = config.output_gcs_uri - is_file = False - last_part = "" - if not output_gcs_uri.endswith("/"): - last_part = output_gcs_uri.split("/")[-1] - if "." in last_part: - is_file = True - - if is_file: - path_parts = output_gcs_uri.split("/") - file_name = path_parts[-1] - base_uri = "/".join(path_parts[:-1]) - name_parts = file_name.rsplit(".", 1) - if len(name_parts) == 2: - name_part, ext = name_parts[0], "." + name_parts[1] - else: - name_part = name_parts[0] - ext = "" - input_gcs_uri = f"{base_uri}/{name_part}_input{ext}" - else: - job_uuid = uuid.uuid4().hex - gcs_path = output_gcs_uri.rstrip("/") - input_gcs_uri = f"{gcs_path}/{job_uuid}_input.json" - output_gcs_uri = f"{gcs_path}/{job_uuid}_output.json" - - storage_client = storage.Client( - project=self._api_client.project, credentials=self._api_client._credentials - ) - - # Handle creating the bucket if it does not exist - bucket_name = config.output_gcs_uri.replace("gs://", "").split("/")[0] - bucket = storage_client.bucket(bucket_name) - - try: - bucket_exists = bucket.exists() - except exceptions.Forbidden as e: - raise ValueError( - f"Permission denied to check existence of bucket '{bucket_name}'. " - "The service account may lack 'storage.buckets.get' permission." - ) from e - - if not bucket_exists: - try: - bucket.create() - except exceptions.Forbidden as e: - raise ValueError( - f"Permission denied to create bucket '{bucket_name}'. " - "The service account may lack 'storage.buckets.create' permission." - ) from e - - input_blob_name = input_gcs_uri.replace(f"gs://{bucket_name}/", "") - blob = bucket.blob(input_blob_name) - blob.upload_from_string(config.query) - - new_config = types._RunQueryJobAgentEngineConfig( - input_gcs_uri=input_gcs_uri, - output_gcs_uri=output_gcs_uri, - ) - - # Proceed with sending the async query via the auto-generated method - operation = self._run_query_job(name=name, config=new_config) - - return types.RunQueryJobResult( - job_name=operation.name, - input_gcs_uri=input_gcs_uri, - output_gcs_uri=output_gcs_uri, - ) - - def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineConfigOrDict] = None, - ) -> types.AgentEngine: - """Gets an agent engine. - - Args: - name (str): - Required. A fully-qualified resource name or ID such as - "projects/123/locations/us-central1/reasoningEngines/456" or - a shortened name such as "reasoningEngines/456". - """ - api_resource = self._get(name=name, config=config) - agent_engine = types.AgentEngine( - api_client=self, - api_async_client=AsyncAgentEngines(api_client_=self._api_client), - api_resource=api_resource, - ) - if api_resource.spec: - self._register_api_methods(agent_engine=agent_engine) - return agent_engine - - def delete( - self, - *, - name: str, - force: Optional[bool] = None, - config: Optional[types.DeleteAgentEngineConfigOrDict] = None, - ) -> types.DeleteAgentEngineOperation: - """ - Delete an Agent Engine resource. - - Args: - name (str): - Required. The name of the Agent Engine to be deleted. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` - or `reasoningEngines/{resource_id}`. - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error when - the Agent Engine has undeleted child resources. Defaults to False. - config (DeleteAgentEngineConfig): - Optional. Additional configurations for deleting the Agent Engine. - - """ - logger.info(f"Deleting AgentEngine resource: {name}") - operation = self._delete(name=name, force=force, config=config) - logger.info(f"Started AgentEngine delete operation: {operation.name}") - return operation - - def create( - self, - *, - agent_engine: Any = None, - agent: Any = None, - config: Optional[types.AgentEngineConfigOrDict] = None, - ) -> types.AgentEngine: - """Creates an agent engine. - - The Agent Engine will be an instance of the `agent_engine` that - was passed in, running remotely on Vertex AI. - - Sample ``src_dir`` contents (e.g. ``./user_src_dir``): - - .. code-block:: python - - user_src_dir/ - |-- main.py - |-- requirements.txt - |-- user_code/ - | |-- utils.py - | |-- ... - |-- ... - - To build an Agent Engine with the above files, run: - - .. code-block:: python - - client = agentplatform.Client( - project="your-project", - location="us-central1", - ) - remote_agent = client.agent_engines.create( - agent=local_agent, - config=dict( - requirements=[ - # I.e. the PyPI dependencies listed in requirements.txt - "google-cloud-aiplatform[agent_engines,adk]", - ... - ], - extra_packages=[ - "./user_src_dir/main.py", # a single file - "./user_src_dir/user_code", # a directory - ... - ], - ), - ) - - Args: - agent (Any): - Optional. The Agent to be created. If not specified, this will - correspond to a lightweight instance that cannot be queried - (but can be updated to future instances that can be queried). - agent_engine (Any): - Optional. This is deprecated. Please use `agent` instead. - config (AgentEngineConfig): - Optional. The configurations to use for creating the Agent Engine. - - Returns: - AgentEngine: The created Agent Engine instance. - - Raises: - ValueError: If the `project` was not set using `client.Client`. - ValueError: If the `location` was not set using `client.Client`. - ValueError: If `config.staging_bucket` was not set when `agent` - is specified. - ValueError: If `config.staging_bucket` does not start with "gs://". - ValueError: If `config.extra_packages` is specified but `agent` - is None. - ValueError: If `config.requirements` is specified but `agent` is None. - ValueError: If `config.env_vars` has a dictionary entry that does not - correspond to an environment variable value or a SecretRef. - TypeError: If `config.env_vars` is not a dictionary. - FileNotFoundError: If `config.extra_packages` includes a file or - directory that does not exist. - IOError: If ``config.requirements` is a string that corresponds to a - nonexistent file. - """ - if config is None: - config = {} - if isinstance(config, dict): - config = types.AgentEngineConfig.model_validate(config) - elif not isinstance(config, types.AgentEngineConfig): - raise TypeError( - f"config must be a dict or AgentEngineConfig, but got {type(config)}." - ) - context_spec = config.context_spec - if context_spec is not None: - # Conversion to a dict for _create_config - context_spec = json.loads(context_spec.model_dump_json()) - developer_connect_source = config.developer_connect_source - if developer_connect_source is not None: - developer_connect_source = json.loads( - developer_connect_source.model_dump_json() - ) - agent_config_source = config.agent_config_source - if agent_config_source is not None: - agent_config_source = json.loads(agent_config_source.model_dump_json()) - keep_alive_probe = config.keep_alive_probe - if keep_alive_probe is not None: - keep_alive_probe = json.loads( - keep_alive_probe.model_dump_json(exclude_none=True) - ) - if agent and agent_engine: - raise ValueError("Please specify only one of `agent` or `agent_engine`.") - elif agent_engine: - raise DeprecationWarning( - "The `agent_engine` argument is deprecated. Please use `agent` instead." - ) - agent = agent or agent_engine - api_config = self._create_config( - mode="create", - agent=agent, - identity_type=config.identity_type, - staging_bucket=config.staging_bucket, - requirements=config.requirements, - display_name=config.display_name, - description=config.description, - gcs_dir_name=config.gcs_dir_name, - extra_packages=config.extra_packages, - env_vars=config.env_vars, - service_account=config.service_account, - context_spec=context_spec, - psc_interface_config=config.psc_interface_config, - agent_gateway_config=config.agent_gateway_config, - min_instances=config.min_instances, - max_instances=config.max_instances, - resource_limits=config.resource_limits, - container_concurrency=config.container_concurrency, - encryption_spec=config.encryption_spec, - agent_server_mode=config.agent_server_mode, - labels=config.labels, - class_methods=config.class_methods, - source_packages=config.source_packages, - developer_connect_source=developer_connect_source, - entrypoint_module=config.entrypoint_module, - entrypoint_object=config.entrypoint_object, - requirements_file=config.requirements_file, - agent_framework=config.agent_framework, - python_version=config.python_version, - build_options=config.build_options, - image_spec=config.image_spec, - agent_config_source=agent_config_source, - container_spec=config.container_spec, - keep_alive_probe=keep_alive_probe, - ) - operation = self._create(config=api_config) - reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( - operation_name=operation.name - ) - logger.info( - "View progress and logs at https://console.cloud.google.com/logs/query?" - f"project={self._api_client.project}" - "&query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22%0A" - f"resource.labels.reasoning_engine_id%3D%22{reasoning_engine_id}%22." - ) - if not self._is_lightweight_creation(agent, config): - poll_interval_seconds = 10 - else: - poll_interval_seconds = 1 # Lightweight agent engine resource creation. - operation = _agent_engines_utils._await_operation( - operation_name=operation.name, - get_operation_fn=self._get_agent_operation, - poll_interval_seconds=poll_interval_seconds, - ) - - agent_engine = types.AgentEngine( - api_client=self, - api_async_client=AsyncAgentEngines(api_client_=self._api_client), - api_resource=operation.response, - ) - if agent_engine.api_resource: - logger.info("Agent Engine created. To use it in another session:") - logger.info( - f"agent_engine=client.agent_engines.get(name='{agent_engine.api_resource.name}')" - ) - elif operation.error: - raise RuntimeError(f"Failed to create Agent Engine: {operation.error}") - else: - logger.warning("The operation returned an empty response.") - if not self._is_lightweight_creation(agent, config): - # If the user did not provide an agent_engine (e.g. lightweight - # provisioning), it will not have any API methods registered. - agent_engine = self._register_api_methods(agent_engine=agent_engine) - return agent_engine # type: ignore[no-any-return] - - def _set_source_code_spec( - self, - *, - spec: types.ReasoningEngineSpecDict, - update_masks: builtins.list[str], - source_packages: Optional[Sequence[str]] = None, - developer_connect_source: Optional[ - types.ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict - ] = None, - class_methods: Optional[Sequence[dict[str, Any]]] = None, - entrypoint_module: Optional[str] = None, - entrypoint_object: Optional[str] = None, - requirements_file: Optional[str] = None, - sys_version: str, - build_options: Optional[dict[str, builtins.list[str]]] = None, - image_spec: Optional[ - types.ReasoningEngineSpecSourceCodeSpecImageSpecDict - ] = None, - agent_config_source: Optional[ - types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict - ] = None, - ) -> None: - """Sets source_code_spec for agent engine inside the `spec`.""" - source_code_spec = types.ReasoningEngineSpecSourceCodeSpecDict() - if source_packages and not agent_config_source: - source_packages = _agent_engines_utils._validate_packages_or_raise( - packages=source_packages, - build_options=build_options, - ) - update_masks.append("spec.source_code_spec.inline_source.source_archive") - source_code_spec["inline_source"] = { # type: ignore[typeddict-item] - "source_archive": _agent_engines_utils._create_base64_encoded_tarball( - source_packages=source_packages - ) - } - elif developer_connect_source: - update_masks.append("spec.source_code_spec.developer_connect_source") - source_code_spec["developer_connect_source"] = { - "config": developer_connect_source - } - elif not agent_config_source: - raise ValueError( - "Please specify one of `source_packages`, `developer_connect_source`, " - "or `agent_config_source`." - ) - if class_methods is not None: - update_masks.append("spec.class_methods") - class_methods_spec_list = ( - _agent_engines_utils._class_methods_to_class_methods_spec( - class_methods=class_methods - ) - ) - spec["class_methods"] = [ - _agent_engines_utils._to_dict(class_method_spec) - for class_method_spec in class_methods_spec_list - ] - elif image_spec is None: - raise ValueError( - "`class_methods` must be specified if `source_packages`, " - "`developer_connect_source`, or `agent_config_source` is " - "specified without a Dockerfile or `image_spec`." - ) - if image_spec is not None: - if entrypoint_module or entrypoint_object or requirements_file: - raise ValueError( - "`image_spec` cannot be specified alongside `entrypoint_module`, " - "`entrypoint_object`, or `requirements_file`, as they are " - "mutually exclusive." - ) - if agent_config_source: - raise ValueError( - "`image_spec` cannot be specified alongside `agent_config_source`, " - "as they are mutually exclusive." - ) - update_masks.append("spec.source_code_spec.image_spec") - source_code_spec["image_spec"] = image_spec - spec["source_code_spec"] = source_code_spec - return - - update_masks.append("spec.source_code_spec.python_spec.version") - python_spec: types.ReasoningEngineSpecSourceCodeSpecPythonSpecDict = { - "version": sys_version, - } - if agent_config_source is not None: - if entrypoint_module or entrypoint_object: - logger.warning( - "`entrypoint_module` and `entrypoint_object` are ignored when " - "`agent_config_source` is specified, as they are pre-defined." - ) - if source_packages: - source_packages = _agent_engines_utils._validate_packages_or_raise( - packages=source_packages, - build_options=build_options, - ) - update_masks.append( - "spec.source_code_spec.agent_config_source.inline_source.source_archive" - ) - agent_config_source["inline_source"] = { # type: ignore[typeddict-item] - "source_archive": _agent_engines_utils._create_base64_encoded_tarball( - source_packages=source_packages - ) - } - update_masks.append("spec.source_code_spec.agent_config_source") - source_code_spec["agent_config_source"] = agent_config_source - - if requirements_file is not None: - update_masks.append( - "spec.source_code_spec.python_spec.requirements_file" - ) - python_spec["requirements_file"] = requirements_file - source_code_spec["python_spec"] = python_spec - - spec["source_code_spec"] = source_code_spec - return - - if not entrypoint_module: - raise ValueError( - "`entrypoint_module` must be specified if `source_packages` or `developer_connect_source` is specified." - ) - update_masks.append("spec.source_code_spec.python_spec.entrypoint_module") - python_spec["entrypoint_module"] = entrypoint_module - if not entrypoint_object: - raise ValueError( - "`entrypoint_object` must be specified if `source_packages` or `developer_connect_source` is specified." - ) - update_masks.append("spec.source_code_spec.python_spec.entrypoint_object") - python_spec["entrypoint_object"] = entrypoint_object - if requirements_file is not None: - update_masks.append("spec.source_code_spec.python_spec.requirements_file") - python_spec["requirements_file"] = requirements_file - source_code_spec["python_spec"] = python_spec - spec["source_code_spec"] = source_code_spec - - def _set_package_spec( - self, - *, - spec: types.ReasoningEngineSpecDict, - update_masks: builtins.list[str], - agent: Any, - staging_bucket: Optional[str] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - class_methods: Optional[Sequence[dict[str, Any]]] = None, - sys_version: str, - build_options: Optional[dict[str, builtins.list[str]]] = None, - ) -> None: - """Sets package spec for agent engine.""" - project = self._api_client.project - if project is None: - raise ValueError("project must be set using `agentplatform.Client`.") - location = self._api_client.location - if location is None: - raise ValueError("location must be set using `agentplatform.Client`.") - gcs_dir_name = gcs_dir_name or _agent_engines_utils._DEFAULT_GCS_DIR_NAME - staging_bucket = _agent_engines_utils._validate_staging_bucket_or_raise( - staging_bucket=staging_bucket, - ) - requirements = _agent_engines_utils._validate_requirements_or_raise( - agent=agent, - requirements=requirements, - ) - extra_packages = _agent_engines_utils._validate_packages_or_raise( - packages=extra_packages, - build_options=build_options, - ) - # Prepares the Agent Engine for creation/update in Vertex AI. This - # involves packaging and uploading the artifacts for agent_engine, - # requirements and extra_packages to `staging_bucket/gcs_dir_name`. - _agent_engines_utils._prepare( - agent=agent, - requirements=requirements, - project=project, - location=location, - staging_bucket=staging_bucket, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - credentials=self._api_client._credentials, - ) - # Update the package spec. - update_masks.append("spec.package_spec.pickle_object_gcs_uri") - package_spec: types.ReasoningEngineSpecPackageSpecDict = { - "python_version": sys_version, - "pickle_object_gcs_uri": "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _agent_engines_utils._BLOB_FILENAME, - ), - } - if extra_packages: - update_masks.append("spec.package_spec.dependency_files_gcs_uri") - package_spec["dependency_files_gcs_uri"] = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _agent_engines_utils._EXTRA_PACKAGES_FILE, - ) - if requirements: - update_masks.append("spec.package_spec.requirements_gcs_uri") - package_spec["requirements_gcs_uri"] = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _agent_engines_utils._REQUIREMENTS_FILE, - ) - spec["package_spec"] = package_spec - - update_masks.append("spec.class_methods") - if class_methods is not None: - class_methods_spec_list = ( - _agent_engines_utils._class_methods_to_class_methods_spec( - class_methods=class_methods - ) - ) - else: - class_methods_spec_list = ( - _agent_engines_utils._generate_class_methods_spec_or_raise( - agent=agent, - operations=_agent_engines_utils._get_registered_operations( - agent=agent - ), - ) - ) - spec["class_methods"] = [ - _agent_engines_utils._to_dict(class_method_spec) - for class_method_spec in class_methods_spec_list - ] - - def _create_config( - self, - *, - mode: str, - agent: Any = None, - identity_type: Optional[types.IdentityType] = None, - staging_bucket: Optional[str] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[dict[str, Union[str, Any]]] = None, - service_account: Optional[str] = None, - context_spec: Optional[types.ReasoningEngineContextSpecDict] = None, - psc_interface_config: Optional[types.PscInterfaceConfigDict] = None, - agent_gateway_config: Optional[ - types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict - ] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[genai_types.EncryptionSpecDict] = None, - labels: Optional[dict[str, str]] = None, - agent_server_mode: Optional[types.AgentServerMode] = None, - class_methods: Optional[Sequence[dict[str, Any]]] = None, - source_packages: Optional[Sequence[str]] = None, - developer_connect_source: Optional[ - types.ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict - ] = None, - entrypoint_module: Optional[str] = None, - entrypoint_object: Optional[str] = None, - requirements_file: Optional[str] = None, - agent_framework: Optional[str] = None, - python_version: Optional[str] = None, - build_options: Optional[dict[str, builtins.list[str]]] = None, - image_spec: Optional[ - types.ReasoningEngineSpecSourceCodeSpecImageSpecDict - ] = None, - agent_config_source: Optional[ - types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict - ] = None, - container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None, - keep_alive_probe: Optional[dict[str, Any]] = None, - traffic_config: Optional[types.ReasoningEngineTrafficConfigDict] = None, - ) -> types.UpdateAgentEngineConfigDict: - import sys - - config: types.UpdateAgentEngineConfigDict = {} - update_masks = [] - if mode not in ["create", "update"]: - raise ValueError(f"Unsupported mode: {mode}") - if agent is None: - if requirements is not None: - raise ValueError("requirements must be None if agent is None.") - if extra_packages is not None: - raise ValueError("extra_packages must be None if agent is None.") - if display_name is not None: - update_masks.append("display_name") - config["display_name"] = display_name - if description is not None: - update_masks.append("description") - config["description"] = description - if context_spec is not None: - update_masks.append("context_spec") - config["context_spec"] = context_spec - if encryption_spec is not None: - update_masks.append("encryption_spec") - config["encryption_spec"] = encryption_spec - if labels is not None: - update_masks.append("labels") - config["labels"] = labels - if traffic_config is not None: - update_masks.append("traffic_config") - config["traffic_config"] = traffic_config - - if agent_framework == "google-adk": - env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) - - if python_version: - sys_version = python_version - else: - sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" - - if agent: - if source_packages: - raise ValueError( - "If you have provided `source_packages` in `config`, please " - "do not specify `agent` in `agent_engines.create()` or " - "`agent_engines.update()`." - ) - if developer_connect_source: - raise ValueError( - "If you have provided `developer_connect_source` in `config`, please " - "do not specify `agent` in `agent_engines.create()` or " - "`agent_engines.update()`." - ) - elif source_packages and developer_connect_source: - raise ValueError( - "Please specify only one of `source_packages` or `developer_connect_source` in `config`." - ) - - if container_spec: - if agent: - raise ValueError( - "If you have provided `container_spec` in `config`, please " - "do not specify `agent` in `agent_engines.create()` or " - "`agent_engines.update()`." - ) - if source_packages or developer_connect_source: - raise ValueError( - "If you have provided `container_spec` in `config`, please " - "do not specify `source_packages` or `developer_connect_source` in `config`." - ) - - agent_engine_spec: Any = None - if agent: - agent_engine_spec = {} - agent = _agent_engines_utils._validate_agent_or_raise(agent=agent) - if _agent_engines_utils._is_adk_agent(agent): - env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) - self._set_package_spec( - spec=agent_engine_spec, - update_masks=update_masks, - agent=agent, - staging_bucket=staging_bucket, - requirements=requirements, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - class_methods=class_methods, - sys_version=sys_version, - build_options=build_options, - ) - elif ( - source_packages - or developer_connect_source - or image_spec - or agent_config_source - ): - agent_engine_spec = {} - self._set_source_code_spec( - spec=agent_engine_spec, - update_masks=update_masks, - source_packages=source_packages, - developer_connect_source=developer_connect_source, - class_methods=class_methods, - entrypoint_module=entrypoint_module, - entrypoint_object=entrypoint_object, - requirements_file=requirements_file, - sys_version=sys_version, - build_options=build_options, - image_spec=image_spec, - agent_config_source=agent_config_source, - ) - elif container_spec: - agent_engine_spec = {} - if class_methods is not None: - update_masks.append("spec.class_methods") - class_methods_spec_list = ( - _agent_engines_utils._class_methods_to_class_methods_spec( - class_methods=class_methods - ) - ) - agent_engine_spec["class_methods"] = [ - _agent_engines_utils._to_dict(class_method_spec) - for class_method_spec in class_methods_spec_list - ] - update_masks.append("spec.container_spec") - agent_engine_spec["container_spec"] = container_spec - - is_deployment_spec_updated = ( - env_vars is not None - or psc_interface_config is not None - or agent_gateway_config is not None - or min_instances is not None - or max_instances is not None - or resource_limits is not None - or container_concurrency is not None - or keep_alive_probe is not None - ) - if agent_engine_spec is None and is_deployment_spec_updated: - raise ValueError( - "To update `env_vars`, `psc_interface_config`, `min_instances`, " - "`max_instances`, `resource_limits`, `container_concurrency`, or " - "`keep_alive_probe`, you must also provide the `agent` variable or " - "the source code options (`source_packages`, " - "`developer_connect_source` or `agent_config_source`)." - ) - - if agent_engine_spec is not None: - if is_deployment_spec_updated: - ( - deployment_spec, - deployment_update_masks, - ) = self._generate_deployment_spec_or_raise( - env_vars=env_vars, - psc_interface_config=psc_interface_config, - agent_gateway_config=agent_gateway_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - keep_alive_probe=keep_alive_probe, - ) - update_masks.extend(deployment_update_masks) - agent_engine_spec["deployment_spec"] = deployment_spec - - if agent_server_mode: - if not agent_engine_spec.get("deployment_spec"): - agent_engine_spec["deployment_spec"] = ( - types.ReasoningEngineSpecDeploymentSpecDict() - ) - agent_engine_spec["deployment_spec"][ - "agent_server_mode" - ] = agent_server_mode - - agent_engine_spec["agent_framework"] = ( - _agent_engines_utils._get_agent_framework( - agent_framework=agent_framework, - agent=agent, - ) - ) - - if hasattr(agent, "agent_card"): - agent_card = getattr(agent, "agent_card") - if agent_card: - try: - from google.protobuf import json_format - - agent_engine_spec["agent_card"] = json_format.MessageToDict( - agent_card - ) - except Exception as e: - raise ValueError( - f"Failed to convert agent card to dict (serialization error): {e}" - ) from e - update_masks.append("spec.agent_framework") - - if identity_type is not None or service_account is not None: - if agent_engine_spec is None: - agent_engine_spec = {} - - if identity_type is not None: - agent_engine_spec["identity_type"] = identity_type - update_masks.append("spec.identity_type") - if service_account is not None: - # Clear the field in case of empty service_account. - if service_account: - agent_engine_spec["service_account"] = service_account - update_masks.append("spec.service_account") - - if agent_engine_spec is not None: - config["spec"] = agent_engine_spec - - if update_masks and mode == "update": - config["update_mask"] = ",".join(update_masks) - return config - - def _generate_deployment_spec_or_raise( - self, - *, - env_vars: Optional[dict[str, Union[str, Any]]] = None, - psc_interface_config: Optional[types.PscInterfaceConfigDict] = None, - agent_gateway_config: Optional[ - types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict - ] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[dict[str, str]] = None, - container_concurrency: Optional[int] = None, - keep_alive_probe: Optional[dict[str, Any]] = None, - ) -> Tuple[dict[str, Any], Sequence[str]]: - deployment_spec: dict[str, Any] = {} - update_masks = [] - if env_vars: - deployment_spec["env"] = [] - deployment_spec["secret_env"] = [] - if isinstance(env_vars, dict): - self._update_deployment_spec_with_env_vars_dict_or_raise( - deployment_spec=deployment_spec, - env_vars=env_vars, - ) - else: - raise TypeError(f"env_vars must be a dict, but got {type(env_vars)}.") - if deployment_spec.get("env"): - update_masks.append("spec.deployment_spec.env") - if deployment_spec.get("secret_env"): - update_masks.append("spec.deployment_spec.secret_env") - if psc_interface_config: - deployment_spec["psc_interface_config"] = psc_interface_config - update_masks.append("spec.deployment_spec.psc_interface_config") - if agent_gateway_config: - deployment_spec["agent_gateway_config"] = agent_gateway_config - update_masks.append("spec.deployment_spec.agent_gateway_config") - if min_instances is not None: - if not 0 <= min_instances <= 10: - raise ValueError( - f"min_instances must be between 0 and 10. Got {min_instances}" - ) - deployment_spec["min_instances"] = min_instances - update_masks.append("spec.deployment_spec.min_instances") - if max_instances is not None: - if psc_interface_config and not 1 <= max_instances <= 100: - raise ValueError( - f"max_instances must be between 1 and 100 when PSC-I is enabled. Got {max_instances}" - ) - elif not psc_interface_config and not 1 <= max_instances <= 1000: - raise ValueError( - f"max_instances must be between 1 and 1000. Got {max_instances}" - ) - deployment_spec["max_instances"] = max_instances - update_masks.append("spec.deployment_spec.max_instances") - if resource_limits: - _agent_engines_utils._validate_resource_limits_or_raise( - resource_limits=resource_limits - ) - deployment_spec["resource_limits"] = resource_limits - update_masks.append("spec.deployment_spec.resource_limits") - if container_concurrency: - deployment_spec["container_concurrency"] = container_concurrency - update_masks.append("spec.deployment_spec.container_concurrency") - if keep_alive_probe is not None: - deployment_spec["keep_alive_probe"] = keep_alive_probe - update_masks.append("spec.deployment_spec.keep_alive_probe") - return deployment_spec, update_masks - - def _update_deployment_spec_with_env_vars_dict_or_raise( - self, - *, - deployment_spec: dict[str, Any], - env_vars: dict[str, Any], - ) -> None: - for key, value in env_vars.items(): - if isinstance(value, dict): - if "secret_env" not in deployment_spec: - deployment_spec["secret_env"] = [] - deployment_spec["secret_env"].append({"name": key, "secret_ref": value}) - elif isinstance(value, str): - if "env" not in deployment_spec: - deployment_spec["env"] = [] - deployment_spec["env"].append({"name": key, "value": value}) - else: - raise TypeError( - f"Unknown value type in env_vars for {key}. " - f"Must be a str or SecretRef: {value}" - ) - - def _register_api_methods( - self, - *, - agent_engine: types.AgentEngine, - ) -> types.AgentEngine: - """Registers the API methods for the agent engine.""" - try: - _agent_engines_utils._register_api_methods_or_raise( - agent_engine=agent_engine, - wrap_operation_fn={ - "": _agent_engines_utils._wrap_query_operation, # type: ignore[dict-item] - "async": _agent_engines_utils._wrap_async_query_operation, # type: ignore[dict-item] - "stream": _agent_engines_utils._wrap_stream_query_operation, # type: ignore[dict-item] - "async_stream": _agent_engines_utils._wrap_async_stream_query_operation, # type: ignore[dict-item] - "a2a_extension": _agent_engines_utils._wrap_a2a_operation, - }, - ) - except Exception as e: - logger.warning( - _agent_engines_utils._FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e - ) - return agent_engine - - def list( - self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None - ) -> Iterator[types.AgentEngine]: - """List all instances of Agent Engine matching the filter. - - Example Usage: - - .. code-block:: python - import agentplatform - - client = agentplatform.Client(project="my_project", location="us-central1") - for agent in client.agent_engines.list( - config={"filter": "'display_name="My Custom Agent"'}, - ): - print(agent.api_resource.name) - - Args: - config (ListAgentEngineConfig): - Optional. The config (e.g. filter) for the agents to be listed. - - Returns: - Iterable[AgentEngine]: An iterable of Agent Engines matching the filter. - """ - - for reasoning_engine in self._list_pager(config=config): - yield types.AgentEngine( - api_client=self, - api_async_client=AsyncAgentEngines(api_client_=self._api_client), - api_resource=reasoning_engine, - ) - - def update( - self, - *, - name: str, - agent: Any = None, - agent_engine: Any = None, - config: types.AgentEngineConfigOrDict, - ) -> types.AgentEngine: - """Updates an existing Agent Engine. - - This method updates the configuration of an existing Agent Engine running - remotely, which is identified by its name. - - Args: - name (str): Required. A fully-qualified resource name or ID such as - "projects/123/locations/us-central1/reasoningEngines/456" or a - shortened name such as "reasoningEngines/456". - agent (Any): - Optional. The instance to be used as the updated Agent Engine. - If it is not specified, the existing instance will be used. - agent_engine (Any): - Optional. This is deprecated. Please use `agent` instead. - config (AgentEngineConfig): - Optional. The configurations to use for updating the Agent Engine. - - Returns: - AgentEngine: The updated Agent Engine. - - Raises: - ValueError: If the `project` was not set using `client.Client`. - ValueError: If the `location` was not set using `client.Client`. - ValueError: If `config.staging_bucket` was not set when `agent_engine` - is specified. - ValueError: If `config.staging_bucket` does not start with "gs://". - ValueError: If `config.extra_packages` is specified but `agent_engine` - is None. - ValueError: If `config.requirements` is specified but `agent_engine` is - None. - ValueError: If `config.env_vars` has a dictionary entry that does not - correspond to an environment variable value or a SecretRef. - TypeError: If `config.env_vars` is not a dictionary. - FileNotFoundError: If `config.extra_packages` includes a file or - directory that does not exist. - IOError: If `config.requirements` is a string that corresponds to a - nonexistent file. - """ - if isinstance(config, dict): - config = types.AgentEngineConfig.model_validate(config) - elif not isinstance(config, types.AgentEngineConfig): - raise TypeError( - f"config must be a dict or AgentEngineConfig, but got {type(config)}." - ) - context_spec = config.context_spec - if context_spec is not None: - # Conversion to a dict for _create_config - context_spec = json.loads(context_spec.model_dump_json()) - developer_connect_source = config.developer_connect_source - if developer_connect_source is not None: - developer_connect_source = json.loads( - developer_connect_source.model_dump_json() - ) - agent_config_source = config.agent_config_source - if agent_config_source is not None: - agent_config_source = json.loads(agent_config_source.model_dump_json()) - keep_alive_probe = config.keep_alive_probe - if keep_alive_probe is not None: - keep_alive_probe = json.loads( - keep_alive_probe.model_dump_json(exclude_none=True) - ) - traffic_config = config.traffic_config - if traffic_config is not None: - traffic_config = json.loads(traffic_config.model_dump_json()) - if agent and agent_engine: - raise ValueError("Please specify only one of `agent` or `agent_engine`.") - elif agent_engine: - raise DeprecationWarning( - "The `agent_engine` argument is deprecated. Please use `agent` instead." - ) - image_spec = config.image_spec - if image_spec is not None: - # Conversion to a dict for _create_config - image_spec = json.loads(image_spec.model_dump_json()) - container_spec = config.container_spec - if container_spec is not None: - # Conversion to a dict for _create_config - container_spec = json.loads(container_spec.model_dump_json()) - agent = agent or agent_engine - api_config = self._create_config( - mode="update", - agent=agent, - identity_type=config.identity_type, - staging_bucket=config.staging_bucket, - requirements=config.requirements, - display_name=config.display_name, - description=config.description, - gcs_dir_name=config.gcs_dir_name, - extra_packages=config.extra_packages, - env_vars=config.env_vars, - service_account=config.service_account, - context_spec=context_spec, - psc_interface_config=config.psc_interface_config, - agent_gateway_config=config.agent_gateway_config, - min_instances=config.min_instances, - max_instances=config.max_instances, - resource_limits=config.resource_limits, - container_concurrency=config.container_concurrency, - labels=config.labels, - class_methods=config.class_methods, - source_packages=config.source_packages, - developer_connect_source=developer_connect_source, - entrypoint_module=config.entrypoint_module, - entrypoint_object=config.entrypoint_object, - requirements_file=config.requirements_file, - agent_framework=config.agent_framework, - python_version=config.python_version, - build_options=config.build_options, - image_spec=image_spec, - agent_config_source=agent_config_source, - container_spec=container_spec, - keep_alive_probe=keep_alive_probe, - traffic_config=traffic_config, - ) - operation = self._update(name=name, config=api_config) - reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( - resource_name=name - ) - logger.info( - "View progress and logs at https://console.cloud.google.com/logs/query?" - f"project={self._api_client.project}" - "&query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22%0A" - f"resource.labels.reasoning_engine_id%3D%22{reasoning_engine_id}%22." - ) - operation = _agent_engines_utils._await_operation( - operation_name=operation.name, - get_operation_fn=self._get_agent_operation, - ) - agent_engine = types.AgentEngine( - api_client=self, - api_async_client=AsyncAgentEngines(api_client_=self._api_client), - api_resource=operation.response, - ) - if agent_engine.api_resource: - logger.info("Agent Engine updated. To use it in another session:") - logger.info( - f"agent_engine=client.agent_engines.get(name='{agent_engine.api_resource.name}')" - ) - elif operation.error: - raise RuntimeError(f"Failed to update Agent Engine: {operation.error}") - if agent_engine.api_resource.spec: - self._register_api_methods(agent_engine=agent_engine) - return agent_engine # type: ignore[no-any-return] - - def _stream_query( - self, *, name: str, config: Optional[types.QueryAgentEngineConfigOrDict] = None - ) -> Iterator[Any]: - """Streams the response of the agent engine.""" - parameter_model = types._QueryAgentEngineRequestParameters( - name=name, - config=config, - ) - request_dict = _QueryAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) - else: - path = "{name}:streamQuery?alt=sse" - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - http_options = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - for response in self._api_client.request_streamed( - "post", path, request_dict, http_options - ): - yield response - - # TODO: b/436704146 - Replace with generated methods - # TODO: b/437129724 - Add replay test for async stream query - async def _async_stream_query( - self, - *, - name: str, - config: Optional[types.QueryAgentEngineConfigOrDict] = None, - ) -> AsyncIterator[Any]: - """Streams the response of the agent engine asynchronously.""" - parameter_model = types._QueryAgentEngineRequestParameters( - name=name, - config=config, - ) - request_dict = _QueryAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) - else: - path = "{name}:streamQuery?alt=sse" - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - http_options = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - async_iterator = await self._api_client.async_request_streamed( - "post", path, request_dict, http_options - ) - async for response in async_iterator: - yield response - - def create_memory( - self, - *, - name: str, - fact: str, - scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - """Deprecated. Use agent_engines.memories.create instead.""" - warnings.warn( - ( - "agent_engines.create_memory is deprecated. " - "Use agent_engines.memories.create instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.create( - name=name, - fact=fact, - scope=scope, - config=config, - ) - - def delete_memory( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, - ) -> types.DeleteAgentEngineMemoryOperation: - """Deprecated. Use agent_engines.memories.delete instead.""" - warnings.warn( - ( - "agent_engines.delete_memory is deprecated. " - "Use agent_engines.memories.delete instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.delete(name=name, config=config) - - def generate_memories( - self, - *, - name: str, - vertex_session_source: Optional[ - types.GenerateMemoriesRequestVertexSessionSourceOrDict - ] = None, - direct_contents_source: Optional[ - types.GenerateMemoriesRequestDirectContentsSourceOrDict - ] = None, - direct_memories_source: Optional[ - types.GenerateMemoriesRequestDirectMemoriesSourceOrDict - ] = None, - scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - """Deprecated. Use agent_engines.memories.generate instead.""" - warnings.warn( - ( - "agent_engines.generate_memories is deprecated. " - "Use agent_engines.memories.generate instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.generate( - name=name, - vertex_session_source=vertex_session_source, - direct_contents_source=direct_contents_source, - direct_memories_source=direct_memories_source, - scope=scope, - config=config, - ) - - def get_memory( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, - ) -> types.Memory: - """Deprecated. Use agent_engines.memories.get instead.""" - warnings.warn( - ( - "agent_engines.get_memory is deprecated. " - "Use agent_engines.memories.get instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.get(name=name, config=config) - - def list_memories( - self, - *, - name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, - ) -> Iterator[types.Memory]: - """Deprecated. Use agent_engines.memories.list instead.""" - warnings.warn( - ( - "agent_engines.list_memories is deprecated. " - "Use agent_engines.memories.list instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.list(name=name, config=config) - - def retrieve_memories( - self, - *, - name: str, - scope: dict[str, str], - similarity_search_params: Optional[ - types.RetrieveMemoriesRequestSimilaritySearchParamsOrDict - ] = None, - simple_retrieval_params: Optional[ - types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict - ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, - ) -> Iterator[types.RetrieveMemoriesResponseRetrievedMemory]: - """Deprecated. Use agent_engines.memories.retrieve instead.""" - warnings.warn( - ( - "agent_engines.retrieve_memories is deprecated. " - "Use agent_engines.memories.retrieve instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.retrieve( - name=name, - scope=scope, - similarity_search_params=similarity_search_params, - simple_retrieval_params=simple_retrieval_params, - config=config, - ) - - def create_session( - self, - *, - name: str, - user_id: str, - config: Optional[types.CreateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: - """Deprecated. Use agent_engines.sessions.create instead.""" - warnings.warn( - ( - "agent_engines.create_session is deprecated. " - "Use agent_engines.sessions.create instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.sessions.create(name=name, user_id=user_id, config=config) - - def delete_session( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineSessionConfigOrDict] = None, - ) -> types.DeleteAgentEngineSessionOperation: - """Deprecated. Use agent_engines.sessions.delete instead.""" - warnings.warn( - ( - "agent_engines.delete_session is deprecated. " - "Use agent_engines.sessions.delete instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.sessions.delete(name=name, config=config) - - def get_session( - self, - *, - name: str, - config: Optional[types.GetAgentEngineSessionConfigOrDict] = None, - ) -> types.Session: - """Deprecated. Use agent_engines.sessions.get instead.""" - warnings.warn( - ( - "agent_engines.get_session is deprecated. " - "Use agent_engines.sessions.get instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.sessions.get(name=name, config=config) - - def list_sessions( - self, - *, - name: str, - config: Optional[types.ListAgentEngineSessionsConfigOrDict] = None, - ) -> Iterator[types.Session]: - """Deprecated. Use agent_engines.sessions.list instead.""" - warnings.warn( - ( - "agent_engines.list_sessions is deprecated. " - "Use agent_engines.sessions.list instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.sessions.list(name=name, config=config) - - def append_session_event( - self, - *, - name: str, - author: str, - invocation_id: str, - timestamp: datetime.datetime, - config: Optional[types.AppendAgentEngineSessionEventConfigOrDict] = None, - ) -> types.AppendAgentEngineSessionEventResponse: - """Deprecated. Use agent_engines.sessions.events.append instead.""" - warnings.warn( - ( - "agent_engines.append_session_event is deprecated. " - "Use agent_engines.sessions.events.append instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.sessions.events.append( - name=name, - author=author, - invocation_id=invocation_id, - timestamp=timestamp, - config=config, - ) - - def list_session_events( - self, - *, - name: str, - config: Optional[types.ListAgentEngineSessionEventsConfigOrDict] = None, - ) -> Iterator[types.SessionEvent]: - """Deprecated. Use agent_engines.sessions.events.list instead.""" - warnings.warn( - ( - "agent_engines.list_session_events is deprecated. " - "Use agent_engines.sessions.events.list instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.sessions.events.list(name=name, config=config) - - -class AsyncAgentEngines(_api_module.BaseModule): - - async def cancel_query_job( - self, - *, - name: str, - config: Optional[types.CancelQueryJobAgentEngineConfigOrDict] = None, - ) -> types.CancelQueryJobResult: - """ - Cancels a long-running query job on an Agent Engine. - - Args: - name (str): - Required. The reasoning engine resource name. - config (CancelQueryJobAgentEngineConfigOrDict): - Optional. The configuration for the cancel_query_job. - - """ - - parameter_model = types._CancelQueryJobAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _CancelQueryJobAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:cancelAsyncQuery".format_map(request_url_dict) - else: - path = "{name}:cancelAsyncQuery" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "post", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.CancelQueryJobResult._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _check_query_job( - self, - *, - name: str, - config: Optional[types.CheckQueryJobAgentEngineConfigOrDict] = None, - ) -> types.CheckQueryJobResult: - """ - Query an Agent Engine asynchronously. - """ - - parameter_model = types._CheckQueryJobAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _CheckQueryJobAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:checkQueryJob".format_map(request_url_dict) - else: - path = "{name}:checkQueryJob" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "post", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _CheckQueryJobResult_from_vertex(response_dict) - - return_value = types.CheckQueryJobResult._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _run_query_job( - self, - *, - name: str, - config: Optional[types._RunQueryJobAgentEngineConfigOrDict] = None, - ) -> types.AgentEngineOperation: - """ - Run a query job on an agent engine. - """ - - parameter_model = types._RunQueryJobAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _RunQueryJobAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:asyncQuery".format_map(request_url_dict) - else: - path = "{name}:asyncQuery" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "post", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _create( - self, *, config: Optional[types.CreateAgentEngineConfigOrDict] = None - ) -> types.AgentEngineOperation: - """ - Creates a new Agent Engine. - """ - - parameter_model = types._CreateAgentEngineRequestParameters( - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _CreateAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "reasoningEngines".format_map(request_url_dict) - else: - path = "reasoningEngines" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "post", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _delete( - self, - *, - name: str, - force: Optional[bool] = None, - config: Optional[types.DeleteAgentEngineConfigOrDict] = None, - ) -> types.DeleteAgentEngineOperation: - """ - Delete an Agent Engine resource. - - Args: - name (str): - Required. The name of the Agent Engine to be deleted. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` - or `reasoningEngines/{resource_id}`. - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error when - the Agent Engine has undeleted child resources. Defaults to False. - config (DeleteAgentEngineConfig): - Optional. Additional configurations for deleting the Agent Engine. - - """ - - parameter_model = types._DeleteAgentEngineRequestParameters( - name=name, - force=force, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _DeleteAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}".format_map(request_url_dict) - else: - path = "{name}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "delete", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.DeleteAgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _get( - self, *, name: str, config: Optional[types.GetAgentEngineConfigOrDict] = None - ) -> types.ReasoningEngine: - """ - Get an Agent Engine instance. - """ - - parameter_model = types._GetAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _GetAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}".format_map(request_url_dict) - else: - path = "{name}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "get", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _ReasoningEngine_from_vertex(response_dict) - - return_value = types.ReasoningEngine._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _list( - self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None - ) -> types.ListReasoningEnginesResponse: - """ - Lists Agent Engines. - """ - - parameter_model = types._ListAgentEngineRequestParameters( - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _ListAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "reasoningEngines".format_map(request_url_dict) - else: - path = "reasoningEngines" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "get", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) - - return_value = types.ListReasoningEnginesResponse._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _get_agent_operation( - self, - *, - operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineOperation: - parameter_model = types._GetAgentEngineOperationParameters( - operation_name=operation_name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _GetAgentEngineOperationParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{operationName}".format_map(request_url_dict) - else: - path = "{operationName}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "get", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _query( - self, *, name: str, config: Optional[types.QueryAgentEngineConfigOrDict] = None - ) -> types.QueryReasoningEngineResponse: - """ - Query an Agent Engine. - """ - - parameter_model = types._QueryAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _QueryAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:query".format_map(request_url_dict) - else: - path = "{name}:query" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "post", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.QueryReasoningEngineResponse._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _update( - self, *, name: str, config: Optional[types.UpdateAgentEngineConfigOrDict] = None - ) -> types.AgentEngineOperation: - """ - Updates an Agent Engine. - """ - - parameter_model = types._UpdateAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _UpdateAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}".format_map(request_url_dict) - else: - path = "{name}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "patch", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - _a2a_tasks = None - _memories = None - _sessions = None - _runtimes = None - - async def delete( - self, - *, - name: str, - force: Optional[bool] = None, - config: Optional[types.DeleteAgentEngineConfigOrDict] = None, - ) -> types.DeleteAgentEngineOperation: - """ - Delete an Agent Engine resource. - - Args: - name (str): - Required. The name of the Agent Engine to be deleted. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` - or `reasoningEngines/{resource_id}`. - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error when - the Agent Engine has undeleted child resources. Defaults to False. - config (DeleteAgentEngineConfig): - Optional. Additional configurations for deleting the Agent Engine. - - """ - logger.info(f"Deleting AgentEngine resource: {name}") - operation = await self._delete(name=name, force=force, config=config) - logger.info(f"Started AgentEngine delete operation: {operation.name}") - return operation - - @property - def runtimes(self) -> "runtimes_module.AsyncRuntimes": - if self._runtimes is None: - try: - # We need to lazy load the runtimes module to handle the - # possibility of ImportError when dependencies are not installed. - self._runtimes = importlib.import_module(".runtimes", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.runtimes' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._runtimes.AsyncRuntimes(self._api_client) # type: ignore[no-any-return] - - @property - def a2a_tasks(self) -> "a2a_tasks_module.AsyncA2aTasks": - if self._a2a_tasks is None: - try: - # We need to lazy load the a2a_tasks module to handle the - # possibility of ImportError when dependencies are not installed. - self._a2a_tasks = importlib.import_module(".a2a_tasks", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.a2a_tasks' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._a2a_tasks.AsyncA2aTasks(self._api_client) # type: ignore[no-any-return] - - @property - def memories(self) -> "memories_module.AsyncMemories": - if self._memories is None: - try: - # We need to lazy load the memories module to handle the - # possibility of ImportError when dependencies are not installed. - self._memories = importlib.import_module(".memories", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.memories' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._memories.AsyncMemories(self._api_client) # type: ignore[no-any-return] - - @property - def sessions(self) -> "sessions_module.AsyncSessions": - if self._sessions is None: - try: - # We need to lazy load the sessions module to handle the - # possibility of ImportError when dependencies are not installed. - self._sessions = importlib.import_module(".sessions", __package__) - except ImportError as e: - raise ImportError( - "The agent_engines.sessions module requires additional packages. " - "Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._sessions.AsyncSessions(self._api_client) # type: ignore[no-any-return] - - async def append_session_event( - self, - *, - name: str, - author: str, - invocation_id: str, - timestamp: datetime.datetime, - config: Optional[types.AppendAgentEngineSessionEventConfigOrDict] = None, - ) -> types.AppendAgentEngineSessionEventResponse: - """Deprecated. Use agent_engines.sessions.events.append instead.""" - warnings.warn( - ( - "agent_engines.append_session_event is deprecated. " - "Use agent_engines.sessions.events.append instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return await self.sessions.events.append( - name=name, - author=author, - invocation_id=invocation_id, - timestamp=timestamp, - config=config, - ) - - async def delete_memory( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, - ) -> types.DeleteAgentEngineMemoryOperation: - """Deprecated. Use agent_engines.memories.delete instead.""" - warnings.warn( - ( - "agent_engines.delete_memory is deprecated. " - "Use agent_engines.memories.delete instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return await self.memories.delete(name=name, config=config) - - async def delete_session( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineSessionConfigOrDict] = None, - ) -> types.DeleteAgentEngineSessionOperation: - """Deprecated. Use agent_engines.sessions.delete instead.""" - warnings.warn( - ( - "agent_engines.delete_session is deprecated. " - "Use agent_engines.sessions.delete instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return await self.sessions.delete(name=name, config=config) - - async def get_memory( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, - ) -> types.Memory: - """Deprecated. Use agent_engines.memories.get instead.""" - warnings.warn( - ( - "agent_engines.get_memory is deprecated. " - "Use agent_engines.memories.get instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return await self.memories.get(name=name, config=config) - - async def get_session( - self, - *, - name: str, - config: Optional[types.GetAgentEngineSessionConfigOrDict] = None, - ) -> types.Session: - """Deprecated. Use agent_engines.sessions.get instead.""" - warnings.warn( - ( - "agent_engines.get_session is deprecated. " - "Use agent_engines.sessions.get instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return await self.sessions.get(name=name, config=config) diff --git a/agentplatform/_genai/client.py b/agentplatform/_genai/client.py index 0dcb0d445e..a9f88ae01a 100644 --- a/agentplatform/_genai/client.py +++ b/agentplatform/_genai/client.py @@ -30,7 +30,7 @@ if TYPE_CHECKING: from agentplatform._genai import ( - agent_engines as agent_engines_module, + runtimes as runtimes_module, ) from agentplatform._genai import datasets as datasets_module from agentplatform._genai import evals as evals_module @@ -47,6 +47,11 @@ from agentplatform._genai import ( feedback_entries as feedback_entries_module, ) + from agentplatform._genai import sessions as sessions_module + from agentplatform._genai import ( + sandboxes as sandboxes_module, + ) + _GENAI_MODULES_TELEMETRY_HEADER = "vertex-genai-modules" @@ -84,7 +89,7 @@ def __init__(self, api_client: genai_client.BaseApiClient): # type: ignore[name self._api_client = api_client self._live = live.AsyncLive(self._api_client) self._evals: Optional[ModuleType] = None - self._agent_engines: Optional[ModuleType] = None + self._runtimes: Optional[ModuleType] = None self._prompt_optimizer: Optional[ModuleType] = None self._prompts: Optional[ModuleType] = None self._datasets: Optional[ModuleType] = None @@ -92,6 +97,8 @@ def __init__(self, api_client: genai_client.BaseApiClient): # type: ignore[name self._rag: Optional[ModuleType] = None self._model_garden: Optional[ModuleType] = None self._feedback_entries: Optional[ModuleType] = None + self._sessions: Optional[ModuleType] = None + self._sandboxes: Optional[ModuleType] = None @property @_common.experimental_warning( @@ -125,22 +132,58 @@ def prompt_optimizer(self) -> "prompt_optimizer_module.AsyncPromptOptimizer": return self._prompt_optimizer.AsyncPromptOptimizer(self._api_client) # type: ignore[no-any-return] @property - def agent_engines(self) -> "agent_engines_module.AsyncAgentEngines": - if self._agent_engines is None: + def runtimes(self) -> "runtimes_module.AsyncRuntimes": + if self._runtimes is None: + try: + # We need to lazy load the runtimes module to handle the + # possibility of ImportError when dependencies are not installed. + self._runtimes = importlib.import_module( + ".runtimes", + __package__, + ) + except ImportError as e: + raise ImportError( + "The 'runtimes' module requires 'additional packages'. " + "Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._runtimes.AsyncRuntimes(self._api_client) # type: ignore[no-any-return] + + @property + def sessions(self) -> "sessions_module.AsyncSessions": + if self._sessions is None: try: - # We need to lazy load the agent_engines module to handle the + # We need to lazy load the sessions module to handle the # possibility of ImportError when dependencies are not installed. - self._agent_engines = importlib.import_module( - ".agent_engines", + self._sessions = importlib.import_module( + ".sessions", __package__, ) except ImportError as e: raise ImportError( - "The 'agent_engines' module requires 'additional packages'. " + "The 'sessions' module requires 'additional packages'. " "Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e - return self._agent_engines.AsyncAgentEngines(self._api_client) # type: ignore[no-any-return] + return self._sessions.AsyncSessions(self._api_client) # type: ignore[no-any-return] + + @property + def sandboxes(self) -> "sandboxes_module.AsyncSandboxes": + if self._sandboxes is None: + try: + # We need to lazy load the sandboxes module to handle the + # possibility of ImportError when dependencies are not installed. + self._sandboxes = importlib.import_module( + ".sandboxes", + __package__, + ) + except ImportError as e: + raise ImportError( + "The 'sandboxes' module requires 'additional packages'. " + "Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._sandboxes.AsyncSandboxes(self._api_client) # type: ignore[no-any-return] @property def prompts(self) -> "prompts_module.AsyncPrompts": @@ -307,13 +350,15 @@ def __init__( self._aio = AsyncClient(self._api_client) self._evals: Optional[ModuleType] = None self._prompt_optimizer: Optional[ModuleType] = None - self._agent_engines: Optional[ModuleType] = None + self._runtimes: Optional[ModuleType] = None self._prompts: Optional[ModuleType] = None self._datasets: Optional[ModuleType] = None self._skills: Optional[ModuleType] = None self._rag: Optional[ModuleType] = None self._model_garden: Optional[ModuleType] = None self._feedback_entries: Optional[ModuleType] = None + self._sessions: Optional[ModuleType] = None + self._sandboxes: Optional[ModuleType] = None @property def evals(self) -> "evals_module.Evals": @@ -371,22 +416,58 @@ def _get_api_client( return None @property - def agent_engines(self) -> "agent_engines_module.AgentEngines": - if self._agent_engines is None: + def runtimes(self) -> "runtimes_module.Runtimes": + if self._runtimes is None: + try: + # We need to lazy load the runtimes module to handle the + # possibility of ImportError when dependencies are not installed. + self._runtimes = importlib.import_module( + ".runtimes", + __package__, + ) + except ImportError as e: + raise ImportError( + "The 'runtimes' module requires 'additional packages'. " + "Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._runtimes.Runtimes(self._api_client) # type: ignore[no-any-return] + + @property + def sessions(self) -> "sessions_module.Sessions": + if self._sessions is None: + try: + # We need to lazy load the sessions module to handle the + # possibility of ImportError when dependencies are not installed. + self._sessions = importlib.import_module( + ".sessions", + __package__, + ) + except ImportError as e: + raise ImportError( + "The 'sessions' module requires 'additional packages'. " + "Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._sessions.Sessions(self._api_client) # type: ignore[no-any-return] + + @property + def sandboxes(self) -> "sandboxes_module.Sandboxes": + if self._sandboxes is None: try: - # We need to lazy load the agent_engines module to handle the + # We need to lazy load the sandboxes module to handle the # possibility of ImportError when dependencies are not installed. - self._agent_engines = importlib.import_module( - ".agent_engines", + self._sandboxes = importlib.import_module( + ".sandboxes", __package__, ) except ImportError as e: raise ImportError( - "The 'agent_engines' module requires 'additional packages'. " + "The 'sandboxes' module requires 'additional packages'. " "Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e - return self._agent_engines.AgentEngines(self._api_client) # type: ignore[no-any-return] + return self._sandboxes.Sandboxes(self._api_client) # type: ignore[no-any-return] @property def prompts(self) -> "prompts_module.Prompts": diff --git a/agentplatform/_genai/evals.py b/agentplatform/_genai/evals.py index 4b600f7054..aef03fac88 100644 --- a/agentplatform/_genai/evals.py +++ b/agentplatform/_genai/evals.py @@ -2712,7 +2712,7 @@ def run_inference( *, src: Union[str, pd.DataFrame, types.EvaluationDataset], model: Optional[Union[str, Callable[[Any], Any]]] = None, - agent: Optional[Union[str, types.AgentEngine, LlmAgent]] = None, + agent: Optional[Union[str, types.Runtime, LlmAgent]] = None, location: Optional[str] = None, config: Optional[types.EvalRunInferenceConfigOrDict] = None, ) -> types.EvaluationDataset: @@ -2735,11 +2735,11 @@ def run_inference( - For custom logic, provide a callable function that accepts a prompt and returns a response. agent: This field is experimental and may change in future versions - The agent engine used or local agent to run agent, optional for non-agent evaluations. - - agent engine resource name in str type, with format + The agent runtime used or local agent to run agent, optional for non-agent evaluations. + - agent runtime resource name in str type, with format `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine_id}`, - run_inference will fetch the agent engine from the resource name. - - Or `types.AgentEngine` object. + run_inference will fetch the agent runtime from the resource name. + - Or `types.Runtime` object. - Or ADK agent in LlMAgent type. location: The location to use for the inference. If not specified, the location configured in the client will be used. If specified, @@ -2771,7 +2771,7 @@ def run_inference( " populated." ) - agent_engine_instance = None + runtime_instance = None agent_instance = None gemini_agent_instance = None if agent: @@ -2779,15 +2779,15 @@ def run_inference( agent ): gemini_agent_instance = agent - elif isinstance(agent, str) or isinstance(agent, types.AgentEngine): - agent_engine_instance = agent + elif isinstance(agent, str) or isinstance(agent, types.Runtime): + runtime_instance = agent else: agent_instance = agent return _evals_common._execute_inference( # type: ignore[no-any-return] api_client=self._api_client, model=model, - agent_engine=agent_engine_instance, + runtime=runtime_instance, agent=agent_instance, gemini_agent=gemini_agent_instance, src=src, @@ -3259,7 +3259,7 @@ def create_evaluation_run( or a Gemini Agent (Vertex AI Agent) resource name `projects/{project}/locations/{location}/agents/{agent}`. When a Gemini Agent resource is provided, the backend scrapes the agent to produce - agent responses. If an Agent Engine resource name is provided, runs + agent responses. If an Agent Runtime resource name is provided, runs inference with the deployed agent to get agent responses for evaluation. The `agent` parameter is required if `agent_info` is provided. user_simulator_config: The user simulator configuration for agent evaluation. @@ -5538,7 +5538,7 @@ async def create_evaluation_run( or a Gemini Agent (Vertex AI Agent) resource name `projects/{project}/locations/{location}/agents/{agent}`. When a Gemini Agent resource is provided, the backend scrapes the agent to produce - agent responses. If an Agent Engine resource name is provided, runs + agent responses. If an Agent Runtime resource name is provided, runs inference with the deployed agent to get agent responses for evaluation. The `agent` parameter is required if `agent_info` is provided. user_simulator_config: The user simulator configuration for agent evaluation. diff --git a/agentplatform/_genai/feedback_contexts.py b/agentplatform/_genai/feedback_contexts.py index 1e00ef64d2..8fd4d04850 100644 --- a/agentplatform/_genai/feedback_contexts.py +++ b/agentplatform/_genai/feedback_contexts.py @@ -25,7 +25,7 @@ from google.genai._common import get_value_by_path as getv from google.genai._common import set_value_by_path as setv -from . import _agent_engines_utils +from . import _runtimes_utils from . import types logger = logging.getLogger("agentplatform_genai.feedbackcontexts") @@ -382,7 +382,7 @@ def update( ) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_feedback_context_operation, poll_interval_seconds=0.5, @@ -688,7 +688,7 @@ async def update( ) if config.wait_for_completion: if not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_feedback_context_operation, poll_interval_seconds=0.5, diff --git a/agentplatform/_genai/feedback_entries.py b/agentplatform/_genai/feedback_entries.py index 79f2ec0040..b214646ad0 100644 --- a/agentplatform/_genai/feedback_entries.py +++ b/agentplatform/_genai/feedback_entries.py @@ -29,7 +29,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import AsyncPager, Pager -from . import _agent_engines_utils +from . import _runtimes_utils from . import types if typing.TYPE_CHECKING: @@ -805,7 +805,7 @@ def create( ) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_feedback_entry_operation, poll_interval_seconds=0.5, @@ -872,7 +872,7 @@ def update( ) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_feedback_entry_operation, poll_interval_seconds=0.5, @@ -911,7 +911,7 @@ def delete( ) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_delete_feedback_entry_operation, poll_interval_seconds=0.5, @@ -1526,7 +1526,7 @@ async def create( ) if config.wait_for_completion: if not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_feedback_entry_operation, poll_interval_seconds=0.5, @@ -1593,7 +1593,7 @@ async def update( ) if config.wait_for_completion: if not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_feedback_entry_operation, poll_interval_seconds=0.5, @@ -1632,7 +1632,7 @@ async def delete( ) if config.wait_for_completion: if not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_delete_feedback_entry_operation, poll_interval_seconds=0.5, diff --git a/agentplatform/_genai/live.py b/agentplatform/_genai/live.py index 1a4bbbf006..ac11533ddb 100644 --- a/agentplatform/_genai/live.py +++ b/agentplatform/_genai/live.py @@ -29,7 +29,7 @@ if TYPE_CHECKING: from agentplatform._genai import ( - live_agent_engines as live_agent_engines_module, + live_runtimes as live_runtimes_module, ) @@ -38,27 +38,27 @@ class AsyncLive(_api_module.BaseModule): def __init__(self, api_client: BaseApiClient): super().__init__(api_client) - self._agent_engines: Optional[ModuleType] = None + self._runtimes: Optional[ModuleType] = None @property @_common.experimental_warning( - "The Vertex SDK GenAI agent engines module is experimental, " + "The Vertex SDK GenAI runtimes module is experimental, " "and may change in future versions." ) - def agent_engines(self) -> "live_agent_engines_module.AsyncLiveAgentEngines": - if self._agent_engines is None: + def runtimes(self) -> "live_runtimes_module.AsyncLiveRuntimes": + if self._runtimes is None: try: - # We need to lazy load the live_agent_engines module to handle + # We need to lazy load the live_runtimes module to handle # the possibility of ImportError when dependencies are not # installed. - self._agent_engines = importlib.import_module( - ".live_agent_engines", + self._runtimes = importlib.import_module( + ".live_runtimes", __package__, ) except ImportError as e: raise ImportError( - "The 'agent_engines' module requires 'additional packages'. " + "The 'runtimes' module requires 'additional packages'. " "Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e - return self._agent_engines.AsyncLiveAgentEngines(self._api_client) # type: ignore[no-any-return] + return self._runtimes.AsyncLiveRuntimes(self._api_client) # type: ignore[no-any-return] diff --git a/agentplatform/_genai/live_agent_engines.py b/agentplatform/_genai/live_runtimes.py similarity index 78% rename from agentplatform/_genai/live_agent_engines.py rename to agentplatform/_genai/live_runtimes.py index ed79ed5d0c..64fd1237ce 100644 --- a/agentplatform/_genai/live_agent_engines.py +++ b/agentplatform/_genai/live_runtimes.py @@ -13,7 +13,7 @@ # limitations under the License. # -"""Live AgentEngine API client.""" +"""Live Runtime API client.""" import contextlib import json @@ -21,7 +21,7 @@ import google.auth from google.genai import _api_module -from .types import QueryAgentEngineConfig, QueryAgentEngineConfigOrDict +from .types import QueryRuntimeConfig, QueryRuntimeConfigOrDict try: @@ -33,8 +33,8 @@ from websockets.client import connect as ws_connect # type: ignore -class AsyncLiveAgentEngineSession: - """AsyncLiveAgentEngineSession.""" +class AsyncLiveRuntimeSession: + """AsyncLiveRuntimeSession.""" def __init__(self, websocket: ClientConnection): self._ws = websocket @@ -50,7 +50,7 @@ async def send(self, query_input: Dict[str, Any]) -> None: json_request = json.dumps({"bidi_stream_input": query_input}) except Exception as exc: raise ValueError( - "Failed to encode query input to JSON in live_agent_engines: " + "Failed to encode query input to JSON in live_runtimes: " f"{str(query_input)}" ) from exc await self._ws.send(json_request) @@ -70,8 +70,7 @@ async def receive(self) -> Any: return json.loads(response) except json.decoder.JSONDecodeError as exc: raise ValueError( - "Failed to parse response to JSON in live_agent_engines: " - f"{str(response)}" + "Failed to parse response to JSON in live_runtimes: " f"{str(response)}" ) from exc async def close(self) -> None: @@ -79,8 +78,8 @@ async def close(self) -> None: await self._ws.close() -class AsyncLiveAgentEngines(_api_module.BaseModule): - """AsyncLiveAgentEngines. +class AsyncLiveRuntimes(_api_module.BaseModule): + """AsyncLiveRuntimes. Example usage: @@ -91,17 +90,17 @@ class AsyncLiveAgentEngines(_api_module.BaseModule): from google import genai from google.genai import types - class MyAgentEngine(client): + class MyRuntime(client): def bidi_stream_query(self, input_queue: asyncio.Queue): while True: input = await input_queue.get() yield {"output": f"Agent received {input}!"} client = agentplatform.Client(project="my-project", location="us-central1") - agent_engine = client.agent_engines.create(agent) + runtime = client.runtimes.create(agent) - async with client.aio.live.agent_engines.connect( - agent_engine=agent_engine.api_resource.name, + async with client.aio.live.runtimes.connect( + runtime=runtime.api_resource.name, setup={"class_method": "bidi_stream_query"}, ) as session: await session.send(input={"input": "Hello world"}) @@ -115,13 +114,13 @@ def bidi_stream_query(self, input_queue: asyncio.Queue): async def connect( self, *, - agent_engine: str, - config: Optional[QueryAgentEngineConfigOrDict] = None, - ) -> AsyncIterator[AsyncLiveAgentEngineSession]: + runtime: str, + config: Optional[QueryRuntimeConfigOrDict] = None, + ) -> AsyncIterator[AsyncLiveRuntimeSession]: """Connect to the agent deployed to Agent Engine in a live (bidirectional streaming) session. Args: - agent_engine: The resource name of the Agent Engine to use for the + runtime: The resource name of the Agent Engine to use for the live session. config: The optional configuration for starting the live Agent Engine session. Custom class_method and an optional initial input could be @@ -129,15 +128,15 @@ async def connect( "bidi_stream_query" will be used by the Agent Engine. Yields: - An AsyncLiveAgentEngineSession object. + An AsyncLiveRuntimeSession object. """ if isinstance(config, dict): - config = QueryAgentEngineConfig(**config) + config = QueryRuntimeConfig(**config) - agent_engine_resource_name = agent_engine - if not agent_engine_resource_name.startswith("projects/"): - agent_engine_resource_name = f"projects/{self._api_client.project}/locations/{self._api_client.location}/reasoningEngines/{agent_engine}" - request_dict = {"setup": {"name": agent_engine_resource_name}} + runtime_resource_name = runtime + if not runtime_resource_name.startswith("projects/"): + runtime_resource_name = f"projects/{self._api_client.project}/locations/{self._api_client.location}/reasoningEngines/{runtime}" + request_dict = {"setup": {"name": runtime_resource_name}} if config is not None and config.class_method: request_dict["setup"]["class_method"] = config.class_method if config is not None and config.input: @@ -176,4 +175,4 @@ async def connect( uri, additional_headers=headers, **self._api_client._websocket_ssl_ctx ) as ws: await ws.send(request) - yield AsyncLiveAgentEngineSession(websocket=ws) + yield AsyncLiveRuntimeSession(websocket=ws) diff --git a/agentplatform/_genai/memories.py b/agentplatform/_genai/memories.py index c83cd240e1..dae85e6cde 100644 --- a/agentplatform/_genai/memories.py +++ b/agentplatform/_genai/memories.py @@ -30,7 +30,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import AsyncPager, Pager -from . import _agent_engines_utils +from . import _runtimes_utils from . import types if typing.TYPE_CHECKING: @@ -44,60 +44,7 @@ logger.setLevel(logging.INFO) -def _AgentEngineMemoryConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["display_name"]) is not None: - setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) - - if getv(from_object, ["description"]) is not None: - setv(parent_object, ["description"], getv(from_object, ["description"])) - - if getv(from_object, ["ttl"]) is not None: - setv(parent_object, ["ttl"], getv(from_object, ["ttl"])) - - if getv(from_object, ["expire_time"]) is not None: - setv(parent_object, ["expireTime"], getv(from_object, ["expire_time"])) - - if getv(from_object, ["revision_expire_time"]) is not None: - setv( - parent_object, - ["revisionExpireTime"], - getv(from_object, ["revision_expire_time"]), - ) - - if getv(from_object, ["revision_ttl"]) is not None: - setv(parent_object, ["revisionTtl"], getv(from_object, ["revision_ttl"])) - - if getv(from_object, ["disable_memory_revisions"]) is not None: - setv( - parent_object, - ["disableMemoryRevisions"], - getv(from_object, ["disable_memory_revisions"]), - ) - - if getv(from_object, ["topics"]) is not None: - setv( - parent_object, ["topics"], [item for item in getv(from_object, ["topics"])] - ) - - if getv(from_object, ["metadata"]) is not None: - setv( - parent_object, - ["metadata"], - {k: v for k, v in getv(from_object, ["metadata"]).items()}, - ) - - if getv(from_object, ["memory_id"]) is not None: - setv(parent_object, ["_query", "memoryId"], getv(from_object, ["memory_id"])) - - return to_object - - -def _CreateAgentEngineMemoryRequestParameters_to_vertex( +def _CreateRuntimeMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -112,12 +59,12 @@ def _CreateAgentEngineMemoryRequestParameters_to_vertex( setv(to_object, ["scope"], getv(from_object, ["scope"])) if getv(from_object, ["config"]) is not None: - _AgentEngineMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) + _RuntimeMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _DeleteAgentEngineMemoryRequestParameters_to_vertex( +def _DeleteRuntimeMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -128,7 +75,7 @@ def _DeleteAgentEngineMemoryRequestParameters_to_vertex( return to_object -def _GenerateAgentEngineMemoriesConfig_to_vertex( +def _GenerateRuntimeMemoriesConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -185,7 +132,7 @@ def _GenerateAgentEngineMemoriesConfig_to_vertex( return to_object -def _GenerateAgentEngineMemoriesRequestParameters_to_vertex( +def _GenerateRuntimeMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -218,14 +165,14 @@ def _GenerateAgentEngineMemoriesRequestParameters_to_vertex( setv(to_object, ["scope"], getv(from_object, ["scope"])) if getv(from_object, ["config"]) is not None: - _GenerateAgentEngineMemoriesConfig_to_vertex( + _GenerateRuntimeMemoriesConfig_to_vertex( getv(from_object, ["config"]), to_object ) return to_object -def _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( +def _GetRuntimeGenerateMemoriesOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -238,7 +185,7 @@ def _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( return to_object -def _GetAgentEngineMemoryOperationParameters_to_vertex( +def _GetRuntimeMemoryOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -251,7 +198,7 @@ def _GetAgentEngineMemoryOperationParameters_to_vertex( return to_object -def _GetAgentEngineMemoryRequestParameters_to_vertex( +def _GetRuntimeMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -342,7 +289,7 @@ def _IngestEventsRequestParameters_to_vertex( return to_object -def _ListAgentEngineMemoryConfig_to_vertex( +def _ListRuntimeMemoryConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -363,7 +310,7 @@ def _ListAgentEngineMemoryConfig_to_vertex( return to_object -def _ListAgentEngineMemoryRequestParameters_to_vertex( +def _ListRuntimeMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -372,12 +319,12 @@ def _ListAgentEngineMemoryRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) + _ListRuntimeMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _PurgeAgentEngineMemoriesRequestParameters_to_vertex( +def _PurgeRuntimeMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -401,7 +348,21 @@ def _PurgeAgentEngineMemoriesRequestParameters_to_vertex( return to_object -def _RetrieveAgentEngineMemoriesConfig_to_vertex( +def _RetrieveMemoryProfilesRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["scope"]) is not None: + setv(to_object, ["scope"], getv(from_object, ["scope"])) + + return to_object + + +def _RetrieveRuntimeMemoriesConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -423,7 +384,7 @@ def _RetrieveAgentEngineMemoriesConfig_to_vertex( return to_object -def _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( +def _RetrieveRuntimeMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -449,14 +410,14 @@ def _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( ) if getv(from_object, ["config"]) is not None: - _RetrieveAgentEngineMemoriesConfig_to_vertex( + _RetrieveRuntimeMemoriesConfig_to_vertex( getv(from_object, ["config"]), to_object ) return to_object -def _RetrieveMemoryProfilesRequestParameters_to_vertex( +def _RollbackRuntimeMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -464,27 +425,66 @@ def _RetrieveMemoryProfilesRequestParameters_to_vertex( if getv(from_object, ["name"]) is not None: setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - if getv(from_object, ["scope"]) is not None: - setv(to_object, ["scope"], getv(from_object, ["scope"])) + if getv(from_object, ["target_revision_id"]) is not None: + setv(to_object, ["targetRevisionId"], getv(from_object, ["target_revision_id"])) return to_object -def _RollbackAgentEngineMemoryRequestParameters_to_vertex( +def _RuntimeMemoryConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - if getv(from_object, ["target_revision_id"]) is not None: - setv(to_object, ["targetRevisionId"], getv(from_object, ["target_revision_id"])) + if getv(from_object, ["display_name"]) is not None: + setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) + + if getv(from_object, ["description"]) is not None: + setv(parent_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["ttl"]) is not None: + setv(parent_object, ["ttl"], getv(from_object, ["ttl"])) + + if getv(from_object, ["expire_time"]) is not None: + setv(parent_object, ["expireTime"], getv(from_object, ["expire_time"])) + + if getv(from_object, ["revision_expire_time"]) is not None: + setv( + parent_object, + ["revisionExpireTime"], + getv(from_object, ["revision_expire_time"]), + ) + + if getv(from_object, ["revision_ttl"]) is not None: + setv(parent_object, ["revisionTtl"], getv(from_object, ["revision_ttl"])) + + if getv(from_object, ["disable_memory_revisions"]) is not None: + setv( + parent_object, + ["disableMemoryRevisions"], + getv(from_object, ["disable_memory_revisions"]), + ) + + if getv(from_object, ["topics"]) is not None: + setv( + parent_object, ["topics"], [item for item in getv(from_object, ["topics"])] + ) + + if getv(from_object, ["metadata"]) is not None: + setv( + parent_object, + ["metadata"], + {k: v for k, v in getv(from_object, ["metadata"]).items()}, + ) + + if getv(from_object, ["memory_id"]) is not None: + setv(parent_object, ["_query", "memoryId"], getv(from_object, ["memory_id"])) return to_object -def _UpdateAgentEngineMemoryConfig_to_vertex( +def _UpdateRuntimeMemoryConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -542,7 +542,7 @@ def _UpdateAgentEngineMemoryConfig_to_vertex( return to_object -def _UpdateAgentEngineMemoryRequestParameters_to_vertex( +def _UpdateRuntimeMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -557,9 +557,7 @@ def _UpdateAgentEngineMemoryRequestParameters_to_vertex( setv(to_object, ["scope"], getv(from_object, ["scope"])) if getv(from_object, ["config"]) is not None: - _UpdateAgentEngineMemoryConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _UpdateRuntimeMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object @@ -572,13 +570,13 @@ def _create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: + config: Optional[types.RuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: """ - Creates a new memory in the Agent Engine. + Creates a new memory in the Agent Runtime. """ - parameter_model = types._CreateAgentEngineMemoryRequestParameters( + parameter_model = types._CreateRuntimeMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -591,7 +589,7 @@ def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _CreateRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -620,7 +618,7 @@ def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.RuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -648,21 +646,21 @@ def delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, - ) -> types.DeleteAgentEngineMemoryOperation: + config: Optional[types.DeleteRuntimeMemoryConfigOrDict] = None, + ) -> types.DeleteRuntimeMemoryOperation: """ - Delete an Agent Engine memory. + Delete an Agent Runtime memory. Args: name (str): - Required. The name of the Agent Engine memory to be deleted. Format: + Required. The name of the Agent Runtime memory to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory}`. - config (DeleteAgentEngineMemoryConfig): - Optional. Additional configurations for deleting the Agent Engine. + config (DeleteRuntimeMemoryConfig): + Optional. Additional configurations for deleting the Agent Runtime. """ - parameter_model = types._DeleteAgentEngineMemoryRequestParameters( + parameter_model = types._DeleteRuntimeMemoryRequestParameters( name=name, config=config, ) @@ -673,7 +671,7 @@ def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _DeleteRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -702,7 +700,7 @@ def delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineMemoryOperation._from_response( + return_value = types.DeleteRuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -740,13 +738,13 @@ def _generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: + config: Optional[types.GenerateRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: """ - Generates memories for an Agent Engine. + Generates memories for an Agent Runtime. """ - parameter_model = types._GenerateAgentEngineMemoriesRequestParameters( + parameter_model = types._GenerateRuntimeMemoriesRequestParameters( name=name, vertex_session_source=vertex_session_source, direct_contents_source=direct_contents_source, @@ -761,7 +759,7 @@ def _generate( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GenerateAgentEngineMemoriesRequestParameters_to_vertex( + request_dict = _GenerateRuntimeMemoriesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -790,7 +788,7 @@ def _generate( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineGenerateMemoriesOperation._from_response( + return_value = types.RuntimeGenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -815,13 +813,10 @@ def _generate( return return_value def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeMemoryConfigOrDict] = None ) -> types.Memory: """ - Gets an agent engine memory. + Gets an agent runtime memory. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -830,7 +825,7 @@ def get( """ - parameter_model = types._GetAgentEngineMemoryRequestParameters( + parameter_model = types._GetRuntimeMemoryRequestParameters( name=name, config=config, ) @@ -841,9 +836,7 @@ def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _GetRuntimeMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -979,16 +972,13 @@ def _ingest_events( return return_value def _list( - self, - *, - name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, + self, *, name: str, config: Optional[types.ListRuntimeMemoryConfigOrDict] = None ) -> types.ListReasoningEnginesMemoriesResponse: """ - Lists Agent Engine memories. + Lists Agent Runtime memories. """ - parameter_model = types._ListAgentEngineMemoryRequestParameters( + parameter_model = types._ListRuntimeMemoryRequestParameters( name=name, config=config, ) @@ -999,7 +989,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _ListRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1056,9 +1046,9 @@ def _get_memory_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - parameter_model = types._GetAgentEngineMemoryOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: + parameter_model = types._GetRuntimeMemoryOperationParameters( operation_name=operation_name, config=config, ) @@ -1069,7 +1059,7 @@ def _get_memory_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryOperationParameters_to_vertex( + request_dict = _GetRuntimeMemoryOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1098,7 +1088,7 @@ def _get_memory_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.RuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1126,9 +1116,9 @@ def _get_generate_memories_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - parameter_model = types._GetAgentEngineGenerateMemoriesOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: + parameter_model = types._GetRuntimeGenerateMemoriesOperationParameters( operation_name=operation_name, config=config, ) @@ -1139,7 +1129,7 @@ def _get_generate_memories_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( + request_dict = _GetRuntimeGenerateMemoriesOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1168,7 +1158,7 @@ def _get_generate_memories_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineGenerateMemoriesOperation._from_response( + return_value = types.RuntimeGenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -1203,13 +1193,13 @@ def _retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveRuntimeMemoriesConfigOrDict] = None, ) -> types.RetrieveMemoriesResponse: """ - Retrieves memories for an Agent Engine. + Retrieves memories for an Agent Runtime. """ - parameter_model = types._RetrieveAgentEngineMemoriesRequestParameters( + parameter_model = types._RetrieveRuntimeMemoriesRequestParameters( name=name, scope=scope, similarity_search_params=similarity_search_params, @@ -1223,7 +1213,7 @@ def _retrieve( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( + request_dict = _RetrieveRuntimeMemoriesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1284,13 +1274,13 @@ def retrieve_profiles( config: Optional[types.RetrieveMemoryProfilesConfigOrDict] = None, ) -> types.RetrieveProfilesResponse: """ - Retrieves memory profiles for an Agent Engine. + Retrieves memory profiles for an Agent Runtime. For example, you can use the following code to retrieve all memory profiles for scope `{'user_id': '123'}`: ```python - result = client.agent_engines.memories.retrieve_profiles( + result = client.runtimes.memories.retrieve_profiles( name="projects/123/locations/us-central1/reasoningEngines/456", scope={"user_id": "123"} ) @@ -1383,13 +1373,13 @@ def _rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineRollbackMemoryOperation: + config: Optional[types.RollbackRuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeRollbackMemoryOperation: """ Rollback a memory to a previous revision. """ - parameter_model = types._RollbackAgentEngineMemoryRequestParameters( + parameter_model = types._RollbackRuntimeMemoryRequestParameters( name=name, target_revision_id=target_revision_id, config=config, @@ -1401,7 +1391,7 @@ def _rollback( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RollbackAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _RollbackRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1430,7 +1420,7 @@ def _rollback( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineRollbackMemoryOperation._from_response( + return_value = types.RuntimeRollbackMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1460,13 +1450,13 @@ def _update( name: str, fact: Optional[str] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.UpdateAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: + config: Optional[types.UpdateRuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: """ - Updates an Agent Engine memory. + Updates an Agent Runtime memory. """ - parameter_model = types._UpdateAgentEngineMemoryRequestParameters( + parameter_model = types._UpdateRuntimeMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -1479,7 +1469,7 @@ def _update( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _UpdateAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _UpdateRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1508,7 +1498,7 @@ def _update( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.RuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1541,13 +1531,13 @@ def _purge( builtins.list[types.MemoryConjunctionFilterOrDict] ] = None, force: Optional[bool] = None, - config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEnginePurgeMemoriesOperation: + config: Optional[types.PurgeRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimePurgeMemoriesOperation: """ - Purges memories from an Agent Engine. + Purges memories from an Agent Runtime. """ - parameter_model = types._PurgeAgentEngineMemoriesRequestParameters( + parameter_model = types._PurgeRuntimeMemoriesRequestParameters( name=name, filter=filter, filter_groups=filter_groups, @@ -1561,7 +1551,7 @@ def _purge( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _PurgeAgentEngineMemoriesRequestParameters_to_vertex( + request_dict = _PurgeRuntimeMemoriesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1590,7 +1580,7 @@ def _purge( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEnginePurgeMemoriesOperation._from_response( + return_value = types.RuntimePurgeMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -1627,7 +1617,7 @@ def revisions(self) -> "memory_revisions_module.MemoryRevisions": ) except ImportError as e: raise ImportError( - "The 'agent_engines.memories.revisions' module requires " + "The 'runtimes.memories.revisions' module requires " "additional packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e @@ -1639,9 +1629,9 @@ def create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - """Creates a new memory in the Agent Engine. + config: Optional[types.RuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: + """Creates a new memory in the Agent Runtime. Args: name (str): @@ -1650,16 +1640,16 @@ def create( Required. The fact to be stored in the memory. scope (dict[str, str]): Required. The scope of the memory. For example, {"user_id": "123"}. - config (AgentEngineMemoryConfigOrDict): + config (RuntimeMemoryConfigOrDict): Optional. The configuration for the memory. Returns: - AgentEngineMemoryOperation: The operation for creating the memory. + RuntimeMemoryOperation: The operation for creating the memory. """ if config is None: - config = types.AgentEngineMemoryConfig() + config = types.RuntimeMemoryConfig() elif isinstance(config, dict): - config = types.AgentEngineMemoryConfig.model_validate(config) + config = types.RuntimeMemoryConfig.model_validate(config) operation = self._create( name=name, fact=fact, @@ -1668,7 +1658,7 @@ def create( ) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -1697,13 +1687,13 @@ def generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - """Generates memories for the agent engine. + config: Optional[types.GenerateRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: + """Generates memories for the agent runtime. Args: name (str): - Required. The name of the agent engine to generate memories for. + Required. The name of the agent runtime to generate memories for. vertex_session_source (GenerateMemoriesRequestVertexSessionSource): Optional. The vertex session source to use for generating memories. Only one of vertex_session_source, @@ -1724,13 +1714,13 @@ def generate( Optional. The configuration for the memories to generate. Returns: - AgentEngineGenerateMemoriesOperation: + RuntimeGenerateMemoriesOperation: The operation for generating the memories. """ if config is None: - config = types.GenerateAgentEngineMemoriesConfig() + config = types.GenerateRuntimeMemoriesConfig() elif isinstance(config, dict): - config = types.GenerateAgentEngineMemoriesConfig.model_validate(config) + config = types.GenerateRuntimeMemoriesConfig.model_validate(config) operation = self._generate( name=name, vertex_session_source=vertex_session_source, @@ -1740,7 +1730,7 @@ def generate( config=config, ) if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_generate_memories_operation, poll_interval_seconds=0.5, @@ -1753,14 +1743,14 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, + config: Optional[types.ListRuntimeMemoryConfigOrDict] = None, ) -> Iterator[types.Memory]: - """Lists Agent Engine memories. + """Lists Agent Runtime memories. Args: name (str): - Required. The name of the agent engine to list memories for. - config (ListAgentEngineMemoryConfig): + Required. The name of the agent runtime to list memories for. + config (ListRuntimeMemoryConfig): Optional. The configuration for the memories to list. Returns: @@ -1785,13 +1775,13 @@ def retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveRuntimeMemoriesConfigOrDict] = None, ) -> Iterator[types.RetrieveMemoriesResponseRetrievedMemory]: """Retrieves memories for the agent. Args: name (str): - Required. The name of the agent engine to retrieve memories for. + Required. The name of the agent runtime to retrieve memories for. scope (dict[str, str]): Required. The scope of the memories to retrieve. For example, {"user_id": "123"}. @@ -1801,7 +1791,7 @@ def retrieve( simple_retrieval_params (RetrieveMemoriesRequestSimpleRetrievalParams): Optional. The simple retrieval parameters to use for retrieving memories. - config (RetrieveAgentEngineMemoriesConfig): + config (RetrieveRuntimeMemoriesConfig): Optional. The configuration for the memories to retrieve. Returns: @@ -1832,8 +1822,8 @@ def rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineRollbackMemoryOperation: + config: Optional[types.RollbackRuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeRollbackMemoryOperation: """Rolls back a memory to a previous revision. Args: @@ -1841,24 +1831,24 @@ def rollback( Required. The name of the memory to rollback. target_revision_id (str): Required. The revision ID to roll back to - config (RollbackAgentEngineMemoryConfig): + config (RollbackRuntimeMemoryConfig): Optional. The configuration for the rollback. Returns: - AgentEngineRollbackMemoryOperation: + RuntimeRollbackMemoryOperation: The operation for rolling back the memory. """ if config is None: - config = types.RollbackAgentEngineMemoryConfig() + config = types.RollbackRuntimeMemoryConfig() elif isinstance(config, dict): - config = types.RollbackAgentEngineMemoryConfig.model_validate(config) + config = types.RollbackRuntimeMemoryConfig.model_validate(config) operation = self._rollback( name=name, target_revision_id=target_revision_id, config=config, ) if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -1874,13 +1864,13 @@ def purge( filter: Optional[str] = None, filter_groups: Optional[List[types.MemoryConjunctionFilter]] = None, force: bool = False, - config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEnginePurgeMemoriesOperation: - """Purges memories from an Agent Engine. + config: Optional[types.PurgeRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimePurgeMemoriesOperation: + """Purges memories from an Agent Runtime. Args: name (str): - Required. The name of the Agent Engine to purge memories from. + Required. The name of the Agent Runtime to purge memories from. filter (str): Optional. The standard list filter to determine which memories to purge. filter_groups (list[MemoryConjunctionFilter]): @@ -1890,17 +1880,17 @@ def purge( force (bool): Optional. Whether to force the purge operation. If false, the operation will be staged but not executed. - config (PurgeAgentEngineMemoriesConfig): + config (PurgeRuntimeMemoriesConfig): Optional. The configuration for the purge operation. Returns: - AgentEnginePurgeMemoriesOperation: + RuntimePurgeMemoriesOperation: The operation for purging the memories. """ if config is None: - config = types.PurgeAgentEngineMemoriesConfig() + config = types.PurgeRuntimeMemoriesConfig() elif isinstance(config, dict): - config = types.PurgeAgentEngineMemoriesConfig.model_validate(config) + config = types.PurgeRuntimeMemoriesConfig.model_validate(config) operation = self._purge( name=name, filter=filter, @@ -1909,7 +1899,7 @@ def purge( config=config, ) if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -1932,11 +1922,11 @@ def ingest_events( ] = None, config: Optional[types.IngestEventsConfigOrDict] = None, ) -> types.MemoryBankIngestEventsOperation: - """Ingests events into an Agent Engine. + """Ingests events into an Agent Runtime. Example usage: ``` - client.agent_engines.memories.ingest_events( + client.runtimes.memories.ingest_events( name="projects/test-project/locations/us-central1/reasoningEngines/test-agent-engine", scope={"user_id": "test-user-id"}, direct_contents_source={ @@ -1961,7 +1951,7 @@ def ingest_events( Args: name (str): - Required. The name of the Agent Engine to ingest events into. + Required. The name of the Agent Runtime to ingest events into. scope (dict[str, str]): Required. The scope of the events to ingest. For example, {"user_id": "123"}. @@ -1976,7 +1966,7 @@ def ingest_events( Optional. The configuration for the ingest events operation. Returns: - AgentEngineIngestEventsOperation: + RuntimeIngestEventsOperation: The operation for ingesting the events. """ if config is None: @@ -1992,7 +1982,7 @@ def ingest_events( config=config, ) if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -2010,13 +2000,13 @@ async def _create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: + config: Optional[types.RuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: """ - Creates a new memory in the Agent Engine. + Creates a new memory in the Agent Runtime. """ - parameter_model = types._CreateAgentEngineMemoryRequestParameters( + parameter_model = types._CreateRuntimeMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -2029,7 +2019,7 @@ async def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _CreateRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2060,7 +2050,7 @@ async def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.RuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2088,21 +2078,21 @@ async def delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, - ) -> types.DeleteAgentEngineMemoryOperation: + config: Optional[types.DeleteRuntimeMemoryConfigOrDict] = None, + ) -> types.DeleteRuntimeMemoryOperation: """ - Delete an Agent Engine memory. + Delete an Agent Runtime memory. Args: name (str): - Required. The name of the Agent Engine memory to be deleted. Format: + Required. The name of the Agent Runtime memory to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory}`. - config (DeleteAgentEngineMemoryConfig): - Optional. Additional configurations for deleting the Agent Engine. + config (DeleteRuntimeMemoryConfig): + Optional. Additional configurations for deleting the Agent Runtime. """ - parameter_model = types._DeleteAgentEngineMemoryRequestParameters( + parameter_model = types._DeleteRuntimeMemoryRequestParameters( name=name, config=config, ) @@ -2113,7 +2103,7 @@ async def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _DeleteRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2144,7 +2134,7 @@ async def delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineMemoryOperation._from_response( + return_value = types.DeleteRuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2182,13 +2172,13 @@ async def _generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: + config: Optional[types.GenerateRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: """ - Generates memories for an Agent Engine. + Generates memories for an Agent Runtime. """ - parameter_model = types._GenerateAgentEngineMemoriesRequestParameters( + parameter_model = types._GenerateRuntimeMemoriesRequestParameters( name=name, vertex_session_source=vertex_session_source, direct_contents_source=direct_contents_source, @@ -2203,7 +2193,7 @@ async def _generate( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GenerateAgentEngineMemoriesRequestParameters_to_vertex( + request_dict = _GenerateRuntimeMemoriesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2234,7 +2224,7 @@ async def _generate( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineGenerateMemoriesOperation._from_response( + return_value = types.RuntimeGenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -2259,13 +2249,10 @@ async def _generate( return return_value async def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeMemoryConfigOrDict] = None ) -> types.Memory: """ - Gets an agent engine memory. + Gets an agent runtime memory. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -2274,7 +2261,7 @@ async def get( """ - parameter_model = types._GetAgentEngineMemoryRequestParameters( + parameter_model = types._GetRuntimeMemoryRequestParameters( name=name, config=config, ) @@ -2285,9 +2272,7 @@ async def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _GetRuntimeMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -2427,16 +2412,13 @@ async def _ingest_events( return return_value async def _list( - self, - *, - name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, + self, *, name: str, config: Optional[types.ListRuntimeMemoryConfigOrDict] = None ) -> types.ListReasoningEnginesMemoriesResponse: """ - Lists Agent Engine memories. + Lists Agent Runtime memories. """ - parameter_model = types._ListAgentEngineMemoryRequestParameters( + parameter_model = types._ListRuntimeMemoryRequestParameters( name=name, config=config, ) @@ -2447,7 +2429,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _ListRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2506,9 +2488,9 @@ async def _get_memory_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - parameter_model = types._GetAgentEngineMemoryOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: + parameter_model = types._GetRuntimeMemoryOperationParameters( operation_name=operation_name, config=config, ) @@ -2519,7 +2501,7 @@ async def _get_memory_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryOperationParameters_to_vertex( + request_dict = _GetRuntimeMemoryOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2550,7 +2532,7 @@ async def _get_memory_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.RuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2578,9 +2560,9 @@ async def _get_generate_memories_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - parameter_model = types._GetAgentEngineGenerateMemoriesOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: + parameter_model = types._GetRuntimeGenerateMemoriesOperationParameters( operation_name=operation_name, config=config, ) @@ -2591,7 +2573,7 @@ async def _get_generate_memories_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( + request_dict = _GetRuntimeGenerateMemoriesOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2622,7 +2604,7 @@ async def _get_generate_memories_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineGenerateMemoriesOperation._from_response( + return_value = types.RuntimeGenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -2657,13 +2639,13 @@ async def _retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveRuntimeMemoriesConfigOrDict] = None, ) -> types.RetrieveMemoriesResponse: """ - Retrieves memories for an Agent Engine. + Retrieves memories for an Agent Runtime. """ - parameter_model = types._RetrieveAgentEngineMemoriesRequestParameters( + parameter_model = types._RetrieveRuntimeMemoriesRequestParameters( name=name, scope=scope, similarity_search_params=similarity_search_params, @@ -2677,7 +2659,7 @@ async def _retrieve( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( + request_dict = _RetrieveRuntimeMemoriesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2740,13 +2722,13 @@ async def retrieve_profiles( config: Optional[types.RetrieveMemoryProfilesConfigOrDict] = None, ) -> types.RetrieveProfilesResponse: """ - Retrieves memory profiles for an Agent Engine. + Retrieves memory profiles for an Agent Runtime. For example, you can use the following code to retrieve all memory profiles for scope `{'user_id': '123'}`: ```python - result = client.agent_engines.memories.retrieve_profiles( + result = client.runtimes.memories.retrieve_profiles( name="projects/123/locations/us-central1/reasoningEngines/456", scope={"user_id": "123"} ) @@ -2841,13 +2823,13 @@ async def _rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineRollbackMemoryOperation: + config: Optional[types.RollbackRuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeRollbackMemoryOperation: """ Rollback a memory to a previous revision. """ - parameter_model = types._RollbackAgentEngineMemoryRequestParameters( + parameter_model = types._RollbackRuntimeMemoryRequestParameters( name=name, target_revision_id=target_revision_id, config=config, @@ -2859,7 +2841,7 @@ async def _rollback( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RollbackAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _RollbackRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2890,7 +2872,7 @@ async def _rollback( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineRollbackMemoryOperation._from_response( + return_value = types.RuntimeRollbackMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2920,13 +2902,13 @@ async def _update( name: str, fact: Optional[str] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.UpdateAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: + config: Optional[types.UpdateRuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: """ - Updates an Agent Engine memory. + Updates an Agent Runtime memory. """ - parameter_model = types._UpdateAgentEngineMemoryRequestParameters( + parameter_model = types._UpdateRuntimeMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -2939,7 +2921,7 @@ async def _update( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _UpdateAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _UpdateRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2970,7 +2952,7 @@ async def _update( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.RuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -3003,13 +2985,13 @@ async def _purge( builtins.list[types.MemoryConjunctionFilterOrDict] ] = None, force: Optional[bool] = None, - config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEnginePurgeMemoriesOperation: + config: Optional[types.PurgeRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimePurgeMemoriesOperation: """ - Purges memories from an Agent Engine. + Purges memories from an Agent Runtime. """ - parameter_model = types._PurgeAgentEngineMemoriesRequestParameters( + parameter_model = types._PurgeRuntimeMemoriesRequestParameters( name=name, filter=filter, filter_groups=filter_groups, @@ -3023,7 +3005,7 @@ async def _purge( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _PurgeAgentEngineMemoriesRequestParameters_to_vertex( + request_dict = _PurgeRuntimeMemoriesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -3054,7 +3036,7 @@ async def _purge( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEnginePurgeMemoriesOperation._from_response( + return_value = types.RuntimePurgeMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -3091,7 +3073,7 @@ def revisions(self) -> "memory_revisions_module.AsyncMemoryRevisions": ) except ImportError as e: raise ImportError( - "The 'agent_engines.memories.revisions' module requires " + "The 'runtimes.memories.revisions' module requires " "additional packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e @@ -3103,9 +3085,9 @@ async def create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - """Creates a new memory in the Agent Engine. + config: Optional[types.RuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: + """Creates a new memory in the Agent Runtime. Args: name (str): @@ -3114,16 +3096,16 @@ async def create( Required. The fact to be stored in the memory. scope (dict[str, str]): Required. The scope of the memory. For example, {"user_id": "123"}. - config (AgentEngineMemoryConfigOrDict): + config (RuntimeMemoryConfigOrDict): Optional. The configuration for the memory. Returns: - AgentEngineMemoryOperation: The operation for creating the memory. + RuntimeMemoryOperation: The operation for creating the memory. """ if config is None: - config = types.AgentEngineMemoryConfig() + config = types.RuntimeMemoryConfig() elif isinstance(config, dict): - config = types.AgentEngineMemoryConfig.model_validate(config) + config = types.RuntimeMemoryConfig.model_validate(config) operation = await self._create( name=name, fact=fact, @@ -3132,7 +3114,7 @@ async def create( ) if config.wait_for_completion: if not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -3161,13 +3143,13 @@ async def generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - """Generates memories for the agent engine. + config: Optional[types.GenerateRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: + """Generates memories for the agent runtime. Args: name (str): - Required. The name of the agent engine to generate memories for. + Required. The name of the agent runtime to generate memories for. vertex_session_source (GenerateMemoriesRequestVertexSessionSource): Optional. The vertex session source to use for generating memories. Only one of vertex_session_source, @@ -3188,13 +3170,13 @@ async def generate( Optional. The configuration for the memories to generate. Returns: - AgentEngineGenerateMemoriesOperation: + RuntimeGenerateMemoriesOperation: The operation for generating the memories. """ if config is None: - config = types.GenerateAgentEngineMemoriesConfig() + config = types.GenerateRuntimeMemoriesConfig() elif isinstance(config, dict): - config = types.GenerateAgentEngineMemoriesConfig.model_validate(config) + config = types.GenerateRuntimeMemoriesConfig.model_validate(config) operation = await self._generate( name=name, vertex_session_source=vertex_session_source, @@ -3204,7 +3186,7 @@ async def generate( config=config, ) if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_generate_memories_operation, poll_interval_seconds=0.5, @@ -3217,14 +3199,14 @@ async def list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, + config: Optional[types.ListRuntimeMemoryConfigOrDict] = None, ) -> AsyncPager[types.Memory]: - """Lists Agent Engine memories. + """Lists Agent Runtime memories. Args: name (str): - Required. The name of the agent engine to list memories for. - config (ListAgentEngineMemoryConfig): + Required. The name of the agent runtime to list memories for. + config (ListRuntimeMemoryConfig): Optional. The configuration for the memories to list. Returns: @@ -3249,13 +3231,13 @@ async def retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveRuntimeMemoriesConfigOrDict] = None, ) -> AsyncPager[types.RetrieveMemoriesResponseRetrievedMemory]: """Retrieves memories for the agent. Args: name (str): - Required. The name of the agent engine to retrieve memories for. + Required. The name of the agent runtime to retrieve memories for. scope (dict[str, str]): Required. The scope of the memories to retrieve. For example, {"user_id": "123"}. @@ -3265,7 +3247,7 @@ async def retrieve( simple_retrieval_params (RetrieveMemoriesRequestSimpleRetrievalParams): Optional. The simple retrieval parameters to use for retrieving memories. - config (RetrieveAgentEngineMemoriesConfig): + config (RetrieveRuntimeMemoriesConfig): Optional. The configuration for the memories to retrieve. Returns: @@ -3296,8 +3278,8 @@ async def rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineRollbackMemoryOperation: + config: Optional[types.RollbackRuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeRollbackMemoryOperation: """Rolls back a memory to a previous revision. Args: @@ -3305,24 +3287,24 @@ async def rollback( Required. The name of the memory to rollback. target_revision_id (str): Required. The revision ID to roll back to - config (RollbackAgentEngineMemoryConfig): + config (RollbackRuntimeMemoryConfig): Optional. The configuration for the rollback. Returns: - AgentEngineRollbackMemoryOperation: + RuntimeRollbackMemoryOperation: The operation for rolling back the memory. """ if config is None: - config = types.RollbackAgentEngineMemoryConfig() + config = types.RollbackRuntimeMemoryConfig() elif isinstance(config, dict): - config = types.RollbackAgentEngineMemoryConfig.model_validate(config) + config = types.RollbackRuntimeMemoryConfig.model_validate(config) operation = await self._rollback( name=name, target_revision_id=target_revision_id, config=config, ) if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -3338,13 +3320,13 @@ async def purge( filter: Optional[str] = None, filter_groups: Optional[List[types.MemoryConjunctionFilter]] = None, force: bool = False, - config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEnginePurgeMemoriesOperation: - """Purges memories from an Agent Engine. + config: Optional[types.PurgeRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimePurgeMemoriesOperation: + """Purges memories from an Agent Runtime. Args: name (str): - Required. The name of the Agent Engine to purge memories from. + Required. The name of the Agent Runtime to purge memories from. filter (str): Optional. The standard list filter to determine which memories to purge. filter_groups (list[MemoryConjunctionFilter]): @@ -3354,17 +3336,17 @@ async def purge( force (bool): Optional. Whether to force the purge operation. If false, the operation will be staged but not executed. - config (PurgeAgentEngineMemoriesConfig): + config (PurgeRuntimeMemoriesConfig): Optional. The configuration for the purge operation. Returns: - AgentEnginePurgeMemoriesOperation: + RuntimePurgeMemoriesOperation: The operation for purging the memories. """ if config is None: - config = types.PurgeAgentEngineMemoriesConfig() + config = types.PurgeRuntimeMemoriesConfig() elif isinstance(config, dict): - config = types.PurgeAgentEngineMemoriesConfig.model_validate(config) + config = types.PurgeRuntimeMemoriesConfig.model_validate(config) operation = await self._purge( name=name, filter=filter, @@ -3373,7 +3355,7 @@ async def purge( config=config, ) if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -3396,11 +3378,11 @@ async def ingest_events( ] = None, config: Optional[types.IngestEventsConfigOrDict] = None, ) -> types.MemoryBankIngestEventsOperation: - """Ingests events into an Agent Engine. + """Ingests events into an Agent Runtime. Example usage: ``` - await client.aio.agent_engines.memories.ingest_events( + await client.aio.runtimes.memories.ingest_events( name="projects/test-project/locations/us-central1/reasoningEngines/test-agent-engine", scope={"user_id": "test-user-id"}, direct_contents_source={ @@ -3425,7 +3407,7 @@ async def ingest_events( Args: name (str): - Required. The name of the Agent Engine to ingest events into. + Required. The name of the Agent Runtime to ingest events into. scope (dict[str, str]): Required. The scope of the events to ingest. For example, {"user_id": "123"}. @@ -3440,7 +3422,7 @@ async def ingest_events( Optional. The configuration for the ingest events operation. Returns: - AgentEngineIngestEventsOperation: + RuntimeIngestEventsOperation: The operation for ingesting the events. """ if config is None: @@ -3456,7 +3438,7 @@ async def ingest_events( config=config, ) if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, diff --git a/agentplatform/_genai/memory_revisions.py b/agentplatform/_genai/memory_revisions.py index c1a1e307b1..910b44ae53 100644 --- a/agentplatform/_genai/memory_revisions.py +++ b/agentplatform/_genai/memory_revisions.py @@ -34,7 +34,7 @@ logger.setLevel(logging.INFO) -def _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( +def _GetRuntimeMemoryRevisionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -45,7 +45,7 @@ def _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( return to_object -def _ListAgentEngineMemoryRevisionsConfig_to_vertex( +def _ListRuntimeMemoryRevisionsConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -63,7 +63,7 @@ def _ListAgentEngineMemoryRevisionsConfig_to_vertex( return to_object -def _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( +def _ListRuntimeMemoryRevisionsRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -72,7 +72,7 @@ def _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineMemoryRevisionsConfig_to_vertex( + _ListRuntimeMemoryRevisionsConfig_to_vertex( getv(from_object, ["config"]), to_object ) @@ -85,23 +85,23 @@ def get( self, *, name: str, - config: Optional[types.GetAgentEngineMemoryRevisionConfigOrDict] = None, + config: Optional[types.GetRuntimeMemoryRevisionConfigOrDict] = None, ) -> types.MemoryRevision: """ - Gets an agent engine memory revision. + Gets an agent runtime memory revision. Args: - name (str): Required. The name of the Agent Engine memory revision to get. Format: + name (str): Required. The name of the Agent Runtime memory revision to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}/revisions/{revision_id}`. - config (GetAgentEngineMemoryRevisionConfig): - Optional. Additional configurations for getting the Agent Engine memory revision. + config (GetRuntimeMemoryRevisionConfig): + Optional. Additional configurations for getting the Agent Runtime memory revision. Returns: - AgentEngineMemoryRevision: The requested Agent Engine memory revision. + RuntimeMemoryRevision: The requested Agent Runtime memory revision. """ - parameter_model = types._GetAgentEngineMemoryRevisionRequestParameters( + parameter_model = types._GetRuntimeMemoryRevisionRequestParameters( name=name, config=config, ) @@ -112,7 +112,7 @@ def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( + request_dict = _GetRuntimeMemoryRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -169,23 +169,23 @@ def _list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, - ) -> types.ListAgentEngineMemoryRevisionsResponse: + config: Optional[types.ListRuntimeMemoryRevisionsConfigOrDict] = None, + ) -> types.ListRuntimeMemoryRevisionsResponse: """ - Lists Agent Engine memory revisions. + Lists Agent Runtime memory revisions. Args: - name (str): Required. The name of the Agent Engine memory to list revisions for. Format: + name (str): Required. The name of the Agent Runtime memory to list revisions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}`. - config (ListAgentEngineMemoryRevisionsConfig): - Optional. Additional configurations for listing the Agent Engine memory revisions. + config (ListRuntimeMemoryRevisionsConfig): + Optional. Additional configurations for listing the Agent Runtime memory revisions. Returns: - ListAgentEngineMemoryRevisionsResponse: The requested Agent Engine memory revisions. + ListRuntimeMemoryRevisionsResponse: The requested Agent Runtime memory revisions. """ - parameter_model = types._ListAgentEngineMemoryRevisionsRequestParameters( + parameter_model = types._ListRuntimeMemoryRevisionsRequestParameters( name=name, config=config, ) @@ -196,7 +196,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( + request_dict = _ListRuntimeMemoryRevisionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -225,7 +225,7 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineMemoryRevisionsResponse._from_response( + return_value = types.ListRuntimeMemoryRevisionsResponse._from_response( response=response_dict, kwargs=( { @@ -253,14 +253,14 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, + config: Optional[types.ListRuntimeMemoryRevisionsConfigOrDict] = None, ) -> Iterator[types.MemoryRevision]: - """Lists Agent Engine memory revisions. + """Lists Agent Runtime memory revisions. Args: name (str): Required. The name of the Memory to list revisions for. - config (ListAgentEngineMemoryRevisionsConfigOrDict): + config (ListRuntimeMemoryRevisionsConfigOrDict): Optional. The configuration for the memories to list revisions. Returns: @@ -281,23 +281,23 @@ async def get( self, *, name: str, - config: Optional[types.GetAgentEngineMemoryRevisionConfigOrDict] = None, + config: Optional[types.GetRuntimeMemoryRevisionConfigOrDict] = None, ) -> types.MemoryRevision: """ - Gets an agent engine memory revision. + Gets an agent runtime memory revision. Args: - name (str): Required. The name of the Agent Engine memory revision to get. Format: + name (str): Required. The name of the Agent Runtime memory revision to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}/revisions/{revision_id}`. - config (GetAgentEngineMemoryRevisionConfig): - Optional. Additional configurations for getting the Agent Engine memory revision. + config (GetRuntimeMemoryRevisionConfig): + Optional. Additional configurations for getting the Agent Runtime memory revision. Returns: - AgentEngineMemoryRevision: The requested Agent Engine memory revision. + RuntimeMemoryRevision: The requested Agent Runtime memory revision. """ - parameter_model = types._GetAgentEngineMemoryRevisionRequestParameters( + parameter_model = types._GetRuntimeMemoryRevisionRequestParameters( name=name, config=config, ) @@ -308,7 +308,7 @@ async def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( + request_dict = _GetRuntimeMemoryRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -367,23 +367,23 @@ async def _list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, - ) -> types.ListAgentEngineMemoryRevisionsResponse: + config: Optional[types.ListRuntimeMemoryRevisionsConfigOrDict] = None, + ) -> types.ListRuntimeMemoryRevisionsResponse: """ - Lists Agent Engine memory revisions. + Lists Agent Runtime memory revisions. Args: - name (str): Required. The name of the Agent Engine memory to list revisions for. Format: + name (str): Required. The name of the Agent Runtime memory to list revisions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}`. - config (ListAgentEngineMemoryRevisionsConfig): - Optional. Additional configurations for listing the Agent Engine memory revisions. + config (ListRuntimeMemoryRevisionsConfig): + Optional. Additional configurations for listing the Agent Runtime memory revisions. Returns: - ListAgentEngineMemoryRevisionsResponse: The requested Agent Engine memory revisions. + ListRuntimeMemoryRevisionsResponse: The requested Agent Runtime memory revisions. """ - parameter_model = types._ListAgentEngineMemoryRevisionsRequestParameters( + parameter_model = types._ListRuntimeMemoryRevisionsRequestParameters( name=name, config=config, ) @@ -394,7 +394,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( + request_dict = _ListRuntimeMemoryRevisionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -425,7 +425,7 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineMemoryRevisionsResponse._from_response( + return_value = types.ListRuntimeMemoryRevisionsResponse._from_response( response=response_dict, kwargs=( { @@ -453,14 +453,14 @@ async def list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, + config: Optional[types.ListRuntimeMemoryRevisionsConfigOrDict] = None, ) -> AsyncPager[types.MemoryRevision]: - """Lists Agent Engine memory revisions. + """Lists Agent Runtime memory revisions. Args: name (str): Required. The name of the Memory to list revisions for. - config (ListAgentEngineMemoryRevisionsConfigOrDict): + config (ListRuntimeMemoryRevisionsConfigOrDict): Optional. The configuration for the memories to list revisions. Returns: diff --git a/agentplatform/_genai/runtime_revisions.py b/agentplatform/_genai/runtime_revisions.py index e4a8d49fd1..bb7476b159 100644 --- a/agentplatform/_genai/runtime_revisions.py +++ b/agentplatform/_genai/runtime_revisions.py @@ -27,7 +27,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import AsyncPager, Pager -from . import _agent_engines_utils +from . import _runtimes_utils from . import types logger = logging.getLogger("agentplatform_genai.runtimerevisions") @@ -35,7 +35,7 @@ logger.setLevel(logging.INFO) -def _DeleteAgentEngineRuntimeRevisionRequestParameters_to_vertex( +def _DeleteRuntimeRevisionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -46,31 +46,31 @@ def _DeleteAgentEngineRuntimeRevisionRequestParameters_to_vertex( return to_object -def _GetAgentEngineRuntimeRevisionRequestParameters_to_vertex( +def _GetDeleteRuntimeRevisionOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + if getv(from_object, ["operation_name"]) is not None: + setv( + to_object, ["_url", "operationName"], getv(from_object, ["operation_name"]) + ) return to_object -def _GetDeleteAgentEngineRuntimeRevisionOperationParameters_to_vertex( +def _GetRuntimeRevisionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: to_object: dict[str, Any] = {} - if getv(from_object, ["operation_name"]) is not None: - setv( - to_object, ["_url", "operationName"], getv(from_object, ["operation_name"]) - ) + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) return to_object -def _ListAgentEngineRuntimeRevisionsConfig_to_vertex( +def _ListRuntimeRevisionsConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -88,7 +88,7 @@ def _ListAgentEngineRuntimeRevisionsConfig_to_vertex( return to_object -def _ListAgentEngineRuntimeRevisionsRequestParameters_to_vertex( +def _ListRuntimeRevisionsRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -97,14 +97,12 @@ def _ListAgentEngineRuntimeRevisionsRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineRuntimeRevisionsConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _ListRuntimeRevisionsConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _QueryAgentEngineRuntimeRevisionConfig_to_vertex( +def _QueryRuntimeRevisionConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -122,7 +120,7 @@ def _QueryAgentEngineRuntimeRevisionConfig_to_vertex( return to_object -def _QueryAgentEngineRuntimeRevisionRequestParameters_to_vertex( +def _QueryRuntimeRevisionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -131,9 +129,7 @@ def _QueryAgentEngineRuntimeRevisionRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _QueryAgentEngineRuntimeRevisionConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _QueryRuntimeRevisionConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object @@ -144,13 +140,13 @@ def _get( self, *, name: str, - config: Optional[types.GetAgentEngineRuntimeRevisionConfigOrDict] = None, + config: Optional[types.GetRuntimeRevisionConfigOrDict] = None, ) -> types.ReasoningEngineRuntimeRevision: """ - Get an agent engine runtime revision instance. + Get an agent runtime runtime revision instance. """ - parameter_model = types._GetAgentEngineRuntimeRevisionRequestParameters( + parameter_model = types._GetRuntimeRevisionRequestParameters( name=name, config=config, ) @@ -161,7 +157,7 @@ def _get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineRuntimeRevisionRequestParameters_to_vertex( + request_dict = _GetRuntimeRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -218,7 +214,7 @@ def _list( self, *, name: str, - config: Optional[types.ListAgentEngineRuntimeRevisionsConfigOrDict] = None, + config: Optional[types.ListRuntimeRevisionsConfigOrDict] = None, ) -> types.ListReasoningEnginesRuntimeRevisionsResponse: """ Lists reasoning engine runtime revisions. @@ -226,7 +222,7 @@ def _list( Args: name (str): Required. The name of the reasoning engine to list runtime revisions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineRuntimeRevisionsConfig): + config (ListRuntimeRevisionsConfig): Optional. Additional configurations for listing the reasoning engine runtime revisions. Returns: @@ -234,7 +230,7 @@ def _list( """ - parameter_model = types._ListAgentEngineRuntimeRevisionsRequestParameters( + parameter_model = types._ListRuntimeRevisionsRequestParameters( name=name, config=config, ) @@ -245,7 +241,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineRuntimeRevisionsRequestParameters_to_vertex( + request_dict = _ListRuntimeRevisionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -304,23 +300,23 @@ def _delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, - ) -> types.DeleteAgentEngineRuntimeRevisionOperation: + config: Optional[types.DeleteRuntimeRevisionConfigOrDict] = None, + ) -> types.DeleteRuntimeRevisionOperation: """ - Delete an Agent Engine runtime revision. + Delete an Agent Runtime runtime revision. Args: - name (str): Required. The name of the Agent Engine runtime revision to be deleted. Format: + name (str): Required. The name of the Agent Runtime runtime revision to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/runtimeRevisions/{runtime_revision_id}`. - config (DeleteAgentEngineRuntimeRevisionConfig): - Optional. Additional configurations for deleting the Agent Engine runtime revision. + config (DeleteRuntimeRevisionConfig): + Optional. Additional configurations for deleting the Agent Runtime runtime revision. Returns: - DeleteAgentEngineRuntimeRevisionOperation: The operation for deleting the Agent Engine runtime revision. + DeleteRuntimeRevisionOperation: The operation for deleting the Agent Runtime runtime revision. """ - parameter_model = types._DeleteAgentEngineRuntimeRevisionRequestParameters( + parameter_model = types._DeleteRuntimeRevisionRequestParameters( name=name, config=config, ) @@ -331,7 +327,7 @@ def _delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineRuntimeRevisionRequestParameters_to_vertex( + request_dict = _DeleteRuntimeRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -360,7 +356,7 @@ def _delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineRuntimeRevisionOperation._from_response( + return_value = types.DeleteRuntimeRevisionOperation._from_response( response=response_dict, kwargs=( { @@ -388,11 +384,9 @@ def _get_delete_runtime_revision_operation( self, *, operation_name: str, - config: Optional[ - types.GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict - ] = None, - ) -> types.DeleteAgentEngineRuntimeRevisionOperation: - parameter_model = types._GetDeleteAgentEngineRuntimeRevisionOperationParameters( + config: Optional[types.GetDeleteRuntimeRevisionOperationConfigOrDict] = None, + ) -> types.DeleteRuntimeRevisionOperation: + parameter_model = types._GetDeleteRuntimeRevisionOperationParameters( operation_name=operation_name, config=config, ) @@ -403,10 +397,8 @@ def _get_delete_runtime_revision_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = ( - _GetDeleteAgentEngineRuntimeRevisionOperationParameters_to_vertex( - parameter_model - ) + request_dict = _GetDeleteRuntimeRevisionOperationParameters_to_vertex( + parameter_model ) request_url_dict = request_dict.get("_url") if request_url_dict: @@ -434,7 +426,7 @@ def _get_delete_runtime_revision_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineRuntimeRevisionOperation._from_response( + return_value = types.DeleteRuntimeRevisionOperation._from_response( response=response_dict, kwargs=( { @@ -462,13 +454,13 @@ def _query( self, *, name: str, - config: Optional[types.QueryAgentEngineRuntimeRevisionConfigOrDict] = None, + config: Optional[types.QueryRuntimeRevisionConfigOrDict] = None, ) -> types.QueryReasoningEngineResponse: """ - Query an Agent Engine runtime revision. + Query an Agent runtime revision. """ - parameter_model = types._QueryAgentEngineRuntimeRevisionRequestParameters( + parameter_model = types._QueryRuntimeRevisionRequestParameters( name=name, config=config, ) @@ -479,7 +471,7 @@ def _query( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _QueryAgentEngineRuntimeRevisionRequestParameters_to_vertex( + request_dict = _QueryRuntimeRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -536,47 +528,45 @@ def get( self, *, name: str, - config: Optional[types.GetAgentEngineRuntimeRevisionConfigOrDict] = None, - ) -> types.AgentEngineRuntimeRevision: - """Gets an agent engine runtime revision. + config: Optional[types.GetRuntimeRevisionConfigOrDict] = None, + ) -> types.RuntimeRevision: + """Gets an agent runtime revision. Args: - name (str): Required. The name of the Agent Engine runtime revision to get. Format: + name (str): Required. The name of the Agent Runtime revision to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/runtimeRevisions/{runtime_revision_id}`. - config (GetAgentEngineRuntimeRevisionConfigOrDict): - Optional. Additional configurations for getting the Agent Engine runtime revision. + config (GetRuntimeRevisionConfigOrDict): + Optional. Additional configurations for getting the Agent Runtime revision. Returns: - AgentEngineRuntimeRevision: The requested Agent Engine runtime revision instance. + RuntimeRevision: The requested Agent Runtime revision instance. """ api_resource = self._get(name=name, config=config) - agent_engine_runtime_revision = types.AgentEngineRuntimeRevision( + runtime_revision = types.RuntimeRevision( api_client=self, api_async_client=AsyncRuntimeRevisions(api_client_=self._api_client), api_resource=api_resource, ) if api_resource.spec: - self._register_api_methods( - agent_engine_runtime_revision=agent_engine_runtime_revision - ) - return agent_engine_runtime_revision + self._register_api_methods(runtime_revision=runtime_revision) + return runtime_revision def list( self, *, name: str, - config: Optional[types.ListAgentEngineRuntimeRevisionsConfigOrDict] = None, - ) -> Iterator[types.AgentEngineRuntimeRevision]: + config: Optional[types.ListRuntimeRevisionsConfigOrDict] = None, + ) -> Iterator[types.RuntimeRevision]: """Lists all reasoning engine runtime revision instances matching the given query. Args: name (str): Required. The name of the reasoning engine to list runtime revisions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineRuntimeRevisionsConfig): + config (ListRuntimeRevisionsConfig): Optional. Additional configurations for listing the reasoning engine runtime revisions. Returns: - Iterable[AgentEngineRuntimeRevision]: An iterable of runtime revisions. + Iterable[RuntimeRevision]: An iterable of runtime revisions. """ list_pager: Pager[types.ReasoningEngineRuntimeRevision] = Pager( "reasoning_engine_runtime_revisions", @@ -586,7 +576,7 @@ def list( ) return ( - types.AgentEngineRuntimeRevision( + types.RuntimeRevision( api_client=self, api_async_client=AsyncRuntimeRevisions(api_client_=self._api_client), api_resource=runtime_revision, @@ -598,29 +588,29 @@ def delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, - ) -> types.DeleteAgentEngineRuntimeRevisionOperation: - """Delete an Agent Engine runtime revision. + config: Optional[types.DeleteRuntimeRevisionConfigOrDict] = None, + ) -> types.DeleteRuntimeRevisionOperation: + """Delete an Agent Runtime revision. Args: - name (str): Required. The name of the Agent Engine runtime revision to be deleted. Format: + name (str): Required. The name of the Agent Runtime revision to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/runtimeRevisions/{runtime_revision_id}`. - config (DeleteAgentEngineRuntimeRevisionConfig): - Optional. Additional configurations for deleting the Agent Engine runtime revision. + config (DeleteRuntimeRevisionConfig): + Optional. Additional configurations for deleting the Agent Runtime revision. Returns: - DeleteAgentEngineRuntimeRevisionOperation: The operation for deleting the Agent Engine runtime revision. + DeleteRuntimeRevisionOperation: The operation for deleting the Agent Runtime revision. """ if config is None: - config = types.DeleteAgentEngineRuntimeRevisionConfig() + config = types.DeleteRuntimeRevisionConfig() elif isinstance(config, dict): - config = types.DeleteAgentEngineRuntimeRevisionConfig.model_validate(config) + config = types.DeleteRuntimeRevisionConfig.model_validate(config) operation = self._delete( name=name, config=config, ) if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_delete_runtime_revision_operation, poll_interval_seconds=0.5, @@ -634,40 +624,38 @@ def delete( def _register_api_methods( self, *, - agent_engine_runtime_revision: types.AgentEngineRuntimeRevision, - ) -> types.AgentEngineRuntimeRevision: - """Registers the API methods for the agent engine runtime revision.""" + runtime_revision: types.RuntimeRevision, + ) -> types.RuntimeRevision: + """Registers the API methods for the agent runtime revision.""" try: - _agent_engines_utils._register_api_methods_or_raise( - agent_engine=agent_engine_runtime_revision, + _runtimes_utils._register_api_methods_or_raise( + runtime=runtime_revision, wrap_operation_fn={ - "": _agent_engines_utils._wrap_query_operation, # type: ignore[dict-item] - "async": _agent_engines_utils._wrap_async_query_operation, # type: ignore[dict-item] - "stream": _agent_engines_utils._wrap_stream_query_operation, # type: ignore[dict-item] - "async_stream": _agent_engines_utils._wrap_async_stream_query_operation, # type: ignore[dict-item] - "a2a_extension": _agent_engines_utils._wrap_a2a_operation, + "": _runtimes_utils._wrap_query_operation, # type: ignore[dict-item] + "async": _runtimes_utils._wrap_async_query_operation, # type: ignore[dict-item] + "stream": _runtimes_utils._wrap_stream_query_operation, # type: ignore[dict-item] + "async_stream": _runtimes_utils._wrap_async_stream_query_operation, # type: ignore[dict-item] + "a2a_extension": _runtimes_utils._wrap_a2a_operation, }, ) except Exception as e: logger.warning( - _agent_engines_utils._FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e + _runtimes_utils._FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e ) - return agent_engine_runtime_revision + return runtime_revision def _stream_query( self, *, name: str, - config: Optional[types.QueryAgentEngineRuntimeRevisionConfigOrDict] = None, + config: Optional[types.QueryRuntimeRevisionConfigOrDict] = None, ) -> Iterator[Any]: - """Streams the response of the agent engine.""" - parameter_model = types._QueryAgentEngineRuntimeRevisionRequestParameters( + """Streams the response of the agent runtime.""" + parameter_model = types._QueryRuntimeRevisionRequestParameters( name=name, config=config, ) - request_dict = _QueryAgentEngineRuntimeRevisionRequestParameters_to_vertex( - parameter_model - ) + request_dict = _QueryRuntimeRevisionRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) @@ -696,16 +684,14 @@ async def _async_stream_query( self, *, name: str, - config: Optional[types.QueryAgentEngineRuntimeRevisionConfigOrDict] = None, + config: Optional[types.QueryRuntimeRevisionConfigOrDict] = None, ) -> AsyncIterator[Any]: - """Streams the response of the agent engine.""" - parameter_model = types._QueryAgentEngineRuntimeRevisionRequestParameters( + """Streams the response of the agent runtime.""" + parameter_model = types._QueryRuntimeRevisionRequestParameters( name=name, config=config, ) - request_dict = _QueryAgentEngineRuntimeRevisionRequestParameters_to_vertex( - parameter_model - ) + request_dict = _QueryRuntimeRevisionRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) @@ -738,13 +724,13 @@ async def _get( self, *, name: str, - config: Optional[types.GetAgentEngineRuntimeRevisionConfigOrDict] = None, + config: Optional[types.GetRuntimeRevisionConfigOrDict] = None, ) -> types.ReasoningEngineRuntimeRevision: """ - Get an agent engine runtime revision instance. + Get an agent runtime runtime revision instance. """ - parameter_model = types._GetAgentEngineRuntimeRevisionRequestParameters( + parameter_model = types._GetRuntimeRevisionRequestParameters( name=name, config=config, ) @@ -755,7 +741,7 @@ async def _get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineRuntimeRevisionRequestParameters_to_vertex( + request_dict = _GetRuntimeRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -814,7 +800,7 @@ async def _list( self, *, name: str, - config: Optional[types.ListAgentEngineRuntimeRevisionsConfigOrDict] = None, + config: Optional[types.ListRuntimeRevisionsConfigOrDict] = None, ) -> types.ListReasoningEnginesRuntimeRevisionsResponse: """ Lists reasoning engine runtime revisions. @@ -822,7 +808,7 @@ async def _list( Args: name (str): Required. The name of the reasoning engine to list runtime revisions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineRuntimeRevisionsConfig): + config (ListRuntimeRevisionsConfig): Optional. Additional configurations for listing the reasoning engine runtime revisions. Returns: @@ -830,7 +816,7 @@ async def _list( """ - parameter_model = types._ListAgentEngineRuntimeRevisionsRequestParameters( + parameter_model = types._ListRuntimeRevisionsRequestParameters( name=name, config=config, ) @@ -841,7 +827,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineRuntimeRevisionsRequestParameters_to_vertex( + request_dict = _ListRuntimeRevisionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -902,23 +888,23 @@ async def _delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, - ) -> types.DeleteAgentEngineRuntimeRevisionOperation: + config: Optional[types.DeleteRuntimeRevisionConfigOrDict] = None, + ) -> types.DeleteRuntimeRevisionOperation: """ - Delete an Agent Engine runtime revision. + Delete an Agent Runtime runtime revision. Args: - name (str): Required. The name of the Agent Engine runtime revision to be deleted. Format: + name (str): Required. The name of the Agent Runtime runtime revision to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/runtimeRevisions/{runtime_revision_id}`. - config (DeleteAgentEngineRuntimeRevisionConfig): - Optional. Additional configurations for deleting the Agent Engine runtime revision. + config (DeleteRuntimeRevisionConfig): + Optional. Additional configurations for deleting the Agent Runtime runtime revision. Returns: - DeleteAgentEngineRuntimeRevisionOperation: The operation for deleting the Agent Engine runtime revision. + DeleteRuntimeRevisionOperation: The operation for deleting the Agent Runtime runtime revision. """ - parameter_model = types._DeleteAgentEngineRuntimeRevisionRequestParameters( + parameter_model = types._DeleteRuntimeRevisionRequestParameters( name=name, config=config, ) @@ -929,7 +915,7 @@ async def _delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineRuntimeRevisionRequestParameters_to_vertex( + request_dict = _DeleteRuntimeRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -960,7 +946,7 @@ async def _delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineRuntimeRevisionOperation._from_response( + return_value = types.DeleteRuntimeRevisionOperation._from_response( response=response_dict, kwargs=( { @@ -988,11 +974,9 @@ async def _get_delete_runtime_revision_operation( self, *, operation_name: str, - config: Optional[ - types.GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict - ] = None, - ) -> types.DeleteAgentEngineRuntimeRevisionOperation: - parameter_model = types._GetDeleteAgentEngineRuntimeRevisionOperationParameters( + config: Optional[types.GetDeleteRuntimeRevisionOperationConfigOrDict] = None, + ) -> types.DeleteRuntimeRevisionOperation: + parameter_model = types._GetDeleteRuntimeRevisionOperationParameters( operation_name=operation_name, config=config, ) @@ -1003,10 +987,8 @@ async def _get_delete_runtime_revision_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = ( - _GetDeleteAgentEngineRuntimeRevisionOperationParameters_to_vertex( - parameter_model - ) + request_dict = _GetDeleteRuntimeRevisionOperationParameters_to_vertex( + parameter_model ) request_url_dict = request_dict.get("_url") if request_url_dict: @@ -1036,7 +1018,7 @@ async def _get_delete_runtime_revision_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineRuntimeRevisionOperation._from_response( + return_value = types.DeleteRuntimeRevisionOperation._from_response( response=response_dict, kwargs=( { @@ -1064,13 +1046,13 @@ async def _query( self, *, name: str, - config: Optional[types.QueryAgentEngineRuntimeRevisionConfigOrDict] = None, + config: Optional[types.QueryRuntimeRevisionConfigOrDict] = None, ) -> types.QueryReasoningEngineResponse: """ - Query an Agent Engine runtime revision. + Query an Agent runtime revision. """ - parameter_model = types._QueryAgentEngineRuntimeRevisionRequestParameters( + parameter_model = types._QueryRuntimeRevisionRequestParameters( name=name, config=config, ) @@ -1081,7 +1063,7 @@ async def _query( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _QueryAgentEngineRuntimeRevisionRequestParameters_to_vertex( + request_dict = _QueryRuntimeRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1140,47 +1122,45 @@ async def get( self, *, name: str, - config: Optional[types.GetAgentEngineRuntimeRevisionConfigOrDict] = None, - ) -> types.AgentEngineRuntimeRevision: - """Gets an agent engine runtime revision. + config: Optional[types.GetRuntimeRevisionConfigOrDict] = None, + ) -> types.RuntimeRevision: + """Gets an agent runtime revision. Args: - name (str): Required. The name of the Agent Engine runtime revision to get. Format: + name (str): Required. The name of the Agent Runtime revision to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/runtimeRevisions/{runtime_revision_id}`. - config (GetAgentEngineRuntimeRevisionConfigOrDict): - Optional. Additional configurations for getting the Agent Engine runtime revision. + config (GetRuntimeRevisionConfigOrDict): + Optional. Additional configurations for getting the Agent Runtime revision. Returns: - AgentEngineRuntimeRevision: The requested Agent Engine runtime revision instance. + RuntimeRevision: The requested Agent Runtime revision instance. """ api_resource = await self._get(name=name, config=config) - agent_engine_runtime_revision = types.AgentEngineRuntimeRevision( + runtime_revision = types.RuntimeRevision( api_client=self, api_async_client=AsyncRuntimeRevisions(api_client_=self._api_client), api_resource=api_resource, ) if api_resource.spec: - self._register_api_methods( - agent_engine_runtime_revision=agent_engine_runtime_revision - ) - return agent_engine_runtime_revision + self._register_api_methods(runtime_revision=runtime_revision) + return runtime_revision async def list( self, *, name: str, - config: Optional[types.ListAgentEngineRuntimeRevisionsConfigOrDict] = None, - ) -> AsyncIterator[types.AgentEngineRuntimeRevision]: + config: Optional[types.ListRuntimeRevisionsConfigOrDict] = None, + ) -> AsyncIterator[types.RuntimeRevision]: """Lists reasoning engine runtime revisions. Args: name (str): Required. The name of the reasoning engine to list runtime revisions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineRuntimeRevisionsConfig): + config (ListRuntimeRevisionsConfig): Optional. Additional configurations for listing the reasoning engine runtime revisions. Returns: - AsyncIterator[AgentEngineRuntimeRevision]: An async iterator of runtime revisions. + AsyncIterator[RuntimeRevision]: An async iterator of runtime revisions. """ list_pager: AsyncPager[types.ReasoningEngineRuntimeRevision] = AsyncPager( "reasoning_engine_runtime_revisions", @@ -1190,7 +1170,7 @@ async def list( ) async for runtime_revision in list_pager: - yield types.AgentEngineRuntimeRevision( + yield types.RuntimeRevision( api_client=self, api_async_client=AsyncRuntimeRevisions(api_client_=self._api_client), api_resource=runtime_revision, @@ -1200,29 +1180,29 @@ async def delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, - ) -> types.DeleteAgentEngineRuntimeRevisionOperation: - """Delete an Agent Engine runtime revision. + config: Optional[types.DeleteRuntimeRevisionConfigOrDict] = None, + ) -> types.DeleteRuntimeRevisionOperation: + """Delete an Agent Runtime revision. Args: - name (str): Required. The name of the Agent Engine runtime revision to be deleted. Format: + name (str): Required. The name of the Agent Runtime revision to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/runtimeRevisions/{runtime_revision_id}`. - config (DeleteAgentEngineRuntimeRevisionConfig): - Optional. Additional configurations for deleting the Agent Engine runtime revision. + config (DeleteRuntimeRevisionConfig): + Optional. Additional configurations for deleting the Agent Runtime revision. Returns: - DeleteAgentEngineRuntimeRevisionOperation: The operation for deleting the Agent Engine runtime revision. + DeleteRuntimeRevisionOperation: The operation for deleting the Agent Runtime revision. """ if config is None: - config = types.DeleteAgentEngineRuntimeRevisionConfig() + config = types.DeleteRuntimeRevisionConfig() elif isinstance(config, dict): - config = types.DeleteAgentEngineRuntimeRevisionConfig.model_validate(config) + config = types.DeleteRuntimeRevisionConfig.model_validate(config) operation = await self._delete( name=name, config=config, ) if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_delete_runtime_revision_operation, poll_interval_seconds=0.5, @@ -1236,22 +1216,22 @@ async def delete( def _register_api_methods( self, *, - agent_engine_runtime_revision: types.AgentEngineRuntimeRevision, - ) -> types.AgentEngineRuntimeRevision: - """Registers the API methods for the agent engine runtime revision.""" + runtime_revision: types.RuntimeRevision, + ) -> types.RuntimeRevision: + """Registers the API methods for the agent runtime revision.""" try: - _agent_engines_utils._register_api_methods_or_raise( - agent_engine=agent_engine_runtime_revision, + _runtimes_utils._register_api_methods_or_raise( + runtime=runtime_revision, wrap_operation_fn={ - "": _agent_engines_utils._wrap_query_operation, # type: ignore[dict-item] - "async": _agent_engines_utils._wrap_async_query_operation, # type: ignore[dict-item] - "stream": _agent_engines_utils._wrap_stream_query_operation, # type: ignore[dict-item] - "async_stream": _agent_engines_utils._wrap_async_stream_query_operation, # type: ignore[dict-item] - "a2a_extension": _agent_engines_utils._wrap_a2a_operation, + "": _runtimes_utils._wrap_query_operation, # type: ignore[dict-item] + "async": _runtimes_utils._wrap_async_query_operation, # type: ignore[dict-item] + "stream": _runtimes_utils._wrap_stream_query_operation, # type: ignore[dict-item] + "async_stream": _runtimes_utils._wrap_async_stream_query_operation, # type: ignore[dict-item] + "a2a_extension": _runtimes_utils._wrap_a2a_operation, }, ) except Exception as e: logger.warning( - _agent_engines_utils._FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e + _runtimes_utils._FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e ) - return agent_engine_runtime_revision + return runtime_revision diff --git a/agentplatform/_genai/runtimes.py b/agentplatform/_genai/runtimes.py index 9f41816bf5..137af12f45 100644 --- a/agentplatform/_genai/runtimes.py +++ b/agentplatform/_genai/runtimes.py @@ -13,20 +13,33 @@ # limitations under the License. # -# Handwritten placeholder code for the runtimes.py file. -# Should be replaced by the generated file. +# Code generated by the Google Gen AI SDK generator DO NOT EDIT. +import builtins import importlib +import json import logging import typing +from typing import Any, AsyncIterator, Iterator, Optional, Sequence, Tuple, Union +from urllib.parse import urlencode +import warnings from google.genai import _api_module +from google.genai import _common +from google.genai import types as genai_types +from google.genai._common import get_value_by_path as getv +from google.genai._common import set_value_by_path as setv +from google.genai.pagers import Pager +from . import _runtimes_utils +from . import types if typing.TYPE_CHECKING: + from . import memories as memories_module from . import runtime_revisions as runtime_revisions_module - _ = runtime_revisions_module + _ = memories_module + __ = runtime_revisions_module logger = logging.getLogger("agentplatform_genai.runtimes") @@ -34,8 +47,1461 @@ logger.setLevel(logging.INFO) +def _CancelQueryJobRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["operation_name"]) is not None: + setv(parent_object, ["operationName"], getv(from_object, ["operation_name"])) + + return to_object + + +def _CancelQueryJobRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + setv( + to_object, + ["config"], + _CancelQueryJobRuntimeConfig_to_vertex( + getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _CheckQueryJobResult_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(parent_object, ["operationName"]) is not None: + setv(to_object, ["operation_name"], getv(parent_object, ["operationName"])) + + if getv(parent_object, ["outputGcsUri"]) is not None: + setv(to_object, ["output_gcs_uri"], getv(parent_object, ["outputGcsUri"])) + + if getv(parent_object, ["status"]) is not None: + setv(to_object, ["status"], getv(parent_object, ["status"])) + + if getv(parent_object, ["result"]) is not None: + setv(to_object, ["result"], getv(parent_object, ["result"])) + + return to_object + + +def _CheckQueryJobRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["retrieve_result"]) is not None: + setv(parent_object, ["retrieveResult"], getv(from_object, ["retrieve_result"])) + + return to_object + + +def _CheckQueryJobRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + setv( + to_object, + ["config"], + _CheckQueryJobRuntimeConfig_to_vertex( + getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _CreateRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["display_name"]) is not None: + setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) + + if getv(from_object, ["description"]) is not None: + setv(parent_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["spec"]) is not None: + setv(parent_object, ["spec"], getv(from_object, ["spec"])) + + if getv(from_object, ["context_spec"]) is not None: + setv( + parent_object, + ["contextSpec"], + _ReasoningEngineContextSpec_to_vertex( + getv(from_object, ["context_spec"]), to_object + ), + ) + + if getv(from_object, ["psc_interface_config"]) is not None: + setv( + parent_object, + ["pscInterfaceConfig"], + getv(from_object, ["psc_interface_config"]), + ) + + if getv(from_object, ["agent_gateway_config"]) is not None: + setv( + parent_object, + ["agentGatewayConfig"], + getv(from_object, ["agent_gateway_config"]), + ) + + if getv(from_object, ["encryption_spec"]) is not None: + setv(parent_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"])) + + if getv(from_object, ["labels"]) is not None: + setv(parent_object, ["labels"], getv(from_object, ["labels"])) + + if getv(from_object, ["source_packages"]) is not None: + setv(parent_object, ["sourcePackages"], getv(from_object, ["source_packages"])) + + if getv(from_object, ["entrypoint_module"]) is not None: + setv( + parent_object, + ["entrypointModule"], + getv(from_object, ["entrypoint_module"]), + ) + + if getv(from_object, ["entrypoint_object"]) is not None: + setv( + parent_object, + ["entrypointObject"], + getv(from_object, ["entrypoint_object"]), + ) + + if getv(from_object, ["requirements_file"]) is not None: + setv( + parent_object, + ["requirementsFile"], + getv(from_object, ["requirements_file"]), + ) + + if getv(from_object, ["agent_framework"]) is not None: + setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) + + if getv(from_object, ["python_version"]) is not None: + setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) + + return to_object + + +def _CreateRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["config"]) is not None: + _CreateRuntimeConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _DeleteRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["force"]) is not None: + setv(to_object, ["force"], getv(from_object, ["force"])) + + return to_object + + +def _GetRuntimeOperationParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["operation_name"]) is not None: + setv( + to_object, ["_url", "operationName"], getv(from_object, ["operation_name"]) + ) + + return to_object + + +def _GetRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + return to_object + + +def _ListReasoningEnginesResponse_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["sdkHttpResponse"]) is not None: + setv(to_object, ["sdk_http_response"], getv(from_object, ["sdkHttpResponse"])) + + if getv(from_object, ["nextPageToken"]) is not None: + setv(to_object, ["next_page_token"], getv(from_object, ["nextPageToken"])) + + if getv(from_object, ["reasoningEngines"]) is not None: + setv( + to_object, + ["reasoning_engines"], + [ + _ReasoningEngine_from_vertex(item, to_object) + for item in getv(from_object, ["reasoningEngines"]) + ], + ) + + return to_object + + +def _ListRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["page_size"]) is not None: + setv(parent_object, ["_query", "pageSize"], getv(from_object, ["page_size"])) + + if getv(from_object, ["page_token"]) is not None: + setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"])) + + if getv(from_object, ["filter"]) is not None: + setv(parent_object, ["_query", "filter"], getv(from_object, ["filter"])) + + return to_object + + +def _ListRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["config"]) is not None: + _ListRuntimeConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _QueryRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["class_method"]) is not None: + setv(parent_object, ["classMethod"], getv(from_object, ["class_method"])) + + if getv(from_object, ["input"]) is not None: + setv(parent_object, ["input"], getv(from_object, ["input"])) + + if getv(from_object, ["include_all_fields"]) is not None: + setv(to_object, ["includeAllFields"], getv(from_object, ["include_all_fields"])) + + return to_object + + +def _QueryRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + _QueryRuntimeConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _ReasoningEngineContextSpecMemoryBankConfig_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["customizationConfigs"]) is not None: + setv( + to_object, + ["customization_configs"], + [item for item in getv(from_object, ["customizationConfigs"])], + ) + + if getv(from_object, ["disableMemoryRevisions"]) is not None: + setv( + to_object, + ["disable_memory_revisions"], + getv(from_object, ["disableMemoryRevisions"]), + ) + + if getv(from_object, ["generationConfig"]) is not None: + setv(to_object, ["generation_config"], getv(from_object, ["generationConfig"])) + + if getv(from_object, ["similaritySearchConfig"]) is not None: + setv( + to_object, + ["similarity_search_config"], + getv(from_object, ["similaritySearchConfig"]), + ) + + if getv(from_object, ["ttlConfig"]) is not None: + setv(to_object, ["ttl_config"], getv(from_object, ["ttlConfig"])) + + if getv(from_object, ["structuredMemoryConfigs"]) is not None: + setv( + to_object, + ["structured_memory_configs"], + [ + _StructuredMemoryConfig_from_vertex(item, to_object) + for item in getv(from_object, ["structuredMemoryConfigs"]) + ], + ) + + return to_object + + +def _ReasoningEngineContextSpecMemoryBankConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["customization_configs"]) is not None: + setv( + to_object, + ["customizationConfigs"], + [item for item in getv(from_object, ["customization_configs"])], + ) + + if getv(from_object, ["disable_memory_revisions"]) is not None: + setv( + to_object, + ["disableMemoryRevisions"], + getv(from_object, ["disable_memory_revisions"]), + ) + + if getv(from_object, ["generation_config"]) is not None: + setv(to_object, ["generationConfig"], getv(from_object, ["generation_config"])) + + if getv(from_object, ["similarity_search_config"]) is not None: + setv( + to_object, + ["similaritySearchConfig"], + getv(from_object, ["similarity_search_config"]), + ) + + if getv(from_object, ["ttl_config"]) is not None: + setv(to_object, ["ttlConfig"], getv(from_object, ["ttl_config"])) + + if getv(from_object, ["structured_memory_configs"]) is not None: + setv( + to_object, + ["structuredMemoryConfigs"], + [ + _StructuredMemoryConfig_to_vertex(item, to_object) + for item in getv(from_object, ["structured_memory_configs"]) + ], + ) + + return to_object + + +def _ReasoningEngineContextSpec_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["memoryBankConfig"]) is not None: + setv( + to_object, + ["memory_bank_config"], + _ReasoningEngineContextSpecMemoryBankConfig_from_vertex( + getv(from_object, ["memoryBankConfig"]), to_object + ), + ) + + return to_object + + +def _ReasoningEngineContextSpec_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["memory_bank_config"]) is not None: + setv( + to_object, + ["memoryBankConfig"], + _ReasoningEngineContextSpecMemoryBankConfig_to_vertex( + getv(from_object, ["memory_bank_config"]), to_object + ), + ) + + return to_object + + +def _ReasoningEngine_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["encryptionSpec"]) is not None: + setv(to_object, ["encryption_spec"], getv(from_object, ["encryptionSpec"])) + + if getv(from_object, ["contextSpec"]) is not None: + setv( + to_object, + ["context_spec"], + _ReasoningEngineContextSpec_from_vertex( + getv(from_object, ["contextSpec"]), to_object + ), + ) + + if getv(from_object, ["createTime"]) is not None: + setv(to_object, ["create_time"], getv(from_object, ["createTime"])) + + if getv(from_object, ["description"]) is not None: + setv(to_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["displayName"]) is not None: + setv(to_object, ["display_name"], getv(from_object, ["displayName"])) + + if getv(from_object, ["etag"]) is not None: + setv(to_object, ["etag"], getv(from_object, ["etag"])) + + if getv(from_object, ["labels"]) is not None: + setv(to_object, ["labels"], getv(from_object, ["labels"])) + + if getv(from_object, ["name"]) is not None: + setv(to_object, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["spec"]) is not None: + setv(to_object, ["spec"], getv(from_object, ["spec"])) + + if getv(from_object, ["updateTime"]) is not None: + setv(to_object, ["update_time"], getv(from_object, ["updateTime"])) + + if getv(from_object, ["trafficConfig"]) is not None: + setv(to_object, ["traffic_config"], getv(from_object, ["trafficConfig"])) + + return to_object + + +def _RunQueryJobRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["input_gcs_uri"]) is not None: + setv(parent_object, ["inputGcsUri"], getv(from_object, ["input_gcs_uri"])) + + if getv(from_object, ["output_gcs_uri"]) is not None: + setv(parent_object, ["outputGcsUri"], getv(from_object, ["output_gcs_uri"])) + + return to_object + + +def _RunQueryJobRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + setv( + to_object, + ["config"], + _RunQueryJobRuntimeConfig_to_vertex( + getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _RuntimeOperation_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["metadata"]) is not None: + setv(to_object, ["metadata"], getv(from_object, ["metadata"])) + + if getv(from_object, ["done"]) is not None: + setv(to_object, ["done"], getv(from_object, ["done"])) + + if getv(from_object, ["error"]) is not None: + setv(to_object, ["error"], getv(from_object, ["error"])) + + if getv(from_object, ["response"]) is not None: + setv( + to_object, + ["response"], + _ReasoningEngine_from_vertex(getv(from_object, ["response"]), to_object), + ) + + return to_object + + +def _StructuredMemoryConfig_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["schemaConfigs"]) is not None: + setv( + to_object, + ["schema_configs"], + [ + _StructuredMemorySchemaConfig_from_vertex(item, to_object) + for item in getv(from_object, ["schemaConfigs"]) + ], + ) + + if getv(from_object, ["scopeKeys"]) is not None: + setv(to_object, ["scope_keys"], getv(from_object, ["scopeKeys"])) + + return to_object + + +def _StructuredMemoryConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["schema_configs"]) is not None: + setv( + to_object, + ["schemaConfigs"], + [ + _StructuredMemorySchemaConfig_to_vertex(item, to_object) + for item in getv(from_object, ["schema_configs"]) + ], + ) + + if getv(from_object, ["scope_keys"]) is not None: + setv(to_object, ["scopeKeys"], getv(from_object, ["scope_keys"])) + + return to_object + + +def _StructuredMemorySchemaConfig_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["schema"]) is not None: + setv(to_object, ["memory_schema"], getv(from_object, ["schema"])) + + if getv(from_object, ["id"]) is not None: + setv(to_object, ["id"], getv(from_object, ["id"])) + + if getv(from_object, ["memoryType"]) is not None: + setv(to_object, ["memory_type"], getv(from_object, ["memoryType"])) + + return to_object + + +def _StructuredMemorySchemaConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["memory_schema"]) is not None: + setv(to_object, ["schema"], getv(from_object, ["memory_schema"])) + + if getv(from_object, ["id"]) is not None: + setv(to_object, ["id"], getv(from_object, ["id"])) + + if getv(from_object, ["memory_type"]) is not None: + setv(to_object, ["memoryType"], getv(from_object, ["memory_type"])) + + return to_object + + +def _UpdateRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["display_name"]) is not None: + setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) + + if getv(from_object, ["description"]) is not None: + setv(parent_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["spec"]) is not None: + setv(parent_object, ["spec"], getv(from_object, ["spec"])) + + if getv(from_object, ["context_spec"]) is not None: + setv( + parent_object, + ["contextSpec"], + _ReasoningEngineContextSpec_to_vertex( + getv(from_object, ["context_spec"]), to_object + ), + ) + + if getv(from_object, ["psc_interface_config"]) is not None: + setv( + parent_object, + ["pscInterfaceConfig"], + getv(from_object, ["psc_interface_config"]), + ) + + if getv(from_object, ["agent_gateway_config"]) is not None: + setv( + parent_object, + ["agentGatewayConfig"], + getv(from_object, ["agent_gateway_config"]), + ) + + if getv(from_object, ["encryption_spec"]) is not None: + setv(parent_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"])) + + if getv(from_object, ["labels"]) is not None: + setv(parent_object, ["labels"], getv(from_object, ["labels"])) + + if getv(from_object, ["source_packages"]) is not None: + setv(parent_object, ["sourcePackages"], getv(from_object, ["source_packages"])) + + if getv(from_object, ["entrypoint_module"]) is not None: + setv( + parent_object, + ["entrypointModule"], + getv(from_object, ["entrypoint_module"]), + ) + + if getv(from_object, ["entrypoint_object"]) is not None: + setv( + parent_object, + ["entrypointObject"], + getv(from_object, ["entrypoint_object"]), + ) + + if getv(from_object, ["requirements_file"]) is not None: + setv( + parent_object, + ["requirementsFile"], + getv(from_object, ["requirements_file"]), + ) + + if getv(from_object, ["agent_framework"]) is not None: + setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) + + if getv(from_object, ["python_version"]) is not None: + setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) + + if getv(from_object, ["update_mask"]) is not None: + setv( + parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"]) + ) + + if getv(from_object, ["traffic_config"]) is not None: + setv(parent_object, ["trafficConfig"], getv(from_object, ["traffic_config"])) + + return to_object + + +def _UpdateRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + _UpdateRuntimeConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + class Runtimes(_api_module.BaseModule): + def cancel_query_job( + self, + *, + name: str, + config: Optional[types.CancelQueryJobRuntimeConfigOrDict] = None, + ) -> types.CancelQueryJobResult: + """ + Cancels a long-running query job on an Agent Runtime. + + Args: + name (str): + Required. The reasoning engine resource name. + config (CancelQueryJobRuntimeConfigOrDict): + Optional. The configuration for the cancel_query_job. + + """ + + parameter_model = types._CancelQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CancelQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:cancelAsyncQuery".format_map(request_url_dict) + else: + path = "{name}:cancelAsyncQuery" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.CancelQueryJobResult._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _check_query_job( + self, + *, + name: str, + config: Optional[types.CheckQueryJobRuntimeConfigOrDict] = None, + ) -> types.CheckQueryJobResult: + """ + Query an Agent Runtime asynchronously. + """ + + parameter_model = types._CheckQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CheckQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:checkQueryJob".format_map(request_url_dict) + else: + path = "{name}:checkQueryJob" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _CheckQueryJobResult_from_vertex(response_dict) + + return_value = types.CheckQueryJobResult._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _run_query_job( + self, + *, + name: str, + config: Optional[types._RunQueryJobRuntimeConfigOrDict] = None, + ) -> types.RuntimeOperation: + """ + Run a query job on an agent runtime. + """ + + parameter_model = types._RunQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _RunQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:asyncQuery".format_map(request_url_dict) + else: + path = "{name}:asyncQuery" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _create( + self, *, config: Optional[types.CreateRuntimeConfigOrDict] = None + ) -> types.RuntimeOperation: + """ + Creates a new Agent Runtime. + """ + + parameter_model = types._CreateRuntimeRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CreateRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _delete( + self, + *, + name: str, + force: Optional[bool] = None, + config: Optional[types.DeleteRuntimeConfigOrDict] = None, + ) -> types.DeleteRuntimeOperation: + """ + Delete an Agent Runtime resource. + + Args: + name (str): + Required. The name of the Agent Runtime to be deleted. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` + or `reasoningEngines/{resource_id}`. + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Agent Runtime has undeleted child resources. Defaults to False. + config (DeleteRuntimeConfig): + Optional. Additional configurations for deleting the Agent Runtime. + + """ + + parameter_model = types._DeleteRuntimeRequestParameters( + name=name, + force=force, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _DeleteRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("delete", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteRuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _get( + self, *, name: str, config: Optional[types.GetRuntimeConfigOrDict] = None + ) -> types.ReasoningEngine: + """ + Get an Agent Runtime instance. + """ + + parameter_model = types._GetRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ReasoningEngine_from_vertex(response_dict) + + return_value = types.ReasoningEngine._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _list( + self, *, config: Optional[types.ListRuntimeConfigOrDict] = None + ) -> types.ListReasoningEnginesResponse: + """ + Lists Agent Runtimes. + """ + + parameter_model = types._ListRuntimeRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _ListRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) + + return_value = types.ListReasoningEnginesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _get_agent_operation( + self, + *, + operation_name: str, + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeOperation: + parameter_model = types._GetRuntimeOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRuntimeOperationParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _query( + self, *, name: str, config: Optional[types.QueryRuntimeConfigOrDict] = None + ) -> types.QueryReasoningEngineResponse: + """ + Query an Agent Runtime. + """ + + parameter_model = types._QueryRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _QueryRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:query".format_map(request_url_dict) + else: + path = "{name}:query" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.QueryReasoningEngineResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _update( + self, *, name: str, config: Optional[types.UpdateRuntimeConfigOrDict] = None + ) -> types.RuntimeOperation: + """ + Updates an Agent Runtime. + """ + + parameter_model = types._UpdateRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpdateRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("patch", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + _memories = None _revisions = None @property @@ -49,17 +1515,2268 @@ def revisions(self) -> "runtime_revisions_module.RuntimeRevisions": ) except ImportError as e: raise ImportError( - "The 'agent_engines.runtimes.revisions' module requires " - "additional packages. Please install them using pip install " + "The 'runtimes.revisions' module requires additional " + "packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e return self._revisions.RuntimeRevisions(self._api_client) # type: ignore[no-any-return] + @property + def memories(self) -> "memories_module.Memories": + if self._memories is None: + try: + # We need to lazy load the memories module to handle the + # possibility of ImportError when dependencies are not installed. + self._memories = importlib.import_module(".memories", __package__) + except ImportError as e: + raise ImportError( + "The 'runtimes.memories' module requires additional " + "packages. Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._memories.Memories(self._api_client) # type: ignore[no-any-return] + + def _list_pager( + self, *, config: Optional[types.ListRuntimeConfigOrDict] = None + ) -> Pager[types.ReasoningEngine]: + return Pager( + "reasoning_engines", + self._list, + self._list(config=config), + config, + ) + + def check_query_job( + self, + *, + name: str, + config: Optional[types.CheckQueryJobRuntimeConfigOrDict] = None, + ) -> types.CheckQueryJobResult: + """Checks a query job on an agent runtime and optionally returns the results. + + Args: + name (str): + Required. A fully-qualified resource name or ID. + config (CheckQueryJobRuntimeConfigOrDict): + Optional. The configuration for the check_query_job. If not provided, + the default configuration will be used. This can be used to specify + the following fields: + - retrieve_result: Whether to retrieve the results of the query job. + """ + from google.cloud import storage # type: ignore[attr-defined] + import json + + if config is None: + config = types.CheckQueryJobRuntimeConfig() + elif isinstance(config, dict): + config = types.CheckQueryJobRuntimeConfig(**config) + + raw_response = self._api_client.request("get", name, {}) + if hasattr(raw_response, "body"): + operation = ( + json.loads(raw_response.body) + if isinstance(raw_response.body, str) + else raw_response.body + ) + else: + operation = raw_response + + status = "RUNNING" + if isinstance(operation, dict): + if operation.get("done"): + status = "FAILED" if operation.get("error") else "SUCCESS" + + response_dict = operation.get("response", {}) + output_gcs_uri = response_dict.get("outputGcsUri") or response_dict.get( + "output_gcs_uri" + ) + error = operation.get("error") + else: + if getattr(operation, "done", False): + status = "FAILED" if getattr(operation, "error", None) else "SUCCESS" + + response_obj = getattr(operation, "response", None) + if isinstance(response_obj, dict): + output_gcs_uri = response_obj.get("outputGcsUri") or response_obj.get( + "output_gcs_uri" + ) + else: + output_gcs_uri = ( + getattr( + response_obj, + "output_gcs_uri", + getattr(response_obj, "outputGcsUri", None), + ) + if response_obj + else None + ) + error = getattr(operation, "error", None) + + result_str = None + if status == "SUCCESS" and config.retrieve_result and output_gcs_uri: + storage_client = storage.Client( + project=self._api_client.project, + credentials=self._api_client._credentials, + ) + bucket_name = output_gcs_uri.replace("gs://", "").split("/")[0] + blob_name = output_gcs_uri.replace(f"gs://{bucket_name}/", "") + bucket = storage_client.bucket(bucket_name) + blob = bucket.blob(blob_name) + if blob.exists(): + result_str = blob.download_as_string().decode("utf-8") + else: + raise ValueError( + f"Failed to retrieve blob results for {output_gcs_uri}" + ) + + elif status == "FAILED" and error: + result_str = str(error) + + return types.CheckQueryJobResult( + operation_name=name, + output_gcs_uri=output_gcs_uri, + status=status, + result=result_str, + ) + + def _is_lightweight_creation( + self, agent: Any, config: types.AgentRuntimeConfig + ) -> bool: + if ( + agent + or config.source_packages + or config.developer_connect_source + or config.agent_config_source + or config.container_spec + ): + return False + return True + + def run_query_job( + self, + *, + name: str, + config: Optional[types.RunQueryJobRuntimeConfigOrDict] = None, + ) -> types.RunQueryJobResult: + """Launches a long-running query job on an Agent Runtime + + Args: + name (str): + Required. A fully-qualified resource name or ID. + config (RunQueryJobRuntimeConfigOrDict): + Optional. The configuration for the async query. If not provided, + the default configuration will be used. This can be used to specify + the following fields: + - query: The query to send to the agent runtime. + - output_gcs_uri: The GCS URI to use for the output. + """ + from google.cloud import storage # type: ignore[attr-defined] + from google.api_core import exceptions + import uuid + + if config is None: + config = types.RunQueryJobRuntimeConfig() + elif isinstance(config, dict): + config = types.RunQueryJobRuntimeConfig(**config) + + if not config.query: + raise ValueError("`query` is required in the config object.") + if not config.output_gcs_uri: + raise ValueError("`output_gcs_uri` is required in the config object.") + + output_gcs_uri = config.output_gcs_uri + is_file = False + last_part = "" + if not output_gcs_uri.endswith("/"): + last_part = output_gcs_uri.split("/")[-1] + if "." in last_part: + is_file = True + + if is_file: + path_parts = output_gcs_uri.split("/") + file_name = path_parts[-1] + base_uri = "/".join(path_parts[:-1]) + name_parts = file_name.rsplit(".", 1) + if len(name_parts) == 2: + name_part, ext = name_parts[0], "." + name_parts[1] + else: + name_part = name_parts[0] + ext = "" + input_gcs_uri = f"{base_uri}/{name_part}_input{ext}" + else: + job_uuid = uuid.uuid4().hex + gcs_path = output_gcs_uri.rstrip("/") + input_gcs_uri = f"{gcs_path}/{job_uuid}_input.json" + output_gcs_uri = f"{gcs_path}/{job_uuid}_output.json" + + storage_client = storage.Client( + project=self._api_client.project, credentials=self._api_client._credentials + ) + + # Handle creating the bucket if it does not exist + bucket_name = config.output_gcs_uri.replace("gs://", "").split("/")[0] + bucket = storage_client.bucket(bucket_name) + + try: + bucket_exists = bucket.exists() + except exceptions.Forbidden as e: + raise ValueError( + f"Permission denied to check existence of bucket '{bucket_name}'. " + "The service account may lack 'storage.buckets.get' permission." + ) from e + + if not bucket_exists: + try: + bucket.create() + except exceptions.Forbidden as e: + raise ValueError( + f"Permission denied to create bucket '{bucket_name}'. " + "The service account may lack 'storage.buckets.create' permission." + ) from e + + input_blob_name = input_gcs_uri.replace(f"gs://{bucket_name}/", "") + blob = bucket.blob(input_blob_name) + blob.upload_from_string(config.query) + + new_config = types._RunQueryJobRuntimeConfig( + input_gcs_uri=input_gcs_uri, + output_gcs_uri=output_gcs_uri, + ) + + # Proceed with sending the async query via the auto-generated method + operation = self._run_query_job(name=name, config=new_config) + + return types.RunQueryJobResult( + job_name=operation.name, + input_gcs_uri=input_gcs_uri, + output_gcs_uri=output_gcs_uri, + ) + + def get( + self, + *, + name: str, + config: Optional[types.GetRuntimeConfigOrDict] = None, + ) -> types.Runtime: + """Gets an agent runtime. + + Args: + name (str): + Required. A fully-qualified resource name or ID such as + "projects/123/locations/us-central1/reasoningEngines/456" or + a shortened name such as "reasoningEngines/456". + """ + api_resource = self._get(name=name, config=config) + runtime = types.Runtime( + api_client=self, + api_async_client=AsyncRuntimes(api_client_=self._api_client), + api_resource=api_resource, + ) + if api_resource.spec: + self._register_api_methods(runtime=runtime) + return runtime + + def delete( + self, + *, + name: str, + force: Optional[bool] = None, + config: Optional[types.DeleteRuntimeConfigOrDict] = None, + ) -> types.DeleteRuntimeOperation: + """ + Delete an Agent Runtime resource. + + Args: + name (str): + Required. The name of the Agent Runtime to be deleted. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` + or `reasoningEngines/{resource_id}`. + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Agent Runtime has undeleted child resources. Defaults to False. + config (DeleteRuntimeConfig): + Optional. Additional configurations for deleting the Agent Runtime. + + """ + logger.info(f"Deleting Runtime resource: {name}") + operation = self._delete(name=name, force=force, config=config) + logger.info(f"Started Runtime delete operation: {operation.name}") + return operation + + def create( + self, + *, + runtime: Any = None, + agent: Any = None, + config: Optional[types.AgentRuntimeConfigOrDict] = None, + ) -> types.Runtime: + """Creates an agent runtime. + + The Agent Runtime will be an instance of the `runtime` that + was passed in, running remotely on Vertex AI. + + Sample ``src_dir`` contents (e.g. ``./user_src_dir``): + + .. code-block:: python + + user_src_dir/ + |-- main.py + |-- requirements.txt + |-- user_code/ + | |-- utils.py + | |-- ... + |-- ... + + To build an Agent Runtime with the above files, run: + + .. code-block:: python + + client = agentplatform.Client( + project="your-project", + location="us-central1", + ) + remote_agent = client.runtimes.create( + agent=local_agent, + config=dict( + requirements=[ + # I.e. the PyPI dependencies listed in requirements.txt + "google-cloud-aiplatform[agent_engines,adk]", + ... + ], + extra_packages=[ + "./user_src_dir/main.py", # a single file + "./user_src_dir/user_code", # a directory + ... + ], + ), + ) + + Args: + agent (Any): + Optional. The Agent to be created. If not specified, this will + correspond to a lightweight instance that cannot be queried + (but can be updated to future instances that can be queried). + runtime (Any): + Optional. This is deprecated. Please use `agent` instead. + config (AgentRuntimeConfig): + Optional. The configurations to use for creating the Agent Runtime. + + Returns: + Runtime: The created Agent Runtime instance. + + Raises: + ValueError: If the `project` was not set using `client.Client`. + ValueError: If the `location` was not set using `client.Client`. + ValueError: If `config.staging_bucket` was not set when `agent` + is specified. + ValueError: If `config.staging_bucket` does not start with "gs://". + ValueError: If `config.extra_packages` is specified but `agent` + is None. + ValueError: If `config.requirements` is specified but `agent` is None. + ValueError: If `config.env_vars` has a dictionary entry that does not + correspond to an environment variable value or a SecretRef. + TypeError: If `config.env_vars` is not a dictionary. + FileNotFoundError: If `config.extra_packages` includes a file or + directory that does not exist. + IOError: If ``config.requirements` is a string that corresponds to a + nonexistent file. + """ + if config is None: + config = {} + if isinstance(config, dict): + config = types.AgentRuntimeConfig.model_validate(config) + elif not isinstance(config, types.AgentRuntimeConfig): + raise TypeError( + f"config must be a dict or AgentRuntimeConfig, but got {type(config)}." + ) + context_spec = config.context_spec + if context_spec is not None: + # Conversion to a dict for _create_config + context_spec = json.loads(context_spec.model_dump_json()) + developer_connect_source = config.developer_connect_source + if developer_connect_source is not None: + developer_connect_source = json.loads( + developer_connect_source.model_dump_json() + ) + agent_config_source = config.agent_config_source + if agent_config_source is not None: + agent_config_source = json.loads(agent_config_source.model_dump_json()) + keep_alive_probe = config.keep_alive_probe + if keep_alive_probe is not None: + keep_alive_probe = json.loads( + keep_alive_probe.model_dump_json(exclude_none=True) + ) + if agent and runtime: + raise ValueError("Please specify only one of `agent` or `runtime`.") + elif runtime: + raise DeprecationWarning( + "The `runtime` argument is deprecated. Please use `agent` instead." + ) + agent = agent or runtime + api_config = self._create_config( + mode="create", + agent=agent, + identity_type=config.identity_type, + staging_bucket=config.staging_bucket, + requirements=config.requirements, + display_name=config.display_name, + description=config.description, + gcs_dir_name=config.gcs_dir_name, + extra_packages=config.extra_packages, + env_vars=config.env_vars, + service_account=config.service_account, + context_spec=context_spec, + psc_interface_config=config.psc_interface_config, + agent_gateway_config=config.agent_gateway_config, + min_instances=config.min_instances, + max_instances=config.max_instances, + resource_limits=config.resource_limits, + container_concurrency=config.container_concurrency, + encryption_spec=config.encryption_spec, + agent_server_mode=config.agent_server_mode, + labels=config.labels, + class_methods=config.class_methods, + source_packages=config.source_packages, + developer_connect_source=developer_connect_source, + entrypoint_module=config.entrypoint_module, + entrypoint_object=config.entrypoint_object, + requirements_file=config.requirements_file, + agent_framework=config.agent_framework, + python_version=config.python_version, + build_options=config.build_options, + image_spec=config.image_spec, + agent_config_source=agent_config_source, + container_spec=config.container_spec, + keep_alive_probe=keep_alive_probe, + ) + operation = self._create(config=api_config) + reasoning_engine_id = _runtimes_utils._get_reasoning_engine_id( + operation_name=operation.name + ) + logger.info( + "View progress and logs at https://console.cloud.google.com/logs/query?" + f"project={self._api_client.project}" + "&query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22%0A" + f"resource.labels.reasoning_engine_id%3D%22{reasoning_engine_id}%22." + ) + if not self._is_lightweight_creation(agent, config): + poll_interval_seconds = 10 + else: + poll_interval_seconds = 1 # Lightweight agent runtime resource creation. + operation = _runtimes_utils._await_operation( + operation_name=operation.name, + get_operation_fn=self._get_agent_operation, + poll_interval_seconds=poll_interval_seconds, + ) + + runtime = types.Runtime( + api_client=self, + api_async_client=AsyncRuntimes(api_client_=self._api_client), + api_resource=operation.response, + ) + if runtime.api_resource: + logger.info("Agent Runtime created. To use it in another session:") + logger.info( + f"runtime=client.runtimes.get(name='{runtime.api_resource.name}')" + ) + elif operation.error: + raise RuntimeError(f"Failed to create Agent Runtime: {operation.error}") + else: + logger.warning("The operation returned an empty response.") + if not self._is_lightweight_creation(agent, config): + # If the user did not provide an runtime (e.g. lightweight + # provisioning), it will not have any API methods registered. + runtime = self._register_api_methods(runtime=runtime) + return runtime # type: ignore[no-any-return] + + def _set_source_code_spec( + self, + *, + spec: types.ReasoningEngineSpecDict, + update_masks: builtins.list[str], + source_packages: Optional[Sequence[str]] = None, + developer_connect_source: Optional[ + types.ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict + ] = None, + class_methods: Optional[Sequence[dict[str, Any]]] = None, + entrypoint_module: Optional[str] = None, + entrypoint_object: Optional[str] = None, + requirements_file: Optional[str] = None, + sys_version: str, + build_options: Optional[dict[str, builtins.list[str]]] = None, + image_spec: Optional[ + types.ReasoningEngineSpecSourceCodeSpecImageSpecDict + ] = None, + agent_config_source: Optional[ + types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict + ] = None, + ) -> None: + """Sets source_code_spec for agent runtime inside the `spec`.""" + source_code_spec = types.ReasoningEngineSpecSourceCodeSpecDict() + if source_packages and not agent_config_source: + source_packages = _runtimes_utils._validate_packages_or_raise( + packages=source_packages, + build_options=build_options, + ) + update_masks.append("spec.source_code_spec.inline_source.source_archive") + source_code_spec["inline_source"] = { # type: ignore[typeddict-item] + "source_archive": _runtimes_utils._create_base64_encoded_tarball( + source_packages=source_packages + ) + } + elif developer_connect_source: + update_masks.append("spec.source_code_spec.developer_connect_source") + source_code_spec["developer_connect_source"] = { + "config": developer_connect_source + } + elif not agent_config_source: + raise ValueError( + "Please specify one of `source_packages`, `developer_connect_source`, " + "or `agent_config_source`." + ) + if class_methods is not None: + update_masks.append("spec.class_methods") + class_methods_spec_list = ( + _runtimes_utils._class_methods_to_class_methods_spec( + class_methods=class_methods + ) + ) + spec["class_methods"] = [ + _runtimes_utils._to_dict(class_method_spec) + for class_method_spec in class_methods_spec_list + ] + elif image_spec is None: + raise ValueError( + "`class_methods` must be specified if `source_packages`, " + "`developer_connect_source`, or `agent_config_source` is " + "specified without a Dockerfile or `image_spec`." + ) + if image_spec is not None: + if entrypoint_module or entrypoint_object or requirements_file: + raise ValueError( + "`image_spec` cannot be specified alongside `entrypoint_module`, " + "`entrypoint_object`, or `requirements_file`, as they are " + "mutually exclusive." + ) + if agent_config_source: + raise ValueError( + "`image_spec` cannot be specified alongside `agent_config_source`, " + "as they are mutually exclusive." + ) + update_masks.append("spec.source_code_spec.image_spec") + source_code_spec["image_spec"] = image_spec + spec["source_code_spec"] = source_code_spec + return + + update_masks.append("spec.source_code_spec.python_spec.version") + python_spec: types.ReasoningEngineSpecSourceCodeSpecPythonSpecDict = { + "version": sys_version, + } + if agent_config_source is not None: + if entrypoint_module or entrypoint_object: + logger.warning( + "`entrypoint_module` and `entrypoint_object` are ignored when " + "`agent_config_source` is specified, as they are pre-defined." + ) + if source_packages: + source_packages = _runtimes_utils._validate_packages_or_raise( + packages=source_packages, + build_options=build_options, + ) + update_masks.append( + "spec.source_code_spec.agent_config_source.inline_source.source_archive" + ) + agent_config_source["inline_source"] = { # type: ignore[typeddict-item] + "source_archive": _runtimes_utils._create_base64_encoded_tarball( + source_packages=source_packages + ) + } + update_masks.append("spec.source_code_spec.agent_config_source") + source_code_spec["agent_config_source"] = agent_config_source + + if requirements_file is not None: + update_masks.append( + "spec.source_code_spec.python_spec.requirements_file" + ) + python_spec["requirements_file"] = requirements_file + source_code_spec["python_spec"] = python_spec + + spec["source_code_spec"] = source_code_spec + return + + if not entrypoint_module: + raise ValueError( + "`entrypoint_module` must be specified if `source_packages` or `developer_connect_source` is specified." + ) + update_masks.append("spec.source_code_spec.python_spec.entrypoint_module") + python_spec["entrypoint_module"] = entrypoint_module + if not entrypoint_object: + raise ValueError( + "`entrypoint_object` must be specified if `source_packages` or `developer_connect_source` is specified." + ) + update_masks.append("spec.source_code_spec.python_spec.entrypoint_object") + python_spec["entrypoint_object"] = entrypoint_object + if requirements_file is not None: + update_masks.append("spec.source_code_spec.python_spec.requirements_file") + python_spec["requirements_file"] = requirements_file + source_code_spec["python_spec"] = python_spec + spec["source_code_spec"] = source_code_spec + + def _set_package_spec( + self, + *, + spec: types.ReasoningEngineSpecDict, + update_masks: builtins.list[str], + agent: Any, + staging_bucket: Optional[str] = None, + requirements: Optional[Union[str, Sequence[str]]] = None, + gcs_dir_name: Optional[str] = None, + extra_packages: Optional[Sequence[str]] = None, + class_methods: Optional[Sequence[dict[str, Any]]] = None, + sys_version: str, + build_options: Optional[dict[str, builtins.list[str]]] = None, + ) -> None: + """Sets package spec for agent runtime.""" + project = self._api_client.project + if project is None: + raise ValueError("project must be set using `agentplatform.Client`.") + location = self._api_client.location + if location is None: + raise ValueError("location must be set using `agentplatform.Client`.") + gcs_dir_name = gcs_dir_name or _runtimes_utils._DEFAULT_GCS_DIR_NAME + staging_bucket = _runtimes_utils._validate_staging_bucket_or_raise( + staging_bucket=staging_bucket, + ) + requirements = _runtimes_utils._validate_requirements_or_raise( + agent=agent, + requirements=requirements, + ) + extra_packages = _runtimes_utils._validate_packages_or_raise( + packages=extra_packages, + build_options=build_options, + ) + # Prepares the Agent Runtime for creation/update in Vertex AI. This + # involves packaging and uploading the artifacts for runtime, + # requirements and extra_packages to `staging_bucket/gcs_dir_name`. + _runtimes_utils._prepare( + agent=agent, + requirements=requirements, + project=project, + location=location, + staging_bucket=staging_bucket, + gcs_dir_name=gcs_dir_name, + extra_packages=extra_packages, + credentials=self._api_client._credentials, + ) + # Update the package spec. + update_masks.append("spec.package_spec.pickle_object_gcs_uri") + package_spec: types.ReasoningEngineSpecPackageSpecDict = { + "python_version": sys_version, + "pickle_object_gcs_uri": "{}/{}/{}".format( + staging_bucket, + gcs_dir_name, + _runtimes_utils._BLOB_FILENAME, + ), + } + if extra_packages: + update_masks.append("spec.package_spec.dependency_files_gcs_uri") + package_spec["dependency_files_gcs_uri"] = "{}/{}/{}".format( + staging_bucket, + gcs_dir_name, + _runtimes_utils._EXTRA_PACKAGES_FILE, + ) + if requirements: + update_masks.append("spec.package_spec.requirements_gcs_uri") + package_spec["requirements_gcs_uri"] = "{}/{}/{}".format( + staging_bucket, + gcs_dir_name, + _runtimes_utils._REQUIREMENTS_FILE, + ) + spec["package_spec"] = package_spec + + update_masks.append("spec.class_methods") + if class_methods is not None: + class_methods_spec_list = ( + _runtimes_utils._class_methods_to_class_methods_spec( + class_methods=class_methods + ) + ) + else: + class_methods_spec_list = ( + _runtimes_utils._generate_class_methods_spec_or_raise( + agent=agent, + operations=_runtimes_utils._get_registered_operations(agent=agent), + ) + ) + spec["class_methods"] = [ + _runtimes_utils._to_dict(class_method_spec) + for class_method_spec in class_methods_spec_list + ] + + def _create_config( + self, + *, + mode: str, + agent: Any = None, + identity_type: Optional[types.IdentityType] = None, + staging_bucket: Optional[str] = None, + requirements: Optional[Union[str, Sequence[str]]] = None, + display_name: Optional[str] = None, + description: Optional[str] = None, + gcs_dir_name: Optional[str] = None, + extra_packages: Optional[Sequence[str]] = None, + env_vars: Optional[dict[str, Union[str, Any]]] = None, + service_account: Optional[str] = None, + context_spec: Optional[types.ReasoningEngineContextSpecDict] = None, + psc_interface_config: Optional[types.PscInterfaceConfigDict] = None, + agent_gateway_config: Optional[ + types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] = None, + min_instances: Optional[int] = None, + max_instances: Optional[int] = None, + resource_limits: Optional[dict[str, str]] = None, + container_concurrency: Optional[int] = None, + encryption_spec: Optional[genai_types.EncryptionSpecDict] = None, + labels: Optional[dict[str, str]] = None, + agent_server_mode: Optional[types.AgentServerMode] = None, + class_methods: Optional[Sequence[dict[str, Any]]] = None, + source_packages: Optional[Sequence[str]] = None, + developer_connect_source: Optional[ + types.ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict + ] = None, + entrypoint_module: Optional[str] = None, + entrypoint_object: Optional[str] = None, + requirements_file: Optional[str] = None, + agent_framework: Optional[str] = None, + python_version: Optional[str] = None, + build_options: Optional[dict[str, builtins.list[str]]] = None, + image_spec: Optional[ + types.ReasoningEngineSpecSourceCodeSpecImageSpecDict + ] = None, + agent_config_source: Optional[ + types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict + ] = None, + container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None, + keep_alive_probe: Optional[dict[str, Any]] = None, + traffic_config: Optional[types.ReasoningEngineTrafficConfigDict] = None, + ) -> types.UpdateRuntimeConfigDict: + import sys + + config: types.UpdateRuntimeConfigDict = {} + update_masks = [] + if mode not in ["create", "update"]: + raise ValueError(f"Unsupported mode: {mode}") + if agent is None: + if requirements is not None: + raise ValueError("requirements must be None if agent is None.") + if extra_packages is not None: + raise ValueError("extra_packages must be None if agent is None.") + if display_name is not None: + update_masks.append("display_name") + config["display_name"] = display_name + if description is not None: + update_masks.append("description") + config["description"] = description + if context_spec is not None: + update_masks.append("context_spec") + config["context_spec"] = context_spec + if encryption_spec is not None: + update_masks.append("encryption_spec") + config["encryption_spec"] = encryption_spec + if labels is not None: + update_masks.append("labels") + config["labels"] = labels + if traffic_config is not None: + update_masks.append("traffic_config") + config["traffic_config"] = traffic_config + + if agent_framework == "google-adk": + env_vars = _runtimes_utils._add_telemetry_enablement_env(env_vars) + + if python_version: + sys_version = python_version + else: + sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" + + if agent: + if source_packages: + raise ValueError( + "If you have provided `source_packages` in `config`, please " + "do not specify `agent` in `runtimes.create()` or " + "`runtimes.update()`." + ) + if developer_connect_source: + raise ValueError( + "If you have provided `developer_connect_source` in `config`, please " + "do not specify `agent` in `runtimes.create()` or " + "`runtimes.update()`." + ) + elif source_packages and developer_connect_source: + raise ValueError( + "Please specify only one of `source_packages` or `developer_connect_source` in `config`." + ) + + if container_spec: + if agent: + raise ValueError( + "If you have provided `container_spec` in `config`, please " + "do not specify `agent` in `runtimes.create()` or " + "`runtimes.update()`." + ) + if source_packages or developer_connect_source: + raise ValueError( + "If you have provided `container_spec` in `config`, please " + "do not specify `source_packages` or `developer_connect_source` in `config`." + ) + + runtime_spec: Any = None + if agent: + runtime_spec = {} + agent = _runtimes_utils._validate_agent_or_raise(agent=agent) + if _runtimes_utils._is_adk_agent(agent): + env_vars = _runtimes_utils._add_telemetry_enablement_env(env_vars) + self._set_package_spec( + spec=runtime_spec, + update_masks=update_masks, + agent=agent, + staging_bucket=staging_bucket, + requirements=requirements, + gcs_dir_name=gcs_dir_name, + extra_packages=extra_packages, + class_methods=class_methods, + sys_version=sys_version, + build_options=build_options, + ) + elif ( + source_packages + or developer_connect_source + or image_spec + or agent_config_source + ): + runtime_spec = {} + self._set_source_code_spec( + spec=runtime_spec, + update_masks=update_masks, + source_packages=source_packages, + developer_connect_source=developer_connect_source, + class_methods=class_methods, + entrypoint_module=entrypoint_module, + entrypoint_object=entrypoint_object, + requirements_file=requirements_file, + sys_version=sys_version, + build_options=build_options, + image_spec=image_spec, + agent_config_source=agent_config_source, + ) + elif container_spec: + runtime_spec = {} + if class_methods is not None: + update_masks.append("spec.class_methods") + class_methods_spec_list = ( + _runtimes_utils._class_methods_to_class_methods_spec( + class_methods=class_methods + ) + ) + runtime_spec["class_methods"] = [ + _runtimes_utils._to_dict(class_method_spec) + for class_method_spec in class_methods_spec_list + ] + update_masks.append("spec.container_spec") + runtime_spec["container_spec"] = container_spec + + is_deployment_spec_updated = ( + env_vars is not None + or psc_interface_config is not None + or agent_gateway_config is not None + or min_instances is not None + or max_instances is not None + or resource_limits is not None + or container_concurrency is not None + or keep_alive_probe is not None + ) + if runtime_spec is None and is_deployment_spec_updated: + raise ValueError( + "To update `env_vars`, `psc_interface_config`, `min_instances`, " + "`max_instances`, `resource_limits`, `container_concurrency`, or " + "`keep_alive_probe`, you must also provide the `agent` variable or " + "the source code options (`source_packages`, " + "`developer_connect_source` or `agent_config_source`)." + ) + + if runtime_spec is not None: + if is_deployment_spec_updated: + ( + deployment_spec, + deployment_update_masks, + ) = self._generate_deployment_spec_or_raise( + env_vars=env_vars, + psc_interface_config=psc_interface_config, + agent_gateway_config=agent_gateway_config, + min_instances=min_instances, + max_instances=max_instances, + resource_limits=resource_limits, + container_concurrency=container_concurrency, + keep_alive_probe=keep_alive_probe, + ) + update_masks.extend(deployment_update_masks) + runtime_spec["deployment_spec"] = deployment_spec + + if agent_server_mode: + if not runtime_spec.get("deployment_spec"): + runtime_spec["deployment_spec"] = ( + types.ReasoningEngineSpecDeploymentSpecDict() + ) + runtime_spec["deployment_spec"]["agent_server_mode"] = agent_server_mode + + runtime_spec["agent_framework"] = _runtimes_utils._get_agent_framework( + agent_framework=agent_framework, + agent=agent, + ) + + if hasattr(agent, "agent_card"): + agent_card = getattr(agent, "agent_card") + if agent_card: + try: + from google.protobuf import json_format + + runtime_spec["agent_card"] = json_format.MessageToDict( + agent_card + ) + except Exception as e: + raise ValueError( + f"Failed to convert agent card to dict (serialization error): {e}" + ) from e + update_masks.append("spec.agent_framework") + + if identity_type is not None or service_account is not None: + if runtime_spec is None: + runtime_spec = {} + + if identity_type is not None: + runtime_spec["identity_type"] = identity_type + update_masks.append("spec.identity_type") + if service_account is not None: + # Clear the field in case of empty service_account. + if service_account: + runtime_spec["service_account"] = service_account + update_masks.append("spec.service_account") + + if runtime_spec is not None: + config["spec"] = runtime_spec + + if update_masks and mode == "update": + config["update_mask"] = ",".join(update_masks) + return config + + def _generate_deployment_spec_or_raise( + self, + *, + env_vars: Optional[dict[str, Union[str, Any]]] = None, + psc_interface_config: Optional[types.PscInterfaceConfigDict] = None, + agent_gateway_config: Optional[ + types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] = None, + min_instances: Optional[int] = None, + max_instances: Optional[int] = None, + resource_limits: Optional[dict[str, str]] = None, + container_concurrency: Optional[int] = None, + keep_alive_probe: Optional[dict[str, Any]] = None, + ) -> Tuple[dict[str, Any], Sequence[str]]: + deployment_spec: dict[str, Any] = {} + update_masks = [] + if env_vars: + deployment_spec["env"] = [] + deployment_spec["secret_env"] = [] + if isinstance(env_vars, dict): + self._update_deployment_spec_with_env_vars_dict_or_raise( + deployment_spec=deployment_spec, + env_vars=env_vars, + ) + else: + raise TypeError(f"env_vars must be a dict, but got {type(env_vars)}.") + if deployment_spec.get("env"): + update_masks.append("spec.deployment_spec.env") + if deployment_spec.get("secret_env"): + update_masks.append("spec.deployment_spec.secret_env") + if psc_interface_config: + deployment_spec["psc_interface_config"] = psc_interface_config + update_masks.append("spec.deployment_spec.psc_interface_config") + if agent_gateway_config: + deployment_spec["agent_gateway_config"] = agent_gateway_config + update_masks.append("spec.deployment_spec.agent_gateway_config") + if min_instances is not None: + if not 0 <= min_instances <= 10: + raise ValueError( + f"min_instances must be between 0 and 10. Got {min_instances}" + ) + deployment_spec["min_instances"] = min_instances + update_masks.append("spec.deployment_spec.min_instances") + if max_instances is not None: + if psc_interface_config and not 1 <= max_instances <= 100: + raise ValueError( + f"max_instances must be between 1 and 100 when PSC-I is enabled. Got {max_instances}" + ) + elif not psc_interface_config and not 1 <= max_instances <= 1000: + raise ValueError( + f"max_instances must be between 1 and 1000. Got {max_instances}" + ) + deployment_spec["max_instances"] = max_instances + update_masks.append("spec.deployment_spec.max_instances") + if resource_limits: + _runtimes_utils._validate_resource_limits_or_raise( + resource_limits=resource_limits + ) + deployment_spec["resource_limits"] = resource_limits + update_masks.append("spec.deployment_spec.resource_limits") + if container_concurrency: + deployment_spec["container_concurrency"] = container_concurrency + update_masks.append("spec.deployment_spec.container_concurrency") + if keep_alive_probe is not None: + deployment_spec["keep_alive_probe"] = keep_alive_probe + update_masks.append("spec.deployment_spec.keep_alive_probe") + return deployment_spec, update_masks + + def _update_deployment_spec_with_env_vars_dict_or_raise( + self, + *, + deployment_spec: dict[str, Any], + env_vars: dict[str, Any], + ) -> None: + for key, value in env_vars.items(): + if isinstance(value, dict): + if "secret_env" not in deployment_spec: + deployment_spec["secret_env"] = [] + deployment_spec["secret_env"].append({"name": key, "secret_ref": value}) + elif isinstance(value, str): + if "env" not in deployment_spec: + deployment_spec["env"] = [] + deployment_spec["env"].append({"name": key, "value": value}) + else: + raise TypeError( + f"Unknown value type in env_vars for {key}. " + f"Must be a str or SecretRef: {value}" + ) + + def _register_api_methods( + self, + *, + runtime: types.Runtime, + ) -> types.Runtime: + """Registers the API methods for the agent runtime.""" + try: + _runtimes_utils._register_api_methods_or_raise( + runtime=runtime, + wrap_operation_fn={ + "": _runtimes_utils._wrap_query_operation, # type: ignore[dict-item] + "async": _runtimes_utils._wrap_async_query_operation, # type: ignore[dict-item] + "stream": _runtimes_utils._wrap_stream_query_operation, # type: ignore[dict-item] + "async_stream": _runtimes_utils._wrap_async_stream_query_operation, # type: ignore[dict-item] + "a2a_extension": _runtimes_utils._wrap_a2a_operation, + }, + ) + except Exception as e: + logger.warning( + _runtimes_utils._FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e + ) + return runtime + + def list( + self, *, config: Optional[types.ListRuntimeConfigOrDict] = None + ) -> Iterator[types.Runtime]: + """List all instances of Agent Runtime matching the filter. + + Example Usage: + + .. code-block:: python + import agentplatform + + client = agentplatform.Client(project="my_project", location="us-central1") + for agent in client.runtimes.list( + config={"filter": "'display_name="My Custom Agent"'}, + ): + print(agent.api_resource.name) + + Args: + config (ListRuntimeConfig): + Optional. The config (e.g. filter) for the agents to be listed. + + Returns: + Iterable[Runtime]: An iterable of Agent Runtimes matching the filter. + """ + + for reasoning_engine in self._list_pager(config=config): + yield types.Runtime( + api_client=self, + api_async_client=AsyncRuntimes(api_client_=self._api_client), + api_resource=reasoning_engine, + ) + + def update( + self, + *, + name: str, + agent: Any = None, + runtime: Any = None, + config: types.AgentRuntimeConfigOrDict, + ) -> types.Runtime: + """Updates an existing Agent Runtime. + + This method updates the configuration of an existing Agent Runtime running + remotely, which is identified by its name. + + Args: + name (str): Required. A fully-qualified resource name or ID such as + "projects/123/locations/us-central1/reasoningEngines/456" or a + shortened name such as "reasoningEngines/456". + agent (Any): + Optional. The instance to be used as the updated Agent Runtime. + If it is not specified, the existing instance will be used. + runtime (Any): + Optional. This is deprecated. Please use `agent` instead. + config (AgentRuntimeConfig): + Optional. The configurations to use for updating the Agent Runtime. + + Returns: + Runtime: The updated Agent Runtime. + + Raises: + ValueError: If the `project` was not set using `client.Client`. + ValueError: If the `location` was not set using `client.Client`. + ValueError: If `config.staging_bucket` was not set when `runtime` + is specified. + ValueError: If `config.staging_bucket` does not start with "gs://". + ValueError: If `config.extra_packages` is specified but `runtime` + is None. + ValueError: If `config.requirements` is specified but `runtime` is + None. + ValueError: If `config.env_vars` has a dictionary entry that does not + correspond to an environment variable value or a SecretRef. + TypeError: If `config.env_vars` is not a dictionary. + FileNotFoundError: If `config.extra_packages` includes a file or + directory that does not exist. + IOError: If `config.requirements` is a string that corresponds to a + nonexistent file. + """ + if isinstance(config, dict): + config = types.AgentRuntimeConfig.model_validate(config) + elif not isinstance(config, types.AgentRuntimeConfig): + raise TypeError( + f"config must be a dict or AgentRuntimeConfig, but got {type(config)}." + ) + context_spec = config.context_spec + if context_spec is not None: + # Conversion to a dict for _create_config + context_spec = json.loads(context_spec.model_dump_json()) + developer_connect_source = config.developer_connect_source + if developer_connect_source is not None: + developer_connect_source = json.loads( + developer_connect_source.model_dump_json() + ) + agent_config_source = config.agent_config_source + if agent_config_source is not None: + agent_config_source = json.loads(agent_config_source.model_dump_json()) + keep_alive_probe = config.keep_alive_probe + if keep_alive_probe is not None: + keep_alive_probe = json.loads( + keep_alive_probe.model_dump_json(exclude_none=True) + ) + traffic_config = config.traffic_config + if traffic_config is not None: + traffic_config = json.loads(traffic_config.model_dump_json()) + if agent and runtime: + raise ValueError("Please specify only one of `agent` or `runtime`.") + elif runtime: + raise DeprecationWarning( + "The `runtime` argument is deprecated. Please use `agent` instead." + ) + image_spec = config.image_spec + if image_spec is not None: + # Conversion to a dict for _create_config + image_spec = json.loads(image_spec.model_dump_json()) + container_spec = config.container_spec + if container_spec is not None: + # Conversion to a dict for _create_config + container_spec = json.loads(container_spec.model_dump_json()) + agent = agent or runtime + api_config = self._create_config( + mode="update", + agent=agent, + identity_type=config.identity_type, + staging_bucket=config.staging_bucket, + requirements=config.requirements, + display_name=config.display_name, + description=config.description, + gcs_dir_name=config.gcs_dir_name, + extra_packages=config.extra_packages, + env_vars=config.env_vars, + service_account=config.service_account, + context_spec=context_spec, + psc_interface_config=config.psc_interface_config, + agent_gateway_config=config.agent_gateway_config, + min_instances=config.min_instances, + max_instances=config.max_instances, + resource_limits=config.resource_limits, + container_concurrency=config.container_concurrency, + labels=config.labels, + class_methods=config.class_methods, + source_packages=config.source_packages, + developer_connect_source=developer_connect_source, + entrypoint_module=config.entrypoint_module, + entrypoint_object=config.entrypoint_object, + requirements_file=config.requirements_file, + agent_framework=config.agent_framework, + python_version=config.python_version, + build_options=config.build_options, + image_spec=image_spec, + agent_config_source=agent_config_source, + container_spec=container_spec, + keep_alive_probe=keep_alive_probe, + traffic_config=traffic_config, + ) + operation = self._update(name=name, config=api_config) + reasoning_engine_id = _runtimes_utils._get_reasoning_engine_id( + resource_name=name + ) + logger.info( + "View progress and logs at https://console.cloud.google.com/logs/query?" + f"project={self._api_client.project}" + "&query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22%0A" + f"resource.labels.reasoning_engine_id%3D%22{reasoning_engine_id}%22." + ) + operation = _runtimes_utils._await_operation( + operation_name=operation.name, + get_operation_fn=self._get_agent_operation, + ) + runtime = types.Runtime( + api_client=self, + api_async_client=AsyncRuntimes(api_client_=self._api_client), + api_resource=operation.response, + ) + if runtime.api_resource: + logger.info("Agent Runtime updated. To use it in another session:") + logger.info( + f"runtime=client.runtimes.get(name='{runtime.api_resource.name}')" + ) + elif operation.error: + raise RuntimeError(f"Failed to update Agent Runtime: {operation.error}") + if runtime.api_resource.spec: + self._register_api_methods(runtime=runtime) + return runtime # type: ignore[no-any-return] + + def _stream_query( + self, *, name: str, config: Optional[types.QueryRuntimeConfigOrDict] = None + ) -> Iterator[Any]: + """Streams the response of the agent runtime.""" + parameter_model = types._QueryRuntimeRequestParameters( + name=name, + config=config, + ) + request_dict = _QueryRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) + else: + path = "{name}:streamQuery?alt=sse" + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + http_options = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + for response in self._api_client.request_streamed( + "post", path, request_dict, http_options + ): + yield response + + # TODO: b/436704146 - Replace with generated methods + # TODO: b/437129724 - Add replay test for async stream query + async def _async_stream_query( + self, + *, + name: str, + config: Optional[types.QueryRuntimeConfigOrDict] = None, + ) -> AsyncIterator[Any]: + """Streams the response of the agent runtime asynchronously.""" + parameter_model = types._QueryRuntimeRequestParameters( + name=name, + config=config, + ) + request_dict = _QueryRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) + else: + path = "{name}:streamQuery?alt=sse" + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + http_options = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + async_iterator = await self._api_client.async_request_streamed( + "post", path, request_dict, http_options + ) + async for response in async_iterator: + yield response + + def create_memory( + self, + *, + name: str, + fact: str, + scope: dict[str, str], + config: Optional[types.RuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: + """Deprecated. Use runtimes.memories.create instead.""" + warnings.warn( + ( + "runtimes.create_memory is deprecated. " + "Use runtimes.memories.create instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return self.memories.create( + name=name, + fact=fact, + scope=scope, + config=config, + ) + + def delete_memory( + self, + *, + name: str, + config: Optional[types.DeleteRuntimeMemoryConfigOrDict] = None, + ) -> types.DeleteRuntimeMemoryOperation: + """Deprecated. Use runtimes.memories.delete instead.""" + warnings.warn( + ( + "runtimes.delete_memory is deprecated. " + "Use runtimes.memories.delete instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return self.memories.delete(name=name, config=config) + + def generate_memories( + self, + *, + name: str, + vertex_session_source: Optional[ + types.GenerateMemoriesRequestVertexSessionSourceOrDict + ] = None, + direct_contents_source: Optional[ + types.GenerateMemoriesRequestDirectContentsSourceOrDict + ] = None, + direct_memories_source: Optional[ + types.GenerateMemoriesRequestDirectMemoriesSourceOrDict + ] = None, + scope: Optional[dict[str, str]] = None, + config: Optional[types.GenerateRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: + """Deprecated. Use runtimes.memories.generate instead.""" + warnings.warn( + ( + "runtimes.generate_memories is deprecated. " + "Use runtimes.memories.generate instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return self.memories.generate( + name=name, + vertex_session_source=vertex_session_source, + direct_contents_source=direct_contents_source, + direct_memories_source=direct_memories_source, + scope=scope, + config=config, + ) + + def get_memory( + self, + *, + name: str, + config: Optional[types.GetRuntimeMemoryConfigOrDict] = None, + ) -> types.Memory: + """Deprecated. Use runtimes.memories.get instead.""" + warnings.warn( + ( + "runtimes.get_memory is deprecated. " + "Use runtimes.memories.get instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return self.memories.get(name=name, config=config) + + def list_memories( + self, + *, + name: str, + config: Optional[types.ListRuntimeMemoryConfigOrDict] = None, + ) -> Iterator[types.Memory]: + """Deprecated. Use runtimes.memories.list instead.""" + warnings.warn( + ( + "runtimes.list_memories is deprecated. " + "Use runtimes.memories.list instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return self.memories.list(name=name, config=config) + + def retrieve_memories( + self, + *, + name: str, + scope: dict[str, str], + similarity_search_params: Optional[ + types.RetrieveMemoriesRequestSimilaritySearchParamsOrDict + ] = None, + simple_retrieval_params: Optional[ + types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict + ] = None, + config: Optional[types.RetrieveRuntimeMemoriesConfigOrDict] = None, + ) -> Iterator[types.RetrieveMemoriesResponseRetrievedMemory]: + """Deprecated. Use runtimes.memories.retrieve instead.""" + warnings.warn( + ( + "runtimes.retrieve_memories is deprecated. " + "Use runtimes.memories.retrieve instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return self.memories.retrieve( + name=name, + scope=scope, + similarity_search_params=similarity_search_params, + simple_retrieval_params=simple_retrieval_params, + config=config, + ) + class AsyncRuntimes(_api_module.BaseModule): + async def cancel_query_job( + self, + *, + name: str, + config: Optional[types.CancelQueryJobRuntimeConfigOrDict] = None, + ) -> types.CancelQueryJobResult: + """ + Cancels a long-running query job on an Agent Runtime. + + Args: + name (str): + Required. The reasoning engine resource name. + config (CancelQueryJobRuntimeConfigOrDict): + Optional. The configuration for the cancel_query_job. + + """ + + parameter_model = types._CancelQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CancelQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:cancelAsyncQuery".format_map(request_url_dict) + else: + path = "{name}:cancelAsyncQuery" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.CancelQueryJobResult._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _check_query_job( + self, + *, + name: str, + config: Optional[types.CheckQueryJobRuntimeConfigOrDict] = None, + ) -> types.CheckQueryJobResult: + """ + Query an Agent Runtime asynchronously. + """ + + parameter_model = types._CheckQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CheckQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:checkQueryJob".format_map(request_url_dict) + else: + path = "{name}:checkQueryJob" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _CheckQueryJobResult_from_vertex(response_dict) + + return_value = types.CheckQueryJobResult._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _run_query_job( + self, + *, + name: str, + config: Optional[types._RunQueryJobRuntimeConfigOrDict] = None, + ) -> types.RuntimeOperation: + """ + Run a query job on an agent runtime. + """ + + parameter_model = types._RunQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _RunQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:asyncQuery".format_map(request_url_dict) + else: + path = "{name}:asyncQuery" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _create( + self, *, config: Optional[types.CreateRuntimeConfigOrDict] = None + ) -> types.RuntimeOperation: + """ + Creates a new Agent Runtime. + """ + + parameter_model = types._CreateRuntimeRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CreateRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _delete( + self, + *, + name: str, + force: Optional[bool] = None, + config: Optional[types.DeleteRuntimeConfigOrDict] = None, + ) -> types.DeleteRuntimeOperation: + """ + Delete an Agent Runtime resource. + + Args: + name (str): + Required. The name of the Agent Runtime to be deleted. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` + or `reasoningEngines/{resource_id}`. + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Agent Runtime has undeleted child resources. Defaults to False. + config (DeleteRuntimeConfig): + Optional. Additional configurations for deleting the Agent Runtime. + + """ + + parameter_model = types._DeleteRuntimeRequestParameters( + name=name, + force=force, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _DeleteRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "delete", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteRuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _get( + self, *, name: str, config: Optional[types.GetRuntimeConfigOrDict] = None + ) -> types.ReasoningEngine: + """ + Get an Agent Runtime instance. + """ + + parameter_model = types._GetRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ReasoningEngine_from_vertex(response_dict) + + return_value = types.ReasoningEngine._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _list( + self, *, config: Optional[types.ListRuntimeConfigOrDict] = None + ) -> types.ListReasoningEnginesResponse: + """ + Lists Agent Runtimes. + """ + + parameter_model = types._ListRuntimeRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _ListRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) + + return_value = types.ListReasoningEnginesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _get_agent_operation( + self, + *, + operation_name: str, + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeOperation: + parameter_model = types._GetRuntimeOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRuntimeOperationParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _query( + self, *, name: str, config: Optional[types.QueryRuntimeConfigOrDict] = None + ) -> types.QueryReasoningEngineResponse: + """ + Query an Agent Runtime. + """ + + parameter_model = types._QueryRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _QueryRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:query".format_map(request_url_dict) + else: + path = "{name}:query" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.QueryReasoningEngineResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _update( + self, *, name: str, config: Optional[types.UpdateRuntimeConfigOrDict] = None + ) -> types.RuntimeOperation: + """ + Updates an Agent Runtime. + """ + + parameter_model = types._UpdateRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpdateRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "patch", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + _memories = None _revisions = None + async def delete( + self, + *, + name: str, + force: Optional[bool] = None, + config: Optional[types.DeleteRuntimeConfigOrDict] = None, + ) -> types.DeleteRuntimeOperation: + """ + Delete an Agent Runtime resource. + + Args: + name (str): + Required. The name of the Agent Runtime to be deleted. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` + or `reasoningEngines/{resource_id}`. + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Agent Runtime has undeleted child resources. Defaults to False. + config (DeleteRuntimeConfig): + Optional. Additional configurations for deleting the Agent Runtime. + + """ + logger.info(f"Deleting Runtime resource: {name}") + operation = await self._delete(name=name, force=force, config=config) + logger.info(f"Started Runtime delete operation: {operation.name}") + return operation + @property def revisions(self) -> "runtime_revisions_module.AsyncRuntimeRevisions": if self._revisions is None: @@ -71,8 +3788,57 @@ def revisions(self) -> "runtime_revisions_module.AsyncRuntimeRevisions": ) except ImportError as e: raise ImportError( - "The 'agent_engines.runtimes.revisions' module requires " - "additional packages. Please install them using pip install " + "The 'runtimes.revisions' module requires additional " + "packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e return self._revisions.AsyncRuntimeRevisions(self._api_client) # type: ignore[no-any-return] + + @property + def memories(self) -> "memories_module.AsyncMemories": + if self._memories is None: + try: + # We need to lazy load the memories module to handle the + # possibility of ImportError when dependencies are not installed. + self._memories = importlib.import_module(".memories", __package__) + except ImportError as e: + raise ImportError( + "The 'runtimes.memories' module requires additional " + "packages. Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._memories.AsyncMemories(self._api_client) # type: ignore[no-any-return] + + async def delete_memory( + self, + *, + name: str, + config: Optional[types.DeleteRuntimeMemoryConfigOrDict] = None, + ) -> types.DeleteRuntimeMemoryOperation: + """Deprecated. Use runtimes.memories.delete instead.""" + warnings.warn( + ( + "runtimes.delete_memory is deprecated. " + "Use runtimes.memories.delete instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return await self.memories.delete(name=name, config=config) + + async def get_memory( + self, + *, + name: str, + config: Optional[types.GetRuntimeMemoryConfigOrDict] = None, + ) -> types.Memory: + """Deprecated. Use runtimes.memories.get instead.""" + warnings.warn( + ( + "runtimes.get_memory is deprecated. " + "Use runtimes.memories.get instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return await self.memories.get(name=name, config=config) diff --git a/agentplatform/_genai/sandbox_snapshots.py b/agentplatform/_genai/sandbox_snapshots.py index aaeb01da1e..6e4731d482 100644 --- a/agentplatform/_genai/sandbox_snapshots.py +++ b/agentplatform/_genai/sandbox_snapshots.py @@ -27,7 +27,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import Pager -from . import _agent_engines_utils +from . import _runtimes_utils from . import types logger = logging.getLogger("agentplatform_genai.sandboxsnapshots") @@ -35,7 +35,7 @@ logger.setLevel(logging.INFO) -def _CreateAgentEngineSandboxSnapshotConfig_to_vertex( +def _CreateRuntimeSandboxSnapshotConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -69,7 +69,7 @@ def _CreateSandboxEnvironmentSnapshotRequestParameters_to_vertex( setv( to_object, ["config"], - _CreateAgentEngineSandboxSnapshotConfig_to_vertex( + _CreateRuntimeSandboxSnapshotConfig_to_vertex( getv(from_object, ["config"]), to_object ), ) @@ -88,7 +88,7 @@ def _DeleteSandboxEnvironmentSnapshotRequestParameters_to_vertex( return to_object -def _GetAgentEngineSandboxSnapshotOperationParameters_to_vertex( +def _GetRuntimeSandboxSnapshotOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -153,8 +153,8 @@ def _create( self, *, source_sandbox_environment_name: str, - config: Optional[types.CreateAgentEngineSandboxSnapshotConfigOrDict] = None, - ) -> types.AgentEngineSandboxSnapshotOperation: + config: Optional[types.CreateRuntimeSandboxSnapshotConfigOrDict] = None, + ) -> types.RuntimeSandboxSnapshotOperation: """ Snapshots an existing sandbox environment. @@ -162,7 +162,7 @@ def _create( source_sandbox_environment_name (str): Required. The name of the sandbox environment to snapshot. projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_environment_id} - config (CreateAgentEngineSandboxSnapshotConfig): + config (CreateRuntimeSandboxSnapshotConfig): Optional. The configuration for the sandbox snapshot. """ @@ -207,7 +207,7 @@ def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxSnapshotOperation._from_response( + return_value = types.RuntimeSandboxSnapshotOperation._from_response( response=response_dict, kwargs=( { @@ -460,9 +460,9 @@ def get_sandbox_snapshot_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineSandboxSnapshotOperation: - parameter_model = types._GetAgentEngineSandboxSnapshotOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeSandboxSnapshotOperation: + parameter_model = types._GetRuntimeSandboxSnapshotOperationParameters( operation_name=operation_name, config=config, ) @@ -473,7 +473,7 @@ def get_sandbox_snapshot_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSandboxSnapshotOperationParameters_to_vertex( + request_dict = _GetRuntimeSandboxSnapshotOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -502,7 +502,7 @@ def get_sandbox_snapshot_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxSnapshotOperation._from_response( + return_value = types.RuntimeSandboxSnapshotOperation._from_response( response=response_dict, kwargs=( { @@ -530,34 +530,34 @@ def create( self, *, source_sandbox_environment_name: str, - config: Optional[types.CreateAgentEngineSandboxSnapshotConfigOrDict] = None, + config: Optional[types.CreateRuntimeSandboxSnapshotConfigOrDict] = None, poll_interval_seconds: float = 0.1, - ) -> types.AgentEngineSandboxSnapshotOperation: + ) -> types.RuntimeSandboxSnapshotOperation: """Snapshots an existing sandbox environment. Args: source_sandbox_environment_name (str): Required. The name of the sandbox environment to snapshot. projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_environment_id} - config (CreateAgentEngineSandboxSnapshotConfig): + config (CreateRuntimeSandboxSnapshotConfig): Optional. The configuration for the sandbox snapshot. poll_interval_seconds (int): Optional. Seconds to wait between polling for operation status. Defaults to 0.1. Returns: - AgentEngineSandboxSnapshotOperation: The operation for creating the sandbox snapshot. + RuntimeSandboxSnapshotOperation: The operation for creating the sandbox snapshot. """ operation = self._create( source_sandbox_environment_name=source_sandbox_environment_name, config=config, ) if config is None: - config = types.CreateAgentEngineSandboxSnapshotConfig() + config = types.CreateRuntimeSandboxSnapshotConfig() elif isinstance(config, dict): - config = types.CreateAgentEngineSandboxSnapshotConfig.model_validate(config) + config = types.CreateRuntimeSandboxSnapshotConfig.model_validate(config) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self.get_sandbox_snapshot_operation, poll_interval_seconds=poll_interval_seconds, @@ -575,17 +575,17 @@ def list( name: str, config: Optional[types.ListSandboxEnvironmentSnapshotsConfigOrDict] = None, ) -> Iterator[types.SandboxEnvironmentSnapshot]: - """Lists Agent Engine sandbox snapshots. + """Lists Agent Runtime sandbox snapshots. Args: name (str): - Required. The name of the agent engine to list sandbox snapshots for. + Required. The name of the agent runtime to list sandbox snapshots for. projects/{project}/locations/{location}/reasoningEngines/{resource_id} config (ListSandboxEnvironmentSnapshotsConfig): Optional. The configuration for the sandbox snapshots to list. Returns: - Iterable[SandboxEnvironmentSnapshot]: An iterable of agent engine sandbox snapshots. + Iterable[SandboxEnvironmentSnapshot]: An iterable of agent runtime sandbox snapshots. """ return Pager( "sandbox_environment_snapshots", @@ -600,7 +600,7 @@ def get( name: str, config: Optional[types.GetSandboxEnvironmentSnapshotConfigOrDict] = None, ) -> types.SandboxEnvironmentSnapshot: - """Gets a sandbox snapshot in the Agent Engine. + """Gets a sandbox snapshot in the Agent Runtime. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -617,7 +617,7 @@ def delete( name: str, config: Optional[types.DeleteSandboxEnvironmentSnapshotConfigOrDict] = None, ) -> types.DeleteSandboxEnvironmentSnapshotOperation: - """Deletes a sandbox snapshot in the Agent Engine. + """Deletes a sandbox snapshot in the Agent Runtime. Args: name (str): Required. The name of the sandbox snapshot to delete. @@ -635,8 +635,8 @@ async def _create( self, *, source_sandbox_environment_name: str, - config: Optional[types.CreateAgentEngineSandboxSnapshotConfigOrDict] = None, - ) -> types.AgentEngineSandboxSnapshotOperation: + config: Optional[types.CreateRuntimeSandboxSnapshotConfigOrDict] = None, + ) -> types.RuntimeSandboxSnapshotOperation: """ Snapshots an existing sandbox environment. @@ -644,7 +644,7 @@ async def _create( source_sandbox_environment_name (str): Required. The name of the sandbox environment to snapshot. projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_environment_id} - config (CreateAgentEngineSandboxSnapshotConfig): + config (CreateRuntimeSandboxSnapshotConfig): Optional. The configuration for the sandbox snapshot. """ @@ -691,7 +691,7 @@ async def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxSnapshotOperation._from_response( + return_value = types.RuntimeSandboxSnapshotOperation._from_response( response=response_dict, kwargs=( { @@ -950,9 +950,9 @@ async def get_sandbox_snapshot_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineSandboxSnapshotOperation: - parameter_model = types._GetAgentEngineSandboxSnapshotOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeSandboxSnapshotOperation: + parameter_model = types._GetRuntimeSandboxSnapshotOperationParameters( operation_name=operation_name, config=config, ) @@ -963,7 +963,7 @@ async def get_sandbox_snapshot_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSandboxSnapshotOperationParameters_to_vertex( + request_dict = _GetRuntimeSandboxSnapshotOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -994,7 +994,7 @@ async def get_sandbox_snapshot_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxSnapshotOperation._from_response( + return_value = types.RuntimeSandboxSnapshotOperation._from_response( response=response_dict, kwargs=( { diff --git a/agentplatform/_genai/sandbox_templates.py b/agentplatform/_genai/sandbox_templates.py index a503195be4..55d0575ba3 100644 --- a/agentplatform/_genai/sandbox_templates.py +++ b/agentplatform/_genai/sandbox_templates.py @@ -27,7 +27,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import Pager -from . import _agent_engines_utils +from . import _runtimes_utils from . import types logger = logging.getLogger("agentplatform_genai.sandboxtemplates") @@ -164,11 +164,11 @@ def _create( display_name: str, ) -> types.SandboxEnvironmentTemplateOperation: """ - Creates a new sandbox template in the Agent Engine. + Creates a new sandbox template in the Agent Runtime. Args: name (str): - Required. The name of the agent engine to create the template under. + Required. The name of the agent runtime to create the template under. Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id} display_name (str): Required. The display name of the sandbox template. @@ -249,7 +249,7 @@ def _delete( config: Optional[types.DeleteSandboxEnvironmentTemplateConfigOrDict] = None, ) -> types.DeleteSandboxEnvironmentTemplateOperation: """ - Delete an Agent Engine sandbox template. + Delete an Agent Runtime sandbox template. Args: name (str): @@ -331,7 +331,7 @@ def _get( config: Optional[types.GetSandboxEnvironmentTemplateConfigOrDict] = None, ) -> types.SandboxEnvironmentTemplate: """ - Gets an agent engine sandbox template. + Gets an agent runtime sandbox template. Args: name (str): The resource name of the SandboxEnvironmentTemplate. @@ -417,10 +417,10 @@ def _list( config: Optional[types.ListSandboxEnvironmentTemplatesConfigOrDict] = None, ) -> types.ListSandboxEnvironmentTemplatesResponse: """ - Lists Agent Engine sandbox templates. + Lists Agent Runtime sandbox templates. Args: - name (str): Name of the agent engine. Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id} + name (str): Name of the agent runtime. Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id} config (ListSandboxEnvironmentTemplatesConfig): Configuration for listing sandbox templates. Returns: @@ -496,7 +496,7 @@ def get_sandbox_environment_template_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, ) -> types.SandboxEnvironmentTemplateOperation: parameter_model = types._GetSandboxEnvironmentTemplateOperationParameters( operation_name=operation_name, @@ -570,11 +570,11 @@ def create( config: Optional[types.CreateSandboxEnvironmentTemplateConfigOrDict] = None, poll_interval_seconds: float = 0.1, ) -> types.SandboxEnvironmentTemplateOperation: - """Creates a new sandbox template in the Agent Engine. + """Creates a new sandbox template in the Agent Runtime. Args: name (str): - Required. The name of the agent engine to create sandbox template for. + Required. The name of the agent runtime to create sandbox template for. projects/{project}/locations/{location}/reasoningEngines/{resource_id} display_name (str): Required. The display name of the sandbox template. @@ -597,7 +597,7 @@ def create( config = types.CreateSandboxEnvironmentTemplateConfig.model_validate(config) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self.get_sandbox_environment_template_operation, poll_interval_seconds=poll_interval_seconds, @@ -615,17 +615,17 @@ def list( name: str, config: Optional[types.ListSandboxEnvironmentTemplatesConfigOrDict] = None, ) -> Iterator[types.SandboxEnvironmentTemplate]: - """Lists Agent Engine sandbox templates. + """Lists Agent Runtime sandbox templates. Args: name (str): - Required. The name of the agent engine to list sandbox templates for. + Required. The name of the agent runtime to list sandbox templates for. projects/{project}/locations/{location}/reasoningEngines/{resource_id} config (ListSandboxEnvironmentTemplatesConfig): Optional. The configuration for the sandbox templates to list. Returns: - Iterable[SandboxEnvironmentTemplate]: An iterable of agent engine sandbox templates. + Iterable[SandboxEnvironmentTemplate]: An iterable of agent runtime sandbox templates. """ return Pager( "sandbox_environment_templates", @@ -640,7 +640,7 @@ def get( name: str, config: Optional[types.GetSandboxEnvironmentTemplateConfigOrDict] = None, ) -> types.SandboxEnvironmentTemplate: - """Gets a sandbox template in the Agent Engine. + """Gets a sandbox template in the Agent Runtime. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -657,7 +657,7 @@ def delete( name: str, config: Optional[types.DeleteSandboxEnvironmentTemplateConfigOrDict] = None, ) -> types.DeleteSandboxEnvironmentTemplateOperation: - """Deletes a sandbox template in the Agent Engine. + """Deletes a sandbox template in the Agent Runtime. Args: name (str): Required. The name of the sandbox template to delete. @@ -679,11 +679,11 @@ async def _create( display_name: str, ) -> types.SandboxEnvironmentTemplateOperation: """ - Creates a new sandbox template in the Agent Engine. + Creates a new sandbox template in the Agent Runtime. Args: name (str): - Required. The name of the agent engine to create the template under. + Required. The name of the agent runtime to create the template under. Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id} display_name (str): Required. The display name of the sandbox template. @@ -766,7 +766,7 @@ async def _delete( config: Optional[types.DeleteSandboxEnvironmentTemplateConfigOrDict] = None, ) -> types.DeleteSandboxEnvironmentTemplateOperation: """ - Delete an Agent Engine sandbox template. + Delete an Agent Runtime sandbox template. Args: name (str): @@ -850,7 +850,7 @@ async def _get( config: Optional[types.GetSandboxEnvironmentTemplateConfigOrDict] = None, ) -> types.SandboxEnvironmentTemplate: """ - Gets an agent engine sandbox template. + Gets an agent runtime sandbox template. Args: name (str): The resource name of the SandboxEnvironmentTemplate. @@ -938,10 +938,10 @@ async def _list( config: Optional[types.ListSandboxEnvironmentTemplatesConfigOrDict] = None, ) -> types.ListSandboxEnvironmentTemplatesResponse: """ - Lists Agent Engine sandbox templates. + Lists Agent Runtime sandbox templates. Args: - name (str): Name of the agent engine. Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id} + name (str): Name of the agent runtime. Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id} config (ListSandboxEnvironmentTemplatesConfig): Configuration for listing sandbox templates. Returns: @@ -1019,7 +1019,7 @@ async def get_sandbox_environment_template_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, ) -> types.SandboxEnvironmentTemplateOperation: parameter_model = types._GetSandboxEnvironmentTemplateOperationParameters( operation_name=operation_name, diff --git a/agentplatform/_genai/sandboxes.py b/agentplatform/_genai/sandboxes.py index 033f033359..248dbdfdb7 100644 --- a/agentplatform/_genai/sandboxes.py +++ b/agentplatform/_genai/sandboxes.py @@ -34,7 +34,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import Pager -from . import _agent_engines_utils +from . import _runtimes_utils from . import types logger = logging.getLogger("agentplatform_genai.sandboxes") @@ -42,7 +42,7 @@ logger.setLevel(logging.INFO) -def _CreateAgentEngineSandboxConfig_to_vertex( +def _CreateRuntimeSandboxConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -77,7 +77,7 @@ def _CreateAgentEngineSandboxConfig_to_vertex( return to_object -def _CreateAgentEngineSandboxRequestParameters_to_vertex( +def _CreateRuntimeSandboxRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -89,14 +89,12 @@ def _CreateAgentEngineSandboxRequestParameters_to_vertex( setv(to_object, ["spec"], getv(from_object, ["spec"])) if getv(from_object, ["config"]) is not None: - _CreateAgentEngineSandboxConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _CreateRuntimeSandboxConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _DeleteAgentEngineSandboxRequestParameters_to_vertex( +def _DeleteRuntimeSandboxRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -107,7 +105,7 @@ def _DeleteAgentEngineSandboxRequestParameters_to_vertex( return to_object -def _ExecuteCodeAgentEngineSandboxRequestParameters_to_vertex( +def _ExecuteCodeRuntimeSandboxRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -121,7 +119,7 @@ def _ExecuteCodeAgentEngineSandboxRequestParameters_to_vertex( return to_object -def _GetAgentEngineSandboxOperationParameters_to_vertex( +def _GetRuntimeSandboxOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -134,7 +132,7 @@ def _GetAgentEngineSandboxOperationParameters_to_vertex( return to_object -def _GetAgentEngineSandboxRequestParameters_to_vertex( +def _GetRuntimeSandboxRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -145,7 +143,7 @@ def _GetAgentEngineSandboxRequestParameters_to_vertex( return to_object -def _ListAgentEngineSandboxesConfig_to_vertex( +def _ListRuntimeSandboxesConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -163,7 +161,7 @@ def _ListAgentEngineSandboxesConfig_to_vertex( return to_object -def _ListAgentEngineSandboxesRequestParameters_to_vertex( +def _ListRuntimeSandboxesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -172,9 +170,7 @@ def _ListAgentEngineSandboxesRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineSandboxesConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _ListRuntimeSandboxesConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object @@ -186,13 +182,13 @@ def _create( *, name: str, spec: Optional[types.SandboxEnvironmentSpecOrDict] = None, - config: Optional[types.CreateAgentEngineSandboxConfigOrDict] = None, - ) -> types.AgentEngineSandboxOperation: + config: Optional[types.CreateRuntimeSandboxConfigOrDict] = None, + ) -> types.RuntimeSandboxOperation: """ - Creates a new sandbox in the Agent Engine. + Creates a new sandbox in the Agent Runtime. """ - parameter_model = types._CreateAgentEngineSandboxRequestParameters( + parameter_model = types._CreateRuntimeSandboxRequestParameters( name=name, spec=spec, config=config, @@ -204,7 +200,7 @@ def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _CreateRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -233,7 +229,7 @@ def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxOperation._from_response( + return_value = types.RuntimeSandboxOperation._from_response( response=response_dict, kwargs=( { @@ -261,19 +257,19 @@ def _delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineSandboxConfigOrDict] = None, - ) -> types.DeleteAgentEngineSandboxOperation: + config: Optional[types.DeleteRuntimeSandboxConfigOrDict] = None, + ) -> types.DeleteRuntimeSandboxOperation: """ - Delete an Agent Engine sandbox. + Delete an Agent Runtime sandbox. Args: name (str): - Required. The name of the Agent Engine sandbox to be deleted. Format: + Required. The name of the Agent Runtime sandbox to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox}`. """ - parameter_model = types._DeleteAgentEngineSandboxRequestParameters( + parameter_model = types._DeleteRuntimeSandboxRequestParameters( name=name, config=config, ) @@ -284,7 +280,7 @@ def _delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _DeleteRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -313,7 +309,7 @@ def _delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineSandboxOperation._from_response( + return_value = types.DeleteRuntimeSandboxOperation._from_response( response=response_dict, kwargs=( { @@ -342,13 +338,13 @@ def _execute_code( *, name: str, inputs: Optional[builtins.list[types.ChunkOrDict]] = None, - config: Optional[types.ExecuteCodeAgentEngineSandboxConfigOrDict] = None, + config: Optional[types.ExecuteCodeRuntimeSandboxConfigOrDict] = None, ) -> types.ExecuteSandboxEnvironmentResponse: """ - Execute code in an Agent Engine sandbox. + Execute code in an Agent Runtime sandbox. """ - parameter_model = types._ExecuteCodeAgentEngineSandboxRequestParameters( + parameter_model = types._ExecuteCodeRuntimeSandboxRequestParameters( name=name, inputs=inputs, config=config, @@ -360,7 +356,7 @@ def _execute_code( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ExecuteCodeAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _ExecuteCodeRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -414,13 +410,10 @@ def _execute_code( return return_value def _get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineSandboxConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeSandboxConfigOrDict] = None ) -> types.SandboxEnvironment: """ - Gets an agent engine sandbox. + Gets an agent runtime sandbox. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -429,7 +422,7 @@ def _get( """ - parameter_model = types._GetAgentEngineSandboxRequestParameters( + parameter_model = types._GetRuntimeSandboxRequestParameters( name=name, config=config, ) @@ -440,7 +433,7 @@ def _get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _GetRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -497,23 +490,23 @@ def _list( self, *, name: str, - config: Optional[types.ListAgentEngineSandboxesConfigOrDict] = None, - ) -> types.ListAgentEngineSandboxesResponse: + config: Optional[types.ListRuntimeSandboxesConfigOrDict] = None, + ) -> types.ListRuntimeSandboxesResponse: """ - Lists Agent Engine sandboxes. + Lists Agent Runtime sandboxes. Args: - name (str): Required. The name of the Agent Engine to list sessions for. Format: + name (str): Required. The name of the Agent Runtime to list sessions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineSandboxesConfig): - Optional. Additional configurations for listing the Agent Engine sandboxes. + config (ListRuntimeSandboxesConfig): + Optional. Additional configurations for listing the Agent Runtime sandboxes. Returns: - ListReasoningEnginesSandboxesResponse: The requested Agent Engine sandboxes. + ListReasoningEnginesSandboxesResponse: The requested Agent Runtime sandboxes. """ - parameter_model = types._ListAgentEngineSandboxesRequestParameters( + parameter_model = types._ListRuntimeSandboxesRequestParameters( name=name, config=config, ) @@ -524,7 +517,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineSandboxesRequestParameters_to_vertex( + request_dict = _ListRuntimeSandboxesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -553,7 +546,7 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineSandboxesResponse._from_response( + return_value = types.ListRuntimeSandboxesResponse._from_response( response=response_dict, kwargs=( { @@ -581,9 +574,9 @@ def _get_sandbox_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineSandboxOperation: - parameter_model = types._GetAgentEngineSandboxOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeSandboxOperation: + parameter_model = types._GetRuntimeSandboxOperationParameters( operation_name=operation_name, config=config, ) @@ -594,7 +587,7 @@ def _get_sandbox_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSandboxOperationParameters_to_vertex( + request_dict = _GetRuntimeSandboxOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -623,7 +616,7 @@ def _get_sandbox_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxOperation._from_response( + return_value = types.RuntimeSandboxOperation._from_response( response=response_dict, kwargs=( { @@ -659,7 +652,7 @@ def templates(self) -> Any: ) except ImportError as e: raise ImportError( - "The 'agent_engines.sandboxes.templates' module requires " + "The 'runtimes.sandboxes.templates' module requires " "additional packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e @@ -674,7 +667,7 @@ def snapshots(self) -> Any: ) except ImportError as e: raise ImportError( - "The 'agent_engines.sandboxes.snapshots' module requires " + "The 'runtimes.sandboxes.snapshots' module requires " "additional packages. Please install them using pip install " "google-cloud-aiplatform[sandbox_snapshots]" ) from e @@ -686,24 +679,24 @@ def create( name: str, poll_interval_seconds: float = 0.1, spec: Optional[types.SandboxEnvironmentSpecOrDict] = None, - config: Optional[types.CreateAgentEngineSandboxConfigOrDict] = None, - ) -> types.AgentEngineSandboxOperation: - """Creates a new sandbox in the Agent Engine. + config: Optional[types.CreateRuntimeSandboxConfigOrDict] = None, + ) -> types.RuntimeSandboxOperation: + """Creates a new sandbox in the Agent Runtime. Args: name (str): - Required. The name of the agent engine to create sandbox for. + Required. The name of the agent runtime to create sandbox for. projects/{project}/locations/{location}/reasoningEngines/{resource_id} poll_interval_seconds (float): Optional. The interval in seconds to poll for sandbox creation completion. spec (SandboxEnvironmentSpec): Optional. The specification for the sandbox to create. - config (CreateAgentEngineSandboxConfigOrDict): + config (CreateRuntimeSandboxConfigOrDict): Optional. The configuration for the sandbox. Returns: - AgentEngineSandboxOperation: The operation for creating the sandbox. + RuntimeSandboxOperation: The operation for creating the sandbox. """ if spec: computer_use = False @@ -722,12 +715,12 @@ def create( config=config, ) if config is None: - config = types.CreateAgentEngineSandboxConfig() + config = types.CreateRuntimeSandboxConfig() elif isinstance(config, dict): - config = types.CreateAgentEngineSandboxConfig.model_validate(config) + config = types.CreateRuntimeSandboxConfig.model_validate(config) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_sandbox_operation, poll_interval_seconds=poll_interval_seconds, @@ -743,19 +736,19 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineSandboxesConfigOrDict] = None, + config: Optional[types.ListRuntimeSandboxesConfigOrDict] = None, ) -> Iterator[types.SandboxEnvironment]: - """Lists Agent Engine sandboxes. + """Lists Agent Runtime sandboxes. Args: name (str): - Required. The name of the agent engine to list sandboxes for. + Required. The name of the agent runtime to list sandboxes for. projects/{project}/locations/{location}/reasoningEngines/{resource_id} - config (ListAgentEngineSandboxConfig): + config (ListRuntimeSandboxConfig): Optional. The configuration for the sandboxes to list. Returns: - Iterable[SandboxEnvironment]: An iterable of agent engine sandboxes. + Iterable[SandboxEnvironment]: An iterable of agent runtime sandboxes. """ return Pager( "sandbox_environments", @@ -769,17 +762,17 @@ def execute_code( *, name: str, input_data: dict[str, Any], - config: Optional[types.ExecuteCodeAgentEngineSandboxConfigOrDict] = None, + config: Optional[types.ExecuteCodeRuntimeSandboxConfigOrDict] = None, ) -> types.ExecuteSandboxEnvironmentResponse: - """Executes code in the Agent Engine sandbox. + """Executes code in the Agent Runtime sandbox. Args: name (str): - Required. The name of the agent engine sandbox to run code in. + Required. The name of the agent runtime sandbox to run code in. projects/{project}/locations/{location}/reasoningEngines/{resource_id}/SandboxEnvironments/{sandbox_id} input_data (dict[str, Any]): Required. The input to the code to execute. - config (ExecuteCodeAgentEngineSandboxConfigOrDict): + config (ExecuteCodeRuntimeSandboxConfigOrDict): Optional. The configuration for the sandboxes to run code in. Returns: @@ -837,15 +830,15 @@ def get( self, *, name: str, - config: Optional[types.GetAgentEngineSandboxConfigOrDict] = None, + config: Optional[types.GetRuntimeSandboxConfigOrDict] = None, ) -> types.SandboxEnvironment: - """Gets an agent engine sandbox. + """Gets an agent runtime sandbox. Args: name (str): Required. A fully-qualified resource name or ID such as projects/{project}/locations/{location}/reasoningEngines/{resource_id}/SandboxEnvironments/{sandbox_id} or a shortened name such as "reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_id}". - config (GetAgentEngineSandboxConfigOrDict): + config (GetRuntimeSandboxConfigOrDict): Optional. The configuration for the sandbox to get. """ @@ -855,15 +848,15 @@ def delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineSandboxConfigOrDict] = None, - ) -> types.DeleteAgentEngineSandboxOperation: - """Deletes an agent engine sandbox. + config: Optional[types.DeleteRuntimeSandboxConfigOrDict] = None, + ) -> types.DeleteRuntimeSandboxOperation: + """Deletes an agent runtime sandbox. Args: name (str): Required. A fully-qualified resource name or ID such as projects/{project}/locations/{location}/reasoningEngines/{resource_id}/SandboxEnvironments/{sandbox_id} or a shortened name such as "reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_id}". - config (DeleteAgentEngineSandboxConfigOrDict): + config (DeleteRuntimeSandboxConfigOrDict): Optional. The configuration for the sandbox to delete. """ return self._delete(name=name, config=config) @@ -1036,13 +1029,13 @@ async def _create( *, name: str, spec: Optional[types.SandboxEnvironmentSpecOrDict] = None, - config: Optional[types.CreateAgentEngineSandboxConfigOrDict] = None, - ) -> types.AgentEngineSandboxOperation: + config: Optional[types.CreateRuntimeSandboxConfigOrDict] = None, + ) -> types.RuntimeSandboxOperation: """ - Creates a new sandbox in the Agent Engine. + Creates a new sandbox in the Agent Runtime. """ - parameter_model = types._CreateAgentEngineSandboxRequestParameters( + parameter_model = types._CreateRuntimeSandboxRequestParameters( name=name, spec=spec, config=config, @@ -1054,7 +1047,7 @@ async def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _CreateRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1085,7 +1078,7 @@ async def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxOperation._from_response( + return_value = types.RuntimeSandboxOperation._from_response( response=response_dict, kwargs=( { @@ -1113,19 +1106,19 @@ async def _delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineSandboxConfigOrDict] = None, - ) -> types.DeleteAgentEngineSandboxOperation: + config: Optional[types.DeleteRuntimeSandboxConfigOrDict] = None, + ) -> types.DeleteRuntimeSandboxOperation: """ - Delete an Agent Engine sandbox. + Delete an Agent Runtime sandbox. Args: name (str): - Required. The name of the Agent Engine sandbox to be deleted. Format: + Required. The name of the Agent Runtime sandbox to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox}`. """ - parameter_model = types._DeleteAgentEngineSandboxRequestParameters( + parameter_model = types._DeleteRuntimeSandboxRequestParameters( name=name, config=config, ) @@ -1136,7 +1129,7 @@ async def _delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _DeleteRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1167,7 +1160,7 @@ async def _delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineSandboxOperation._from_response( + return_value = types.DeleteRuntimeSandboxOperation._from_response( response=response_dict, kwargs=( { @@ -1196,13 +1189,13 @@ async def _execute_code( *, name: str, inputs: Optional[builtins.list[types.ChunkOrDict]] = None, - config: Optional[types.ExecuteCodeAgentEngineSandboxConfigOrDict] = None, + config: Optional[types.ExecuteCodeRuntimeSandboxConfigOrDict] = None, ) -> types.ExecuteSandboxEnvironmentResponse: """ - Execute code in an Agent Engine sandbox. + Execute code in an Agent Runtime sandbox. """ - parameter_model = types._ExecuteCodeAgentEngineSandboxRequestParameters( + parameter_model = types._ExecuteCodeRuntimeSandboxRequestParameters( name=name, inputs=inputs, config=config, @@ -1214,7 +1207,7 @@ async def _execute_code( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ExecuteCodeAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _ExecuteCodeRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1270,13 +1263,10 @@ async def _execute_code( return return_value async def _get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineSandboxConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeSandboxConfigOrDict] = None ) -> types.SandboxEnvironment: """ - Gets an agent engine sandbox. + Gets an agent runtime sandbox. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -1285,7 +1275,7 @@ async def _get( """ - parameter_model = types._GetAgentEngineSandboxRequestParameters( + parameter_model = types._GetRuntimeSandboxRequestParameters( name=name, config=config, ) @@ -1296,7 +1286,7 @@ async def _get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _GetRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1355,23 +1345,23 @@ async def _list( self, *, name: str, - config: Optional[types.ListAgentEngineSandboxesConfigOrDict] = None, - ) -> types.ListAgentEngineSandboxesResponse: + config: Optional[types.ListRuntimeSandboxesConfigOrDict] = None, + ) -> types.ListRuntimeSandboxesResponse: """ - Lists Agent Engine sandboxes. + Lists Agent Runtime sandboxes. Args: - name (str): Required. The name of the Agent Engine to list sessions for. Format: + name (str): Required. The name of the Agent Runtime to list sessions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineSandboxesConfig): - Optional. Additional configurations for listing the Agent Engine sandboxes. + config (ListRuntimeSandboxesConfig): + Optional. Additional configurations for listing the Agent Runtime sandboxes. Returns: - ListReasoningEnginesSandboxesResponse: The requested Agent Engine sandboxes. + ListReasoningEnginesSandboxesResponse: The requested Agent Runtime sandboxes. """ - parameter_model = types._ListAgentEngineSandboxesRequestParameters( + parameter_model = types._ListRuntimeSandboxesRequestParameters( name=name, config=config, ) @@ -1382,7 +1372,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineSandboxesRequestParameters_to_vertex( + request_dict = _ListRuntimeSandboxesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1413,7 +1403,7 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineSandboxesResponse._from_response( + return_value = types.ListRuntimeSandboxesResponse._from_response( response=response_dict, kwargs=( { @@ -1441,9 +1431,9 @@ async def _get_sandbox_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineSandboxOperation: - parameter_model = types._GetAgentEngineSandboxOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeSandboxOperation: + parameter_model = types._GetRuntimeSandboxOperationParameters( operation_name=operation_name, config=config, ) @@ -1454,7 +1444,7 @@ async def _get_sandbox_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSandboxOperationParameters_to_vertex( + request_dict = _GetRuntimeSandboxOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1485,7 +1475,7 @@ async def _get_sandbox_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxOperation._from_response( + return_value = types.RuntimeSandboxOperation._from_response( response=response_dict, kwargs=( { diff --git a/agentplatform/_genai/session_events.py b/agentplatform/_genai/session_events.py index 9333f29549..a4f46f5d49 100644 --- a/agentplatform/_genai/session_events.py +++ b/agentplatform/_genai/session_events.py @@ -35,7 +35,7 @@ logger.setLevel(logging.INFO) -def _AppendAgentEngineSessionEventConfig_to_vertex( +def _AppendRuntimeSessionEventConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -62,7 +62,7 @@ def _AppendAgentEngineSessionEventConfig_to_vertex( return to_object -def _AppendAgentEngineSessionEventRequestParameters_to_vertex( +def _AppendRuntimeSessionEventRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -80,14 +80,14 @@ def _AppendAgentEngineSessionEventRequestParameters_to_vertex( setv(to_object, ["timestamp"], getv(from_object, ["timestamp"])) if getv(from_object, ["config"]) is not None: - _AppendAgentEngineSessionEventConfig_to_vertex( + _AppendRuntimeSessionEventConfig_to_vertex( getv(from_object, ["config"]), to_object ) return to_object -def _ListAgentEngineSessionEventsConfig_to_vertex( +def _ListRuntimeSessionEventsConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -105,7 +105,7 @@ def _ListAgentEngineSessionEventsConfig_to_vertex( return to_object -def _ListAgentEngineSessionEventsRequestParameters_to_vertex( +def _ListRuntimeSessionEventsRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -114,7 +114,7 @@ def _ListAgentEngineSessionEventsRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineSessionEventsConfig_to_vertex( + _ListRuntimeSessionEventsConfig_to_vertex( getv(from_object, ["config"]), to_object ) @@ -130,26 +130,26 @@ def append( author: str, invocation_id: str, timestamp: datetime.datetime, - config: Optional[types.AppendAgentEngineSessionEventConfigOrDict] = None, - ) -> types.AppendAgentEngineSessionEventResponse: + config: Optional[types.AppendRuntimeSessionEventConfigOrDict] = None, + ) -> types.AppendRuntimeSessionEventResponse: """ - Appends Agent Engine session event. + Appends Agent Runtime session event. Args: - name (str): Required. The name of the Agent Engine session to append the event to. Format: + name (str): Required. The name of the Agent Runtime session to append the event to. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - author (str): Required. The author of the Agent Engine session event. - invocation_id (str): Required. The invocation ID of the Agent Engine session event. - timestamp (datetime.datetime): Required. The timestamp of the Agent Engine session event. - config (AppendAgentEngineSessionEventConfig): - Optional. Additional configurations for appending the Agent Engine session event. + author (str): Required. The author of the Agent Runtime session event. + invocation_id (str): Required. The invocation ID of the Agent Runtime session event. + timestamp (datetime.datetime): Required. The timestamp of the Agent Runtime session event. + config (AppendRuntimeSessionEventConfig): + Optional. Additional configurations for appending the Agent Runtime session event. Returns: - AppendAgentEngineSessionEventResponse: The requested Agent Engine session event. + AppendRuntimeSessionEventResponse: The requested Agent Runtime session event. """ - parameter_model = types._AppendAgentEngineSessionEventRequestParameters( + parameter_model = types._AppendRuntimeSessionEventRequestParameters( name=name, author=author, invocation_id=invocation_id, @@ -163,7 +163,7 @@ def append( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _AppendAgentEngineSessionEventRequestParameters_to_vertex( + request_dict = _AppendRuntimeSessionEventRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -192,7 +192,7 @@ def append( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AppendAgentEngineSessionEventResponse._from_response( + return_value = types.AppendRuntimeSessionEventResponse._from_response( response=response_dict, kwargs=( { @@ -220,23 +220,23 @@ def _list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionEventsConfigOrDict] = None, - ) -> types.ListAgentEngineSessionEventsResponse: + config: Optional[types.ListRuntimeSessionEventsConfigOrDict] = None, + ) -> types.ListRuntimeSessionEventsResponse: """ - Lists Agent Engine session events. + Lists Agent Runtime session events. Args: - name (str): Required. The name of the Agent Engine session to list events for. Format: + name (str): Required. The name of the Agent Runtime session to list events for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (ListAgentEngineSessionEventsConfig): - Optional. Additional configurations for listing the Agent Engine session events. + config (ListRuntimeSessionEventsConfig): + Optional. Additional configurations for listing the Agent Runtime session events. Returns: - ListAgentEngineSessionEventsResponse: The requested Agent Engine session events. + ListRuntimeSessionEventsResponse: The requested Agent Runtime session events. """ - parameter_model = types._ListAgentEngineSessionEventsRequestParameters( + parameter_model = types._ListRuntimeSessionEventsRequestParameters( name=name, config=config, ) @@ -247,7 +247,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineSessionEventsRequestParameters_to_vertex( + request_dict = _ListRuntimeSessionEventsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -276,7 +276,7 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineSessionEventsResponse._from_response( + return_value = types.ListRuntimeSessionEventsResponse._from_response( response=response_dict, kwargs=( { @@ -304,14 +304,14 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionEventsConfigOrDict] = None, + config: Optional[types.ListRuntimeSessionEventsConfigOrDict] = None, ) -> Iterator[types.SessionEvent]: - """Lists Agent Engine session events. + """Lists Agent Runtime session events. Args: - name (str): Required. The name of the agent engine to list session + name (str): Required. The name of the agent runtime to list session events for. - config (ListAgentEngineSessionEventsConfig): Optional. The configuration + config (ListRuntimeSessionEventsConfig): Optional. The configuration for the session events to list. Currently, the `filter` field in `config` only supports filtering by `timestamp`. The timestamp value must be enclosed in double quotes and include the time zone @@ -339,26 +339,26 @@ async def append( author: str, invocation_id: str, timestamp: datetime.datetime, - config: Optional[types.AppendAgentEngineSessionEventConfigOrDict] = None, - ) -> types.AppendAgentEngineSessionEventResponse: + config: Optional[types.AppendRuntimeSessionEventConfigOrDict] = None, + ) -> types.AppendRuntimeSessionEventResponse: """ - Appends Agent Engine session event. + Appends Agent Runtime session event. Args: - name (str): Required. The name of the Agent Engine session to append the event to. Format: + name (str): Required. The name of the Agent Runtime session to append the event to. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - author (str): Required. The author of the Agent Engine session event. - invocation_id (str): Required. The invocation ID of the Agent Engine session event. - timestamp (datetime.datetime): Required. The timestamp of the Agent Engine session event. - config (AppendAgentEngineSessionEventConfig): - Optional. Additional configurations for appending the Agent Engine session event. + author (str): Required. The author of the Agent Runtime session event. + invocation_id (str): Required. The invocation ID of the Agent Runtime session event. + timestamp (datetime.datetime): Required. The timestamp of the Agent Runtime session event. + config (AppendRuntimeSessionEventConfig): + Optional. Additional configurations for appending the Agent Runtime session event. Returns: - AppendAgentEngineSessionEventResponse: The requested Agent Engine session event. + AppendRuntimeSessionEventResponse: The requested Agent Runtime session event. """ - parameter_model = types._AppendAgentEngineSessionEventRequestParameters( + parameter_model = types._AppendRuntimeSessionEventRequestParameters( name=name, author=author, invocation_id=invocation_id, @@ -372,7 +372,7 @@ async def append( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _AppendAgentEngineSessionEventRequestParameters_to_vertex( + request_dict = _AppendRuntimeSessionEventRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -403,7 +403,7 @@ async def append( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AppendAgentEngineSessionEventResponse._from_response( + return_value = types.AppendRuntimeSessionEventResponse._from_response( response=response_dict, kwargs=( { @@ -431,23 +431,23 @@ async def _list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionEventsConfigOrDict] = None, - ) -> types.ListAgentEngineSessionEventsResponse: + config: Optional[types.ListRuntimeSessionEventsConfigOrDict] = None, + ) -> types.ListRuntimeSessionEventsResponse: """ - Lists Agent Engine session events. + Lists Agent Runtime session events. Args: - name (str): Required. The name of the Agent Engine session to list events for. Format: + name (str): Required. The name of the Agent Runtime session to list events for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (ListAgentEngineSessionEventsConfig): - Optional. Additional configurations for listing the Agent Engine session events. + config (ListRuntimeSessionEventsConfig): + Optional. Additional configurations for listing the Agent Runtime session events. Returns: - ListAgentEngineSessionEventsResponse: The requested Agent Engine session events. + ListRuntimeSessionEventsResponse: The requested Agent Runtime session events. """ - parameter_model = types._ListAgentEngineSessionEventsRequestParameters( + parameter_model = types._ListRuntimeSessionEventsRequestParameters( name=name, config=config, ) @@ -458,7 +458,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineSessionEventsRequestParameters_to_vertex( + request_dict = _ListRuntimeSessionEventsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -489,7 +489,7 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineSessionEventsResponse._from_response( + return_value = types.ListRuntimeSessionEventsResponse._from_response( response=response_dict, kwargs=( { @@ -517,14 +517,14 @@ async def list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionEventsConfigOrDict] = None, + config: Optional[types.ListRuntimeSessionEventsConfigOrDict] = None, ) -> AsyncPager[types.SessionEvent]: - """Lists Agent Engine session events. + """Lists Agent Runtime session events. Args: - name (str): Required. The name of the agent engine to list session + name (str): Required. The name of the agent runtime to list session events for. - config (ListAgentEngineSessionEventsConfig): Optional. The configuration + config (ListRuntimeSessionEventsConfig): Optional. The configuration for the session events to list. Currently, the `filter` field in `config` only supports filtering by `timestamp`. The timestamp value must be enclosed in double quotes and include the time zone diff --git a/agentplatform/_genai/sessions.py b/agentplatform/_genai/sessions.py index bcd4862494..61ade3fa4d 100644 --- a/agentplatform/_genai/sessions.py +++ b/agentplatform/_genai/sessions.py @@ -29,7 +29,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import AsyncPager, Pager -from . import _agent_engines_utils +from . import _runtimes_utils from . import types if typing.TYPE_CHECKING: @@ -43,7 +43,7 @@ logger.setLevel(logging.INFO) -def _CreateAgentEngineSessionConfig_to_vertex( +def _CreateRuntimeSessionConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -70,7 +70,7 @@ def _CreateAgentEngineSessionConfig_to_vertex( return to_object -def _CreateAgentEngineSessionRequestParameters_to_vertex( +def _CreateRuntimeSessionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -82,14 +82,12 @@ def _CreateAgentEngineSessionRequestParameters_to_vertex( setv(to_object, ["userId"], getv(from_object, ["user_id"])) if getv(from_object, ["config"]) is not None: - _CreateAgentEngineSessionConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _CreateRuntimeSessionConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _DeleteAgentEngineSessionRequestParameters_to_vertex( +def _DeleteRuntimeSessionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -100,7 +98,7 @@ def _DeleteAgentEngineSessionRequestParameters_to_vertex( return to_object -def _GetAgentEngineSessionOperationParameters_to_vertex( +def _GetRuntimeSessionOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -113,7 +111,7 @@ def _GetAgentEngineSessionOperationParameters_to_vertex( return to_object -def _GetAgentEngineSessionRequestParameters_to_vertex( +def _GetRuntimeSessionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -124,7 +122,7 @@ def _GetAgentEngineSessionRequestParameters_to_vertex( return to_object -def _ListAgentEngineSessionsConfig_to_vertex( +def _ListRuntimeSessionsConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -142,7 +140,7 @@ def _ListAgentEngineSessionsConfig_to_vertex( return to_object -def _ListAgentEngineSessionsRequestParameters_to_vertex( +def _ListRuntimeSessionsRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -151,14 +149,12 @@ def _ListAgentEngineSessionsRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineSessionsConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _ListRuntimeSessionsConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _UpdateAgentEngineSessionConfig_to_vertex( +def _UpdateRuntimeSessionConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -193,7 +189,7 @@ def _UpdateAgentEngineSessionConfig_to_vertex( return to_object -def _UpdateAgentEngineSessionRequestParameters_to_vertex( +def _UpdateRuntimeSessionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -202,9 +198,7 @@ def _UpdateAgentEngineSessionRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _UpdateAgentEngineSessionConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _UpdateRuntimeSessionConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object @@ -216,24 +210,24 @@ def _create( *, name: str, user_id: str, - config: Optional[types.CreateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: + config: Optional[types.CreateRuntimeSessionConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: """ - Creates a new session in the Agent Engine. + Creates a new session in the Agent Runtime. Args: - name (str): Required. The name of the Agent Engine to create the session under. Format: + name (str): Required. The name of the Agent Runtime to create the session under. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. user_id (str): Required. The user ID of the session. - config (CreateAgentEngineSessionConfig): - Optional. Additional configurations for creating the Agent Engine session. + config (CreateRuntimeSessionConfig): + Optional. Additional configurations for creating the Agent Runtime session. Returns: - AgentEngineSessionOperation: The operation for creating the Agent Engine session. + RuntimeSessionOperation: The operation for creating the Agent Runtime session. """ - parameter_model = types._CreateAgentEngineSessionRequestParameters( + parameter_model = types._CreateRuntimeSessionRequestParameters( name=name, user_id=user_id, config=config, @@ -245,7 +239,7 @@ def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineSessionRequestParameters_to_vertex( + request_dict = _CreateRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -274,7 +268,7 @@ def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSessionOperation._from_response( + return_value = types.RuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -302,23 +296,23 @@ def delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineSessionConfigOrDict] = None, - ) -> types.DeleteAgentEngineSessionOperation: + config: Optional[types.DeleteRuntimeSessionConfigOrDict] = None, + ) -> types.DeleteRuntimeSessionOperation: """ - Delete an Agent Engine session. + Delete an Agent Runtime session. Args: - name (str): Required. The name of the Agent Engine session to be deleted. Format: + name (str): Required. The name of the Agent Runtime session to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (DeleteAgentEngineSessionConfig): - Optional. Additional configurations for deleting the Agent Engine session. + config (DeleteRuntimeSessionConfig): + Optional. Additional configurations for deleting the Agent Runtime session. Returns: - DeleteAgentEngineSessionOperation: The operation for deleting the Agent Engine session. + DeleteRuntimeSessionOperation: The operation for deleting the Agent Runtime session. """ - parameter_model = types._DeleteAgentEngineSessionRequestParameters( + parameter_model = types._DeleteRuntimeSessionRequestParameters( name=name, config=config, ) @@ -329,7 +323,7 @@ def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineSessionRequestParameters_to_vertex( + request_dict = _DeleteRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -358,7 +352,7 @@ def delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineSessionOperation._from_response( + return_value = types.DeleteRuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -383,26 +377,23 @@ def delete( return return_value def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineSessionConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeSessionConfigOrDict] = None ) -> types.Session: """ - Gets an agent engine session. + Gets an agent runtime session. Args: - name (str): Required. The name of the Agent Engine session to get. Format: + name (str): Required. The name of the Agent Runtime session to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (GetAgentEngineSessionConfig): - Optional. Additional configurations for getting the Agent Engine session. + config (GetRuntimeSessionConfig): + Optional. Additional configurations for getting the Agent Runtime session. Returns: - AgentEngineSession: The requested Agent Engine session. + RuntimeSession: The requested Agent Runtime session. """ - parameter_model = types._GetAgentEngineSessionRequestParameters( + parameter_model = types._GetRuntimeSessionRequestParameters( name=name, config=config, ) @@ -413,7 +404,7 @@ def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSessionRequestParameters_to_vertex( + request_dict = _GetRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -470,23 +461,23 @@ def _list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionsConfigOrDict] = None, + config: Optional[types.ListRuntimeSessionsConfigOrDict] = None, ) -> types.ListReasoningEnginesSessionsResponse: """ - Lists Agent Engine sessions. + Lists Agent Runtime sessions. Args: - name (str): Required. The name of the Agent Engine to list sessions for. Format: + name (str): Required. The name of the Agent Runtime to list sessions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineSessionsConfig): - Optional. Additional configurations for listing the Agent Engine sessions. + config (ListRuntimeSessionsConfig): + Optional. Additional configurations for listing the Agent Runtime sessions. Returns: - ListReasoningEnginesSessionsResponse: The requested Agent Engine sessions. + ListReasoningEnginesSessionsResponse: The requested Agent Runtime sessions. """ - parameter_model = types._ListAgentEngineSessionsRequestParameters( + parameter_model = types._ListRuntimeSessionsRequestParameters( name=name, config=config, ) @@ -497,7 +488,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineSessionsRequestParameters_to_vertex( + request_dict = _ListRuntimeSessionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -554,9 +545,9 @@ def _get_session_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: - parameter_model = types._GetAgentEngineSessionOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: + parameter_model = types._GetRuntimeSessionOperationParameters( operation_name=operation_name, config=config, ) @@ -567,7 +558,7 @@ def _get_session_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSessionOperationParameters_to_vertex( + request_dict = _GetRuntimeSessionOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -596,7 +587,7 @@ def _get_session_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSessionOperation._from_response( + return_value = types.RuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -624,23 +615,23 @@ def _update( self, *, name: str, - config: Optional[types.UpdateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: + config: Optional[types.UpdateRuntimeSessionConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: """ - Updates an Agent Engine session. + Updates an Agent Runtime session. Args: - name (str): Required. The name of the Agent Engine session to be updated. Format: + name (str): Required. The name of the Agent Runtime session to be updated. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (UpdateAgentEngineSessionConfig): - Optional. Additional configurations for updating the Agent Engine session. + config (UpdateRuntimeSessionConfig): + Optional. Additional configurations for updating the Agent Runtime session. Returns: - AgentEngineSessionOperation: The operation for updating the Agent Engine session. + RuntimeSessionOperation: The operation for updating the Agent Runtime session. """ - parameter_model = types._UpdateAgentEngineSessionRequestParameters( + parameter_model = types._UpdateRuntimeSessionRequestParameters( name=name, config=config, ) @@ -651,7 +642,7 @@ def _update( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _UpdateAgentEngineSessionRequestParameters_to_vertex( + request_dict = _UpdateRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -680,7 +671,7 @@ def _update( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSessionOperation._from_response( + return_value = types.RuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -715,7 +706,7 @@ def events(self) -> "session_events_module.SessionEvents": self._events = importlib.import_module(".session_events", __package__) except ImportError as e: raise ImportError( - "The 'agent_engines.sessions.events' module requires" + "The 'runtimes.sessions.events' module requires" "additional packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e @@ -726,25 +717,25 @@ def create( *, name: str, user_id: str, - config: Optional[types.CreateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: - """Creates a new session in the Agent Engine. + config: Optional[types.CreateRuntimeSessionConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: + """Creates a new session in the Agent Runtime. Args: name (str): - Required. The name of the agent engine to create the session for. + Required. The name of the agent runtime to create the session for. user_id (str): Required. The user ID of the session. - config (CreateAgentEngineSessionConfig): + config (CreateRuntimeSessionConfig): Optional. The configuration for the session to create. Returns: - AgentEngineSessionOperation: The operation for creating the session. + RuntimeSessionOperation: The operation for creating the session. """ if config is None: - config = types.CreateAgentEngineSessionConfig() + config = types.CreateRuntimeSessionConfig() elif isinstance(config, dict): - config = types.CreateAgentEngineSessionConfig.model_validate(config) + config = types.CreateRuntimeSessionConfig.model_validate(config) operation = self._create( name=name, user_id=user_id, @@ -752,7 +743,7 @@ def create( ) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_session_operation, poll_interval_seconds=0.5, @@ -774,14 +765,14 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionsConfigOrDict] = None, + config: Optional[types.ListRuntimeSessionsConfigOrDict] = None, ) -> Iterator[types.Session]: - """Lists Agent Engine sessions. + """Lists Agent Runtime sessions. Args: - name (str): Required. The name of the agent engine to list sessions + name (str): Required. The name of the agent runtime to list sessions for. - config (ListAgentEngineSessionConfig): Optional. The configuration + config (ListRuntimeSessionConfig): Optional. The configuration for the sessions to list. Returns: @@ -803,24 +794,24 @@ async def _create( *, name: str, user_id: str, - config: Optional[types.CreateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: + config: Optional[types.CreateRuntimeSessionConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: """ - Creates a new session in the Agent Engine. + Creates a new session in the Agent Runtime. Args: - name (str): Required. The name of the Agent Engine to create the session under. Format: + name (str): Required. The name of the Agent Runtime to create the session under. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. user_id (str): Required. The user ID of the session. - config (CreateAgentEngineSessionConfig): - Optional. Additional configurations for creating the Agent Engine session. + config (CreateRuntimeSessionConfig): + Optional. Additional configurations for creating the Agent Runtime session. Returns: - AgentEngineSessionOperation: The operation for creating the Agent Engine session. + RuntimeSessionOperation: The operation for creating the Agent Runtime session. """ - parameter_model = types._CreateAgentEngineSessionRequestParameters( + parameter_model = types._CreateRuntimeSessionRequestParameters( name=name, user_id=user_id, config=config, @@ -832,7 +823,7 @@ async def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineSessionRequestParameters_to_vertex( + request_dict = _CreateRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -863,7 +854,7 @@ async def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSessionOperation._from_response( + return_value = types.RuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -891,23 +882,23 @@ async def delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineSessionConfigOrDict] = None, - ) -> types.DeleteAgentEngineSessionOperation: + config: Optional[types.DeleteRuntimeSessionConfigOrDict] = None, + ) -> types.DeleteRuntimeSessionOperation: """ - Delete an Agent Engine session. + Delete an Agent Runtime session. Args: - name (str): Required. The name of the Agent Engine session to be deleted. Format: + name (str): Required. The name of the Agent Runtime session to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (DeleteAgentEngineSessionConfig): - Optional. Additional configurations for deleting the Agent Engine session. + config (DeleteRuntimeSessionConfig): + Optional. Additional configurations for deleting the Agent Runtime session. Returns: - DeleteAgentEngineSessionOperation: The operation for deleting the Agent Engine session. + DeleteRuntimeSessionOperation: The operation for deleting the Agent Runtime session. """ - parameter_model = types._DeleteAgentEngineSessionRequestParameters( + parameter_model = types._DeleteRuntimeSessionRequestParameters( name=name, config=config, ) @@ -918,7 +909,7 @@ async def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineSessionRequestParameters_to_vertex( + request_dict = _DeleteRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -949,7 +940,7 @@ async def delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineSessionOperation._from_response( + return_value = types.DeleteRuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -974,26 +965,23 @@ async def delete( return return_value async def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineSessionConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeSessionConfigOrDict] = None ) -> types.Session: """ - Gets an agent engine session. + Gets an agent runtime session. Args: - name (str): Required. The name of the Agent Engine session to get. Format: + name (str): Required. The name of the Agent Runtime session to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (GetAgentEngineSessionConfig): - Optional. Additional configurations for getting the Agent Engine session. + config (GetRuntimeSessionConfig): + Optional. Additional configurations for getting the Agent Runtime session. Returns: - AgentEngineSession: The requested Agent Engine session. + RuntimeSession: The requested Agent Runtime session. """ - parameter_model = types._GetAgentEngineSessionRequestParameters( + parameter_model = types._GetRuntimeSessionRequestParameters( name=name, config=config, ) @@ -1004,7 +992,7 @@ async def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSessionRequestParameters_to_vertex( + request_dict = _GetRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1063,23 +1051,23 @@ async def _list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionsConfigOrDict] = None, + config: Optional[types.ListRuntimeSessionsConfigOrDict] = None, ) -> types.ListReasoningEnginesSessionsResponse: """ - Lists Agent Engine sessions. + Lists Agent Runtime sessions. Args: - name (str): Required. The name of the Agent Engine to list sessions for. Format: + name (str): Required. The name of the Agent Runtime to list sessions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineSessionsConfig): - Optional. Additional configurations for listing the Agent Engine sessions. + config (ListRuntimeSessionsConfig): + Optional. Additional configurations for listing the Agent Runtime sessions. Returns: - ListReasoningEnginesSessionsResponse: The requested Agent Engine sessions. + ListReasoningEnginesSessionsResponse: The requested Agent Runtime sessions. """ - parameter_model = types._ListAgentEngineSessionsRequestParameters( + parameter_model = types._ListRuntimeSessionsRequestParameters( name=name, config=config, ) @@ -1090,7 +1078,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineSessionsRequestParameters_to_vertex( + request_dict = _ListRuntimeSessionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1149,9 +1137,9 @@ async def _get_session_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: - parameter_model = types._GetAgentEngineSessionOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: + parameter_model = types._GetRuntimeSessionOperationParameters( operation_name=operation_name, config=config, ) @@ -1162,7 +1150,7 @@ async def _get_session_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSessionOperationParameters_to_vertex( + request_dict = _GetRuntimeSessionOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1193,7 +1181,7 @@ async def _get_session_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSessionOperation._from_response( + return_value = types.RuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -1221,23 +1209,23 @@ async def _update( self, *, name: str, - config: Optional[types.UpdateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: + config: Optional[types.UpdateRuntimeSessionConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: """ - Updates an Agent Engine session. + Updates an Agent Runtime session. Args: - name (str): Required. The name of the Agent Engine session to be updated. Format: + name (str): Required. The name of the Agent Runtime session to be updated. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (UpdateAgentEngineSessionConfig): - Optional. Additional configurations for updating the Agent Engine session. + config (UpdateRuntimeSessionConfig): + Optional. Additional configurations for updating the Agent Runtime session. Returns: - AgentEngineSessionOperation: The operation for updating the Agent Engine session. + RuntimeSessionOperation: The operation for updating the Agent Runtime session. """ - parameter_model = types._UpdateAgentEngineSessionRequestParameters( + parameter_model = types._UpdateRuntimeSessionRequestParameters( name=name, config=config, ) @@ -1248,7 +1236,7 @@ async def _update( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _UpdateAgentEngineSessionRequestParameters_to_vertex( + request_dict = _UpdateRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1279,7 +1267,7 @@ async def _update( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSessionOperation._from_response( + return_value = types.RuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -1314,7 +1302,7 @@ def events(self) -> "session_events_module.AsyncSessionEvents": self._events = importlib.import_module(".session_events", __package__) except ImportError as e: raise ImportError( - "The 'agent_engines.sessions.events' module requires" + "The 'runtimes.sessions.events' module requires" "additional packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e @@ -1325,25 +1313,25 @@ async def create( *, name: str, user_id: str, - config: Optional[types.CreateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: - """Creates a new session in the Agent Engine. + config: Optional[types.CreateRuntimeSessionConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: + """Creates a new session in the Agent Runtime. Args: name (str): - Required. The name of the agent engine to create the session for. + Required. The name of the agent runtime to create the session for. user_id (str): Required. The user ID of the session. - config (CreateAgentEngineSessionConfig): + config (CreateRuntimeSessionConfig): Optional. The configuration for the session to create. Returns: - AgentEngineSessionOperation: The operation for creating the session. + RuntimeSessionOperation: The operation for creating the session. """ if config is None: - config = types.CreateAgentEngineSessionConfig() + config = types.CreateRuntimeSessionConfig() elif isinstance(config, dict): - config = types.CreateAgentEngineSessionConfig.model_validate(config) + config = types.CreateRuntimeSessionConfig.model_validate(config) operation = await self._create( name=name, user_id=user_id, @@ -1351,7 +1339,7 @@ async def create( ) if config.wait_for_completion: if not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_session_operation, poll_interval_seconds=0.5, @@ -1373,14 +1361,14 @@ async def list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionsConfigOrDict] = None, + config: Optional[types.ListRuntimeSessionsConfigOrDict] = None, ) -> AsyncPager[types.Session]: - """Lists Agent Engine sessions. + """Lists Agent Runtime sessions. Args: - name (str): Required. The name of the agent engine to list sessions + name (str): Required. The name of the agent runtime to list sessions for. - config (ListAgentEngineSessionConfig): Optional. The configuration + config (ListRuntimeSessionConfig): Optional. The configuration for the sessions to list. Returns: diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 6d719ddd84..bf8830ddc2 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -22,18 +22,12 @@ from . import agent_engines from . import evals from . import prompts -from .common import _AppendAgentEngineSessionEventRequestParameters -from .common import _AppendAgentEngineTaskEventRequestParameters +from .common import _AppendRuntimeSessionEventRequestParameters from .common import _AskContextsRequestParameters from .common import _AssembleDatasetParameters from .common import _AssessDatasetParameters -from .common import _CancelQueryJobAgentEngineRequestParameters -from .common import _CheckQueryJobAgentEngineRequestParameters -from .common import _CreateAgentEngineMemoryRequestParameters -from .common import _CreateAgentEngineRequestParameters -from .common import _CreateAgentEngineSandboxRequestParameters -from .common import _CreateAgentEngineSessionRequestParameters -from .common import _CreateAgentEngineTaskRequestParameters +from .common import _CancelQueryJobRuntimeRequestParameters +from .common import _CheckQueryJobRuntimeRequestParameters from .common import _CreateDatasetParameters from .common import _CreateDatasetVersionParameters from .common import _CreateEvaluationExperimentParameters @@ -44,17 +38,15 @@ from .common import _CreateMultimodalDatasetParameters from .common import _CreateRagCorpusRequestParameters from .common import _CreateRuntimeFeedbackEntryRequestParameters +from .common import _CreateRuntimeMemoryRequestParameters +from .common import _CreateRuntimeRequestParameters +from .common import _CreateRuntimeSandboxRequestParameters +from .common import _CreateRuntimeSessionRequestParameters from .common import _CreateSandboxEnvironmentSnapshotRequestParameters from .common import _CreateSandboxEnvironmentTemplateRequestParameters from .common import _CreateSkillRequestParameters from .common import _CustomJobParameters from .common import _CustomJobParameters -from .common import _DeleteAgentEngineMemoryRequestParameters -from .common import _DeleteAgentEngineRequestParameters -from .common import _DeleteAgentEngineRuntimeRevisionRequestParameters -from .common import _DeleteAgentEngineSandboxRequestParameters -from .common import _DeleteAgentEngineSessionRequestParameters -from .common import _DeleteAgentEngineTaskRequestParameters from .common import _DeleteDatasetRequestParameters from .common import _DeleteEvaluationExperimentParameters from .common import _DeleteEvaluationMetricParameters @@ -63,36 +55,28 @@ from .common import _DeleteRagCorpusRequestParameters from .common import _DeleteRagFileRequestParameters from .common import _DeleteRuntimeFeedbackEntryRequestParameters +from .common import _DeleteRuntimeMemoryRequestParameters +from .common import _DeleteRuntimeRequestParameters +from .common import _DeleteRuntimeRevisionRequestParameters +from .common import _DeleteRuntimeSandboxRequestParameters +from .common import _DeleteRuntimeSessionRequestParameters from .common import _DeleteSandboxEnvironmentSnapshotRequestParameters from .common import _DeleteSandboxEnvironmentTemplateRequestParameters from .common import _DeleteSkillRequestParameters from .common import _EvaluateInstancesRequestParameters -from .common import _ExecuteCodeAgentEngineSandboxRequestParameters +from .common import _ExecuteCodeRuntimeSandboxRequestParameters from .common import _ExportPublisherModelRequestParameters -from .common import _GenerateAgentEngineMemoriesRequestParameters from .common import _GenerateInstanceRubricsRequest from .common import _GenerateLossClustersParameters +from .common import _GenerateRuntimeMemoriesRequestParameters from .common import _GenerateUserScenariosParameters -from .common import _GetAgentEngineGenerateMemoriesOperationParameters -from .common import _GetAgentEngineMemoryOperationParameters -from .common import _GetAgentEngineMemoryRequestParameters -from .common import _GetAgentEngineMemoryRevisionRequestParameters -from .common import _GetAgentEngineOperationParameters -from .common import _GetAgentEngineRequestParameters -from .common import _GetAgentEngineRuntimeRevisionRequestParameters -from .common import _GetAgentEngineSandboxOperationParameters -from .common import _GetAgentEngineSandboxRequestParameters -from .common import _GetAgentEngineSandboxSnapshotOperationParameters -from .common import _GetAgentEngineSessionOperationParameters -from .common import _GetAgentEngineSessionRequestParameters -from .common import _GetAgentEngineTaskRequestParameters from .common import _GetCorpusOperationParameters from .common import _GetCustomJobParameters from .common import _GetCustomJobParameters from .common import _GetDatasetOperationParameters from .common import _GetDatasetParameters from .common import _GetDatasetVersionParameters -from .common import _GetDeleteAgentEngineRuntimeRevisionOperationParameters +from .common import _GetDeleteRuntimeRevisionOperationParameters from .common import _GetEvaluationExperimentParameters from .common import _GetEvaluationItemParameters from .common import _GetEvaluationMetricParameters @@ -112,6 +96,18 @@ from .common import _GetRuntimeFeedbackOperationParameters from .common import _GetRuntimeFeedbackOperationParameters from .common import _GetRuntimeFeedbackRequestParameters +from .common import _GetRuntimeGenerateMemoriesOperationParameters +from .common import _GetRuntimeMemoryOperationParameters +from .common import _GetRuntimeMemoryRequestParameters +from .common import _GetRuntimeMemoryRevisionRequestParameters +from .common import _GetRuntimeOperationParameters +from .common import _GetRuntimeRequestParameters +from .common import _GetRuntimeRevisionRequestParameters +from .common import _GetRuntimeSandboxOperationParameters +from .common import _GetRuntimeSandboxRequestParameters +from .common import _GetRuntimeSandboxSnapshotOperationParameters +from .common import _GetRuntimeSessionOperationParameters +from .common import _GetRuntimeSessionRequestParameters from .common import _GetSandboxEnvironmentSnapshotRequestParameters from .common import _GetSandboxEnvironmentTemplateOperationParameters from .common import _GetSandboxEnvironmentTemplateRequestParameters @@ -120,15 +116,6 @@ from .common import _GetSkillRevisionRequestParameters from .common import _ImportRagFilesRequestParameters from .common import _IngestEventsRequestParameters -from .common import _ListAgentEngineMemoryRequestParameters -from .common import _ListAgentEngineMemoryRevisionsRequestParameters -from .common import _ListAgentEngineRequestParameters -from .common import _ListAgentEngineRuntimeRevisionsRequestParameters -from .common import _ListAgentEngineSandboxesRequestParameters -from .common import _ListAgentEngineSessionEventsRequestParameters -from .common import _ListAgentEngineSessionsRequestParameters -from .common import _ListAgentEngineTaskEventsRequestParameters -from .common import _ListAgentEngineTasksRequestParameters from .common import _ListDatasetsRequestParameters from .common import _ListDatasetVersionsRequestParameters from .common import _ListEvaluationExperimentsParameters @@ -138,29 +125,33 @@ from .common import _ListRagCorporaRequestParameters from .common import _ListRagFilesRequestParameters from .common import _ListRuntimeFeedbackEntriesRequestParameters +from .common import _ListRuntimeMemoryRequestParameters +from .common import _ListRuntimeMemoryRevisionsRequestParameters +from .common import _ListRuntimeRequestParameters +from .common import _ListRuntimeRevisionsRequestParameters +from .common import _ListRuntimeSandboxesRequestParameters +from .common import _ListRuntimeSessionEventsRequestParameters +from .common import _ListRuntimeSessionsRequestParameters from .common import _ListSandboxEnvironmentSnapshotsRequestParameters from .common import _ListSandboxEnvironmentTemplatesRequestParameters from .common import _ListSkillRevisionsRequestParameters from .common import _ListSkillsRequestParameters from .common import _OptimizeRequestParameters from .common import _OptimizeRequestParameters -from .common import _PurgeAgentEngineMemoriesRequestParameters -from .common import _QueryAgentEngineRequestParameters -from .common import _QueryAgentEngineRuntimeRevisionRequestParameters +from .common import _PurgeRuntimeMemoriesRequestParameters +from .common import _QueryRuntimeRequestParameters +from .common import _QueryRuntimeRevisionRequestParameters from .common import _RecommendSpecRequestParameters from .common import _RestoreVersionRequestParameters -from .common import _RetrieveAgentEngineMemoriesRequestParameters from .common import _RetrieveMemoryProfilesRequestParameters from .common import _RetrieveRagContextsRequestParameters +from .common import _RetrieveRuntimeMemoriesRequestParameters from .common import _RetrieveSkillsRequestParameters -from .common import _RollbackAgentEngineMemoryRequestParameters -from .common import _RunQueryJobAgentEngineConfig -from .common import _RunQueryJobAgentEngineConfigDict -from .common import _RunQueryJobAgentEngineConfigOrDict -from .common import _RunQueryJobAgentEngineRequestParameters -from .common import _UpdateAgentEngineMemoryRequestParameters -from .common import _UpdateAgentEngineRequestParameters -from .common import _UpdateAgentEngineSessionRequestParameters +from .common import _RollbackRuntimeMemoryRequestParameters +from .common import _RunQueryJobRuntimeConfig +from .common import _RunQueryJobRuntimeConfigDict +from .common import _RunQueryJobRuntimeConfigOrDict +from .common import _RunQueryJobRuntimeRequestParameters from .common import _UpdateDatasetParameters from .common import _UpdateEvaluationExperimentParameters from .common import _UpdateMultimodalDatasetParameters @@ -168,52 +159,18 @@ from .common import _UpdateRagCorpusRequestParameters from .common import _UpdateRuntimeFeedbackContextRequestParameters from .common import _UpdateRuntimeFeedbackEntryRequestParameters +from .common import _UpdateRuntimeMemoryRequestParameters +from .common import _UpdateRuntimeRequestParameters +from .common import _UpdateRuntimeSessionRequestParameters from .common import _UpdateSkillRequestParameters from .common import _UploadRagFileParameters -from .common import A2aTask -from .common import A2aTaskDict -from .common import A2aTaskOrDict -from .common import A2aTaskState from .common import AcceleratorType -from .common import AgentEngine -from .common import AgentEngineConfig -from .common import AgentEngineConfigDict -from .common import AgentEngineConfigOrDict -from .common import AgentEngineDict -from .common import AgentEngineGenerateMemoriesOperation -from .common import AgentEngineGenerateMemoriesOperationDict -from .common import AgentEngineGenerateMemoriesOperationOrDict -from .common import AgentEngineMemoryConfig -from .common import AgentEngineMemoryConfigDict -from .common import AgentEngineMemoryConfigOrDict -from .common import AgentEngineMemoryOperation -from .common import AgentEngineMemoryOperationDict -from .common import AgentEngineMemoryOperationOrDict -from .common import AgentEngineOperation -from .common import AgentEngineOperationDict -from .common import AgentEngineOperationOrDict -from .common import AgentEngineOrDict -from .common import AgentEnginePurgeMemoriesOperation -from .common import AgentEnginePurgeMemoriesOperationDict -from .common import AgentEnginePurgeMemoriesOperationOrDict -from .common import AgentEngineRollbackMemoryOperation -from .common import AgentEngineRollbackMemoryOperationDict -from .common import AgentEngineRollbackMemoryOperationOrDict -from .common import AgentEngineRuntimeRevision -from .common import AgentEngineRuntimeRevisionDict -from .common import AgentEngineRuntimeRevisionOrDict -from .common import AgentEngineSandboxOperation -from .common import AgentEngineSandboxOperationDict -from .common import AgentEngineSandboxOperationOrDict -from .common import AgentEngineSandboxSnapshotOperation -from .common import AgentEngineSandboxSnapshotOperationDict -from .common import AgentEngineSandboxSnapshotOperationOrDict -from .common import AgentEngineSessionOperation -from .common import AgentEngineSessionOperationDict -from .common import AgentEngineSessionOperationOrDict from .common import AgentRunConfig from .common import AgentRunConfigDict from .common import AgentRunConfigOrDict +from .common import AgentRuntimeConfig +from .common import AgentRuntimeConfigDict +from .common import AgentRuntimeConfigOrDict from .common import AgentServerMode from .common import AggregatedMetricResult from .common import AggregatedMetricResultDict @@ -221,18 +178,12 @@ from .common import AnalysisConfig from .common import AnalysisConfigDict from .common import AnalysisConfigOrDict -from .common import AppendAgentEngineSessionEventConfig -from .common import AppendAgentEngineSessionEventConfigDict -from .common import AppendAgentEngineSessionEventConfigOrDict -from .common import AppendAgentEngineSessionEventResponse -from .common import AppendAgentEngineSessionEventResponseDict -from .common import AppendAgentEngineSessionEventResponseOrDict -from .common import AppendAgentEngineTaskEventConfig -from .common import AppendAgentEngineTaskEventConfigDict -from .common import AppendAgentEngineTaskEventConfigOrDict -from .common import AppendAgentEngineTaskEventResponse -from .common import AppendAgentEngineTaskEventResponseDict -from .common import AppendAgentEngineTaskEventResponseOrDict +from .common import AppendRuntimeSessionEventConfig +from .common import AppendRuntimeSessionEventConfigDict +from .common import AppendRuntimeSessionEventConfigOrDict +from .common import AppendRuntimeSessionEventResponse +from .common import AppendRuntimeSessionEventResponseDict +from .common import AppendRuntimeSessionEventResponseOrDict from .common import AskContextsConfig from .common import AskContextsConfigDict from .common import AskContextsConfigOrDict @@ -284,26 +235,26 @@ from .common import BleuResults from .common import BleuResultsDict from .common import BleuResultsOrDict -from .common import CancelQueryJobAgentEngineConfig -from .common import CancelQueryJobAgentEngineConfigDict -from .common import CancelQueryJobAgentEngineConfigOrDict from .common import CancelQueryJobResult from .common import CancelQueryJobResultDict from .common import CancelQueryJobResultOrDict +from .common import CancelQueryJobRuntimeConfig +from .common import CancelQueryJobRuntimeConfigDict +from .common import CancelQueryJobRuntimeConfigOrDict from .common import CandidateResponse from .common import CandidateResponseDict from .common import CandidateResponseOrDict from .common import CandidateResult from .common import CandidateResultDict -from .common import CheckQueryJobAgentEngineConfig -from .common import CheckQueryJobAgentEngineConfigDict -from .common import CheckQueryJobAgentEngineConfigOrDict from .common import CheckQueryJobResponse from .common import CheckQueryJobResponseDict from .common import CheckQueryJobResponseOrDict from .common import CheckQueryJobResult from .common import CheckQueryJobResultDict from .common import CheckQueryJobResultOrDict +from .common import CheckQueryJobRuntimeConfig +from .common import CheckQueryJobRuntimeConfigDict +from .common import CheckQueryJobRuntimeConfigOrDict from .common import Chunk from .common import ChunkDict from .common import ChunkOrDict @@ -326,21 +277,6 @@ from .common import CorpusStatus from .common import CorpusStatusDict from .common import CorpusStatusOrDict -from .common import CreateAgentEngineConfig -from .common import CreateAgentEngineConfigDict -from .common import CreateAgentEngineConfigOrDict -from .common import CreateAgentEngineSandboxConfig -from .common import CreateAgentEngineSandboxConfigDict -from .common import CreateAgentEngineSandboxConfigOrDict -from .common import CreateAgentEngineSandboxSnapshotConfig -from .common import CreateAgentEngineSandboxSnapshotConfigDict -from .common import CreateAgentEngineSandboxSnapshotConfigOrDict -from .common import CreateAgentEngineSessionConfig -from .common import CreateAgentEngineSessionConfigDict -from .common import CreateAgentEngineSessionConfigOrDict -from .common import CreateAgentEngineTaskConfig -from .common import CreateAgentEngineTaskConfigDict -from .common import CreateAgentEngineTaskConfigOrDict from .common import CreateDatasetConfig from .common import CreateDatasetConfigDict from .common import CreateDatasetConfigOrDict @@ -377,9 +313,21 @@ from .common import CreateRagCorpusOperation from .common import CreateRagCorpusOperationDict from .common import CreateRagCorpusOperationOrDict +from .common import CreateRuntimeConfig +from .common import CreateRuntimeConfigDict +from .common import CreateRuntimeConfigOrDict from .common import CreateRuntimeFeedbackEntryConfig from .common import CreateRuntimeFeedbackEntryConfigDict from .common import CreateRuntimeFeedbackEntryConfigOrDict +from .common import CreateRuntimeSandboxConfig +from .common import CreateRuntimeSandboxConfigDict +from .common import CreateRuntimeSandboxConfigOrDict +from .common import CreateRuntimeSandboxSnapshotConfig +from .common import CreateRuntimeSandboxSnapshotConfigDict +from .common import CreateRuntimeSandboxSnapshotConfigOrDict +from .common import CreateRuntimeSessionConfig +from .common import CreateRuntimeSessionConfigDict +from .common import CreateRuntimeSessionConfigOrDict from .common import CreateSandboxEnvironmentTemplateConfig from .common import CreateSandboxEnvironmentTemplateConfigDict from .common import CreateSandboxEnvironmentTemplateConfigOrDict @@ -411,39 +359,6 @@ from .common import DedicatedResourcesScaleToZeroSpecDict from .common import DedicatedResourcesScaleToZeroSpecOrDict from .common import DefaultContainerCategory -from .common import DeleteAgentEngineConfig -from .common import DeleteAgentEngineConfigDict -from .common import DeleteAgentEngineConfigOrDict -from .common import DeleteAgentEngineMemoryConfig -from .common import DeleteAgentEngineMemoryConfigDict -from .common import DeleteAgentEngineMemoryConfigOrDict -from .common import DeleteAgentEngineMemoryOperation -from .common import DeleteAgentEngineMemoryOperationDict -from .common import DeleteAgentEngineMemoryOperationOrDict -from .common import DeleteAgentEngineOperation -from .common import DeleteAgentEngineOperationDict -from .common import DeleteAgentEngineOperationOrDict -from .common import DeleteAgentEngineRuntimeRevisionConfig -from .common import DeleteAgentEngineRuntimeRevisionConfigDict -from .common import DeleteAgentEngineRuntimeRevisionConfigOrDict -from .common import DeleteAgentEngineRuntimeRevisionOperation -from .common import DeleteAgentEngineRuntimeRevisionOperationDict -from .common import DeleteAgentEngineRuntimeRevisionOperationOrDict -from .common import DeleteAgentEngineSandboxConfig -from .common import DeleteAgentEngineSandboxConfigDict -from .common import DeleteAgentEngineSandboxConfigOrDict -from .common import DeleteAgentEngineSandboxOperation -from .common import DeleteAgentEngineSandboxOperationDict -from .common import DeleteAgentEngineSandboxOperationOrDict -from .common import DeleteAgentEngineSessionConfig -from .common import DeleteAgentEngineSessionConfigDict -from .common import DeleteAgentEngineSessionConfigOrDict -from .common import DeleteAgentEngineSessionOperation -from .common import DeleteAgentEngineSessionOperationDict -from .common import DeleteAgentEngineSessionOperationOrDict -from .common import DeleteAgentEngineTaskConfig -from .common import DeleteAgentEngineTaskConfigDict -from .common import DeleteAgentEngineTaskConfigOrDict from .common import DeleteEvaluationExperimentConfig from .common import DeleteEvaluationExperimentConfigDict from .common import DeleteEvaluationExperimentConfigOrDict @@ -477,12 +392,42 @@ from .common import DeleteRagFileOperation from .common import DeleteRagFileOperationDict from .common import DeleteRagFileOperationOrDict +from .common import DeleteRuntimeConfig +from .common import DeleteRuntimeConfigDict +from .common import DeleteRuntimeConfigOrDict from .common import DeleteRuntimeFeedbackEntryConfig from .common import DeleteRuntimeFeedbackEntryConfigDict from .common import DeleteRuntimeFeedbackEntryConfigOrDict from .common import DeleteRuntimeFeedbackEntryOperation from .common import DeleteRuntimeFeedbackEntryOperationDict from .common import DeleteRuntimeFeedbackEntryOperationOrDict +from .common import DeleteRuntimeMemoryConfig +from .common import DeleteRuntimeMemoryConfigDict +from .common import DeleteRuntimeMemoryConfigOrDict +from .common import DeleteRuntimeMemoryOperation +from .common import DeleteRuntimeMemoryOperationDict +from .common import DeleteRuntimeMemoryOperationOrDict +from .common import DeleteRuntimeOperation +from .common import DeleteRuntimeOperationDict +from .common import DeleteRuntimeOperationOrDict +from .common import DeleteRuntimeRevisionConfig +from .common import DeleteRuntimeRevisionConfigDict +from .common import DeleteRuntimeRevisionConfigOrDict +from .common import DeleteRuntimeRevisionOperation +from .common import DeleteRuntimeRevisionOperationDict +from .common import DeleteRuntimeRevisionOperationOrDict +from .common import DeleteRuntimeSandboxConfig +from .common import DeleteRuntimeSandboxConfigDict +from .common import DeleteRuntimeSandboxConfigOrDict +from .common import DeleteRuntimeSandboxOperation +from .common import DeleteRuntimeSandboxOperationDict +from .common import DeleteRuntimeSandboxOperationOrDict +from .common import DeleteRuntimeSessionConfig +from .common import DeleteRuntimeSessionConfigDict +from .common import DeleteRuntimeSessionConfigOrDict +from .common import DeleteRuntimeSessionOperation +from .common import DeleteRuntimeSessionOperationDict +from .common import DeleteRuntimeSessionOperationOrDict from .common import DeleteSandboxEnvironmentSnapshotConfig from .common import DeleteSandboxEnvironmentSnapshotConfigDict from .common import DeleteSandboxEnvironmentSnapshotConfigOrDict @@ -626,9 +571,9 @@ from .common import ExactMatchSpec from .common import ExactMatchSpecDict from .common import ExactMatchSpecOrDict -from .common import ExecuteCodeAgentEngineSandboxConfig -from .common import ExecuteCodeAgentEngineSandboxConfigDict -from .common import ExecuteCodeAgentEngineSandboxConfigOrDict +from .common import ExecuteCodeRuntimeSandboxConfig +from .common import ExecuteCodeRuntimeSandboxConfigDict +from .common import ExecuteCodeRuntimeSandboxConfigOrDict from .common import ExecuteSandboxEnvironmentResponse from .common import ExecuteSandboxEnvironmentResponseDict from .common import ExecuteSandboxEnvironmentResponseOrDict @@ -670,9 +615,6 @@ from .common import GeminiTemplateConfig from .common import GeminiTemplateConfigDict from .common import GeminiTemplateConfigOrDict -from .common import GenerateAgentEngineMemoriesConfig -from .common import GenerateAgentEngineMemoriesConfigDict -from .common import GenerateAgentEngineMemoriesConfigOrDict from .common import GenerateInstanceRubricsResponse from .common import GenerateInstanceRubricsResponseDict from .common import GenerateInstanceRubricsResponseOrDict @@ -707,45 +649,24 @@ from .common import GenerateMemoriesResponseGeneratedMemoryDict from .common import GenerateMemoriesResponseGeneratedMemoryOrDict from .common import GenerateMemoriesResponseOrDict +from .common import GenerateRuntimeMemoriesConfig +from .common import GenerateRuntimeMemoriesConfigDict +from .common import GenerateRuntimeMemoriesConfigOrDict from .common import GenerateUserScenariosConfig from .common import GenerateUserScenariosConfigDict from .common import GenerateUserScenariosConfigOrDict from .common import GenerateUserScenariosResponse from .common import GenerateUserScenariosResponseDict from .common import GenerateUserScenariosResponseOrDict -from .common import GetAgentEngineConfig -from .common import GetAgentEngineConfigDict -from .common import GetAgentEngineConfigOrDict -from .common import GetAgentEngineMemoryConfig -from .common import GetAgentEngineMemoryConfigDict -from .common import GetAgentEngineMemoryConfigOrDict -from .common import GetAgentEngineMemoryRevisionConfig -from .common import GetAgentEngineMemoryRevisionConfigDict -from .common import GetAgentEngineMemoryRevisionConfigOrDict -from .common import GetAgentEngineOperationConfig -from .common import GetAgentEngineOperationConfigDict -from .common import GetAgentEngineOperationConfigOrDict -from .common import GetAgentEngineRuntimeRevisionConfig -from .common import GetAgentEngineRuntimeRevisionConfigDict -from .common import GetAgentEngineRuntimeRevisionConfigOrDict -from .common import GetAgentEngineSandboxConfig -from .common import GetAgentEngineSandboxConfigDict -from .common import GetAgentEngineSandboxConfigOrDict -from .common import GetAgentEngineSessionConfig -from .common import GetAgentEngineSessionConfigDict -from .common import GetAgentEngineSessionConfigOrDict -from .common import GetAgentEngineTaskConfig -from .common import GetAgentEngineTaskConfigDict -from .common import GetAgentEngineTaskConfigOrDict from .common import GetCorpusOperationConfig from .common import GetCorpusOperationConfigDict from .common import GetCorpusOperationConfigOrDict from .common import GetDatasetOperationConfig from .common import GetDatasetOperationConfigDict from .common import GetDatasetOperationConfigOrDict -from .common import GetDeleteAgentEngineRuntimeRevisionOperationConfig -from .common import GetDeleteAgentEngineRuntimeRevisionOperationConfigDict -from .common import GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict +from .common import GetDeleteRuntimeRevisionOperationConfig +from .common import GetDeleteRuntimeRevisionOperationConfigDict +from .common import GetDeleteRuntimeRevisionOperationConfigOrDict from .common import GetEvaluationExperimentConfig from .common import GetEvaluationExperimentConfigDict from .common import GetEvaluationExperimentConfigOrDict @@ -788,6 +709,9 @@ from .common import GetRagFileConfig from .common import GetRagFileConfigDict from .common import GetRagFileConfigOrDict +from .common import GetRuntimeConfig +from .common import GetRuntimeConfigDict +from .common import GetRuntimeConfigOrDict from .common import GetRuntimeFeedbackConfig from .common import GetRuntimeFeedbackConfigDict from .common import GetRuntimeFeedbackConfigOrDict @@ -800,6 +724,24 @@ from .common import GetRuntimeFeedbackEntryConfig from .common import GetRuntimeFeedbackEntryConfigDict from .common import GetRuntimeFeedbackEntryConfigOrDict +from .common import GetRuntimeMemoryConfig +from .common import GetRuntimeMemoryConfigDict +from .common import GetRuntimeMemoryConfigOrDict +from .common import GetRuntimeMemoryRevisionConfig +from .common import GetRuntimeMemoryRevisionConfigDict +from .common import GetRuntimeMemoryRevisionConfigOrDict +from .common import GetRuntimeOperationConfig +from .common import GetRuntimeOperationConfigDict +from .common import GetRuntimeOperationConfigOrDict +from .common import GetRuntimeRevisionConfig +from .common import GetRuntimeRevisionConfigDict +from .common import GetRuntimeRevisionConfigOrDict +from .common import GetRuntimeSandboxConfig +from .common import GetRuntimeSandboxConfigDict +from .common import GetRuntimeSandboxConfigOrDict +from .common import GetRuntimeSessionConfig +from .common import GetRuntimeSessionConfigDict +from .common import GetRuntimeSessionConfigOrDict from .common import GetSandboxEnvironmentSnapshotConfig from .common import GetSandboxEnvironmentSnapshotConfigDict from .common import GetSandboxEnvironmentSnapshotConfigOrDict @@ -871,48 +813,6 @@ from .common import LargeModelReferenceDict from .common import LargeModelReferenceOrDict from .common import LaunchStage -from .common import ListAgentEngineConfig -from .common import ListAgentEngineConfigDict -from .common import ListAgentEngineConfigOrDict -from .common import ListAgentEngineMemoryConfig -from .common import ListAgentEngineMemoryConfigDict -from .common import ListAgentEngineMemoryConfigOrDict -from .common import ListAgentEngineMemoryRevisionsConfig -from .common import ListAgentEngineMemoryRevisionsConfigDict -from .common import ListAgentEngineMemoryRevisionsConfigOrDict -from .common import ListAgentEngineMemoryRevisionsResponse -from .common import ListAgentEngineMemoryRevisionsResponseDict -from .common import ListAgentEngineMemoryRevisionsResponseOrDict -from .common import ListAgentEngineRuntimeRevisionsConfig -from .common import ListAgentEngineRuntimeRevisionsConfigDict -from .common import ListAgentEngineRuntimeRevisionsConfigOrDict -from .common import ListAgentEngineSandboxesConfig -from .common import ListAgentEngineSandboxesConfigDict -from .common import ListAgentEngineSandboxesConfigOrDict -from .common import ListAgentEngineSandboxesResponse -from .common import ListAgentEngineSandboxesResponseDict -from .common import ListAgentEngineSandboxesResponseOrDict -from .common import ListAgentEngineSessionEventsConfig -from .common import ListAgentEngineSessionEventsConfigDict -from .common import ListAgentEngineSessionEventsConfigOrDict -from .common import ListAgentEngineSessionEventsResponse -from .common import ListAgentEngineSessionEventsResponseDict -from .common import ListAgentEngineSessionEventsResponseOrDict -from .common import ListAgentEngineSessionsConfig -from .common import ListAgentEngineSessionsConfigDict -from .common import ListAgentEngineSessionsConfigOrDict -from .common import ListAgentEngineTaskEventsConfig -from .common import ListAgentEngineTaskEventsConfigDict -from .common import ListAgentEngineTaskEventsConfigOrDict -from .common import ListAgentEngineTaskEventsResponse -from .common import ListAgentEngineTaskEventsResponseDict -from .common import ListAgentEngineTaskEventsResponseOrDict -from .common import ListAgentEngineTasksConfig -from .common import ListAgentEngineTasksConfigDict -from .common import ListAgentEngineTasksConfigOrDict -from .common import ListAgentEngineTasksResponse -from .common import ListAgentEngineTasksResponseDict -from .common import ListAgentEngineTasksResponseOrDict from .common import ListCustomModelDeployOptionsConfig from .common import ListCustomModelDeployOptionsConfigDict from .common import ListCustomModelDeployOptionsConfigOrDict @@ -982,12 +882,42 @@ from .common import ListReasoningEnginesSessionsResponse from .common import ListReasoningEnginesSessionsResponseDict from .common import ListReasoningEnginesSessionsResponseOrDict +from .common import ListRuntimeConfig +from .common import ListRuntimeConfigDict +from .common import ListRuntimeConfigOrDict from .common import ListRuntimeFeedbackEntriesConfig from .common import ListRuntimeFeedbackEntriesConfigDict from .common import ListRuntimeFeedbackEntriesConfigOrDict from .common import ListRuntimeFeedbackEntriesResponse from .common import ListRuntimeFeedbackEntriesResponseDict from .common import ListRuntimeFeedbackEntriesResponseOrDict +from .common import ListRuntimeMemoryConfig +from .common import ListRuntimeMemoryConfigDict +from .common import ListRuntimeMemoryConfigOrDict +from .common import ListRuntimeMemoryRevisionsConfig +from .common import ListRuntimeMemoryRevisionsConfigDict +from .common import ListRuntimeMemoryRevisionsConfigOrDict +from .common import ListRuntimeMemoryRevisionsResponse +from .common import ListRuntimeMemoryRevisionsResponseDict +from .common import ListRuntimeMemoryRevisionsResponseOrDict +from .common import ListRuntimeRevisionsConfig +from .common import ListRuntimeRevisionsConfigDict +from .common import ListRuntimeRevisionsConfigOrDict +from .common import ListRuntimeSandboxesConfig +from .common import ListRuntimeSandboxesConfigDict +from .common import ListRuntimeSandboxesConfigOrDict +from .common import ListRuntimeSandboxesResponse +from .common import ListRuntimeSandboxesResponseDict +from .common import ListRuntimeSandboxesResponseOrDict +from .common import ListRuntimeSessionEventsConfig +from .common import ListRuntimeSessionEventsConfigDict +from .common import ListRuntimeSessionEventsConfigOrDict +from .common import ListRuntimeSessionEventsResponse +from .common import ListRuntimeSessionEventsResponseDict +from .common import ListRuntimeSessionEventsResponseOrDict +from .common import ListRuntimeSessionsConfig +from .common import ListRuntimeSessionsConfigDict +from .common import ListRuntimeSessionsConfigOrDict from .common import ListSandboxEnvironmentSnapshotsConfig from .common import ListSandboxEnvironmentSnapshotsConfigDict from .common import ListSandboxEnvironmentSnapshotsConfigOrDict @@ -1269,24 +1199,24 @@ from .common import PublisherModelResourceReference from .common import PublisherModelResourceReferenceDict from .common import PublisherModelResourceReferenceOrDict -from .common import PurgeAgentEngineMemoriesConfig -from .common import PurgeAgentEngineMemoriesConfigDict -from .common import PurgeAgentEngineMemoriesConfigOrDict from .common import PurgeMemoriesResponse from .common import PurgeMemoriesResponseDict from .common import PurgeMemoriesResponseOrDict +from .common import PurgeRuntimeMemoriesConfig +from .common import PurgeRuntimeMemoriesConfigDict +from .common import PurgeRuntimeMemoriesConfigOrDict from .common import PythonPackageSpec from .common import PythonPackageSpecDict from .common import PythonPackageSpecOrDict -from .common import QueryAgentEngineConfig -from .common import QueryAgentEngineConfigDict -from .common import QueryAgentEngineConfigOrDict -from .common import QueryAgentEngineRuntimeRevisionConfig -from .common import QueryAgentEngineRuntimeRevisionConfigDict -from .common import QueryAgentEngineRuntimeRevisionConfigOrDict from .common import QueryReasoningEngineResponse from .common import QueryReasoningEngineResponseDict from .common import QueryReasoningEngineResponseOrDict +from .common import QueryRuntimeConfig +from .common import QueryRuntimeConfigDict +from .common import QueryRuntimeConfigOrDict +from .common import QueryRuntimeRevisionConfig +from .common import QueryRuntimeRevisionConfigDict +from .common import QueryRuntimeRevisionConfigOrDict from .common import QuotaState from .common import RagContexts from .common import RagContextsContext @@ -1546,9 +1476,6 @@ from .common import RestoreVersionOperation from .common import RestoreVersionOperationDict from .common import RestoreVersionOperationOrDict -from .common import RetrieveAgentEngineMemoriesConfig -from .common import RetrieveAgentEngineMemoriesConfigDict -from .common import RetrieveAgentEngineMemoriesConfigOrDict from .common import RetrieveContextsConfig from .common import RetrieveContextsConfigDict from .common import RetrieveContextsConfigOrDict @@ -1576,15 +1503,18 @@ from .common import RetrieveProfilesResponse from .common import RetrieveProfilesResponseDict from .common import RetrieveProfilesResponseOrDict +from .common import RetrieveRuntimeMemoriesConfig +from .common import RetrieveRuntimeMemoriesConfigDict +from .common import RetrieveRuntimeMemoriesConfigOrDict from .common import RetrieveSkillsConfig from .common import RetrieveSkillsConfigDict from .common import RetrieveSkillsConfigOrDict from .common import RetrieveSkillsResponse from .common import RetrieveSkillsResponseDict from .common import RetrieveSkillsResponseOrDict -from .common import RollbackAgentEngineMemoryConfig -from .common import RollbackAgentEngineMemoryConfigDict -from .common import RollbackAgentEngineMemoryConfigOrDict +from .common import RollbackRuntimeMemoryConfig +from .common import RollbackRuntimeMemoryConfigDict +from .common import RollbackRuntimeMemoryConfigOrDict from .common import RougeInput from .common import RougeInputDict from .common import RougeInputOrDict @@ -1623,18 +1553,51 @@ from .common import RubricGroupOrDict from .common import RubricVerdict from .common import RubricVerdictDict -from .common import RunQueryJobAgentEngineConfig -from .common import RunQueryJobAgentEngineConfigDict -from .common import RunQueryJobAgentEngineConfigOrDict from .common import RunQueryJobResult from .common import RunQueryJobResultDict from .common import RunQueryJobResultOrDict +from .common import RunQueryJobRuntimeConfig +from .common import RunQueryJobRuntimeConfigDict +from .common import RunQueryJobRuntimeConfigOrDict +from .common import Runtime +from .common import RuntimeDict from .common import RuntimeFeedbackContextOperation from .common import RuntimeFeedbackContextOperationDict from .common import RuntimeFeedbackContextOperationOrDict from .common import RuntimeFeedbackEntryOperation from .common import RuntimeFeedbackEntryOperationDict from .common import RuntimeFeedbackEntryOperationOrDict +from .common import RuntimeGenerateMemoriesOperation +from .common import RuntimeGenerateMemoriesOperationDict +from .common import RuntimeGenerateMemoriesOperationOrDict +from .common import RuntimeMemoryConfig +from .common import RuntimeMemoryConfigDict +from .common import RuntimeMemoryConfigOrDict +from .common import RuntimeMemoryOperation +from .common import RuntimeMemoryOperationDict +from .common import RuntimeMemoryOperationOrDict +from .common import RuntimeOperation +from .common import RuntimeOperationDict +from .common import RuntimeOperationOrDict +from .common import RuntimeOrDict +from .common import RuntimePurgeMemoriesOperation +from .common import RuntimePurgeMemoriesOperationDict +from .common import RuntimePurgeMemoriesOperationOrDict +from .common import RuntimeRevision +from .common import RuntimeRevisionDict +from .common import RuntimeRevisionOrDict +from .common import RuntimeRollbackMemoryOperation +from .common import RuntimeRollbackMemoryOperationDict +from .common import RuntimeRollbackMemoryOperationOrDict +from .common import RuntimeSandboxOperation +from .common import RuntimeSandboxOperationDict +from .common import RuntimeSandboxOperationOrDict +from .common import RuntimeSandboxSnapshotOperation +from .common import RuntimeSandboxSnapshotOperationDict +from .common import RuntimeSandboxSnapshotOperationOrDict +from .common import RuntimeSessionOperation +from .common import RuntimeSessionOperationDict +from .common import RuntimeSessionOperationOrDict from .common import SamplingConfig from .common import SamplingConfigDict from .common import SamplingConfigOrDict @@ -1810,39 +1773,6 @@ from .common import SummaryMetric from .common import SummaryMetricDict from .common import SummaryMetricOrDict -from .common import TaskArtifact -from .common import TaskArtifactChange -from .common import TaskArtifactChangeDict -from .common import TaskArtifactChangeOrDict -from .common import TaskArtifactDict -from .common import TaskArtifactOrDict -from .common import TaskEvent -from .common import TaskEventData -from .common import TaskEventDataDict -from .common import TaskEventDataOrDict -from .common import TaskEventDict -from .common import TaskEventOrDict -from .common import TaskMessage -from .common import TaskMessageDict -from .common import TaskMessageOrDict -from .common import TaskMetadataChange -from .common import TaskMetadataChangeDict -from .common import TaskMetadataChangeOrDict -from .common import TaskOutput -from .common import TaskOutputChange -from .common import TaskOutputChangeDict -from .common import TaskOutputChangeOrDict -from .common import TaskOutputDict -from .common import TaskOutputOrDict -from .common import TaskStateChange -from .common import TaskStateChangeDict -from .common import TaskStateChangeOrDict -from .common import TaskStatusDetails -from .common import TaskStatusDetailsChange -from .common import TaskStatusDetailsChangeDict -from .common import TaskStatusDetailsChangeOrDict -from .common import TaskStatusDetailsDict -from .common import TaskStatusDetailsOrDict from .common import ToolCallValidInput from .common import ToolCallValidInputDict from .common import ToolCallValidInputOrDict @@ -1919,15 +1849,6 @@ from .common import UnifiedMetric from .common import UnifiedMetricDict from .common import UnifiedMetricOrDict -from .common import UpdateAgentEngineConfig -from .common import UpdateAgentEngineConfigDict -from .common import UpdateAgentEngineConfigOrDict -from .common import UpdateAgentEngineMemoryConfig -from .common import UpdateAgentEngineMemoryConfigDict -from .common import UpdateAgentEngineMemoryConfigOrDict -from .common import UpdateAgentEngineSessionConfig -from .common import UpdateAgentEngineSessionConfigDict -from .common import UpdateAgentEngineSessionConfigOrDict from .common import UpdateEvaluationExperimentConfig from .common import UpdateEvaluationExperimentConfigDict from .common import UpdateEvaluationExperimentConfigOrDict @@ -1946,12 +1867,21 @@ from .common import UpdateRagCorpusOperation from .common import UpdateRagCorpusOperationDict from .common import UpdateRagCorpusOperationOrDict +from .common import UpdateRuntimeConfig +from .common import UpdateRuntimeConfigDict +from .common import UpdateRuntimeConfigOrDict from .common import UpdateRuntimeFeedbackContextConfig from .common import UpdateRuntimeFeedbackContextConfigDict from .common import UpdateRuntimeFeedbackContextConfigOrDict from .common import UpdateRuntimeFeedbackEntryConfig from .common import UpdateRuntimeFeedbackEntryConfigDict from .common import UpdateRuntimeFeedbackEntryConfigOrDict +from .common import UpdateRuntimeMemoryConfig +from .common import UpdateRuntimeMemoryConfigDict +from .common import UpdateRuntimeMemoryConfigOrDict +from .common import UpdateRuntimeSessionConfig +from .common import UpdateRuntimeSessionConfigDict +from .common import UpdateRuntimeSessionConfigOrDict from .common import UpdateSkillConfig from .common import UpdateSkillConfigDict from .common import UpdateSkillConfigOrDict @@ -1982,69 +1912,6 @@ from .common import WorkerPoolSpecOrDict __all__ = [ - "DeleteAgentEngineTaskConfig", - "DeleteAgentEngineTaskConfigDict", - "DeleteAgentEngineTaskConfigOrDict", - "GetAgentEngineTaskConfig", - "GetAgentEngineTaskConfigDict", - "GetAgentEngineTaskConfigOrDict", - "TaskArtifact", - "TaskArtifactDict", - "TaskArtifactOrDict", - "TaskOutput", - "TaskOutputDict", - "TaskOutputOrDict", - "TaskMessage", - "TaskMessageDict", - "TaskMessageOrDict", - "TaskStatusDetails", - "TaskStatusDetailsDict", - "TaskStatusDetailsOrDict", - "A2aTask", - "A2aTaskDict", - "A2aTaskOrDict", - "ListAgentEngineTasksConfig", - "ListAgentEngineTasksConfigDict", - "ListAgentEngineTasksConfigOrDict", - "ListAgentEngineTasksResponse", - "ListAgentEngineTasksResponseDict", - "ListAgentEngineTasksResponseOrDict", - "CreateAgentEngineTaskConfig", - "CreateAgentEngineTaskConfigDict", - "CreateAgentEngineTaskConfigOrDict", - "TaskMetadataChange", - "TaskMetadataChangeDict", - "TaskMetadataChangeOrDict", - "TaskArtifactChange", - "TaskArtifactChangeDict", - "TaskArtifactChangeOrDict", - "TaskOutputChange", - "TaskOutputChangeDict", - "TaskOutputChangeOrDict", - "TaskStateChange", - "TaskStateChangeDict", - "TaskStateChangeOrDict", - "TaskStatusDetailsChange", - "TaskStatusDetailsChangeDict", - "TaskStatusDetailsChangeOrDict", - "TaskEventData", - "TaskEventDataDict", - "TaskEventDataOrDict", - "TaskEvent", - "TaskEventDict", - "TaskEventOrDict", - "AppendAgentEngineTaskEventConfig", - "AppendAgentEngineTaskEventConfigDict", - "AppendAgentEngineTaskEventConfigOrDict", - "AppendAgentEngineTaskEventResponse", - "AppendAgentEngineTaskEventResponseDict", - "AppendAgentEngineTaskEventResponseOrDict", - "ListAgentEngineTaskEventsConfig", - "ListAgentEngineTaskEventsConfigDict", - "ListAgentEngineTaskEventsConfigOrDict", - "ListAgentEngineTaskEventsResponse", - "ListAgentEngineTaskEventsResponseDict", - "ListAgentEngineTaskEventsResponseOrDict", "CreateEvaluationExperimentConfig", "CreateEvaluationExperimentConfigDict", "CreateEvaluationExperimentConfigOrDict", @@ -2459,21 +2326,21 @@ "VertexBaseConfig", "VertexBaseConfigDict", "VertexBaseConfigOrDict", - "CancelQueryJobAgentEngineConfig", - "CancelQueryJobAgentEngineConfigDict", - "CancelQueryJobAgentEngineConfigOrDict", + "CancelQueryJobRuntimeConfig", + "CancelQueryJobRuntimeConfigDict", + "CancelQueryJobRuntimeConfigOrDict", "CancelQueryJobResult", "CancelQueryJobResultDict", "CancelQueryJobResultOrDict", - "CheckQueryJobAgentEngineConfig", - "CheckQueryJobAgentEngineConfigDict", - "CheckQueryJobAgentEngineConfigOrDict", + "CheckQueryJobRuntimeConfig", + "CheckQueryJobRuntimeConfigDict", + "CheckQueryJobRuntimeConfigOrDict", "CheckQueryJobResult", "CheckQueryJobResultDict", "CheckQueryJobResultOrDict", - "_RunQueryJobAgentEngineConfig", - "_RunQueryJobAgentEngineConfigDict", - "_RunQueryJobAgentEngineConfigOrDict", + "_RunQueryJobRuntimeConfig", + "_RunQueryJobRuntimeConfigDict", + "_RunQueryJobRuntimeConfigOrDict", "MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent", "MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict", "MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventOrDict", @@ -2606,60 +2473,84 @@ "ReasoningEngine", "ReasoningEngineDict", "ReasoningEngineOrDict", - "AgentEngineOperation", - "AgentEngineOperationDict", - "AgentEngineOperationOrDict", - "CreateAgentEngineConfig", - "CreateAgentEngineConfigDict", - "CreateAgentEngineConfigOrDict", - "DeleteAgentEngineConfig", - "DeleteAgentEngineConfigDict", - "DeleteAgentEngineConfigOrDict", - "DeleteAgentEngineOperation", - "DeleteAgentEngineOperationDict", - "DeleteAgentEngineOperationOrDict", - "GetAgentEngineConfig", - "GetAgentEngineConfigDict", - "GetAgentEngineConfigOrDict", - "ListAgentEngineConfig", - "ListAgentEngineConfigDict", - "ListAgentEngineConfigOrDict", + "RuntimeOperation", + "RuntimeOperationDict", + "RuntimeOperationOrDict", + "CreateRuntimeConfig", + "CreateRuntimeConfigDict", + "CreateRuntimeConfigOrDict", + "DeleteRuntimeConfig", + "DeleteRuntimeConfigDict", + "DeleteRuntimeConfigOrDict", + "DeleteRuntimeOperation", + "DeleteRuntimeOperationDict", + "DeleteRuntimeOperationOrDict", + "GetRuntimeConfig", + "GetRuntimeConfigDict", + "GetRuntimeConfigOrDict", + "ListRuntimeConfig", + "ListRuntimeConfigDict", + "ListRuntimeConfigOrDict", "ListReasoningEnginesResponse", "ListReasoningEnginesResponseDict", "ListReasoningEnginesResponseOrDict", - "GetAgentEngineOperationConfig", - "GetAgentEngineOperationConfigDict", - "GetAgentEngineOperationConfigOrDict", - "QueryAgentEngineConfig", - "QueryAgentEngineConfigDict", - "QueryAgentEngineConfigOrDict", + "GetRuntimeOperationConfig", + "GetRuntimeOperationConfigDict", + "GetRuntimeOperationConfigOrDict", + "QueryRuntimeConfig", + "QueryRuntimeConfigDict", + "QueryRuntimeConfigOrDict", "QueryReasoningEngineResponse", "QueryReasoningEngineResponseDict", "QueryReasoningEngineResponseOrDict", - "UpdateAgentEngineConfig", - "UpdateAgentEngineConfigDict", - "UpdateAgentEngineConfigOrDict", + "UpdateRuntimeConfig", + "UpdateRuntimeConfigDict", + "UpdateRuntimeConfigOrDict", + "GetRuntimeRevisionConfig", + "GetRuntimeRevisionConfigDict", + "GetRuntimeRevisionConfigOrDict", + "ReasoningEngineRuntimeRevision", + "ReasoningEngineRuntimeRevisionDict", + "ReasoningEngineRuntimeRevisionOrDict", + "ListRuntimeRevisionsConfig", + "ListRuntimeRevisionsConfigDict", + "ListRuntimeRevisionsConfigOrDict", + "ListReasoningEnginesRuntimeRevisionsResponse", + "ListReasoningEnginesRuntimeRevisionsResponseDict", + "ListReasoningEnginesRuntimeRevisionsResponseOrDict", + "DeleteRuntimeRevisionConfig", + "DeleteRuntimeRevisionConfigDict", + "DeleteRuntimeRevisionConfigOrDict", + "DeleteRuntimeRevisionOperation", + "DeleteRuntimeRevisionOperationDict", + "DeleteRuntimeRevisionOperationOrDict", + "GetDeleteRuntimeRevisionOperationConfig", + "GetDeleteRuntimeRevisionOperationConfigDict", + "GetDeleteRuntimeRevisionOperationConfigOrDict", + "QueryRuntimeRevisionConfig", + "QueryRuntimeRevisionConfigDict", + "QueryRuntimeRevisionConfigOrDict", "MemoryMetadataValue", "MemoryMetadataValueDict", "MemoryMetadataValueOrDict", - "AgentEngineMemoryConfig", - "AgentEngineMemoryConfigDict", - "AgentEngineMemoryConfigOrDict", + "RuntimeMemoryConfig", + "RuntimeMemoryConfigDict", + "RuntimeMemoryConfigOrDict", "MemoryStructuredContent", "MemoryStructuredContentDict", "MemoryStructuredContentOrDict", "Memory", "MemoryDict", "MemoryOrDict", - "AgentEngineMemoryOperation", - "AgentEngineMemoryOperationDict", - "AgentEngineMemoryOperationOrDict", - "DeleteAgentEngineMemoryConfig", - "DeleteAgentEngineMemoryConfigDict", - "DeleteAgentEngineMemoryConfigOrDict", - "DeleteAgentEngineMemoryOperation", - "DeleteAgentEngineMemoryOperationDict", - "DeleteAgentEngineMemoryOperationOrDict", + "RuntimeMemoryOperation", + "RuntimeMemoryOperationDict", + "RuntimeMemoryOperationOrDict", + "DeleteRuntimeMemoryConfig", + "DeleteRuntimeMemoryConfigDict", + "DeleteRuntimeMemoryConfigOrDict", + "DeleteRuntimeMemoryOperation", + "DeleteRuntimeMemoryOperationDict", + "DeleteRuntimeMemoryOperationOrDict", "GenerateMemoriesRequestVertexSessionSource", "GenerateMemoriesRequestVertexSessionSourceDict", "GenerateMemoriesRequestVertexSessionSourceOrDict", @@ -2675,21 +2566,21 @@ "GenerateMemoriesRequestDirectMemoriesSource", "GenerateMemoriesRequestDirectMemoriesSourceDict", "GenerateMemoriesRequestDirectMemoriesSourceOrDict", - "GenerateAgentEngineMemoriesConfig", - "GenerateAgentEngineMemoriesConfigDict", - "GenerateAgentEngineMemoriesConfigOrDict", + "GenerateRuntimeMemoriesConfig", + "GenerateRuntimeMemoriesConfigDict", + "GenerateRuntimeMemoriesConfigOrDict", "GenerateMemoriesResponseGeneratedMemory", "GenerateMemoriesResponseGeneratedMemoryDict", "GenerateMemoriesResponseGeneratedMemoryOrDict", "GenerateMemoriesResponse", "GenerateMemoriesResponseDict", "GenerateMemoriesResponseOrDict", - "AgentEngineGenerateMemoriesOperation", - "AgentEngineGenerateMemoriesOperationDict", - "AgentEngineGenerateMemoriesOperationOrDict", - "GetAgentEngineMemoryConfig", - "GetAgentEngineMemoryConfigDict", - "GetAgentEngineMemoryConfigOrDict", + "RuntimeGenerateMemoriesOperation", + "RuntimeGenerateMemoriesOperationDict", + "RuntimeGenerateMemoriesOperationOrDict", + "GetRuntimeMemoryConfig", + "GetRuntimeMemoryConfigDict", + "GetRuntimeMemoryConfigOrDict", "IngestionDirectContentsSourceEvent", "IngestionDirectContentsSourceEventDict", "IngestionDirectContentsSourceEventOrDict", @@ -2702,9 +2593,9 @@ "MemoryBankIngestEventsOperation", "MemoryBankIngestEventsOperationDict", "MemoryBankIngestEventsOperationOrDict", - "ListAgentEngineMemoryConfig", - "ListAgentEngineMemoryConfigDict", - "ListAgentEngineMemoryConfigOrDict", + "ListRuntimeMemoryConfig", + "ListRuntimeMemoryConfigDict", + "ListRuntimeMemoryConfigOrDict", "ListReasoningEnginesMemoriesResponse", "ListReasoningEnginesMemoriesResponseDict", "ListReasoningEnginesMemoriesResponseOrDict", @@ -2720,9 +2611,9 @@ "MemoryConjunctionFilter", "MemoryConjunctionFilterDict", "MemoryConjunctionFilterOrDict", - "RetrieveAgentEngineMemoriesConfig", - "RetrieveAgentEngineMemoriesConfigDict", - "RetrieveAgentEngineMemoriesConfigOrDict", + "RetrieveRuntimeMemoriesConfig", + "RetrieveRuntimeMemoriesConfigDict", + "RetrieveRuntimeMemoriesConfigOrDict", "RetrieveMemoriesResponseRetrievedMemory", "RetrieveMemoriesResponseRetrievedMemoryDict", "RetrieveMemoriesResponseRetrievedMemoryOrDict", @@ -2738,39 +2629,39 @@ "RetrieveProfilesResponse", "RetrieveProfilesResponseDict", "RetrieveProfilesResponseOrDict", - "RollbackAgentEngineMemoryConfig", - "RollbackAgentEngineMemoryConfigDict", - "RollbackAgentEngineMemoryConfigOrDict", - "AgentEngineRollbackMemoryOperation", - "AgentEngineRollbackMemoryOperationDict", - "AgentEngineRollbackMemoryOperationOrDict", - "UpdateAgentEngineMemoryConfig", - "UpdateAgentEngineMemoryConfigDict", - "UpdateAgentEngineMemoryConfigOrDict", - "PurgeAgentEngineMemoriesConfig", - "PurgeAgentEngineMemoriesConfigDict", - "PurgeAgentEngineMemoriesConfigOrDict", + "RollbackRuntimeMemoryConfig", + "RollbackRuntimeMemoryConfigDict", + "RollbackRuntimeMemoryConfigOrDict", + "RuntimeRollbackMemoryOperation", + "RuntimeRollbackMemoryOperationDict", + "RuntimeRollbackMemoryOperationOrDict", + "UpdateRuntimeMemoryConfig", + "UpdateRuntimeMemoryConfigDict", + "UpdateRuntimeMemoryConfigOrDict", + "PurgeRuntimeMemoriesConfig", + "PurgeRuntimeMemoriesConfigDict", + "PurgeRuntimeMemoriesConfigOrDict", "PurgeMemoriesResponse", "PurgeMemoriesResponseDict", "PurgeMemoriesResponseOrDict", - "AgentEnginePurgeMemoriesOperation", - "AgentEnginePurgeMemoriesOperationDict", - "AgentEnginePurgeMemoriesOperationOrDict", - "GetAgentEngineMemoryRevisionConfig", - "GetAgentEngineMemoryRevisionConfigDict", - "GetAgentEngineMemoryRevisionConfigOrDict", + "RuntimePurgeMemoriesOperation", + "RuntimePurgeMemoriesOperationDict", + "RuntimePurgeMemoriesOperationOrDict", + "GetRuntimeMemoryRevisionConfig", + "GetRuntimeMemoryRevisionConfigDict", + "GetRuntimeMemoryRevisionConfigOrDict", "IntermediateExtractedMemory", "IntermediateExtractedMemoryDict", "IntermediateExtractedMemoryOrDict", "MemoryRevision", "MemoryRevisionDict", "MemoryRevisionOrDict", - "ListAgentEngineMemoryRevisionsConfig", - "ListAgentEngineMemoryRevisionsConfigDict", - "ListAgentEngineMemoryRevisionsConfigOrDict", - "ListAgentEngineMemoryRevisionsResponse", - "ListAgentEngineMemoryRevisionsResponseDict", - "ListAgentEngineMemoryRevisionsResponseOrDict", + "ListRuntimeMemoryRevisionsConfig", + "ListRuntimeMemoryRevisionsConfigDict", + "ListRuntimeMemoryRevisionsConfigOrDict", + "ListRuntimeMemoryRevisionsResponse", + "ListRuntimeMemoryRevisionsResponseDict", + "ListRuntimeMemoryRevisionsResponseOrDict", "AskContextsConfig", "AskContextsConfigDict", "AskContextsConfigOrDict", @@ -3032,30 +2923,6 @@ "UploadRagFileResponse", "UploadRagFileResponseDict", "UploadRagFileResponseOrDict", - "GetAgentEngineRuntimeRevisionConfig", - "GetAgentEngineRuntimeRevisionConfigDict", - "GetAgentEngineRuntimeRevisionConfigOrDict", - "ReasoningEngineRuntimeRevision", - "ReasoningEngineRuntimeRevisionDict", - "ReasoningEngineRuntimeRevisionOrDict", - "ListAgentEngineRuntimeRevisionsConfig", - "ListAgentEngineRuntimeRevisionsConfigDict", - "ListAgentEngineRuntimeRevisionsConfigOrDict", - "ListReasoningEnginesRuntimeRevisionsResponse", - "ListReasoningEnginesRuntimeRevisionsResponseDict", - "ListReasoningEnginesRuntimeRevisionsResponseOrDict", - "DeleteAgentEngineRuntimeRevisionConfig", - "DeleteAgentEngineRuntimeRevisionConfigDict", - "DeleteAgentEngineRuntimeRevisionConfigOrDict", - "DeleteAgentEngineRuntimeRevisionOperation", - "DeleteAgentEngineRuntimeRevisionOperationDict", - "DeleteAgentEngineRuntimeRevisionOperationOrDict", - "GetDeleteAgentEngineRuntimeRevisionOperationConfig", - "GetDeleteAgentEngineRuntimeRevisionOperationConfigDict", - "GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict", - "QueryAgentEngineRuntimeRevisionConfig", - "QueryAgentEngineRuntimeRevisionConfigDict", - "QueryAgentEngineRuntimeRevisionConfigOrDict", "SandboxEnvironmentSpecCodeExecutionEnvironment", "SandboxEnvironmentSpecCodeExecutionEnvironmentDict", "SandboxEnvironmentSpecCodeExecutionEnvironmentOrDict", @@ -3065,45 +2932,45 @@ "SandboxEnvironmentSpec", "SandboxEnvironmentSpecDict", "SandboxEnvironmentSpecOrDict", - "CreateAgentEngineSandboxConfig", - "CreateAgentEngineSandboxConfigDict", - "CreateAgentEngineSandboxConfigOrDict", + "CreateRuntimeSandboxConfig", + "CreateRuntimeSandboxConfigDict", + "CreateRuntimeSandboxConfigOrDict", "SandboxEnvironmentConnectionInfo", "SandboxEnvironmentConnectionInfoDict", "SandboxEnvironmentConnectionInfoOrDict", "SandboxEnvironment", "SandboxEnvironmentDict", "SandboxEnvironmentOrDict", - "AgentEngineSandboxOperation", - "AgentEngineSandboxOperationDict", - "AgentEngineSandboxOperationOrDict", - "DeleteAgentEngineSandboxConfig", - "DeleteAgentEngineSandboxConfigDict", - "DeleteAgentEngineSandboxConfigOrDict", - "DeleteAgentEngineSandboxOperation", - "DeleteAgentEngineSandboxOperationDict", - "DeleteAgentEngineSandboxOperationOrDict", + "RuntimeSandboxOperation", + "RuntimeSandboxOperationDict", + "RuntimeSandboxOperationOrDict", + "DeleteRuntimeSandboxConfig", + "DeleteRuntimeSandboxConfigDict", + "DeleteRuntimeSandboxConfigOrDict", + "DeleteRuntimeSandboxOperation", + "DeleteRuntimeSandboxOperationDict", + "DeleteRuntimeSandboxOperationOrDict", "Metadata", "MetadataDict", "MetadataOrDict", "Chunk", "ChunkDict", "ChunkOrDict", - "ExecuteCodeAgentEngineSandboxConfig", - "ExecuteCodeAgentEngineSandboxConfigDict", - "ExecuteCodeAgentEngineSandboxConfigOrDict", + "ExecuteCodeRuntimeSandboxConfig", + "ExecuteCodeRuntimeSandboxConfigDict", + "ExecuteCodeRuntimeSandboxConfigOrDict", "ExecuteSandboxEnvironmentResponse", "ExecuteSandboxEnvironmentResponseDict", "ExecuteSandboxEnvironmentResponseOrDict", - "GetAgentEngineSandboxConfig", - "GetAgentEngineSandboxConfigDict", - "GetAgentEngineSandboxConfigOrDict", - "ListAgentEngineSandboxesConfig", - "ListAgentEngineSandboxesConfigDict", - "ListAgentEngineSandboxesConfigOrDict", - "ListAgentEngineSandboxesResponse", - "ListAgentEngineSandboxesResponseDict", - "ListAgentEngineSandboxesResponseOrDict", + "GetRuntimeSandboxConfig", + "GetRuntimeSandboxConfigDict", + "GetRuntimeSandboxConfigOrDict", + "ListRuntimeSandboxesConfig", + "ListRuntimeSandboxesConfigDict", + "ListRuntimeSandboxesConfigOrDict", + "ListRuntimeSandboxesResponse", + "ListRuntimeSandboxesResponseDict", + "ListRuntimeSandboxesResponseOrDict", "SandboxEnvironmentTemplateCustomContainerSpec", "SandboxEnvironmentTemplateCustomContainerSpecDict", "SandboxEnvironmentTemplateCustomContainerSpecOrDict", @@ -3146,15 +3013,15 @@ "ListSandboxEnvironmentTemplatesResponse", "ListSandboxEnvironmentTemplatesResponseDict", "ListSandboxEnvironmentTemplatesResponseOrDict", - "CreateAgentEngineSandboxSnapshotConfig", - "CreateAgentEngineSandboxSnapshotConfigDict", - "CreateAgentEngineSandboxSnapshotConfigOrDict", + "CreateRuntimeSandboxSnapshotConfig", + "CreateRuntimeSandboxSnapshotConfigDict", + "CreateRuntimeSandboxSnapshotConfigOrDict", "SandboxEnvironmentSnapshot", "SandboxEnvironmentSnapshotDict", "SandboxEnvironmentSnapshotOrDict", - "AgentEngineSandboxSnapshotOperation", - "AgentEngineSandboxSnapshotOperationDict", - "AgentEngineSandboxSnapshotOperationOrDict", + "RuntimeSandboxSnapshotOperation", + "RuntimeSandboxSnapshotOperationDict", + "RuntimeSandboxSnapshotOperationOrDict", "DeleteSandboxEnvironmentSnapshotConfig", "DeleteSandboxEnvironmentSnapshotConfigDict", "DeleteSandboxEnvironmentSnapshotConfigOrDict", @@ -3170,54 +3037,54 @@ "ListSandboxEnvironmentSnapshotsResponse", "ListSandboxEnvironmentSnapshotsResponseDict", "ListSandboxEnvironmentSnapshotsResponseOrDict", - "CreateAgentEngineSessionConfig", - "CreateAgentEngineSessionConfigDict", - "CreateAgentEngineSessionConfigOrDict", + "CreateRuntimeSessionConfig", + "CreateRuntimeSessionConfigDict", + "CreateRuntimeSessionConfigOrDict", "Session", "SessionDict", "SessionOrDict", - "AgentEngineSessionOperation", - "AgentEngineSessionOperationDict", - "AgentEngineSessionOperationOrDict", - "DeleteAgentEngineSessionConfig", - "DeleteAgentEngineSessionConfigDict", - "DeleteAgentEngineSessionConfigOrDict", - "DeleteAgentEngineSessionOperation", - "DeleteAgentEngineSessionOperationDict", - "DeleteAgentEngineSessionOperationOrDict", - "GetAgentEngineSessionConfig", - "GetAgentEngineSessionConfigDict", - "GetAgentEngineSessionConfigOrDict", - "ListAgentEngineSessionsConfig", - "ListAgentEngineSessionsConfigDict", - "ListAgentEngineSessionsConfigOrDict", + "RuntimeSessionOperation", + "RuntimeSessionOperationDict", + "RuntimeSessionOperationOrDict", + "DeleteRuntimeSessionConfig", + "DeleteRuntimeSessionConfigDict", + "DeleteRuntimeSessionConfigOrDict", + "DeleteRuntimeSessionOperation", + "DeleteRuntimeSessionOperationDict", + "DeleteRuntimeSessionOperationOrDict", + "GetRuntimeSessionConfig", + "GetRuntimeSessionConfigDict", + "GetRuntimeSessionConfigOrDict", + "ListRuntimeSessionsConfig", + "ListRuntimeSessionsConfigDict", + "ListRuntimeSessionsConfigOrDict", "ListReasoningEnginesSessionsResponse", "ListReasoningEnginesSessionsResponseDict", "ListReasoningEnginesSessionsResponseOrDict", - "UpdateAgentEngineSessionConfig", - "UpdateAgentEngineSessionConfigDict", - "UpdateAgentEngineSessionConfigOrDict", + "UpdateRuntimeSessionConfig", + "UpdateRuntimeSessionConfigDict", + "UpdateRuntimeSessionConfigOrDict", "EventActions", "EventActionsDict", "EventActionsOrDict", "EventMetadata", "EventMetadataDict", "EventMetadataOrDict", - "AppendAgentEngineSessionEventConfig", - "AppendAgentEngineSessionEventConfigDict", - "AppendAgentEngineSessionEventConfigOrDict", - "AppendAgentEngineSessionEventResponse", - "AppendAgentEngineSessionEventResponseDict", - "AppendAgentEngineSessionEventResponseOrDict", - "ListAgentEngineSessionEventsConfig", - "ListAgentEngineSessionEventsConfigDict", - "ListAgentEngineSessionEventsConfigOrDict", + "AppendRuntimeSessionEventConfig", + "AppendRuntimeSessionEventConfigDict", + "AppendRuntimeSessionEventConfigOrDict", + "AppendRuntimeSessionEventResponse", + "AppendRuntimeSessionEventResponseDict", + "AppendRuntimeSessionEventResponseOrDict", + "ListRuntimeSessionEventsConfig", + "ListRuntimeSessionEventsConfigDict", + "ListRuntimeSessionEventsConfigOrDict", "SessionEvent", "SessionEventDict", "SessionEventOrDict", - "ListAgentEngineSessionEventsResponse", - "ListAgentEngineSessionEventsResponseDict", - "ListAgentEngineSessionEventsResponseOrDict", + "ListRuntimeSessionEventsResponse", + "ListRuntimeSessionEventsResponseDict", + "ListRuntimeSessionEventsResponseOrDict", "GeminiExample", "GeminiExampleDict", "GeminiExampleOrDict", @@ -3626,21 +3493,6 @@ "EvalRunInferenceConfig", "EvalRunInferenceConfigDict", "EvalRunInferenceConfigOrDict", - "AgentEngine", - "AgentEngineDict", - "AgentEngineOrDict", - "AgentEngineConfig", - "AgentEngineConfigDict", - "AgentEngineConfigOrDict", - "RunQueryJobAgentEngineConfig", - "RunQueryJobAgentEngineConfigDict", - "RunQueryJobAgentEngineConfigOrDict", - "RunQueryJobResult", - "RunQueryJobResultDict", - "RunQueryJobResultOrDict", - "CheckQueryJobResponse", - "CheckQueryJobResponseDict", - "CheckQueryJobResponseOrDict", "AssembleDataset", "AssembleDatasetDict", "AssembleDatasetOrDict", @@ -3680,9 +3532,6 @@ "OptimizeJobConfig", "OptimizeJobConfigDict", "OptimizeJobConfigOrDict", - "AgentEngineRuntimeRevision", - "AgentEngineRuntimeRevisionDict", - "AgentEngineRuntimeRevisionOrDict", "ListDeployableModelsConfig", "ListDeployableModelsConfigDict", "ListDeployableModelsConfigOrDict", @@ -3701,8 +3550,24 @@ "DeployOption", "DeployOptionDict", "DeployOptionOrDict", - "A2aTaskState", - "State", + "Runtime", + "RuntimeDict", + "RuntimeOrDict", + "AgentRuntimeConfig", + "AgentRuntimeConfigDict", + "AgentRuntimeConfigOrDict", + "RunQueryJobRuntimeConfig", + "RunQueryJobRuntimeConfigDict", + "RunQueryJobRuntimeConfigOrDict", + "RunQueryJobResult", + "RunQueryJobResultDict", + "RunQueryJobResultOrDict", + "CheckQueryJobResponse", + "CheckQueryJobResponseDict", + "CheckQueryJobResponseOrDict", + "RuntimeRevision", + "RuntimeRevisionDict", + "RuntimeRevisionOrDict", "Strategy", "AcceleratorType", "Type", @@ -3710,6 +3575,7 @@ "ManagedTopicEnum", "IdentityType", "AgentServerMode", + "State", "MemoryType", "Operator", "RagFileType", @@ -3761,12 +3627,6 @@ "MessageDict", "Importance", "ParsedResponseUnion", - "_DeleteAgentEngineTaskRequestParameters", - "_GetAgentEngineTaskRequestParameters", - "_ListAgentEngineTasksRequestParameters", - "_CreateAgentEngineTaskRequestParameters", - "_AppendAgentEngineTaskEventRequestParameters", - "_ListAgentEngineTaskEventsRequestParameters", "_CreateEvaluationExperimentParameters", "_CreateEvaluationItemParameters", "_CreateEvaluationMetricParameters", @@ -3789,31 +3649,36 @@ "_OptimizeRequestParameters", "_CustomJobParameters", "_GetCustomJobParameters", - "_CancelQueryJobAgentEngineRequestParameters", - "_CheckQueryJobAgentEngineRequestParameters", - "_RunQueryJobAgentEngineRequestParameters", - "_CreateAgentEngineRequestParameters", - "_DeleteAgentEngineRequestParameters", - "_GetAgentEngineRequestParameters", - "_ListAgentEngineRequestParameters", - "_GetAgentEngineOperationParameters", - "_QueryAgentEngineRequestParameters", - "_UpdateAgentEngineRequestParameters", - "_CreateAgentEngineMemoryRequestParameters", - "_DeleteAgentEngineMemoryRequestParameters", - "_GenerateAgentEngineMemoriesRequestParameters", - "_GetAgentEngineMemoryRequestParameters", + "_CancelQueryJobRuntimeRequestParameters", + "_CheckQueryJobRuntimeRequestParameters", + "_RunQueryJobRuntimeRequestParameters", + "_CreateRuntimeRequestParameters", + "_DeleteRuntimeRequestParameters", + "_GetRuntimeRequestParameters", + "_ListRuntimeRequestParameters", + "_GetRuntimeOperationParameters", + "_QueryRuntimeRequestParameters", + "_UpdateRuntimeRequestParameters", + "_GetRuntimeRevisionRequestParameters", + "_ListRuntimeRevisionsRequestParameters", + "_DeleteRuntimeRevisionRequestParameters", + "_GetDeleteRuntimeRevisionOperationParameters", + "_QueryRuntimeRevisionRequestParameters", + "_CreateRuntimeMemoryRequestParameters", + "_DeleteRuntimeMemoryRequestParameters", + "_GenerateRuntimeMemoriesRequestParameters", + "_GetRuntimeMemoryRequestParameters", "_IngestEventsRequestParameters", - "_ListAgentEngineMemoryRequestParameters", - "_GetAgentEngineMemoryOperationParameters", - "_GetAgentEngineGenerateMemoriesOperationParameters", - "_RetrieveAgentEngineMemoriesRequestParameters", + "_ListRuntimeMemoryRequestParameters", + "_GetRuntimeMemoryOperationParameters", + "_GetRuntimeGenerateMemoriesOperationParameters", + "_RetrieveRuntimeMemoriesRequestParameters", "_RetrieveMemoryProfilesRequestParameters", - "_RollbackAgentEngineMemoryRequestParameters", - "_UpdateAgentEngineMemoryRequestParameters", - "_PurgeAgentEngineMemoriesRequestParameters", - "_GetAgentEngineMemoryRevisionRequestParameters", - "_ListAgentEngineMemoryRevisionsRequestParameters", + "_RollbackRuntimeMemoryRequestParameters", + "_UpdateRuntimeMemoryRequestParameters", + "_PurgeRuntimeMemoriesRequestParameters", + "_GetRuntimeMemoryRevisionRequestParameters", + "_ListRuntimeMemoryRevisionsRequestParameters", "_AskContextsRequestParameters", "_CreateRagCorpusRequestParameters", "_GetCorpusOperationParameters", @@ -3831,17 +3696,12 @@ "_ImportRagFilesRequestParameters", "_GetImportFilesOperationParameters", "_UploadRagFileParameters", - "_GetAgentEngineRuntimeRevisionRequestParameters", - "_ListAgentEngineRuntimeRevisionsRequestParameters", - "_DeleteAgentEngineRuntimeRevisionRequestParameters", - "_GetDeleteAgentEngineRuntimeRevisionOperationParameters", - "_QueryAgentEngineRuntimeRevisionRequestParameters", - "_CreateAgentEngineSandboxRequestParameters", - "_DeleteAgentEngineSandboxRequestParameters", - "_ExecuteCodeAgentEngineSandboxRequestParameters", - "_GetAgentEngineSandboxRequestParameters", - "_ListAgentEngineSandboxesRequestParameters", - "_GetAgentEngineSandboxOperationParameters", + "_CreateRuntimeSandboxRequestParameters", + "_DeleteRuntimeSandboxRequestParameters", + "_ExecuteCodeRuntimeSandboxRequestParameters", + "_GetRuntimeSandboxRequestParameters", + "_ListRuntimeSandboxesRequestParameters", + "_GetRuntimeSandboxOperationParameters", "_CreateSandboxEnvironmentTemplateRequestParameters", "_DeleteSandboxEnvironmentTemplateRequestParameters", "_GetSandboxEnvironmentTemplateRequestParameters", @@ -3851,15 +3711,15 @@ "_DeleteSandboxEnvironmentSnapshotRequestParameters", "_GetSandboxEnvironmentSnapshotRequestParameters", "_ListSandboxEnvironmentSnapshotsRequestParameters", - "_GetAgentEngineSandboxSnapshotOperationParameters", - "_CreateAgentEngineSessionRequestParameters", - "_DeleteAgentEngineSessionRequestParameters", - "_GetAgentEngineSessionRequestParameters", - "_ListAgentEngineSessionsRequestParameters", - "_GetAgentEngineSessionOperationParameters", - "_UpdateAgentEngineSessionRequestParameters", - "_AppendAgentEngineSessionEventRequestParameters", - "_ListAgentEngineSessionEventsRequestParameters", + "_GetRuntimeSandboxSnapshotOperationParameters", + "_CreateRuntimeSessionRequestParameters", + "_DeleteRuntimeSessionRequestParameters", + "_GetRuntimeSessionRequestParameters", + "_ListRuntimeSessionsRequestParameters", + "_GetRuntimeSessionOperationParameters", + "_UpdateRuntimeSessionRequestParameters", + "_AppendRuntimeSessionEventRequestParameters", + "_ListRuntimeSessionEventsRequestParameters", "_AssembleDatasetParameters", "_AssessDatasetParameters", "_CreateMultimodalDatasetParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 24265d7dab..114fbe43de 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -91,56 +91,6 @@ def _camel_key_to_snake(message: Any) -> Any: MetricSubclass = TypeVar("MetricSubclass", bound="Metric") -class A2aTaskState(_common.CaseInSensitiveEnum): - """Output only. The state of the task. The state of a new task is SUBMITTED by default. The state of a task can only be updated via AppendA2aTaskEvents API.""" - - STATE_UNSPECIFIED = "STATE_UNSPECIFIED" - """Task state unspecified. Default value if not set.""" - SUBMITTED = "SUBMITTED" - """Task is submitted and waiting to be processed.""" - WORKING = "WORKING" - """Task is actively being processed.""" - COMPLETED = "COMPLETED" - """Task is finished.""" - CANCELLED = "CANCELLED" - """Task is cancelled.""" - FAILED = "FAILED" - """Task has failed.""" - REJECTED = "REJECTED" - """Task is rejected by the system.""" - INPUT_REQUIRED = "INPUT_REQUIRED" - """Task requires input from the user.""" - AUTH_REQUIRED = "AUTH_REQUIRED" - """Task requires auth (e.g. OAuth) from the user.""" - PAUSED = "PAUSED" - """Task is paused.""" - - -class State(_common.CaseInSensitiveEnum): - """The new state of the task.""" - - STATE_UNSPECIFIED = "STATE_UNSPECIFIED" - """Task state unspecified. Default value if not set.""" - SUBMITTED = "SUBMITTED" - """Task is submitted and waiting to be processed.""" - WORKING = "WORKING" - """Task is actively being processed.""" - COMPLETED = "COMPLETED" - """Task is finished.""" - CANCELLED = "CANCELLED" - """Task is cancelled.""" - FAILED = "FAILED" - """Task has failed.""" - REJECTED = "REJECTED" - """Task is rejected by the system.""" - INPUT_REQUIRED = "INPUT_REQUIRED" - """Task requires input from the user.""" - AUTH_REQUIRED = "AUTH_REQUIRED" - """Task requires auth (e.g. OAuth) from the user.""" - PAUSED = "PAUSED" - """Task is paused.""" - - class Strategy(_common.CaseInSensitiveEnum): """This determines which type of scheduling strategy to use.""" @@ -280,6 +230,17 @@ class AgentServerMode(_common.CaseInSensitiveEnum): """Experimental agent server mode. This mode contains experimental features.""" +class State(_common.CaseInSensitiveEnum): + """Output only. The state of the revision.""" + + STATE_UNSPECIFIED = "STATE_UNSPECIFIED" + """The unspecified state.""" + ACTIVE = "ACTIVE" + """Is deployed and ready to be used.""" + DEPRECATED = "DEPRECATED" + """Is deprecated, may not be used, only preserved for historical purposes.""" + + class MemoryType(_common.CaseInSensitiveEnum): """The type of the memory.""" @@ -633,3695 +594,3229 @@ class OptimizationMethod(_common.CaseInSensitiveEnum): """The data driven prompt optimizer designer for prompts from Android core API.""" -class DeleteAgentEngineTaskConfig(_common.BaseModel): - """Config for deleting an Agent Engine Task.""" +class CreateEvaluationExperimentConfig(_common.BaseModel): + """Config to create an evaluation experiment.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class DeleteAgentEngineTaskConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Task.""" +class CreateEvaluationExperimentConfigDict(TypedDict, total=False): + """Config to create an evaluation experiment.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteAgentEngineTaskConfigOrDict = Union[ - DeleteAgentEngineTaskConfig, DeleteAgentEngineTaskConfigDict +CreateEvaluationExperimentConfigOrDict = Union[ + CreateEvaluationExperimentConfig, CreateEvaluationExperimentConfigDict ] -class _DeleteAgentEngineTaskRequestParameters(_common.BaseModel): - """Parameters for deleting an agent engine task.""" +class _CreateEvaluationExperimentParameters(_common.BaseModel): + """Parameters for creating an evaluation experiment.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine task.""" + display_name: Optional[str] = Field(default=None, description="""""") + labels: Optional[dict[str, str]] = Field(default=None, description="""""") + merge_strategy: Optional[EvaluationExperimentMergeStrategy] = Field( + default=None, description="""""" ) - config: Optional[DeleteAgentEngineTaskConfig] = Field( + metadata: Optional[dict[str, Any]] = Field(default=None, description="""""") + config: Optional[CreateEvaluationExperimentConfig] = Field( default=None, description="""""" ) -class _DeleteAgentEngineTaskRequestParametersDict(TypedDict, total=False): - """Parameters for deleting an agent engine task.""" - - name: Optional[str] - """Name of the agent engine task.""" +class _CreateEvaluationExperimentParametersDict(TypedDict, total=False): + """Parameters for creating an evaluation experiment.""" - config: Optional[DeleteAgentEngineTaskConfigDict] + display_name: Optional[str] """""" + labels: Optional[dict[str, str]] + """""" -_DeleteAgentEngineTaskRequestParametersOrDict = Union[ - _DeleteAgentEngineTaskRequestParameters, _DeleteAgentEngineTaskRequestParametersDict -] - - -class GetAgentEngineTaskConfig(_common.BaseModel): - """Config for getting an Agent Engine Task.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - - -class GetAgentEngineTaskConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Task.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -GetAgentEngineTaskConfigOrDict = Union[ - GetAgentEngineTaskConfig, GetAgentEngineTaskConfigDict -] - - -class _GetAgentEngineTaskRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine task.""" - - name: Optional[str] = Field( - default=None, description="""Name of the agent engine task.""" - ) - config: Optional[GetAgentEngineTaskConfig] = Field(default=None, description="""""") - - -class _GetAgentEngineTaskRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine task.""" + merge_strategy: Optional[EvaluationExperimentMergeStrategy] + """""" - name: Optional[str] - """Name of the agent engine task.""" + metadata: Optional[dict[str, Any]] + """""" - config: Optional[GetAgentEngineTaskConfigDict] + config: Optional[CreateEvaluationExperimentConfigDict] """""" -_GetAgentEngineTaskRequestParametersOrDict = Union[ - _GetAgentEngineTaskRequestParameters, _GetAgentEngineTaskRequestParametersDict +_CreateEvaluationExperimentParametersOrDict = Union[ + _CreateEvaluationExperimentParameters, _CreateEvaluationExperimentParametersDict ] -class TaskArtifact(_common.BaseModel): - """The artifact of the task event.""" +class EvaluationExperiment(_common.BaseModel): + """Represents an experiment for iterating on and visualizing evaluation runs.""" - artifact_id: Optional[str] = Field( - default=None, - description="""Required. The unique identifier of the artifact within the task. This id is provided by the creator of the artifact.""", - ) - description: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Optional. A human readable description of the artifact.""", + description="""The resource name of the EvaluationExperiment. Format: + `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""", ) display_name: Optional[str] = Field( + default=None, description="""The display name of the evaluation experiment.""" + ) + evaluation_runs: Optional[list[str]] = Field( default=None, - description="""Optional. The human-readable name of the artifact provided by the creator.""", + description="""The EvaluationRuns that are part of this experiment.""", + ) + labels: Optional[dict[str, str]] = Field( + default=None, description="""Labels for the evaluation experiment.""" + ) + merge_strategy: Optional[EvaluationExperimentMergeStrategy] = Field( + default=None, description="""Merge strategy for the evaluation experiment.""" ) metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. Additional metadata for the artifact. For A2A, the URIs of the extensions that were used to produce this artifact will be stored here.""", + description="""Metadata about the evaluation experiment, can be used by the caller + to store additional tracking information about the experiment.""", ) - parts: Optional[list[genai_types.Part]] = Field( - default=None, description="""The parts of the artifact.""" + create_time: Optional[datetime.datetime] = Field( + default=None, description="""Timestamp when this experiment was created.""" + ) + update_time: Optional[datetime.datetime] = Field( + default=None, description="""Timestamp when this experiment was last updated.""" ) -class TaskArtifactDict(TypedDict, total=False): - """The artifact of the task event.""" - - artifact_id: Optional[str] - """Required. The unique identifier of the artifact within the task. This id is provided by the creator of the artifact.""" +class EvaluationExperimentDict(TypedDict, total=False): + """Represents an experiment for iterating on and visualizing evaluation runs.""" - description: Optional[str] - """Optional. A human readable description of the artifact.""" + name: Optional[str] + """The resource name of the EvaluationExperiment. Format: + `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""" display_name: Optional[str] - """Optional. The human-readable name of the artifact provided by the creator.""" + """The display name of the evaluation experiment.""" + + evaluation_runs: Optional[list[str]] + """The EvaluationRuns that are part of this experiment.""" + + labels: Optional[dict[str, str]] + """Labels for the evaluation experiment.""" + + merge_strategy: Optional[EvaluationExperimentMergeStrategy] + """Merge strategy for the evaluation experiment.""" metadata: Optional[dict[str, Any]] - """Optional. Additional metadata for the artifact. For A2A, the URIs of the extensions that were used to produce this artifact will be stored here.""" + """Metadata about the evaluation experiment, can be used by the caller + to store additional tracking information about the experiment.""" - parts: Optional[list[genai_types.Part]] - """The parts of the artifact.""" + create_time: Optional[datetime.datetime] + """Timestamp when this experiment was created.""" + update_time: Optional[datetime.datetime] + """Timestamp when this experiment was last updated.""" -TaskArtifactOrDict = Union[TaskArtifact, TaskArtifactDict] + +EvaluationExperimentOrDict = Union[EvaluationExperiment, EvaluationExperimentDict] -class TaskOutput(_common.BaseModel): - """The output of the task event.""" +class CreateEvaluationItemConfig(_common.BaseModel): + """Config to create an evaluation item.""" - artifacts: Optional[list[TaskArtifact]] = Field( - default=None, description="""The artifacts of the task event.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class TaskOutputDict(TypedDict, total=False): - """The output of the task event.""" +class CreateEvaluationItemConfigDict(TypedDict, total=False): + """Config to create an evaluation item.""" - artifacts: Optional[list[TaskArtifactDict]] - """The artifacts of the task event.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -TaskOutputOrDict = Union[TaskOutput, TaskOutputDict] +CreateEvaluationItemConfigOrDict = Union[ + CreateEvaluationItemConfig, CreateEvaluationItemConfigDict +] -class TaskMessage(_common.BaseModel): - """The message of the task event.""" +class _CreateEvaluationItemParameters(_common.BaseModel): + """Represents a job that creates an evaluation item.""" - message_id: Optional[str] = Field( - default=None, description="""Required. The unique identifier of the message.""" - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, - description="""Optional. A2A message may have extension_uris or reference_task_ids. They will be stored under metadata.""", - ) - parts: Optional[list[genai_types.Part]] = Field( - default=None, description="""The parts of the message.""" - ) - role: Optional[str] = Field( - default=None, - description="""Required. The role of the sender of the message. e.g. "user", "agent".""", + evaluation_item_type: Optional[str] = Field(default=None, description="""""") + gcs_uri: Optional[str] = Field(default=None, description="""""") + display_name: Optional[str] = Field(default=None, description="""""") + config: Optional[CreateEvaluationItemConfig] = Field( + default=None, description="""""" ) -class TaskMessageDict(TypedDict, total=False): - """The message of the task event.""" +class _CreateEvaluationItemParametersDict(TypedDict, total=False): + """Represents a job that creates an evaluation item.""" + + evaluation_item_type: Optional[str] + """""" - message_id: Optional[str] - """Required. The unique identifier of the message.""" + gcs_uri: Optional[str] + """""" - metadata: Optional[dict[str, Any]] - """Optional. A2A message may have extension_uris or reference_task_ids. They will be stored under metadata.""" + display_name: Optional[str] + """""" - parts: Optional[list[genai_types.Part]] - """The parts of the message.""" + config: Optional[CreateEvaluationItemConfigDict] + """""" - role: Optional[str] - """Required. The role of the sender of the message. e.g. "user", "agent".""" +_CreateEvaluationItemParametersOrDict = Union[ + _CreateEvaluationItemParameters, _CreateEvaluationItemParametersDict +] -TaskMessageOrDict = Union[TaskMessage, TaskMessageDict] +class PromptTemplateData(_common.BaseModel): + """Holds data for a prompt template. -class TaskStatusDetails(_common.BaseModel): - """The status details of the task event.""" + Message to hold a prompt template and the values to populate the template. + """ - task_message: Optional[TaskMessage] = Field( - default=None, description="""The status of the task event.""" + values: Optional[dict[str, genai_types.Content]] = Field( + default=None, description="""The values for fields in the prompt template.""" ) -class TaskStatusDetailsDict(TypedDict, total=False): - """The status details of the task event.""" +class PromptTemplateDataDict(TypedDict, total=False): + """Holds data for a prompt template. + + Message to hold a prompt template and the values to populate the template. + """ - task_message: Optional[TaskMessageDict] - """The status of the task event.""" + values: Optional[dict[str, genai_types.Content]] + """The values for fields in the prompt template.""" -TaskStatusDetailsOrDict = Union[TaskStatusDetails, TaskStatusDetailsDict] +PromptTemplateDataOrDict = Union[PromptTemplateData, PromptTemplateDataDict] -class A2aTask(_common.BaseModel): - """A task.""" +class EvaluationPrompt(_common.BaseModel): + """Represents the prompt to be evaluated.""" - context_id: Optional[str] = Field( + text: Optional[str] = Field(default=None, description="""Text prompt.""") + value: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. A generic identifier for grouping related tasks (e.g., session_id, workflow_id).""", - ) - create_time: Optional[datetime.datetime] = Field( - default=None, description="""Output only. The creation timestamp of the task.""" - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, description="""Optional. Arbitrary, user-defined metadata.""" + description="""Fields and values that can be used to populate the prompt template.""", ) - name: Optional[str] = Field( - default=None, - description="""Identifier. The resource name of the task. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/a2aTasks/{a2a_task}`""", + prompt_template_data: Optional[PromptTemplateData] = Field( + default=None, description="""Prompt template data.""" ) - next_event_sequence_number: Optional[int] = Field( + user_scenario: Optional[evals_types.UserScenario] = Field( default=None, - description="""Output only. The next event sequence number to be appended to the task. This value starts at 1 and is guaranteed to be monotonically increasing.""", - ) - output: Optional[TaskOutput] = Field( - default=None, description="""Optional. The final output of the task.""" + description="""User scenario to help simulate multi-turn agent run results.""", ) - state: Optional[A2aTaskState] = Field( + + +class EvaluationPromptDict(TypedDict, total=False): + """Represents the prompt to be evaluated.""" + + text: Optional[str] + """Text prompt.""" + + value: Optional[dict[str, Any]] + """Fields and values that can be used to populate the prompt template.""" + + prompt_template_data: Optional[PromptTemplateDataDict] + """Prompt template data.""" + + user_scenario: Optional[evals_types.UserScenario] + """User scenario to help simulate multi-turn agent run results.""" + + +EvaluationPromptOrDict = Union[EvaluationPrompt, EvaluationPromptDict] + + +class CandidateResponse(_common.BaseModel): + """Responses from model or agent.""" + + candidate: Optional[str] = Field( default=None, - description="""Output only. The state of the task. The state of a new task is SUBMITTED by default. The state of a task can only be updated via AppendA2aTaskEvents API.""", - ) - status_details: Optional[TaskStatusDetails] = Field( - default=None, description="""Optional. The status details of the task.""" + description="""The name of the candidate that produced the response.""", ) - update_time: Optional[datetime.datetime] = Field( + text: Optional[str] = Field(default=None, description="""The text response.""") + value: Optional[dict[str, Any]] = Field( default=None, - description="""Output only. The last update timestamp of the task.""", + description="""Fields and values that can be used to populate the response template.""", ) - expire_time: Optional[datetime.datetime] = Field( + events: Optional[list[genai_types.Content]] = Field( default=None, - description="""Optional. Timestamp of when this task is considered expired. This is *always* provided on output, and is calculated based on the `ttl` if set on the request""", + description="""Intermediate events (such as tool calls and responses) that led to the final response.""", ) - ttl: Optional[str] = Field( + agent_data: Optional[evals_types.AgentData] = Field( default=None, - description="""Optional. Input only. The TTL (Time To Live) for the task. If not set, the task will expire in 24 hours by default. Valid range: (0 seconds, 1000 days]""", + description="""Represents the complete execution trace of an agent conversation, + which can involve single or multiple agents. This field is used to + provide the full output of an agent's run, including all turns and + events, for direct evaluation.""", ) -class A2aTaskDict(TypedDict, total=False): - """A task.""" - - context_id: Optional[str] - """Optional. A generic identifier for grouping related tasks (e.g., session_id, workflow_id).""" +class CandidateResponseDict(TypedDict, total=False): + """Responses from model or agent.""" - create_time: Optional[datetime.datetime] - """Output only. The creation timestamp of the task.""" + candidate: Optional[str] + """The name of the candidate that produced the response.""" - metadata: Optional[dict[str, Any]] - """Optional. Arbitrary, user-defined metadata.""" + text: Optional[str] + """The text response.""" - name: Optional[str] - """Identifier. The resource name of the task. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/a2aTasks/{a2a_task}`""" + value: Optional[dict[str, Any]] + """Fields and values that can be used to populate the response template.""" - next_event_sequence_number: Optional[int] - """Output only. The next event sequence number to be appended to the task. This value starts at 1 and is guaranteed to be monotonically increasing.""" + events: Optional[list[genai_types.Content]] + """Intermediate events (such as tool calls and responses) that led to the final response.""" - output: Optional[TaskOutputDict] - """Optional. The final output of the task.""" + agent_data: Optional[evals_types.AgentData] + """Represents the complete execution trace of an agent conversation, + which can involve single or multiple agents. This field is used to + provide the full output of an agent's run, including all turns and + events, for direct evaluation.""" - state: Optional[A2aTaskState] - """Output only. The state of the task. The state of a new task is SUBMITTED by default. The state of a task can only be updated via AppendA2aTaskEvents API.""" - status_details: Optional[TaskStatusDetailsDict] - """Optional. The status details of the task.""" +CandidateResponseOrDict = Union[CandidateResponse, CandidateResponseDict] - update_time: Optional[datetime.datetime] - """Output only. The last update timestamp of the task.""" - expire_time: Optional[datetime.datetime] - """Optional. Timestamp of when this task is considered expired. This is *always* provided on output, and is calculated based on the `ttl` if set on the request""" - - ttl: Optional[str] - """Optional. Input only. The TTL (Time To Live) for the task. If not set, the task will expire in 24 hours by default. Valid range: (0 seconds, 1000 days]""" - - -A2aTaskOrDict = Union[A2aTask, A2aTaskDict] - - -class ListAgentEngineTasksConfig(_common.BaseModel): - """Config for listing agent engine tasks.""" +class EvaluationItemRequest(_common.BaseModel): + """Single evaluation request.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + prompt: Optional[EvaluationPrompt] = Field( + default=None, description="""The request/prompt to evaluate.""" ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( + golden_response: Optional[CandidateResponse] = Field( + default=None, description="""The ideal response or ground truth.""" + ) + rubrics: Optional[dict[str, "RubricGroup"]] = Field( default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""", + description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", ) - order_by: Optional[str] = Field( + candidate_responses: Optional[list[CandidateResponse]] = Field( default=None, - description="""A comma-separated list of fields to order by, sorted in ascending order. - Use "desc" after a field name for descending. - If this field is omitted, the default ordering is `create_time` descending. - More detail in [AIP-132](https://google.aip.dev/132). - - Supported fields: - * `create_time` - * `update_time` - - Example: `create_time desc`.""", + description="""Responses from model under test and other baseline models for comparison.""", ) -class ListAgentEngineTasksConfigDict(TypedDict, total=False): - """Config for listing agent engine tasks.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - page_size: Optional[int] - """""" - - page_token: Optional[str] - """""" +class EvaluationItemRequestDict(TypedDict, total=False): + """Single evaluation request.""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" + prompt: Optional[EvaluationPromptDict] + """The request/prompt to evaluate.""" - order_by: Optional[str] - """A comma-separated list of fields to order by, sorted in ascending order. - Use "desc" after a field name for descending. - If this field is omitted, the default ordering is `create_time` descending. - More detail in [AIP-132](https://google.aip.dev/132). + golden_response: Optional[CandidateResponseDict] + """The ideal response or ground truth.""" - Supported fields: - * `create_time` - * `update_time` + rubrics: Optional[dict[str, "RubricGroupDict"]] + """Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""" - Example: `create_time desc`.""" + candidate_responses: Optional[list[CandidateResponseDict]] + """Responses from model under test and other baseline models for comparison.""" -ListAgentEngineTasksConfigOrDict = Union[ - ListAgentEngineTasksConfig, ListAgentEngineTasksConfigDict -] +EvaluationItemRequestOrDict = Union[EvaluationItemRequest, EvaluationItemRequestDict] -class _ListAgentEngineTasksRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" +class EvaluationItemResult(_common.BaseModel): + """Represents the result of an evaluation item.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + evaluation_request: Optional[str] = Field( + default=None, description="""The request item that was evaluated.""" ) - config: Optional[ListAgentEngineTasksConfig] = Field( - default=None, description="""""" + evaluation_run: Optional[str] = Field( + default=None, + description="""The evaluation run that was used to generate the result.""", + ) + request: Optional[EvaluationItemRequest] = Field( + default=None, description="""The request that was evaluated.""" + ) + metric: Optional[str] = Field( + default=None, description="""The metric that was evaluated.""" + ) + candidate_results: Optional[list[evals_types.CandidateResult]] = Field( + default=None, description="""The results for the metric.""" + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, description="""Metadata about the evaluation result.""" ) -class _ListAgentEngineTasksRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" - - name: Optional[str] - """Name of the agent engine.""" - - config: Optional[ListAgentEngineTasksConfigDict] - """""" - - -_ListAgentEngineTasksRequestParametersOrDict = Union[ - _ListAgentEngineTasksRequestParameters, _ListAgentEngineTasksRequestParametersDict -] - +class EvaluationItemResultDict(TypedDict, total=False): + """Represents the result of an evaluation item.""" -class ListAgentEngineTasksResponse(_common.BaseModel): - """Response for listing agent engine tasks.""" + evaluation_request: Optional[str] + """The request item that was evaluated.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" - ) - next_page_token: Optional[str] = Field(default=None, description="""""") - a2aTasks: Optional[list[A2aTask]] = Field( - default=None, description="""List of agent engine tasks.""" - ) + evaluation_run: Optional[str] + """The evaluation run that was used to generate the result.""" + request: Optional[EvaluationItemRequestDict] + """The request that was evaluated.""" -class ListAgentEngineTasksResponseDict(TypedDict, total=False): - """Response for listing agent engine tasks.""" + metric: Optional[str] + """The metric that was evaluated.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + candidate_results: Optional[list[evals_types.CandidateResult]] + """The results for the metric.""" - next_page_token: Optional[str] - """""" + metadata: Optional[dict[str, Any]] + """Metadata about the evaluation result.""" - a2aTasks: Optional[list[A2aTaskDict]] - """List of agent engine tasks.""" +EvaluationItemResultOrDict = Union[EvaluationItemResult, EvaluationItemResultDict] -ListAgentEngineTasksResponseOrDict = Union[ - ListAgentEngineTasksResponse, ListAgentEngineTasksResponseDict -] +class EvaluationItem(_common.BaseModel): + """EvaluationItem is a single evaluation request or result. -class CreateAgentEngineTaskConfig(_common.BaseModel): - """Config for creating a Session.""" + The content of an EvaluationItem is immutable - it cannot be updated once + created. EvaluationItems can be deleted when no longer needed. + """ - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + name: Optional[str] = Field( + default=None, description="""The resource name of the EvaluationItem.""" ) - context_id: Optional[str] = Field( - default=None, description="""The context id of the task to create.""" + display_name: Optional[str] = Field( + default=None, description="""The display name of the EvaluationItem.""" ) metadata: Optional[dict[str, Any]] = Field( - default=None, description="""The metadata of the task to create.""" + default=None, description="""Metadata for the EvaluationItem.""" + ) + labels: Optional[dict[str, str]] = Field( + default=None, description="""Labels for the EvaluationItem.""" ) - status_details: Optional[TaskStatusDetails] = Field( - default=None, description="""The status details of the task to create.""" + evaluation_item_type: Optional[EvaluationItemType] = Field( + default=None, description="""The type of the EvaluationItem.""" ) - output: Optional[TaskOutput] = Field( - default=None, description="""The output of the task to create.""" + evaluation_request: Optional[EvaluationItemRequest] = Field( + default=None, description="""The request to evaluate.""" ) - - -class CreateAgentEngineTaskConfigDict(TypedDict, total=False): - """Config for creating a Session.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - context_id: Optional[str] - """The context id of the task to create.""" - - metadata: Optional[dict[str, Any]] - """The metadata of the task to create.""" - - status_details: Optional[TaskStatusDetailsDict] - """The status details of the task to create.""" - - output: Optional[TaskOutputDict] - """The output of the task to create.""" - - -CreateAgentEngineTaskConfigOrDict = Union[ - CreateAgentEngineTaskConfig, CreateAgentEngineTaskConfigDict -] - - -class _CreateAgentEngineTaskRequestParameters(_common.BaseModel): - """Parameters for creating Agent Engine Tasks.""" - - name: Optional[str] = Field( + evaluation_response: Optional[EvaluationItemResult] = Field( + default=None, description="""The response from evaluation.""" + ) + gcs_uri: Optional[str] = Field( default=None, - description="""Name of the agent engine to create the task under.""", + description="""The Cloud Storage object where the request or response is stored.""", ) - a2a_task_id: Optional[str] = Field( - default=None, description="""The ID of the task.""" + create_time: Optional[datetime.datetime] = Field( + default=None, description="""Timestamp when this item was created.""" ) - config: Optional[CreateAgentEngineTaskConfig] = Field( - default=None, description="""""" + error: Optional[genai_types.GoogleRpcStatus] = Field( + default=None, description="""Error for the evaluation item.""" ) + # TODO(b/448806531): Remove all the overridden _from_response methods once the + # ticket is resolved and published. + @classmethod + def _from_response( + cls: typing.Type["EvaluationItem"], + *, + response: dict[str, object], + kwargs: dict[str, object], + ) -> "EvaluationItem": + """Converts a dictionary response into a EvaluationItem object.""" -class _CreateAgentEngineTaskRequestParametersDict(TypedDict, total=False): - """Parameters for creating Agent Engine Tasks.""" + response = _camel_key_to_snake(response) + result = super()._from_response(response=response, kwargs=kwargs) + return result - name: Optional[str] - """Name of the agent engine to create the task under.""" - a2a_task_id: Optional[str] - """The ID of the task.""" +class EvaluationItemDict(TypedDict, total=False): + """EvaluationItem is a single evaluation request or result. - config: Optional[CreateAgentEngineTaskConfigDict] - """""" + The content of an EvaluationItem is immutable - it cannot be updated once + created. EvaluationItems can be deleted when no longer needed. + """ + name: Optional[str] + """The resource name of the EvaluationItem.""" -_CreateAgentEngineTaskRequestParametersOrDict = Union[ - _CreateAgentEngineTaskRequestParameters, _CreateAgentEngineTaskRequestParametersDict -] + display_name: Optional[str] + """The display name of the EvaluationItem.""" + metadata: Optional[dict[str, Any]] + """Metadata for the EvaluationItem.""" -class TaskMetadataChange(_common.BaseModel): - """An event representing a change to the task's top-level metadata. example: metadata_change: { new_metadata: { "name": "My task", } update_mask: { paths: "name" } }""" + labels: Optional[dict[str, str]] + """Labels for the EvaluationItem.""" - new_metadata: Optional[dict[str, Any]] = Field( - default=None, - description="""Required. The complete state of the metadata object *after* the change.""", - ) - update_mask: Optional[str] = Field( - default=None, - description="""Optional. A field mask indicating which paths in the Struct were changed. If not set, all fields will be updated. go/aip-internal/cloud-standard/2412""", - ) + evaluation_item_type: Optional[EvaluationItemType] + """The type of the EvaluationItem.""" + + evaluation_request: Optional[EvaluationItemRequestDict] + """The request to evaluate.""" + evaluation_response: Optional[EvaluationItemResultDict] + """The response from evaluation.""" -class TaskMetadataChangeDict(TypedDict, total=False): - """An event representing a change to the task's top-level metadata. example: metadata_change: { new_metadata: { "name": "My task", } update_mask: { paths: "name" } }""" + gcs_uri: Optional[str] + """The Cloud Storage object where the request or response is stored.""" - new_metadata: Optional[dict[str, Any]] - """Required. The complete state of the metadata object *after* the change.""" + create_time: Optional[datetime.datetime] + """Timestamp when this item was created.""" - update_mask: Optional[str] - """Optional. A field mask indicating which paths in the Struct were changed. If not set, all fields will be updated. go/aip-internal/cloud-standard/2412""" + error: Optional[genai_types.GoogleRpcStatus] + """Error for the evaluation item.""" -TaskMetadataChangeOrDict = Union[TaskMetadataChange, TaskMetadataChangeDict] +EvaluationItemOrDict = Union[EvaluationItem, EvaluationItemDict] -class TaskArtifactChange(_common.BaseModel): - """Describes changes to the artifact list.""" +class Metric(_common.BaseModel): + """The metric used for evaluation.""" - added_artifacts: Optional[list[TaskArtifact]] = Field( + name: Optional[str] = Field(default=None, description="""The name of the metric.""") + custom_function: Optional[Union[str, Callable[..., Any]]] = Field( default=None, - description="""Optional. A list of brand-new artifacts created in this event.""", + description="""The custom function that defines the end-to-end logic for metric computation.""", ) - deleted_artifact_ids: Optional[list[str]] = Field( - default=None, - description="""Optional. A list of artifact IDs that were removed in this event.""", + prompt_template: Optional[str] = Field( + default=None, description="""The prompt template for the metric.""" + ) + judge_model_system_instruction: Optional[str] = Field( + default=None, description="""The system instruction for the judge model.""" ) - updated_artifacts: Optional[list[TaskArtifact]] = Field( + return_raw_output: Optional[bool] = Field( default=None, - description="""Optional. A list of existing artifacts that were modified in this event.""", + description="""Whether to return the raw output from the judge model.""", ) - - -class TaskArtifactChangeDict(TypedDict, total=False): - """Describes changes to the artifact list.""" - - added_artifacts: Optional[list[TaskArtifactDict]] - """Optional. A list of brand-new artifacts created in this event.""" - - deleted_artifact_ids: Optional[list[str]] - """Optional. A list of artifact IDs that were removed in this event.""" - - updated_artifacts: Optional[list[TaskArtifactDict]] - """Optional. A list of existing artifacts that were modified in this event.""" - - -TaskArtifactChangeOrDict = Union[TaskArtifactChange, TaskArtifactChangeDict] - - -class TaskOutputChange(_common.BaseModel): - """An event representing a change to the task's outputs.""" - - task_artifact_change: Optional[TaskArtifactChange] = Field( + parse_and_reduce_fn: Optional[Callable[..., Any]] = Field( default=None, - description="""Required. A granular change to the list of artifacts.""", + description="""The parse and reduce function for the judge model.""", ) - - -class TaskOutputChangeDict(TypedDict, total=False): - """An event representing a change to the task's outputs.""" - - task_artifact_change: Optional[TaskArtifactChangeDict] - """Required. A granular change to the list of artifacts.""" - - -TaskOutputChangeOrDict = Union[TaskOutputChange, TaskOutputChangeDict] - - -class TaskStateChange(_common.BaseModel): - """A message representing a change in a task's state.""" - - new_state: Optional[State] = Field( - default=None, description="""Required. The new state of the task.""" + aggregate_summary_fn: Optional[Callable[..., Any]] = Field( + default=None, + description="""The aggregate summary function for the judge model.""", ) - - -class TaskStateChangeDict(TypedDict, total=False): - """A message representing a change in a task's state.""" - - new_state: Optional[State] - """Required. The new state of the task.""" - - -TaskStateChangeOrDict = Union[TaskStateChange, TaskStateChangeDict] - - -class TaskStatusDetailsChange(_common.BaseModel): - """Represents a change to the task's status details.""" - - new_task_status: Optional[TaskStatusDetails] = Field( + remote_custom_function: Optional[str] = Field( default=None, - description="""Required. The complete state of the task's status *after* the change.""", + description="""The evaluation function for the custom code execution metric. This custom code is run remotely in the evaluation service.""", + ) + judge_model: Optional[str] = Field( + default=None, description="""The judge model for the metric.""" + ) + judge_model_generation_config: Optional[genai_types.GenerationConfig] = Field( + default=None, + description="""The generation config for the judge LLM (temperature, top_k, top_p, etc).""", + ) + judge_model_sampling_count: Optional[int] = Field( + default=None, description="""The sampling count for the judge model.""" + ) + rubric_group_name: Optional[str] = Field( + default=None, + description="""The rubric group name for the rubric-based metric.""", + ) + metric_spec_parameters: Optional[dict[str, Any]] = Field( + default=None, + description="""Optional steering instruction parameters for the automated predefined metric.""", + ) + metric_resource_name: Optional[str] = Field( + default=None, + description="""The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""", + ) + result_parsing_function: Optional[str] = Field( + default=None, + description="""Optional. A Python function string used to parse the raw output of the LLM judge model. The function must be named `parse_results` and accept a list of model response strings. It should return a dictionary with `score` (float) and `explanation` (str) keys.""", ) + # Allow extra fields to support metric-specific config fields. + model_config = ConfigDict(extra="allow") -class TaskStatusDetailsChangeDict(TypedDict, total=False): - """Represents a change to the task's status details.""" + _is_predefined: bool = PrivateAttr(default=False) + """A boolean indicating whether the metric is predefined.""" - new_task_status: Optional[TaskStatusDetailsDict] - """Required. The complete state of the task's status *after* the change.""" + _config_source: Optional[str] = PrivateAttr(default=None) + """An optional string indicating the source of the metric configuration.""" + _version: Optional[str] = PrivateAttr(default=None) + """An optional string indicating the version of the metric.""" -TaskStatusDetailsChangeOrDict = Union[ - TaskStatusDetailsChange, TaskStatusDetailsChangeDict -] + @model_validator(mode="after") + @classmethod + def validate_name(cls, model: "Metric") -> "Metric": + if not model.name: + raise ValueError("Metric name cannot be empty.") + model.name = model.name.lower() + return model + def to_yaml_file(self, file_path: str, version: Optional[str] = None) -> None: + """Dumps the metric object to a YAML file. -class TaskEventData(_common.BaseModel): - """Data for a TaskEvent.""" + Args: + file_path: The path to the YAML file. + version: Optional version string to include in the YAML output. - metadata_change: Optional[TaskMetadataChange] = Field( - default=None, description="""Optional. A change to the task's metadata.""" - ) - output_change: Optional[TaskOutputChange] = Field( - default=None, description="""Optional. A change to the task's final outputs.""" - ) - state_change: Optional[TaskStateChange] = Field( - default=None, description="""Optional. A change in the task's state.""" - ) - status_details_change: Optional[TaskStatusDetailsChange] = Field( - default=None, - description="""Optional. A change to the framework-specific status details.""", - ) + Raises: + ImportError: If the pyyaml library is not installed. + """ + if yaml is None: + raise ImportError( + "YAML serialization requires the pyyaml library. Please install" + " it using 'pip install google-cloud-aiplatform[evaluation]'." + ) + fields_to_exclude = { + field_name + for field_name, field_info in self.model_fields.items() + if self.__getattribute__(field_name) is not None + and isinstance(self.__getattribute__(field_name), Callable) + } -class TaskEventDataDict(TypedDict, total=False): - """Data for a TaskEvent.""" + data_to_dump = self.model_dump( + exclude_unset=True, + exclude_none=True, + mode="json", + exclude=fields_to_exclude if fields_to_exclude else None, + ) - metadata_change: Optional[TaskMetadataChangeDict] - """Optional. A change to the task's metadata.""" + if version: + data_to_dump["version"] = version - output_change: Optional[TaskOutputChangeDict] - """Optional. A change to the task's final outputs.""" + with open(file_path, "w", encoding="utf-8") as f: + yaml.dump(data_to_dump, f, sort_keys=False, allow_unicode=True) - state_change: Optional[TaskStateChangeDict] - """Optional. A change in the task's state.""" - status_details_change: Optional[TaskStatusDetailsChangeDict] - """Optional. A change to the framework-specific status details.""" +class CodeExecutionMetric(Metric): + """A metric that executes custom Python code for evaluation.""" + # You can use standard Pydantic Field syntax here because this is raw Python code + custom_function: Optional[str] = Field( + default=None, + description="""The Python function code to be executed on the server side.""", + ) -TaskEventDataOrDict = Union[TaskEventData, TaskEventDataDict] + # You can also add hand-written validators or methods here + @field_validator("custom_function") + @classmethod + def validate_code(cls, value: Optional[str]) -> Optional[str]: + if value and "def evaluate" not in value: + raise ValueError( + "custom_function must contain a 'def evaluate(instance):' signature." + ) + return value -class TaskEvent(_common.BaseModel): - """A task event.""" +class LLMMetric(Metric): + """A metric that uses LLM-as-a-judge for evaluation.""" - create_time: Optional[datetime.datetime] = Field( - default=None, description="""Output only. The create time of the event.""" - ) - event_data: Optional[TaskEventData] = Field( - default=None, description="""Required. The delta associated with the event.""" + rubric_group_name: Optional[str] = Field( + default=None, + description="""Optional. The name of the column in the EvaluationDataset containing the list of rubrics to use for this metric.""", ) - event_sequence_number: Optional[int] = Field( + + result_parsing_function: Optional[str] = Field( default=None, - description="""Required. The sequence number of the event. This is used to uniquely identify the event within the task and order events chronologically. This is a id generated by the SDK.""", + description="""Optional. A Python function string used to parse the raw output of the LLM judge model. The function must be named `parse_results` and accept a list of model response strings. It should return a dictionary with `score` (float) and `explanation` (str) keys.""", ) + @field_validator("prompt_template", mode="before") + @classmethod + def validate_prompt_template(cls, value: Union[str, "MetricPromptBuilder"]) -> str: + """Validates prompt template to be a non-empty string.""" + if value is None: + raise ValueError("Prompt template cannot be empty.") + if isinstance(value, MetricPromptBuilder): + value = str(value) + if not value.strip(): + raise ValueError("Prompt template cannot be an empty string.") + return value -class TaskEventDict(TypedDict, total=False): - """A task event.""" + @field_validator("judge_model_sampling_count") + @classmethod + def validate_judge_model_sampling_count(cls, value: Optional[int]) -> Optional[int]: + """Validates judge_model_sampling_count to be between 1 and 32.""" + if value is not None and (value < 1 or value > 32): + raise ValueError("judge_model_sampling_count must be between 1 and 32.") + return value - create_time: Optional[datetime.datetime] - """Output only. The create time of the event.""" + @classmethod + def load(cls, config_path: str, client: Optional[Any] = None) -> "LLMMetric": + """Loads a metric configuration from a YAML or JSON file. - event_data: Optional[TaskEventDataDict] - """Required. The delta associated with the event.""" + This method allows for the creation of an LLMMetric instance from a + local file path or a Google Cloud Storage (GCS) URI. It will automatically + detect the file type (.yaml, .yml, or .json) and parse it accordingly. - event_sequence_number: Optional[int] - """Required. The sequence number of the event. This is used to uniquely identify the event within the task and order events chronologically. This is a id generated by the SDK.""" + Args: + config_path: The local path or GCS URI (e.g., 'gs://bucket/metric.yaml') + to the metric configuration file. + client: Optional. The Vertex AI client instance to use for authentication. + If not provided, Application Default Credentials (ADC) will be used. + Returns: + An instance of LLMMetric configured with the loaded data. -TaskEventOrDict = Union[TaskEvent, TaskEventDict] + Raises: + ValueError: If the file path is invalid or the file content cannot be parsed. + ImportError: If a required library like 'PyYAML' or 'google-cloud-storage' is not installed. + IOError: If the file cannot be read from the specified path. + """ + file_extension = os.path.splitext(config_path)[1].lower() + if file_extension not in [".yaml", ".yml", ".json"]: + raise ValueError( + "Unsupported file extension for metric config. Must be .yaml, .yml, or .json" + ) + content_str: str + if config_path.startswith("gs://"): + try: + from google.cloud import storage # type: ignore[attr-defined] -class AppendAgentEngineTaskEventConfig(_common.BaseModel): - """Config for appending Agent Engine task events.""" + storage_client = storage.Client( + credentials=client._api_client._credentials if client else None + ) + path_without_prefix = config_path[len("gs://") :] + bucket_name, blob_path = path_without_prefix.split("/", 1) - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + bucket = storage_client.bucket(bucket_name) + blob = bucket.blob(blob_path) + content_str = blob.download_as_bytes().decode("utf-8") + except ImportError as e: + raise ImportError( + "Reading from GCS requires the 'google-cloud-storage' library. Please install it with 'pip install google-cloud-aiplatform[evaluation]'." + ) from e + except Exception as e: + raise IOError(f"Failed to read from GCS path {config_path}: {e}") from e + else: + try: + with open(config_path, "r", encoding="utf-8") as f: + content_str = f.read() + except FileNotFoundError: + raise FileNotFoundError( + f"Local configuration file not found at: {config_path}" + ) + except Exception as e: + raise IOError(f"Failed to read local file {config_path}: {e}") from e + data: Dict[str, Any] -class AppendAgentEngineTaskEventConfigDict(TypedDict, total=False): - """Config for appending Agent Engine task events.""" + if file_extension in [".yaml", ".yml"]: + if yaml is None: + raise ImportError( + "YAML parsing requires the pyyaml library. Please install it with 'pip install google-cloud-aiplatform[evaluation]'." + ) + data = yaml.safe_load(content_str) + elif file_extension == ".json": + data = json.loads(content_str) - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + if not isinstance(data, dict): + raise ValueError("Metric config content did not parse into a dictionary.") + return cls.model_validate(data) -AppendAgentEngineTaskEventConfigOrDict = Union[ - AppendAgentEngineTaskEventConfig, AppendAgentEngineTaskEventConfigDict -] +class MetricDict(TypedDict, total=False): + """The metric used for evaluation.""" -class _AppendAgentEngineTaskEventRequestParameters(_common.BaseModel): - """Parameters for appending Agent Engine task events.""" + name: Optional[str] + """The name of the metric.""" - name: Optional[str] = Field( - default=None, - description="""Name of the Agent Engine task to append the events to.""", - ) - task_events: Optional[list[TaskEvent]] = Field( - default=None, description="""The events to append to the task.""" - ) - config: Optional[AppendAgentEngineTaskEventConfig] = Field( - default=None, description="""""" - ) + custom_function: Optional[Union[str, Callable[..., Any]]] + """The custom function that defines the end-to-end logic for metric computation.""" + prompt_template: Optional[str] + """The prompt template for the metric.""" -class _AppendAgentEngineTaskEventRequestParametersDict(TypedDict, total=False): - """Parameters for appending Agent Engine task events.""" + judge_model_system_instruction: Optional[str] + """The system instruction for the judge model.""" - name: Optional[str] - """Name of the Agent Engine task to append the events to.""" + return_raw_output: Optional[bool] + """Whether to return the raw output from the judge model.""" - task_events: Optional[list[TaskEventDict]] - """The events to append to the task.""" + parse_and_reduce_fn: Optional[Callable[..., Any]] + """The parse and reduce function for the judge model.""" - config: Optional[AppendAgentEngineTaskEventConfigDict] - """""" + aggregate_summary_fn: Optional[Callable[..., Any]] + """The aggregate summary function for the judge model.""" + remote_custom_function: Optional[str] + """The evaluation function for the custom code execution metric. This custom code is run remotely in the evaluation service.""" -_AppendAgentEngineTaskEventRequestParametersOrDict = Union[ - _AppendAgentEngineTaskEventRequestParameters, - _AppendAgentEngineTaskEventRequestParametersDict, -] + judge_model: Optional[str] + """The judge model for the metric.""" + judge_model_generation_config: Optional[genai_types.GenerationConfig] + """The generation config for the judge LLM (temperature, top_k, top_p, etc).""" -class AppendAgentEngineTaskEventResponse(_common.BaseModel): - """Response for appending Agent Engine task events.""" + judge_model_sampling_count: Optional[int] + """The sampling count for the judge model.""" - pass + rubric_group_name: Optional[str] + """The rubric group name for the rubric-based metric.""" + metric_spec_parameters: Optional[dict[str, Any]] + """Optional steering instruction parameters for the automated predefined metric.""" -class AppendAgentEngineTaskEventResponseDict(TypedDict, total=False): - """Response for appending Agent Engine task events.""" + metric_resource_name: Optional[str] + """The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" - pass + result_parsing_function: Optional[str] + """Optional. A Python function string used to parse the raw output of the LLM judge model. The function must be named `parse_results` and accept a list of model response strings. It should return a dictionary with `score` (float) and `explanation` (str) keys.""" -AppendAgentEngineTaskEventResponseOrDict = Union[ - AppendAgentEngineTaskEventResponse, AppendAgentEngineTaskEventResponseDict -] +MetricOrDict = Union[Metric, MetricDict] -class ListAgentEngineTaskEventsConfig(_common.BaseModel): - """Config for listing agent engine tasks.""" +class CreateEvaluationMetricConfig(_common.BaseModel): + """Config for creating an evaluation metric.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( - default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""", - ) - order_by: Optional[str] = Field( - default=None, - description="""A comma-separated list of fields to order by, sorted in ascending order. - Use "desc" after a field name for descending. - If this field is omitted, the default ordering is `create_time` descending. - More detail in [AIP-132](https://google.aip.dev/132). - - Supported fields: - * `create_time` - * `update_time` - - Example: `create_time desc`.""", - ) -class ListAgentEngineTaskEventsConfigDict(TypedDict, total=False): - """Config for listing agent engine tasks.""" +class CreateEvaluationMetricConfigDict(TypedDict, total=False): + """Config for creating an evaluation metric.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - page_size: Optional[int] - """""" - page_token: Optional[str] - """""" - - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" - - order_by: Optional[str] - """A comma-separated list of fields to order by, sorted in ascending order. - Use "desc" after a field name for descending. - If this field is omitted, the default ordering is `create_time` descending. - More detail in [AIP-132](https://google.aip.dev/132). - - Supported fields: - * `create_time` - * `update_time` - - Example: `create_time desc`.""" - - -ListAgentEngineTaskEventsConfigOrDict = Union[ - ListAgentEngineTaskEventsConfig, ListAgentEngineTaskEventsConfigDict +CreateEvaluationMetricConfigOrDict = Union[ + CreateEvaluationMetricConfig, CreateEvaluationMetricConfigDict ] -class _ListAgentEngineTaskEventsRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" +class _CreateEvaluationMetricParameters(_common.BaseModel): + """Parameters for creating an evaluation metric.""" - name: Optional[str] = Field( - default=None, description="""Name of the Agent Engine task.""" + display_name: Optional[str] = Field( + default=None, + description="""The user-defined name of the evaluation metric. + + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """, + ) + description: Optional[str] = Field( + default=None, description="""The description of the evaluation metric.""" + ) + metric: Optional[Metric] = Field( + default=None, + description="""The metric configuration of the evaluation metric.""", ) - config: Optional[ListAgentEngineTaskEventsConfig] = Field( + config: Optional[CreateEvaluationMetricConfig] = Field( default=None, description="""""" ) -class _ListAgentEngineTaskEventsRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" +class _CreateEvaluationMetricParametersDict(TypedDict, total=False): + """Parameters for creating an evaluation metric.""" - name: Optional[str] - """Name of the Agent Engine task.""" + display_name: Optional[str] + """The user-defined name of the evaluation metric. + + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """ + + description: Optional[str] + """The description of the evaluation metric.""" + + metric: Optional[MetricDict] + """The metric configuration of the evaluation metric.""" - config: Optional[ListAgentEngineTaskEventsConfigDict] + config: Optional[CreateEvaluationMetricConfigDict] """""" -_ListAgentEngineTaskEventsRequestParametersOrDict = Union[ - _ListAgentEngineTaskEventsRequestParameters, - _ListAgentEngineTaskEventsRequestParametersDict, +_CreateEvaluationMetricParametersOrDict = Union[ + _CreateEvaluationMetricParameters, _CreateEvaluationMetricParametersDict ] -class ListAgentEngineTaskEventsResponse(_common.BaseModel): - """Response for listing Agent Engine tasks events.""" +class CustomCodeExecutionSpec(_common.BaseModel): + """Specifies a metric using remote Python function execution. - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + This metric is computed by running user-defined Python functions remotely. + """ + + evaluation_function: Optional[str] = Field( + default=None, + description="""Required. Python function. Expected user to define the following function, e.g.: def evaluate(instance: dict[str, Any]) -> float: Please include this function signature in the code snippet. Instance is the evaluation instance, any fields populated in the instance are available to the function as instance[field_name]. Example: Example input: ``` instance= EvaluationInstance( response=EvaluationInstance.InstanceData(text="The answer is 4."), reference=EvaluationInstance.InstanceData(text="4") ) ``` Example converted input: ``` { 'response': {'text': 'The answer is 4.'}, 'reference': {'text': '4'} } ``` Example python function: ``` def evaluate(instance: dict[str, Any]) -> float: if instance'response' == instance'reference': return 1.0 return 0.0 ``` CustomCodeExecutionSpec is also supported in Batch Evaluation (EvalDataset RPC) and Tuning Evaluation. Each line in the input jsonl file will be converted to dict[str, Any] and passed to the evaluation function.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - taskEvents: Optional[list[TaskEvent]] = Field( - default=None, description="""List of Agent Engine task events.""" + remote_custom_function: Optional[str] = Field( + default=None, + description="""A string representing a user-defined function for evaluation. + Expected user to define the following function, e.g.: + def evaluate(instance: dict[str, Any]) -> float: + Please include this function signature in the code snippet. + Instance is the evaluation instance, any fields populated in the instance + are available to the function as instance[field_name].""", ) -class ListAgentEngineTaskEventsResponseDict(TypedDict, total=False): - """Response for listing Agent Engine tasks events.""" +class CustomCodeExecutionSpecDict(TypedDict, total=False): + """Specifies a metric using remote Python function execution. - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + This metric is computed by running user-defined Python functions remotely. + """ - next_page_token: Optional[str] - """""" + evaluation_function: Optional[str] + """Required. Python function. Expected user to define the following function, e.g.: def evaluate(instance: dict[str, Any]) -> float: Please include this function signature in the code snippet. Instance is the evaluation instance, any fields populated in the instance are available to the function as instance[field_name]. Example: Example input: ``` instance= EvaluationInstance( response=EvaluationInstance.InstanceData(text="The answer is 4."), reference=EvaluationInstance.InstanceData(text="4") ) ``` Example converted input: ``` { 'response': {'text': 'The answer is 4.'}, 'reference': {'text': '4'} } ``` Example python function: ``` def evaluate(instance: dict[str, Any]) -> float: if instance'response' == instance'reference': return 1.0 return 0.0 ``` CustomCodeExecutionSpec is also supported in Batch Evaluation (EvalDataset RPC) and Tuning Evaluation. Each line in the input jsonl file will be converted to dict[str, Any] and passed to the evaluation function.""" - taskEvents: Optional[list[TaskEventDict]] - """List of Agent Engine task events.""" + remote_custom_function: Optional[str] + """A string representing a user-defined function for evaluation. + Expected user to define the following function, e.g.: + def evaluate(instance: dict[str, Any]) -> float: + Please include this function signature in the code snippet. + Instance is the evaluation instance, any fields populated in the instance + are available to the function as instance[field_name].""" -ListAgentEngineTaskEventsResponseOrDict = Union[ - ListAgentEngineTaskEventsResponse, ListAgentEngineTaskEventsResponseDict +CustomCodeExecutionSpecOrDict = Union[ + CustomCodeExecutionSpec, CustomCodeExecutionSpecDict ] -class CreateEvaluationExperimentConfig(_common.BaseModel): - """Config to create an evaluation experiment.""" +class UnifiedMetric(_common.BaseModel): + """The unified metric used for evaluation.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + bleu_spec: Optional[genai_types.BleuSpec] = Field( + default=None, description="""The Bleu metric spec.""" ) - - -class CreateEvaluationExperimentConfigDict(TypedDict, total=False): - """Config to create an evaluation experiment.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -CreateEvaluationExperimentConfigOrDict = Union[ - CreateEvaluationExperimentConfig, CreateEvaluationExperimentConfigDict -] - - -class _CreateEvaluationExperimentParameters(_common.BaseModel): - """Parameters for creating an evaluation experiment.""" - - display_name: Optional[str] = Field(default=None, description="""""") - labels: Optional[dict[str, str]] = Field(default=None, description="""""") - merge_strategy: Optional[EvaluationExperimentMergeStrategy] = Field( - default=None, description="""""" + rouge_spec: Optional[genai_types.RougeSpec] = Field( + default=None, description="""The rouge metric spec.""" ) - metadata: Optional[dict[str, Any]] = Field(default=None, description="""""") - config: Optional[CreateEvaluationExperimentConfig] = Field( - default=None, description="""""" + pointwise_metric_spec: Optional[genai_types.PointwiseMetricSpec] = Field( + default=None, description="""The pointwise metric spec.""" + ) + llm_based_metric_spec: Optional[genai_types.LLMBasedMetricSpec] = Field( + default=None, description="""The spec for an LLM based metric.""" + ) + custom_code_execution_spec: Optional[CustomCodeExecutionSpec] = Field( + default=None, description="""The spec for a custom code execution metric.""" + ) + predefined_metric_spec: Optional[genai_types.PredefinedMetricSpec] = Field( + default=None, description="""The spec for a pre-defined metric.""" + ) + computation_based_metric_spec: Optional[genai_types.ComputationBasedMetricSpec] = ( + Field(default=None, description="""The spec for a computation based metric.""") ) -class _CreateEvaluationExperimentParametersDict(TypedDict, total=False): - """Parameters for creating an evaluation experiment.""" +class UnifiedMetricDict(TypedDict, total=False): + """The unified metric used for evaluation.""" - display_name: Optional[str] - """""" + bleu_spec: Optional[genai_types.BleuSpec] + """The Bleu metric spec.""" - labels: Optional[dict[str, str]] - """""" + rouge_spec: Optional[genai_types.RougeSpec] + """The rouge metric spec.""" - merge_strategy: Optional[EvaluationExperimentMergeStrategy] - """""" + pointwise_metric_spec: Optional[genai_types.PointwiseMetricSpec] + """The pointwise metric spec.""" - metadata: Optional[dict[str, Any]] - """""" + llm_based_metric_spec: Optional[genai_types.LLMBasedMetricSpec] + """The spec for an LLM based metric.""" - config: Optional[CreateEvaluationExperimentConfigDict] - """""" + custom_code_execution_spec: Optional[CustomCodeExecutionSpecDict] + """The spec for a custom code execution metric.""" + predefined_metric_spec: Optional[genai_types.PredefinedMetricSpec] + """The spec for a pre-defined metric.""" -_CreateEvaluationExperimentParametersOrDict = Union[ - _CreateEvaluationExperimentParameters, _CreateEvaluationExperimentParametersDict -] + computation_based_metric_spec: Optional[genai_types.ComputationBasedMetricSpec] + """The spec for a computation based metric.""" -class EvaluationExperiment(_common.BaseModel): - """Represents an experiment for iterating on and visualizing evaluation runs.""" +UnifiedMetricOrDict = Union[UnifiedMetric, UnifiedMetricDict] + + +class EvaluationMetric(_common.BaseModel): + """Represents an evaluation metric.""" name: Optional[str] = Field( - default=None, - description="""The resource name of the EvaluationExperiment. Format: - `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""", + default=None, description="""The resource name of the evaluation metric.""" ) display_name: Optional[str] = Field( - default=None, description="""The display name of the evaluation experiment.""" - ) - evaluation_runs: Optional[list[str]] = Field( default=None, - description="""The EvaluationRuns that are part of this experiment.""", - ) - labels: Optional[dict[str, str]] = Field( - default=None, description="""Labels for the evaluation experiment.""" + description="""The user-friendly display name for the EvaluationMetric.""", ) - merge_strategy: Optional[EvaluationExperimentMergeStrategy] = Field( - default=None, description="""Merge strategy for the evaluation experiment.""" + description: Optional[str] = Field( + default=None, description="""The description of the EvaluationMetric.""" ) - metadata: Optional[dict[str, Any]] = Field( + metric: Optional[UnifiedMetric] = Field( default=None, - description="""Metadata about the evaluation experiment, can be used by the caller - to store additional tracking information about the experiment.""", - ) - create_time: Optional[datetime.datetime] = Field( - default=None, description="""Timestamp when this experiment was created.""" - ) - update_time: Optional[datetime.datetime] = Field( - default=None, description="""Timestamp when this experiment was last updated.""" + description="""The metric configuration of the evaluation metric.""", ) -class EvaluationExperimentDict(TypedDict, total=False): - """Represents an experiment for iterating on and visualizing evaluation runs.""" +class EvaluationMetricDict(TypedDict, total=False): + """Represents an evaluation metric.""" name: Optional[str] - """The resource name of the EvaluationExperiment. Format: - `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""" + """The resource name of the evaluation metric.""" display_name: Optional[str] - """The display name of the evaluation experiment.""" + """The user-friendly display name for the EvaluationMetric.""" - evaluation_runs: Optional[list[str]] - """The EvaluationRuns that are part of this experiment.""" + description: Optional[str] + """The description of the EvaluationMetric.""" - labels: Optional[dict[str, str]] - """Labels for the evaluation experiment.""" + metric: Optional[UnifiedMetricDict] + """The metric configuration of the evaluation metric.""" - merge_strategy: Optional[EvaluationExperimentMergeStrategy] - """Merge strategy for the evaluation experiment.""" - metadata: Optional[dict[str, Any]] - """Metadata about the evaluation experiment, can be used by the caller - to store additional tracking information about the experiment.""" +EvaluationMetricOrDict = Union[EvaluationMetric, EvaluationMetricDict] - create_time: Optional[datetime.datetime] - """Timestamp when this experiment was created.""" - update_time: Optional[datetime.datetime] - """Timestamp when this experiment was last updated.""" +class SamplingConfig(_common.BaseModel): + """Sampling config for a BigQuery request set.""" + sampling_count: Optional[int] = Field(default=None, description="""""") + sampling_method: Optional[SamplingMethod] = Field(default=None, description="""""") + sampling_duration: Optional[str] = Field(default=None, description="""""") -EvaluationExperimentOrDict = Union[EvaluationExperiment, EvaluationExperimentDict] +class SamplingConfigDict(TypedDict, total=False): + """Sampling config for a BigQuery request set.""" -class CreateEvaluationItemConfig(_common.BaseModel): - """Config to create an evaluation item.""" + sampling_count: Optional[int] + """""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + sampling_method: Optional[SamplingMethod] + """""" + sampling_duration: Optional[str] + """""" -class CreateEvaluationItemConfigDict(TypedDict, total=False): - """Config to create an evaluation item.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +SamplingConfigOrDict = Union[SamplingConfig, SamplingConfigDict] -CreateEvaluationItemConfigOrDict = Union[ - CreateEvaluationItemConfig, CreateEvaluationItemConfigDict -] - - -class _CreateEvaluationItemParameters(_common.BaseModel): - """Represents a job that creates an evaluation item.""" +class BigQueryRequestSet(_common.BaseModel): + """Represents a BigQuery request set.""" - evaluation_item_type: Optional[str] = Field(default=None, description="""""") - gcs_uri: Optional[str] = Field(default=None, description="""""") - display_name: Optional[str] = Field(default=None, description="""""") - config: Optional[CreateEvaluationItemConfig] = Field( - default=None, description="""""" + uri: Optional[str] = Field(default=None, description="""""") + prompt_column: Optional[str] = Field( + default=None, + description="""The column name of the prompt in the BigQuery table. Used for EvaluationRun only.""", + ) + rubrics_column: Optional[str] = Field( + default=None, + description="""The column name of the rubrics in the BigQuery table. Used for EvaluationRun only.""", + ) + candidate_response_columns: Optional[dict[str, str]] = Field( + default=None, + description="""The column name of the response candidates in the BigQuery table. Used for EvaluationRun only.""", + ) + sampling_config: Optional[SamplingConfig] = Field( + default=None, + description="""The sampling config for the BigQuery request set. Used for EvaluationRun only.""", ) -class _CreateEvaluationItemParametersDict(TypedDict, total=False): - """Represents a job that creates an evaluation item.""" +class BigQueryRequestSetDict(TypedDict, total=False): + """Represents a BigQuery request set.""" - evaluation_item_type: Optional[str] + uri: Optional[str] """""" - gcs_uri: Optional[str] - """""" + prompt_column: Optional[str] + """The column name of the prompt in the BigQuery table. Used for EvaluationRun only.""" - display_name: Optional[str] - """""" + rubrics_column: Optional[str] + """The column name of the rubrics in the BigQuery table. Used for EvaluationRun only.""" - config: Optional[CreateEvaluationItemConfigDict] - """""" + candidate_response_columns: Optional[dict[str, str]] + """The column name of the response candidates in the BigQuery table. Used for EvaluationRun only.""" + sampling_config: Optional[SamplingConfigDict] + """The sampling config for the BigQuery request set. Used for EvaluationRun only.""" -_CreateEvaluationItemParametersOrDict = Union[ - _CreateEvaluationItemParameters, _CreateEvaluationItemParametersDict -] +BigQueryRequestSetOrDict = Union[BigQueryRequestSet, BigQueryRequestSetDict] -class PromptTemplateData(_common.BaseModel): - """Holds data for a prompt template. - Message to hold a prompt template and the values to populate the template. - """ +class EvaluationRunDataSource(_common.BaseModel): + """Represents an evaluation run data source.""" - values: Optional[dict[str, genai_types.Content]] = Field( - default=None, description="""The values for fields in the prompt template.""" + evaluation_set: Optional[str] = Field(default=None, description="""""") + bigquery_request_set: Optional[BigQueryRequestSet] = Field( + default=None, description="""""" ) -class PromptTemplateDataDict(TypedDict, total=False): - """Holds data for a prompt template. +class EvaluationRunDataSourceDict(TypedDict, total=False): + """Represents an evaluation run data source.""" - Message to hold a prompt template and the values to populate the template. - """ + evaluation_set: Optional[str] + """""" - values: Optional[dict[str, genai_types.Content]] - """The values for fields in the prompt template.""" + bigquery_request_set: Optional[BigQueryRequestSetDict] + """""" -PromptTemplateDataOrDict = Union[PromptTemplateData, PromptTemplateDataDict] +EvaluationRunDataSourceOrDict = Union[ + EvaluationRunDataSource, EvaluationRunDataSourceDict +] -class EvaluationPrompt(_common.BaseModel): - """Represents the prompt to be evaluated.""" +class EvaluationRunMetric(_common.BaseModel): + """The metric used for evaluation run.""" - text: Optional[str] = Field(default=None, description="""Text prompt.""") - value: Optional[dict[str, Any]] = Field( - default=None, - description="""Fields and values that can be used to populate the prompt template.""", - ) - prompt_template_data: Optional[PromptTemplateData] = Field( - default=None, description="""Prompt template data.""" + metric: Optional[str] = Field( + default=None, description="""The name of the metric.""" ) - user_scenario: Optional[evals_types.UserScenario] = Field( + metric_resource_name: Optional[str] = Field( default=None, - description="""User scenario to help simulate multi-turn agent run results.""", + description="""The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""", + ) + metric_config: Optional[UnifiedMetric] = Field( + default=None, description="""The unified metric used for evaluation run.""" ) -class EvaluationPromptDict(TypedDict, total=False): - """Represents the prompt to be evaluated.""" +class EvaluationRunMetricDict(TypedDict, total=False): + """The metric used for evaluation run.""" - text: Optional[str] - """Text prompt.""" + metric: Optional[str] + """The name of the metric.""" - value: Optional[dict[str, Any]] - """Fields and values that can be used to populate the prompt template.""" + metric_resource_name: Optional[str] + """The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" - prompt_template_data: Optional[PromptTemplateDataDict] - """Prompt template data.""" + metric_config: Optional[UnifiedMetricDict] + """The unified metric used for evaluation run.""" - user_scenario: Optional[evals_types.UserScenario] - """User scenario to help simulate multi-turn agent run results.""" +EvaluationRunMetricOrDict = Union[EvaluationRunMetric, EvaluationRunMetricDict] -EvaluationPromptOrDict = Union[EvaluationPrompt, EvaluationPromptDict] +class EvaluationRunPromptTemplate(_common.BaseModel): + """Prompt template used for inference. -class CandidateResponse(_common.BaseModel): - """Responses from model or agent.""" + Only one of `prompt_template` or `gcs_uri` should be set. If both are + provided, an error will be raised. + """ - candidate: Optional[str] = Field( - default=None, - description="""The name of the candidate that produced the response.""", - ) - text: Optional[str] = Field(default=None, description="""The text response.""") - value: Optional[dict[str, Any]] = Field( - default=None, - description="""Fields and values that can be used to populate the response template.""", - ) - events: Optional[list[genai_types.Content]] = Field( + prompt_template: Optional[str] = Field( default=None, - description="""Intermediate events (such as tool calls and responses) that led to the final response.""", + description="""Inline prompt template. Template variables should be in the format + "{var_name}". Only one of `prompt_template` or `gcs_uri` should be set.""", ) - agent_data: Optional[evals_types.AgentData] = Field( + gcs_uri: Optional[str] = Field( default=None, - description="""Represents the complete execution trace of an agent conversation, - which can involve single or multiple agents. This field is used to - provide the full output of an agent's run, including all turns and - events, for direct evaluation.""", + description="""Prompt template stored in Cloud Storage. Format: + "gs://my-bucket/file-name.txt". Only one of `prompt_template` or `gcs_uri` + should be set.""", ) -class CandidateResponseDict(TypedDict, total=False): - """Responses from model or agent.""" - - candidate: Optional[str] - """The name of the candidate that produced the response.""" - - text: Optional[str] - """The text response.""" +class EvaluationRunPromptTemplateDict(TypedDict, total=False): + """Prompt template used for inference. - value: Optional[dict[str, Any]] - """Fields and values that can be used to populate the response template.""" + Only one of `prompt_template` or `gcs_uri` should be set. If both are + provided, an error will be raised. + """ - events: Optional[list[genai_types.Content]] - """Intermediate events (such as tool calls and responses) that led to the final response.""" + prompt_template: Optional[str] + """Inline prompt template. Template variables should be in the format + "{var_name}". Only one of `prompt_template` or `gcs_uri` should be set.""" - agent_data: Optional[evals_types.AgentData] - """Represents the complete execution trace of an agent conversation, - which can involve single or multiple agents. This field is used to - provide the full output of an agent's run, including all turns and - events, for direct evaluation.""" + gcs_uri: Optional[str] + """Prompt template stored in Cloud Storage. Format: + "gs://my-bucket/file-name.txt". Only one of `prompt_template` or `gcs_uri` + should be set.""" -CandidateResponseOrDict = Union[CandidateResponse, CandidateResponseDict] +EvaluationRunPromptTemplateOrDict = Union[ + EvaluationRunPromptTemplate, EvaluationRunPromptTemplateDict +] -class EvaluationItemRequest(_common.BaseModel): - """Single evaluation request.""" +class LossAnalysisConfig(_common.BaseModel): + """Configuration for the loss analysis job.""" - prompt: Optional[EvaluationPrompt] = Field( - default=None, description="""The request/prompt to evaluate.""" + metric: Optional[str] = Field( + default=None, + description="""Required. The metric to analyze (e.g., "multi_turn_tool_use_quality_v1").""", ) - golden_response: Optional[CandidateResponse] = Field( - default=None, description="""The ideal response or ground truth.""" + candidate: Optional[str] = Field( + default=None, + description="""Required. The candidate model/agent to analyze (e.g., "gemini-3.1-pro-preview"). This targets the specific CandidateResult within the EvaluationResult.""", ) - rubrics: Optional[dict[str, "RubricGroup"]] = Field( + predefined_taxonomy: Optional[str] = Field( default=None, - description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", + description="""Optional. The identifier for the pre-defined taxonomy to use (e.g., "agent_taxonomy_v1", "tool_use_v2"). If not specified, the service may select a default based on the metric.""", ) - candidate_responses: Optional[list[CandidateResponse]] = Field( + max_top_cluster_count: Optional[int] = Field( default=None, - description="""Responses from model under test and other baseline models for comparison.""", + description="""Optional. Limits the analysis to the top N clusters. If not specified or set to 0, all clusters are returned.""", ) -class EvaluationItemRequestDict(TypedDict, total=False): - """Single evaluation request.""" +class LossAnalysisConfigDict(TypedDict, total=False): + """Configuration for the loss analysis job.""" - prompt: Optional[EvaluationPromptDict] - """The request/prompt to evaluate.""" + metric: Optional[str] + """Required. The metric to analyze (e.g., "multi_turn_tool_use_quality_v1").""" - golden_response: Optional[CandidateResponseDict] - """The ideal response or ground truth.""" + candidate: Optional[str] + """Required. The candidate model/agent to analyze (e.g., "gemini-3.1-pro-preview"). This targets the specific CandidateResult within the EvaluationResult.""" - rubrics: Optional[dict[str, "RubricGroupDict"]] - """Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""" + predefined_taxonomy: Optional[str] + """Optional. The identifier for the pre-defined taxonomy to use (e.g., "agent_taxonomy_v1", "tool_use_v2"). If not specified, the service may select a default based on the metric.""" - candidate_responses: Optional[list[CandidateResponseDict]] - """Responses from model under test and other baseline models for comparison.""" + max_top_cluster_count: Optional[int] + """Optional. Limits the analysis to the top N clusters. If not specified or set to 0, all clusters are returned.""" -EvaluationItemRequestOrDict = Union[EvaluationItemRequest, EvaluationItemRequestDict] +LossAnalysisConfigOrDict = Union[LossAnalysisConfig, LossAnalysisConfigDict] -class EvaluationItemResult(_common.BaseModel): - """Represents the result of an evaluation item.""" +class EvaluationRunConfig(_common.BaseModel): + """The evaluation configuration used for the evaluation run.""" - evaluation_request: Optional[str] = Field( - default=None, description="""The request item that was evaluated.""" - ) - evaluation_run: Optional[str] = Field( + metrics: Optional[list[EvaluationRunMetric]] = Field( default=None, - description="""The evaluation run that was used to generate the result.""", + description="""The metrics to be calculated in the evaluation run.""", ) - request: Optional[EvaluationItemRequest] = Field( - default=None, description="""The request that was evaluated.""" + output_config: Optional[genai_types.OutputConfig] = Field( + default=None, description="""The output config for the evaluation run.""" ) - metric: Optional[str] = Field( - default=None, description="""The metric that was evaluated.""" + autorater_config: Optional[genai_types.AutoraterConfig] = Field( + default=None, + description="""The autorater config for the evaluation run. Not applicable for predefined metrics (PredefinedMetricSpec); the server uses its own model configuration for predefined metrics and this field is ignored.""", ) - candidate_results: Optional[list[evals_types.CandidateResult]] = Field( - default=None, description="""The results for the metric.""" + prompt_template: Optional[EvaluationRunPromptTemplate] = Field( + default=None, description="""The prompt template used for inference.""" ) - metadata: Optional[dict[str, Any]] = Field( - default=None, description="""Metadata about the evaluation result.""" + loss_analysis_config: Optional[list[LossAnalysisConfig]] = Field( + default=None, + description="""Specifications for loss analysis. Each config specifies a metric and candidate to analyze for loss patterns.""", + ) + allow_cross_region_model: Optional[bool] = Field( + default=None, + description="""Allows the evaluation run to use cross region models. When this + flag is set, the service may route traffic to other regions if a model is + unavailable in the current region (e.g., to a `global`endpoint). If a + fully-qualified model endpoint resource name with a different region than + the run location is provided elsewhere in the run config, this flag must + be set to true or the request will fail.""", ) -class EvaluationItemResultDict(TypedDict, total=False): - """Represents the result of an evaluation item.""" - - evaluation_request: Optional[str] - """The request item that was evaluated.""" +class EvaluationRunConfigDict(TypedDict, total=False): + """The evaluation configuration used for the evaluation run.""" - evaluation_run: Optional[str] - """The evaluation run that was used to generate the result.""" + metrics: Optional[list[EvaluationRunMetricDict]] + """The metrics to be calculated in the evaluation run.""" - request: Optional[EvaluationItemRequestDict] - """The request that was evaluated.""" + output_config: Optional[genai_types.OutputConfig] + """The output config for the evaluation run.""" - metric: Optional[str] - """The metric that was evaluated.""" + autorater_config: Optional[genai_types.AutoraterConfig] + """The autorater config for the evaluation run. Not applicable for predefined metrics (PredefinedMetricSpec); the server uses its own model configuration for predefined metrics and this field is ignored.""" - candidate_results: Optional[list[evals_types.CandidateResult]] - """The results for the metric.""" + prompt_template: Optional[EvaluationRunPromptTemplateDict] + """The prompt template used for inference.""" - metadata: Optional[dict[str, Any]] - """Metadata about the evaluation result.""" + loss_analysis_config: Optional[list[LossAnalysisConfigDict]] + """Specifications for loss analysis. Each config specifies a metric and candidate to analyze for loss patterns.""" + allow_cross_region_model: Optional[bool] + """Allows the evaluation run to use cross region models. When this + flag is set, the service may route traffic to other regions if a model is + unavailable in the current region (e.g., to a `global`endpoint). If a + fully-qualified model endpoint resource name with a different region than + the run location is provided elsewhere in the run config, this flag must + be set to true or the request will fail.""" -EvaluationItemResultOrDict = Union[EvaluationItemResult, EvaluationItemResultDict] +EvaluationRunConfigOrDict = Union[EvaluationRunConfig, EvaluationRunConfigDict] -class EvaluationItem(_common.BaseModel): - """EvaluationItem is a single evaluation request or result. - The content of an EvaluationItem is immutable - it cannot be updated once - created. EvaluationItems can be deleted when no longer needed. - """ +class EvaluationRunAgentConfig(_common.BaseModel): + """Agent config for an evaluation run.""" - name: Optional[str] = Field( - default=None, description="""The resource name of the EvaluationItem.""" - ) - display_name: Optional[str] = Field( - default=None, description="""The display name of the EvaluationItem.""" - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, description="""Metadata for the EvaluationItem.""" - ) - labels: Optional[dict[str, str]] = Field( - default=None, description="""Labels for the EvaluationItem.""" - ) - evaluation_item_type: Optional[EvaluationItemType] = Field( - default=None, description="""The type of the EvaluationItem.""" - ) - evaluation_request: Optional[EvaluationItemRequest] = Field( - default=None, description="""The request to evaluate.""" - ) - evaluation_response: Optional[EvaluationItemResult] = Field( - default=None, description="""The response from evaluation.""" - ) - gcs_uri: Optional[str] = Field( - default=None, - description="""The Cloud Storage object where the request or response is stored.""", - ) - create_time: Optional[datetime.datetime] = Field( - default=None, description="""Timestamp when this item was created.""" + developer_instruction: Optional[genai_types.Content] = Field( + default=None, description="""The developer instruction for the agent.""" ) - error: Optional[genai_types.GoogleRpcStatus] = Field( - default=None, description="""Error for the evaluation item.""" + tools: Optional[list[genai_types.Tool]] = Field( + default=None, description="""The tools available to the agent.""" ) - # TODO(b/448806531): Remove all the overridden _from_response methods once the - # ticket is resolved and published. - @classmethod - def _from_response( - cls: typing.Type["EvaluationItem"], - *, - response: dict[str, object], - kwargs: dict[str, object], - ) -> "EvaluationItem": - """Converts a dictionary response into a EvaluationItem object.""" - - response = _camel_key_to_snake(response) - result = super()._from_response(response=response, kwargs=kwargs) - return result +class EvaluationRunAgentConfigDict(TypedDict, total=False): + """Agent config for an evaluation run.""" -class EvaluationItemDict(TypedDict, total=False): - """EvaluationItem is a single evaluation request or result. + developer_instruction: Optional[genai_types.Content] + """The developer instruction for the agent.""" - The content of an EvaluationItem is immutable - it cannot be updated once - created. EvaluationItems can be deleted when no longer needed. - """ + tools: Optional[list[genai_types.Tool]] + """The tools available to the agent.""" - name: Optional[str] - """The resource name of the EvaluationItem.""" - display_name: Optional[str] - """The display name of the EvaluationItem.""" +EvaluationRunAgentConfigOrDict = Union[ + EvaluationRunAgentConfig, EvaluationRunAgentConfigDict +] - metadata: Optional[dict[str, Any]] - """Metadata for the EvaluationItem.""" - labels: Optional[dict[str, str]] - """Labels for the EvaluationItem.""" +class GeminiAgentConfig(_common.BaseModel): + """Config for scraping a Gemini Agent. - evaluation_item_type: Optional[EvaluationItemType] - """The type of the EvaluationItem.""" + A Gemini Agent is a Vertex AI Agent resource scraped via the Vertex + Interactions API. + """ - evaluation_request: Optional[EvaluationItemRequestDict] - """The request to evaluate.""" + gemini_agent: Optional[str] = Field( + default=None, + description="""The resource name of the Gemini Agent. + Format: `projects/{project}/locations/{location}/agents/{agent}`.""", + ) - evaluation_response: Optional[EvaluationItemResultDict] - """The response from evaluation.""" - gcs_uri: Optional[str] - """The Cloud Storage object where the request or response is stored.""" +class GeminiAgentConfigDict(TypedDict, total=False): + """Config for scraping a Gemini Agent. - create_time: Optional[datetime.datetime] - """Timestamp when this item was created.""" + A Gemini Agent is a Vertex AI Agent resource scraped via the Vertex + Interactions API. + """ - error: Optional[genai_types.GoogleRpcStatus] - """Error for the evaluation item.""" + gemini_agent: Optional[str] + """The resource name of the Gemini Agent. + Format: `projects/{project}/locations/{location}/agents/{agent}`.""" -EvaluationItemOrDict = Union[EvaluationItem, EvaluationItemDict] +GeminiAgentConfigOrDict = Union[GeminiAgentConfig, GeminiAgentConfigDict] -class Metric(_common.BaseModel): - """The metric used for evaluation.""" +class AgentRunConfig(_common.BaseModel): + """Configuration for an Agent Run.""" - name: Optional[str] = Field(default=None, description="""The name of the metric.""") - custom_function: Optional[Union[str, Callable[..., Any]]] = Field( - default=None, - description="""The custom function that defines the end-to-end logic for metric computation.""", - ) - prompt_template: Optional[str] = Field( - default=None, description="""The prompt template for the metric.""" - ) - judge_model_system_instruction: Optional[str] = Field( - default=None, description="""The system instruction for the judge model.""" - ) - return_raw_output: Optional[bool] = Field( - default=None, - description="""Whether to return the raw output from the judge model.""", - ) - parse_and_reduce_fn: Optional[Callable[..., Any]] = Field( - default=None, - description="""The parse and reduce function for the judge model.""", + session_input: Optional[evals_types.SessionInput] = Field( + default=None, description="""The session input to get agent running results.""" ) - aggregate_summary_fn: Optional[Callable[..., Any]] = Field( - default=None, - description="""The aggregate summary function for the judge model.""", + agent_engine: Optional[str] = Field( + default=None, description="""The resource name of the Agent Engine.""" ) - remote_custom_function: Optional[str] = Field( + user_simulator_config: Optional[evals_types.UserSimulatorConfig] = Field( default=None, - description="""The evaluation function for the custom code execution metric. This custom code is run remotely in the evaluation service.""", - ) - judge_model: Optional[str] = Field( - default=None, description="""The judge model for the metric.""" + description="""Used for multi-turn agent run. + Contains configuration for a user simulator that + uses an LLM to generate messages on behalf of the user.""", ) - judge_model_generation_config: Optional[genai_types.GenerationConfig] = Field( + gemini_agent_config: Optional[GeminiAgentConfig] = Field( default=None, - description="""The generation config for the judge LLM (temperature, top_k, top_p, etc).""", + description="""Config for scraping a Gemini Agent (Vertex AI Agent resource). + Used to target a Gemini agent for an evaluation run.""", ) - judge_model_sampling_count: Optional[int] = Field( - default=None, description="""The sampling count for the judge model.""" + + +class AgentRunConfigDict(TypedDict, total=False): + """Configuration for an Agent Run.""" + + session_input: Optional[evals_types.SessionInput] + """The session input to get agent running results.""" + + agent_engine: Optional[str] + """The resource name of the Agent Engine.""" + + user_simulator_config: Optional[evals_types.UserSimulatorConfig] + """Used for multi-turn agent run. + Contains configuration for a user simulator that + uses an LLM to generate messages on behalf of the user.""" + + gemini_agent_config: Optional[GeminiAgentConfigDict] + """Config for scraping a Gemini Agent (Vertex AI Agent resource). + Used to target a Gemini agent for an evaluation run.""" + + +AgentRunConfigOrDict = Union[AgentRunConfig, AgentRunConfigDict] + + +class EvaluationRunInferenceConfig(_common.BaseModel): + """Configuration that describes an agent.""" + + agent_config: Optional[EvaluationRunAgentConfig] = Field( + default=None, description="""The agent config.""" ) - rubric_group_name: Optional[str] = Field( + model: Optional[str] = Field( default=None, - description="""The rubric group name for the rubric-based metric.""", + description="""The model to use for inference. Accepts a short Gemini model name (e.g. `gemini-2.5-flash`), which is automatically expanded to a fully-qualified resource name using the client's project and location, or an already fully-qualified publisher-model or endpoint resource name (e.g. `projects/{project}/locations/{location}/publishers/google/models/gemini-2.5-flash`).""", ) - metric_spec_parameters: Optional[dict[str, Any]] = Field( - default=None, - description="""Optional steering instruction parameters for the automated predefined metric.""", + prompt_template: Optional[EvaluationRunPromptTemplate] = Field( + default=None, description="""The prompt template used for inference.""" ) - metric_resource_name: Optional[str] = Field( + agent_run_config: Optional[AgentRunConfig] = Field( default=None, - description="""The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""", + description="""Configuration for Agent Run in evaluation management service.""", ) - result_parsing_function: Optional[str] = Field( + agent_configs: Optional[dict[str, evals_types.AgentConfig]] = Field( default=None, - description="""Optional. A Python function string used to parse the raw output of the LLM judge model. The function must be named `parse_results` and accept a list of model response strings. It should return a dictionary with `score` (float) and `explanation` (str) keys.""", + description="""A map of agent IDs to their respective agent config.""", ) - # Allow extra fields to support metric-specific config fields. - model_config = ConfigDict(extra="allow") - _is_predefined: bool = PrivateAttr(default=False) - """A boolean indicating whether the metric is predefined.""" +class EvaluationRunInferenceConfigDict(TypedDict, total=False): + """Configuration that describes an agent.""" - _config_source: Optional[str] = PrivateAttr(default=None) - """An optional string indicating the source of the metric configuration.""" + agent_config: Optional[EvaluationRunAgentConfigDict] + """The agent config.""" - _version: Optional[str] = PrivateAttr(default=None) - """An optional string indicating the version of the metric.""" + model: Optional[str] + """The model to use for inference. Accepts a short Gemini model name (e.g. `gemini-2.5-flash`), which is automatically expanded to a fully-qualified resource name using the client's project and location, or an already fully-qualified publisher-model or endpoint resource name (e.g. `projects/{project}/locations/{location}/publishers/google/models/gemini-2.5-flash`).""" - @model_validator(mode="after") - @classmethod - def validate_name(cls, model: "Metric") -> "Metric": - if not model.name: - raise ValueError("Metric name cannot be empty.") - model.name = model.name.lower() - return model + prompt_template: Optional[EvaluationRunPromptTemplateDict] + """The prompt template used for inference.""" - def to_yaml_file(self, file_path: str, version: Optional[str] = None) -> None: - """Dumps the metric object to a YAML file. + agent_run_config: Optional[AgentRunConfigDict] + """Configuration for Agent Run in evaluation management service.""" - Args: - file_path: The path to the YAML file. - version: Optional version string to include in the YAML output. + agent_configs: Optional[dict[str, evals_types.AgentConfig]] + """A map of agent IDs to their respective agent config.""" - Raises: - ImportError: If the pyyaml library is not installed. - """ - if yaml is None: - raise ImportError( - "YAML serialization requires the pyyaml library. Please install" - " it using 'pip install google-cloud-aiplatform[evaluation]'." - ) - fields_to_exclude = { - field_name - for field_name, field_info in self.model_fields.items() - if self.__getattribute__(field_name) is not None - and isinstance(self.__getattribute__(field_name), Callable) - } +EvaluationRunInferenceConfigOrDict = Union[ + EvaluationRunInferenceConfig, EvaluationRunInferenceConfigDict +] - data_to_dump = self.model_dump( - exclude_unset=True, - exclude_none=True, - mode="json", - exclude=fields_to_exclude if fields_to_exclude else None, - ) - if version: - data_to_dump["version"] = version +class VulnerableTool(_common.BaseModel): + """A tool considered high risk for prompt injection.""" - with open(file_path, "w", encoding="utf-8") as f: - yaml.dump(data_to_dump, f, sort_keys=False, allow_unicode=True) + tool_name: Optional[str] = Field( + default=None, + description="""Optional. The name of the vulnerable function/tool (e.g., "search_flights").""", + ) + json_paths: Optional[list[str]] = Field( + default=None, + description="""Optional. JSON Paths within the tool's FunctionResponse where malicious content could be injected.""", + ) -class CodeExecutionMetric(Metric): - """A metric that executes custom Python code for evaluation.""" +class VulnerableToolDict(TypedDict, total=False): + """A tool considered high risk for prompt injection.""" - # You can use standard Pydantic Field syntax here because this is raw Python code - custom_function: Optional[str] = Field( - default=None, - description="""The Python function code to be executed on the server side.""", - ) + tool_name: Optional[str] + """Optional. The name of the vulnerable function/tool (e.g., "search_flights").""" - # You can also add hand-written validators or methods here - @field_validator("custom_function") - @classmethod - def validate_code(cls, value: Optional[str]) -> Optional[str]: - if value and "def evaluate" not in value: - raise ValueError( - "custom_function must contain a 'def evaluate(instance):' signature." - ) - return value + json_paths: Optional[list[str]] + """Optional. JSON Paths within the tool's FunctionResponse where malicious content could be injected.""" -class LLMMetric(Metric): - """A metric that uses LLM-as-a-judge for evaluation.""" +VulnerableToolOrDict = Union[VulnerableTool, VulnerableToolDict] - rubric_group_name: Optional[str] = Field( + +class RedTeamingAnalysisConfig(_common.BaseModel): + """Configuration for the automated Agent Red Teaming analysis.""" + + attack_categories: Optional[list[str]] = Field( default=None, - description="""Optional. The name of the column in the EvaluationDataset containing the list of rubrics to use for this metric.""", + description="""Optional. Specific attack categories to test against.""", ) - - result_parsing_function: Optional[str] = Field( + vulnerable_tools: Optional[list[VulnerableTool]] = Field( default=None, - description="""Optional. A Python function string used to parse the raw output of the LLM judge model. The function must be named `parse_results` and accept a list of model response strings. It should return a dictionary with `score` (float) and `explanation` (str) keys.""", + description="""Optional. Manually defined vulnerable tools and their injection paths.""", ) - @field_validator("prompt_template", mode="before") - @classmethod - def validate_prompt_template(cls, value: Union[str, "MetricPromptBuilder"]) -> str: - """Validates prompt template to be a non-empty string.""" - if value is None: - raise ValueError("Prompt template cannot be empty.") - if isinstance(value, MetricPromptBuilder): - value = str(value) - if not value.strip(): - raise ValueError("Prompt template cannot be an empty string.") - return value - - @field_validator("judge_model_sampling_count") - @classmethod - def validate_judge_model_sampling_count(cls, value: Optional[int]) -> Optional[int]: - """Validates judge_model_sampling_count to be between 1 and 32.""" - if value is not None and (value < 1 or value > 32): - raise ValueError("judge_model_sampling_count must be between 1 and 32.") - return value - @classmethod - def load(cls, config_path: str, client: Optional[Any] = None) -> "LLMMetric": - """Loads a metric configuration from a YAML or JSON file. +class RedTeamingAnalysisConfigDict(TypedDict, total=False): + """Configuration for the automated Agent Red Teaming analysis.""" - This method allows for the creation of an LLMMetric instance from a - local file path or a Google Cloud Storage (GCS) URI. It will automatically - detect the file type (.yaml, .yml, or .json) and parse it accordingly. + attack_categories: Optional[list[str]] + """Optional. Specific attack categories to test against.""" - Args: - config_path: The local path or GCS URI (e.g., 'gs://bucket/metric.yaml') - to the metric configuration file. - client: Optional. The Vertex AI client instance to use for authentication. - If not provided, Application Default Credentials (ADC) will be used. + vulnerable_tools: Optional[list[VulnerableToolDict]] + """Optional. Manually defined vulnerable tools and their injection paths.""" - Returns: - An instance of LLMMetric configured with the loaded data. - Raises: - ValueError: If the file path is invalid or the file content cannot be parsed. - ImportError: If a required library like 'PyYAML' or 'google-cloud-storage' is not installed. - IOError: If the file cannot be read from the specified path. - """ - file_extension = os.path.splitext(config_path)[1].lower() - if file_extension not in [".yaml", ".yml", ".json"]: - raise ValueError( - "Unsupported file extension for metric config. Must be .yaml, .yml, or .json" - ) +RedTeamingAnalysisConfigOrDict = Union[ + RedTeamingAnalysisConfig, RedTeamingAnalysisConfigDict +] - content_str: str - if config_path.startswith("gs://"): - try: - from google.cloud import storage # type: ignore[attr-defined] - storage_client = storage.Client( - credentials=client._api_client._credentials if client else None - ) - path_without_prefix = config_path[len("gs://") :] - bucket_name, blob_path = path_without_prefix.split("/", 1) +class AnalysisConfig(_common.BaseModel): + """Configuration for an analysis to be performed on an evaluation run.""" - bucket = storage_client.bucket(bucket_name) - blob = bucket.blob(blob_path) - content_str = blob.download_as_bytes().decode("utf-8") - except ImportError as e: - raise ImportError( - "Reading from GCS requires the 'google-cloud-storage' library. Please install it with 'pip install google-cloud-aiplatform[evaluation]'." - ) from e - except Exception as e: - raise IOError(f"Failed to read from GCS path {config_path}: {e}") from e - else: - try: - with open(config_path, "r", encoding="utf-8") as f: - content_str = f.read() - except FileNotFoundError: - raise FileNotFoundError( - f"Local configuration file not found at: {config_path}" - ) - except Exception as e: - raise IOError(f"Failed to read local file {config_path}: {e}") from e + analysis_name: Optional[str] = Field( + default=None, description="""Optional. A name for this analysis.""" + ) + red_teaming_analysis_config: Optional[RedTeamingAnalysisConfig] = Field( + default=None, + description="""Configuration for the automated Agent Red Teaming analysis.""", + ) - data: Dict[str, Any] - if file_extension in [".yaml", ".yml"]: - if yaml is None: - raise ImportError( - "YAML parsing requires the pyyaml library. Please install it with 'pip install google-cloud-aiplatform[evaluation]'." - ) - data = yaml.safe_load(content_str) - elif file_extension == ".json": - data = json.loads(content_str) +class AnalysisConfigDict(TypedDict, total=False): + """Configuration for an analysis to be performed on an evaluation run.""" - if not isinstance(data, dict): - raise ValueError("Metric config content did not parse into a dictionary.") + analysis_name: Optional[str] + """Optional. A name for this analysis.""" - return cls.model_validate(data) + red_teaming_analysis_config: Optional[RedTeamingAnalysisConfigDict] + """Configuration for the automated Agent Red Teaming analysis.""" -class MetricDict(TypedDict, total=False): - """The metric used for evaluation.""" +AnalysisConfigOrDict = Union[AnalysisConfig, AnalysisConfigDict] - name: Optional[str] - """The name of the metric.""" - custom_function: Optional[Union[str, Callable[..., Any]]] - """The custom function that defines the end-to-end logic for metric computation.""" +class CreateEvaluationRunConfig(_common.BaseModel): + """Config to create an evaluation run.""" - prompt_template: Optional[str] - """The prompt template for the metric.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + allow_cross_region_model: Optional[bool] = Field( + default=None, + description="""Allows the evaluation run to use cross region models. When this + flag is set, the service may route traffic to other regions if a model is + unavailable in the current region (e.g., to a `global`endpoint). If a + fully-qualified model endpoint resource name with a different region than + the run location is provided elsewhere in the run config, this flag must + be set to true or the request will fail.""", + ) - judge_model_system_instruction: Optional[str] - """The system instruction for the judge model.""" - return_raw_output: Optional[bool] - """Whether to return the raw output from the judge model.""" +class CreateEvaluationRunConfigDict(TypedDict, total=False): + """Config to create an evaluation run.""" - parse_and_reduce_fn: Optional[Callable[..., Any]] - """The parse and reduce function for the judge model.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - aggregate_summary_fn: Optional[Callable[..., Any]] - """The aggregate summary function for the judge model.""" + allow_cross_region_model: Optional[bool] + """Allows the evaluation run to use cross region models. When this + flag is set, the service may route traffic to other regions if a model is + unavailable in the current region (e.g., to a `global`endpoint). If a + fully-qualified model endpoint resource name with a different region than + the run location is provided elsewhere in the run config, this flag must + be set to true or the request will fail.""" - remote_custom_function: Optional[str] - """The evaluation function for the custom code execution metric. This custom code is run remotely in the evaluation service.""" - judge_model: Optional[str] - """The judge model for the metric.""" +CreateEvaluationRunConfigOrDict = Union[ + CreateEvaluationRunConfig, CreateEvaluationRunConfigDict +] - judge_model_generation_config: Optional[genai_types.GenerationConfig] - """The generation config for the judge LLM (temperature, top_k, top_p, etc).""" - judge_model_sampling_count: Optional[int] - """The sampling count for the judge model.""" +class _CreateEvaluationRunParameters(_common.BaseModel): + """Represents a job that creates an evaluation run.""" - rubric_group_name: Optional[str] - """The rubric group name for the rubric-based metric.""" + name: Optional[str] = Field(default=None, description="""""") + display_name: Optional[str] = Field(default=None, description="""""") + data_source: Optional[EvaluationRunDataSource] = Field( + default=None, description="""""" + ) + evaluation_config: Optional[EvaluationRunConfig] = Field( + default=None, description="""""" + ) + labels: Optional[dict[str, str]] = Field(default=None, description="""""") + inference_configs: Optional[dict[str, EvaluationRunInferenceConfig]] = Field( + default=None, description="""""" + ) + config: Optional[CreateEvaluationRunConfig] = Field( + default=None, description="""""" + ) + analysis_configs: Optional[list[AnalysisConfig]] = Field( + default=None, description="""""" + ) + evaluation_experiment: Optional[str] = Field( + default=None, + description="""The resource name of the parent EvaluationExperiment that this run + belongs to. Format: + `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""", + ) - metric_spec_parameters: Optional[dict[str, Any]] - """Optional steering instruction parameters for the automated predefined metric.""" - metric_resource_name: Optional[str] - """The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" +class _CreateEvaluationRunParametersDict(TypedDict, total=False): + """Represents a job that creates an evaluation run.""" - result_parsing_function: Optional[str] - """Optional. A Python function string used to parse the raw output of the LLM judge model. The function must be named `parse_results` and accept a list of model response strings. It should return a dictionary with `score` (float) and `explanation` (str) keys.""" + name: Optional[str] + """""" + display_name: Optional[str] + """""" -MetricOrDict = Union[Metric, MetricDict] + data_source: Optional[EvaluationRunDataSourceDict] + """""" + evaluation_config: Optional[EvaluationRunConfigDict] + """""" -class CreateEvaluationMetricConfig(_common.BaseModel): - """Config for creating an evaluation metric.""" + labels: Optional[dict[str, str]] + """""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + inference_configs: Optional[dict[str, EvaluationRunInferenceConfigDict]] + """""" + config: Optional[CreateEvaluationRunConfigDict] + """""" -class CreateEvaluationMetricConfigDict(TypedDict, total=False): - """Config for creating an evaluation metric.""" + analysis_configs: Optional[list[AnalysisConfigDict]] + """""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + evaluation_experiment: Optional[str] + """The resource name of the parent EvaluationExperiment that this run + belongs to. Format: + `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""" -CreateEvaluationMetricConfigOrDict = Union[ - CreateEvaluationMetricConfig, CreateEvaluationMetricConfigDict +_CreateEvaluationRunParametersOrDict = Union[ + _CreateEvaluationRunParameters, _CreateEvaluationRunParametersDict ] -class _CreateEvaluationMetricParameters(_common.BaseModel): - """Parameters for creating an evaluation metric.""" - - display_name: Optional[str] = Field( - default=None, - description="""The user-defined name of the evaluation metric. +class SummaryMetric(_common.BaseModel): + """Represents a summary metric for an evaluation run.""" - The display name can be up to 128 characters long and can comprise any - UTF-8 characters. - """, - ) - description: Optional[str] = Field( - default=None, description="""The description of the evaluation metric.""" + metrics: Optional[dict[str, Any]] = Field( + default=None, description="""Map of metric name to metric value.""" ) - metric: Optional[Metric] = Field( - default=None, - description="""The metric configuration of the evaluation metric.""", + total_items: Optional[int] = Field( + default=None, description="""The total number of items that were evaluated.""" ) - config: Optional[CreateEvaluationMetricConfig] = Field( - default=None, description="""""" + failed_items: Optional[int] = Field( + default=None, description="""The number of items that failed to be evaluated.""" ) -class _CreateEvaluationMetricParametersDict(TypedDict, total=False): - """Parameters for creating an evaluation metric.""" - - display_name: Optional[str] - """The user-defined name of the evaluation metric. - - The display name can be up to 128 characters long and can comprise any - UTF-8 characters. - """ - - description: Optional[str] - """The description of the evaluation metric.""" +class SummaryMetricDict(TypedDict, total=False): + """Represents a summary metric for an evaluation run.""" - metric: Optional[MetricDict] - """The metric configuration of the evaluation metric.""" + metrics: Optional[dict[str, Any]] + """Map of metric name to metric value.""" - config: Optional[CreateEvaluationMetricConfigDict] - """""" + total_items: Optional[int] + """The total number of items that were evaluated.""" + failed_items: Optional[int] + """The number of items that failed to be evaluated.""" -_CreateEvaluationMetricParametersOrDict = Union[ - _CreateEvaluationMetricParameters, _CreateEvaluationMetricParametersDict -] +SummaryMetricOrDict = Union[SummaryMetric, SummaryMetricDict] -class CustomCodeExecutionSpec(_common.BaseModel): - """Specifies a metric using remote Python function execution. - This metric is computed by running user-defined Python functions remotely. - """ +class AttackCategoryResult(_common.BaseModel): + """The red teaming outcome for a specific attack category.""" - evaluation_function: Optional[str] = Field( - default=None, - description="""Required. Python function. Expected user to define the following function, e.g.: def evaluate(instance: dict[str, Any]) -> float: Please include this function signature in the code snippet. Instance is the evaluation instance, any fields populated in the instance are available to the function as instance[field_name]. Example: Example input: ``` instance= EvaluationInstance( response=EvaluationInstance.InstanceData(text="The answer is 4."), reference=EvaluationInstance.InstanceData(text="4") ) ``` Example converted input: ``` { 'response': {'text': 'The answer is 4.'}, 'reference': {'text': '4'} } ``` Example python function: ``` def evaluate(instance: dict[str, Any]) -> float: if instance'response' == instance'reference': return 1.0 return 0.0 ``` CustomCodeExecutionSpec is also supported in Batch Evaluation (EvalDataset RPC) and Tuning Evaluation. Each line in the input jsonl file will be converted to dict[str, Any] and passed to the evaluation function.""", + attack_category: Optional[str] = Field( + default=None, description="""The category of the attack evaluated.""" ) - remote_custom_function: Optional[str] = Field( + attack_success_rate: Optional[float] = Field( default=None, - description="""A string representing a user-defined function for evaluation. - Expected user to define the following function, e.g.: - def evaluate(instance: dict[str, Any]) -> float: - Please include this function signature in the code snippet. - Instance is the evaluation instance, any fields populated in the instance - are available to the function as instance[field_name].""", + description="""The ratio of successful attacks given a fixed budget.""", + ) + vulnerability_insight: Optional[str] = Field( + default=None, description="""Insights into why an attack succeeded or failed.""" ) -class CustomCodeExecutionSpecDict(TypedDict, total=False): - """Specifies a metric using remote Python function execution. +class AttackCategoryResultDict(TypedDict, total=False): + """The red teaming outcome for a specific attack category.""" - This metric is computed by running user-defined Python functions remotely. - """ + attack_category: Optional[str] + """The category of the attack evaluated.""" - evaluation_function: Optional[str] - """Required. Python function. Expected user to define the following function, e.g.: def evaluate(instance: dict[str, Any]) -> float: Please include this function signature in the code snippet. Instance is the evaluation instance, any fields populated in the instance are available to the function as instance[field_name]. Example: Example input: ``` instance= EvaluationInstance( response=EvaluationInstance.InstanceData(text="The answer is 4."), reference=EvaluationInstance.InstanceData(text="4") ) ``` Example converted input: ``` { 'response': {'text': 'The answer is 4.'}, 'reference': {'text': '4'} } ``` Example python function: ``` def evaluate(instance: dict[str, Any]) -> float: if instance'response' == instance'reference': return 1.0 return 0.0 ``` CustomCodeExecutionSpec is also supported in Batch Evaluation (EvalDataset RPC) and Tuning Evaluation. Each line in the input jsonl file will be converted to dict[str, Any] and passed to the evaluation function.""" + attack_success_rate: Optional[float] + """The ratio of successful attacks given a fixed budget.""" - remote_custom_function: Optional[str] - """A string representing a user-defined function for evaluation. - Expected user to define the following function, e.g.: - def evaluate(instance: dict[str, Any]) -> float: - Please include this function signature in the code snippet. - Instance is the evaluation instance, any fields populated in the instance - are available to the function as instance[field_name].""" + vulnerability_insight: Optional[str] + """Insights into why an attack succeeded or failed.""" -CustomCodeExecutionSpecOrDict = Union[ - CustomCodeExecutionSpec, CustomCodeExecutionSpecDict -] +AttackCategoryResultOrDict = Union[AttackCategoryResult, AttackCategoryResultDict] -class UnifiedMetric(_common.BaseModel): - """The unified metric used for evaluation.""" +class RedTeamingAnalysisResult(_common.BaseModel): + """The top-level result for Red Teaming analysis.""" - bleu_spec: Optional[genai_types.BleuSpec] = Field( - default=None, description="""The Bleu metric spec.""" - ) - rouge_spec: Optional[genai_types.RougeSpec] = Field( - default=None, description="""The rouge metric spec.""" - ) - pointwise_metric_spec: Optional[genai_types.PointwiseMetricSpec] = Field( - default=None, description="""The pointwise metric spec.""" - ) - llm_based_metric_spec: Optional[genai_types.LLMBasedMetricSpec] = Field( - default=None, description="""The spec for an LLM based metric.""" - ) - custom_code_execution_spec: Optional[CustomCodeExecutionSpec] = Field( - default=None, description="""The spec for a custom code execution metric.""" + config: Optional[RedTeamingAnalysisConfig] = Field( + default=None, + description="""The configuration used to generate this analysis.""", ) - predefined_metric_spec: Optional[genai_types.PredefinedMetricSpec] = Field( - default=None, description="""The spec for a pre-defined metric.""" + analysis_time: Optional[str] = Field( + default=None, description="""The timestamp when this analysis was performed.""" ) - computation_based_metric_spec: Optional[genai_types.ComputationBasedMetricSpec] = ( - Field(default=None, description="""The spec for a computation based metric.""") + category_results: Optional[list[AttackCategoryResult]] = Field( + default=None, description="""Detailed results by attack category.""" ) -class UnifiedMetricDict(TypedDict, total=False): - """The unified metric used for evaluation.""" - - bleu_spec: Optional[genai_types.BleuSpec] - """The Bleu metric spec.""" - - rouge_spec: Optional[genai_types.RougeSpec] - """The rouge metric spec.""" - - pointwise_metric_spec: Optional[genai_types.PointwiseMetricSpec] - """The pointwise metric spec.""" - - llm_based_metric_spec: Optional[genai_types.LLMBasedMetricSpec] - """The spec for an LLM based metric.""" +class RedTeamingAnalysisResultDict(TypedDict, total=False): + """The top-level result for Red Teaming analysis.""" - custom_code_execution_spec: Optional[CustomCodeExecutionSpecDict] - """The spec for a custom code execution metric.""" + config: Optional[RedTeamingAnalysisConfigDict] + """The configuration used to generate this analysis.""" - predefined_metric_spec: Optional[genai_types.PredefinedMetricSpec] - """The spec for a pre-defined metric.""" + analysis_time: Optional[str] + """The timestamp when this analysis was performed.""" - computation_based_metric_spec: Optional[genai_types.ComputationBasedMetricSpec] - """The spec for a computation based metric.""" + category_results: Optional[list[AttackCategoryResultDict]] + """Detailed results by attack category.""" -UnifiedMetricOrDict = Union[UnifiedMetric, UnifiedMetricDict] +RedTeamingAnalysisResultOrDict = Union[ + RedTeamingAnalysisResult, RedTeamingAnalysisResultDict +] -class EvaluationMetric(_common.BaseModel): - """Represents an evaluation metric.""" +class LossTaxonomyEntry(_common.BaseModel): + """A specific entry in the loss pattern taxonomy.""" - name: Optional[str] = Field( - default=None, description="""The resource name of the evaluation metric.""" + l1_category: Optional[str] = Field( + default=None, + description="""The primary category of the loss (e.g., "Hallucination", "Tool Calling").""", ) - display_name: Optional[str] = Field( + l2_category: Optional[str] = Field( default=None, - description="""The user-friendly display name for the EvaluationMetric.""", + description="""The secondary category of the loss (e.g., "Hallucination of Action", "Incorrect Tool Selection").""", ) description: Optional[str] = Field( - default=None, description="""The description of the EvaluationMetric.""" - ) - metric: Optional[UnifiedMetric] = Field( default=None, - description="""The metric configuration of the evaluation metric.""", + description="""A detailed description of this loss pattern. Example: "The agent verbally confirms an action without executing the tool." """, ) -class EvaluationMetricDict(TypedDict, total=False): - """Represents an evaluation metric.""" +class LossTaxonomyEntryDict(TypedDict, total=False): + """A specific entry in the loss pattern taxonomy.""" - name: Optional[str] - """The resource name of the evaluation metric.""" + l1_category: Optional[str] + """The primary category of the loss (e.g., "Hallucination", "Tool Calling").""" - display_name: Optional[str] - """The user-friendly display name for the EvaluationMetric.""" + l2_category: Optional[str] + """The secondary category of the loss (e.g., "Hallucination of Action", "Incorrect Tool Selection").""" description: Optional[str] - """The description of the EvaluationMetric.""" - - metric: Optional[UnifiedMetricDict] - """The metric configuration of the evaluation metric.""" - + """A detailed description of this loss pattern. Example: "The agent verbally confirms an action without executing the tool." """ -EvaluationMetricOrDict = Union[EvaluationMetric, EvaluationMetricDict] +LossTaxonomyEntryOrDict = Union[LossTaxonomyEntry, LossTaxonomyEntryDict] -class SamplingConfig(_common.BaseModel): - """Sampling config for a BigQuery request set.""" - sampling_count: Optional[int] = Field(default=None, description="""""") - sampling_method: Optional[SamplingMethod] = Field(default=None, description="""""") - sampling_duration: Optional[str] = Field(default=None, description="""""") +class FailedRubric(_common.BaseModel): + """A specific failed rubric and the associated analysis.""" + rubric_id: Optional[str] = Field( + default=None, + description="""The unique ID of the rubric (if available from the metric source).""", + ) + classification_rationale: Optional[str] = Field( + default=None, + description="""The rationale provided by the Loss Analysis Classifier for why this failure maps to this specific Loss Cluster.""", + ) -class SamplingConfigDict(TypedDict, total=False): - """Sampling config for a BigQuery request set.""" - sampling_count: Optional[int] - """""" +class FailedRubricDict(TypedDict, total=False): + """A specific failed rubric and the associated analysis.""" - sampling_method: Optional[SamplingMethod] - """""" + rubric_id: Optional[str] + """The unique ID of the rubric (if available from the metric source).""" - sampling_duration: Optional[str] - """""" + classification_rationale: Optional[str] + """The rationale provided by the Loss Analysis Classifier for why this failure maps to this specific Loss Cluster.""" -SamplingConfigOrDict = Union[SamplingConfig, SamplingConfigDict] +FailedRubricOrDict = Union[FailedRubric, FailedRubricDict] -class BigQueryRequestSet(_common.BaseModel): - """Represents a BigQuery request set.""" +class LossExample(_common.BaseModel): + """A specific example of a loss pattern.""" - uri: Optional[str] = Field(default=None, description="""""") - prompt_column: Optional[str] = Field( - default=None, - description="""The column name of the prompt in the BigQuery table. Used for EvaluationRun only.""", - ) - rubrics_column: Optional[str] = Field( + evaluation_item: Optional[str] = Field( default=None, - description="""The column name of the rubrics in the BigQuery table. Used for EvaluationRun only.""", + description="""Reference to the persisted EvalItem resource name. Format: projects/.../locations/.../evaluationItems/{item_id}.""", ) - candidate_response_columns: Optional[dict[str, str]] = Field( + evaluation_result: Optional[dict[str, Any]] = Field( default=None, - description="""The column name of the response candidates in the BigQuery table. Used for EvaluationRun only.""", + description="""The full evaluation result object provided inline. Used when the analysis is performed on ephemeral data.""", ) - sampling_config: Optional[SamplingConfig] = Field( + failed_rubrics: Optional[list[FailedRubric]] = Field( default=None, - description="""The sampling config for the BigQuery request set. Used for EvaluationRun only.""", + description="""The specific rubric(s) that failed and caused this example to be classified here. An example might fail multiple rubrics, but only specific ones trigger this loss pattern.""", ) -class BigQueryRequestSetDict(TypedDict, total=False): - """Represents a BigQuery request set.""" +class LossExampleDict(TypedDict, total=False): + """A specific example of a loss pattern.""" - uri: Optional[str] - """""" + evaluation_item: Optional[str] + """Reference to the persisted EvalItem resource name. Format: projects/.../locations/.../evaluationItems/{item_id}.""" - prompt_column: Optional[str] - """The column name of the prompt in the BigQuery table. Used for EvaluationRun only.""" + evaluation_result: Optional[dict[str, Any]] + """The full evaluation result object provided inline. Used when the analysis is performed on ephemeral data.""" - rubrics_column: Optional[str] - """The column name of the rubrics in the BigQuery table. Used for EvaluationRun only.""" + failed_rubrics: Optional[list[FailedRubricDict]] + """The specific rubric(s) that failed and caused this example to be classified here. An example might fail multiple rubrics, but only specific ones trigger this loss pattern.""" - candidate_response_columns: Optional[dict[str, str]] - """The column name of the response candidates in the BigQuery table. Used for EvaluationRun only.""" - sampling_config: Optional[SamplingConfigDict] - """The sampling config for the BigQuery request set. Used for EvaluationRun only.""" +LossExampleOrDict = Union[LossExample, LossExampleDict] -BigQueryRequestSetOrDict = Union[BigQueryRequestSet, BigQueryRequestSetDict] +class LossCluster(_common.BaseModel): + """A semantic grouping of failures (e.g., "Hallucination of Action").""" + cluster_id: Optional[str] = Field( + default=None, + description="""Unique identifier for the loss cluster within the scope of the analysis result.""", + ) + taxonomy_entry: Optional[LossTaxonomyEntry] = Field( + default=None, + description="""The structured definition of the loss taxonomy for this cluster.""", + ) + item_count: Optional[int] = Field( + default=None, + description="""The total number of EvaluationItems falling into this cluster.""", + ) + examples: Optional[list[LossExample]] = Field( + default=None, + description="""A list of examples that belong to this cluster. This links the cluster back to the specific EvaluationItems and Rubrics.""", + ) -class EvaluationRunDataSource(_common.BaseModel): - """Represents an evaluation run data source.""" - evaluation_set: Optional[str] = Field(default=None, description="""""") - bigquery_request_set: Optional[BigQueryRequestSet] = Field( - default=None, description="""""" - ) +class LossClusterDict(TypedDict, total=False): + """A semantic grouping of failures (e.g., "Hallucination of Action").""" + cluster_id: Optional[str] + """Unique identifier for the loss cluster within the scope of the analysis result.""" -class EvaluationRunDataSourceDict(TypedDict, total=False): - """Represents an evaluation run data source.""" + taxonomy_entry: Optional[LossTaxonomyEntryDict] + """The structured definition of the loss taxonomy for this cluster.""" - evaluation_set: Optional[str] - """""" + item_count: Optional[int] + """The total number of EvaluationItems falling into this cluster.""" - bigquery_request_set: Optional[BigQueryRequestSetDict] - """""" + examples: Optional[list[LossExampleDict]] + """A list of examples that belong to this cluster. This links the cluster back to the specific EvaluationItems and Rubrics.""" -EvaluationRunDataSourceOrDict = Union[ - EvaluationRunDataSource, EvaluationRunDataSourceDict -] +LossClusterOrDict = Union[LossCluster, LossClusterDict] -class EvaluationRunMetric(_common.BaseModel): - """The metric used for evaluation run.""" +class LossAnalysisResult(_common.BaseModel): + """The top-level result for loss analysis.""" - metric: Optional[str] = Field( - default=None, description="""The name of the metric.""" - ) - metric_resource_name: Optional[str] = Field( + config: Optional[LossAnalysisConfig] = Field( default=None, - description="""The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""", + description="""The configuration used to generate this analysis.""", ) - metric_config: Optional[UnifiedMetric] = Field( - default=None, description="""The unified metric used for evaluation run.""" + analysis_time: Optional[str] = Field( + default=None, description="""The timestamp when this analysis was performed.""" + ) + clusters: Optional[list[LossCluster]] = Field( + default=None, description="""The list of identified loss clusters.""" ) + def show(self) -> None: + """Shows the loss analysis result with rich HTML visualization.""" + from .. import _evals_visualization -class EvaluationRunMetricDict(TypedDict, total=False): - """The metric used for evaluation run.""" + _evals_visualization.display_loss_analysis_result(self) - metric: Optional[str] - """The name of the metric.""" - metric_resource_name: Optional[str] - """The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" +class LossAnalysisResultDict(TypedDict, total=False): + """The top-level result for loss analysis.""" - metric_config: Optional[UnifiedMetricDict] - """The unified metric used for evaluation run.""" + config: Optional[LossAnalysisConfigDict] + """The configuration used to generate this analysis.""" + analysis_time: Optional[str] + """The timestamp when this analysis was performed.""" -EvaluationRunMetricOrDict = Union[EvaluationRunMetric, EvaluationRunMetricDict] + clusters: Optional[list[LossClusterDict]] + """The list of identified loss clusters.""" -class EvaluationRunPromptTemplate(_common.BaseModel): - """Prompt template used for inference. +LossAnalysisResultOrDict = Union[LossAnalysisResult, LossAnalysisResultDict] - Only one of `prompt_template` or `gcs_uri` should be set. If both are - provided, an error will be raised. - """ - prompt_template: Optional[str] = Field( +class EvaluationRunResults(_common.BaseModel): + """Represents the results of an evaluation run.""" + + evaluation_set: Optional[str] = Field( default=None, - description="""Inline prompt template. Template variables should be in the format - "{var_name}". Only one of `prompt_template` or `gcs_uri` should be set.""", + description="""The evaluation set where item level results are stored.""", ) - gcs_uri: Optional[str] = Field( + summary_metrics: Optional[SummaryMetric] = Field( + default=None, description="""The summary metrics for the evaluation run.""" + ) + loss_analysis_results: Optional[list[LossAnalysisResult]] = Field( default=None, - description="""Prompt template stored in Cloud Storage. Format: - "gs://my-bucket/file-name.txt". Only one of `prompt_template` or `gcs_uri` - should be set.""", + description="""The loss analysis results for the evaluation run.""", + ) + red_teaming_analysis_results: Optional[list[RedTeamingAnalysisResult]] = Field( + default=None, description="""The Red Teaming analysis results.""" ) -class EvaluationRunPromptTemplateDict(TypedDict, total=False): - """Prompt template used for inference. +class EvaluationRunResultsDict(TypedDict, total=False): + """Represents the results of an evaluation run.""" - Only one of `prompt_template` or `gcs_uri` should be set. If both are - provided, an error will be raised. - """ + evaluation_set: Optional[str] + """The evaluation set where item level results are stored.""" - prompt_template: Optional[str] - """Inline prompt template. Template variables should be in the format - "{var_name}". Only one of `prompt_template` or `gcs_uri` should be set.""" + summary_metrics: Optional[SummaryMetricDict] + """The summary metrics for the evaluation run.""" - gcs_uri: Optional[str] - """Prompt template stored in Cloud Storage. Format: - "gs://my-bucket/file-name.txt". Only one of `prompt_template` or `gcs_uri` - should be set.""" + loss_analysis_results: Optional[list[LossAnalysisResultDict]] + """The loss analysis results for the evaluation run.""" + red_teaming_analysis_results: Optional[list[RedTeamingAnalysisResultDict]] + """The Red Teaming analysis results.""" -EvaluationRunPromptTemplateOrDict = Union[ - EvaluationRunPromptTemplate, EvaluationRunPromptTemplateDict -] +EvaluationRunResultsOrDict = Union[EvaluationRunResults, EvaluationRunResultsDict] -class LossAnalysisConfig(_common.BaseModel): - """Configuration for the loss analysis job.""" - metric: Optional[str] = Field( - default=None, - description="""Required. The metric to analyze (e.g., "multi_turn_tool_use_quality_v1").""", +class EvalCaseMetricResult(_common.BaseModel): + """Evaluation result for a single evaluation case for a single metric.""" + + metric_name: Optional[str] = Field( + default=None, description="""Name of the metric.""" ) - candidate: Optional[str] = Field( - default=None, - description="""Required. The candidate model/agent to analyze (e.g., "gemini-3.1-pro-preview"). This targets the specific CandidateResult within the EvaluationResult.""", + score: Optional[float] = Field(default=None, description="""Score of the metric.""") + explanation: Optional[str] = Field( + default=None, description="""Explanation of the metric.""" ) - predefined_taxonomy: Optional[str] = Field( + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( default=None, - description="""Optional. The identifier for the pre-defined taxonomy to use (e.g., "agent_taxonomy_v1", "tool_use_v2"). If not specified, the service may select a default based on the metric.""", + description="""The details of all the rubrics and their verdicts for rubric-based metrics.""", ) - max_top_cluster_count: Optional[int] = Field( - default=None, - description="""Optional. Limits the analysis to the top N clusters. If not specified or set to 0, all clusters are returned.""", + raw_output: Optional[list[str]] = Field( + default=None, description="""Raw output of the metric.""" + ) + error_message: Optional[str] = Field( + default=None, description="""Error message for the metric.""" ) -class LossAnalysisConfigDict(TypedDict, total=False): - """Configuration for the loss analysis job.""" - - metric: Optional[str] - """Required. The metric to analyze (e.g., "multi_turn_tool_use_quality_v1").""" +class EvalCaseMetricResultDict(TypedDict, total=False): + """Evaluation result for a single evaluation case for a single metric.""" - candidate: Optional[str] - """Required. The candidate model/agent to analyze (e.g., "gemini-3.1-pro-preview"). This targets the specific CandidateResult within the EvaluationResult.""" + metric_name: Optional[str] + """Name of the metric.""" - predefined_taxonomy: Optional[str] - """Optional. The identifier for the pre-defined taxonomy to use (e.g., "agent_taxonomy_v1", "tool_use_v2"). If not specified, the service may select a default based on the metric.""" + score: Optional[float] + """Score of the metric.""" - max_top_cluster_count: Optional[int] - """Optional. Limits the analysis to the top N clusters. If not specified or set to 0, all clusters are returned.""" + explanation: Optional[str] + """Explanation of the metric.""" + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] + """The details of all the rubrics and their verdicts for rubric-based metrics.""" -LossAnalysisConfigOrDict = Union[LossAnalysisConfig, LossAnalysisConfigDict] + raw_output: Optional[list[str]] + """Raw output of the metric.""" + error_message: Optional[str] + """Error message for the metric.""" -class EvaluationRunConfig(_common.BaseModel): - """The evaluation configuration used for the evaluation run.""" - metrics: Optional[list[EvaluationRunMetric]] = Field( - default=None, - description="""The metrics to be calculated in the evaluation run.""", - ) - output_config: Optional[genai_types.OutputConfig] = Field( - default=None, description="""The output config for the evaluation run.""" - ) - autorater_config: Optional[genai_types.AutoraterConfig] = Field( - default=None, - description="""The autorater config for the evaluation run. Not applicable for predefined metrics (PredefinedMetricSpec); the server uses its own model configuration for predefined metrics and this field is ignored.""", - ) - prompt_template: Optional[EvaluationRunPromptTemplate] = Field( - default=None, description="""The prompt template used for inference.""" - ) - loss_analysis_config: Optional[list[LossAnalysisConfig]] = Field( +EvalCaseMetricResultOrDict = Union[EvalCaseMetricResult, EvalCaseMetricResultDict] + + +class ResponseCandidateResult(_common.BaseModel): + """Aggregated metric results for a single response candidate.""" + + response_index: Optional[int] = Field( default=None, - description="""Specifications for loss analysis. Each config specifies a metric and candidate to analyze for loss patterns.""", + description="""Index of the response candidate this result pertains to.""", ) - allow_cross_region_model: Optional[bool] = Field( + metric_results: Optional[dict[str, EvalCaseMetricResult]] = Field( default=None, - description="""Allows the evaluation run to use cross region models. When this - flag is set, the service may route traffic to other regions if a model is - unavailable in the current region (e.g., to a `global`endpoint). If a - fully-qualified model endpoint resource name with a different region than - the run location is provided elsewhere in the run config, this flag must - be set to true or the request will fail.""", + description="""A dictionary of metric results for this response candidate, keyed by metric name.""", ) -class EvaluationRunConfigDict(TypedDict, total=False): - """The evaluation configuration used for the evaluation run.""" +class ResponseCandidateResultDict(TypedDict, total=False): + """Aggregated metric results for a single response candidate.""" - metrics: Optional[list[EvaluationRunMetricDict]] - """The metrics to be calculated in the evaluation run.""" + response_index: Optional[int] + """Index of the response candidate this result pertains to.""" - output_config: Optional[genai_types.OutputConfig] - """The output config for the evaluation run.""" + metric_results: Optional[dict[str, EvalCaseMetricResultDict]] + """A dictionary of metric results for this response candidate, keyed by metric name.""" - autorater_config: Optional[genai_types.AutoraterConfig] - """The autorater config for the evaluation run. Not applicable for predefined metrics (PredefinedMetricSpec); the server uses its own model configuration for predefined metrics and this field is ignored.""" - prompt_template: Optional[EvaluationRunPromptTemplateDict] - """The prompt template used for inference.""" +ResponseCandidateResultOrDict = Union[ + ResponseCandidateResult, ResponseCandidateResultDict +] - loss_analysis_config: Optional[list[LossAnalysisConfigDict]] - """Specifications for loss analysis. Each config specifies a metric and candidate to analyze for loss patterns.""" - allow_cross_region_model: Optional[bool] - """Allows the evaluation run to use cross region models. When this - flag is set, the service may route traffic to other regions if a model is - unavailable in the current region (e.g., to a `global`endpoint). If a - fully-qualified model endpoint resource name with a different region than - the run location is provided elsewhere in the run config, this flag must - be set to true or the request will fail.""" +class EvalCaseResult(_common.BaseModel): + """Eval result for a single evaluation case.""" + eval_case_index: Optional[int] = Field( + default=None, description="""Index of the evaluation case.""" + ) + response_candidate_results: Optional[list[ResponseCandidateResult]] = Field( + default=None, + description="""A list of results, one for each response candidate of the EvalCase.""", + ) -EvaluationRunConfigOrDict = Union[EvaluationRunConfig, EvaluationRunConfigDict] +class EvalCaseResultDict(TypedDict, total=False): + """Eval result for a single evaluation case.""" -class EvaluationRunAgentConfig(_common.BaseModel): - """Agent config for an evaluation run.""" + eval_case_index: Optional[int] + """Index of the evaluation case.""" - developer_instruction: Optional[genai_types.Content] = Field( - default=None, description="""The developer instruction for the agent.""" - ) - tools: Optional[list[genai_types.Tool]] = Field( - default=None, description="""The tools available to the agent.""" - ) + response_candidate_results: Optional[list[ResponseCandidateResultDict]] + """A list of results, one for each response candidate of the EvalCase.""" -class EvaluationRunAgentConfigDict(TypedDict, total=False): - """Agent config for an evaluation run.""" +EvalCaseResultOrDict = Union[EvalCaseResult, EvalCaseResultDict] - developer_instruction: Optional[genai_types.Content] - """The developer instruction for the agent.""" - tools: Optional[list[genai_types.Tool]] - """The tools available to the agent.""" +class AggregatedMetricResult(_common.BaseModel): + """Evaluation result for a single metric for an evaluation dataset.""" + + metric_name: Optional[str] = Field( + default=None, description="""Name of the metric.""" + ) + num_cases_total: Optional[int] = Field( + default=None, description="""Total number of cases in the dataset.""" + ) + num_cases_valid: Optional[int] = Field( + default=None, description="""Number of valid cases in the dataset.""" + ) + num_cases_error: Optional[int] = Field( + default=None, description="""Number of cases with errors in the dataset.""" + ) + mean_score: Optional[float] = Field( + default=None, description="""Mean score of the metric.""" + ) + stdev_score: Optional[float] = Field( + default=None, description="""Standard deviation of the metric.""" + ) + pass_rate: Optional[float] = Field( + default=None, + description="""Pass rate of the adaptive rubric metric. Calculated as the number of cases where all criteria passed divided by the total number of valid cases. A case is passing if it has a score of 1.0.""", + ) + # Allow extra fields to support custom aggregation stats. + model_config = ConfigDict(extra="allow") -EvaluationRunAgentConfigOrDict = Union[ - EvaluationRunAgentConfig, EvaluationRunAgentConfigDict -] +class AggregatedMetricResultDict(TypedDict, total=False): + """Evaluation result for a single metric for an evaluation dataset.""" -class GeminiAgentConfig(_common.BaseModel): - """Config for scraping a Gemini Agent. + metric_name: Optional[str] + """Name of the metric.""" - A Gemini Agent is a Vertex AI Agent resource scraped via the Vertex - Interactions API. - """ + num_cases_total: Optional[int] + """Total number of cases in the dataset.""" - gemini_agent: Optional[str] = Field( - default=None, - description="""The resource name of the Gemini Agent. - Format: `projects/{project}/locations/{location}/agents/{agent}`.""", - ) + num_cases_valid: Optional[int] + """Number of valid cases in the dataset.""" + num_cases_error: Optional[int] + """Number of cases with errors in the dataset.""" -class GeminiAgentConfigDict(TypedDict, total=False): - """Config for scraping a Gemini Agent. + mean_score: Optional[float] + """Mean score of the metric.""" - A Gemini Agent is a Vertex AI Agent resource scraped via the Vertex - Interactions API. - """ + stdev_score: Optional[float] + """Standard deviation of the metric.""" - gemini_agent: Optional[str] - """The resource name of the Gemini Agent. - Format: `projects/{project}/locations/{location}/agents/{agent}`.""" + pass_rate: Optional[float] + """Pass rate of the adaptive rubric metric. Calculated as the number of cases where all criteria passed divided by the total number of valid cases. A case is passing if it has a score of 1.0.""" -GeminiAgentConfigOrDict = Union[GeminiAgentConfig, GeminiAgentConfigDict] +AggregatedMetricResultOrDict = Union[AggregatedMetricResult, AggregatedMetricResultDict] -class AgentRunConfig(_common.BaseModel): - """Configuration for an Agent Run.""" +class WinRateStats(_common.BaseModel): + """Statistics for win rates for a single metric.""" - session_input: Optional[evals_types.SessionInput] = Field( - default=None, description="""The session input to get agent running results.""" - ) - agent_engine: Optional[str] = Field( - default=None, description="""The resource name of the Agent Engine.""" - ) - user_simulator_config: Optional[evals_types.UserSimulatorConfig] = Field( + win_rates: Optional[list[float]] = Field( default=None, - description="""Used for multi-turn agent run. - Contains configuration for a user simulator that - uses an LLM to generate messages on behalf of the user.""", + description="""Win rates for the metric, one for each candidate.""", ) - gemini_agent_config: Optional[GeminiAgentConfig] = Field( - default=None, - description="""Config for scraping a Gemini Agent (Vertex AI Agent resource). - Used to target a Gemini agent for an evaluation run.""", + tie_rate: Optional[float] = Field( + default=None, description="""Tie rate for the metric.""" ) -class AgentRunConfigDict(TypedDict, total=False): - """Configuration for an Agent Run.""" - - session_input: Optional[evals_types.SessionInput] - """The session input to get agent running results.""" - - agent_engine: Optional[str] - """The resource name of the Agent Engine.""" +class WinRateStatsDict(TypedDict, total=False): + """Statistics for win rates for a single metric.""" - user_simulator_config: Optional[evals_types.UserSimulatorConfig] - """Used for multi-turn agent run. - Contains configuration for a user simulator that - uses an LLM to generate messages on behalf of the user.""" + win_rates: Optional[list[float]] + """Win rates for the metric, one for each candidate.""" - gemini_agent_config: Optional[GeminiAgentConfigDict] - """Config for scraping a Gemini Agent (Vertex AI Agent resource). - Used to target a Gemini agent for an evaluation run.""" + tie_rate: Optional[float] + """Tie rate for the metric.""" -AgentRunConfigOrDict = Union[AgentRunConfig, AgentRunConfigDict] +WinRateStatsOrDict = Union[WinRateStats, WinRateStatsDict] -class EvaluationRunInferenceConfig(_common.BaseModel): - """Configuration that describes an agent.""" +class ResponseCandidate(_common.BaseModel): + """A model-generated content to the prompt.""" - agent_config: Optional[EvaluationRunAgentConfig] = Field( - default=None, description="""The agent config.""" - ) - model: Optional[str] = Field( - default=None, - description="""The model to use for inference. Accepts a short Gemini model name (e.g. `gemini-2.5-flash`), which is automatically expanded to a fully-qualified resource name using the client's project and location, or an already fully-qualified publisher-model or endpoint resource name (e.g. `projects/{project}/locations/{location}/publishers/google/models/gemini-2.5-flash`).""", - ) - prompt_template: Optional[EvaluationRunPromptTemplate] = Field( - default=None, description="""The prompt template used for inference.""" - ) - agent_run_config: Optional[AgentRunConfig] = Field( - default=None, - description="""Configuration for Agent Run in evaluation management service.""", - ) - agent_configs: Optional[dict[str, evals_types.AgentConfig]] = Field( + response: Optional[genai_types.Content] = Field( default=None, - description="""A map of agent IDs to their respective agent config.""", + description="""The final model-generated response to the `prompt`.""", ) -class EvaluationRunInferenceConfigDict(TypedDict, total=False): - """Configuration that describes an agent.""" +class ResponseCandidateDict(TypedDict, total=False): + """A model-generated content to the prompt.""" - agent_config: Optional[EvaluationRunAgentConfigDict] - """The agent config.""" + response: Optional[genai_types.Content] + """The final model-generated response to the `prompt`.""" - model: Optional[str] - """The model to use for inference. Accepts a short Gemini model name (e.g. `gemini-2.5-flash`), which is automatically expanded to a fully-qualified resource name using the client's project and location, or an already fully-qualified publisher-model or endpoint resource name (e.g. `projects/{project}/locations/{location}/publishers/google/models/gemini-2.5-flash`).""" - prompt_template: Optional[EvaluationRunPromptTemplateDict] - """The prompt template used for inference.""" - - agent_run_config: Optional[AgentRunConfigDict] - """Configuration for Agent Run in evaluation management service.""" - - agent_configs: Optional[dict[str, evals_types.AgentConfig]] - """A map of agent IDs to their respective agent config.""" - - -EvaluationRunInferenceConfigOrDict = Union[ - EvaluationRunInferenceConfig, EvaluationRunInferenceConfigDict -] +ResponseCandidateOrDict = Union[ResponseCandidate, ResponseCandidateDict] -class VulnerableTool(_common.BaseModel): - """A tool considered high risk for prompt injection.""" +class InteractionsDataSource(_common.BaseModel): + """Source for populating agent data from an Interactions API interaction.""" - tool_name: Optional[str] = Field( + gemini_agent_config: Optional[GeminiAgentConfig] = Field( default=None, - description="""Optional. The name of the vulnerable function/tool (e.g., "search_flights").""", + description="""The Gemini Agent (Vertex AI Agent resource) that produced the + interaction.""", ) - json_paths: Optional[list[str]] = Field( + interaction: Optional[str] = Field( default=None, - description="""Optional. JSON Paths within the tool's FunctionResponse where malicious content could be injected.""", + description="""The interaction to evaluate. Required by the backend. + Format: + `projects/{project}/locations/{location}/interactions/{interaction}`.""", ) -class VulnerableToolDict(TypedDict, total=False): - """A tool considered high risk for prompt injection.""" +class InteractionsDataSourceDict(TypedDict, total=False): + """Source for populating agent data from an Interactions API interaction.""" - tool_name: Optional[str] - """Optional. The name of the vulnerable function/tool (e.g., "search_flights").""" + gemini_agent_config: Optional[GeminiAgentConfigDict] + """The Gemini Agent (Vertex AI Agent resource) that produced the + interaction.""" - json_paths: Optional[list[str]] - """Optional. JSON Paths within the tool's FunctionResponse where malicious content could be injected.""" + interaction: Optional[str] + """The interaction to evaluate. Required by the backend. + Format: + `projects/{project}/locations/{location}/interactions/{interaction}`.""" -VulnerableToolOrDict = Union[VulnerableTool, VulnerableToolDict] +InteractionsDataSourceOrDict = Union[InteractionsDataSource, InteractionsDataSourceDict] -class RedTeamingAnalysisConfig(_common.BaseModel): - """Configuration for the automated Agent Red Teaming analysis.""" +class EvalCase(_common.BaseModel): + """A comprehensive representation of a GenAI interaction for evaluation.""" - attack_categories: Optional[list[str]] = Field( + prompt: Optional[genai_types.Content] = Field( + default=None, description="""The most recent user message (current input).""" + ) + responses: Optional[list[ResponseCandidate]] = Field( default=None, - description="""Optional. Specific attack categories to test against.""", + description="""Model-generated replies to the last user message in a conversation. Multiple responses are allowed to support use cases such as comparing different model outputs.""", ) - vulnerable_tools: Optional[list[VulnerableTool]] = Field( + reference: Optional[ResponseCandidate] = Field( default=None, - description="""Optional. Manually defined vulnerable tools and their injection paths.""", + description="""User-provided, golden reference model reply to prompt in context of chat history; Reference for last response in a conversation.""", ) - - -class RedTeamingAnalysisConfigDict(TypedDict, total=False): - """Configuration for the automated Agent Red Teaming analysis.""" - - attack_categories: Optional[list[str]] - """Optional. Specific attack categories to test against.""" - - vulnerable_tools: Optional[list[VulnerableToolDict]] - """Optional. Manually defined vulnerable tools and their injection paths.""" - - -RedTeamingAnalysisConfigOrDict = Union[ - RedTeamingAnalysisConfig, RedTeamingAnalysisConfigDict -] - - -class AnalysisConfig(_common.BaseModel): - """Configuration for an analysis to be performed on an evaluation run.""" - - analysis_name: Optional[str] = Field( - default=None, description="""Optional. A name for this analysis.""" + system_instruction: Optional[genai_types.Content] = Field( + default=None, description="""System instruction for the model.""" ) - red_teaming_analysis_config: Optional[RedTeamingAnalysisConfig] = Field( + conversation_history: Optional[list[evals_types.Message]] = Field( default=None, - description="""Configuration for the automated Agent Red Teaming analysis.""", + description="""List of all prior messages in the conversation (chat history).""", + ) + rubric_groups: Optional[dict[str, "RubricGroup"]] = Field( + default=None, + description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", + ) + eval_case_id: Optional[str] = Field( + default=None, description="""Unique identifier for the evaluation case.""" ) + intermediate_events: Optional[list[evals_types.Event]] = Field( + default=None, + description="""Intermediate events of a single turn in an agent run or intermediate events of the last turn for multi-turn an agent run.""", + ) + agent_info: Optional[evals_types.AgentInfo] = Field( + default=None, + description="""The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""", + ) + agent_data: Optional[evals_types.AgentData] = Field( + default=None, description="""The agent data of the agent under evaluation.""" + ) + user_scenario: Optional[evals_types.UserScenario] = Field( + default=None, description="""The user scenario for the evaluation case.""" + ) + interactions_data_source: Optional[InteractionsDataSource] = Field( + default=None, + description="""Source for populating agent data from an Interactions API interaction. When set, the backend fetches the interaction (and the Gemini Agent config) and parses it into agent data for evaluation; agent_data must not also be set.""", + ) + # Allow extra fields to support custom metric prompts and stay backward compatible. + model_config = ConfigDict(frozen=True, extra="allow") -class AnalysisConfigDict(TypedDict, total=False): - """Configuration for an analysis to be performed on an evaluation run.""" +class EvalCaseDict(TypedDict, total=False): + """A comprehensive representation of a GenAI interaction for evaluation.""" - analysis_name: Optional[str] - """Optional. A name for this analysis.""" + prompt: Optional[genai_types.Content] + """The most recent user message (current input).""" - red_teaming_analysis_config: Optional[RedTeamingAnalysisConfigDict] - """Configuration for the automated Agent Red Teaming analysis.""" + responses: Optional[list[ResponseCandidateDict]] + """Model-generated replies to the last user message in a conversation. Multiple responses are allowed to support use cases such as comparing different model outputs.""" + reference: Optional[ResponseCandidateDict] + """User-provided, golden reference model reply to prompt in context of chat history; Reference for last response in a conversation.""" -AnalysisConfigOrDict = Union[AnalysisConfig, AnalysisConfigDict] + system_instruction: Optional[genai_types.Content] + """System instruction for the model.""" + conversation_history: Optional[list[evals_types.Message]] + """List of all prior messages in the conversation (chat history).""" -class CreateEvaluationRunConfig(_common.BaseModel): - """Config to create an evaluation run.""" + rubric_groups: Optional[dict[str, "RubricGroupDict"]] + """Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - allow_cross_region_model: Optional[bool] = Field( - default=None, - description="""Allows the evaluation run to use cross region models. When this - flag is set, the service may route traffic to other regions if a model is - unavailable in the current region (e.g., to a `global`endpoint). If a - fully-qualified model endpoint resource name with a different region than - the run location is provided elsewhere in the run config, this flag must - be set to true or the request will fail.""", - ) + eval_case_id: Optional[str] + """Unique identifier for the evaluation case.""" + intermediate_events: Optional[list[evals_types.Event]] + """Intermediate events of a single turn in an agent run or intermediate events of the last turn for multi-turn an agent run.""" -class CreateEvaluationRunConfigDict(TypedDict, total=False): - """Config to create an evaluation run.""" + agent_info: Optional[evals_types.AgentInfo] + """The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + agent_data: Optional[evals_types.AgentData] + """The agent data of the agent under evaluation.""" - allow_cross_region_model: Optional[bool] - """Allows the evaluation run to use cross region models. When this - flag is set, the service may route traffic to other regions if a model is - unavailable in the current region (e.g., to a `global`endpoint). If a - fully-qualified model endpoint resource name with a different region than - the run location is provided elsewhere in the run config, this flag must - be set to true or the request will fail.""" + user_scenario: Optional[evals_types.UserScenario] + """The user scenario for the evaluation case.""" + interactions_data_source: Optional[InteractionsDataSourceDict] + """Source for populating agent data from an Interactions API interaction. When set, the backend fetches the interaction (and the Gemini Agent config) and parses it into agent data for evaluation; agent_data must not also be set.""" -CreateEvaluationRunConfigOrDict = Union[ - CreateEvaluationRunConfig, CreateEvaluationRunConfigDict -] +EvalCaseOrDict = Union[EvalCase, EvalCaseDict] -class _CreateEvaluationRunParameters(_common.BaseModel): - """Represents a job that creates an evaluation run.""" - name: Optional[str] = Field(default=None, description="""""") - display_name: Optional[str] = Field(default=None, description="""""") - data_source: Optional[EvaluationRunDataSource] = Field( - default=None, description="""""" - ) - evaluation_config: Optional[EvaluationRunConfig] = Field( - default=None, description="""""" +class EvaluationDataset(_common.BaseModel): + """The dataset used for evaluation.""" + + bigquery_source: Optional[genai_types.BigQuerySource] = Field( + default=None, description="""The BigQuery source for the evaluation dataset.""" ) - labels: Optional[dict[str, str]] = Field(default=None, description="""""") - inference_configs: Optional[dict[str, EvaluationRunInferenceConfig]] = Field( - default=None, description="""""" + gcs_source: Optional[genai_types.GcsSource] = Field( + default=None, description="""The GCS source for the evaluation dataset.""" ) - config: Optional[CreateEvaluationRunConfig] = Field( - default=None, description="""""" + eval_cases: Optional[list[EvalCase]] = Field( + default=None, description="""The evaluation cases to be evaluated.""" ) - analysis_configs: Optional[list[AnalysisConfig]] = Field( - default=None, description="""""" + eval_dataset_df: Optional[PandasDataFrame] = Field( + default=None, + description="""The evaluation dataset in the form of a Pandas DataFrame.""", ) - evaluation_experiment: Optional[str] = Field( + candidate_name: Optional[str] = Field( default=None, - description="""The resource name of the parent EvaluationExperiment that this run - belongs to. Format: - `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""", + description="""The name of the candidate model or agent for this evaluation dataset.""", ) + @model_validator(mode="before") + @classmethod + def _check_pandas_installed(cls, data: Any) -> Any: + if isinstance(data, dict) and data.get("eval_dataset_df") is not None: + if pd is None: + logger.warning( + "Pandas is not installed, some evals features are not available." + " Please install it with `pip install" + " google-cloud-aiplatform[evaluation]`." + ) + return data -class _CreateEvaluationRunParametersDict(TypedDict, total=False): - """Represents a job that creates an evaluation run.""" + @classmethod + def load_from_observability_eval_cases( + cls, cases: list["ObservabilityEvalCase"] + ) -> "EvaluationDataset": + """Fetches GenAI Observability data from GCS and parses into a DataFrame.""" + try: + import pandas as pd + from .. import _gcs_utils - name: Optional[str] - """""" + formats = [] + requests = [] + responses = [] + system_instructions = [] - display_name: Optional[str] - """""" + for case in cases: + gcs_utils = _gcs_utils.GcsUtils( + case.api_client._api_client if case.api_client else None + ) - data_source: Optional[EvaluationRunDataSourceDict] - """""" + # Associate "observability" data format for given sources + formats.append("observability") - evaluation_config: Optional[EvaluationRunConfigDict] - """""" + # Input source + request_data = gcs_utils.read_file_contents(case.input_src) + requests.append(request_data) - labels: Optional[dict[str, str]] - """""" - - inference_configs: Optional[dict[str, EvaluationRunInferenceConfigDict]] - """""" + # Output source + response_data = gcs_utils.read_file_contents(case.output_src) + responses.append(response_data) - config: Optional[CreateEvaluationRunConfigDict] - """""" + # System instruction source + system_instruction_data = "" + if case.system_instruction_src is not None: + system_instruction_data = gcs_utils.read_file_contents( + case.system_instruction_src + ) + system_instructions.append(system_instruction_data) - analysis_configs: Optional[list[AnalysisConfigDict]] - """""" + eval_dataset_df = pd.DataFrame( + { + "format": formats, + "request": requests, + "response": responses, + "system_instruction": system_instructions, + } + ) - evaluation_experiment: Optional[str] - """The resource name of the parent EvaluationExperiment that this run - belongs to. Format: - `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""" + except ImportError as e: + raise ImportError("Pandas DataFrame library is required.") from e + return EvaluationDataset(eval_dataset_df=eval_dataset_df) -_CreateEvaluationRunParametersOrDict = Union[ - _CreateEvaluationRunParameters, _CreateEvaluationRunParametersDict -] + def show(self) -> None: + """Shows the evaluation dataset.""" + from .. import _evals_visualization + _evals_visualization.display_evaluation_dataset(self) -class SummaryMetric(_common.BaseModel): - """Represents a summary metric for an evaluation run.""" - metrics: Optional[dict[str, Any]] = Field( - default=None, description="""Map of metric name to metric value.""" - ) - total_items: Optional[int] = Field( - default=None, description="""The total number of items that were evaluated.""" - ) - failed_items: Optional[int] = Field( - default=None, description="""The number of items that failed to be evaluated.""" - ) +class EvaluationDatasetDict(TypedDict, total=False): + """The dataset used for evaluation.""" + bigquery_source: Optional[genai_types.BigQuerySource] + """The BigQuery source for the evaluation dataset.""" -class SummaryMetricDict(TypedDict, total=False): - """Represents a summary metric for an evaluation run.""" + gcs_source: Optional[genai_types.GcsSource] + """The GCS source for the evaluation dataset.""" - metrics: Optional[dict[str, Any]] - """Map of metric name to metric value.""" + eval_cases: Optional[list[EvalCaseDict]] + """The evaluation cases to be evaluated.""" - total_items: Optional[int] - """The total number of items that were evaluated.""" + eval_dataset_df: Optional[PandasDataFrame] + """The evaluation dataset in the form of a Pandas DataFrame.""" - failed_items: Optional[int] - """The number of items that failed to be evaluated.""" + candidate_name: Optional[str] + """The name of the candidate model or agent for this evaluation dataset.""" -SummaryMetricOrDict = Union[SummaryMetric, SummaryMetricDict] +EvaluationDatasetOrDict = Union[EvaluationDataset, EvaluationDatasetDict] -class AttackCategoryResult(_common.BaseModel): - """The red teaming outcome for a specific attack category.""" +class EvaluationRunMetadata(_common.BaseModel): + """Metadata for an evaluation run.""" - attack_category: Optional[str] = Field( - default=None, description="""The category of the attack evaluated.""" + candidate_names: Optional[list[str]] = Field( + default=None, + description="""Name of the candidate(s) being evaluated in the evaluation run.""", ) - attack_success_rate: Optional[float] = Field( + dataset_name: Optional[str] = Field( default=None, - description="""The ratio of successful attacks given a fixed budget.""", + description="""Name of the evaluation dataset used for the evaluation run.""", ) - vulnerability_insight: Optional[str] = Field( - default=None, description="""Insights into why an attack succeeded or failed.""" + dataset_id: Optional[str] = Field( + default=None, + description="""Unique identifier for the evaluation dataset used for the evaluation run.""", + ) + creation_timestamp: Optional[datetime.datetime] = Field( + default=None, description="""Creation timestamp of the evaluation run.""" ) -class AttackCategoryResultDict(TypedDict, total=False): - """The red teaming outcome for a specific attack category.""" +class EvaluationRunMetadataDict(TypedDict, total=False): + """Metadata for an evaluation run.""" - attack_category: Optional[str] - """The category of the attack evaluated.""" + candidate_names: Optional[list[str]] + """Name of the candidate(s) being evaluated in the evaluation run.""" - attack_success_rate: Optional[float] - """The ratio of successful attacks given a fixed budget.""" + dataset_name: Optional[str] + """Name of the evaluation dataset used for the evaluation run.""" - vulnerability_insight: Optional[str] - """Insights into why an attack succeeded or failed.""" + dataset_id: Optional[str] + """Unique identifier for the evaluation dataset used for the evaluation run.""" + creation_timestamp: Optional[datetime.datetime] + """Creation timestamp of the evaluation run.""" -AttackCategoryResultOrDict = Union[AttackCategoryResult, AttackCategoryResultDict] +EvaluationRunMetadataOrDict = Union[EvaluationRunMetadata, EvaluationRunMetadataDict] -class RedTeamingAnalysisResult(_common.BaseModel): - """The top-level result for Red Teaming analysis.""" - config: Optional[RedTeamingAnalysisConfig] = Field( +class EvaluationResult(_common.BaseModel): + """Result of an evaluation run for an evaluation dataset.""" + + eval_case_results: Optional[list[EvalCaseResult]] = Field( default=None, - description="""The configuration used to generate this analysis.""", - ) - analysis_time: Optional[str] = Field( - default=None, description="""The timestamp when this analysis was performed.""" + description="""A list of evaluation results for each evaluation case.""", ) - category_results: Optional[list[AttackCategoryResult]] = Field( - default=None, description="""Detailed results by attack category.""" + summary_metrics: Optional[list[AggregatedMetricResult]] = Field( + default=None, + description="""A list of summary-level evaluation results for each metric.""", ) - - -class RedTeamingAnalysisResultDict(TypedDict, total=False): - """The top-level result for Red Teaming analysis.""" - - config: Optional[RedTeamingAnalysisConfigDict] - """The configuration used to generate this analysis.""" - - analysis_time: Optional[str] - """The timestamp when this analysis was performed.""" - - category_results: Optional[list[AttackCategoryResultDict]] - """Detailed results by attack category.""" - - -RedTeamingAnalysisResultOrDict = Union[ - RedTeamingAnalysisResult, RedTeamingAnalysisResultDict -] - - -class LossTaxonomyEntry(_common.BaseModel): - """A specific entry in the loss pattern taxonomy.""" - - l1_category: Optional[str] = Field( + win_rates: Optional[dict[str, WinRateStats]] = Field( default=None, - description="""The primary category of the loss (e.g., "Hallucination", "Tool Calling").""", + description="""A dictionary of win rates for each metric, only populated for multi-response evaluation runs.""", ) - l2_category: Optional[str] = Field( + evaluation_dataset: Optional[list[EvaluationDataset]] = Field( default=None, - description="""The secondary category of the loss (e.g., "Hallucination of Action", "Incorrect Tool Selection").""", + description="""The input evaluation dataset(s) for the evaluation run.""", ) - description: Optional[str] = Field( + metadata: Optional[EvaluationRunMetadata] = Field( + default=None, description="""Metadata for the evaluation run.""" + ) + agent_info: Optional[evals_types.AgentInfo] = Field( default=None, - description="""A detailed description of this loss pattern. Example: "The agent verbally confirms an action without executing the tool." """, + description="""The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""", ) + def show(self, candidate_names: Optional[List[str]] = None) -> None: + """Shows the evaluation result. -class LossTaxonomyEntryDict(TypedDict, total=False): - """A specific entry in the loss pattern taxonomy.""" - - l1_category: Optional[str] - """The primary category of the loss (e.g., "Hallucination", "Tool Calling").""" - - l2_category: Optional[str] - """The secondary category of the loss (e.g., "Hallucination of Action", "Incorrect Tool Selection").""" - - description: Optional[str] - """A detailed description of this loss pattern. Example: "The agent verbally confirms an action without executing the tool." """ + Args: + candidate_names: list of names for the evaluated candidates, used in + comparison reports. + """ + from .. import _evals_visualization + _evals_visualization.display_evaluation_result(self, candidate_names) -LossTaxonomyEntryOrDict = Union[LossTaxonomyEntry, LossTaxonomyEntryDict] +class EvaluationResultDict(TypedDict, total=False): + """Result of an evaluation run for an evaluation dataset.""" -class FailedRubric(_common.BaseModel): - """A specific failed rubric and the associated analysis.""" + eval_case_results: Optional[list[EvalCaseResultDict]] + """A list of evaluation results for each evaluation case.""" - rubric_id: Optional[str] = Field( - default=None, - description="""The unique ID of the rubric (if available from the metric source).""", - ) - classification_rationale: Optional[str] = Field( - default=None, - description="""The rationale provided by the Loss Analysis Classifier for why this failure maps to this specific Loss Cluster.""", - ) + summary_metrics: Optional[list[AggregatedMetricResultDict]] + """A list of summary-level evaluation results for each metric.""" + win_rates: Optional[dict[str, WinRateStatsDict]] + """A dictionary of win rates for each metric, only populated for multi-response evaluation runs.""" -class FailedRubricDict(TypedDict, total=False): - """A specific failed rubric and the associated analysis.""" + evaluation_dataset: Optional[list[EvaluationDatasetDict]] + """The input evaluation dataset(s) for the evaluation run.""" - rubric_id: Optional[str] - """The unique ID of the rubric (if available from the metric source).""" + metadata: Optional[EvaluationRunMetadataDict] + """Metadata for the evaluation run.""" - classification_rationale: Optional[str] - """The rationale provided by the Loss Analysis Classifier for why this failure maps to this specific Loss Cluster.""" + agent_info: Optional[evals_types.AgentInfo] + """The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""" -FailedRubricOrDict = Union[FailedRubric, FailedRubricDict] +EvaluationResultOrDict = Union[EvaluationResult, EvaluationResultDict] -class LossExample(_common.BaseModel): - """A specific example of a loss pattern.""" +class EvaluationRun(_common.BaseModel): + """Represents an evaluation run.""" - evaluation_item: Optional[str] = Field( - default=None, - description="""Reference to the persisted EvalItem resource name. Format: projects/.../locations/.../evaluationItems/{item_id}.""", + name: Optional[str] = Field(default=None, description="""""") + display_name: Optional[str] = Field(default=None, description="""""") + metadata: Optional[dict[str, Any]] = Field(default=None, description="""""") + create_time: Optional[datetime.datetime] = Field(default=None, description="""""") + completion_time: Optional[datetime.datetime] = Field( + default=None, description="""""" ) - evaluation_result: Optional[dict[str, Any]] = Field( - default=None, - description="""The full evaluation result object provided inline. Used when the analysis is performed on ephemeral data.""", + state: Optional[EvaluationRunState] = Field(default=None, description="""""") + evaluation_set_snapshot: Optional[str] = Field(default=None, description="""""") + error: Optional[genai_types.GoogleRpcStatus] = Field( + default=None, description="""""" ) - failed_rubrics: Optional[list[FailedRubric]] = Field( + data_source: Optional[EvaluationRunDataSource] = Field( + default=None, description="""""" + ) + evaluation_run_results: Optional[EvaluationRunResults] = Field( + default=None, description="""The evaluation run formatted results.""" + ) + evaluation_item_results: Optional[EvaluationResult] = Field( default=None, - description="""The specific rubric(s) that failed and caused this example to be classified here. An example might fail multiple rubrics, but only specific ones trigger this loss pattern.""", + description="""The parsed EvaluationItem results for the evaluation run. This is only populated when include_evaluation_items is set to True.""", ) - - -class LossExampleDict(TypedDict, total=False): - """A specific example of a loss pattern.""" - - evaluation_item: Optional[str] - """Reference to the persisted EvalItem resource name. Format: projects/.../locations/.../evaluationItems/{item_id}.""" - - evaluation_result: Optional[dict[str, Any]] - """The full evaluation result object provided inline. Used when the analysis is performed on ephemeral data.""" - - failed_rubrics: Optional[list[FailedRubricDict]] - """The specific rubric(s) that failed and caused this example to be classified here. An example might fail multiple rubrics, but only specific ones trigger this loss pattern.""" - - -LossExampleOrDict = Union[LossExample, LossExampleDict] - - -class LossCluster(_common.BaseModel): - """A semantic grouping of failures (e.g., "Hallucination of Action").""" - - cluster_id: Optional[str] = Field( - default=None, - description="""Unique identifier for the loss cluster within the scope of the analysis result.""", - ) - taxonomy_entry: Optional[LossTaxonomyEntry] = Field( - default=None, - description="""The structured definition of the loss taxonomy for this cluster.""", + evaluation_config: Optional[EvaluationRunConfig] = Field( + default=None, description="""The evaluation config for the evaluation run.""" ) - item_count: Optional[int] = Field( - default=None, - description="""The total number of EvaluationItems falling into this cluster.""", + inference_configs: Optional[dict[str, EvaluationRunInferenceConfig]] = Field( + default=None, description="""The inference configs for the evaluation run.""" ) - examples: Optional[list[LossExample]] = Field( + labels: Optional[dict[str, str]] = Field(default=None, description="""""") + analysis_configs: Optional[list[AnalysisConfig]] = Field( default=None, - description="""A list of examples that belong to this cluster. This links the cluster back to the specific EvaluationItems and Rubrics.""", + description="""The analysis configurations for the evaluation run.""", ) + # TODO(b/448806531): Remove all the overridden _from_response methods once the + # ticket is resolved and published. + @classmethod + def _from_response( + cls: typing.Type["EvaluationRun"], + *, + response: dict[str, object], + kwargs: dict[str, object], + ) -> "EvaluationRun": + """Converts a dictionary response into a EvaluationRun object.""" -class LossClusterDict(TypedDict, total=False): - """A semantic grouping of failures (e.g., "Hallucination of Action").""" - - cluster_id: Optional[str] - """Unique identifier for the loss cluster within the scope of the analysis result.""" + snaked_response = _camel_key_to_snake(response) - taxonomy_entry: Optional[LossTaxonomyEntryDict] - """The structured definition of the loss taxonomy for this cluster.""" + evaluation_run_results = response.get("evaluation_run_results") - item_count: Optional[int] - """The total number of EvaluationItems falling into this cluster.""" + if ( + isinstance(evaluation_run_results, dict) + and "summaryMetrics" in evaluation_run_results + ): + snaked_response["evaluation_run_results"]["summary_metrics"] = ( + evaluation_run_results["summaryMetrics"] + ) + result = super()._from_response(response=snaked_response, kwargs=kwargs) + return result - examples: Optional[list[LossExampleDict]] - """A list of examples that belong to this cluster. This links the cluster back to the specific EvaluationItems and Rubrics.""" + def show(self) -> None: + """Shows the evaluation result.""" + from .. import _evals_visualization + if self.state == "SUCCEEDED": + if self.evaluation_item_results is not None: + _evals_visualization.display_evaluation_result( + self.evaluation_item_results, None + ) + else: + logger.warning( + "Evaluation Run succeeded but no evaluation item results found. To display results, please set include_evaluation_items to True when calling get_evaluation_run()." + ) + # Show loss analysis results if present on the evaluation run. + # Pass the eval item map so the visualization can enrich + # loss examples with scenario/rubric data. + if ( + self.evaluation_run_results + and self.evaluation_run_results.loss_analysis_results + ): + eval_item_map = getattr(self, "_eval_item_map", None) + _evals_visualization.display_loss_analysis_results( + self.evaluation_run_results.loss_analysis_results, + eval_item_map=eval_item_map, + ) + else: + _evals_visualization.display_evaluation_run_status(self) -LossClusterOrDict = Union[LossCluster, LossClusterDict] +class EvaluationRunDict(TypedDict, total=False): + """Represents an evaluation run.""" -class LossAnalysisResult(_common.BaseModel): - """The top-level result for loss analysis.""" + name: Optional[str] + """""" - config: Optional[LossAnalysisConfig] = Field( - default=None, - description="""The configuration used to generate this analysis.""", - ) - analysis_time: Optional[str] = Field( - default=None, description="""The timestamp when this analysis was performed.""" - ) - clusters: Optional[list[LossCluster]] = Field( - default=None, description="""The list of identified loss clusters.""" - ) + display_name: Optional[str] + """""" - def show(self) -> None: - """Shows the loss analysis result with rich HTML visualization.""" - from .. import _evals_visualization + metadata: Optional[dict[str, Any]] + """""" - _evals_visualization.display_loss_analysis_result(self) + create_time: Optional[datetime.datetime] + """""" + completion_time: Optional[datetime.datetime] + """""" -class LossAnalysisResultDict(TypedDict, total=False): - """The top-level result for loss analysis.""" + state: Optional[EvaluationRunState] + """""" - config: Optional[LossAnalysisConfigDict] - """The configuration used to generate this analysis.""" + evaluation_set_snapshot: Optional[str] + """""" - analysis_time: Optional[str] - """The timestamp when this analysis was performed.""" + error: Optional[genai_types.GoogleRpcStatus] + """""" - clusters: Optional[list[LossClusterDict]] - """The list of identified loss clusters.""" + data_source: Optional[EvaluationRunDataSourceDict] + """""" + evaluation_run_results: Optional[EvaluationRunResultsDict] + """The evaluation run formatted results.""" -LossAnalysisResultOrDict = Union[LossAnalysisResult, LossAnalysisResultDict] + evaluation_item_results: Optional[EvaluationResultDict] + """The parsed EvaluationItem results for the evaluation run. This is only populated when include_evaluation_items is set to True.""" + evaluation_config: Optional[EvaluationRunConfigDict] + """The evaluation config for the evaluation run.""" -class EvaluationRunResults(_common.BaseModel): - """Represents the results of an evaluation run.""" + inference_configs: Optional[dict[str, EvaluationRunInferenceConfigDict]] + """The inference configs for the evaluation run.""" - evaluation_set: Optional[str] = Field( - default=None, - description="""The evaluation set where item level results are stored.""", - ) - summary_metrics: Optional[SummaryMetric] = Field( - default=None, description="""The summary metrics for the evaluation run.""" - ) - loss_analysis_results: Optional[list[LossAnalysisResult]] = Field( - default=None, - description="""The loss analysis results for the evaluation run.""", - ) - red_teaming_analysis_results: Optional[list[RedTeamingAnalysisResult]] = Field( - default=None, description="""The Red Teaming analysis results.""" - ) + labels: Optional[dict[str, str]] + """""" + analysis_configs: Optional[list[AnalysisConfigDict]] + """The analysis configurations for the evaluation run.""" -class EvaluationRunResultsDict(TypedDict, total=False): - """Represents the results of an evaluation run.""" - evaluation_set: Optional[str] - """The evaluation set where item level results are stored.""" +EvaluationRunOrDict = Union[EvaluationRun, EvaluationRunDict] - summary_metrics: Optional[SummaryMetricDict] - """The summary metrics for the evaluation run.""" - loss_analysis_results: Optional[list[LossAnalysisResultDict]] - """The loss analysis results for the evaluation run.""" +class CreateEvaluationSetConfig(_common.BaseModel): + """Config to create an evaluation set.""" - red_teaming_analysis_results: Optional[list[RedTeamingAnalysisResultDict]] - """The Red Teaming analysis results.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -EvaluationRunResultsOrDict = Union[EvaluationRunResults, EvaluationRunResultsDict] +class CreateEvaluationSetConfigDict(TypedDict, total=False): + """Config to create an evaluation set.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -class EvalCaseMetricResult(_common.BaseModel): - """Evaluation result for a single evaluation case for a single metric.""" - metric_name: Optional[str] = Field( - default=None, description="""Name of the metric.""" - ) - score: Optional[float] = Field(default=None, description="""Score of the metric.""") - explanation: Optional[str] = Field( - default=None, description="""Explanation of the metric.""" - ) - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( - default=None, - description="""The details of all the rubrics and their verdicts for rubric-based metrics.""", - ) - raw_output: Optional[list[str]] = Field( - default=None, description="""Raw output of the metric.""" - ) - error_message: Optional[str] = Field( - default=None, description="""Error message for the metric.""" - ) +CreateEvaluationSetConfigOrDict = Union[ + CreateEvaluationSetConfig, CreateEvaluationSetConfigDict +] -class EvalCaseMetricResultDict(TypedDict, total=False): - """Evaluation result for a single evaluation case for a single metric.""" +class _CreateEvaluationSetParameters(_common.BaseModel): + """Represents a job that creates an evaluation set.""" - metric_name: Optional[str] - """Name of the metric.""" + evaluation_items: Optional[list[str]] = Field(default=None, description="""""") + display_name: Optional[str] = Field(default=None, description="""""") + config: Optional[CreateEvaluationSetConfig] = Field( + default=None, description="""""" + ) - score: Optional[float] - """Score of the metric.""" - explanation: Optional[str] - """Explanation of the metric.""" +class _CreateEvaluationSetParametersDict(TypedDict, total=False): + """Represents a job that creates an evaluation set.""" - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] - """The details of all the rubrics and their verdicts for rubric-based metrics.""" + evaluation_items: Optional[list[str]] + """""" - raw_output: Optional[list[str]] - """Raw output of the metric.""" + display_name: Optional[str] + """""" - error_message: Optional[str] - """Error message for the metric.""" + config: Optional[CreateEvaluationSetConfigDict] + """""" -EvalCaseMetricResultOrDict = Union[EvalCaseMetricResult, EvalCaseMetricResultDict] +_CreateEvaluationSetParametersOrDict = Union[ + _CreateEvaluationSetParameters, _CreateEvaluationSetParametersDict +] -class ResponseCandidateResult(_common.BaseModel): - """Aggregated metric results for a single response candidate.""" +class EvaluationSet(_common.BaseModel): + """Represents an evaluation set.""" - response_index: Optional[int] = Field( - default=None, - description="""Index of the response candidate this result pertains to.""", + name: Optional[str] = Field( + default=None, description="""The resource name of the evaluation set.""" ) - metric_results: Optional[dict[str, EvalCaseMetricResult]] = Field( - default=None, - description="""A dictionary of metric results for this response candidate, keyed by metric name.""", + display_name: Optional[str] = Field( + default=None, description="""The display name of the evaluation set.""" + ) + evaluation_items: Optional[list[str]] = Field( + default=None, + description="""The EvaluationItems that are part of this dataset.""", + ) + create_time: Optional[datetime.datetime] = Field( + default=None, description="""The create time of the evaluation set.""" + ) + update_time: Optional[datetime.datetime] = Field( + default=None, description="""The update time of the evaluation set.""" + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, description="""The metadata of the evaluation set.""" ) -class ResponseCandidateResultDict(TypedDict, total=False): - """Aggregated metric results for a single response candidate.""" - - response_index: Optional[int] - """Index of the response candidate this result pertains to.""" - - metric_results: Optional[dict[str, EvalCaseMetricResultDict]] - """A dictionary of metric results for this response candidate, keyed by metric name.""" - +class EvaluationSetDict(TypedDict, total=False): + """Represents an evaluation set.""" -ResponseCandidateResultOrDict = Union[ - ResponseCandidateResult, ResponseCandidateResultDict -] + name: Optional[str] + """The resource name of the evaluation set.""" + display_name: Optional[str] + """The display name of the evaluation set.""" -class EvalCaseResult(_common.BaseModel): - """Eval result for a single evaluation case.""" + evaluation_items: Optional[list[str]] + """The EvaluationItems that are part of this dataset.""" - eval_case_index: Optional[int] = Field( - default=None, description="""Index of the evaluation case.""" - ) - response_candidate_results: Optional[list[ResponseCandidateResult]] = Field( - default=None, - description="""A list of results, one for each response candidate of the EvalCase.""", - ) + create_time: Optional[datetime.datetime] + """The create time of the evaluation set.""" + update_time: Optional[datetime.datetime] + """The update time of the evaluation set.""" -class EvalCaseResultDict(TypedDict, total=False): - """Eval result for a single evaluation case.""" + metadata: Optional[dict[str, Any]] + """The metadata of the evaluation set.""" - eval_case_index: Optional[int] - """Index of the evaluation case.""" - response_candidate_results: Optional[list[ResponseCandidateResultDict]] - """A list of results, one for each response candidate of the EvalCase.""" +EvaluationSetOrDict = Union[EvaluationSet, EvaluationSetDict] -EvalCaseResultOrDict = Union[EvalCaseResult, EvalCaseResultDict] +class DeleteEvaluationExperimentConfig(_common.BaseModel): + """Config for deleting an evaluation experiment.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -class AggregatedMetricResult(_common.BaseModel): - """Evaluation result for a single metric for an evaluation dataset.""" - metric_name: Optional[str] = Field( - default=None, description="""Name of the metric.""" - ) - num_cases_total: Optional[int] = Field( - default=None, description="""Total number of cases in the dataset.""" - ) - num_cases_valid: Optional[int] = Field( - default=None, description="""Number of valid cases in the dataset.""" - ) - num_cases_error: Optional[int] = Field( - default=None, description="""Number of cases with errors in the dataset.""" - ) - mean_score: Optional[float] = Field( - default=None, description="""Mean score of the metric.""" - ) - stdev_score: Optional[float] = Field( - default=None, description="""Standard deviation of the metric.""" - ) - pass_rate: Optional[float] = Field( - default=None, - description="""Pass rate of the adaptive rubric metric. Calculated as the number of cases where all criteria passed divided by the total number of valid cases. A case is passing if it has a score of 1.0.""", - ) +class DeleteEvaluationExperimentConfigDict(TypedDict, total=False): + """Config for deleting an evaluation experiment.""" - # Allow extra fields to support custom aggregation stats. - model_config = ConfigDict(extra="allow") + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -class AggregatedMetricResultDict(TypedDict, total=False): - """Evaluation result for a single metric for an evaluation dataset.""" +DeleteEvaluationExperimentConfigOrDict = Union[ + DeleteEvaluationExperimentConfig, DeleteEvaluationExperimentConfigDict +] - metric_name: Optional[str] - """Name of the metric.""" - num_cases_total: Optional[int] - """Total number of cases in the dataset.""" +class _DeleteEvaluationExperimentParameters(_common.BaseModel): + """Parameters for deleting an evaluation experiment.""" - num_cases_valid: Optional[int] - """Number of valid cases in the dataset.""" + name: Optional[str] = Field(default=None, description="""""") + config: Optional[DeleteEvaluationExperimentConfig] = Field( + default=None, description="""""" + ) - num_cases_error: Optional[int] - """Number of cases with errors in the dataset.""" - mean_score: Optional[float] - """Mean score of the metric.""" +class _DeleteEvaluationExperimentParametersDict(TypedDict, total=False): + """Parameters for deleting an evaluation experiment.""" - stdev_score: Optional[float] - """Standard deviation of the metric.""" + name: Optional[str] + """""" - pass_rate: Optional[float] - """Pass rate of the adaptive rubric metric. Calculated as the number of cases where all criteria passed divided by the total number of valid cases. A case is passing if it has a score of 1.0.""" + config: Optional[DeleteEvaluationExperimentConfigDict] + """""" -AggregatedMetricResultOrDict = Union[AggregatedMetricResult, AggregatedMetricResultDict] +_DeleteEvaluationExperimentParametersOrDict = Union[ + _DeleteEvaluationExperimentParameters, _DeleteEvaluationExperimentParametersDict +] -class WinRateStats(_common.BaseModel): - """Statistics for win rates for a single metric.""" +class DeleteEvaluationExperimentOperation(_common.BaseModel): + """Operation for deleting an evaluation experiment.""" - win_rates: Optional[list[float]] = Field( + name: Optional[str] = Field( default=None, - description="""Win rates for the metric, one for each candidate.""", + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - tie_rate: Optional[float] = Field( - default=None, description="""Tie rate for the metric.""" + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", ) -class WinRateStatsDict(TypedDict, total=False): - """Statistics for win rates for a single metric.""" +class DeleteEvaluationExperimentOperationDict(TypedDict, total=False): + """Operation for deleting an evaluation experiment.""" - win_rates: Optional[list[float]] - """Win rates for the metric, one for each candidate.""" + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - tie_rate: Optional[float] - """Tie rate for the metric.""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -WinRateStatsOrDict = Union[WinRateStats, WinRateStatsDict] +DeleteEvaluationExperimentOperationOrDict = Union[ + DeleteEvaluationExperimentOperation, DeleteEvaluationExperimentOperationDict +] -class ResponseCandidate(_common.BaseModel): - """A model-generated content to the prompt.""" +class DeleteEvaluationMetricConfig(_common.BaseModel): + """Config for deleting an evaluation metric.""" - response: Optional[genai_types.Content] = Field( - default=None, - description="""The final model-generated response to the `prompt`.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class ResponseCandidateDict(TypedDict, total=False): - """A model-generated content to the prompt.""" +class DeleteEvaluationMetricConfigDict(TypedDict, total=False): + """Config for deleting an evaluation metric.""" - response: Optional[genai_types.Content] - """The final model-generated response to the `prompt`.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ResponseCandidateOrDict = Union[ResponseCandidate, ResponseCandidateDict] +DeleteEvaluationMetricConfigOrDict = Union[ + DeleteEvaluationMetricConfig, DeleteEvaluationMetricConfigDict +] -class InteractionsDataSource(_common.BaseModel): - """Source for populating agent data from an Interactions API interaction.""" +class _DeleteEvaluationMetricParameters(_common.BaseModel): + """Parameters for deleting an evaluation metric.""" - gemini_agent_config: Optional[GeminiAgentConfig] = Field( - default=None, - description="""The Gemini Agent (Vertex AI Agent resource) that produced the - interaction.""", - ) - interaction: Optional[str] = Field( - default=None, - description="""The interaction to evaluate. Required by the backend. - Format: - `projects/{project}/locations/{location}/interactions/{interaction}`.""", + metric_resource_name: Optional[str] = Field(default=None, description="""""") + config: Optional[DeleteEvaluationMetricConfig] = Field( + default=None, description="""""" ) -class InteractionsDataSourceDict(TypedDict, total=False): - """Source for populating agent data from an Interactions API interaction.""" +class _DeleteEvaluationMetricParametersDict(TypedDict, total=False): + """Parameters for deleting an evaluation metric.""" - gemini_agent_config: Optional[GeminiAgentConfigDict] - """The Gemini Agent (Vertex AI Agent resource) that produced the - interaction.""" + metric_resource_name: Optional[str] + """""" - interaction: Optional[str] - """The interaction to evaluate. Required by the backend. - Format: - `projects/{project}/locations/{location}/interactions/{interaction}`.""" + config: Optional[DeleteEvaluationMetricConfigDict] + """""" -InteractionsDataSourceOrDict = Union[InteractionsDataSource, InteractionsDataSourceDict] +_DeleteEvaluationMetricParametersOrDict = Union[ + _DeleteEvaluationMetricParameters, _DeleteEvaluationMetricParametersDict +] -class EvalCase(_common.BaseModel): - """A comprehensive representation of a GenAI interaction for evaluation.""" +class DeleteEvaluationMetricOperation(_common.BaseModel): + """Operation for deleting an evaluation metric.""" - prompt: Optional[genai_types.Content] = Field( - default=None, description="""The most recent user message (current input).""" - ) - responses: Optional[list[ResponseCandidate]] = Field( + name: Optional[str] = Field( default=None, - description="""Model-generated replies to the last user message in a conversation. Multiple responses are allowed to support use cases such as comparing different model outputs.""", + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - reference: Optional[ResponseCandidate] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""User-provided, golden reference model reply to prompt in context of chat history; Reference for last response in a conversation.""", - ) - system_instruction: Optional[genai_types.Content] = Field( - default=None, description="""System instruction for the model.""" + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", ) - conversation_history: Optional[list[evals_types.Message]] = Field( + done: Optional[bool] = Field( default=None, - description="""List of all prior messages in the conversation (chat history).""", + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", ) - rubric_groups: Optional[dict[str, "RubricGroup"]] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", - ) - eval_case_id: Optional[str] = Field( - default=None, description="""Unique identifier for the evaluation case.""" + description="""The error result of the operation in case of failure or cancellation.""", ) - intermediate_events: Optional[list[evals_types.Event]] = Field( - default=None, - description="""Intermediate events of a single turn in an agent run or intermediate events of the last turn for multi-turn an agent run.""", - ) - agent_info: Optional[evals_types.AgentInfo] = Field( - default=None, - description="""The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""", - ) - agent_data: Optional[evals_types.AgentData] = Field( - default=None, description="""The agent data of the agent under evaluation.""" - ) - user_scenario: Optional[evals_types.UserScenario] = Field( - default=None, description="""The user scenario for the evaluation case.""" - ) - interactions_data_source: Optional[InteractionsDataSource] = Field( - default=None, - description="""Source for populating agent data from an Interactions API interaction. When set, the backend fetches the interaction (and the Gemini Agent config) and parses it into agent data for evaluation; agent_data must not also be set.""", - ) - # Allow extra fields to support custom metric prompts and stay backward compatible. - model_config = ConfigDict(frozen=True, extra="allow") -class EvalCaseDict(TypedDict, total=False): - """A comprehensive representation of a GenAI interaction for evaluation.""" +class DeleteEvaluationMetricOperationDict(TypedDict, total=False): + """Operation for deleting an evaluation metric.""" - prompt: Optional[genai_types.Content] - """The most recent user message (current input).""" + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - responses: Optional[list[ResponseCandidateDict]] - """Model-generated replies to the last user message in a conversation. Multiple responses are allowed to support use cases such as comparing different model outputs.""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - reference: Optional[ResponseCandidateDict] - """User-provided, golden reference model reply to prompt in context of chat history; Reference for last response in a conversation.""" + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - system_instruction: Optional[genai_types.Content] - """System instruction for the model.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - conversation_history: Optional[list[evals_types.Message]] - """List of all prior messages in the conversation (chat history).""" - rubric_groups: Optional[dict[str, "RubricGroupDict"]] - """Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""" +DeleteEvaluationMetricOperationOrDict = Union[ + DeleteEvaluationMetricOperation, DeleteEvaluationMetricOperationDict +] - eval_case_id: Optional[str] - """Unique identifier for the evaluation case.""" - intermediate_events: Optional[list[evals_types.Event]] - """Intermediate events of a single turn in an agent run or intermediate events of the last turn for multi-turn an agent run.""" +class BleuInstance(_common.BaseModel): + """Bleu instance.""" - agent_info: Optional[evals_types.AgentInfo] - """The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""" + prediction: Optional[str] = Field( + default=None, description="""Required. Output of the evaluated model.""" + ) + reference: Optional[str] = Field( + default=None, + description="""Required. Ground truth used to compare against the prediction.""", + ) - agent_data: Optional[evals_types.AgentData] - """The agent data of the agent under evaluation.""" - user_scenario: Optional[evals_types.UserScenario] - """The user scenario for the evaluation case.""" +class BleuInstanceDict(TypedDict, total=False): + """Bleu instance.""" - interactions_data_source: Optional[InteractionsDataSourceDict] - """Source for populating agent data from an Interactions API interaction. When set, the backend fetches the interaction (and the Gemini Agent config) and parses it into agent data for evaluation; agent_data must not also be set.""" + prediction: Optional[str] + """Required. Output of the evaluated model.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" -EvalCaseOrDict = Union[EvalCase, EvalCaseDict] +BleuInstanceOrDict = Union[BleuInstance, BleuInstanceDict] -class EvaluationDataset(_common.BaseModel): - """The dataset used for evaluation.""" - bigquery_source: Optional[genai_types.BigQuerySource] = Field( - default=None, description="""The BigQuery source for the evaluation dataset.""" - ) - gcs_source: Optional[genai_types.GcsSource] = Field( - default=None, description="""The GCS source for the evaluation dataset.""" +class BleuInput(_common.BaseModel): + + instances: Optional[list[BleuInstance]] = Field( + default=None, description="""Required. Repeated bleu instances.""" ) - eval_cases: Optional[list[EvalCase]] = Field( - default=None, description="""The evaluation cases to be evaluated.""" + metric_spec: Optional[genai_types.BleuSpec] = Field( + default=None, description="""Required. Spec for bleu score metric.""" ) - eval_dataset_df: Optional[PandasDataFrame] = Field( - default=None, - description="""The evaluation dataset in the form of a Pandas DataFrame.""", + + +class BleuInputDict(TypedDict, total=False): + + instances: Optional[list[BleuInstanceDict]] + """Required. Repeated bleu instances.""" + + metric_spec: Optional[genai_types.BleuSpec] + """Required. Spec for bleu score metric.""" + + +BleuInputOrDict = Union[BleuInput, BleuInputDict] + + +class ExactMatchInstance(_common.BaseModel): + """Exact match instance.""" + + prediction: Optional[str] = Field( + default=None, description="""Required. Output of the evaluated model.""" ) - candidate_name: Optional[str] = Field( + reference: Optional[str] = Field( default=None, - description="""The name of the candidate model or agent for this evaluation dataset.""", + description="""Required. Ground truth used to compare against the prediction.""", ) - @model_validator(mode="before") - @classmethod - def _check_pandas_installed(cls, data: Any) -> Any: - if isinstance(data, dict) and data.get("eval_dataset_df") is not None: - if pd is None: - logger.warning( - "Pandas is not installed, some evals features are not available." - " Please install it with `pip install" - " google-cloud-aiplatform[evaluation]`." - ) - return data - @classmethod - def load_from_observability_eval_cases( - cls, cases: list["ObservabilityEvalCase"] - ) -> "EvaluationDataset": - """Fetches GenAI Observability data from GCS and parses into a DataFrame.""" - try: - import pandas as pd - from .. import _gcs_utils +class ExactMatchInstanceDict(TypedDict, total=False): + """Exact match instance.""" - formats = [] - requests = [] - responses = [] - system_instructions = [] + prediction: Optional[str] + """Required. Output of the evaluated model.""" - for case in cases: - gcs_utils = _gcs_utils.GcsUtils( - case.api_client._api_client if case.api_client else None - ) + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" - # Associate "observability" data format for given sources - formats.append("observability") - # Input source - request_data = gcs_utils.read_file_contents(case.input_src) - requests.append(request_data) +ExactMatchInstanceOrDict = Union[ExactMatchInstance, ExactMatchInstanceDict] - # Output source - response_data = gcs_utils.read_file_contents(case.output_src) - responses.append(response_data) - # System instruction source - system_instruction_data = "" - if case.system_instruction_src is not None: - system_instruction_data = gcs_utils.read_file_contents( - case.system_instruction_src - ) - system_instructions.append(system_instruction_data) +class ExactMatchSpec(_common.BaseModel): + """Spec for exact match metric.""" - eval_dataset_df = pd.DataFrame( - { - "format": formats, - "request": requests, - "response": responses, - "system_instruction": system_instructions, - } - ) + pass - except ImportError as e: - raise ImportError("Pandas DataFrame library is required.") from e - return EvaluationDataset(eval_dataset_df=eval_dataset_df) +class ExactMatchSpecDict(TypedDict, total=False): + """Spec for exact match metric.""" - def show(self) -> None: - """Shows the evaluation dataset.""" - from .. import _evals_visualization + pass - _evals_visualization.display_evaluation_dataset(self) +ExactMatchSpecOrDict = Union[ExactMatchSpec, ExactMatchSpecDict] -class EvaluationDatasetDict(TypedDict, total=False): - """The dataset used for evaluation.""" - bigquery_source: Optional[genai_types.BigQuerySource] - """The BigQuery source for the evaluation dataset.""" +class ExactMatchInput(_common.BaseModel): - gcs_source: Optional[genai_types.GcsSource] - """The GCS source for the evaluation dataset.""" + instances: Optional[list[ExactMatchInstance]] = Field( + default=None, description="""Required. Repeated exact match instances.""" + ) + metric_spec: Optional[ExactMatchSpec] = Field( + default=None, description="""Required. Spec for exact match metric.""" + ) - eval_cases: Optional[list[EvalCaseDict]] - """The evaluation cases to be evaluated.""" - eval_dataset_df: Optional[PandasDataFrame] - """The evaluation dataset in the form of a Pandas DataFrame.""" +class ExactMatchInputDict(TypedDict, total=False): - candidate_name: Optional[str] - """The name of the candidate model or agent for this evaluation dataset.""" + instances: Optional[list[ExactMatchInstanceDict]] + """Required. Repeated exact match instances.""" + metric_spec: Optional[ExactMatchSpecDict] + """Required. Spec for exact match metric.""" -EvaluationDatasetOrDict = Union[EvaluationDataset, EvaluationDatasetDict] +ExactMatchInputOrDict = Union[ExactMatchInput, ExactMatchInputDict] -class EvaluationRunMetadata(_common.BaseModel): - """Metadata for an evaluation run.""" - candidate_names: Optional[list[str]] = Field( - default=None, - description="""Name of the candidate(s) being evaluated in the evaluation run.""", - ) - dataset_name: Optional[str] = Field( - default=None, - description="""Name of the evaluation dataset used for the evaluation run.""", +class RougeInstance(_common.BaseModel): + """Rouge instance.""" + + prediction: Optional[str] = Field( + default=None, description="""Required. Output of the evaluated model.""" ) - dataset_id: Optional[str] = Field( + reference: Optional[str] = Field( default=None, - description="""Unique identifier for the evaluation dataset used for the evaluation run.""", - ) - creation_timestamp: Optional[datetime.datetime] = Field( - default=None, description="""Creation timestamp of the evaluation run.""" + description="""Required. Ground truth used to compare against the prediction.""", ) -class EvaluationRunMetadataDict(TypedDict, total=False): - """Metadata for an evaluation run.""" - - candidate_names: Optional[list[str]] - """Name of the candidate(s) being evaluated in the evaluation run.""" +class RougeInstanceDict(TypedDict, total=False): + """Rouge instance.""" - dataset_name: Optional[str] - """Name of the evaluation dataset used for the evaluation run.""" + prediction: Optional[str] + """Required. Output of the evaluated model.""" - dataset_id: Optional[str] - """Unique identifier for the evaluation dataset used for the evaluation run.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" - creation_timestamp: Optional[datetime.datetime] - """Creation timestamp of the evaluation run.""" +RougeInstanceOrDict = Union[RougeInstance, RougeInstanceDict] -EvaluationRunMetadataOrDict = Union[EvaluationRunMetadata, EvaluationRunMetadataDict] +class RougeInput(_common.BaseModel): + """Rouge input.""" -class EvaluationResult(_common.BaseModel): - """Result of an evaluation run for an evaluation dataset.""" - - eval_case_results: Optional[list[EvalCaseResult]] = Field( - default=None, - description="""A list of evaluation results for each evaluation case.""", - ) - summary_metrics: Optional[list[AggregatedMetricResult]] = Field( - default=None, - description="""A list of summary-level evaluation results for each metric.""", - ) - win_rates: Optional[dict[str, WinRateStats]] = Field( - default=None, - description="""A dictionary of win rates for each metric, only populated for multi-response evaluation runs.""", - ) - evaluation_dataset: Optional[list[EvaluationDataset]] = Field( - default=None, - description="""The input evaluation dataset(s) for the evaluation run.""", - ) - metadata: Optional[EvaluationRunMetadata] = Field( - default=None, description="""Metadata for the evaluation run.""" + instances: Optional[list[RougeInstance]] = Field( + default=None, description="""Required. Repeated rouge instances.""" ) - agent_info: Optional[evals_types.AgentInfo] = Field( - default=None, - description="""The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""", + metric_spec: Optional[genai_types.RougeSpec] = Field( + default=None, description="""Required. Spec for rouge score metric.""" ) - def show(self, candidate_names: Optional[List[str]] = None) -> None: - """Shows the evaluation result. - Args: - candidate_names: list of names for the evaluated candidates, used in - comparison reports. - """ - from .. import _evals_visualization +class RougeInputDict(TypedDict, total=False): + """Rouge input.""" - _evals_visualization.display_evaluation_result(self, candidate_names) + instances: Optional[list[RougeInstanceDict]] + """Required. Repeated rouge instances.""" + metric_spec: Optional[genai_types.RougeSpec] + """Required. Spec for rouge score metric.""" -class EvaluationResultDict(TypedDict, total=False): - """Result of an evaluation run for an evaluation dataset.""" - eval_case_results: Optional[list[EvalCaseResultDict]] - """A list of evaluation results for each evaluation case.""" +RougeInputOrDict = Union[RougeInput, RougeInputDict] - summary_metrics: Optional[list[AggregatedMetricResultDict]] - """A list of summary-level evaluation results for each metric.""" - win_rates: Optional[dict[str, WinRateStatsDict]] - """A dictionary of win rates for each metric, only populated for multi-response evaluation runs.""" +class ContentMap(_common.BaseModel): + """Map of placeholder in metric prompt template to contents of model input.""" - evaluation_dataset: Optional[list[EvaluationDatasetDict]] - """The input evaluation dataset(s) for the evaluation run.""" + values: Optional[dict[str, "ContentMapContents"]] = Field( + default=None, description="""Map of placeholder to contents.""" + ) - metadata: Optional[EvaluationRunMetadataDict] - """Metadata for the evaluation run.""" - agent_info: Optional[evals_types.AgentInfo] - """The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""" +class ContentMapDict(TypedDict, total=False): + """Map of placeholder in metric prompt template to contents of model input.""" + + values: Optional[dict[str, "ContentMapContents"]] + """Map of placeholder to contents.""" -EvaluationResultOrDict = Union[EvaluationResult, EvaluationResultDict] +ContentMapOrDict = Union[ContentMap, ContentMapDict] -class EvaluationRun(_common.BaseModel): - """Represents an evaluation run.""" +class PointwiseMetricInstance(_common.BaseModel): + """Pointwise metric instance.""" - name: Optional[str] = Field(default=None, description="""""") - display_name: Optional[str] = Field(default=None, description="""""") - metadata: Optional[dict[str, Any]] = Field(default=None, description="""""") - create_time: Optional[datetime.datetime] = Field(default=None, description="""""") - completion_time: Optional[datetime.datetime] = Field( - default=None, description="""""" - ) - state: Optional[EvaluationRunState] = Field(default=None, description="""""") - evaluation_set_snapshot: Optional[str] = Field(default=None, description="""""") - error: Optional[genai_types.GoogleRpcStatus] = Field( - default=None, description="""""" - ) - data_source: Optional[EvaluationRunDataSource] = Field( - default=None, description="""""" - ) - evaluation_run_results: Optional[EvaluationRunResults] = Field( - default=None, description="""The evaluation run formatted results.""" - ) - evaluation_item_results: Optional[EvaluationResult] = Field( + json_instance: Optional[str] = Field( default=None, - description="""The parsed EvaluationItem results for the evaluation run. This is only populated when include_evaluation_items is set to True.""", - ) - evaluation_config: Optional[EvaluationRunConfig] = Field( - default=None, description="""The evaluation config for the evaluation run.""" - ) - inference_configs: Optional[dict[str, EvaluationRunInferenceConfig]] = Field( - default=None, description="""The inference configs for the evaluation run.""" + description="""Instance specified as a json string. String key-value pairs are expected in the json_instance to render PointwiseMetricSpec.instance_prompt_template.""", ) - labels: Optional[dict[str, str]] = Field(default=None, description="""""") - analysis_configs: Optional[list[AnalysisConfig]] = Field( + content_map_instance: Optional[ContentMap] = Field( default=None, - description="""The analysis configurations for the evaluation run.""", + description="""Key-value contents for the mutlimodality input, including text, image, video, audio, and pdf, etc. The key is placeholder in metric prompt template, and the value is the multimodal content.""", ) - # TODO(b/448806531): Remove all the overridden _from_response methods once the - # ticket is resolved and published. - @classmethod - def _from_response( - cls: typing.Type["EvaluationRun"], - *, - response: dict[str, object], - kwargs: dict[str, object], - ) -> "EvaluationRun": - """Converts a dictionary response into a EvaluationRun object.""" - - snaked_response = _camel_key_to_snake(response) - - evaluation_run_results = response.get("evaluation_run_results") - if ( - isinstance(evaluation_run_results, dict) - and "summaryMetrics" in evaluation_run_results - ): - snaked_response["evaluation_run_results"]["summary_metrics"] = ( - evaluation_run_results["summaryMetrics"] - ) - result = super()._from_response(response=snaked_response, kwargs=kwargs) - return result +class PointwiseMetricInstanceDict(TypedDict, total=False): + """Pointwise metric instance.""" - def show(self) -> None: - """Shows the evaluation result.""" - from .. import _evals_visualization + json_instance: Optional[str] + """Instance specified as a json string. String key-value pairs are expected in the json_instance to render PointwiseMetricSpec.instance_prompt_template.""" - if self.state == "SUCCEEDED": - if self.evaluation_item_results is not None: - _evals_visualization.display_evaluation_result( - self.evaluation_item_results, None - ) - else: - logger.warning( - "Evaluation Run succeeded but no evaluation item results found. To display results, please set include_evaluation_items to True when calling get_evaluation_run()." - ) - # Show loss analysis results if present on the evaluation run. - # Pass the eval item map so the visualization can enrich - # loss examples with scenario/rubric data. - if ( - self.evaluation_run_results - and self.evaluation_run_results.loss_analysis_results - ): - eval_item_map = getattr(self, "_eval_item_map", None) - _evals_visualization.display_loss_analysis_results( - self.evaluation_run_results.loss_analysis_results, - eval_item_map=eval_item_map, - ) - else: - _evals_visualization.display_evaluation_run_status(self) + content_map_instance: Optional[ContentMapDict] + """Key-value contents for the mutlimodality input, including text, image, video, audio, and pdf, etc. The key is placeholder in metric prompt template, and the value is the multimodal content.""" -class EvaluationRunDict(TypedDict, total=False): - """Represents an evaluation run.""" +PointwiseMetricInstanceOrDict = Union[ + PointwiseMetricInstance, PointwiseMetricInstanceDict +] - name: Optional[str] - """""" - display_name: Optional[str] - """""" +class PointwiseMetricInput(_common.BaseModel): + """Pointwise metric input.""" - metadata: Optional[dict[str, Any]] - """""" + instance: Optional[PointwiseMetricInstance] = Field( + default=None, description="""Required. Pointwise metric instance.""" + ) + metric_spec: Optional[genai_types.PointwiseMetricSpec] = Field( + default=None, description="""Required. Spec for pointwise metric.""" + ) - create_time: Optional[datetime.datetime] - """""" - completion_time: Optional[datetime.datetime] - """""" +class PointwiseMetricInputDict(TypedDict, total=False): + """Pointwise metric input.""" - state: Optional[EvaluationRunState] - """""" + instance: Optional[PointwiseMetricInstanceDict] + """Required. Pointwise metric instance.""" - evaluation_set_snapshot: Optional[str] - """""" + metric_spec: Optional[genai_types.PointwiseMetricSpec] + """Required. Spec for pointwise metric.""" - error: Optional[genai_types.GoogleRpcStatus] - """""" - data_source: Optional[EvaluationRunDataSourceDict] - """""" +PointwiseMetricInputOrDict = Union[PointwiseMetricInput, PointwiseMetricInputDict] - evaluation_run_results: Optional[EvaluationRunResultsDict] - """The evaluation run formatted results.""" - evaluation_item_results: Optional[EvaluationResultDict] - """The parsed EvaluationItem results for the evaluation run. This is only populated when include_evaluation_items is set to True.""" +class PairwiseMetricInstance(_common.BaseModel): + """Pairwise metric instance.""" - evaluation_config: Optional[EvaluationRunConfigDict] - """The evaluation config for the evaluation run.""" + json_instance: Optional[str] = Field( + default=None, + description="""Instance specified as a json string. String key-value pairs are expected in the json_instance to render PairwiseMetricSpec.instance_prompt_template.""", + ) - inference_configs: Optional[dict[str, EvaluationRunInferenceConfigDict]] - """The inference configs for the evaluation run.""" - labels: Optional[dict[str, str]] - """""" +class PairwiseMetricInstanceDict(TypedDict, total=False): + """Pairwise metric instance.""" - analysis_configs: Optional[list[AnalysisConfigDict]] - """The analysis configurations for the evaluation run.""" + json_instance: Optional[str] + """Instance specified as a json string. String key-value pairs are expected in the json_instance to render PairwiseMetricSpec.instance_prompt_template.""" -EvaluationRunOrDict = Union[EvaluationRun, EvaluationRunDict] +PairwiseMetricInstanceOrDict = Union[PairwiseMetricInstance, PairwiseMetricInstanceDict] -class CreateEvaluationSetConfig(_common.BaseModel): - """Config to create an evaluation set.""" +class PairwiseMetricInput(_common.BaseModel): + """Pairwise metric instance.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + instance: Optional[PairwiseMetricInstance] = Field( + default=None, description="""Required. Pairwise metric instance.""" + ) + metric_spec: Optional[genai_types.PairwiseMetricSpec] = Field( + default=None, description="""Required. Spec for pairwise metric.""" ) -class CreateEvaluationSetConfigDict(TypedDict, total=False): - """Config to create an evaluation set.""" +class PairwiseMetricInputDict(TypedDict, total=False): + """Pairwise metric instance.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + instance: Optional[PairwiseMetricInstanceDict] + """Required. Pairwise metric instance.""" + metric_spec: Optional[genai_types.PairwiseMetricSpec] + """Required. Spec for pairwise metric.""" -CreateEvaluationSetConfigOrDict = Union[ - CreateEvaluationSetConfig, CreateEvaluationSetConfigDict -] +PairwiseMetricInputOrDict = Union[PairwiseMetricInput, PairwiseMetricInputDict] -class _CreateEvaluationSetParameters(_common.BaseModel): - """Represents a job that creates an evaluation set.""" - evaluation_items: Optional[list[str]] = Field(default=None, description="""""") - display_name: Optional[str] = Field(default=None, description="""""") - config: Optional[CreateEvaluationSetConfig] = Field( - default=None, description="""""" - ) +class ToolCallValidInstance(_common.BaseModel): + """Tool call valid instance.""" - -class _CreateEvaluationSetParametersDict(TypedDict, total=False): - """Represents a job that creates an evaluation set.""" - - evaluation_items: Optional[list[str]] - """""" - - display_name: Optional[str] - """""" - - config: Optional[CreateEvaluationSetConfigDict] - """""" - - -_CreateEvaluationSetParametersOrDict = Union[ - _CreateEvaluationSetParameters, _CreateEvaluationSetParametersDict -] - - -class EvaluationSet(_common.BaseModel): - """Represents an evaluation set.""" - - name: Optional[str] = Field( - default=None, description="""The resource name of the evaluation set.""" - ) - display_name: Optional[str] = Field( - default=None, description="""The display name of the evaluation set.""" + prediction: Optional[str] = Field( + default=None, description="""Required. Output of the evaluated model.""" ) - evaluation_items: Optional[list[str]] = Field( + reference: Optional[str] = Field( default=None, - description="""The EvaluationItems that are part of this dataset.""", - ) - create_time: Optional[datetime.datetime] = Field( - default=None, description="""The create time of the evaluation set.""" - ) - update_time: Optional[datetime.datetime] = Field( - default=None, description="""The update time of the evaluation set.""" - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, description="""The metadata of the evaluation set.""" + description="""Required. Ground truth used to compare against the prediction.""", ) -class EvaluationSetDict(TypedDict, total=False): - """Represents an evaluation set.""" - - name: Optional[str] - """The resource name of the evaluation set.""" - - display_name: Optional[str] - """The display name of the evaluation set.""" - - evaluation_items: Optional[list[str]] - """The EvaluationItems that are part of this dataset.""" - - create_time: Optional[datetime.datetime] - """The create time of the evaluation set.""" +class ToolCallValidInstanceDict(TypedDict, total=False): + """Tool call valid instance.""" - update_time: Optional[datetime.datetime] - """The update time of the evaluation set.""" + prediction: Optional[str] + """Required. Output of the evaluated model.""" - metadata: Optional[dict[str, Any]] - """The metadata of the evaluation set.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" -EvaluationSetOrDict = Union[EvaluationSet, EvaluationSetDict] +ToolCallValidInstanceOrDict = Union[ToolCallValidInstance, ToolCallValidInstanceDict] -class DeleteEvaluationExperimentConfig(_common.BaseModel): - """Config for deleting an evaluation experiment.""" +class ToolCallValidSpec(_common.BaseModel): + """Spec for tool call valid metric.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + pass -class DeleteEvaluationExperimentConfigDict(TypedDict, total=False): - """Config for deleting an evaluation experiment.""" +class ToolCallValidSpecDict(TypedDict, total=False): + """Spec for tool call valid metric.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + pass -DeleteEvaluationExperimentConfigOrDict = Union[ - DeleteEvaluationExperimentConfig, DeleteEvaluationExperimentConfigDict -] +ToolCallValidSpecOrDict = Union[ToolCallValidSpec, ToolCallValidSpecDict] -class _DeleteEvaluationExperimentParameters(_common.BaseModel): - """Parameters for deleting an evaluation experiment.""" +class ToolCallValidInput(_common.BaseModel): + """Tool call valid input.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[DeleteEvaluationExperimentConfig] = Field( - default=None, description="""""" + instances: Optional[list[ToolCallValidInstance]] = Field( + default=None, description="""Required. Repeated tool call valid instances.""" + ) + metric_spec: Optional[ToolCallValidSpec] = Field( + default=None, description="""Required. Spec for tool call valid metric.""" ) -class _DeleteEvaluationExperimentParametersDict(TypedDict, total=False): - """Parameters for deleting an evaluation experiment.""" +class ToolCallValidInputDict(TypedDict, total=False): + """Tool call valid input.""" - name: Optional[str] - """""" + instances: Optional[list[ToolCallValidInstanceDict]] + """Required. Repeated tool call valid instances.""" - config: Optional[DeleteEvaluationExperimentConfigDict] - """""" + metric_spec: Optional[ToolCallValidSpecDict] + """Required. Spec for tool call valid metric.""" -_DeleteEvaluationExperimentParametersOrDict = Union[ - _DeleteEvaluationExperimentParameters, _DeleteEvaluationExperimentParametersDict -] +ToolCallValidInputOrDict = Union[ToolCallValidInput, ToolCallValidInputDict] -class DeleteEvaluationExperimentOperation(_common.BaseModel): - """Operation for deleting an evaluation experiment.""" +class ToolNameMatchInstance(_common.BaseModel): + """Tool name match instance.""" - name: Optional[str] = Field( - default=None, - description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, - description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", - ) - done: Optional[bool] = Field( - default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + prediction: Optional[str] = Field( + default=None, description="""Required. Output of the evaluated model.""" ) - error: Optional[dict[str, Any]] = Field( + reference: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", - ) - - -class DeleteEvaluationExperimentOperationDict(TypedDict, total=False): - """Operation for deleting an evaluation experiment.""" - - name: Optional[str] - """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - - metadata: Optional[dict[str, Any]] - """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - - done: Optional[bool] - """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" - - -DeleteEvaluationExperimentOperationOrDict = Union[ - DeleteEvaluationExperimentOperation, DeleteEvaluationExperimentOperationDict -] - - -class DeleteEvaluationMetricConfig(_common.BaseModel): - """Config for deleting an evaluation metric.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + description="""Required. Ground truth used to compare against the prediction.""", ) -class DeleteEvaluationMetricConfigDict(TypedDict, total=False): - """Config for deleting an evaluation metric.""" +class ToolNameMatchInstanceDict(TypedDict, total=False): + """Tool name match instance.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + prediction: Optional[str] + """Required. Output of the evaluated model.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" -DeleteEvaluationMetricConfigOrDict = Union[ - DeleteEvaluationMetricConfig, DeleteEvaluationMetricConfigDict -] +ToolNameMatchInstanceOrDict = Union[ToolNameMatchInstance, ToolNameMatchInstanceDict] -class _DeleteEvaluationMetricParameters(_common.BaseModel): - """Parameters for deleting an evaluation metric.""" - metric_resource_name: Optional[str] = Field(default=None, description="""""") - config: Optional[DeleteEvaluationMetricConfig] = Field( - default=None, description="""""" - ) +class ToolNameMatchSpec(_common.BaseModel): + """Spec for tool name match metric.""" + pass -class _DeleteEvaluationMetricParametersDict(TypedDict, total=False): - """Parameters for deleting an evaluation metric.""" - metric_resource_name: Optional[str] - """""" +class ToolNameMatchSpecDict(TypedDict, total=False): + """Spec for tool name match metric.""" - config: Optional[DeleteEvaluationMetricConfigDict] - """""" + pass -_DeleteEvaluationMetricParametersOrDict = Union[ - _DeleteEvaluationMetricParameters, _DeleteEvaluationMetricParametersDict -] +ToolNameMatchSpecOrDict = Union[ToolNameMatchSpec, ToolNameMatchSpecDict] -class DeleteEvaluationMetricOperation(_common.BaseModel): - """Operation for deleting an evaluation metric.""" +class ToolNameMatchInput(_common.BaseModel): + """Tool name match input.""" - name: Optional[str] = Field( - default=None, - description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, - description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", - ) - done: Optional[bool] = Field( - default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + instances: Optional[list[ToolNameMatchInstance]] = Field( + default=None, description="""Required. Repeated tool name match instances.""" ) - error: Optional[dict[str, Any]] = Field( - default=None, - description="""The error result of the operation in case of failure or cancellation.""", + metric_spec: Optional[ToolNameMatchSpec] = Field( + default=None, description="""Required. Spec for tool name match metric.""" ) -class DeleteEvaluationMetricOperationDict(TypedDict, total=False): - """Operation for deleting an evaluation metric.""" - - name: Optional[str] - """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - - metadata: Optional[dict[str, Any]] - """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" +class ToolNameMatchInputDict(TypedDict, total=False): + """Tool name match input.""" - done: Optional[bool] - """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + instances: Optional[list[ToolNameMatchInstanceDict]] + """Required. Repeated tool name match instances.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + metric_spec: Optional[ToolNameMatchSpecDict] + """Required. Spec for tool name match metric.""" -DeleteEvaluationMetricOperationOrDict = Union[ - DeleteEvaluationMetricOperation, DeleteEvaluationMetricOperationDict -] +ToolNameMatchInputOrDict = Union[ToolNameMatchInput, ToolNameMatchInputDict] -class BleuInstance(_common.BaseModel): - """Bleu instance.""" +class ToolParameterKeyMatchInstance(_common.BaseModel): + """Tool parameter key match instance.""" prediction: Optional[str] = Field( default=None, description="""Required. Output of the evaluated model.""" @@ -4332,8 +3827,8 @@ class BleuInstance(_common.BaseModel): ) -class BleuInstanceDict(TypedDict, total=False): - """Bleu instance.""" +class ToolParameterKeyMatchInstanceDict(TypedDict, total=False): + """Tool parameter key match instance.""" prediction: Optional[str] """Required. Output of the evaluated model.""" @@ -4342,33 +3837,58 @@ class BleuInstanceDict(TypedDict, total=False): """Required. Ground truth used to compare against the prediction.""" -BleuInstanceOrDict = Union[BleuInstance, BleuInstanceDict] - - -class BleuInput(_common.BaseModel): +ToolParameterKeyMatchInstanceOrDict = Union[ + ToolParameterKeyMatchInstance, ToolParameterKeyMatchInstanceDict +] - instances: Optional[list[BleuInstance]] = Field( - default=None, description="""Required. Repeated bleu instances.""" + +class ToolParameterKeyMatchSpec(_common.BaseModel): + """Spec for tool parameter key match metric.""" + + pass + + +class ToolParameterKeyMatchSpecDict(TypedDict, total=False): + """Spec for tool parameter key match metric.""" + + pass + + +ToolParameterKeyMatchSpecOrDict = Union[ + ToolParameterKeyMatchSpec, ToolParameterKeyMatchSpecDict +] + + +class ToolParameterKeyMatchInput(_common.BaseModel): + """Tool parameter key match input.""" + + instances: Optional[list[ToolParameterKeyMatchInstance]] = Field( + default=None, + description="""Required. Repeated tool parameter key match instances.""", ) - metric_spec: Optional[genai_types.BleuSpec] = Field( - default=None, description="""Required. Spec for bleu score metric.""" + metric_spec: Optional[ToolParameterKeyMatchSpec] = Field( + default=None, + description="""Required. Spec for tool parameter key match metric.""", ) -class BleuInputDict(TypedDict, total=False): +class ToolParameterKeyMatchInputDict(TypedDict, total=False): + """Tool parameter key match input.""" - instances: Optional[list[BleuInstanceDict]] - """Required. Repeated bleu instances.""" + instances: Optional[list[ToolParameterKeyMatchInstanceDict]] + """Required. Repeated tool parameter key match instances.""" - metric_spec: Optional[genai_types.BleuSpec] - """Required. Spec for bleu score metric.""" + metric_spec: Optional[ToolParameterKeyMatchSpecDict] + """Required. Spec for tool parameter key match metric.""" -BleuInputOrDict = Union[BleuInput, BleuInputDict] +ToolParameterKeyMatchInputOrDict = Union[ + ToolParameterKeyMatchInput, ToolParameterKeyMatchInputDict +] -class ExactMatchInstance(_common.BaseModel): - """Exact match instance.""" +class ToolParameterKVMatchInstance(_common.BaseModel): + """Tool parameter kv match instance.""" prediction: Optional[str] = Field( default=None, description="""Required. Output of the evaluated model.""" @@ -4379,8 +3899,8 @@ class ExactMatchInstance(_common.BaseModel): ) -class ExactMatchInstanceDict(TypedDict, total=False): - """Exact match instance.""" +class ToolParameterKVMatchInstanceDict(TypedDict, total=False): + """Tool parameter kv match instance.""" prediction: Optional[str] """Required. Output of the evaluated model.""" @@ -4389,4745 +3909,4734 @@ class ExactMatchInstanceDict(TypedDict, total=False): """Required. Ground truth used to compare against the prediction.""" -ExactMatchInstanceOrDict = Union[ExactMatchInstance, ExactMatchInstanceDict] +ToolParameterKVMatchInstanceOrDict = Union[ + ToolParameterKVMatchInstance, ToolParameterKVMatchInstanceDict +] -class ExactMatchSpec(_common.BaseModel): - """Spec for exact match metric.""" +class ToolParameterKVMatchSpec(_common.BaseModel): + """Spec for tool parameter kv match metric.""" - pass + use_strict_string_match: Optional[bool] = Field( + default=None, + description="""Optional. Whether to use STRICT string match on parameter values.""", + ) -class ExactMatchSpecDict(TypedDict, total=False): - """Spec for exact match metric.""" +class ToolParameterKVMatchSpecDict(TypedDict, total=False): + """Spec for tool parameter kv match metric.""" - pass + use_strict_string_match: Optional[bool] + """Optional. Whether to use STRICT string match on parameter values.""" -ExactMatchSpecOrDict = Union[ExactMatchSpec, ExactMatchSpecDict] +ToolParameterKVMatchSpecOrDict = Union[ + ToolParameterKVMatchSpec, ToolParameterKVMatchSpecDict +] -class ExactMatchInput(_common.BaseModel): +class ToolParameterKVMatchInput(_common.BaseModel): + """Tool parameter kv match input.""" - instances: Optional[list[ExactMatchInstance]] = Field( - default=None, description="""Required. Repeated exact match instances.""" + instances: Optional[list[ToolParameterKVMatchInstance]] = Field( + default=None, + description="""Required. Repeated tool parameter key value match instances.""", ) - metric_spec: Optional[ExactMatchSpec] = Field( - default=None, description="""Required. Spec for exact match metric.""" + metric_spec: Optional[ToolParameterKVMatchSpec] = Field( + default=None, + description="""Required. Spec for tool parameter key value match metric.""", ) -class ExactMatchInputDict(TypedDict, total=False): +class ToolParameterKVMatchInputDict(TypedDict, total=False): + """Tool parameter kv match input.""" - instances: Optional[list[ExactMatchInstanceDict]] - """Required. Repeated exact match instances.""" + instances: Optional[list[ToolParameterKVMatchInstanceDict]] + """Required. Repeated tool parameter key value match instances.""" - metric_spec: Optional[ExactMatchSpecDict] - """Required. Spec for exact match metric.""" + metric_spec: Optional[ToolParameterKVMatchSpecDict] + """Required. Spec for tool parameter key value match metric.""" -ExactMatchInputOrDict = Union[ExactMatchInput, ExactMatchInputDict] +ToolParameterKVMatchInputOrDict = Union[ + ToolParameterKVMatchInput, ToolParameterKVMatchInputDict +] -class RougeInstance(_common.BaseModel): - """Rouge instance.""" +class MapInstance(_common.BaseModel): + """Instance data specified as a map.""" - prediction: Optional[str] = Field( - default=None, description="""Required. Output of the evaluated model.""" - ) - reference: Optional[str] = Field( - default=None, - description="""Required. Ground truth used to compare against the prediction.""", + map_instance: Optional[dict[str, evals_types.InstanceData]] = Field( + default=None, description="""Map of instance data.""" ) -class RougeInstanceDict(TypedDict, total=False): - """Rouge instance.""" - - prediction: Optional[str] - """Required. Output of the evaluated model.""" +class MapInstanceDict(TypedDict, total=False): + """Instance data specified as a map.""" - reference: Optional[str] - """Required. Ground truth used to compare against the prediction.""" + map_instance: Optional[dict[str, evals_types.InstanceData]] + """Map of instance data.""" -RougeInstanceOrDict = Union[RougeInstance, RougeInstanceDict] +MapInstanceOrDict = Union[MapInstance, MapInstanceDict] -class RougeInput(_common.BaseModel): - """Rouge input.""" +class EvaluationInstance(_common.BaseModel): + """A single instance to be evaluated.""" - instances: Optional[list[RougeInstance]] = Field( - default=None, description="""Required. Repeated rouge instances.""" + prompt: Optional[evals_types.InstanceData] = Field( + default=None, + description="""Data used to populate placeholder `prompt` in a metric prompt template.""", ) - metric_spec: Optional[genai_types.RougeSpec] = Field( - default=None, description="""Required. Spec for rouge score metric.""" + response: Optional[evals_types.InstanceData] = Field( + default=None, + description="""Data used to populate placeholder `response` in a metric prompt template.""", + ) + reference: Optional[evals_types.InstanceData] = Field( + default=None, + description="""Data used to populate placeholder `reference` in a metric prompt template.""", + ) + other_data: Optional[MapInstance] = Field( + default=None, + description="""Other data used to populate placeholders based on their key.""", + ) + agent_data: Optional[evals_types.AgentData] = Field( + default=None, description="""Data used for agent evaluation.""" + ) + rubric_groups: Optional[dict[str, "RubricGroup"]] = Field( + default=None, + description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", + ) + interactions_data_source: Optional[InteractionsDataSource] = Field( + default=None, + description="""Source for populating agent data from an Interactions API + interaction. If set, no other agent data source may be set. The backend + fetches the interaction (and the agent that produced it) and parses it + into agent data for grading.""", ) -class RougeInputDict(TypedDict, total=False): - """Rouge input.""" - - instances: Optional[list[RougeInstanceDict]] - """Required. Repeated rouge instances.""" - - metric_spec: Optional[genai_types.RougeSpec] - """Required. Spec for rouge score metric.""" - +class EvaluationInstanceDict(TypedDict, total=False): + """A single instance to be evaluated.""" -RougeInputOrDict = Union[RougeInput, RougeInputDict] + prompt: Optional[evals_types.InstanceData] + """Data used to populate placeholder `prompt` in a metric prompt template.""" + response: Optional[evals_types.InstanceData] + """Data used to populate placeholder `response` in a metric prompt template.""" -class ContentMap(_common.BaseModel): - """Map of placeholder in metric prompt template to contents of model input.""" + reference: Optional[evals_types.InstanceData] + """Data used to populate placeholder `reference` in a metric prompt template.""" - values: Optional[dict[str, "ContentMapContents"]] = Field( - default=None, description="""Map of placeholder to contents.""" - ) + other_data: Optional[MapInstanceDict] + """Other data used to populate placeholders based on their key.""" + agent_data: Optional[evals_types.AgentData] + """Data used for agent evaluation.""" -class ContentMapDict(TypedDict, total=False): - """Map of placeholder in metric prompt template to contents of model input.""" + rubric_groups: Optional[dict[str, "RubricGroupDict"]] + """Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""" - values: Optional[dict[str, "ContentMapContents"]] - """Map of placeholder to contents.""" + interactions_data_source: Optional[InteractionsDataSourceDict] + """Source for populating agent data from an Interactions API + interaction. If set, no other agent data source may be set. The backend + fetches the interaction (and the agent that produced it) and parses it + into agent data for grading.""" -ContentMapOrDict = Union[ContentMap, ContentMapDict] +EvaluationInstanceOrDict = Union[EvaluationInstance, EvaluationInstanceDict] -class PointwiseMetricInstance(_common.BaseModel): - """Pointwise metric instance.""" +class EvaluateInstancesConfig(_common.BaseModel): + """Config for evaluate instances.""" - json_instance: Optional[str] = Field( - default=None, - description="""Instance specified as a json string. String key-value pairs are expected in the json_instance to render PointwiseMetricSpec.instance_prompt_template.""", - ) - content_map_instance: Optional[ContentMap] = Field( - default=None, - description="""Key-value contents for the mutlimodality input, including text, image, video, audio, and pdf, etc. The key is placeholder in metric prompt template, and the value is the multimodal content.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class PointwiseMetricInstanceDict(TypedDict, total=False): - """Pointwise metric instance.""" - - json_instance: Optional[str] - """Instance specified as a json string. String key-value pairs are expected in the json_instance to render PointwiseMetricSpec.instance_prompt_template.""" +class EvaluateInstancesConfigDict(TypedDict, total=False): + """Config for evaluate instances.""" - content_map_instance: Optional[ContentMapDict] - """Key-value contents for the mutlimodality input, including text, image, video, audio, and pdf, etc. The key is placeholder in metric prompt template, and the value is the multimodal content.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -PointwiseMetricInstanceOrDict = Union[ - PointwiseMetricInstance, PointwiseMetricInstanceDict +EvaluateInstancesConfigOrDict = Union[ + EvaluateInstancesConfig, EvaluateInstancesConfigDict ] -class PointwiseMetricInput(_common.BaseModel): - """Pointwise metric input.""" +class RubricBasedMetricSpec(_common.BaseModel): + """Specification for a metric that is based on rubrics.""" - instance: Optional[PointwiseMetricInstance] = Field( - default=None, description="""Required. Pointwise metric instance.""" + metric_prompt_template: Optional[str] = Field( + default=None, + description="""Template for the prompt used by the judge model to evaluate against + rubrics.""", ) - metric_spec: Optional[genai_types.PointwiseMetricSpec] = Field( - default=None, description="""Required. Spec for pointwise metric.""" + judge_autorater_config: Optional[genai_types.AutoraterConfig] = Field( + default=None, + description="""Optional configuration for the judge LLM (Autorater).""", ) - - -class PointwiseMetricInputDict(TypedDict, total=False): - """Pointwise metric input.""" - - instance: Optional[PointwiseMetricInstanceDict] - """Required. Pointwise metric instance.""" - - metric_spec: Optional[genai_types.PointwiseMetricSpec] - """Required. Spec for pointwise metric.""" - - -PointwiseMetricInputOrDict = Union[PointwiseMetricInput, PointwiseMetricInputDict] - - -class PairwiseMetricInstance(_common.BaseModel): - """Pairwise metric instance.""" - - json_instance: Optional[str] = Field( + inline_rubrics: Optional[list[evals_types.Rubric]] = Field( + default=None, description="""Use rubrics provided directly in the spec.""" + ) + rubric_group_key: Optional[str] = Field( default=None, - description="""Instance specified as a json string. String key-value pairs are expected in the json_instance to render PairwiseMetricSpec.instance_prompt_template.""", + description="""Use a pre-defined group of rubrics associated with the input content. + This refers to a key in the `rubric_groups` map of + `RubricEnhancedContents`.""", + ) + rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] = Field( + default=None, + description="""Dynamically generate rubrics for evaluation using this specification.""", ) -class PairwiseMetricInstanceDict(TypedDict, total=False): - """Pairwise metric instance.""" - - json_instance: Optional[str] - """Instance specified as a json string. String key-value pairs are expected in the json_instance to render PairwiseMetricSpec.instance_prompt_template.""" - - -PairwiseMetricInstanceOrDict = Union[PairwiseMetricInstance, PairwiseMetricInstanceDict] - - -class PairwiseMetricInput(_common.BaseModel): - """Pairwise metric instance.""" +class RubricBasedMetricSpecDict(TypedDict, total=False): + """Specification for a metric that is based on rubrics.""" - instance: Optional[PairwiseMetricInstance] = Field( - default=None, description="""Required. Pairwise metric instance.""" - ) - metric_spec: Optional[genai_types.PairwiseMetricSpec] = Field( - default=None, description="""Required. Spec for pairwise metric.""" - ) + metric_prompt_template: Optional[str] + """Template for the prompt used by the judge model to evaluate against + rubrics.""" + judge_autorater_config: Optional[genai_types.AutoraterConfig] + """Optional configuration for the judge LLM (Autorater).""" -class PairwiseMetricInputDict(TypedDict, total=False): - """Pairwise metric instance.""" + inline_rubrics: Optional[list[evals_types.Rubric]] + """Use rubrics provided directly in the spec.""" - instance: Optional[PairwiseMetricInstanceDict] - """Required. Pairwise metric instance.""" + rubric_group_key: Optional[str] + """Use a pre-defined group of rubrics associated with the input content. + This refers to a key in the `rubric_groups` map of + `RubricEnhancedContents`.""" - metric_spec: Optional[genai_types.PairwiseMetricSpec] - """Required. Spec for pairwise metric.""" + rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] + """Dynamically generate rubrics for evaluation using this specification.""" -PairwiseMetricInputOrDict = Union[PairwiseMetricInput, PairwiseMetricInputDict] +RubricBasedMetricSpecOrDict = Union[RubricBasedMetricSpec, RubricBasedMetricSpecDict] -class ToolCallValidInstance(_common.BaseModel): - """Tool call valid instance.""" +class RubricEnhancedContents(_common.BaseModel): + """Rubric-enhanced contents for evaluation.""" - prediction: Optional[str] = Field( - default=None, description="""Required. Output of the evaluated model.""" + prompt: Optional[list[genai_types.Content]] = Field( + default=None, + description="""User prompt, using the standard Content type from the Gen AI SDK.""", ) - reference: Optional[str] = Field( + rubric_groups: Optional[dict[str, "RubricGroup"]] = Field( default=None, - description="""Required. Ground truth used to compare against the prediction.""", + description="""Named groups of rubrics associated with this prompt. + The key is a user-defined name for the rubric group.""", + ) + response: Optional[list[genai_types.Content]] = Field( + default=None, + description="""Response, using the standard Content type from the Gen AI SDK.""", + ) + other_content: Optional[ContentMap] = Field( + default=None, + description="""Other contents needed for the metric. + For example, if `reference` is needed for the metric, it can be provided + here.""", ) -class ToolCallValidInstanceDict(TypedDict, total=False): - """Tool call valid instance.""" - - prediction: Optional[str] - """Required. Output of the evaluated model.""" - - reference: Optional[str] - """Required. Ground truth used to compare against the prediction.""" - - -ToolCallValidInstanceOrDict = Union[ToolCallValidInstance, ToolCallValidInstanceDict] - +class RubricEnhancedContentsDict(TypedDict, total=False): + """Rubric-enhanced contents for evaluation.""" -class ToolCallValidSpec(_common.BaseModel): - """Spec for tool call valid metric.""" + prompt: Optional[list[genai_types.Content]] + """User prompt, using the standard Content type from the Gen AI SDK.""" - pass + rubric_groups: Optional[dict[str, "RubricGroup"]] + """Named groups of rubrics associated with this prompt. + The key is a user-defined name for the rubric group.""" + response: Optional[list[genai_types.Content]] + """Response, using the standard Content type from the Gen AI SDK.""" -class ToolCallValidSpecDict(TypedDict, total=False): - """Spec for tool call valid metric.""" + other_content: Optional[ContentMapDict] + """Other contents needed for the metric. + For example, if `reference` is needed for the metric, it can be provided + here.""" - pass +RubricEnhancedContentsOrDict = Union[RubricEnhancedContents, RubricEnhancedContentsDict] -ToolCallValidSpecOrDict = Union[ToolCallValidSpec, ToolCallValidSpecDict] +class RubricBasedMetricInstance(_common.BaseModel): + """Defines an instance for Rubric-based metrics. -class ToolCallValidInput(_common.BaseModel): - """Tool call valid input.""" + This class allows various input formats. + """ - instances: Optional[list[ToolCallValidInstance]] = Field( - default=None, description="""Required. Repeated tool call valid instances.""" - ) - metric_spec: Optional[ToolCallValidSpec] = Field( - default=None, description="""Required. Spec for tool call valid metric.""" + json_instance: Optional[str] = Field( + default=None, + description="""Specify evaluation fields and their string values in JSON format.""", ) - - -class ToolCallValidInputDict(TypedDict, total=False): - """Tool call valid input.""" - - instances: Optional[list[ToolCallValidInstanceDict]] - """Required. Repeated tool call valid instances.""" - - metric_spec: Optional[ToolCallValidSpecDict] - """Required. Spec for tool call valid metric.""" - - -ToolCallValidInputOrDict = Union[ToolCallValidInput, ToolCallValidInputDict] - - -class ToolNameMatchInstance(_common.BaseModel): - """Tool name match instance.""" - - prediction: Optional[str] = Field( - default=None, description="""Required. Output of the evaluated model.""" + content_map_instance: Optional[ContentMap] = Field( + default=None, + description="""Specify evaluation fields and their content values using a ContentMap.""", ) - reference: Optional[str] = Field( + rubric_enhanced_contents: Optional[RubricEnhancedContents] = Field( default=None, - description="""Required. Ground truth used to compare against the prediction.""", + description="""Provide input as Gemini Content along with one or more + associated rubric groups.""", ) -class ToolNameMatchInstanceDict(TypedDict, total=False): - """Tool name match instance.""" - - prediction: Optional[str] - """Required. Output of the evaluated model.""" - - reference: Optional[str] - """Required. Ground truth used to compare against the prediction.""" - - -ToolNameMatchInstanceOrDict = Union[ToolNameMatchInstance, ToolNameMatchInstanceDict] - - -class ToolNameMatchSpec(_common.BaseModel): - """Spec for tool name match metric.""" +class RubricBasedMetricInstanceDict(TypedDict, total=False): + """Defines an instance for Rubric-based metrics. - pass + This class allows various input formats. + """ + json_instance: Optional[str] + """Specify evaluation fields and their string values in JSON format.""" -class ToolNameMatchSpecDict(TypedDict, total=False): - """Spec for tool name match metric.""" + content_map_instance: Optional[ContentMapDict] + """Specify evaluation fields and their content values using a ContentMap.""" - pass + rubric_enhanced_contents: Optional[RubricEnhancedContentsDict] + """Provide input as Gemini Content along with one or more + associated rubric groups.""" -ToolNameMatchSpecOrDict = Union[ToolNameMatchSpec, ToolNameMatchSpecDict] +RubricBasedMetricInstanceOrDict = Union[ + RubricBasedMetricInstance, RubricBasedMetricInstanceDict +] -class ToolNameMatchInput(_common.BaseModel): - """Tool name match input.""" +class RubricBasedMetricInput(_common.BaseModel): + """Input for a rubric-based metrics.""" - instances: Optional[list[ToolNameMatchInstance]] = Field( - default=None, description="""Required. Repeated tool name match instances.""" + metric_spec: Optional[RubricBasedMetricSpec] = Field( + default=None, description="""Specification for the rubric-based metric.""" ) - metric_spec: Optional[ToolNameMatchSpec] = Field( - default=None, description="""Required. Spec for tool name match metric.""" + instance: Optional[RubricBasedMetricInstance] = Field( + default=None, description="""The instance to be evaluated.""" ) -class ToolNameMatchInputDict(TypedDict, total=False): - """Tool name match input.""" +class RubricBasedMetricInputDict(TypedDict, total=False): + """Input for a rubric-based metrics.""" - instances: Optional[list[ToolNameMatchInstanceDict]] - """Required. Repeated tool name match instances.""" + metric_spec: Optional[RubricBasedMetricSpecDict] + """Specification for the rubric-based metric.""" - metric_spec: Optional[ToolNameMatchSpecDict] - """Required. Spec for tool name match metric.""" + instance: Optional[RubricBasedMetricInstanceDict] + """The instance to be evaluated.""" -ToolNameMatchInputOrDict = Union[ToolNameMatchInput, ToolNameMatchInputDict] +RubricBasedMetricInputOrDict = Union[RubricBasedMetricInput, RubricBasedMetricInputDict] -class ToolParameterKeyMatchInstance(_common.BaseModel): - """Tool parameter key match instance.""" +class MetricSource(_common.BaseModel): + """The metric source used for evaluation.""" - prediction: Optional[str] = Field( - default=None, description="""Required. Output of the evaluated model.""" + metric: Optional[Metric] = Field( + default=None, description="""Inline metric config.""" ) - reference: Optional[str] = Field( + metric_resource_name: Optional[str] = Field( default=None, - description="""Required. Ground truth used to compare against the prediction.""", + description="""Resource name for registered metric. Example: + projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""", ) -class ToolParameterKeyMatchInstanceDict(TypedDict, total=False): - """Tool parameter key match instance.""" +class MetricSourceDict(TypedDict, total=False): + """The metric source used for evaluation.""" - prediction: Optional[str] - """Required. Output of the evaluated model.""" + metric: Optional[MetricDict] + """Inline metric config.""" - reference: Optional[str] - """Required. Ground truth used to compare against the prediction.""" + metric_resource_name: Optional[str] + """Resource name for registered metric. Example: + projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" -ToolParameterKeyMatchInstanceOrDict = Union[ - ToolParameterKeyMatchInstance, ToolParameterKeyMatchInstanceDict -] +MetricSourceOrDict = Union[MetricSource, MetricSourceDict] -class ToolParameterKeyMatchSpec(_common.BaseModel): - """Spec for tool parameter key match metric.""" +class _EvaluateInstancesRequestParameters(_common.BaseModel): + """Parameters for evaluating instances.""" - pass + bleu_input: Optional[BleuInput] = Field(default=None, description="""""") + exact_match_input: Optional[ExactMatchInput] = Field( + default=None, description="""""" + ) + rouge_input: Optional[RougeInput] = Field(default=None, description="""""") + pointwise_metric_input: Optional[PointwiseMetricInput] = Field( + default=None, description="""""" + ) + pairwise_metric_input: Optional[PairwiseMetricInput] = Field( + default=None, description="""""" + ) + tool_call_valid_input: Optional[ToolCallValidInput] = Field( + default=None, description="""""" + ) + tool_name_match_input: Optional[ToolNameMatchInput] = Field( + default=None, description="""""" + ) + tool_parameter_key_match_input: Optional[ToolParameterKeyMatchInput] = Field( + default=None, description="""""" + ) + tool_parameter_kv_match_input: Optional[ToolParameterKVMatchInput] = Field( + default=None, description="""""" + ) + rubric_based_metric_input: Optional[RubricBasedMetricInput] = Field( + default=None, description="""""" + ) + autorater_config: Optional[genai_types.AutoraterConfig] = Field( + default=None, + description="""Autorater config used for evaluation. Not applicable for predefined metrics (PredefinedMetricSpec); the server uses its own model configuration for predefined metrics and this field is ignored.""", + ) + metrics: Optional[list[Metric]] = Field( + default=None, + description="""The metrics used for evaluation. + Currently, we only support evaluating a single metric. If multiple metrics + are provided, only the first one will be evaluated.""", + ) + instance: Optional[EvaluationInstance] = Field( + default=None, description="""The instance to be evaluated.""" + ) + metric_sources: Optional[list[MetricSource]] = Field( + default=None, description="""The metrics used for evaluation.""" + ) + config: Optional[EvaluateInstancesConfig] = Field(default=None, description="""""") -class ToolParameterKeyMatchSpecDict(TypedDict, total=False): - """Spec for tool parameter key match metric.""" +class _EvaluateInstancesRequestParametersDict(TypedDict, total=False): + """Parameters for evaluating instances.""" - pass + bleu_input: Optional[BleuInputDict] + """""" + exact_match_input: Optional[ExactMatchInputDict] + """""" -ToolParameterKeyMatchSpecOrDict = Union[ - ToolParameterKeyMatchSpec, ToolParameterKeyMatchSpecDict -] + rouge_input: Optional[RougeInputDict] + """""" + pointwise_metric_input: Optional[PointwiseMetricInputDict] + """""" -class ToolParameterKeyMatchInput(_common.BaseModel): - """Tool parameter key match input.""" + pairwise_metric_input: Optional[PairwiseMetricInputDict] + """""" - instances: Optional[list[ToolParameterKeyMatchInstance]] = Field( - default=None, - description="""Required. Repeated tool parameter key match instances.""", - ) - metric_spec: Optional[ToolParameterKeyMatchSpec] = Field( - default=None, - description="""Required. Spec for tool parameter key match metric.""", - ) + tool_call_valid_input: Optional[ToolCallValidInputDict] + """""" + tool_name_match_input: Optional[ToolNameMatchInputDict] + """""" -class ToolParameterKeyMatchInputDict(TypedDict, total=False): - """Tool parameter key match input.""" + tool_parameter_key_match_input: Optional[ToolParameterKeyMatchInputDict] + """""" - instances: Optional[list[ToolParameterKeyMatchInstanceDict]] - """Required. Repeated tool parameter key match instances.""" + tool_parameter_kv_match_input: Optional[ToolParameterKVMatchInputDict] + """""" - metric_spec: Optional[ToolParameterKeyMatchSpecDict] - """Required. Spec for tool parameter key match metric.""" + rubric_based_metric_input: Optional[RubricBasedMetricInputDict] + """""" + + autorater_config: Optional[genai_types.AutoraterConfig] + """Autorater config used for evaluation. Not applicable for predefined metrics (PredefinedMetricSpec); the server uses its own model configuration for predefined metrics and this field is ignored.""" + metrics: Optional[list[MetricDict]] + """The metrics used for evaluation. + Currently, we only support evaluating a single metric. If multiple metrics + are provided, only the first one will be evaluated.""" -ToolParameterKeyMatchInputOrDict = Union[ - ToolParameterKeyMatchInput, ToolParameterKeyMatchInputDict + instance: Optional[EvaluationInstanceDict] + """The instance to be evaluated.""" + + metric_sources: Optional[list[MetricSourceDict]] + """The metrics used for evaluation.""" + + config: Optional[EvaluateInstancesConfigDict] + """""" + + +_EvaluateInstancesRequestParametersOrDict = Union[ + _EvaluateInstancesRequestParameters, _EvaluateInstancesRequestParametersDict ] -class ToolParameterKVMatchInstance(_common.BaseModel): - """Tool parameter kv match instance.""" +class MetricResult(_common.BaseModel): + """Result for a single metric on a single instance.""" - prediction: Optional[str] = Field( - default=None, description="""Required. Output of the evaluated model.""" + score: Optional[float] = Field( + default=None, + description="""The score for the metric. Please refer to each metric's documentation for the meaning of the score.""", ) - reference: Optional[str] = Field( + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( default=None, - description="""Required. Ground truth used to compare against the prediction.""", + description="""For rubric-based metrics, the verdicts for each rubric.""", + ) + explanation: Optional[str] = Field( + default=None, description="""The explanation for the metric result.""" + ) + error: Optional[genai_types.GoogleRpcStatus] = Field( + default=None, description="""The error status for the metric result.""" ) -class ToolParameterKVMatchInstanceDict(TypedDict, total=False): - """Tool parameter kv match instance.""" +class MetricResultDict(TypedDict, total=False): + """Result for a single metric on a single instance.""" - prediction: Optional[str] - """Required. Output of the evaluated model.""" + score: Optional[float] + """The score for the metric. Please refer to each metric's documentation for the meaning of the score.""" - reference: Optional[str] - """Required. Ground truth used to compare against the prediction.""" + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] + """For rubric-based metrics, the verdicts for each rubric.""" + explanation: Optional[str] + """The explanation for the metric result.""" -ToolParameterKVMatchInstanceOrDict = Union[ - ToolParameterKVMatchInstance, ToolParameterKVMatchInstanceDict -] + error: Optional[genai_types.GoogleRpcStatus] + """The error status for the metric result.""" -class ToolParameterKVMatchSpec(_common.BaseModel): - """Spec for tool parameter kv match metric.""" +MetricResultOrDict = Union[MetricResult, MetricResultDict] - use_strict_string_match: Optional[bool] = Field( - default=None, - description="""Optional. Whether to use STRICT string match on parameter values.""", + +class BleuResults(_common.BaseModel): + """Result of evaluating a bleu metric.""" + + bleu_metric_values: Optional[list[genai_types.BleuMetricValue]] = Field( + default=None, description="""Output only. Bleu metric values.""" ) -class ToolParameterKVMatchSpecDict(TypedDict, total=False): - """Spec for tool parameter kv match metric.""" +class BleuResultsDict(TypedDict, total=False): + """Result of evaluating a bleu metric.""" - use_strict_string_match: Optional[bool] - """Optional. Whether to use STRICT string match on parameter values.""" + bleu_metric_values: Optional[list[genai_types.BleuMetricValue]] + """Output only. Bleu metric values.""" -ToolParameterKVMatchSpecOrDict = Union[ - ToolParameterKVMatchSpec, ToolParameterKVMatchSpecDict -] +BleuResultsOrDict = Union[BleuResults, BleuResultsDict] -class ToolParameterKVMatchInput(_common.BaseModel): - """Tool parameter kv match input.""" +class ExactMatchResults(_common.BaseModel): + """Result of evaluating an exact match metric.""" - instances: Optional[list[ToolParameterKVMatchInstance]] = Field( - default=None, - description="""Required. Repeated tool parameter key value match instances.""", - ) - metric_spec: Optional[ToolParameterKVMatchSpec] = Field( - default=None, - description="""Required. Spec for tool parameter key value match metric.""", + exact_match_metric_values: Optional[list[genai_types.ExactMatchMetricValue]] = ( + Field(default=None, description="""Output only. Exact match metric values.""") ) -class ToolParameterKVMatchInputDict(TypedDict, total=False): - """Tool parameter kv match input.""" - - instances: Optional[list[ToolParameterKVMatchInstanceDict]] - """Required. Repeated tool parameter key value match instances.""" +class ExactMatchResultsDict(TypedDict, total=False): + """Result of evaluating an exact match metric.""" - metric_spec: Optional[ToolParameterKVMatchSpecDict] - """Required. Spec for tool parameter key value match metric.""" + exact_match_metric_values: Optional[list[genai_types.ExactMatchMetricValue]] + """Output only. Exact match metric values.""" -ToolParameterKVMatchInputOrDict = Union[ - ToolParameterKVMatchInput, ToolParameterKVMatchInputDict -] +ExactMatchResultsOrDict = Union[ExactMatchResults, ExactMatchResultsDict] -class MapInstance(_common.BaseModel): - """Instance data specified as a map.""" +class RougeResults(_common.BaseModel): + """Result of evaluating a rouge metric.""" - map_instance: Optional[dict[str, evals_types.InstanceData]] = Field( - default=None, description="""Map of instance data.""" + rouge_metric_values: Optional[list[genai_types.RougeMetricValue]] = Field( + default=None, description="""Output only. Rouge metric values.""" ) -class MapInstanceDict(TypedDict, total=False): - """Instance data specified as a map.""" +class RougeResultsDict(TypedDict, total=False): + """Result of evaluating a rouge metric.""" - map_instance: Optional[dict[str, evals_types.InstanceData]] - """Map of instance data.""" + rouge_metric_values: Optional[list[genai_types.RougeMetricValue]] + """Output only. Rouge metric values.""" -MapInstanceOrDict = Union[MapInstance, MapInstanceDict] +RougeResultsOrDict = Union[RougeResults, RougeResultsDict] -class EvaluationInstance(_common.BaseModel): - """A single instance to be evaluated.""" +class RubricBasedMetricResult(_common.BaseModel): + """Result for a rubric-based metric.""" - prompt: Optional[evals_types.InstanceData] = Field( - default=None, - description="""Data used to populate placeholder `prompt` in a metric prompt template.""", - ) - response: Optional[evals_types.InstanceData] = Field( - default=None, - description="""Data used to populate placeholder `response` in a metric prompt template.""", - ) - reference: Optional[evals_types.InstanceData] = Field( - default=None, - description="""Data used to populate placeholder `reference` in a metric prompt template.""", - ) - other_data: Optional[MapInstance] = Field( - default=None, - description="""Other data used to populate placeholders based on their key.""", - ) - agent_data: Optional[evals_types.AgentData] = Field( - default=None, description="""Data used for agent evaluation.""" + score: Optional[float] = Field( + default=None, description="""Passing rate of all the rubrics.""" ) - rubric_groups: Optional[dict[str, "RubricGroup"]] = Field( + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( default=None, - description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", + description="""The details of all the rubrics and their verdicts.""", ) - interactions_data_source: Optional[InteractionsDataSource] = Field( - default=None, - description="""Source for populating agent data from an Interactions API - interaction. If set, no other agent data source may be set. The backend - fetches the interaction (and the agent that produced it) and parses it - into agent data for grading.""", - ) - - -class EvaluationInstanceDict(TypedDict, total=False): - """A single instance to be evaluated.""" - - prompt: Optional[evals_types.InstanceData] - """Data used to populate placeholder `prompt` in a metric prompt template.""" - - response: Optional[evals_types.InstanceData] - """Data used to populate placeholder `response` in a metric prompt template.""" - - reference: Optional[evals_types.InstanceData] - """Data used to populate placeholder `reference` in a metric prompt template.""" - other_data: Optional[MapInstanceDict] - """Other data used to populate placeholders based on their key.""" - agent_data: Optional[evals_types.AgentData] - """Data used for agent evaluation.""" +class RubricBasedMetricResultDict(TypedDict, total=False): + """Result for a rubric-based metric.""" - rubric_groups: Optional[dict[str, "RubricGroupDict"]] - """Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""" + score: Optional[float] + """Passing rate of all the rubrics.""" - interactions_data_source: Optional[InteractionsDataSourceDict] - """Source for populating agent data from an Interactions API - interaction. If set, no other agent data source may be set. The backend - fetches the interaction (and the agent that produced it) and parses it - into agent data for grading.""" + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] + """The details of all the rubrics and their verdicts.""" -EvaluationInstanceOrDict = Union[EvaluationInstance, EvaluationInstanceDict] +RubricBasedMetricResultOrDict = Union[ + RubricBasedMetricResult, RubricBasedMetricResultDict +] -class EvaluateInstancesConfig(_common.BaseModel): - """Config for evaluate instances.""" +class CometResult(_common.BaseModel): + """Spec for Comet result - calculates the comet score for the given instance using the version specified in the spec.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + score: Optional[float] = Field( + default=None, + description="""Output only. Comet score. Range depends on version.""", ) -class EvaluateInstancesConfigDict(TypedDict, total=False): - """Config for evaluate instances.""" +class CometResultDict(TypedDict, total=False): + """Spec for Comet result - calculates the comet score for the given instance using the version specified in the spec.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + score: Optional[float] + """Output only. Comet score. Range depends on version.""" -EvaluateInstancesConfigOrDict = Union[ - EvaluateInstancesConfig, EvaluateInstancesConfigDict -] +CometResultOrDict = Union[CometResult, CometResultDict] -class RubricBasedMetricSpec(_common.BaseModel): - """Specification for a metric that is based on rubrics.""" +class MetricxResult(_common.BaseModel): + """Spec for MetricX result - calculates the MetricX score for the given instance using the version specified in the spec.""" - metric_prompt_template: Optional[str] = Field( - default=None, - description="""Template for the prompt used by the judge model to evaluate against - rubrics.""", - ) - judge_autorater_config: Optional[genai_types.AutoraterConfig] = Field( - default=None, - description="""Optional configuration for the judge LLM (Autorater).""", - ) - inline_rubrics: Optional[list[evals_types.Rubric]] = Field( - default=None, description="""Use rubrics provided directly in the spec.""" - ) - rubric_group_key: Optional[str] = Field( - default=None, - description="""Use a pre-defined group of rubrics associated with the input content. - This refers to a key in the `rubric_groups` map of - `RubricEnhancedContents`.""", - ) - rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] = Field( + score: Optional[float] = Field( default=None, - description="""Dynamically generate rubrics for evaluation using this specification.""", + description="""Output only. MetricX score. Range depends on version.""", ) -class RubricBasedMetricSpecDict(TypedDict, total=False): - """Specification for a metric that is based on rubrics.""" +class MetricxResultDict(TypedDict, total=False): + """Spec for MetricX result - calculates the MetricX score for the given instance using the version specified in the spec.""" - metric_prompt_template: Optional[str] - """Template for the prompt used by the judge model to evaluate against - rubrics.""" + score: Optional[float] + """Output only. MetricX score. Range depends on version.""" - judge_autorater_config: Optional[genai_types.AutoraterConfig] - """Optional configuration for the judge LLM (Autorater).""" - inline_rubrics: Optional[list[evals_types.Rubric]] - """Use rubrics provided directly in the spec.""" +MetricxResultOrDict = Union[MetricxResult, MetricxResultDict] - rubric_group_key: Optional[str] - """Use a pre-defined group of rubrics associated with the input content. - This refers to a key in the `rubric_groups` map of - `RubricEnhancedContents`.""" - rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] - """Dynamically generate rubrics for evaluation using this specification.""" +class ToolCallValidMetricValue(_common.BaseModel): + """Tool call valid metric value for an instance.""" + score: Optional[float] = Field( + default=None, description="""Output only. Tool call valid score.""" + ) -RubricBasedMetricSpecOrDict = Union[RubricBasedMetricSpec, RubricBasedMetricSpecDict] +class ToolCallValidMetricValueDict(TypedDict, total=False): + """Tool call valid metric value for an instance.""" -class RubricEnhancedContents(_common.BaseModel): - """Rubric-enhanced contents for evaluation.""" + score: Optional[float] + """Output only. Tool call valid score.""" - prompt: Optional[list[genai_types.Content]] = Field( - default=None, - description="""User prompt, using the standard Content type from the Gen AI SDK.""", - ) - rubric_groups: Optional[dict[str, "RubricGroup"]] = Field( - default=None, - description="""Named groups of rubrics associated with this prompt. - The key is a user-defined name for the rubric group.""", - ) - response: Optional[list[genai_types.Content]] = Field( - default=None, - description="""Response, using the standard Content type from the Gen AI SDK.""", - ) - other_content: Optional[ContentMap] = Field( - default=None, - description="""Other contents needed for the metric. - For example, if `reference` is needed for the metric, it can be provided - here.""", - ) +ToolCallValidMetricValueOrDict = Union[ + ToolCallValidMetricValue, ToolCallValidMetricValueDict +] -class RubricEnhancedContentsDict(TypedDict, total=False): - """Rubric-enhanced contents for evaluation.""" - prompt: Optional[list[genai_types.Content]] - """User prompt, using the standard Content type from the Gen AI SDK.""" +class ToolCallValidResults(_common.BaseModel): + """Results for tool call valid metric.""" - rubric_groups: Optional[dict[str, "RubricGroup"]] - """Named groups of rubrics associated with this prompt. - The key is a user-defined name for the rubric group.""" + tool_call_valid_metric_values: Optional[list[ToolCallValidMetricValue]] = Field( + default=None, description="""Output only. Tool call valid metric values.""" + ) - response: Optional[list[genai_types.Content]] - """Response, using the standard Content type from the Gen AI SDK.""" - other_content: Optional[ContentMapDict] - """Other contents needed for the metric. - For example, if `reference` is needed for the metric, it can be provided - here.""" +class ToolCallValidResultsDict(TypedDict, total=False): + """Results for tool call valid metric.""" + tool_call_valid_metric_values: Optional[list[ToolCallValidMetricValueDict]] + """Output only. Tool call valid metric values.""" -RubricEnhancedContentsOrDict = Union[RubricEnhancedContents, RubricEnhancedContentsDict] +ToolCallValidResultsOrDict = Union[ToolCallValidResults, ToolCallValidResultsDict] -class RubricBasedMetricInstance(_common.BaseModel): - """Defines an instance for Rubric-based metrics. - This class allows various input formats. - """ +class ToolNameMatchMetricValue(_common.BaseModel): + """Tool name match metric value for an instance.""" - json_instance: Optional[str] = Field( - default=None, - description="""Specify evaluation fields and their string values in JSON format.""", - ) - content_map_instance: Optional[ContentMap] = Field( - default=None, - description="""Specify evaluation fields and their content values using a ContentMap.""", - ) - rubric_enhanced_contents: Optional[RubricEnhancedContents] = Field( - default=None, - description="""Provide input as Gemini Content along with one or more - associated rubric groups.""", + score: Optional[float] = Field( + default=None, description="""Output only. Tool name match score.""" ) -class RubricBasedMetricInstanceDict(TypedDict, total=False): - """Defines an instance for Rubric-based metrics. - - This class allows various input formats. - """ - - json_instance: Optional[str] - """Specify evaluation fields and their string values in JSON format.""" - - content_map_instance: Optional[ContentMapDict] - """Specify evaluation fields and their content values using a ContentMap.""" +class ToolNameMatchMetricValueDict(TypedDict, total=False): + """Tool name match metric value for an instance.""" - rubric_enhanced_contents: Optional[RubricEnhancedContentsDict] - """Provide input as Gemini Content along with one or more - associated rubric groups.""" + score: Optional[float] + """Output only. Tool name match score.""" -RubricBasedMetricInstanceOrDict = Union[ - RubricBasedMetricInstance, RubricBasedMetricInstanceDict +ToolNameMatchMetricValueOrDict = Union[ + ToolNameMatchMetricValue, ToolNameMatchMetricValueDict ] -class RubricBasedMetricInput(_common.BaseModel): - """Input for a rubric-based metrics.""" +class ToolNameMatchResults(_common.BaseModel): + """Results for tool name match metric.""" - metric_spec: Optional[RubricBasedMetricSpec] = Field( - default=None, description="""Specification for the rubric-based metric.""" - ) - instance: Optional[RubricBasedMetricInstance] = Field( - default=None, description="""The instance to be evaluated.""" + tool_name_match_metric_values: Optional[list[ToolNameMatchMetricValue]] = Field( + default=None, description="""Output only. Tool name match metric values.""" ) -class RubricBasedMetricInputDict(TypedDict, total=False): - """Input for a rubric-based metrics.""" - - metric_spec: Optional[RubricBasedMetricSpecDict] - """Specification for the rubric-based metric.""" +class ToolNameMatchResultsDict(TypedDict, total=False): + """Results for tool name match metric.""" - instance: Optional[RubricBasedMetricInstanceDict] - """The instance to be evaluated.""" + tool_name_match_metric_values: Optional[list[ToolNameMatchMetricValueDict]] + """Output only. Tool name match metric values.""" -RubricBasedMetricInputOrDict = Union[RubricBasedMetricInput, RubricBasedMetricInputDict] +ToolNameMatchResultsOrDict = Union[ToolNameMatchResults, ToolNameMatchResultsDict] -class MetricSource(_common.BaseModel): - """The metric source used for evaluation.""" +class ToolParameterKeyMatchMetricValue(_common.BaseModel): + """Tool parameter key match metric value for an instance.""" - metric: Optional[Metric] = Field( - default=None, description="""Inline metric config.""" - ) - metric_resource_name: Optional[str] = Field( - default=None, - description="""Resource name for registered metric. Example: - projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""", + score: Optional[float] = Field( + default=None, description="""Output only. Tool parameter key match score.""" ) -class MetricSourceDict(TypedDict, total=False): - """The metric source used for evaluation.""" +class ToolParameterKeyMatchMetricValueDict(TypedDict, total=False): + """Tool parameter key match metric value for an instance.""" - metric: Optional[MetricDict] - """Inline metric config.""" + score: Optional[float] + """Output only. Tool parameter key match score.""" - metric_resource_name: Optional[str] - """Resource name for registered metric. Example: - projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" +ToolParameterKeyMatchMetricValueOrDict = Union[ + ToolParameterKeyMatchMetricValue, ToolParameterKeyMatchMetricValueDict +] -MetricSourceOrDict = Union[MetricSource, MetricSourceDict] +class ToolParameterKeyMatchResults(_common.BaseModel): + """Results for tool parameter key match metric.""" -class _EvaluateInstancesRequestParameters(_common.BaseModel): - """Parameters for evaluating instances.""" - - bleu_input: Optional[BleuInput] = Field(default=None, description="""""") - exact_match_input: Optional[ExactMatchInput] = Field( - default=None, description="""""" - ) - rouge_input: Optional[RougeInput] = Field(default=None, description="""""") - pointwise_metric_input: Optional[PointwiseMetricInput] = Field( - default=None, description="""""" - ) - pairwise_metric_input: Optional[PairwiseMetricInput] = Field( - default=None, description="""""" - ) - tool_call_valid_input: Optional[ToolCallValidInput] = Field( - default=None, description="""""" - ) - tool_name_match_input: Optional[ToolNameMatchInput] = Field( - default=None, description="""""" - ) - tool_parameter_key_match_input: Optional[ToolParameterKeyMatchInput] = Field( - default=None, description="""""" - ) - tool_parameter_kv_match_input: Optional[ToolParameterKVMatchInput] = Field( - default=None, description="""""" - ) - rubric_based_metric_input: Optional[RubricBasedMetricInput] = Field( - default=None, description="""""" - ) - autorater_config: Optional[genai_types.AutoraterConfig] = Field( - default=None, - description="""Autorater config used for evaluation. Not applicable for predefined metrics (PredefinedMetricSpec); the server uses its own model configuration for predefined metrics and this field is ignored.""", - ) - metrics: Optional[list[Metric]] = Field( + tool_parameter_key_match_metric_values: Optional[ + list[ToolParameterKeyMatchMetricValue] + ] = Field( default=None, - description="""The metrics used for evaluation. - Currently, we only support evaluating a single metric. If multiple metrics - are provided, only the first one will be evaluated.""", - ) - instance: Optional[EvaluationInstance] = Field( - default=None, description="""The instance to be evaluated.""" - ) - metric_sources: Optional[list[MetricSource]] = Field( - default=None, description="""The metrics used for evaluation.""" + description="""Output only. Tool parameter key match metric values.""", ) - config: Optional[EvaluateInstancesConfig] = Field(default=None, description="""""") -class _EvaluateInstancesRequestParametersDict(TypedDict, total=False): - """Parameters for evaluating instances.""" +class ToolParameterKeyMatchResultsDict(TypedDict, total=False): + """Results for tool parameter key match metric.""" - bleu_input: Optional[BleuInputDict] - """""" + tool_parameter_key_match_metric_values: Optional[ + list[ToolParameterKeyMatchMetricValueDict] + ] + """Output only. Tool parameter key match metric values.""" - exact_match_input: Optional[ExactMatchInputDict] - """""" - rouge_input: Optional[RougeInputDict] - """""" +ToolParameterKeyMatchResultsOrDict = Union[ + ToolParameterKeyMatchResults, ToolParameterKeyMatchResultsDict +] - pointwise_metric_input: Optional[PointwiseMetricInputDict] - """""" - pairwise_metric_input: Optional[PairwiseMetricInputDict] - """""" +class ToolParameterKVMatchMetricValue(_common.BaseModel): + """Tool parameter key value match metric value for an instance.""" - tool_call_valid_input: Optional[ToolCallValidInputDict] - """""" + score: Optional[float] = Field( + default=None, + description="""Output only. Tool parameter key value match score.""", + ) - tool_name_match_input: Optional[ToolNameMatchInputDict] - """""" - tool_parameter_key_match_input: Optional[ToolParameterKeyMatchInputDict] - """""" +class ToolParameterKVMatchMetricValueDict(TypedDict, total=False): + """Tool parameter key value match metric value for an instance.""" - tool_parameter_kv_match_input: Optional[ToolParameterKVMatchInputDict] - """""" + score: Optional[float] + """Output only. Tool parameter key value match score.""" - rubric_based_metric_input: Optional[RubricBasedMetricInputDict] - """""" - autorater_config: Optional[genai_types.AutoraterConfig] - """Autorater config used for evaluation. Not applicable for predefined metrics (PredefinedMetricSpec); the server uses its own model configuration for predefined metrics and this field is ignored.""" +ToolParameterKVMatchMetricValueOrDict = Union[ + ToolParameterKVMatchMetricValue, ToolParameterKVMatchMetricValueDict +] - metrics: Optional[list[MetricDict]] - """The metrics used for evaluation. - Currently, we only support evaluating a single metric. If multiple metrics - are provided, only the first one will be evaluated.""" - instance: Optional[EvaluationInstanceDict] - """The instance to be evaluated.""" +class ToolParameterKVMatchResults(_common.BaseModel): + """Results for tool parameter key value match metric.""" - metric_sources: Optional[list[MetricSourceDict]] - """The metrics used for evaluation.""" + tool_parameter_kv_match_metric_values: Optional[ + list[ToolParameterKVMatchMetricValue] + ] = Field( + default=None, + description="""Output only. Tool parameter key value match metric values.""", + ) - config: Optional[EvaluateInstancesConfigDict] - """""" +class ToolParameterKVMatchResultsDict(TypedDict, total=False): + """Results for tool parameter key value match metric.""" -_EvaluateInstancesRequestParametersOrDict = Union[ - _EvaluateInstancesRequestParameters, _EvaluateInstancesRequestParametersDict + tool_parameter_kv_match_metric_values: Optional[ + list[ToolParameterKVMatchMetricValueDict] + ] + """Output only. Tool parameter key value match metric values.""" + + +ToolParameterKVMatchResultsOrDict = Union[ + ToolParameterKVMatchResults, ToolParameterKVMatchResultsDict ] -class MetricResult(_common.BaseModel): - """Result for a single metric on a single instance.""" +class EvaluateInstancesResponse(_common.BaseModel): + """Result of evaluating an LLM metric.""" - score: Optional[float] = Field( + rubric_based_metric_result: Optional[RubricBasedMetricResult] = Field( + default=None, description="""Result for rubric based metric.""" + ) + metric_results: Optional[list[MetricResult]] = Field( default=None, - description="""The score for the metric. Please refer to each metric's documentation for the meaning of the score.""", + description="""A list of metric results for each evaluation case. The order of the metric results is guaranteed to be the same as the order of the instances in the request.""", ) - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( + bleu_results: Optional[BleuResults] = Field( + default=None, description="""Results for bleu metric.""" + ) + comet_result: Optional[CometResult] = Field( + default=None, description="""Translation metrics. Result for Comet metric.""" + ) + exact_match_results: Optional[ExactMatchResults] = Field( default=None, - description="""For rubric-based metrics, the verdicts for each rubric.""", + description="""Auto metric evaluation results. Results for exact match metric.""", ) - explanation: Optional[str] = Field( - default=None, description="""The explanation for the metric result.""" + metricx_result: Optional[MetricxResult] = Field( + default=None, description="""Result for Metricx metric.""" ) - error: Optional[genai_types.GoogleRpcStatus] = Field( - default=None, description="""The error status for the metric result.""" + pairwise_metric_result: Optional[genai_types.PairwiseMetricResult] = Field( + default=None, description="""Result for pairwise metric.""" + ) + pointwise_metric_result: Optional[genai_types.PointwiseMetricResult] = Field( + default=None, description="""Generic metrics. Result for pointwise metric.""" + ) + rouge_results: Optional[RougeResults] = Field( + default=None, description="""Results for rouge metric.""" + ) + tool_call_valid_results: Optional[ToolCallValidResults] = Field( + default=None, + description="""Tool call metrics. Results for tool call valid metric.""", + ) + tool_name_match_results: Optional[ToolNameMatchResults] = Field( + default=None, description="""Results for tool name match metric.""" + ) + tool_parameter_key_match_results: Optional[ToolParameterKeyMatchResults] = Field( + default=None, description="""Results for tool parameter key match metric.""" + ) + tool_parameter_kv_match_results: Optional[ToolParameterKVMatchResults] = Field( + default=None, + description="""Results for tool parameter key value match metric.""", ) -class MetricResultDict(TypedDict, total=False): - """Result for a single metric on a single instance.""" +class EvaluateInstancesResponseDict(TypedDict, total=False): + """Result of evaluating an LLM metric.""" - score: Optional[float] - """The score for the metric. Please refer to each metric's documentation for the meaning of the score.""" + rubric_based_metric_result: Optional[RubricBasedMetricResultDict] + """Result for rubric based metric.""" - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] - """For rubric-based metrics, the verdicts for each rubric.""" + metric_results: Optional[list[MetricResultDict]] + """A list of metric results for each evaluation case. The order of the metric results is guaranteed to be the same as the order of the instances in the request.""" - explanation: Optional[str] - """The explanation for the metric result.""" + bleu_results: Optional[BleuResultsDict] + """Results for bleu metric.""" - error: Optional[genai_types.GoogleRpcStatus] - """The error status for the metric result.""" + comet_result: Optional[CometResultDict] + """Translation metrics. Result for Comet metric.""" + exact_match_results: Optional[ExactMatchResultsDict] + """Auto metric evaluation results. Results for exact match metric.""" -MetricResultOrDict = Union[MetricResult, MetricResultDict] + metricx_result: Optional[MetricxResultDict] + """Result for Metricx metric.""" + pairwise_metric_result: Optional[genai_types.PairwiseMetricResult] + """Result for pairwise metric.""" -class BleuResults(_common.BaseModel): - """Result of evaluating a bleu metric.""" + pointwise_metric_result: Optional[genai_types.PointwiseMetricResult] + """Generic metrics. Result for pointwise metric.""" - bleu_metric_values: Optional[list[genai_types.BleuMetricValue]] = Field( - default=None, description="""Output only. Bleu metric values.""" - ) + rouge_results: Optional[RougeResultsDict] + """Results for rouge metric.""" + tool_call_valid_results: Optional[ToolCallValidResultsDict] + """Tool call metrics. Results for tool call valid metric.""" -class BleuResultsDict(TypedDict, total=False): - """Result of evaluating a bleu metric.""" + tool_name_match_results: Optional[ToolNameMatchResultsDict] + """Results for tool name match metric.""" - bleu_metric_values: Optional[list[genai_types.BleuMetricValue]] - """Output only. Bleu metric values.""" + tool_parameter_key_match_results: Optional[ToolParameterKeyMatchResultsDict] + """Results for tool parameter key match metric.""" + tool_parameter_kv_match_results: Optional[ToolParameterKVMatchResultsDict] + """Results for tool parameter key value match metric.""" -BleuResultsOrDict = Union[BleuResults, BleuResultsDict] +EvaluateInstancesResponseOrDict = Union[ + EvaluateInstancesResponse, EvaluateInstancesResponseDict +] -class ExactMatchResults(_common.BaseModel): - """Result of evaluating an exact match metric.""" - exact_match_metric_values: Optional[list[genai_types.ExactMatchMetricValue]] = ( - Field(default=None, description="""Output only. Exact match metric values.""") +class GenerateUserScenariosConfig(_common.BaseModel): + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class ExactMatchResultsDict(TypedDict, total=False): - """Result of evaluating an exact match metric.""" +class GenerateUserScenariosConfigDict(TypedDict, total=False): - exact_match_metric_values: Optional[list[genai_types.ExactMatchMetricValue]] - """Output only. Exact match metric values.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ExactMatchResultsOrDict = Union[ExactMatchResults, ExactMatchResultsDict] +GenerateUserScenariosConfigOrDict = Union[ + GenerateUserScenariosConfig, GenerateUserScenariosConfigDict +] -class RougeResults(_common.BaseModel): - """Result of evaluating a rouge metric.""" +class _GenerateUserScenariosParameters(_common.BaseModel): + """Parameters for GenerateUserScenarios.""" - rouge_metric_values: Optional[list[genai_types.RougeMetricValue]] = Field( - default=None, description="""Output only. Rouge metric values.""" + location: Optional[str] = Field(default=None, description="""""") + agents: Optional[dict[str, evals_types.AgentConfig]] = Field( + default=None, description="""""" + ) + root_agent_id: Optional[str] = Field(default=None, description="""""") + user_scenario_generation_config: Optional[ + evals_types.UserScenarioGenerationConfig + ] = Field(default=None, description="""""") + 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.""", + ) + 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 RougeResultsDict(TypedDict, total=False): - """Result of evaluating a rouge metric.""" - - rouge_metric_values: Optional[list[genai_types.RougeMetricValue]] - """Output only. Rouge metric values.""" - - -RougeResultsOrDict = Union[RougeResults, RougeResultsDict] +class _GenerateUserScenariosParametersDict(TypedDict, total=False): + """Parameters for GenerateUserScenarios.""" + location: Optional[str] + """""" -class RubricBasedMetricResult(_common.BaseModel): - """Result for a rubric-based metric.""" + agents: Optional[dict[str, evals_types.AgentConfig]] + """""" - score: Optional[float] = Field( - default=None, description="""Passing rate of all the rubrics.""" - ) - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( - default=None, - description="""The details of all the rubrics and their verdicts.""", - ) + root_agent_id: Optional[str] + """""" + user_scenario_generation_config: Optional[evals_types.UserScenarioGenerationConfig] + """""" -class RubricBasedMetricResultDict(TypedDict, total=False): - """Result for a rubric-based metric.""" + config: Optional[GenerateUserScenariosConfigDict] + """""" - score: Optional[float] - """Passing rate of all the rubrics.""" + allow_cross_region_model: Optional[bool] + """Opt-in flag to authorize cross-region routing for LLM models.""" - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] - """The details of all the rubrics and their verdicts.""" + gemini_agent_config: Optional[GeminiAgentConfigDict] + """If set, the server derives the agents map and root_agent_id + from the referenced Gemini Agent server-side.""" -RubricBasedMetricResultOrDict = Union[ - RubricBasedMetricResult, RubricBasedMetricResultDict +_GenerateUserScenariosParametersOrDict = Union[ + _GenerateUserScenariosParameters, _GenerateUserScenariosParametersDict ] -class CometResult(_common.BaseModel): - """Spec for Comet result - calculates the comet score for the given instance using the version specified in the spec.""" +class GenerateUserScenariosResponse(_common.BaseModel): + """Response message for DataFoundryService.GenerateUserScenarios.""" - score: Optional[float] = Field( - default=None, - description="""Output only. Comet score. Range depends on version.""", + user_scenarios: Optional[list[evals_types.UserScenario]] = Field( + default=None, description="""""" ) -class CometResultDict(TypedDict, total=False): - """Spec for Comet result - calculates the comet score for the given instance using the version specified in the spec.""" +class GenerateUserScenariosResponseDict(TypedDict, total=False): + """Response message for DataFoundryService.GenerateUserScenarios.""" - score: Optional[float] - """Output only. Comet score. Range depends on version.""" + user_scenarios: Optional[list[evals_types.UserScenario]] + """""" -CometResultOrDict = Union[CometResult, CometResultDict] +GenerateUserScenariosResponseOrDict = Union[ + GenerateUserScenariosResponse, GenerateUserScenariosResponseDict +] -class MetricxResult(_common.BaseModel): - """Spec for MetricX result - calculates the MetricX score for the given instance using the version specified in the spec.""" +class GenerateLossClustersConfig(_common.BaseModel): + """Config for generating loss clusters.""" - score: Optional[float] = Field( - default=None, - description="""Output only. MetricX score. Range depends on version.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class MetricxResultDict(TypedDict, total=False): - """Spec for MetricX result - calculates the MetricX score for the given instance using the version specified in the spec.""" +class GenerateLossClustersConfigDict(TypedDict, total=False): + """Config for generating loss clusters.""" - score: Optional[float] - """Output only. MetricX score. Range depends on version.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -MetricxResultOrDict = Union[MetricxResult, MetricxResultDict] +GenerateLossClustersConfigOrDict = Union[ + GenerateLossClustersConfig, GenerateLossClustersConfigDict +] -class ToolCallValidMetricValue(_common.BaseModel): - """Tool call valid metric value for an instance.""" +class _GenerateLossClustersParameters(_common.BaseModel): + """Parameters for GenerateLossClusters.""" - score: Optional[float] = Field( - default=None, description="""Output only. Tool call valid score.""" + location: Optional[str] = Field( + default=None, + description="""The resource name of the Location. Format: `projects/{project}/locations/{location}`.""", + ) + evaluation_set: Optional[str] = Field( + default=None, + description="""Reference to a persisted EvaluationSet. The service will read items from this set.""", + ) + inline_results: Optional[list[EvaluationResult]] = Field( + default=None, + description="""Inline evaluation results. Useful for ephemeral analysis in notebooks/SDKs where data isn't persisted.""", + ) + configs: Optional[list[LossAnalysisConfig]] = Field( + default=None, + description="""Configuration for the analysis algorithm. Analysis for multiple metrics and multiple candidates could be specified.""", + ) + config: Optional[GenerateLossClustersConfig] = Field( + default=None, description="""Config for generating loss clusters.""" ) -class ToolCallValidMetricValueDict(TypedDict, total=False): - """Tool call valid metric value for an instance.""" - - score: Optional[float] - """Output only. Tool call valid score.""" +class _GenerateLossClustersParametersDict(TypedDict, total=False): + """Parameters for GenerateLossClusters.""" + location: Optional[str] + """The resource name of the Location. Format: `projects/{project}/locations/{location}`.""" -ToolCallValidMetricValueOrDict = Union[ - ToolCallValidMetricValue, ToolCallValidMetricValueDict -] + evaluation_set: Optional[str] + """Reference to a persisted EvaluationSet. The service will read items from this set.""" + inline_results: Optional[list[EvaluationResultDict]] + """Inline evaluation results. Useful for ephemeral analysis in notebooks/SDKs where data isn't persisted.""" -class ToolCallValidResults(_common.BaseModel): - """Results for tool call valid metric.""" + configs: Optional[list[LossAnalysisConfigDict]] + """Configuration for the analysis algorithm. Analysis for multiple metrics and multiple candidates could be specified.""" - tool_call_valid_metric_values: Optional[list[ToolCallValidMetricValue]] = Field( - default=None, description="""Output only. Tool call valid metric values.""" - ) + config: Optional[GenerateLossClustersConfigDict] + """Config for generating loss clusters.""" -class ToolCallValidResultsDict(TypedDict, total=False): - """Results for tool call valid metric.""" +_GenerateLossClustersParametersOrDict = Union[ + _GenerateLossClustersParameters, _GenerateLossClustersParametersDict +] - tool_call_valid_metric_values: Optional[list[ToolCallValidMetricValueDict]] - """Output only. Tool call valid metric values.""" +class GenerateLossClustersResponse(_common.BaseModel): + """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" -ToolCallValidResultsOrDict = Union[ToolCallValidResults, ToolCallValidResultsDict] + analysis_time: Optional[str] = Field( + default=None, description="""The timestamp when this analysis was completed.""" + ) + results: Optional[list[LossAnalysisResult]] = Field( + default=None, + description="""The analysis results, one per config provided in the request.""", + ) + def show(self) -> None: + """Shows the loss pattern analysis report with rich HTML visualization.""" + from .. import _evals_visualization -class ToolNameMatchMetricValue(_common.BaseModel): - """Tool name match metric value for an instance.""" + _evals_visualization.display_loss_clusters_response(self) - score: Optional[float] = Field( - default=None, description="""Output only. Tool name match score.""" - ) +class GenerateLossClustersResponseDict(TypedDict, total=False): + """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" -class ToolNameMatchMetricValueDict(TypedDict, total=False): - """Tool name match metric value for an instance.""" + analysis_time: Optional[str] + """The timestamp when this analysis was completed.""" - score: Optional[float] - """Output only. Tool name match score.""" + results: Optional[list[LossAnalysisResultDict]] + """The analysis results, one per config provided in the request.""" -ToolNameMatchMetricValueOrDict = Union[ - ToolNameMatchMetricValue, ToolNameMatchMetricValueDict +GenerateLossClustersResponseOrDict = Union[ + GenerateLossClustersResponse, GenerateLossClustersResponseDict ] -class ToolNameMatchResults(_common.BaseModel): - """Results for tool name match metric.""" +class GenerateLossClustersOperation(_common.BaseModel): + """Long-running operation for generating loss clusters.""" - tool_name_match_metric_values: Optional[list[ToolNameMatchMetricValue]] = Field( - default=None, description="""Output only. Tool name match metric values.""" + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + response: Optional[GenerateLossClustersResponse] = Field( + default=None, + description="""Response message for EvaluationAnalyticsService.GenerateLossClusters.""", ) -class ToolNameMatchResultsDict(TypedDict, total=False): - """Results for tool name match metric.""" +class GenerateLossClustersOperationDict(TypedDict, total=False): + """Long-running operation for generating loss clusters.""" - tool_name_match_metric_values: Optional[list[ToolNameMatchMetricValueDict]] - """Output only. Tool name match metric values.""" + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" -ToolNameMatchResultsOrDict = Union[ToolNameMatchResults, ToolNameMatchResultsDict] + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -class ToolParameterKeyMatchMetricValue(_common.BaseModel): - """Tool parameter key match metric value for an instance.""" + response: Optional[GenerateLossClustersResponseDict] + """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" - score: Optional[float] = Field( - default=None, description="""Output only. Tool parameter key match score.""" - ) +GenerateLossClustersOperationOrDict = Union[ + GenerateLossClustersOperation, GenerateLossClustersOperationDict +] -class ToolParameterKeyMatchMetricValueDict(TypedDict, total=False): - """Tool parameter key match metric value for an instance.""" - score: Optional[float] - """Output only. Tool parameter key match score.""" +class RubricGenerationConfig(_common.BaseModel): + """Config for generating rubrics.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -ToolParameterKeyMatchMetricValueOrDict = Union[ - ToolParameterKeyMatchMetricValue, ToolParameterKeyMatchMetricValueDict -] +class RubricGenerationConfigDict(TypedDict, total=False): + """Config for generating rubrics.""" -class ToolParameterKeyMatchResults(_common.BaseModel): - """Results for tool parameter key match metric.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - tool_parameter_key_match_metric_values: Optional[ - list[ToolParameterKeyMatchMetricValue] - ] = Field( - default=None, - description="""Output only. Tool parameter key match metric values.""", - ) +RubricGenerationConfigOrDict = Union[RubricGenerationConfig, RubricGenerationConfigDict] -class ToolParameterKeyMatchResultsDict(TypedDict, total=False): - """Results for tool parameter key match metric.""" - tool_parameter_key_match_metric_values: Optional[ - list[ToolParameterKeyMatchMetricValueDict] - ] - """Output only. Tool parameter key match metric values.""" +class _GenerateInstanceRubricsRequest(_common.BaseModel): + """Parameters for generating rubrics.""" + contents: Optional[list[genai_types.Content]] = Field( + default=None, + description="""The prompt to generate rubrics from. For single-turn queries, this is a single instance. For multi-turn queries, this is a repeated field that contains conversation history + latest request.""", + ) + predefined_rubric_generation_spec: Optional[genai_types.PredefinedMetricSpec] = ( + Field( + default=None, + description="""Specification for using the rubric generation configs of a pre-defined + metric, e.g. "generic_quality_v1" and "instruction_following_v1". + Some of the configs may be only used in rubric generation and not + supporting evaluation, e.g. "fully_customized_generic_quality_v1". + If this field is set, the `rubric_generation_spec` field will be ignored. + """, + ) + ) + rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] = Field( + default=None, + description="""Specification for how the rubrics should be generated.""", + ) + metric_resource_name: Optional[str] = Field( + default=None, + description="""Registered metric resource name. If this field is set, the configuration provided in this field is used for rubric generation. The `predefined_rubric_generation_spec` and `rubric_generation_spec` fields will be ignored.""", + ) + config: Optional[RubricGenerationConfig] = Field(default=None, description="""""") -ToolParameterKeyMatchResultsOrDict = Union[ - ToolParameterKeyMatchResults, ToolParameterKeyMatchResultsDict -] +class _GenerateInstanceRubricsRequestDict(TypedDict, total=False): + """Parameters for generating rubrics.""" -class ToolParameterKVMatchMetricValue(_common.BaseModel): - """Tool parameter key value match metric value for an instance.""" + contents: Optional[list[genai_types.Content]] + """The prompt to generate rubrics from. For single-turn queries, this is a single instance. For multi-turn queries, this is a repeated field that contains conversation history + latest request.""" - score: Optional[float] = Field( - default=None, - description="""Output only. Tool parameter key value match score.""", - ) + predefined_rubric_generation_spec: Optional[genai_types.PredefinedMetricSpec] + """Specification for using the rubric generation configs of a pre-defined + metric, e.g. "generic_quality_v1" and "instruction_following_v1". + Some of the configs may be only used in rubric generation and not + supporting evaluation, e.g. "fully_customized_generic_quality_v1". + If this field is set, the `rubric_generation_spec` field will be ignored. + """ + rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] + """Specification for how the rubrics should be generated.""" -class ToolParameterKVMatchMetricValueDict(TypedDict, total=False): - """Tool parameter key value match metric value for an instance.""" + metric_resource_name: Optional[str] + """Registered metric resource name. If this field is set, the configuration provided in this field is used for rubric generation. The `predefined_rubric_generation_spec` and `rubric_generation_spec` fields will be ignored.""" - score: Optional[float] - """Output only. Tool parameter key value match score.""" + config: Optional[RubricGenerationConfigDict] + """""" -ToolParameterKVMatchMetricValueOrDict = Union[ - ToolParameterKVMatchMetricValue, ToolParameterKVMatchMetricValueDict +_GenerateInstanceRubricsRequestOrDict = Union[ + _GenerateInstanceRubricsRequest, _GenerateInstanceRubricsRequestDict ] -class ToolParameterKVMatchResults(_common.BaseModel): - """Results for tool parameter key value match metric.""" +class GenerateInstanceRubricsResponse(_common.BaseModel): + """Response for generating rubrics.""" - tool_parameter_kv_match_metric_values: Optional[ - list[ToolParameterKVMatchMetricValue] - ] = Field( - default=None, - description="""Output only. Tool parameter key value match metric values.""", + generated_rubrics: Optional[list[evals_types.Rubric]] = Field( + default=None, description="""A list of generated rubrics.""" ) -class ToolParameterKVMatchResultsDict(TypedDict, total=False): - """Results for tool parameter key value match metric.""" +class GenerateInstanceRubricsResponseDict(TypedDict, total=False): + """Response for generating rubrics.""" - tool_parameter_kv_match_metric_values: Optional[ - list[ToolParameterKVMatchMetricValueDict] - ] - """Output only. Tool parameter key value match metric values.""" + generated_rubrics: Optional[list[evals_types.Rubric]] + """A list of generated rubrics.""" -ToolParameterKVMatchResultsOrDict = Union[ - ToolParameterKVMatchResults, ToolParameterKVMatchResultsDict +GenerateInstanceRubricsResponseOrDict = Union[ + GenerateInstanceRubricsResponse, GenerateInstanceRubricsResponseDict ] -class EvaluateInstancesResponse(_common.BaseModel): - """Result of evaluating an LLM metric.""" +class GetEvaluationExperimentConfig(_common.BaseModel): + """Config for getting an evaluation experiment.""" - rubric_based_metric_result: Optional[RubricBasedMetricResult] = Field( - default=None, description="""Result for rubric based metric.""" - ) - metric_results: Optional[list[MetricResult]] = Field( - default=None, - description="""A list of metric results for each evaluation case. The order of the metric results is guaranteed to be the same as the order of the instances in the request.""", - ) - bleu_results: Optional[BleuResults] = Field( - default=None, description="""Results for bleu metric.""" - ) - comet_result: Optional[CometResult] = Field( - default=None, description="""Translation metrics. Result for Comet metric.""" - ) - exact_match_results: Optional[ExactMatchResults] = Field( - default=None, - description="""Auto metric evaluation results. Results for exact match metric.""", - ) - metricx_result: Optional[MetricxResult] = Field( - default=None, description="""Result for Metricx metric.""" - ) - pairwise_metric_result: Optional[genai_types.PairwiseMetricResult] = Field( - default=None, description="""Result for pairwise metric.""" - ) - pointwise_metric_result: Optional[genai_types.PointwiseMetricResult] = Field( - default=None, description="""Generic metrics. Result for pointwise metric.""" - ) - rouge_results: Optional[RougeResults] = Field( - default=None, description="""Results for rouge metric.""" - ) - tool_call_valid_results: Optional[ToolCallValidResults] = Field( - default=None, - description="""Tool call metrics. Results for tool call valid metric.""", - ) - tool_name_match_results: Optional[ToolNameMatchResults] = Field( - default=None, description="""Results for tool name match metric.""" - ) - tool_parameter_key_match_results: Optional[ToolParameterKeyMatchResults] = Field( - default=None, description="""Results for tool parameter key match metric.""" - ) - tool_parameter_kv_match_results: Optional[ToolParameterKVMatchResults] = Field( - default=None, - description="""Results for tool parameter key value match metric.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class EvaluateInstancesResponseDict(TypedDict, total=False): - """Result of evaluating an LLM metric.""" - - rubric_based_metric_result: Optional[RubricBasedMetricResultDict] - """Result for rubric based metric.""" - - metric_results: Optional[list[MetricResultDict]] - """A list of metric results for each evaluation case. The order of the metric results is guaranteed to be the same as the order of the instances in the request.""" - - bleu_results: Optional[BleuResultsDict] - """Results for bleu metric.""" +class GetEvaluationExperimentConfigDict(TypedDict, total=False): + """Config for getting an evaluation experiment.""" - comet_result: Optional[CometResultDict] - """Translation metrics. Result for Comet metric.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - exact_match_results: Optional[ExactMatchResultsDict] - """Auto metric evaluation results. Results for exact match metric.""" - metricx_result: Optional[MetricxResultDict] - """Result for Metricx metric.""" +GetEvaluationExperimentConfigOrDict = Union[ + GetEvaluationExperimentConfig, GetEvaluationExperimentConfigDict +] - pairwise_metric_result: Optional[genai_types.PairwiseMetricResult] - """Result for pairwise metric.""" - pointwise_metric_result: Optional[genai_types.PointwiseMetricResult] - """Generic metrics. Result for pointwise metric.""" +class _GetEvaluationExperimentParameters(_common.BaseModel): + """Parameters for getting an evaluation experiment.""" - rouge_results: Optional[RougeResultsDict] - """Results for rouge metric.""" + name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationExperimentConfig] = Field( + default=None, description="""""" + ) - tool_call_valid_results: Optional[ToolCallValidResultsDict] - """Tool call metrics. Results for tool call valid metric.""" - tool_name_match_results: Optional[ToolNameMatchResultsDict] - """Results for tool name match metric.""" +class _GetEvaluationExperimentParametersDict(TypedDict, total=False): + """Parameters for getting an evaluation experiment.""" - tool_parameter_key_match_results: Optional[ToolParameterKeyMatchResultsDict] - """Results for tool parameter key match metric.""" + name: Optional[str] + """""" - tool_parameter_kv_match_results: Optional[ToolParameterKVMatchResultsDict] - """Results for tool parameter key value match metric.""" + config: Optional[GetEvaluationExperimentConfigDict] + """""" -EvaluateInstancesResponseOrDict = Union[ - EvaluateInstancesResponse, EvaluateInstancesResponseDict +_GetEvaluationExperimentParametersOrDict = Union[ + _GetEvaluationExperimentParameters, _GetEvaluationExperimentParametersDict ] -class GenerateUserScenariosConfig(_common.BaseModel): +class GetEvaluationMetricConfig(_common.BaseModel): + """Config for getting an evaluation metric.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GenerateUserScenariosConfigDict(TypedDict, total=False): +class GetEvaluationMetricConfigDict(TypedDict, total=False): + """Config for getting an evaluation metric.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GenerateUserScenariosConfigOrDict = Union[ - GenerateUserScenariosConfig, GenerateUserScenariosConfigDict +GetEvaluationMetricConfigOrDict = Union[ + GetEvaluationMetricConfig, GetEvaluationMetricConfigDict ] -class _GenerateUserScenariosParameters(_common.BaseModel): - """Parameters for GenerateUserScenarios.""" +class _GetEvaluationMetricParameters(_common.BaseModel): + """Parameters for getting an evaluation metric.""" - location: Optional[str] = Field(default=None, description="""""") - agents: Optional[dict[str, evals_types.AgentConfig]] = Field( - default=None, description="""""" - ) - root_agent_id: Optional[str] = Field(default=None, description="""""") - user_scenario_generation_config: Optional[ - evals_types.UserScenarioGenerationConfig - ] = Field(default=None, description="""""") - config: Optional[GenerateUserScenariosConfig] = Field( + metric_resource_name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationMetricConfig] = 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.""", - ) - 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): - """Parameters for GenerateUserScenarios.""" - - location: Optional[str] - """""" +class _GetEvaluationMetricParametersDict(TypedDict, total=False): + """Parameters for getting an evaluation metric.""" - agents: Optional[dict[str, evals_types.AgentConfig]] + metric_resource_name: Optional[str] """""" - root_agent_id: Optional[str] + config: Optional[GetEvaluationMetricConfigDict] """""" - user_scenario_generation_config: Optional[evals_types.UserScenarioGenerationConfig] - """""" - config: Optional[GenerateUserScenariosConfigDict] - """""" +_GetEvaluationMetricParametersOrDict = Union[ + _GetEvaluationMetricParameters, _GetEvaluationMetricParametersDict +] - 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.""" +class GetEvaluationRunConfig(_common.BaseModel): + """Config for get evaluation run.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -_GenerateUserScenariosParametersOrDict = Union[ - _GenerateUserScenariosParameters, _GenerateUserScenariosParametersDict -] +class GetEvaluationRunConfigDict(TypedDict, total=False): + """Config for get evaluation run.""" -class GenerateUserScenariosResponse(_common.BaseModel): - """Response message for DataFoundryService.GenerateUserScenarios.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - user_scenarios: Optional[list[evals_types.UserScenario]] = Field( - default=None, description="""""" - ) +GetEvaluationRunConfigOrDict = Union[GetEvaluationRunConfig, GetEvaluationRunConfigDict] -class GenerateUserScenariosResponseDict(TypedDict, total=False): - """Response message for DataFoundryService.GenerateUserScenarios.""" - user_scenarios: Optional[list[evals_types.UserScenario]] +class _GetEvaluationRunParameters(_common.BaseModel): + """Represents a job that runs evaluation.""" + + name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationRunConfig] = Field(default=None, description="""""") + + +class _GetEvaluationRunParametersDict(TypedDict, total=False): + """Represents a job that runs evaluation.""" + + name: Optional[str] + """""" + + config: Optional[GetEvaluationRunConfigDict] """""" -GenerateUserScenariosResponseOrDict = Union[ - GenerateUserScenariosResponse, GenerateUserScenariosResponseDict +_GetEvaluationRunParametersOrDict = Union[ + _GetEvaluationRunParameters, _GetEvaluationRunParametersDict ] -class GenerateLossClustersConfig(_common.BaseModel): - """Config for generating loss clusters.""" +class GetEvaluationSetConfig(_common.BaseModel): + """Config for get evaluation set.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GenerateLossClustersConfigDict(TypedDict, total=False): - """Config for generating loss clusters.""" +class GetEvaluationSetConfigDict(TypedDict, total=False): + """Config for get evaluation set.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GenerateLossClustersConfigOrDict = Union[ - GenerateLossClustersConfig, GenerateLossClustersConfigDict -] +GetEvaluationSetConfigOrDict = Union[GetEvaluationSetConfig, GetEvaluationSetConfigDict] -class _GenerateLossClustersParameters(_common.BaseModel): - """Parameters for GenerateLossClusters.""" +class _GetEvaluationSetParameters(_common.BaseModel): + """Represents a job that gets an evaluation set.""" - location: Optional[str] = Field( - default=None, - description="""The resource name of the Location. Format: `projects/{project}/locations/{location}`.""", - ) - evaluation_set: Optional[str] = Field( - default=None, - description="""Reference to a persisted EvaluationSet. The service will read items from this set.""", - ) - inline_results: Optional[list[EvaluationResult]] = Field( - default=None, - description="""Inline evaluation results. Useful for ephemeral analysis in notebooks/SDKs where data isn't persisted.""", - ) - configs: Optional[list[LossAnalysisConfig]] = Field( - default=None, - description="""Configuration for the analysis algorithm. Analysis for multiple metrics and multiple candidates could be specified.""", - ) - config: Optional[GenerateLossClustersConfig] = Field( - default=None, description="""Config for generating loss clusters.""" - ) + name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationSetConfig] = Field(default=None, description="""""") -class _GenerateLossClustersParametersDict(TypedDict, total=False): - """Parameters for GenerateLossClusters.""" +class _GetEvaluationSetParametersDict(TypedDict, total=False): + """Represents a job that gets an evaluation set.""" - location: Optional[str] - """The resource name of the Location. Format: `projects/{project}/locations/{location}`.""" + name: Optional[str] + """""" - evaluation_set: Optional[str] - """Reference to a persisted EvaluationSet. The service will read items from this set.""" + config: Optional[GetEvaluationSetConfigDict] + """""" - inline_results: Optional[list[EvaluationResultDict]] - """Inline evaluation results. Useful for ephemeral analysis in notebooks/SDKs where data isn't persisted.""" - configs: Optional[list[LossAnalysisConfigDict]] - """Configuration for the analysis algorithm. Analysis for multiple metrics and multiple candidates could be specified.""" +_GetEvaluationSetParametersOrDict = Union[ + _GetEvaluationSetParameters, _GetEvaluationSetParametersDict +] - config: Optional[GenerateLossClustersConfigDict] - """Config for generating loss clusters.""" +class GetEvaluationItemConfig(_common.BaseModel): + """Config for get evaluation item.""" -_GenerateLossClustersParametersOrDict = Union[ - _GenerateLossClustersParameters, _GenerateLossClustersParametersDict -] + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -class GenerateLossClustersResponse(_common.BaseModel): - """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" +class GetEvaluationItemConfigDict(TypedDict, total=False): + """Config for get evaluation item.""" - analysis_time: Optional[str] = Field( - default=None, description="""The timestamp when this analysis was completed.""" - ) - results: Optional[list[LossAnalysisResult]] = Field( - default=None, - description="""The analysis results, one per config provided in the request.""", - ) + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - def show(self) -> None: - """Shows the loss pattern analysis report with rich HTML visualization.""" - from .. import _evals_visualization - _evals_visualization.display_loss_clusters_response(self) +GetEvaluationItemConfigOrDict = Union[ + GetEvaluationItemConfig, GetEvaluationItemConfigDict +] -class GenerateLossClustersResponseDict(TypedDict, total=False): - """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" +class _GetEvaluationItemParameters(_common.BaseModel): + """Represents a job that gets an evaluation item.""" - analysis_time: Optional[str] - """The timestamp when this analysis was completed.""" + name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationItemConfig] = Field(default=None, description="""""") - results: Optional[list[LossAnalysisResultDict]] - """The analysis results, one per config provided in the request.""" +class _GetEvaluationItemParametersDict(TypedDict, total=False): + """Represents a job that gets an evaluation item.""" -GenerateLossClustersResponseOrDict = Union[ - GenerateLossClustersResponse, GenerateLossClustersResponseDict + name: Optional[str] + """""" + + config: Optional[GetEvaluationItemConfigDict] + """""" + + +_GetEvaluationItemParametersOrDict = Union[ + _GetEvaluationItemParameters, _GetEvaluationItemParametersDict ] -class GenerateLossClustersOperation(_common.BaseModel): - """Long-running operation for generating loss clusters.""" +class ListEvaluationExperimentsConfig(_common.BaseModel): + """Config for listing evaluation experiments.""" - name: Optional[str] = Field( - default=None, - description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, - description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", - ) - done: Optional[bool] = Field( - default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - error: Optional[dict[str, Any]] = Field( + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported. + For more information about filter syntax, see + `AIP-160 `_.""", ) - response: Optional[GenerateLossClustersResponse] = Field( + order_by: Optional[str] = Field( default=None, - description="""Response message for EvaluationAnalyticsService.GenerateLossClusters.""", + description="""A comma-separated list of fields to order by, sorted in ascending + order by default. Use ``desc`` after a field name for descending. + Example: ``"create_time desc"``.""", ) -class GenerateLossClustersOperationDict(TypedDict, total=False): - """Long-running operation for generating loss clusters.""" +class ListEvaluationExperimentsConfigDict(TypedDict, total=False): + """Config for listing evaluation experiments.""" - name: Optional[str] - """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - metadata: Optional[dict[str, Any]] - """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + page_size: Optional[int] + """""" - done: Optional[bool] - """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + page_token: Optional[str] + """""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported. + For more information about filter syntax, see + `AIP-160 `_.""" - response: Optional[GenerateLossClustersResponseDict] - """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" + order_by: Optional[str] + """A comma-separated list of fields to order by, sorted in ascending + order by default. Use ``desc`` after a field name for descending. + Example: ``"create_time desc"``.""" -GenerateLossClustersOperationOrDict = Union[ - GenerateLossClustersOperation, GenerateLossClustersOperationDict +ListEvaluationExperimentsConfigOrDict = Union[ + ListEvaluationExperimentsConfig, ListEvaluationExperimentsConfigDict ] -class RubricGenerationConfig(_common.BaseModel): - """Config for generating rubrics.""" +class _ListEvaluationExperimentsParameters(_common.BaseModel): + """Parameters for listing evaluation experiments.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + config: Optional[ListEvaluationExperimentsConfig] = Field( + default=None, description="""""" ) -class RubricGenerationConfigDict(TypedDict, total=False): - """Config for generating rubrics.""" +class _ListEvaluationExperimentsParametersDict(TypedDict, total=False): + """Parameters for listing evaluation experiments.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + config: Optional[ListEvaluationExperimentsConfigDict] + """""" -RubricGenerationConfigOrDict = Union[RubricGenerationConfig, RubricGenerationConfigDict] +_ListEvaluationExperimentsParametersOrDict = Union[ + _ListEvaluationExperimentsParameters, _ListEvaluationExperimentsParametersDict +] -class _GenerateInstanceRubricsRequest(_common.BaseModel): - """Parameters for generating rubrics.""" +class ListEvaluationExperimentsResponse(_common.BaseModel): + """Response for listing evaluation experiments.""" - contents: Optional[list[genai_types.Content]] = Field( + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" + ) + next_page_token: Optional[str] = Field(default=None, description="""""") + evaluation_experiments: Optional[list[EvaluationExperiment]] = Field( default=None, - description="""The prompt to generate rubrics from. For single-turn queries, this is a single instance. For multi-turn queries, this is a repeated field that contains conversation history + latest request.""", + description="""List of evaluation experiments. + """, ) - predefined_rubric_generation_spec: Optional[genai_types.PredefinedMetricSpec] = ( - Field( - default=None, - description="""Specification for using the rubric generation configs of a pre-defined - metric, e.g. "generic_quality_v1" and "instruction_following_v1". - Some of the configs may be only used in rubric generation and not - supporting evaluation, e.g. "fully_customized_generic_quality_v1". - If this field is set, the `rubric_generation_spec` field will be ignored. - """, - ) - ) - rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] = Field( - default=None, - description="""Specification for how the rubrics should be generated.""", - ) - metric_resource_name: Optional[str] = Field( - default=None, - description="""Registered metric resource name. If this field is set, the configuration provided in this field is used for rubric generation. The `predefined_rubric_generation_spec` and `rubric_generation_spec` fields will be ignored.""", - ) - config: Optional[RubricGenerationConfig] = Field(default=None, description="""""") - - -class _GenerateInstanceRubricsRequestDict(TypedDict, total=False): - """Parameters for generating rubrics.""" - - contents: Optional[list[genai_types.Content]] - """The prompt to generate rubrics from. For single-turn queries, this is a single instance. For multi-turn queries, this is a repeated field that contains conversation history + latest request.""" - predefined_rubric_generation_spec: Optional[genai_types.PredefinedMetricSpec] - """Specification for using the rubric generation configs of a pre-defined - metric, e.g. "generic_quality_v1" and "instruction_following_v1". - Some of the configs may be only used in rubric generation and not - supporting evaluation, e.g. "fully_customized_generic_quality_v1". - If this field is set, the `rubric_generation_spec` field will be ignored. - """ - rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] - """Specification for how the rubrics should be generated.""" +class ListEvaluationExperimentsResponseDict(TypedDict, total=False): + """Response for listing evaluation experiments.""" - metric_resource_name: Optional[str] - """Registered metric resource name. If this field is set, the configuration provided in this field is used for rubric generation. The `predefined_rubric_generation_spec` and `rubric_generation_spec` fields will be ignored.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - config: Optional[RubricGenerationConfigDict] + next_page_token: Optional[str] """""" + evaluation_experiments: Optional[list[EvaluationExperimentDict]] + """List of evaluation experiments. + """ + -_GenerateInstanceRubricsRequestOrDict = Union[ - _GenerateInstanceRubricsRequest, _GenerateInstanceRubricsRequestDict +ListEvaluationExperimentsResponseOrDict = Union[ + ListEvaluationExperimentsResponse, ListEvaluationExperimentsResponseDict ] -class GenerateInstanceRubricsResponse(_common.BaseModel): - """Response for generating rubrics.""" +class ListEvaluationMetricsConfig(_common.BaseModel): + """Config for listing evaluation metrics.""" - generated_rubrics: Optional[list[evals_types.Rubric]] = Field( - default=None, description="""A list of generated rubrics.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( + default=None, + description="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported. + For more information about filter syntax, see + `AIP-160 `_.""", + ) + order_by: Optional[str] = Field( + default=None, + description="""A comma-separated list of fields to order by, sorted in ascending + order by default. Use ``desc`` after a field name for descending. + Example: ``"create_time desc"``.""", ) -class GenerateInstanceRubricsResponseDict(TypedDict, total=False): - """Response for generating rubrics.""" +class ListEvaluationMetricsConfigDict(TypedDict, total=False): + """Config for listing evaluation metrics.""" - generated_rubrics: Optional[list[evals_types.Rubric]] - """A list of generated rubrics.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + page_size: Optional[int] + """""" -GenerateInstanceRubricsResponseOrDict = Union[ - GenerateInstanceRubricsResponse, GenerateInstanceRubricsResponseDict + page_token: Optional[str] + """""" + + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported. + For more information about filter syntax, see + `AIP-160 `_.""" + + order_by: Optional[str] + """A comma-separated list of fields to order by, sorted in ascending + order by default. Use ``desc`` after a field name for descending. + Example: ``"create_time desc"``.""" + + +ListEvaluationMetricsConfigOrDict = Union[ + ListEvaluationMetricsConfig, ListEvaluationMetricsConfigDict ] -class GetEvaluationExperimentConfig(_common.BaseModel): - """Config for getting an evaluation experiment.""" +class _ListEvaluationMetricsParameters(_common.BaseModel): + """Parameters for listing evaluation metrics.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + config: Optional[ListEvaluationMetricsConfig] = Field( + default=None, description="""""" ) -class GetEvaluationExperimentConfigDict(TypedDict, total=False): - """Config for getting an evaluation experiment.""" +class _ListEvaluationMetricsParametersDict(TypedDict, total=False): + """Parameters for listing evaluation metrics.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + config: Optional[ListEvaluationMetricsConfigDict] + """""" -GetEvaluationExperimentConfigOrDict = Union[ - GetEvaluationExperimentConfig, GetEvaluationExperimentConfigDict +_ListEvaluationMetricsParametersOrDict = Union[ + _ListEvaluationMetricsParameters, _ListEvaluationMetricsParametersDict ] -class _GetEvaluationExperimentParameters(_common.BaseModel): - """Parameters for getting an evaluation experiment.""" +class ListEvaluationMetricsResponse(_common.BaseModel): + """Response for listing evaluation metrics.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationExperimentConfig] = Field( - default=None, description="""""" + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" + ) + next_page_token: Optional[str] = Field(default=None, description="""""") + evaluation_metrics: Optional[list[EvaluationMetric]] = Field( + default=None, + description="""List of evaluation metrics. + """, ) -class _GetEvaluationExperimentParametersDict(TypedDict, total=False): - """Parameters for getting an evaluation experiment.""" +class ListEvaluationMetricsResponseDict(TypedDict, total=False): + """Response for listing evaluation metrics.""" - name: Optional[str] - """""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - config: Optional[GetEvaluationExperimentConfigDict] + next_page_token: Optional[str] """""" + evaluation_metrics: Optional[list[EvaluationMetricDict]] + """List of evaluation metrics. + """ + -_GetEvaluationExperimentParametersOrDict = Union[ - _GetEvaluationExperimentParameters, _GetEvaluationExperimentParametersDict +ListEvaluationMetricsResponseOrDict = Union[ + ListEvaluationMetricsResponse, ListEvaluationMetricsResponseDict ] -class GetEvaluationMetricConfig(_common.BaseModel): - """Config for getting an evaluation metric.""" +class UpdateEvaluationExperimentConfig(_common.BaseModel): + """Config for updating an evaluation experiment.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) + update_mask: Optional[str] = Field( + default=None, + description="""The update mask to apply. For the `FieldMask` definition, see + https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""", + ) + display_name: Optional[str] = Field( + default=None, description="""The display name of the evaluation experiment.""" + ) + labels: Optional[dict[str, str]] = Field( + default=None, description="""Labels for the evaluation experiment.""" + ) + merge_strategy: Optional[EvaluationExperimentMergeStrategy] = Field( + default=None, description="""Merge strategy for the evaluation experiment.""" + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, description="""Metadata about the evaluation experiment.""" + ) -class GetEvaluationMetricConfigDict(TypedDict, total=False): - """Config for getting an evaluation metric.""" +class UpdateEvaluationExperimentConfigDict(TypedDict, total=False): + """Config for updating an evaluation experiment.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" + update_mask: Optional[str] + """The update mask to apply. For the `FieldMask` definition, see + https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" + + display_name: Optional[str] + """The display name of the evaluation experiment.""" -GetEvaluationMetricConfigOrDict = Union[ - GetEvaluationMetricConfig, GetEvaluationMetricConfigDict + labels: Optional[dict[str, str]] + """Labels for the evaluation experiment.""" + + merge_strategy: Optional[EvaluationExperimentMergeStrategy] + """Merge strategy for the evaluation experiment.""" + + metadata: Optional[dict[str, Any]] + """Metadata about the evaluation experiment.""" + + +UpdateEvaluationExperimentConfigOrDict = Union[ + UpdateEvaluationExperimentConfig, UpdateEvaluationExperimentConfigDict ] -class _GetEvaluationMetricParameters(_common.BaseModel): - """Parameters for getting an evaluation metric.""" +class _UpdateEvaluationExperimentParameters(_common.BaseModel): + """Parameters for updating an evaluation experiment.""" - metric_resource_name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationMetricConfig] = Field( + name: Optional[str] = Field( + default=None, description="""The resource name of the EvaluationExperiment.""" + ) + config: Optional[UpdateEvaluationExperimentConfig] = Field( default=None, description="""""" ) -class _GetEvaluationMetricParametersDict(TypedDict, total=False): - """Parameters for getting an evaluation metric.""" +class _UpdateEvaluationExperimentParametersDict(TypedDict, total=False): + """Parameters for updating an evaluation experiment.""" - metric_resource_name: Optional[str] - """""" + name: Optional[str] + """The resource name of the EvaluationExperiment.""" - config: Optional[GetEvaluationMetricConfigDict] + config: Optional[UpdateEvaluationExperimentConfigDict] """""" -_GetEvaluationMetricParametersOrDict = Union[ - _GetEvaluationMetricParameters, _GetEvaluationMetricParametersDict +_UpdateEvaluationExperimentParametersOrDict = Union[ + _UpdateEvaluationExperimentParameters, _UpdateEvaluationExperimentParametersDict ] -class GetEvaluationRunConfig(_common.BaseModel): - """Config for get evaluation run.""" +class OptimizeConfig(_common.BaseModel): + """Config for Prompt Optimizer.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) + optimization_target: Optional[OptimizeTarget] = Field( + default=None, + description="""The optimization target for the prompt optimizer. It must be one of the OptimizeTarget enum values: OPTIMIZATION_TARGET_GEMINI_NANO for the prompts from Android core API, OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS for the few-shot prompt optimizer with rubrics, OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE for the few-shot prompt optimizer with target responses.""", + ) + examples_dataframe: Optional[PandasDataFrame] = Field( + default=None, + description="""The examples dataframe for the few-shot prompt optimizer. It must contain "prompt" and "model_response" columns. Depending on which optimization target is used, it also needs to contain "rubrics" and "rubrics_evaluations" or "target_response" columns.""", + ) -class GetEvaluationRunConfigDict(TypedDict, total=False): - """Config for get evaluation run.""" +class OptimizeConfigDict(TypedDict, total=False): + """Config for Prompt Optimizer.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" + optimization_target: Optional[OptimizeTarget] + """The optimization target for the prompt optimizer. It must be one of the OptimizeTarget enum values: OPTIMIZATION_TARGET_GEMINI_NANO for the prompts from Android core API, OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS for the few-shot prompt optimizer with rubrics, OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE for the few-shot prompt optimizer with target responses.""" -GetEvaluationRunConfigOrDict = Union[GetEvaluationRunConfig, GetEvaluationRunConfigDict] - + examples_dataframe: Optional[PandasDataFrame] + """The examples dataframe for the few-shot prompt optimizer. It must contain "prompt" and "model_response" columns. Depending on which optimization target is used, it also needs to contain "rubrics" and "rubrics_evaluations" or "target_response" columns.""" -class _GetEvaluationRunParameters(_common.BaseModel): - """Represents a job that runs evaluation.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationRunConfig] = Field(default=None, description="""""") +OptimizeConfigOrDict = Union[OptimizeConfig, OptimizeConfigDict] -class _GetEvaluationRunParametersDict(TypedDict, total=False): - """Represents a job that runs evaluation.""" +class _OptimizeRequestParameters(_common.BaseModel): + """Request for the optimize_prompt method.""" - name: Optional[str] - """""" + content: Optional[genai_types.Content] = Field(default=None, description="""""") + config: Optional[OptimizeConfig] = Field(default=None, description="""""") - config: Optional[GetEvaluationRunConfigDict] - """""" +class _OptimizeRequestParametersDict(TypedDict, total=False): + """Request for the optimize_prompt method.""" -_GetEvaluationRunParametersOrDict = Union[ - _GetEvaluationRunParameters, _GetEvaluationRunParametersDict -] - - -class GetEvaluationSetConfig(_common.BaseModel): - """Config for get evaluation set.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - - -class GetEvaluationSetConfigDict(TypedDict, total=False): - """Config for get evaluation set.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + content: Optional[genai_types.Content] + """""" + config: Optional[OptimizeConfigDict] + """""" -GetEvaluationSetConfigOrDict = Union[GetEvaluationSetConfig, GetEvaluationSetConfigDict] +_OptimizeRequestParametersOrDict = Union[ + _OptimizeRequestParameters, _OptimizeRequestParametersDict +] -class _GetEvaluationSetParameters(_common.BaseModel): - """Represents a job that gets an evaluation set.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationSetConfig] = Field(default=None, description="""""") +class OptimizeResponseEndpoint(_common.BaseModel): + """Response for the optimize_prompt method.""" + content: Optional[genai_types.Content] = Field(default=None, description="""""") -class _GetEvaluationSetParametersDict(TypedDict, total=False): - """Represents a job that gets an evaluation set.""" - name: Optional[str] - """""" +class OptimizeResponseEndpointDict(TypedDict, total=False): + """Response for the optimize_prompt method.""" - config: Optional[GetEvaluationSetConfigDict] + content: Optional[genai_types.Content] """""" -_GetEvaluationSetParametersOrDict = Union[ - _GetEvaluationSetParameters, _GetEvaluationSetParametersDict +OptimizeResponseEndpointOrDict = Union[ + OptimizeResponseEndpoint, OptimizeResponseEndpointDict ] -class GetEvaluationItemConfig(_common.BaseModel): - """Config for get evaluation item.""" +class DnsPeeringConfig(_common.BaseModel): + """DNS peering configuration. These configurations are used to create DNS peering zones in the Vertex tenant project VPC, enabling resolution of records within the specified domain hosted in the target network's Cloud DNS.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + domain: Optional[str] = Field( + default=None, + description="""Required. The DNS name suffix of the zone being peered to, e.g., "my-internal-domain.corp.". Must end with a dot.""", + ) + target_network: Optional[str] = Field( + default=None, + description="""Required. The VPC network name in the target_project where the DNS zone specified by 'domain' is visible.""", + ) + target_project: Optional[str] = Field( + default=None, + description="""Required. The project ID hosting the Cloud DNS managed zone that contains the 'domain'. The Vertex AI Service Agent requires the dns.peer role on this project.""", ) -class GetEvaluationItemConfigDict(TypedDict, total=False): - """Config for get evaluation item.""" +class DnsPeeringConfigDict(TypedDict, total=False): + """DNS peering configuration. These configurations are used to create DNS peering zones in the Vertex tenant project VPC, enabling resolution of records within the specified domain hosted in the target network's Cloud DNS.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + domain: Optional[str] + """Required. The DNS name suffix of the zone being peered to, e.g., "my-internal-domain.corp.". Must end with a dot.""" + target_network: Optional[str] + """Required. The VPC network name in the target_project where the DNS zone specified by 'domain' is visible.""" -GetEvaluationItemConfigOrDict = Union[ - GetEvaluationItemConfig, GetEvaluationItemConfigDict -] + target_project: Optional[str] + """Required. The project ID hosting the Cloud DNS managed zone that contains the 'domain'. The Vertex AI Service Agent requires the dns.peer role on this project.""" -class _GetEvaluationItemParameters(_common.BaseModel): - """Represents a job that gets an evaluation item.""" +DnsPeeringConfigOrDict = Union[DnsPeeringConfig, DnsPeeringConfigDict] - name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationItemConfig] = Field(default=None, description="""""") +class PscInterfaceConfig(_common.BaseModel): + """Configuration for PSC-I.""" -class _GetEvaluationItemParametersDict(TypedDict, total=False): - """Represents a job that gets an evaluation item.""" + dns_peering_configs: Optional[list[DnsPeeringConfig]] = Field( + default=None, + description="""Optional. DNS peering configurations. When specified, Vertex AI will attempt to configure DNS peering zones in the tenant project VPC to resolve the specified domains using the target network's Cloud DNS. The user must grant the dns.peer role to the Vertex AI Service Agent on the target project.""", + ) + network_attachment: Optional[str] = Field( + default=None, + description="""Optional. The name of the Compute Engine [network attachment](https://cloud.google.com/vpc/docs/about-network-attachments) to attach to the resource within the region and user project. To specify this field, you must have already [created a network attachment] (https://cloud.google.com/vpc/docs/create-manage-network-attachments#create-network-attachments). This field is only used for resources using PSC-I.""", + ) - name: Optional[str] - """""" - config: Optional[GetEvaluationItemConfigDict] - """""" +class PscInterfaceConfigDict(TypedDict, total=False): + """Configuration for PSC-I.""" + dns_peering_configs: Optional[list[DnsPeeringConfigDict]] + """Optional. DNS peering configurations. When specified, Vertex AI will attempt to configure DNS peering zones in the tenant project VPC to resolve the specified domains using the target network's Cloud DNS. The user must grant the dns.peer role to the Vertex AI Service Agent on the target project.""" -_GetEvaluationItemParametersOrDict = Union[ - _GetEvaluationItemParameters, _GetEvaluationItemParametersDict -] + network_attachment: Optional[str] + """Optional. The name of the Compute Engine [network attachment](https://cloud.google.com/vpc/docs/about-network-attachments) to attach to the resource within the region and user project. To specify this field, you must have already [created a network attachment] (https://cloud.google.com/vpc/docs/create-manage-network-attachments#create-network-attachments). This field is only used for resources using PSC-I.""" -class ListEvaluationExperimentsConfig(_common.BaseModel): - """Config for listing evaluation experiments.""" +PscInterfaceConfigOrDict = Union[PscInterfaceConfig, PscInterfaceConfigDict] - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + +class Scheduling(_common.BaseModel): + """All parameters related to queuing and scheduling of custom jobs.""" + + disable_retries: Optional[bool] = Field( + default=None, + description="""Optional. Indicates if the job should retry for internal errors after the job starts running. If true, overrides `Scheduling.restart_job_on_worker_restart` to false.""", ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( + max_wait_duration: Optional[str] = Field( default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported. - For more information about filter syntax, see - `AIP-160 `_.""", + description="""Optional. This is the maximum duration that a job will wait for the requested resources to be provisioned if the scheduling strategy is set to [Strategy.DWS_FLEX_START]. If set to 0, the job will wait indefinitely. The default is 24 hours.""", ) - order_by: Optional[str] = Field( + restart_job_on_worker_restart: Optional[bool] = Field( default=None, - description="""A comma-separated list of fields to order by, sorted in ascending - order by default. Use ``desc`` after a field name for descending. - Example: ``"create_time desc"``.""", + description="""Optional. Restarts the entire CustomJob if a worker gets restarted. This feature can be used by distributed training jobs that are not resilient to workers leaving and joining a job.""", + ) + strategy: Optional[Strategy] = Field( + default=None, + description="""Optional. This determines which type of scheduling strategy to use.""", + ) + timeout: Optional[str] = Field( + default=None, + description="""Optional. The maximum job running time. The default is 7 days.""", ) -class ListEvaluationExperimentsConfigDict(TypedDict, total=False): - """Config for listing evaluation experiments.""" +class SchedulingDict(TypedDict, total=False): + """All parameters related to queuing and scheduling of custom jobs.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + disable_retries: Optional[bool] + """Optional. Indicates if the job should retry for internal errors after the job starts running. If true, overrides `Scheduling.restart_job_on_worker_restart` to false.""" - page_size: Optional[int] - """""" + max_wait_duration: Optional[str] + """Optional. This is the maximum duration that a job will wait for the requested resources to be provisioned if the scheduling strategy is set to [Strategy.DWS_FLEX_START]. If set to 0, the job will wait indefinitely. The default is 24 hours.""" - page_token: Optional[str] - """""" + restart_job_on_worker_restart: Optional[bool] + """Optional. Restarts the entire CustomJob if a worker gets restarted. This feature can be used by distributed training jobs that are not resilient to workers leaving and joining a job.""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported. - For more information about filter syntax, see - `AIP-160 `_.""" + strategy: Optional[Strategy] + """Optional. This determines which type of scheduling strategy to use.""" - order_by: Optional[str] - """A comma-separated list of fields to order by, sorted in ascending - order by default. Use ``desc`` after a field name for descending. - Example: ``"create_time desc"``.""" + timeout: Optional[str] + """Optional. The maximum job running time. The default is 7 days.""" -ListEvaluationExperimentsConfigOrDict = Union[ - ListEvaluationExperimentsConfig, ListEvaluationExperimentsConfigDict -] +SchedulingOrDict = Union[Scheduling, SchedulingDict] -class _ListEvaluationExperimentsParameters(_common.BaseModel): - """Parameters for listing evaluation experiments.""" +class EnvVar(_common.BaseModel): + """Represents an environment variable present in a Container or Python Module.""" - config: Optional[ListEvaluationExperimentsConfig] = Field( - default=None, description="""""" + name: Optional[str] = Field( + default=None, + description="""Required. Name of the environment variable. Must be a valid C identifier.""", + ) + value: Optional[str] = Field( + default=None, + description="""Required. Variables that reference a $(VAR_NAME) are expanded using the previous defined environment variables in the container and any service environment variables. If a variable cannot be resolved, the reference in the input string will be unchanged. The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped references will never be expanded, regardless of whether the variable exists or not.""", ) -class _ListEvaluationExperimentsParametersDict(TypedDict, total=False): - """Parameters for listing evaluation experiments.""" +class EnvVarDict(TypedDict, total=False): + """Represents an environment variable present in a Container or Python Module.""" - config: Optional[ListEvaluationExperimentsConfigDict] - """""" + name: Optional[str] + """Required. Name of the environment variable. Must be a valid C identifier.""" + value: Optional[str] + """Required. Variables that reference a $(VAR_NAME) are expanded using the previous defined environment variables in the container and any service environment variables. If a variable cannot be resolved, the reference in the input string will be unchanged. The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped references will never be expanded, regardless of whether the variable exists or not.""" -_ListEvaluationExperimentsParametersOrDict = Union[ - _ListEvaluationExperimentsParameters, _ListEvaluationExperimentsParametersDict -] +EnvVarOrDict = Union[EnvVar, EnvVarDict] -class ListEvaluationExperimentsResponse(_common.BaseModel): - """Response for listing evaluation experiments.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" +class ContainerSpec(_common.BaseModel): + """The spec of a Container.""" + + args: Optional[list[str]] = Field( + default=None, + description="""The arguments to be passed when starting the container.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - evaluation_experiments: Optional[list[EvaluationExperiment]] = Field( + command: Optional[list[str]] = Field( default=None, - description="""List of evaluation experiments. - """, + description="""The command to be invoked when the container is started. It overrides the entrypoint instruction in Dockerfile when provided.""", + ) + env: Optional[list[EnvVar]] = Field( + default=None, + description="""Environment variables to be passed to the container. Maximum limit is 100.""", + ) + image_uri: Optional[str] = Field( + default=None, + description="""Required. The URI of a container image in the Container Registry that is to be run on each worker replica.""", ) -class ListEvaluationExperimentsResponseDict(TypedDict, total=False): - """Response for listing evaluation experiments.""" +class ContainerSpecDict(TypedDict, total=False): + """The spec of a Container.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + args: Optional[list[str]] + """The arguments to be passed when starting the container.""" - next_page_token: Optional[str] - """""" + command: Optional[list[str]] + """The command to be invoked when the container is started. It overrides the entrypoint instruction in Dockerfile when provided.""" - evaluation_experiments: Optional[list[EvaluationExperimentDict]] - """List of evaluation experiments. - """ + env: Optional[list[EnvVarDict]] + """Environment variables to be passed to the container. Maximum limit is 100.""" + image_uri: Optional[str] + """Required. The URI of a container image in the Container Registry that is to be run on each worker replica.""" -ListEvaluationExperimentsResponseOrDict = Union[ - ListEvaluationExperimentsResponse, ListEvaluationExperimentsResponseDict -] +ContainerSpecOrDict = Union[ContainerSpec, ContainerSpecDict] -class ListEvaluationMetricsConfig(_common.BaseModel): - """Config for listing evaluation metrics.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( - default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported. - For more information about filter syntax, see - `AIP-160 `_.""", +class DiskSpec(_common.BaseModel): + """Represents the spec of disk options.""" + + boot_disk_size_gb: Optional[int] = Field( + default=None, description="""Size in GB of the boot disk (default is 100GB).""" ) - order_by: Optional[str] = Field( + boot_disk_type: Optional[str] = Field( default=None, - description="""A comma-separated list of fields to order by, sorted in ascending - order by default. Use ``desc`` after a field name for descending. - Example: ``"create_time desc"``.""", + description="""Type of the boot disk. For non-A3U machines, the default value is "pd-ssd", for A3U machines, the default value is "hyperdisk-balanced". Valid values: "pd-ssd" (Persistent Disk Solid State Drive), "pd-standard" (Persistent Disk Hard Disk Drive) or "hyperdisk-balanced".""", ) -class ListEvaluationMetricsConfigDict(TypedDict, total=False): - """Config for listing evaluation metrics.""" +class DiskSpecDict(TypedDict, total=False): + """Represents the spec of disk options.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + boot_disk_size_gb: Optional[int] + """Size in GB of the boot disk (default is 100GB).""" - page_size: Optional[int] - """""" + boot_disk_type: Optional[str] + """Type of the boot disk. For non-A3U machines, the default value is "pd-ssd", for A3U machines, the default value is "hyperdisk-balanced". Valid values: "pd-ssd" (Persistent Disk Solid State Drive), "pd-standard" (Persistent Disk Hard Disk Drive) or "hyperdisk-balanced".""" - page_token: Optional[str] - """""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported. - For more information about filter syntax, see - `AIP-160 `_.""" +DiskSpecOrDict = Union[DiskSpec, DiskSpecDict] - order_by: Optional[str] - """A comma-separated list of fields to order by, sorted in ascending - order by default. Use ``desc`` after a field name for descending. - Example: ``"create_time desc"``.""" +class LustreMount(_common.BaseModel): + """Represents a mount configuration for Lustre file system.""" -ListEvaluationMetricsConfigOrDict = Union[ - ListEvaluationMetricsConfig, ListEvaluationMetricsConfigDict -] + filesystem: Optional[str] = Field( + default=None, description="""Required. The name of the Lustre filesystem.""" + ) + instance_ip: Optional[str] = Field( + default=None, description="""Required. IP address of the Lustre instance.""" + ) + mount_point: Optional[str] = Field( + default=None, + description="""Required. Destination mount path. The Lustre file system will be mounted for the user under /mnt/lustre/""", + ) + volume_handle: Optional[str] = Field( + default=None, + description="""Required. The unique identifier of the Lustre volume.""", + ) -class _ListEvaluationMetricsParameters(_common.BaseModel): - """Parameters for listing evaluation metrics.""" +class LustreMountDict(TypedDict, total=False): + """Represents a mount configuration for Lustre file system.""" - config: Optional[ListEvaluationMetricsConfig] = Field( - default=None, description="""""" - ) + filesystem: Optional[str] + """Required. The name of the Lustre filesystem.""" + instance_ip: Optional[str] + """Required. IP address of the Lustre instance.""" -class _ListEvaluationMetricsParametersDict(TypedDict, total=False): - """Parameters for listing evaluation metrics.""" + mount_point: Optional[str] + """Required. Destination mount path. The Lustre file system will be mounted for the user under /mnt/lustre/""" - config: Optional[ListEvaluationMetricsConfigDict] - """""" + volume_handle: Optional[str] + """Required. The unique identifier of the Lustre volume.""" -_ListEvaluationMetricsParametersOrDict = Union[ - _ListEvaluationMetricsParameters, _ListEvaluationMetricsParametersDict -] +LustreMountOrDict = Union[LustreMount, LustreMountDict] -class ListEvaluationMetricsResponse(_common.BaseModel): - """Response for listing evaluation metrics.""" +class ReservationAffinity(_common.BaseModel): + """A ReservationAffinity can be used to configure a Vertex AI resource (e.g., a DeployedModel) to draw its Compute Engine resources from a Shared Reservation, or exclusively from on-demand capacity.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + key: Optional[str] = Field( + default=None, + description="""Optional. Corresponds to the label key of a reservation resource. To target a SPECIFIC_RESERVATION by name, use `compute.googleapis.com/reservation-name` as the key and specify the name of your reservation as its value.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - evaluation_metrics: Optional[list[EvaluationMetric]] = Field( + reservation_affinity_type: Optional[Type] = Field( default=None, - description="""List of evaluation metrics. - """, + description="""Required. Specifies the reservation affinity type.""", + ) + values: Optional[list[str]] = Field( + default=None, + description="""Optional. Corresponds to the label values of a reservation resource. This must be the full resource name of the reservation or reservation block.""", ) -class ListEvaluationMetricsResponseDict(TypedDict, total=False): - """Response for listing evaluation metrics.""" +class ReservationAffinityDict(TypedDict, total=False): + """A ReservationAffinity can be used to configure a Vertex AI resource (e.g., a DeployedModel) to draw its Compute Engine resources from a Shared Reservation, or exclusively from on-demand capacity.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + key: Optional[str] + """Optional. Corresponds to the label key of a reservation resource. To target a SPECIFIC_RESERVATION by name, use `compute.googleapis.com/reservation-name` as the key and specify the name of your reservation as its value.""" - next_page_token: Optional[str] - """""" + reservation_affinity_type: Optional[Type] + """Required. Specifies the reservation affinity type.""" - evaluation_metrics: Optional[list[EvaluationMetricDict]] - """List of evaluation metrics. - """ + values: Optional[list[str]] + """Optional. Corresponds to the label values of a reservation resource. This must be the full resource name of the reservation or reservation block.""" -ListEvaluationMetricsResponseOrDict = Union[ - ListEvaluationMetricsResponse, ListEvaluationMetricsResponseDict -] +ReservationAffinityOrDict = Union[ReservationAffinity, ReservationAffinityDict] -class UpdateEvaluationExperimentConfig(_common.BaseModel): - """Config for updating an evaluation experiment.""" +class MachineSpec(_common.BaseModel): + """Specification of a single machine.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + accelerator_count: Optional[int] = Field( + default=None, + description="""The number of accelerators to attach to the machine. For [accelerator optimized machine types](https://cloud.google.com/compute/docs/accelerator-optimized-machines), One may set the accelerator_count from 1 to N for machine with N GPUs. If accelerator_count is less than or equal to N / 2, Agent Platform co-schedules the replicas of the model into the same VM to save cost. For example, if the machine type is a3-highgpu-8g, which has 8 H100 GPUs, one can set accelerator_count to 1 to 8. If accelerator_count is 1, 2, 3, or 4, Agent Platform co-schedules 8, 4, 2, or 2 replicas of the model into the same VM to save cost. When co-scheduling, CPU, memory and storage on the VM will be distributed to replicas on the VM. For example, one can expect a co-scheduled replica requesting 2 GPUs out of a 8-GPU VM will receive 25% of the CPU, memory and storage of the VM. Note that the feature is not compatible with multihost_gpu_node_count. When multihost_gpu_node_count is set, the co-scheduling will not be enabled.""", ) - update_mask: Optional[str] = Field( + accelerator_type: Optional[AcceleratorType] = Field( default=None, - description="""The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""", + description="""Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""", ) - display_name: Optional[str] = Field( - default=None, description="""The display name of the evaluation experiment.""" + gpu_partition_size: Optional[str] = Field( + default=None, + description="""Optional. Immutable. The Nvidia GPU partition size. When specified, the requested accelerators will be partitioned into smaller GPU partitions. For example, if the request is for 8 units of NVIDIA A100 GPUs, and gpu_partition_size="1g.10gb", the service will create 8 * 7 = 56 partitioned MIG instances. The partition size must be a value supported by the requested accelerator. Refer to [Nvidia GPU Partitioning](https://cloud.google.com/kubernetes-engine/docs/how-to/gpus-multi#multi-instance_gpu_partitions) for the available partition sizes. If set, the accelerator_count should be set to 1.""", ) - labels: Optional[dict[str, str]] = Field( - default=None, description="""Labels for the evaluation experiment.""" + machine_type: Optional[str] = Field( + default=None, + description="""Immutable. The type of the machine. See the [list of machine types supported for prediction](https://cloud.google.com/gemini-enterprise-agent-platform/machine-learning/predictions/configure-compute#machine-types) See the [list of machine types supported for custom training](https://cloud.google.com/gemini-enterprise-agent-platform/machine-learning/training/configure-compute#machine-types). For DeployedModel this field is optional, and the default value is `n1-standard-2`. For BatchPredictionJob or as part of WorkerPoolSpec this field is required.""", ) - merge_strategy: Optional[EvaluationExperimentMergeStrategy] = Field( - default=None, description="""Merge strategy for the evaluation experiment.""" + min_gpu_driver_version: Optional[str] = Field( + default=None, + description="""Optional. Immutable. The minimum GPU driver version that this machine requires. For example, "535.104.06". If not specified, the default GPU driver version will be used by the underlying infrastructure.""", ) - metadata: Optional[dict[str, Any]] = Field( - default=None, description="""Metadata about the evaluation experiment.""" + multihost_gpu_node_count: Optional[int] = Field( + default=None, + description="""Optional. Immutable. The number of nodes per replica for multihost GPU deployments.""", + ) + reservation_affinity: Optional[ReservationAffinity] = Field( + default=None, + description="""Optional. Immutable. Configuration controlling how this resource pool consumes reservation.""", + ) + tpu_topology: Optional[str] = Field( + default=None, + description="""Immutable. The topology of the TPUs. Corresponds to the TPU topologies available from GKE. (Example: tpu_topology: "2x2x1").""", ) -class UpdateEvaluationExperimentConfigDict(TypedDict, total=False): - """Config for updating an evaluation experiment.""" +class MachineSpecDict(TypedDict, total=False): + """Specification of a single machine.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + accelerator_count: Optional[int] + """The number of accelerators to attach to the machine. For [accelerator optimized machine types](https://cloud.google.com/compute/docs/accelerator-optimized-machines), One may set the accelerator_count from 1 to N for machine with N GPUs. If accelerator_count is less than or equal to N / 2, Agent Platform co-schedules the replicas of the model into the same VM to save cost. For example, if the machine type is a3-highgpu-8g, which has 8 H100 GPUs, one can set accelerator_count to 1 to 8. If accelerator_count is 1, 2, 3, or 4, Agent Platform co-schedules 8, 4, 2, or 2 replicas of the model into the same VM to save cost. When co-scheduling, CPU, memory and storage on the VM will be distributed to replicas on the VM. For example, one can expect a co-scheduled replica requesting 2 GPUs out of a 8-GPU VM will receive 25% of the CPU, memory and storage of the VM. Note that the feature is not compatible with multihost_gpu_node_count. When multihost_gpu_node_count is set, the co-scheduling will not be enabled.""" - update_mask: Optional[str] - """The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" + accelerator_type: Optional[AcceleratorType] + """Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""" - display_name: Optional[str] - """The display name of the evaluation experiment.""" + gpu_partition_size: Optional[str] + """Optional. Immutable. The Nvidia GPU partition size. When specified, the requested accelerators will be partitioned into smaller GPU partitions. For example, if the request is for 8 units of NVIDIA A100 GPUs, and gpu_partition_size="1g.10gb", the service will create 8 * 7 = 56 partitioned MIG instances. The partition size must be a value supported by the requested accelerator. Refer to [Nvidia GPU Partitioning](https://cloud.google.com/kubernetes-engine/docs/how-to/gpus-multi#multi-instance_gpu_partitions) for the available partition sizes. If set, the accelerator_count should be set to 1.""" - labels: Optional[dict[str, str]] - """Labels for the evaluation experiment.""" + machine_type: Optional[str] + """Immutable. The type of the machine. See the [list of machine types supported for prediction](https://cloud.google.com/gemini-enterprise-agent-platform/machine-learning/predictions/configure-compute#machine-types) See the [list of machine types supported for custom training](https://cloud.google.com/gemini-enterprise-agent-platform/machine-learning/training/configure-compute#machine-types). For DeployedModel this field is optional, and the default value is `n1-standard-2`. For BatchPredictionJob or as part of WorkerPoolSpec this field is required.""" - merge_strategy: Optional[EvaluationExperimentMergeStrategy] - """Merge strategy for the evaluation experiment.""" + min_gpu_driver_version: Optional[str] + """Optional. Immutable. The minimum GPU driver version that this machine requires. For example, "535.104.06". If not specified, the default GPU driver version will be used by the underlying infrastructure.""" - metadata: Optional[dict[str, Any]] - """Metadata about the evaluation experiment.""" + multihost_gpu_node_count: Optional[int] + """Optional. Immutable. The number of nodes per replica for multihost GPU deployments.""" + reservation_affinity: Optional[ReservationAffinityDict] + """Optional. Immutable. Configuration controlling how this resource pool consumes reservation.""" -UpdateEvaluationExperimentConfigOrDict = Union[ - UpdateEvaluationExperimentConfig, UpdateEvaluationExperimentConfigDict -] + tpu_topology: Optional[str] + """Immutable. The topology of the TPUs. Corresponds to the TPU topologies available from GKE. (Example: tpu_topology: "2x2x1").""" -class _UpdateEvaluationExperimentParameters(_common.BaseModel): - """Parameters for updating an evaluation experiment.""" +MachineSpecOrDict = Union[MachineSpec, MachineSpecDict] - name: Optional[str] = Field( - default=None, description="""The resource name of the EvaluationExperiment.""" + +class NfsMount(_common.BaseModel): + """Represents a mount configuration for Network File System (NFS) to mount.""" + + mount_point: Optional[str] = Field( + default=None, + description="""Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""", ) - config: Optional[UpdateEvaluationExperimentConfig] = Field( - default=None, description="""""" + path: Optional[str] = Field( + default=None, + description="""Required. Source path exported from NFS server. Has to start with '/', and combined with the ip address, it indicates the source mount path in the form of `server:path`""", + ) + server: Optional[str] = Field( + default=None, description="""Required. IP address of the NFS server.""" ) -class _UpdateEvaluationExperimentParametersDict(TypedDict, total=False): - """Parameters for updating an evaluation experiment.""" +class NfsMountDict(TypedDict, total=False): + """Represents a mount configuration for Network File System (NFS) to mount.""" - name: Optional[str] - """The resource name of the EvaluationExperiment.""" + mount_point: Optional[str] + """Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""" - config: Optional[UpdateEvaluationExperimentConfigDict] - """""" + path: Optional[str] + """Required. Source path exported from NFS server. Has to start with '/', and combined with the ip address, it indicates the source mount path in the form of `server:path`""" + server: Optional[str] + """Required. IP address of the NFS server.""" -_UpdateEvaluationExperimentParametersOrDict = Union[ - _UpdateEvaluationExperimentParameters, _UpdateEvaluationExperimentParametersDict -] +NfsMountOrDict = Union[NfsMount, NfsMountDict] -class OptimizeConfig(_common.BaseModel): - """Config for Prompt Optimizer.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - optimization_target: Optional[OptimizeTarget] = Field( +class PythonPackageSpec(_common.BaseModel): + """The spec of a Python packaged code.""" + + args: Optional[list[str]] = Field( default=None, - description="""The optimization target for the prompt optimizer. It must be one of the OptimizeTarget enum values: OPTIMIZATION_TARGET_GEMINI_NANO for the prompts from Android core API, OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS for the few-shot prompt optimizer with rubrics, OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE for the few-shot prompt optimizer with target responses.""", + description="""Command line arguments to be passed to the Python task.""", ) - examples_dataframe: Optional[PandasDataFrame] = Field( + env: Optional[list[EnvVar]] = Field( default=None, - description="""The examples dataframe for the few-shot prompt optimizer. It must contain "prompt" and "model_response" columns. Depending on which optimization target is used, it also needs to contain "rubrics" and "rubrics_evaluations" or "target_response" columns.""", + description="""Environment variables to be passed to the python module. Maximum limit is 100.""", + ) + executor_image_uri: Optional[str] = Field( + default=None, + description="""Required. The URI of a container image in Artifact Registry that will run the provided Python package. Vertex AI provides a wide range of executor images with pre-installed packages to meet users' various use cases. See the list of [pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers). You must use an image from this list.""", + ) + package_uris: Optional[list[str]] = Field( + default=None, + description="""Required. The Google Cloud Storage location of the Python package files which are the training program and its dependent packages. The maximum number of package URIs is 100.""", + ) + python_module: Optional[str] = Field( + default=None, + description="""Required. The Python module name to run after installing the packages.""", ) -class OptimizeConfigDict(TypedDict, total=False): - """Config for Prompt Optimizer.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class PythonPackageSpecDict(TypedDict, total=False): + """The spec of a Python packaged code.""" - optimization_target: Optional[OptimizeTarget] - """The optimization target for the prompt optimizer. It must be one of the OptimizeTarget enum values: OPTIMIZATION_TARGET_GEMINI_NANO for the prompts from Android core API, OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS for the few-shot prompt optimizer with rubrics, OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE for the few-shot prompt optimizer with target responses.""" + args: Optional[list[str]] + """Command line arguments to be passed to the Python task.""" - examples_dataframe: Optional[PandasDataFrame] - """The examples dataframe for the few-shot prompt optimizer. It must contain "prompt" and "model_response" columns. Depending on which optimization target is used, it also needs to contain "rubrics" and "rubrics_evaluations" or "target_response" columns.""" + env: Optional[list[EnvVarDict]] + """Environment variables to be passed to the python module. Maximum limit is 100.""" + executor_image_uri: Optional[str] + """Required. The URI of a container image in Artifact Registry that will run the provided Python package. Vertex AI provides a wide range of executor images with pre-installed packages to meet users' various use cases. See the list of [pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers). You must use an image from this list.""" -OptimizeConfigOrDict = Union[OptimizeConfig, OptimizeConfigDict] + package_uris: Optional[list[str]] + """Required. The Google Cloud Storage location of the Python package files which are the training program and its dependent packages. The maximum number of package URIs is 100.""" + python_module: Optional[str] + """Required. The Python module name to run after installing the packages.""" -class _OptimizeRequestParameters(_common.BaseModel): - """Request for the optimize_prompt method.""" - content: Optional[genai_types.Content] = Field(default=None, description="""""") - config: Optional[OptimizeConfig] = Field(default=None, description="""""") +PythonPackageSpecOrDict = Union[PythonPackageSpec, PythonPackageSpecDict] -class _OptimizeRequestParametersDict(TypedDict, total=False): - """Request for the optimize_prompt method.""" +class WorkerPoolSpec(_common.BaseModel): + """Represents the spec of a worker pool in a job.""" - content: Optional[genai_types.Content] - """""" + container_spec: Optional[ContainerSpec] = Field( + default=None, description="""The custom container task.""" + ) + disk_spec: Optional[DiskSpec] = Field(default=None, description="""Disk spec.""") + lustre_mounts: Optional[list[LustreMount]] = Field( + default=None, description="""Optional. List of Lustre mounts.""" + ) + machine_spec: Optional[MachineSpec] = Field( + default=None, + description="""Optional. Immutable. The specification of a single machine.""", + ) + nfs_mounts: Optional[list[NfsMount]] = Field( + default=None, description="""Optional. List of NFS mount spec.""" + ) + python_package_spec: Optional[PythonPackageSpec] = Field( + default=None, description="""The Python packaged task.""" + ) + replica_count: Optional[int] = Field( + default=None, + description="""Optional. The number of worker replicas to use for this worker pool.""", + ) - config: Optional[OptimizeConfigDict] - """""" +class WorkerPoolSpecDict(TypedDict, total=False): + """Represents the spec of a worker pool in a job.""" -_OptimizeRequestParametersOrDict = Union[ - _OptimizeRequestParameters, _OptimizeRequestParametersDict -] + container_spec: Optional[ContainerSpecDict] + """The custom container task.""" + disk_spec: Optional[DiskSpecDict] + """Disk spec.""" -class OptimizeResponseEndpoint(_common.BaseModel): - """Response for the optimize_prompt method.""" + lustre_mounts: Optional[list[LustreMountDict]] + """Optional. List of Lustre mounts.""" - content: Optional[genai_types.Content] = Field(default=None, description="""""") + machine_spec: Optional[MachineSpecDict] + """Optional. Immutable. The specification of a single machine.""" + nfs_mounts: Optional[list[NfsMountDict]] + """Optional. List of NFS mount spec.""" -class OptimizeResponseEndpointDict(TypedDict, total=False): - """Response for the optimize_prompt method.""" + python_package_spec: Optional[PythonPackageSpecDict] + """The Python packaged task.""" - content: Optional[genai_types.Content] - """""" + replica_count: Optional[int] + """Optional. The number of worker replicas to use for this worker pool.""" -OptimizeResponseEndpointOrDict = Union[ - OptimizeResponseEndpoint, OptimizeResponseEndpointDict -] +WorkerPoolSpecOrDict = Union[WorkerPoolSpec, WorkerPoolSpecDict] -class DnsPeeringConfig(_common.BaseModel): - """DNS peering configuration. These configurations are used to create DNS peering zones in the Vertex tenant project VPC, enabling resolution of records within the specified domain hosted in the target network's Cloud DNS.""" +class CustomJobSpec(_common.BaseModel): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - domain: Optional[str] = Field( + base_output_directory: Optional[genai_types.GcsDestination] = Field( default=None, - description="""Required. The DNS name suffix of the zone being peered to, e.g., "my-internal-domain.corp.". Must end with a dot.""", + description="""The Cloud Storage location to store the output of this CustomJob or HyperparameterTuningJob. For HyperparameterTuningJob, the baseOutputDirectory of each child CustomJob backing a Trial is set to a subdirectory of name id under its parent HyperparameterTuningJob's baseOutputDirectory. The following Vertex AI environment variables will be passed to containers or python modules when this field is set: For CustomJob: * AIP_MODEL_DIR = `/model/` * AIP_CHECKPOINT_DIR = `/checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `/logs/` For CustomJob backing a Trial of HyperparameterTuningJob: * AIP_MODEL_DIR = `//model/` * AIP_CHECKPOINT_DIR = `//checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `//logs/`""", ) - target_network: Optional[str] = Field( + enable_dashboard_access: Optional[bool] = Field( default=None, - description="""Required. The VPC network name in the target_project where the DNS zone specified by 'domain' is visible.""", + description="""Optional. Whether you want Vertex AI to enable access to the customized dashboard in training chief container. If set to `true`, you can access the dashboard at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""", ) - target_project: Optional[str] = Field( + enable_web_access: Optional[bool] = Field( default=None, - description="""Required. The project ID hosting the Cloud DNS managed zone that contains the 'domain'. The Vertex AI Service Agent requires the dns.peer role on this project.""", + description="""Optional. Whether you want Vertex AI to enable [interactive shell access](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) to training containers. If set to `true`, you can access interactive shells at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""", ) - - -class DnsPeeringConfigDict(TypedDict, total=False): - """DNS peering configuration. These configurations are used to create DNS peering zones in the Vertex tenant project VPC, enabling resolution of records within the specified domain hosted in the target network's Cloud DNS.""" - - domain: Optional[str] - """Required. The DNS name suffix of the zone being peered to, e.g., "my-internal-domain.corp.". Must end with a dot.""" - - target_network: Optional[str] - """Required. The VPC network name in the target_project where the DNS zone specified by 'domain' is visible.""" - - target_project: Optional[str] - """Required. The project ID hosting the Cloud DNS managed zone that contains the 'domain'. The Vertex AI Service Agent requires the dns.peer role on this project.""" - - -DnsPeeringConfigOrDict = Union[DnsPeeringConfig, DnsPeeringConfigDict] - - -class PscInterfaceConfig(_common.BaseModel): - """Configuration for PSC-I.""" - - dns_peering_configs: Optional[list[DnsPeeringConfig]] = Field( + experiment: Optional[str] = Field( default=None, - description="""Optional. DNS peering configurations. When specified, Vertex AI will attempt to configure DNS peering zones in the tenant project VPC to resolve the specified domains using the target network's Cloud DNS. The user must grant the dns.peer role to the Vertex AI Service Agent on the target project.""", + description="""Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`""", ) - network_attachment: Optional[str] = Field( + experiment_run: Optional[str] = Field( default=None, - description="""Optional. The name of the Compute Engine [network attachment](https://cloud.google.com/vpc/docs/about-network-attachments) to attach to the resource within the region and user project. To specify this field, you must have already [created a network attachment] (https://cloud.google.com/vpc/docs/create-manage-network-attachments#create-network-attachments). This field is only used for resources using PSC-I.""", + description="""Optional. The Experiment Run associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}-{experiment-run-name}`""", ) - - -class PscInterfaceConfigDict(TypedDict, total=False): - """Configuration for PSC-I.""" - - dns_peering_configs: Optional[list[DnsPeeringConfigDict]] - """Optional. DNS peering configurations. When specified, Vertex AI will attempt to configure DNS peering zones in the tenant project VPC to resolve the specified domains using the target network's Cloud DNS. The user must grant the dns.peer role to the Vertex AI Service Agent on the target project.""" - - network_attachment: Optional[str] - """Optional. The name of the Compute Engine [network attachment](https://cloud.google.com/vpc/docs/about-network-attachments) to attach to the resource within the region and user project. To specify this field, you must have already [created a network attachment] (https://cloud.google.com/vpc/docs/create-manage-network-attachments#create-network-attachments). This field is only used for resources using PSC-I.""" - - -PscInterfaceConfigOrDict = Union[PscInterfaceConfig, PscInterfaceConfigDict] - - -class Scheduling(_common.BaseModel): - """All parameters related to queuing and scheduling of custom jobs.""" - - disable_retries: Optional[bool] = Field( + models: Optional[list[str]] = Field( default=None, - description="""Optional. Indicates if the job should retry for internal errors after the job starts running. If true, overrides `Scheduling.restart_job_on_worker_restart` to false.""", + description="""Optional. The name of the Model resources for which to generate a mapping to artifact URIs. Applicable only to some of the Google-provided custom jobs. Format: `projects/{project}/locations/{location}/models/{model}` In order to retrieve a specific version of the model, also provide the version ID or version alias. Example: `projects/{project}/locations/{location}/models/{model}@2` or `projects/{project}/locations/{location}/models/{model}@golden` If no version ID or alias is specified, the "default" version will be returned. The "default" version alias is created for the first version of the model, and can be moved to other versions later on. There will be exactly one default version.""", ) - max_wait_duration: Optional[str] = Field( + network: Optional[str] = Field( default=None, - description="""Optional. This is the maximum duration that a job will wait for the requested resources to be provisioned if the scheduling strategy is set to [Strategy.DWS_FLEX_START]. If set to 0, the job will wait indefinitely. The default is 24 hours.""", + description="""Optional. The full name of the Compute Engine [network](/compute/docs/networks-and-firewalls#networks) to which the Job should be peered. For example, `projects/12345/global/networks/myVPC`. [Format](/compute/docs/reference/rest/v1/networks/insert) is of the form `projects/{project}/global/networks/{network}`. Where {project} is a project number, as in `12345`, and {network} is a network name. To specify this field, you must have already [configured VPC Network Peering for Vertex AI](https://cloud.google.com/vertex-ai/docs/general/vpc-peering). If this field is left unspecified, the job is not peered with any network.""", ) - restart_job_on_worker_restart: Optional[bool] = Field( + persistent_resource_id: Optional[str] = Field( default=None, - description="""Optional. Restarts the entire CustomJob if a worker gets restarted. This feature can be used by distributed training jobs that are not resilient to workers leaving and joining a job.""", + description="""Optional. The ID of the PersistentResource in the same Project and Location which to run If this is specified, the job will be run on existing machines held by the PersistentResource instead of on-demand short-live machines. The network and CMEK configs on the job should be consistent with those on the PersistentResource, otherwise, the job will be rejected.""", ) - strategy: Optional[Strategy] = Field( + protected_artifact_location_id: Optional[str] = Field( default=None, - description="""Optional. This determines which type of scheduling strategy to use.""", + description="""The ID of the location to store protected artifacts. e.g. us-central1. Populate only when the location is different than CustomJob location. List of supported locations: https://cloud.google.com/vertex-ai/docs/general/locations""", ) - timeout: Optional[str] = Field( + psc_interface_config: Optional[PscInterfaceConfig] = Field( + default=None, description="""Optional. Configuration for PSC-I for CustomJob.""" + ) + reserved_ip_ranges: Optional[list[str]] = Field( default=None, - description="""Optional. The maximum job running time. The default is 7 days.""", + description="""Optional. A list of names for the reserved ip ranges under the VPC network that can be used for this job. If set, we will deploy the job within the provided ip ranges. Otherwise, the job will be deployed to any ip ranges under the provided VPC network. Example: ['vertex-ai-ip-range'].""", + ) + scheduling: Optional[Scheduling] = Field( + default=None, description="""Scheduling options for a CustomJob.""" + ) + service_account: Optional[str] = Field( + default=None, + description="""Specifies the service account for workload run-as account. Users submitting jobs must have act-as permission on this run-as account. If unspecified, the [Vertex AI Custom Code Service Agent](https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents) for the CustomJob's project is used.""", + ) + tensorboard: Optional[str] = Field( + default=None, + description="""Optional. The name of a Vertex AI Tensorboard resource to which this CustomJob will upload Tensorboard logs. Format: `projects/{project}/locations/{location}/tensorboards/{tensorboard}`""", + ) + worker_pool_specs: Optional[list[WorkerPoolSpec]] = Field( + default=None, + description="""Required. The spec of the worker pools including machine type and Docker image. All worker pools except the first one are optional and can be skipped by providing an empty value.""", ) -class SchedulingDict(TypedDict, total=False): - """All parameters related to queuing and scheduling of custom jobs.""" +class CustomJobSpecDict(TypedDict, total=False): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - disable_retries: Optional[bool] - """Optional. Indicates if the job should retry for internal errors after the job starts running. If true, overrides `Scheduling.restart_job_on_worker_restart` to false.""" + base_output_directory: Optional[genai_types.GcsDestination] + """The Cloud Storage location to store the output of this CustomJob or HyperparameterTuningJob. For HyperparameterTuningJob, the baseOutputDirectory of each child CustomJob backing a Trial is set to a subdirectory of name id under its parent HyperparameterTuningJob's baseOutputDirectory. The following Vertex AI environment variables will be passed to containers or python modules when this field is set: For CustomJob: * AIP_MODEL_DIR = `/model/` * AIP_CHECKPOINT_DIR = `/checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `/logs/` For CustomJob backing a Trial of HyperparameterTuningJob: * AIP_MODEL_DIR = `//model/` * AIP_CHECKPOINT_DIR = `//checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `//logs/`""" - max_wait_duration: Optional[str] - """Optional. This is the maximum duration that a job will wait for the requested resources to be provisioned if the scheduling strategy is set to [Strategy.DWS_FLEX_START]. If set to 0, the job will wait indefinitely. The default is 24 hours.""" + enable_dashboard_access: Optional[bool] + """Optional. Whether you want Vertex AI to enable access to the customized dashboard in training chief container. If set to `true`, you can access the dashboard at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""" - restart_job_on_worker_restart: Optional[bool] - """Optional. Restarts the entire CustomJob if a worker gets restarted. This feature can be used by distributed training jobs that are not resilient to workers leaving and joining a job.""" + enable_web_access: Optional[bool] + """Optional. Whether you want Vertex AI to enable [interactive shell access](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) to training containers. If set to `true`, you can access interactive shells at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""" - strategy: Optional[Strategy] - """Optional. This determines which type of scheduling strategy to use.""" + experiment: Optional[str] + """Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`""" - timeout: Optional[str] - """Optional. The maximum job running time. The default is 7 days.""" + experiment_run: Optional[str] + """Optional. The Experiment Run associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}-{experiment-run-name}`""" + models: Optional[list[str]] + """Optional. The name of the Model resources for which to generate a mapping to artifact URIs. Applicable only to some of the Google-provided custom jobs. Format: `projects/{project}/locations/{location}/models/{model}` In order to retrieve a specific version of the model, also provide the version ID or version alias. Example: `projects/{project}/locations/{location}/models/{model}@2` or `projects/{project}/locations/{location}/models/{model}@golden` If no version ID or alias is specified, the "default" version will be returned. The "default" version alias is created for the first version of the model, and can be moved to other versions later on. There will be exactly one default version.""" -SchedulingOrDict = Union[Scheduling, SchedulingDict] + network: Optional[str] + """Optional. The full name of the Compute Engine [network](/compute/docs/networks-and-firewalls#networks) to which the Job should be peered. For example, `projects/12345/global/networks/myVPC`. [Format](/compute/docs/reference/rest/v1/networks/insert) is of the form `projects/{project}/global/networks/{network}`. Where {project} is a project number, as in `12345`, and {network} is a network name. To specify this field, you must have already [configured VPC Network Peering for Vertex AI](https://cloud.google.com/vertex-ai/docs/general/vpc-peering). If this field is left unspecified, the job is not peered with any network.""" + persistent_resource_id: Optional[str] + """Optional. The ID of the PersistentResource in the same Project and Location which to run If this is specified, the job will be run on existing machines held by the PersistentResource instead of on-demand short-live machines. The network and CMEK configs on the job should be consistent with those on the PersistentResource, otherwise, the job will be rejected.""" -class EnvVar(_common.BaseModel): - """Represents an environment variable present in a Container or Python Module.""" + protected_artifact_location_id: Optional[str] + """The ID of the location to store protected artifacts. e.g. us-central1. Populate only when the location is different than CustomJob location. List of supported locations: https://cloud.google.com/vertex-ai/docs/general/locations""" - name: Optional[str] = Field( - default=None, - description="""Required. Name of the environment variable. Must be a valid C identifier.""", - ) - value: Optional[str] = Field( - default=None, - description="""Required. Variables that reference a $(VAR_NAME) are expanded using the previous defined environment variables in the container and any service environment variables. If a variable cannot be resolved, the reference in the input string will be unchanged. The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped references will never be expanded, regardless of whether the variable exists or not.""", - ) + psc_interface_config: Optional[PscInterfaceConfigDict] + """Optional. Configuration for PSC-I for CustomJob.""" + reserved_ip_ranges: Optional[list[str]] + """Optional. A list of names for the reserved ip ranges under the VPC network that can be used for this job. If set, we will deploy the job within the provided ip ranges. Otherwise, the job will be deployed to any ip ranges under the provided VPC network. Example: ['vertex-ai-ip-range'].""" -class EnvVarDict(TypedDict, total=False): - """Represents an environment variable present in a Container or Python Module.""" + scheduling: Optional[SchedulingDict] + """Scheduling options for a CustomJob.""" - name: Optional[str] - """Required. Name of the environment variable. Must be a valid C identifier.""" + service_account: Optional[str] + """Specifies the service account for workload run-as account. Users submitting jobs must have act-as permission on this run-as account. If unspecified, the [Vertex AI Custom Code Service Agent](https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents) for the CustomJob's project is used.""" - value: Optional[str] - """Required. Variables that reference a $(VAR_NAME) are expanded using the previous defined environment variables in the container and any service environment variables. If a variable cannot be resolved, the reference in the input string will be unchanged. The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped references will never be expanded, regardless of whether the variable exists or not.""" + tensorboard: Optional[str] + """Optional. The name of a Vertex AI Tensorboard resource to which this CustomJob will upload Tensorboard logs. Format: `projects/{project}/locations/{location}/tensorboards/{tensorboard}`""" + worker_pool_specs: Optional[list[WorkerPoolSpecDict]] + """Required. The spec of the worker pools including machine type and Docker image. All worker pools except the first one are optional and can be skipped by providing an empty value.""" -EnvVarOrDict = Union[EnvVar, EnvVarDict] +CustomJobSpecOrDict = Union[CustomJobSpec, CustomJobSpecDict] -class ContainerSpec(_common.BaseModel): - """The spec of a Container.""" - args: Optional[list[str]] = Field( +class CustomJob(_common.BaseModel): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" + + display_name: Optional[str] = Field( default=None, - description="""The arguments to be passed when starting the container.""", + description="""Required. The display name of the CustomJob. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", ) - command: Optional[list[str]] = Field( + job_spec: Optional[CustomJobSpec] = Field( + default=None, description="""Required. Job spec.""" + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, - description="""The command to be invoked when the container is started. It overrides the entrypoint instruction in Dockerfile when provided.""", + description="""Customer-managed encryption key options for a CustomJob. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key.""", ) - env: Optional[list[EnvVar]] = Field( + state: Optional[genai_types.JobState] = Field( + default=None, description="""Output only. The detailed state of the job.""" + ) + error: Optional[genai_types.GoogleRpcStatus] = Field( default=None, - description="""Environment variables to be passed to the container. Maximum limit is 100.""", + description="""Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`.""", ) - image_uri: Optional[str] = Field( + create_time: Optional[datetime.datetime] = Field( default=None, - description="""Required. The URI of a container image in the Container Registry that is to be run on each worker replica.""", + description="""Output only. Time when the CustomJob was created.""", + ) + end_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Time when the CustomJob entered any of the following states: `JOB_STATE_SUCCEEDED`, `JOB_STATE_FAILED`, `JOB_STATE_CANCELLED`.""", + ) + labels: Optional[dict[str, str]] = Field( + default=None, + description="""The labels with user-defined metadata to organize CustomJobs. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", + ) + name: Optional[str] = Field( + default=None, description="""Output only. Resource name of a CustomJob.""" + ) + satisfies_pzi: Optional[bool] = Field( + default=None, description="""Output only. Reserved for future use.""" + ) + satisfies_pzs: Optional[bool] = Field( + default=None, description="""Output only. Reserved for future use.""" + ) + start_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state.""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Time when the CustomJob was most recently updated.""", + ) + web_access_uris: Optional[dict[str, str]] = Field( + default=None, + description="""Output only. URIs for accessing [interactive shells](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) (one URI for each training node). Only available if job_spec.enable_web_access is `true`. The keys are names of each node in the training job; for example, `workerpool0-0` for the primary node, `workerpool1-0` for the first node in the second worker pool, and `workerpool1-1` for the second node in the second worker pool. The values are the URIs for each node's interactive shell.""", ) -class ContainerSpecDict(TypedDict, total=False): - """The spec of a Container.""" +class CustomJobDict(TypedDict, total=False): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - args: Optional[list[str]] - """The arguments to be passed when starting the container.""" + display_name: Optional[str] + """Required. The display name of the CustomJob. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" - command: Optional[list[str]] - """The command to be invoked when the container is started. It overrides the entrypoint instruction in Dockerfile when provided.""" + job_spec: Optional[CustomJobSpecDict] + """Required. Job spec.""" - env: Optional[list[EnvVarDict]] - """Environment variables to be passed to the container. Maximum limit is 100.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key options for a CustomJob. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key.""" - image_uri: Optional[str] - """Required. The URI of a container image in the Container Registry that is to be run on each worker replica.""" + state: Optional[genai_types.JobState] + """Output only. The detailed state of the job.""" + error: Optional[genai_types.GoogleRpcStatus] + """Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`.""" -ContainerSpecOrDict = Union[ContainerSpec, ContainerSpecDict] + create_time: Optional[datetime.datetime] + """Output only. Time when the CustomJob was created.""" + end_time: Optional[datetime.datetime] + """Output only. Time when the CustomJob entered any of the following states: `JOB_STATE_SUCCEEDED`, `JOB_STATE_FAILED`, `JOB_STATE_CANCELLED`.""" -class DiskSpec(_common.BaseModel): - """Represents the spec of disk options.""" + labels: Optional[dict[str, str]] + """The labels with user-defined metadata to organize CustomJobs. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""" - boot_disk_size_gb: Optional[int] = Field( - default=None, description="""Size in GB of the boot disk (default is 100GB).""" - ) - boot_disk_type: Optional[str] = Field( - default=None, - description="""Type of the boot disk. For non-A3U machines, the default value is "pd-ssd", for A3U machines, the default value is "hyperdisk-balanced". Valid values: "pd-ssd" (Persistent Disk Solid State Drive), "pd-standard" (Persistent Disk Hard Disk Drive) or "hyperdisk-balanced".""", - ) + name: Optional[str] + """Output only. Resource name of a CustomJob.""" + satisfies_pzi: Optional[bool] + """Output only. Reserved for future use.""" -class DiskSpecDict(TypedDict, total=False): - """Represents the spec of disk options.""" + satisfies_pzs: Optional[bool] + """Output only. Reserved for future use.""" - boot_disk_size_gb: Optional[int] - """Size in GB of the boot disk (default is 100GB).""" + start_time: Optional[datetime.datetime] + """Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state.""" - boot_disk_type: Optional[str] - """Type of the boot disk. For non-A3U machines, the default value is "pd-ssd", for A3U machines, the default value is "hyperdisk-balanced". Valid values: "pd-ssd" (Persistent Disk Solid State Drive), "pd-standard" (Persistent Disk Hard Disk Drive) or "hyperdisk-balanced".""" + update_time: Optional[datetime.datetime] + """Output only. Time when the CustomJob was most recently updated.""" + web_access_uris: Optional[dict[str, str]] + """Output only. URIs for accessing [interactive shells](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) (one URI for each training node). Only available if job_spec.enable_web_access is `true`. The keys are names of each node in the training job; for example, `workerpool0-0` for the primary node, `workerpool1-0` for the first node in the second worker pool, and `workerpool1-1` for the second node in the second worker pool. The values are the URIs for each node's interactive shell.""" -DiskSpecOrDict = Union[DiskSpec, DiskSpecDict] +CustomJobOrDict = Union[CustomJob, CustomJobDict] -class LustreMount(_common.BaseModel): - """Represents a mount configuration for Lustre file system.""" - filesystem: Optional[str] = Field( - default=None, description="""Required. The name of the Lustre filesystem.""" - ) - instance_ip: Optional[str] = Field( - default=None, description="""Required. IP address of the Lustre instance.""" - ) - mount_point: Optional[str] = Field( - default=None, - description="""Required. Destination mount path. The Lustre file system will be mounted for the user under /mnt/lustre/""", - ) - volume_handle: Optional[str] = Field( - default=None, - description="""Required. The unique identifier of the Lustre volume.""", +class VertexBaseConfig(_common.BaseModel): + """Base config for Vertex AI.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class LustreMountDict(TypedDict, total=False): - """Represents a mount configuration for Lustre file system.""" +class VertexBaseConfigDict(TypedDict, total=False): + """Base config for Vertex AI.""" - filesystem: Optional[str] - """Required. The name of the Lustre filesystem.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - instance_ip: Optional[str] - """Required. IP address of the Lustre instance.""" - mount_point: Optional[str] - """Required. Destination mount path. The Lustre file system will be mounted for the user under /mnt/lustre/""" +VertexBaseConfigOrDict = Union[VertexBaseConfig, VertexBaseConfigDict] - volume_handle: Optional[str] - """Required. The unique identifier of the Lustre volume.""" +class _CustomJobParameters(_common.BaseModel): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" -LustreMountOrDict = Union[LustreMount, LustreMountDict] + custom_job: Optional[CustomJob] = Field(default=None, description="""""") + config: Optional[VertexBaseConfig] = Field(default=None, description="""""") -class ReservationAffinity(_common.BaseModel): - """A ReservationAffinity can be used to configure a Vertex AI resource (e.g., a DeployedModel) to draw its Compute Engine resources from a Shared Reservation, or exclusively from on-demand capacity.""" +class _CustomJobParametersDict(TypedDict, total=False): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - key: Optional[str] = Field( - default=None, - description="""Optional. Corresponds to the label key of a reservation resource. To target a SPECIFIC_RESERVATION by name, use `compute.googleapis.com/reservation-name` as the key and specify the name of your reservation as its value.""", - ) - reservation_affinity_type: Optional[Type] = Field( - default=None, - description="""Required. Specifies the reservation affinity type.""", - ) - values: Optional[list[str]] = Field( - default=None, - description="""Optional. Corresponds to the label values of a reservation resource. This must be the full resource name of the reservation or reservation block.""", - ) + custom_job: Optional[CustomJobDict] + """""" + config: Optional[VertexBaseConfigDict] + """""" -class ReservationAffinityDict(TypedDict, total=False): - """A ReservationAffinity can be used to configure a Vertex AI resource (e.g., a DeployedModel) to draw its Compute Engine resources from a Shared Reservation, or exclusively from on-demand capacity.""" - key: Optional[str] - """Optional. Corresponds to the label key of a reservation resource. To target a SPECIFIC_RESERVATION by name, use `compute.googleapis.com/reservation-name` as the key and specify the name of your reservation as its value.""" +_CustomJobParametersOrDict = Union[_CustomJobParameters, _CustomJobParametersDict] - reservation_affinity_type: Optional[Type] - """Required. Specifies the reservation affinity type.""" - values: Optional[list[str]] - """Optional. Corresponds to the label values of a reservation resource. This must be the full resource name of the reservation or reservation block.""" +class _GetCustomJobParameters(_common.BaseModel): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" + name: Optional[str] = Field(default=None, description="""""") + config: Optional[VertexBaseConfig] = Field(default=None, description="""""") -ReservationAffinityOrDict = Union[ReservationAffinity, ReservationAffinityDict] +class _GetCustomJobParametersDict(TypedDict, total=False): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" -class MachineSpec(_common.BaseModel): - """Specification of a single machine.""" + name: Optional[str] + """""" - accelerator_count: Optional[int] = Field( - default=None, - description="""The number of accelerators to attach to the machine. For [accelerator optimized machine types](https://cloud.google.com/compute/docs/accelerator-optimized-machines), One may set the accelerator_count from 1 to N for machine with N GPUs. If accelerator_count is less than or equal to N / 2, Agent Platform co-schedules the replicas of the model into the same VM to save cost. For example, if the machine type is a3-highgpu-8g, which has 8 H100 GPUs, one can set accelerator_count to 1 to 8. If accelerator_count is 1, 2, 3, or 4, Agent Platform co-schedules 8, 4, 2, or 2 replicas of the model into the same VM to save cost. When co-scheduling, CPU, memory and storage on the VM will be distributed to replicas on the VM. For example, one can expect a co-scheduled replica requesting 2 GPUs out of a 8-GPU VM will receive 25% of the CPU, memory and storage of the VM. Note that the feature is not compatible with multihost_gpu_node_count. When multihost_gpu_node_count is set, the co-scheduling will not be enabled.""", - ) - accelerator_type: Optional[AcceleratorType] = Field( - default=None, - description="""Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""", - ) - gpu_partition_size: Optional[str] = Field( - default=None, - description="""Optional. Immutable. The Nvidia GPU partition size. When specified, the requested accelerators will be partitioned into smaller GPU partitions. For example, if the request is for 8 units of NVIDIA A100 GPUs, and gpu_partition_size="1g.10gb", the service will create 8 * 7 = 56 partitioned MIG instances. The partition size must be a value supported by the requested accelerator. Refer to [Nvidia GPU Partitioning](https://cloud.google.com/kubernetes-engine/docs/how-to/gpus-multi#multi-instance_gpu_partitions) for the available partition sizes. If set, the accelerator_count should be set to 1.""", - ) - machine_type: Optional[str] = Field( - default=None, - description="""Immutable. The type of the machine. See the [list of machine types supported for prediction](https://cloud.google.com/gemini-enterprise-agent-platform/machine-learning/predictions/configure-compute#machine-types) See the [list of machine types supported for custom training](https://cloud.google.com/gemini-enterprise-agent-platform/machine-learning/training/configure-compute#machine-types). For DeployedModel this field is optional, and the default value is `n1-standard-2`. For BatchPredictionJob or as part of WorkerPoolSpec this field is required.""", - ) - min_gpu_driver_version: Optional[str] = Field( - default=None, - description="""Optional. Immutable. The minimum GPU driver version that this machine requires. For example, "535.104.06". If not specified, the default GPU driver version will be used by the underlying infrastructure.""", - ) - multihost_gpu_node_count: Optional[int] = Field( - default=None, - description="""Optional. Immutable. The number of nodes per replica for multihost GPU deployments.""", - ) - reservation_affinity: Optional[ReservationAffinity] = Field( - default=None, - description="""Optional. Immutable. Configuration controlling how this resource pool consumes reservation.""", + config: Optional[VertexBaseConfigDict] + """""" + + +_GetCustomJobParametersOrDict = Union[ + _GetCustomJobParameters, _GetCustomJobParametersDict +] + + +class CancelQueryJobRuntimeConfig(_common.BaseModel): + """Config for canceling async querying agent runtimes.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - tpu_topology: Optional[str] = Field( + operation_name: Optional[str] = Field( default=None, - description="""Immutable. The topology of the TPUs. Corresponds to the TPU topologies available from GKE. (Example: tpu_topology: "2x2x1").""", + description="""Name of the longrunning operation returned from run_query_job.""", ) -class MachineSpecDict(TypedDict, total=False): - """Specification of a single machine.""" +class CancelQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for canceling async querying agent runtimes.""" - accelerator_count: Optional[int] - """The number of accelerators to attach to the machine. For [accelerator optimized machine types](https://cloud.google.com/compute/docs/accelerator-optimized-machines), One may set the accelerator_count from 1 to N for machine with N GPUs. If accelerator_count is less than or equal to N / 2, Agent Platform co-schedules the replicas of the model into the same VM to save cost. For example, if the machine type is a3-highgpu-8g, which has 8 H100 GPUs, one can set accelerator_count to 1 to 8. If accelerator_count is 1, 2, 3, or 4, Agent Platform co-schedules 8, 4, 2, or 2 replicas of the model into the same VM to save cost. When co-scheduling, CPU, memory and storage on the VM will be distributed to replicas on the VM. For example, one can expect a co-scheduled replica requesting 2 GPUs out of a 8-GPU VM will receive 25% of the CPU, memory and storage of the VM. Note that the feature is not compatible with multihost_gpu_node_count. When multihost_gpu_node_count is set, the co-scheduling will not be enabled.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - accelerator_type: Optional[AcceleratorType] - """Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""" + operation_name: Optional[str] + """Name of the longrunning operation returned from run_query_job.""" - gpu_partition_size: Optional[str] - """Optional. Immutable. The Nvidia GPU partition size. When specified, the requested accelerators will be partitioned into smaller GPU partitions. For example, if the request is for 8 units of NVIDIA A100 GPUs, and gpu_partition_size="1g.10gb", the service will create 8 * 7 = 56 partitioned MIG instances. The partition size must be a value supported by the requested accelerator. Refer to [Nvidia GPU Partitioning](https://cloud.google.com/kubernetes-engine/docs/how-to/gpus-multi#multi-instance_gpu_partitions) for the available partition sizes. If set, the accelerator_count should be set to 1.""" - machine_type: Optional[str] - """Immutable. The type of the machine. See the [list of machine types supported for prediction](https://cloud.google.com/gemini-enterprise-agent-platform/machine-learning/predictions/configure-compute#machine-types) See the [list of machine types supported for custom training](https://cloud.google.com/gemini-enterprise-agent-platform/machine-learning/training/configure-compute#machine-types). For DeployedModel this field is optional, and the default value is `n1-standard-2`. For BatchPredictionJob or as part of WorkerPoolSpec this field is required.""" +CancelQueryJobRuntimeConfigOrDict = Union[ + CancelQueryJobRuntimeConfig, CancelQueryJobRuntimeConfigDict +] - min_gpu_driver_version: Optional[str] - """Optional. Immutable. The minimum GPU driver version that this machine requires. For example, "535.104.06". If not specified, the default GPU driver version will be used by the underlying infrastructure.""" - multihost_gpu_node_count: Optional[int] - """Optional. Immutable. The number of nodes per replica for multihost GPU deployments.""" +class _CancelQueryJobRuntimeRequestParameters(_common.BaseModel): + """Parameters for canceling async querying agent runtimes.""" - reservation_affinity: Optional[ReservationAffinityDict] - """Optional. Immutable. Configuration controlling how this resource pool consumes reservation.""" + name: Optional[str] = Field( + default=None, description="""Name of the reasoning engine resource.""" + ) + config: Optional[CancelQueryJobRuntimeConfig] = Field( + default=None, description="""""" + ) - tpu_topology: Optional[str] - """Immutable. The topology of the TPUs. Corresponds to the TPU topologies available from GKE. (Example: tpu_topology: "2x2x1").""" +class _CancelQueryJobRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for canceling async querying agent runtimes.""" -MachineSpecOrDict = Union[MachineSpec, MachineSpecDict] + name: Optional[str] + """Name of the reasoning engine resource.""" + config: Optional[CancelQueryJobRuntimeConfigDict] + """""" -class NfsMount(_common.BaseModel): - """Represents a mount configuration for Network File System (NFS) to mount.""" - mount_point: Optional[str] = Field( - default=None, - description="""Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""", - ) - path: Optional[str] = Field( - default=None, - description="""Required. Source path exported from NFS server. Has to start with '/', and combined with the ip address, it indicates the source mount path in the form of `server:path`""", - ) - server: Optional[str] = Field( - default=None, description="""Required. IP address of the NFS server.""" - ) +_CancelQueryJobRuntimeRequestParametersOrDict = Union[ + _CancelQueryJobRuntimeRequestParameters, _CancelQueryJobRuntimeRequestParametersDict +] -class NfsMountDict(TypedDict, total=False): - """Represents a mount configuration for Network File System (NFS) to mount.""" +class CancelQueryJobResult(_common.BaseModel): + """Result of canceling a query job.""" - mount_point: Optional[str] - """Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) - path: Optional[str] - """Required. Source path exported from NFS server. Has to start with '/', and combined with the ip address, it indicates the source mount path in the form of `server:path`""" - server: Optional[str] - """Required. IP address of the NFS server.""" +class CancelQueryJobResultDict(TypedDict, total=False): + """Result of canceling a query job.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -NfsMountOrDict = Union[NfsMount, NfsMountDict] +CancelQueryJobResultOrDict = Union[CancelQueryJobResult, CancelQueryJobResultDict] -class PythonPackageSpec(_common.BaseModel): - """The spec of a Python packaged code.""" - args: Optional[list[str]] = Field( - default=None, - description="""Command line arguments to be passed to the Python task.""", - ) - env: Optional[list[EnvVar]] = Field( - default=None, - description="""Environment variables to be passed to the python module. Maximum limit is 100.""", - ) - executor_image_uri: Optional[str] = Field( - default=None, - description="""Required. The URI of a container image in Artifact Registry that will run the provided Python package. Vertex AI provides a wide range of executor images with pre-installed packages to meet users' various use cases. See the list of [pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers). You must use an image from this list.""", - ) - package_uris: Optional[list[str]] = Field( - default=None, - description="""Required. The Google Cloud Storage location of the Python package files which are the training program and its dependent packages. The maximum number of package URIs is 100.""", +class CheckQueryJobRuntimeConfig(_common.BaseModel): + """Config for async querying agent runtimes.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - python_module: Optional[str] = Field( + retrieve_result: Optional[bool] = Field( default=None, - description="""Required. The Python module name to run after installing the packages.""", + description="""Whether to retrieve the results of the query job.""", ) -class PythonPackageSpecDict(TypedDict, total=False): - """The spec of a Python packaged code.""" +class CheckQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for async querying agent runtimes.""" - args: Optional[list[str]] - """Command line arguments to be passed to the Python task.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - env: Optional[list[EnvVarDict]] - """Environment variables to be passed to the python module. Maximum limit is 100.""" + retrieve_result: Optional[bool] + """Whether to retrieve the results of the query job.""" - executor_image_uri: Optional[str] - """Required. The URI of a container image in Artifact Registry that will run the provided Python package. Vertex AI provides a wide range of executor images with pre-installed packages to meet users' various use cases. See the list of [pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers). You must use an image from this list.""" - package_uris: Optional[list[str]] - """Required. The Google Cloud Storage location of the Python package files which are the training program and its dependent packages. The maximum number of package URIs is 100.""" +CheckQueryJobRuntimeConfigOrDict = Union[ + CheckQueryJobRuntimeConfig, CheckQueryJobRuntimeConfigDict +] - python_module: Optional[str] - """Required. The Python module name to run after installing the packages.""" +class _CheckQueryJobRuntimeRequestParameters(_common.BaseModel): + """Parameters for async querying agent runtimes.""" -PythonPackageSpecOrDict = Union[PythonPackageSpec, PythonPackageSpecDict] + name: Optional[str] = Field(default=None, description="""Name of the query job.""") + config: Optional[CheckQueryJobRuntimeConfig] = Field( + default=None, description="""""" + ) -class WorkerPoolSpec(_common.BaseModel): - """Represents the spec of a worker pool in a job.""" +class _CheckQueryJobRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for async querying agent runtimes.""" - container_spec: Optional[ContainerSpec] = Field( - default=None, description="""The custom container task.""" - ) - disk_spec: Optional[DiskSpec] = Field(default=None, description="""Disk spec.""") - lustre_mounts: Optional[list[LustreMount]] = Field( - default=None, description="""Optional. List of Lustre mounts.""" + name: Optional[str] + """Name of the query job.""" + + config: Optional[CheckQueryJobRuntimeConfigDict] + """""" + + +_CheckQueryJobRuntimeRequestParametersOrDict = Union[ + _CheckQueryJobRuntimeRequestParameters, _CheckQueryJobRuntimeRequestParametersDict +] + + +class CheckQueryJobResult(_common.BaseModel): + """Result of checking a query job.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - machine_spec: Optional[MachineSpec] = Field( - default=None, - description="""Optional. Immutable. The specification of a single machine.""", + operation_name: Optional[str] = Field( + default=None, description="""Name of the agent runtime operation.""" ) - nfs_mounts: Optional[list[NfsMount]] = Field( - default=None, description="""Optional. List of NFS mount spec.""" + output_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the output file.""" ) - python_package_spec: Optional[PythonPackageSpec] = Field( - default=None, description="""The Python packaged task.""" + status: Optional[str] = Field( + default=None, description="""Status of the operation.""" ) - replica_count: Optional[int] = Field( - default=None, - description="""Optional. The number of worker replicas to use for this worker pool.""", + result: Optional[str] = Field( + default=None, description="""JSON result of the operation.""" ) -class WorkerPoolSpecDict(TypedDict, total=False): - """Represents the spec of a worker pool in a job.""" - - container_spec: Optional[ContainerSpecDict] - """The custom container task.""" - - disk_spec: Optional[DiskSpecDict] - """Disk spec.""" +class CheckQueryJobResultDict(TypedDict, total=False): + """Result of checking a query job.""" - lustre_mounts: Optional[list[LustreMountDict]] - """Optional. List of Lustre mounts.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - machine_spec: Optional[MachineSpecDict] - """Optional. Immutable. The specification of a single machine.""" + operation_name: Optional[str] + """Name of the agent runtime operation.""" - nfs_mounts: Optional[list[NfsMountDict]] - """Optional. List of NFS mount spec.""" + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" - python_package_spec: Optional[PythonPackageSpecDict] - """The Python packaged task.""" + status: Optional[str] + """Status of the operation.""" - replica_count: Optional[int] - """Optional. The number of worker replicas to use for this worker pool.""" + result: Optional[str] + """JSON result of the operation.""" -WorkerPoolSpecOrDict = Union[WorkerPoolSpec, WorkerPoolSpecDict] +CheckQueryJobResultOrDict = Union[CheckQueryJobResult, CheckQueryJobResultDict] -class CustomJobSpec(_common.BaseModel): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +class _RunQueryJobRuntimeConfig(_common.BaseModel): + """Config for running a query job on an agent runtime.""" - base_output_directory: Optional[genai_types.GcsDestination] = Field( - default=None, - description="""The Cloud Storage location to store the output of this CustomJob or HyperparameterTuningJob. For HyperparameterTuningJob, the baseOutputDirectory of each child CustomJob backing a Trial is set to a subdirectory of name id under its parent HyperparameterTuningJob's baseOutputDirectory. The following Vertex AI environment variables will be passed to containers or python modules when this field is set: For CustomJob: * AIP_MODEL_DIR = `/model/` * AIP_CHECKPOINT_DIR = `/checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `/logs/` For CustomJob backing a Trial of HyperparameterTuningJob: * AIP_MODEL_DIR = `//model/` * AIP_CHECKPOINT_DIR = `//checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `//logs/`""", - ) - enable_dashboard_access: Optional[bool] = Field( - default=None, - description="""Optional. Whether you want Vertex AI to enable access to the customized dashboard in training chief container. If set to `true`, you can access the dashboard at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""", - ) - enable_web_access: Optional[bool] = Field( - default=None, - description="""Optional. Whether you want Vertex AI to enable [interactive shell access](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) to training containers. If set to `true`, you can access interactive shells at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""", - ) - experiment: Optional[str] = Field( - default=None, - description="""Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`""", - ) - experiment_run: Optional[str] = Field( - default=None, - description="""Optional. The Experiment Run associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}-{experiment-run-name}`""", - ) - models: Optional[list[str]] = Field( - default=None, - description="""Optional. The name of the Model resources for which to generate a mapping to artifact URIs. Applicable only to some of the Google-provided custom jobs. Format: `projects/{project}/locations/{location}/models/{model}` In order to retrieve a specific version of the model, also provide the version ID or version alias. Example: `projects/{project}/locations/{location}/models/{model}@2` or `projects/{project}/locations/{location}/models/{model}@golden` If no version ID or alias is specified, the "default" version will be returned. The "default" version alias is created for the first version of the model, and can be moved to other versions later on. There will be exactly one default version.""", - ) - network: Optional[str] = Field( - default=None, - description="""Optional. The full name of the Compute Engine [network](/compute/docs/networks-and-firewalls#networks) to which the Job should be peered. For example, `projects/12345/global/networks/myVPC`. [Format](/compute/docs/reference/rest/v1/networks/insert) is of the form `projects/{project}/global/networks/{network}`. Where {project} is a project number, as in `12345`, and {network} is a network name. To specify this field, you must have already [configured VPC Network Peering for Vertex AI](https://cloud.google.com/vertex-ai/docs/general/vpc-peering). If this field is left unspecified, the job is not peered with any network.""", - ) - persistent_resource_id: Optional[str] = Field( - default=None, - description="""Optional. The ID of the PersistentResource in the same Project and Location which to run If this is specified, the job will be run on existing machines held by the PersistentResource instead of on-demand short-live machines. The network and CMEK configs on the job should be consistent with those on the PersistentResource, otherwise, the job will be rejected.""", - ) - protected_artifact_location_id: Optional[str] = Field( - default=None, - description="""The ID of the location to store protected artifacts. e.g. us-central1. Populate only when the location is different than CustomJob location. List of supported locations: https://cloud.google.com/vertex-ai/docs/general/locations""", - ) - psc_interface_config: Optional[PscInterfaceConfig] = Field( - default=None, description="""Optional. Configuration for PSC-I for CustomJob.""" - ) - reserved_ip_ranges: Optional[list[str]] = Field( - default=None, - description="""Optional. A list of names for the reserved ip ranges under the VPC network that can be used for this job. If set, we will deploy the job within the provided ip ranges. Otherwise, the job will be deployed to any ip ranges under the provided VPC network. Example: ['vertex-ai-ip-range'].""", - ) - scheduling: Optional[Scheduling] = Field( - default=None, description="""Scheduling options for a CustomJob.""" - ) - service_account: Optional[str] = Field( - default=None, - description="""Specifies the service account for workload run-as account. Users submitting jobs must have act-as permission on this run-as account. If unspecified, the [Vertex AI Custom Code Service Agent](https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents) for the CustomJob's project is used.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - tensorboard: Optional[str] = Field( - default=None, - description="""Optional. The name of a Vertex AI Tensorboard resource to which this CustomJob will upload Tensorboard logs. Format: `projects/{project}/locations/{location}/tensorboards/{tensorboard}`""", + input_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the input file.""" ) - worker_pool_specs: Optional[list[WorkerPoolSpec]] = Field( - default=None, - description="""Required. The spec of the worker pools including machine type and Docker image. All worker pools except the first one are optional and can be skipped by providing an empty value.""", + output_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the output file.""" ) -class CustomJobSpecDict(TypedDict, total=False): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +class _RunQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for running a query job on an agent runtime.""" - base_output_directory: Optional[genai_types.GcsDestination] - """The Cloud Storage location to store the output of this CustomJob or HyperparameterTuningJob. For HyperparameterTuningJob, the baseOutputDirectory of each child CustomJob backing a Trial is set to a subdirectory of name id under its parent HyperparameterTuningJob's baseOutputDirectory. The following Vertex AI environment variables will be passed to containers or python modules when this field is set: For CustomJob: * AIP_MODEL_DIR = `/model/` * AIP_CHECKPOINT_DIR = `/checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `/logs/` For CustomJob backing a Trial of HyperparameterTuningJob: * AIP_MODEL_DIR = `//model/` * AIP_CHECKPOINT_DIR = `//checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `//logs/`""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - enable_dashboard_access: Optional[bool] - """Optional. Whether you want Vertex AI to enable access to the customized dashboard in training chief container. If set to `true`, you can access the dashboard at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""" + input_gcs_uri: Optional[str] + """The GCS URI of the input file.""" - enable_web_access: Optional[bool] - """Optional. Whether you want Vertex AI to enable [interactive shell access](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) to training containers. If set to `true`, you can access interactive shells at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""" + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" - experiment: Optional[str] - """Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`""" - experiment_run: Optional[str] - """Optional. The Experiment Run associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}-{experiment-run-name}`""" +_RunQueryJobRuntimeConfigOrDict = Union[ + _RunQueryJobRuntimeConfig, _RunQueryJobRuntimeConfigDict +] - models: Optional[list[str]] - """Optional. The name of the Model resources for which to generate a mapping to artifact URIs. Applicable only to some of the Google-provided custom jobs. Format: `projects/{project}/locations/{location}/models/{model}` In order to retrieve a specific version of the model, also provide the version ID or version alias. Example: `projects/{project}/locations/{location}/models/{model}@2` or `projects/{project}/locations/{location}/models/{model}@golden` If no version ID or alias is specified, the "default" version will be returned. The "default" version alias is created for the first version of the model, and can be moved to other versions later on. There will be exactly one default version.""" - network: Optional[str] - """Optional. The full name of the Compute Engine [network](/compute/docs/networks-and-firewalls#networks) to which the Job should be peered. For example, `projects/12345/global/networks/myVPC`. [Format](/compute/docs/reference/rest/v1/networks/insert) is of the form `projects/{project}/global/networks/{network}`. Where {project} is a project number, as in `12345`, and {network} is a network name. To specify this field, you must have already [configured VPC Network Peering for Vertex AI](https://cloud.google.com/vertex-ai/docs/general/vpc-peering). If this field is left unspecified, the job is not peered with any network.""" +class _RunQueryJobRuntimeRequestParameters(_common.BaseModel): + """Parameters for running a query job on an agent runtime.""" - persistent_resource_id: Optional[str] - """Optional. The ID of the PersistentResource in the same Project and Location which to run If this is specified, the job will be run on existing machines held by the PersistentResource instead of on-demand short-live machines. The network and CMEK configs on the job should be consistent with those on the PersistentResource, otherwise, the job will be rejected.""" - - protected_artifact_location_id: Optional[str] - """The ID of the location to store protected artifacts. e.g. us-central1. Populate only when the location is different than CustomJob location. List of supported locations: https://cloud.google.com/vertex-ai/docs/general/locations""" - - psc_interface_config: Optional[PscInterfaceConfigDict] - """Optional. Configuration for PSC-I for CustomJob.""" - - reserved_ip_ranges: Optional[list[str]] - """Optional. A list of names for the reserved ip ranges under the VPC network that can be used for this job. If set, we will deploy the job within the provided ip ranges. Otherwise, the job will be deployed to any ip ranges under the provided VPC network. Example: ['vertex-ai-ip-range'].""" - - scheduling: Optional[SchedulingDict] - """Scheduling options for a CustomJob.""" - - service_account: Optional[str] - """Specifies the service account for workload run-as account. Users submitting jobs must have act-as permission on this run-as account. If unspecified, the [Vertex AI Custom Code Service Agent](https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents) for the CustomJob's project is used.""" - - tensorboard: Optional[str] - """Optional. The name of a Vertex AI Tensorboard resource to which this CustomJob will upload Tensorboard logs. Format: `projects/{project}/locations/{location}/tensorboards/{tensorboard}`""" - - worker_pool_specs: Optional[list[WorkerPoolSpecDict]] - """Required. The spec of the worker pools including machine type and Docker image. All worker pools except the first one are optional and can be skipped by providing an empty value.""" - - -CustomJobSpecOrDict = Union[CustomJobSpec, CustomJobSpecDict] - - -class CustomJob(_common.BaseModel): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - - display_name: Optional[str] = Field( - default=None, - description="""Required. The display name of the CustomJob. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", - ) - job_spec: Optional[CustomJobSpec] = Field( - default=None, description="""Required. Job spec.""" - ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, - description="""Customer-managed encryption key options for a CustomJob. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key.""", - ) - state: Optional[genai_types.JobState] = Field( - default=None, description="""Output only. The detailed state of the job.""" - ) - error: Optional[genai_types.GoogleRpcStatus] = Field( - default=None, - description="""Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`.""", - ) - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Time when the CustomJob was created.""", - ) - end_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Time when the CustomJob entered any of the following states: `JOB_STATE_SUCCEEDED`, `JOB_STATE_FAILED`, `JOB_STATE_CANCELLED`.""", - ) - labels: Optional[dict[str, str]] = Field( - default=None, - description="""The labels with user-defined metadata to organize CustomJobs. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", - ) name: Optional[str] = Field( - default=None, description="""Output only. Resource name of a CustomJob.""" - ) - satisfies_pzi: Optional[bool] = Field( - default=None, description="""Output only. Reserved for future use.""" - ) - satisfies_pzs: Optional[bool] = Field( - default=None, description="""Output only. Reserved for future use.""" - ) - start_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state.""", - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Time when the CustomJob was most recently updated.""", + default=None, description="""Name of the agent runtime.""" ) - web_access_uris: Optional[dict[str, str]] = Field( - default=None, - description="""Output only. URIs for accessing [interactive shells](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) (one URI for each training node). Only available if job_spec.enable_web_access is `true`. The keys are names of each node in the training job; for example, `workerpool0-0` for the primary node, `workerpool1-0` for the first node in the second worker pool, and `workerpool1-1` for the second node in the second worker pool. The values are the URIs for each node's interactive shell.""", + config: Optional[_RunQueryJobRuntimeConfig] = Field( + default=None, description="""""" ) -class CustomJobDict(TypedDict, total=False): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - - display_name: Optional[str] - """Required. The display name of the CustomJob. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" - - job_spec: Optional[CustomJobSpecDict] - """Required. Job spec.""" - - encryption_spec: Optional[genai_types.EncryptionSpec] - """Customer-managed encryption key options for a CustomJob. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key.""" - - state: Optional[genai_types.JobState] - """Output only. The detailed state of the job.""" +class _RunQueryJobRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for running a query job on an agent runtime.""" - error: Optional[genai_types.GoogleRpcStatus] - """Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`.""" + name: Optional[str] + """Name of the agent runtime.""" - create_time: Optional[datetime.datetime] - """Output only. Time when the CustomJob was created.""" + config: Optional[_RunQueryJobRuntimeConfigDict] + """""" - end_time: Optional[datetime.datetime] - """Output only. Time when the CustomJob entered any of the following states: `JOB_STATE_SUCCEEDED`, `JOB_STATE_FAILED`, `JOB_STATE_CANCELLED`.""" - labels: Optional[dict[str, str]] - """The labels with user-defined metadata to organize CustomJobs. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""" +_RunQueryJobRuntimeRequestParametersOrDict = Union[ + _RunQueryJobRuntimeRequestParameters, _RunQueryJobRuntimeRequestParametersDict +] - name: Optional[str] - """Output only. Resource name of a CustomJob.""" - satisfies_pzi: Optional[bool] - """Output only. Reserved for future use.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent( + _common.BaseModel +): + """The conversation source event for generating memories.""" - satisfies_pzs: Optional[bool] - """Output only. Reserved for future use.""" + content: Optional[genai_types.Content] = Field( + default=None, description="""Required. Represents the content of the event.""" + ) - start_time: Optional[datetime.datetime] - """Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state.""" - update_time: Optional[datetime.datetime] - """Output only. Time when the CustomJob was most recently updated.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict( + TypedDict, total=False +): + """The conversation source event for generating memories.""" - web_access_uris: Optional[dict[str, str]] - """Output only. URIs for accessing [interactive shells](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) (one URI for each training node). Only available if job_spec.enable_web_access is `true`. The keys are names of each node in the training job; for example, `workerpool0-0` for the primary node, `workerpool1-0` for the first node in the second worker pool, and `workerpool1-1` for the second node in the second worker pool. The values are the URIs for each node's interactive shell.""" + content: Optional[genai_types.Content] + """Required. Represents the content of the event.""" -CustomJobOrDict = Union[CustomJob, CustomJobDict] +MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventOrDict = ( + Union[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent, + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict, + ] +) -class VertexBaseConfig(_common.BaseModel): - """Base config for Vertex AI.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource( + _common.BaseModel +): + """A conversation source for the example. This is similar to `DirectContentsSource`.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + events: Optional[ + list[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent + ] + ] = Field( + default=None, + description="""Optional. Represents the input conversation events for the example.""", ) -class VertexBaseConfigDict(TypedDict, total=False): - """Base config for Vertex AI.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict( + TypedDict, total=False +): + """A conversation source for the example. This is similar to `DirectContentsSource`.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + events: Optional[ + list[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict + ] + ] + """Optional. Represents the input conversation events for the example.""" -VertexBaseConfigOrDict = Union[VertexBaseConfig, VertexBaseConfigDict] +MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceOrDict = Union[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource, + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict, +] -class _CustomJobParameters(_common.BaseModel): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +class MemoryTopicId(_common.BaseModel): + """The topic ID for a memory.""" - custom_job: Optional[CustomJob] = Field(default=None, description="""""") - config: Optional[VertexBaseConfig] = Field(default=None, description="""""") + custom_memory_topic_label: Optional[str] = Field( + default=None, + description="""Optional. Represents the custom memory topic label.""", + ) + managed_memory_topic: Optional[ManagedTopicEnum] = Field( + default=None, description="""Optional. Represents the managed memory topic.""" + ) -class _CustomJobParametersDict(TypedDict, total=False): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +class MemoryTopicIdDict(TypedDict, total=False): + """The topic ID for a memory.""" - custom_job: Optional[CustomJobDict] - """""" + custom_memory_topic_label: Optional[str] + """Optional. Represents the custom memory topic label.""" - config: Optional[VertexBaseConfigDict] - """""" + managed_memory_topic: Optional[ManagedTopicEnum] + """Optional. Represents the managed memory topic.""" -_CustomJobParametersOrDict = Union[_CustomJobParameters, _CustomJobParametersDict] +MemoryTopicIdOrDict = Union[MemoryTopicId, MemoryTopicIdDict] -class _GetCustomJobParameters(_common.BaseModel): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory( + _common.BaseModel +): + """A memory generated by the operation.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[VertexBaseConfig] = Field(default=None, description="""""") + fact: Optional[str] = Field( + default=None, + description="""Required. Represents the fact to generate a memory from.""", + ) + topics: Optional[list[MemoryTopicId]] = Field( + default=None, + description="""Optional. Represents the list of topics that the memory should be associated with. For example, use `custom_memory_topic_label = "jargon"` if the extracted memory is an example of memory extraction for the custom topic `jargon`.""", + ) -class _GetCustomJobParametersDict(TypedDict, total=False): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict( + TypedDict, total=False +): + """A memory generated by the operation.""" - name: Optional[str] - """""" + fact: Optional[str] + """Required. Represents the fact to generate a memory from.""" - config: Optional[VertexBaseConfigDict] - """""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. Represents the list of topics that the memory should be associated with. For example, use `custom_memory_topic_label = "jargon"` if the extracted memory is an example of memory extraction for the custom topic `jargon`.""" -_GetCustomJobParametersOrDict = Union[ - _GetCustomJobParameters, _GetCustomJobParametersDict +MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryOrDict = Union[ + MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory, + MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict, ] -class CancelQueryJobAgentEngineConfig(_common.BaseModel): - """Config for canceling async querying agent engines.""" +class MemoryBankCustomizationConfigGenerateMemoriesExample(_common.BaseModel): + """An example of how to generate memories for a particular scope.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - operation_name: Optional[str] = Field( + conversation_source: Optional[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource + ] = Field(default=None, description="""A conversation source for the example.""") + generated_memories: Optional[ + list[MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory] + ] = Field( default=None, - description="""Name of the longrunning operation returned from run_query_job.""", + description="""Optional. Represents the memories that are expected to be generated from the input conversation. An empty list indicates that no memories are expected to be generated for the input conversation.""", ) -class CancelQueryJobAgentEngineConfigDict(TypedDict, total=False): - """Config for canceling async querying agent engines.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleDict(TypedDict, total=False): + """An example of how to generate memories for a particular scope.""" - operation_name: Optional[str] - """Name of the longrunning operation returned from run_query_job.""" + conversation_source: Optional[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict + ] + """A conversation source for the example.""" + generated_memories: Optional[ + list[MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict] + ] + """Optional. Represents the memories that are expected to be generated from the input conversation. An empty list indicates that no memories are expected to be generated for the input conversation.""" -CancelQueryJobAgentEngineConfigOrDict = Union[ - CancelQueryJobAgentEngineConfig, CancelQueryJobAgentEngineConfigDict + +MemoryBankCustomizationConfigGenerateMemoriesExampleOrDict = Union[ + MemoryBankCustomizationConfigGenerateMemoriesExample, + MemoryBankCustomizationConfigGenerateMemoriesExampleDict, ] -class _CancelQueryJobAgentEngineRequestParameters(_common.BaseModel): - """Parameters for canceling async querying agent engines.""" +class MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic(_common.BaseModel): + """A custom memory topic defined by the developer.""" - name: Optional[str] = Field( - default=None, description="""Name of the reasoning engine resource.""" + label: Optional[str] = Field( + default=None, description="""Required. Represents the label of the topic.""" ) - config: Optional[CancelQueryJobAgentEngineConfig] = Field( - default=None, description="""""" + description: Optional[str] = Field( + default=None, + description="""Required. Represents the description of the memory topic. This should explain what information should be extracted for this topic.""", ) -class _CancelQueryJobAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for canceling async querying agent engines.""" +class MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict( + TypedDict, total=False +): + """A custom memory topic defined by the developer.""" - name: Optional[str] - """Name of the reasoning engine resource.""" + label: Optional[str] + """Required. Represents the label of the topic.""" - config: Optional[CancelQueryJobAgentEngineConfigDict] - """""" + description: Optional[str] + """Required. Represents the description of the memory topic. This should explain what information should be extracted for this topic.""" -_CancelQueryJobAgentEngineRequestParametersOrDict = Union[ - _CancelQueryJobAgentEngineRequestParameters, - _CancelQueryJobAgentEngineRequestParametersDict, +MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicOrDict = Union[ + MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic, + MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict, ] -class CancelQueryJobResult(_common.BaseModel): - """Result of canceling a query job.""" +class MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic(_common.BaseModel): + """A managed memory topic defined by the system.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + managed_topic_enum: Optional[ManagedTopicEnum] = Field( + default=None, description="""Required. Represents the managed topic.""" ) -class CancelQueryJobResultDict(TypedDict, total=False): - """Result of canceling a query job.""" +class MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict( + TypedDict, total=False +): + """A managed memory topic defined by the system.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + managed_topic_enum: Optional[ManagedTopicEnum] + """Required. Represents the managed topic.""" -CancelQueryJobResultOrDict = Union[CancelQueryJobResult, CancelQueryJobResultDict] +MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicOrDict = Union[ + MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic, + MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict, +] -class CheckQueryJobAgentEngineConfig(_common.BaseModel): - """Config for async querying agent engines.""" +class MemoryBankCustomizationConfigMemoryTopic(_common.BaseModel): + """A topic of information that should be extracted from conversations and stored as memories.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + custom_memory_topic: Optional[ + MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic + ] = Field( + default=None, description="""A custom memory topic defined by the developer.""" ) - retrieve_result: Optional[bool] = Field( - default=None, - description="""Whether to retrieve the results of the query job.""", + managed_memory_topic: Optional[ + MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic + ] = Field( + default=None, description="""A managed memory topic defined by Memory Bank.""" ) -class CheckQueryJobAgentEngineConfigDict(TypedDict, total=False): - """Config for async querying agent engines.""" +class MemoryBankCustomizationConfigMemoryTopicDict(TypedDict, total=False): + """A topic of information that should be extracted from conversations and stored as memories.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + custom_memory_topic: Optional[ + MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict + ] + """A custom memory topic defined by the developer.""" - retrieve_result: Optional[bool] - """Whether to retrieve the results of the query job.""" + managed_memory_topic: Optional[ + MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict + ] + """A managed memory topic defined by Memory Bank.""" -CheckQueryJobAgentEngineConfigOrDict = Union[ - CheckQueryJobAgentEngineConfig, CheckQueryJobAgentEngineConfigDict +MemoryBankCustomizationConfigMemoryTopicOrDict = Union[ + MemoryBankCustomizationConfigMemoryTopic, + MemoryBankCustomizationConfigMemoryTopicDict, ] -class _CheckQueryJobAgentEngineRequestParameters(_common.BaseModel): - """Parameters for async querying agent engines.""" +class MemoryBankCustomizationConfigConsolidationConfig(_common.BaseModel): + """Represents configuration for customizing how memories are consolidated.""" - name: Optional[str] = Field(default=None, description="""Name of the query job.""") - config: Optional[CheckQueryJobAgentEngineConfig] = Field( - default=None, description="""""" + revisions_per_candidate_count: Optional[int] = Field( + default=None, + description="""Optional. Represents the maximum number of revisions to consider for each candidate memory. If not set, then the default value (1) will be used, which means that only the latest revision will be considered.""", ) -class _CheckQueryJobAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for async querying agent engines.""" - - name: Optional[str] - """Name of the query job.""" +class MemoryBankCustomizationConfigConsolidationConfigDict(TypedDict, total=False): + """Represents configuration for customizing how memories are consolidated.""" - config: Optional[CheckQueryJobAgentEngineConfigDict] - """""" + revisions_per_candidate_count: Optional[int] + """Optional. Represents the maximum number of revisions to consider for each candidate memory. If not set, then the default value (1) will be used, which means that only the latest revision will be considered.""" -_CheckQueryJobAgentEngineRequestParametersOrDict = Union[ - _CheckQueryJobAgentEngineRequestParameters, - _CheckQueryJobAgentEngineRequestParametersDict, +MemoryBankCustomizationConfigConsolidationConfigOrDict = Union[ + MemoryBankCustomizationConfigConsolidationConfig, + MemoryBankCustomizationConfigConsolidationConfigDict, ] -class CheckQueryJobResult(_common.BaseModel): - """Result of checking a query job.""" +class MemoryBankCustomizationConfig(_common.BaseModel): + """Represents configuration for organizing natural language memories for a particular scope.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + enable_third_person_memories: Optional[bool] = Field( + default=None, + description="""Optional. Indicates whether the memories will be generated in the third person (i.e. "The user generates memories with Memory Bank."). By default, the memories will be generated in the first person (i.e. "I generate memories with Memory Bank.")""", ) - operation_name: Optional[str] = Field( - default=None, description="""Name of the agent engine operation.""" + generate_memories_examples: Optional[ + list[MemoryBankCustomizationConfigGenerateMemoriesExample] + ] = Field( + default=None, + description="""Optional. Provides examples of how to generate memories for a particular scope.""", ) - output_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the output file.""" + memory_topics: Optional[list[MemoryBankCustomizationConfigMemoryTopic]] = Field( + default=None, + description="""Optional. Represents topics of information that should be extracted from conversations and stored as memories. If not set, then Memory Bank's default topics will be used.""", ) - status: Optional[str] = Field( - default=None, description="""Status of the operation.""" + scope_keys: Optional[list[str]] = Field( + default=None, + description="""Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""", ) - result: Optional[str] = Field( - default=None, description="""JSON result of the operation.""" + consolidation_config: Optional[MemoryBankCustomizationConfigConsolidationConfig] = ( + Field( + default=None, + description="""Optional. Represents configuration for customizing how memories are consolidated together.""", + ) + ) + disable_natural_language_memories: Optional[bool] = Field( + default=None, + description="""Optional. Indicates whether natural language memory generation should be disabled for all requests. By default, natural language memory generation is enabled. Set this to `true` when you only want to generate structured memories.""", ) -class CheckQueryJobResultDict(TypedDict, total=False): - """Result of checking a query job.""" +class MemoryBankCustomizationConfigDict(TypedDict, total=False): + """Represents configuration for organizing natural language memories for a particular scope.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + enable_third_person_memories: Optional[bool] + """Optional. Indicates whether the memories will be generated in the third person (i.e. "The user generates memories with Memory Bank."). By default, the memories will be generated in the first person (i.e. "I generate memories with Memory Bank.")""" - operation_name: Optional[str] - """Name of the agent engine operation.""" + generate_memories_examples: Optional[ + list[MemoryBankCustomizationConfigGenerateMemoriesExampleDict] + ] + """Optional. Provides examples of how to generate memories for a particular scope.""" - output_gcs_uri: Optional[str] - """The GCS URI of the output file.""" + memory_topics: Optional[list[MemoryBankCustomizationConfigMemoryTopicDict]] + """Optional. Represents topics of information that should be extracted from conversations and stored as memories. If not set, then Memory Bank's default topics will be used.""" - status: Optional[str] - """Status of the operation.""" + scope_keys: Optional[list[str]] + """Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""" - result: Optional[str] - """JSON result of the operation.""" + consolidation_config: Optional[MemoryBankCustomizationConfigConsolidationConfigDict] + """Optional. Represents configuration for customizing how memories are consolidated together.""" + disable_natural_language_memories: Optional[bool] + """Optional. Indicates whether natural language memory generation should be disabled for all requests. By default, natural language memory generation is enabled. Set this to `true` when you only want to generate structured memories.""" -CheckQueryJobResultOrDict = Union[CheckQueryJobResult, CheckQueryJobResultDict] +MemoryBankCustomizationConfigOrDict = Union[ + MemoryBankCustomizationConfig, MemoryBankCustomizationConfigDict +] -class _RunQueryJobAgentEngineConfig(_common.BaseModel): - """Config for running a query job on an agent engine.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" +class MemoryGenerationTriggerConfigGenerationTriggerRule(_common.BaseModel): + """Represents the active rule that determines when to flush the buffer.""" + + event_count: Optional[int] = Field( + default=None, + description="""Optional. Specifies to trigger generation when the event count reaches this limit.""", ) - input_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the input file.""" + fixed_interval: Optional[str] = Field( + default=None, + description="""Optional. Specifies to trigger generation at a fixed interval. The duration must have a minute-level granularity.""", ) - output_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the output file.""" + idle_duration: Optional[str] = Field( + default=None, + description="""Optional. Specifies to trigger generation if the stream is inactive for the specified duration after the most recent event. The duration must have a minute-level granularity.""", + ) + overlap_event_count: Optional[int] = Field( + default=None, + description="""Optional. Re-include the last N already-processed events in the next window.""", ) -class _RunQueryJobAgentEngineConfigDict(TypedDict, total=False): - """Config for running a query job on an agent engine.""" +class MemoryGenerationTriggerConfigGenerationTriggerRuleDict(TypedDict, total=False): + """Represents the active rule that determines when to flush the buffer.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + event_count: Optional[int] + """Optional. Specifies to trigger generation when the event count reaches this limit.""" - input_gcs_uri: Optional[str] - """The GCS URI of the input file.""" + fixed_interval: Optional[str] + """Optional. Specifies to trigger generation at a fixed interval. The duration must have a minute-level granularity.""" - output_gcs_uri: Optional[str] - """The GCS URI of the output file.""" + idle_duration: Optional[str] + """Optional. Specifies to trigger generation if the stream is inactive for the specified duration after the most recent event. The duration must have a minute-level granularity.""" + overlap_event_count: Optional[int] + """Optional. Re-include the last N already-processed events in the next window.""" -_RunQueryJobAgentEngineConfigOrDict = Union[ - _RunQueryJobAgentEngineConfig, _RunQueryJobAgentEngineConfigDict + +MemoryGenerationTriggerConfigGenerationTriggerRuleOrDict = Union[ + MemoryGenerationTriggerConfigGenerationTriggerRule, + MemoryGenerationTriggerConfigGenerationTriggerRuleDict, ] -class _RunQueryJobAgentEngineRequestParameters(_common.BaseModel): - """Parameters for running a query job on an agent engine.""" +class MemoryGenerationTriggerConfig(_common.BaseModel): + """The configuration for triggering memory generation for ingested events.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[_RunQueryJobAgentEngineConfig] = Field( - default=None, description="""""" + generation_rule: Optional[MemoryGenerationTriggerConfigGenerationTriggerRule] = ( + Field( + default=None, + description="""Optional. Represents the active rule that determines when to flush the buffer. If not set, then the stream will be force flushed immediately.""", + ) ) -class _RunQueryJobAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for running a query job on an agent engine.""" - - name: Optional[str] - """Name of the agent engine.""" +class MemoryGenerationTriggerConfigDict(TypedDict, total=False): + """The configuration for triggering memory generation for ingested events.""" - config: Optional[_RunQueryJobAgentEngineConfigDict] - """""" + generation_rule: Optional[MemoryGenerationTriggerConfigGenerationTriggerRuleDict] + """Optional. Represents the active rule that determines when to flush the buffer. If not set, then the stream will be force flushed immediately.""" -_RunQueryJobAgentEngineRequestParametersOrDict = Union[ - _RunQueryJobAgentEngineRequestParameters, - _RunQueryJobAgentEngineRequestParametersDict, +MemoryGenerationTriggerConfigOrDict = Union[ + MemoryGenerationTriggerConfig, MemoryGenerationTriggerConfigDict ] -class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent( - _common.BaseModel -): - """The conversation source event for generating memories.""" +class ReasoningEngineContextSpecMemoryBankConfigGenerationConfig(_common.BaseModel): + """Configuration for how to generate memories.""" - content: Optional[genai_types.Content] = Field( - default=None, description="""Required. Represents the content of the event.""" + model: Optional[str] = Field( + default=None, + description="""Optional. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", + ) + generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + default=None, + description="""Optional. Specifies the default trigger configuration for generating memories using `IngestEvents`.""", ) -class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict( +class ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict( TypedDict, total=False ): - """The conversation source event for generating memories.""" + """Configuration for how to generate memories.""" - content: Optional[genai_types.Content] - """Required. Represents the content of the event.""" + model: Optional[str] + """Optional. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] + """Optional. Specifies the default trigger configuration for generating memories using `IngestEvents`.""" -MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventOrDict = ( - Union[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent, - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict, - ] -) +ReasoningEngineContextSpecMemoryBankConfigGenerationConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigGenerationConfig, + ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict, +] -class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource( + +class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig( _common.BaseModel ): - """A conversation source for the example. This is similar to `DirectContentsSource`.""" + """Configuration for how to perform similarity search on memories.""" - events: Optional[ - list[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent - ] - ] = Field( + embedding_model: Optional[str] = Field( default=None, - description="""Optional. Represents the input conversation events for the example.""", + description="""Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", ) -class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict( +class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict( TypedDict, total=False ): - """A conversation source for the example. This is similar to `DirectContentsSource`.""" + """Configuration for how to perform similarity search on memories.""" - events: Optional[ - list[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict - ] - ] - """Optional. Represents the input conversation events for the example.""" + embedding_model: Optional[str] + """Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" -MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceOrDict = Union[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource, - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict, +ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig, + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict, ] -class MemoryTopicId(_common.BaseModel): - """The topic ID for a memory.""" - - custom_memory_topic_label: Optional[str] = Field( - default=None, - description="""Optional. Represents the custom memory topic label.""", - ) - managed_memory_topic: Optional[ManagedTopicEnum] = Field( - default=None, description="""Optional. Represents the managed memory topic.""" - ) - - -class MemoryTopicIdDict(TypedDict, total=False): - """The topic ID for a memory.""" - - custom_memory_topic_label: Optional[str] - """Optional. Represents the custom memory topic label.""" - - managed_memory_topic: Optional[ManagedTopicEnum] - """Optional. Represents the managed memory topic.""" - - -MemoryTopicIdOrDict = Union[MemoryTopicId, MemoryTopicIdDict] - - -class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory( +class ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig( _common.BaseModel ): - """A memory generated by the operation.""" + """Configuration for TTL of the memories in the Memory Bank based on the action that created or updated the memory.""" - fact: Optional[str] = Field( + create_ttl: Optional[str] = Field( default=None, - description="""Required. Represents the fact to generate a memory from.""", + description="""Optional. The TTL duration for memories uploaded via CreateMemory.""", ) - topics: Optional[list[MemoryTopicId]] = Field( + generate_created_ttl: Optional[str] = Field( default=None, - description="""Optional. Represents the list of topics that the memory should be associated with. For example, use `custom_memory_topic_label = "jargon"` if the extracted memory is an example of memory extraction for the custom topic `jargon`.""", + description="""Optional. The TTL duration for memories newly generated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.CREATED).""", + ) + generate_updated_ttl: Optional[str] = Field( + default=None, + description="""Optional. The TTL duration for memories updated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.UPDATED). In the case of an UPDATE action, the `expire_time` of the existing memory will be updated to the new value (now + TTL).""", ) -class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict( +class ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict( TypedDict, total=False ): - """A memory generated by the operation.""" + """Configuration for TTL of the memories in the Memory Bank based on the action that created or updated the memory.""" - fact: Optional[str] - """Required. Represents the fact to generate a memory from.""" + create_ttl: Optional[str] + """Optional. The TTL duration for memories uploaded via CreateMemory.""" - topics: Optional[list[MemoryTopicIdDict]] - """Optional. Represents the list of topics that the memory should be associated with. For example, use `custom_memory_topic_label = "jargon"` if the extracted memory is an example of memory extraction for the custom topic `jargon`.""" + generate_created_ttl: Optional[str] + """Optional. The TTL duration for memories newly generated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.CREATED).""" + generate_updated_ttl: Optional[str] + """Optional. The TTL duration for memories updated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.UPDATED). In the case of an UPDATE action, the `expire_time` of the existing memory will be updated to the new value (now + TTL).""" -MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryOrDict = Union[ - MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory, - MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict, + +ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig, + ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict, ] -class MemoryBankCustomizationConfigGenerateMemoriesExample(_common.BaseModel): - """An example of how to generate memories for a particular scope.""" +class ReasoningEngineContextSpecMemoryBankConfigTtlConfig(_common.BaseModel): + """Configuration for automatically setting the TTL ("time-to-live") of the memories in the Memory Bank.""" - conversation_source: Optional[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource - ] = Field(default=None, description="""A conversation source for the example.""") - generated_memories: Optional[ - list[MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory] + default_ttl: Optional[str] = Field( + default=None, + description="""Optional. The default TTL duration of the memories in the Memory Bank. This applies to all operations that create or update a memory.""", + ) + granular_ttl_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig ] = Field( default=None, - description="""Optional. Represents the memories that are expected to be generated from the input conversation. An empty list indicates that no memories are expected to be generated for the input conversation.""", + description="""Optional. The granular TTL configuration of the memories in the Memory Bank.""", + ) + memory_revision_default_ttl: Optional[str] = Field( + default=None, + description="""Optional. The default TTL duration of the memory revisions in the Memory Bank. This applies to all operations that create a memory revision. If not set, a default TTL of 365 days will be used.""", ) -class MemoryBankCustomizationConfigGenerateMemoriesExampleDict(TypedDict, total=False): - """An example of how to generate memories for a particular scope.""" +class ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict(TypedDict, total=False): + """Configuration for automatically setting the TTL ("time-to-live") of the memories in the Memory Bank.""" - conversation_source: Optional[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict - ] - """A conversation source for the example.""" + default_ttl: Optional[str] + """Optional. The default TTL duration of the memories in the Memory Bank. This applies to all operations that create or update a memory.""" - generated_memories: Optional[ - list[MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict] + granular_ttl_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict ] - """Optional. Represents the memories that are expected to be generated from the input conversation. An empty list indicates that no memories are expected to be generated for the input conversation.""" + """Optional. The granular TTL configuration of the memories in the Memory Bank.""" + + memory_revision_default_ttl: Optional[str] + """Optional. The default TTL duration of the memory revisions in the Memory Bank. This applies to all operations that create a memory revision. If not set, a default TTL of 365 days will be used.""" -MemoryBankCustomizationConfigGenerateMemoriesExampleOrDict = Union[ - MemoryBankCustomizationConfigGenerateMemoriesExample, - MemoryBankCustomizationConfigGenerateMemoriesExampleDict, +ReasoningEngineContextSpecMemoryBankConfigTtlConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigTtlConfig, + ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict, ] -class MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic(_common.BaseModel): - """A custom memory topic defined by the developer.""" +class StructuredMemorySchemaConfig(_common.BaseModel): + """Represents the OpenAPI schema of the structured memories.""" - label: Optional[str] = Field( - default=None, description="""Required. Represents the label of the topic.""" + memory_schema: Optional[genai_types.Schema] = Field( + default=None, + description="""Required. Represents the OpenAPI schema of the structured memories.""", ) - description: Optional[str] = Field( + id: Optional[str] = Field( default=None, - description="""Required. Represents the description of the memory topic. This should explain what information should be extracted for this topic.""", + description="""Required. Represents the ID of the schema. Must be 1-63 characters, start with a lowercase letter, and consist of lowercase letters, numbers, and hyphens.""", + ) + memory_type: Optional[MemoryType] = Field( + default=None, + description="""Optional. Represents the type of the structured memories associated with the schema. If not set, then `STRUCTURED_PROFILE` will be used.""", ) -class MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict( - TypedDict, total=False -): - """A custom memory topic defined by the developer.""" +class StructuredMemorySchemaConfigDict(TypedDict, total=False): + """Represents the OpenAPI schema of the structured memories.""" - label: Optional[str] - """Required. Represents the label of the topic.""" + memory_schema: Optional[genai_types.Schema] + """Required. Represents the OpenAPI schema of the structured memories.""" - description: Optional[str] - """Required. Represents the description of the memory topic. This should explain what information should be extracted for this topic.""" + id: Optional[str] + """Required. Represents the ID of the schema. Must be 1-63 characters, start with a lowercase letter, and consist of lowercase letters, numbers, and hyphens.""" + + memory_type: Optional[MemoryType] + """Optional. Represents the type of the structured memories associated with the schema. If not set, then `STRUCTURED_PROFILE` will be used.""" -MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicOrDict = Union[ - MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic, - MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict, +StructuredMemorySchemaConfigOrDict = Union[ + StructuredMemorySchemaConfig, StructuredMemorySchemaConfigDict ] -class MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic(_common.BaseModel): - """A managed memory topic defined by the system.""" +class StructuredMemoryConfig(_common.BaseModel): + """Configuration for organizing structured memories within a scope.""" - managed_topic_enum: Optional[ManagedTopicEnum] = Field( - default=None, description="""Required. Represents the managed topic.""" + schema_configs: Optional[list[StructuredMemorySchemaConfig]] = Field( + default=None, + description="""Optional. Represents configuration of the structured memories' schemas.""", + ) + scope_keys: Optional[list[str]] = Field( + default=None, + description="""Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""", ) -class MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict( - TypedDict, total=False -): - """A managed memory topic defined by the system.""" +class StructuredMemoryConfigDict(TypedDict, total=False): + """Configuration for organizing structured memories within a scope.""" - managed_topic_enum: Optional[ManagedTopicEnum] - """Required. Represents the managed topic.""" + schema_configs: Optional[list[StructuredMemorySchemaConfigDict]] + """Optional. Represents configuration of the structured memories' schemas.""" + scope_keys: Optional[list[str]] + """Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""" -MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicOrDict = Union[ - MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic, - MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict, -] +StructuredMemoryConfigOrDict = Union[StructuredMemoryConfig, StructuredMemoryConfigDict] -class MemoryBankCustomizationConfigMemoryTopic(_common.BaseModel): - """A topic of information that should be extracted from conversations and stored as memories.""" - custom_memory_topic: Optional[ - MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic +class ReasoningEngineContextSpecMemoryBankConfig(_common.BaseModel): + """Specification for a Memory Bank.""" + + customization_configs: Optional[list[MemoryBankCustomizationConfig]] = Field( + default=None, + description="""Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""", + ) + disable_memory_revisions: Optional[bool] = Field( + default=None, + description="""If true, no memory revisions will be created for any requests to the Memory Bank.""", + ) + generation_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigGenerationConfig ] = Field( - default=None, description="""A custom memory topic defined by the developer.""" + default=None, + description="""Optional. Configuration for how to generate memories for the Memory Bank.""", ) - managed_memory_topic: Optional[ - MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic + similarity_search_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig ] = Field( - default=None, description="""A managed memory topic defined by Memory Bank.""" + default=None, + description="""Optional. Configuration for how to perform similarity search on memories. If not set, the Memory Bank will use the default embedding model `text-embedding-005`.""", + ) + ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfig] = Field( + default=None, + description="""Optional. Configuration for automatic TTL ("time-to-live") of the memories in the Memory Bank. If not set, TTL will not be applied automatically. The TTL can be explicitly set by modifying the `expire_time` of each Memory resource.""", + ) + structured_memory_configs: Optional[list[StructuredMemoryConfig]] = Field( + default=None, + description="""Optional. Configuration for organizing structured memories for a particular scope.""", ) -class MemoryBankCustomizationConfigMemoryTopicDict(TypedDict, total=False): - """A topic of information that should be extracted from conversations and stored as memories.""" +class ReasoningEngineContextSpecMemoryBankConfigDict(TypedDict, total=False): + """Specification for a Memory Bank.""" - custom_memory_topic: Optional[ - MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict + customization_configs: Optional[list[MemoryBankCustomizationConfigDict]] + """Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""" + + disable_memory_revisions: Optional[bool] + """If true, no memory revisions will be created for any requests to the Memory Bank.""" + + generation_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict ] - """A custom memory topic defined by the developer.""" + """Optional. Configuration for how to generate memories for the Memory Bank.""" - managed_memory_topic: Optional[ - MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict + similarity_search_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict ] - """A managed memory topic defined by Memory Bank.""" + """Optional. Configuration for how to perform similarity search on memories. If not set, the Memory Bank will use the default embedding model `text-embedding-005`.""" + ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict] + """Optional. Configuration for automatic TTL ("time-to-live") of the memories in the Memory Bank. If not set, TTL will not be applied automatically. The TTL can be explicitly set by modifying the `expire_time` of each Memory resource.""" -MemoryBankCustomizationConfigMemoryTopicOrDict = Union[ - MemoryBankCustomizationConfigMemoryTopic, - MemoryBankCustomizationConfigMemoryTopicDict, + structured_memory_configs: Optional[list[StructuredMemoryConfigDict]] + """Optional. Configuration for organizing structured memories for a particular scope.""" + + +ReasoningEngineContextSpecMemoryBankConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfig, + ReasoningEngineContextSpecMemoryBankConfigDict, ] -class MemoryBankCustomizationConfigConsolidationConfig(_common.BaseModel): - """Represents configuration for customizing how memories are consolidated.""" +class ReasoningEngineContextSpec(_common.BaseModel): + """Configuration for how Agent Engine sub-resources should manage context.""" - revisions_per_candidate_count: Optional[int] = Field( + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfig] = Field( default=None, - description="""Optional. Represents the maximum number of revisions to consider for each candidate memory. If not set, then the default value (1) will be used, which means that only the latest revision will be considered.""", + description="""Optional. Specification for a Memory Bank, which manages memories for the Agent Engine.""", ) -class MemoryBankCustomizationConfigConsolidationConfigDict(TypedDict, total=False): - """Represents configuration for customizing how memories are consolidated.""" +class ReasoningEngineContextSpecDict(TypedDict, total=False): + """Configuration for how Agent Engine sub-resources should manage context.""" - revisions_per_candidate_count: Optional[int] - """Optional. Represents the maximum number of revisions to consider for each candidate memory. If not set, then the default value (1) will be used, which means that only the latest revision will be considered.""" + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfigDict] + """Optional. Specification for a Memory Bank, which manages memories for the Agent Engine.""" -MemoryBankCustomizationConfigConsolidationConfigOrDict = Union[ - MemoryBankCustomizationConfigConsolidationConfig, - MemoryBankCustomizationConfigConsolidationConfigDict, +ReasoningEngineContextSpecOrDict = Union[ + ReasoningEngineContextSpec, ReasoningEngineContextSpecDict ] -class MemoryBankCustomizationConfig(_common.BaseModel): - """Represents configuration for organizing natural language memories for a particular scope.""" +class SecretRef(_common.BaseModel): + """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" - enable_third_person_memories: Optional[bool] = Field( - default=None, - description="""Optional. Indicates whether the memories will be generated in the third person (i.e. "The user generates memories with Memory Bank."). By default, the memories will be generated in the first person (i.e. "I generate memories with Memory Bank.")""", - ) - generate_memories_examples: Optional[ - list[MemoryBankCustomizationConfigGenerateMemoriesExample] - ] = Field( - default=None, - description="""Optional. Provides examples of how to generate memories for a particular scope.""", - ) - memory_topics: Optional[list[MemoryBankCustomizationConfigMemoryTopic]] = Field( - default=None, - description="""Optional. Represents topics of information that should be extracted from conversations and stored as memories. If not set, then Memory Bank's default topics will be used.""", - ) - scope_keys: Optional[list[str]] = Field( + secret: Optional[str] = Field( default=None, - description="""Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""", - ) - consolidation_config: Optional[MemoryBankCustomizationConfigConsolidationConfig] = ( - Field( - default=None, - description="""Optional. Represents configuration for customizing how memories are consolidated together.""", - ) + description="""Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""", ) - disable_natural_language_memories: Optional[bool] = Field( + version: Optional[str] = Field( default=None, - description="""Optional. Indicates whether natural language memory generation should be disabled for all requests. By default, natural language memory generation is enabled. Set this to `true` when you only want to generate structured memories.""", + description="""The Cloud Secret Manager secret version. Can be 'latest' for the latest version, an integer for a specific version, or a version alias.""", ) -class MemoryBankCustomizationConfigDict(TypedDict, total=False): - """Represents configuration for organizing natural language memories for a particular scope.""" - - enable_third_person_memories: Optional[bool] - """Optional. Indicates whether the memories will be generated in the third person (i.e. "The user generates memories with Memory Bank."). By default, the memories will be generated in the first person (i.e. "I generate memories with Memory Bank.")""" - - generate_memories_examples: Optional[ - list[MemoryBankCustomizationConfigGenerateMemoriesExampleDict] - ] - """Optional. Provides examples of how to generate memories for a particular scope.""" - - memory_topics: Optional[list[MemoryBankCustomizationConfigMemoryTopicDict]] - """Optional. Represents topics of information that should be extracted from conversations and stored as memories. If not set, then Memory Bank's default topics will be used.""" - - scope_keys: Optional[list[str]] - """Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""" +class SecretRefDict(TypedDict, total=False): + """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" - consolidation_config: Optional[MemoryBankCustomizationConfigConsolidationConfigDict] - """Optional. Represents configuration for customizing how memories are consolidated together.""" + secret: Optional[str] + """Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""" - disable_natural_language_memories: Optional[bool] - """Optional. Indicates whether natural language memory generation should be disabled for all requests. By default, natural language memory generation is enabled. Set this to `true` when you only want to generate structured memories.""" + version: Optional[str] + """The Cloud Secret Manager secret version. Can be 'latest' for the latest version, an integer for a specific version, or a version alias.""" -MemoryBankCustomizationConfigOrDict = Union[ - MemoryBankCustomizationConfig, MemoryBankCustomizationConfigDict -] +SecretRefOrDict = Union[SecretRef, SecretRefDict] -class MemoryGenerationTriggerConfigGenerationTriggerRule(_common.BaseModel): - """Represents the active rule that determines when to flush the buffer.""" +class SecretEnvVar(_common.BaseModel): + """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" - event_count: Optional[int] = Field( - default=None, - description="""Optional. Specifies to trigger generation when the event count reaches this limit.""", - ) - fixed_interval: Optional[str] = Field( - default=None, - description="""Optional. Specifies to trigger generation at a fixed interval. The duration must have a minute-level granularity.""", - ) - idle_duration: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Optional. Specifies to trigger generation if the stream is inactive for the specified duration after the most recent event. The duration must have a minute-level granularity.""", + description="""Required. Name of the secret environment variable.""", ) - overlap_event_count: Optional[int] = Field( + secret_ref: Optional[SecretRef] = Field( default=None, - description="""Optional. Re-include the last N already-processed events in the next window.""", + description="""Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""", ) -class MemoryGenerationTriggerConfigGenerationTriggerRuleDict(TypedDict, total=False): - """Represents the active rule that determines when to flush the buffer.""" - - event_count: Optional[int] - """Optional. Specifies to trigger generation when the event count reaches this limit.""" - - fixed_interval: Optional[str] - """Optional. Specifies to trigger generation at a fixed interval. The duration must have a minute-level granularity.""" - - idle_duration: Optional[str] - """Optional. Specifies to trigger generation if the stream is inactive for the specified duration after the most recent event. The duration must have a minute-level granularity.""" - - overlap_event_count: Optional[int] - """Optional. Re-include the last N already-processed events in the next window.""" - - -MemoryGenerationTriggerConfigGenerationTriggerRuleOrDict = Union[ - MemoryGenerationTriggerConfigGenerationTriggerRule, - MemoryGenerationTriggerConfigGenerationTriggerRuleDict, -] - - -class MemoryGenerationTriggerConfig(_common.BaseModel): - """The configuration for triggering memory generation for ingested events.""" - - generation_rule: Optional[MemoryGenerationTriggerConfigGenerationTriggerRule] = ( - Field( - default=None, - description="""Optional. Represents the active rule that determines when to flush the buffer. If not set, then the stream will be force flushed immediately.""", - ) - ) - +class SecretEnvVarDict(TypedDict, total=False): + """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" -class MemoryGenerationTriggerConfigDict(TypedDict, total=False): - """The configuration for triggering memory generation for ingested events.""" + name: Optional[str] + """Required. Name of the secret environment variable.""" - generation_rule: Optional[MemoryGenerationTriggerConfigGenerationTriggerRuleDict] - """Optional. Represents the active rule that determines when to flush the buffer. If not set, then the stream will be force flushed immediately.""" + secret_ref: Optional[SecretRefDict] + """Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" -MemoryGenerationTriggerConfigOrDict = Union[ - MemoryGenerationTriggerConfig, MemoryGenerationTriggerConfigDict -] +SecretEnvVarOrDict = Union[SecretEnvVar, SecretEnvVarDict] -class ReasoningEngineContextSpecMemoryBankConfigGenerationConfig(_common.BaseModel): - """Configuration for how to generate memories.""" +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig( + _common.BaseModel +): + """Configuration for traffic originating from a Reasoning Engine.""" - model: Optional[str] = Field( - default=None, - description="""Optional. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", - ) - generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + agent_gateway: Optional[str] = Field( default=None, - description="""Optional. Specifies the default trigger configuration for generating memories using `IngestEvents`.""", + description="""Required. The resource name of the Agent Gateway for outbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `AGENT_TO_ANYWHERE`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`""", ) -class ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict( +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict( TypedDict, total=False ): - """Configuration for how to generate memories.""" - - model: Optional[str] - """Optional. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + """Configuration for traffic originating from a Reasoning Engine.""" - generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] - """Optional. Specifies the default trigger configuration for generating memories using `IngestEvents`.""" + agent_gateway: Optional[str] + """Required. The resource name of the Agent Gateway for outbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `AGENT_TO_ANYWHERE`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`""" -ReasoningEngineContextSpecMemoryBankConfigGenerationConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigGenerationConfig, - ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict, +ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigOrDict = Union[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig, + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict, ] -class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig( +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig( _common.BaseModel ): - """Configuration for how to perform similarity search on memories.""" + """Configuration for traffic targeting a Reasoning Engine.""" - embedding_model: Optional[str] = Field( + agent_gateway: Optional[str] = Field( default=None, - description="""Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", + description="""Required. The resource name of the Agent Gateway to use for inbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `CLIENT_TO_AGENT`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`""", ) -class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict( +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict( TypedDict, total=False ): - """Configuration for how to perform similarity search on memories.""" + """Configuration for traffic targeting a Reasoning Engine.""" - embedding_model: Optional[str] - """Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + agent_gateway: Optional[str] + """Required. The resource name of the Agent Gateway to use for inbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `CLIENT_TO_AGENT`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`""" -ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig, - ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict, +ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigOrDict = Union[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig, + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict, ] -class ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig( - _common.BaseModel -): - """Configuration for TTL of the memories in the Memory Bank based on the action that created or updated the memory.""" +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfig(_common.BaseModel): + """Agent Gateway configuration for a Reasoning Engine deployment.""" - create_ttl: Optional[str] = Field( + agent_to_anywhere_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig + ] = Field( default=None, - description="""Optional. The TTL duration for memories uploaded via CreateMemory.""", + description="""Optional. Configuration for traffic originating from the Reasoning Engine. When unset, outgoing traffic is not routed through an Agent Gateway.""", ) - generate_created_ttl: Optional[str] = Field( + client_to_agent_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig + ] = Field( default=None, - description="""Optional. The TTL duration for memories newly generated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.CREATED).""", + description="""Optional. Configuration for traffic targeting the Reasoning Engine. When unset, incoming traffic is not routed through an Agent Gateway.""", ) - generate_updated_ttl: Optional[str] = Field( - default=None, - description="""Optional. The TTL duration for memories updated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.UPDATED). In the case of an UPDATE action, the `expire_time` of the existing memory will be updated to the new value (now + TTL).""", - ) - -class ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict( - TypedDict, total=False -): - """Configuration for TTL of the memories in the Memory Bank based on the action that created or updated the memory.""" - create_ttl: Optional[str] - """Optional. The TTL duration for memories uploaded via CreateMemory.""" +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict(TypedDict, total=False): + """Agent Gateway configuration for a Reasoning Engine deployment.""" - generate_created_ttl: Optional[str] - """Optional. The TTL duration for memories newly generated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.CREATED).""" + agent_to_anywhere_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict + ] + """Optional. Configuration for traffic originating from the Reasoning Engine. When unset, outgoing traffic is not routed through an Agent Gateway.""" - generate_updated_ttl: Optional[str] - """Optional. The TTL duration for memories updated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.UPDATED). In the case of an UPDATE action, the `expire_time` of the existing memory will be updated to the new value (now + TTL).""" + client_to_agent_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict + ] + """Optional. Configuration for traffic targeting the Reasoning Engine. When unset, incoming traffic is not routed through an Agent Gateway.""" -ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig, - ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict, +ReasoningEngineSpecDeploymentSpecAgentGatewayConfigOrDict = Union[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig, + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict, ] -class ReasoningEngineContextSpecMemoryBankConfigTtlConfig(_common.BaseModel): - """Configuration for automatically setting the TTL ("time-to-live") of the memories in the Memory Bank.""" +class KeepAliveProbeHttpGet(_common.BaseModel): + """Specifies the HTTP GET configuration for the probe.""" - default_ttl: Optional[str] = Field( - default=None, - description="""Optional. The default TTL duration of the memories in the Memory Bank. This applies to all operations that create or update a memory.""", - ) - granular_ttl_config: Optional[ - ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig - ] = Field( + path: Optional[str] = Field( default=None, - description="""Optional. The granular TTL configuration of the memories in the Memory Bank.""", + description="""Required. Specifies the path of the HTTP GET request (e.g., `"/is_busy"`).""", ) - memory_revision_default_ttl: Optional[str] = Field( + port: Optional[int] = Field( default=None, - description="""Optional. The default TTL duration of the memory revisions in the Memory Bank. This applies to all operations that create a memory revision. If not set, a default TTL of 365 days will be used.""", + description="""Optional. Specifies the port number on the container to which the request is sent.""", ) -class ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict(TypedDict, total=False): - """Configuration for automatically setting the TTL ("time-to-live") of the memories in the Memory Bank.""" - - default_ttl: Optional[str] - """Optional. The default TTL duration of the memories in the Memory Bank. This applies to all operations that create or update a memory.""" +class KeepAliveProbeHttpGetDict(TypedDict, total=False): + """Specifies the HTTP GET configuration for the probe.""" - granular_ttl_config: Optional[ - ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict - ] - """Optional. The granular TTL configuration of the memories in the Memory Bank.""" + path: Optional[str] + """Required. Specifies the path of the HTTP GET request (e.g., `"/is_busy"`).""" - memory_revision_default_ttl: Optional[str] - """Optional. The default TTL duration of the memory revisions in the Memory Bank. This applies to all operations that create a memory revision. If not set, a default TTL of 365 days will be used.""" + port: Optional[int] + """Optional. Specifies the port number on the container to which the request is sent.""" -ReasoningEngineContextSpecMemoryBankConfigTtlConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigTtlConfig, - ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict, -] +KeepAliveProbeHttpGetOrDict = Union[KeepAliveProbeHttpGet, KeepAliveProbeHttpGetDict] -class StructuredMemorySchemaConfig(_common.BaseModel): - """Represents the OpenAPI schema of the structured memories.""" +class KeepAliveProbe(_common.BaseModel): + """Represents the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""" - memory_schema: Optional[genai_types.Schema] = Field( - default=None, - description="""Required. Represents the OpenAPI schema of the structured memories.""", - ) - id: Optional[str] = Field( + http_get: Optional[KeepAliveProbeHttpGet] = Field( default=None, - description="""Required. Represents the ID of the schema. Must be 1-63 characters, start with a lowercase letter, and consist of lowercase letters, numbers, and hyphens.""", + description="""Optional. Specifies the HTTP GET configuration for the probe.""", ) - memory_type: Optional[MemoryType] = Field( + max_seconds: Optional[int] = Field( default=None, - description="""Optional. Represents the type of the structured memories associated with the schema. If not set, then `STRUCTURED_PROFILE` will be used.""", + description="""Optional. Specifies the maximum duration (in seconds) to keep the instance alive via this probe. Can be a maximum of 3600 seconds (1 hour).""", ) -class StructuredMemorySchemaConfigDict(TypedDict, total=False): - """Represents the OpenAPI schema of the structured memories.""" - - memory_schema: Optional[genai_types.Schema] - """Required. Represents the OpenAPI schema of the structured memories.""" +class KeepAliveProbeDict(TypedDict, total=False): + """Represents the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""" - id: Optional[str] - """Required. Represents the ID of the schema. Must be 1-63 characters, start with a lowercase letter, and consist of lowercase letters, numbers, and hyphens.""" + http_get: Optional[KeepAliveProbeHttpGetDict] + """Optional. Specifies the HTTP GET configuration for the probe.""" - memory_type: Optional[MemoryType] - """Optional. Represents the type of the structured memories associated with the schema. If not set, then `STRUCTURED_PROFILE` will be used.""" + max_seconds: Optional[int] + """Optional. Specifies the maximum duration (in seconds) to keep the instance alive via this probe. Can be a maximum of 3600 seconds (1 hour).""" -StructuredMemorySchemaConfigOrDict = Union[ - StructuredMemorySchemaConfig, StructuredMemorySchemaConfigDict -] +KeepAliveProbeOrDict = Union[KeepAliveProbe, KeepAliveProbeDict] -class StructuredMemoryConfig(_common.BaseModel): - """Configuration for organizing structured memories within a scope.""" +class ReasoningEngineSpecDeploymentSpec(_common.BaseModel): + """The specification of a Reasoning Engine deployment.""" - schema_configs: Optional[list[StructuredMemorySchemaConfig]] = Field( + agent_server_mode: Optional[AgentServerMode] = Field( + default=None, description="""The agent server mode.""" + ) + container_concurrency: Optional[int] = Field( default=None, - description="""Optional. Represents configuration of the structured memories' schemas.""", + description="""Optional. Concurrency for each container and agent server. Recommended value: 2 * cpu + 1. Defaults to 9.""", ) - scope_keys: Optional[list[str]] = Field( + env: Optional[list[EnvVar]] = Field( default=None, - description="""Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""", + description="""Optional. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API.""", ) - - -class StructuredMemoryConfigDict(TypedDict, total=False): - """Configuration for organizing structured memories within a scope.""" - - schema_configs: Optional[list[StructuredMemorySchemaConfigDict]] - """Optional. Represents configuration of the structured memories' schemas.""" - - scope_keys: Optional[list[str]] - """Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""" - - -StructuredMemoryConfigOrDict = Union[StructuredMemoryConfig, StructuredMemoryConfigDict] - - -class ReasoningEngineContextSpecMemoryBankConfig(_common.BaseModel): - """Specification for a Memory Bank.""" - - customization_configs: Optional[list[MemoryBankCustomizationConfig]] = Field( + max_instances: Optional[int] = Field( default=None, - description="""Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""", + description="""Optional. The maximum number of application instances that can be launched to handle increased traffic. Defaults to 100. Range: [1, 1000]. If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100].""", ) - disable_memory_revisions: Optional[bool] = Field( + min_instances: Optional[int] = Field( default=None, - description="""If true, no memory revisions will be created for any requests to the Memory Bank.""", + description="""Optional. The minimum number of application instances that will be kept running at all times. Defaults to 1. Range: [0, 75].""", ) - generation_config: Optional[ - ReasoningEngineContextSpecMemoryBankConfigGenerationConfig - ] = Field( + psc_interface_config: Optional[PscInterfaceConfig] = Field( + default=None, description="""Optional. Configuration for PSC-I.""" + ) + resource_limits: Optional[dict[str, str]] = Field( default=None, - description="""Optional. Configuration for how to generate memories for the Memory Bank.""", + description="""Optional. Resource limits for each container. Only 'cpu' and 'memory' keys are supported. Defaults to {"cpu": "4", "memory": "4Gi"}. * The only supported values for CPU are '1', '2', '4', '6' and '8'. For more information, go to https://cloud.google.com/run/docs/configuring/cpu. * The only supported values for memory are '1Gi', '2Gi', ... '32 Gi'. * For required cpu on different memory values, go to https://cloud.google.com/run/docs/configuring/memory-limits""", ) - similarity_search_config: Optional[ - ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig - ] = Field( + secret_env: Optional[list[SecretEnvVar]] = Field( default=None, - description="""Optional. Configuration for how to perform similarity search on memories. If not set, the Memory Bank will use the default embedding model `text-embedding-005`.""", + description="""Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""", ) - ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfig] = Field( + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + ] = Field( default=None, - description="""Optional. Configuration for automatic TTL ("time-to-live") of the memories in the Memory Bank. If not set, TTL will not be applied automatically. The TTL can be explicitly set by modifying the `expire_time` of each Memory resource.""", + description="""Optional. Agent Gateway configuration for the Reasoning Engine deployment.""", ) - structured_memory_configs: Optional[list[StructuredMemoryConfig]] = Field( + keep_alive_probe: Optional[KeepAliveProbe] = Field( default=None, - description="""Optional. Configuration for organizing structured memories for a particular scope.""", + description="""Optional. Specifies the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""", ) -class ReasoningEngineContextSpecMemoryBankConfigDict(TypedDict, total=False): - """Specification for a Memory Bank.""" - - customization_configs: Optional[list[MemoryBankCustomizationConfigDict]] - """Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""" - - disable_memory_revisions: Optional[bool] - """If true, no memory revisions will be created for any requests to the Memory Bank.""" - - generation_config: Optional[ - ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict - ] - """Optional. Configuration for how to generate memories for the Memory Bank.""" - - similarity_search_config: Optional[ - ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict - ] - """Optional. Configuration for how to perform similarity search on memories. If not set, the Memory Bank will use the default embedding model `text-embedding-005`.""" +class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False): + """The specification of a Reasoning Engine deployment.""" - ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict] - """Optional. Configuration for automatic TTL ("time-to-live") of the memories in the Memory Bank. If not set, TTL will not be applied automatically. The TTL can be explicitly set by modifying the `expire_time` of each Memory resource.""" + agent_server_mode: Optional[AgentServerMode] + """The agent server mode.""" - structured_memory_configs: Optional[list[StructuredMemoryConfigDict]] - """Optional. Configuration for organizing structured memories for a particular scope.""" + container_concurrency: Optional[int] + """Optional. Concurrency for each container and agent server. Recommended value: 2 * cpu + 1. Defaults to 9.""" + env: Optional[list[EnvVarDict]] + """Optional. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API.""" -ReasoningEngineContextSpecMemoryBankConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfig, - ReasoningEngineContextSpecMemoryBankConfigDict, -] + max_instances: Optional[int] + """Optional. The maximum number of application instances that can be launched to handle increased traffic. Defaults to 100. Range: [1, 1000]. If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100].""" + min_instances: Optional[int] + """Optional. The minimum number of application instances that will be kept running at all times. Defaults to 1. Range: [0, 75].""" -class ReasoningEngineContextSpec(_common.BaseModel): - """Configuration for how Agent Engine sub-resources should manage context.""" + psc_interface_config: Optional[PscInterfaceConfigDict] + """Optional. Configuration for PSC-I.""" - memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfig] = Field( - default=None, - description="""Optional. Specification for a Memory Bank, which manages memories for the Agent Engine.""", - ) + resource_limits: Optional[dict[str, str]] + """Optional. Resource limits for each container. Only 'cpu' and 'memory' keys are supported. Defaults to {"cpu": "4", "memory": "4Gi"}. * The only supported values for CPU are '1', '2', '4', '6' and '8'. For more information, go to https://cloud.google.com/run/docs/configuring/cpu. * The only supported values for memory are '1Gi', '2Gi', ... '32 Gi'. * For required cpu on different memory values, go to https://cloud.google.com/run/docs/configuring/memory-limits""" + secret_env: Optional[list[SecretEnvVarDict]] + """Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""" -class ReasoningEngineContextSpecDict(TypedDict, total=False): - """Configuration for how Agent Engine sub-resources should manage context.""" + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] + """Optional. Agent Gateway configuration for the Reasoning Engine deployment.""" - memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfigDict] - """Optional. Specification for a Memory Bank, which manages memories for the Agent Engine.""" + keep_alive_probe: Optional[KeepAliveProbeDict] + """Optional. Specifies the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""" -ReasoningEngineContextSpecOrDict = Union[ - ReasoningEngineContextSpec, ReasoningEngineContextSpecDict +ReasoningEngineSpecDeploymentSpecOrDict = Union[ + ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict ] -class SecretRef(_common.BaseModel): - """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" +class ReasoningEngineSpecPackageSpec(_common.BaseModel): + """User-provided package specification, containing pickled object and package requirements.""" - secret: Optional[str] = Field( + dependency_files_gcs_uri: Optional[str] = Field( default=None, - description="""Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""", + description="""Optional. The Cloud Storage URI of the dependency files in tar.gz format.""", ) - version: Optional[str] = Field( + pickle_object_gcs_uri: Optional[str] = Field( default=None, - description="""The Cloud Secret Manager secret version. Can be 'latest' for the latest version, an integer for a specific version, or a version alias.""", + description="""Optional. The Cloud Storage URI of the pickled python object.""", + ) + python_version: Optional[str] = Field( + default=None, + description="""Optional. The Python version. Supported values are 3.10, 3.11, 3.12, 3.13, 3.14. If not specified, the default value is 3.10.""", + ) + requirements_gcs_uri: Optional[str] = Field( + default=None, + description="""Optional. The Cloud Storage URI of the `requirements.txt` file""", ) -class SecretRefDict(TypedDict, total=False): - """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" - - secret: Optional[str] - """Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""" +class ReasoningEngineSpecPackageSpecDict(TypedDict, total=False): + """User-provided package specification, containing pickled object and package requirements.""" - version: Optional[str] - """The Cloud Secret Manager secret version. Can be 'latest' for the latest version, an integer for a specific version, or a version alias.""" + dependency_files_gcs_uri: Optional[str] + """Optional. The Cloud Storage URI of the dependency files in tar.gz format.""" + pickle_object_gcs_uri: Optional[str] + """Optional. The Cloud Storage URI of the pickled python object.""" -SecretRefOrDict = Union[SecretRef, SecretRefDict] + python_version: Optional[str] + """Optional. The Python version. Supported values are 3.10, 3.11, 3.12, 3.13, 3.14. If not specified, the default value is 3.10.""" + requirements_gcs_uri: Optional[str] + """Optional. The Cloud Storage URI of the `requirements.txt` file""" -class SecretEnvVar(_common.BaseModel): - """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" - name: Optional[str] = Field( - default=None, - description="""Required. Name of the secret environment variable.""", - ) - secret_ref: Optional[SecretRef] = Field( +ReasoningEngineSpecPackageSpecOrDict = Union[ + ReasoningEngineSpecPackageSpec, ReasoningEngineSpecPackageSpecDict +] + + +class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig(_common.BaseModel): + """Configuration for the Agent Development Kit (ADK).""" + + json_config: Optional[dict[str, Any]] = Field( default=None, - description="""Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""", + description="""Required. The value of the ADK config in JSON format.""", ) -class SecretEnvVarDict(TypedDict, total=False): - """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" - - name: Optional[str] - """Required. Name of the secret environment variable.""" +class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict( + TypedDict, total=False +): + """Configuration for the Agent Development Kit (ADK).""" - secret_ref: Optional[SecretRefDict] - """Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" + json_config: Optional[dict[str, Any]] + """Required. The value of the ADK config in JSON format.""" -SecretEnvVarOrDict = Union[SecretEnvVar, SecretEnvVarDict] +ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig, + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict, +] -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig( - _common.BaseModel -): - """Configuration for traffic originating from a Reasoning Engine.""" +class ReasoningEngineSpecSourceCodeSpecInlineSource(_common.BaseModel): + """Specifies source code provided as a byte stream.""" - agent_gateway: Optional[str] = Field( + source_archive: Optional[bytes] = Field( default=None, - description="""Required. The resource name of the Agent Gateway for outbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `AGENT_TO_ANYWHERE`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`""", + description="""Required. Input only. The application source code archive. It must be a compressed tarball (.tar.gz) file.""", ) -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict( - TypedDict, total=False -): - """Configuration for traffic originating from a Reasoning Engine.""" +class ReasoningEngineSpecSourceCodeSpecInlineSourceDict(TypedDict, total=False): + """Specifies source code provided as a byte stream.""" - agent_gateway: Optional[str] - """Required. The resource name of the Agent Gateway for outbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `AGENT_TO_ANYWHERE`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`""" + source_archive: Optional[bytes] + """Required. Input only. The application source code archive. It must be a compressed tarball (.tar.gz) file.""" -ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigOrDict = Union[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig, - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict, +ReasoningEngineSpecSourceCodeSpecInlineSourceOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecInlineSource, + ReasoningEngineSpecSourceCodeSpecInlineSourceDict, ] -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig( - _common.BaseModel -): - """Configuration for traffic targeting a Reasoning Engine.""" +class ReasoningEngineSpecSourceCodeSpecAgentConfigSource(_common.BaseModel): + """Specification for the deploying from agent config.""" - agent_gateway: Optional[str] = Field( + adk_config: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig + ] = Field(default=None, description="""Required. The ADK configuration.""") + inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSource] = Field( default=None, - description="""Required. The resource name of the Agent Gateway to use for inbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `CLIENT_TO_AGENT`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`""", + description="""Optional. Any additional files needed to interpret the config. If a `requirements.txt` file is present in the `inline_source`, the corresponding packages will be installed. If no `requirements.txt` file is present in `inline_source`, then the latest version of `google-adk` will be installed for interpreting the ADK config.""", ) -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict( - TypedDict, total=False -): - """Configuration for traffic targeting a Reasoning Engine.""" +class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict(TypedDict, total=False): + """Specification for the deploying from agent config.""" - agent_gateway: Optional[str] - """Required. The resource name of the Agent Gateway to use for inbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `CLIENT_TO_AGENT`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`""" + adk_config: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict + ] + """Required. The ADK configuration.""" + + inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSourceDict] + """Optional. Any additional files needed to interpret the config. If a `requirements.txt` file is present in the `inline_source`, the corresponding packages will be installed. If no `requirements.txt` file is present in `inline_source`, then the latest version of `google-adk` will be installed for interpreting the ADK config.""" -ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigOrDict = Union[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig, - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict, +ReasoningEngineSpecSourceCodeSpecAgentConfigSourceOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSource, + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict, ] -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfig(_common.BaseModel): - """Agent Gateway configuration for a Reasoning Engine deployment.""" +class ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig(_common.BaseModel): + """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. - agent_to_anywhere_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig - ] = Field( + This includes the repository, revision, and directory to use. + """ + + git_repository_link: Optional[str] = Field( default=None, - description="""Optional. Configuration for traffic originating from the Reasoning Engine. When unset, outgoing traffic is not routed through an Agent Gateway.""", + description="""Required. The Developer Connect Git repository link, formatted as `projects/{project_id}/locations/{location_id}/connections/{connection_id}/gitRepositoryLink/{repository_link_id}`.""", ) - client_to_agent_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig - ] = Field( + dir: Optional[str] = Field( default=None, - description="""Optional. Configuration for traffic targeting the Reasoning Engine. When unset, incoming traffic is not routed through an Agent Gateway.""", + description="""Required. Directory, relative to the source root, in which to run the build.""", + ) + revision: Optional[str] = Field( + default=None, + description="""Required. The revision to fetch from the Git repository such as a branch, a tag, a commit SHA, or any Git ref.""", ) -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict(TypedDict, total=False): - """Agent Gateway configuration for a Reasoning Engine deployment.""" +class ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict( + TypedDict, total=False +): + """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. - agent_to_anywhere_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict - ] - """Optional. Configuration for traffic originating from the Reasoning Engine. When unset, outgoing traffic is not routed through an Agent Gateway.""" + This includes the repository, revision, and directory to use. + """ - client_to_agent_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict - ] - """Optional. Configuration for traffic targeting the Reasoning Engine. When unset, incoming traffic is not routed through an Agent Gateway.""" + git_repository_link: Optional[str] + """Required. The Developer Connect Git repository link, formatted as `projects/{project_id}/locations/{location_id}/connections/{connection_id}/gitRepositoryLink/{repository_link_id}`.""" + dir: Optional[str] + """Required. Directory, relative to the source root, in which to run the build.""" -ReasoningEngineSpecDeploymentSpecAgentGatewayConfigOrDict = Union[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig, - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict, + revision: Optional[str] + """Required. The revision to fetch from the Git repository such as a branch, a tag, a commit SHA, or any Git ref.""" + + +ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig, + ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict, ] -class KeepAliveProbeHttpGet(_common.BaseModel): - """Specifies the HTTP GET configuration for the probe.""" +class ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource(_common.BaseModel): + """Specifies source code to be fetched from a Git repository managed through the Developer Connect service.""" - path: Optional[str] = Field( - default=None, - description="""Required. Specifies the path of the HTTP GET request (e.g., `"/is_busy"`).""", - ) - port: Optional[int] = Field( + config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig] = Field( default=None, - description="""Optional. Specifies the port number on the container to which the request is sent.""", + description="""Required. The Developer Connect configuration that defines the specific repository, revision, and directory to use as the source code root.""", ) -class KeepAliveProbeHttpGetDict(TypedDict, total=False): - """Specifies the HTTP GET configuration for the probe.""" - - path: Optional[str] - """Required. Specifies the path of the HTTP GET request (e.g., `"/is_busy"`).""" +class ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict( + TypedDict, total=False +): + """Specifies source code to be fetched from a Git repository managed through the Developer Connect service.""" - port: Optional[int] - """Optional. Specifies the port number on the container to which the request is sent.""" + config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict] + """Required. The Developer Connect configuration that defines the specific repository, revision, and directory to use as the source code root.""" -KeepAliveProbeHttpGetOrDict = Union[KeepAliveProbeHttpGet, KeepAliveProbeHttpGetDict] +ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource, + ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict, +] -class KeepAliveProbe(_common.BaseModel): - """Represents the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""" +class ReasoningEngineSpecSourceCodeSpecImageSpec(_common.BaseModel): + """The image spec for building an image (within a single build step), based on the config file (i.e. Dockerfile) in the source directory.""" - http_get: Optional[KeepAliveProbeHttpGet] = Field( - default=None, - description="""Optional. Specifies the HTTP GET configuration for the probe.""", - ) - max_seconds: Optional[int] = Field( + build_args: Optional[dict[str, str]] = Field( default=None, - description="""Optional. Specifies the maximum duration (in seconds) to keep the instance alive via this probe. Can be a maximum of 3600 seconds (1 hour).""", + description="""Optional. Build arguments to be used. They will be passed through --build-arg flags.""", ) -class KeepAliveProbeDict(TypedDict, total=False): - """Represents the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""" - - http_get: Optional[KeepAliveProbeHttpGetDict] - """Optional. Specifies the HTTP GET configuration for the probe.""" +class ReasoningEngineSpecSourceCodeSpecImageSpecDict(TypedDict, total=False): + """The image spec for building an image (within a single build step), based on the config file (i.e. Dockerfile) in the source directory.""" - max_seconds: Optional[int] - """Optional. Specifies the maximum duration (in seconds) to keep the instance alive via this probe. Can be a maximum of 3600 seconds (1 hour).""" + build_args: Optional[dict[str, str]] + """Optional. Build arguments to be used. They will be passed through --build-arg flags.""" -KeepAliveProbeOrDict = Union[KeepAliveProbe, KeepAliveProbeDict] +ReasoningEngineSpecSourceCodeSpecImageSpecOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecImageSpec, + ReasoningEngineSpecSourceCodeSpecImageSpecDict, +] -class ReasoningEngineSpecDeploymentSpec(_common.BaseModel): - """The specification of a Reasoning Engine deployment.""" +class ReasoningEngineSpecSourceCodeSpecPythonSpec(_common.BaseModel): + """Specification for running a Python application from source.""" - agent_server_mode: Optional[AgentServerMode] = Field( - default=None, description="""The agent server mode.""" - ) - container_concurrency: Optional[int] = Field( - default=None, - description="""Optional. Concurrency for each container and agent server. Recommended value: 2 * cpu + 1. Defaults to 9.""", - ) - env: Optional[list[EnvVar]] = Field( - default=None, - description="""Optional. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API.""", - ) - max_instances: Optional[int] = Field( + entrypoint_module: Optional[str] = Field( default=None, - description="""Optional. The maximum number of application instances that can be launched to handle increased traffic. Defaults to 100. Range: [1, 1000]. If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100].""", + description="""Optional. The Python module to load as the entrypoint, specified as a fully qualified module name. For example: path.to.agent. If not specified, defaults to "agent". The project root will be added to Python sys.path, allowing imports to be specified relative to the root. This field should not be set if the source is `agent_config_source`.""", ) - min_instances: Optional[int] = Field( + entrypoint_object: Optional[str] = Field( default=None, - description="""Optional. The minimum number of application instances that will be kept running at all times. Defaults to 1. Range: [0, 75].""", - ) - psc_interface_config: Optional[PscInterfaceConfig] = Field( - default=None, description="""Optional. Configuration for PSC-I.""" + description="""Optional. The name of the callable object within the `entrypoint_module` to use as the application If not specified, defaults to "root_agent". This field should not be set if the source is `agent_config_source`.""", ) - resource_limits: Optional[dict[str, str]] = Field( + requirements_file: Optional[str] = Field( default=None, - description="""Optional. Resource limits for each container. Only 'cpu' and 'memory' keys are supported. Defaults to {"cpu": "4", "memory": "4Gi"}. * The only supported values for CPU are '1', '2', '4', '6' and '8'. For more information, go to https://cloud.google.com/run/docs/configuring/cpu. * The only supported values for memory are '1Gi', '2Gi', ... '32 Gi'. * For required cpu on different memory values, go to https://cloud.google.com/run/docs/configuring/memory-limits""", + description="""Optional. The path to the requirements file, relative to the source root. If not specified, defaults to "requirements.txt".""", ) - secret_env: Optional[list[SecretEnvVar]] = Field( + version: Optional[str] = Field( default=None, - description="""Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""", + description="""Optional. The version of Python to use. Supported versions include 3.10, 3.11, 3.12, 3.13, 3.14. If not specified, default value is 3.10.""", ) - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + + +class ReasoningEngineSpecSourceCodeSpecPythonSpecDict(TypedDict, total=False): + """Specification for running a Python application from source.""" + + entrypoint_module: Optional[str] + """Optional. The Python module to load as the entrypoint, specified as a fully qualified module name. For example: path.to.agent. If not specified, defaults to "agent". The project root will be added to Python sys.path, allowing imports to be specified relative to the root. This field should not be set if the source is `agent_config_source`.""" + + entrypoint_object: Optional[str] + """Optional. The name of the callable object within the `entrypoint_module` to use as the application If not specified, defaults to "root_agent". This field should not be set if the source is `agent_config_source`.""" + + requirements_file: Optional[str] + """Optional. The path to the requirements file, relative to the source root. If not specified, defaults to "requirements.txt".""" + + version: Optional[str] + """Optional. The version of Python to use. Supported versions include 3.10, 3.11, 3.12, 3.13, 3.14. If not specified, default value is 3.10.""" + + +ReasoningEngineSpecSourceCodeSpecPythonSpecOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecPythonSpec, + ReasoningEngineSpecSourceCodeSpecPythonSpecDict, +] + + +class ReasoningEngineSpecSourceCodeSpec(_common.BaseModel): + """Specification for deploying from source code.""" + + agent_config_source: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSource + ] = Field( + default=None, description="""Source code is generated from the agent config.""" + ) + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource ] = Field( default=None, - description="""Optional. Agent Gateway configuration for the Reasoning Engine deployment.""", + description="""Source code is in a Git repository managed by Developer Connect.""", ) - keep_alive_probe: Optional[KeepAliveProbe] = Field( + image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpec] = Field( default=None, - description="""Optional. Specifies the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""", + description="""Optional. Configuration for building an image with custom config file.""", + ) + inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSource] = Field( + default=None, description="""Source code is provided directly in the request.""" + ) + python_spec: Optional[ReasoningEngineSpecSourceCodeSpecPythonSpec] = Field( + default=None, description="""Configuration for a Python application.""" ) -class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False): - """The specification of a Reasoning Engine deployment.""" +class ReasoningEngineSpecSourceCodeSpecDict(TypedDict, total=False): + """Specification for deploying from source code.""" - agent_server_mode: Optional[AgentServerMode] - """The agent server mode.""" + agent_config_source: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict + ] + """Source code is generated from the agent config.""" - container_concurrency: Optional[int] - """Optional. Concurrency for each container and agent server. Recommended value: 2 * cpu + 1. Defaults to 9.""" + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict + ] + """Source code is in a Git repository managed by Developer Connect.""" - env: Optional[list[EnvVarDict]] - """Optional. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API.""" + image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict] + """Optional. Configuration for building an image with custom config file.""" - max_instances: Optional[int] - """Optional. The maximum number of application instances that can be launched to handle increased traffic. Defaults to 100. Range: [1, 1000]. If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100].""" + inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSourceDict] + """Source code is provided directly in the request.""" - min_instances: Optional[int] - """Optional. The minimum number of application instances that will be kept running at all times. Defaults to 1. Range: [0, 75].""" + python_spec: Optional[ReasoningEngineSpecSourceCodeSpecPythonSpecDict] + """Configuration for a Python application.""" - psc_interface_config: Optional[PscInterfaceConfigDict] - """Optional. Configuration for PSC-I.""" - resource_limits: Optional[dict[str, str]] - """Optional. Resource limits for each container. Only 'cpu' and 'memory' keys are supported. Defaults to {"cpu": "4", "memory": "4Gi"}. * The only supported values for CPU are '1', '2', '4', '6' and '8'. For more information, go to https://cloud.google.com/run/docs/configuring/cpu. * The only supported values for memory are '1Gi', '2Gi', ... '32 Gi'. * For required cpu on different memory values, go to https://cloud.google.com/run/docs/configuring/memory-limits""" +ReasoningEngineSpecSourceCodeSpecOrDict = Union[ + ReasoningEngineSpecSourceCodeSpec, ReasoningEngineSpecSourceCodeSpecDict +] - secret_env: Optional[list[SecretEnvVarDict]] - """Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""" - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict - ] - """Optional. Agent Gateway configuration for the Reasoning Engine deployment.""" +class ReasoningEngineSpecContainerSpec(_common.BaseModel): + """Specification for deploying from a container image.""" - keep_alive_probe: Optional[KeepAliveProbeDict] - """Optional. Specifies the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""" + image_uri: Optional[str] = Field( + default=None, + description="""Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""", + ) -ReasoningEngineSpecDeploymentSpecOrDict = Union[ - ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict +class ReasoningEngineSpecContainerSpecDict(TypedDict, total=False): + """Specification for deploying from a container image.""" + + image_uri: Optional[str] + """Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""" + + +ReasoningEngineSpecContainerSpecOrDict = Union[ + ReasoningEngineSpecContainerSpec, ReasoningEngineSpecContainerSpecDict ] -class ReasoningEngineSpecPackageSpec(_common.BaseModel): - """User-provided package specification, containing pickled object and package requirements.""" +class ReasoningEngineSpec(_common.BaseModel): + """The specification of an agent runtime.""" - dependency_files_gcs_uri: Optional[str] = Field( + agent_card: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. The Cloud Storage URI of the dependency files in tar.gz format.""", + description="""Optional. The A2A Agent Card for the agent (if available). It follows the specification at https://a2a-protocol.org/latest/specification/#5-agent-discovery-the-agent-card.""", ) - pickle_object_gcs_uri: Optional[str] = Field( + agent_framework: Optional[str] = Field( default=None, - description="""Optional. The Cloud Storage URI of the pickled python object.""", + description="""Optional. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom".""", ) - python_version: Optional[str] = Field( + class_methods: Optional[list[dict[str, Any]]] = Field( default=None, - description="""Optional. The Python version. Supported values are 3.10, 3.11, 3.12, 3.13, 3.14. If not specified, the default value is 3.10.""", + description="""Optional. Declarations for object class methods in OpenAPI specification format.""", ) - requirements_gcs_uri: Optional[str] = Field( + deployment_spec: Optional[ReasoningEngineSpecDeploymentSpec] = Field( default=None, - description="""Optional. The Cloud Storage URI of the `requirements.txt` file""", + description="""Optional. The specification of a Reasoning Engine deployment.""", + ) + effective_identity: Optional[str] = Field( + default=None, + description="""Output only. The identity to use for the Reasoning Engine. It can contain one of the following values: * service-{project}@gcp-sa-aiplatform-re.googleapis.com (for SERVICE_AGENT identity type) * {name}@{project}.gserviceaccount.com (for SERVICE_ACCOUNT identity type) * agents.global.{org}.system.id.goog/resources/aiplatform/projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine} (for AGENT_IDENTITY identity type)""", + ) + identity_type: Optional[IdentityType] = Field( + default=None, + description="""Optional. The identity type to use for the Reasoning Engine. If not specified, the `service_account` field will be used if set, otherwise the default Vertex AI Reasoning Engine Service Agent in the project will be used.""", + ) + package_spec: Optional[ReasoningEngineSpecPackageSpec] = Field( + default=None, + description="""Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes. The `deployment_source` field should not be set if `package_spec` is specified.""", + ) + service_account: Optional[str] = Field( + default=None, + description="""Optional. The service account that the Reasoning Engine artifact runs as. It should have "roles/storage.objectViewer" for reading the user project's Cloud Storage and "roles/aiplatform.user" for using Vertex extensions. If not specified, the Vertex AI Reasoning Engine Service Agent in the project will be used.""", + ) + source_code_spec: Optional[ReasoningEngineSpecSourceCodeSpec] = Field( + default=None, + description="""Deploy from source code files with a defined entrypoint.""", + ) + container_spec: Optional[ReasoningEngineSpecContainerSpec] = Field( + default=None, + description="""Deploy from a container image with a defined entrypoint and commands.""", ) -class ReasoningEngineSpecPackageSpecDict(TypedDict, total=False): - """User-provided package specification, containing pickled object and package requirements.""" - - dependency_files_gcs_uri: Optional[str] - """Optional. The Cloud Storage URI of the dependency files in tar.gz format.""" - - pickle_object_gcs_uri: Optional[str] - """Optional. The Cloud Storage URI of the pickled python object.""" +class ReasoningEngineSpecDict(TypedDict, total=False): + """The specification of an agent runtime.""" - python_version: Optional[str] - """Optional. The Python version. Supported values are 3.10, 3.11, 3.12, 3.13, 3.14. If not specified, the default value is 3.10.""" + agent_card: Optional[dict[str, Any]] + """Optional. The A2A Agent Card for the agent (if available). It follows the specification at https://a2a-protocol.org/latest/specification/#5-agent-discovery-the-agent-card.""" - requirements_gcs_uri: Optional[str] - """Optional. The Cloud Storage URI of the `requirements.txt` file""" + agent_framework: Optional[str] + """Optional. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom".""" + class_methods: Optional[list[dict[str, Any]]] + """Optional. Declarations for object class methods in OpenAPI specification format.""" -ReasoningEngineSpecPackageSpecOrDict = Union[ - ReasoningEngineSpecPackageSpec, ReasoningEngineSpecPackageSpecDict -] + deployment_spec: Optional[ReasoningEngineSpecDeploymentSpecDict] + """Optional. The specification of a Reasoning Engine deployment.""" + effective_identity: Optional[str] + """Output only. The identity to use for the Reasoning Engine. It can contain one of the following values: * service-{project}@gcp-sa-aiplatform-re.googleapis.com (for SERVICE_AGENT identity type) * {name}@{project}.gserviceaccount.com (for SERVICE_ACCOUNT identity type) * agents.global.{org}.system.id.goog/resources/aiplatform/projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine} (for AGENT_IDENTITY identity type)""" -class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig(_common.BaseModel): - """Configuration for the Agent Development Kit (ADK).""" + identity_type: Optional[IdentityType] + """Optional. The identity type to use for the Reasoning Engine. If not specified, the `service_account` field will be used if set, otherwise the default Vertex AI Reasoning Engine Service Agent in the project will be used.""" - json_config: Optional[dict[str, Any]] = Field( - default=None, - description="""Required. The value of the ADK config in JSON format.""", - ) + package_spec: Optional[ReasoningEngineSpecPackageSpecDict] + """Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes. The `deployment_source` field should not be set if `package_spec` is specified.""" + service_account: Optional[str] + """Optional. The service account that the Reasoning Engine artifact runs as. It should have "roles/storage.objectViewer" for reading the user project's Cloud Storage and "roles/aiplatform.user" for using Vertex extensions. If not specified, the Vertex AI Reasoning Engine Service Agent in the project will be used.""" -class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict( - TypedDict, total=False -): - """Configuration for the Agent Development Kit (ADK).""" + source_code_spec: Optional[ReasoningEngineSpecSourceCodeSpecDict] + """Deploy from source code files with a defined entrypoint.""" - json_config: Optional[dict[str, Any]] - """Required. The value of the ADK config in JSON format.""" + container_spec: Optional[ReasoningEngineSpecContainerSpecDict] + """Deploy from a container image with a defined entrypoint and commands.""" -ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig, - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict, -] +ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict] -class ReasoningEngineSpecSourceCodeSpecInlineSource(_common.BaseModel): - """Specifies source code provided as a byte stream.""" +class ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest(_common.BaseModel): + """Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" - source_archive: Optional[bytes] = Field( - default=None, - description="""Required. Input only. The application source code archive. It must be a compressed tarball (.tar.gz) file.""", - ) + pass -class ReasoningEngineSpecSourceCodeSpecInlineSourceDict(TypedDict, total=False): - """Specifies source code provided as a byte stream.""" +class ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict(TypedDict, total=False): + """Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" - source_archive: Optional[bytes] - """Required. Input only. The application source code archive. It must be a compressed tarball (.tar.gz) file.""" + pass -ReasoningEngineSpecSourceCodeSpecInlineSourceOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecInlineSource, - ReasoningEngineSpecSourceCodeSpecInlineSourceDict, +ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestOrDict = Union[ + ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest, + ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict, ] -class ReasoningEngineSpecSourceCodeSpecAgentConfigSource(_common.BaseModel): - """Specification for the deploying from agent config.""" +class ReasoningEngineTrafficConfigTrafficSplitManualTarget(_common.BaseModel): + """A single target for the traffic split, specifying a Runtime Revision and the percentage of traffic to send to it.""" - adk_config: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig - ] = Field(default=None, description="""Required. The ADK configuration.""") - inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSource] = Field( + percent: Optional[int] = Field( default=None, - description="""Optional. Any additional files needed to interpret the config. If a `requirements.txt` file is present in the `inline_source`, the corresponding packages will be installed. If no `requirements.txt` file is present in `inline_source`, then the latest version of `google-adk` will be installed for interpreting the ADK config.""", + description="""Required. Specifies percent of the traffic to this Runtime Revision.""", + ) + runtime_revision_name: Optional[str] = Field( + default=None, + description="""Required. The Runtime Revision name to which to send this portion of traffic, if traffic allocation is by Runtime Revision.""", ) -class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict(TypedDict, total=False): - """Specification for the deploying from agent config.""" +class ReasoningEngineTrafficConfigTrafficSplitManualTargetDict(TypedDict, total=False): + """A single target for the traffic split, specifying a Runtime Revision and the percentage of traffic to send to it.""" - adk_config: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict - ] - """Required. The ADK configuration.""" + percent: Optional[int] + """Required. Specifies percent of the traffic to this Runtime Revision.""" - inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSourceDict] - """Optional. Any additional files needed to interpret the config. If a `requirements.txt` file is present in the `inline_source`, the corresponding packages will be installed. If no `requirements.txt` file is present in `inline_source`, then the latest version of `google-adk` will be installed for interpreting the ADK config.""" + runtime_revision_name: Optional[str] + """Required. The Runtime Revision name to which to send this portion of traffic, if traffic allocation is by Runtime Revision.""" -ReasoningEngineSpecSourceCodeSpecAgentConfigSourceOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSource, - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict, +ReasoningEngineTrafficConfigTrafficSplitManualTargetOrDict = Union[ + ReasoningEngineTrafficConfigTrafficSplitManualTarget, + ReasoningEngineTrafficConfigTrafficSplitManualTargetDict, ] -class ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig(_common.BaseModel): - """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. - - This includes the repository, revision, and directory to use. - """ +class ReasoningEngineTrafficConfigTrafficSplitManual(_common.BaseModel): + """Manual traffic distribution configuration, where the user specifies the Runtime Revision IDs and the percentage of traffic to send to each.""" - git_repository_link: Optional[str] = Field( - default=None, - description="""Required. The Developer Connect Git repository link, formatted as `projects/{project_id}/locations/{location_id}/connections/{connection_id}/gitRepositoryLink/{repository_link_id}`.""", - ) - dir: Optional[str] = Field( - default=None, - description="""Required. Directory, relative to the source root, in which to run the build.""", - ) - revision: Optional[str] = Field( - default=None, - description="""Required. The revision to fetch from the Git repository such as a branch, a tag, a commit SHA, or any Git ref.""", + targets: Optional[list[ReasoningEngineTrafficConfigTrafficSplitManualTarget]] = ( + Field( + default=None, + description="""A list of traffic targets for the Runtimes Revisions. The sum of percentages must equal to 100.""", + ) ) -class ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict( - TypedDict, total=False -): - """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. - - This includes the repository, revision, and directory to use. - """ - - git_repository_link: Optional[str] - """Required. The Developer Connect Git repository link, formatted as `projects/{project_id}/locations/{location_id}/connections/{connection_id}/gitRepositoryLink/{repository_link_id}`.""" - - dir: Optional[str] - """Required. Directory, relative to the source root, in which to run the build.""" +class ReasoningEngineTrafficConfigTrafficSplitManualDict(TypedDict, total=False): + """Manual traffic distribution configuration, where the user specifies the Runtime Revision IDs and the percentage of traffic to send to each.""" - revision: Optional[str] - """Required. The revision to fetch from the Git repository such as a branch, a tag, a commit SHA, or any Git ref.""" + targets: Optional[list[ReasoningEngineTrafficConfigTrafficSplitManualTargetDict]] + """A list of traffic targets for the Runtimes Revisions. The sum of percentages must equal to 100.""" -ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig, - ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict, +ReasoningEngineTrafficConfigTrafficSplitManualOrDict = Union[ + ReasoningEngineTrafficConfigTrafficSplitManual, + ReasoningEngineTrafficConfigTrafficSplitManualDict, ] -class ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource(_common.BaseModel): - """Specifies source code to be fetched from a Git repository managed through the Developer Connect service.""" +class ReasoningEngineTrafficConfig(_common.BaseModel): + """Traffic distribution configuration.""" - config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig] = Field( + traffic_split_always_latest: Optional[ + ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest + ] = Field( default=None, - description="""Required. The Developer Connect configuration that defines the specific repository, revision, and directory to use as the source code root.""", + description="""Optional. Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""", + ) + traffic_split_manual: Optional[ReasoningEngineTrafficConfigTrafficSplitManual] = ( + Field( + default=None, + description="""Optional. Manual traffic distribution configuration, where the user specifies the Runtime Revision IDs and the percentage of traffic to send to each.""", + ) ) -class ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict( - TypedDict, total=False -): - """Specifies source code to be fetched from a Git repository managed through the Developer Connect service.""" +class ReasoningEngineTrafficConfigDict(TypedDict, total=False): + """Traffic distribution configuration.""" - config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict] - """Required. The Developer Connect configuration that defines the specific repository, revision, and directory to use as the source code root.""" + traffic_split_always_latest: Optional[ + ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict + ] + """Optional. Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" + traffic_split_manual: Optional[ReasoningEngineTrafficConfigTrafficSplitManualDict] + """Optional. Manual traffic distribution configuration, where the user specifies the Runtime Revision IDs and the percentage of traffic to send to each.""" -ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource, - ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict, + +ReasoningEngineTrafficConfigOrDict = Union[ + ReasoningEngineTrafficConfig, ReasoningEngineTrafficConfigDict ] -class ReasoningEngineSpecSourceCodeSpecImageSpec(_common.BaseModel): - """The image spec for building an image (within a single build step), based on the config file (i.e. Dockerfile) in the source directory.""" +class ReasoningEngine(_common.BaseModel): + """An agent runtime.""" - build_args: Optional[dict[str, str]] = Field( + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, - description="""Optional. Build arguments to be used. They will be passed through --build-arg flags.""", + description="""Customer-managed encryption key spec for a ReasoningEngine. If set, this ReasoningEngine and all sub-resources of this ReasoningEngine will be secured by this key.""", ) - - -class ReasoningEngineSpecSourceCodeSpecImageSpecDict(TypedDict, total=False): - """The image spec for building an image (within a single build step), based on the config file (i.e. Dockerfile) in the source directory.""" - - build_args: Optional[dict[str, str]] - """Optional. Build arguments to be used. They will be passed through --build-arg flags.""" - - -ReasoningEngineSpecSourceCodeSpecImageSpecOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecImageSpec, - ReasoningEngineSpecSourceCodeSpecImageSpecDict, -] - - -class ReasoningEngineSpecSourceCodeSpecPythonSpec(_common.BaseModel): - """Specification for running a Python application from source.""" - - entrypoint_module: Optional[str] = Field( + context_spec: Optional[ReasoningEngineContextSpec] = Field( default=None, - description="""Optional. The Python module to load as the entrypoint, specified as a fully qualified module name. For example: path.to.agent. If not specified, defaults to "agent". The project root will be added to Python sys.path, allowing imports to be specified relative to the root. This field should not be set if the source is `agent_config_source`.""", + description="""Optional. Configuration for how Agent Engine sub-resources should manage context.""", ) - entrypoint_object: Optional[str] = Field( + create_time: Optional[datetime.datetime] = Field( default=None, - description="""Optional. The name of the callable object within the `entrypoint_module` to use as the application If not specified, defaults to "root_agent". This field should not be set if the source is `agent_config_source`.""", + description="""Output only. Timestamp when this ReasoningEngine was created.""", ) - requirements_file: Optional[str] = Field( + description: Optional[str] = Field( default=None, - description="""Optional. The path to the requirements file, relative to the source root. If not specified, defaults to "requirements.txt".""", + description="""Optional. The description of the ReasoningEngine.""", ) - version: Optional[str] = Field( + display_name: Optional[str] = Field( default=None, - description="""Optional. The version of Python to use. Supported versions include 3.10, 3.11, 3.12, 3.13, 3.14. If not specified, default value is 3.10.""", + description="""Required. The display name of the ReasoningEngine.""", + ) + etag: Optional[str] = Field( + default=None, + description="""Optional. Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""", + ) + labels: Optional[dict[str, str]] = Field( + default=None, description="""Labels for the ReasoningEngine.""" + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. The resource name of the ReasoningEngine. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`""", + ) + spec: Optional[ReasoningEngineSpec] = Field( + default=None, description="""Optional. Configurations of the ReasoningEngine""" + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this ReasoningEngine was most recently updated.""", + ) + traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( + default=None, + description="""Optional. Traffic distribution configuration for the Reasoning Engine.""", ) -class ReasoningEngineSpecSourceCodeSpecPythonSpecDict(TypedDict, total=False): - """Specification for running a Python application from source.""" +class ReasoningEngineDict(TypedDict, total=False): + """An agent runtime.""" - entrypoint_module: Optional[str] - """Optional. The Python module to load as the entrypoint, specified as a fully qualified module name. For example: path.to.agent. If not specified, defaults to "agent". The project root will be added to Python sys.path, allowing imports to be specified relative to the root. This field should not be set if the source is `agent_config_source`.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for a ReasoningEngine. If set, this ReasoningEngine and all sub-resources of this ReasoningEngine will be secured by this key.""" - entrypoint_object: Optional[str] - """Optional. The name of the callable object within the `entrypoint_module` to use as the application If not specified, defaults to "root_agent". This field should not be set if the source is `agent_config_source`.""" + context_spec: Optional[ReasoningEngineContextSpecDict] + """Optional. Configuration for how Agent Engine sub-resources should manage context.""" - requirements_file: Optional[str] - """Optional. The path to the requirements file, relative to the source root. If not specified, defaults to "requirements.txt".""" + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this ReasoningEngine was created.""" - version: Optional[str] - """Optional. The version of Python to use. Supported versions include 3.10, 3.11, 3.12, 3.13, 3.14. If not specified, default value is 3.10.""" + description: Optional[str] + """Optional. The description of the ReasoningEngine.""" + display_name: Optional[str] + """Required. The display name of the ReasoningEngine.""" -ReasoningEngineSpecSourceCodeSpecPythonSpecOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecPythonSpec, - ReasoningEngineSpecSourceCodeSpecPythonSpecDict, -] + etag: Optional[str] + """Optional. Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""" + labels: Optional[dict[str, str]] + """Labels for the ReasoningEngine.""" -class ReasoningEngineSpecSourceCodeSpec(_common.BaseModel): - """Specification for deploying from source code.""" + name: Optional[str] + """Identifier. The resource name of the ReasoningEngine. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`""" - agent_config_source: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSource - ] = Field( - default=None, description="""Source code is generated from the agent config.""" + spec: Optional[ReasoningEngineSpecDict] + """Optional. Configurations of the ReasoningEngine""" + + update_time: Optional[datetime.datetime] + """Output only. Timestamp when this ReasoningEngine was most recently updated.""" + + traffic_config: Optional[ReasoningEngineTrafficConfigDict] + """Optional. Traffic distribution configuration for the Reasoning Engine.""" + + +ReasoningEngineOrDict = Union[ReasoningEngine, ReasoningEngineDict] + + +class RuntimeOperation(_common.BaseModel): + """Operation that has an agent runtime as a response.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - developer_connect_source: Optional[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource - ] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Source code is in a Git repository managed by Developer Connect.""", + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", ) - image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpec] = Field( + done: Optional[bool] = Field( default=None, - description="""Optional. Configuration for building an image with custom config file.""", + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", ) - inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSource] = Field( - default=None, description="""Source code is provided directly in the request.""" + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", ) - python_spec: Optional[ReasoningEngineSpecSourceCodeSpecPythonSpec] = Field( - default=None, description="""Configuration for a Python application.""" + response: Optional[ReasoningEngine] = Field( + default=None, description="""The created Agent Runtime.""" ) -class ReasoningEngineSpecSourceCodeSpecDict(TypedDict, total=False): - """Specification for deploying from source code.""" +class RuntimeOperationDict(TypedDict, total=False): + """Operation that has an agent runtime as a response.""" - agent_config_source: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict - ] - """Source code is generated from the agent config.""" + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - developer_connect_source: Optional[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict - ] - """Source code is in a Git repository managed by Developer Connect.""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict] - """Optional. Configuration for building an image with custom config file.""" + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSourceDict] - """Source code is provided directly in the request.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - python_spec: Optional[ReasoningEngineSpecSourceCodeSpecPythonSpecDict] - """Configuration for a Python application.""" + response: Optional[ReasoningEngineDict] + """The created Agent Runtime.""" -ReasoningEngineSpecSourceCodeSpecOrDict = Union[ - ReasoningEngineSpecSourceCodeSpec, ReasoningEngineSpecSourceCodeSpecDict -] +RuntimeOperationOrDict = Union[RuntimeOperation, RuntimeOperationDict] -class ReasoningEngineSpecContainerSpec(_common.BaseModel): - """Specification for deploying from a container image.""" +class CreateRuntimeConfig(_common.BaseModel): + """Config for create agent runtime.""" - image_uri: Optional[str] = Field( - default=None, - description="""Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) + display_name: Optional[str] = Field( + default=None, + description="""The user-defined name of the Agent Runtime. - -class ReasoningEngineSpecContainerSpecDict(TypedDict, total=False): - """Specification for deploying from a container image.""" - - image_uri: Optional[str] - """Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""" - - -ReasoningEngineSpecContainerSpecOrDict = Union[ - ReasoningEngineSpecContainerSpec, ReasoningEngineSpecContainerSpecDict -] - - -class ReasoningEngineSpec(_common.BaseModel): - """The specification of an agent engine.""" - - agent_card: Optional[dict[str, Any]] = Field( + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """, + ) + description: Optional[str] = Field( + default=None, description="""The description of the Agent Runtime.""" + ) + spec: Optional[ReasoningEngineSpec] = Field( + default=None, description="""Optional. Configurations of the Agent Runtime.""" + ) + context_spec: Optional[ReasoningEngineContextSpec] = Field( default=None, - description="""Optional. The A2A Agent Card for the agent (if available). It follows the specification at https://a2a-protocol.org/latest/specification/#5-agent-discovery-the-agent-card.""", + description="""Optional. The context spec to be used for the Agent Runtime.""", ) - agent_framework: Optional[str] = Field( + psc_interface_config: Optional[PscInterfaceConfig] = Field( default=None, - description="""Optional. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom".""", + description="""Optional. The PSC interface config for PSC-I to be used for the + Agent Runtime.""", + ) + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + ] = Field( + default=None, + description="""Agent Gateway configuration for a Reasoning Engine deployment.""", + ) + min_instances: Optional[int] = Field( + default=None, + description="""The minimum number of instances to run for the Agent Runtime. + Defaults to 1. Range: [0, 10]. + """, + ) + max_instances: Optional[int] = Field( + default=None, + description="""The maximum number of instances to run for the Agent Runtime. + Defaults to 100. Range: [1, 1000]. + If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. + """, + ) + resource_limits: Optional[dict[str, str]] = Field( + default=None, + description="""The resource limits to be applied to the Agent Runtime. + Required keys: 'cpu' and 'memory'. + Supported values for 'cpu': '1', '2', '4', '6', '8'. + Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. + """, + ) + container_concurrency: Optional[int] = Field( + default=None, + description="""The container concurrency to be used for the Agent Runtime. + Recommended value: 2 * cpu + 1. Defaults to 9. + """, + ) + keep_alive_probe: Optional[KeepAliveProbe] = Field( + default=None, + description="""Optional. Specifies the configuration for keep-alive probe. + Contains configuration on a specified endpoint that a deployment host + should use to keep the container alive based on the probe settings.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""The encryption spec to be used for the Agent Runtime.""", + ) + labels: Optional[dict[str, str]] = Field( + default=None, description="""The labels to be used for the Agent Runtime.""" ) class_methods: Optional[list[dict[str, Any]]] = Field( default=None, - description="""Optional. Declarations for object class methods in OpenAPI specification format.""", + description="""The class methods to be used for the Agent Runtime. + If specified, they'll override the class methods that are autogenerated by + default. By default, methods are generated by inspecting the agent object + and generating a corresponding method for each method defined on the + agent class. + """, ) - deployment_spec: Optional[ReasoningEngineSpecDeploymentSpec] = Field( + source_packages: Optional[list[str]] = Field( default=None, - description="""Optional. The specification of a Reasoning Engine deployment.""", + description="""The user-provided paths to the source packages (if any). + If specified, the files in the source packages will be packed into a + a tarball file, uploaded to Agent Runtime's API, and deployed to the + Agent Runtime. + The following fields will be ignored: + - agent + - extra_packages + - staging_bucket + - requirements + The following fields will be used to install and use the agent from the + source packages: + - entrypoint_module (required) + - entrypoint_object (required) + - requirements_file (optional) + - class_methods (required) + """, ) - effective_identity: Optional[str] = Field( + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig + ] = Field( default=None, - description="""Output only. The identity to use for the Reasoning Engine. It can contain one of the following values: * service-{project}@gcp-sa-aiplatform-re.googleapis.com (for SERVICE_AGENT identity type) * {name}@{project}.gserviceaccount.com (for SERVICE_ACCOUNT identity type) * agents.global.{org}.system.id.goog/resources/aiplatform/projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine} (for AGENT_IDENTITY identity type)""", + description="""Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""", ) - identity_type: Optional[IdentityType] = Field( + entrypoint_module: Optional[str] = Field( default=None, - description="""Optional. The identity type to use for the Reasoning Engine. If not specified, the `service_account` field will be used if set, otherwise the default Vertex AI Reasoning Engine Service Agent in the project will be used.""", + description="""The entrypoint module to be used for the Agent Runtime + This field only used when source_packages is specified.""", ) - package_spec: Optional[ReasoningEngineSpecPackageSpec] = Field( + entrypoint_object: Optional[str] = Field( default=None, - description="""Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes. The `deployment_source` field should not be set if `package_spec` is specified.""", + description="""The entrypoint object to be used for the Agent Runtime. + This field only used when source_packages is specified.""", ) - service_account: Optional[str] = Field( + requirements_file: Optional[str] = Field( default=None, - description="""Optional. The service account that the Reasoning Engine artifact runs as. It should have "roles/storage.objectViewer" for reading the user project's Cloud Storage and "roles/aiplatform.user" for using Vertex extensions. If not specified, the Vertex AI Reasoning Engine Service Agent in the project will be used.""", + description="""The user-provided path to the requirements file (if any). + This field is only used when source_packages is specified. + If not specified, agent runtime will find and use the `requirements.txt` in + the source package. + """, ) - source_code_spec: Optional[ReasoningEngineSpecSourceCodeSpec] = Field( + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + ] = Field( default=None, - description="""Deploy from source code files with a defined entrypoint.""", + description="""The agent framework to be used for the Agent Runtime. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""", ) - container_spec: Optional[ReasoningEngineSpecContainerSpec] = Field( + python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] = Field( default=None, - description="""Deploy from a container image with a defined entrypoint and commands.""", + description="""The Python version to be used for the Agent Runtime. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". + """, + ) + build_options: Optional[dict[str, list[str]]] = Field( + default=None, + description="""The build options for the Agent Runtime. + The following keys are supported: + - installation_scripts: + Optional. The paths to the installation scripts to be + executed in the Docker image. + The scripts must be located in the `installation_scripts` + subdirectory and the path must be added to `extra_packages`. + """, ) -class ReasoningEngineSpecDict(TypedDict, total=False): - """The specification of an agent engine.""" - - agent_card: Optional[dict[str, Any]] - """Optional. The A2A Agent Card for the agent (if available). It follows the specification at https://a2a-protocol.org/latest/specification/#5-agent-discovery-the-agent-card.""" +class CreateRuntimeConfigDict(TypedDict, total=False): + """Config for create agent runtime.""" - agent_framework: Optional[str] - """Optional. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom".""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - class_methods: Optional[list[dict[str, Any]]] - """Optional. Declarations for object class methods in OpenAPI specification format.""" + display_name: Optional[str] + """The user-defined name of the Agent Runtime. - deployment_spec: Optional[ReasoningEngineSpecDeploymentSpecDict] - """Optional. The specification of a Reasoning Engine deployment.""" + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """ - effective_identity: Optional[str] - """Output only. The identity to use for the Reasoning Engine. It can contain one of the following values: * service-{project}@gcp-sa-aiplatform-re.googleapis.com (for SERVICE_AGENT identity type) * {name}@{project}.gserviceaccount.com (for SERVICE_ACCOUNT identity type) * agents.global.{org}.system.id.goog/resources/aiplatform/projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine} (for AGENT_IDENTITY identity type)""" + description: Optional[str] + """The description of the Agent Runtime.""" - identity_type: Optional[IdentityType] - """Optional. The identity type to use for the Reasoning Engine. If not specified, the `service_account` field will be used if set, otherwise the default Vertex AI Reasoning Engine Service Agent in the project will be used.""" + spec: Optional[ReasoningEngineSpecDict] + """Optional. Configurations of the Agent Runtime.""" - package_spec: Optional[ReasoningEngineSpecPackageSpecDict] - """Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes. The `deployment_source` field should not be set if `package_spec` is specified.""" + context_spec: Optional[ReasoningEngineContextSpecDict] + """Optional. The context spec to be used for the Agent Runtime.""" - service_account: Optional[str] - """Optional. The service account that the Reasoning Engine artifact runs as. It should have "roles/storage.objectViewer" for reading the user project's Cloud Storage and "roles/aiplatform.user" for using Vertex extensions. If not specified, the Vertex AI Reasoning Engine Service Agent in the project will be used.""" + psc_interface_config: Optional[PscInterfaceConfigDict] + """Optional. The PSC interface config for PSC-I to be used for the + Agent Runtime.""" - source_code_spec: Optional[ReasoningEngineSpecSourceCodeSpecDict] - """Deploy from source code files with a defined entrypoint.""" + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] + """Agent Gateway configuration for a Reasoning Engine deployment.""" - container_spec: Optional[ReasoningEngineSpecContainerSpecDict] - """Deploy from a container image with a defined entrypoint and commands.""" + min_instances: Optional[int] + """The minimum number of instances to run for the Agent Runtime. + Defaults to 1. Range: [0, 10]. + """ + max_instances: Optional[int] + """The maximum number of instances to run for the Agent Runtime. + Defaults to 100. Range: [1, 1000]. + If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. + """ -ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict] + resource_limits: Optional[dict[str, str]] + """The resource limits to be applied to the Agent Runtime. + Required keys: 'cpu' and 'memory'. + Supported values for 'cpu': '1', '2', '4', '6', '8'. + Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. + """ + container_concurrency: Optional[int] + """The container concurrency to be used for the Agent Runtime. + Recommended value: 2 * cpu + 1. Defaults to 9. + """ -class ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest(_common.BaseModel): - """Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" + keep_alive_probe: Optional[KeepAliveProbeDict] + """Optional. Specifies the configuration for keep-alive probe. + Contains configuration on a specified endpoint that a deployment host + should use to keep the container alive based on the probe settings.""" - pass + encryption_spec: Optional[genai_types.EncryptionSpec] + """The encryption spec to be used for the Agent Runtime.""" + labels: Optional[dict[str, str]] + """The labels to be used for the Agent Runtime.""" -class ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict(TypedDict, total=False): - """Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" + class_methods: Optional[list[dict[str, Any]]] + """The class methods to be used for the Agent Runtime. + If specified, they'll override the class methods that are autogenerated by + default. By default, methods are generated by inspecting the agent object + and generating a corresponding method for each method defined on the + agent class. + """ - pass + source_packages: Optional[list[str]] + """The user-provided paths to the source packages (if any). + If specified, the files in the source packages will be packed into a + a tarball file, uploaded to Agent Runtime's API, and deployed to the + Agent Runtime. + The following fields will be ignored: + - agent + - extra_packages + - staging_bucket + - requirements + The following fields will be used to install and use the agent from the + source packages: + - entrypoint_module (required) + - entrypoint_object (required) + - requirements_file (optional) + - class_methods (required) + """ + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict + ] + """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""" -ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestOrDict = Union[ - ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest, - ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict, -] + entrypoint_module: Optional[str] + """The entrypoint module to be used for the Agent Runtime + This field only used when source_packages is specified.""" + entrypoint_object: Optional[str] + """The entrypoint object to be used for the Agent Runtime. + This field only used when source_packages is specified.""" -class ReasoningEngineTrafficConfigTrafficSplitManualTarget(_common.BaseModel): - """A single target for the traffic split, specifying a Runtime Revision and the percentage of traffic to send to it.""" + requirements_file: Optional[str] + """The user-provided path to the requirements file (if any). + This field is only used when source_packages is specified. + If not specified, agent runtime will find and use the `requirements.txt` in + the source package. + """ - percent: Optional[int] = Field( - default=None, - description="""Required. Specifies percent of the traffic to this Runtime Revision.""", - ) - runtime_revision_name: Optional[str] = Field( - default=None, - description="""Required. The Runtime Revision name to which to send this portion of traffic, if traffic allocation is by Runtime Revision.""", - ) - - -class ReasoningEngineTrafficConfigTrafficSplitManualTargetDict(TypedDict, total=False): - """A single target for the traffic split, specifying a Runtime Revision and the percentage of traffic to send to it.""" + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + ] + """The agent framework to be used for the Agent Runtime. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""" - percent: Optional[int] - """Required. Specifies percent of the traffic to this Runtime Revision.""" + python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] + """The Python version to be used for the Agent Runtime. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". + """ - runtime_revision_name: Optional[str] - """Required. The Runtime Revision name to which to send this portion of traffic, if traffic allocation is by Runtime Revision.""" + build_options: Optional[dict[str, list[str]]] + """The build options for the Agent Runtime. + The following keys are supported: + - installation_scripts: + Optional. The paths to the installation scripts to be + executed in the Docker image. + The scripts must be located in the `installation_scripts` + subdirectory and the path must be added to `extra_packages`. + """ -ReasoningEngineTrafficConfigTrafficSplitManualTargetOrDict = Union[ - ReasoningEngineTrafficConfigTrafficSplitManualTarget, - ReasoningEngineTrafficConfigTrafficSplitManualTargetDict, -] +CreateRuntimeConfigOrDict = Union[CreateRuntimeConfig, CreateRuntimeConfigDict] -class ReasoningEngineTrafficConfigTrafficSplitManual(_common.BaseModel): - """Manual traffic distribution configuration, where the user specifies the Runtime Revision IDs and the percentage of traffic to send to each.""" +class _CreateRuntimeRequestParameters(_common.BaseModel): + """Parameters for creating agent runtimes.""" - targets: Optional[list[ReasoningEngineTrafficConfigTrafficSplitManualTarget]] = ( - Field( - default=None, - description="""A list of traffic targets for the Runtimes Revisions. The sum of percentages must equal to 100.""", - ) - ) + config: Optional[CreateRuntimeConfig] = Field(default=None, description="""""") -class ReasoningEngineTrafficConfigTrafficSplitManualDict(TypedDict, total=False): - """Manual traffic distribution configuration, where the user specifies the Runtime Revision IDs and the percentage of traffic to send to each.""" +class _CreateRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for creating agent runtimes.""" - targets: Optional[list[ReasoningEngineTrafficConfigTrafficSplitManualTargetDict]] - """A list of traffic targets for the Runtimes Revisions. The sum of percentages must equal to 100.""" + config: Optional[CreateRuntimeConfigDict] + """""" -ReasoningEngineTrafficConfigTrafficSplitManualOrDict = Union[ - ReasoningEngineTrafficConfigTrafficSplitManual, - ReasoningEngineTrafficConfigTrafficSplitManualDict, +_CreateRuntimeRequestParametersOrDict = Union[ + _CreateRuntimeRequestParameters, _CreateRuntimeRequestParametersDict ] -class ReasoningEngineTrafficConfig(_common.BaseModel): - """Traffic distribution configuration.""" +class DeleteRuntimeConfig(_common.BaseModel): + """Config for deleting agent runtime.""" - traffic_split_always_latest: Optional[ - ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest - ] = Field( - default=None, - description="""Optional. Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""", - ) - traffic_split_manual: Optional[ReasoningEngineTrafficConfigTrafficSplitManual] = ( - Field( - default=None, - description="""Optional. Manual traffic distribution configuration, where the user specifies the Runtime Revision IDs and the percentage of traffic to send to each.""", - ) + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class ReasoningEngineTrafficConfigDict(TypedDict, total=False): - """Traffic distribution configuration.""" - - traffic_split_always_latest: Optional[ - ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict - ] - """Optional. Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" +class DeleteRuntimeConfigDict(TypedDict, total=False): + """Config for deleting agent runtime.""" - traffic_split_manual: Optional[ReasoningEngineTrafficConfigTrafficSplitManualDict] - """Optional. Manual traffic distribution configuration, where the user specifies the Runtime Revision IDs and the percentage of traffic to send to each.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ReasoningEngineTrafficConfigOrDict = Union[ - ReasoningEngineTrafficConfig, ReasoningEngineTrafficConfigDict -] +DeleteRuntimeConfigOrDict = Union[DeleteRuntimeConfig, DeleteRuntimeConfigDict] -class ReasoningEngine(_common.BaseModel): - """An agent engine.""" +class _DeleteRuntimeRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtimes.""" - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, - description="""Customer-managed encryption key spec for a ReasoningEngine. If set, this ReasoningEngine and all sub-resources of this ReasoningEngine will be secured by this key.""", - ) - context_spec: Optional[ReasoningEngineContextSpec] = Field( - default=None, - description="""Optional. Configuration for how Agent Engine sub-resources should manage context.""", - ) - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when this ReasoningEngine was created.""", - ) - description: Optional[str] = Field( - default=None, - description="""Optional. The description of the ReasoningEngine.""", - ) - display_name: Optional[str] = Field( - default=None, - description="""Required. The display name of the ReasoningEngine.""", - ) - etag: Optional[str] = Field( - default=None, - description="""Optional. Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""", - ) - labels: Optional[dict[str, str]] = Field( - default=None, description="""Labels for the ReasoningEngine.""" - ) name: Optional[str] = Field( - default=None, - description="""Identifier. The resource name of the ReasoningEngine. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`""", - ) - spec: Optional[ReasoningEngineSpec] = Field( - default=None, description="""Optional. Configurations of the ReasoningEngine""" - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when this ReasoningEngine was most recently updated.""", + default=None, description="""Name of the agent runtime.""" ) - traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( - default=None, - description="""Optional. Traffic distribution configuration for the Reasoning Engine.""", + force: Optional[bool] = Field( + default=False, + description="""If set to true, any child resources will also be deleted.""", ) + config: Optional[DeleteRuntimeConfig] = Field(default=None, description="""""") -class ReasoningEngineDict(TypedDict, total=False): - """An agent engine.""" - - encryption_spec: Optional[genai_types.EncryptionSpec] - """Customer-managed encryption key spec for a ReasoningEngine. If set, this ReasoningEngine and all sub-resources of this ReasoningEngine will be secured by this key.""" - - context_spec: Optional[ReasoningEngineContextSpecDict] - """Optional. Configuration for how Agent Engine sub-resources should manage context.""" - - create_time: Optional[datetime.datetime] - """Output only. Timestamp when this ReasoningEngine was created.""" - - description: Optional[str] - """Optional. The description of the ReasoningEngine.""" - - display_name: Optional[str] - """Required. The display name of the ReasoningEngine.""" - - etag: Optional[str] - """Optional. Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""" - - labels: Optional[dict[str, str]] - """Labels for the ReasoningEngine.""" +class _DeleteRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtimes.""" name: Optional[str] - """Identifier. The resource name of the ReasoningEngine. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`""" - - spec: Optional[ReasoningEngineSpecDict] - """Optional. Configurations of the ReasoningEngine""" + """Name of the agent runtime.""" - update_time: Optional[datetime.datetime] - """Output only. Timestamp when this ReasoningEngine was most recently updated.""" + force: Optional[bool] + """If set to true, any child resources will also be deleted.""" - traffic_config: Optional[ReasoningEngineTrafficConfigDict] - """Optional. Traffic distribution configuration for the Reasoning Engine.""" + config: Optional[DeleteRuntimeConfigDict] + """""" -ReasoningEngineOrDict = Union[ReasoningEngine, ReasoningEngineDict] +_DeleteRuntimeRequestParametersOrDict = Union[ + _DeleteRuntimeRequestParameters, _DeleteRuntimeRequestParametersDict +] -class AgentEngineOperation(_common.BaseModel): - """Operation that has an agent engine as a response.""" +class DeleteRuntimeOperation(_common.BaseModel): + """Operation for deleting agent runtimes.""" name: Optional[str] = Field( default=None, @@ -9145,13 +8654,10 @@ class AgentEngineOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) - response: Optional[ReasoningEngine] = Field( - default=None, description="""The created Agent Engine.""" - ) -class AgentEngineOperationDict(TypedDict, total=False): - """Operation that has an agent engine as a response.""" +class DeleteRuntimeOperationDict(TypedDict, total=False): + """Operation for deleting agent runtimes.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -9165,79 +8671,339 @@ class AgentEngineOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" - response: Optional[ReasoningEngineDict] - """The created Agent Engine.""" - -AgentEngineOperationOrDict = Union[AgentEngineOperation, AgentEngineOperationDict] +DeleteRuntimeOperationOrDict = Union[DeleteRuntimeOperation, DeleteRuntimeOperationDict] -class CreateAgentEngineConfig(_common.BaseModel): - """Config for create agent engine.""" +class GetRuntimeConfig(_common.BaseModel): + """Config for create agent runtime.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - display_name: Optional[str] = Field( - default=None, - description="""The user-defined name of the Agent Engine. - The display name can be up to 128 characters long and can comprise any - UTF-8 characters. - """, - ) - description: Optional[str] = Field( - default=None, description="""The description of the Agent Engine.""" - ) - spec: Optional[ReasoningEngineSpec] = Field( - default=None, description="""Optional. Configurations of the Agent Engine.""" - ) - context_spec: Optional[ReasoningEngineContextSpec] = Field( - default=None, - description="""Optional. The context spec to be used for the Agent Engine.""", + +class GetRuntimeConfigDict(TypedDict, total=False): + """Config for create agent runtime.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetRuntimeConfigOrDict = Union[GetRuntimeConfig, GetRuntimeConfigDict] + + +class _GetRuntimeRequestParameters(_common.BaseModel): + """Parameters for getting agent runtimes.""" + + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime.""" ) - psc_interface_config: Optional[PscInterfaceConfig] = Field( - default=None, - description="""Optional. The PSC interface config for PSC-I to be used for the - Agent Engine.""", + config: Optional[GetRuntimeConfig] = Field(default=None, description="""""") + + +class _GetRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for getting agent runtimes.""" + + name: Optional[str] + """Name of the agent runtime.""" + + config: Optional[GetRuntimeConfigDict] + """""" + + +_GetRuntimeRequestParametersOrDict = Union[ + _GetRuntimeRequestParameters, _GetRuntimeRequestParametersDict +] + + +class ListRuntimeConfig(_common.BaseModel): + """Config for listing agent runtimes.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - min_instances: Optional[int] = Field( + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( default=None, - description="""The minimum number of instances to run for the Agent Engine. - Defaults to 1. Range: [0, 10]. - """, + description="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""", ) - max_instances: Optional[int] = Field( - default=None, - description="""The maximum number of instances to run for the Agent Engine. - Defaults to 100. Range: [1, 1000]. - If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. - """, + + +class ListRuntimeConfigDict(TypedDict, total=False): + """Config for listing agent runtimes.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + page_size: Optional[int] + """""" + + page_token: Optional[str] + """""" + + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""" + + +ListRuntimeConfigOrDict = Union[ListRuntimeConfig, ListRuntimeConfigDict] + + +class _ListRuntimeRequestParameters(_common.BaseModel): + """Parameters for listing agent runtimes.""" + + config: Optional[ListRuntimeConfig] = Field(default=None, description="""""") + + +class _ListRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtimes.""" + + config: Optional[ListRuntimeConfigDict] + """""" + + +_ListRuntimeRequestParametersOrDict = Union[ + _ListRuntimeRequestParameters, _ListRuntimeRequestParametersDict +] + + +class ListReasoningEnginesResponse(_common.BaseModel): + """Response for listing agent runtimes.""" + + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - resource_limits: Optional[dict[str, str]] = Field( + next_page_token: Optional[str] = Field(default=None, description="""""") + reasoning_engines: Optional[list[ReasoningEngine]] = Field( default=None, - description="""The resource limits to be applied to the Agent Engine. - Required keys: 'cpu' and 'memory'. - Supported values for 'cpu': '1', '2', '4', '6', '8'. - Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. + description="""List of agent runtimes. + """, + ) + + +class ListReasoningEnginesResponseDict(TypedDict, total=False): + """Response for listing agent runtimes.""" + + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + + next_page_token: Optional[str] + """""" + + reasoning_engines: Optional[list[ReasoningEngineDict]] + """List of agent runtimes. + """ + + +ListReasoningEnginesResponseOrDict = Union[ + ListReasoningEnginesResponse, ListReasoningEnginesResponseDict +] + + +class GetRuntimeOperationConfig(_common.BaseModel): + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetRuntimeOperationConfigDict(TypedDict, total=False): + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetRuntimeOperationConfigOrDict = Union[ + GetRuntimeOperationConfig, GetRuntimeOperationConfigDict +] + + +class _GetRuntimeOperationParameters(_common.BaseModel): + """Parameters for getting an operation with an agent runtime as a response.""" + + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetRuntimeOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" + ) + + +class _GetRuntimeOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with an agent runtime as a response.""" + + operation_name: Optional[str] + """The server-assigned name for the operation.""" + + config: Optional[GetRuntimeOperationConfigDict] + """Used to override the default configuration.""" + + +_GetRuntimeOperationParametersOrDict = Union[ + _GetRuntimeOperationParameters, _GetRuntimeOperationParametersDict +] + + +class QueryRuntimeConfig(_common.BaseModel): + """Config for querying agent runtimes.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + class_method: Optional[str] = Field( + default=None, description="""The class method to call.""" + ) + input: Optional[dict[str, Any]] = Field( + default=None, description="""The input to the class method.""" + ) + include_all_fields: Optional[bool] = Field(default=False, description="""""") + + +class QueryRuntimeConfigDict(TypedDict, total=False): + """Config for querying agent runtimes.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + class_method: Optional[str] + """The class method to call.""" + + input: Optional[dict[str, Any]] + """The input to the class method.""" + + include_all_fields: Optional[bool] + """""" + + +QueryRuntimeConfigOrDict = Union[QueryRuntimeConfig, QueryRuntimeConfigDict] + + +class _QueryRuntimeRequestParameters(_common.BaseModel): + """Parameters for querying agent runtimes.""" + + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime.""" + ) + config: Optional[QueryRuntimeConfig] = Field(default=None, description="""""") + + +class _QueryRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for querying agent runtimes.""" + + name: Optional[str] + """Name of the agent runtime.""" + + config: Optional[QueryRuntimeConfigDict] + """""" + + +_QueryRuntimeRequestParametersOrDict = Union[ + _QueryRuntimeRequestParameters, _QueryRuntimeRequestParametersDict +] + + +class QueryReasoningEngineResponse(_common.BaseModel): + """The response for querying an agent runtime.""" + + output: Optional[Any] = Field( + default=None, + description="""Response provided by users in JSON object format.""", + ) + + +class QueryReasoningEngineResponseDict(TypedDict, total=False): + """The response for querying an agent runtime.""" + + output: Optional[Any] + """Response provided by users in JSON object format.""" + + +QueryReasoningEngineResponseOrDict = Union[ + QueryReasoningEngineResponse, QueryReasoningEngineResponseDict +] + + +class UpdateRuntimeConfig(_common.BaseModel): + """Config for updating agent runtime.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + display_name: Optional[str] = Field( + default=None, + description="""The user-defined name of the Agent Runtime. + + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """, + ) + description: Optional[str] = Field( + default=None, description="""The description of the Agent Runtime.""" + ) + spec: Optional[ReasoningEngineSpec] = Field( + default=None, description="""Optional. Configurations of the Agent Runtime.""" + ) + context_spec: Optional[ReasoningEngineContextSpec] = Field( + default=None, + description="""Optional. The context spec to be used for the Agent Runtime.""", + ) + psc_interface_config: Optional[PscInterfaceConfig] = Field( + default=None, + description="""Optional. The PSC interface config for PSC-I to be used for the + Agent Runtime.""", + ) + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + ] = Field( + default=None, + description="""Agent Gateway configuration for a Reasoning Engine deployment.""", + ) + min_instances: Optional[int] = Field( + default=None, + description="""The minimum number of instances to run for the Agent Runtime. + Defaults to 1. Range: [0, 10]. + """, + ) + max_instances: Optional[int] = Field( + default=None, + description="""The maximum number of instances to run for the Agent Runtime. + Defaults to 100. Range: [1, 1000]. + If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. + """, + ) + resource_limits: Optional[dict[str, str]] = Field( + default=None, + description="""The resource limits to be applied to the Agent Runtime. + Required keys: 'cpu' and 'memory'. + Supported values for 'cpu': '1', '2', '4', '6', '8'. + Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. """, ) container_concurrency: Optional[int] = Field( default=None, - description="""The container concurrency to be used for the Agent Engine. + description="""The container concurrency to be used for the Agent Runtime. Recommended value: 2 * cpu + 1. Defaults to 9. """, ) + keep_alive_probe: Optional[KeepAliveProbe] = Field( + default=None, + description="""Optional. Specifies the configuration for keep-alive probe. + Contains configuration on a specified endpoint that a deployment host + should use to keep the container alive based on the probe settings.""", + ) encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, - description="""The encryption spec to be used for the Agent Engine.""", + description="""The encryption spec to be used for the Agent Runtime.""", ) labels: Optional[dict[str, str]] = Field( - default=None, description="""The labels to be used for the Agent Engine.""" + default=None, description="""The labels to be used for the Agent Runtime.""" ) class_methods: Optional[list[dict[str, Any]]] = Field( default=None, - description="""The class methods to be used for the Agent Engine. + description="""The class methods to be used for the Agent Runtime. If specified, they'll override the class methods that are autogenerated by default. By default, methods are generated by inspecting the agent object and generating a corresponding method for each method defined on the @@ -9248,8 +9014,8 @@ class CreateAgentEngineConfig(_common.BaseModel): default=None, description="""The user-provided paths to the source packages (if any). If specified, the files in the source packages will be packed into a - a tarball file, uploaded to Agent Engine's API, and deployed to the - Agent Engine. + a tarball file, uploaded to Agent Runtime's API, and deployed to the + Agent Runtime. The following fields will be ignored: - agent - extra_packages @@ -9271,19 +9037,19 @@ class CreateAgentEngineConfig(_common.BaseModel): ) entrypoint_module: Optional[str] = Field( default=None, - description="""The entrypoint module to be used for the Agent Engine + description="""The entrypoint module to be used for the Agent Runtime This field only used when source_packages is specified.""", ) entrypoint_object: Optional[str] = Field( default=None, - description="""The entrypoint object to be used for the Agent Engine. + description="""The entrypoint object to be used for the Agent Runtime. This field only used when source_packages is specified.""", ) requirements_file: Optional[str] = Field( default=None, description="""The user-provided path to the requirements file (if any). This field is only used when source_packages is specified. - If not specified, agent engine will find and use the `requirements.txt` in + If not specified, agent runtime will find and use the `requirements.txt` in the source package. """, ) @@ -9291,7 +9057,7 @@ class CreateAgentEngineConfig(_common.BaseModel): Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] ] = Field( default=None, - description="""The agent framework to be used for the Agent Engine. + description="""The agent framework to be used for the Agent Runtime. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom". @@ -9302,14 +9068,14 @@ class CreateAgentEngineConfig(_common.BaseModel): ) python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] = Field( default=None, - description="""The Python version to be used for the Agent Engine. + description="""The Python version to be used for the Agent Runtime. If not specified, it will use the current Python version of the environment. Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". """, ) build_options: Optional[dict[str, list[str]]] = Field( default=None, - description="""The build options for the Agent Engine. + description="""The build options for the Agent Runtime. The following keys are supported: - installation_scripts: Optional. The paths to the installation scripts to be @@ -9318,77 +9084,84 @@ class CreateAgentEngineConfig(_common.BaseModel): subdirectory and the path must be added to `extra_packages`. """, ) - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig - ] = Field( + update_mask: Optional[str] = Field( default=None, - description="""Agent Gateway configuration for a Reasoning Engine deployment.""", + description="""The update mask to apply. For the `FieldMask` definition, see + https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""", ) - keep_alive_probe: Optional[KeepAliveProbe] = Field( + traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( default=None, - description="""Optional. Specifies the configuration for keep-alive probe. - Contains configuration on a specified endpoint that a deployment host - should use to keep the container alive based on the probe settings.""", + description="""Traffic distribution configuration for the Reasoning Engine.""", ) -class CreateAgentEngineConfigDict(TypedDict, total=False): - """Config for create agent engine.""" +class UpdateRuntimeConfigDict(TypedDict, total=False): + """Config for updating agent runtime.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" display_name: Optional[str] - """The user-defined name of the Agent Engine. + """The user-defined name of the Agent Runtime. The display name can be up to 128 characters long and can comprise any UTF-8 characters. """ description: Optional[str] - """The description of the Agent Engine.""" + """The description of the Agent Runtime.""" spec: Optional[ReasoningEngineSpecDict] - """Optional. Configurations of the Agent Engine.""" + """Optional. Configurations of the Agent Runtime.""" context_spec: Optional[ReasoningEngineContextSpecDict] - """Optional. The context spec to be used for the Agent Engine.""" + """Optional. The context spec to be used for the Agent Runtime.""" psc_interface_config: Optional[PscInterfaceConfigDict] """Optional. The PSC interface config for PSC-I to be used for the - Agent Engine.""" + Agent Runtime.""" + + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] + """Agent Gateway configuration for a Reasoning Engine deployment.""" min_instances: Optional[int] - """The minimum number of instances to run for the Agent Engine. + """The minimum number of instances to run for the Agent Runtime. Defaults to 1. Range: [0, 10]. """ max_instances: Optional[int] - """The maximum number of instances to run for the Agent Engine. + """The maximum number of instances to run for the Agent Runtime. Defaults to 100. Range: [1, 1000]. If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. """ resource_limits: Optional[dict[str, str]] - """The resource limits to be applied to the Agent Engine. + """The resource limits to be applied to the Agent Runtime. Required keys: 'cpu' and 'memory'. Supported values for 'cpu': '1', '2', '4', '6', '8'. Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. """ container_concurrency: Optional[int] - """The container concurrency to be used for the Agent Engine. + """The container concurrency to be used for the Agent Runtime. Recommended value: 2 * cpu + 1. Defaults to 9. """ + keep_alive_probe: Optional[KeepAliveProbeDict] + """Optional. Specifies the configuration for keep-alive probe. + Contains configuration on a specified endpoint that a deployment host + should use to keep the container alive based on the probe settings.""" + encryption_spec: Optional[genai_types.EncryptionSpec] - """The encryption spec to be used for the Agent Engine.""" + """The encryption spec to be used for the Agent Runtime.""" labels: Optional[dict[str, str]] - """The labels to be used for the Agent Engine.""" + """The labels to be used for the Agent Runtime.""" class_methods: Optional[list[dict[str, Any]]] - """The class methods to be used for the Agent Engine. + """The class methods to be used for the Agent Runtime. If specified, they'll override the class methods that are autogenerated by default. By default, methods are generated by inspecting the agent object and generating a corresponding method for each method defined on the @@ -9398,8 +9171,8 @@ class CreateAgentEngineConfigDict(TypedDict, total=False): source_packages: Optional[list[str]] """The user-provided paths to the source packages (if any). If specified, the files in the source packages will be packed into a - a tarball file, uploaded to Agent Engine's API, and deployed to the - Agent Engine. + a tarball file, uploaded to Agent Runtime's API, and deployed to the + Agent Runtime. The following fields will be ignored: - agent - extra_packages @@ -9419,24 +9192,24 @@ class CreateAgentEngineConfigDict(TypedDict, total=False): """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""" entrypoint_module: Optional[str] - """The entrypoint module to be used for the Agent Engine + """The entrypoint module to be used for the Agent Runtime This field only used when source_packages is specified.""" entrypoint_object: Optional[str] - """The entrypoint object to be used for the Agent Engine. + """The entrypoint object to be used for the Agent Runtime. This field only used when source_packages is specified.""" requirements_file: Optional[str] """The user-provided path to the requirements file (if any). This field is only used when source_packages is specified. - If not specified, agent engine will find and use the `requirements.txt` in + If not specified, agent runtime will find and use the `requirements.txt` in the source package. """ agent_framework: Optional[ Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] ] - """The agent framework to be used for the Agent Engine. + """The agent framework to be used for the Agent Runtime. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom". @@ -9446,13 +9219,13 @@ class CreateAgentEngineConfigDict(TypedDict, total=False): default to "custom".""" python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] - """The Python version to be used for the Agent Engine. + """The Python version to be used for the Agent Runtime. If not specified, it will use the current Python version of the environment. Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". """ build_options: Optional[dict[str, list[str]]] - """The build options for the Agent Engine. + """The build options for the Agent Runtime. The following keys are supported: - installation_scripts: Optional. The paths to the installation scripts to be @@ -9461,307 +9234,365 @@ class CreateAgentEngineConfigDict(TypedDict, total=False): subdirectory and the path must be added to `extra_packages`. """ - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict - ] - """Agent Gateway configuration for a Reasoning Engine deployment.""" + update_mask: Optional[str] + """The update mask to apply. For the `FieldMask` definition, see + https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" - keep_alive_probe: Optional[KeepAliveProbeDict] - """Optional. Specifies the configuration for keep-alive probe. - Contains configuration on a specified endpoint that a deployment host - should use to keep the container alive based on the probe settings.""" + traffic_config: Optional[ReasoningEngineTrafficConfigDict] + """Traffic distribution configuration for the Reasoning Engine.""" -CreateAgentEngineConfigOrDict = Union[ - CreateAgentEngineConfig, CreateAgentEngineConfigDict -] +UpdateRuntimeConfigOrDict = Union[UpdateRuntimeConfig, UpdateRuntimeConfigDict] -class _CreateAgentEngineRequestParameters(_common.BaseModel): - """Parameters for creating agent engines.""" +class _UpdateRuntimeRequestParameters(_common.BaseModel): + """Parameters for updating agent runtimes.""" + + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime.""" + ) + config: Optional[UpdateRuntimeConfig] = Field(default=None, description="""""") - config: Optional[CreateAgentEngineConfig] = Field(default=None, description="""""") +class _UpdateRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for updating agent runtimes.""" -class _CreateAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for creating agent engines.""" + name: Optional[str] + """Name of the agent runtime.""" - config: Optional[CreateAgentEngineConfigDict] + config: Optional[UpdateRuntimeConfigDict] """""" -_CreateAgentEngineRequestParametersOrDict = Union[ - _CreateAgentEngineRequestParameters, _CreateAgentEngineRequestParametersDict +_UpdateRuntimeRequestParametersOrDict = Union[ + _UpdateRuntimeRequestParameters, _UpdateRuntimeRequestParametersDict ] -class DeleteAgentEngineConfig(_common.BaseModel): - """Config for deleting agent engine.""" +class GetRuntimeRevisionConfig(_common.BaseModel): + """Config for getting an Agent Runtime Runtime Revision.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class DeleteAgentEngineConfigDict(TypedDict, total=False): - """Config for deleting agent engine.""" +class GetRuntimeRevisionConfigDict(TypedDict, total=False): + """Config for getting an Agent Runtime Runtime Revision.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteAgentEngineConfigOrDict = Union[ - DeleteAgentEngineConfig, DeleteAgentEngineConfigDict +GetRuntimeRevisionConfigOrDict = Union[ + GetRuntimeRevisionConfig, GetRuntimeRevisionConfigDict ] -class _DeleteAgentEngineRequestParameters(_common.BaseModel): - """Parameters for deleting agent engines.""" +class _GetRuntimeRevisionRequestParameters(_common.BaseModel): + """Parameters for getting an agent runtime runtime revision.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - force: Optional[bool] = Field( - default=False, - description="""If set to true, any child resources will also be deleted.""", + default=None, description="""Name of the agent runtime runtime revision.""" ) - config: Optional[DeleteAgentEngineConfig] = Field(default=None, description="""""") + config: Optional[GetRuntimeRevisionConfig] = Field(default=None, description="""""") -class _DeleteAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engines.""" +class _GetRuntimeRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for getting an agent runtime runtime revision.""" name: Optional[str] - """Name of the agent engine.""" - - force: Optional[bool] - """If set to true, any child resources will also be deleted.""" + """Name of the agent runtime runtime revision.""" - config: Optional[DeleteAgentEngineConfigDict] + config: Optional[GetRuntimeRevisionConfigDict] """""" -_DeleteAgentEngineRequestParametersOrDict = Union[ - _DeleteAgentEngineRequestParameters, _DeleteAgentEngineRequestParametersDict +_GetRuntimeRevisionRequestParametersOrDict = Union[ + _GetRuntimeRevisionRequestParameters, _GetRuntimeRevisionRequestParametersDict ] -class DeleteAgentEngineOperation(_common.BaseModel): - """Operation for deleting agent engines.""" +class ReasoningEngineRuntimeRevision(_common.BaseModel): + """A runtime revision.""" - name: Optional[str] = Field( + create_time: Optional[datetime.datetime] = Field( default=None, - description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + description="""Output only. Timestamp when this ReasoningEngineRuntimeRevision was created.""", ) - metadata: Optional[dict[str, Any]] = Field( + name: Optional[str] = Field( default=None, - description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + description="""Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""", ) - done: Optional[bool] = Field( + spec: Optional[ReasoningEngineSpec] = Field( default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + description="""Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""", ) - error: Optional[dict[str, Any]] = Field( - default=None, - description="""The error result of the operation in case of failure or cancellation.""", + state: Optional[State] = Field( + default=None, description="""Output only. The state of the revision.""" ) -class DeleteAgentEngineOperationDict(TypedDict, total=False): - """Operation for deleting agent engines.""" +class ReasoningEngineRuntimeRevisionDict(TypedDict, total=False): + """A runtime revision.""" - name: Optional[str] - """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this ReasoningEngineRuntimeRevision was created.""" - metadata: Optional[dict[str, Any]] - """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + name: Optional[str] + """Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""" - done: Optional[bool] - """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + spec: Optional[ReasoningEngineSpecDict] + """Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + state: Optional[State] + """Output only. The state of the revision.""" -DeleteAgentEngineOperationOrDict = Union[ - DeleteAgentEngineOperation, DeleteAgentEngineOperationDict +ReasoningEngineRuntimeRevisionOrDict = Union[ + ReasoningEngineRuntimeRevision, ReasoningEngineRuntimeRevisionDict ] -class GetAgentEngineConfig(_common.BaseModel): - """Config for create agent engine.""" +class ListRuntimeRevisionsConfig(_common.BaseModel): + """Config for listing agent runtime revisions.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - - -class GetAgentEngineConfigDict(TypedDict, total=False): - """Config for create agent engine.""" - - http_options: Optional[genai_types.HttpOptions] + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( + default=None, + description="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""", + ) + + +class ListRuntimeRevisionsConfigDict(TypedDict, total=False): + """Config for listing agent runtime revisions.""" + + http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" + page_size: Optional[int] + """""" + + page_token: Optional[str] + """""" + + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""" + -GetAgentEngineConfigOrDict = Union[GetAgentEngineConfig, GetAgentEngineConfigDict] +ListRuntimeRevisionsConfigOrDict = Union[ + ListRuntimeRevisionsConfig, ListRuntimeRevisionsConfigDict +] -class _GetAgentEngineRequestParameters(_common.BaseModel): - """Parameters for getting agent engines.""" +class _ListRuntimeRevisionsRequestParameters(_common.BaseModel): + """Parameters for listing agent runtime revisions.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the reasoning engine.""" + ) + config: Optional[ListRuntimeRevisionsConfig] = Field( + default=None, description="""""" ) - config: Optional[GetAgentEngineConfig] = Field(default=None, description="""""") -class _GetAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for getting agent engines.""" +class _ListRuntimeRevisionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtime revisions.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the reasoning engine.""" + + config: Optional[ListRuntimeRevisionsConfigDict] + """""" + + +_ListRuntimeRevisionsRequestParametersOrDict = Union[ + _ListRuntimeRevisionsRequestParameters, _ListRuntimeRevisionsRequestParametersDict +] + - config: Optional[GetAgentEngineConfigDict] +class ListReasoningEnginesRuntimeRevisionsResponse(_common.BaseModel): + """Response for listing agent runtime runtime revisions.""" + + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" + ) + next_page_token: Optional[str] = Field(default=None, description="""""") + reasoning_engine_runtime_revisions: Optional[ + list[ReasoningEngineRuntimeRevision] + ] = Field(default=None, description="""List of agent runtime revisions.""") + + +class ListReasoningEnginesRuntimeRevisionsResponseDict(TypedDict, total=False): + """Response for listing agent runtime runtime revisions.""" + + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + + next_page_token: Optional[str] """""" + reasoning_engine_runtime_revisions: Optional[ + list[ReasoningEngineRuntimeRevisionDict] + ] + """List of agent runtime revisions.""" + -_GetAgentEngineRequestParametersOrDict = Union[ - _GetAgentEngineRequestParameters, _GetAgentEngineRequestParametersDict +ListReasoningEnginesRuntimeRevisionsResponseOrDict = Union[ + ListReasoningEnginesRuntimeRevisionsResponse, + ListReasoningEnginesRuntimeRevisionsResponseDict, ] -class ListAgentEngineConfig(_common.BaseModel): - """Config for listing agent engines.""" +class DeleteRuntimeRevisionConfig(_common.BaseModel): + """Config for deleting an Agent Runtime Runtime Revision.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( - default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""", + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) -class ListAgentEngineConfigDict(TypedDict, total=False): - """Config for listing agent engines.""" +class DeleteRuntimeRevisionConfigDict(TypedDict, total=False): + """Config for deleting an Agent Runtime Runtime Revision.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - page_size: Optional[int] - """""" - - page_token: Optional[str] - """""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" +DeleteRuntimeRevisionConfigOrDict = Union[ + DeleteRuntimeRevisionConfig, DeleteRuntimeRevisionConfigDict +] -ListAgentEngineConfigOrDict = Union[ListAgentEngineConfig, ListAgentEngineConfigDict] +class _DeleteRuntimeRevisionRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtime runtime revisions.""" -class _ListAgentEngineRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" + name: Optional[str] = Field( + default=None, + description="""Name of the agent runtime runtime revision to delete.""", + ) + config: Optional[DeleteRuntimeRevisionConfig] = Field( + default=None, description="""""" + ) - config: Optional[ListAgentEngineConfig] = Field(default=None, description="""""") +class _DeleteRuntimeRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtime runtime revisions.""" -class _ListAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" + name: Optional[str] + """Name of the agent runtime runtime revision to delete.""" - config: Optional[ListAgentEngineConfigDict] + config: Optional[DeleteRuntimeRevisionConfigDict] """""" -_ListAgentEngineRequestParametersOrDict = Union[ - _ListAgentEngineRequestParameters, _ListAgentEngineRequestParametersDict +_DeleteRuntimeRevisionRequestParametersOrDict = Union[ + _DeleteRuntimeRevisionRequestParameters, _DeleteRuntimeRevisionRequestParametersDict ] -class ListReasoningEnginesResponse(_common.BaseModel): - """Response for listing agent engines.""" +class DeleteRuntimeRevisionOperation(_common.BaseModel): + """Operation for deleting agent runtime revisions.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - reasoning_engines: Optional[list[ReasoningEngine]] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""List of agent engines. - """, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", ) -class ListReasoningEnginesResponseDict(TypedDict, total=False): - """Response for listing agent engines.""" +class DeleteRuntimeRevisionOperationDict(TypedDict, total=False): + """Operation for deleting agent runtime revisions.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - next_page_token: Optional[str] - """""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - reasoning_engines: Optional[list[ReasoningEngineDict]] - """List of agent engines. - """ + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -ListReasoningEnginesResponseOrDict = Union[ - ListReasoningEnginesResponse, ListReasoningEnginesResponseDict + +DeleteRuntimeRevisionOperationOrDict = Union[ + DeleteRuntimeRevisionOperation, DeleteRuntimeRevisionOperationDict ] -class GetAgentEngineOperationConfig(_common.BaseModel): +class GetDeleteRuntimeRevisionOperationConfig(_common.BaseModel): http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetAgentEngineOperationConfigDict(TypedDict, total=False): +class GetDeleteRuntimeRevisionOperationConfigDict(TypedDict, total=False): http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetAgentEngineOperationConfigOrDict = Union[ - GetAgentEngineOperationConfig, GetAgentEngineOperationConfigDict +GetDeleteRuntimeRevisionOperationConfigOrDict = Union[ + GetDeleteRuntimeRevisionOperationConfig, GetDeleteRuntimeRevisionOperationConfigDict ] -class _GetAgentEngineOperationParameters(_common.BaseModel): - """Parameters for getting an operation with an agent engine as a response.""" +class _GetDeleteRuntimeRevisionOperationParameters(_common.BaseModel): + """Parameters for getting an operation that deletes an agent runtime revision.""" operation_name: Optional[str] = Field( default=None, description="""The server-assigned name for the operation.""" ) - config: Optional[GetAgentEngineOperationConfig] = Field( + config: Optional[GetDeleteRuntimeRevisionOperationConfig] = Field( default=None, description="""Used to override the default configuration.""" ) -class _GetAgentEngineOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with an agent engine as a response.""" +class _GetDeleteRuntimeRevisionOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation that deletes an agent runtime revision.""" operation_name: Optional[str] """The server-assigned name for the operation.""" - config: Optional[GetAgentEngineOperationConfigDict] + config: Optional[GetDeleteRuntimeRevisionOperationConfigDict] """Used to override the default configuration.""" -_GetAgentEngineOperationParametersOrDict = Union[ - _GetAgentEngineOperationParameters, _GetAgentEngineOperationParametersDict +_GetDeleteRuntimeRevisionOperationParametersOrDict = Union[ + _GetDeleteRuntimeRevisionOperationParameters, + _GetDeleteRuntimeRevisionOperationParametersDict, ] -class QueryAgentEngineConfig(_common.BaseModel): - """Config for querying agent engines.""" +class QueryRuntimeRevisionConfig(_common.BaseModel): + """Config for querying agent runtime revisions.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -9775,8 +9606,8 @@ class QueryAgentEngineConfig(_common.BaseModel): include_all_fields: Optional[bool] = Field(default=False, description="""""") -class QueryAgentEngineConfigDict(TypedDict, total=False): - """Config for querying agent engines.""" +class QueryRuntimeRevisionConfigDict(TypedDict, total=False): + """Config for querying agent runtime revisions.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -9791,620 +9622,255 @@ class QueryAgentEngineConfigDict(TypedDict, total=False): """""" -QueryAgentEngineConfigOrDict = Union[QueryAgentEngineConfig, QueryAgentEngineConfigDict] +QueryRuntimeRevisionConfigOrDict = Union[ + QueryRuntimeRevisionConfig, QueryRuntimeRevisionConfigDict +] -class _QueryAgentEngineRequestParameters(_common.BaseModel): - """Parameters for querying agent engines.""" +class _QueryRuntimeRevisionRequestParameters(_common.BaseModel): + """Parameters for querying agent runtime revisions.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime revision.""" + ) + config: Optional[QueryRuntimeRevisionConfig] = Field( + default=None, description="""""" ) - config: Optional[QueryAgentEngineConfig] = Field(default=None, description="""""") -class _QueryAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for querying agent engines.""" +class _QueryRuntimeRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for querying agent runtime revisions.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime revision.""" - config: Optional[QueryAgentEngineConfigDict] + config: Optional[QueryRuntimeRevisionConfigDict] """""" -_QueryAgentEngineRequestParametersOrDict = Union[ - _QueryAgentEngineRequestParameters, _QueryAgentEngineRequestParametersDict +_QueryRuntimeRevisionRequestParametersOrDict = Union[ + _QueryRuntimeRevisionRequestParameters, _QueryRuntimeRevisionRequestParametersDict ] -class QueryReasoningEngineResponse(_common.BaseModel): - """The response for querying an agent engine.""" +class MemoryMetadataValue(_common.BaseModel): + """The metadata values for memories.""" - output: Optional[Any] = Field( + bool_value: Optional[bool] = Field( + default=None, description="""Represents a boolean value.""" + ) + double_value: Optional[float] = Field( + default=None, description="""Represents a double value.""" + ) + string_value: Optional[str] = Field( + default=None, description="""Represents a string value.""" + ) + timestamp_value: Optional[datetime.datetime] = Field( default=None, - description="""Response provided by users in JSON object format.""", + description="""Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""", ) -class QueryReasoningEngineResponseDict(TypedDict, total=False): - """The response for querying an agent engine.""" +class MemoryMetadataValueDict(TypedDict, total=False): + """The metadata values for memories.""" - output: Optional[Any] - """Response provided by users in JSON object format.""" + bool_value: Optional[bool] + """Represents a boolean value.""" + double_value: Optional[float] + """Represents a double value.""" -QueryReasoningEngineResponseOrDict = Union[ - QueryReasoningEngineResponse, QueryReasoningEngineResponseDict -] + string_value: Optional[str] + """Represents a string value.""" + + timestamp_value: Optional[datetime.datetime] + """Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""" + + +MemoryMetadataValueOrDict = Union[MemoryMetadataValue, MemoryMetadataValueDict] -class UpdateAgentEngineConfig(_common.BaseModel): - """Config for updating agent engine.""" +class RuntimeMemoryConfig(_common.BaseModel): + """Config for creating a Memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) display_name: Optional[str] = Field( - default=None, - description="""The user-defined name of the Agent Engine. - - The display name can be up to 128 characters long and can comprise any - UTF-8 characters. - """, + default=None, description="""The display name of the memory.""" ) description: Optional[str] = Field( - default=None, description="""The description of the Agent Engine.""" - ) - spec: Optional[ReasoningEngineSpec] = Field( - default=None, description="""Optional. Configurations of the Agent Engine.""" - ) - context_spec: Optional[ReasoningEngineContextSpec] = Field( - default=None, - description="""Optional. The context spec to be used for the Agent Engine.""", - ) - psc_interface_config: Optional[PscInterfaceConfig] = Field( - default=None, - description="""Optional. The PSC interface config for PSC-I to be used for the - Agent Engine.""", - ) - min_instances: Optional[int] = Field( - default=None, - description="""The minimum number of instances to run for the Agent Engine. - Defaults to 1. Range: [0, 10]. - """, - ) - max_instances: Optional[int] = Field( - default=None, - description="""The maximum number of instances to run for the Agent Engine. - Defaults to 100. Range: [1, 1000]. - If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. - """, + default=None, description="""The description of the memory.""" ) - resource_limits: Optional[dict[str, str]] = Field( - default=None, - description="""The resource limits to be applied to the Agent Engine. - Required keys: 'cpu' and 'memory'. - Supported values for 'cpu': '1', '2', '4', '6', '8'. - Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. - """, + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) - container_concurrency: Optional[int] = Field( + ttl: Optional[str] = Field( default=None, - description="""The container concurrency to be used for the Agent Engine. - Recommended value: 2 * cpu + 1. Defaults to 9. - """, + description="""Optional. Input only. The TTL for this resource. + + The expiration time is computed: now + TTL.""", ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + expire_time: Optional[datetime.datetime] = Field( default=None, - description="""The encryption spec to be used for the Agent Engine.""", - ) - labels: Optional[dict[str, str]] = Field( - default=None, description="""The labels to be used for the Agent Engine.""" + description="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", ) - class_methods: Optional[list[dict[str, Any]]] = Field( + revision_expire_time: Optional[datetime.datetime] = Field( default=None, - description="""The class methods to be used for the Agent Engine. - If specified, they'll override the class methods that are autogenerated by - default. By default, methods are generated by inspecting the agent object - and generating a corresponding method for each method defined on the - agent class. - """, + description="""Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", ) - source_packages: Optional[list[str]] = Field( + revision_ttl: Optional[str] = Field( default=None, - description="""The user-provided paths to the source packages (if any). - If specified, the files in the source packages will be packed into a - a tarball file, uploaded to Agent Engine's API, and deployed to the - Agent Engine. - The following fields will be ignored: - - agent - - extra_packages - - staging_bucket - - requirements - The following fields will be used to install and use the agent from the - source packages: - - entrypoint_module (required) - - entrypoint_object (required) - - requirements_file (optional) - - class_methods (required) - """, + description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", ) - developer_connect_source: Optional[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig - ] = Field( + disable_memory_revisions: Optional[bool] = Field( default=None, - description="""Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""", + description="""Optional. Input only. If true, no revision will be created for this request.""", ) - entrypoint_module: Optional[str] = Field( - default=None, - description="""The entrypoint module to be used for the Agent Engine - This field only used when source_packages is specified.""", + topics: Optional[list[MemoryTopicId]] = Field( + default=None, description="""Optional. The topics of the memory.""" ) - entrypoint_object: Optional[str] = Field( + metadata: Optional[dict[str, MemoryMetadataValue]] = Field( default=None, - description="""The entrypoint object to be used for the Agent Engine. - This field only used when source_packages is specified.""", + description="""Optional. User-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""", ) - requirements_file: Optional[str] = Field( + memory_id: Optional[str] = Field( default=None, - description="""The user-provided path to the requirements file (if any). - This field is only used when source_packages is specified. - If not specified, agent engine will find and use the `requirements.txt` in - the source package. - """, - ) - agent_framework: Optional[ - Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] - ] = Field( - default=None, - description="""The agent framework to be used for the Agent Engine. - The OSS agent framework used to develop the agent. - Currently supported values: "google-adk", "langchain", "langgraph", - "ag2", "llama-index", "custom". - If not specified: - - If `agent` is specified, the agent framework will be auto-detected. - - If `source_packages` is specified, the agent framework will - default to "custom".""", - ) - python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] = Field( - default=None, - description="""The Python version to be used for the Agent Engine. - If not specified, it will use the current Python version of the environment. - Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". - """, - ) - build_options: Optional[dict[str, list[str]]] = Field( - default=None, - description="""The build options for the Agent Engine. - The following keys are supported: - - installation_scripts: - Optional. The paths to the installation scripts to be - executed in the Docker image. - The scripts must be located in the `installation_scripts` - subdirectory and the path must be added to `extra_packages`. - """, - ) - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig - ] = Field( - default=None, - description="""Agent Gateway configuration for a Reasoning Engine deployment.""", - ) - keep_alive_probe: Optional[KeepAliveProbe] = Field( - default=None, - description="""Optional. Specifies the configuration for keep-alive probe. - Contains configuration on a specified endpoint that a deployment host - should use to keep the container alive based on the probe settings.""", - ) - update_mask: Optional[str] = Field( - default=None, - description="""The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""", - ) - traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( - default=None, - description="""Traffic distribution configuration for the Reasoning Engine.""", + description="""Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""", ) -class UpdateAgentEngineConfigDict(TypedDict, total=False): - """Config for updating agent engine.""" +class RuntimeMemoryConfigDict(TypedDict, total=False): + """Config for creating a Memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" display_name: Optional[str] - """The user-defined name of the Agent Engine. - - The display name can be up to 128 characters long and can comprise any - UTF-8 characters. - """ + """The display name of the memory.""" description: Optional[str] - """The description of the Agent Engine.""" - - spec: Optional[ReasoningEngineSpecDict] - """Optional. Configurations of the Agent Engine.""" - - context_spec: Optional[ReasoningEngineContextSpecDict] - """Optional. The context spec to be used for the Agent Engine.""" - - psc_interface_config: Optional[PscInterfaceConfigDict] - """Optional. The PSC interface config for PSC-I to be used for the - Agent Engine.""" - - min_instances: Optional[int] - """The minimum number of instances to run for the Agent Engine. - Defaults to 1. Range: [0, 10]. - """ - - max_instances: Optional[int] - """The maximum number of instances to run for the Agent Engine. - Defaults to 100. Range: [1, 1000]. - If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. - """ - - resource_limits: Optional[dict[str, str]] - """The resource limits to be applied to the Agent Engine. - Required keys: 'cpu' and 'memory'. - Supported values for 'cpu': '1', '2', '4', '6', '8'. - Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. - """ + """The description of the memory.""" - container_concurrency: Optional[int] - """The container concurrency to be used for the Agent Engine. - Recommended value: 2 * cpu + 1. Defaults to 9. - """ + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """The encryption spec to be used for the Agent Engine.""" + ttl: Optional[str] + """Optional. Input only. The TTL for this resource. - labels: Optional[dict[str, str]] - """The labels to be used for the Agent Engine.""" + The expiration time is computed: now + TTL.""" - class_methods: Optional[list[dict[str, Any]]] - """The class methods to be used for the Agent Engine. - If specified, they'll override the class methods that are autogenerated by - default. By default, methods are generated by inspecting the agent object - and generating a corresponding method for each method defined on the - agent class. - """ + expire_time: Optional[datetime.datetime] + """Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""" - source_packages: Optional[list[str]] - """The user-provided paths to the source packages (if any). - If specified, the files in the source packages will be packed into a - a tarball file, uploaded to Agent Engine's API, and deployed to the - Agent Engine. - The following fields will be ignored: - - agent - - extra_packages - - staging_bucket - - requirements - The following fields will be used to install and use the agent from the - source packages: - - entrypoint_module (required) - - entrypoint_object (required) - - requirements_file (optional) - - class_methods (required) - """ + revision_expire_time: Optional[datetime.datetime] + """Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" - developer_connect_source: Optional[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict - ] - """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""" + revision_ttl: Optional[str] + """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" - entrypoint_module: Optional[str] - """The entrypoint module to be used for the Agent Engine - This field only used when source_packages is specified.""" + disable_memory_revisions: Optional[bool] + """Optional. Input only. If true, no revision will be created for this request.""" - entrypoint_object: Optional[str] - """The entrypoint object to be used for the Agent Engine. - This field only used when source_packages is specified.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The topics of the memory.""" - requirements_file: Optional[str] - """The user-provided path to the requirements file (if any). - This field is only used when source_packages is specified. - If not specified, agent engine will find and use the `requirements.txt` in - the source package. - """ + metadata: Optional[dict[str, MemoryMetadataValueDict]] + """Optional. User-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""" - agent_framework: Optional[ - Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] - ] - """The agent framework to be used for the Agent Engine. - The OSS agent framework used to develop the agent. - Currently supported values: "google-adk", "langchain", "langgraph", - "ag2", "llama-index", "custom". - If not specified: - - If `agent` is specified, the agent framework will be auto-detected. - - If `source_packages` is specified, the agent framework will - default to "custom".""" + memory_id: Optional[str] + """Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""" - python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] - """The Python version to be used for the Agent Engine. - If not specified, it will use the current Python version of the environment. - Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". - """ - build_options: Optional[dict[str, list[str]]] - """The build options for the Agent Engine. - The following keys are supported: - - installation_scripts: - Optional. The paths to the installation scripts to be - executed in the Docker image. - The scripts must be located in the `installation_scripts` - subdirectory and the path must be added to `extra_packages`. - """ +RuntimeMemoryConfigOrDict = Union[RuntimeMemoryConfig, RuntimeMemoryConfigDict] - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict - ] - """Agent Gateway configuration for a Reasoning Engine deployment.""" - keep_alive_probe: Optional[KeepAliveProbeDict] - """Optional. Specifies the configuration for keep-alive probe. - Contains configuration on a specified endpoint that a deployment host - should use to keep the container alive based on the probe settings.""" +class _CreateRuntimeMemoryRequestParameters(_common.BaseModel): + """Parameters for creating Agent Runtime Memories.""" - update_mask: Optional[str] - """The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" + name: Optional[str] = Field( + default=None, + description="""Name of the agent runtime to create the memory under.""", + ) + fact: Optional[str] = Field( + default=None, + description="""The fact of the memory. - traffic_config: Optional[ReasoningEngineTrafficConfigDict] - """Traffic distribution configuration for the Reasoning Engine.""" + This is the semantic knowledge extracted from the source content).""", + ) + scope: Optional[dict[str, str]] = Field( + default=None, + description="""The scope of the memory. + Memories are isolated within their scope. The scope is defined when + creating or generating memories. Up to 5 key-value pairs are accepted, + and scope values cannot contain the wildcard character '*'.""", + ) + config: Optional[RuntimeMemoryConfig] = Field(default=None, description="""""") -UpdateAgentEngineConfigOrDict = Union[ - UpdateAgentEngineConfig, UpdateAgentEngineConfigDict -] +class _CreateRuntimeMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for creating Agent Runtime Memories.""" -class _UpdateAgentEngineRequestParameters(_common.BaseModel): - """Parameters for updating agent engines.""" + name: Optional[str] + """Name of the agent runtime to create the memory under.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[UpdateAgentEngineConfig] = Field(default=None, description="""""") + fact: Optional[str] + """The fact of the memory. + This is the semantic knowledge extracted from the source content).""" -class _UpdateAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for updating agent engines.""" + scope: Optional[dict[str, str]] + """The scope of the memory. - name: Optional[str] - """Name of the agent engine.""" + Memories are isolated within their scope. The scope is defined when + creating or generating memories. Up to 5 key-value pairs are accepted, + and scope values cannot contain the wildcard character '*'.""" - config: Optional[UpdateAgentEngineConfigDict] + config: Optional[RuntimeMemoryConfigDict] """""" -_UpdateAgentEngineRequestParametersOrDict = Union[ - _UpdateAgentEngineRequestParameters, _UpdateAgentEngineRequestParametersDict +_CreateRuntimeMemoryRequestParametersOrDict = Union[ + _CreateRuntimeMemoryRequestParameters, _CreateRuntimeMemoryRequestParametersDict ] -class MemoryMetadataValue(_common.BaseModel): - """The metadata values for memories.""" +class MemoryStructuredContent(_common.BaseModel): + """Represents the structured value of the memory.""" - bool_value: Optional[bool] = Field( - default=None, description="""Represents a boolean value.""" - ) - double_value: Optional[float] = Field( - default=None, description="""Represents a double value.""" - ) - string_value: Optional[str] = Field( - default=None, description="""Represents a string value.""" + data: Optional[dict[str, Any]] = Field( + default=None, + description="""Required. Represents the structured value of the memory.""", ) - timestamp_value: Optional[datetime.datetime] = Field( + schema_id: Optional[str] = Field( default=None, - description="""Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""", + description="""Required. Represents the schema ID for which this structured memory belongs to.""", ) -class MemoryMetadataValueDict(TypedDict, total=False): - """The metadata values for memories.""" - - bool_value: Optional[bool] - """Represents a boolean value.""" - - double_value: Optional[float] - """Represents a double value.""" +class MemoryStructuredContentDict(TypedDict, total=False): + """Represents the structured value of the memory.""" - string_value: Optional[str] - """Represents a string value.""" + data: Optional[dict[str, Any]] + """Required. Represents the structured value of the memory.""" - timestamp_value: Optional[datetime.datetime] - """Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""" + schema_id: Optional[str] + """Required. Represents the schema ID for which this structured memory belongs to.""" -MemoryMetadataValueOrDict = Union[MemoryMetadataValue, MemoryMetadataValueDict] +MemoryStructuredContentOrDict = Union[ + MemoryStructuredContent, MemoryStructuredContentDict +] -class AgentEngineMemoryConfig(_common.BaseModel): - """Config for creating a Memory.""" +class Memory(_common.BaseModel): + """A memory.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - display_name: Optional[str] = Field( - default=None, description="""The display name of the memory.""" - ) - description: Optional[str] = Field( - default=None, description="""The description of the memory.""" - ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", - ) - ttl: Optional[str] = Field( - default=None, - description="""Optional. Input only. The TTL for this resource. - - The expiration time is computed: now + TTL.""", - ) - expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", - ) - revision_expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", - ) - revision_ttl: Optional[str] = Field( - default=None, - description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", - ) - disable_memory_revisions: Optional[bool] = Field( - default=None, - description="""Optional. Input only. If true, no revision will be created for this request.""", - ) - topics: Optional[list[MemoryTopicId]] = Field( - default=None, description="""Optional. The topics of the memory.""" - ) - metadata: Optional[dict[str, MemoryMetadataValue]] = Field( - default=None, - description="""Optional. User-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""", - ) - memory_id: Optional[str] = Field( - default=None, - description="""Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""", - ) - - -class AgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for creating a Memory.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - display_name: Optional[str] - """The display name of the memory.""" - - description: Optional[str] - """The description of the memory.""" - - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" - - ttl: Optional[str] - """Optional. Input only. The TTL for this resource. - - The expiration time is computed: now + TTL.""" - - expire_time: Optional[datetime.datetime] - """Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""" - - revision_expire_time: Optional[datetime.datetime] - """Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" - - revision_ttl: Optional[str] - """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" - - disable_memory_revisions: Optional[bool] - """Optional. Input only. If true, no revision will be created for this request.""" - - topics: Optional[list[MemoryTopicIdDict]] - """Optional. The topics of the memory.""" - - metadata: Optional[dict[str, MemoryMetadataValueDict]] - """Optional. User-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""" - - memory_id: Optional[str] - """Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""" - - -AgentEngineMemoryConfigOrDict = Union[ - AgentEngineMemoryConfig, AgentEngineMemoryConfigDict -] - - -class _CreateAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for creating Agent Engine Memories.""" - - name: Optional[str] = Field( - default=None, - description="""Name of the agent engine to create the memory under.""", - ) - fact: Optional[str] = Field( - default=None, - description="""The fact of the memory. - - This is the semantic knowledge extracted from the source content).""", - ) - scope: Optional[dict[str, str]] = Field( - default=None, - description="""The scope of the memory. - - Memories are isolated within their scope. The scope is defined when - creating or generating memories. Up to 5 key-value pairs are accepted, - and scope values cannot contain the wildcard character '*'.""", - ) - config: Optional[AgentEngineMemoryConfig] = Field(default=None, description="""""") - - -class _CreateAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for creating Agent Engine Memories.""" - - name: Optional[str] - """Name of the agent engine to create the memory under.""" - - fact: Optional[str] - """The fact of the memory. - - This is the semantic knowledge extracted from the source content).""" - - scope: Optional[dict[str, str]] - """The scope of the memory. - - Memories are isolated within their scope. The scope is defined when - creating or generating memories. Up to 5 key-value pairs are accepted, - and scope values cannot contain the wildcard character '*'.""" - - config: Optional[AgentEngineMemoryConfigDict] - """""" - - -_CreateAgentEngineMemoryRequestParametersOrDict = Union[ - _CreateAgentEngineMemoryRequestParameters, - _CreateAgentEngineMemoryRequestParametersDict, -] - - -class MemoryStructuredContent(_common.BaseModel): - """Represents the structured value of the memory.""" - - data: Optional[dict[str, Any]] = Field( - default=None, - description="""Required. Represents the structured value of the memory.""", - ) - schema_id: Optional[str] = Field( - default=None, - description="""Required. Represents the schema ID for which this structured memory belongs to.""", - ) - - -class MemoryStructuredContentDict(TypedDict, total=False): - """Represents the structured value of the memory.""" - - data: Optional[dict[str, Any]] - """Required. Represents the structured value of the memory.""" - - schema_id: Optional[str] - """Required. Represents the schema ID for which this structured memory belongs to.""" - - -MemoryStructuredContentOrDict = Union[ - MemoryStructuredContent, MemoryStructuredContentDict -] - - -class Memory(_common.BaseModel): - """A memory.""" - - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Represents the timestamp when this Memory was created.""", + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Represents the timestamp when this Memory was created.""", ) description: Optional[str] = Field( default=None, @@ -10529,8 +9995,8 @@ class MemoryDict(TypedDict, total=False): MemoryOrDict = Union[Memory, MemoryDict] -class AgentEngineMemoryOperation(_common.BaseModel): - """Operation that has an agent engine memory as a response.""" +class RuntimeMemoryOperation(_common.BaseModel): + """Operation that has an agent runtime memory as a response.""" name: Optional[str] = Field( default=None, @@ -10549,12 +10015,12 @@ class AgentEngineMemoryOperation(_common.BaseModel): description="""The error result of the operation in case of failure or cancellation.""", ) response: Optional[Memory] = Field( - default=None, description="""The Agent Engine Memory.""" + default=None, description="""The Agent Runtime Memory.""" ) -class AgentEngineMemoryOperationDict(TypedDict, total=False): - """Operation that has an agent engine memory as a response.""" +class RuntimeMemoryOperationDict(TypedDict, total=False): + """Operation that has an agent runtime memory as a response.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -10569,63 +10035,60 @@ class AgentEngineMemoryOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" response: Optional[MemoryDict] - """The Agent Engine Memory.""" + """The Agent Runtime Memory.""" -AgentEngineMemoryOperationOrDict = Union[ - AgentEngineMemoryOperation, AgentEngineMemoryOperationDict -] +RuntimeMemoryOperationOrDict = Union[RuntimeMemoryOperation, RuntimeMemoryOperationDict] -class DeleteAgentEngineMemoryConfig(_common.BaseModel): - """Config for deleting an Agent Engine Memory.""" +class DeleteRuntimeMemoryConfig(_common.BaseModel): + """Config for deleting an Agent Runtime Memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class DeleteAgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Memory.""" +class DeleteRuntimeMemoryConfigDict(TypedDict, total=False): + """Config for deleting an Agent Runtime Memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteAgentEngineMemoryConfigOrDict = Union[ - DeleteAgentEngineMemoryConfig, DeleteAgentEngineMemoryConfigDict +DeleteRuntimeMemoryConfigOrDict = Union[ + DeleteRuntimeMemoryConfig, DeleteRuntimeMemoryConfigDict ] -class _DeleteAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for deleting agent engines.""" +class _DeleteRuntimeMemoryRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtimes.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine memory to delete.""" + default=None, description="""Name of the agent runtime memory to delete.""" ) - config: Optional[DeleteAgentEngineMemoryConfig] = Field( + config: Optional[DeleteRuntimeMemoryConfig] = Field( default=None, description="""""" ) -class _DeleteAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engines.""" +class _DeleteRuntimeMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtimes.""" name: Optional[str] - """Name of the agent engine memory to delete.""" + """Name of the agent runtime memory to delete.""" - config: Optional[DeleteAgentEngineMemoryConfigDict] + config: Optional[DeleteRuntimeMemoryConfigDict] """""" -_DeleteAgentEngineMemoryRequestParametersOrDict = Union[ - _DeleteAgentEngineMemoryRequestParameters, - _DeleteAgentEngineMemoryRequestParametersDict, +_DeleteRuntimeMemoryRequestParametersOrDict = Union[ + _DeleteRuntimeMemoryRequestParameters, _DeleteRuntimeMemoryRequestParametersDict ] -class DeleteAgentEngineMemoryOperation(_common.BaseModel): - """Operation for deleting agent engines.""" +class DeleteRuntimeMemoryOperation(_common.BaseModel): + """Operation for deleting agent runtimes.""" name: Optional[str] = Field( default=None, @@ -10645,8 +10108,8 @@ class DeleteAgentEngineMemoryOperation(_common.BaseModel): ) -class DeleteAgentEngineMemoryOperationDict(TypedDict, total=False): - """Operation for deleting agent engines.""" +class DeleteRuntimeMemoryOperationDict(TypedDict, total=False): + """Operation for deleting agent runtimes.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -10661,8 +10124,8 @@ class DeleteAgentEngineMemoryOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -DeleteAgentEngineMemoryOperationOrDict = Union[ - DeleteAgentEngineMemoryOperation, DeleteAgentEngineMemoryOperationDict +DeleteRuntimeMemoryOperationOrDict = Union[ + DeleteRuntimeMemoryOperation, DeleteRuntimeMemoryOperationDict ] @@ -10801,7 +10264,7 @@ class GenerateMemoriesRequestDirectMemoriesSourceDict(TypedDict, total=False): ] -class GenerateAgentEngineMemoriesConfig(_common.BaseModel): +class GenerateRuntimeMemoriesConfig(_common.BaseModel): """Config for generating memories.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -10850,7 +10313,7 @@ class GenerateAgentEngineMemoriesConfig(_common.BaseModel): ) -class GenerateAgentEngineMemoriesConfigDict(TypedDict, total=False): +class GenerateRuntimeMemoriesConfigDict(TypedDict, total=False): """Config for generating memories.""" http_options: Optional[genai_types.HttpOptions] @@ -10889,17 +10352,17 @@ class GenerateAgentEngineMemoriesConfigDict(TypedDict, total=False): """Optional. Restricts memory generation to a subset of memory topics.""" -GenerateAgentEngineMemoriesConfigOrDict = Union[ - GenerateAgentEngineMemoriesConfig, GenerateAgentEngineMemoriesConfigDict +GenerateRuntimeMemoriesConfigOrDict = Union[ + GenerateRuntimeMemoriesConfig, GenerateRuntimeMemoriesConfigDict ] -class _GenerateAgentEngineMemoriesRequestParameters(_common.BaseModel): - """Parameters for generating agent engine memories.""" +class _GenerateRuntimeMemoriesRequestParameters(_common.BaseModel): + """Parameters for generating agent runtime memories.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to generate memories for.""", + description="""Name of the agent runtime to generate memories for.""", ) vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSource] = Field( default=None, @@ -10926,16 +10389,16 @@ class _GenerateAgentEngineMemoriesRequestParameters(_common.BaseModel): provided, it will override the scope defined in the source content. Scope values cannot contain the wildcard character '*'.""", ) - config: Optional[GenerateAgentEngineMemoriesConfig] = Field( + config: Optional[GenerateRuntimeMemoriesConfig] = Field( default=None, description="""""" ) -class _GenerateAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for generating agent engine memories.""" +class _GenerateRuntimeMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for generating agent runtime memories.""" name: Optional[str] - """Name of the agent engine to generate memories for.""" + """Name of the agent runtime to generate memories for.""" vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSourceDict] """The vertex session source of the memories that should be generated.""" @@ -10954,13 +10417,13 @@ class _GenerateAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): provided, it will override the scope defined in the source content. Scope values cannot contain the wildcard character '*'.""" - config: Optional[GenerateAgentEngineMemoriesConfigDict] + config: Optional[GenerateRuntimeMemoriesConfigDict] """""" -_GenerateAgentEngineMemoriesRequestParametersOrDict = Union[ - _GenerateAgentEngineMemoriesRequestParameters, - _GenerateAgentEngineMemoriesRequestParametersDict, +_GenerateRuntimeMemoriesRequestParametersOrDict = Union[ + _GenerateRuntimeMemoriesRequestParameters, + _GenerateRuntimeMemoriesRequestParametersDict, ] @@ -11025,8 +10488,8 @@ class GenerateMemoriesResponseDict(TypedDict, total=False): ] -class AgentEngineGenerateMemoriesOperation(_common.BaseModel): - """Operation that generates memories for an agent engine.""" +class RuntimeGenerateMemoriesOperation(_common.BaseModel): + """Operation that generates memories for an agent runtime.""" name: Optional[str] = Field( default=None, @@ -11049,8 +10512,8 @@ class AgentEngineGenerateMemoriesOperation(_common.BaseModel): ) -class AgentEngineGenerateMemoriesOperationDict(TypedDict, total=False): - """Operation that generates memories for an agent engine.""" +class RuntimeGenerateMemoriesOperationDict(TypedDict, total=False): + """Operation that generates memories for an agent runtime.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -11068,54 +10531,50 @@ class AgentEngineGenerateMemoriesOperationDict(TypedDict, total=False): """The response for generating memories.""" -AgentEngineGenerateMemoriesOperationOrDict = Union[ - AgentEngineGenerateMemoriesOperation, AgentEngineGenerateMemoriesOperationDict +RuntimeGenerateMemoriesOperationOrDict = Union[ + RuntimeGenerateMemoriesOperation, RuntimeGenerateMemoriesOperationDict ] -class GetAgentEngineMemoryConfig(_common.BaseModel): - """Config for getting an Agent Engine Memory.""" +class GetRuntimeMemoryConfig(_common.BaseModel): + """Config for getting an Agent Runtime Memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetAgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Memory.""" +class GetRuntimeMemoryConfigDict(TypedDict, total=False): + """Config for getting an Agent Runtime Memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetAgentEngineMemoryConfigOrDict = Union[ - GetAgentEngineMemoryConfig, GetAgentEngineMemoryConfigDict -] +GetRuntimeMemoryConfigOrDict = Union[GetRuntimeMemoryConfig, GetRuntimeMemoryConfigDict] -class _GetAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine.""" +class _GetRuntimeMemoryRequestParameters(_common.BaseModel): + """Parameters for getting an agent runtime.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[GetAgentEngineMemoryConfig] = Field( - default=None, description="""""" + default=None, description="""Name of the agent runtime.""" ) + config: Optional[GetRuntimeMemoryConfig] = Field(default=None, description="""""") -class _GetAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine.""" +class _GetRuntimeMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for getting an agent runtime.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime.""" - config: Optional[GetAgentEngineMemoryConfigDict] + config: Optional[GetRuntimeMemoryConfigDict] """""" -_GetAgentEngineMemoryRequestParametersOrDict = Union[ - _GetAgentEngineMemoryRequestParameters, _GetAgentEngineMemoryRequestParametersDict +_GetRuntimeMemoryRequestParametersOrDict = Union[ + _GetRuntimeMemoryRequestParameters, _GetRuntimeMemoryRequestParametersDict ] @@ -11250,10 +10709,10 @@ class IngestEventsConfigDict(TypedDict, total=False): class _IngestEventsRequestParameters(_common.BaseModel): - """Parameters for purging agent engine memories.""" + """Parameters for purging agent runtime memories.""" name: Optional[str] = Field( - default=None, description="""Name of the Agent Engine to ingest events into.""" + default=None, description="""Name of the Agent Runtime to ingest events into.""" ) stream_id: Optional[str] = Field( default=None, description="""The ID of the stream to ingest events into.""" @@ -11277,10 +10736,10 @@ class _IngestEventsRequestParameters(_common.BaseModel): class _IngestEventsRequestParametersDict(TypedDict, total=False): - """Parameters for purging agent engine memories.""" + """Parameters for purging agent runtime memories.""" name: Optional[str] - """Name of the Agent Engine to ingest events into.""" + """Name of the Agent Runtime to ingest events into.""" stream_id: Optional[str] """The ID of the stream to ingest events into.""" @@ -11348,8 +10807,8 @@ class MemoryBankIngestEventsOperationDict(TypedDict, total=False): ] -class ListAgentEngineMemoryConfig(_common.BaseModel): - """Config for listing agent engine memories.""" +class ListRuntimeMemoryConfig(_common.BaseModel): + """Config for listing agent runtime memories.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -11374,8 +10833,8 @@ class ListAgentEngineMemoryConfig(_common.BaseModel): ) -class ListAgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for listing agent engine memories.""" +class ListRuntimeMemoryConfigDict(TypedDict, total=False): + """Config for listing agent runtime memories.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -11401,51 +10860,49 @@ class ListAgentEngineMemoryConfigDict(TypedDict, total=False): * `update_time`""" -ListAgentEngineMemoryConfigOrDict = Union[ - ListAgentEngineMemoryConfig, ListAgentEngineMemoryConfigDict +ListRuntimeMemoryConfigOrDict = Union[ + ListRuntimeMemoryConfig, ListRuntimeMemoryConfigDict ] -class _ListAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" +class _ListRuntimeMemoryRequestParameters(_common.BaseModel): + """Parameters for listing agent runtimes.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[ListAgentEngineMemoryConfig] = Field( - default=None, description="""""" + default=None, description="""Name of the agent runtime.""" ) + config: Optional[ListRuntimeMemoryConfig] = Field(default=None, description="""""") -class _ListAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" +class _ListRuntimeMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtimes.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime.""" - config: Optional[ListAgentEngineMemoryConfigDict] + config: Optional[ListRuntimeMemoryConfigDict] """""" -_ListAgentEngineMemoryRequestParametersOrDict = Union[ - _ListAgentEngineMemoryRequestParameters, _ListAgentEngineMemoryRequestParametersDict +_ListRuntimeMemoryRequestParametersOrDict = Union[ + _ListRuntimeMemoryRequestParameters, _ListRuntimeMemoryRequestParametersDict ] class ListReasoningEnginesMemoriesResponse(_common.BaseModel): - """Response for listing agent engine memories.""" + """Response for listing agent runtime memories.""" sdk_http_response: Optional[genai_types.HttpResponse] = Field( default=None, description="""Used to retain the full HTTP response.""" ) next_page_token: Optional[str] = Field(default=None, description="""""") memories: Optional[list[Memory]] = Field( - default=None, description="""List of agent engine memories.""" + default=None, description="""List of agent runtime memories.""" ) class ListReasoningEnginesMemoriesResponseDict(TypedDict, total=False): - """Response for listing agent engine memories.""" + """Response for listing agent runtime memories.""" sdk_http_response: Optional[genai_types.HttpResponse] """Used to retain the full HTTP response.""" @@ -11454,7 +10911,7 @@ class ListReasoningEnginesMemoriesResponseDict(TypedDict, total=False): """""" memories: Optional[list[MemoryDict]] - """List of agent engine memories.""" + """List of agent runtime memories.""" ListReasoningEnginesMemoriesResponseOrDict = Union[ @@ -11462,57 +10919,56 @@ class ListReasoningEnginesMemoriesResponseDict(TypedDict, total=False): ] -class _GetAgentEngineMemoryOperationParameters(_common.BaseModel): +class _GetRuntimeMemoryOperationParameters(_common.BaseModel): """Parameters for getting an operation with a memory as a response.""" operation_name: Optional[str] = Field( default=None, description="""The server-assigned name for the operation.""" ) - config: Optional[GetAgentEngineOperationConfig] = Field( + config: Optional[GetRuntimeOperationConfig] = Field( default=None, description="""Used to override the default configuration.""" ) -class _GetAgentEngineMemoryOperationParametersDict(TypedDict, total=False): +class _GetRuntimeMemoryOperationParametersDict(TypedDict, total=False): """Parameters for getting an operation with a memory as a response.""" operation_name: Optional[str] """The server-assigned name for the operation.""" - config: Optional[GetAgentEngineOperationConfigDict] + config: Optional[GetRuntimeOperationConfigDict] """Used to override the default configuration.""" -_GetAgentEngineMemoryOperationParametersOrDict = Union[ - _GetAgentEngineMemoryOperationParameters, - _GetAgentEngineMemoryOperationParametersDict, +_GetRuntimeMemoryOperationParametersOrDict = Union[ + _GetRuntimeMemoryOperationParameters, _GetRuntimeMemoryOperationParametersDict ] -class _GetAgentEngineGenerateMemoriesOperationParameters(_common.BaseModel): +class _GetRuntimeGenerateMemoriesOperationParameters(_common.BaseModel): """Parameters for getting an operation with generated memories as a response.""" operation_name: Optional[str] = Field( default=None, description="""The server-assigned name for the operation.""" ) - config: Optional[GetAgentEngineOperationConfig] = Field( + config: Optional[GetRuntimeOperationConfig] = Field( default=None, description="""Used to override the default configuration.""" ) -class _GetAgentEngineGenerateMemoriesOperationParametersDict(TypedDict, total=False): +class _GetRuntimeGenerateMemoriesOperationParametersDict(TypedDict, total=False): """Parameters for getting an operation with generated memories as a response.""" operation_name: Optional[str] """The server-assigned name for the operation.""" - config: Optional[GetAgentEngineOperationConfigDict] + config: Optional[GetRuntimeOperationConfigDict] """Used to override the default configuration.""" -_GetAgentEngineGenerateMemoriesOperationParametersOrDict = Union[ - _GetAgentEngineGenerateMemoriesOperationParameters, - _GetAgentEngineGenerateMemoriesOperationParametersDict, +_GetRuntimeGenerateMemoriesOperationParametersOrDict = Union[ + _GetRuntimeGenerateMemoriesOperationParameters, + _GetRuntimeGenerateMemoriesOperationParametersDict, ] @@ -11633,7 +11089,7 @@ class MemoryConjunctionFilterDict(TypedDict, total=False): ] -class RetrieveAgentEngineMemoriesConfig(_common.BaseModel): +class RetrieveRuntimeMemoriesConfig(_common.BaseModel): """Config for retrieving memories.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -11676,7 +11132,7 @@ class RetrieveAgentEngineMemoriesConfig(_common.BaseModel): ) -class RetrieveAgentEngineMemoriesConfigDict(TypedDict, total=False): +class RetrieveRuntimeMemoriesConfigDict(TypedDict, total=False): """Config for retrieving memories.""" http_options: Optional[genai_types.HttpOptions] @@ -11715,17 +11171,17 @@ class RetrieveAgentEngineMemoriesConfigDict(TypedDict, total=False): retrieve memories matching any of the specified `MemoryType` values.""" -RetrieveAgentEngineMemoriesConfigOrDict = Union[ - RetrieveAgentEngineMemoriesConfig, RetrieveAgentEngineMemoriesConfigDict +RetrieveRuntimeMemoriesConfigOrDict = Union[ + RetrieveRuntimeMemoriesConfig, RetrieveRuntimeMemoriesConfigDict ] -class _RetrieveAgentEngineMemoriesRequestParameters(_common.BaseModel): - """Parameters for retrieving agent engine memories.""" +class _RetrieveRuntimeMemoriesRequestParameters(_common.BaseModel): + """Parameters for retrieving agent runtime memories.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to retrieve memories from.""", + description="""Name of the agent runtime to retrieve memories from.""", ) scope: Optional[dict[str, str]] = Field( default=None, @@ -11747,16 +11203,16 @@ class _RetrieveAgentEngineMemoriesRequestParameters(_common.BaseModel): description="""Parameters for simple (non-similarity search) retrieval.""", ) ) - config: Optional[RetrieveAgentEngineMemoriesConfig] = Field( + config: Optional[RetrieveRuntimeMemoriesConfig] = Field( default=None, description="""""" ) -class _RetrieveAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving agent engine memories.""" +class _RetrieveRuntimeMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for retrieving agent runtime memories.""" name: Optional[str] - """Name of the agent engine to retrieve memories from.""" + """Name of the agent runtime to retrieve memories from.""" scope: Optional[dict[str, str]] """The scope of the memories to retrieve. @@ -11773,13 +11229,13 @@ class _RetrieveAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): simple_retrieval_params: Optional[RetrieveMemoriesRequestSimpleRetrievalParamsDict] """Parameters for simple (non-similarity search) retrieval.""" - config: Optional[RetrieveAgentEngineMemoriesConfigDict] + config: Optional[RetrieveRuntimeMemoriesConfigDict] """""" -_RetrieveAgentEngineMemoriesRequestParametersOrDict = Union[ - _RetrieveAgentEngineMemoriesRequestParameters, - _RetrieveAgentEngineMemoriesRequestParametersDict, +_RetrieveRuntimeMemoriesRequestParametersOrDict = Union[ + _RetrieveRuntimeMemoriesRequestParameters, + _RetrieveRuntimeMemoriesRequestParametersDict, ] @@ -11858,11 +11314,11 @@ class RetrieveMemoryProfilesConfigDict(TypedDict, total=False): class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): - """Parameters for retrieving agent engine memory profiles.""" + """Parameters for retrieving agent runtime memory profiles.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to retrieve memory profiles from.""", + description="""Name of the agent runtime to retrieve memory profiles from.""", ) scope: Optional[dict[str, str]] = Field( default=None, @@ -11878,10 +11334,10 @@ class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): class _RetrieveMemoryProfilesRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving agent engine memory profiles.""" + """Parameters for retrieving agent runtime memory profiles.""" name: Optional[str] - """Name of the agent engine to retrieve memory profiles from.""" + """Name of the agent runtime to retrieve memory profiles from.""" scope: Optional[dict[str, str]] """The scope of the memories to retrieve. @@ -11952,7 +11408,7 @@ class RetrieveProfilesResponseDict(TypedDict, total=False): ] -class RollbackAgentEngineMemoryConfig(_common.BaseModel): +class RollbackRuntimeMemoryConfig(_common.BaseModel): """Config for rolling back a memory.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -11964,7 +11420,7 @@ class RollbackAgentEngineMemoryConfig(_common.BaseModel): ) -class RollbackAgentEngineMemoryConfigDict(TypedDict, total=False): +class RollbackRuntimeMemoryConfigDict(TypedDict, total=False): """Config for rolling back a memory.""" http_options: Optional[genai_types.HttpOptions] @@ -11974,45 +11430,44 @@ class RollbackAgentEngineMemoryConfigDict(TypedDict, total=False): """Waits for the operation to complete before returning.""" -RollbackAgentEngineMemoryConfigOrDict = Union[ - RollbackAgentEngineMemoryConfig, RollbackAgentEngineMemoryConfigDict +RollbackRuntimeMemoryConfigOrDict = Union[ + RollbackRuntimeMemoryConfig, RollbackRuntimeMemoryConfigDict ] -class _RollbackAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for generating agent engine memories.""" +class _RollbackRuntimeMemoryRequestParameters(_common.BaseModel): + """Parameters for generating agent runtime memories.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine memory to rollback.""" + default=None, description="""Name of the agent runtime memory to rollback.""" ) target_revision_id: Optional[str] = Field( default=None, description="""The ID of the revision to rollback to.""" ) - config: Optional[RollbackAgentEngineMemoryConfig] = Field( + config: Optional[RollbackRuntimeMemoryConfig] = Field( default=None, description="""""" ) -class _RollbackAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for generating agent engine memories.""" +class _RollbackRuntimeMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for generating agent runtime memories.""" name: Optional[str] - """Name of the agent engine memory to rollback.""" + """Name of the agent runtime memory to rollback.""" target_revision_id: Optional[str] """The ID of the revision to rollback to.""" - config: Optional[RollbackAgentEngineMemoryConfigDict] + config: Optional[RollbackRuntimeMemoryConfigDict] """""" -_RollbackAgentEngineMemoryRequestParametersOrDict = Union[ - _RollbackAgentEngineMemoryRequestParameters, - _RollbackAgentEngineMemoryRequestParametersDict, +_RollbackRuntimeMemoryRequestParametersOrDict = Union[ + _RollbackRuntimeMemoryRequestParameters, _RollbackRuntimeMemoryRequestParametersDict ] -class AgentEngineRollbackMemoryOperation(_common.BaseModel): +class RuntimeRollbackMemoryOperation(_common.BaseModel): """Operation that rolls back a memory.""" name: Optional[str] = Field( @@ -12033,7 +11488,7 @@ class AgentEngineRollbackMemoryOperation(_common.BaseModel): ) -class AgentEngineRollbackMemoryOperationDict(TypedDict, total=False): +class RuntimeRollbackMemoryOperationDict(TypedDict, total=False): """Operation that rolls back a memory.""" name: Optional[str] @@ -12049,13 +11504,13 @@ class AgentEngineRollbackMemoryOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -AgentEngineRollbackMemoryOperationOrDict = Union[ - AgentEngineRollbackMemoryOperation, AgentEngineRollbackMemoryOperationDict +RuntimeRollbackMemoryOperationOrDict = Union[ + RuntimeRollbackMemoryOperation, RuntimeRollbackMemoryOperationDict ] -class UpdateAgentEngineMemoryConfig(_common.BaseModel): - """Config for updating agent engine memory.""" +class UpdateRuntimeMemoryConfig(_common.BaseModel): + """Config for updating agent runtime memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -12110,8 +11565,8 @@ class UpdateAgentEngineMemoryConfig(_common.BaseModel): ) -class UpdateAgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for updating agent engine memory.""" +class UpdateRuntimeMemoryConfigDict(TypedDict, total=False): + """Config for updating agent runtime memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -12156,16 +11611,16 @@ class UpdateAgentEngineMemoryConfigDict(TypedDict, total=False): https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" -UpdateAgentEngineMemoryConfigOrDict = Union[ - UpdateAgentEngineMemoryConfig, UpdateAgentEngineMemoryConfigDict +UpdateRuntimeMemoryConfigOrDict = Union[ + UpdateRuntimeMemoryConfig, UpdateRuntimeMemoryConfigDict ] -class _UpdateAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for updating agent engine memories.""" +class _UpdateRuntimeMemoryRequestParameters(_common.BaseModel): + """Parameters for updating agent runtime memories.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine memory to update.""" + default=None, description="""Name of the agent runtime memory to update.""" ) fact: Optional[str] = Field( default=None, @@ -12181,16 +11636,16 @@ class _UpdateAgentEngineMemoryRequestParameters(_common.BaseModel): creating or generating memories. Up to 5 key-value pairs are accepted, and scope values cannot contain the wildcard character '*'.""", ) - config: Optional[UpdateAgentEngineMemoryConfig] = Field( + config: Optional[UpdateRuntimeMemoryConfig] = Field( default=None, description="""""" ) -class _UpdateAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for updating agent engine memories.""" +class _UpdateRuntimeMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for updating agent runtime memories.""" name: Optional[str] - """Name of the agent engine memory to update.""" + """Name of the agent runtime memory to update.""" fact: Optional[str] """The updated fact of the memory. @@ -12204,17 +11659,16 @@ class _UpdateAgentEngineMemoryRequestParametersDict(TypedDict, total=False): creating or generating memories. Up to 5 key-value pairs are accepted, and scope values cannot contain the wildcard character '*'.""" - config: Optional[UpdateAgentEngineMemoryConfigDict] + config: Optional[UpdateRuntimeMemoryConfigDict] """""" -_UpdateAgentEngineMemoryRequestParametersOrDict = Union[ - _UpdateAgentEngineMemoryRequestParameters, - _UpdateAgentEngineMemoryRequestParametersDict, +_UpdateRuntimeMemoryRequestParametersOrDict = Union[ + _UpdateRuntimeMemoryRequestParameters, _UpdateRuntimeMemoryRequestParametersDict ] -class PurgeAgentEngineMemoriesConfig(_common.BaseModel): +class PurgeRuntimeMemoriesConfig(_common.BaseModel): """Config for purging memories.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -12226,7 +11680,7 @@ class PurgeAgentEngineMemoriesConfig(_common.BaseModel): ) -class PurgeAgentEngineMemoriesConfigDict(TypedDict, total=False): +class PurgeRuntimeMemoriesConfigDict(TypedDict, total=False): """Config for purging memories.""" http_options: Optional[genai_types.HttpOptions] @@ -12236,16 +11690,17 @@ class PurgeAgentEngineMemoriesConfigDict(TypedDict, total=False): """Waits for the operation to complete before returning.""" -PurgeAgentEngineMemoriesConfigOrDict = Union[ - PurgeAgentEngineMemoriesConfig, PurgeAgentEngineMemoriesConfigDict +PurgeRuntimeMemoriesConfigOrDict = Union[ + PurgeRuntimeMemoriesConfig, PurgeRuntimeMemoriesConfigDict ] -class _PurgeAgentEngineMemoriesRequestParameters(_common.BaseModel): - """Parameters for purging agent engine memories.""" +class _PurgeRuntimeMemoriesRequestParameters(_common.BaseModel): + """Parameters for purging agent runtime memories.""" name: Optional[str] = Field( - default=None, description="""Name of the Agent Engine to purge memories from.""" + default=None, + description="""Name of the Agent Runtime to purge memories from.""", ) filter: Optional[str] = Field( default=None, @@ -12273,16 +11728,16 @@ class _PurgeAgentEngineMemoriesRequestParameters(_common.BaseModel): default=None, description="""If true, the memories will actually be purged. If false, the purge request will be validated but not executed.""", ) - config: Optional[PurgeAgentEngineMemoriesConfig] = Field( + config: Optional[PurgeRuntimeMemoriesConfig] = Field( default=None, description="""""" ) -class _PurgeAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for purging agent engine memories.""" +class _PurgeRuntimeMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for purging agent runtime memories.""" name: Optional[str] - """Name of the Agent Engine to purge memories from.""" + """Name of the Agent Runtime to purge memories from.""" filter: Optional[str] """The standard list filter to determine which memories to purge. @@ -12307,13 +11762,12 @@ class _PurgeAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): force: Optional[bool] """If true, the memories will actually be purged. If false, the purge request will be validated but not executed.""" - config: Optional[PurgeAgentEngineMemoriesConfigDict] + config: Optional[PurgeRuntimeMemoriesConfigDict] """""" -_PurgeAgentEngineMemoriesRequestParametersOrDict = Union[ - _PurgeAgentEngineMemoriesRequestParameters, - _PurgeAgentEngineMemoriesRequestParametersDict, +_PurgeRuntimeMemoriesRequestParametersOrDict = Union[ + _PurgeRuntimeMemoriesRequestParameters, _PurgeRuntimeMemoriesRequestParametersDict ] @@ -12335,8 +11789,8 @@ class PurgeMemoriesResponseDict(TypedDict, total=False): PurgeMemoriesResponseOrDict = Union[PurgeMemoriesResponse, PurgeMemoriesResponseDict] -class AgentEnginePurgeMemoriesOperation(_common.BaseModel): - """Operation that purges memories from an agent engine.""" +class RuntimePurgeMemoriesOperation(_common.BaseModel): + """Operation that purges memories from an agent runtime.""" name: Optional[str] = Field( default=None, @@ -12359,8 +11813,8 @@ class AgentEnginePurgeMemoriesOperation(_common.BaseModel): ) -class AgentEnginePurgeMemoriesOperationDict(TypedDict, total=False): - """Operation that purges memories from an agent engine.""" +class RuntimePurgeMemoriesOperationDict(TypedDict, total=False): + """Operation that purges memories from an agent runtime.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -12378,55 +11832,55 @@ class AgentEnginePurgeMemoriesOperationDict(TypedDict, total=False): """The response for purging memories.""" -AgentEnginePurgeMemoriesOperationOrDict = Union[ - AgentEnginePurgeMemoriesOperation, AgentEnginePurgeMemoriesOperationDict +RuntimePurgeMemoriesOperationOrDict = Union[ + RuntimePurgeMemoriesOperation, RuntimePurgeMemoriesOperationDict ] -class GetAgentEngineMemoryRevisionConfig(_common.BaseModel): - """Config for getting an Agent Engine Memory Revision.""" +class GetRuntimeMemoryRevisionConfig(_common.BaseModel): + """Config for getting an Agent Runtime Memory Revision.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetAgentEngineMemoryRevisionConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Memory Revision.""" +class GetRuntimeMemoryRevisionConfigDict(TypedDict, total=False): + """Config for getting an Agent Runtime Memory Revision.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetAgentEngineMemoryRevisionConfigOrDict = Union[ - GetAgentEngineMemoryRevisionConfig, GetAgentEngineMemoryRevisionConfigDict +GetRuntimeMemoryRevisionConfigOrDict = Union[ + GetRuntimeMemoryRevisionConfig, GetRuntimeMemoryRevisionConfigDict ] -class _GetAgentEngineMemoryRevisionRequestParameters(_common.BaseModel): - """Parameters for getting an Agent Engine memory revision.""" +class _GetRuntimeMemoryRevisionRequestParameters(_common.BaseModel): + """Parameters for getting an Agent Runtime memory revision.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime.""" ) - config: Optional[GetAgentEngineMemoryRevisionConfig] = Field( + config: Optional[GetRuntimeMemoryRevisionConfig] = Field( default=None, description="""""" ) -class _GetAgentEngineMemoryRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for getting an Agent Engine memory revision.""" +class _GetRuntimeMemoryRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for getting an Agent Runtime memory revision.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime.""" - config: Optional[GetAgentEngineMemoryRevisionConfigDict] + config: Optional[GetRuntimeMemoryRevisionConfigDict] """""" -_GetAgentEngineMemoryRevisionRequestParametersOrDict = Union[ - _GetAgentEngineMemoryRevisionRequestParameters, - _GetAgentEngineMemoryRevisionRequestParametersDict, +_GetRuntimeMemoryRevisionRequestParametersOrDict = Union[ + _GetRuntimeMemoryRevisionRequestParameters, + _GetRuntimeMemoryRevisionRequestParametersDict, ] @@ -12526,8 +11980,8 @@ class MemoryRevisionDict(TypedDict, total=False): MemoryRevisionOrDict = Union[MemoryRevision, MemoryRevisionDict] -class ListAgentEngineMemoryRevisionsConfig(_common.BaseModel): - """Config for listing Agent Engine memory revisions.""" +class ListRuntimeMemoryRevisionsConfig(_common.BaseModel): + """Config for listing Agent Runtime memory revisions.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -12541,8 +11995,8 @@ class ListAgentEngineMemoryRevisionsConfig(_common.BaseModel): ) -class ListAgentEngineMemoryRevisionsConfigDict(TypedDict, total=False): - """Config for listing Agent Engine memory revisions.""" +class ListRuntimeMemoryRevisionsConfigDict(TypedDict, total=False): + """Config for listing Agent Runtime memory revisions.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -12558,40 +12012,40 @@ class ListAgentEngineMemoryRevisionsConfigDict(TypedDict, total=False): For field names both snake_case and camelCase are supported.""" -ListAgentEngineMemoryRevisionsConfigOrDict = Union[ - ListAgentEngineMemoryRevisionsConfig, ListAgentEngineMemoryRevisionsConfigDict +ListRuntimeMemoryRevisionsConfigOrDict = Union[ + ListRuntimeMemoryRevisionsConfig, ListRuntimeMemoryRevisionsConfigDict ] -class _ListAgentEngineMemoryRevisionsRequestParameters(_common.BaseModel): - """Parameters for listing Agent Engine memory revisions.""" +class _ListRuntimeMemoryRevisionsRequestParameters(_common.BaseModel): + """Parameters for listing Agent Runtime memory revisions.""" name: Optional[str] = Field( - default=None, description="""Name of the Agent Engine memory""" + default=None, description="""Name of the Agent Runtime memory""" ) - config: Optional[ListAgentEngineMemoryRevisionsConfig] = Field( + config: Optional[ListRuntimeMemoryRevisionsConfig] = Field( default=None, description="""""" ) -class _ListAgentEngineMemoryRevisionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing Agent Engine memory revisions.""" +class _ListRuntimeMemoryRevisionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing Agent Runtime memory revisions.""" name: Optional[str] - """Name of the Agent Engine memory""" + """Name of the Agent Runtime memory""" - config: Optional[ListAgentEngineMemoryRevisionsConfigDict] + config: Optional[ListRuntimeMemoryRevisionsConfigDict] """""" -_ListAgentEngineMemoryRevisionsRequestParametersOrDict = Union[ - _ListAgentEngineMemoryRevisionsRequestParameters, - _ListAgentEngineMemoryRevisionsRequestParametersDict, +_ListRuntimeMemoryRevisionsRequestParametersOrDict = Union[ + _ListRuntimeMemoryRevisionsRequestParameters, + _ListRuntimeMemoryRevisionsRequestParametersDict, ] -class ListAgentEngineMemoryRevisionsResponse(_common.BaseModel): - """Response for listing agent engine memory revisions.""" +class ListRuntimeMemoryRevisionsResponse(_common.BaseModel): + """Response for listing agent runtime memory revisions.""" sdk_http_response: Optional[genai_types.HttpResponse] = Field( default=None, description="""Used to retain the full HTTP response.""" @@ -12602,8 +12056,8 @@ class ListAgentEngineMemoryRevisionsResponse(_common.BaseModel): ) -class ListAgentEngineMemoryRevisionsResponseDict(TypedDict, total=False): - """Response for listing agent engine memory revisions.""" +class ListRuntimeMemoryRevisionsResponseDict(TypedDict, total=False): + """Response for listing agent runtime memory revisions.""" sdk_http_response: Optional[genai_types.HttpResponse] """Used to retain the full HTTP response.""" @@ -12615,8 +12069,8 @@ class ListAgentEngineMemoryRevisionsResponseDict(TypedDict, total=False): """List of memory revisions.""" -ListAgentEngineMemoryRevisionsResponseOrDict = Union[ - ListAgentEngineMemoryRevisionsResponse, ListAgentEngineMemoryRevisionsResponseDict +ListRuntimeMemoryRevisionsResponseOrDict = Union[ + ListRuntimeMemoryRevisionsResponse, ListRuntimeMemoryRevisionsResponseDict ] @@ -15734,250 +15188,330 @@ class UploadRagFileResponseDict(TypedDict, total=False): UploadRagFileResponseOrDict = Union[UploadRagFileResponse, UploadRagFileResponseDict] -class GetAgentEngineRuntimeRevisionConfig(_common.BaseModel): - """Config for getting an Agent Engine Runtime Revision.""" +class SandboxEnvironmentSpecCodeExecutionEnvironment(_common.BaseModel): + """The code execution environment with customized settings.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + code_language: Optional[Language] = Field( + default=None, + description="""The coding language supported in this environment.""", + ) + machine_config: Optional[MachineConfig] = Field( + default=None, + description="""The machine config of the code execution environment.""", ) -class GetAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Runtime Revision.""" +class SandboxEnvironmentSpecCodeExecutionEnvironmentDict(TypedDict, total=False): + """The code execution environment with customized settings.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + code_language: Optional[Language] + """The coding language supported in this environment.""" + machine_config: Optional[MachineConfig] + """The machine config of the code execution environment.""" -GetAgentEngineRuntimeRevisionConfigOrDict = Union[ - GetAgentEngineRuntimeRevisionConfig, GetAgentEngineRuntimeRevisionConfigDict -] +SandboxEnvironmentSpecCodeExecutionEnvironmentOrDict = Union[ + SandboxEnvironmentSpecCodeExecutionEnvironment, + SandboxEnvironmentSpecCodeExecutionEnvironmentDict, +] -class _GetAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine runtime revision.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine runtime revision.""" - ) - config: Optional[GetAgentEngineRuntimeRevisionConfig] = Field( - default=None, description="""""" - ) +class SandboxEnvironmentSpecComputerUseEnvironment(_common.BaseModel): + """The computer use environment with customized settings.""" + pass -class _GetAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine runtime revision.""" - name: Optional[str] - """Name of the agent engine runtime revision.""" +class SandboxEnvironmentSpecComputerUseEnvironmentDict(TypedDict, total=False): + """The computer use environment with customized settings.""" - config: Optional[GetAgentEngineRuntimeRevisionConfigDict] - """""" + pass -_GetAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ - _GetAgentEngineRuntimeRevisionRequestParameters, - _GetAgentEngineRuntimeRevisionRequestParametersDict, +SandboxEnvironmentSpecComputerUseEnvironmentOrDict = Union[ + SandboxEnvironmentSpecComputerUseEnvironment, + SandboxEnvironmentSpecComputerUseEnvironmentDict, ] -class ReasoningEngineRuntimeRevision(_common.BaseModel): - """A runtime revision.""" +class SandboxEnvironmentSpec(_common.BaseModel): + """The specification of a sandbox environment.""" - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when this ReasoningEngineRuntimeRevision was created.""", - ) - name: Optional[str] = Field( - default=None, - description="""Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""", - ) - spec: Optional[ReasoningEngineSpec] = Field( - default=None, - description="""Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""", - ) - state: Optional[State] = Field( - default=None, description="""Output only. The state of the revision.""" + code_execution_environment: Optional[ + SandboxEnvironmentSpecCodeExecutionEnvironment + ] = Field(default=None, description="""Optional. The code execution environment.""") + computer_use_environment: Optional[SandboxEnvironmentSpecComputerUseEnvironment] = ( + Field(default=None, description="""Optional. The computer use environment.""") ) -class ReasoningEngineRuntimeRevisionDict(TypedDict, total=False): - """A runtime revision.""" - - create_time: Optional[datetime.datetime] - """Output only. Timestamp when this ReasoningEngineRuntimeRevision was created.""" - - name: Optional[str] - """Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""" +class SandboxEnvironmentSpecDict(TypedDict, total=False): + """The specification of a sandbox environment.""" - spec: Optional[ReasoningEngineSpecDict] - """Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""" + code_execution_environment: Optional[ + SandboxEnvironmentSpecCodeExecutionEnvironmentDict + ] + """Optional. The code execution environment.""" - state: Optional[State] - """Output only. The state of the revision.""" + computer_use_environment: Optional[SandboxEnvironmentSpecComputerUseEnvironmentDict] + """Optional. The computer use environment.""" -ReasoningEngineRuntimeRevisionOrDict = Union[ - ReasoningEngineRuntimeRevision, ReasoningEngineRuntimeRevisionDict -] +SandboxEnvironmentSpecOrDict = Union[SandboxEnvironmentSpec, SandboxEnvironmentSpecDict] -class ListAgentEngineRuntimeRevisionsConfig(_common.BaseModel): - """Config for listing reasoning engine runtime revisions.""" +class CreateRuntimeSandboxConfig(_common.BaseModel): + """Config for creating a Sandbox.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( - default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""", + display_name: Optional[str] = Field( + default=None, description="""The display name of the sandbox.""" + ) + description: Optional[str] = Field( + default=None, description="""The description of the sandbox.""" + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", + ) + ttl: Optional[str] = Field( + default=None, + description="""The TTL for this resource. The expiration time is computed: now + TTL.""", + ) + sandbox_environment_template: Optional[str] = Field( + default=None, + description="""The name of the sandbox environment template to create the sandbox from. The sandbox environment template should be in the format: + projects/{project}/locations/{location}/agentEngines/{agent_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}""", + ) + sandbox_environment_snapshot: Optional[str] = Field( + default=None, + description="""The name of the sandbox environment snapshot to restore the sandbox from. The sandbox environment snapshot should be in the format: + projects/{project}/locations/{location}/agentEngines/{agent_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}""", + ) + owner: Optional[str] = Field( + default=None, + description="""Owner information for this sandbox environment. A sandbox can only be restored from a snapshot belonging to the same owner.""", ) -class ListAgentEngineRuntimeRevisionsConfigDict(TypedDict, total=False): - """Config for listing reasoning engine runtime revisions.""" +class CreateRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for creating a Sandbox.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - page_size: Optional[int] - """""" + display_name: Optional[str] + """The display name of the sandbox.""" - page_token: Optional[str] - """""" + description: Optional[str] + """The description of the sandbox.""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" + + ttl: Optional[str] + """The TTL for this resource. The expiration time is computed: now + TTL.""" + + sandbox_environment_template: Optional[str] + """The name of the sandbox environment template to create the sandbox from. The sandbox environment template should be in the format: + projects/{project}/locations/{location}/agentEngines/{agent_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}""" + + sandbox_environment_snapshot: Optional[str] + """The name of the sandbox environment snapshot to restore the sandbox from. The sandbox environment snapshot should be in the format: + projects/{project}/locations/{location}/agentEngines/{agent_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}""" + + owner: Optional[str] + """Owner information for this sandbox environment. A sandbox can only be restored from a snapshot belonging to the same owner.""" -ListAgentEngineRuntimeRevisionsConfigOrDict = Union[ - ListAgentEngineRuntimeRevisionsConfig, ListAgentEngineRuntimeRevisionsConfigDict +CreateRuntimeSandboxConfigOrDict = Union[ + CreateRuntimeSandboxConfig, CreateRuntimeSandboxConfigDict ] -class _ListAgentEngineRuntimeRevisionsRequestParameters(_common.BaseModel): - """Parameters for listing reasoning engine runtime revisions.""" +class _CreateRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for creating Agent Runtime Sandboxes.""" name: Optional[str] = Field( - default=None, description="""Name of the reasoning engine.""" + default=None, + description="""Name of the Agent Runtime to create the sandbox under.""", + ) + spec: Optional[SandboxEnvironmentSpec] = Field( + default=None, description="""The specification of the sandbox.""" ) - config: Optional[ListAgentEngineRuntimeRevisionsConfig] = Field( + config: Optional[CreateRuntimeSandboxConfig] = Field( default=None, description="""""" ) -class _ListAgentEngineRuntimeRevisionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing reasoning engine runtime revisions.""" +class _CreateRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for creating Agent Runtime Sandboxes.""" name: Optional[str] - """Name of the reasoning engine.""" + """Name of the Agent Runtime to create the sandbox under.""" + + spec: Optional[SandboxEnvironmentSpecDict] + """The specification of the sandbox.""" - config: Optional[ListAgentEngineRuntimeRevisionsConfigDict] + config: Optional[CreateRuntimeSandboxConfigDict] """""" -_ListAgentEngineRuntimeRevisionsRequestParametersOrDict = Union[ - _ListAgentEngineRuntimeRevisionsRequestParameters, - _ListAgentEngineRuntimeRevisionsRequestParametersDict, +_CreateRuntimeSandboxRequestParametersOrDict = Union[ + _CreateRuntimeSandboxRequestParameters, _CreateRuntimeSandboxRequestParametersDict ] -class ListReasoningEnginesRuntimeRevisionsResponse(_common.BaseModel): - """Response for listing agent engine runtime revisions.""" +class SandboxEnvironmentConnectionInfo(_common.BaseModel): + """The connection information of the SandboxEnvironment.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + load_balancer_hostname: Optional[str] = Field( + default=None, description="""Output only. The hostname of the load balancer.""" ) - next_page_token: Optional[str] = Field(default=None, description="""""") - reasoning_engine_runtime_revisions: Optional[ - list[ReasoningEngineRuntimeRevision] - ] = Field( - default=None, description="""List of reasoning engine runtime revisions.""" + load_balancer_ip: Optional[str] = Field( + default=None, + description="""Output only. The IP address of the load balancer.""", + ) + sandbox_internal_ip: Optional[str] = Field( + default=None, + description="""Output only. The internal IP address of the SandboxEnvironment.""", + ) + routing_token: Optional[str] = Field( + default=None, + description="""Output only. The routing token for the SandboxEnvironment.""", ) -class ListReasoningEnginesRuntimeRevisionsResponseDict(TypedDict, total=False): - """Response for listing agent engine runtime revisions.""" +class SandboxEnvironmentConnectionInfoDict(TypedDict, total=False): + """The connection information of the SandboxEnvironment.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + load_balancer_hostname: Optional[str] + """Output only. The hostname of the load balancer.""" - next_page_token: Optional[str] - """""" + load_balancer_ip: Optional[str] + """Output only. The IP address of the load balancer.""" - reasoning_engine_runtime_revisions: Optional[ - list[ReasoningEngineRuntimeRevisionDict] - ] - """List of reasoning engine runtime revisions.""" + sandbox_internal_ip: Optional[str] + """Output only. The internal IP address of the SandboxEnvironment.""" + routing_token: Optional[str] + """Output only. The routing token for the SandboxEnvironment.""" -ListReasoningEnginesRuntimeRevisionsResponseOrDict = Union[ - ListReasoningEnginesRuntimeRevisionsResponse, - ListReasoningEnginesRuntimeRevisionsResponseDict, + +SandboxEnvironmentConnectionInfoOrDict = Union[ + SandboxEnvironmentConnectionInfo, SandboxEnvironmentConnectionInfoDict ] -class DeleteAgentEngineRuntimeRevisionConfig(_common.BaseModel): - """Config for deleting an Agent Engine Runtime Revision.""" +class SandboxEnvironment(_common.BaseModel): + """A sandbox environment.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + expire_time: Optional[datetime.datetime] = Field( + default=None, + description="""Expiration time of the sandbox environment. + """, ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", + connection_info: Optional[SandboxEnvironmentConnectionInfo] = Field( + default=None, + description="""Output only. The connection information of the SandboxEnvironment.""", + ) + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. The timestamp when this SandboxEnvironment was created.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""Required. The display name of the SandboxEnvironment.""", + ) + name: Optional[str] = Field( + default=None, description="""Identifier. The name of the SandboxEnvironment.""" + ) + spec: Optional[SandboxEnvironmentSpec] = Field( + default=None, + description="""Optional. The configuration of the SandboxEnvironment.""", + ) + state: Optional[SandboxState] = Field( + default=None, + description="""Output only. The runtime state of the SandboxEnvironment.""", + ) + ttl: Optional[str] = Field( + default=None, + description="""Optional. Input only. The TTL for the sandbox environment. The expiration time is computed: now + TTL.""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. The timestamp when this SandboxEnvironment was most recently updated.""", + ) + latest_sandbox_environment_snapshot: Optional[str] = Field( + default=None, + description="""Output only. The resource name of the latest snapshot taken for this SandboxEnvironment.""", + ) + owner: Optional[str] = Field( + default=None, + description="""Optional. Owner information for this sandbox environment. A Sandbox can only be restored from a snapshot that belongs to the same owner. If not set, sandbox will be created as the default owner.""", + ) + sandbox_environment_snapshot: Optional[str] = Field( + default=None, + description="""Optional. The resource name of the SandboxEnvironmentSnapshot to use for creating this SandboxEnvironment. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}`""", + ) + sandbox_environment_template: Optional[str] = Field( + default=None, + description="""Optional. The name of the SandboxEnvironmentTemplate specified in the parent Agent Engine resource that this SandboxEnvironment is created from.""", ) -class DeleteAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Runtime Revision.""" +class SandboxEnvironmentDict(TypedDict, total=False): + """A sandbox environment.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + expire_time: Optional[datetime.datetime] + """Expiration time of the sandbox environment. + """ - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + connection_info: Optional[SandboxEnvironmentConnectionInfoDict] + """Output only. The connection information of the SandboxEnvironment.""" + create_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironment was created.""" -DeleteAgentEngineRuntimeRevisionConfigOrDict = Union[ - DeleteAgentEngineRuntimeRevisionConfig, DeleteAgentEngineRuntimeRevisionConfigDict -] + display_name: Optional[str] + """Required. The display name of the SandboxEnvironment.""" + name: Optional[str] + """Identifier. The name of the SandboxEnvironment.""" -class _DeleteAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): - """Parameters for deleting agent engine runtime revisions.""" + spec: Optional[SandboxEnvironmentSpecDict] + """Optional. The configuration of the SandboxEnvironment.""" - name: Optional[str] = Field( - default=None, - description="""Name of the agent engine runtime revision to delete.""", - ) - config: Optional[DeleteAgentEngineRuntimeRevisionConfig] = Field( - default=None, description="""""" - ) + state: Optional[SandboxState] + """Output only. The runtime state of the SandboxEnvironment.""" + ttl: Optional[str] + """Optional. Input only. The TTL for the sandbox environment. The expiration time is computed: now + TTL.""" -class _DeleteAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engine runtime revisions.""" + update_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironment was most recently updated.""" - name: Optional[str] - """Name of the agent engine runtime revision to delete.""" + latest_sandbox_environment_snapshot: Optional[str] + """Output only. The resource name of the latest snapshot taken for this SandboxEnvironment.""" - config: Optional[DeleteAgentEngineRuntimeRevisionConfigDict] - """""" + owner: Optional[str] + """Optional. Owner information for this sandbox environment. A Sandbox can only be restored from a snapshot that belongs to the same owner. If not set, sandbox will be created as the default owner.""" + + sandbox_environment_snapshot: Optional[str] + """Optional. The resource name of the SandboxEnvironmentSnapshot to use for creating this SandboxEnvironment. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}`""" + sandbox_environment_template: Optional[str] + """Optional. The name of the SandboxEnvironmentTemplate specified in the parent Agent Engine resource that this SandboxEnvironment is created from.""" -_DeleteAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ - _DeleteAgentEngineRuntimeRevisionRequestParameters, - _DeleteAgentEngineRuntimeRevisionRequestParametersDict, -] + +SandboxEnvironmentOrDict = Union[SandboxEnvironment, SandboxEnvironmentDict] -class DeleteAgentEngineRuntimeRevisionOperation(_common.BaseModel): - """Operation for deleting agent engine runtime revisions.""" +class RuntimeSandboxOperation(_common.BaseModel): + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] = Field( default=None, @@ -15995,10 +15529,13 @@ class DeleteAgentEngineRuntimeRevisionOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) + response: Optional[SandboxEnvironment] = Field( + default=None, description="""The Agent Runtime Sandbox.""" + ) -class DeleteAgentEngineRuntimeRevisionOperationDict(TypedDict, total=False): - """Operation for deleting agent engine runtime revisions.""" +class RuntimeSandboxOperationDict(TypedDict, total=False): + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -16012,1233 +15549,1218 @@ class DeleteAgentEngineRuntimeRevisionOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" + response: Optional[SandboxEnvironmentDict] + """The Agent Runtime Sandbox.""" + -DeleteAgentEngineRuntimeRevisionOperationOrDict = Union[ - DeleteAgentEngineRuntimeRevisionOperation, - DeleteAgentEngineRuntimeRevisionOperationDict, +RuntimeSandboxOperationOrDict = Union[ + RuntimeSandboxOperation, RuntimeSandboxOperationDict ] -class GetDeleteAgentEngineRuntimeRevisionOperationConfig(_common.BaseModel): +class DeleteRuntimeSandboxConfig(_common.BaseModel): + """Config for deleting an Agent Runtime Sandbox.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetDeleteAgentEngineRuntimeRevisionOperationConfigDict(TypedDict, total=False): +class DeleteRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for deleting an Agent Runtime Sandbox.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict = Union[ - GetDeleteAgentEngineRuntimeRevisionOperationConfig, - GetDeleteAgentEngineRuntimeRevisionOperationConfigDict, +DeleteRuntimeSandboxConfigOrDict = Union[ + DeleteRuntimeSandboxConfig, DeleteRuntimeSandboxConfigDict ] -class _GetDeleteAgentEngineRuntimeRevisionOperationParameters(_common.BaseModel): - """Parameters for getting an operation that deletes a agent engine runtime revision.""" +class _DeleteRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtimes.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime sandbox to delete.""" ) - config: Optional[GetDeleteAgentEngineRuntimeRevisionOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + config: Optional[DeleteRuntimeSandboxConfig] = Field( + default=None, description="""""" ) -class _GetDeleteAgentEngineRuntimeRevisionOperationParametersDict( - TypedDict, total=False -): - """Parameters for getting an operation that deletes a agent engine runtime revision.""" +class _DeleteRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtimes.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + name: Optional[str] + """Name of the agent runtime sandbox to delete.""" - config: Optional[GetDeleteAgentEngineRuntimeRevisionOperationConfigDict] - """Used to override the default configuration.""" + config: Optional[DeleteRuntimeSandboxConfigDict] + """""" -_GetDeleteAgentEngineRuntimeRevisionOperationParametersOrDict = Union[ - _GetDeleteAgentEngineRuntimeRevisionOperationParameters, - _GetDeleteAgentEngineRuntimeRevisionOperationParametersDict, +_DeleteRuntimeSandboxRequestParametersOrDict = Union[ + _DeleteRuntimeSandboxRequestParameters, _DeleteRuntimeSandboxRequestParametersDict ] -class QueryAgentEngineRuntimeRevisionConfig(_common.BaseModel): - """Config for querying agent engine runtime revisions.""" +class DeleteRuntimeSandboxOperation(_common.BaseModel): + """Operation for deleting agent runtimes.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - class_method: Optional[str] = Field( - default=None, description="""The class method to call.""" + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", ) - input: Optional[dict[str, Any]] = Field( - default=None, description="""The input to the class method.""" + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", ) - include_all_fields: Optional[bool] = Field(default=False, description="""""") -class QueryAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): - """Config for querying agent engine runtime revisions.""" +class DeleteRuntimeSandboxOperationDict(TypedDict, total=False): + """Operation for deleting agent runtimes.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - class_method: Optional[str] - """The class method to call.""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - input: Optional[dict[str, Any]] - """The input to the class method.""" + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - include_all_fields: Optional[bool] - """""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +DeleteRuntimeSandboxOperationOrDict = Union[ + DeleteRuntimeSandboxOperation, DeleteRuntimeSandboxOperationDict +] + + +class Metadata(_common.BaseModel): + """Metadata for a chunk.""" + + attributes: Optional[dict[str, bytes]] = Field( + default=None, + description="""Optional. Attributes attached to the data. The keys have semantic conventions and the consumers of the attributes should know how to deserialize the value bytes based on the keys.""", + ) + + +class MetadataDict(TypedDict, total=False): + """Metadata for a chunk.""" + + attributes: Optional[dict[str, bytes]] + """Optional. Attributes attached to the data. The keys have semantic conventions and the consumers of the attributes should know how to deserialize the value bytes based on the keys.""" + + +MetadataOrDict = Union[Metadata, MetadataDict] + + +class Chunk(_common.BaseModel): + """A chunk of data.""" + + data: Optional[bytes] = Field( + default=None, description="""Required. The data in the chunk.""" + ) + metadata: Optional[Metadata] = Field( + default=None, + description="""Optional. Metadata that is associated with the data in the payload.""", + ) + mime_type: Optional[str] = Field( + default=None, + description="""Required. Mime type of the chunk data. See https://www.iana.org/assignments/media-types/media-types.xhtml for the full list.""", + ) + + +class ChunkDict(TypedDict, total=False): + """A chunk of data.""" + + data: Optional[bytes] + """Required. The data in the chunk.""" + + metadata: Optional[MetadataDict] + """Optional. Metadata that is associated with the data in the payload.""" + + mime_type: Optional[str] + """Required. Mime type of the chunk data. See https://www.iana.org/assignments/media-types/media-types.xhtml for the full list.""" + + +ChunkOrDict = Union[Chunk, ChunkDict] + + +class ExecuteCodeRuntimeSandboxConfig(_common.BaseModel): + """Config for executing code in an Agent Runtime sandbox.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class ExecuteCodeRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for executing code in an Agent Runtime sandbox.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -QueryAgentEngineRuntimeRevisionConfigOrDict = Union[ - QueryAgentEngineRuntimeRevisionConfig, QueryAgentEngineRuntimeRevisionConfigDict + +ExecuteCodeRuntimeSandboxConfigOrDict = Union[ + ExecuteCodeRuntimeSandboxConfig, ExecuteCodeRuntimeSandboxConfigDict ] -class _QueryAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): - """Parameters for querying agent engine runtime revisions.""" +class _ExecuteCodeRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for executing code in an agent runtime sandbox.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine runtime revision.""" + default=None, + description="""Name of the agent runtime sandbox to execute code in.""", ) - config: Optional[QueryAgentEngineRuntimeRevisionConfig] = Field( + inputs: Optional[list[Chunk]] = Field( + default=None, description="""Inputs to the code execution.""" + ) + config: Optional[ExecuteCodeRuntimeSandboxConfig] = Field( default=None, description="""""" ) -class _QueryAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for querying agent engine runtime revisions.""" +class _ExecuteCodeRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for executing code in an agent runtime sandbox.""" name: Optional[str] - """Name of the agent engine runtime revision.""" + """Name of the agent runtime sandbox to execute code in.""" + + inputs: Optional[list[ChunkDict]] + """Inputs to the code execution.""" - config: Optional[QueryAgentEngineRuntimeRevisionConfigDict] + config: Optional[ExecuteCodeRuntimeSandboxConfigDict] """""" -_QueryAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ - _QueryAgentEngineRuntimeRevisionRequestParameters, - _QueryAgentEngineRuntimeRevisionRequestParametersDict, +_ExecuteCodeRuntimeSandboxRequestParametersOrDict = Union[ + _ExecuteCodeRuntimeSandboxRequestParameters, + _ExecuteCodeRuntimeSandboxRequestParametersDict, ] -class SandboxEnvironmentSpecCodeExecutionEnvironment(_common.BaseModel): - """The code execution environment with customized settings.""" +class ExecuteSandboxEnvironmentResponse(_common.BaseModel): + """The response for executing a sandbox environment.""" - code_language: Optional[Language] = Field( - default=None, - description="""The coding language supported in this environment.""", - ) - machine_config: Optional[MachineConfig] = Field( - default=None, - description="""The machine config of the code execution environment.""", + outputs: Optional[list[Chunk]] = Field( + default=None, description="""The outputs from the sandbox environment.""" ) -class SandboxEnvironmentSpecCodeExecutionEnvironmentDict(TypedDict, total=False): - """The code execution environment with customized settings.""" - - code_language: Optional[Language] - """The coding language supported in this environment.""" +class ExecuteSandboxEnvironmentResponseDict(TypedDict, total=False): + """The response for executing a sandbox environment.""" - machine_config: Optional[MachineConfig] - """The machine config of the code execution environment.""" + outputs: Optional[list[ChunkDict]] + """The outputs from the sandbox environment.""" -SandboxEnvironmentSpecCodeExecutionEnvironmentOrDict = Union[ - SandboxEnvironmentSpecCodeExecutionEnvironment, - SandboxEnvironmentSpecCodeExecutionEnvironmentDict, +ExecuteSandboxEnvironmentResponseOrDict = Union[ + ExecuteSandboxEnvironmentResponse, ExecuteSandboxEnvironmentResponseDict ] -class SandboxEnvironmentSpecComputerUseEnvironment(_common.BaseModel): - """The computer use environment with customized settings.""" +class GetRuntimeSandboxConfig(_common.BaseModel): + """Config for getting an Agent Runtime Memory.""" - pass + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -class SandboxEnvironmentSpecComputerUseEnvironmentDict(TypedDict, total=False): - """The computer use environment with customized settings.""" +class GetRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for getting an Agent Runtime Memory.""" - pass + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -SandboxEnvironmentSpecComputerUseEnvironmentOrDict = Union[ - SandboxEnvironmentSpecComputerUseEnvironment, - SandboxEnvironmentSpecComputerUseEnvironmentDict, +GetRuntimeSandboxConfigOrDict = Union[ + GetRuntimeSandboxConfig, GetRuntimeSandboxConfigDict ] -class SandboxEnvironmentSpec(_common.BaseModel): - """The specification of a sandbox environment.""" +class _GetRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for getting an agent runtime sandbox.""" - code_execution_environment: Optional[ - SandboxEnvironmentSpecCodeExecutionEnvironment - ] = Field(default=None, description="""Optional. The code execution environment.""") - computer_use_environment: Optional[SandboxEnvironmentSpecComputerUseEnvironment] = ( - Field(default=None, description="""Optional. The computer use environment.""") + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime sandbox.""" ) + config: Optional[GetRuntimeSandboxConfig] = Field(default=None, description="""""") -class SandboxEnvironmentSpecDict(TypedDict, total=False): - """The specification of a sandbox environment.""" +class _GetRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for getting an agent runtime sandbox.""" - code_execution_environment: Optional[ - SandboxEnvironmentSpecCodeExecutionEnvironmentDict - ] - """Optional. The code execution environment.""" + name: Optional[str] + """Name of the agent runtime sandbox.""" - computer_use_environment: Optional[SandboxEnvironmentSpecComputerUseEnvironmentDict] - """Optional. The computer use environment.""" + config: Optional[GetRuntimeSandboxConfigDict] + """""" -SandboxEnvironmentSpecOrDict = Union[SandboxEnvironmentSpec, SandboxEnvironmentSpecDict] +_GetRuntimeSandboxRequestParametersOrDict = Union[ + _GetRuntimeSandboxRequestParameters, _GetRuntimeSandboxRequestParametersDict +] -class CreateAgentEngineSandboxConfig(_common.BaseModel): - """Config for creating a Sandbox.""" +class ListRuntimeSandboxesConfig(_common.BaseModel): + """Config for listing agent runtime sandboxes.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - display_name: Optional[str] = Field( - default=None, description="""The display name of the sandbox.""" - ) - description: Optional[str] = Field( - default=None, description="""The description of the sandbox.""" - ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", - ) - ttl: Optional[str] = Field( - default=None, - description="""The TTL for this resource. The expiration time is computed: now + TTL.""", - ) - sandbox_environment_template: Optional[str] = Field( - default=None, - description="""The name of the sandbox environment template to create the sandbox from. The sandbox environment template should be in the format: - projects/{project}/locations/{location}/agentEngines/{agent_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}""", - ) - sandbox_environment_snapshot: Optional[str] = Field( - default=None, - description="""The name of the sandbox environment snapshot to restore the sandbox from. The sandbox environment snapshot should be in the format: - projects/{project}/locations/{location}/agentEngines/{agent_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}""", - ) - owner: Optional[str] = Field( + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( default=None, - description="""Owner information for this sandbox environment. A sandbox can only be restored from a snapshot belonging to the same owner.""", + description="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""", ) -class CreateAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for creating a Sandbox.""" +class ListRuntimeSandboxesConfigDict(TypedDict, total=False): + """Config for listing agent runtime sandboxes.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - display_name: Optional[str] - """The display name of the sandbox.""" - - description: Optional[str] - """The description of the sandbox.""" - - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" - - ttl: Optional[str] - """The TTL for this resource. The expiration time is computed: now + TTL.""" - - sandbox_environment_template: Optional[str] - """The name of the sandbox environment template to create the sandbox from. The sandbox environment template should be in the format: - projects/{project}/locations/{location}/agentEngines/{agent_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}""" + page_size: Optional[int] + """""" - sandbox_environment_snapshot: Optional[str] - """The name of the sandbox environment snapshot to restore the sandbox from. The sandbox environment snapshot should be in the format: - projects/{project}/locations/{location}/agentEngines/{agent_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}""" + page_token: Optional[str] + """""" - owner: Optional[str] - """Owner information for this sandbox environment. A sandbox can only be restored from a snapshot belonging to the same owner.""" + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""" -CreateAgentEngineSandboxConfigOrDict = Union[ - CreateAgentEngineSandboxConfig, CreateAgentEngineSandboxConfigDict +ListRuntimeSandboxesConfigOrDict = Union[ + ListRuntimeSandboxesConfig, ListRuntimeSandboxesConfigDict ] -class _CreateAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for creating Agent Engine Sandboxes.""" +class _ListRuntimeSandboxesRequestParameters(_common.BaseModel): + """Parameters for listing agent runtime sandboxes.""" name: Optional[str] = Field( - default=None, - description="""Name of the agent engine to create the sandbox under.""", - ) - spec: Optional[SandboxEnvironmentSpec] = Field( - default=None, description="""The specification of the sandbox.""" + default=None, description="""Name of the agent runtime.""" ) - config: Optional[CreateAgentEngineSandboxConfig] = Field( + config: Optional[ListRuntimeSandboxesConfig] = Field( default=None, description="""""" ) -class _CreateAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for creating Agent Engine Sandboxes.""" +class _ListRuntimeSandboxesRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtime sandboxes.""" name: Optional[str] - """Name of the agent engine to create the sandbox under.""" - - spec: Optional[SandboxEnvironmentSpecDict] - """The specification of the sandbox.""" + """Name of the agent runtime.""" - config: Optional[CreateAgentEngineSandboxConfigDict] + config: Optional[ListRuntimeSandboxesConfigDict] """""" -_CreateAgentEngineSandboxRequestParametersOrDict = Union[ - _CreateAgentEngineSandboxRequestParameters, - _CreateAgentEngineSandboxRequestParametersDict, +_ListRuntimeSandboxesRequestParametersOrDict = Union[ + _ListRuntimeSandboxesRequestParameters, _ListRuntimeSandboxesRequestParametersDict ] -class SandboxEnvironmentConnectionInfo(_common.BaseModel): - """The connection information of the SandboxEnvironment.""" +class ListRuntimeSandboxesResponse(_common.BaseModel): + """Response for listing agent runtime sandboxes.""" - load_balancer_hostname: Optional[str] = Field( - default=None, description="""Output only. The hostname of the load balancer.""" - ) - load_balancer_ip: Optional[str] = Field( - default=None, - description="""Output only. The IP address of the load balancer.""", - ) - sandbox_internal_ip: Optional[str] = Field( - default=None, - description="""Output only. The internal IP address of the SandboxEnvironment.""", + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - routing_token: Optional[str] = Field( - default=None, - description="""Output only. The routing token for the SandboxEnvironment.""", + next_page_token: Optional[str] = Field(default=None, description="""""") + sandbox_environments: Optional[list[SandboxEnvironment]] = Field( + default=None, description="""List of agent runtime sandboxes.""" ) -class SandboxEnvironmentConnectionInfoDict(TypedDict, total=False): - """The connection information of the SandboxEnvironment.""" - - load_balancer_hostname: Optional[str] - """Output only. The hostname of the load balancer.""" +class ListRuntimeSandboxesResponseDict(TypedDict, total=False): + """Response for listing agent runtime sandboxes.""" - load_balancer_ip: Optional[str] - """Output only. The IP address of the load balancer.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - sandbox_internal_ip: Optional[str] - """Output only. The internal IP address of the SandboxEnvironment.""" + next_page_token: Optional[str] + """""" - routing_token: Optional[str] - """Output only. The routing token for the SandboxEnvironment.""" + sandbox_environments: Optional[list[SandboxEnvironmentDict]] + """List of agent runtime sandboxes.""" -SandboxEnvironmentConnectionInfoOrDict = Union[ - SandboxEnvironmentConnectionInfo, SandboxEnvironmentConnectionInfoDict +ListRuntimeSandboxesResponseOrDict = Union[ + ListRuntimeSandboxesResponse, ListRuntimeSandboxesResponseDict ] -class SandboxEnvironment(_common.BaseModel): - """A sandbox environment.""" +class _GetRuntimeSandboxOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a sandbox as a response.""" - expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Expiration time of the sandbox environment. - """, + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - connection_info: Optional[SandboxEnvironmentConnectionInfo] = Field( - default=None, - description="""Output only. The connection information of the SandboxEnvironment.""", - ) - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. The timestamp when this SandboxEnvironment was created.""", - ) - display_name: Optional[str] = Field( - default=None, - description="""Required. The display name of the SandboxEnvironment.""", - ) - name: Optional[str] = Field( - default=None, description="""Identifier. The name of the SandboxEnvironment.""" - ) - spec: Optional[SandboxEnvironmentSpec] = Field( - default=None, - description="""Optional. The configuration of the SandboxEnvironment.""", - ) - state: Optional[SandboxState] = Field( - default=None, - description="""Output only. The runtime state of the SandboxEnvironment.""", - ) - ttl: Optional[str] = Field( - default=None, - description="""Optional. Input only. The TTL for the sandbox environment. The expiration time is computed: now + TTL.""", - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. The timestamp when this SandboxEnvironment was most recently updated.""", - ) - latest_sandbox_environment_snapshot: Optional[str] = Field( - default=None, - description="""Output only. The resource name of the latest snapshot taken for this SandboxEnvironment.""", - ) - owner: Optional[str] = Field( - default=None, - description="""Optional. Owner information for this sandbox environment. A Sandbox can only be restored from a snapshot that belongs to the same owner. If not set, sandbox will be created as the default owner.""", - ) - sandbox_environment_snapshot: Optional[str] = Field( - default=None, - description="""Optional. The resource name of the SandboxEnvironmentSnapshot to use for creating this SandboxEnvironment. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}`""", - ) - sandbox_environment_template: Optional[str] = Field( - default=None, - description="""Optional. The name of the SandboxEnvironmentTemplate specified in the parent Agent Engine resource that this SandboxEnvironment is created from.""", + config: Optional[GetRuntimeOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -class SandboxEnvironmentDict(TypedDict, total=False): - """A sandbox environment.""" - - expire_time: Optional[datetime.datetime] - """Expiration time of the sandbox environment. - """ - - connection_info: Optional[SandboxEnvironmentConnectionInfoDict] - """Output only. The connection information of the SandboxEnvironment.""" - - create_time: Optional[datetime.datetime] - """Output only. The timestamp when this SandboxEnvironment was created.""" +class _GetRuntimeSandboxOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a sandbox as a response.""" - display_name: Optional[str] - """Required. The display name of the SandboxEnvironment.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - name: Optional[str] - """Identifier. The name of the SandboxEnvironment.""" + config: Optional[GetRuntimeOperationConfigDict] + """Used to override the default configuration.""" - spec: Optional[SandboxEnvironmentSpecDict] - """Optional. The configuration of the SandboxEnvironment.""" - state: Optional[SandboxState] - """Output only. The runtime state of the SandboxEnvironment.""" +_GetRuntimeSandboxOperationParametersOrDict = Union[ + _GetRuntimeSandboxOperationParameters, _GetRuntimeSandboxOperationParametersDict +] - ttl: Optional[str] - """Optional. Input only. The TTL for the sandbox environment. The expiration time is computed: now + TTL.""" - update_time: Optional[datetime.datetime] - """Output only. The timestamp when this SandboxEnvironment was most recently updated.""" +class SandboxEnvironmentTemplateCustomContainerSpec(_common.BaseModel): + """Specification for deploying from a custom container image.""" - latest_sandbox_environment_snapshot: Optional[str] - """Output only. The resource name of the latest snapshot taken for this SandboxEnvironment.""" + image_uri: Optional[str] = Field( + default=None, + description="""Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""", + ) - owner: Optional[str] - """Optional. Owner information for this sandbox environment. A Sandbox can only be restored from a snapshot that belongs to the same owner. If not set, sandbox will be created as the default owner.""" - sandbox_environment_snapshot: Optional[str] - """Optional. The resource name of the SandboxEnvironmentSnapshot to use for creating this SandboxEnvironment. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}`""" +class SandboxEnvironmentTemplateCustomContainerSpecDict(TypedDict, total=False): + """Specification for deploying from a custom container image.""" - sandbox_environment_template: Optional[str] - """Optional. The name of the SandboxEnvironmentTemplate specified in the parent Agent Engine resource that this SandboxEnvironment is created from.""" + image_uri: Optional[str] + """Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""" -SandboxEnvironmentOrDict = Union[SandboxEnvironment, SandboxEnvironmentDict] +SandboxEnvironmentTemplateCustomContainerSpecOrDict = Union[ + SandboxEnvironmentTemplateCustomContainerSpec, + SandboxEnvironmentTemplateCustomContainerSpecDict, +] -class AgentEngineSandboxOperation(_common.BaseModel): - """Operation that has an agent engine sandbox as a response.""" +class SandboxEnvironmentTemplateNetworkPort(_common.BaseModel): + """Represents a network port in a container.""" - name: Optional[str] = Field( - default=None, - description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, - description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", - ) - done: Optional[bool] = Field( + port: Optional[int] = Field( default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + description="""Optional. Port number to expose. This must be a valid port number, between 1 and 65535.""", ) - error: Optional[dict[str, Any]] = Field( + protocol: Optional[Protocol] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", - ) - response: Optional[SandboxEnvironment] = Field( - default=None, description="""The Agent Engine Sandbox.""" + description="""Optional. Protocol for port. Defaults to TCP if not specified.""", ) -class AgentEngineSandboxOperationDict(TypedDict, total=False): - """Operation that has an agent engine sandbox as a response.""" - - name: Optional[str] - """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - - metadata: Optional[dict[str, Any]] - """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - - done: Optional[bool] - """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" +class SandboxEnvironmentTemplateNetworkPortDict(TypedDict, total=False): + """Represents a network port in a container.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + port: Optional[int] + """Optional. Port number to expose. This must be a valid port number, between 1 and 65535.""" - response: Optional[SandboxEnvironmentDict] - """The Agent Engine Sandbox.""" + protocol: Optional[Protocol] + """Optional. Protocol for port. Defaults to TCP if not specified.""" -AgentEngineSandboxOperationOrDict = Union[ - AgentEngineSandboxOperation, AgentEngineSandboxOperationDict +SandboxEnvironmentTemplateNetworkPortOrDict = Union[ + SandboxEnvironmentTemplateNetworkPort, SandboxEnvironmentTemplateNetworkPortDict ] -class DeleteAgentEngineSandboxConfig(_common.BaseModel): - """Config for deleting an Agent Engine Sandbox.""" +class SandboxEnvironmentTemplateResourceRequirements(_common.BaseModel): + """Message to define resource requests and limits (mirroring Kubernetes) for each sandbox instance created from this template.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + limits: Optional[dict[str, str]] = Field( + default=None, + description="""Optional. The maximum amounts of compute resources allowed. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "500m", "1Gi").""", + ) + requests: Optional[dict[str, str]] = Field( + default=None, + description="""Optional. The requested amounts of compute resources. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "250m", "512Mi").""", ) -class DeleteAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Sandbox.""" +class SandboxEnvironmentTemplateResourceRequirementsDict(TypedDict, total=False): + """Message to define resource requests and limits (mirroring Kubernetes) for each sandbox instance created from this template.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + limits: Optional[dict[str, str]] + """Optional. The maximum amounts of compute resources allowed. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "500m", "1Gi").""" + + requests: Optional[dict[str, str]] + """Optional. The requested amounts of compute resources. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "250m", "512Mi").""" -DeleteAgentEngineSandboxConfigOrDict = Union[ - DeleteAgentEngineSandboxConfig, DeleteAgentEngineSandboxConfigDict +SandboxEnvironmentTemplateResourceRequirementsOrDict = Union[ + SandboxEnvironmentTemplateResourceRequirements, + SandboxEnvironmentTemplateResourceRequirementsDict, ] -class _DeleteAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for deleting agent engines.""" +class SandboxEnvironmentTemplateCustomContainerEnvironment(_common.BaseModel): + """The customized sandbox runtime environment for BYOC.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine sandbox to delete.""" + custom_container_spec: Optional[SandboxEnvironmentTemplateCustomContainerSpec] = ( + Field( + default=None, + description="""The specification of the custom container environment.""", + ) ) - config: Optional[DeleteAgentEngineSandboxConfig] = Field( - default=None, description="""""" + ports: Optional[list[SandboxEnvironmentTemplateNetworkPort]] = Field( + default=None, description="""Ports to expose from the container.""" + ) + resources: Optional[SandboxEnvironmentTemplateResourceRequirements] = Field( + default=None, description="""Resource requests and limits for the container.""" ) -class _DeleteAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engines.""" +class SandboxEnvironmentTemplateCustomContainerEnvironmentDict(TypedDict, total=False): + """The customized sandbox runtime environment for BYOC.""" + + custom_container_spec: Optional[SandboxEnvironmentTemplateCustomContainerSpecDict] + """The specification of the custom container environment.""" - name: Optional[str] - """Name of the agent engine sandbox to delete.""" + ports: Optional[list[SandboxEnvironmentTemplateNetworkPortDict]] + """Ports to expose from the container.""" - config: Optional[DeleteAgentEngineSandboxConfigDict] - """""" + resources: Optional[SandboxEnvironmentTemplateResourceRequirementsDict] + """Resource requests and limits for the container.""" -_DeleteAgentEngineSandboxRequestParametersOrDict = Union[ - _DeleteAgentEngineSandboxRequestParameters, - _DeleteAgentEngineSandboxRequestParametersDict, +SandboxEnvironmentTemplateCustomContainerEnvironmentOrDict = Union[ + SandboxEnvironmentTemplateCustomContainerEnvironment, + SandboxEnvironmentTemplateCustomContainerEnvironmentDict, ] -class DeleteAgentEngineSandboxOperation(_common.BaseModel): - """Operation for deleting agent engines.""" +class SandboxEnvironmentTemplateDefaultContainerEnvironment(_common.BaseModel): + """The default sandbox runtime environment for default container workloads.""" - name: Optional[str] = Field( - default=None, - description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, - description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", - ) - done: Optional[bool] = Field( + default_container_category: Optional[DefaultContainerCategory] = Field( default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + description="""Required. The category of the default container image.""", ) - error: Optional[dict[str, Any]] = Field( + resources: Optional[SandboxEnvironmentTemplateResourceRequirements] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Optional. Resource requests and limits for the default container.""", ) -class DeleteAgentEngineSandboxOperationDict(TypedDict, total=False): - """Operation for deleting agent engines.""" - - name: Optional[str] - """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - - metadata: Optional[dict[str, Any]] - """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" +class SandboxEnvironmentTemplateDefaultContainerEnvironmentDict(TypedDict, total=False): + """The default sandbox runtime environment for default container workloads.""" - done: Optional[bool] - """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + default_container_category: Optional[DefaultContainerCategory] + """Required. The category of the default container image.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + resources: Optional[SandboxEnvironmentTemplateResourceRequirementsDict] + """Optional. Resource requests and limits for the default container.""" -DeleteAgentEngineSandboxOperationOrDict = Union[ - DeleteAgentEngineSandboxOperation, DeleteAgentEngineSandboxOperationDict +SandboxEnvironmentTemplateDefaultContainerEnvironmentOrDict = Union[ + SandboxEnvironmentTemplateDefaultContainerEnvironment, + SandboxEnvironmentTemplateDefaultContainerEnvironmentDict, ] -class Metadata(_common.BaseModel): - """Metadata for a chunk.""" +class SandboxEnvironmentTemplateEgressControlConfig(_common.BaseModel): + """Configuration for egress control of sandbox instances.""" - attributes: Optional[dict[str, bytes]] = Field( - default=None, - description="""Optional. Attributes attached to the data. The keys have semantic conventions and the consumers of the attributes should know how to deserialize the value bytes based on the keys.""", + internet_access: Optional[bool] = Field( + default=None, description="""Optional. Whether to allow internet access.""" ) -class MetadataDict(TypedDict, total=False): - """Metadata for a chunk.""" +class SandboxEnvironmentTemplateEgressControlConfigDict(TypedDict, total=False): + """Configuration for egress control of sandbox instances.""" - attributes: Optional[dict[str, bytes]] - """Optional. Attributes attached to the data. The keys have semantic conventions and the consumers of the attributes should know how to deserialize the value bytes based on the keys.""" + internet_access: Optional[bool] + """Optional. Whether to allow internet access.""" -MetadataOrDict = Union[Metadata, MetadataDict] +SandboxEnvironmentTemplateEgressControlConfigOrDict = Union[ + SandboxEnvironmentTemplateEgressControlConfig, + SandboxEnvironmentTemplateEgressControlConfigDict, +] -class Chunk(_common.BaseModel): - """A chunk of data.""" +class CreateSandboxEnvironmentTemplateConfig(_common.BaseModel): + """Config for creating a Sandbox Template.""" - data: Optional[bytes] = Field( - default=None, description="""Required. The data in the chunk.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - metadata: Optional[Metadata] = Field( + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", + ) + custom_container_environment: Optional[ + SandboxEnvironmentTemplateCustomContainerEnvironment + ] = Field( default=None, - description="""Optional. Metadata that is associated with the data in the payload.""", + description="""The custom container environment for the sandbox template.""", ) - mime_type: Optional[str] = Field( + default_container_environment: Optional[ + SandboxEnvironmentTemplateDefaultContainerEnvironment + ] = Field( default=None, - description="""Required. Mime type of the chunk data. See https://www.iana.org/assignments/media-types/media-types.xhtml for the full list.""", + description="""The default container environment for the sandbox template.""", + ) + egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = ( + Field( + default=None, + description="""The egress control config for the sandbox template.""", + ) ) -class ChunkDict(TypedDict, total=False): - """A chunk of data.""" - - data: Optional[bytes] - """Required. The data in the chunk.""" - - metadata: Optional[MetadataDict] - """Optional. Metadata that is associated with the data in the payload.""" - - mime_type: Optional[str] - """Required. Mime type of the chunk data. See https://www.iana.org/assignments/media-types/media-types.xhtml for the full list.""" - - -ChunkOrDict = Union[Chunk, ChunkDict] - +class CreateSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): + """Config for creating a Sandbox Template.""" -class ExecuteCodeAgentEngineSandboxConfig(_common.BaseModel): - """Config for executing code in an Agent Engine sandbox.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" + custom_container_environment: Optional[ + SandboxEnvironmentTemplateCustomContainerEnvironmentDict + ] + """The custom container environment for the sandbox template.""" -class ExecuteCodeAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for executing code in an Agent Engine sandbox.""" + default_container_environment: Optional[ + SandboxEnvironmentTemplateDefaultContainerEnvironmentDict + ] + """The default container environment for the sandbox template.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfigDict] + """The egress control config for the sandbox template.""" -ExecuteCodeAgentEngineSandboxConfigOrDict = Union[ - ExecuteCodeAgentEngineSandboxConfig, ExecuteCodeAgentEngineSandboxConfigDict +CreateSandboxEnvironmentTemplateConfigOrDict = Union[ + CreateSandboxEnvironmentTemplateConfig, CreateSandboxEnvironmentTemplateConfigDict ] -class _ExecuteCodeAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for executing code in an agent engine sandbox.""" +class _CreateSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): + """Parameters for creating Sandbox Environment Templates.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine sandbox to execute code in.""", - ) - inputs: Optional[list[Chunk]] = Field( - default=None, description="""Inputs to the code execution.""" + description="""Name of the agent runtime to create the template under.""", ) - config: Optional[ExecuteCodeAgentEngineSandboxConfig] = Field( + config: Optional[CreateSandboxEnvironmentTemplateConfig] = Field( default=None, description="""""" ) + display_name: Optional[str] = Field( + default=None, description="""The display name of the sandbox template.""" + ) -class _ExecuteCodeAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for executing code in an agent engine sandbox.""" +class _CreateSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): + """Parameters for creating Sandbox Environment Templates.""" name: Optional[str] - """Name of the agent engine sandbox to execute code in.""" - - inputs: Optional[list[ChunkDict]] - """Inputs to the code execution.""" + """Name of the agent runtime to create the template under.""" - config: Optional[ExecuteCodeAgentEngineSandboxConfigDict] + config: Optional[CreateSandboxEnvironmentTemplateConfigDict] """""" + display_name: Optional[str] + """The display name of the sandbox template.""" + -_ExecuteCodeAgentEngineSandboxRequestParametersOrDict = Union[ - _ExecuteCodeAgentEngineSandboxRequestParameters, - _ExecuteCodeAgentEngineSandboxRequestParametersDict, +_CreateSandboxEnvironmentTemplateRequestParametersOrDict = Union[ + _CreateSandboxEnvironmentTemplateRequestParameters, + _CreateSandboxEnvironmentTemplateRequestParametersDict, ] -class ExecuteSandboxEnvironmentResponse(_common.BaseModel): - """The response for executing a sandbox environment.""" +class SandboxEnvironmentTemplate(_common.BaseModel): + """A sandbox environment template.""" - outputs: Optional[list[Chunk]] = Field( - default=None, description="""The outputs from the sandbox environment.""" + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. The timestamp when this SandboxEnvironmentTemplate was created.""", + ) + custom_container_environment: Optional[ + SandboxEnvironmentTemplateCustomContainerEnvironment + ] = Field( + default=None, + description="""The sandbox environment for custom container workloads.""", + ) + default_container_environment: Optional[ + SandboxEnvironmentTemplateDefaultContainerEnvironment + ] = Field( + default=None, + description="""The sandbox environment for default container workloads.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""Required. The display name of the SandboxEnvironmentTemplate.""", + ) + egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = ( + Field( + default=None, + description="""Optional. The configuration for egress control of this template.""", + ) + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. The resource name of the SandboxEnvironmentTemplate. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}`""", + ) + state: Optional[ + Literal[ + "UNSPECIFIED", + "PROVISIONING", + "ACTIVE", + "DEPROVISIONING", + "DELETED", + "FAILED", + ] + ] = Field( + default=None, + description="""Output only. The state of the sandbox environment template.""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. The timestamp when this SandboxEnvironmentTemplate was most recently updated.""", ) -class ExecuteSandboxEnvironmentResponseDict(TypedDict, total=False): - """The response for executing a sandbox environment.""" - - outputs: Optional[list[ChunkDict]] - """The outputs from the sandbox environment.""" +class SandboxEnvironmentTemplateDict(TypedDict, total=False): + """A sandbox environment template.""" + create_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironmentTemplate was created.""" -ExecuteSandboxEnvironmentResponseOrDict = Union[ - ExecuteSandboxEnvironmentResponse, ExecuteSandboxEnvironmentResponseDict -] + custom_container_environment: Optional[ + SandboxEnvironmentTemplateCustomContainerEnvironmentDict + ] + """The sandbox environment for custom container workloads.""" + default_container_environment: Optional[ + SandboxEnvironmentTemplateDefaultContainerEnvironmentDict + ] + """The sandbox environment for default container workloads.""" -class GetAgentEngineSandboxConfig(_common.BaseModel): - """Config for getting an Agent Engine Memory.""" + display_name: Optional[str] + """Required. The display name of the SandboxEnvironmentTemplate.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfigDict] + """Optional. The configuration for egress control of this template.""" + name: Optional[str] + """Identifier. The resource name of the SandboxEnvironmentTemplate. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}`""" -class GetAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Memory.""" + state: Optional[ + Literal[ + "UNSPECIFIED", + "PROVISIONING", + "ACTIVE", + "DEPROVISIONING", + "DELETED", + "FAILED", + ] + ] + """Output only. The state of the sandbox environment template.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + update_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironmentTemplate was most recently updated.""" -GetAgentEngineSandboxConfigOrDict = Union[ - GetAgentEngineSandboxConfig, GetAgentEngineSandboxConfigDict +SandboxEnvironmentTemplateOrDict = Union[ + SandboxEnvironmentTemplate, SandboxEnvironmentTemplateDict ] -class _GetAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine sandbox.""" +class SandboxEnvironmentTemplateOperation(_common.BaseModel): + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine sandbox.""" + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - config: Optional[GetAgentEngineSandboxConfig] = Field( - default=None, description="""""" + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + response: Optional[SandboxEnvironmentTemplate] = Field( + default=None, description="""The Agent Runtime Sandbox Template.""" ) -class _GetAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine sandbox.""" +class SandboxEnvironmentTemplateOperationDict(TypedDict, total=False): + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] - """Name of the agent engine sandbox.""" + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - config: Optional[GetAgentEngineSandboxConfigDict] - """""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + response: Optional[SandboxEnvironmentTemplateDict] + """The Agent Runtime Sandbox Template.""" -_GetAgentEngineSandboxRequestParametersOrDict = Union[ - _GetAgentEngineSandboxRequestParameters, _GetAgentEngineSandboxRequestParametersDict +SandboxEnvironmentTemplateOperationOrDict = Union[ + SandboxEnvironmentTemplateOperation, SandboxEnvironmentTemplateOperationDict ] -class ListAgentEngineSandboxesConfig(_common.BaseModel): - """Config for listing agent engine sandboxes.""" +class DeleteSandboxEnvironmentTemplateConfig(_common.BaseModel): + """Config for deleting a Sandbox Template.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( - default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""", - ) -class ListAgentEngineSandboxesConfigDict(TypedDict, total=False): - """Config for listing agent engine sandboxes.""" +class DeleteSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): + """Config for deleting a Sandbox Template.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - page_size: Optional[int] - """""" - - page_token: Optional[str] - """""" - - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" - -ListAgentEngineSandboxesConfigOrDict = Union[ - ListAgentEngineSandboxesConfig, ListAgentEngineSandboxesConfigDict +DeleteSandboxEnvironmentTemplateConfigOrDict = Union[ + DeleteSandboxEnvironmentTemplateConfig, DeleteSandboxEnvironmentTemplateConfigDict ] -class _ListAgentEngineSandboxesRequestParameters(_common.BaseModel): - """Parameters for listing agent engine sandboxes.""" +class _DeleteSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): + """Parameters for deleting sandbox templates.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the sandbox template to delete.""" ) - config: Optional[ListAgentEngineSandboxesConfig] = Field( + config: Optional[DeleteSandboxEnvironmentTemplateConfig] = Field( default=None, description="""""" ) -class _ListAgentEngineSandboxesRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engine sandboxes.""" +class _DeleteSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): + """Parameters for deleting sandbox templates.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the sandbox template to delete.""" - config: Optional[ListAgentEngineSandboxesConfigDict] + config: Optional[DeleteSandboxEnvironmentTemplateConfigDict] """""" -_ListAgentEngineSandboxesRequestParametersOrDict = Union[ - _ListAgentEngineSandboxesRequestParameters, - _ListAgentEngineSandboxesRequestParametersDict, +_DeleteSandboxEnvironmentTemplateRequestParametersOrDict = Union[ + _DeleteSandboxEnvironmentTemplateRequestParameters, + _DeleteSandboxEnvironmentTemplateRequestParametersDict, ] -class ListAgentEngineSandboxesResponse(_common.BaseModel): - """Response for listing agent engine sandboxes.""" +class DeleteSandboxEnvironmentTemplateOperation(_common.BaseModel): + """Operation for deleting sandbox templates.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - sandbox_environments: Optional[list[SandboxEnvironment]] = Field( - default=None, description="""List of agent engine sandboxes.""" + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", ) -class ListAgentEngineSandboxesResponseDict(TypedDict, total=False): - """Response for listing agent engine sandboxes.""" +class DeleteSandboxEnvironmentTemplateOperationDict(TypedDict, total=False): + """Operation for deleting sandbox templates.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - next_page_token: Optional[str] - """""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - sandbox_environments: Optional[list[SandboxEnvironmentDict]] - """List of agent engine sandboxes.""" + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -ListAgentEngineSandboxesResponseOrDict = Union[ - ListAgentEngineSandboxesResponse, ListAgentEngineSandboxesResponseDict +DeleteSandboxEnvironmentTemplateOperationOrDict = Union[ + DeleteSandboxEnvironmentTemplateOperation, + DeleteSandboxEnvironmentTemplateOperationDict, ] -class _GetAgentEngineSandboxOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a sandbox as a response.""" +class GetSandboxEnvironmentTemplateConfig(_common.BaseModel): + """Config for getting a Sandbox Template.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" - ) - config: Optional[GetAgentEngineOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class _GetAgentEngineSandboxOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a sandbox as a response.""" - - operation_name: Optional[str] - """The server-assigned name for the operation.""" +class GetSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): + """Config for getting a Sandbox Template.""" - config: Optional[GetAgentEngineOperationConfigDict] - """Used to override the default configuration.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -_GetAgentEngineSandboxOperationParametersOrDict = Union[ - _GetAgentEngineSandboxOperationParameters, - _GetAgentEngineSandboxOperationParametersDict, +GetSandboxEnvironmentTemplateConfigOrDict = Union[ + GetSandboxEnvironmentTemplateConfig, GetSandboxEnvironmentTemplateConfigDict ] -class SandboxEnvironmentTemplateCustomContainerSpec(_common.BaseModel): - """Specification for deploying from a custom container image.""" +class _GetSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): + """Parameters for getting a sandbox template.""" - image_uri: Optional[str] = Field( - default=None, - description="""Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""", + name: Optional[str] = Field( + default=None, description="""Name of the sandbox template.""" + ) + config: Optional[GetSandboxEnvironmentTemplateConfig] = Field( + default=None, description="""""" ) -class SandboxEnvironmentTemplateCustomContainerSpecDict(TypedDict, total=False): - """Specification for deploying from a custom container image.""" +class _GetSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): + """Parameters for getting a sandbox template.""" - image_uri: Optional[str] - """Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""" + name: Optional[str] + """Name of the sandbox template.""" + config: Optional[GetSandboxEnvironmentTemplateConfigDict] + """""" -SandboxEnvironmentTemplateCustomContainerSpecOrDict = Union[ - SandboxEnvironmentTemplateCustomContainerSpec, - SandboxEnvironmentTemplateCustomContainerSpecDict, + +_GetSandboxEnvironmentTemplateRequestParametersOrDict = Union[ + _GetSandboxEnvironmentTemplateRequestParameters, + _GetSandboxEnvironmentTemplateRequestParametersDict, ] -class SandboxEnvironmentTemplateNetworkPort(_common.BaseModel): - """Represents a network port in a container.""" +class ListSandboxEnvironmentTemplatesConfig(_common.BaseModel): + """Config for listing sandbox templates.""" - port: Optional[int] = Field( - default=None, - description="""Optional. Port number to expose. This must be a valid port number, between 1 and 65535.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - protocol: Optional[Protocol] = Field( + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( default=None, - description="""Optional. Protocol for port. Defaults to TCP if not specified.""", + description="""An expression for filtering the results of the request.""", ) -class SandboxEnvironmentTemplateNetworkPortDict(TypedDict, total=False): - """Represents a network port in a container.""" +class ListSandboxEnvironmentTemplatesConfigDict(TypedDict, total=False): + """Config for listing sandbox templates.""" - port: Optional[int] - """Optional. Port number to expose. This must be a valid port number, between 1 and 65535.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - protocol: Optional[Protocol] - """Optional. Protocol for port. Defaults to TCP if not specified.""" + page_size: Optional[int] + """""" + + page_token: Optional[str] + """""" + + filter: Optional[str] + """An expression for filtering the results of the request.""" -SandboxEnvironmentTemplateNetworkPortOrDict = Union[ - SandboxEnvironmentTemplateNetworkPort, SandboxEnvironmentTemplateNetworkPortDict +ListSandboxEnvironmentTemplatesConfigOrDict = Union[ + ListSandboxEnvironmentTemplatesConfig, ListSandboxEnvironmentTemplatesConfigDict ] -class SandboxEnvironmentTemplateResourceRequirements(_common.BaseModel): - """Message to define resource requests and limits (mirroring Kubernetes) for each sandbox instance created from this template.""" +class _ListSandboxEnvironmentTemplatesRequestParameters(_common.BaseModel): + """Parameters for listing sandbox templates.""" - limits: Optional[dict[str, str]] = Field( - default=None, - description="""Optional. The maximum amounts of compute resources allowed. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "500m", "1Gi").""", + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime.""" ) - requests: Optional[dict[str, str]] = Field( - default=None, - description="""Optional. The requested amounts of compute resources. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "250m", "512Mi").""", + config: Optional[ListSandboxEnvironmentTemplatesConfig] = Field( + default=None, description="""""" ) -class SandboxEnvironmentTemplateResourceRequirementsDict(TypedDict, total=False): - """Message to define resource requests and limits (mirroring Kubernetes) for each sandbox instance created from this template.""" +class _ListSandboxEnvironmentTemplatesRequestParametersDict(TypedDict, total=False): + """Parameters for listing sandbox templates.""" - limits: Optional[dict[str, str]] - """Optional. The maximum amounts of compute resources allowed. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "500m", "1Gi").""" + name: Optional[str] + """Name of the agent runtime.""" - requests: Optional[dict[str, str]] - """Optional. The requested amounts of compute resources. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "250m", "512Mi").""" + config: Optional[ListSandboxEnvironmentTemplatesConfigDict] + """""" -SandboxEnvironmentTemplateResourceRequirementsOrDict = Union[ - SandboxEnvironmentTemplateResourceRequirements, - SandboxEnvironmentTemplateResourceRequirementsDict, +_ListSandboxEnvironmentTemplatesRequestParametersOrDict = Union[ + _ListSandboxEnvironmentTemplatesRequestParameters, + _ListSandboxEnvironmentTemplatesRequestParametersDict, ] -class SandboxEnvironmentTemplateCustomContainerEnvironment(_common.BaseModel): - """The customized sandbox runtime environment for BYOC.""" +class ListSandboxEnvironmentTemplatesResponse(_common.BaseModel): + """Response for listing sandbox templates.""" - custom_container_spec: Optional[SandboxEnvironmentTemplateCustomContainerSpec] = ( - Field( - default=None, - description="""The specification of the custom container environment.""", - ) - ) - ports: Optional[list[SandboxEnvironmentTemplateNetworkPort]] = Field( - default=None, description="""Ports to expose from the container.""" + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - resources: Optional[SandboxEnvironmentTemplateResourceRequirements] = Field( - default=None, description="""Resource requests and limits for the container.""" + next_page_token: Optional[str] = Field(default=None, description="""""") + sandbox_environment_templates: Optional[list[SandboxEnvironmentTemplate]] = Field( + default=None, description="""List of sandbox templates.""" ) -class SandboxEnvironmentTemplateCustomContainerEnvironmentDict(TypedDict, total=False): - """The customized sandbox runtime environment for BYOC.""" +class ListSandboxEnvironmentTemplatesResponseDict(TypedDict, total=False): + """Response for listing sandbox templates.""" - custom_container_spec: Optional[SandboxEnvironmentTemplateCustomContainerSpecDict] - """The specification of the custom container environment.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - ports: Optional[list[SandboxEnvironmentTemplateNetworkPortDict]] - """Ports to expose from the container.""" + next_page_token: Optional[str] + """""" - resources: Optional[SandboxEnvironmentTemplateResourceRequirementsDict] - """Resource requests and limits for the container.""" + sandbox_environment_templates: Optional[list[SandboxEnvironmentTemplateDict]] + """List of sandbox templates.""" -SandboxEnvironmentTemplateCustomContainerEnvironmentOrDict = Union[ - SandboxEnvironmentTemplateCustomContainerEnvironment, - SandboxEnvironmentTemplateCustomContainerEnvironmentDict, +ListSandboxEnvironmentTemplatesResponseOrDict = Union[ + ListSandboxEnvironmentTemplatesResponse, ListSandboxEnvironmentTemplatesResponseDict ] -class SandboxEnvironmentTemplateDefaultContainerEnvironment(_common.BaseModel): - """The default sandbox runtime environment for default container workloads.""" +class _GetSandboxEnvironmentTemplateOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a sandbox template as a response.""" - default_container_category: Optional[DefaultContainerCategory] = Field( - default=None, - description="""Required. The category of the default container image.""", + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - resources: Optional[SandboxEnvironmentTemplateResourceRequirements] = Field( - default=None, - description="""Optional. Resource requests and limits for the default container.""", + config: Optional[GetRuntimeOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -class SandboxEnvironmentTemplateDefaultContainerEnvironmentDict(TypedDict, total=False): - """The default sandbox runtime environment for default container workloads.""" +class _GetSandboxEnvironmentTemplateOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a sandbox template as a response.""" - default_container_category: Optional[DefaultContainerCategory] - """Required. The category of the default container image.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - resources: Optional[SandboxEnvironmentTemplateResourceRequirementsDict] - """Optional. Resource requests and limits for the default container.""" + config: Optional[GetRuntimeOperationConfigDict] + """Used to override the default configuration.""" -SandboxEnvironmentTemplateDefaultContainerEnvironmentOrDict = Union[ - SandboxEnvironmentTemplateDefaultContainerEnvironment, - SandboxEnvironmentTemplateDefaultContainerEnvironmentDict, +_GetSandboxEnvironmentTemplateOperationParametersOrDict = Union[ + _GetSandboxEnvironmentTemplateOperationParameters, + _GetSandboxEnvironmentTemplateOperationParametersDict, ] -class SandboxEnvironmentTemplateEgressControlConfig(_common.BaseModel): - """Configuration for egress control of sandbox instances.""" +class CreateRuntimeSandboxSnapshotConfig(_common.BaseModel): + """Config for creating a Sandbox Environment Snapshot.""" - internet_access: Optional[bool] = Field( - default=None, description="""Optional. Whether to allow internet access.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + display_name: Optional[str] = Field( + default=None, description="""The display name of the sandbox snapshot.""" + ) + owner: Optional[str] = Field( + default=None, description="""The owner of the sandbox snapshot.""" + ) + ttl: Optional[str] = Field( + default=None, + description="""The TTL for this resource. The expiration time is computed: now + TTL.""", + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) -class SandboxEnvironmentTemplateEgressControlConfigDict(TypedDict, total=False): - """Configuration for egress control of sandbox instances.""" +class CreateRuntimeSandboxSnapshotConfigDict(TypedDict, total=False): + """Config for creating a Sandbox Environment Snapshot.""" - internet_access: Optional[bool] - """Optional. Whether to allow internet access.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + display_name: Optional[str] + """The display name of the sandbox snapshot.""" -SandboxEnvironmentTemplateEgressControlConfigOrDict = Union[ - SandboxEnvironmentTemplateEgressControlConfig, - SandboxEnvironmentTemplateEgressControlConfigDict, -] + owner: Optional[str] + """The owner of the sandbox snapshot.""" + ttl: Optional[str] + """The TTL for this resource. The expiration time is computed: now + TTL.""" -class CreateSandboxEnvironmentTemplateConfig(_common.BaseModel): - """Config for creating a Sandbox Template.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", - ) - custom_container_environment: Optional[ - SandboxEnvironmentTemplateCustomContainerEnvironment - ] = Field( - default=None, - description="""The custom container environment for the sandbox template.""", - ) - default_container_environment: Optional[ - SandboxEnvironmentTemplateDefaultContainerEnvironment - ] = Field( - default=None, - description="""The default container environment for the sandbox template.""", - ) - egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = ( - Field( - default=None, - description="""The egress control config for the sandbox template.""", - ) - ) - - -class CreateSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): - """Config for creating a Sandbox Template.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" - - custom_container_environment: Optional[ - SandboxEnvironmentTemplateCustomContainerEnvironmentDict - ] - """The custom container environment for the sandbox template.""" - - default_container_environment: Optional[ - SandboxEnvironmentTemplateDefaultContainerEnvironmentDict - ] - """The default container environment for the sandbox template.""" - - egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfigDict] - """The egress control config for the sandbox template.""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" -CreateSandboxEnvironmentTemplateConfigOrDict = Union[ - CreateSandboxEnvironmentTemplateConfig, CreateSandboxEnvironmentTemplateConfigDict +CreateRuntimeSandboxSnapshotConfigOrDict = Union[ + CreateRuntimeSandboxSnapshotConfig, CreateRuntimeSandboxSnapshotConfigDict ] -class _CreateSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): - """Parameters for creating Sandbox Environment Templates.""" +class _CreateSandboxEnvironmentSnapshotRequestParameters(_common.BaseModel): + """Parameters for creating a sandbox environment snapshot.""" - name: Optional[str] = Field( - default=None, - description="""Name of the agent engine to create the template under.""", + source_sandbox_environment_name: Optional[str] = Field( + default=None, description="""Name of the sandbox environment to snapshot.""" ) - config: Optional[CreateSandboxEnvironmentTemplateConfig] = Field( + config: Optional[CreateRuntimeSandboxSnapshotConfig] = Field( default=None, description="""""" ) - display_name: Optional[str] = Field( - default=None, description="""The display name of the sandbox template.""" - ) -class _CreateSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): - """Parameters for creating Sandbox Environment Templates.""" +class _CreateSandboxEnvironmentSnapshotRequestParametersDict(TypedDict, total=False): + """Parameters for creating a sandbox environment snapshot.""" - name: Optional[str] - """Name of the agent engine to create the template under.""" + source_sandbox_environment_name: Optional[str] + """Name of the sandbox environment to snapshot.""" - config: Optional[CreateSandboxEnvironmentTemplateConfigDict] + config: Optional[CreateRuntimeSandboxSnapshotConfigDict] """""" - display_name: Optional[str] - """The display name of the sandbox template.""" - -_CreateSandboxEnvironmentTemplateRequestParametersOrDict = Union[ - _CreateSandboxEnvironmentTemplateRequestParameters, - _CreateSandboxEnvironmentTemplateRequestParametersDict, +_CreateSandboxEnvironmentSnapshotRequestParametersOrDict = Union[ + _CreateSandboxEnvironmentSnapshotRequestParameters, + _CreateSandboxEnvironmentSnapshotRequestParametersDict, ] -class SandboxEnvironmentTemplate(_common.BaseModel): - """A sandbox environment template.""" +class SandboxEnvironmentSnapshot(_common.BaseModel): + """A sandbox environment snapshot.""" + display_name: Optional[str] = Field( + default=None, + description="""The display name of the sandbox environment snapshot.""", + ) + expire_time: Optional[datetime.datetime] = Field( + default=None, + description="""Expiration time of the sandbox environment snapshot. + """, + ) create_time: Optional[datetime.datetime] = Field( default=None, - description="""Output only. The timestamp when this SandboxEnvironmentTemplate was created.""", + description="""Output only. The timestamp when this SandboxEnvironmentSnapshot was created.""", ) - custom_container_environment: Optional[ - SandboxEnvironmentTemplateCustomContainerEnvironment - ] = Field( + name: Optional[str] = Field( default=None, - description="""The sandbox environment for custom container workloads.""", + description="""Identifier. The resource name of the SandboxEnvironmentSnapshot. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}`""", ) - default_container_environment: Optional[ - SandboxEnvironmentTemplateDefaultContainerEnvironment - ] = Field( + owner: Optional[str] = Field( default=None, - description="""The sandbox environment for default container workloads.""", + description="""Optional. Owner information for this sandbox snapshot. Different owners will have isolations on snapshot storage and identity. If not set, snapshot will be created as the default owner.""", ) - display_name: Optional[str] = Field( + parent_snapshot: Optional[str] = Field( default=None, - description="""Required. The display name of the SandboxEnvironmentTemplate.""", + description="""Output only. The resource name of the parent SandboxEnvironmentSnapshot. Empty if this is a root Snapshot (the first snapshot from a newly created sandbox). Can be used to reconstruct the whole ancestry tree of snapshots.""", ) - egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = ( - Field( - default=None, - description="""Optional. The configuration for egress control of this template.""", - ) + post_snapshot_action: Optional[PostSnapshotAction] = Field( + default=None, + description="""Optional. Input only. Action to take on the source SandboxEnvironment after the snapshot is taken. This field is only used in CreateSandboxEnvironmentSnapshotRequest and it is not stored in the resource.""", ) - name: Optional[str] = Field( + size_bytes: Optional[int] = Field( default=None, - description="""Identifier. The resource name of the SandboxEnvironmentTemplate. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}`""", + description="""Optional. Output only. Size of the snapshot data in bytes.""", ) - state: Optional[ - Literal[ - "UNSPECIFIED", - "PROVISIONING", - "ACTIVE", - "DEPROVISIONING", - "DELETED", - "FAILED", - ] - ] = Field( + source_sandbox_environment: Optional[str] = Field( default=None, - description="""Output only. The state of the sandbox environment template.""", + description="""Required. The resource name of the source SandboxEnvironment this snapshot was taken from.""", + ) + ttl: Optional[str] = Field( + default=None, + description="""Optional. Input only. The TTL for the sandbox environment snapshot. The expiration time is computed: now + TTL.""", ) update_time: Optional[datetime.datetime] = Field( default=None, - description="""Output only. The timestamp when this SandboxEnvironmentTemplate was most recently updated.""", + description="""Output only. The timestamp when this SandboxEnvironment was most recently updated.""", ) -class SandboxEnvironmentTemplateDict(TypedDict, total=False): - """A sandbox environment template.""" +class SandboxEnvironmentSnapshotDict(TypedDict, total=False): + """A sandbox environment snapshot.""" + + display_name: Optional[str] + """The display name of the sandbox environment snapshot.""" + + expire_time: Optional[datetime.datetime] + """Expiration time of the sandbox environment snapshot. + """ create_time: Optional[datetime.datetime] - """Output only. The timestamp when this SandboxEnvironmentTemplate was created.""" + """Output only. The timestamp when this SandboxEnvironmentSnapshot was created.""" - custom_container_environment: Optional[ - SandboxEnvironmentTemplateCustomContainerEnvironmentDict - ] - """The sandbox environment for custom container workloads.""" + name: Optional[str] + """Identifier. The resource name of the SandboxEnvironmentSnapshot. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}`""" - default_container_environment: Optional[ - SandboxEnvironmentTemplateDefaultContainerEnvironmentDict - ] - """The sandbox environment for default container workloads.""" + owner: Optional[str] + """Optional. Owner information for this sandbox snapshot. Different owners will have isolations on snapshot storage and identity. If not set, snapshot will be created as the default owner.""" - display_name: Optional[str] - """Required. The display name of the SandboxEnvironmentTemplate.""" + parent_snapshot: Optional[str] + """Output only. The resource name of the parent SandboxEnvironmentSnapshot. Empty if this is a root Snapshot (the first snapshot from a newly created sandbox). Can be used to reconstruct the whole ancestry tree of snapshots.""" - egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfigDict] - """Optional. The configuration for egress control of this template.""" + post_snapshot_action: Optional[PostSnapshotAction] + """Optional. Input only. Action to take on the source SandboxEnvironment after the snapshot is taken. This field is only used in CreateSandboxEnvironmentSnapshotRequest and it is not stored in the resource.""" - name: Optional[str] - """Identifier. The resource name of the SandboxEnvironmentTemplate. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}`""" + size_bytes: Optional[int] + """Optional. Output only. Size of the snapshot data in bytes.""" - state: Optional[ - Literal[ - "UNSPECIFIED", - "PROVISIONING", - "ACTIVE", - "DEPROVISIONING", - "DELETED", - "FAILED", - ] - ] - """Output only. The state of the sandbox environment template.""" + source_sandbox_environment: Optional[str] + """Required. The resource name of the source SandboxEnvironment this snapshot was taken from.""" + + ttl: Optional[str] + """Optional. Input only. The TTL for the sandbox environment snapshot. The expiration time is computed: now + TTL.""" update_time: Optional[datetime.datetime] - """Output only. The timestamp when this SandboxEnvironmentTemplate was most recently updated.""" + """Output only. The timestamp when this SandboxEnvironment was most recently updated.""" -SandboxEnvironmentTemplateOrDict = Union[ - SandboxEnvironmentTemplate, SandboxEnvironmentTemplateDict +SandboxEnvironmentSnapshotOrDict = Union[ + SandboxEnvironmentSnapshot, SandboxEnvironmentSnapshotDict ] -class SandboxEnvironmentTemplateOperation(_common.BaseModel): - """Operation that has an agent engine sandbox as a response.""" +class RuntimeSandboxSnapshotOperation(_common.BaseModel): + """Operation that has an agent runtime sandbox snapshot as a response.""" name: Optional[str] = Field( default=None, @@ -17256,13 +16778,13 @@ class SandboxEnvironmentTemplateOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) - response: Optional[SandboxEnvironmentTemplate] = Field( - default=None, description="""The Agent Engine Sandbox Template.""" + response: Optional[SandboxEnvironmentSnapshot] = Field( + default=None, description="""The Agent Runtime Sandbox Snapshot.""" ) -class SandboxEnvironmentTemplateOperationDict(TypedDict, total=False): - """Operation that has an agent engine sandbox as a response.""" +class RuntimeSandboxSnapshotOperationDict(TypedDict, total=False): + """Operation that has an agent runtime sandbox snapshot as a response.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -17276,64 +16798,65 @@ class SandboxEnvironmentTemplateOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" - response: Optional[SandboxEnvironmentTemplateDict] - """The Agent Engine Sandbox Template.""" + response: Optional[SandboxEnvironmentSnapshotDict] + """The Agent Runtime Sandbox Snapshot.""" -SandboxEnvironmentTemplateOperationOrDict = Union[ - SandboxEnvironmentTemplateOperation, SandboxEnvironmentTemplateOperationDict +RuntimeSandboxSnapshotOperationOrDict = Union[ + RuntimeSandboxSnapshotOperation, RuntimeSandboxSnapshotOperationDict ] -class DeleteSandboxEnvironmentTemplateConfig(_common.BaseModel): - """Config for deleting a Sandbox Template.""" +class DeleteSandboxEnvironmentSnapshotConfig(_common.BaseModel): + """Config for deleting a Sandbox Environment Snapshot.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class DeleteSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): - """Config for deleting a Sandbox Template.""" +class DeleteSandboxEnvironmentSnapshotConfigDict(TypedDict, total=False): + """Config for deleting a Sandbox Environment Snapshot.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteSandboxEnvironmentTemplateConfigOrDict = Union[ - DeleteSandboxEnvironmentTemplateConfig, DeleteSandboxEnvironmentTemplateConfigDict +DeleteSandboxEnvironmentSnapshotConfigOrDict = Union[ + DeleteSandboxEnvironmentSnapshotConfig, DeleteSandboxEnvironmentSnapshotConfigDict ] -class _DeleteSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): - """Parameters for deleting sandbox templates.""" +class _DeleteSandboxEnvironmentSnapshotRequestParameters(_common.BaseModel): + """Parameters for deleting sandbox environment snapshots.""" name: Optional[str] = Field( - default=None, description="""Name of the sandbox template to delete.""" + default=None, + description="""Name of the sandbox environment snapshot to delete.""", ) - config: Optional[DeleteSandboxEnvironmentTemplateConfig] = Field( + config: Optional[DeleteSandboxEnvironmentSnapshotConfig] = Field( default=None, description="""""" ) -class _DeleteSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): - """Parameters for deleting sandbox templates.""" +class _DeleteSandboxEnvironmentSnapshotRequestParametersDict(TypedDict, total=False): + """Parameters for deleting sandbox environment snapshots.""" name: Optional[str] - """Name of the sandbox template to delete.""" + """Name of the sandbox environment snapshot to delete.""" - config: Optional[DeleteSandboxEnvironmentTemplateConfigDict] + config: Optional[DeleteSandboxEnvironmentSnapshotConfigDict] """""" -_DeleteSandboxEnvironmentTemplateRequestParametersOrDict = Union[ - _DeleteSandboxEnvironmentTemplateRequestParameters, - _DeleteSandboxEnvironmentTemplateRequestParametersDict, +_DeleteSandboxEnvironmentSnapshotRequestParametersOrDict = Union[ + _DeleteSandboxEnvironmentSnapshotRequestParameters, + _DeleteSandboxEnvironmentSnapshotRequestParametersDict, ] -class DeleteSandboxEnvironmentTemplateOperation(_common.BaseModel): - """Operation for deleting sandbox templates.""" +class DeleteSandboxEnvironmentSnapshotOperation(_common.BaseModel): + """Operation for deleting sandbox environment snapshots.""" name: Optional[str] = Field( default=None, @@ -17353,8 +16876,8 @@ class DeleteSandboxEnvironmentTemplateOperation(_common.BaseModel): ) -class DeleteSandboxEnvironmentTemplateOperationDict(TypedDict, total=False): - """Operation for deleting sandbox templates.""" +class DeleteSandboxEnvironmentSnapshotOperationDict(TypedDict, total=False): + """Operation for deleting sandbox environment snapshots.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -17369,61 +16892,61 @@ class DeleteSandboxEnvironmentTemplateOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -DeleteSandboxEnvironmentTemplateOperationOrDict = Union[ - DeleteSandboxEnvironmentTemplateOperation, - DeleteSandboxEnvironmentTemplateOperationDict, -] - +DeleteSandboxEnvironmentSnapshotOperationOrDict = Union[ + DeleteSandboxEnvironmentSnapshotOperation, + DeleteSandboxEnvironmentSnapshotOperationDict, +] -class GetSandboxEnvironmentTemplateConfig(_common.BaseModel): - """Config for getting a Sandbox Template.""" + +class GetSandboxEnvironmentSnapshotConfig(_common.BaseModel): + """Config for getting a Sandbox Environment Snapshot.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): - """Config for getting a Sandbox Template.""" +class GetSandboxEnvironmentSnapshotConfigDict(TypedDict, total=False): + """Config for getting a Sandbox Environment Snapshot.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetSandboxEnvironmentTemplateConfigOrDict = Union[ - GetSandboxEnvironmentTemplateConfig, GetSandboxEnvironmentTemplateConfigDict +GetSandboxEnvironmentSnapshotConfigOrDict = Union[ + GetSandboxEnvironmentSnapshotConfig, GetSandboxEnvironmentSnapshotConfigDict ] -class _GetSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): - """Parameters for getting a sandbox template.""" +class _GetSandboxEnvironmentSnapshotRequestParameters(_common.BaseModel): + """Parameters for getting a sandbox environment snapshot.""" name: Optional[str] = Field( - default=None, description="""Name of the sandbox template.""" + default=None, description="""Name of the sandbox environment snapshot.""" ) - config: Optional[GetSandboxEnvironmentTemplateConfig] = Field( + config: Optional[GetSandboxEnvironmentSnapshotConfig] = Field( default=None, description="""""" ) -class _GetSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): - """Parameters for getting a sandbox template.""" +class _GetSandboxEnvironmentSnapshotRequestParametersDict(TypedDict, total=False): + """Parameters for getting a sandbox environment snapshot.""" name: Optional[str] - """Name of the sandbox template.""" + """Name of the sandbox environment snapshot.""" - config: Optional[GetSandboxEnvironmentTemplateConfigDict] + config: Optional[GetSandboxEnvironmentSnapshotConfigDict] """""" -_GetSandboxEnvironmentTemplateRequestParametersOrDict = Union[ - _GetSandboxEnvironmentTemplateRequestParameters, - _GetSandboxEnvironmentTemplateRequestParametersDict, +_GetSandboxEnvironmentSnapshotRequestParametersOrDict = Union[ + _GetSandboxEnvironmentSnapshotRequestParameters, + _GetSandboxEnvironmentSnapshotRequestParametersDict, ] -class ListSandboxEnvironmentTemplatesConfig(_common.BaseModel): - """Config for listing sandbox templates.""" +class ListSandboxEnvironmentSnapshotsConfig(_common.BaseModel): + """Config for listing sandbox environment snapshots.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -17436,8 +16959,8 @@ class ListSandboxEnvironmentTemplatesConfig(_common.BaseModel): ) -class ListSandboxEnvironmentTemplatesConfigDict(TypedDict, total=False): - """Config for listing sandbox templates.""" +class ListSandboxEnvironmentSnapshotsConfigDict(TypedDict, total=False): + """Config for listing sandbox environment snapshots.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -17452,52 +16975,53 @@ class ListSandboxEnvironmentTemplatesConfigDict(TypedDict, total=False): """An expression for filtering the results of the request.""" -ListSandboxEnvironmentTemplatesConfigOrDict = Union[ - ListSandboxEnvironmentTemplatesConfig, ListSandboxEnvironmentTemplatesConfigDict +ListSandboxEnvironmentSnapshotsConfigOrDict = Union[ + ListSandboxEnvironmentSnapshotsConfig, ListSandboxEnvironmentSnapshotsConfigDict ] -class _ListSandboxEnvironmentTemplatesRequestParameters(_common.BaseModel): - """Parameters for listing sandbox templates.""" +class _ListSandboxEnvironmentSnapshotsRequestParameters(_common.BaseModel): + """Parameters for listing sandbox environment snapshots.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, + description="""Name of the reasoning engine to list snapshots from.""", ) - config: Optional[ListSandboxEnvironmentTemplatesConfig] = Field( + config: Optional[ListSandboxEnvironmentSnapshotsConfig] = Field( default=None, description="""""" ) -class _ListSandboxEnvironmentTemplatesRequestParametersDict(TypedDict, total=False): - """Parameters for listing sandbox templates.""" +class _ListSandboxEnvironmentSnapshotsRequestParametersDict(TypedDict, total=False): + """Parameters for listing sandbox environment snapshots.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the reasoning engine to list snapshots from.""" - config: Optional[ListSandboxEnvironmentTemplatesConfigDict] + config: Optional[ListSandboxEnvironmentSnapshotsConfigDict] """""" -_ListSandboxEnvironmentTemplatesRequestParametersOrDict = Union[ - _ListSandboxEnvironmentTemplatesRequestParameters, - _ListSandboxEnvironmentTemplatesRequestParametersDict, +_ListSandboxEnvironmentSnapshotsRequestParametersOrDict = Union[ + _ListSandboxEnvironmentSnapshotsRequestParameters, + _ListSandboxEnvironmentSnapshotsRequestParametersDict, ] -class ListSandboxEnvironmentTemplatesResponse(_common.BaseModel): - """Response for listing sandbox templates.""" +class ListSandboxEnvironmentSnapshotsResponse(_common.BaseModel): + """Response for listing sandbox environment snapshots.""" sdk_http_response: Optional[genai_types.HttpResponse] = Field( default=None, description="""Used to retain the full HTTP response.""" ) next_page_token: Optional[str] = Field(default=None, description="""""") - sandbox_environment_templates: Optional[list[SandboxEnvironmentTemplate]] = Field( - default=None, description="""List of sandbox templates.""" + sandbox_environment_snapshots: Optional[list[SandboxEnvironmentSnapshot]] = Field( + default=None, description="""List of sandbox environment snapshots.""" ) -class ListSandboxEnvironmentTemplatesResponseDict(TypedDict, total=False): - """Response for listing sandbox templates.""" +class ListSandboxEnvironmentSnapshotsResponseDict(TypedDict, total=False): + """Response for listing sandbox environment snapshots.""" sdk_http_response: Optional[genai_types.HttpResponse] """Used to retain the full HTTP response.""" @@ -17505,210 +17029,223 @@ class ListSandboxEnvironmentTemplatesResponseDict(TypedDict, total=False): next_page_token: Optional[str] """""" - sandbox_environment_templates: Optional[list[SandboxEnvironmentTemplateDict]] - """List of sandbox templates.""" + sandbox_environment_snapshots: Optional[list[SandboxEnvironmentSnapshotDict]] + """List of sandbox environment snapshots.""" -ListSandboxEnvironmentTemplatesResponseOrDict = Union[ - ListSandboxEnvironmentTemplatesResponse, ListSandboxEnvironmentTemplatesResponseDict +ListSandboxEnvironmentSnapshotsResponseOrDict = Union[ + ListSandboxEnvironmentSnapshotsResponse, ListSandboxEnvironmentSnapshotsResponseDict ] -class _GetSandboxEnvironmentTemplateOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a sandbox template as a response.""" +class _GetRuntimeSandboxSnapshotOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a sandbox snapshot as a response.""" operation_name: Optional[str] = Field( default=None, description="""The server-assigned name for the operation.""" ) - config: Optional[GetAgentEngineOperationConfig] = Field( + config: Optional[GetRuntimeOperationConfig] = Field( default=None, description="""Used to override the default configuration.""" ) -class _GetSandboxEnvironmentTemplateOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a sandbox template as a response.""" +class _GetRuntimeSandboxSnapshotOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a sandbox snapshot as a response.""" operation_name: Optional[str] """The server-assigned name for the operation.""" - config: Optional[GetAgentEngineOperationConfigDict] + config: Optional[GetRuntimeOperationConfigDict] """Used to override the default configuration.""" -_GetSandboxEnvironmentTemplateOperationParametersOrDict = Union[ - _GetSandboxEnvironmentTemplateOperationParameters, - _GetSandboxEnvironmentTemplateOperationParametersDict, +_GetRuntimeSandboxSnapshotOperationParametersOrDict = Union[ + _GetRuntimeSandboxSnapshotOperationParameters, + _GetRuntimeSandboxSnapshotOperationParametersDict, ] -class CreateAgentEngineSandboxSnapshotConfig(_common.BaseModel): - """Config for creating a Sandbox Environment Snapshot.""" +class CreateRuntimeSessionConfig(_common.BaseModel): + """Config for creating a Session.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) display_name: Optional[str] = Field( - default=None, description="""The display name of the sandbox snapshot.""" - ) - owner: Optional[str] = Field( - default=None, description="""The owner of the sandbox snapshot.""" + default=None, description="""The display name of the session.""" ) - ttl: Optional[str] = Field( + session_state: Optional[dict[str, Any]] = Field( default=None, - description="""The TTL for this resource. The expiration time is computed: now + TTL.""", + description="""Session state which stores key conversation points.""", ) wait_for_completion: Optional[bool] = Field( default=True, description="""Waits for the operation to complete before returning.""", ) + ttl: Optional[str] = Field( + default=None, + description="""Optional. Input only. The TTL for this resource. + + The expiration time is computed: now + TTL.""", + ) + expire_time: Optional[datetime.datetime] = Field( + default=None, + description="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", + ) + labels: Optional[dict[str, str]] = Field( + default=None, + description="""Optional. The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", + ) + session_id: Optional[str] = Field( + default=None, + description="""Optional. The user defined ID to use for session, which will become the final component of the session resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""", + ) -class CreateAgentEngineSandboxSnapshotConfigDict(TypedDict, total=False): - """Config for creating a Sandbox Environment Snapshot.""" +class CreateRuntimeSessionConfigDict(TypedDict, total=False): + """Config for creating a Session.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" display_name: Optional[str] - """The display name of the sandbox snapshot.""" - - owner: Optional[str] - """The owner of the sandbox snapshot.""" + """The display name of the session.""" - ttl: Optional[str] - """The TTL for this resource. The expiration time is computed: now + TTL.""" + session_state: Optional[dict[str, Any]] + """Session state which stores key conversation points.""" wait_for_completion: Optional[bool] """Waits for the operation to complete before returning.""" + ttl: Optional[str] + """Optional. Input only. The TTL for this resource. + + The expiration time is computed: now + TTL.""" + + expire_time: Optional[datetime.datetime] + """Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""" + + labels: Optional[dict[str, str]] + """Optional. The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""" + + session_id: Optional[str] + """Optional. The user defined ID to use for session, which will become the final component of the session resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""" -CreateAgentEngineSandboxSnapshotConfigOrDict = Union[ - CreateAgentEngineSandboxSnapshotConfig, CreateAgentEngineSandboxSnapshotConfigDict + +CreateRuntimeSessionConfigOrDict = Union[ + CreateRuntimeSessionConfig, CreateRuntimeSessionConfigDict ] -class _CreateSandboxEnvironmentSnapshotRequestParameters(_common.BaseModel): - """Parameters for creating a sandbox environment snapshot.""" +class _CreateRuntimeSessionRequestParameters(_common.BaseModel): + """Parameters for creating Agent Runtime Sessions.""" - source_sandbox_environment_name: Optional[str] = Field( - default=None, description="""Name of the sandbox environment to snapshot.""" + name: Optional[str] = Field( + default=None, + description="""Name of the agent runtime to create the session under.""", + ) + user_id: Optional[str] = Field( + default=None, description="""The user ID of the session.""" ) - config: Optional[CreateAgentEngineSandboxSnapshotConfig] = Field( + config: Optional[CreateRuntimeSessionConfig] = Field( default=None, description="""""" ) -class _CreateSandboxEnvironmentSnapshotRequestParametersDict(TypedDict, total=False): - """Parameters for creating a sandbox environment snapshot.""" +class _CreateRuntimeSessionRequestParametersDict(TypedDict, total=False): + """Parameters for creating Agent Runtime Sessions.""" - source_sandbox_environment_name: Optional[str] - """Name of the sandbox environment to snapshot.""" + name: Optional[str] + """Name of the agent runtime to create the session under.""" + + user_id: Optional[str] + """The user ID of the session.""" - config: Optional[CreateAgentEngineSandboxSnapshotConfigDict] + config: Optional[CreateRuntimeSessionConfigDict] """""" -_CreateSandboxEnvironmentSnapshotRequestParametersOrDict = Union[ - _CreateSandboxEnvironmentSnapshotRequestParameters, - _CreateSandboxEnvironmentSnapshotRequestParametersDict, +_CreateRuntimeSessionRequestParametersOrDict = Union[ + _CreateRuntimeSessionRequestParameters, _CreateRuntimeSessionRequestParametersDict ] -class SandboxEnvironmentSnapshot(_common.BaseModel): - """A sandbox environment snapshot.""" +class Session(_common.BaseModel): + """A session.""" - display_name: Optional[str] = Field( - default=None, - description="""The display name of the sandbox environment snapshot.""", - ) - expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Expiration time of the sandbox environment snapshot. - """, - ) create_time: Optional[datetime.datetime] = Field( default=None, - description="""Output only. The timestamp when this SandboxEnvironmentSnapshot was created.""", - ) - name: Optional[str] = Field( - default=None, - description="""Identifier. The resource name of the SandboxEnvironmentSnapshot. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}`""", + description="""Output only. Timestamp when the session was created.""", ) - owner: Optional[str] = Field( - default=None, - description="""Optional. Owner information for this sandbox snapshot. Different owners will have isolations on snapshot storage and identity. If not set, snapshot will be created as the default owner.""", + display_name: Optional[str] = Field( + default=None, description="""Optional. The display name of the session.""" ) - parent_snapshot: Optional[str] = Field( + expire_time: Optional[datetime.datetime] = Field( default=None, - description="""Output only. The resource name of the parent SandboxEnvironmentSnapshot. Empty if this is a root Snapshot (the first snapshot from a newly created sandbox). Can be used to reconstruct the whole ancestry tree of snapshots.""", + description="""Optional. Timestamp of when this session is considered expired. This is *always* provided on output, regardless of what was sent on input. The minimum value is 24 hours from the time of creation.""", ) - post_snapshot_action: Optional[PostSnapshotAction] = Field( + labels: Optional[dict[str, str]] = Field( default=None, - description="""Optional. Input only. Action to take on the source SandboxEnvironment after the snapshot is taken. This field is only used in CreateSandboxEnvironmentSnapshotRequest and it is not stored in the resource.""", + description="""The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", ) - size_bytes: Optional[int] = Field( + name: Optional[str] = Field( default=None, - description="""Optional. Output only. Size of the snapshot data in bytes.""", + description="""Identifier. The resource name of the session. Format: 'projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}'.""", ) - source_sandbox_environment: Optional[str] = Field( + session_state: Optional[dict[str, Any]] = Field( default=None, - description="""Required. The resource name of the source SandboxEnvironment this snapshot was taken from.""", + description="""Optional. Session specific memory which stores key conversation points.""", ) ttl: Optional[str] = Field( default=None, - description="""Optional. Input only. The TTL for the sandbox environment snapshot. The expiration time is computed: now + TTL.""", + description="""Optional. Input only. The TTL for this session. The minimum value is 24 hours.""", ) update_time: Optional[datetime.datetime] = Field( default=None, - description="""Output only. The timestamp when this SandboxEnvironment was most recently updated.""", + description="""Output only. Timestamp when the session was updated.""", + ) + user_id: Optional[str] = Field( + default=None, + description="""Required. Immutable. String id provided by the user""", ) -class SandboxEnvironmentSnapshotDict(TypedDict, total=False): - """A sandbox environment snapshot.""" +class SessionDict(TypedDict, total=False): + """A session.""" + + create_time: Optional[datetime.datetime] + """Output only. Timestamp when the session was created.""" display_name: Optional[str] - """The display name of the sandbox environment snapshot.""" + """Optional. The display name of the session.""" expire_time: Optional[datetime.datetime] - """Expiration time of the sandbox environment snapshot. - """ + """Optional. Timestamp of when this session is considered expired. This is *always* provided on output, regardless of what was sent on input. The minimum value is 24 hours from the time of creation.""" - create_time: Optional[datetime.datetime] - """Output only. The timestamp when this SandboxEnvironmentSnapshot was created.""" + labels: Optional[dict[str, str]] + """The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""" name: Optional[str] - """Identifier. The resource name of the SandboxEnvironmentSnapshot. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}`""" - - owner: Optional[str] - """Optional. Owner information for this sandbox snapshot. Different owners will have isolations on snapshot storage and identity. If not set, snapshot will be created as the default owner.""" + """Identifier. The resource name of the session. Format: 'projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}'.""" - parent_snapshot: Optional[str] - """Output only. The resource name of the parent SandboxEnvironmentSnapshot. Empty if this is a root Snapshot (the first snapshot from a newly created sandbox). Can be used to reconstruct the whole ancestry tree of snapshots.""" + session_state: Optional[dict[str, Any]] + """Optional. Session specific memory which stores key conversation points.""" - post_snapshot_action: Optional[PostSnapshotAction] - """Optional. Input only. Action to take on the source SandboxEnvironment after the snapshot is taken. This field is only used in CreateSandboxEnvironmentSnapshotRequest and it is not stored in the resource.""" + ttl: Optional[str] + """Optional. Input only. The TTL for this session. The minimum value is 24 hours.""" - size_bytes: Optional[int] - """Optional. Output only. Size of the snapshot data in bytes.""" + update_time: Optional[datetime.datetime] + """Output only. Timestamp when the session was updated.""" - source_sandbox_environment: Optional[str] - """Required. The resource name of the source SandboxEnvironment this snapshot was taken from.""" - - ttl: Optional[str] - """Optional. Input only. The TTL for the sandbox environment snapshot. The expiration time is computed: now + TTL.""" - - update_time: Optional[datetime.datetime] - """Output only. The timestamp when this SandboxEnvironment was most recently updated.""" + user_id: Optional[str] + """Required. Immutable. String id provided by the user""" -SandboxEnvironmentSnapshotOrDict = Union[ - SandboxEnvironmentSnapshot, SandboxEnvironmentSnapshotDict -] +SessionOrDict = Union[Session, SessionDict] -class AgentEngineSandboxSnapshotOperation(_common.BaseModel): - """Operation that has an agent engine sandbox snapshot as a response.""" +class RuntimeSessionOperation(_common.BaseModel): + """Operation that has an agent runtime session as a response.""" name: Optional[str] = Field( default=None, @@ -17726,13 +17263,13 @@ class AgentEngineSandboxSnapshotOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) - response: Optional[SandboxEnvironmentSnapshot] = Field( - default=None, description="""The Agent Engine Sandbox Snapshot.""" + response: Optional[Session] = Field( + default=None, description="""The Agent Runtime Session.""" ) -class AgentEngineSandboxSnapshotOperationDict(TypedDict, total=False): - """Operation that has an agent engine sandbox snapshot as a response.""" +class RuntimeSessionOperationDict(TypedDict, total=False): + """Operation that has an agent runtime session as a response.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -17746,65 +17283,63 @@ class AgentEngineSandboxSnapshotOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" - response: Optional[SandboxEnvironmentSnapshotDict] - """The Agent Engine Sandbox Snapshot.""" + response: Optional[SessionDict] + """The Agent Runtime Session.""" -AgentEngineSandboxSnapshotOperationOrDict = Union[ - AgentEngineSandboxSnapshotOperation, AgentEngineSandboxSnapshotOperationDict +RuntimeSessionOperationOrDict = Union[ + RuntimeSessionOperation, RuntimeSessionOperationDict ] -class DeleteSandboxEnvironmentSnapshotConfig(_common.BaseModel): - """Config for deleting a Sandbox Environment Snapshot.""" +class DeleteRuntimeSessionConfig(_common.BaseModel): + """Config for deleting an Agent Runtime Session.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class DeleteSandboxEnvironmentSnapshotConfigDict(TypedDict, total=False): - """Config for deleting a Sandbox Environment Snapshot.""" +class DeleteRuntimeSessionConfigDict(TypedDict, total=False): + """Config for deleting an Agent Runtime Session.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteSandboxEnvironmentSnapshotConfigOrDict = Union[ - DeleteSandboxEnvironmentSnapshotConfig, DeleteSandboxEnvironmentSnapshotConfigDict +DeleteRuntimeSessionConfigOrDict = Union[ + DeleteRuntimeSessionConfig, DeleteRuntimeSessionConfigDict ] -class _DeleteSandboxEnvironmentSnapshotRequestParameters(_common.BaseModel): - """Parameters for deleting sandbox environment snapshots.""" +class _DeleteRuntimeSessionRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtime sessions.""" name: Optional[str] = Field( - default=None, - description="""Name of the sandbox environment snapshot to delete.""", + default=None, description="""Name of the agent runtime session to delete.""" ) - config: Optional[DeleteSandboxEnvironmentSnapshotConfig] = Field( + config: Optional[DeleteRuntimeSessionConfig] = Field( default=None, description="""""" ) -class _DeleteSandboxEnvironmentSnapshotRequestParametersDict(TypedDict, total=False): - """Parameters for deleting sandbox environment snapshots.""" +class _DeleteRuntimeSessionRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtime sessions.""" name: Optional[str] - """Name of the sandbox environment snapshot to delete.""" + """Name of the agent runtime session to delete.""" - config: Optional[DeleteSandboxEnvironmentSnapshotConfigDict] + config: Optional[DeleteRuntimeSessionConfigDict] """""" -_DeleteSandboxEnvironmentSnapshotRequestParametersOrDict = Union[ - _DeleteSandboxEnvironmentSnapshotRequestParameters, - _DeleteSandboxEnvironmentSnapshotRequestParametersDict, +_DeleteRuntimeSessionRequestParametersOrDict = Union[ + _DeleteRuntimeSessionRequestParameters, _DeleteRuntimeSessionRequestParametersDict ] -class DeleteSandboxEnvironmentSnapshotOperation(_common.BaseModel): - """Operation for deleting sandbox environment snapshots.""" +class DeleteRuntimeSessionOperation(_common.BaseModel): + """Operation for deleting agent runtime sessions.""" name: Optional[str] = Field( default=None, @@ -17824,8 +17359,8 @@ class DeleteSandboxEnvironmentSnapshotOperation(_common.BaseModel): ) -class DeleteSandboxEnvironmentSnapshotOperationDict(TypedDict, total=False): - """Operation for deleting sandbox environment snapshots.""" +class DeleteRuntimeSessionOperationDict(TypedDict, total=False): + """Operation for deleting agent runtime sessions.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -17840,61 +17375,57 @@ class DeleteSandboxEnvironmentSnapshotOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -DeleteSandboxEnvironmentSnapshotOperationOrDict = Union[ - DeleteSandboxEnvironmentSnapshotOperation, - DeleteSandboxEnvironmentSnapshotOperationDict, +DeleteRuntimeSessionOperationOrDict = Union[ + DeleteRuntimeSessionOperation, DeleteRuntimeSessionOperationDict ] -class GetSandboxEnvironmentSnapshotConfig(_common.BaseModel): - """Config for getting a Sandbox Environment Snapshot.""" +class GetRuntimeSessionConfig(_common.BaseModel): + """Config for getting an Agent Runtime Session.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetSandboxEnvironmentSnapshotConfigDict(TypedDict, total=False): - """Config for getting a Sandbox Environment Snapshot.""" +class GetRuntimeSessionConfigDict(TypedDict, total=False): + """Config for getting an Agent Runtime Session.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetSandboxEnvironmentSnapshotConfigOrDict = Union[ - GetSandboxEnvironmentSnapshotConfig, GetSandboxEnvironmentSnapshotConfigDict +GetRuntimeSessionConfigOrDict = Union[ + GetRuntimeSessionConfig, GetRuntimeSessionConfigDict ] -class _GetSandboxEnvironmentSnapshotRequestParameters(_common.BaseModel): - """Parameters for getting a sandbox environment snapshot.""" +class _GetRuntimeSessionRequestParameters(_common.BaseModel): + """Parameters for getting an agent runtime session.""" name: Optional[str] = Field( - default=None, description="""Name of the sandbox environment snapshot.""" - ) - config: Optional[GetSandboxEnvironmentSnapshotConfig] = Field( - default=None, description="""""" + default=None, description="""Name of the agent runtime session.""" ) + config: Optional[GetRuntimeSessionConfig] = Field(default=None, description="""""") -class _GetSandboxEnvironmentSnapshotRequestParametersDict(TypedDict, total=False): - """Parameters for getting a sandbox environment snapshot.""" +class _GetRuntimeSessionRequestParametersDict(TypedDict, total=False): + """Parameters for getting an agent runtime session.""" name: Optional[str] - """Name of the sandbox environment snapshot.""" + """Name of the agent runtime session.""" - config: Optional[GetSandboxEnvironmentSnapshotConfigDict] + config: Optional[GetRuntimeSessionConfigDict] """""" -_GetSandboxEnvironmentSnapshotRequestParametersOrDict = Union[ - _GetSandboxEnvironmentSnapshotRequestParameters, - _GetSandboxEnvironmentSnapshotRequestParametersDict, +_GetRuntimeSessionRequestParametersOrDict = Union[ + _GetRuntimeSessionRequestParameters, _GetRuntimeSessionRequestParametersDict ] -class ListSandboxEnvironmentSnapshotsConfig(_common.BaseModel): - """Config for listing sandbox environment snapshots.""" +class ListRuntimeSessionsConfig(_common.BaseModel): + """Config for listing agent runtime sessions.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -17903,12 +17434,13 @@ class ListSandboxEnvironmentSnapshotsConfig(_common.BaseModel): page_token: Optional[str] = Field(default=None, description="""""") filter: Optional[str] = Field( default=None, - description="""An expression for filtering the results of the request.""", + description="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""", ) -class ListSandboxEnvironmentSnapshotsConfigDict(TypedDict, total=False): - """Config for listing sandbox environment snapshots.""" +class ListRuntimeSessionsConfigDict(TypedDict, total=False): + """Config for listing agent runtime sessions.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -17920,56 +17452,55 @@ class ListSandboxEnvironmentSnapshotsConfigDict(TypedDict, total=False): """""" filter: Optional[str] - """An expression for filtering the results of the request.""" + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""" -ListSandboxEnvironmentSnapshotsConfigOrDict = Union[ - ListSandboxEnvironmentSnapshotsConfig, ListSandboxEnvironmentSnapshotsConfigDict +ListRuntimeSessionsConfigOrDict = Union[ + ListRuntimeSessionsConfig, ListRuntimeSessionsConfigDict ] -class _ListSandboxEnvironmentSnapshotsRequestParameters(_common.BaseModel): - """Parameters for listing sandbox environment snapshots.""" +class _ListRuntimeSessionsRequestParameters(_common.BaseModel): + """Parameters for listing agent runtimes.""" name: Optional[str] = Field( - default=None, - description="""Name of the reasoning engine to list snapshots from.""", + default=None, description="""Name of the agent runtime.""" ) - config: Optional[ListSandboxEnvironmentSnapshotsConfig] = Field( + config: Optional[ListRuntimeSessionsConfig] = Field( default=None, description="""""" ) -class _ListSandboxEnvironmentSnapshotsRequestParametersDict(TypedDict, total=False): - """Parameters for listing sandbox environment snapshots.""" +class _ListRuntimeSessionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtimes.""" name: Optional[str] - """Name of the reasoning engine to list snapshots from.""" + """Name of the agent runtime.""" - config: Optional[ListSandboxEnvironmentSnapshotsConfigDict] + config: Optional[ListRuntimeSessionsConfigDict] """""" -_ListSandboxEnvironmentSnapshotsRequestParametersOrDict = Union[ - _ListSandboxEnvironmentSnapshotsRequestParameters, - _ListSandboxEnvironmentSnapshotsRequestParametersDict, +_ListRuntimeSessionsRequestParametersOrDict = Union[ + _ListRuntimeSessionsRequestParameters, _ListRuntimeSessionsRequestParametersDict ] -class ListSandboxEnvironmentSnapshotsResponse(_common.BaseModel): - """Response for listing sandbox environment snapshots.""" +class ListReasoningEnginesSessionsResponse(_common.BaseModel): + """Response for listing agent runtime sessions.""" sdk_http_response: Optional[genai_types.HttpResponse] = Field( default=None, description="""Used to retain the full HTTP response.""" ) next_page_token: Optional[str] = Field(default=None, description="""""") - sandbox_environment_snapshots: Optional[list[SandboxEnvironmentSnapshot]] = Field( - default=None, description="""List of sandbox environment snapshots.""" + sessions: Optional[list[Session]] = Field( + default=None, description="""List of agent runtime sessions.""" ) -class ListSandboxEnvironmentSnapshotsResponseDict(TypedDict, total=False): - """Response for listing sandbox environment snapshots.""" +class ListReasoningEnginesSessionsResponseDict(TypedDict, total=False): + """Response for listing agent runtime sessions.""" sdk_http_response: Optional[genai_types.HttpResponse] """Used to retain the full HTTP response.""" @@ -17977,44 +17508,43 @@ class ListSandboxEnvironmentSnapshotsResponseDict(TypedDict, total=False): next_page_token: Optional[str] """""" - sandbox_environment_snapshots: Optional[list[SandboxEnvironmentSnapshotDict]] - """List of sandbox environment snapshots.""" + sessions: Optional[list[SessionDict]] + """List of agent runtime sessions.""" -ListSandboxEnvironmentSnapshotsResponseOrDict = Union[ - ListSandboxEnvironmentSnapshotsResponse, ListSandboxEnvironmentSnapshotsResponseDict +ListReasoningEnginesSessionsResponseOrDict = Union[ + ListReasoningEnginesSessionsResponse, ListReasoningEnginesSessionsResponseDict ] -class _GetAgentEngineSandboxSnapshotOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a sandbox snapshot as a response.""" +class _GetRuntimeSessionOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a session as a response.""" operation_name: Optional[str] = Field( default=None, description="""The server-assigned name for the operation.""" ) - config: Optional[GetAgentEngineOperationConfig] = Field( + config: Optional[GetRuntimeOperationConfig] = Field( default=None, description="""Used to override the default configuration.""" ) -class _GetAgentEngineSandboxSnapshotOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a sandbox snapshot as a response.""" +class _GetRuntimeSessionOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a session as a response.""" operation_name: Optional[str] """The server-assigned name for the operation.""" - config: Optional[GetAgentEngineOperationConfigDict] + config: Optional[GetRuntimeOperationConfigDict] """Used to override the default configuration.""" -_GetAgentEngineSandboxSnapshotOperationParametersOrDict = Union[ - _GetAgentEngineSandboxSnapshotOperationParameters, - _GetAgentEngineSandboxSnapshotOperationParametersDict, +_GetRuntimeSessionOperationParametersOrDict = Union[ + _GetRuntimeSessionOperationParameters, _GetRuntimeSessionOperationParametersDict ] -class CreateAgentEngineSessionConfig(_common.BaseModel): - """Config for creating a Session.""" +class UpdateRuntimeSessionConfig(_common.BaseModel): + """Config for updating agent runtime session.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -18048,10 +17578,18 @@ class CreateAgentEngineSessionConfig(_common.BaseModel): default=None, description="""Optional. The user defined ID to use for session, which will become the final component of the session resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""", ) + update_mask: Optional[str] = Field( + default=None, + description="""The update mask to apply. For the `FieldMask` definition, see + https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""", + ) + user_id: Optional[str] = Field( + default=None, description="""User ID of the agent runtime session to update.""" + ) -class CreateAgentEngineSessionConfigDict(TypedDict, total=False): - """Config for creating a Session.""" +class UpdateRuntimeSessionConfigDict(TypedDict, total=False): + """Config for updating agent runtime session.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -18079,305 +17617,293 @@ class CreateAgentEngineSessionConfigDict(TypedDict, total=False): session_id: Optional[str] """Optional. The user defined ID to use for session, which will become the final component of the session resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""" + update_mask: Optional[str] + """The update mask to apply. For the `FieldMask` definition, see + https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" + + user_id: Optional[str] + """User ID of the agent runtime session to update.""" + -CreateAgentEngineSessionConfigOrDict = Union[ - CreateAgentEngineSessionConfig, CreateAgentEngineSessionConfigDict +UpdateRuntimeSessionConfigOrDict = Union[ + UpdateRuntimeSessionConfig, UpdateRuntimeSessionConfigDict ] -class _CreateAgentEngineSessionRequestParameters(_common.BaseModel): - """Parameters for creating Agent Engine Sessions.""" +class _UpdateRuntimeSessionRequestParameters(_common.BaseModel): + """Parameters for updating agent runtime sessions.""" name: Optional[str] = Field( - default=None, - description="""Name of the agent engine to create the session under.""", - ) - user_id: Optional[str] = Field( - default=None, description="""The user ID of the session.""" + default=None, description="""Name of the agent runtime session to update.""" ) - config: Optional[CreateAgentEngineSessionConfig] = Field( + config: Optional[UpdateRuntimeSessionConfig] = Field( default=None, description="""""" ) -class _CreateAgentEngineSessionRequestParametersDict(TypedDict, total=False): - """Parameters for creating Agent Engine Sessions.""" +class _UpdateRuntimeSessionRequestParametersDict(TypedDict, total=False): + """Parameters for updating agent runtime sessions.""" name: Optional[str] - """Name of the agent engine to create the session under.""" - - user_id: Optional[str] - """The user ID of the session.""" + """Name of the agent runtime session to update.""" - config: Optional[CreateAgentEngineSessionConfigDict] + config: Optional[UpdateRuntimeSessionConfigDict] """""" -_CreateAgentEngineSessionRequestParametersOrDict = Union[ - _CreateAgentEngineSessionRequestParameters, - _CreateAgentEngineSessionRequestParametersDict, +_UpdateRuntimeSessionRequestParametersOrDict = Union[ + _UpdateRuntimeSessionRequestParameters, _UpdateRuntimeSessionRequestParametersDict ] -class Session(_common.BaseModel): - """A session.""" +class EventActions(_common.BaseModel): + """Actions are parts of events that are executed by the agent.""" - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when the session was created.""", - ) - display_name: Optional[str] = Field( - default=None, description="""Optional. The display name of the session.""" - ) - expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. Timestamp of when this session is considered expired. This is *always* provided on output, regardless of what was sent on input. The minimum value is 24 hours from the time of creation.""", - ) - labels: Optional[dict[str, str]] = Field( + artifact_delta: Optional[dict[str, int]] = Field( default=None, - description="""The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", + description="""Optional. Indicates that the event is updating an artifact. key is the filename, value is the version.""", ) - name: Optional[str] = Field( + escalate: Optional[bool] = Field( default=None, - description="""Identifier. The resource name of the session. Format: 'projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}'.""", + description="""Optional. The agent is escalating to a higher level agent.""", ) - session_state: Optional[dict[str, Any]] = Field( + requested_auth_configs: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. Session specific memory which stores key conversation points.""", + description="""Optional. Will only be set by a tool response indicating tool request euc. Struct key is the function call id since one function call response (from model) could correspond to multiple function calls. Struct value is the required auth config, which can be another struct.""", ) - ttl: Optional[str] = Field( + skip_summarization: Optional[bool] = Field( default=None, - description="""Optional. Input only. The TTL for this session. The minimum value is 24 hours.""", + description="""Optional. If true, it won't call model to summarize function response. Only used for function_response event.""", ) - update_time: Optional[datetime.datetime] = Field( + state_delta: Optional[dict[str, Any]] = Field( default=None, - description="""Output only. Timestamp when the session was updated.""", + description="""Optional. Indicates that the event is updating the state with the given delta.""", ) - user_id: Optional[str] = Field( + transfer_agent: Optional[str] = Field( default=None, - description="""Required. Immutable. String id provided by the user""", + description="""Optional. If set, the event transfers to the specified agent.""", ) -class SessionDict(TypedDict, total=False): - """A session.""" - - create_time: Optional[datetime.datetime] - """Output only. Timestamp when the session was created.""" - - display_name: Optional[str] - """Optional. The display name of the session.""" - - expire_time: Optional[datetime.datetime] - """Optional. Timestamp of when this session is considered expired. This is *always* provided on output, regardless of what was sent on input. The minimum value is 24 hours from the time of creation.""" +class EventActionsDict(TypedDict, total=False): + """Actions are parts of events that are executed by the agent.""" - labels: Optional[dict[str, str]] - """The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""" + artifact_delta: Optional[dict[str, int]] + """Optional. Indicates that the event is updating an artifact. key is the filename, value is the version.""" - name: Optional[str] - """Identifier. The resource name of the session. Format: 'projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}'.""" + escalate: Optional[bool] + """Optional. The agent is escalating to a higher level agent.""" - session_state: Optional[dict[str, Any]] - """Optional. Session specific memory which stores key conversation points.""" + requested_auth_configs: Optional[dict[str, Any]] + """Optional. Will only be set by a tool response indicating tool request euc. Struct key is the function call id since one function call response (from model) could correspond to multiple function calls. Struct value is the required auth config, which can be another struct.""" - ttl: Optional[str] - """Optional. Input only. The TTL for this session. The minimum value is 24 hours.""" + skip_summarization: Optional[bool] + """Optional. If true, it won't call model to summarize function response. Only used for function_response event.""" - update_time: Optional[datetime.datetime] - """Output only. Timestamp when the session was updated.""" + state_delta: Optional[dict[str, Any]] + """Optional. Indicates that the event is updating the state with the given delta.""" - user_id: Optional[str] - """Required. Immutable. String id provided by the user""" + transfer_agent: Optional[str] + """Optional. If set, the event transfers to the specified agent.""" -SessionOrDict = Union[Session, SessionDict] +EventActionsOrDict = Union[EventActions, EventActionsDict] -class AgentEngineSessionOperation(_common.BaseModel): - """Operation that has an agent engine session as a response.""" +class EventMetadata(_common.BaseModel): + """Metadata relating to a LLM response event.""" - name: Optional[str] = Field( + grounding_metadata: Optional[genai_types.GroundingMetadata] = Field( default=None, - description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + description="""Optional. Metadata returned to client when grounding is enabled.""", ) - metadata: Optional[dict[str, Any]] = Field( + branch: Optional[str] = Field( default=None, - description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + description="""Optional. The branch of the event. The format is like agent_1.agent_2.agent_3, where agent_1 is the parent of agent_2, and agent_2 is the parent of agent_3. Branch is used when multiple child agents shouldn't see their siblings' conversation history.""", ) - done: Optional[bool] = Field( + custom_metadata: Optional[dict[str, Any]] = Field( + default=None, description="""The custom metadata of the LlmResponse.""" + ) + interrupted: Optional[bool] = Field( default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + description="""Optional. Flag indicating that LLM was interrupted when generating the content. Usually it's due to user interruption during a bidi streaming.""", ) - error: Optional[dict[str, Any]] = Field( + long_running_tool_ids: Optional[list[str]] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Optional. Set of ids of the long running function calls. Agent client will know from this field about which function call is long running. Only valid for function call event.""", ) - response: Optional[Session] = Field( - default=None, description="""The Agent Engine Session.""" + partial: Optional[bool] = Field( + default=None, + description="""Optional. Indicates whether the text content is part of a unfinished text stream. Only used for streaming mode and when the content is plain text.""", + ) + turn_complete: Optional[bool] = Field( + default=None, + description="""Optional. Indicates whether the response from the model is complete. Only used for streaming mode.""", + ) + input_transcription: Optional[genai_types.Transcription] = Field( + default=None, description="""Optional. Audio transcription of user input.""" + ) + output_transcription: Optional[genai_types.Transcription] = Field( + default=None, description="""Optional. Audio transcription of model output.""" ) -class AgentEngineSessionOperationDict(TypedDict, total=False): - """Operation that has an agent engine session as a response.""" +class EventMetadataDict(TypedDict, total=False): + """Metadata relating to a LLM response event.""" - name: Optional[str] - """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + grounding_metadata: Optional[genai_types.GroundingMetadata] + """Optional. Metadata returned to client when grounding is enabled.""" - metadata: Optional[dict[str, Any]] - """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + branch: Optional[str] + """Optional. The branch of the event. The format is like agent_1.agent_2.agent_3, where agent_1 is the parent of agent_2, and agent_2 is the parent of agent_3. Branch is used when multiple child agents shouldn't see their siblings' conversation history.""" - done: Optional[bool] - """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + custom_metadata: Optional[dict[str, Any]] + """The custom metadata of the LlmResponse.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + interrupted: Optional[bool] + """Optional. Flag indicating that LLM was interrupted when generating the content. Usually it's due to user interruption during a bidi streaming.""" - response: Optional[SessionDict] - """The Agent Engine Session.""" + long_running_tool_ids: Optional[list[str]] + """Optional. Set of ids of the long running function calls. Agent client will know from this field about which function call is long running. Only valid for function call event.""" + + partial: Optional[bool] + """Optional. Indicates whether the text content is part of a unfinished text stream. Only used for streaming mode and when the content is plain text.""" + turn_complete: Optional[bool] + """Optional. Indicates whether the response from the model is complete. Only used for streaming mode.""" -AgentEngineSessionOperationOrDict = Union[ - AgentEngineSessionOperation, AgentEngineSessionOperationDict -] + input_transcription: Optional[genai_types.Transcription] + """Optional. Audio transcription of user input.""" + + output_transcription: Optional[genai_types.Transcription] + """Optional. Audio transcription of model output.""" + + +EventMetadataOrDict = Union[EventMetadata, EventMetadataDict] -class DeleteAgentEngineSessionConfig(_common.BaseModel): - """Config for deleting an Agent Engine Session.""" +class AppendRuntimeSessionEventConfig(_common.BaseModel): + """Config for appending agent runtime session event.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) + content: Optional[genai_types.Content] = Field( + default=None, description="""The content of the session event.""" + ) + actions: Optional[EventActions] = Field( + default=None, + description="""Actions are parts of events that are related to the session event.""", + ) + error_code: Optional[str] = Field( + default=None, description="""The error code of the session event.""" + ) + error_message: Optional[str] = Field( + default=None, description="""The error message of the session event.""" + ) + event_metadata: Optional[EventMetadata] = Field( + default=None, description="""Metadata relating to the session event.""" + ) + raw_event: Optional[dict[str, Any]] = Field( + default=None, + description="""Weakly typed raw event data in proto struct format.""", + ) -class DeleteAgentEngineSessionConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Session.""" +class AppendRuntimeSessionEventConfigDict(TypedDict, total=False): + """Config for appending agent runtime session event.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" + content: Optional[genai_types.Content] + """The content of the session event.""" -DeleteAgentEngineSessionConfigOrDict = Union[ - DeleteAgentEngineSessionConfig, DeleteAgentEngineSessionConfigDict -] - - -class _DeleteAgentEngineSessionRequestParameters(_common.BaseModel): - """Parameters for deleting agent engine sessions.""" - - name: Optional[str] = Field( - default=None, description="""Name of the agent engine session to delete.""" - ) - config: Optional[DeleteAgentEngineSessionConfig] = Field( - default=None, description="""""" - ) + actions: Optional[EventActionsDict] + """Actions are parts of events that are related to the session event.""" + error_code: Optional[str] + """The error code of the session event.""" -class _DeleteAgentEngineSessionRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engine sessions.""" + error_message: Optional[str] + """The error message of the session event.""" - name: Optional[str] - """Name of the agent engine session to delete.""" + event_metadata: Optional[EventMetadataDict] + """Metadata relating to the session event.""" - config: Optional[DeleteAgentEngineSessionConfigDict] - """""" + raw_event: Optional[dict[str, Any]] + """Weakly typed raw event data in proto struct format.""" -_DeleteAgentEngineSessionRequestParametersOrDict = Union[ - _DeleteAgentEngineSessionRequestParameters, - _DeleteAgentEngineSessionRequestParametersDict, +AppendRuntimeSessionEventConfigOrDict = Union[ + AppendRuntimeSessionEventConfig, AppendRuntimeSessionEventConfigDict ] -class DeleteAgentEngineSessionOperation(_common.BaseModel): - """Operation for deleting agent engine sessions.""" +class _AppendRuntimeSessionEventRequestParameters(_common.BaseModel): + """Parameters for appending agent runtimes.""" name: Optional[str] = Field( - default=None, - description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + default=None, description="""Name of the agent runtime session.""" ) - metadata: Optional[dict[str, Any]] = Field( - default=None, - description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + author: Optional[str] = Field( + default=None, description="""Author of the agent runtime session event.""" ) - done: Optional[bool] = Field( - default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + invocation_id: Optional[str] = Field( + default=None, description="""Invocation ID of the agent runtime.""" ) - error: Optional[dict[str, Any]] = Field( - default=None, - description="""The error result of the operation in case of failure or cancellation.""", + timestamp: Optional[datetime.datetime] = Field( + default=None, description="""Timestamp indicating when the event was created.""" + ) + config: Optional[AppendRuntimeSessionEventConfig] = Field( + default=None, description="""""" ) -class DeleteAgentEngineSessionOperationDict(TypedDict, total=False): - """Operation for deleting agent engine sessions.""" +class _AppendRuntimeSessionEventRequestParametersDict(TypedDict, total=False): + """Parameters for appending agent runtimes.""" name: Optional[str] - """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + """Name of the agent runtime session.""" - metadata: Optional[dict[str, Any]] - """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - - done: Optional[bool] - """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" - - -DeleteAgentEngineSessionOperationOrDict = Union[ - DeleteAgentEngineSessionOperation, DeleteAgentEngineSessionOperationDict -] - - -class GetAgentEngineSessionConfig(_common.BaseModel): - """Config for getting an Agent Engine Session.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + author: Optional[str] + """Author of the agent runtime session event.""" + invocation_id: Optional[str] + """Invocation ID of the agent runtime.""" -class GetAgentEngineSessionConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Session.""" + timestamp: Optional[datetime.datetime] + """Timestamp indicating when the event was created.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + config: Optional[AppendRuntimeSessionEventConfigDict] + """""" -GetAgentEngineSessionConfigOrDict = Union[ - GetAgentEngineSessionConfig, GetAgentEngineSessionConfigDict +_AppendRuntimeSessionEventRequestParametersOrDict = Union[ + _AppendRuntimeSessionEventRequestParameters, + _AppendRuntimeSessionEventRequestParametersDict, ] -class _GetAgentEngineSessionRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine session.""" - - name: Optional[str] = Field( - default=None, description="""Name of the agent engine session.""" - ) - config: Optional[GetAgentEngineSessionConfig] = Field( - default=None, description="""""" - ) +class AppendRuntimeSessionEventResponse(_common.BaseModel): + """Response for appending agent runtime session event.""" + pass -class _GetAgentEngineSessionRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine session.""" - name: Optional[str] - """Name of the agent engine session.""" +class AppendRuntimeSessionEventResponseDict(TypedDict, total=False): + """Response for appending agent runtime session event.""" - config: Optional[GetAgentEngineSessionConfigDict] - """""" + pass -_GetAgentEngineSessionRequestParametersOrDict = Union[ - _GetAgentEngineSessionRequestParameters, _GetAgentEngineSessionRequestParametersDict +AppendRuntimeSessionEventResponseOrDict = Union[ + AppendRuntimeSessionEventResponse, AppendRuntimeSessionEventResponseDict ] -class ListAgentEngineSessionsConfig(_common.BaseModel): - """Config for listing agent engine sessions.""" +class ListRuntimeSessionEventsConfig(_common.BaseModel): + """Config for listing agent runtime session events.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -18391,8 +17917,8 @@ class ListAgentEngineSessionsConfig(_common.BaseModel): ) -class ListAgentEngineSessionsConfigDict(TypedDict, total=False): - """Config for listing agent engine sessions.""" +class ListRuntimeSessionEventsConfigDict(TypedDict, total=False): + """Config for listing agent runtime session events.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -18408,874 +17934,586 @@ class ListAgentEngineSessionsConfigDict(TypedDict, total=False): For field names both snake_case and camelCase are supported.""" -ListAgentEngineSessionsConfigOrDict = Union[ - ListAgentEngineSessionsConfig, ListAgentEngineSessionsConfigDict +ListRuntimeSessionEventsConfigOrDict = Union[ + ListRuntimeSessionEventsConfig, ListRuntimeSessionEventsConfigDict ] -class _ListAgentEngineSessionsRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" +class _ListRuntimeSessionEventsRequestParameters(_common.BaseModel): + """Parameters for listing agent runtime session events.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime session.""" ) - config: Optional[ListAgentEngineSessionsConfig] = Field( + config: Optional[ListRuntimeSessionEventsConfig] = Field( default=None, description="""""" ) -class _ListAgentEngineSessionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" +class _ListRuntimeSessionEventsRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtime session events.""" name: Optional[str] - """Name of the agent engine.""" - - config: Optional[ListAgentEngineSessionsConfigDict] - """""" - - -_ListAgentEngineSessionsRequestParametersOrDict = Union[ - _ListAgentEngineSessionsRequestParameters, - _ListAgentEngineSessionsRequestParametersDict, -] - - -class ListReasoningEnginesSessionsResponse(_common.BaseModel): - """Response for listing agent engine sessions.""" - - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" - ) - next_page_token: Optional[str] = Field(default=None, description="""""") - sessions: Optional[list[Session]] = Field( - default=None, description="""List of agent engine sessions.""" - ) + """Name of the agent runtime session.""" - -class ListReasoningEnginesSessionsResponseDict(TypedDict, total=False): - """Response for listing agent engine sessions.""" - - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" - - next_page_token: Optional[str] + config: Optional[ListRuntimeSessionEventsConfigDict] """""" - sessions: Optional[list[SessionDict]] - """List of agent engine sessions.""" - - -ListReasoningEnginesSessionsResponseOrDict = Union[ - ListReasoningEnginesSessionsResponse, ListReasoningEnginesSessionsResponseDict -] - - -class _GetAgentEngineSessionOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a session as a response.""" - - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" - ) - config: Optional[GetAgentEngineOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" - ) - - -class _GetAgentEngineSessionOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a session as a response.""" - - operation_name: Optional[str] - """The server-assigned name for the operation.""" - config: Optional[GetAgentEngineOperationConfigDict] - """Used to override the default configuration.""" - - -_GetAgentEngineSessionOperationParametersOrDict = Union[ - _GetAgentEngineSessionOperationParameters, - _GetAgentEngineSessionOperationParametersDict, +_ListRuntimeSessionEventsRequestParametersOrDict = Union[ + _ListRuntimeSessionEventsRequestParameters, + _ListRuntimeSessionEventsRequestParametersDict, ] -class UpdateAgentEngineSessionConfig(_common.BaseModel): - """Config for updating agent engine session.""" +class SessionEvent(_common.BaseModel): + """A session event.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + content: Optional[genai_types.Content] = Field( + default=None, + description="""Optional. Content of the event provided by the author.""", ) - display_name: Optional[str] = Field( - default=None, description="""The display name of the session.""" + actions: Optional[EventActions] = Field( + default=None, description="""Optional. Actions executed by the agent.""" ) - session_state: Optional[dict[str, Any]] = Field( + author: Optional[str] = Field( default=None, - description="""Session state which stores key conversation points.""", - ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", + description="""Required. The name of the agent that sent the event, or user.""", ) - ttl: Optional[str] = Field( + error_code: Optional[str] = Field( default=None, - description="""Optional. Input only. The TTL for this resource. - - The expiration time is computed: now + TTL.""", + description="""Optional. Error code if the response is an error. Code varies by model.""", ) - expire_time: Optional[datetime.datetime] = Field( + error_message: Optional[str] = Field( default=None, - description="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", + description="""Optional. Error message if the response is an error.""", ) - labels: Optional[dict[str, str]] = Field( + event_metadata: Optional[EventMetadata] = Field( + default=None, description="""Optional. Metadata relating to this event.""" + ) + invocation_id: Optional[str] = Field( default=None, - description="""Optional. The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", + description="""Required. The invocation id of the event, multiple events can have the same invocation id.""", ) - session_id: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Optional. The user defined ID to use for session, which will become the final component of the session resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""", + description="""Identifier. The resource name of the event. Format:`projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}/events/{event}`.""", ) - update_mask: Optional[str] = Field( + timestamp: Optional[datetime.datetime] = Field( default=None, - description="""The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""", + description="""Required. Timestamp when the event was created on client side.""", ) - user_id: Optional[str] = Field( - default=None, description="""User ID of the agent engine session to update.""" + raw_event: Optional[dict[str, Any]] = Field( + default=None, + description="""Optional. Weakly typed raw event data in proto struct format.""", ) -class UpdateAgentEngineSessionConfigDict(TypedDict, total=False): - """Config for updating agent engine session.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class SessionEventDict(TypedDict, total=False): + """A session event.""" - display_name: Optional[str] - """The display name of the session.""" + content: Optional[genai_types.Content] + """Optional. Content of the event provided by the author.""" - session_state: Optional[dict[str, Any]] - """Session state which stores key conversation points.""" + actions: Optional[EventActionsDict] + """Optional. Actions executed by the agent.""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + author: Optional[str] + """Required. The name of the agent that sent the event, or user.""" - ttl: Optional[str] - """Optional. Input only. The TTL for this resource. + error_code: Optional[str] + """Optional. Error code if the response is an error. Code varies by model.""" - The expiration time is computed: now + TTL.""" + error_message: Optional[str] + """Optional. Error message if the response is an error.""" - expire_time: Optional[datetime.datetime] - """Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""" + event_metadata: Optional[EventMetadataDict] + """Optional. Metadata relating to this event.""" - labels: Optional[dict[str, str]] - """Optional. The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""" + invocation_id: Optional[str] + """Required. The invocation id of the event, multiple events can have the same invocation id.""" - session_id: Optional[str] - """Optional. The user defined ID to use for session, which will become the final component of the session resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""" + name: Optional[str] + """Identifier. The resource name of the event. Format:`projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}/events/{event}`.""" - update_mask: Optional[str] - """The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" + timestamp: Optional[datetime.datetime] + """Required. Timestamp when the event was created on client side.""" - user_id: Optional[str] - """User ID of the agent engine session to update.""" + raw_event: Optional[dict[str, Any]] + """Optional. Weakly typed raw event data in proto struct format.""" -UpdateAgentEngineSessionConfigOrDict = Union[ - UpdateAgentEngineSessionConfig, UpdateAgentEngineSessionConfigDict -] +SessionEventOrDict = Union[SessionEvent, SessionEventDict] -class _UpdateAgentEngineSessionRequestParameters(_common.BaseModel): - """Parameters for updating agent engine sessions.""" +class ListRuntimeSessionEventsResponse(_common.BaseModel): + """Response for listing agent runtime session events.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine session to update.""" - ) - config: Optional[UpdateAgentEngineSessionConfig] = Field( - default=None, description="""""" + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" + ) + next_page_token: Optional[str] = Field(default=None, description="""""") + session_events: Optional[list[SessionEvent]] = Field( + default=None, description="""List of session events.""" ) -class _UpdateAgentEngineSessionRequestParametersDict(TypedDict, total=False): - """Parameters for updating agent engine sessions.""" +class ListRuntimeSessionEventsResponseDict(TypedDict, total=False): + """Response for listing agent runtime session events.""" - name: Optional[str] - """Name of the agent engine session to update.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - config: Optional[UpdateAgentEngineSessionConfigDict] + next_page_token: Optional[str] """""" + session_events: Optional[list[SessionEventDict]] + """List of session events.""" + -_UpdateAgentEngineSessionRequestParametersOrDict = Union[ - _UpdateAgentEngineSessionRequestParameters, - _UpdateAgentEngineSessionRequestParametersDict, +ListRuntimeSessionEventsResponseOrDict = Union[ + ListRuntimeSessionEventsResponse, ListRuntimeSessionEventsResponseDict ] -class EventActions(_common.BaseModel): - """Actions are parts of events that are executed by the agent.""" +class GeminiExample(_common.BaseModel): + """Represents a Gemini example.""" - artifact_delta: Optional[dict[str, int]] = Field( - default=None, - description="""Optional. Indicates that the event is updating an artifact. key is the filename, value is the version.""", - ) - escalate: Optional[bool] = Field( - default=None, - description="""Optional. The agent is escalating to a higher level agent.""", - ) - requested_auth_configs: Optional[dict[str, Any]] = Field( - default=None, - description="""Optional. Will only be set by a tool response indicating tool request euc. Struct key is the function call id since one function call response (from model) could correspond to multiple function calls. Struct value is the required auth config, which can be another struct.""", - ) - skip_summarization: Optional[bool] = Field( - default=None, - description="""Optional. If true, it won't call model to summarize function response. Only used for function_response event.""", - ) - state_delta: Optional[dict[str, Any]] = Field( - default=None, - description="""Optional. Indicates that the event is updating the state with the given delta.""", + model: Optional[str] = Field( + default=None, description="""The model used to generate the Gemini example.""" ) - transfer_agent: Optional[str] = Field( - default=None, - description="""Optional. If set, the event transfers to the specified agent.""", + contents: Optional[list[genai_types.Content]] = Field( + default=None, description="""Contents of the Gemini example.""" ) - - -class EventActionsDict(TypedDict, total=False): - """Actions are parts of events that are executed by the agent.""" - - artifact_delta: Optional[dict[str, int]] - """Optional. Indicates that the event is updating an artifact. key is the filename, value is the version.""" - - escalate: Optional[bool] - """Optional. The agent is escalating to a higher level agent.""" - - requested_auth_configs: Optional[dict[str, Any]] - """Optional. Will only be set by a tool response indicating tool request euc. Struct key is the function call id since one function call response (from model) could correspond to multiple function calls. Struct value is the required auth config, which can be another struct.""" - - skip_summarization: Optional[bool] - """Optional. If true, it won't call model to summarize function response. Only used for function_response event.""" - - state_delta: Optional[dict[str, Any]] - """Optional. Indicates that the event is updating the state with the given delta.""" - - transfer_agent: Optional[str] - """Optional. If set, the event transfers to the specified agent.""" - - -EventActionsOrDict = Union[EventActions, EventActionsDict] - - -class EventMetadata(_common.BaseModel): - """Metadata relating to a LLM response event.""" - - grounding_metadata: Optional[genai_types.GroundingMetadata] = Field( - default=None, - description="""Optional. Metadata returned to client when grounding is enabled.""", + system_instruction: Optional[genai_types.Content] = Field( + default=None, description="""System instruction for the Gemini example.""" ) - branch: Optional[str] = Field( - default=None, - description="""Optional. The branch of the event. The format is like agent_1.agent_2.agent_3, where agent_1 is the parent of agent_2, and agent_2 is the parent of agent_3. Branch is used when multiple child agents shouldn't see their siblings' conversation history.""", + cached_content: Optional[str] = Field( + default=None, description="""Cached content for the Gemini example.""" ) - custom_metadata: Optional[dict[str, Any]] = Field( - default=None, description="""The custom metadata of the LlmResponse.""" + tools: Optional[list[genai_types.Tool]] = Field( + default=None, description="""Tools for the Gemini example.""" ) - interrupted: Optional[bool] = Field( - default=None, - description="""Optional. Flag indicating that LLM was interrupted when generating the content. Usually it's due to user interruption during a bidi streaming.""", + tool_config: Optional[genai_types.ToolConfig] = Field( + default=None, description="""Tools for the Gemini example.""" ) - long_running_tool_ids: Optional[list[str]] = Field( - default=None, - description="""Optional. Set of ids of the long running function calls. Agent client will know from this field about which function call is long running. Only valid for function call event.""", + safety_settings: Optional[list[genai_types.SafetySetting]] = Field( + default=None, description="""Safety settings for the Gemini example.""" ) - partial: Optional[bool] = Field( - default=None, - description="""Optional. Indicates whether the text content is part of a unfinished text stream. Only used for streaming mode and when the content is plain text.""", + generation_config: Optional[genai_types.GenerationConfig] = Field( + default=None, description="""Generation config for the Gemini example.""" ) - turn_complete: Optional[bool] = Field( + model_armor_config: Optional[genai_types.ModelArmorConfig] = Field( default=None, - description="""Optional. Indicates whether the response from the model is complete. Only used for streaming mode.""", - ) - input_transcription: Optional[genai_types.Transcription] = Field( - default=None, description="""Optional. Audio transcription of user input.""" - ) - output_transcription: Optional[genai_types.Transcription] = Field( - default=None, description="""Optional. Audio transcription of model output.""" + description="""Optional. Settings for prompt and response sanitization using the Model Armor service. If supplied, safety_settings must not be supplied.""", ) -class EventMetadataDict(TypedDict, total=False): - """Metadata relating to a LLM response event.""" +class GeminiExampleDict(TypedDict, total=False): + """Represents a Gemini example.""" - grounding_metadata: Optional[genai_types.GroundingMetadata] - """Optional. Metadata returned to client when grounding is enabled.""" + model: Optional[str] + """The model used to generate the Gemini example.""" - branch: Optional[str] - """Optional. The branch of the event. The format is like agent_1.agent_2.agent_3, where agent_1 is the parent of agent_2, and agent_2 is the parent of agent_3. Branch is used when multiple child agents shouldn't see their siblings' conversation history.""" + contents: Optional[list[genai_types.Content]] + """Contents of the Gemini example.""" - custom_metadata: Optional[dict[str, Any]] - """The custom metadata of the LlmResponse.""" + system_instruction: Optional[genai_types.Content] + """System instruction for the Gemini example.""" - interrupted: Optional[bool] - """Optional. Flag indicating that LLM was interrupted when generating the content. Usually it's due to user interruption during a bidi streaming.""" + cached_content: Optional[str] + """Cached content for the Gemini example.""" - long_running_tool_ids: Optional[list[str]] - """Optional. Set of ids of the long running function calls. Agent client will know from this field about which function call is long running. Only valid for function call event.""" + tools: Optional[list[genai_types.Tool]] + """Tools for the Gemini example.""" - partial: Optional[bool] - """Optional. Indicates whether the text content is part of a unfinished text stream. Only used for streaming mode and when the content is plain text.""" + tool_config: Optional[genai_types.ToolConfig] + """Tools for the Gemini example.""" - turn_complete: Optional[bool] - """Optional. Indicates whether the response from the model is complete. Only used for streaming mode.""" + safety_settings: Optional[list[genai_types.SafetySetting]] + """Safety settings for the Gemini example.""" - input_transcription: Optional[genai_types.Transcription] - """Optional. Audio transcription of user input.""" + generation_config: Optional[genai_types.GenerationConfig] + """Generation config for the Gemini example.""" - output_transcription: Optional[genai_types.Transcription] - """Optional. Audio transcription of model output.""" + model_armor_config: Optional[genai_types.ModelArmorConfig] + """Optional. Settings for prompt and response sanitization using the Model Armor service. If supplied, safety_settings must not be supplied.""" -EventMetadataOrDict = Union[EventMetadata, EventMetadataDict] +GeminiExampleOrDict = Union[GeminiExample, GeminiExampleDict] -class AppendAgentEngineSessionEventConfig(_common.BaseModel): - """Config for appending agent engine session event.""" +class GeminiTemplateConfig(_common.BaseModel): + """Represents a Gemini template config.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - content: Optional[genai_types.Content] = Field( - default=None, description="""The content of the session event.""" - ) - actions: Optional[EventActions] = Field( + gemini_example: Optional[GeminiExample] = Field( default=None, - description="""Actions are parts of events that are related to the session event.""", - ) - error_code: Optional[str] = Field( - default=None, description="""The error code of the session event.""" - ) - error_message: Optional[str] = Field( - default=None, description="""The error message of the session event.""" - ) - event_metadata: Optional[EventMetadata] = Field( - default=None, description="""Metadata relating to the session event.""" + description="""Required. The template that will be used for assembling the request to use for downstream applications.""", ) - raw_event: Optional[dict[str, Any]] = Field( + field_mapping: Optional[dict[str, str]] = Field( default=None, - description="""Weakly typed raw event data in proto struct format.""", + description="""Required. Map of template parameters to the columns in the dataset table.""", ) -class AppendAgentEngineSessionEventConfigDict(TypedDict, total=False): - """Config for appending agent engine session event.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - content: Optional[genai_types.Content] - """The content of the session event.""" - - actions: Optional[EventActionsDict] - """Actions are parts of events that are related to the session event.""" - - error_code: Optional[str] - """The error code of the session event.""" - - error_message: Optional[str] - """The error message of the session event.""" +class GeminiTemplateConfigDict(TypedDict, total=False): + """Represents a Gemini template config.""" - event_metadata: Optional[EventMetadataDict] - """Metadata relating to the session event.""" + gemini_example: Optional[GeminiExampleDict] + """Required. The template that will be used for assembling the request to use for downstream applications.""" - raw_event: Optional[dict[str, Any]] - """Weakly typed raw event data in proto struct format.""" + field_mapping: Optional[dict[str, str]] + """Required. Map of template parameters to the columns in the dataset table.""" -AppendAgentEngineSessionEventConfigOrDict = Union[ - AppendAgentEngineSessionEventConfig, AppendAgentEngineSessionEventConfigDict -] +GeminiTemplateConfigOrDict = Union[GeminiTemplateConfig, GeminiTemplateConfigDict] -class _AppendAgentEngineSessionEventRequestParameters(_common.BaseModel): - """Parameters for appending agent engines.""" +class GeminiRequestReadConfig(_common.BaseModel): + """Represents the config for reading Gemini requests.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine session.""" - ) - author: Optional[str] = Field( - default=None, description="""Author of the agent engine session event.""" - ) - invocation_id: Optional[str] = Field( - default=None, description="""Invocation ID of the agent engine.""" - ) - timestamp: Optional[datetime.datetime] = Field( - default=None, description="""Timestamp indicating when the event was created.""" + template_config: Optional[GeminiTemplateConfig] = Field( + default=None, description="""Gemini request template with placeholders.""" ) - config: Optional[AppendAgentEngineSessionEventConfig] = Field( - default=None, description="""""" + assembled_request_column_name: Optional[str] = Field( + default=None, + description="""Column name in the underlying BigQuery table that contains already fully assembled Gemini requests.""", ) + @classmethod + def single_turn_template( + cls, + *, + prompt: str, + response: Optional[str] = None, + system_instruction: Optional[str] = None, + model: Optional[str] = None, + cached_content: Optional[str] = None, + tools: Optional[list[Union[genai_types.Tool, dict[str, Any]]]] = None, + tool_config: Optional[Union[genai_types.ToolConfig, dict[str, Any]]] = None, + safety_settings: Optional[ + list[Union[genai_types.SafetySetting, dict[str, Any]]] + ] = None, + generation_config: Optional[ + Union[genai_types.GenerationConfig, dict[str, Any]] + ] = None, + field_mapping: Optional[dict[str, str]] = None, + ) -> "GeminiRequestReadConfig": + """Constructs a GeminiRequestReadConfig object for single-turn cases. -class _AppendAgentEngineSessionEventRequestParametersDict(TypedDict, total=False): - """Parameters for appending agent engines.""" - - name: Optional[str] - """Name of the agent engine session.""" - - author: Optional[str] - """Author of the agent engine session event.""" - - invocation_id: Optional[str] - """Invocation ID of the agent engine.""" - - timestamp: Optional[datetime.datetime] - """Timestamp indicating when the event was created.""" - - config: Optional[AppendAgentEngineSessionEventConfigDict] - """""" + Example: + read_config = GeminiRequestReadConfig.single_turn_template( + prompt="Which flower is this {flower_image}?", + response="This is a {label}.", + system_instruction="You are a botanical classifier." + ) + Args: + prompt: Required. User input. + response: Optional. Model response to user input. + system_instruction: Optional. System instructions for the model. + model: Optional. The model to use for the GeminiExample. + cached_content: Optional. The cached content to use for the GeminiExample. + tools: Optional. The tools to use for the GeminiExample. + tool_config: Optional. The tool config to use for the GeminiExample. + safety_settings: Optional. The safety settings to use for the GeminiExample. + generation_config: Optional. The generation config to use for the GeminiExample. + field_mapping: Optional. Mapping of placeholders to dataset columns. -_AppendAgentEngineSessionEventRequestParametersOrDict = Union[ - _AppendAgentEngineSessionEventRequestParameters, - _AppendAgentEngineSessionEventRequestParametersDict, -] + Returns: + A GeminiRequestReadConfig object. + """ + contents = [] + contents.append( + genai_types.Content( + role="user", + parts=[ + genai_types.Part.from_text(text=prompt), + ], + ) + ) + if response: + contents.append( + genai_types.Content( + role="model", + parts=[ + genai_types.Part.from_text(text=response), + ], + ) + ) + system_instruction_content = None + if system_instruction: + system_instruction_content = genai_types.Content( + parts=[ + genai_types.Part.from_text(text=system_instruction), + ], + ) -class AppendAgentEngineSessionEventResponse(_common.BaseModel): - """Response for appending agent engine session event.""" + return cls( + template_config=GeminiTemplateConfig( + gemini_example=GeminiExample( + model=model, + contents=contents, + system_instruction=system_instruction_content, + cached_content=cached_content, + tools=tools, + tool_config=tool_config, + safety_settings=safety_settings, + generation_config=generation_config, + ), + field_mapping=field_mapping, + ), + ) - pass +class GeminiRequestReadConfigDict(TypedDict, total=False): + """Represents the config for reading Gemini requests.""" -class AppendAgentEngineSessionEventResponseDict(TypedDict, total=False): - """Response for appending agent engine session event.""" + template_config: Optional[GeminiTemplateConfigDict] + """Gemini request template with placeholders.""" - pass + assembled_request_column_name: Optional[str] + """Column name in the underlying BigQuery table that contains already fully assembled Gemini requests.""" -AppendAgentEngineSessionEventResponseOrDict = Union[ - AppendAgentEngineSessionEventResponse, AppendAgentEngineSessionEventResponseDict +GeminiRequestReadConfigOrDict = Union[ + GeminiRequestReadConfig, GeminiRequestReadConfigDict ] -class ListAgentEngineSessionEventsConfig(_common.BaseModel): - """Config for listing agent engine session events.""" +class AssembleDatasetConfig(_common.BaseModel): + """Config for assembling a multimodal dataset resource.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( - default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""", + timeout: Optional[int] = Field( + default=90, + description="""The timeout for the assemble dataset request in seconds. If not + set, the default timeout is 90 seconds.""", ) -class ListAgentEngineSessionEventsConfigDict(TypedDict, total=False): - """Config for listing agent engine session events.""" +class AssembleDatasetConfigDict(TypedDict, total=False): + """Config for assembling a multimodal dataset resource.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - page_size: Optional[int] - """""" - - page_token: Optional[str] - """""" - - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" + timeout: Optional[int] + """The timeout for the assemble dataset request in seconds. If not + set, the default timeout is 90 seconds.""" -ListAgentEngineSessionEventsConfigOrDict = Union[ - ListAgentEngineSessionEventsConfig, ListAgentEngineSessionEventsConfigDict -] +AssembleDatasetConfigOrDict = Union[AssembleDatasetConfig, AssembleDatasetConfigDict] -class _ListAgentEngineSessionEventsRequestParameters(_common.BaseModel): - """Parameters for listing agent engine session events.""" +class _AssembleDatasetParameters(_common.BaseModel): + """Parameters for assembling a multimodal dataset resource.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine session.""" - ) - config: Optional[ListAgentEngineSessionEventsConfig] = Field( + name: Optional[str] = Field(default=None, description="""""") + gemini_request_read_config: Optional[GeminiRequestReadConfig] = Field( default=None, description="""""" ) + config: Optional[AssembleDatasetConfig] = Field(default=None, description="""""") -class _ListAgentEngineSessionEventsRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engine session events.""" +class _AssembleDatasetParametersDict(TypedDict, total=False): + """Parameters for assembling a multimodal dataset resource.""" name: Optional[str] - """Name of the agent engine session.""" + """""" + + gemini_request_read_config: Optional[GeminiRequestReadConfigDict] + """""" - config: Optional[ListAgentEngineSessionEventsConfigDict] + config: Optional[AssembleDatasetConfigDict] """""" -_ListAgentEngineSessionEventsRequestParametersOrDict = Union[ - _ListAgentEngineSessionEventsRequestParameters, - _ListAgentEngineSessionEventsRequestParametersDict, +_AssembleDatasetParametersOrDict = Union[ + _AssembleDatasetParameters, _AssembleDatasetParametersDict ] -class SessionEvent(_common.BaseModel): - """A session event.""" +class MultimodalDatasetOperation(_common.BaseModel): + """Represents the create dataset operation.""" - content: Optional[genai_types.Content] = Field( - default=None, - description="""Optional. Content of the event provided by the author.""", - ) - actions: Optional[EventActions] = Field( - default=None, description="""Optional. Actions executed by the agent.""" - ) - author: Optional[str] = Field( - default=None, - description="""Required. The name of the agent that sent the event, or user.""", - ) - error_code: Optional[str] = Field( - default=None, - description="""Optional. Error code if the response is an error. Code varies by model.""", - ) - error_message: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Optional. Error message if the response is an error.""", - ) - event_metadata: Optional[EventMetadata] = Field( - default=None, description="""Optional. Metadata relating to this event.""" + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - invocation_id: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Required. The invocation id of the event, multiple events can have the same invocation id.""", + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", ) - name: Optional[str] = Field( + done: Optional[bool] = Field( default=None, - description="""Identifier. The resource name of the event. Format:`projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}/events/{event}`.""", + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", ) - timestamp: Optional[datetime.datetime] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Required. Timestamp when the event was created on client side.""", + description="""The error result of the operation in case of failure or cancellation.""", ) - raw_event: Optional[dict[str, Any]] = Field( - default=None, - description="""Optional. Weakly typed raw event data in proto struct format.""", + response: Optional[dict[str, Any]] = Field( + default=None, description="""The result of the dataset operation.""" ) -class SessionEventDict(TypedDict, total=False): - """A session event.""" +class MultimodalDatasetOperationDict(TypedDict, total=False): + """Represents the create dataset operation.""" - content: Optional[genai_types.Content] - """Optional. Content of the event provided by the author.""" + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - actions: Optional[EventActionsDict] - """Optional. Actions executed by the agent.""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - author: Optional[str] - """Required. The name of the agent that sent the event, or user.""" + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - error_code: Optional[str] - """Optional. Error code if the response is an error. Code varies by model.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - error_message: Optional[str] - """Optional. Error message if the response is an error.""" + response: Optional[dict[str, Any]] + """The result of the dataset operation.""" - event_metadata: Optional[EventMetadataDict] - """Optional. Metadata relating to this event.""" - invocation_id: Optional[str] - """Required. The invocation id of the event, multiple events can have the same invocation id.""" +MultimodalDatasetOperationOrDict = Union[ + MultimodalDatasetOperation, MultimodalDatasetOperationDict +] - name: Optional[str] - """Identifier. The resource name of the event. Format:`projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}/events/{event}`.""" - timestamp: Optional[datetime.datetime] - """Required. Timestamp when the event was created on client side.""" +class TuningResourceUsageAssessmentConfig(_common.BaseModel): + """Config for tuning resource usage assessment.""" - raw_event: Optional[dict[str, Any]] - """Optional. Weakly typed raw event data in proto struct format.""" + model_name: Optional[str] = Field(default=None, description="""""") -SessionEventOrDict = Union[SessionEvent, SessionEventDict] +class TuningResourceUsageAssessmentConfigDict(TypedDict, total=False): + """Config for tuning resource usage assessment.""" + model_name: Optional[str] + """""" -class ListAgentEngineSessionEventsResponse(_common.BaseModel): - """Response for listing agent engine session events.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" - ) - next_page_token: Optional[str] = Field(default=None, description="""""") - session_events: Optional[list[SessionEvent]] = Field( - default=None, description="""List of session events.""" - ) +TuningResourceUsageAssessmentConfigOrDict = Union[ + TuningResourceUsageAssessmentConfig, TuningResourceUsageAssessmentConfigDict +] -class ListAgentEngineSessionEventsResponseDict(TypedDict, total=False): - """Response for listing agent engine session events.""" +class TuningValidationAssessmentConfig(_common.BaseModel): + """Config for tuning validation assessment.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + model_name: Optional[str] = Field(default=None, description="""""") + dataset_usage: Optional[str] = Field(default=None, description="""""") - next_page_token: Optional[str] + +class TuningValidationAssessmentConfigDict(TypedDict, total=False): + """Config for tuning validation assessment.""" + + model_name: Optional[str] """""" - session_events: Optional[list[SessionEventDict]] - """List of session events.""" + dataset_usage: Optional[str] + """""" -ListAgentEngineSessionEventsResponseOrDict = Union[ - ListAgentEngineSessionEventsResponse, ListAgentEngineSessionEventsResponseDict +TuningValidationAssessmentConfigOrDict = Union[ + TuningValidationAssessmentConfig, TuningValidationAssessmentConfigDict ] -class GeminiExample(_common.BaseModel): - """Represents a Gemini example.""" +class BatchPredictionResourceUsageAssessmentConfig(_common.BaseModel): + """Config for batch prediction resource usage assessment.""" - model: Optional[str] = Field( - default=None, description="""The model used to generate the Gemini example.""" - ) - contents: Optional[list[genai_types.Content]] = Field( - default=None, description="""Contents of the Gemini example.""" - ) - system_instruction: Optional[genai_types.Content] = Field( - default=None, description="""System instruction for the Gemini example.""" - ) - cached_content: Optional[str] = Field( - default=None, description="""Cached content for the Gemini example.""" - ) - tools: Optional[list[genai_types.Tool]] = Field( - default=None, description="""Tools for the Gemini example.""" - ) - tool_config: Optional[genai_types.ToolConfig] = Field( - default=None, description="""Tools for the Gemini example.""" - ) - safety_settings: Optional[list[genai_types.SafetySetting]] = Field( - default=None, description="""Safety settings for the Gemini example.""" - ) - generation_config: Optional[genai_types.GenerationConfig] = Field( - default=None, description="""Generation config for the Gemini example.""" - ) - model_armor_config: Optional[genai_types.ModelArmorConfig] = Field( - default=None, - description="""Optional. Settings for prompt and response sanitization using the Model Armor service. If supplied, safety_settings must not be supplied.""", - ) - - -class GeminiExampleDict(TypedDict, total=False): - """Represents a Gemini example.""" - - model: Optional[str] - """The model used to generate the Gemini example.""" - - contents: Optional[list[genai_types.Content]] - """Contents of the Gemini example.""" - - system_instruction: Optional[genai_types.Content] - """System instruction for the Gemini example.""" - - cached_content: Optional[str] - """Cached content for the Gemini example.""" - - tools: Optional[list[genai_types.Tool]] - """Tools for the Gemini example.""" - - tool_config: Optional[genai_types.ToolConfig] - """Tools for the Gemini example.""" - - safety_settings: Optional[list[genai_types.SafetySetting]] - """Safety settings for the Gemini example.""" - - generation_config: Optional[genai_types.GenerationConfig] - """Generation config for the Gemini example.""" - - model_armor_config: Optional[genai_types.ModelArmorConfig] - """Optional. Settings for prompt and response sanitization using the Model Armor service. If supplied, safety_settings must not be supplied.""" - - -GeminiExampleOrDict = Union[GeminiExample, GeminiExampleDict] - - -class GeminiTemplateConfig(_common.BaseModel): - """Represents a Gemini template config.""" - - gemini_example: Optional[GeminiExample] = Field( - default=None, - description="""Required. The template that will be used for assembling the request to use for downstream applications.""", - ) - field_mapping: Optional[dict[str, str]] = Field( - default=None, - description="""Required. Map of template parameters to the columns in the dataset table.""", - ) - - -class GeminiTemplateConfigDict(TypedDict, total=False): - """Represents a Gemini template config.""" - - gemini_example: Optional[GeminiExampleDict] - """Required. The template that will be used for assembling the request to use for downstream applications.""" - - field_mapping: Optional[dict[str, str]] - """Required. Map of template parameters to the columns in the dataset table.""" - - -GeminiTemplateConfigOrDict = Union[GeminiTemplateConfig, GeminiTemplateConfigDict] - - -class GeminiRequestReadConfig(_common.BaseModel): - """Represents the config for reading Gemini requests.""" + model_name: Optional[str] = Field(default=None, description="""""") - template_config: Optional[GeminiTemplateConfig] = Field( - default=None, description="""Gemini request template with placeholders.""" - ) - assembled_request_column_name: Optional[str] = Field( - default=None, - description="""Column name in the underlying BigQuery table that contains already fully assembled Gemini requests.""", - ) - @classmethod - def single_turn_template( - cls, - *, - prompt: str, - response: Optional[str] = None, - system_instruction: Optional[str] = None, - model: Optional[str] = None, - cached_content: Optional[str] = None, - tools: Optional[list[Union[genai_types.Tool, dict[str, Any]]]] = None, - tool_config: Optional[Union[genai_types.ToolConfig, dict[str, Any]]] = None, - safety_settings: Optional[ - list[Union[genai_types.SafetySetting, dict[str, Any]]] - ] = None, - generation_config: Optional[ - Union[genai_types.GenerationConfig, dict[str, Any]] - ] = None, - field_mapping: Optional[dict[str, str]] = None, - ) -> "GeminiRequestReadConfig": - """Constructs a GeminiRequestReadConfig object for single-turn cases. +class BatchPredictionResourceUsageAssessmentConfigDict(TypedDict, total=False): + """Config for batch prediction resource usage assessment.""" - Example: - read_config = GeminiRequestReadConfig.single_turn_template( - prompt="Which flower is this {flower_image}?", - response="This is a {label}.", - system_instruction="You are a botanical classifier." - ) + model_name: Optional[str] + """""" - Args: - prompt: Required. User input. - response: Optional. Model response to user input. - system_instruction: Optional. System instructions for the model. - model: Optional. The model to use for the GeminiExample. - cached_content: Optional. The cached content to use for the GeminiExample. - tools: Optional. The tools to use for the GeminiExample. - tool_config: Optional. The tool config to use for the GeminiExample. - safety_settings: Optional. The safety settings to use for the GeminiExample. - generation_config: Optional. The generation config to use for the GeminiExample. - field_mapping: Optional. Mapping of placeholders to dataset columns. - Returns: - A GeminiRequestReadConfig object. - """ - contents = [] - contents.append( - genai_types.Content( - role="user", - parts=[ - genai_types.Part.from_text(text=prompt), - ], - ) - ) - if response: - contents.append( - genai_types.Content( - role="model", - parts=[ - genai_types.Part.from_text(text=response), - ], - ) - ) +BatchPredictionResourceUsageAssessmentConfigOrDict = Union[ + BatchPredictionResourceUsageAssessmentConfig, + BatchPredictionResourceUsageAssessmentConfigDict, +] - system_instruction_content = None - if system_instruction: - system_instruction_content = genai_types.Content( - parts=[ - genai_types.Part.from_text(text=system_instruction), - ], - ) - return cls( - template_config=GeminiTemplateConfig( - gemini_example=GeminiExample( - model=model, - contents=contents, - system_instruction=system_instruction_content, - cached_content=cached_content, - tools=tools, - tool_config=tool_config, - safety_settings=safety_settings, - generation_config=generation_config, - ), - field_mapping=field_mapping, - ), - ) +class BatchPredictionValidationAssessmentConfig(_common.BaseModel): + """Config for batch prediction validation assessment.""" + model_name: Optional[str] = Field(default=None, description="""""") -class GeminiRequestReadConfigDict(TypedDict, total=False): - """Represents the config for reading Gemini requests.""" - template_config: Optional[GeminiTemplateConfigDict] - """Gemini request template with placeholders.""" +class BatchPredictionValidationAssessmentConfigDict(TypedDict, total=False): + """Config for batch prediction validation assessment.""" - assembled_request_column_name: Optional[str] - """Column name in the underlying BigQuery table that contains already fully assembled Gemini requests.""" + model_name: Optional[str] + """""" -GeminiRequestReadConfigOrDict = Union[ - GeminiRequestReadConfig, GeminiRequestReadConfigDict +BatchPredictionValidationAssessmentConfigOrDict = Union[ + BatchPredictionValidationAssessmentConfig, + BatchPredictionValidationAssessmentConfigDict, ] -class AssembleDatasetConfig(_common.BaseModel): - """Config for assembling a multimodal dataset resource.""" +class AssessDatasetConfig(_common.BaseModel): + """Config for assessing a multimodal dataset resource.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) timeout: Optional[int] = Field( default=90, - description="""The timeout for the assemble dataset request in seconds. If not - set, the default timeout is 90 seconds.""", + description="""The timeout for the assess dataset request in seconds. If not set, + the default timeout is 90 seconds.""", ) -class AssembleDatasetConfigDict(TypedDict, total=False): - """Config for assembling a multimodal dataset resource.""" +class AssessDatasetConfigDict(TypedDict, total=False): + """Config for assessing a multimodal dataset resource.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" timeout: Optional[int] - """The timeout for the assemble dataset request in seconds. If not - set, the default timeout is 90 seconds.""" + """The timeout for the assess dataset request in seconds. If not set, + the default timeout is 90 seconds.""" -AssembleDatasetConfigOrDict = Union[AssembleDatasetConfig, AssembleDatasetConfigDict] +AssessDatasetConfigOrDict = Union[AssessDatasetConfig, AssessDatasetConfigDict] -class _AssembleDatasetParameters(_common.BaseModel): - """Parameters for assembling a multimodal dataset resource.""" +class _AssessDatasetParameters(_common.BaseModel): + """Parameters for assessing a multimodal dataset resource.""" name: Optional[str] = Field(default=None, description="""""") gemini_request_read_config: Optional[GeminiRequestReadConfig] = Field( default=None, description="""""" ) - config: Optional[AssembleDatasetConfig] = Field(default=None, description="""""") + tuning_resource_usage_assessment_config: Optional[ + TuningResourceUsageAssessmentConfig + ] = Field(default=None, description="""""") + tuning_validation_assessment_config: Optional[TuningValidationAssessmentConfig] = ( + Field(default=None, description="""""") + ) + batch_prediction_resource_usage_assessment_config: Optional[ + BatchPredictionResourceUsageAssessmentConfig + ] = Field(default=None, description="""""") + batch_prediction_validation_assessment_config: Optional[ + BatchPredictionValidationAssessmentConfig + ] = Field(default=None, description="""""") + config: Optional[AssessDatasetConfig] = Field(default=None, description="""""") -class _AssembleDatasetParametersDict(TypedDict, total=False): - """Parameters for assembling a multimodal dataset resource.""" +class _AssessDatasetParametersDict(TypedDict, total=False): + """Parameters for assessing a multimodal dataset resource.""" name: Optional[str] """""" @@ -19283,232 +18521,39 @@ class _AssembleDatasetParametersDict(TypedDict, total=False): gemini_request_read_config: Optional[GeminiRequestReadConfigDict] """""" - config: Optional[AssembleDatasetConfigDict] + tuning_resource_usage_assessment_config: Optional[ + TuningResourceUsageAssessmentConfigDict + ] """""" + tuning_validation_assessment_config: Optional[TuningValidationAssessmentConfigDict] + """""" -_AssembleDatasetParametersOrDict = Union[ - _AssembleDatasetParameters, _AssembleDatasetParametersDict + batch_prediction_resource_usage_assessment_config: Optional[ + BatchPredictionResourceUsageAssessmentConfigDict + ] + """""" + + batch_prediction_validation_assessment_config: Optional[ + BatchPredictionValidationAssessmentConfigDict + ] + """""" + + config: Optional[AssessDatasetConfigDict] + """""" + + +_AssessDatasetParametersOrDict = Union[ + _AssessDatasetParameters, _AssessDatasetParametersDict ] -class MultimodalDatasetOperation(_common.BaseModel): - """Represents the create dataset operation.""" +class SchemaTablesDatasetMetadataBigQuerySource(_common.BaseModel): + """Represents the BigQuery source for multimodal dataset metadata.""" - name: Optional[str] = Field( + uri: Optional[str] = Field( default=None, - description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, - description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", - ) - done: Optional[bool] = Field( - default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", - ) - error: Optional[dict[str, Any]] = Field( - default=None, - description="""The error result of the operation in case of failure or cancellation.""", - ) - response: Optional[dict[str, Any]] = Field( - default=None, description="""The result of the dataset operation.""" - ) - - -class MultimodalDatasetOperationDict(TypedDict, total=False): - """Represents the create dataset operation.""" - - name: Optional[str] - """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - - metadata: Optional[dict[str, Any]] - """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - - done: Optional[bool] - """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" - - response: Optional[dict[str, Any]] - """The result of the dataset operation.""" - - -MultimodalDatasetOperationOrDict = Union[ - MultimodalDatasetOperation, MultimodalDatasetOperationDict -] - - -class TuningResourceUsageAssessmentConfig(_common.BaseModel): - """Config for tuning resource usage assessment.""" - - model_name: Optional[str] = Field(default=None, description="""""") - - -class TuningResourceUsageAssessmentConfigDict(TypedDict, total=False): - """Config for tuning resource usage assessment.""" - - model_name: Optional[str] - """""" - - -TuningResourceUsageAssessmentConfigOrDict = Union[ - TuningResourceUsageAssessmentConfig, TuningResourceUsageAssessmentConfigDict -] - - -class TuningValidationAssessmentConfig(_common.BaseModel): - """Config for tuning validation assessment.""" - - model_name: Optional[str] = Field(default=None, description="""""") - dataset_usage: Optional[str] = Field(default=None, description="""""") - - -class TuningValidationAssessmentConfigDict(TypedDict, total=False): - """Config for tuning validation assessment.""" - - model_name: Optional[str] - """""" - - dataset_usage: Optional[str] - """""" - - -TuningValidationAssessmentConfigOrDict = Union[ - TuningValidationAssessmentConfig, TuningValidationAssessmentConfigDict -] - - -class BatchPredictionResourceUsageAssessmentConfig(_common.BaseModel): - """Config for batch prediction resource usage assessment.""" - - model_name: Optional[str] = Field(default=None, description="""""") - - -class BatchPredictionResourceUsageAssessmentConfigDict(TypedDict, total=False): - """Config for batch prediction resource usage assessment.""" - - model_name: Optional[str] - """""" - - -BatchPredictionResourceUsageAssessmentConfigOrDict = Union[ - BatchPredictionResourceUsageAssessmentConfig, - BatchPredictionResourceUsageAssessmentConfigDict, -] - - -class BatchPredictionValidationAssessmentConfig(_common.BaseModel): - """Config for batch prediction validation assessment.""" - - model_name: Optional[str] = Field(default=None, description="""""") - - -class BatchPredictionValidationAssessmentConfigDict(TypedDict, total=False): - """Config for batch prediction validation assessment.""" - - model_name: Optional[str] - """""" - - -BatchPredictionValidationAssessmentConfigOrDict = Union[ - BatchPredictionValidationAssessmentConfig, - BatchPredictionValidationAssessmentConfigDict, -] - - -class AssessDatasetConfig(_common.BaseModel): - """Config for assessing a multimodal dataset resource.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - timeout: Optional[int] = Field( - default=90, - description="""The timeout for the assess dataset request in seconds. If not set, - the default timeout is 90 seconds.""", - ) - - -class AssessDatasetConfigDict(TypedDict, total=False): - """Config for assessing a multimodal dataset resource.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - timeout: Optional[int] - """The timeout for the assess dataset request in seconds. If not set, - the default timeout is 90 seconds.""" - - -AssessDatasetConfigOrDict = Union[AssessDatasetConfig, AssessDatasetConfigDict] - - -class _AssessDatasetParameters(_common.BaseModel): - """Parameters for assessing a multimodal dataset resource.""" - - name: Optional[str] = Field(default=None, description="""""") - gemini_request_read_config: Optional[GeminiRequestReadConfig] = Field( - default=None, description="""""" - ) - tuning_resource_usage_assessment_config: Optional[ - TuningResourceUsageAssessmentConfig - ] = Field(default=None, description="""""") - tuning_validation_assessment_config: Optional[TuningValidationAssessmentConfig] = ( - Field(default=None, description="""""") - ) - batch_prediction_resource_usage_assessment_config: Optional[ - BatchPredictionResourceUsageAssessmentConfig - ] = Field(default=None, description="""""") - batch_prediction_validation_assessment_config: Optional[ - BatchPredictionValidationAssessmentConfig - ] = Field(default=None, description="""""") - config: Optional[AssessDatasetConfig] = Field(default=None, description="""""") - - -class _AssessDatasetParametersDict(TypedDict, total=False): - """Parameters for assessing a multimodal dataset resource.""" - - name: Optional[str] - """""" - - gemini_request_read_config: Optional[GeminiRequestReadConfigDict] - """""" - - tuning_resource_usage_assessment_config: Optional[ - TuningResourceUsageAssessmentConfigDict - ] - """""" - - tuning_validation_assessment_config: Optional[TuningValidationAssessmentConfigDict] - """""" - - batch_prediction_resource_usage_assessment_config: Optional[ - BatchPredictionResourceUsageAssessmentConfigDict - ] - """""" - - batch_prediction_validation_assessment_config: Optional[ - BatchPredictionValidationAssessmentConfigDict - ] - """""" - - config: Optional[AssessDatasetConfigDict] - """""" - - -_AssessDatasetParametersOrDict = Union[ - _AssessDatasetParameters, _AssessDatasetParametersDict -] - - -class SchemaTablesDatasetMetadataBigQuerySource(_common.BaseModel): - """Represents the BigQuery source for multimodal dataset metadata.""" - - uri: Optional[str] = Field( - default=None, - description="""The URI of the BigQuery table. This accepts the table name with or without the bq:// prefix.""", + description="""The URI of the BigQuery table. This accepts the table name with or without the bq:// prefix.""", ) @@ -26062,1437 +25107,1435 @@ class EvalRunInferenceConfigDict(TypedDict, total=False): EvalRunInferenceConfigOrDict = Union[EvalRunInferenceConfig, EvalRunInferenceConfigDict] -class AgentEngine(_common.BaseModel): - """An agent engine instance.""" +class AssembleDataset(_common.BaseModel): + """Represents the assembled dataset.""" - api_client: Optional[Any] = Field( - default=None, description="""The underlying API client.""" - ) - api_async_client: Optional[Any] = Field( - default=None, - description="""The underlying API client for asynchronous operations.""", - ) - api_resource: Optional[ReasoningEngine] = Field( + bigquery_destination: Optional[str] = Field( default=None, - description="""The underlying API resource (i.e. ReasoningEngine).""", + description="""The BigQuery destination of the assembled dataset.""", ) - # Allows dynamic binding of methods based on the registered operations. - model_config = ConfigDict(extra="allow") - def __repr__(self) -> str: - return ( - f"AgentEngine(api_resource.name='{self.api_resource.name}')" - if self.api_resource is not None - else "AgentEngine(api_resource.name=None)" - ) +class AssembleDatasetDict(TypedDict, total=False): + """Represents the assembled dataset.""" - def operation_schemas(self) -> Optional[list[Dict[str, Any]]]: - """Returns the schemas of all registered operations for the agent.""" - if not isinstance(self.api_resource, ReasoningEngine): - raise ValueError("api_resource is not initialized.") - if not self.api_resource.spec: - raise ValueError("api_resource.spec is not initialized.") - return self.api_resource.spec.class_methods + bigquery_destination: Optional[str] + """The BigQuery destination of the assembled dataset.""" - def delete( - self, - force: bool = False, - config: Optional[DeleteAgentEngineConfigOrDict] = None, - ) -> None: - """Deletes the agent engine. - Args: - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error when - the Agent Engine has undeleted child resources. Defaults to False. - config (DeleteAgentEngineConfig): - Optional. Additional configurations for deleting the Agent Engine. - """ - if not isinstance(self.api_resource, ReasoningEngine): - raise ValueError("api_resource is not initialized.") - self.api_client.delete(name=self.api_resource.name, force=force, config=config) # type: ignore[union-attr] +AssembleDatasetOrDict = Union[AssembleDataset, AssembleDatasetDict] -RubricContentProperty = evals_types.RubricContentProperty -RubricContentPropertyDict = evals_types.RubricContentPropertyDict -RubricContentPropertyDictOrDict = evals_types.RubricContentPropertyOrDict +class BatchPredictionResourceUsageAssessmentResult(_common.BaseModel): + """Result of batch prediction resource usage assessment.""" -RubricContent = evals_types.RubricContent -RubricContentDict = evals_types.RubricContentDict -RubricContentDictOrDict = evals_types.RubricContentOrDict + token_count: Optional[int] = Field( + default=None, description="""Number of tokens in the dataset.""" + ) + audio_token_count: Optional[int] = Field( + default=None, description="""Number of audio tokens in the dataset.""" + ) -Rubric = evals_types.Rubric -RubricDict = evals_types.RubricDict -RubricDictOrDict = evals_types.RubricOrDict -RubricVerdict = evals_types.RubricVerdict -RubricVerdictDict = evals_types.RubricVerdictDict -RubricVerdictDictOrDict = evals_types.RubricVerdictOrDict +class BatchPredictionResourceUsageAssessmentResultDict(TypedDict, total=False): + """Result of batch prediction resource usage assessment.""" -CandidateResult = evals_types.CandidateResult -CandidateResultDict = evals_types.CandidateResultDict -CandidateResultDictOrDict = evals_types.CandidateResultOrDict + token_count: Optional[int] + """Number of tokens in the dataset.""" -Event = evals_types.Event -EventDict = evals_types.EventDict -EventDictOrDict = evals_types.EventOrDict + audio_token_count: Optional[int] + """Number of audio tokens in the dataset.""" -Message = evals_types.Message -MessageDict = evals_types.MessageDict -MessageDictOrDict = evals_types.MessageOrDict -Importance = evals_types.Importance +BatchPredictionResourceUsageAssessmentResultOrDict = Union[ + BatchPredictionResourceUsageAssessmentResult, + BatchPredictionResourceUsageAssessmentResultDict, +] -class AgentEngineDict(TypedDict, total=False): - """An agent engine instance.""" +class BatchPredictionValidationAssessmentResult(_common.BaseModel): + """Result of batch prediction validation assessment.""" - api_client: Optional[Any] - """The underlying API client.""" + errors: Optional[list[str]] = Field( + default=None, description="""The list of errors found in the dataset.""" + ) - api_async_client: Optional[Any] - """The underlying API client for asynchronous operations.""" - api_resource: Optional[ReasoningEngineDict] - """The underlying API resource (i.e. ReasoningEngine).""" +class BatchPredictionValidationAssessmentResultDict(TypedDict, total=False): + """Result of batch prediction validation assessment.""" + errors: Optional[list[str]] + """The list of errors found in the dataset.""" -AgentEngineOrDict = Union[AgentEngine, AgentEngineDict] +BatchPredictionValidationAssessmentResultOrDict = Union[ + BatchPredictionValidationAssessmentResult, + BatchPredictionValidationAssessmentResultDict, +] -class AgentEngineConfig(_common.BaseModel): - """Config for agent engine methods.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - staging_bucket: Optional[str] = Field( - default=None, - description="""The GCS bucket to use for staging the artifacts needed. +class TuningResourceUsageAssessmentResult(_common.BaseModel): + """Result of tuning resource usage assessment.""" - It must be a valid GCS bucket name, e.g. "gs://bucket-name". It is - required if `agent_engine` is specified.""", + token_count: Optional[int] = Field( + default=None, description="""The number of tokens in the dataset.""" ) - requirements: Optional[Any] = Field( + billable_character_count: Optional[int] = Field( default=None, - description="""The set of PyPI dependencies needed. - - It can either be the path to a single file (requirements.txt), or an - ordered list of strings corresponding to each line of the requirements - file.""", + description="""The number of billable characters in the dataset.""", ) - display_name: Optional[str] = Field( - default=None, - description="""The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any UTF-8 - character.""", - ) - description: Optional[str] = Field( - default=None, description="""The description of the Agent Engine.""" - ) - gcs_dir_name: Optional[str] = Field( - default=None, - description="""The GCS bucket directory under `staging_bucket` to use for staging - the artifacts needed.""", - ) - extra_packages: Optional[list[str]] = Field( - default=None, - description="""The set of extra user-provided packages (if any).""", - ) - env_vars: Optional[Any] = Field( - default=None, - description="""The environment variables to be set when running the Agent Engine. - If it is a dictionary, the keys are the environment variable names, and - the values are the corresponding values.""", - ) - service_account: Optional[str] = Field( - default=None, - description="""The service account to be used for the Agent Engine. +class TuningResourceUsageAssessmentResultDict(TypedDict, total=False): + """Result of tuning resource usage assessment.""" - If not specified, the default Reasoning Engine P6SA service agent will be used.""", - ) - identity_type: Optional[IdentityType] = Field( - default=None, description="""The identity type to use for the Agent Engine.""" - ) - context_spec: Optional[ReasoningEngineContextSpec] = Field( - default=None, - description="""The context spec to be used for the Agent Engine.""", - ) - psc_interface_config: Optional[PscInterfaceConfig] = Field( - default=None, - description="""The PSC interface config for PSC-I to be used for the Agent Engine.""", - ) - min_instances: Optional[int] = Field( - default=None, - description="""The minimum number of instances to run for the Agent Engine. - Defaults to 1. Range: [0, 10]. - """, - ) - max_instances: Optional[int] = Field( - default=None, - description="""The maximum number of instances to run for the Agent Engine. - Defaults to 100. Range: [1, 1000]. - If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. - """, - ) - resource_limits: Optional[dict[str, str]] = Field( - default=None, - description="""The resource limits to be applied to the Agent Engine. - Required keys: 'cpu' and 'memory'. - Supported values for 'cpu': '1', '2', '4', '6', '8'. - Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. - """, - ) - container_concurrency: Optional[int] = Field( - default=None, - description="""The container concurrency to be used for the Agent Engine. - Recommended value: 2 * cpu + 1. Defaults to 9. - """, - ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, - description="""The encryption spec to be used for the Agent Engine.""", - ) - labels: Optional[dict[str, str]] = Field( - default=None, description="""The labels to be used for the Agent Engine.""" - ) - agent_server_mode: Optional[AgentServerMode] = Field( - default=None, description="""The agent server mode to use for deployment.""" - ) - class_methods: Optional[list[dict[str, Any]]] = Field( - default=None, - description="""The class methods to be used for the Agent Engine. - If specified, they'll override the class methods that are autogenerated by - default. By default, methods are generated by inspecting the agent object - and generating a corresponding method for each method defined on the - agent class. - """, - ) - source_packages: Optional[list[str]] = Field( - default=None, - description="""The user-provided paths to the source packages (if any). - If specified, the files in the source packages will be packed into a - a tarball file, uploaded to Agent Engine's API, and deployed to the - Agent Engine. - The following fields will be ignored: - - agent - - extra_packages - - staging_bucket - - requirements - The following fields will be used to install and use the agent from the - source packages: - - entrypoint_module (required) - - entrypoint_object (required) - - requirements_file (optional) - - class_methods (required) - """, - ) - developer_connect_source: Optional[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig - ] = Field( - default=None, - description="""Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""", - ) - entrypoint_module: Optional[str] = Field( - default=None, - description="""The entrypoint module to be used for the Agent Engine - This field only used when source_packages is specified.""", - ) - entrypoint_object: Optional[str] = Field( - default=None, - description="""The entrypoint object to be used for the Agent Engine. - This field only used when source_packages is specified.""", - ) - requirements_file: Optional[str] = Field( - default=None, - description="""The user-provided path to the requirements file (if any). - This field is only used when source_packages is specified. - If not specified, agent engine will find and use the `requirements.txt` in - the source package. - """, - ) - agent_framework: Optional[ - Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] - ] = Field( - default=None, - description="""The agent framework to be used for the Agent Engine. - The OSS agent framework used to develop the agent. - Currently supported values: "google-adk", "langchain", "langgraph", - "ag2", "llama-index", "custom". - If not specified: - - If `agent` is specified, the agent framework will be auto-detected. - - If `source_packages` is specified, the agent framework will - default to "custom".""", - ) - python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] = Field( - default=None, - description="""The Python version to be used for the Agent Engine. - If not specified, it will use the current Python version of the environment. - Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". - """, - ) - build_options: Optional[dict[str, list[str]]] = Field( - default=None, - description="""The build options for the Agent Engine. - The following keys are supported: - - installation_scripts: - Optional. The paths to the installation scripts to be - executed in the Docker image. - The scripts must be located in the `installation_scripts` - subdirectory and the path must be added to `extra_packages`. - """, - ) - image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpec] = Field( - default=None, description="""The image spec for the Agent Engine.""" - ) - agent_config_source: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSource - ] = Field( - default=None, description="""The agent config source for the Agent Engine.""" - ) - container_spec: Optional[ReasoningEngineSpecContainerSpec] = Field( - default=None, description="""The container spec for the Agent Engine.""" - ) - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig - ] = Field( - default=None, - description="""Agent Gateway configuration for a Reasoning Engine deployment.""", - ) - keep_alive_probe: Optional[KeepAliveProbe] = Field( - default=None, - description="""Optional. Specifies the configuration for keep-alive probe. - Contains configuration on a specified endpoint that a deployment host - should use to keep the container alive based on the probe settings.""", - ) - traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( - default=None, description="""The traffic config for the Agent Engine.""" + token_count: Optional[int] + """The number of tokens in the dataset.""" + + billable_character_count: Optional[int] + """The number of billable characters in the dataset.""" + + +TuningResourceUsageAssessmentResultOrDict = Union[ + TuningResourceUsageAssessmentResult, TuningResourceUsageAssessmentResultDict +] + + +class TuningValidationAssessmentResult(_common.BaseModel): + """The result of a tuning validation assessment.""" + + errors: Optional[list[str]] = Field( + default=None, description="""The list of errors found in the dataset.""" ) -class AgentEngineConfigDict(TypedDict, total=False): - """Config for agent engine methods.""" +class TuningValidationAssessmentResultDict(TypedDict, total=False): + """The result of a tuning validation assessment.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + errors: Optional[list[str]] + """The list of errors found in the dataset.""" - staging_bucket: Optional[str] - """The GCS bucket to use for staging the artifacts needed. - It must be a valid GCS bucket name, e.g. "gs://bucket-name". It is - required if `agent_engine` is specified.""" +TuningValidationAssessmentResultOrDict = Union[ + TuningValidationAssessmentResult, TuningValidationAssessmentResultDict +] - requirements: Optional[Any] - """The set of PyPI dependencies needed. - It can either be the path to a single file (requirements.txt), or an - ordered list of strings corresponding to each line of the requirements - file.""" +class Prompt(_common.BaseModel): + """Represents a prompt.""" - display_name: Optional[str] - """The user-defined name of the Agent Engine. + prompt_data: Optional["PromptData"] = Field(default=None, description="""""") + _dataset: Optional["Dataset"] = PrivateAttr(default=None) + _dataset_version: Optional["DatasetVersion"] = PrivateAttr(default=None) - The name can be up to 128 characters long and can comprise any UTF-8 - character.""" + @property + def dataset(self) -> "Dataset": + return self._dataset # type: ignore[return-value] - description: Optional[str] - """The description of the Agent Engine.""" + @property + def dataset_version(self) -> "DatasetVersion": + return self._dataset_version # type: ignore[return-value] - gcs_dir_name: Optional[str] - """The GCS bucket directory under `staging_bucket` to use for staging - the artifacts needed.""" + @property + def prompt_id(self) -> Optional[str]: + """Returns the ID associated with the prompt resource.""" + if self._dataset and self._dataset.name: + return self._dataset.name.split("/")[-1] + elif not self._dataset and ( + self._dataset_version and self._dataset_version.name + ): + return self._dataset_version.name.split("datasets/")[1].split("/")[0] + return None - extra_packages: Optional[list[str]] - """The set of extra user-provided packages (if any).""" + @property + def version_id(self) -> Optional[str]: + """Returns the ID associated with the prompt version resource.""" + if self._dataset_version and self._dataset_version.name: + return self._dataset_version.name.split("/")[-1] + return None - env_vars: Optional[Any] - """The environment variables to be set when running the Agent Engine. + def assemble_contents(self) -> list[genai_types.Content]: + """Transforms a Prompt object into a list with a single genai_types.Content object. - If it is a dictionary, the keys are the environment variable names, and - the values are the corresponding values.""" + This method replaces the variables in the prompt template with the values provided in prompt.prompt_data.variables. + If no variables are provided, prompt.prompt_data.contents is returned as is. Only single-turn prompts are supported. - service_account: Optional[str] - """The service account to be used for the Agent Engine. + This can be used to call generate_content() in the Gen AI SDK. - If not specified, the default Reasoning Engine P6SA service agent will be used.""" + Example usage: - identity_type: Optional[IdentityType] - """The identity type to use for the Agent Engine.""" + my_prompt = types.Prompt( + prompt_data=types.PromptData( + model="gemini-2.0-flash-001", + contents=[ + genai_types.Content( + parts=[ + genai_types.Part(text="Hello {name}!"), + ], + ), + ], + variables=[ + { + "name": genai_types.Part(text="Alice"), + }, + ], + ), + ) - context_spec: Optional[ReasoningEngineContextSpecDict] - """The context spec to be used for the Agent Engine.""" + from google import genai - psc_interface_config: Optional[PscInterfaceConfigDict] - """The PSC interface config for PSC-I to be used for the Agent Engine.""" + genai_client = genai.Client(vertexai=True, project="my-project", location="us-central1") + genai_client.models.generate_content( + model=my_prompt.prompt_data.model, + contents=my_prompt.assemble_contents(), + ) - min_instances: Optional[int] - """The minimum number of instances to run for the Agent Engine. - Defaults to 1. Range: [0, 10]. - """ + Returns: + A list with a single Content object that can be used to call + generate_content(). + """ + if not self.prompt_data or not self.prompt_data.contents: + return [] + + if not self.prompt_data.variables: + return self.prompt_data.contents + + if len(self.prompt_data.contents) > 1: + raise ValueError( + "Multiple contents are not supported. Use assemble_contents() for a prompt with a single Content item." + ) - max_instances: Optional[int] - """The maximum number of instances to run for the Agent Engine. - Defaults to 100. Range: [1, 1000]. - If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. - """ + parts_to_process = self.prompt_data.contents[0].parts + if parts_to_process is None: + return [] + if not isinstance(parts_to_process, list): + parts_to_process = [parts_to_process] - resource_limits: Optional[dict[str, str]] - """The resource limits to be applied to the Agent Engine. - Required keys: 'cpu' and 'memory'. - Supported values for 'cpu': '1', '2', '4', '6', '8'. - Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. - """ + has_placeholders = False + variable_regex = r"\{.*?\}" + for item in parts_to_process: + part = ( + item + if isinstance(item, genai_types.Part) + else genai_types.Part(text=str(item)) + ) + if part.text and re.search(variable_regex, part.text): + has_placeholders = True + break - container_concurrency: Optional[int] - """The container concurrency to be used for the Agent Engine. - Recommended value: 2 * cpu + 1. Defaults to 9. - """ + if not has_placeholders: + return [genai_types.Content(parts=parts_to_process)] - encryption_spec: Optional[genai_types.EncryptionSpec] - """The encryption spec to be used for the Agent Engine.""" + all_rendered_parts: list[genai_types.Part] = [] - labels: Optional[dict[str, str]] - """The labels to be used for the Agent Engine.""" + for var_dict in self.prompt_data.variables: + for template_item in parts_to_process: + template_part = ( + template_item + if isinstance(template_item, genai_types.Part) + else genai_types.Part(text=str(template_item)) + ) + if template_part.text: + rendered_text = template_part.text - agent_server_mode: Optional[AgentServerMode] - """The agent server mode to use for deployment.""" + for key, value in var_dict.items(): + placeholder = f"{{{key}}}" + replacement_text = None - class_methods: Optional[list[dict[str, Any]]] - """The class methods to be used for the Agent Engine. - If specified, they'll override the class methods that are autogenerated by - default. By default, methods are generated by inspecting the agent object - and generating a corresponding method for each method defined on the - agent class. - """ + if isinstance(value, str): + replacement_text = value + elif isinstance(value, genai_types.Part): + if value.text: + replacement_text = value.text + else: + all_rendered_parts.append(value) + if ( + replacement_text is not None + and placeholder in rendered_text + ): + rendered_text = rendered_text.replace( + placeholder, replacement_text + ) + all_rendered_parts.append(genai_types.Part(text=rendered_text)) + else: + all_rendered_parts.append(template_part) + return [genai_types.Content(parts=all_rendered_parts, role="user")] - source_packages: Optional[list[str]] - """The user-provided paths to the source packages (if any). - If specified, the files in the source packages will be packed into a - a tarball file, uploaded to Agent Engine's API, and deployed to the - Agent Engine. - The following fields will be ignored: - - agent - - extra_packages - - staging_bucket - - requirements - The following fields will be used to install and use the agent from the - source packages: - - entrypoint_module (required) - - entrypoint_object (required) - - requirements_file (optional) - - class_methods (required) - """ - developer_connect_source: Optional[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict - ] - """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""" +PromptData = SchemaPromptSpecPromptMessage +PromptDataDict = SchemaPromptSpecPromptMessageDict +PromptDataOrDict = Union[PromptData, PromptDataDict] - entrypoint_module: Optional[str] - """The entrypoint module to be used for the Agent Engine - This field only used when source_packages is specified.""" +ParsedResponseUnion = Union[ + prompts_types.ParsedResponse, prompts_types.ParsedResponseFewShot +] +ParsedResponseUnionDict = Union[ + prompts_types.ParsedResponseDict, prompts_types.ParsedResponseFewShotDict +] - entrypoint_object: Optional[str] - """The entrypoint object to be used for the Agent Engine. - This field only used when source_packages is specified.""" - requirements_file: Optional[str] - """The user-provided path to the requirements file (if any). - This field is only used when source_packages is specified. - If not specified, agent engine will find and use the `requirements.txt` in - the source package. - """ +class PromptDict(TypedDict, total=False): + """Represents a prompt.""" - agent_framework: Optional[ - Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] - ] - """The agent framework to be used for the Agent Engine. - The OSS agent framework used to develop the agent. - Currently supported values: "google-adk", "langchain", "langgraph", - "ag2", "llama-index", "custom". - If not specified: - - If `agent` is specified, the agent framework will be auto-detected. - - If `source_packages` is specified, the agent framework will - default to "custom".""" + prompt_data: Optional["PromptDataDict"] + """""" - python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] - """The Python version to be used for the Agent Engine. - If not specified, it will use the current Python version of the environment. - Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". - """ - build_options: Optional[dict[str, list[str]]] - """The build options for the Agent Engine. - The following keys are supported: - - installation_scripts: - Optional. The paths to the installation scripts to be - executed in the Docker image. - The scripts must be located in the `installation_scripts` - subdirectory and the path must be added to `extra_packages`. - """ +PromptOrDict = Union[Prompt, PromptDict] - image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict] - """The image spec for the Agent Engine.""" - agent_config_source: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict - ] - """The agent config source for the Agent Engine.""" +class SchemaPromptInstanceVariableValue(_common.BaseModel): + """Represents a prompt instance variable.""" - container_spec: Optional[ReasoningEngineSpecContainerSpecDict] - """The container spec for the Agent Engine.""" + part_list: Optional[SchemaPromptSpecPartList] = Field( + default=None, description="""The parts of the variable value.""" + ) - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict - ] - """Agent Gateway configuration for a Reasoning Engine deployment.""" - keep_alive_probe: Optional[KeepAliveProbeDict] - """Optional. Specifies the configuration for keep-alive probe. - Contains configuration on a specified endpoint that a deployment host - should use to keep the container alive based on the probe settings.""" +class SchemaPromptInstanceVariableValueDict(TypedDict, total=False): + """Represents a prompt instance variable.""" - traffic_config: Optional[ReasoningEngineTrafficConfigDict] - """The traffic config for the Agent Engine.""" + part_list: Optional[SchemaPromptSpecPartListDict] + """The parts of the variable value.""" -AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict] +SchemaPromptInstanceVariableValueOrDict = Union[ + SchemaPromptInstanceVariableValue, SchemaPromptInstanceVariableValueDict +] -class RunQueryJobAgentEngineConfig(_common.BaseModel): - """Config for checking a query job on an agent engine.""" +class CreatePromptConfig(_common.BaseModel): + """Config for creating a prompt.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - query: Optional[str] = Field( - default=None, description="""The query to send to the agent engine.""" + prompt_display_name: Optional[str] = Field( + default=None, + description="""The display name for the prompt. If not set, a default name with a timestamp will be used.""", ) - output_gcs_uri: Optional[str] = Field( + timeout: Optional[int] = Field( + default=90, + description="""The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, - description="""The GCS URI to use for the output. - If it is a file, the system use this file to store the response. - If it represents a directory, the system automatically generate a file - for the response. - In both cases, the input query will be stored in the same directory under - the same file name prefix as the output file.""", + description="""Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""", + ) + version_display_name: Optional[str] = Field( + default=None, + description="""The display name for the prompt version. If not set, a default name with a timestamp will be used.""", + ) + max_wait_time: Optional[int] = Field( + default=60, + description="""The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""", ) -class RunQueryJobAgentEngineConfigDict(TypedDict, total=False): - """Config for checking a query job on an agent engine.""" +class CreatePromptConfigDict(TypedDict, total=False): + """Config for creating a prompt.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - query: Optional[str] - """The query to send to the agent engine.""" + prompt_display_name: Optional[str] + """The display name for the prompt. If not set, a default name with a timestamp will be used.""" - output_gcs_uri: Optional[str] - """The GCS URI to use for the output. - If it is a file, the system use this file to store the response. - If it represents a directory, the system automatically generate a file - for the response. - In both cases, the input query will be stored in the same directory under - the same file name prefix as the output file.""" + timeout: Optional[int] + """The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""" -RunQueryJobAgentEngineConfigOrDict = Union[ - RunQueryJobAgentEngineConfig, RunQueryJobAgentEngineConfigDict -] + version_display_name: Optional[str] + """The display name for the prompt version. If not set, a default name with a timestamp will be used.""" + max_wait_time: Optional[int] + """The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""" + + +CreatePromptConfigOrDict = Union[CreatePromptConfig, CreatePromptConfigDict] -class RunQueryJobResult(_common.BaseModel): - """Result of running a query job.""" + +class CreatePromptVersionConfig(_common.BaseModel): + """Config for creating a prompt version.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - job_name: Optional[str] = Field( + version_display_name: Optional[str] = Field( default=None, - description="""Name of the agent engine operation to later check for status.""", + description="""The display name for the prompt version. If not set, a default name with a timestamp will be used.""", ) - input_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the input file.""" + timeout: Optional[int] = Field( + default=90, + description="""The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""", ) - output_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the output file.""" + prompt_display_name: Optional[str] = Field( + default=None, + description="""The display name for the prompt. If not set, a default name with a timestamp will be used.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""", + ) + max_wait_time: Optional[int] = Field( + default=60, + description="""The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""", ) -class RunQueryJobResultDict(TypedDict, total=False): - """Result of running a query job.""" +class CreatePromptVersionConfigDict(TypedDict, total=False): + """Config for creating a prompt version.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - job_name: Optional[str] - """Name of the agent engine operation to later check for status.""" + version_display_name: Optional[str] + """The display name for the prompt version. If not set, a default name with a timestamp will be used.""" - input_gcs_uri: Optional[str] - """The GCS URI of the input file.""" + timeout: Optional[int] + """The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""" + + prompt_display_name: Optional[str] + """The display name for the prompt. If not set, a default name with a timestamp will be used.""" + + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""" - output_gcs_uri: Optional[str] - """The GCS URI of the output file.""" + max_wait_time: Optional[int] + """The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""" -RunQueryJobResultOrDict = Union[RunQueryJobResult, RunQueryJobResultDict] +CreatePromptVersionConfigOrDict = Union[ + CreatePromptVersionConfig, CreatePromptVersionConfigDict +] -class CheckQueryJobResponse(_common.BaseModel): - """Response from LRO.""" +class GetPromptConfig(_common.BaseModel): + """Config for getting a prompt.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - output_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the output file.""" - ) -class CheckQueryJobResponseDict(TypedDict, total=False): - """Response from LRO.""" +class GetPromptConfigDict(TypedDict, total=False): + """Config for getting a prompt.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - output_gcs_uri: Optional[str] - """The GCS URI of the output file.""" - - -CheckQueryJobResponseOrDict = Union[CheckQueryJobResponse, CheckQueryJobResponseDict] - - -class AssembleDataset(_common.BaseModel): - """Represents the assembled dataset.""" - - bigquery_destination: Optional[str] = Field( - default=None, - description="""The BigQuery destination of the assembled dataset.""", - ) - - -class AssembleDatasetDict(TypedDict, total=False): - """Represents the assembled dataset.""" - bigquery_destination: Optional[str] - """The BigQuery destination of the assembled dataset.""" +GetPromptConfigOrDict = Union[GetPromptConfig, GetPromptConfigDict] -AssembleDatasetOrDict = Union[AssembleDataset, AssembleDatasetDict] +class PromptRef(_common.BaseModel): + """Reference to a prompt.""" + prompt_id: Optional[str] = Field(default=None, description="""""") + model: Optional[str] = Field(default=None, description="""""") -class BatchPredictionResourceUsageAssessmentResult(_common.BaseModel): - """Result of batch prediction resource usage assessment.""" - token_count: Optional[int] = Field( - default=None, description="""Number of tokens in the dataset.""" - ) - audio_token_count: Optional[int] = Field( - default=None, description="""Number of audio tokens in the dataset.""" - ) +class PromptRefDict(TypedDict, total=False): + """Reference to a prompt.""" + prompt_id: Optional[str] + """""" -class BatchPredictionResourceUsageAssessmentResultDict(TypedDict, total=False): - """Result of batch prediction resource usage assessment.""" + model: Optional[str] + """""" - token_count: Optional[int] - """Number of tokens in the dataset.""" - audio_token_count: Optional[int] - """Number of audio tokens in the dataset.""" +PromptRefOrDict = Union[PromptRef, PromptRefDict] -BatchPredictionResourceUsageAssessmentResultOrDict = Union[ - BatchPredictionResourceUsageAssessmentResult, - BatchPredictionResourceUsageAssessmentResultDict, -] +class PromptVersionRef(_common.BaseModel): + """Reference to a prompt version.""" + prompt_id: Optional[str] = Field(default=None, description="""""") + version_id: Optional[str] = Field(default=None, description="""""") + model: Optional[str] = Field(default=None, description="""""") -class BatchPredictionValidationAssessmentResult(_common.BaseModel): - """Result of batch prediction validation assessment.""" - errors: Optional[list[str]] = Field( - default=None, description="""The list of errors found in the dataset.""" - ) +class PromptVersionRefDict(TypedDict, total=False): + """Reference to a prompt version.""" + prompt_id: Optional[str] + """""" -class BatchPredictionValidationAssessmentResultDict(TypedDict, total=False): - """Result of batch prediction validation assessment.""" + version_id: Optional[str] + """""" - errors: Optional[list[str]] - """The list of errors found in the dataset.""" + model: Optional[str] + """""" -BatchPredictionValidationAssessmentResultOrDict = Union[ - BatchPredictionValidationAssessmentResult, - BatchPredictionValidationAssessmentResultDict, -] +PromptVersionRefOrDict = Union[PromptVersionRef, PromptVersionRefDict] -class TuningResourceUsageAssessmentResult(_common.BaseModel): - """Result of tuning resource usage assessment.""" +class OptimizeJobConfig(_common.BaseModel): + """VAPO Prompt Optimizer Config.""" - token_count: Optional[int] = Field( - default=None, description="""The number of tokens in the dataset.""" + config_path: Optional[str] = Field( + default=None, + description="""The gcs path to the config file, e.g. gs://bucket/config.json.""", ) - billable_character_count: Optional[int] = Field( + service_account: Optional[str] = Field( default=None, - description="""The number of billable characters in the dataset.""", + description="""The service account to use for the custom job. Cannot be provided at the same time as service_account_project_number.""", + ) + service_account_project_number: Optional[Union[int, str]] = Field( + default=None, + description="""The project number used to construct the default service account:{service_account_project_number}-compute@developer.gserviceaccount.comCannot be provided at the same time as "service_account".""", + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to wait for the job tocomplete. Ignored for async jobs.""", + ) + optimizer_job_display_name: Optional[str] = Field( + default=None, + description="""The display name of the optimization job. If not provided, a display name in the format of "vapo-optimizer-{timestamp}" will be used.""", ) -class TuningResourceUsageAssessmentResultDict(TypedDict, total=False): - """Result of tuning resource usage assessment.""" - - token_count: Optional[int] - """The number of tokens in the dataset.""" - - billable_character_count: Optional[int] - """The number of billable characters in the dataset.""" - - -TuningResourceUsageAssessmentResultOrDict = Union[ - TuningResourceUsageAssessmentResult, TuningResourceUsageAssessmentResultDict -] +class OptimizeJobConfigDict(TypedDict, total=False): + """VAPO Prompt Optimizer Config.""" + config_path: Optional[str] + """The gcs path to the config file, e.g. gs://bucket/config.json.""" -class TuningValidationAssessmentResult(_common.BaseModel): - """The result of a tuning validation assessment.""" + service_account: Optional[str] + """The service account to use for the custom job. Cannot be provided at the same time as service_account_project_number.""" - errors: Optional[list[str]] = Field( - default=None, description="""The list of errors found in the dataset.""" - ) + service_account_project_number: Optional[Union[int, str]] + """The project number used to construct the default service account:{service_account_project_number}-compute@developer.gserviceaccount.comCannot be provided at the same time as "service_account".""" + wait_for_completion: Optional[bool] + """Whether to wait for the job tocomplete. Ignored for async jobs.""" -class TuningValidationAssessmentResultDict(TypedDict, total=False): - """The result of a tuning validation assessment.""" + optimizer_job_display_name: Optional[str] + """The display name of the optimization job. If not provided, a display name in the format of "vapo-optimizer-{timestamp}" will be used.""" - errors: Optional[list[str]] - """The list of errors found in the dataset.""" +OptimizeJobConfigOrDict = Union[OptimizeJobConfig, OptimizeJobConfigDict] -TuningValidationAssessmentResultOrDict = Union[ - TuningValidationAssessmentResult, TuningValidationAssessmentResultDict -] +class ListDeployableModelsConfig(_common.BaseModel): + """Config for listing deployable models.""" -class Prompt(_common.BaseModel): - """Represents a prompt.""" + include_hugging_face_models: Optional[bool] = Field( + default=None, description="""Whether to list Hugging Face models.""" + ) + model_filter: Optional[str] = Field( + default=None, description="""Optional. A string to filter the models by.""" + ) - prompt_data: Optional["PromptData"] = Field(default=None, description="""""") - _dataset: Optional["Dataset"] = PrivateAttr(default=None) - _dataset_version: Optional["DatasetVersion"] = PrivateAttr(default=None) - @property - def dataset(self) -> "Dataset": - return self._dataset # type: ignore[return-value] +class ListDeployableModelsConfigDict(TypedDict, total=False): + """Config for listing deployable models.""" - @property - def dataset_version(self) -> "DatasetVersion": - return self._dataset_version # type: ignore[return-value] + include_hugging_face_models: Optional[bool] + """Whether to list Hugging Face models.""" - @property - def prompt_id(self) -> Optional[str]: - """Returns the ID associated with the prompt resource.""" - if self._dataset and self._dataset.name: - return self._dataset.name.split("/")[-1] - elif not self._dataset and ( - self._dataset_version and self._dataset_version.name - ): - return self._dataset_version.name.split("datasets/")[1].split("/")[0] - return None + model_filter: Optional[str] + """Optional. A string to filter the models by.""" - @property - def version_id(self) -> Optional[str]: - """Returns the ID associated with the prompt version resource.""" - if self._dataset_version and self._dataset_version.name: - return self._dataset_version.name.split("/")[-1] - return None - def assemble_contents(self) -> list[genai_types.Content]: - """Transforms a Prompt object into a list with a single genai_types.Content object. +ListDeployableModelsConfigOrDict = Union[ + ListDeployableModelsConfig, ListDeployableModelsConfigDict +] - This method replaces the variables in the prompt template with the values provided in prompt.prompt_data.variables. - If no variables are provided, prompt.prompt_data.contents is returned as is. Only single-turn prompts are supported. - This can be used to call generate_content() in the Gen AI SDK. +class ListModelGardenModelsConfig(_common.BaseModel): + """Config for listing Model Garden models.""" - Example usage: + include_hugging_face_models: Optional[bool] = Field( + default=None, description="""Whether to list Hugging Face models.""" + ) + model_filter: Optional[str] = Field( + default=None, description="""Optional. A string to filter the models by.""" + ) - my_prompt = types.Prompt( - prompt_data=types.PromptData( - model="gemini-2.0-flash-001", - contents=[ - genai_types.Content( - parts=[ - genai_types.Part(text="Hello {name}!"), - ], - ), - ], - variables=[ - { - "name": genai_types.Part(text="Alice"), - }, - ], - ), - ) - from google import genai +class ListModelGardenModelsConfigDict(TypedDict, total=False): + """Config for listing Model Garden models.""" - genai_client = genai.Client(vertexai=True, project="my-project", location="us-central1") - genai_client.models.generate_content( - model=my_prompt.prompt_data.model, - contents=my_prompt.assemble_contents(), - ) + include_hugging_face_models: Optional[bool] + """Whether to list Hugging Face models.""" - Returns: - A list with a single Content object that can be used to call - generate_content(). - """ - if not self.prompt_data or not self.prompt_data.contents: - return [] + model_filter: Optional[str] + """Optional. A string to filter the models by.""" - if not self.prompt_data.variables: - return self.prompt_data.contents - if len(self.prompt_data.contents) > 1: - raise ValueError( - "Multiple contents are not supported. Use assemble_contents() for a prompt with a single Content item." - ) +ListModelGardenModelsConfigOrDict = Union[ + ListModelGardenModelsConfig, ListModelGardenModelsConfigDict +] - parts_to_process = self.prompt_data.contents[0].parts - if parts_to_process is None: - return [] - if not isinstance(parts_to_process, list): - parts_to_process = [parts_to_process] - has_placeholders = False - variable_regex = r"\{.*?\}" - for item in parts_to_process: - part = ( - item - if isinstance(item, genai_types.Part) - else genai_types.Part(text=str(item)) - ) - if part.text and re.search(variable_regex, part.text): - has_placeholders = True - break +class ListPublisherModelDeployOptionsConfig(_common.BaseModel): + """Config for listing the deploy options of a publisher model.""" - if not has_placeholders: - return [genai_types.Content(parts=parts_to_process)] + machine_type_filter: Optional[Union[str, list[str]]] = Field( + default=None, + description="""Optional. Case-insensitive substring filter on the machine type. + Accepts a single keyword (e.g. ``'g2'``) or a list of keywords (e.g. + ``['n1', 'g2']``); an option matches if it contains any of them.""", + ) + accelerator_type_filter: Optional[Union[str, list[str]]] = Field( + default=None, + description="""Optional. Case-insensitive substring filter on the accelerator + type. Accepts a single keyword (e.g. ``'L4'``) or a list of keywords + (e.g. ``['T4', 'L4']``); an option matches if it contains any of them.""", + ) + serving_container_image_uri_filter: Optional[Union[str, list[str]]] = Field( + default=None, + description="""Optional. Case-insensitive substring filter on the serving + container image URI. Accepts a single keyword (e.g. ``'vllm'``) or a list + of keywords (e.g. ``['vllm', 'tgi']``); an option matches if it contains + any of them.""", + ) + concise: Optional[bool] = Field( + default=None, + description="""Optional. If True, returns a human-readable string describing the + deploy options (container and machine specs) instead of a list of + ``DeployOption`` objects.""", + ) - all_rendered_parts: list[genai_types.Part] = [] - for var_dict in self.prompt_data.variables: - for template_item in parts_to_process: - template_part = ( - template_item - if isinstance(template_item, genai_types.Part) - else genai_types.Part(text=str(template_item)) - ) - if template_part.text: - rendered_text = template_part.text +class ListPublisherModelDeployOptionsConfigDict(TypedDict, total=False): + """Config for listing the deploy options of a publisher model.""" - for key, value in var_dict.items(): - placeholder = f"{{{key}}}" - replacement_text = None + machine_type_filter: Optional[Union[str, list[str]]] + """Optional. Case-insensitive substring filter on the machine type. + Accepts a single keyword (e.g. ``'g2'``) or a list of keywords (e.g. + ``['n1', 'g2']``); an option matches if it contains any of them.""" - if isinstance(value, str): - replacement_text = value - elif isinstance(value, genai_types.Part): - if value.text: - replacement_text = value.text - else: - all_rendered_parts.append(value) - if ( - replacement_text is not None - and placeholder in rendered_text - ): - rendered_text = rendered_text.replace( - placeholder, replacement_text - ) - all_rendered_parts.append(genai_types.Part(text=rendered_text)) - else: - all_rendered_parts.append(template_part) - return [genai_types.Content(parts=all_rendered_parts, role="user")] + accelerator_type_filter: Optional[Union[str, list[str]]] + """Optional. Case-insensitive substring filter on the accelerator + type. Accepts a single keyword (e.g. ``'L4'``) or a list of keywords + (e.g. ``['T4', 'L4']``); an option matches if it contains any of them.""" + + serving_container_image_uri_filter: Optional[Union[str, list[str]]] + """Optional. Case-insensitive substring filter on the serving + container image URI. Accepts a single keyword (e.g. ``'vllm'``) or a list + of keywords (e.g. ``['vllm', 'tgi']``); an option matches if it contains + any of them.""" + concise: Optional[bool] + """Optional. If True, returns a human-readable string describing the + deploy options (container and machine specs) instead of a list of + ``DeployOption`` objects.""" -PromptData = SchemaPromptSpecPromptMessage -PromptDataDict = SchemaPromptSpecPromptMessageDict -PromptDataOrDict = Union[PromptData, PromptDataDict] -ParsedResponseUnion = Union[ - prompts_types.ParsedResponse, prompts_types.ParsedResponseFewShot -] -ParsedResponseUnionDict = Union[ - prompts_types.ParsedResponseDict, prompts_types.ParsedResponseFewShotDict +ListPublisherModelDeployOptionsConfigOrDict = Union[ + ListPublisherModelDeployOptionsConfig, ListPublisherModelDeployOptionsConfigDict ] -class PromptDict(TypedDict, total=False): - """Represents a prompt.""" +class ListCustomModelDeployOptionsConfig(_common.BaseModel): + """Config for listing custom model deploy options.""" - prompt_data: Optional["PromptDataDict"] - """""" + filter_by_user_quota: Optional[bool] = Field( + default=True, + description="""Whether to filter recommendations to regions with user quota. + Only takes effect when ``check_machine_availability=True``; the specs + fallback returned when ``check_machine_availability=False`` carries no + per-region quota information, so this flag is ignored in that mode. + """, + ) + check_machine_availability: Optional[bool] = Field( + default=True, + description="""Whether to check per-region machine availability. -PromptOrDict = Union[Prompt, PromptDict] + When True (the default), the API returns per-region recommendations + that include the machine spec, region and user quota state. When + False, the API returns a flat list of specs without per-region or + quota information (and ``filter_by_user_quota`` has no effect). + """, + ) -class SchemaPromptInstanceVariableValue(_common.BaseModel): - """Represents a prompt instance variable.""" +class ListCustomModelDeployOptionsConfigDict(TypedDict, total=False): + """Config for listing custom model deploy options.""" - part_list: Optional[SchemaPromptSpecPartList] = Field( - default=None, description="""The parts of the variable value.""" - ) + filter_by_user_quota: Optional[bool] + """Whether to filter recommendations to regions with user quota. + Only takes effect when ``check_machine_availability=True``; the specs + fallback returned when ``check_machine_availability=False`` carries no + per-region quota information, so this flag is ignored in that mode. + """ -class SchemaPromptInstanceVariableValueDict(TypedDict, total=False): - """Represents a prompt instance variable.""" + check_machine_availability: Optional[bool] + """Whether to check per-region machine availability. - part_list: Optional[SchemaPromptSpecPartListDict] - """The parts of the variable value.""" + When True (the default), the API returns per-region recommendations + that include the machine spec, region and user quota state. When + False, the API returns a flat list of specs without per-region or + quota information (and ``filter_by_user_quota`` has no effect). + """ -SchemaPromptInstanceVariableValueOrDict = Union[ - SchemaPromptInstanceVariableValue, SchemaPromptInstanceVariableValueDict +ListCustomModelDeployOptionsConfigOrDict = Union[ + ListCustomModelDeployOptionsConfig, ListCustomModelDeployOptionsConfigDict ] -class CreatePromptConfig(_common.BaseModel): - """Config for creating a prompt.""" +class ExportOpenModelConfig(_common.BaseModel): + """Config for ``export_open_model``.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - prompt_display_name: Optional[str] = Field( - default=None, - description="""The display name for the prompt. If not set, a default name with a timestamp will be used.""", - ) - timeout: Optional[int] = Field( - default=90, - description="""The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""", + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to block on the export long-running operation. When + ``True`` (default), returns the destination URI on completion. When + ``False``, returns the ``ExportModelOperation`` for the caller to + poll.""", ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + poll_interval_seconds: Optional[float] = Field( default=None, - description="""Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""", + description="""Seconds between LRO polls when ``wait_for_completion=True``. + Defaults to 30. Ignored when ``wait_for_completion=False``.""", ) - version_display_name: Optional[str] = Field( + timeout_seconds: Optional[float] = Field( default=None, - description="""The display name for the prompt version. If not set, a default name with a timestamp will be used.""", - ) - max_wait_time: Optional[int] = Field( - default=60, - description="""The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""", + description="""Total wall-clock seconds to wait for the export to complete + when ``wait_for_completion=True``. Defaults to 2 hours to + accommodate large model weights. Ignored when + ``wait_for_completion=False``.""", ) -class CreatePromptConfigDict(TypedDict, total=False): - """Config for creating a prompt.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - prompt_display_name: Optional[str] - """The display name for the prompt. If not set, a default name with a timestamp will be used.""" - - timeout: Optional[int] - """The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""" +class ExportOpenModelConfigDict(TypedDict, total=False): + """Config for ``export_open_model``.""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""" + wait_for_completion: Optional[bool] + """Whether to block on the export long-running operation. When + ``True`` (default), returns the destination URI on completion. When + ``False``, returns the ``ExportModelOperation`` for the caller to + poll.""" - version_display_name: Optional[str] - """The display name for the prompt version. If not set, a default name with a timestamp will be used.""" + poll_interval_seconds: Optional[float] + """Seconds between LRO polls when ``wait_for_completion=True``. + Defaults to 30. Ignored when ``wait_for_completion=False``.""" - max_wait_time: Optional[int] - """The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""" + timeout_seconds: Optional[float] + """Total wall-clock seconds to wait for the export to complete + when ``wait_for_completion=True``. Defaults to 2 hours to + accommodate large model weights. Ignored when + ``wait_for_completion=False``.""" -CreatePromptConfigOrDict = Union[CreatePromptConfig, CreatePromptConfigDict] +ExportOpenModelConfigOrDict = Union[ExportOpenModelConfig, ExportOpenModelConfigDict] -class CreatePromptVersionConfig(_common.BaseModel): - """Config for creating a prompt version.""" +class DeployOption(_common.BaseModel): + """A verified deploy option for a model.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - version_display_name: Optional[str] = Field( - default=None, - description="""The display name for the prompt version. If not set, a default name with a timestamp will be used.""", + option_name: Optional[str] = Field( + default=None, description="""The name of the deploy task.""" ) - timeout: Optional[int] = Field( - default=90, - description="""The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""", + serving_container_image_uri: Optional[str] = Field( + default=None, description="""The URI of the serving container.""" ) - prompt_display_name: Optional[str] = Field( - default=None, - description="""The display name for the prompt. If not set, a default name with a timestamp will be used.""", + machine_type: Optional[str] = Field( + default=None, description="""The machine type.""" ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, - description="""Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""", + accelerator_type: Optional[str] = Field( + default=None, description="""The accelerator type.""" ) - max_wait_time: Optional[int] = Field( - default=60, - description="""The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""", + accelerator_count: Optional[int] = Field( + default=None, description="""The number of accelerators.""" ) -class CreatePromptVersionConfigDict(TypedDict, total=False): - """Config for creating a prompt version.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class DeployOptionDict(TypedDict, total=False): + """A verified deploy option for a model.""" - version_display_name: Optional[str] - """The display name for the prompt version. If not set, a default name with a timestamp will be used.""" + option_name: Optional[str] + """The name of the deploy task.""" - timeout: Optional[int] - """The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""" + serving_container_image_uri: Optional[str] + """The URI of the serving container.""" - prompt_display_name: Optional[str] - """The display name for the prompt. If not set, a default name with a timestamp will be used.""" + machine_type: Optional[str] + """The machine type.""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""" + accelerator_type: Optional[str] + """The accelerator type.""" - max_wait_time: Optional[int] - """The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""" + accelerator_count: Optional[int] + """The number of accelerators.""" -CreatePromptVersionConfigOrDict = Union[ - CreatePromptVersionConfig, CreatePromptVersionConfigDict -] +DeployOptionOrDict = Union[DeployOption, DeployOptionDict] -class GetPromptConfig(_common.BaseModel): - """Config for getting a prompt.""" +class Runtime(_common.BaseModel): + """An agent runtime instance.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + api_client: Optional[Any] = Field( + default=None, description="""The underlying API client.""" + ) + api_async_client: Optional[Any] = Field( + default=None, + description="""The underlying API client for asynchronous operations.""", + ) + api_resource: Optional[ReasoningEngine] = Field( + default=None, + description="""The underlying API resource (i.e. ReasoningEngine).""", ) + # Allows dynamic binding of methods based on the registered operations. + model_config = ConfigDict(extra="allow") -class GetPromptConfigDict(TypedDict, total=False): - """Config for getting a prompt.""" + def __repr__(self) -> str: + return ( + f"Runtime(api_resource.name='{self.api_resource.name}')" + if self.api_resource is not None + else "Runtime(api_resource.name=None)" + ) - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + def operation_schemas(self) -> Optional[list[Dict[str, Any]]]: + """Returns the schemas of all registered operations for the agent.""" + if not isinstance(self.api_resource, ReasoningEngine): + raise ValueError("api_resource is not initialized.") + if not self.api_resource.spec: + raise ValueError("api_resource.spec is not initialized.") + return self.api_resource.spec.class_methods + def delete( + self, + force: bool = False, + config: Optional[DeleteRuntimeConfigOrDict] = None, + ) -> None: + """Deletes the agent engine. -GetPromptConfigOrDict = Union[GetPromptConfig, GetPromptConfigDict] + Args: + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Agent Engine has undeleted child resources. Defaults to False. + config (DeleteRuntimeConfig): + Optional. Additional configurations for deleting the Agent Engine. + """ + if not isinstance(self.api_resource, ReasoningEngine): + raise ValueError("api_resource is not initialized.") + self.api_client.delete(name=self.api_resource.name, force=force, config=config) # type: ignore[union-attr] -class PromptRef(_common.BaseModel): - """Reference to a prompt.""" +RubricContentProperty = evals_types.RubricContentProperty +RubricContentPropertyDict = evals_types.RubricContentPropertyDict +RubricContentPropertyDictOrDict = evals_types.RubricContentPropertyOrDict - prompt_id: Optional[str] = Field(default=None, description="""""") - model: Optional[str] = Field(default=None, description="""""") +RubricContent = evals_types.RubricContent +RubricContentDict = evals_types.RubricContentDict +RubricContentDictOrDict = evals_types.RubricContentOrDict +Rubric = evals_types.Rubric +RubricDict = evals_types.RubricDict +RubricDictOrDict = evals_types.RubricOrDict -class PromptRefDict(TypedDict, total=False): - """Reference to a prompt.""" +RubricVerdict = evals_types.RubricVerdict +RubricVerdictDict = evals_types.RubricVerdictDict +RubricVerdictDictOrDict = evals_types.RubricVerdictOrDict - prompt_id: Optional[str] - """""" +CandidateResult = evals_types.CandidateResult +CandidateResultDict = evals_types.CandidateResultDict +CandidateResultDictOrDict = evals_types.CandidateResultOrDict - model: Optional[str] - """""" +Event = evals_types.Event +EventDict = evals_types.EventDict +EventDictOrDict = evals_types.EventOrDict +Message = evals_types.Message +MessageDict = evals_types.MessageDict +MessageDictOrDict = evals_types.MessageOrDict -PromptRefOrDict = Union[PromptRef, PromptRefDict] +Importance = evals_types.Importance -class PromptVersionRef(_common.BaseModel): - """Reference to a prompt version.""" +class RuntimeDict(TypedDict, total=False): + """An agent runtime instance.""" - prompt_id: Optional[str] = Field(default=None, description="""""") - version_id: Optional[str] = Field(default=None, description="""""") - model: Optional[str] = Field(default=None, description="""""") + api_client: Optional[Any] + """The underlying API client.""" + api_async_client: Optional[Any] + """The underlying API client for asynchronous operations.""" -class PromptVersionRefDict(TypedDict, total=False): - """Reference to a prompt version.""" + api_resource: Optional[ReasoningEngineDict] + """The underlying API resource (i.e. ReasoningEngine).""" - prompt_id: Optional[str] - """""" - version_id: Optional[str] - """""" +RuntimeOrDict = Union[Runtime, RuntimeDict] - model: Optional[str] - """""" +class AgentRuntimeConfig(_common.BaseModel): + """Config for agent runtime methods.""" -PromptVersionRefOrDict = Union[PromptVersionRef, PromptVersionRefDict] + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + staging_bucket: Optional[str] = Field( + default=None, + description="""The GCS bucket to use for staging the artifacts needed. + It must be a valid GCS bucket name, e.g. "gs://bucket-name". It is + required if `agent_engine` is specified.""", + ) + requirements: Optional[Any] = Field( + default=None, + description="""The set of PyPI dependencies needed. -class OptimizeJobConfig(_common.BaseModel): - """VAPO Prompt Optimizer Config.""" + It can either be the path to a single file (requirements.txt), or an + ordered list of strings corresponding to each line of the requirements + file.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""The user-defined name of the Agent Runtime. - config_path: Optional[str] = Field( + The name can be up to 128 characters long and can comprise any UTF-8 + character.""", + ) + description: Optional[str] = Field( + default=None, description="""The description of the Agent Runtime.""" + ) + gcs_dir_name: Optional[str] = Field( + default=None, + description="""The GCS bucket directory under `staging_bucket` to use for staging + the artifacts needed.""", + ) + extra_packages: Optional[list[str]] = Field( + default=None, + description="""The set of extra user-provided packages (if any).""", + ) + env_vars: Optional[Any] = Field( + default=None, + description="""The environment variables to be set when running the Agent Runtime. + + If it is a dictionary, the keys are the environment variable names, and + the values are the corresponding values.""", + ) + service_account: Optional[str] = Field( + default=None, + description="""The service account to be used for the Agent Runtime. + + If not specified, the default Reasoning Engine P6SA service agent will be used.""", + ) + identity_type: Optional[IdentityType] = Field( + default=None, description="""The identity type to use for the Agent Runtime.""" + ) + context_spec: Optional[ReasoningEngineContextSpec] = Field( + default=None, + description="""The context spec to be used for the Agent Runtime.""", + ) + psc_interface_config: Optional[PscInterfaceConfig] = Field( + default=None, + description="""The PSC interface config for PSC-I to be used for the Agent Runtime.""", + ) + min_instances: Optional[int] = Field( + default=None, + description="""The minimum number of instances to run for the Agent Runtime. + Defaults to 1. Range: [0, 10]. + """, + ) + max_instances: Optional[int] = Field( + default=None, + description="""The maximum number of instances to run for the Agent Runtime. + Defaults to 100. Range: [1, 1000]. + If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. + """, + ) + resource_limits: Optional[dict[str, str]] = Field( + default=None, + description="""The resource limits to be applied to the Agent Runtime. + Required keys: 'cpu' and 'memory'. + Supported values for 'cpu': '1', '2', '4', '6', '8'. + Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. + """, + ) + container_concurrency: Optional[int] = Field( + default=None, + description="""The container concurrency to be used for the Agent Runtime. + Recommended value: 2 * cpu + 1. Defaults to 9. + """, + ) + keep_alive_probe: Optional[KeepAliveProbe] = Field( + default=None, + description="""Optional. Specifies the configuration for keep-alive probe. + Contains configuration on a specified endpoint that a deployment host + should use to keep the container alive based on the probe settings.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""The encryption spec to be used for the Agent Runtime.""", + ) + labels: Optional[dict[str, str]] = Field( + default=None, description="""The labels to be used for the Agent Runtime.""" + ) + agent_server_mode: Optional[AgentServerMode] = Field( + default=None, description="""The agent server mode to use for deployment.""" + ) + class_methods: Optional[list[dict[str, Any]]] = Field( + default=None, + description="""The class methods to be used for the Agent Runtime. + If specified, they'll override the class methods that are autogenerated by + default. By default, methods are generated by inspecting the agent object + and generating a corresponding method for each method defined on the + agent class. + """, + ) + source_packages: Optional[list[str]] = Field( + default=None, + description="""The user-provided paths to the source packages (if any). + If specified, the files in the source packages will be packed into a + a tarball file, uploaded to Agent Runtime's API, and deployed to the + Agent Runtime. + The following fields will be ignored: + - agent + - extra_packages + - staging_bucket + - requirements + The following fields will be used to install and use the agent from the + source packages: + - entrypoint_module (required) + - entrypoint_object (required) + - requirements_file (optional) + - class_methods (required) + """, + ) + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig + ] = Field( default=None, - description="""The gcs path to the config file, e.g. gs://bucket/config.json.""", + description="""Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""", ) - service_account: Optional[str] = Field( + entrypoint_module: Optional[str] = Field( default=None, - description="""The service account to use for the custom job. Cannot be provided at the same time as service_account_project_number.""", + description="""The entrypoint module to be used for the Agent Runtime + This field only used when source_packages is specified.""", ) - service_account_project_number: Optional[Union[int, str]] = Field( + entrypoint_object: Optional[str] = Field( default=None, - description="""The project number used to construct the default service account:{service_account_project_number}-compute@developer.gserviceaccount.comCannot be provided at the same time as "service_account".""", + description="""The entrypoint object to be used for the Agent Runtime. + This field only used when source_packages is specified.""", ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Whether to wait for the job tocomplete. Ignored for async jobs.""", + requirements_file: Optional[str] = Field( + default=None, + description="""The user-provided path to the requirements file (if any). + This field is only used when source_packages is specified. + If not specified, agent runtime will find and use the `requirements.txt` in + the source package. + """, ) - optimizer_job_display_name: Optional[str] = Field( + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + ] = Field( default=None, - description="""The display name of the optimization job. If not provided, a display name in the format of "vapo-optimizer-{timestamp}" will be used.""", + description="""The agent framework to be used for the Agent Runtime. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""", + ) + python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] = Field( + default=None, + description="""The Python version to be used for the Agent Runtime. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". + """, + ) + build_options: Optional[dict[str, list[str]]] = Field( + default=None, + description="""The build options for the Agent Runtime. + The following keys are supported: + - installation_scripts: + Optional. The paths to the installation scripts to be + executed in the Docker image. + The scripts must be located in the `installation_scripts` + subdirectory and the path must be added to `extra_packages`. + """, + ) + image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpec] = Field( + default=None, description="""The image spec for the Agent Runtime.""" + ) + agent_config_source: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSource + ] = Field( + default=None, description="""The agent config source for the Agent Runtime.""" + ) + traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( + default=None, description="""The traffic config for the Agent Runtime.""" + ) + container_spec: Optional[ReasoningEngineSpecContainerSpec] = Field( + default=None, description="""The container spec for the Agent Runtime.""" + ) + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + ] = Field( + default=None, + description="""Agent Gateway configuration for a Agent Runtime deployment.""", ) -class OptimizeJobConfigDict(TypedDict, total=False): - """VAPO Prompt Optimizer Config.""" +class AgentRuntimeConfigDict(TypedDict, total=False): + """Config for agent runtime methods.""" - config_path: Optional[str] - """The gcs path to the config file, e.g. gs://bucket/config.json.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - service_account: Optional[str] - """The service account to use for the custom job. Cannot be provided at the same time as service_account_project_number.""" + staging_bucket: Optional[str] + """The GCS bucket to use for staging the artifacts needed. - service_account_project_number: Optional[Union[int, str]] - """The project number used to construct the default service account:{service_account_project_number}-compute@developer.gserviceaccount.comCannot be provided at the same time as "service_account".""" + It must be a valid GCS bucket name, e.g. "gs://bucket-name". It is + required if `agent_engine` is specified.""" - wait_for_completion: Optional[bool] - """Whether to wait for the job tocomplete. Ignored for async jobs.""" + requirements: Optional[Any] + """The set of PyPI dependencies needed. - optimizer_job_display_name: Optional[str] - """The display name of the optimization job. If not provided, a display name in the format of "vapo-optimizer-{timestamp}" will be used.""" + It can either be the path to a single file (requirements.txt), or an + ordered list of strings corresponding to each line of the requirements + file.""" + display_name: Optional[str] + """The user-defined name of the Agent Runtime. -OptimizeJobConfigOrDict = Union[OptimizeJobConfig, OptimizeJobConfigDict] + The name can be up to 128 characters long and can comprise any UTF-8 + character.""" + description: Optional[str] + """The description of the Agent Runtime.""" -class AgentEngineRuntimeRevision(_common.BaseModel): - """An agent engine runtime revision instance.""" + gcs_dir_name: Optional[str] + """The GCS bucket directory under `staging_bucket` to use for staging + the artifacts needed.""" - api_client: Optional[Any] = Field( - default=None, description="""The underlying API client.""" - ) - api_async_client: Optional[Any] = Field( - default=None, - description="""The underlying API client for asynchronous operations.""", - ) - api_resource: Optional[ReasoningEngineRuntimeRevision] = Field( - default=None, - description="""The underlying API resource (i.e. ReasoningEngineRuntimeRevision).""", - ) + extra_packages: Optional[list[str]] + """The set of extra user-provided packages (if any).""" - # Allows dynamic binding of methods based on the registered operations. - model_config = ConfigDict(extra="allow") + env_vars: Optional[Any] + """The environment variables to be set when running the Agent Runtime. - def __repr__(self) -> str: - return ( - f"AgentEngineRuntimeRevision(api_resource.name='{self.api_resource.name}')" - if self.api_resource is not None - else "AgentEngineRuntimeRevision(api_resource.name=None)" - ) + If it is a dictionary, the keys are the environment variable names, and + the values are the corresponding values.""" - def operation_schemas(self) -> Optional[list[Dict[str, Any]]]: - """Returns the schemas of all registered operations for the agent.""" - if not isinstance(self.api_resource, ReasoningEngineRuntimeRevision): - raise ValueError("api_resource is not initialized.") - if not self.api_resource.spec: - raise ValueError("api_resource.spec is not initialized.") - return self.api_resource.spec.class_methods + service_account: Optional[str] + """The service account to be used for the Agent Runtime. - def delete( - self, - config: Optional[DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, - ) -> None: - """Deletes the agent engine runtime revision. + If not specified, the default Reasoning Engine P6SA service agent will be used.""" - Args: - config (DeleteAgentEngineRuntimeRevisionConfig): - Optional. Additional configurations for deleting the Agent Engine Runtime Revision. - """ - if not isinstance(self.api_resource, ReasoningEngineRuntimeRevision): - raise ValueError("api_resource is not initialized.") - self.api_client.delete(name=self.api_resource.name, config=config) # type: ignore[union-attr] + identity_type: Optional[IdentityType] + """The identity type to use for the Agent Runtime.""" + context_spec: Optional[ReasoningEngineContextSpecDict] + """The context spec to be used for the Agent Runtime.""" -class AgentEngineRuntimeRevisionDict(TypedDict, total=False): - """An agent engine runtime revision instance.""" + psc_interface_config: Optional[PscInterfaceConfigDict] + """The PSC interface config for PSC-I to be used for the Agent Runtime.""" - api_client: Optional[Any] - """The underlying API client.""" + min_instances: Optional[int] + """The minimum number of instances to run for the Agent Runtime. + Defaults to 1. Range: [0, 10]. + """ - api_async_client: Optional[Any] - """The underlying API client for asynchronous operations.""" + max_instances: Optional[int] + """The maximum number of instances to run for the Agent Runtime. + Defaults to 100. Range: [1, 1000]. + If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. + """ - api_resource: Optional[ReasoningEngineRuntimeRevisionDict] - """The underlying API resource (i.e. ReasoningEngineRuntimeRevision).""" + resource_limits: Optional[dict[str, str]] + """The resource limits to be applied to the Agent Runtime. + Required keys: 'cpu' and 'memory'. + Supported values for 'cpu': '1', '2', '4', '6', '8'. + Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. + """ + container_concurrency: Optional[int] + """The container concurrency to be used for the Agent Runtime. + Recommended value: 2 * cpu + 1. Defaults to 9. + """ -AgentEngineRuntimeRevisionOrDict = Union[ - AgentEngineRuntimeRevision, AgentEngineRuntimeRevisionDict -] + keep_alive_probe: Optional[KeepAliveProbeDict] + """Optional. Specifies the configuration for keep-alive probe. + Contains configuration on a specified endpoint that a deployment host + should use to keep the container alive based on the probe settings.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """The encryption spec to be used for the Agent Runtime.""" -class ListDeployableModelsConfig(_common.BaseModel): - """Config for listing deployable models.""" + labels: Optional[dict[str, str]] + """The labels to be used for the Agent Runtime.""" - include_hugging_face_models: Optional[bool] = Field( - default=None, description="""Whether to list Hugging Face models.""" - ) - model_filter: Optional[str] = Field( - default=None, description="""Optional. A string to filter the models by.""" - ) + agent_server_mode: Optional[AgentServerMode] + """The agent server mode to use for deployment.""" + class_methods: Optional[list[dict[str, Any]]] + """The class methods to be used for the Agent Runtime. + If specified, they'll override the class methods that are autogenerated by + default. By default, methods are generated by inspecting the agent object + and generating a corresponding method for each method defined on the + agent class. + """ -class ListDeployableModelsConfigDict(TypedDict, total=False): - """Config for listing deployable models.""" + source_packages: Optional[list[str]] + """The user-provided paths to the source packages (if any). + If specified, the files in the source packages will be packed into a + a tarball file, uploaded to Agent Runtime's API, and deployed to the + Agent Runtime. + The following fields will be ignored: + - agent + - extra_packages + - staging_bucket + - requirements + The following fields will be used to install and use the agent from the + source packages: + - entrypoint_module (required) + - entrypoint_object (required) + - requirements_file (optional) + - class_methods (required) + """ - include_hugging_face_models: Optional[bool] - """Whether to list Hugging Face models.""" + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict + ] + """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""" - model_filter: Optional[str] - """Optional. A string to filter the models by.""" + entrypoint_module: Optional[str] + """The entrypoint module to be used for the Agent Runtime + This field only used when source_packages is specified.""" + entrypoint_object: Optional[str] + """The entrypoint object to be used for the Agent Runtime. + This field only used when source_packages is specified.""" -ListDeployableModelsConfigOrDict = Union[ - ListDeployableModelsConfig, ListDeployableModelsConfigDict -] + requirements_file: Optional[str] + """The user-provided path to the requirements file (if any). + This field is only used when source_packages is specified. + If not specified, agent runtime will find and use the `requirements.txt` in + the source package. + """ + + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + ] + """The agent framework to be used for the Agent Runtime. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""" + python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] + """The Python version to be used for the Agent Runtime. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". + """ -class ListModelGardenModelsConfig(_common.BaseModel): - """Config for listing Model Garden models.""" + build_options: Optional[dict[str, list[str]]] + """The build options for the Agent Runtime. + The following keys are supported: + - installation_scripts: + Optional. The paths to the installation scripts to be + executed in the Docker image. + The scripts must be located in the `installation_scripts` + subdirectory and the path must be added to `extra_packages`. + """ - include_hugging_face_models: Optional[bool] = Field( - default=None, description="""Whether to list Hugging Face models.""" - ) - model_filter: Optional[str] = Field( - default=None, description="""Optional. A string to filter the models by.""" - ) + image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict] + """The image spec for the Agent Runtime.""" + agent_config_source: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict + ] + """The agent config source for the Agent Runtime.""" -class ListModelGardenModelsConfigDict(TypedDict, total=False): - """Config for listing Model Garden models.""" + traffic_config: Optional[ReasoningEngineTrafficConfigDict] + """The traffic config for the Agent Runtime.""" - include_hugging_face_models: Optional[bool] - """Whether to list Hugging Face models.""" + container_spec: Optional[ReasoningEngineSpecContainerSpecDict] + """The container spec for the Agent Runtime.""" - model_filter: Optional[str] - """Optional. A string to filter the models by.""" + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] + """Agent Gateway configuration for a Agent Runtime deployment.""" -ListModelGardenModelsConfigOrDict = Union[ - ListModelGardenModelsConfig, ListModelGardenModelsConfigDict -] +AgentRuntimeConfigOrDict = Union[AgentRuntimeConfig, AgentRuntimeConfigDict] -class ListPublisherModelDeployOptionsConfig(_common.BaseModel): - """Config for listing the deploy options of a publisher model.""" +class RunQueryJobRuntimeConfig(_common.BaseModel): + """Config for checking a query job on an agent runtime.""" - machine_type_filter: Optional[Union[str, list[str]]] = Field( - default=None, - description="""Optional. Case-insensitive substring filter on the machine type. - Accepts a single keyword (e.g. ``'g2'``) or a list of keywords (e.g. - ``['n1', 'g2']``); an option matches if it contains any of them.""", - ) - accelerator_type_filter: Optional[Union[str, list[str]]] = Field( - default=None, - description="""Optional. Case-insensitive substring filter on the accelerator - type. Accepts a single keyword (e.g. ``'L4'``) or a list of keywords - (e.g. ``['T4', 'L4']``); an option matches if it contains any of them.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - serving_container_image_uri_filter: Optional[Union[str, list[str]]] = Field( - default=None, - description="""Optional. Case-insensitive substring filter on the serving - container image URI. Accepts a single keyword (e.g. ``'vllm'``) or a list - of keywords (e.g. ``['vllm', 'tgi']``); an option matches if it contains - any of them.""", + query: Optional[str] = Field( + default=None, description="""The query to send to the agent runtime.""" ) - concise: Optional[bool] = Field( + output_gcs_uri: Optional[str] = Field( default=None, - description="""Optional. If True, returns a human-readable string describing the - deploy options (container and machine specs) instead of a list of - ``DeployOption`` objects.""", + description="""The GCS URI to use for the output. + If it is a file, the system use this file to store the response. + If it represents a directory, the system automatically generate a file + for the response. + In both cases, the input query will be stored in the same directory under + the same file name prefix as the output file.""", ) -class ListPublisherModelDeployOptionsConfigDict(TypedDict, total=False): - """Config for listing the deploy options of a publisher model.""" - - machine_type_filter: Optional[Union[str, list[str]]] - """Optional. Case-insensitive substring filter on the machine type. - Accepts a single keyword (e.g. ``'g2'``) or a list of keywords (e.g. - ``['n1', 'g2']``); an option matches if it contains any of them.""" +class RunQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for checking a query job on an agent runtime.""" - accelerator_type_filter: Optional[Union[str, list[str]]] - """Optional. Case-insensitive substring filter on the accelerator - type. Accepts a single keyword (e.g. ``'L4'``) or a list of keywords - (e.g. ``['T4', 'L4']``); an option matches if it contains any of them.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - serving_container_image_uri_filter: Optional[Union[str, list[str]]] - """Optional. Case-insensitive substring filter on the serving - container image URI. Accepts a single keyword (e.g. ``'vllm'``) or a list - of keywords (e.g. ``['vllm', 'tgi']``); an option matches if it contains - any of them.""" + query: Optional[str] + """The query to send to the agent runtime.""" - concise: Optional[bool] - """Optional. If True, returns a human-readable string describing the - deploy options (container and machine specs) instead of a list of - ``DeployOption`` objects.""" + output_gcs_uri: Optional[str] + """The GCS URI to use for the output. + If it is a file, the system use this file to store the response. + If it represents a directory, the system automatically generate a file + for the response. + In both cases, the input query will be stored in the same directory under + the same file name prefix as the output file.""" -ListPublisherModelDeployOptionsConfigOrDict = Union[ - ListPublisherModelDeployOptionsConfig, ListPublisherModelDeployOptionsConfigDict +RunQueryJobRuntimeConfigOrDict = Union[ + RunQueryJobRuntimeConfig, RunQueryJobRuntimeConfigDict ] -class ListCustomModelDeployOptionsConfig(_common.BaseModel): - """Config for listing custom model deploy options.""" - - filter_by_user_quota: Optional[bool] = Field( - default=True, - description="""Whether to filter recommendations to regions with user quota. +class RunQueryJobResult(_common.BaseModel): + """Result of running a query job.""" - Only takes effect when ``check_machine_availability=True``; the specs - fallback returned when ``check_machine_availability=False`` carries no - per-region quota information, so this flag is ignored in that mode. - """, + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - check_machine_availability: Optional[bool] = Field( - default=True, - description="""Whether to check per-region machine availability. - - When True (the default), the API returns per-region recommendations - that include the machine spec, region and user quota state. When - False, the API returns a flat list of specs without per-region or - quota information (and ``filter_by_user_quota`` has no effect). - """, + job_name: Optional[str] = Field( + default=None, + description="""Name of the agent runtime operation to later check for status.""", + ) + input_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the input file.""" + ) + output_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the output file.""" ) -class ListCustomModelDeployOptionsConfigDict(TypedDict, total=False): - """Config for listing custom model deploy options.""" +class RunQueryJobResultDict(TypedDict, total=False): + """Result of running a query job.""" - filter_by_user_quota: Optional[bool] - """Whether to filter recommendations to regions with user quota. + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - Only takes effect when ``check_machine_availability=True``; the specs - fallback returned when ``check_machine_availability=False`` carries no - per-region quota information, so this flag is ignored in that mode. - """ + job_name: Optional[str] + """Name of the agent runtime operation to later check for status.""" - check_machine_availability: Optional[bool] - """Whether to check per-region machine availability. + input_gcs_uri: Optional[str] + """The GCS URI of the input file.""" - When True (the default), the API returns per-region recommendations - that include the machine spec, region and user quota state. When - False, the API returns a flat list of specs without per-region or - quota information (and ``filter_by_user_quota`` has no effect). - """ + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" -ListCustomModelDeployOptionsConfigOrDict = Union[ - ListCustomModelDeployOptionsConfig, ListCustomModelDeployOptionsConfigDict -] +RunQueryJobResultOrDict = Union[RunQueryJobResult, RunQueryJobResultDict] -class ExportOpenModelConfig(_common.BaseModel): - """Config for ``export_open_model``.""" +class CheckQueryJobResponse(_common.BaseModel): + """Response from LRO.""" - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Whether to block on the export long-running operation. When - ``True`` (default), returns the destination URI on completion. When - ``False``, returns the ``ExportModelOperation`` for the caller to - poll.""", - ) - poll_interval_seconds: Optional[float] = Field( - default=None, - description="""Seconds between LRO polls when ``wait_for_completion=True``. - Defaults to 30. Ignored when ``wait_for_completion=False``.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - timeout_seconds: Optional[float] = Field( - default=None, - description="""Total wall-clock seconds to wait for the export to complete - when ``wait_for_completion=True``. Defaults to 2 hours to - accommodate large model weights. Ignored when - ``wait_for_completion=False``.""", + output_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the output file.""" ) -class ExportOpenModelConfigDict(TypedDict, total=False): - """Config for ``export_open_model``.""" - - wait_for_completion: Optional[bool] - """Whether to block on the export long-running operation. When - ``True`` (default), returns the destination URI on completion. When - ``False``, returns the ``ExportModelOperation`` for the caller to - poll.""" +class CheckQueryJobResponseDict(TypedDict, total=False): + """Response from LRO.""" - poll_interval_seconds: Optional[float] - """Seconds between LRO polls when ``wait_for_completion=True``. - Defaults to 30. Ignored when ``wait_for_completion=False``.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - timeout_seconds: Optional[float] - """Total wall-clock seconds to wait for the export to complete - when ``wait_for_completion=True``. Defaults to 2 hours to - accommodate large model weights. Ignored when - ``wait_for_completion=False``.""" + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" -ExportOpenModelConfigOrDict = Union[ExportOpenModelConfig, ExportOpenModelConfigDict] +CheckQueryJobResponseOrDict = Union[CheckQueryJobResponse, CheckQueryJobResponseDict] -class DeployOption(_common.BaseModel): - """A verified deploy option for a model.""" +class RuntimeRevision(_common.BaseModel): + """An agent runtime revision instance.""" - option_name: Optional[str] = Field( - default=None, description="""The name of the deploy task.""" - ) - serving_container_image_uri: Optional[str] = Field( - default=None, description="""The URI of the serving container.""" - ) - machine_type: Optional[str] = Field( - default=None, description="""The machine type.""" + api_client: Optional[Any] = Field( + default=None, description="""The underlying API client.""" ) - accelerator_type: Optional[str] = Field( - default=None, description="""The accelerator type.""" + api_async_client: Optional[Any] = Field( + default=None, + description="""The underlying API client for asynchronous operations.""", ) - accelerator_count: Optional[int] = Field( - default=None, description="""The number of accelerators.""" + api_resource: Optional[ReasoningEngineRuntimeRevision] = Field( + default=None, + description="""The underlying API resource (i.e. ReasoningEngineRuntimeRevision).""", ) + # Allows dynamic binding of methods based on the registered operations. + model_config = ConfigDict(extra="allow") -class DeployOptionDict(TypedDict, total=False): - """A verified deploy option for a model.""" + def __repr__(self) -> str: + return ( + f"RuntimeRevision(api_resource.name='{self.api_resource.name}')" + if self.api_resource is not None + else "RuntimeRevision(api_resource.name=None)" + ) - option_name: Optional[str] - """The name of the deploy task.""" + def operation_schemas(self) -> Optional[list[Dict[str, Any]]]: + """Returns the schemas of all registered operations for the agent.""" + if not isinstance(self.api_resource, ReasoningEngineRuntimeRevision): + raise ValueError("api_resource is not initialized.") + if not self.api_resource.spec: + raise ValueError("api_resource.spec is not initialized.") + return self.api_resource.spec.class_methods - serving_container_image_uri: Optional[str] - """The URI of the serving container.""" + def delete( + self, + config: Optional[DeleteRuntimeRevisionConfigOrDict] = None, + ) -> None: + """Deletes the agent engine runtime revision. - machine_type: Optional[str] - """The machine type.""" + Args: + config (DeleteRuntimeRevisionConfig): + Optional. Additional configurations for deleting the Agent Engine Runtime Revision. + """ + if not isinstance(self.api_resource, ReasoningEngineRuntimeRevision): + raise ValueError("api_resource is not initialized.") + self.api_client.delete(name=self.api_resource.name, config=config) # type: ignore[union-attr] - accelerator_type: Optional[str] - """The accelerator type.""" - accelerator_count: Optional[int] - """The number of accelerators.""" +class RuntimeRevisionDict(TypedDict, total=False): + """An agent runtime revision instance.""" + + api_client: Optional[Any] + """The underlying API client.""" + api_async_client: Optional[Any] + """The underlying API client for asynchronous operations.""" -DeployOptionOrDict = Union[DeployOption, DeployOptionDict] + api_resource: Optional[ReasoningEngineRuntimeRevisionDict] + """The underlying API resource (i.e. ReasoningEngineRuntimeRevision).""" + + +RuntimeRevisionOrDict = Union[RuntimeRevision, RuntimeRevisionDict] diff --git a/agentplatform/agent_engines/__init__.py b/agentplatform/agent_engines/__init__.py deleted file mode 100644 index ac4802a9b2..0000000000 --- a/agentplatform/agent_engines/__init__.py +++ /dev/null @@ -1,428 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -"""Classes and functions for working with agent engines.""" - -from typing import Dict, Iterable, Optional, Sequence, Union - -from google.cloud.aiplatform import base -from google.cloud.aiplatform import initializer -from google.cloud.aiplatform import utils as aip_utils -from google.cloud.aiplatform_v1 import types as aip_types - -# We just want to re-export certain classes -# pylint: disable=g-multiple-import,g-importing-member -from agentplatform.agent_engines._agent_engines import ( - _AgentEngineInterface, - AgentEngine, - Cloneable, - ModuleAgent, - OperationRegistrable, - Queryable, - AsyncQueryable, - StreamQueryable, - AsyncStreamQueryable, -) -from agentplatform.agent_engines.templates.adk import ( - AdkApp, -) -from agentplatform.agent_engines.templates.ag2 import ( - AG2Agent, -) -from agentplatform.agent_engines.templates.langchain import ( - LangchainAgent, -) -from agentplatform.agent_engines.templates.langgraph import ( - LanggraphAgent, -) -from agentplatform.agent_engines.templates.llama_index import ( - LlamaIndexQueryPipelineAgent, -) - - -_LOGGER = base.Logger(__name__) - - -def get(resource_name: str) -> AgentEngine: - """Retrieves an Agent Engine resource. - - Args: - resource_name (str): - Required. A fully-qualified resource name or ID such as - "projects/123/locations/us-central1/reasoningEngines/456" or - "456" when project and location are initialized or passed. - """ - return AgentEngine(resource_name) - - -def create( - agent_engine: Optional[_AgentEngineInterface] = None, - *, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - build_options: Optional[Dict[str, Sequence[str]]] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[aip_types.EncryptionSpec] = None, -) -> AgentEngine: - """Creates a new Agent Engine. - - The Agent Engine will be an instance of the `agent_engine` that - was passed in, running remotely on Vertex AI. - - Sample ``src_dir`` contents (e.g. ``./user_src_dir``): - - .. code-block:: python - - user_src_dir/ - |-- main.py - |-- requirements.txt - |-- user_code/ - | |-- utils.py - | |-- ... - |-- installation_scripts/ - | |-- install_package.sh - | |-- ... - |-- ... - - To build an Agent Engine with the above files, run: - - .. code-block:: python - - remote_agent = agent_engines.create( - agent_engine=local_agent, - requirements=[ - # I.e. the PyPI dependencies listed in requirements.txt - "google-cloud-aiplatform==1.25.0", - "langchain==0.0.242", - ... - ], - extra_packages=[ - "./user_src_dir/main.py", # a single file - "./user_src_dir/user_code", # a directory - ... - ], - build_options={ - "installation": [ - "./user_src_dir/installation_scripts/install_package.sh", - ... - ], - }, - ) - - Args: - agent_engine (AgentEngineInterface): - Required. The Agent Engine to be created. - requirements (Union[str, Sequence[str]]): - Optional. The set of PyPI dependencies needed. It can either be - the path to a single file (requirements.txt), or an ordered list - of strings corresponding to each line of the requirements file. - display_name (str): - Optional. The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any - UTF-8 character. - description (str): - Optional. The description of the Agent Engine. - gcs_dir_name (str): - Optional. The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - extra_packages (Sequence[str]): - Optional. The set of extra user-provided packages (if any). - env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): - Optional. The environment variables to be set when running the - Agent Engine. If it is a list of strings, each string should be - a valid key to `os.environ`. If it is a dictionary, the keys are - the environment variable names, and the values are the - corresponding values. - build_options (Dict[str, Sequence[str]]): - Optional. The build options for the Agent Engine. This includes - options such as installation scripts. - service_account (str): - Optional. The service account to be used for the Agent Engine. If - not specified, the default reasoning engine service agent service - account will be used. - psc_interface_config (PscInterfaceConfig): - Optional. The PSC interface config for the Agent Engine. If not - specified, the default PSC interface config will be used. - min_instances (int): - Optional. The minimum number of instances to run the Agent Engine. - If not specified, the default value will be used. - max_instances (int): - Optional. The maximum number of instances to run the Agent Engine. - If not specified, the default value will be used. - resource_limits (Dict[str, str]): - Optional. The resource limits for the Agent Engine. If not - specified, the default value will be used. - container_concurrency (int): - Optional. The container concurrency for the Agent Engine. If not - specified, the default value will be used. - encryption_spec (EncryptionSpec): - Optional. The encryption spec for the Agent Engine. If not - specified, the default encryption spec will be used. - - Returns: - AgentEngine: The Agent Engine that was created. - - Raises: - ValueError: If the `project` was not set using `agentplatform.init`. - ValueError: If the `location` was not set using `agentplatform.init`. - ValueError: If the `staging_bucket` was not set using agentplatform.init. - ValueError: If the `staging_bucket` does not start with "gs://". - FileNotFoundError: If `extra_packages` includes a file or directory - that does not exist. - IOError: If requirements is a string that corresponds to a - nonexistent file. - """ - return AgentEngine.create( - agent_engine=agent_engine, - requirements=requirements, - display_name=display_name, - description=description, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - env_vars=env_vars, - build_options=build_options, - service_account=service_account, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - encryption_spec=encryption_spec, - ) - - -def list(*, filter: str = "") -> Iterable[AgentEngine]: - """List all instances of Agent Engine matching the filter. - - Example Usage: - - .. code-block:: python - import agentplatform - from agentplatform import agent_engines - - agentplatform.init(project="my_project", location="us-central1") - agent_engines.list(filter='display_name="My Custom Agent"') - - Args: - filter (str): - Optional. An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported. - - Returns: - Iterable[AgentEngine]: An iterable of Agent Engines matching the filter. - """ - api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineClientWithOverride, - ) - for agent in api_client.list_reasoning_engines( - request=aip_types.ListReasoningEnginesRequest( - parent=initializer.global_config.common_location_path(), - filter=filter, - ) - ): - yield AgentEngine(agent.name) - - -def delete( - resource_name: str, - *, - force: bool = False, - **kwargs, -) -> None: - """Delete an Agent Engine resource. - - Args: - resource_name (str): - Required. The name of the Agent Engine to be deleted. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error - when the Agent Engine has undeleted child resources. Defaults to - False. - **kwargs (dict[str, Any]): - Optional. Additional keyword arguments to pass to the - delete_reasoning_engine method. - """ - api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineClientWithOverride, - ) - _LOGGER.info(f"Deleting AgentEngine resource: {resource_name}") - operation_future = api_client.delete_reasoning_engine( - request=aip_types.DeleteReasoningEngineRequest( - name=resource_name, - force=force, - **(kwargs or {}), - ) - ) - _LOGGER.info(f"Delete AgentEngine backing LRO: {operation_future.operation.name}") - operation_future.result() - _LOGGER.info(f"AgentEngine resource deleted: {resource_name}") - - -def update( - resource_name: str, - *, - agent_engine: Optional[Union[Queryable, OperationRegistrable]] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - build_options: Optional[Dict[str, Sequence[str]]] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[aip_types.EncryptionSpec] = None, -) -> "AgentEngine": - """Updates an existing Agent Engine. - - This method updates the configuration of a deployed Agent Engine, identified - by its resource name. Unlike the `create` function which requires an - `agent_engine` object, all arguments in this method are optional. This - method allows you to modify individual aspects of the configuration by - providing any of the optional arguments. - - Args: - resource_name (str): - Required. The name of the Agent Engine to be updated. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - agent_engine (AgentEngineInterface): - Optional. The instance to be used as the updated Agent Engine. If it - is not specified, the existing instance will be used. - requirements (Union[str, Sequence[str]]): - Optional. The set of PyPI dependencies needed. It can either be - the path to a single file (requirements.txt), or an ordered list - of strings corresponding to each line of the requirements file. - If it is not specified, the existing requirements will be used. - If it is set to an empty string or list, the existing - requirements will be removed. - display_name (str): - Optional. The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any - UTF-8 character. - description (str): - Optional. The description of the Agent Engine. - gcs_dir_name (str): - Optional. The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - extra_packages (Sequence[str]): - Optional. The set of extra user-provided packages (if any). If - it is not specified, the existing extra packages will be used. - If it is set to an empty list, the existing extra packages will - be removed. - env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): - Optional. The environment variables to be set when running the - Agent Engine. If it is a list of strings, each string should be - a valid key to `os.environ`. If it is a dictionary, the keys are - the environment variable names, and the values are the - corresponding values. - build_options (Dict[str, Sequence[str]]): - Optional. The build options for the Agent Engine. This includes - options such as installation scripts. - service_account (str): - Optional. The service account to be used for the Agent Engine. If - not specified, the default reasoning engine service agent service - account will be used. - min_instances (int): - Optional. The minimum number of instances to run the Agent Engine. - If not specified, the default value will be used. - max_instances (int): - Optional. The maximum number of instances to run the Agent Engine. - If not specified, the default value will be used. - resource_limits (Dict[str, str]): - Optional. The resource limits for the Agent Engine. If not - specified, the default value will be used. - container_concurrency (int): - Optional. The container concurrency for the Agent Engine. If not - specified, the default value will be used. - encryption_spec (EncryptionSpec): - Optional. The encryption spec for the Agent Engine. If not - specified, the default encryption spec will be used. - - Returns: - AgentEngine: The Agent Engine that was updated. - - Raises: - ValueError: If the `staging_bucket` was not set using agentplatform.init. - ValueError: If the `staging_bucket` does not start with "gs://". - FileNotFoundError: If `extra_packages` includes a file or directory - that does not exist. - ValueError: if none of `display_name`, `description`, - `requirements`, `extra_packages`, `agent_engine`, or `build_options` - were specified. - IOError: If requirements is a string that corresponds to a - nonexistent file. - """ - agent = get(resource_name) - return agent.update( - agent_engine=agent_engine, - requirements=requirements, - display_name=display_name, - description=description, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - env_vars=env_vars, - build_options=build_options, - service_account=service_account, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - encryption_spec=encryption_spec, - ) - - -__all__ = ( - # Resources - "AgentEngine", - # Protocols - "Cloneable", - "OperationRegistrable", - "Queryable", - "AsyncQueryable", - "StreamQueryable", - "AsyncStreamQueryable", - # Methods - "create", - "delete", - "get", - "list", - "update", - # Templates - "AdkApp", - "ModuleAgent", - "LangchainAgent", - "LanggraphAgent", - "AG2Agent", - "LlamaIndexQueryPipelineAgent", -) diff --git a/agentplatform/agent_engines/_agent_engines.py b/agentplatform/agent_engines/_agent_engines.py deleted file mode 100644 index 35333b82a4..0000000000 --- a/agentplatform/agent_engines/_agent_engines.py +++ /dev/null @@ -1,2028 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import abc -import inspect -import io -import json -import logging -import os -import sys -import tarfile -import types -import typing -from typing import ( - Any, - AsyncIterable, - Callable, - Coroutine, - Dict, - Iterable, - List, - Optional, - Protocol, - Sequence, - Tuple, - Union, -) - -from google.api_core import exceptions -from google.cloud import storage -from google.cloud.aiplatform import base -from google.cloud.aiplatform import initializer -from google.cloud.aiplatform import utils as aip_utils -from google.cloud.aiplatform_v1 import types as aip_types -from google.cloud.aiplatform_v1.types import reasoning_engine_service -from agentplatform._genai import _agent_engines_utils -import httpx -import proto - -from google.protobuf import field_mask_pb2 - - -_LOGGER = base.Logger("agentplatform.agent_engines") - -_SUPPORTED_PYTHON_VERSIONS = ("3.10", "3.11", "3.12", "3.13", "3.14") -_DEFAULT_GCS_DIR_NAME = "agent_engine" -_BLOB_FILENAME = "agent_engine.pkl" -_REQUIREMENTS_FILE = "requirements.txt" -_EXTRA_PACKAGES_FILE = "dependencies.tar.gz" -_STANDARD_API_MODE = "" -_ASYNC_API_MODE = "async" -_STREAM_API_MODE = "stream" -_ASYNC_STREAM_API_MODE = "async_stream" -_BIDI_STREAM_API_MODE = "bidi_stream" -_A2A_EXTENSION_MODE = "a2a_extension" -_A2A_AGENT_CARD = "a2a_agent_card" -_MODE_KEY_IN_SCHEMA = "api_mode" -_METHOD_NAME_KEY_IN_SCHEMA = "name" -_DEFAULT_METHOD_NAME = "query" -_DEFAULT_ASYNC_METHOD_NAME = "async_query" -_DEFAULT_STREAM_METHOD_NAME = "stream_query" -_DEFAULT_ASYNC_STREAM_METHOD_NAME = "async_stream_query" -_DEFAULT_METHOD_RETURN_TYPE = "dict[str, Any]" -_DEFAULT_ASYNC_METHOD_RETURN_TYPE = "Coroutine[Any, Any, Any]" -_DEFAULT_STREAM_METHOD_RETURN_TYPE = "Iterable[Any]" -_DEFAULT_ASYNC_STREAM_METHOD_RETURN_TYPE = "AsyncIterable[Any]" -_DEFAULT_METHOD_DOCSTRING_TEMPLATE = """ - Runs the Agent Engine to serve the user request. - This will be based on the `.{method_name}(...)` of the python object that - was passed in when creating the Agent Engine. The method will invoke the - `{default_method_name}` API client of the python object. - Args: - **kwargs: - Optional. The arguments of the `.{method_name}(...)` method. - Returns: - {return_type}: The response from serving the user request. -""" -_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE = ( - "Failed to register API methods. Please follow the guide to " - "register the API methods: " - "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#custom-methods. " - "Error: {%s}" -) -_AGENT_FRAMEWORK_ATTR = "agent_framework" -_DEFAULT_AGENT_FRAMEWORK = "custom" -_BUILD_OPTIONS_INSTALLATION = "installation_scripts" -_DEFAULT_METHOD_NAME_MAP = { - _STANDARD_API_MODE: _DEFAULT_METHOD_NAME, - _ASYNC_API_MODE: _DEFAULT_ASYNC_METHOD_NAME, - _STREAM_API_MODE: _DEFAULT_STREAM_METHOD_NAME, - _ASYNC_STREAM_API_MODE: _DEFAULT_ASYNC_STREAM_METHOD_NAME, -} -_DEFAULT_METHOD_RETURN_TYPE_MAP = { - _STANDARD_API_MODE: _DEFAULT_METHOD_RETURN_TYPE, - _ASYNC_API_MODE: _DEFAULT_ASYNC_METHOD_RETURN_TYPE, - _STREAM_API_MODE: _DEFAULT_STREAM_METHOD_RETURN_TYPE, - _ASYNC_STREAM_API_MODE: _DEFAULT_ASYNC_STREAM_METHOD_RETURN_TYPE, -} - - -try: - from google.adk.agents import BaseAgent - - ADKAgent = BaseAgent -except (ImportError, AttributeError): - ADKAgent = None - -try: - from a2a.types import ( - AgentCard, - AgentInterface, - Message, - TaskIdParams, - TaskQueryParams, - ) - from a2a.utils.constants import TransportProtocol, PROTOCOL_VERSION_CURRENT - from a2a.client import ClientConfig, ClientFactory - - AgentCard = AgentCard - AgentInterface = AgentInterface - TransportProtocol = TransportProtocol - PROTOCOL_VERSION_CURRENT = PROTOCOL_VERSION_CURRENT - Message = Message - ClientConfig = ClientConfig - ClientFactory = ClientFactory - TaskIdParams = TaskIdParams - TaskQueryParams = TaskQueryParams -except (ImportError, AttributeError): - AgentCard = None - AgentInterface = None - TransportProtocol = None - PROTOCOL_VERSION_CURRENT = None - Message = None - ClientConfig = None - ClientFactory = None - TaskIdParams = None - TaskQueryParams = None - - -@typing.runtime_checkable -class Queryable(Protocol): - """Protocol for Agent Engines that can be queried.""" - - @abc.abstractmethod - def query(self, **kwargs) -> Any: - """Runs the Agent Engine to serve the user query.""" - - -@typing.runtime_checkable -class AsyncQueryable(Protocol): - """Protocol for Agent Engines that can be queried asynchronously.""" - - @abc.abstractmethod - def async_query(self, **kwargs) -> Coroutine[Any, Any, Any]: - """Runs the Agent Engine to serve the user query asynchronously.""" - - -@typing.runtime_checkable -class AsyncStreamQueryable(Protocol): - """Protocol for Agent Engines that can stream responses asynchronously.""" - - @abc.abstractmethod - async def async_stream_query(self, **kwargs) -> AsyncIterable[Any]: - """Asynchronously stream responses to serve the user query.""" - - -@typing.runtime_checkable -class StreamQueryable(Protocol): - """Protocol for Agent Engines that can stream responses.""" - - @abc.abstractmethod - def stream_query(self, **kwargs) -> Iterable[Any]: - """Stream responses to serve the user query.""" - - -@typing.runtime_checkable -class BidiStreamQueryable(Protocol): - """Protocol for Agent Engines that can stream requests and responses.""" - - @abc.abstractmethod - async def bidi_stream_query(self, **kwargs) -> AsyncIterable[Any]: - """Asynchronously stream requests and responses to serve the user query.""" - - -@typing.runtime_checkable -class Cloneable(Protocol): - """Protocol for Agent Engines that can be cloned.""" - - @abc.abstractmethod - def clone(self) -> Any: - """Return a clone of the object.""" - - -@typing.runtime_checkable -class OperationRegistrable(Protocol): - """Protocol for agents that have registered operations.""" - - @abc.abstractmethod - def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]: - """Register the user provided operations (modes and methods).""" - - -_AgentEngineInterface = Union[ - ADKAgent, - AsyncQueryable, - AsyncStreamQueryable, - BidiStreamQueryable, - OperationRegistrable, - Queryable, - StreamQueryable, -] - - -def _wrap_agent_operation(agent: Any, operation: str): - """Wraps an agent operation into a method (works for all API modes).""" - - def _method(self, **kwargs): - if not self._tmpl_attrs.get("agent"): - self.set_up() - return getattr(self._tmpl_attrs["agent"], operation)(**kwargs) - - _method.__name__ = operation - _method.__doc__ = getattr(agent, operation).__doc__ - return _method - - -class ModuleAgent(Cloneable, OperationRegistrable): - """Agent that is defined by a module and an agent name. - - This agent is instantiated by importing a module and instantiating an agent - from that module. It also allows to register operations that are defined in - the agent. - """ - - def __init__( - self, - *, - module_name: str, - agent_name: str, - register_operations: Dict[str, Sequence[str]], - sys_paths: Optional[Sequence[str]] = None, - agent_framework: Optional[str] = None, - ): - """Initializes a module-based agent. - - Args: - module_name (str): - Required. The name of the module to import. - agent_name (str): - Required. The name of the agent in the module to instantiate. - register_operations (Dict[str, Sequence[str]]): - Required. A dictionary of API modes to a list of method names. - sys_paths (Sequence[str]): - Optional. The system paths to search for the module. It should - be relative to the directory where the code will be running. - I.e. it should correspond to the directory being passed to - `extra_packages=...` in the create method. It will be appended - to the system path in the sequence being specified here, and - only be appended if it is not already in the system path. - """ - self.agent_framework = agent_framework - self._tmpl_attrs = { - "module_name": module_name, - "agent_name": agent_name, - "register_operations": register_operations, - "sys_paths": sys_paths, - } - - def clone(self): - """Return a clone of the agent.""" - return ModuleAgent( - module_name=self._tmpl_attrs.get("module_name"), - agent_name=self._tmpl_attrs.get("agent_name"), - register_operations=self._tmpl_attrs.get("register_operations"), - sys_paths=self._tmpl_attrs.get("sys_paths"), - agent_framework=self.agent_framework, - ) - - def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]: - return self._tmpl_attrs.get("register_operations") - - def set_up(self) -> None: - """Sets up the agent for execution of queries at runtime. - - It runs the code to import the agent from the module, and registers the - operations of the agent. - """ - if self._tmpl_attrs.get("sys_paths"): - import sys - - for sys_path in self._tmpl_attrs.get("sys_paths"): - abs_path = os.path.abspath(sys_path) - if abs_path not in sys.path: - sys.path.append(abs_path) - - import importlib - - module = importlib.import_module(self._tmpl_attrs.get("module_name")) - try: - importlib.reload(module) - except Exception as e: - _LOGGER.warning( - f"Failed to reload module {self._tmpl_attrs.get('module_name')}: {e}" - ) - agent_name = self._tmpl_attrs.get("agent_name") - try: - agent = getattr(module, agent_name) - except AttributeError as e: - raise AttributeError( - f"Agent {agent_name} not found in module " - f"{self._tmpl_attrs.get('module_name')}" - ) from e - if not self.agent_framework: - self.agent_framework = _get_agent_framework(agent) - self._tmpl_attrs["agent"] = agent - if hasattr(agent, "set_up"): - agent.set_up() - for operations in self.register_operations().values(): - for operation in operations: - op = _wrap_agent_operation(agent, operation) - setattr(self, operation, types.MethodType(op, self)) - - -class AgentEngine(base.VertexAiResourceNounWithFutureManager): - """Represents a Vertex AI Agent Engine resource.""" - - client_class = aip_utils.AgentEngineClientWithOverride - _resource_noun = "reasoning_engine" - _getter_method = "get_reasoning_engine" - _list_method = "list_reasoning_engines" - _delete_method = "delete_reasoning_engine" - _parse_resource_name_method = "parse_reasoning_engine_path" - _format_resource_name_method = "reasoning_engine_path" - - def __init__(self, resource_name: str): - """Retrieves an Agent Engine resource. - - Args: - resource_name (str): - Required. A fully-qualified resource name or ID such as - "projects/123/locations/us-central1/reasoningEngines/456" or - "456" when project and location are initialized or passed. - """ - super().__init__(resource_name=resource_name) - self.execution_api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionClientWithOverride, - ) - self.execution_async_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionAsyncClientWithOverride, - ) - self._gca_resource = self._get_gca_resource(resource_name=resource_name) - try: - _register_api_methods_or_raise(self) - except Exception as e: - _LOGGER.warning(_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e) - self._operation_schemas = None - - @property - def resource_name(self) -> str: - """Fully-qualified resource name.""" - return self._gca_resource.name - - @classmethod - def create( - cls, - agent_engine: Optional[_AgentEngineInterface] = None, - *, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - build_options: Optional[Dict[str, Sequence[str]]] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[aip_types.EncryptionSpec] = None, - ) -> "AgentEngine": - """Creates a new Agent Engine. - - The Agent Engine will be an instance of the `agent_engine` that - was passed in, running remotely on Vertex AI. - - Sample `src_dir` contents (e.g. `./user_src_dir`): - - .. code-block:: python - - user_src_dir/ - |-- main.py - |-- requirements.txt - |-- user_code/ - | |-- utils.py - | |-- ... - |-- installation_scripts/ - | |-- install_package.sh - | |-- ... - |-- ... - - To build an Agent Engine with the above files, run: - - .. code-block:: python - - remote_agent = agent_engines.create( - agent_engine=local_agent, - requirements=[ - # I.e. the PyPI dependencies listed in requirements.txt - "google-cloud-aiplatform==1.25.0", - "langchain==0.0.242", - ... - ], - extra_packages=[ - "./user_src_dir/main.py", # a single file - "./user_src_dir/user_code", # a directory - ... - ], - build_options={ - "installation_scripts": [ - "./user_src_dir/installation_scripts/install_package.sh", - ... - ], - }, - ) - - Args: - agent_engine (AgentEngineInterface): - Optional. The Agent Engine to be created. - requirements (Union[str, Sequence[str]]): - Optional. The set of PyPI dependencies needed. It can either be - the path to a single file (requirements.txt), or an ordered list - of strings corresponding to each line of the requirements file. - display_name (str): - Optional. The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any - UTF-8 character. - description (str): - Optional. The description of the Agent Engine. - gcs_dir_name (str): - Optional. The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - extra_packages (Sequence[str]): - Optional. The set of extra user-provided packages (if any). - env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): - Optional. The environment variables to be set when running the - Agent Engine. If it is a list of strings, each string should be - a valid key to `os.environ`. If it is a dictionary, the keys are - the environment variable names, and the values are the - corresponding values. - build_options (Dict[str, Sequence[str]]): - Optional. The build options for the Agent Engine. - The following keys are supported: - - installation_scripts: - Optional. The paths to the installation scripts to be - executed in the Docker image. - The scripts must be located in the `installation_scripts` - subdirectory and the path must be added to `extra_packages`. - service_account (str): - Optional. The service account to be used for the Agent Engine. - If not specified, the default reasoning engine service agent - service account will be used. - psc_interface_config (aip_types.PscInterfaceConfig): - Optional. The Private Service Connect interface config for the - Agent Engine. - min_instances (int): - Optional. The minimum number of instances to be running for the - Agent Engine. - max_instances (int): - Optional. The maximum number of instances to be running for the - Agent Engine. - resource_limits (Dict[str, str]): - Optional. The resource limits for the Agent Engine. - container_concurrency (int): - Optional. The container concurrency for the Agent Engine. - encryption_spec (aip_types.EncryptionSpec): - Optional. The Cloud KMS resource identifier of the customer - managed encryption key used to protect the model. Has the - form: - `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`. - The key needs to be in the same region as the model. - - Returns: - AgentEngine: The Agent Engine that was created. - - Raises: - ValueError: If the `project` was not set using `agentplatform.init`. - ValueError: If the `location` was not set using `agentplatform.init`. - ValueError: If the `staging_bucket` was not set using agentplatform.init. - ValueError: If the `staging_bucket` does not start with "gs://". - ValueError: If `extra_packages` is specified but `agent_engine` is None. - ValueError: If `requirements` is specified but `agent_engine` is None. - ValueError: If `env_vars` has a dictionary entry that does not - correspond to a SecretRef. - ValueError: If `env_vars` is a list which contains a string that - does not exist in `os.environ`. - TypeError: If `env_vars` is not a list of strings or a dictionary. - TypeError: If `env_vars` has a value that is not a string or SecretRef. - FileNotFoundError: If `extra_packages` includes a file or directory - that does not exist. - IOError: If requirements is a string that corresponds to a - nonexistent file. - """ - sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" - _validate_sys_version_or_raise(sys_version) - gcs_dir_name = gcs_dir_name or _DEFAULT_GCS_DIR_NAME - staging_bucket = initializer.global_config.staging_bucket - - if agent_engine is not None: - agent_engine = _validate_agent_engine_or_raise(agent_engine) - staging_bucket = _validate_staging_bucket_or_raise(staging_bucket) - if _is_adk_agent(None, agent_engine): - env_vars = _add_telemetry_enablement_env(env_vars=env_vars) - - if agent_engine is None: - if requirements is not None: - raise ValueError("requirements must be None if agent_engine is None.") - if extra_packages is not None: - raise ValueError("extra_packages must be None if agent_engine is None.") - requirements = _validate_requirements_or_raise( - agent_engine=agent_engine, - requirements=requirements, - ) - extra_packages = _validate_extra_packages_or_raise( - extra_packages=extra_packages, - build_options=build_options, - ) - - sdk_resource = cls.__new__(cls) - base.VertexAiResourceNounWithFutureManager.__init__(sdk_resource) - - # Prepares the Agent Engine for creation in Vertex AI. - # This involves packaging and uploading the artifacts for - # agent_engine, requirements and extra_packages to - # `staging_bucket/gcs_dir_name`. - _prepare( - agent_engine=agent_engine, - requirements=requirements, - project=sdk_resource.project, - location=sdk_resource.location, - staging_bucket=staging_bucket, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - ) - reasoning_engine = aip_types.ReasoningEngine( - display_name=display_name, - description=description, - encryption_spec=encryption_spec, - ) - if agent_engine is not None: - # Update the package spec. - package_spec = aip_types.ReasoningEngineSpec.PackageSpec( - python_version=sys_version, - pickle_object_gcs_uri="{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _BLOB_FILENAME, - ), - ) - if extra_packages: - package_spec.dependency_files_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _EXTRA_PACKAGES_FILE, - ) - if requirements: - package_spec.requirements_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _REQUIREMENTS_FILE, - ) - agent_engine_spec = aip_types.ReasoningEngineSpec( - package_spec=package_spec, - ) - if ( - env_vars - or psc_interface_config - or min_instances is not None - or max_instances is not None - or resource_limits - or container_concurrency is not None - ): - deployment_spec, _ = _generate_deployment_spec_or_raise( - env_vars=env_vars, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - ) - agent_engine_spec.deployment_spec = deployment_spec - class_methods_spec = _generate_class_methods_spec_or_raise( - agent_engine=agent_engine, - operations=_get_registered_operations(agent_engine), - ) - agent_engine_spec.class_methods.extend(class_methods_spec) - if service_account: - agent_engine_spec.service_account = service_account - reasoning_engine.spec = agent_engine_spec - reasoning_engine.spec.agent_framework = _get_agent_framework(agent_engine) - operation_future = sdk_resource.api_client.create_reasoning_engine( - parent=initializer.global_config.common_location_path( - project=sdk_resource.project, location=sdk_resource.location - ), - reasoning_engine=reasoning_engine, - ) - _LOGGER.log_create_with_lro(cls, operation_future) - _LOGGER.info( - f"View progress and logs at https://console.cloud.google.com/logs/query?project={sdk_resource.project}" - ) - created_resource = operation_future.result() - _LOGGER.info(f"{cls.__name__} created. Resource name: {created_resource.name}") - _LOGGER.info(f"To use this {cls.__name__} in another session:") - _LOGGER.info( - f"agent_engine = agentplatform.agent_engines.get('{created_resource.name}')" - ) - # We use `._get_gca_resource(...)` instead of `created_resource` to - # fully instantiate the attributes of the agent engine. - sdk_resource._gca_resource = sdk_resource._get_gca_resource( - resource_name=created_resource.name - ) - sdk_resource.execution_api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionClientWithOverride, - credentials=sdk_resource.credentials, - location_override=sdk_resource.location, - ) - sdk_resource.execution_async_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionAsyncClientWithOverride, - credentials=sdk_resource.credentials, - location_override=sdk_resource.location, - ) - if agent_engine is not None: - try: - _register_api_methods_or_raise(sdk_resource) - except Exception as e: - _LOGGER.warning(_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e) - sdk_resource._operation_schemas = None - return sdk_resource - - def update( - self, - *, - agent_engine: Optional[_AgentEngineInterface] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - build_options: Optional[Dict[str, Sequence[str]]] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[aip_types.EncryptionSpec] = None, - ) -> "AgentEngine": - """Updates an existing Agent Engine. - - This method updates the configuration of an existing Agent Engine - running remotely, which is identified by its resource name. - Unlike the `create` function which requires a `agent_engine` object, - all arguments in this method are optional. - This method allows you to modify individual aspects of the configuration - by providing any of the optional arguments. - - Args: - agent_engine (AgentEngineInterface): - Optional. The instance to be used as the updated Agent Engine. - If it is not specified, the existing instance will be used. - requirements (Union[str, Sequence[str]]): - Optional. The set of PyPI dependencies needed. It can either be - the path to a single file (requirements.txt), or an ordered list - of strings corresponding to each line of the requirements file. - If it is not specified, the existing requirements will be used. - If it is set to an empty string or list, the existing - requirements will be removed. - display_name (str): - Optional. The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any - UTF-8 character. - description (str): - Optional. The description of the Agent Engine. - gcs_dir_name (str): - Optional. The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - extra_packages (Sequence[str]): - Optional. The set of extra user-provided packages (if any). If - it is not specified, the existing extra packages will be used. - If it is set to an empty list, the existing extra packages will - be removed. - env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): - Optional. The environment variables to be set when running the - Agent Engine. If it is a list of strings, each string should be - a valid key to `os.environ`. If it is a dictionary, the keys are - the environment variable names, and the values are the - corresponding values. - build_options (Dict[str, Sequence[str]]): - Optional. The build options for the Agent Engine. - The following keys are supported: - - installation_scripts: - Optional. The paths to the installation scripts to be - executed in the Docker image. - The scripts must be located in the `installation_scripts` - subdirectory and the path must be added to `extra_packages`. - service_account (str): - Optional. The service account to be used for the Agent Engine. - If not specified, the default reasoning engine service agent - service account will be used. - psc_interface_config (aip_types.PscInterfaceConfig): - Optional. The Private Service Connect interface config for the - Agent Engine. - min_instances (int): - Optional. The minimum number of instances to be running for the - Agent Engine. - max_instances (int): - Optional. The maximum number of instances to be running for the - Agent Engine. - resource_limits (Dict[str, str]): - Optional. The resource limits for the Agent Engine. - container_concurrency (int): - Optional. The container concurrency for the Agent Engine. - encryption_spec (aip_types.EncryptionSpec): - Optional. The Cloud KMS resource identifier of the customer - managed encryption key used to protect the model. Has the - form: - `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`. - The key needs to be in the same region as the model. - - Returns: - AgentEngine: The Agent Engine that was updated. - - Raises: - ValueError: If the `staging_bucket` was not set using agentplatform.init. - ValueError: If the `staging_bucket` does not start with "gs://". - ValueError: If `env_vars` has a dictionary entry that does not - correspond to a SecretRef. - ValueError: If `env_vars` is a list which contains a string that - does not exist in `os.environ`. - TypeError: If `env_vars` is not a list of strings or a dictionary. - TypeError: If `env_vars` has a value that is not a string or SecretRef. - FileNotFoundError: If `extra_packages` includes a file or directory - that does not exist. - ValueError: if none of `display_name`, `description`, `requirements`, - `extra_packages`, `env_vars`, or `agent_engine` were specified. - IOError: If requirements is a string that corresponds to a - nonexistent file. - """ - staging_bucket = initializer.global_config.staging_bucket - staging_bucket = _validate_staging_bucket_or_raise(staging_bucket) - historical_operation_schemas = self.operation_schemas() - gcs_dir_name = gcs_dir_name or _DEFAULT_GCS_DIR_NAME - - # Validate the arguments. - if not any( - [ - agent_engine, - requirements, - extra_packages, - display_name, - description, - env_vars, - build_options, - service_account, - psc_interface_config, - min_instances is not None, - max_instances is not None, - resource_limits, - container_concurrency is not None, - encryption_spec, - ] - ): - raise ValueError( - "At least one of `agent_engine`, `requirements`, " - "`extra_packages`, `display_name`, `description`, " - "`env_vars`, `build_options`, `service_account`, " - "`psc_interface_config`, `min_instances`, `max_instances`, " - "`resource_limits`, `container_concurrency`, or " - "`encryption_spec` must be specified." - ) - if requirements is not None: - requirements = _validate_requirements_or_raise( - agent_engine=agent_engine, - requirements=requirements, - ) - if extra_packages is not None: - extra_packages = _validate_extra_packages_or_raise( - extra_packages=extra_packages, - build_options=build_options, - ) - if agent_engine is not None: - agent_engine = _validate_agent_engine_or_raise(agent_engine) - - if _is_adk_agent(self, agent_engine): - env_vars = _add_telemetry_enablement_env(env_vars=env_vars) - - # Prepares the Agent Engine for update in Vertex AI. This involves - # packaging and uploading the artifacts for agent_engine, requirements - # and extra_packages to `staging_bucket/gcs_dir_name`. - _prepare( - agent_engine=agent_engine, - requirements=requirements, - project=self.project, - location=self.location, - staging_bucket=staging_bucket, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - ) - update_request = _generate_update_request_or_raise( - resource_name=self.resource_name, - staging_bucket=staging_bucket, - gcs_dir_name=gcs_dir_name, - agent_engine=agent_engine, - requirements=requirements, - extra_packages=extra_packages, - display_name=display_name, - description=description, - env_vars=env_vars, - service_account=service_account, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - encryption_spec=encryption_spec, - ) - operation_future = self.api_client.update_reasoning_engine( - request=update_request - ) - _LOGGER.info( - f"Update Agent Engine backing LRO: {operation_future.operation.name}" - ) - created_resource = operation_future.result() - _LOGGER.info(f"Agent Engine updated. Resource name: {created_resource.name}") - self._operation_schemas = None - self.execution_api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionClientWithOverride, - ) - # We use `._get_gca_resource(...)` instead of `created_resource` to - # fully instantiate the attributes of the agent engine. - self._gca_resource = self._get_gca_resource(resource_name=self.resource_name) - - if ( - agent_engine is None - or historical_operation_schemas == self.operation_schemas() - ): - # The operations of the agent engine are unchanged, so we return it. - return self - - # If the agent engine has changed and the historical operation - # schemas are different from the current operation schemas, we need to - # unregister the historical operation schemas and register the current - # operation schemas. - _unregister_api_methods(self, historical_operation_schemas) - try: - _register_api_methods_or_raise(self) - except Exception as e: - _LOGGER.warning(_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e) - return self - - def delete( - self, - *, - force: bool = False, - **kwargs, - ) -> None: - """Deletes the ReasoningEngine. - - Args: - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error - when the Agent Engine has undeleted child resources. Defaults to - False. - **kwargs (dict[str, Any]): - Optional. Additional keyword arguments to pass to the - delete_reasoning_engine method. - """ - kwargs = kwargs or {} - operation_future = self.api_client.delete_reasoning_engine( - request=aip_types.DeleteReasoningEngineRequest( - name=self.resource_name, - force=force, - **kwargs, - ), - ) - _LOGGER.info( - f"Delete Agent Engine backing LRO: {operation_future.operation.name}" - ) - operation_future.result() - _LOGGER.info(f"Agent Engine deleted. Resource name: {self.resource_name}") - - def operation_schemas(self) -> Sequence[_agent_engines_utils.JsonDict]: - """Returns the (Open)API schemas for the Agent Engine.""" - spec = _agent_engines_utils._to_dict(self._gca_resource.spec) - if not hasattr(self, "_operation_schemas") or self._operation_schemas is None: - self._operation_schemas = spec.get("class_methods", []) - return self._operation_schemas - - -def _validate_sys_version_or_raise(sys_version: str) -> None: - """Tries to validate the python system version.""" - if sys_version not in _SUPPORTED_PYTHON_VERSIONS: - raise ValueError( - f"Unsupported python version: {sys_version}. AgentEngine " - f"only supports {_SUPPORTED_PYTHON_VERSIONS} at the moment." - ) - if sys_version != f"{sys.version_info.major}.{sys.version_info.minor}": - _LOGGER.warning( - f"{sys_version=} is inconsistent with {sys.version_info=}. " - "This might result in issues with deployment, and should only " - "be used as a workaround for advanced cases." - ) - - -def _validate_staging_bucket_or_raise(staging_bucket: Optional[str]) -> str: - """Tries to validate the staging bucket.""" - if not staging_bucket: - raise ValueError( - "Please provide a `staging_bucket` in `agentplatform.init(...)`" - ) - if not staging_bucket.startswith("gs://"): - raise ValueError(f"{staging_bucket=} must start with `gs://`") - return staging_bucket - - -def _validate_agent_engine_or_raise( - agent_engine: _AgentEngineInterface, - logger: base.Logger = _LOGGER, -) -> _AgentEngineInterface: - """Tries to validate the agent engine. - - The agent engine must have one of the following: - * a callable method named `query` - * a callable method named `stream_query` - * a callable method named `async_stream_query` - * a callable method named `bidi_stream_query` - * a callable method named `register_operations` - - Args: - agent_engine: The agent engine to be validated. - logger: The logger to use for logging. - - Returns: - The validated agent engine. - - Raises: - TypeError: If `agent_engine` has no callable method named `query`, - `stream_query` or `register_operations`. - ValueError: If `agent_engine` has an invalid `query`, `stream_query` or - `register_operations` signature. - """ - try: - from google.adk.agents import BaseAgent - - if isinstance(agent_engine, BaseAgent): - logger.info("Deploying google.adk.agents.Agent as an application.") - from agentplatform import agent_engines - - agent_engine = agent_engines.AdkApp(agent=agent_engine) - except Exception: - pass - is_queryable = isinstance(agent_engine, Queryable) and callable(agent_engine.query) - is_async_queryable = isinstance(agent_engine, AsyncQueryable) and callable( - agent_engine.async_query - ) - is_stream_queryable = isinstance(agent_engine, StreamQueryable) and callable( - agent_engine.stream_query - ) - is_async_stream_queryable = isinstance( - agent_engine, AsyncStreamQueryable - ) and callable(agent_engine.async_stream_query) - is_bidi_stream_queryable = isinstance( - agent_engine, BidiStreamQueryable - ) and callable(agent_engine.bidi_stream_query) - is_operation_registrable = isinstance( - agent_engine, OperationRegistrable - ) and callable(agent_engine.register_operations) - - if not ( - is_queryable - or is_async_queryable - or is_stream_queryable - or is_operation_registrable - or is_async_stream_queryable - or is_bidi_stream_queryable - ): - raise TypeError( - "agent_engine has none of the following callable methods: " - "`query`, `async_query`, `stream_query`, `async_stream_query`, " - "`bidi_stream_query` or `register_operations`." - ) - - if is_queryable: - try: - inspect.signature(getattr(agent_engine, "query")) - except ValueError as err: - raise ValueError( - "Invalid query signature. This might be due to a missing " - "`self` argument in the agent_engine.query method." - ) from err - - if is_async_queryable: - try: - inspect.signature(getattr(agent_engine, "async_query")) - except ValueError as err: - raise ValueError( - "Invalid async_query signature. This might be due to a missing " - "`self` argument in the agent_engine.async_query method." - ) from err - - if is_stream_queryable: - try: - inspect.signature(getattr(agent_engine, "stream_query")) - except ValueError as err: - raise ValueError( - "Invalid stream_query signature. This might be due to a missing" - " `self` argument in the agent_engine.stream_query method." - ) from err - - if is_async_stream_queryable: - try: - inspect.signature(getattr(agent_engine, "async_stream_query")) - except ValueError as err: - raise ValueError( - "Invalid async_stream_query signature. This might be due to a " - " missing `self` argument in the " - "agent_engine.async_stream_query method." - ) from err - - if is_bidi_stream_queryable: - try: - inspect.signature(getattr(agent_engine, "bidi_stream_query")) - except ValueError as err: - raise ValueError( - "Invalid bidi_stream_query signature. This might be due to a " - " missing `self` argument in the " - "agent_engine.bidi_stream_query method." - ) from err - - if is_operation_registrable: - try: - inspect.signature(getattr(agent_engine, "register_operations")) - except ValueError as err: - raise ValueError( - "Invalid register_operations signature. This might be due to a " - "missing `self` argument in the " - "agent_engine.register_operations method." - ) from err - - if isinstance(agent_engine, Cloneable): - # Avoid undeployable states. - agent_engine = agent_engine.clone() - return agent_engine - - -def _is_adk_agent( - agent_engine_to_update: Optional[AgentEngine], - new_agent_engine: Optional[_AgentEngineInterface], -) -> bool: - """Checks if the agent engine is an ADK agent. - - Args: - agent_engine_to_update: Existing agent engine, None if creating new one. - new_agent_engine: The new agent engine to deploy. Can be None during an update, if the Python agent implementation is not provided, and should remain unchanged. - - Returns: - True if the agent after the create/update operation, will be an ADK agent. - """ - - from agentplatform.agent_engines.templates import adk - - if new_agent_engine is not None: - return ( - getattr(new_agent_engine, "agent_framework", None) - == adk.AdkApp.agent_framework - ) - if agent_engine_to_update is not None: - return ( - agent_engine_to_update.gca_resource.spec.agent_framework - == adk.AdkApp.agent_framework - ) - return False - - -EnvVars = Optional[Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]] - - -def _add_telemetry_enablement_env(*, env_vars: EnvVars) -> EnvVars: - """Adds telemetry enablement env var to the env vars. - - This is in order to achieve default-on telemetry. - If the telemetry enablement env var is already set, we do not override it. - - Args: - env_vars: The env vars to add the telemetry enablement env var to. - - Returns: - The env vars with the telemetry enablement env var added. - """ - - GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = ( - "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" - ) - - if env_vars is None: - return {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"} - if isinstance(env_vars, dict): - return ( - env_vars - if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in env_vars - else env_vars | {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"} - ) - if isinstance(env_vars, list) or isinstance(env_vars, tuple): - if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY not in os.environ: - os.environ[GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY] = "unspecified" - - if isinstance(env_vars, list): - return env_vars + [GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY] - else: - return env_vars + (GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY,) - - raise TypeError( - f"env_vars must be a list, tuple or a dict, but got {type(env_vars)}." - ) - - -def _validate_requirements_or_raise( - *, - agent_engine: _AgentEngineInterface, - requirements: Optional[Sequence[str]] = None, - logger: logging.getLoggerClass() = _LOGGER, -) -> Sequence[str]: - """Tries to validate the requirements.""" - if requirements is None: - requirements = [] - elif isinstance(requirements, str): - try: - logger.info(f"Reading requirements from {requirements=}") - with open(requirements) as f: - requirements = f.read().splitlines() - logger.info(f"Read the following lines: {requirements}") - except IOError as err: - raise IOError(f"Failed to read requirements from {requirements=}") from err - requirements = _agent_engines_utils._validate_requirements_or_warn( - obj=agent_engine, - requirements=requirements, - ) - logger.info(f"The final list of requirements: {requirements}") - return requirements - - -def _validate_extra_packages_or_raise( - extra_packages: Optional[Sequence[str]], - build_options: Optional[Dict[str, Sequence[str]]] = None, -) -> Sequence[str]: - """Tries to validates the extra packages.""" - extra_packages = extra_packages or [] - if build_options and _BUILD_OPTIONS_INSTALLATION in build_options: - _agent_engines_utils._validate_installation_scripts_or_raise( - script_paths=build_options[_BUILD_OPTIONS_INSTALLATION], - packages=extra_packages, - ) - for extra_package in extra_packages: - if not os.path.exists(extra_package): - raise FileNotFoundError( - f"Extra package specified but not found: {extra_package=}" - ) - return extra_packages - - -def _get_gcs_bucket( - *, - project: str, - location: str, - staging_bucket: str, - logger: base.Logger = _LOGGER, -) -> storage.Bucket: - """Gets or creates the GCS bucket.""" - storage = _agent_engines_utils._import_cloud_storage_or_raise() - storage_client = storage.Client(project=project) - staging_bucket = staging_bucket.replace("gs://", "") - try: - gcs_bucket = storage_client.get_bucket(staging_bucket) - logger.info(f"Using bucket {staging_bucket}") - except exceptions.NotFound: - new_bucket = storage_client.bucket(staging_bucket) - gcs_bucket = storage_client.create_bucket(new_bucket, location=location) - logger.info(f"Creating bucket {staging_bucket} in {location=}") - return gcs_bucket - - -def _upload_agent_engine( - *, - agent_engine: _AgentEngineInterface, - gcs_bucket: storage.Bucket, - gcs_dir_name: str, - logger: base.Logger = _LOGGER, -) -> None: - """Uploads the agent engine to GCS.""" - cloudpickle = _agent_engines_utils._import_cloudpickle_or_raise() - blob = gcs_bucket.blob(f"{gcs_dir_name}/{_BLOB_FILENAME}") - with blob.open("wb") as f: - try: - cloudpickle.dump(agent_engine, f) - except Exception as e: - url = "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#deployment-considerations" - error_msg = f"Failed to serialize agent engine. Visit {url} for details." - if "google._upb._message" in str(e) or "Descriptor" in str(e): - error_msg += ( - " This is often caused by protobuf objects (like Part, AgentCard) " - "being imported at the global module level. Please move these " - "imports inside the functions or methods where they are used. " - "Alternatively, you can import the entire module: " - "`from a2a import types as a2a_types`." - ) - raise TypeError(error_msg) from e - with blob.open("rb") as f: - try: - _ = cloudpickle.load(f) - except Exception as e: - raise TypeError("Agent engine serialized to an invalid format") from e - dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" - logger.info(f"Wrote to {dir_name}/{_BLOB_FILENAME}") - - -def _upload_requirements( - *, - requirements: Sequence[str], - gcs_bucket: storage.Bucket, - gcs_dir_name: str, - logger: base.Logger = _LOGGER, -) -> None: - """Uploads the requirements file to GCS.""" - blob = gcs_bucket.blob(f"{gcs_dir_name}/{_REQUIREMENTS_FILE}") - blob.upload_from_string("\n".join(requirements)) - dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" - logger.info(f"Writing to {dir_name}/{_REQUIREMENTS_FILE}") - - -def _upload_extra_packages( - *, - extra_packages: Sequence[str], - gcs_bucket: storage.Bucket, - gcs_dir_name: str, - logger: base.Logger = _LOGGER, -) -> None: - """Uploads extra packages to GCS.""" - logger.info("Creating in-memory tarfile of extra_packages") - tar_fileobj = io.BytesIO() - with tarfile.open(fileobj=tar_fileobj, mode="w|gz") as tar: - for file in extra_packages: - tar.add(file) - tar_fileobj.seek(0) - blob = gcs_bucket.blob(f"{gcs_dir_name}/{_EXTRA_PACKAGES_FILE}") - blob.upload_from_string(tar_fileobj.read()) - dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" - logger.info(f"Writing to {dir_name}/{_EXTRA_PACKAGES_FILE}") - - -def _prepare( - agent_engine: Optional[_AgentEngineInterface], - requirements: Optional[Sequence[str]], - extra_packages: Optional[Sequence[str]], - project: str, - location: str, - staging_bucket: str, - gcs_dir_name: str, - logger: base.Logger = _LOGGER, -) -> None: - """Prepares the agent engine for creation or updates in Vertex AI. - - This involves packaging and uploading artifacts to Cloud Storage. Note that - 1. This does not actually update the Agent Engine in Vertex AI. - 2. This will only generate and upload a pickled object if specified. - 3. This will only generate and upload the dependencies.tar.gz file if - extra_packages is non-empty. - - Args: - agent_engine: The agent engine to be prepared. - requirements (Sequence[str]): The set of PyPI dependencies needed. - extra_packages (Sequence[str]): The set of extra user-provided packages. - project (str): The project for the staging bucket. - location (str): The location for the staging bucket. - staging_bucket (str): The staging bucket name in the form "gs://...". - gcs_dir_name (str): The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - """ - if agent_engine is None: - return - gcs_bucket = _get_gcs_bucket( - project=project, - location=location, - staging_bucket=staging_bucket, - logger=logger, - ) - _upload_agent_engine( - agent_engine=agent_engine, - gcs_bucket=gcs_bucket, - gcs_dir_name=gcs_dir_name, - logger=logger, - ) - if requirements is not None: - _upload_requirements( - requirements=requirements, - gcs_bucket=gcs_bucket, - gcs_dir_name=gcs_dir_name, - logger=logger, - ) - if extra_packages is not None: - _upload_extra_packages( - extra_packages=extra_packages, - gcs_bucket=gcs_bucket, - gcs_dir_name=gcs_dir_name, - logger=logger, - ) - - -def _update_deployment_spec_with_env_vars_dict_or_raise( - *, - deployment_spec: aip_types.ReasoningEngineSpec.DeploymentSpec, - env_vars: Dict[str, Union[str, aip_types.SecretRef]], -) -> None: - for key, value in env_vars.items(): - if isinstance(value, Dict): - try: - secret_ref = _agent_engines_utils._to_proto( - value, aip_types.SecretRef() - ) - except Exception as e: - raise ValueError(f"Failed to convert to secret ref: {value}") from e - deployment_spec.secret_env.append( - aip_types.SecretEnvVar(name=key, secret_ref=secret_ref) - ) - elif isinstance(value, aip_types.SecretRef): - deployment_spec.secret_env.append( - aip_types.SecretEnvVar(name=key, secret_ref=value) - ) - elif isinstance(value, str): - deployment_spec.env.append(aip_types.EnvVar(name=key, value=value)) - else: - raise TypeError( - f"Unknown value type in env_vars for {key}. " - f"Must be a str or SecretRef: {value}" - ) - - -def _update_deployment_spec_with_env_vars_list_or_raise( - *, - deployment_spec: aip_types.ReasoningEngineSpec.DeploymentSpec, - env_vars: Sequence[str], -) -> None: - for env_var in env_vars: - if env_var not in os.environ: - raise ValueError(f"Env var not found in os.environ: {env_var}.") - deployment_spec.env.append( - aip_types.EnvVar(name=env_var, value=os.environ[env_var]) - ) - - -def _generate_deployment_spec_or_raise( - *, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, -) -> Tuple[aip_types.ReasoningEngineSpec.DeploymentSpec, List[str]]: - deployment_spec = aip_types.ReasoningEngineSpec.DeploymentSpec() - update_masks = [] - if env_vars: - deployment_spec.env = [] - deployment_spec.secret_env = [] - if isinstance(env_vars, Dict): - _update_deployment_spec_with_env_vars_dict_or_raise( - deployment_spec=deployment_spec, - env_vars=env_vars, - ) - elif isinstance(env_vars, Sequence): - _update_deployment_spec_with_env_vars_list_or_raise( - deployment_spec=deployment_spec, - env_vars=env_vars, - ) - else: - raise TypeError( - f"env_vars must be a list, tuple or a dict, but got {type(env_vars)}." - ) - if deployment_spec.env: - update_masks.append("spec.deployment_spec.env") - if deployment_spec.secret_env: - update_masks.append("spec.deployment_spec.secret_env") - if psc_interface_config: - deployment_spec.psc_interface_config = psc_interface_config - update_masks.append("spec.deployment_spec.psc_interface_config") - if min_instances is not None: - deployment_spec.min_instances = min_instances - update_masks.append("spec.deployment_spec.min_instances") - if max_instances is not None: - deployment_spec.max_instances = max_instances - update_masks.append("spec.deployment_spec.max_instances") - if resource_limits: - deployment_spec.resource_limits = resource_limits - update_masks.append("spec.deployment_spec.resource_limits") - if container_concurrency is not None: - deployment_spec.container_concurrency = container_concurrency - update_masks.append("spec.deployment_spec.container_concurrency") - return deployment_spec, update_masks - - -def _get_agent_framework( - agent_engine: _AgentEngineInterface, -) -> str: - if ( - hasattr(agent_engine, _AGENT_FRAMEWORK_ATTR) - and getattr(agent_engine, _AGENT_FRAMEWORK_ATTR) is not None - and isinstance(getattr(agent_engine, _AGENT_FRAMEWORK_ATTR), str) - ): - return getattr(agent_engine, _AGENT_FRAMEWORK_ATTR) - return _DEFAULT_AGENT_FRAMEWORK - - -def _generate_update_request_or_raise( - *, - resource_name: str, - staging_bucket: str, - gcs_dir_name: str = _DEFAULT_GCS_DIR_NAME, - agent_engine: Optional[_AgentEngineInterface] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - extra_packages: Optional[Sequence[str]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[str] = None, -) -> reasoning_engine_service.UpdateReasoningEngineRequest: - """Tries to generate the update request for the agent engine.""" - is_spec_update = False - update_masks: List[str] = [] - agent_engine_spec = aip_types.ReasoningEngineSpec() - package_spec = aip_types.ReasoningEngineSpec.PackageSpec() - if requirements is not None: - is_spec_update = True - update_masks.append("spec.package_spec.requirements_gcs_uri") - package_spec.requirements_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _REQUIREMENTS_FILE, - ) - if extra_packages is not None: - is_spec_update = True - update_masks.append("spec.package_spec.dependency_files_gcs_uri") - package_spec.dependency_files_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _EXTRA_PACKAGES_FILE, - ) - if agent_engine is not None: - is_spec_update = True - update_masks.append("spec.package_spec.pickle_object_gcs_uri") - package_spec.pickle_object_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _BLOB_FILENAME, - ) - class_methods_spec = _generate_class_methods_spec_or_raise( - agent_engine=agent_engine, - operations=_get_registered_operations(agent_engine), - ) - agent_engine_spec.class_methods.extend(class_methods_spec) - update_masks.append("spec.class_methods") - agent_engine_spec.agent_framework = _get_agent_framework(agent_engine) - update_masks.append("spec.agent_framework") - if ( - env_vars is not None - or psc_interface_config - or min_instances is not None - or max_instances is not None - or resource_limits - or container_concurrency is not None - ): - is_spec_update = True - deployment_spec, deployment_update_masks = _generate_deployment_spec_or_raise( - env_vars=env_vars, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - ) - update_masks.extend(deployment_update_masks) - agent_engine_spec.deployment_spec = deployment_spec - if service_account is not None: - is_spec_update = True - update_masks.append("spec.service_account") - agent_engine_spec.service_account = service_account - - agent_engine_message = aip_types.ReasoningEngine(name=resource_name) - if is_spec_update: - if package_spec: - agent_engine_spec.package_spec = package_spec - agent_engine_message.spec = agent_engine_spec - if display_name: - agent_engine_message.display_name = display_name - update_masks.append("display_name") - if description: - agent_engine_message.description = description - update_masks.append("description") - if encryption_spec: - agent_engine_message.encryption_spec = encryption_spec - update_masks.append("encryption_spec") - if not update_masks: - raise ValueError( - "At least one of `agent_engine`, `requirements`, `extra_packages`, " - "`display_name`, `description`, `env_vars`, or " - "`encryption_spec` must be specified." - ) - return reasoning_engine_service.UpdateReasoningEngineRequest( - reasoning_engine=agent_engine_message, - update_mask=field_mask_pb2.FieldMask(paths=update_masks), - ) - - -def _wrap_query_operation( - method_name: str, -) -> Callable[..., _agent_engines_utils.JsonDict]: - """Wraps an Agent Engine method, creating a callable for `query` API. - - This function creates a callable object that executes the specified - Agent Engine method using the `query` API. It handles the creation of - the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - doc: Documentation string for the method. - - Returns: - A callable object that executes the method on the Agent Engine via - the `query` API. - """ - - def _method(self, **kwargs) -> _agent_engines_utils.JsonDict: - response = self.execution_api_client.query_reasoning_engine( - request=aip_types.QueryReasoningEngineRequest( - name=self.resource_name, - input=kwargs, - class_method=method_name, - ), - ) - output = _agent_engines_utils._to_dict(response) - return output.get("output", output) - - return _method - - -def _wrap_async_query_operation(method_name: str) -> Callable[..., Coroutine]: - """Wraps an Agent Engine method, creating an async callable for `query` API. - - This function creates a callable object that executes the specified - Agent Engine method asynchronously using the `query` API. It handles the - creation of the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - doc: Documentation string for the method. - - Returns: - A callable object that executes the method on the Agent Engine via - the `query` API. - """ - - async def _method(self, **kwargs) -> _agent_engines_utils.JsonDict: - response = await self.execution_async_client.query_reasoning_engine( - request=aip_types.QueryReasoningEngineRequest( - name=self.resource_name, - input=kwargs, - class_method=method_name, - ), - ) - output = _agent_engines_utils._to_dict(response) - return output.get("output", output) - - return _method - - -def _wrap_stream_query_operation(*, method_name: str) -> Callable[..., Iterable[Any]]: - """Wraps an Agent Engine method, creating a callable for `stream_query` API. - - This function creates a callable object that executes the specified - Agent Engine method using the `stream_query` API. It handles the - creation of the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - doc: Documentation string for the method. - - Returns: - A callable object that executes the method on the Agent Engine via - the `stream_query` API. - """ - - def _method(self, **kwargs) -> Iterable[Any]: - response = self.execution_api_client.stream_query_reasoning_engine( - request=aip_types.StreamQueryReasoningEngineRequest( - name=self.resource_name, - input=kwargs, - class_method=method_name, - ), - ) - for chunk in response: - for parsed_json in _agent_engines_utils._yield_parsed_json_from_httpbody( - chunk - ): - if parsed_json is not None: - yield parsed_json - - return _method - - -def _wrap_async_stream_query_operation( - *, method_name: str -) -> Callable[..., AsyncIterable[Any]]: - """Wraps an Agent Engine method, creating an async callable for `stream_query` API. - - This function creates a callable object that executes the specified - Agent Engine method using the `stream_query` API. It handles the - creation of the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - doc: Documentation string for the method. - - Returns: - A callable object that executes the method on the Agent Engine via - the `stream_query` API. - """ - - async def _method(self, **kwargs) -> AsyncIterable[Any]: - response = self.execution_api_client.stream_query_reasoning_engine( - request=aip_types.StreamQueryReasoningEngineRequest( - name=self.resource_name, - input=kwargs, - class_method=method_name, - ), - ) - for chunk in response: - for parsed_json in _agent_engines_utils._yield_parsed_json_from_httpbody( - chunk - ): - if parsed_json is not None: - yield parsed_json - - return _method - - -def _wrap_bidi_stream_query_operation( - *, method_name: str -) -> Callable[..., AsyncIterable[Any]]: - """Wraps an Agent Engine method, creating an async callable for `bidi_stream_query` API. - - This function creates a callable object that executes the specified - Agent Engine method using the `bidi_stream_query` API. It handles the - creation of the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - - Returns: - A callable object that executes the method on the Agent Engine via - the `bidi_stream_query` API. - """ - - async def _method(self, **kwargs) -> AsyncIterable[Any]: - # Agent Engine bidi streaming query execution should use GenAI SDK Agent - # Engine live API client directly. - raise NotImplementedError( - f"{method_name} is not implemented, please use GenAI SDK Agent " - "Enginve live API client instead." - ) - - -def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list]: - """Wraps an Agent Engine method, creating a callable for A2A API. - - Args: - method_name: The name of the Agent Engine method to call. - agent_card: The agent card to use for the A2A API call. - Example: - {'additionalInterfaces': None, - 'capabilities': {'extensions': None, - 'pushNotifications': None, - 'stateTransitionHistory': None, - 'streaming': False}, - 'defaultInputModes': ['text'], - 'defaultOutputModes': ['text'], - 'description': ( - 'A helpful assistant agent that can answer questions.' - ), - 'documentationUrl': None, - 'iconUrl': None, - 'name': 'Q&A Agent', - 'preferredTransport': 'JSONRPC', - 'protocolVersion': '0.3.0', - 'provider': None, - 'security': None, - 'securitySchemes': None, - 'signatures': None, - 'skills': [{ - 'description': ( - 'A helpful assistant agent that can answer questions.' - ), - 'examples': ['Who is leading 2025 F1 Standings?', - 'Where can i find an active volcano?'], - 'id': 'question_answer', - 'inputModes': None, - 'name': 'Q&A Agent', - 'outputModes': None, - 'security': None, - 'tags': ['Question-Answer']}], - 'supportsAuthenticatedExtendedCard': True, - 'url': 'http://localhost:8080/', - 'version': '1.0.0'} - Returns: - A callable object that executes the method on the Agent Engine via - the A2A API. - """ - - async def _method(self, **kwargs) -> Any: - """Wraps an Agent Engine method, creating a callable for A2A API.""" - a2a_agent_card = AgentCard(**json.loads(agent_card)) - - # A2A + AE integration currently only supports Rest API. - if ( - a2a_agent_card.supported_interfaces - and a2a_agent_card.supported_interfaces[0].protocol_binding - != TransportProtocol.HTTP_JSON - ): - raise ValueError( - "Only HTTP+JSON is supported for primary interface on agent card " - ) - - # Set primary interface to HTTP+JSON if not set. - if not a2a_agent_card.supported_interfaces: - a2a_agent_card.supported_interfaces = [] - a2a_agent_card.supported_interfaces.append( - AgentInterface( - protocol_binding=TransportProtocol.HTTP_JSON, - protocol_version=PROTOCOL_VERSION_CURRENT, - ) - ) - - # AE cannot support streaming yet. Turn off streaming for now. - if a2a_agent_card.capabilities and a2a_agent_card.capabilities.streaming: - raise ValueError( - "Streaming is not supported in Agent Engine, please change " - "a2a_agent_card.capabilities.streaming to False." - ) - - if not hasattr(a2a_agent_card.capabilities, "streaming"): - a2a_agent_card.capabilities.streaming = False - - # agent_card is set on the class_methods before set_up is invoked. - # Ensure that the agent_card url is set correctly before the client is created. - url = f"https://{initializer.global_config.api_endpoint}/v1beta1/{self.resource_name}/a2a" - a2a_agent_card.supported_interfaces[0].url = url - - # Using a2a client, inject the auth token from the global config. - config = ClientConfig( - supported_transports=[ - TransportProtocol.HTTP_JSON, - ], - use_client_preference=True, - httpx_client=httpx.AsyncClient( - headers={ - "Authorization": ( - f"Bearer {initializer.global_config.credentials.token}" - ) - } - ), - ) - factory = ClientFactory(config) - client = factory.create(a2a_agent_card) - - # kokoro job uses python 3.9, replaced match with if else. - if method_name == "on_message_send": - response = client.send_message(Message(**kwargs)) - chunks = [] - async for chunk in response: - chunks.append(chunk) - return chunks - elif method_name == "on_get_task": - response = await client.get_task(TaskQueryParams(**kwargs)) - elif method_name == "on_cancel_task": - response = await client.cancel_task(TaskIdParams(**kwargs)) - elif method_name == "handle_authenticated_agent_card": - response = await client.get_card() - else: - raise ValueError(f"Unknown method name: {method_name}") - - return response - - return _method - - -def _unregister_api_methods( - obj: "AgentEngine", operation_schemas: Sequence[_agent_engines_utils.JsonDict] -): - """Unregisters Agent Engine API methods based on operation schemas. - - This function iterates through operation schemas provided by the - AgentEngine object. Each schema defines an API mode and method name. - It dynamically unregisters methods on the AgentEngine object. This - should only be used when updating the object. - - Args: - obj: The AgentEngine object to augment with API methods. - operation_schemas: The operation schemas to use for method unregistration. - """ - for operation_schema in operation_schemas: - if "name" in operation_schema: - method_name = operation_schema.get("name") - if hasattr(obj, method_name): - delattr(obj, method_name) - - -def _register_api_methods_or_raise( - obj: "AgentEngine", - wrap_operation_fn: Optional[ - dict[str, Callable[[str, str], Callable[..., Any]]] - ] = None, -): - """Registers Agent Engine API methods based on operation schemas. - - This function iterates through operation schemas provided by the - AgentEngine object. Each schema defines an API mode and method name. - It dynamically creates and registers methods on the AgentEngine object - to handle API calls based on the specified API mode. - Currently, only standard API mode `` is supported. - - Args: - obj: The AgentEngine object to augment with API methods. - wrap_operation_fn: A dictionary of API modes and method wrapping - functions. - - Raises: - ValueError: If the API mode is not supported or if the operation schema - is missing any required fields (e.g. `api_mode` or `name`). - """ - for operation_schema in obj.operation_schemas(): - if _MODE_KEY_IN_SCHEMA not in operation_schema: - raise ValueError( - f"Operation schema {operation_schema} does not" - f" contain an `{_MODE_KEY_IN_SCHEMA}` field." - ) - api_mode = operation_schema.get(_MODE_KEY_IN_SCHEMA) - if _METHOD_NAME_KEY_IN_SCHEMA not in operation_schema: - raise ValueError( - f"Operation schema {operation_schema} does not" - f" contain a `{_METHOD_NAME_KEY_IN_SCHEMA}` field." - ) - method_name = operation_schema.get(_METHOD_NAME_KEY_IN_SCHEMA) - method_description = operation_schema.get( - "description", - _DEFAULT_METHOD_DOCSTRING_TEMPLATE.format( - method_name=method_name, - default_method_name=_DEFAULT_METHOD_NAME_MAP.get( - api_mode, _DEFAULT_METHOD_NAME - ), - return_type=_DEFAULT_METHOD_RETURN_TYPE_MAP.get( - api_mode, - _DEFAULT_METHOD_RETURN_TYPE, - ), - ), - ) - _wrap_operation_map = { - _STANDARD_API_MODE: _wrap_query_operation, - _ASYNC_API_MODE: _wrap_async_query_operation, - _STREAM_API_MODE: _wrap_stream_query_operation, - _ASYNC_STREAM_API_MODE: _wrap_async_stream_query_operation, - _BIDI_STREAM_API_MODE: _wrap_bidi_stream_query_operation, - _A2A_EXTENSION_MODE: _wrap_a2a_operation, - } - if isinstance(wrap_operation_fn, dict) and api_mode in wrap_operation_fn: - # Override the default function with user-specified function if it exists. - _wrap_operation = wrap_operation_fn[api_mode] - elif api_mode in _wrap_operation_map: - _wrap_operation = _wrap_operation_map[api_mode] - else: - supported_api_modes = ", ".join( - f"`{mode}`" for mode in sorted(_wrap_operation_map.keys()) - ) - raise ValueError( - f"Unsupported api mode: `{api_mode}`," - f" Supported modes are: {supported_api_modes}." - ) - - # Bind the method to the object. - if api_mode == _A2A_EXTENSION_MODE: - agent_card = operation_schema.get(_A2A_AGENT_CARD) - method = _wrap_operation(method_name=method_name, agent_card=agent_card) - else: - method = _wrap_operation(method_name=method_name) - method.__name__ = method_name - method.__doc__ = method_description - setattr(obj, method_name, types.MethodType(method, obj)) - - -def _get_registered_operations( - agent_engine: _AgentEngineInterface, -) -> Dict[str, List[str]]: - """Retrieves registered operations for a AgentEngine.""" - if isinstance(agent_engine, OperationRegistrable): - return agent_engine.register_operations() - - operations = {} - if isinstance(agent_engine, Queryable): - operations[_STANDARD_API_MODE] = [_DEFAULT_METHOD_NAME] - if isinstance(agent_engine, AsyncQueryable): - operations[_ASYNC_API_MODE] = [_DEFAULT_ASYNC_METHOD_NAME] - if isinstance(agent_engine, StreamQueryable): - operations[_STREAM_API_MODE] = [_DEFAULT_STREAM_METHOD_NAME] - if isinstance(agent_engine, AsyncStreamQueryable): - operations[_ASYNC_STREAM_API_MODE] = [_DEFAULT_ASYNC_STREAM_METHOD_NAME] - return operations - - -def _generate_class_methods_spec_or_raise( - *, - agent_engine: _AgentEngineInterface, - operations: Dict[str, List[str]], - logger: base.Logger = _LOGGER, -) -> List[proto.Message]: - """Generates a ReasoningEngineSpec based on the registered operations. - - Args: - agent_engine: The AgentEngine instance. - operations: A dictionary of API modes and method names. - - Returns: - A list of ReasoningEngineSpec.ClassMethod messages. - - Raises: - ValueError: If a method defined in `register_operations` is not found on - the AgentEngine. - """ - if isinstance(agent_engine, ModuleAgent): - # We do a dry-run of setting up the agent engine to have the operations - # needed for registration. - agent_engine = agent_engine.clone() - try: - agent_engine.set_up() - except Exception as e: - raise ValueError( - f"Failed to set up agent engine {agent_engine}: {e}" - ) from e - class_methods_spec = [] - for mode, method_names in operations.items(): - for method_name in method_names: - if mode == _BIDI_STREAM_API_MODE: - _LOGGER.warning( - "Bidi stream API mode is not supported yet in Vertex SDK, " - "please use the GenAI SDK instead. Skipping " - f"method {method_name}." - ) - continue - if not hasattr(agent_engine, method_name): - raise ValueError( - f"Method `{method_name}` defined in `register_operations`" - " not found on AgentEngine." - ) - - method = getattr(agent_engine, method_name) - try: - schema_dict = _agent_engines_utils._generate_schema( - method, schema_name=method_name - ) - except Exception as e: - logger.warning(f"failed to generate schema for {method_name}: {e}") - continue - - class_method = _agent_engines_utils._to_proto(schema_dict) - class_method[_MODE_KEY_IN_SCHEMA] = mode - # A2A agent card is a special case, when running in A2A mode, - if hasattr(agent_engine, "agent_card"): - from google.protobuf import json_format - - class_method[_A2A_AGENT_CARD] = json_format.MessageToJson( - getattr(agent_engine, "agent_card") - ) - class_methods_spec.append(class_method) - - return class_methods_spec - - -def _class_methods_to_class_methods_spec( - class_methods: List[dict[str, Any]], -) -> List[proto.Message]: - """Converts a list of class methods to a list of ReasoningEngineSpec.ClassMethod messages.""" - return [ - _agent_engines_utils._to_proto(class_method) for class_method in class_methods - ] diff --git a/agentplatform/frameworks/__init__.py b/agentplatform/frameworks/__init__.py new file mode 100644 index 0000000000..346cdc96c9 --- /dev/null +++ b/agentplatform/frameworks/__init__.py @@ -0,0 +1,39 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Classes for working with agent platform.""" + +from agentplatform.frameworks import a2a +from agentplatform.frameworks import adk +from agentplatform.frameworks import ag2 +from agentplatform.frameworks import langchain +from agentplatform.frameworks import langgraph +from agentplatform.frameworks import llama_index + + +A2aAgent = a2a.A2aAgent +AdkApp = adk.AdkApp +AG2Agent = ag2.AG2Agent +LangchainAgent = langchain.LangchainAgent +LanggraphAgent = langgraph.LanggraphAgent +LlamaIndexQueryPipelineAgent = llama_index.LlamaIndexQueryPipelineAgent + +__all__ = ( + "A2aAgent", + "AdkApp", + "AG2Agent", + "LangchainAgent", + "LanggraphAgent", + "LlamaIndexQueryPipelineAgent", +) diff --git a/agentplatform/agent_engines/templates/a2a.py b/agentplatform/frameworks/a2a.py similarity index 97% rename from agentplatform/agent_engines/templates/a2a.py rename to agentplatform/frameworks/a2a.py index 705d57c428..00b5a67388 100644 --- a/agentplatform/agent_engines/templates/a2a.py +++ b/agentplatform/frameworks/a2a.py @@ -259,7 +259,6 @@ def __init__( ): """Initializes the A2A agent.""" # pylint: disable=g-import-not-at-top - from google.cloud.aiplatform import initializer from a2a.utils.constants import TransportProtocol, PROTOCOL_VERSION_CURRENT if ( @@ -276,8 +275,6 @@ def __init__( ) self._tmpl_attrs: dict[str, Any] = { - "project": initializer.global_config.project, - "location": initializer.global_config.location, "agent_card": agent_card, "agent_executor": None, "agent_executor_kwargs": agent_executor_kwargs or {}, @@ -318,14 +315,21 @@ def set_up(self): from a2a.server.tasks import InMemoryTaskStore os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1" - project = self._tmpl_attrs.get("project") - os.environ["GOOGLE_CLOUD_PROJECT"] = project - location = self._tmpl_attrs.get("location") - os.environ["GOOGLE_CLOUD_LOCATION"] = location - agent_engine_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "test-agent-engine") + + project = os.environ.get("GOOGLE_CLOUD_PROJECT") + location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") or os.getenv( + "GOOGLE_CLOUD_LOCATION" + ) + if location: + if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ: + os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location + if "GOOGLE_CLOUD_LOCATION" not in os.environ: + os.environ["GOOGLE_CLOUD_LOCATION"] = location + + runtime_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "test-agent-engine") version = "v1beta1" - new_url = f"https://{location}-aiplatform.googleapis.com/{version}/projects/{project}/locations/{location}/reasoningEngines/{agent_engine_id}/a2a" + new_url = f"https://{location}-aiplatform.googleapis.com/{version}/projects/{project}/locations/{location}/reasoningEngines/{runtime_id}/a2a" if not self.agent_card.supported_interfaces: from a2a.types import AgentInterface from a2a.utils.constants import TransportProtocol, PROTOCOL_VERSION_CURRENT diff --git a/agentplatform/agent_engines/templates/adk.py b/agentplatform/frameworks/adk.py similarity index 96% rename from agentplatform/agent_engines/templates/adk.py rename to agentplatform/frameworks/adk.py index ff3aecbb90..3cf1e10ec7 100644 --- a/agentplatform/agent_engines/templates/adk.py +++ b/agentplatform/frameworks/adk.py @@ -236,11 +236,6 @@ def __init__(self, **kwargs): ) # The session ID. - self.labels: Optional[Dict[str, str]] = kwargs.get("labels") - # Per-request user labels (e.g. for billing/attribution) to attach to - # the RunConfig for this invocation, so they are propagated onto the - # downstream Vertex API requests. - class _StreamingRunResponse: """Response object for `streaming_agent_run_with_events` method. @@ -257,13 +252,13 @@ def __init__(self, **kwargs): # The session ID. def dump(self) -> Dict[str, Any]: - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils result = {} if self.events: result["events"] = [] for event in self.events: - event_dict = _agent_engines_utils.dump_event_for_json(event) + event_dict = _runtimes_utils.dump_event_for_json(event) event_dict["invocation_id"] = event_dict.get("invocation_id", "") result["events"].append(event_dict) if self.artifacts: @@ -367,9 +362,9 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]: location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", "") or os.getenv( "GOOGLE_CLOUD_LOCATION", "" ) - agent_engine_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID") - if all(v is not None for v in (location, agent_engine_id)): - return f"//aiplatform.googleapis.com/projects/{project_id}/locations/{location}/reasoningEngines/{agent_engine_id}" + runtime_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID") + if all(v is not None for v in (location, runtime_id)): + return f"//aiplatform.googleapis.com/projects/{project_id}/locations/{location}/reasoningEngines/{runtime_id}" return None try: @@ -482,9 +477,9 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]: # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils - if _agent_engines_utils.is_noop_or_proxy_tracer_provider(tracer_provider): + if _runtimes_utils.is_noop_or_proxy_tracer_provider(tracer_provider): tracer_provider = opentelemetry.sdk.trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -741,7 +736,6 @@ def __init__( This parameter is ignored if `enable_tracing` is False. """ import os - from google.cloud.aiplatform import initializer adk_version = get_adk_version() if not is_version_sufficient("1.5.0"): @@ -768,8 +762,6 @@ def __init__( ) self._tmpl_attrs: Dict[str, Any] = { - "project": initializer.global_config.project, - "location": initializer.global_config.location, "agent": agent, "app": app, "app_name": app_name, @@ -780,9 +772,7 @@ def __init__( "memory_service_builder": memory_service_builder, "credential_service_builder": credential_service_builder, "instrumentor_builder": instrumentor_builder, - "express_mode_api_key": ( - initializer.global_config.api_key or os.environ.get("GOOGLE_API_KEY") - ), + "express_mode_api_key": os.environ.get("GOOGLE_API_KEY"), } def _serialize(self, obj: Any) -> Any: @@ -952,19 +942,16 @@ def set_up(self): ) os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1" - project = self._tmpl_attrs.get("project") - if project: - os.environ["GOOGLE_CLOUD_PROJECT"] = project - location = self._tmpl_attrs.get("location") + project = os.environ.get("GOOGLE_CLOUD_PROJECT") + location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") or os.getenv( + "GOOGLE_CLOUD_LOCATION" + ) if location: if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ: os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location if "GOOGLE_CLOUD_LOCATION" not in os.environ: os.environ["GOOGLE_CLOUD_LOCATION"] = location - agent_engine_location = os.environ.get( - "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", # the runtime env var (if set) - location, # the location set in the AdkApp template - ) + runtime_location = location express_mode_api_key = self._tmpl_attrs.get("express_mode_api_key") if express_mode_api_key and not project: os.environ["GOOGLE_API_KEY"] = express_mode_api_key @@ -1005,7 +992,7 @@ def set_up(self): "You can then use the " "'GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY' " "environment variable:\n" - "agent_engines.create(\n" + "runtimes.create(\n" " env_vars={\n" ' "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY": true|false\n' " }\n" @@ -1015,12 +1002,13 @@ def set_up(self): ), ) + project_id = self._get_project_id(project) if custom_instrumentor and self._tracing_enabled(): - self._tmpl_attrs["instrumentor"] = custom_instrumentor(self.project_id()) + self._tmpl_attrs["instrumentor"] = custom_instrumentor(project_id) if not custom_instrumentor: self._tmpl_attrs["instrumentor"] = _default_instrumentor_builder( - self.project_id(), + project_id, enable_tracing=self._tracing_enabled(), enable_logging=enable_logging, ) @@ -1052,7 +1040,7 @@ def set_up(self): # environment variable when initializing the session service. self._tmpl_attrs["session_service"] = VertexAiSessionService( project=project, - location=agent_engine_location, + location=runtime_location, agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"), ) except (ImportError, AttributeError): @@ -1064,7 +1052,7 @@ def set_up(self): # environment variable when initializing the session service. self._tmpl_attrs["session_service"] = VertexAiSessionService( project=project, - location=agent_engine_location, + location=runtime_location, agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"), ) @@ -1086,7 +1074,7 @@ def set_up(self): # environment variable when initializing the memory service. self._tmpl_attrs["memory_service"] = VertexAiMemoryBankService( project=project, - location=agent_engine_location, + location=runtime_location, agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"), ) except (ImportError, AttributeError): @@ -1098,7 +1086,7 @@ def set_up(self): # environment variable when initializing the memory service. self._tmpl_attrs["memory_service"] = VertexAiMemoryBankService( project=project, - location=agent_engine_location, + location=runtime_location, agent_engine_id=os.environ.get("GOOGLE_CLOUD_AGENT_ENGINE_ID"), ) else: @@ -1190,7 +1178,7 @@ async def async_stream_query( a Content object. ValueError: If both session_id and session_events are specified. """ - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils from google.genai import types if isinstance(message, Dict): @@ -1250,7 +1238,7 @@ async def async_stream_query( try: async for event in events_async: # Yield the event data as a dictionary - yield _agent_engines_utils.dump_event_for_json(event) + yield _runtimes_utils.dump_event_for_json(event) finally: # Avoid telemetry data loss having to do with CPU throttling on instance turndown _ = await _force_flush_otel( @@ -1300,7 +1288,7 @@ def stream_query( DeprecationWarning, stacklevel=2, ) - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils from google.genai import types if isinstance(message, Dict): @@ -1327,7 +1315,7 @@ def stream_query( run_config=run_config, **kwargs, ): - yield _agent_engines_utils.dump_event_for_json(event) + yield _runtimes_utils.dump_event_for_json(event) else: for event in self._tmpl_attrs.get("runner").run( user_id=user_id, @@ -1335,7 +1323,7 @@ def stream_query( new_message=content, **kwargs, ): - yield _agent_engines_utils.dump_event_for_json(event) + yield _runtimes_utils.dump_event_for_json(event) async def streaming_agent_run_with_events(self, request_json: str): """Streams responses asynchronously from the ADK application. @@ -1420,23 +1408,12 @@ async def streaming_agent_run_with_events(self, request_json: str): # Run the agent message_for_agent = types.Content(**request.message) - # Propagate per-request user labels (e.g. billing/attribution) onto a - # RunConfig so the ADK flow forwards them to downstream Vertex requests. - # Only attach them if the installed google-adk RunConfig supports the - # `labels` field; older versions forbid unknown fields. - run_config = None - if request.labels: - from google.adk.agents.run_config import RunConfig - - if "labels" in RunConfig.model_fields: - run_config = RunConfig(labels=request.labels) try: async for event in runner.run_async( user_id=request.user_id, session_id=session.id, new_message=message_for_agent, state_delta=state_delta, - run_config=run_config, ): converted_event = await self._convert_response_events( user_id=request.user_id, @@ -2155,8 +2132,8 @@ def _tracing_enabled(self) -> bool: and is_version_sufficient("1.17.0") ) - def project_id(self) -> Optional[str]: - if project := self._tmpl_attrs.get("project"): + def _get_project_id(self, project: str) -> Optional[str]: + if project: try: from google.cloud.aiplatform.utils import ( resource_manager_utils, diff --git a/agentplatform/agent_engines/templates/ag2.py b/agentplatform/frameworks/ag2.py similarity index 89% rename from agentplatform/agent_engines/templates/ag2.py rename to agentplatform/frameworks/ag2.py index ee12e49d8f..2d456b7e75 100644 --- a/agentplatform/agent_engines/templates/ag2.py +++ b/agentplatform/frameworks/ag2.py @@ -23,6 +23,8 @@ Sequence, Union, ) +import os +import copy if TYPE_CHECKING: try: @@ -89,15 +91,13 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): - from agentplatform._genai import _agent_engines_utils - - cloud_trace_exporter = _agent_engines_utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _agent_engines_utils._import_cloud_trace_v2_or_warn() - openinference_autogen = _agent_engines_utils._import_openinference_autogen_or_warn() - opentelemetry = _agent_engines_utils._import_opentelemetry_or_warn() - opentelemetry_sdk_trace = ( - _agent_engines_utils._import_opentelemetry_sdk_trace_or_warn() - ) + from agentplatform._genai import _runtimes_utils + + cloud_trace_exporter = _runtimes_utils._import_cloud_trace_exporter_or_warn() + cloud_trace_v2 = _runtimes_utils._import_cloud_trace_v2_or_warn() + openinference_autogen = _runtimes_utils._import_openinference_autogen_or_warn() + opentelemetry = _runtimes_utils._import_opentelemetry_or_warn() + opentelemetry_sdk_trace = _runtimes_utils._import_opentelemetry_sdk_trace_or_warn() if all( ( cloud_trace_exporter, @@ -144,7 +144,7 @@ def _default_instrumentor_builder(project_id: str): # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. - if _agent_engines_utils.is_noop_or_proxy_tracer_provider(tracer_provider): + if _runtimes_utils.is_noop_or_proxy_tracer_provider(tracer_provider): tracer_provider = opentelemetry_sdk_trace.TracerProvider() opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -337,11 +337,7 @@ def __init__( If not provided, a default instrumentor builder will be used. This parameter is ignored if `enable_tracing` is False. """ - from google.cloud.aiplatform import initializer - self._tmpl_attrs: dict[str, Any] = { - "project": initializer.global_config.project, - "location": initializer.global_config.location, "model_name": model, "api_type": api_type or "google", "system_instruction": system_instruction, @@ -353,45 +349,51 @@ def __init__( "instrumentor": None, "instrumentor_builder": instrumentor_builder, "enable_tracing": enable_tracing, + "provided_llm_config": copy.deepcopy(llm_config), + "provided_runnable_kwargs": copy.deepcopy(runnable_kwargs), } - self._tmpl_attrs["llm_config"] = llm_config or { + if tools: + _validate_tools(tools) + self._tmpl_attrs["tools"] = tools + + def set_up(self): + """Sets up the agent for execution of queries at runtime. + + It initializes the runnable, binds the runnable with tools. + Project and Location are sourced from environment variables. + """ + project = os.environ.get("GOOGLE_CLOUD_PROJECT") + location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") or os.getenv( + "GOOGLE_CLOUD_LOCATION" + ) + + llm_config = { "config_list": [ { - "project_id": self._tmpl_attrs.get("project"), - "location": self._tmpl_attrs.get("location"), + "project_id": project, + "location": location, "model": self._tmpl_attrs.get("model_name"), "api_type": self._tmpl_attrs.get("api_type"), } ] } - self._tmpl_attrs["runnable_kwargs"] = _prepare_runnable_kwargs( - runnable_kwargs=runnable_kwargs, - llm_config=self._tmpl_attrs.get("llm_config"), + if self._tmpl_attrs.get("provided_llm_config"): + llm_config = self._tmpl_attrs.get("provided_llm_config") + + runnable_kwargs = _prepare_runnable_kwargs( + runnable_kwargs=self._tmpl_attrs.get("provided_runnable_kwargs"), + llm_config=llm_config, system_instruction=self._tmpl_attrs.get("system_instruction"), runnable_name=self._tmpl_attrs.get("runnable_name"), ) - if tools: - # We validate tools at initialization for actionable feedback before - # they are deployed. - _validate_tools(tools) - self._tmpl_attrs["tools"] = tools - def set_up(self): - """Sets up the agent for execution of queries at runtime. - - It initializes the runnable, binds the runnable with tools. - - This method should not be called for an object that being passed to - the ReasoningEngine service for deployment, as it initializes clients - that can not be serialized. - """ if self._tmpl_attrs.get("enable_tracing"): instrumentor_builder = ( self._tmpl_attrs.get("instrumentor_builder") or _default_instrumentor_builder ) self._tmpl_attrs["instrumentor"] = instrumentor_builder( - project_id=self._tmpl_attrs.get("project") + project_id=project, ) # Set up tools. @@ -399,10 +401,10 @@ def set_up(self): ag2_tool_objects = self._tmpl_attrs.get("ag2_tool_objects") if tools and not ag2_tool_objects: from agentplatform._genai import ( - _agent_engines_utils, + _runtimes_utils, ) - autogen_tools = _agent_engines_utils._import_autogen_tools_or_warn() + autogen_tools = _runtimes_utils._import_autogen_tools_or_warn() if autogen_tools: for tool in tools: ag2_tool_objects.append(autogen_tools.Tool(func_or_tool=tool)) @@ -411,22 +413,21 @@ def set_up(self): runnable_builder = ( self._tmpl_attrs.get("runnable_builder") or _default_runnable_builder ) - self._tmpl_attrs["runnable"] = runnable_builder( - **self._tmpl_attrs.get("runnable_kwargs") - ) + self._tmpl_attrs["runnable"] = runnable_builder(**runnable_kwargs) def clone(self) -> "AG2Agent": """Returns a clone of the AG2Agent.""" - import copy return AG2Agent( model=self._tmpl_attrs.get("model_name"), api_type=self._tmpl_attrs.get("api_type"), - llm_config=copy.deepcopy(self._tmpl_attrs.get("llm_config")), + llm_config=copy.deepcopy(self._tmpl_attrs.get("provided_llm_config")), system_instruction=self._tmpl_attrs.get("system_instruction"), runnable_name=self._tmpl_attrs.get("runnable_name"), tools=copy.deepcopy(self._tmpl_attrs.get("tools")), - runnable_kwargs=copy.deepcopy(self._tmpl_attrs.get("runnable_kwargs")), + runnable_kwargs=copy.deepcopy( + self._tmpl_attrs.get("provided_runnable_kwargs") + ), runnable_builder=self._tmpl_attrs.get("runnable_builder"), enable_tracing=self._tmpl_attrs.get("enable_tracing"), instrumentor_builder=self._tmpl_attrs.get("instrumentor_builder"), @@ -488,6 +489,6 @@ def query( **kwargs, ) - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils - return _agent_engines_utils.to_json_serializable_autogen_object(response) + return _runtimes_utils.to_json_serializable_autogen_object(response) diff --git a/agentplatform/agent_engines/templates/langchain.py b/agentplatform/frameworks/langchain.py similarity index 95% rename from agentplatform/agent_engines/templates/langchain.py rename to agentplatform/frameworks/langchain.py index 17751524ba..188b6f0964 100644 --- a/agentplatform/agent_engines/templates/langchain.py +++ b/agentplatform/frameworks/langchain.py @@ -114,15 +114,11 @@ def _default_model_builder( ) return model except ImportError: - import agentplatform - from google.cloud.aiplatform import initializer from langchain_google_vertexai import ChatVertexAI - current_project = initializer.global_config.project - current_location = initializer.global_config.location - agentplatform.init(project=project, location=location) - model = ChatVertexAI(model_name=model_name, **model_kwargs) - agentplatform.init(project=current_project, location=current_location) + model = ChatVertexAI( + model_name=model_name, project=project, location=location, **model_kwargs + ) return model @@ -192,17 +188,13 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils - cloud_trace_exporter = _agent_engines_utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _agent_engines_utils._import_cloud_trace_v2_or_warn() - openinference_langchain = ( - _agent_engines_utils._import_openinference_langchain_or_warn() - ) - opentelemetry = _agent_engines_utils._import_opentelemetry_or_warn() - opentelemetry_sdk_trace = ( - _agent_engines_utils._import_opentelemetry_sdk_trace_or_warn() - ) + cloud_trace_exporter = _runtimes_utils._import_cloud_trace_exporter_or_warn() + cloud_trace_v2 = _runtimes_utils._import_cloud_trace_v2_or_warn() + openinference_langchain = _runtimes_utils._import_openinference_langchain_or_warn() + opentelemetry = _runtimes_utils._import_opentelemetry_or_warn() + opentelemetry_sdk_trace = _runtimes_utils._import_opentelemetry_sdk_trace_or_warn() if all( ( cloud_trace_exporter, @@ -249,7 +241,7 @@ def _default_instrumentor_builder(project_id: str): # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. - if _agent_engines_utils.is_noop_or_proxy_tracer_provider(tracer_provider): + if _runtimes_utils.is_noop_or_proxy_tracer_provider(tracer_provider): tracer_provider = opentelemetry_sdk_trace.TracerProvider() opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -541,11 +533,7 @@ def __init__( TypeError: If there is an invalid tool (e.g. function with an input that did not specify its type). """ - from google.cloud.aiplatform import initializer - self._tmpl_attrs: dict[str, Any] = { - "project": initializer.global_config.project, - "location": initializer.global_config.location, "tools": [], "model_name": model, "system_instruction": system_instruction, @@ -586,20 +574,26 @@ def set_up(self): service for deployment, as it might initialize clients that can not be serialized. """ + import os + + project = os.environ.get("GOOGLE_CLOUD_PROJECT") + location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") or os.getenv( + "GOOGLE_CLOUD_LOCATION" + ) if self._tmpl_attrs.get("enable_tracing"): instrumentor_builder = ( self._tmpl_attrs.get("instrumentor_builder") or _default_instrumentor_builder ) self._tmpl_attrs["instrumentor"] = instrumentor_builder( - project_id=self._tmpl_attrs.get("project") + project_id=project, ) model_builder = self._tmpl_attrs.get("model_builder") or _default_model_builder self._tmpl_attrs["model"] = model_builder( model_name=self._tmpl_attrs.get("model_name"), model_kwargs=self._tmpl_attrs.get("model_kwargs"), - project=self._tmpl_attrs.get("project"), - location=self._tmpl_attrs.get("location"), + project=project, + location=location, ) runnable_builder = ( self._tmpl_attrs.get("runnable_builder") or _default_runnable_builder diff --git a/agentplatform/agent_engines/templates/langgraph.py b/agentplatform/frameworks/langgraph.py similarity index 95% rename from agentplatform/agent_engines/templates/langgraph.py rename to agentplatform/frameworks/langgraph.py index 0ebdb690cc..c460e0a518 100644 --- a/agentplatform/agent_engines/templates/langgraph.py +++ b/agentplatform/frameworks/langgraph.py @@ -105,15 +105,11 @@ def _default_model_builder( ) return model except ImportError: - import agentplatform - from google.cloud.aiplatform import initializer from langchain_google_vertexai import ChatVertexAI - current_project = initializer.global_config.project - current_location = initializer.global_config.location - agentplatform.init(project=project, location=location) - model = ChatVertexAI(model_name=model_name, **model_kwargs) - agentplatform.init(project=current_project, location=current_location) + model = ChatVertexAI( + model_name=model_name, project=project, location=location, **model_kwargs + ) return model @@ -168,17 +164,13 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils - cloud_trace_exporter = _agent_engines_utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _agent_engines_utils._import_cloud_trace_v2_or_warn() - openinference_langchain = ( - _agent_engines_utils._import_openinference_langchain_or_warn() - ) - opentelemetry = _agent_engines_utils._import_opentelemetry_or_warn() - opentelemetry_sdk_trace = ( - _agent_engines_utils._import_opentelemetry_sdk_trace_or_warn() - ) + cloud_trace_exporter = _runtimes_utils._import_cloud_trace_exporter_or_warn() + cloud_trace_v2 = _runtimes_utils._import_cloud_trace_v2_or_warn() + openinference_langchain = _runtimes_utils._import_openinference_langchain_or_warn() + opentelemetry = _runtimes_utils._import_opentelemetry_or_warn() + opentelemetry_sdk_trace = _runtimes_utils._import_opentelemetry_sdk_trace_or_warn() if all( ( cloud_trace_exporter, @@ -224,7 +216,7 @@ def _default_instrumentor_builder(project_id: str): # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. - if _agent_engines_utils.is_noop_or_proxy_tracer_provider(tracer_provider): + if _runtimes_utils.is_noop_or_proxy_tracer_provider(tracer_provider): tracer_provider = opentelemetry_sdk_trace.TracerProvider() opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -465,11 +457,7 @@ def checkpointer_builder(**kwargs): TypeError: If there is an invalid tool (e.g. function with an input that did not specify its type). """ - from google.cloud.aiplatform import initializer - self._tmpl_attrs: dict[str, Any] = { - "project": initializer.global_config.project, - "location": initializer.global_config.location, "tools": [], "model_name": model, "model_kwargs": model_kwargs, @@ -502,20 +490,26 @@ def set_up(self): the ReasoningEngine service for deployment, as it initializes clients that can not be serialized. """ + import os + + project = os.environ.get("GOOGLE_CLOUD_PROJECT") + location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") or os.getenv( + "GOOGLE_CLOUD_LOCATION" + ) if self._tmpl_attrs.get("enable_tracing"): instrumentor_builder = ( self._tmpl_attrs.get("instrumentor_builder") or _default_instrumentor_builder ) self._tmpl_attrs["instrumentor"] = instrumentor_builder( - project_id=self._tmpl_attrs.get("project") + project_id=project, ) model_builder = self._tmpl_attrs.get("model_builder") or _default_model_builder self._tmpl_attrs["model"] = model_builder( model_name=self._tmpl_attrs.get("model_name"), model_kwargs=self._tmpl_attrs.get("model_kwargs"), - project=self._tmpl_attrs.get("project"), - location=self._tmpl_attrs.get("location"), + project=project, + location=location, ) checkpointer_builder = self._tmpl_attrs.get("checkpointer_builder") if checkpointer_builder: diff --git a/agentplatform/agent_engines/templates/llama_index.py b/agentplatform/frameworks/llama_index.py similarity index 93% rename from agentplatform/agent_engines/templates/llama_index.py rename to agentplatform/frameworks/llama_index.py index d9147fccef..3f68e93e39 100644 --- a/agentplatform/agent_engines/templates/llama_index.py +++ b/agentplatform/frameworks/llama_index.py @@ -58,8 +58,6 @@ def _default_model_builder( model_kwargs: Optional[Mapping[str, Any]] = None, ) -> "FunctionCallingLLM": """Creates a default model builder for LlamaIndex.""" - import agentplatform - from google.cloud.aiplatform import initializer from llama_index.llms import google_genai model_kwargs = model_kwargs or {} @@ -68,9 +66,6 @@ def _default_model_builder( vertexai_config={"project": project, "location": location}, **model_kwargs, ) - current_project = initializer.global_config.project - current_location = initializer.global_config.location - agentplatform.init(project=current_project, location=current_location) return model @@ -346,10 +341,6 @@ def __init__( enable_tracing (bool): Optional. Whether to enable tracing. Defaults to False. """ - from google.cloud.aiplatform import initializer - - self._project = initializer.global_config.project - self._location = initializer.global_config.location self._model_name = model self._system_instruction = system_instruction self._prompt = prompt @@ -383,21 +374,28 @@ def set_up(self): the ReasoningEngine service for deployment, as it initializes clients that can not be serialized. """ + import os + + project = os.environ.get("GOOGLE_CLOUD_PROJECT") + location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") or os.getenv( + "GOOGLE_CLOUD_LOCATION" + ) + if self._enable_tracing: - from agentplatform._genai.agent_engines import ( - _agent_engines_utils, + from agentplatform._genai import ( + _runtimes_utils, ) cloud_trace_exporter = ( - _agent_engines_utils._import_cloud_trace_exporter_or_warn() + _runtimes_utils._import_cloud_trace_exporter_or_warn() ) - cloud_trace_v2 = _agent_engines_utils._import_cloud_trace_v2_or_warn() + cloud_trace_v2 = _runtimes_utils._import_cloud_trace_v2_or_warn() openinference_llama_index = ( - _agent_engines_utils._import_openinference_llama_index_or_warn() + _runtimes_utils._import_openinference_llama_index_or_warn() ) - opentelemetry = _agent_engines_utils._import_opentelemetry_or_warn() + opentelemetry = _runtimes_utils._import_opentelemetry_or_warn() opentelemetry_sdk_trace = ( - _agent_engines_utils._import_opentelemetry_sdk_trace_or_warn() + _runtimes_utils._import_opentelemetry_sdk_trace_or_warn() ) if all( ( @@ -412,9 +410,9 @@ def set_up(self): credentials, _ = google.auth.default() span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=self._project, + project_id=project, client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(self._project), + credentials=credentials.with_quota_project(project), ), ) span_processor: SpanProcessor = ( @@ -447,9 +445,7 @@ def set_up(self): # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. - if _agent_engines_utils.is_noop_or_proxy_tracer_provider( - tracer_provider - ): + if _runtimes_utils.is_noop_or_proxy_tracer_provider(tracer_provider): tracer_provider = opentelemetry_sdk_trace.TracerProvider() opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -482,8 +478,8 @@ def set_up(self): self._model = model_builder( model_name=self._model_name, model_kwargs=self._model_kwargs, - project=self._project, - location=self._location, + project=project, + location=location, ) if self._retriever_builder: @@ -546,8 +542,8 @@ def query( Returns: The output of querying the Agent with the given input and config. """ - from agentplatform._genai.agent_engines import ( - _agent_engines_utils, + from agentplatform._genai import ( + _runtimes_utils, ) if isinstance(input, str): @@ -557,9 +553,9 @@ def query( self.set_up() if kwargs.get("batch"): - nest_asyncio = _agent_engines_utils._import_nest_asyncio_or_warn() + nest_asyncio = _runtimes_utils._import_nest_asyncio_or_warn() nest_asyncio.apply() - return _agent_engines_utils.to_json_serializable_llama_index_object( + return _runtimes_utils.to_json_serializable_llama_index_object( self._runnable.run(**input, **kwargs) ) diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py b/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py index 4fa0d87aff..d7c0bb5028 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py @@ -13,63 +13,25 @@ # limitations under the License. # +import importlib import os -import sys -import tempfile from unittest import mock -import pytest -import cloudpickle -import pydantic from google import auth -from google.api_core import operation as ga_operation from google.auth import credentials as auth_credentials -from google.cloud import storage -from google.cloud import aiplatform -from google.cloud.aiplatform import base - -from google.cloud.aiplatform_v1 import types -from google.cloud.aiplatform_v1.services import reasoning_engine_service -from vertexai import agent_engines -from vertexai.agent_engines import _agent_engines -from vertexai.agent_engines import _utils -from google.protobuf import struct_pb2 - - -class CapitalizeEngine: - """A sample Agent Engine.""" - - def query(self, unused_arbitrary_string_name: str) -> str: - """Runs the engine.""" - return unused_arbitrary_string_name.upper() - - -class CapitalizeEngineWithCard(CapitalizeEngine): - - def __init__(self, card): - self.agent_card = card - - def __getstate__(self): - state = self.__dict__.copy() - if hasattr(self.agent_card, "DESCRIPTOR"): - state["agent_card"] = None - return state +import agentplatform +from agentplatform._genai import _runtimes_utils +from agentplatform._genai import types as _genai_types +from google.genai import types as genai_types +import pytest - def __setstate__(self, state): - self.__dict__.update(state) +from a2a.types import AgentSkill -def _create_empty_fake_package(package_name: str) -> str: - temp_dir = tempfile.mkdtemp() - package_dir = os.path.join(temp_dir, package_name) - os.makedirs(package_dir) - init_path = os.path.join(package_dir, "__init__.py") - open(init_path, "w").close() - return temp_dir +from agentplatform.frameworks.a2a import A2aAgent +from agentplatform.frameworks.a2a import create_agent_card -_TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials()) -_TEST_STAGING_BUCKET = "gs://test-bucket" _TEST_LOCATION = "us-central1" _TEST_PROJECT = "test-project" _TEST_RESOURCE_ID = "1028944691210842416" @@ -77,65 +39,25 @@ def _create_empty_fake_package(package_name: str) -> str: _TEST_AGENT_ENGINE_RESOURCE_NAME = ( f"{_TEST_PARENT}/reasoningEngines/{_TEST_RESOURCE_ID}" ) -_TEST_AGENT_ENGINE_DISPLAY_NAME = "Agent Engine Display Name" -_TEST_GCS_DIR_NAME = _agent_engines._DEFAULT_GCS_DIR_NAME -_TEST_BLOB_FILENAME = _agent_engines._BLOB_FILENAME -_TEST_REQUIREMENTS_FILE = _agent_engines._REQUIREMENTS_FILE -_TEST_EXTRA_PACKAGES_FILE = _agent_engines._EXTRA_PACKAGES_FILE -_TEST_STANDARD_API_MODE = _agent_engines._STANDARD_API_MODE -_TEST_DEFAULT_METHOD_NAME = _agent_engines._DEFAULT_METHOD_NAME -_TEST_MODE_KEY_IN_SCHEMA = _agent_engines._MODE_KEY_IN_SCHEMA - -_TEST_AGENT_ENGINE_EXTRA_PACKAGE = "fake.py" - -_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH = _create_empty_fake_package( - _TEST_AGENT_ENGINE_EXTRA_PACKAGE +_TEST_AGENT_ENGINE_DISPLAY_NAME = "test-a2a-agent" +_TEST_STAGING_BUCKET = "gs://test-bucket" +_TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials()) +_TEST_SKILL = AgentSkill( + id="hello_world", + name="Returns hello world", + description="just returns hello world", + tags=["hello world"], + examples=["hi", "hello world"], ) -_TEST_AGENT_ENGINE_REQUIREMENTS = [ - "google-cloud-aiplatform==1.29.0", - "langchain", -] - -_TEST_AGENT_ENGINE_GCS_URI = "{}/{}/{}".format( - _TEST_STAGING_BUCKET, - _TEST_GCS_DIR_NAME, - _TEST_BLOB_FILENAME, -) -_TEST_AGENT_ENGINE_DEPENDENCY_FILES_GCS_URI = "{}/{}/{}".format( - _TEST_STAGING_BUCKET, - _TEST_GCS_DIR_NAME, - _TEST_EXTRA_PACKAGES_FILE, -) -_TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI = "{}/{}/{}".format( - _TEST_STAGING_BUCKET, - _TEST_GCS_DIR_NAME, - _TEST_REQUIREMENTS_FILE, -) -_TEST_AGENT_ENGINE_QUERY_SCHEMA = _utils.to_proto( - _utils.generate_schema( - CapitalizeEngine().query, - schema_name=_TEST_DEFAULT_METHOD_NAME, +def _make_agent() -> A2aAgent: + card = create_agent_card( + agent_name="Test", + description="Test", + skills=[_TEST_SKILL], ) -) -_TEST_AGENT_ENGINE_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STANDARD_API_MODE - -_TEST_AGENT_ENGINE_PACKAGE_SPEC = types.ReasoningEngineSpec.PackageSpec( - python_version=f"{sys.version_info.major}.{sys.version_info.minor}", - pickle_object_gcs_uri=_TEST_AGENT_ENGINE_GCS_URI, - dependency_files_gcs_uri=_TEST_AGENT_ENGINE_DEPENDENCY_FILES_GCS_URI, - requirements_gcs_uri=_TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI, -) - -_TEST_AGENT_ENGINE_OBJ = types.ReasoningEngine( - name=_TEST_AGENT_ENGINE_RESOURCE_NAME, - spec=types.ReasoningEngineSpec( - package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC, - agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK, - ), -) -_TEST_AGENT_ENGINE_OBJ.spec.class_methods.append(_TEST_AGENT_ENGINE_QUERY_SCHEMA) + return A2aAgent(agent_card=card) @pytest.fixture(scope="module") @@ -148,128 +70,109 @@ def google_auth_mock(): yield google_auth_mock -@pytest.fixture(scope="module") -def cloud_storage_create_bucket_mock(): - with mock.patch.object(storage, "Client") as cloud_storage_mock: - bucket_mock = mock.Mock(spec=storage.Bucket) - bucket_mock.blob.return_value.open.return_value = "blob_file" - bucket_mock.blob.return_value.upload_from_filename.return_value = None - bucket_mock.blob.return_value.upload_from_string.return_value = None - - cloud_storage_mock.get_bucket = mock.Mock( - side_effect=ValueError("bucket not found") - ) - cloud_storage_mock.bucket.return_value = bucket_mock - cloud_storage_mock.create_bucket.return_value = bucket_mock - - yield cloud_storage_mock - - -@pytest.fixture(scope="module") -def cloudpickle_load_mock(): - with mock.patch.object(cloudpickle, "load") as cloudpickle_load_mock: - yield cloudpickle_load_mock - - -@pytest.fixture(scope="module") -def create_agent_engine_mock(): - with mock.patch.object( - reasoning_engine_service.ReasoningEngineServiceClient, - "create_reasoning_engine", - ) as create_agent_engine_mock: - create_agent_engine_lro_mock = mock.Mock(spec=ga_operation.Operation) - create_agent_engine_lro_mock.result.return_value = _TEST_AGENT_ENGINE_OBJ - create_agent_engine_mock.return_value = create_agent_engine_lro_mock - yield create_agent_engine_mock - - -@pytest.fixture(scope="function") -def get_gca_resource_mock(): - with mock.patch.object( - base.VertexAiResourceNoun, - "_get_gca_resource", - ) as get_gca_resource_mock: - get_gca_resource_mock.return_value = _TEST_AGENT_ENGINE_OBJ - yield get_gca_resource_mock - - @pytest.mark.usefixtures("google_auth_mock") -class TestAgentEngineA2A: +class TestA2aAgentCreate: def setup_method(self): - aiplatform.init( + importlib.reload(agentplatform) + os.environ["GOOGLE_CLOUD_PROJECT"] = _TEST_PROJECT + os.environ["GOOGLE_CLOUD_LOCATION"] = _TEST_LOCATION + self.client = agentplatform.Client( project=_TEST_PROJECT, location=_TEST_LOCATION, credentials=_TEST_CREDENTIALS, - staging_bucket=_TEST_STAGING_BUCKET, ) - def test_create_agent_engine_with_protobuf_agent_card( + def teardown_method(self): + for key in [ + "GOOGLE_CLOUD_PROJECT", + "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", + "GOOGLE_CLOUD_LOCATION", + "GOOGLE_GENAI_USE_VERTEXAI", + ]: + os.environ.pop(key, None) + + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") + @mock.patch.object( + _runtimes_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) + def test_create_agent_with_agent_card( self, - create_agent_engine_mock, - cloud_storage_create_bucket_mock, - cloudpickle_load_mock, - get_gca_resource_mock, + mock_get_reasoning_engine_id, + mock_await_operation, + mock_prepare, ): - a2a_pb2 = None - # fmt: off - try: - try: - from a2a.compat.v0_3 import a2a_v0_3_pb2 as a2a_pb2 - except ImportError: - from a2a.grpc import a2a_pb2 - has_a2a_pb2 = True - except (ImportError, TypeError): - has_a2a_pb2 = False - # fmt: on - - if not has_a2a_pb2: - pytest.skip("a2a_pb2 could not be imported.") - - card = a2a_pb2.AgentCard(name="test_agent_card") - agent = CapitalizeEngineWithCard(card) - - agent_engines.create( - agent, - display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, - requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, - extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], + mock_await_operation.return_value = _genai_types.RuntimeOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + ) ) + agent = _make_agent() + with mock.patch.object( + self.client.runtimes._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="") + self.client.runtimes.create( + agent=agent, + config=_genai_types.AgentRuntimeConfig( + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + staging_bucket=_TEST_STAGING_BUCKET, + ), + ) - expected_reasoning_engine = types.ReasoningEngine( - display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, - spec=types.ReasoningEngineSpec( - package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC, - agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK, - ), - ) - from google.protobuf import json_format + request_mock.assert_called_once() + _, _, request_dict, _ = request_mock.call_args[0] + class_methods = request_dict["spec"]["class_methods"] + assert any("a2a_agent_card" in method for method in class_methods) - expected_class_method = struct_pb2.Struct() - expected_class_method.CopyFrom(_TEST_AGENT_ENGINE_QUERY_SCHEMA) - expected_class_method["a2a_agent_card"] = json_format.MessageToJson(card) - expected_reasoning_engine.spec.class_methods.append(expected_class_method) - create_agent_engine_mock.assert_called_with( - parent=_TEST_PARENT, - reasoning_engine=expected_reasoning_engine, - ) +class TestA2aLocationResolution: + @pytest.fixture(autouse=True) + def patch_create_rest_routes(self, monkeypatch): + """Patch to avoid optional dependencies not actually used in these tests.""" + import sys + from unittest.mock import MagicMock - def test_create_agent_engine_with_invalid_agent_card( - self, - create_agent_engine_mock, - cloud_storage_create_bucket_mock, - cloudpickle_load_mock, - get_gca_resource_mock, - ): - agent = CapitalizeEngineWithCard(card="invalid_card_type_string") + mock_rest_routes = MagicMock() + mock_rest_routes.create_rest_routes = MagicMock() + monkeypatch.setitem( + sys.modules, "a2a.server.routes.rest_routes", mock_rest_routes + ) - with pytest.raises( - TypeError, - match="Unsupported AgentCard type", + def test_location_from_runtime_env_var(self): + agent = _make_agent() + with mock.patch.dict( + os.environ, + {"GOOGLE_CLOUD_AGENT_ENGINE_LOCATION": "us-east1"}, + clear=True, ): - agent_engines.create( - agent, - display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, - requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, - extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], - ) + agent.set_up() + assert os.environ.get("GOOGLE_CLOUD_LOCATION") == "us-east1" + + if hasattr(agent.agent_card, "url"): + url = agent.agent_card.url + else: + url = agent.agent_card.supported_interfaces[0].url + assert "us-east1" in url + + def test_location_env_var_precedence(self): + agent = _make_agent() + with mock.patch.dict( + os.environ, + { + "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION": "us-east1", + "GOOGLE_CLOUD_LOCATION": "us-west1", + }, + clear=True, + ): + agent.set_up() + # Should not overwrite existing GOOGLE_CLOUD_LOCATION. + assert os.environ.get("GOOGLE_CLOUD_LOCATION") == "us-west1" + + if hasattr(agent.agent_card, "url"): + url = agent.agent_card.url + else: + url = agent.agent_card.supported_interfaces[0].url + assert "us-east1" in url diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py index d88cd56074..a1b474b6c3 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py @@ -34,9 +34,9 @@ from google.cloud.aiplatform import initializer from google.cloud.aiplatform_v1 import types as aip_types from google.cloud.aiplatform_v1.services import reasoning_engine_service -from agentplatform._genai import agent_engines -from agentplatform._genai import _agent_engines_utils -from agentplatform.agent_engines.templates import ( +from agentplatform._genai import runtimes +from agentplatform._genai import _runtimes_utils +from agentplatform.frameworks import ( adk as adk_template, ) from google.genai import types @@ -414,12 +414,11 @@ def setup_method(self): "GOOGLE_GENAI_USE_VERTEXAI", ]: os.environ.pop(key, None) - importlib.reload(initializer) importlib.reload(agentplatform) - agentplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + os.environ["GOOGLE_CLOUD_PROJECT"] = _TEST_PROJECT + os.environ["GOOGLE_CLOUD_LOCATION"] = _TEST_LOCATION def teardown_method(self): - initializer.global_pool.shutdown(wait=True) for key in [ "GOOGLE_CLOUD_PROJECT", "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", @@ -430,8 +429,6 @@ def teardown_method(self): def test_initialization(self): app = adk_template.AdkApp(agent=_TEST_AGENT) - assert app._tmpl_attrs.get("project") == _TEST_PROJECT - assert app._tmpl_attrs.get("location") == _TEST_LOCATION assert app._tmpl_attrs.get("runner") is None def test_set_up( @@ -547,6 +544,15 @@ async def test_runner_auto_creates_missing_session( default_instrumentor_builder_mock: mock.Mock, get_project_id_mock: mock.Mock, ): + """A missing session_id is auto-created, not rejected (b/537477760). + + Regression test for b/537477760: the interactive playground streams a + query with a session_id that has not been created in this agent's + session service. Because the template constructs the runner with + ``auto_create_session=True``, the real ADK session-resolution logic + (``Runner._get_or_create_session``) creates the session on the fly + instead of raising ``SessionNotFoundError``. + """ from google.adk.runners import Runner from google.adk.sessions.in_memory_session_service import ( InMemorySessionService, @@ -554,8 +560,13 @@ async def test_runner_auto_creates_missing_session( app = adk_template.AdkApp(agent=_TEST_AGENT) app.set_up() + # The template must have opted the runner into auto-creation. assert app._tmpl_attrs.get("runner").auto_create_session is True + # Drive the real ADK session-resolution path (no LLM/backend calls) + # with an in-memory session service. Crucially, the runner inherits + # the auto_create_session value that the template configured, so this + # test fails if the template regresses to auto_create_session=False. app_name = app._tmpl_attrs.get("app_name") session_service = InMemorySessionService() runner = Runner( @@ -576,6 +587,7 @@ async def test_runner_auto_creates_missing_session( is None ) + # This is the call that raised SessionNotFoundError before the fix. session = await runner._get_or_create_session( user_id=_TEST_USER_ID, session_id=missing_session_id, @@ -583,6 +595,7 @@ async def test_runner_auto_creates_missing_session( assert session is not None assert session.id == missing_session_id + # The session is now persisted in the session service. assert ( await session_service.get_session( app_name=app_name, @@ -620,9 +633,7 @@ async def test_async_stream_query_with_session_events( default_instrumentor_builder_mock: mock.Mock, get_project_id_mock: mock.Mock, ): - app = adk_template.AdkApp( - agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL) - ) + app = adk_template.AdkApp(agent=Agent(name=_TEST_AGENT_NAME, model=_TEST_MODEL)) assert app._tmpl_attrs.get("runner") is None app.set_up() app._tmpl_attrs["runner"] = _MockRunner() @@ -716,67 +727,6 @@ async def test_streaming_agent_run_with_events( events.append(event) assert len(events) == 1 - @pytest.mark.asyncio - async def test_streaming_agent_run_with_events_propagates_labels( - self, - default_instrumentor_builder_mock: mock.Mock, - get_project_id_mock: mock.Mock, - ): - from google.adk.agents.run_config import RunConfig - - if "labels" not in RunConfig.model_fields: - pytest.skip("Installed google-adk RunConfig does not support labels.") - - app = adk_template.AdkApp(agent=_TEST_AGENT) - app.set_up() - - # Pre-create a session in the real in-memory session service. - await app.async_create_session( - user_id=_TEST_USER_ID, session_id="test_session_id" - ) - - runner_mock = mock.Mock() - - async def mock_run_async(*args, **kwargs): - from google.adk.events import event - - yield event.Event( - **{ - "author": "currency_exchange_agent", - "content": {"parts": [{"text": "Sweden"}], "role": "model"}, - "id": "9aaItGK9", - "invocation_id": "e-6543c213-6417-484b-9551-b67915d1d5f7", - } - ) - - spy = mock.MagicMock(side_effect=mock_run_async) - runner_mock.run_async = spy - app._tmpl_attrs["runner"] = runner_mock - - labels = {"goog-originating-logical-product-id": "prod1"} - request_json = json.dumps( - { - "user_id": _TEST_USER_ID, - "session_id": "test_session_id", - "message": { - "parts": [{"text": "What is the exchange rate from USD to SEK?"}], - "role": "user", - }, - "labels": labels, - } - ) - - async for _ in app.streaming_agent_run_with_events( - request_json=request_json, - ): - pass - - # The labels are surfaced to run_async via a RunConfig. - spy.assert_called_once() - run_config = spy.call_args.kwargs["run_config"] - assert run_config is not None - assert run_config.labels == labels - @pytest.mark.asyncio @mock.patch.dict( os.environ, @@ -1336,8 +1286,8 @@ def test_tracing_setup( "uuid.uuid4", lambda: uuid.UUID("12345678123456781234567812345678") ) monkeypatch.setattr("os.getpid", lambda: 123123123) - with mock.patch.object(initializer.global_config, "_project", _TEST_PROJECT): - app = adk_template.AdkApp(agent=_TEST_AGENT, enable_tracing=True) + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", _TEST_PROJECT) + app = adk_template.AdkApp(agent=_TEST_AGENT, enable_tracing=True) app.set_up() otlp_span_exporter_mock.assert_called_once_with( @@ -1444,7 +1394,7 @@ def test_dump_event_for_json(): "invocation_id": "test_invocation_id", } ) - dumped_event = _agent_engines_utils.dump_event_for_json(test_event) + dumped_event = _runtimes_utils.dump_event_for_json(test_event) part = dumped_event["content"]["parts"][0] assert "text" in part @@ -1519,27 +1469,27 @@ async def test_async_stream_query_invalid_message_type(self): @pytest.fixture(scope="module") -def create_agent_engine_mock(): +def create_runtime_mock(): with mock.patch.object( reasoning_engine_service.ReasoningEngineServiceClient, "create_reasoning_engine", - ) as create_agent_engine_mock: - create_agent_engine_lro_mock = mock.Mock(ga_operation.Operation) - create_agent_engine_lro_mock.result.return_value = _TEST_AGENT_ENGINE_OBJ - create_agent_engine_mock.return_value = create_agent_engine_lro_mock - yield create_agent_engine_mock + ) as create_runtime_mock: + create_runtime_lro_mock = mock.Mock(ga_operation.Operation) + create_runtime_lro_mock.result.return_value = _TEST_AGENT_ENGINE_OBJ + create_runtime_mock.return_value = create_runtime_lro_mock + yield create_runtime_mock @pytest.fixture(scope="module") -def get_agent_engine_mock(): +def get_runtime_mock(): with mock.patch.object( reasoning_engine_service.ReasoningEngineServiceClient, "get_reasoning_engine", - ) as get_agent_engine_mock: + ) as get_runtime_mock: api_client_mock = mock.Mock() api_client_mock.get_reasoning_engine.return_value = _TEST_AGENT_ENGINE_OBJ - get_agent_engine_mock.return_value = api_client_mock - yield get_agent_engine_mock + get_runtime_mock.return_value = api_client_mock + yield get_runtime_mock @pytest.fixture(scope="module") @@ -1583,12 +1533,12 @@ def get_gca_resource_mock(): # Function scope is required for the pytest parameterized tests. @pytest.fixture(scope="function") -def update_agent_engine_mock(): +def update_runtime_mock(): with mock.patch.object( reasoning_engine_service.ReasoningEngineServiceClient, "update_reasoning_engine", - ) as update_agent_engine_mock: - yield update_agent_engine_mock + ) as update_runtime_mock: + yield update_runtime_mock @pytest.mark.usefixtures( @@ -1597,9 +1547,9 @@ def update_agent_engine_mock(): "cloudpickle_dump_mock", "cloudpickle_load_mock", "get_gca_resource_mock", - "get_agent_engine_mock", + "get_runtime_mock", ) -class TestAgentEngines: +class TestRuntimes: def setup_method(self): importlib.reload(initializer) @@ -1636,8 +1586,8 @@ def teardown_method(self): ), ], ) - @mock.patch.object(agent_engines.AgentEngines, "_create") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(runtimes.Runtimes, "_create") + @mock.patch.object(_runtimes_utils, "_await_operation") def test_create_default_telemetry_enablement( self, mock_await_operation, @@ -1652,13 +1602,13 @@ def test_create_default_telemetry_enablement( "projects/test-project/locations/us-central1/reasoningEngines/123456/operations/789" ) mock_create.return_value = mock_operation - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, ) ) client = agentplatform.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) - client.agent_engines.create( + client.runtimes.create( agent=adk_template.AdkApp(agent=_TEST_AGENT), config={"env_vars": env_vars, "staging_bucket": _TEST_STAGING_BUCKET}, ) @@ -1691,8 +1641,8 @@ def test_create_default_telemetry_enablement( ), ], ) - @mock.patch.object(agent_engines.AgentEngines, "_update") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(runtimes.Runtimes, "_update") + @mock.patch.object(_runtimes_utils, "_await_operation") def test_update_default_telemetry_enablement( self, mock_await_operation, @@ -1707,13 +1657,13 @@ def test_update_default_telemetry_enablement( "projects/test-project/locations/us-central1/reasoningEngines/123456/operations/789" ) mock_update.return_value = mock_operation - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, ) ) client = agentplatform.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) - client.agent_engines.update( + client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, agent=adk_template.AdkApp(agent=_TEST_AGENT), config={ diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_ag2.py b/tests/unit/agentplatform/frameworks/test_frameworks_ag2.py index b8e2337f51..dc85742577 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_ag2.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_ag2.py @@ -15,14 +15,14 @@ import dataclasses import importlib import json +import os from typing import Optional from unittest import mock from google import auth import agentplatform -from google.cloud.aiplatform import initializer -from agentplatform import agent_engines -from agentplatform._genai import _agent_engines_utils +from agentplatform import frameworks +from agentplatform._genai import _runtimes_utils import pytest @@ -94,7 +94,7 @@ def dataclasses_is_dataclass_mock(): @pytest.fixture def to_json_serializable_autogen_object_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "to_json_serializable_autogen_object", ) as to_json_serializable_autogen_object_mock: to_json_serializable_autogen_object_mock.return_value = {} @@ -104,7 +104,7 @@ def to_json_serializable_autogen_object_mock(): @pytest.fixture def cloud_trace_exporter_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_cloud_trace_exporter_or_warn", ) as cloud_trace_exporter_mock: yield cloud_trace_exporter_mock @@ -127,7 +127,7 @@ def simple_span_processor_mock(): @pytest.fixture def autogen_instrumentor_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_autogen_or_warn", ) as autogen_instrumentor_mock: yield autogen_instrumentor_mock @@ -136,7 +136,7 @@ def autogen_instrumentor_mock(): @pytest.fixture def autogen_instrumentor_none_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_autogen_or_warn", ) as autogen_instrumentor_mock: autogen_instrumentor_mock.return_value = None @@ -146,7 +146,7 @@ def autogen_instrumentor_none_mock(): @pytest.fixture def autogen_tools_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_autogen_tools_or_warn", ) as autogen_tools_mock: autogen_tools_mock.return_value = mock.MagicMock() @@ -170,24 +170,24 @@ def model_dump_json(self): @pytest.mark.usefixtures("google_auth_mock") class TestAG2Agent: def setup_method(self): - importlib.reload(initializer) importlib.reload(agentplatform) - agentplatform.init( - project=_TEST_PROJECT, - location=_TEST_LOCATION, - ) + os.environ["GOOGLE_CLOUD_PROJECT"] = _TEST_PROJECT + os.environ["GOOGLE_CLOUD_LOCATION"] = _TEST_LOCATION def teardown_method(self): - initializer.global_pool.shutdown(wait=True) + for key in [ + "GOOGLE_CLOUD_PROJECT", + "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", + "GOOGLE_CLOUD_LOCATION", + ]: + os.environ.pop(key, None) def test_initialization(self): - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME ) assert agent._tmpl_attrs.get("model_name") == _TEST_MODEL assert agent._tmpl_attrs.get("runnable_name") == _TEST_RUNNABLE_NAME - assert agent._tmpl_attrs.get("project") == _TEST_PROJECT - assert agent._tmpl_attrs.get("location") == _TEST_LOCATION assert agent._tmpl_attrs.get("runnable") is None def test_initialization_with_tools(self, autogen_tools_mock): @@ -195,7 +195,7 @@ def test_initialization_with_tools(self, autogen_tools_mock): place_tool_query, place_photo_query, ] - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, system_instruction=_TEST_SYSTEM_INSTRUCTION, @@ -210,7 +210,7 @@ def test_initialization_with_tools(self, autogen_tools_mock): assert agent._tmpl_attrs.get("ag2_tool_objects") def test_set_up(self): - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, runnable_builder=lambda **kwargs: kwargs, @@ -220,7 +220,7 @@ def test_set_up(self): assert agent._tmpl_attrs.get("runnable") is not None def test_clone(self): - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, runnable_builder=lambda **kwargs: kwargs, @@ -234,7 +234,7 @@ def test_clone(self): assert agent_clone._tmpl_attrs.get("runnable") is not None def test_query(self, to_json_serializable_autogen_object_mock): - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, ) @@ -262,7 +262,7 @@ def test_enable_tracing( simple_span_processor_mock, autogen_instrumentor_mock, ): - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, enable_tracing=True, @@ -275,7 +275,7 @@ def test_enable_tracing( @pytest.mark.usefixtures("caplog") def test_enable_tracing_warning(self, caplog, autogen_instrumentor_none_mock): - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, enable_tracing=True, @@ -294,7 +294,7 @@ def _return_input_no_typing(input_): class TestConvertToolsOrRaiseErrors: def test_raise_untyped_input_args(self, agentplatform_init_mock): with pytest.raises(TypeError, match=r"has untyped input_arg"): - agent_engines.AG2Agent( + frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, tools=[_return_input_no_typing], @@ -302,23 +302,23 @@ def test_raise_untyped_input_args(self, agentplatform_init_mock): class TestToJsonSerializableAutoGenObject: - """Tests for `_agent_engines_utils.to_json_serializable_autogen_object`.""" + """Tests for `_runtimes_utils.to_json_serializable_autogen_object`.""" def test_autogen_chat_result( self, dataclasses_asdict_mock, dataclasses_is_dataclass_mock, ): - mock_chat_result: _agent_engines_utils.AutogenChatResult = mock.Mock( - spec=_agent_engines_utils.AutogenChatResult + mock_chat_result: _runtimes_utils.AutogenChatResult = mock.Mock( + spec=_runtimes_utils.AutogenChatResult ) - _agent_engines_utils.to_json_serializable_autogen_object(mock_chat_result) + _runtimes_utils.to_json_serializable_autogen_object(mock_chat_result) dataclasses_is_dataclass_mock.assert_called_once_with(mock_chat_result) dataclasses_asdict_mock.assert_called_once_with(mock_chat_result) def test_autogen_run_response(self): - mock_response: _agent_engines_utils.AutogenRunResponse = mock.Mock( - spec=_agent_engines_utils.AutogenRunResponse + mock_response: _runtimes_utils.AutogenRunResponse = mock.Mock( + spec=_runtimes_utils.AutogenRunResponse ) mock_agent = MockAgent( name="TestAgent", @@ -342,13 +342,13 @@ def test_autogen_run_response(self): }, "cost": {"total_cost": 5.5}, } - got = _agent_engines_utils.to_json_serializable_autogen_object(mock_response) + got = _runtimes_utils.to_json_serializable_autogen_object(mock_response) mock_response.process.assert_called_once() assert got == want def test_autogen_empty_run_response(self): - mock_response: _agent_engines_utils.AutogenRunResponse = mock.Mock( - spec=_agent_engines_utils.AutogenRunResponse + mock_response: _runtimes_utils.AutogenRunResponse = mock.Mock( + spec=_runtimes_utils.AutogenRunResponse ) mock_response.summary = None mock_response.messages = [] @@ -362,7 +362,7 @@ def test_autogen_empty_run_response(self): "last_speaker": None, "cost": None, } - got = _agent_engines_utils.to_json_serializable_autogen_object(mock_response) + got = _runtimes_utils.to_json_serializable_autogen_object(mock_response) assert got == want @@ -377,7 +377,7 @@ class SimpleDataClass: instance = SimpleDataClass(field1="value1", field2=123) want = {"field1": "value1", "field2": 123} - got = _agent_engines_utils._dataclass_to_dict_or_raise(instance) + got = _runtimes_utils._dataclass_to_dict_or_raise(instance) assert got == want def test_not_a_dataclass_raises_type_error(self): @@ -386,4 +386,4 @@ class NotADataclass: instance = NotADataclass() with pytest.raises(TypeError, match="Object is not a dataclass"): - _agent_engines_utils._dataclass_to_dict_or_raise(instance) + _runtimes_utils._dataclass_to_dict_or_raise(instance) diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_langchain.py b/tests/unit/agentplatform/frameworks/test_frameworks_langchain.py index 8d50518aff..d468621562 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_langchain.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_langchain.py @@ -13,15 +13,15 @@ # limitations under the License. # import importlib +import os from typing import Optional from unittest import mock from google import auth import agentplatform -from google.cloud.aiplatform import initializer -from agentplatform import agent_engines +from agentplatform import frameworks -from agentplatform._genai import _agent_engines_utils +from agentplatform._genai import _runtimes_utils import pytest @@ -90,7 +90,7 @@ def langchain_dump_mock(): @pytest.fixture def cloud_trace_exporter_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_cloud_trace_exporter_or_warn", ) as cloud_trace_exporter_mock: yield cloud_trace_exporter_mock @@ -113,7 +113,7 @@ def simple_span_processor_mock(): @pytest.fixture def langchain_instrumentor_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_langchain_or_warn", ) as langchain_instrumentor_mock: yield langchain_instrumentor_mock @@ -122,7 +122,7 @@ def langchain_instrumentor_mock(): @pytest.fixture def langchain_instrumentor_none_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_langchain_or_warn", ) as langchain_instrumentor_mock: langchain_instrumentor_mock.return_value = None @@ -132,12 +132,9 @@ def langchain_instrumentor_none_mock(): @pytest.mark.usefixtures("google_auth_mock") class TestLangchainAgent: def setup_method(self): - importlib.reload(initializer) importlib.reload(agentplatform) - agentplatform.init( - project=_TEST_PROJECT, - location=_TEST_LOCATION, - ) + os.environ["GOOGLE_CLOUD_PROJECT"] = _TEST_PROJECT + os.environ["GOOGLE_CLOUD_LOCATION"] = _TEST_LOCATION self.prompt = { "input": lambda x: x["input"], "agent_scratchpad": ( @@ -152,13 +149,16 @@ def setup_method(self): self.output_parser = mock.Mock() def teardown_method(self): - initializer.global_pool.shutdown(wait=True) + for key in [ + "GOOGLE_CLOUD_PROJECT", + "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", + "GOOGLE_CLOUD_LOCATION", + ]: + os.environ.pop(key, None) def test_initialization(self): - agent = agent_engines.LangchainAgent(model=_TEST_MODEL) + agent = frameworks.LangchainAgent(model=_TEST_MODEL) assert agent._tmpl_attrs.get("model_name") == _TEST_MODEL - assert agent._tmpl_attrs.get("project") == _TEST_PROJECT - assert agent._tmpl_attrs.get("location") == _TEST_LOCATION assert agent._tmpl_attrs.get("runnable") is None def test_initialization_with_tools(self): @@ -166,7 +166,7 @@ def test_initialization_with_tools(self): place_tool_query, StructuredTool.from_function(place_photo_query), ] - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, system_instruction=_TEST_SYSTEM_INSTRUCTION, tools=tools, @@ -180,7 +180,7 @@ def test_initialization_with_tools(self): assert agent._tmpl_attrs.get("runnable") is not None def test_set_up(self): - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, prompt=self.prompt, output_parser=self.output_parser, @@ -192,7 +192,7 @@ def test_set_up(self): assert agent._tmpl_attrs.get("runnable") is not None def test_clone(self): - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, prompt=self.prompt, output_parser=self.output_parser, @@ -208,7 +208,7 @@ def test_clone(self): assert agent_clone._tmpl_attrs.get("runnable") is not None def test_query(self, langchain_dump_mock): - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, prompt=self.prompt, output_parser=self.output_parser, @@ -222,7 +222,7 @@ def test_query(self, langchain_dump_mock): ) def test_stream_query(self, langchain_dump_mock): - agent = agent_engines.LangchainAgent(model=_TEST_MODEL) + agent = frameworks.LangchainAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].stream.return_value = [] list(agent.stream_query(input="test stream query")) @@ -240,7 +240,7 @@ def test_enable_tracing( simple_span_processor_mock, langchain_instrumentor_mock, ): - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, prompt=self.prompt, output_parser=self.output_parser, @@ -257,7 +257,7 @@ def test_enable_tracing( @pytest.mark.usefixtures("caplog") def test_enable_tracing_warning(self, caplog, langchain_instrumentor_none_mock): - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, prompt=self.prompt, output_parser=self.output_parser, @@ -277,7 +277,7 @@ def _return_input_no_typing(input_): class TestConvertToolsOrRaiseErrors: def test_raise_untyped_input_args(self, agentplatform_init_mock): with pytest.raises(TypeError, match=r"has untyped input_arg"): - agent_engines.LangchainAgent( + frameworks.LangchainAgent( model=_TEST_MODEL, tools=[_return_input_no_typing], ) @@ -289,7 +289,7 @@ def test_raise_both_system_instruction_and_prompt_error(self, agentplatform_init ValueError, match=r"Only one of `prompt` or `system_instruction` should be specified.", ): - agent_engines.LangchainAgent( + frameworks.LangchainAgent( model=_TEST_MODEL, system_instruction=_TEST_SYSTEM_INSTRUCTION, prompt=prompts.ChatPromptTemplate.from_messages( diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_langgraph.py b/tests/unit/agentplatform/frameworks/test_frameworks_langgraph.py index 82fa32d53b..099cfeb56b 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_langgraph.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_langgraph.py @@ -13,14 +13,14 @@ # limitations under the License. # import importlib +import os from typing import Any, Dict, List, Optional from unittest import mock from google import auth import agentplatform -from google.cloud.aiplatform import initializer -from agentplatform import agent_engines -from agentplatform._genai import _agent_engines_utils +from agentplatform import frameworks +from agentplatform._genai import _runtimes_utils import pytest from langchain_core import runnables @@ -101,7 +101,7 @@ def langchain_dump_mock(): @pytest.fixture def cloud_trace_exporter_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_cloud_trace_exporter_or_warn", ) as cloud_trace_exporter_mock: yield cloud_trace_exporter_mock @@ -124,7 +124,7 @@ def simple_span_processor_mock(): @pytest.fixture def langchain_instrumentor_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_langchain_or_warn", ) as langchain_instrumentor_mock: yield langchain_instrumentor_mock @@ -133,7 +133,7 @@ def langchain_instrumentor_mock(): @pytest.fixture def langchain_instrumentor_none_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_langchain_or_warn", ) as langchain_instrumentor_mock: langchain_instrumentor_mock.return_value = None @@ -143,21 +143,21 @@ def langchain_instrumentor_none_mock(): @pytest.mark.usefixtures("google_auth_mock") class TestLanggraphAgent: def setup_method(self): - importlib.reload(initializer) importlib.reload(agentplatform) - agentplatform.init( - project=_TEST_PROJECT, - location=_TEST_LOCATION, - ) + os.environ["GOOGLE_CLOUD_PROJECT"] = _TEST_PROJECT + os.environ["GOOGLE_CLOUD_LOCATION"] = _TEST_LOCATION def teardown_method(self): - initializer.global_pool.shutdown(wait=True) + for key in [ + "GOOGLE_CLOUD_PROJECT", + "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", + "GOOGLE_CLOUD_LOCATION", + ]: + os.environ.pop(key, None) def test_initialization(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) assert agent._tmpl_attrs.get("model_name") == _TEST_MODEL - assert agent._tmpl_attrs.get("project") == _TEST_PROJECT - assert agent._tmpl_attrs.get("location") == _TEST_LOCATION assert agent._tmpl_attrs.get("runnable") is None def test_initialization_with_tools(self): @@ -165,7 +165,7 @@ def test_initialization_with_tools(self): place_tool_query, StructuredTool.from_function(place_photo_query), ] - agent = agent_engines.LanggraphAgent( + agent = frameworks.LanggraphAgent( model=_TEST_MODEL, tools=tools, model_builder=lambda **kwargs: kwargs, @@ -178,7 +178,7 @@ def test_initialization_with_tools(self): assert agent._tmpl_attrs.get("runnable") is not None def test_set_up(self): - agent = agent_engines.LanggraphAgent( + agent = frameworks.LanggraphAgent( model=_TEST_MODEL, model_builder=lambda **kwargs: kwargs, runnable_builder=lambda **kwargs: kwargs, @@ -188,7 +188,7 @@ def test_set_up(self): assert agent._tmpl_attrs.get("runnable") is not None def test_clone(self): - agent = agent_engines.LanggraphAgent( + agent = frameworks.LanggraphAgent( model=_TEST_MODEL, model_builder=lambda **kwargs: kwargs, runnable_builder=lambda **kwargs: kwargs, @@ -202,7 +202,7 @@ def test_clone(self): assert agent_clone._tmpl_attrs.get("runnable") is not None def test_query(self, langchain_dump_mock): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() mocks = mock.Mock() mocks.attach_mock(mock=agent._tmpl_attrs.get("runnable"), attribute="invoke") @@ -217,7 +217,7 @@ def test_query(self, langchain_dump_mock): ) def test_stream_query(self, langchain_dump_mock): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].stream.return_value = [] list(agent.stream_query(input="test stream query")) @@ -238,7 +238,7 @@ def test_enable_tracing( simple_span_processor_mock, langchain_instrumentor_mock, ): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL, enable_tracing=True) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL, enable_tracing=True) assert agent._tmpl_attrs.get("instrumentor") is None # TODO(b/384730642): Re-enable this test once the parent issue is fixed. # agent.set_up() @@ -250,21 +250,21 @@ def test_enable_tracing( @pytest.mark.usefixtures("caplog") def test_enable_tracing_warning(self, caplog, langchain_instrumentor_none_mock): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL, enable_tracing=True) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL, enable_tracing=True) assert agent._tmpl_attrs.get("instrumentor") is None # TODO(b/383923584): Re-enable this test once the parent issue is fixed. # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text def test_get_state_history_empty(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].get_state_history.return_value = [] history = list(agent.get_state_history()) assert history == [] def test_get_state_history(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].get_state_history.return_value = [ mock.Mock(), @@ -283,7 +283,7 @@ def test_get_state_history(self): ] def test_get_state_history_with_config(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].get_state_history.return_value = [ mock.Mock(), @@ -302,7 +302,7 @@ def test_get_state_history_with_config(self): ] def test_get_state(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].get_state.return_value = mock.Mock() agent._tmpl_attrs["runnable"].get_state.return_value._asdict.return_value = { @@ -312,7 +312,7 @@ def test_get_state(self): assert state == {"test_key": "test_value"} def test_get_state_with_config(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].get_state.return_value = mock.Mock() agent._tmpl_attrs["runnable"].get_state.return_value._asdict.return_value = { @@ -322,13 +322,13 @@ def test_get_state_with_config(self): assert state == {"test_key": "test_value"} def test_update_state(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent.update_state() agent._tmpl_attrs["runnable"].update_state.assert_called_once() def test_update_state_with_config(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent.update_state(config=_TEST_CONFIG) agent._tmpl_attrs["runnable"].update_state.assert_called_once_with( @@ -336,7 +336,7 @@ def test_update_state_with_config(self): ) def test_update_state_with_config_and_kwargs(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent.update_state(config=_TEST_CONFIG, test_key="test_value") agent._tmpl_attrs["runnable"].update_state.assert_called_once_with( @@ -344,7 +344,7 @@ def test_update_state_with_config_and_kwargs(self): ) def test_register_operations(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) expected_operations = { "": ["query", "get_state", "update_state"], "stream": ["stream_query", "get_state_history"], @@ -360,6 +360,6 @@ def _return_input_no_typing(input_): class TestConvertToolsOrRaiseErrors: def test_raise_untyped_input_args(self, agentplatform_init_mock): with pytest.raises(TypeError, match=r"has untyped input_arg"): - agent_engines.LanggraphAgent( + frameworks.LanggraphAgent( model=_TEST_MODEL, tools=[_return_input_no_typing] ) diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py b/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py index aaef28749b..1f3e002449 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py @@ -14,15 +14,15 @@ # import importlib import json +import os from unittest import mock from google import auth import agentplatform -from google.cloud.aiplatform import initializer -from agentplatform.agent_engines.templates import ( +from agentplatform.frameworks import ( llama_index, ) -from agentplatform._genai.agent_engines import _agent_engines_utils +from agentplatform._genai import _runtimes_utils from llama_index.core import prompts from llama_index.core.base.llms import types @@ -71,7 +71,7 @@ def model_builder_mock(): @pytest.fixture def cloud_trace_exporter_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_cloud_trace_exporter_or_warn", ) as cloud_trace_exporter_mock: yield cloud_trace_exporter_mock @@ -94,7 +94,7 @@ def simple_span_processor_mock(): @pytest.fixture def llama_index_instrumentor_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_llama_index_or_warn", ) as llama_index_instrumentor_mock: yield llama_index_instrumentor_mock @@ -103,7 +103,7 @@ def llama_index_instrumentor_mock(): @pytest.fixture def llama_index_instrumentor_none_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_llama_index_or_warn", ) as llama_index_instrumentor_mock: llama_index_instrumentor_mock.return_value = None @@ -113,7 +113,7 @@ def llama_index_instrumentor_none_mock(): @pytest.fixture def nest_asyncio_apply_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_nest_asyncio_or_warn", ) as nest_asyncio_apply_mock: yield nest_asyncio_apply_mock @@ -122,12 +122,9 @@ def nest_asyncio_apply_mock(): @pytest.mark.usefixtures("google_auth_mock") class TestLlamaIndexQueryPipelineAgent: def setup_method(self): - importlib.reload(initializer) importlib.reload(agentplatform) - agentplatform.init( - project=_TEST_PROJECT, - location=_TEST_LOCATION, - ) + os.environ["GOOGLE_CLOUD_PROJECT"] = _TEST_PROJECT + os.environ["GOOGLE_CLOUD_LOCATION"] = _TEST_LOCATION self.prompt = prompts.ChatPromptTemplate( message_templates=[ types.ChatMessage( @@ -142,13 +139,16 @@ def setup_method(self): ) def teardown_method(self): - initializer.global_pool.shutdown(wait=True) + for key in [ + "GOOGLE_CLOUD_PROJECT", + "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", + "GOOGLE_CLOUD_LOCATION", + ]: + os.environ.pop(key, None) def test_initialization(self): agent = llama_index.LlamaIndexQueryPipelineAgent(model=_TEST_MODEL) assert agent._model_name == _TEST_MODEL - assert agent._project == _TEST_PROJECT - assert agent._location == _TEST_LOCATION assert agent._runnable is None def test_set_up(self): @@ -264,17 +264,17 @@ class TestToJsonSerializableLlamaIndexObject: """Tests for `_utils.to_json_serializable_llama_index_object`.""" def test_llama_index_response(self): - mock_response: _agent_engines_utils.LlamaIndexResponse = mock.Mock( - spec=_agent_engines_utils.LlamaIndexResponse + mock_response: _runtimes_utils.LlamaIndexResponse = mock.Mock( + spec=_runtimes_utils.LlamaIndexResponse ) mock_response.response = "test response" mock_response.source_nodes = [ mock.Mock( - spec=_agent_engines_utils.LlamaIndexBaseModel, + spec=_runtimes_utils.LlamaIndexBaseModel, model_dump_json=lambda: '{"name": "model1"}', ), mock.Mock( - spec=_agent_engines_utils.LlamaIndexBaseModel, + spec=_runtimes_utils.LlamaIndexBaseModel, model_dump_json=lambda: '{"name": "model2"}', ), ] @@ -285,61 +285,61 @@ def test_llama_index_response(self): "source_nodes": ['{"name": "model1"}', '{"name": "model2"}'], "metadata": {"key": "value"}, } - got = _agent_engines_utils.to_json_serializable_llama_index_object(mock_response) + got = _runtimes_utils.to_json_serializable_llama_index_object(mock_response) assert got == want def test_llama_index_chat_response(self): - mock_chat_response: _agent_engines_utils.LlamaIndexChatResponse = mock.Mock( - spec=_agent_engines_utils.LlamaIndexChatResponse + mock_chat_response: _runtimes_utils.LlamaIndexChatResponse = mock.Mock( + spec=_runtimes_utils.LlamaIndexChatResponse ) mock_chat_response.message = mock.Mock( - spec=_agent_engines_utils.LlamaIndexBaseModel, + spec=_runtimes_utils.LlamaIndexBaseModel, model_dump_json=lambda: '{"content": "chat message"}', ) want = {"content": "chat message"} - got = _agent_engines_utils.to_json_serializable_llama_index_object(mock_chat_response) + got = _runtimes_utils.to_json_serializable_llama_index_object(mock_chat_response) assert got == want def test_llama_index_base_model(self): - mock_base_model: _agent_engines_utils.LlamaIndexBaseModel = mock.Mock( - spec=_agent_engines_utils.LlamaIndexBaseModel + mock_base_model: _runtimes_utils.LlamaIndexBaseModel = mock.Mock( + spec=_runtimes_utils.LlamaIndexBaseModel ) mock_base_model.model_dump_json = lambda: '{"name": "test_model"}' want = {"name": "test_model"} - got = _agent_engines_utils.to_json_serializable_llama_index_object(mock_base_model) + got = _runtimes_utils.to_json_serializable_llama_index_object(mock_base_model) assert got == want def test_sequence_of_llama_index_base_model(self): - mock_base_model1: _agent_engines_utils.LlamaIndexBaseModel = mock.Mock( - spec=_agent_engines_utils.LlamaIndexBaseModel + mock_base_model1: _runtimes_utils.LlamaIndexBaseModel = mock.Mock( + spec=_runtimes_utils.LlamaIndexBaseModel ) mock_base_model1.model_dump_json = lambda: '{"name": "test_model1"}' - mock_base_model2: _agent_engines_utils.LlamaIndexBaseModel = mock.Mock( - spec=_agent_engines_utils.LlamaIndexBaseModel + mock_base_model2: _runtimes_utils.LlamaIndexBaseModel = mock.Mock( + spec=_runtimes_utils.LlamaIndexBaseModel ) mock_base_model2.model_dump_json = lambda: '{"name": "test_model2"}' mock_base_model_list = [mock_base_model1, mock_base_model2] want = [{"name": "test_model1"}, {"name": "test_model2"}] - got = _agent_engines_utils.to_json_serializable_llama_index_object(mock_base_model_list) + got = _runtimes_utils.to_json_serializable_llama_index_object(mock_base_model_list) assert got == want def test_sequence_of_mixed_types(self): - mock_base_model: _agent_engines_utils.LlamaIndexBaseModel = mock.Mock( - spec=_agent_engines_utils.LlamaIndexBaseModel + mock_base_model: _runtimes_utils.LlamaIndexBaseModel = mock.Mock( + spec=_runtimes_utils.LlamaIndexBaseModel ) mock_base_model.model_dump_json = lambda: '{"name": "test_model"}' mock_string = "test_string" mock_list = [mock_base_model, mock_string] want = [{"name": "test_model"}, "test_string"] - got = _agent_engines_utils.to_json_serializable_llama_index_object(mock_list) + got = _runtimes_utils.to_json_serializable_llama_index_object(mock_list) assert got == want def test_other_type(self): test_dict = {"name": "test_model"} want = "{'name': 'test_model'}" - got = _agent_engines_utils.to_json_serializable_llama_index_object(test_dict) + got = _runtimes_utils.to_json_serializable_llama_index_object(test_dict) assert got == want diff --git a/tests/unit/agentplatform/genai/replays/conftest.py b/tests/unit/agentplatform/genai/replays/conftest.py index c0ddc6de57..05c74872b5 100644 --- a/tests/unit/agentplatform/genai/replays/conftest.py +++ b/tests/unit/agentplatform/genai/replays/conftest.py @@ -22,7 +22,7 @@ from agentplatform._genai import ( client as agentplatform_genai_client_module, ) -from agentplatform._genai import _agent_engines_utils +from agentplatform._genai import _runtimes_utils from google.cloud import storage, bigquery from google.genai import _replay_api_client from google.genai import types as genai_types @@ -154,30 +154,30 @@ def is_replay_mode(request): @pytest.fixture -def mock_agent_engine_create_path_exists(): +def mock_runtime_create_path_exists(): """Mocks os.path.exists to return True.""" with mock.patch("os.path.exists", return_value=True) as mock_exists: yield mock_exists @pytest.fixture -def mock_agent_engine_create_base64_encoded_tarball(): +def mock_runtime_create_base64_encoded_tarball(): """Mocks the _create_base64_encoded_tarball function.""" with mock.patch.object( - _agent_engines_utils, "_create_base64_encoded_tarball" + _runtimes_utils, "_create_base64_encoded_tarball" ) as mock_create_base64_encoded_tarball: mock_create_base64_encoded_tarball.return_value = "H4sIAAAAAAAAA-3UvWrDMBAHcM9-CpEpGRLkD8VQ6JOUElT7LFxkydEHxG9f2V1CKXSyu_x_i6TjJN2gk6N7HByNZIK_hEfINsCTa82zilcNTyMvRSPKao2vBM8KwZu6vJZXITJepGyRMb5FMT9FH6RjLHsM0mpr1CyN-i1vcsMo3aycjdMede0kV9YqTedW29id5TBpGXrrxjep0pO4kVGDIf-e_3edsI1APtxG20VNl2ne5o6_-r-oRer_Ypk2dd0s_Z82oP_3kLdaes-ensFLzpKOenaP5OajJ92fvoMLRyE6ww7LjrTwkzWeDnm-nmA_PqkN7PX5vOMJnwcAAAAAAAAAAAAAAAAAAADAdr4AI-kzQQAoAAA=" yield mock_create_base64_encoded_tarball @pytest.fixture -def mock_agent_engine_create_docker_base64_encoded_tarball(): +def mock_runtime_create_docker_base64_encoded_tarball(): """Mocks the _create_base64_encoded_tarball function.""" with mock.patch.object( - _agent_engines_utils, "_create_base64_encoded_tarball" - ) as mock_agent_engine_create_docker_base64_encoded_tarball: - mock_agent_engine_create_docker_base64_encoded_tarball.return_value = "H4sIAAAAAAAAA+xdzW8bSXanPYska2DXQRBgMbda7kE2QrW6+pMU1pvYsux11iNrJc1MDEMgWmSJ7HGzm+5uyvYMBtgsctzDHjxIECRAgOSYY475C5L/Yo4Bcshtj6kPypTF7mq+MtX6qoe2aDb796pe1atXVa+qXxlrjTMnk5LvuvyT0ulP/n/sYsfGju9aDr3v+9hvIPfss9ZoTLI8SBFqpEmSy56r+v2SkrFm9JL4MBycoR4sWP+25TiuadsNE2PLdnT910Gz+jeiIMu72SQ9Im+74zQZjXPjbTCKPj4NVsGe45TXP/3tff1j+hz2sEvr3/z4pKvpmtc/r3RR2908HJF1hH2v7fq241qG7fu27dw67zxqOjs63f4n436Qk25vSHovja+yJF5CGlXt3zL9Wfu3LNr+Xfa4bv810DfNuWrndqB5whC0DcelIzPst1DB0yk5CrMwiSnCMi3PxGYHW9h2TPp0nOThYdgLcvp7Rh94sX/MIg4GPB1295tvvz3vYri2NGv/gx45ozR4+59v9yfGf27DMm069Pd9i93H2Dexbv910KMgys6q3jVdfDrd/yd0HBjGNY//Tfu4/3d8z+Tjf6zbfy30zbd6eH+dadb+g14eHtFBHf+21DSq278/a/8OG//bvmfp9l8H9clhMImuomSaFqFZ+x+G/T6Ju4NelEz6UzPQncTUJqQZ6fZJ1kvDcZ6kXTr3C7q9gM7+pk9lRv9Aloa8/dvz/l/L9Vxbt/86aPfXT8OcoMMkHQU5sht/2rhxo/FXtDwajU+m/47pBv33g1Pfq+iThvHlX//o9h8aN2+njdvp7f9eYtY1adKkSZMmTZo0adKkSZOmC0xf3vjjn3z66Y2//WEeHEREeFDE35sbO5v39zbR3v0HTzeRuIfuTD0xQZ6naG/zb/bQ9s6Tz+7vPEe/2nzeQkdBNCHowdNnD+7+xc0/+slffnqjEcZ98iZ7FdF5fTeY5An/fuzQweKTzev/hOXmR/TfjR//V4NemjRp0qRJkyZNmjRp0qRJ0yWg39o3f3D377YGSTKISDAOM6OXjL5pHu9geJiMgpBtUm9++ESzhd4/sztM0nwr4Lvfm+wHOn0ekjif7mCfsWD3Db5RwhDcjjmN0+Qr0su3U3IYvjnmwh/8khzM8EXQftKbjGhip5IqepRlVuy4b2LDNEx2L6NAnu9xGo6C9C3PfpqHh0Ev3yGDMMvTtzOu45cDo0+Omnz3343b/9uglyZNmjRp0qRJkyZNmjRp0qTpKtCffXL3xofej5uz9z+mbwIt+E4HhADvf0zf/3Q839Pvf9RBkvc/bjY+/v2Pm/r9D02aNGnSpEmTJk2aNGnSdD3pAr7/cft/GvTSpEmTJk2aNGnSpEmTJk2aNF16+vMbn/zs014UkjjvBmlv+KbtdT1HvP/xfw16adKkSZMmTZo0adKkSZMmTZquCv34k5/98IQXYPb+R5QMsrM5BJqf8LPY+c+O77Dzny12JKg+/7kGOlX/7ARPw8QG9paoCoD6d03bZec/Wbau/1qovP6xY1i2YWOj08bY9gz6u2Ia0ve/sGNjzz51/rfrYH3+Vy3EKnzVxKvYQ9hZt+x1G7c6bRs93Hzw+WNEicmNTtCjIIxIH+UJSknQR483NhE/Cwod0vvr6MVmmsYJsvbRVoKySW/I76MkRf0wJb08Sd+uo5U1qnM57YJOHj67cmsuL47TopqzSF74UcRoRPKAnU6FMpIekXQd/XySRsmYxIikKc0CO3C4j5JJ/ouStPzytJ7FXNbDNBm9Z88PTy1h1Sln9TQJ+jQfG8loFMS0CNNkMqYltyKO3lppoZVeMhonMS2gbGW/mD+WVBGIP/0mjnMuTcktT2lnEsdhPEAvBGtjxtcQTPfR6zAfoiAd8NhWGc3G6ip7ilbc6vhtPkziddTcSyek2UKrq68mIclnNzaefbb9bGtza2/1ycPddY/+8IJmPSUs0yLF1T4ZcxkOXs1unvi5l/Zsq8dvZJM8jFb2m2VieujJ1qNnXKxut/voydPN7rOtp8/p/9EDMgiFoEIqA+0NwwyN06RHsgyNgrcoD14SqhdULYIIjcJ4kpPMuFWclG0upmYjMqLNRTQvqbK57SUz9KwZQ9qEovDAZi01pg04TOJxkkRoN2fRxWiRxOQ1+uXe3vYumj2B7uC766gfnYhbRvnahYl1On5lYsM8H2fra2tzHFHz8eYeovfXeH1n/ZdrvWFAsVG2lgbjsL82U8lVix9ozzO7Rru4JrJME1mO61sFGXO5OShViaKqZZC2BPI8maSoN0lTmhv0mMuBNli20cbTJ2ga1A2FtJG4rsXiuikk8iSm/RktQ1oxM8lF7U8T+Bju5ZBS/dsYJhk1wP0wG0e0nTziLxqv873PLw6SN608zCNyr7k3JBk5meXXYRShA0J7mVFyRPpGc/9On5r3MMqMKa9uHIzIehQckOgeC9q3HkThIL4XkcO8NZXVmH52szylRTJ9+ItpQYjn03AwzFus2zCy8GvC/9z5mqTJvWazRVvyPWw6bdf37k7Ru/T3k9C7JWUi6QA+tkyEFdJlcrJMQqH4l7BUXImZ+f7d775/95trc/2+pITgJkkG+f7db1ERLaBbhTjGsKaM/9N511Dl9e/Le+zfzlUbGDFjUfYbx00NxPFXhFib1xrxwfWfy3usPo2Aox6Eg19PCB1eH8+7noYxQXts+FpG8EQUIBJpOpKEKFkGNixfAQuHqGYSG230WfigrjzCQWXqt/SykKjfne0oyFmEH7Q7Jr3wMOzdlanfxajZqfq5FzqTP0f4UqqfZ6qUhQz1YBL3mS9um/t0kG1gu1TvpNZPlogCRE0aSkwCA5sXOJPYbhteufotP49wUG3qJ9wou3mSBgOyWA98gWtWkGvYJdbvYmQSW4azXPXDKuonAZWpnzwdhdxVqV+BJYQnolZqCqhK9bsYmZR2vsvPYwnIgqufBCIpCwlqzpe8kaSEqt9BGqQhyQDqJ0lEIV9q0lCAS3tew8bnnUlbXf1kUDikXP0koDL1k6ejkLsP1W9nw7Y20C+DbCiZfdRVako1i9i8ki2QnHMmHfnM15WonwwKh5SrnwRUpn7ydBRyJ5Z6US8KUZ+MSdwnca/E6k27NpUCr0saCqDK5xjYu8CZrLJ+l1X93p23r7Ty+o/lPfYPtRRqGUSyqwS+ZYExXHDvzuIM23XtgfA+3M6z/D0QO5tPN+/vbna3nu1t7p7a/2A7btsuzpXrSkbEj9h+qiFB1MaSLEeHkyhCKYlIkBEUJ/ReC43Ft6MwC/P1W+i9BKePtFtjWZ5iuxx7q0hrWIYk47htkjJjz6ojjGnaQe94odwwjFKGMnPw3bvvv/v9Jbz+XknYf6FY9DB5HUdJ0GdL4KAVDZ7wvyol/I8loKVaCMZQspNQjWH1xqUlWQi/5Z7cO3GWu6TWBHRV7J+jj64evFplOTKx2cEWth3ToEIZg69PGRLcsbDnd4pzb8kmYN8VLftr0MeCZJOOawWSWh8NujAgV4M06AOQZE+2Bl12kGzbowZdR5BsF5MGXUeQbD+MGuifi+aoFGXLVv+ulCegUliQJ+ADL7thGBJPQGXCRZ4ABrKWOXFnDCUvz6kxdOvyBLRbzgKOybPzBMThG+YNcLFpY7ZUYfll3gC76K0qkX1Zm750l4qUhUaIonzZLp0rZYQqhS0wQgVbDMuo3AhVJlxkhBhIarnOvZCXpLZyKUvVVuqfumpqKxV2MbUtW5qWq63U+VWmts4yuzrGcKluecbQq6vv7LScs11nk/adQg2mL73bqxPWk0ZhPHmzKiLxim6V9qjYMzF2SrpVx7V93/KLtkZ1Wh24o5iBwK4ZBgIP5ylIPvorA4F9vgwE9kgwEHi2QUGOikwO2D/KQCoyOSoyyZZZy0EqMknffS0Fqcgk3QJYClJpGp6KTPIxSRlIpWn4KjL5Ks29rdI02ioytVWaRltFpo6KTB2VptFRkakDbRqu2TJNqEwcBJWJg6AycZCKTBja3DlIRSYMbe4cpCITuHfnIBWZwO5bDlKRyVaRyVaRyVaRyYGaZQ5SkcmBmmUOUpHJVZHJhZplDlKRSfoucxlI+npTKUhFJvCQgINUZAIPCThIRSZfxSzLQUWuAo6SlsQVchVUCzvvKoC8xVjqKqhOuMBVwEGS6DzgmT1j2F5mnDPOENfkKqCJWc65udlFqDruDzA9Otd3HQeXudktOrHxOwXuACZCR8HsWCpDyPpAtqlgFRVBCsPiOkEKo4w6QQodUZ0ghSHkhQcpDE3qBCmMMi48SGEAXidIoQOoE6TUAVxsEFbpamoEqXQ1NYJUupoaQSpdTY0glV7jooNUupoaQSq9Ro0glV7jDEAl7gBb7ge9Wu6ACmHh7oDjXQTSXXfVCRe7A2xrybN328LLZlhXlHIX08nque0cmLoDFt955xS8hDcVQdauL91VIqW0lyk0RAwlNf5XyRAxYaVdybwhekgOg0mUo4zkKDl8H3BDGKQMfUDlhqgy4SJDVAk690JeltrC3ekMJe11r5raAt3pi4bFkveflQmXqe1S3emYOzGXzLA2d7rVwp1zc6ezE2MW7j0t2yoIvMMFkHsDLt2lImWhGbJaltyBfZXMUKWwVWbodHi0RXvPyoSLzFAl6NwLeTlqWyFlqdpKB8RXTW2lwkLVdraDXd57ViZcprbL3LfOGS5z3zpnWNe+ddduYa+mM7IKZp8nD1Kb27DOosJY2MXYtktmo66Nba9EKh/sYKegNtjtrUGXBwReNNCgKw4Cr9NokAbNg8DraRqkQUsCgdc9NUiD5kHg9WkNWhRU6CFgqGvjX68Udt5DAAm8Xu4hqEy4yEPAQMs89JozXKrDnjLsmHV5CJwWPr/4sCdOYp/zD2DTwR52bRPjEv+A6RTLU7G77dJdKlIWWiWGui7xNqqF5VZpdug5LGi1xCpVJlxklZyWBd9WyEDgDo2BwDtuKcgCO2UYCDzoZSAVmeDvnVKQrSIT/L1TBgIP2xgIvG+PgcCeUQoCh6LgIJXKBYei4CAVhYW/rMpAKvXkqsgEf1mVgqQHjZWCVBQW/rIqA6k0Qk9FYStiapWAVBQWHPSCg1QKAr6mQUFwXzkDqWgEfObJQCoF0VYpCHB4DQ5SkUm+c6oYpPCWi6PypoGjshndUdm47fB9vgogFZlUhgS2Su9uq/Tutq2UPZUiV+ndbXhUCQZSkQkeVYKCVDpqGxxoioMUmrsNDjTFQPCoEgykIpP0WDI1UMnk0b42wRqrhYVMHj90ask3vVQmXDx5tCuGNpfuUpGyVG2vTeCQSmHn1RYQGbdCbcE7nSmoYiB66S4VKcvUVj7evmJqKxd2IbUtPbRVqrbyhIvU1m858JkaBcHnGuzsNvDgo9Nqw0csFATfYtNpdeCzGhYAViF78HCfDCTvSudAzrpp8uNA4SDg5JiDMDB4iwABZzUCBBzMCxDQ2SRAvkpBACf8HGQBI4kIEHDKKkBAP6wAAcPeCRDQFyZAwEVtDoKGvuGgDvC0JgECenwFCOiqEyBgLFcBAnpzKAiDA4sKkAVuTxQEjGEtQPDmTkFAL7YAwY0l5uHi4KAOuOViFQtLQcAAGgIE9FAJENAFJEDAxRABAvojBAjoZBcg4D43AQKOjTjIAoZdEiCgo1OAgCHrBQgYAFaAZOdml4KA0UgFCN5RUxDcwla+aVwMgkZQEiAVCwv1+AoQcO1TgIBBbRcBzU8eOcqRLwJcmcnjIsLOTR4hsVLLJo+LJDw3eeSgjqfQhjrABUUBAi6+CRC8tVrg6OwCBFzLECC4gaQgeJFbPAYtHATvyCw64FCQCcNtHQXBbR0FwTsyBlKoXAzcdCBAwAMmBAjekbE33BWahgX04AgQVsiewiiFguBDUAu8H0eAgB4cAVJpuRYwKLkAAc/04CBosFUBAi7QCxB86m2Bl2MFCO72oSD4lI6COgrWyAG6HAUIPuqnIKCfUoBUSs+Bj8UpCO72scBOaA5y4b4iCoLPoikIuF9UgFQ6NRe441GAgDsVBAg+9aYguOeaglQ6NRcY6Z2DPGAAWQFSqVwPuLVGgFTqyVMZfHgqjdCHe6UoCLhpexFQ4eSRouTd+1WaPFYKC548vl+FlOzzWCThoskjA13xUFKLSFmmtvIVlyumtvKtbHNqC4nDKFdbecIlaluxtnDpLhUpS9X2OllbubBzartoHMZqaytPuExtz8XaXva8lCm6fJniiim6XNgKRZdH+pMrujzhEkWvcLNcuktFylK1vR4v8i8iLFBtT+zIq7TP8oSL1bZdcbxUncrVrnCSXKC8lCh6W+5Zu1qKXiHsnKJDAlZIFb0i4TJFl0/+L92lImWZ2mK7PPjG5+N+IIJkJDmrnIAHx+DxN6hFKuVXmouyLEi8QNskZfrBMjFOshyN06RHsox9z3IyzsrzIQk7uvmG9CZcsOkUbh29WFnLqdGdj7KxdhDGa0J/V1poZRaPg31jWVqdZmllfz4jmK3KnQiXMlcgy0HYJ7canRa2T7JeGo7zJO1GYZZTUb9ZmcThEUkz8jAZBWG8so5WhNjBOMxYbBIm2/Ezu8MkzbeCEWGPsR+CST6kBTBVhxkLdt/gBXUiygkvpTT5ivTy7ZQchm+OufAHvyQHM3wRtJ/0JiOa2Kmkih5lmaUPsZ+xYRomu0ctUM7zPU7DUZC+5dlPqS4HvXyHDGhxpG9nXMcvB0afHK18W1yT0leLbvG2QlCfKsdPbxUoOmfglevkxjDJSIz6YTaOgrfoEVX6IF+PKbtSXqVV/nDK5FAwQU3GpnmroenikLHGogQdhoO1KBlka6yODRMb2FvDjmE5Bm4bHddxPM+gvyumYVLyHId/Ujr9afoWbmAX25bjuKZtN0zsOthvIHOpkpbQhDbOFKEGU17Zc1W/X1Kab9TtVseVdFpPk6BP+u995o/TZDJm/VZh31RkwRh/iQEC8Z/2fV1Z38fSk4QK25nEMeuCX4gEjBl342S3uo9eh/kQBemAdwQZzVJBYpZPe3x/OdaV82qftXWtav+mbdjYscxOXe0fO6z906GEbv910Pyg1eYHNZ5Z++f8a2z/PL262j8dEZmOpPCAoyvG65zbv9U2PMczffPs2j+26f9n/T9v/w69pdt/DTTf7bRbniPpwhZpn0yfitoi5e07klMgF+bN5lQkL01C4kuYb+6MoUG5FbTw3c2NvSfPtta2d55tb+7sPaeN7L156NIsBgOSrtHGHBxEpDvhM69ub0h6L5st9MX9p59vUkCeTkizLKOSVVgxkeszT8OYpPlb9GKxpPeNglkfT0wyqCu2S32x3FzKbnHTNOXULGblSg422334KyTEQ1y8jFYQQVPB+4aeTH4kVdr/jmF3LLN9huM/bPn2Cfvvcvtv2dr+10HzDbLTsjsSA/pR9r/Doy6cqf2nSbQlXcwy7P8kJ2uDHumOSB5QyxR0UxL0u3k4Iskk72akd8L8m4W2n2VS4sUrtv2VyRabfpHW0kw/Y2cvx/QzVpKRuTb9NVCl/882TB9b9MtZ2X/bdD3ztP/P9rX9r4Xmp512y/Ql9uKj/X92C8vO4oPO/4WFKE1JMpORzPwF04IeYXX11SQkzK7tsXF9C208+2z72dbm1t7qk4e76236w4dZTQn7PHjF/oqD3I9FoOZcrEqRo5C85ktC0XgY8MephV/ZL+o5uEwLHnxALeVL6SkF79mVdkQPyCAUhSRKxEB7wzA7Xv1EI2ro8+AlQRk5ImkQoVEYU7Gyon6IJ7W8ExYEQ9nqqhrD6nMWl3FkA03MaZmWW8+RDauW8VVG83bqpGPH9a3ijDkW+vL+ztaTrccF44Vjtu8HJfsoTlCUxHRGisibMKPtp4QtbE2+EvI8maSoN0lTlpu5TabThVgU0nbruhZbiy1LRDINn20fQbMCFUo1TeBjuJdDgOO1nI2JXhwkb1p5mEfkXnNvSDJyMsuvwyhCBwSltCkc0eFTc/9On5qZMMqMKa9uHIzIehQckOgeW2hfD6JwEN+LyGHemspqTD+7WZ6yHTXi4S+mBSGeT8PBMG+xAaqRhV8T/ufO1yRN7jWbLWog7mHTabu+d3eK3qW/n4TeLSkTsPti8TIRxk2XyckyCYXiX8JSkW5CfPe779/95v/bO5bltpFjzviKKShVJjcg+JZkbZiElmgv17KokFIcl6PijsgRhQgCEDwko1y+5Bv2kEOOOeaYY74mX5LumQEfEAGKFEXZWvS61iYw3dMz09OveeAr+zNnIxVvR4phTlAcaSj/+/nvQjLuMeRShgBlQ6z9c60d+u91FPvXZkeFA06myS/+Tk6Z6CfBSfDtjsx/11FscyOzPJbcYtpEVz4KYjyyAqEVUFbjOAr5Sy+fmsmU8yi/JWXyzni1POIKdSXPrRSkJAFLr2cF7qSAvQKLPCVfKxBarWdWwFooYF8Dk5mAxQTsWOQlxjL2xGO3k3IZwARWwF0eZTUmUwVs/Rwuj5QgYOn1/LxW7+A/6yj2j400PQklZelp+YRULT0VuxrB6qYyXHWtVFr8BdQHZLi6rcNWs9fqH3VOWr1Ydqtaq+9W53NVT/uMH8TJxL9kxKQ+83xyEUBo5DKTUQiZLBueacQRv24Mz/D3FDJuQfw0QBFZlrh9jjt3Q/4ChqZOnhgW1E0HUb4Cz5wkEUw9I/mMDoEtbiw/BDbvs8VxF30uJB4CW1zxnENgHCllRWL5CY0E1/cdZElw8ZeE16QhtrV6qbKZHPidA1UFvtqC3yjeLpVLL8uVcrVWmv+N4t1SKZH5Z36gj7dy6XOoHOuX8eET2djFd9fOU0HxIG5JFbSw4nkqSCCtUWNs8++bPB7Bx1VBO1pt+8lUEK7z3ksD7bzcSeT9uWugha2cq4EQ6xejgRY2NlEDxaP8JTUQVry0BkKklAMgyysMJLhWDRQj+LgaaFerVWtPpYHk7hNUQrVyuVwv12vV0m6CG1RJYj/9+sRv7s8qrZyrhBDrFxOJYWOXu3dmmUAsRQktrHieEkKkZ38dx+JWJolt+pcxnpnYpjc2SWzv5byni216xQlim7af6Kl7eH1im97KRLFNvdLuuYltamOTxPZeHt8CsU2tOElsn/2dt9jK9Durk8S2nrLfb9nLjzi9FFcliYUUP2Wly48k0ae//Kiyo72sTi3L3OmQ+RiVZTGmv4OQXX70GJcf8V5O2xC86PIjTmBNlx9JWt/65UeT8z/ir76HYgXR4xrrWHT/R2nm/D+e/6xVy/Xs/M8mID7+gcunuVdcYx38jpd6/T7jX9upVXD8Kzu1X5H6GnlIhGz854+/VAfy+OaD6lg0/8vV7dj83ylv17L5vwm4e6vEmTLvWgnSIHiThaKMT/0oqWeRoXxJ+eqsXQZx0IvoeelO+Ih1LJj/lUqd3/8ARh/Ufw3P/5Yr5ez+v43AFtm3nZAfp8GNuvXoHN3h4b6ypWyRQ2PALI8NSWANmdga1XQw+IzeaOMTChW9RHJYQJWv1Pz3QCG0A35QFOJXEngMSBgeuTCgDvZpwBwftzWhTjENag2YOHLrT+jrQOKDJGGf+yCshEJ5J8QvoEyVI9TnDCPgYslesXh7e6tTzqxuu6OiKQp6xcP2fuuo1yoAwxzl1DLxNKvL/hYYLjT1PCTUAX4GqAWJSW+J7RI6chm8823k99Y1MHTViGdf+LfUZUBliJGMcR74M50VcQdtni4A3UUtojZ7pN1Tyatmr93TgMb79skPndMTPH3ZbR6dtFs90umS/c7RQRsvooBfr0nz6AN52z460AiDruJnLh0X+QcmDexGvBthi/QYm2HgQm5s8+QN19AuaxTg549GNoRt/JCvwyDK93AwPWBvCFRM49oQgZ93t1G6okCFEJkS21P48toFhdnkGEQ+fg0/m8dtjXSha5nnR8WDG2Ngu5bAGQCCT80+sGL5Ov9/hI9zTjyPMCnru4GFdiZ6Arz77BM1FCX6l25Yhp9TUBBk1NuwPZ1ZN4ZrW/qI+Tn1TafzBgLH/cPO6UH/uNv5sbV/ouY1jmPaIqhORzrs7DdxSBArryggL2DxZHtz8OAP8ESgFS8ZNf1LNa9QL7QGBDwqwq0lti6Xxx2ELvMD1yKf1R940VCFcLHzVv2iSDqY4gBC0FFFECFGr9HcejaOWZ9ZIwjXZqgnlMm5YhT2ouGYrpveUmO6e3X4p6SD4TD1YRa5ucmAaERSg7YqxgXE4ngmEILwRoOo/T4alX5fxQrkYOtAOAeN0QgE2X5DLen8P1UjOIwNw/Jz8Q4/7nRP4P1uabeUz+czZ+IxQC9OjfkjeQEL7H+1LO7/mLL/lVolu/9zIxCpUbwgIPq3aY9AX4ziyt0PHb4LWTxtWqEGltOYoAWBMRxbBGkJ5poFZg1sMCXzbQaYMohHPNwszV/LBCMdXgnb4I0ZmLYLYysQncwXloBjSPU3QRxeNR1H4TbSF+4IWHUwaIErTCa/eAg4QM8G6JBme+KVgEW5MYbibiY0LcXIXCgT7fUxMhVvWkfNdv+01+r/qdU9af252VbPwEyoJ93TlqogB6iuIVzyxFFp5oKuZP3zwDChf9A2oD+jquorfOIJ+y0KE1kYWedOGfAsQvbCOUV/bcgc0w4xb/sCjDm3rUCIE4z3qyTpgeHsiw0ycZairmtb7/j7nnjdk2/JFnFCE1T4XnQ9U2NUECgFcPsK1C/4Nu/xsb2ZTwktp8L7RMgIxJ6BddX37T7KZ47/2kPRy5PC74T0fQQrpZEjCGHP0NpAG1sc1UM3EctjD1HyY69zRMRZdHmrC27bAZ4Z7xYC0W24x+0/1iTqhaGKRJVfoMGDYvlE8JLnGLJJWEQfBteOl5vQyJPfEPUvFtYg3d0W/wtPwUt3g082nUWPc+prCt7xUAglVibaoc5Uhg2Gvpqy+nMNNh/4PTFZx+Zfzjh92g2ATjiI0DzpQWBfNVsRnheJmnDRbgzKhQ5v0WGu7MRzexhCrwlnQuLxrssh8w4Ncf8ZFMBywsYblhP4ah6d189foMw18y/t2SIDk3peX7wQJVXe6qi5UIsbwnRCMRR+mJjhovEN6a4kTLJG0uxDhpOo8XdxdwteMJf6tit8OgJxwo19xUOIBoGGUN93hQMkmiKGU5BB55zdcM/XmuDlvvtOdpmkSKRAYP/MnyGcSF4WBrdMlo/QCQkNZg6jx8rE/xurXb0XjX1XPspJZFAvPnbB3fZqssQ1Gxq0D3aCNVQZQaFEF5E1FQtlTtzXBnpRhr38hivd//QITs7C+z9L8fxPtbyT5X82Aln+J8v/PCD/M7NTgxqO/I7axxnHWwMn80yJIoIo+eOEQ4o7GDKT8JSgF2fyb+tc9h3DPdd/p+P/ermSrf9uAuLj3+9j9rTfX2cuaJH9L1fi9r+2Xcm+/7ARyOx/Zv8fYP95KkmP0kNiqeapRTqDJSCu/0WmcL0LAYvy/3fX/2vl7Wz/30Yg0/+Z/n+o/k9cojk0r8UqDVA6YBcQC8LY4RZ/chFY/GwrT/aPmN+PtNDA8MPcwA4szMdjapVn+uFvmaTuMuhDdsPEYojEIojFW0jJyLhhKFCcgkxLb5EuEzfniNT/wA8AybRH0A05po90jTSP20DNNDWC29lwCQXe21eBg8lUWY1HGuSzeuGinOIS/TF1DU/ViPpXkDILn5zYV6GNTwbUokPK1/F9n95S9csk1RpR49ntiFHTvsUVH41cqD3bdUONtHEf+wufXFn27UxjQe4/S7QvuoorJjMqHJiM+l1kbq/tITMb6ohdg2sPQl8vXJjUu1RF1haX7RvqDAX5Jjq0gHsh1KblAYMezmCX8Ly+kJJzO/DvDgWweGckNJkUh7EMBoKoquLMxq96wGQQzPuX1I8W2eaMMScckRwntTm8v2SoGjiD1LvyxhI/1W9j3L1Z3LJO2kPc+H4RChRRineOWC7Dp0j6hUf4ekOs8opOTj0x4X6Ki/NPQuRBc8AMGE6zFCNS1YnIueMn6hl1zTBacsGaQdn44iRQvFdiZFqfKKgBRv6IfIIEvocufeHFu+I1F+Pfq/Nxo9w/yvQ8PNRnXPz1CQEYTjHG/BBP42O8H85gHh6LkeXMRDoAFKPLBr4Z4kaayQ4TkOMZqcz8yscAvUjdwSXMVHnm/zHqWLT/ezr/U6/U+P6Peub/bQT04oE9uGIuOmSPVcei/E+9Vo7v/93B/F82/o8Pmf+f+f8P8P9fdzvviBP6l7a1V9XL5YIHGIryvtN9e9DukiLFjVb7neMPJL7MTAC7e3pEHMOJPjxACu6dYhJ7NknB6cYSF6LczGZGrILXQR2/AN4IKYTy6szvpx9FtQ8C1xRU5JEIxN9/d0A+qugtE7WAH/ZT5foVL7SHO0QKBdzRSuSGVvjJY59f892rZ5nXkkEGGWSQQQYZZJBBBhlkkEEGGWSQQQYZPBX8H3kASf4A4AEA" - yield mock_agent_engine_create_docker_base64_encoded_tarball + _runtimes_utils, "_create_base64_encoded_tarball" + ) as mock_runtime_create_docker_base64_encoded_tarball: + mock_runtime_create_docker_base64_encoded_tarball.return_value = "H4sIAAAAAAAAA+xdzW8bSXanPYska2DXQRBgMbda7kE2QrW6+pMU1pvYsux11iNrJc1MDEMgWmSJ7HGzm+5uyvYMBtgsctzDHjxIECRAgOSYY475C5L/Yo4Bcshtj6kPypTF7mq+MtX6qoe2aDb796pe1atXVa+qXxlrjTMnk5LvuvyT0ulP/n/sYsfGju9aDr3v+9hvIPfss9ZoTLI8SBFqpEmSy56r+v2SkrFm9JL4MBycoR4sWP+25TiuadsNE2PLdnT910Gz+jeiIMu72SQ9Im+74zQZjXPjbTCKPj4NVsGe45TXP/3tff1j+hz2sEvr3/z4pKvpmtc/r3RR2908HJF1hH2v7fq241qG7fu27dw67zxqOjs63f4n436Qk25vSHovja+yJF5CGlXt3zL9Wfu3LNr+Xfa4bv810DfNuWrndqB5whC0DcelIzPst1DB0yk5CrMwiSnCMi3PxGYHW9h2TPp0nOThYdgLcvp7Rh94sX/MIg4GPB1295tvvz3vYri2NGv/gx45ozR4+59v9yfGf27DMm069Pd9i93H2Dexbv910KMgys6q3jVdfDrd/yd0HBjGNY//Tfu4/3d8z+Tjf6zbfy30zbd6eH+dadb+g14eHtFBHf+21DSq278/a/8OG//bvmfp9l8H9clhMImuomSaFqFZ+x+G/T6Ju4NelEz6UzPQncTUJqQZ6fZJ1kvDcZ6kXTr3C7q9gM7+pk9lRv9Aloa8/dvz/l/L9Vxbt/86aPfXT8OcoMMkHQU5sht/2rhxo/FXtDwajU+m/47pBv33g1Pfq+iThvHlX//o9h8aN2+njdvp7f9eYtY1adKkSZMmTZo0adKkSZOmC0xf3vjjn3z66Y2//WEeHEREeFDE35sbO5v39zbR3v0HTzeRuIfuTD0xQZ6naG/zb/bQ9s6Tz+7vPEe/2nzeQkdBNCHowdNnD+7+xc0/+slffnqjEcZ98iZ7FdF5fTeY5An/fuzQweKTzev/hOXmR/TfjR//V4NemjRp0qRJkyZNmjRp0qRJ0yWg39o3f3D377YGSTKISDAOM6OXjL5pHu9geJiMgpBtUm9++ESzhd4/sztM0nwr4Lvfm+wHOn0ekjif7mCfsWD3Db5RwhDcjjmN0+Qr0su3U3IYvjnmwh/8khzM8EXQftKbjGhip5IqepRlVuy4b2LDNEx2L6NAnu9xGo6C9C3PfpqHh0Ev3yGDMMvTtzOu45cDo0+Omnz3343b/9uglyZNmjRp0qRJkyZNmjRp0qTpKtCffXL3xofej5uz9z+mbwIt+E4HhADvf0zf/3Q839Pvf9RBkvc/bjY+/v2Pm/r9D02aNGnSpEmTJk2aNGnSdD3pAr7/cft/GvTSpEmTJk2aNGnSpEmTJk2aNF16+vMbn/zs014UkjjvBmlv+KbtdT1HvP/xfw16adKkSZMmTZo0adKkSZMmTZquCv34k5/98IQXYPb+R5QMsrM5BJqf8LPY+c+O77Dzny12JKg+/7kGOlX/7ARPw8QG9paoCoD6d03bZec/Wbau/1qovP6xY1i2YWOj08bY9gz6u2Ia0ve/sGNjzz51/rfrYH3+Vy3EKnzVxKvYQ9hZt+x1G7c6bRs93Hzw+WNEicmNTtCjIIxIH+UJSknQR483NhE/Cwod0vvr6MVmmsYJsvbRVoKySW/I76MkRf0wJb08Sd+uo5U1qnM57YJOHj67cmsuL47TopqzSF74UcRoRPKAnU6FMpIekXQd/XySRsmYxIikKc0CO3C4j5JJ/ouStPzytJ7FXNbDNBm9Z88PTy1h1Sln9TQJ+jQfG8loFMS0CNNkMqYltyKO3lppoZVeMhonMS2gbGW/mD+WVBGIP/0mjnMuTcktT2lnEsdhPEAvBGtjxtcQTPfR6zAfoiAd8NhWGc3G6ip7ilbc6vhtPkziddTcSyek2UKrq68mIclnNzaefbb9bGtza2/1ycPddY/+8IJmPSUs0yLF1T4ZcxkOXs1unvi5l/Zsq8dvZJM8jFb2m2VieujJ1qNnXKxut/voydPN7rOtp8/p/9EDMgiFoEIqA+0NwwyN06RHsgyNgrcoD14SqhdULYIIjcJ4kpPMuFWclG0upmYjMqLNRTQvqbK57SUz9KwZQ9qEovDAZi01pg04TOJxkkRoN2fRxWiRxOQ1+uXe3vYumj2B7uC766gfnYhbRvnahYl1On5lYsM8H2fra2tzHFHz8eYeovfXeH1n/ZdrvWFAsVG2lgbjsL82U8lVix9ozzO7Rru4JrJME1mO61sFGXO5OShViaKqZZC2BPI8maSoN0lTmhv0mMuBNli20cbTJ2ga1A2FtJG4rsXiuikk8iSm/RktQ1oxM8lF7U8T+Bju5ZBS/dsYJhk1wP0wG0e0nTziLxqv873PLw6SN608zCNyr7k3JBk5meXXYRShA0J7mVFyRPpGc/9On5r3MMqMKa9uHIzIehQckOgeC9q3HkThIL4XkcO8NZXVmH52szylRTJ9+ItpQYjn03AwzFus2zCy8GvC/9z5mqTJvWazRVvyPWw6bdf37k7Ru/T3k9C7JWUi6QA+tkyEFdJlcrJMQqH4l7BUXImZ+f7d775/95trc/2+pITgJkkG+f7db1ERLaBbhTjGsKaM/9N511Dl9e/Le+zfzlUbGDFjUfYbx00NxPFXhFib1xrxwfWfy3usPo2Aox6Eg19PCB1eH8+7noYxQXts+FpG8EQUIBJpOpKEKFkGNixfAQuHqGYSG230WfigrjzCQWXqt/SykKjfne0oyFmEH7Q7Jr3wMOzdlanfxajZqfq5FzqTP0f4UqqfZ6qUhQz1YBL3mS9um/t0kG1gu1TvpNZPlogCRE0aSkwCA5sXOJPYbhteufotP49wUG3qJ9wou3mSBgOyWA98gWtWkGvYJdbvYmQSW4azXPXDKuonAZWpnzwdhdxVqV+BJYQnolZqCqhK9bsYmZR2vsvPYwnIgqufBCIpCwlqzpe8kaSEqt9BGqQhyQDqJ0lEIV9q0lCAS3tew8bnnUlbXf1kUDikXP0koDL1k6ejkLsP1W9nw7Y20C+DbCiZfdRVako1i9i8ki2QnHMmHfnM15WonwwKh5SrnwRUpn7ydBRyJ5Z6US8KUZ+MSdwnca/E6k27NpUCr0saCqDK5xjYu8CZrLJ+l1X93p23r7Ty+o/lPfYPtRRqGUSyqwS+ZYExXHDvzuIM23XtgfA+3M6z/D0QO5tPN+/vbna3nu1t7p7a/2A7btsuzpXrSkbEj9h+qiFB1MaSLEeHkyhCKYlIkBEUJ/ReC43Ft6MwC/P1W+i9BKePtFtjWZ5iuxx7q0hrWIYk47htkjJjz6ojjGnaQe94odwwjFKGMnPw3bvvv/v9Jbz+XknYf6FY9DB5HUdJ0GdL4KAVDZ7wvyol/I8loKVaCMZQspNQjWH1xqUlWQi/5Z7cO3GWu6TWBHRV7J+jj64evFplOTKx2cEWth3ToEIZg69PGRLcsbDnd4pzb8kmYN8VLftr0MeCZJOOawWSWh8NujAgV4M06AOQZE+2Bl12kGzbowZdR5BsF5MGXUeQbD+MGuifi+aoFGXLVv+ulCegUliQJ+ADL7thGBJPQGXCRZ4ABrKWOXFnDCUvz6kxdOvyBLRbzgKOybPzBMThG+YNcLFpY7ZUYfll3gC76K0qkX1Zm750l4qUhUaIonzZLp0rZYQqhS0wQgVbDMuo3AhVJlxkhBhIarnOvZCXpLZyKUvVVuqfumpqKxV2MbUtW5qWq63U+VWmts4yuzrGcKluecbQq6vv7LScs11nk/adQg2mL73bqxPWk0ZhPHmzKiLxim6V9qjYMzF2SrpVx7V93/KLtkZ1Wh24o5iBwK4ZBgIP5ylIPvorA4F9vgwE9kgwEHi2QUGOikwO2D/KQCoyOSoyyZZZy0EqMknffS0Fqcgk3QJYClJpGp6KTPIxSRlIpWn4KjL5Ks29rdI02ioytVWaRltFpo6KTB2VptFRkakDbRqu2TJNqEwcBJWJg6AycZCKTBja3DlIRSYMbe4cpCITuHfnIBWZwO5bDlKRyVaRyVaRyVaRyYGaZQ5SkcmBmmUOUpHJVZHJhZplDlKRSfoucxlI+npTKUhFJvCQgINUZAIPCThIRSZfxSzLQUWuAo6SlsQVchVUCzvvKoC8xVjqKqhOuMBVwEGS6DzgmT1j2F5mnDPOENfkKqCJWc65udlFqDruDzA9Otd3HQeXudktOrHxOwXuACZCR8HsWCpDyPpAtqlgFRVBCsPiOkEKo4w6QQodUZ0ghSHkhQcpDE3qBCmMMi48SGEAXidIoQOoE6TUAVxsEFbpamoEqXQ1NYJUupoaQSpdTY0glV7jooNUupoaQSq9Ro0glV7jDEAl7gBb7ge9Wu6ACmHh7oDjXQTSXXfVCRe7A2xrybN328LLZlhXlHIX08nque0cmLoDFt955xS8hDcVQdauL91VIqW0lyk0RAwlNf5XyRAxYaVdybwhekgOg0mUo4zkKDl8H3BDGKQMfUDlhqgy4SJDVAk690JeltrC3ekMJe11r5raAt3pi4bFkveflQmXqe1S3emYOzGXzLA2d7rVwp1zc6ezE2MW7j0t2yoIvMMFkHsDLt2lImWhGbJaltyBfZXMUKWwVWbodHi0RXvPyoSLzFAl6NwLeTlqWyFlqdpKB8RXTW2lwkLVdraDXd57ViZcprbL3LfOGS5z3zpnWNe+ddduYa+mM7IKZp8nD1Kb27DOosJY2MXYtktmo66Nba9EKh/sYKegNtjtrUGXBwReNNCgKw4Cr9NokAbNg8DraRqkQUsCgdc9NUiD5kHg9WkNWhRU6CFgqGvjX68Udt5DAAm8Xu4hqEy4yEPAQMs89JozXKrDnjLsmHV5CJwWPr/4sCdOYp/zD2DTwR52bRPjEv+A6RTLU7G77dJdKlIWWiWGui7xNqqF5VZpdug5LGi1xCpVJlxklZyWBd9WyEDgDo2BwDtuKcgCO2UYCDzoZSAVmeDvnVKQrSIT/L1TBgIP2xgIvG+PgcCeUQoCh6LgIJXKBYei4CAVhYW/rMpAKvXkqsgEf1mVgqQHjZWCVBQW/rIqA6k0Qk9FYStiapWAVBQWHPSCg1QKAr6mQUFwXzkDqWgEfObJQCoF0VYpCHB4DQ5SkUm+c6oYpPCWi6PypoGjshndUdm47fB9vgogFZlUhgS2Su9uq/Tutq2UPZUiV+ndbXhUCQZSkQkeVYKCVDpqGxxoioMUmrsNDjTFQPCoEgykIpP0WDI1UMnk0b42wRqrhYVMHj90ask3vVQmXDx5tCuGNpfuUpGyVG2vTeCQSmHn1RYQGbdCbcE7nSmoYiB66S4VKcvUVj7evmJqKxd2IbUtPbRVqrbyhIvU1m858JkaBcHnGuzsNvDgo9Nqw0csFATfYtNpdeCzGhYAViF78HCfDCTvSudAzrpp8uNA4SDg5JiDMDB4iwABZzUCBBzMCxDQ2SRAvkpBACf8HGQBI4kIEHDKKkBAP6wAAcPeCRDQFyZAwEVtDoKGvuGgDvC0JgECenwFCOiqEyBgLFcBAnpzKAiDA4sKkAVuTxQEjGEtQPDmTkFAL7YAwY0l5uHi4KAOuOViFQtLQcAAGgIE9FAJENAFJEDAxRABAvojBAjoZBcg4D43AQKOjTjIAoZdEiCgo1OAgCHrBQgYAFaAZOdml4KA0UgFCN5RUxDcwla+aVwMgkZQEiAVCwv1+AoQcO1TgIBBbRcBzU8eOcqRLwJcmcnjIsLOTR4hsVLLJo+LJDw3eeSgjqfQhjrABUUBAi6+CRC8tVrg6OwCBFzLECC4gaQgeJFbPAYtHATvyCw64FCQCcNtHQXBbR0FwTsyBlKoXAzcdCBAwAMmBAjekbE33BWahgX04AgQVsiewiiFguBDUAu8H0eAgB4cAVJpuRYwKLkAAc/04CBosFUBAi7QCxB86m2Bl2MFCO72oSD4lI6COgrWyAG6HAUIPuqnIKCfUoBUSs+Bj8UpCO72scBOaA5y4b4iCoLPoikIuF9UgFQ6NRe441GAgDsVBAg+9aYguOeaglQ6NRcY6Z2DPGAAWQFSqVwPuLVGgFTqyVMZfHgqjdCHe6UoCLhpexFQ4eSRouTd+1WaPFYKC548vl+FlOzzWCThoskjA13xUFKLSFmmtvIVlyumtvKtbHNqC4nDKFdbecIlaluxtnDpLhUpS9X2OllbubBzartoHMZqaytPuExtz8XaXva8lCm6fJniiim6XNgKRZdH+pMrujzhEkWvcLNcuktFylK1vR4v8i8iLFBtT+zIq7TP8oSL1bZdcbxUncrVrnCSXKC8lCh6W+5Zu1qKXiHsnKJDAlZIFb0i4TJFl0/+L92lImWZ2mK7PPjG5+N+IIJkJDmrnIAHx+DxN6hFKuVXmouyLEi8QNskZfrBMjFOshyN06RHsox9z3IyzsrzIQk7uvmG9CZcsOkUbh29WFnLqdGdj7KxdhDGa0J/V1poZRaPg31jWVqdZmllfz4jmK3KnQiXMlcgy0HYJ7canRa2T7JeGo7zJO1GYZZTUb9ZmcThEUkz8jAZBWG8so5WhNjBOMxYbBIm2/Ezu8MkzbeCEWGPsR+CST6kBTBVhxkLdt/gBXUiygkvpTT5ivTy7ZQchm+OufAHvyQHM3wRtJ/0JiOa2Kmkih5lmaUPsZ+xYRomu0ctUM7zPU7DUZC+5dlPqS4HvXyHDGhxpG9nXMcvB0afHK18W1yT0leLbvG2QlCfKsdPbxUoOmfglevkxjDJSIz6YTaOgrfoEVX6IF+PKbtSXqVV/nDK5FAwQU3GpnmroenikLHGogQdhoO1KBlka6yODRMb2FvDjmE5Bm4bHddxPM+gvyumYVLyHId/Ujr9afoWbmAX25bjuKZtN0zsOthvIHOpkpbQhDbOFKEGU17Zc1W/X1Kab9TtVseVdFpPk6BP+u995o/TZDJm/VZh31RkwRh/iQEC8Z/2fV1Z38fSk4QK25nEMeuCX4gEjBl342S3uo9eh/kQBemAdwQZzVJBYpZPe3x/OdaV82qftXWtav+mbdjYscxOXe0fO6z906GEbv910Pyg1eYHNZ5Z++f8a2z/PL262j8dEZmOpPCAoyvG65zbv9U2PMczffPs2j+26f9n/T9v/w69pdt/DTTf7bRbniPpwhZpn0yfitoi5e07klMgF+bN5lQkL01C4kuYb+6MoUG5FbTw3c2NvSfPtta2d55tb+7sPaeN7L156NIsBgOSrtHGHBxEpDvhM69ub0h6L5st9MX9p59vUkCeTkizLKOSVVgxkeszT8OYpPlb9GKxpPeNglkfT0wyqCu2S32x3FzKbnHTNOXULGblSg422334KyTEQ1y8jFYQQVPB+4aeTH4kVdr/jmF3LLN9huM/bPn2Cfvvcvtv2dr+10HzDbLTsjsSA/pR9r/Doy6cqf2nSbQlXcwy7P8kJ2uDHumOSB5QyxR0UxL0u3k4Iskk72akd8L8m4W2n2VS4sUrtv2VyRabfpHW0kw/Y2cvx/QzVpKRuTb9NVCl/882TB9b9MtZ2X/bdD3ztP/P9rX9r4Xmp512y/Ql9uKj/X92C8vO4oPO/4WFKE1JMpORzPwF04IeYXX11SQkzK7tsXF9C208+2z72dbm1t7qk4e76236w4dZTQn7PHjF/oqD3I9FoOZcrEqRo5C85ktC0XgY8MephV/ZL+o5uEwLHnxALeVL6SkF79mVdkQPyCAUhSRKxEB7wzA7Xv1EI2ro8+AlQRk5ImkQoVEYU7Gyon6IJ7W8ExYEQ9nqqhrD6nMWl3FkA03MaZmWW8+RDauW8VVG83bqpGPH9a3ijDkW+vL+ztaTrccF44Vjtu8HJfsoTlCUxHRGisibMKPtp4QtbE2+EvI8maSoN0lTlpu5TabThVgU0nbruhZbiy1LRDINn20fQbMCFUo1TeBjuJdDgOO1nI2JXhwkb1p5mEfkXnNvSDJyMsuvwyhCBwSltCkc0eFTc/9On5qZMMqMKa9uHIzIehQckOgeW2hfD6JwEN+LyGHemspqTD+7WZ6yHTXi4S+mBSGeT8PBMG+xAaqRhV8T/ufO1yRN7jWbLWog7mHTabu+d3eK3qW/n4TeLSkTsPti8TIRxk2XyckyCYXiX8JSkW5CfPe779/95v/bO5bltpFjzviKKShVJjcg+JZkbZiElmgv17KokFIcl6PijsgRhQgCEDwko1y+5Bv2kEOOOeaYY74mX5LumQEfEAGKFEXZWvS61iYw3dMz09OveeAr+zNnIxVvR4phTlAcaSj/+/nvQjLuMeRShgBlQ6z9c60d+u91FPvXZkeFA06myS/+Tk6Z6CfBSfDtjsx/11FscyOzPJbcYtpEVz4KYjyyAqEVUFbjOAr5Sy+fmsmU8yi/JWXyzni1POIKdSXPrRSkJAFLr2cF7qSAvQKLPCVfKxBarWdWwFooYF8Dk5mAxQTsWOQlxjL2xGO3k3IZwARWwF0eZTUmUwVs/Rwuj5QgYOn1/LxW7+A/6yj2j400PQklZelp+YRULT0VuxrB6qYyXHWtVFr8BdQHZLi6rcNWs9fqH3VOWr1Ydqtaq+9W53NVT/uMH8TJxL9kxKQ+83xyEUBo5DKTUQiZLBueacQRv24Mz/D3FDJuQfw0QBFZlrh9jjt3Q/4ChqZOnhgW1E0HUb4Cz5wkEUw9I/mMDoEtbiw/BDbvs8VxF30uJB4CW1zxnENgHCllRWL5CY0E1/cdZElw8ZeE16QhtrV6qbKZHPidA1UFvtqC3yjeLpVLL8uVcrVWmv+N4t1SKZH5Z36gj7dy6XOoHOuX8eET2djFd9fOU0HxIG5JFbSw4nkqSCCtUWNs8++bPB7Bx1VBO1pt+8lUEK7z3ksD7bzcSeT9uWugha2cq4EQ6xejgRY2NlEDxaP8JTUQVry0BkKklAMgyysMJLhWDRQj+LgaaFerVWtPpYHk7hNUQrVyuVwv12vV0m6CG1RJYj/9+sRv7s8qrZyrhBDrFxOJYWOXu3dmmUAsRQktrHieEkKkZ38dx+JWJolt+pcxnpnYpjc2SWzv5byni216xQlim7af6Kl7eH1im97KRLFNvdLuuYltamOTxPZeHt8CsU2tOElsn/2dt9jK9Durk8S2nrLfb9nLjzi9FFcliYUUP2Wly48k0ae//Kiyo72sTi3L3OmQ+RiVZTGmv4OQXX70GJcf8V5O2xC86PIjTmBNlx9JWt/65UeT8z/ir76HYgXR4xrrWHT/R2nm/D+e/6xVy/Xs/M8mID7+gcunuVdcYx38jpd6/T7jX9upVXD8Kzu1X5H6GnlIhGz854+/VAfy+OaD6lg0/8vV7dj83ylv17L5vwm4e6vEmTLvWgnSIHiThaKMT/0oqWeRoXxJ+eqsXQZx0IvoeelO+Ih1LJj/lUqd3/8ARh/Ufw3P/5Yr5ez+v43AFtm3nZAfp8GNuvXoHN3h4b6ypWyRQ2PALI8NSWANmdga1XQw+IzeaOMTChW9RHJYQJWv1Pz3QCG0A35QFOJXEngMSBgeuTCgDvZpwBwftzWhTjENag2YOHLrT+jrQOKDJGGf+yCshEJ5J8QvoEyVI9TnDCPgYslesXh7e6tTzqxuu6OiKQp6xcP2fuuo1yoAwxzl1DLxNKvL/hYYLjT1PCTUAX4GqAWJSW+J7RI6chm8823k99Y1MHTViGdf+LfUZUBliJGMcR74M50VcQdtni4A3UUtojZ7pN1Tyatmr93TgMb79skPndMTPH3ZbR6dtFs90umS/c7RQRsvooBfr0nz6AN52z460AiDruJnLh0X+QcmDexGvBthi/QYm2HgQm5s8+QN19AuaxTg549GNoRt/JCvwyDK93AwPWBvCFRM49oQgZ93t1G6okCFEJkS21P48toFhdnkGEQ+fg0/m8dtjXSha5nnR8WDG2Ngu5bAGQCCT80+sGL5Ov9/hI9zTjyPMCnru4GFdiZ6Arz77BM1FCX6l25Yhp9TUBBk1NuwPZ1ZN4ZrW/qI+Tn1TafzBgLH/cPO6UH/uNv5sbV/ouY1jmPaIqhORzrs7DdxSBArryggL2DxZHtz8OAP8ESgFS8ZNf1LNa9QL7QGBDwqwq0lti6Xxx2ELvMD1yKf1R940VCFcLHzVv2iSDqY4gBC0FFFECFGr9HcejaOWZ9ZIwjXZqgnlMm5YhT2ouGYrpveUmO6e3X4p6SD4TD1YRa5ucmAaERSg7YqxgXE4ngmEILwRoOo/T4alX5fxQrkYOtAOAeN0QgE2X5DLen8P1UjOIwNw/Jz8Q4/7nRP4P1uabeUz+czZ+IxQC9OjfkjeQEL7H+1LO7/mLL/lVolu/9zIxCpUbwgIPq3aY9AX4ziyt0PHb4LWTxtWqEGltOYoAWBMRxbBGkJ5poFZg1sMCXzbQaYMohHPNwszV/LBCMdXgnb4I0ZmLYLYysQncwXloBjSPU3QRxeNR1H4TbSF+4IWHUwaIErTCa/eAg4QM8G6JBme+KVgEW5MYbibiY0LcXIXCgT7fUxMhVvWkfNdv+01+r/qdU9af252VbPwEyoJ93TlqogB6iuIVzyxFFp5oKuZP3zwDChf9A2oD+jquorfOIJ+y0KE1kYWedOGfAsQvbCOUV/bcgc0w4xb/sCjDm3rUCIE4z3qyTpgeHsiw0ycZairmtb7/j7nnjdk2/JFnFCE1T4XnQ9U2NUECgFcPsK1C/4Nu/xsb2ZTwktp8L7RMgIxJ6BddX37T7KZ47/2kPRy5PC74T0fQQrpZEjCGHP0NpAG1sc1UM3EctjD1HyY69zRMRZdHmrC27bAZ4Z7xYC0W24x+0/1iTqhaGKRJVfoMGDYvlE8JLnGLJJWEQfBteOl5vQyJPfEPUvFtYg3d0W/wtPwUt3g082nUWPc+prCt7xUAglVibaoc5Uhg2Gvpqy+nMNNh/4PTFZx+Zfzjh92g2ATjiI0DzpQWBfNVsRnheJmnDRbgzKhQ5v0WGu7MRzexhCrwlnQuLxrssh8w4Ncf8ZFMBywsYblhP4ah6d189foMw18y/t2SIDk3peX7wQJVXe6qi5UIsbwnRCMRR+mJjhovEN6a4kTLJG0uxDhpOo8XdxdwteMJf6tit8OgJxwo19xUOIBoGGUN93hQMkmiKGU5BB55zdcM/XmuDlvvtOdpmkSKRAYP/MnyGcSF4WBrdMlo/QCQkNZg6jx8rE/xurXb0XjX1XPspJZFAvPnbB3fZqssQ1Gxq0D3aCNVQZQaFEF5E1FQtlTtzXBnpRhr38hivd//QITs7C+z9L8fxPtbyT5X82Aln+J8v/PCD/M7NTgxqO/I7axxnHWwMn80yJIoIo+eOEQ4o7GDKT8JSgF2fyb+tc9h3DPdd/p+P/ermSrf9uAuLj3+9j9rTfX2cuaJH9L1fi9r+2Xcm+/7ARyOx/Zv8fYP95KkmP0kNiqeapRTqDJSCu/0WmcL0LAYvy/3fX/2vl7Wz/30Yg0/+Z/n+o/k9cojk0r8UqDVA6YBcQC8LY4RZ/chFY/GwrT/aPmN+PtNDA8MPcwA4szMdjapVn+uFvmaTuMuhDdsPEYojEIojFW0jJyLhhKFCcgkxLb5EuEzfniNT/wA8AybRH0A05po90jTSP20DNNDWC29lwCQXe21eBg8lUWY1HGuSzeuGinOIS/TF1DU/ViPpXkDILn5zYV6GNTwbUokPK1/F9n95S9csk1RpR49ntiFHTvsUVH41cqD3bdUONtHEf+wufXFn27UxjQe4/S7QvuoorJjMqHJiM+l1kbq/tITMb6ohdg2sPQl8vXJjUu1RF1haX7RvqDAX5Jjq0gHsh1KblAYMezmCX8Ly+kJJzO/DvDgWweGckNJkUh7EMBoKoquLMxq96wGQQzPuX1I8W2eaMMScckRwntTm8v2SoGjiD1LvyxhI/1W9j3L1Z3LJO2kPc+H4RChRRineOWC7Dp0j6hUf4ekOs8opOTj0x4X6Ki/NPQuRBc8AMGE6zFCNS1YnIueMn6hl1zTBacsGaQdn44iRQvFdiZFqfKKgBRv6IfIIEvocufeHFu+I1F+Pfq/Nxo9w/yvQ8PNRnXPz1CQEYTjHG/BBP42O8H85gHh6LkeXMRDoAFKPLBr4Z4kaayQ4TkOMZqcz8yscAvUjdwSXMVHnm/zHqWLT/ezr/U6/U+P6Peub/bQT04oE9uGIuOmSPVcei/E+9Vo7v/93B/F82/o8Pmf+f+f8P8P9fdzvviBP6l7a1V9XL5YIHGIryvtN9e9DukiLFjVb7neMPJL7MTAC7e3pEHMOJPjxACu6dYhJ7NknB6cYSF6LczGZGrILXQR2/AN4IKYTy6szvpx9FtQ8C1xRU5JEIxN9/d0A+qugtE7WAH/ZT5foVL7SHO0QKBdzRSuSGVvjJY59f892rZ5nXkkEGGWSQQQYZZJBBBhlkkEEGGWSQQQYZPBX8H3kASf4A4AEA" + yield mock_runtime_create_docker_base64_encoded_tarball def _get_replay_id(use_vertex: bool, replays_prefix: str) -> str: diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py index e4b33ad2d5..6f7227ec62 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py @@ -19,14 +19,14 @@ def test_delete_memory(client): - ae_memory_operation = client.agent_engines.memories.delete( + ae_memory_operation = client.runtimes.memories.delete( name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/5605466683931099136", ) - assert isinstance(ae_memory_operation, types.DeleteAgentEngineMemoryOperation) + assert isinstance(ae_memory_operation, types.DeleteRuntimeMemoryOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.delete", + test_method="runtimes.memories.delete", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_get.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_get.py index bca6fca8c9..707479f147 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_get.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_get.py @@ -20,7 +20,7 @@ def test_get_memory(client): memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688" - ae_memory = client.agent_engines.memories.get(name=memory_name) + ae_memory = client.runtimes.memories.get(name=memory_name) assert isinstance(ae_memory, types.Memory) assert ae_memory.name == memory_name @@ -28,5 +28,5 @@ def test_get_memory(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.get", + test_method="runtimes.memories.get", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py index 5ca959abe2..9a8ecfb8c3 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py @@ -19,16 +19,16 @@ def test_private_create_memory(client): - ae_memory_operation = client.agent_engines.memories._create( + ae_memory_operation = client.runtimes.memories._create( name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584", fact="memory_fact", scope={"user_id": "123"}, ) - assert isinstance(ae_memory_operation, types.AgentEngineMemoryOperation) + assert isinstance(ae_memory_operation, types.RuntimeMemoryOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._create", + test_method="runtimes.memories._create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_generate.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_generate.py index 7b294a01cf..f9c417ec1d 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_generate.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_generate.py @@ -19,17 +19,17 @@ def test_private_generate_memory(client): - ae_memory_operation = client.agent_engines.memories._generate( + ae_memory_operation = client.runtimes.memories._generate( name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584", vertex_session_source=types.GenerateMemoriesRequestVertexSessionSource( session="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/sessions/6922431337672474624" ), ) - assert isinstance(ae_memory_operation, types.AgentEngineGenerateMemoriesOperation) + assert isinstance(ae_memory_operation, types.RuntimeGenerateMemoriesOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._generate", + test_method="runtimes.memories._generate", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_generate_memories_operation.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_generate_memories_operation.py index bdcd2f14a4..d9feda4c64 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_generate_memories_operation.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_generate_memories_operation.py @@ -19,14 +19,14 @@ def test_private_get_generate_memories_operation(client): - memory_operation = client.agent_engines.memories._get_generate_memories_operation( + memory_operation = client.runtimes.memories._get_generate_memories_operation( operation_name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/operations/5669315676343369728" ) - assert isinstance(memory_operation, types.AgentEngineGenerateMemoriesOperation) + assert isinstance(memory_operation, types.RuntimeGenerateMemoriesOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._get_generate_memories_operation", + test_method="runtimes.memories._get_generate_memories_operation", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py index 23c92f7821..0c9c74a684 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py @@ -19,14 +19,14 @@ def test_private_get_memory_operation(client): - memory_operation = client.agent_engines.memories._get_memory_operation( + memory_operation = client.runtimes.memories._get_memory_operation( operation_name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688/operations/1044963283964002304" ) - assert isinstance(memory_operation, types.AgentEngineMemoryOperation) + assert isinstance(memory_operation, types.RuntimeMemoryOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._get_memory_operation", + test_method="runtimes.memories._get_memory_operation", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py index 662ad020fb..24f608dc6c 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py @@ -20,7 +20,7 @@ def test_private_list_memory(client): ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584" - memory_list = client.agent_engines.memories._list(name=ae_name) + memory_list = client.runtimes.memories._list(name=ae_name) assert isinstance(memory_list, types.ListReasoningEnginesMemoriesResponse) assert isinstance(memory_list.memories[0], types.Memory) @@ -28,5 +28,5 @@ def test_private_list_memory(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._list", + test_method="runtimes.memories._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py index d77cb3a795..025bcfb1a9 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py @@ -19,17 +19,17 @@ def test_private_purge(client): - ae_memory_purge_operation = client.agent_engines.memories._purge( + ae_memory_purge_operation = client.runtimes.memories._purge( name="projects/964831358985/locations/us-central1/reasoningEngines/6086402690647064576", filter="scope.user_id=123", ) assert isinstance( - ae_memory_purge_operation, types.AgentEnginePurgeMemoriesOperation + ae_memory_purge_operation, types.RuntimePurgeMemoriesOperation ) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._purge", + test_method="runtimes.memories._purge", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py index 4098b11d9b..754a3430ee 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py @@ -20,7 +20,7 @@ def test_private_retrieve(client): ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584" - retrieved_memories = client.agent_engines.memories._retrieve( + retrieved_memories = client.runtimes.memories._retrieve( name=ae_name, scope={"user_id": "123"}, ) @@ -34,5 +34,5 @@ def test_private_retrieve(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._retrieve", + test_method="runtimes.memories._retrieve", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_rollback.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_rollback.py index 85329da47e..898ec79aee 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_rollback.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_rollback.py @@ -19,15 +19,15 @@ def test_private_rollback(client): - rollback_operation = client.agent_engines.memories._rollback( + rollback_operation = client.runtimes.memories._rollback( name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688", target_revision_id="3001207491565453312", ) - assert isinstance(rollback_operation, types.AgentEngineRollbackMemoryOperation) + assert isinstance(rollback_operation, types.RuntimeRollbackMemoryOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._rollback", + test_method="runtimes.memories._rollback", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py index c6dbc1f095..40ef680b5b 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py @@ -20,16 +20,16 @@ def test_private_update_memory(client): memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688" - memory_update_operation = client.agent_engines.memories._update( + memory_update_operation = client.runtimes.memories._update( name=memory_name, fact="memory_fact_updated", scope={"user_id": "123"}, ) - assert isinstance(memory_update_operation, types.AgentEngineMemoryOperation) + assert isinstance(memory_update_operation, types.RuntimeMemoryOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._update", + test_method="runtimes.memories._update", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py b/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py index f58e365fe7..54ce8d2bb8 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py @@ -20,7 +20,7 @@ def test_get_memory_revisions(client): memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688/revisions/516064922187071488" - memory_revision = client.agent_engines.memories.revisions.get(name=memory_name) + memory_revision = client.runtimes.memories.revisions.get(name=memory_name) assert isinstance(memory_revision, types.MemoryRevision) assert memory_revision.name == memory_name @@ -28,5 +28,5 @@ def test_get_memory_revisions(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.revisions.get", + test_method="runtimes.memories.revisions.get", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py b/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py index c827a57a06..76eb966f38 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py @@ -20,13 +20,13 @@ def test_private_list_memory_revisions(client): ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688" - memory_list = client.agent_engines.memories.revisions._list(name=ae_name) - assert isinstance(memory_list, types.ListAgentEngineMemoryRevisionsResponse) + memory_list = client.runtimes.memories.revisions._list(name=ae_name) + assert isinstance(memory_list, types.ListRuntimeMemoryRevisionsResponse) assert isinstance(memory_list.memory_revisions[0], types.MemoryRevision) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.revisions._list", + test_method="runtimes.memories.revisions._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_create.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_create.py index fc776d9874..aa46a18a0c 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_create.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_create.py @@ -21,7 +21,7 @@ def test_create_sandbox_snapshot(client): - snapshot = client.agent_engines.sandboxes.snapshots._create( + snapshot = client.sandboxes.snapshots._create( source_sandbox_environment_name="projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironments/525190525100228608", config={ "display_name": "test_snapshot", @@ -30,11 +30,11 @@ def test_create_sandbox_snapshot(client): }, ) - assert isinstance(snapshot, types.AgentEngineSandboxSnapshotOperation) + assert isinstance(snapshot, types.RuntimeSandboxSnapshotOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.snapshots._create", + test_method="sandboxes.snapshots._create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_delete.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_delete.py index 730f9b262a..e6ac2d5d39 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_delete.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_delete.py @@ -21,7 +21,7 @@ def test_delete_sandbox_snapshot(client): - result = client.agent_engines.sandboxes.snapshots._delete( + result = client.sandboxes.snapshots._delete( name="projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironmentSnapshots/421086565159141376", ) @@ -31,5 +31,5 @@ def test_delete_sandbox_snapshot(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.snapshots._delete", + test_method="sandboxes.snapshots._delete", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_get.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_get.py index b9d5075c8b..7f6c35c68c 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_get.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_get.py @@ -22,7 +22,7 @@ def test_get_sandbox_snapshot(client): snapshot_name = "projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironmentSnapshots/2433069698686910464" - fetched_snapshot = client.agent_engines.sandboxes.snapshots._get( + fetched_snapshot = client.sandboxes.snapshots._get( name=snapshot_name, ) @@ -33,5 +33,5 @@ def test_get_sandbox_snapshot(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.snapshots._get", + test_method="sandboxes.snapshots._get", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_list.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_list.py index a22bb1c244..1f2dfc84c3 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_list.py @@ -21,7 +21,7 @@ def test_list_sandbox_snapshots(client): - snapshots_list_operation = client.agent_engines.sandboxes.snapshots._list( + snapshots_list_operation = client.sandboxes.snapshots._list( name="projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472", ) @@ -37,5 +37,5 @@ def test_list_sandbox_snapshots(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.snapshots._list", + test_method="sandboxes.snapshots._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_byoc_create.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_byoc_create.py index 30ec57efb4..8a4cc64d27 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_byoc_create.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_byoc_create.py @@ -48,7 +48,7 @@ def test_sandbox_templates_byoc_create(client): "internet_access": True, }, } - sandbox_template_operation = client.agent_engines.sandboxes.templates.create( + sandbox_template_operation = client.sandboxes.templates.create( name=( "projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472" ), @@ -75,5 +75,5 @@ def test_sandbox_templates_byoc_create(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.templates.create", + test_method="sandboxes.templates.create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_default_create.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_default_create.py index 60087195dc..3e3e28fff0 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_default_create.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_default_create.py @@ -27,7 +27,7 @@ def test_sandbox_templates_default_create(client): "internet_access": True, }, } - sandbox_template_operation = client.agent_engines.sandboxes.templates._create( + sandbox_template_operation = client.sandboxes.templates._create( name=( "projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472" ), @@ -53,5 +53,5 @@ def test_sandbox_templates_default_create(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.templates._create", + test_method="sandboxes.templates._create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_delete.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_delete.py index b55d64b0eb..ca7e607606 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_delete.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_delete.py @@ -19,7 +19,7 @@ def test_sandbox_templates_delete(client): - sandbox_template_delete_operation = client.agent_engines.sandboxes.templates.delete( + sandbox_template_delete_operation = client.sandboxes.templates.delete( name=( "projects/254005681254/locations/us-central1/reasoningEngines/208148546254274560/sandboxEnvironmentTemplates/4632233691727265792" ), @@ -33,5 +33,5 @@ def test_sandbox_templates_delete(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.templates.delete", + test_method="sandboxes.templates.delete", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_get.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_get.py index 9cda01044f..602701369e 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_get.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_get.py @@ -21,7 +21,7 @@ def test_sandbox_templates_get(client): sandbox_template_name = "projects/254005681254/locations/us-central1/reasoningEngines/208148546254274560/sandboxEnvironmentTemplates/4632233691727265792" - sandbox_template = client.agent_engines.sandboxes.templates.get( + sandbox_template = client.sandboxes.templates.get( name=sandbox_template_name ) assert isinstance(sandbox_template, types.SandboxEnvironmentTemplate) @@ -31,5 +31,5 @@ def test_sandbox_templates_get(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.templates.get", + test_method="sandboxes.templates.get", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_get_sandbox_template_operation.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_get_sandbox_template_operation.py index 8022ca6a51..f6ab8c3086 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_get_sandbox_template_operation.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_get_sandbox_template_operation.py @@ -23,7 +23,7 @@ def test_get_sandbox_template_operation(client): "projects/254005681254/locations/us-central1/operations/7252775414349692928" ) - sandbox_template_operation = client.agent_engines.sandboxes.templates.get_sandbox_environment_template_operation( + sandbox_template_operation = client.sandboxes.templates.get_sandbox_environment_template_operation( operation_name=operation_name ) assert isinstance( @@ -35,5 +35,5 @@ def test_get_sandbox_template_operation(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.templates.get_sandbox_environment_template_operation", + test_method="sandboxes.templates.get_sandbox_environment_template_operation", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_list.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_list.py index da9268c1fd..870c908afb 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_list.py @@ -19,7 +19,7 @@ def test_sandbox_templates_list(client): - sandbox_templates_list_operation = client.agent_engines.sandboxes.templates._list( + sandbox_templates_list_operation = client.sandboxes.templates._list( name=( "projects/254005681254/locations/us-central1/reasoningEngines/208148546254274560" ), @@ -36,5 +36,5 @@ def test_sandbox_templates_list(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.templates._list", + test_method="sandboxes.templates._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_create.py b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_create.py index 4a861c99ec..4c63866f61 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_create.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_create.py @@ -22,17 +22,17 @@ def test_private_create(client): spec = { "code_execution_environment": {"machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB"} } - agent_engine_sandbox_operation = client.agent_engines.sandboxes._create( + runtime_sandbox_operation = client.sandboxes._create( name=( "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584" ), spec=spec, ) - assert isinstance(agent_engine_sandbox_operation, types.AgentEngineSandboxOperation) + assert isinstance(runtime_sandbox_operation, types.RuntimeSandboxOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes._create", + test_method="sandboxes._create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_delete.py b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_delete.py index 12b79a9269..2978132f10 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_delete.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_delete.py @@ -19,18 +19,18 @@ def test_private_delete(client): - agent_engine_sandbox_delete_operation = client.agent_engines.sandboxes._delete( + runtime_sandbox_delete_operation = client.sandboxes._delete( name=( "reasoningEngines/2886612747586371584/sandboxEnvironments/6068475153556176896" ), ) assert isinstance( - agent_engine_sandbox_delete_operation, types.DeleteAgentEngineSandboxOperation + runtime_sandbox_delete_operation, types.DeleteRuntimeSandboxOperation ) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes._delete", + test_method="sandboxes._delete", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_execute_code.py b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_execute_code.py index bbc69a9ff2..094a3abbd4 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_execute_code.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_execute_code.py @@ -41,7 +41,7 @@ def test_private_execute_code(client): ), ] - execute_code_response = client.agent_engines.sandboxes._execute_code( + execute_code_response = client.sandboxes._execute_code( name=( "reasoningEngines/2886612747586371584/sandboxEnvironments/6068475153556176896" ), @@ -58,5 +58,5 @@ def test_private_execute_code(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes._execute_code", + test_method="sandboxes._execute_code", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_get.py b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_get.py index 7c6cc2d5b9..158d7a1b90 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_get.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_get.py @@ -22,13 +22,13 @@ def test_private_get(client): sandbox_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/sandboxEnvironments/3186171392039059456" - agent_engine_sandbox = client.agent_engines.sandboxes._get(name=sandbox_name) - assert isinstance(agent_engine_sandbox, types.SandboxEnvironment) - assert agent_engine_sandbox.name == sandbox_name + runtime_sandbox = client.sandboxes._get(name=sandbox_name) + assert isinstance(runtime_sandbox, types.SandboxEnvironment) + assert runtime_sandbox.name == sandbox_name pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes._get", + test_method="sandboxes._get", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_get_sandbox_operation.py b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_get_sandbox_operation.py index 1de905140d..9b6db2217c 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_get_sandbox_operation.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_get_sandbox_operation.py @@ -24,15 +24,15 @@ def test_private_get_operation(client): "projects/964831358985/locations/us-central1/operations/4799455193970245632" ) - agent_engine_sandbox = client.agent_engines.sandboxes._get_sandbox_operation( + runtime_sandbox = client.sandboxes._get_sandbox_operation( operation_name=operation_name ) - assert isinstance(agent_engine_sandbox, types.AgentEngineSandboxOperation) - assert agent_engine_sandbox.name == operation_name + assert isinstance(runtime_sandbox, types.RuntimeSandboxOperation) + assert runtime_sandbox.name == operation_name pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes._get_sandbox_operation", + test_method="sandboxes._get_sandbox_operation", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_list.py b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_list.py index fb9e5ecd3d..8407b4d125 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_list.py @@ -19,20 +19,20 @@ def test_private_list(client): - agent_engine_sandbox_list_operation = client.agent_engines.sandboxes._list( + runtime_sandbox_list_operation = client.sandboxes._list( name=("reasoningEngines/2886612747586371584"), ) assert isinstance( - agent_engine_sandbox_list_operation.sandbox_environments[0], + runtime_sandbox_list_operation.sandbox_environments[0], types.SandboxEnvironment, ) assert isinstance( - agent_engine_sandbox_list_operation, types.ListAgentEngineSandboxesResponse + runtime_sandbox_list_operation, types.ListRuntimeSandboxesResponse ) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes._list", + test_method="sandboxes._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_delete.py b/tests/unit/agentplatform/genai/replays/test_ae_session_delete.py index 73358920ef..7b44a6e775 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_delete.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_delete.py @@ -19,14 +19,14 @@ def test_delete_session_non_blocking(client): - ae_session_operation = client.agent_engines.sessions.delete( + ae_session_operation = client.sessions.delete( name=("reasoningEngines/2886612747586371584/sessions/8521561049109889024"), ) - assert isinstance(ae_session_operation, types.DeleteAgentEngineSessionOperation) + assert isinstance(ae_session_operation, types.DeleteRuntimeSessionOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions.delete", + test_method="sessions.delete", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_events_append.py b/tests/unit/agentplatform/genai/replays/test_ae_session_events_append.py index aef176ea86..62d15a7094 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_events_append.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_events_append.py @@ -21,17 +21,17 @@ def test_append_session_event(client): - session_event = client.agent_engines.sessions.events.append( + session_event = client.sessions.events.append( name="reasoningEngines/2886612747586371584/sessions/6922431337672474624", author="test-user-123", invocation_id="test-invocation-id", timestamp=datetime.datetime.fromtimestamp(1234567860, tz=datetime.timezone.utc), ) - assert isinstance(session_event, types.AppendAgentEngineSessionEventResponse) + assert isinstance(session_event, types.AppendRuntimeSessionEventResponse) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions.events.append", + test_method="sessions.events.append", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_events_private_list.py b/tests/unit/agentplatform/genai/replays/test_ae_session_events_private_list.py index eb3153ab17..f681f3c335 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_events_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_events_private_list.py @@ -19,11 +19,11 @@ def test_private_list_session_events(client): - session_event_list_response = client.agent_engines.sessions.events._list( + session_event_list_response = client.sessions.events._list( name="reasoningEngines/2886612747586371584/sessions/6922431337672474624", ) assert isinstance( - session_event_list_response, types.ListAgentEngineSessionEventsResponse + session_event_list_response, types.ListRuntimeSessionEventsResponse ) assert len(session_event_list_response.session_events) == 1 assert isinstance(session_event_list_response.session_events[0], types.SessionEvent) @@ -32,5 +32,5 @@ def test_private_list_session_events(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions.events._list", + test_method="sessions.events._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_private_create.py b/tests/unit/agentplatform/genai/replays/test_ae_session_private_create.py index 28dd62cac2..f693c6e164 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_private_create.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_private_create.py @@ -19,15 +19,15 @@ def test_private_create_session(client): - ae_session_operation = client.agent_engines.sessions._create( + ae_session_operation = client.sessions._create( name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584", user_id="test-user-id", ) - assert isinstance(ae_session_operation, types.AgentEngineSessionOperation) + assert isinstance(ae_session_operation, types.RuntimeSessionOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions._create", + test_method="sessions._create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_private_get.py b/tests/unit/agentplatform/genai/replays/test_ae_session_private_get.py index 35bfdd7b63..c1100873b6 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_private_get.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_private_get.py @@ -19,14 +19,14 @@ def test_private_get_session_operation(client): - ae_session_operation = client.agent_engines.sessions._get_session_operation( + ae_session_operation = client.sessions._get_session_operation( operation_name="reasoningEngines/2886612747586371584/sessions/3080649749292908544/operations/758783840595476480", ) - assert isinstance(ae_session_operation, types.AgentEngineSessionOperation) + assert isinstance(ae_session_operation, types.RuntimeSessionOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions._get_session_operation", + test_method="sessions._get_session_operation", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_private_list.py b/tests/unit/agentplatform/genai/replays/test_ae_session_private_list.py index f8a2fe338e..b01cf859ef 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_private_list.py @@ -19,7 +19,7 @@ def test_private_list_session(client): - session_list_response = client.agent_engines.sessions._list( + session_list_response = client.sessions._list( name="reasoningEngines/2886612747586371584", ) assert isinstance(session_list_response, types.ListReasoningEnginesSessionsResponse) @@ -29,5 +29,5 @@ def test_private_list_session(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions._list", + test_method="sessions._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_private_update.py b/tests/unit/agentplatform/genai/replays/test_ae_session_private_update.py index 3f62dee2f0..9acd032e67 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_private_update.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_private_update.py @@ -19,18 +19,18 @@ def test_private_update_session(client): - agent_engine_session_operation = client.agent_engines.sessions._update( + runtime_session_operation = client.sessions._update( name="reasoningEngines/2886612747586371584/sessions/3080649749292908544", - config=types.UpdateAgentEngineSessionConfig( + config=types.UpdateRuntimeSessionConfig( display_name="test-agent-engine-session-updated", user_id="test-user-id", ), ) - assert isinstance(agent_engine_session_operation, types.AgentEngineSessionOperation) + assert isinstance(runtime_session_operation, types.RuntimeSessionOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions._update", + test_method="sessions._update", ) diff --git a/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_methods.py b/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_methods.py index f9665236c0..57d66518ed 100644 --- a/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_methods.py +++ b/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_methods.py @@ -35,10 +35,10 @@ @pytest.mark.asyncio async def test_timeout_is_set(client): - agent_engine = client.agent_engines.get( + runtime = client.runtimes.get( name="projects/932854658080/locations/us-central1/reasoningEngines/857830725653626880", ) - assert isinstance(agent_engine, types.AgentEngine) + assert isinstance(runtime, types.Runtime) message_data = { "messageId": "msg-123", @@ -71,10 +71,10 @@ async def test_timeout_is_set(client): class FakeCredentials: token = "fake-token" - agent_engine.api_client._api_client._credentials = FakeCredentials() + runtime.api_client._api_client._credentials = FakeCredentials() try: - await agent_engine.on_message_send(**message_data) + await runtime.on_message_send(**message_data) except a2a_errors.A2AClientHTTPError as e: # Make sure that the authentication error was successfully # propagated, otherwise the test is not valid. @@ -87,6 +87,6 @@ class FakeCredentials: pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.get", + test_method="runtimes.get", http_options=_api_client.HttpOptions(timeout=99000), ) diff --git a/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_v1_methods.py b/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_v1_methods.py index 200e590fc3..48657e4ddb 100644 --- a/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_v1_methods.py +++ b/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_v1_methods.py @@ -49,10 +49,10 @@ def _build_send_message_request() -> "a2a_types.SendMessageRequest": @pytest.mark.asyncio async def test_timeout_is_set(client): - agent_engine = client.agent_engines.get( + runtime = client.runtimes.get( name="projects/964831358985/locations/us-central1/reasoningEngines/6859679872613089280", ) - assert isinstance(agent_engine, types.AgentEngine) + assert isinstance(runtime, types.Runtime) with mock.patch( "httpx.AsyncClient", spec=httpx.AsyncClient @@ -78,13 +78,13 @@ async def test_timeout_is_set(client): class FakeCredentials: token = "fake-token" - agent_engine.api_client._api_client._credentials = FakeCredentials() + runtime.api_client._api_client._credentials = FakeCredentials() # In a2a 1.0 the wrapped operation forwards the `request` kwarg directly # to `client.send_message(request)`, and HTTP failures surface as an # `A2AClientError` (the legacy `A2AClientHTTPError` no longer exists). with pytest.raises(a2a_errors.A2AClientError) as exc_info: - await agent_engine.on_message_send(request=_build_send_message_request()) + await runtime.on_message_send(request=_build_send_message_request()) # Make sure the authentication failure was propagated, otherwise the # test is not validating the request path. @@ -97,6 +97,6 @@ class FakeCredentials: pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.get", + test_method="runtimes.get", http_options=_api_client.HttpOptions(timeout=99000), ) diff --git a/tests/unit/agentplatform/genai/replays/test_agent_engine_private_create.py b/tests/unit/agentplatform/genai/replays/test_agent_engine_private_create.py index 43a3f3bdae..025e904e11 100644 --- a/tests/unit/agentplatform/genai/replays/test_agent_engine_private_create.py +++ b/tests/unit/agentplatform/genai/replays/test_agent_engine_private_create.py @@ -20,14 +20,14 @@ def test_private_create_with_labels(client): labels = {"test-label": "test-value"} - agent_engine_operation = client.agent_engines._create( + runtime_operation = client.runtimes._create( config={"labels": labels}, ) - assert isinstance(agent_engine_operation, types.AgentEngineOperation) + assert isinstance(runtime_operation, types.RuntimeOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines._create", + test_method="runtimes._create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_agent_engine_private_delete.py b/tests/unit/agentplatform/genai/replays/test_agent_engine_private_delete.py index 565ab4eb5c..721cf64c69 100644 --- a/tests/unit/agentplatform/genai/replays/test_agent_engine_private_delete.py +++ b/tests/unit/agentplatform/genai/replays/test_agent_engine_private_delete.py @@ -19,14 +19,14 @@ def test_private_delete(client): - agent_engine_operation = client.agent_engines._delete( + runtime_operation = client.runtimes._delete( name="reasoningEngines/7571341522470174720", ) - assert isinstance(agent_engine_operation, types.DeleteAgentEngineOperation) + assert isinstance(runtime_operation, types.DeleteRuntimeOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines._delete", + test_method="runtimes._delete", ) diff --git a/tests/unit/agentplatform/genai/replays/test_agent_engine_private_get.py b/tests/unit/agentplatform/genai/replays/test_agent_engine_private_get.py index f36a57a70c..ea72853a09 100644 --- a/tests/unit/agentplatform/genai/replays/test_agent_engine_private_get.py +++ b/tests/unit/agentplatform/genai/replays/test_agent_engine_private_get.py @@ -19,15 +19,15 @@ def test_private_get(client): - agent_engine = client.agent_engines._get( + runtime = client.runtimes._get( name="reasoningEngines/2886612747586371584", ) - assert isinstance(agent_engine, types.ReasoningEngine) - assert agent_engine.labels == {"test-label": "test-value"} + assert isinstance(runtime, types.ReasoningEngine) + assert runtime.labels == {"test-label": "test-value"} pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines._get", + test_method="runtimes._get", ) diff --git a/tests/unit/agentplatform/genai/replays/test_agent_engine_private_update.py b/tests/unit/agentplatform/genai/replays/test_agent_engine_private_update.py index ae02bc390e..74da04b0f3 100644 --- a/tests/unit/agentplatform/genai/replays/test_agent_engine_private_update.py +++ b/tests/unit/agentplatform/genai/replays/test_agent_engine_private_update.py @@ -19,15 +19,15 @@ def test_private_update(client): - agent_engine_operation = client.agent_engines._update( + runtime_operation = client.runtimes._update( name="reasoningEngines/2886612747586371584", - config=types.UpdateAgentEngineConfig(display_name="test-agent-engine-updated"), + config=types.UpdateRuntimeConfig(display_name="test-agent-engine-updated"), ) - assert isinstance(agent_engine_operation, types.AgentEngineOperation) + assert isinstance(runtime_operation, types.RuntimeOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines._update", + test_method="runtimes._update", ) diff --git a/tests/unit/agentplatform/genai/replays/test_append_agent_engine_a2a_task_events.py b/tests/unit/agentplatform/genai/replays/test_append_agent_engine_a2a_task_events.py deleted file mode 100644 index 119c82e591..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_append_agent_engine_a2a_task_events.py +++ /dev/null @@ -1,98 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -import pytest - - -def test_append_simple_a2a_task_events(client): - # Use the autopush environment. - client._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client._api_client._http_options.api_version = "internal" - task = client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - assert isinstance(task, types.A2aTask) - - client.agent_engines.a2a_tasks.events.append( - name=task.name, - task_events=[ - types.TaskEvent( - event_data=types.TaskEventData( - metadata_change=types.TaskMetadataChange( - new_metadata={"key1": "value1"} - ) - ), - event_sequence_number=1, - ) - ], - ) - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), -) - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_append_simple_a2a_task_events_async(client): - # Use the autopush environment. - client.aio._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client.aio._api_client._http_options.api_version = "internal" - task = await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - assert isinstance(task, types.A2aTask) - - await client.aio.agent_engines.a2a_tasks.events.append( - name=task.name, - task_events=[ - types.TaskEvent( - event_data=types.TaskEventData( - metadata_change=types.TaskMetadataChange( - new_metadata={"key1": "value1"} - ) - ), - event_sequence_number=1, - ) - ], - ) - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_append_agent_engine_session_event.py b/tests/unit/agentplatform/genai/replays/test_append_agent_engine_session_event.py deleted file mode 100644 index 30794629eb..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_append_agent_engine_session_event.py +++ /dev/null @@ -1,55 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -import datetime - -from tests.unit.agentplatform.genai.replays import pytest_helper - - -def test_append_session_event(client): - agent_engine = client.agent_engines.create() - operation = client.agent_engines.create_session( - name=agent_engine.api_resource.name, - user_id="test-user-123", - ) - session = operation.response - client.agent_engines.append_session_event( - name=session.name, - author="test-user-123", - invocation_id="test-invocation-id", - timestamp=datetime.datetime.fromtimestamp(1234567890, tz=datetime.timezone.utc), - config={ - "content": { - "parts": [ - { - "text": "Hello World", - }, - ], - }, - "error_code": "test-error-code", - "error_message": "test-error-message", - "raw_event": { - "test-key": "test-value", - }, - }, - ) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.append_session_event", -) diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine.py index f19d061fc4..39fbdc025c 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_agent_engine.py +++ b/tests/unit/agentplatform/genai/replays/test_create_agent_engine.py @@ -37,7 +37,7 @@ def test_create_config_lightweight(client): if not os.environ.get("GCS_BUCKET"): raise ValueError("GCS_BUCKET environment variable is not set.") - config = client.agent_engines._create_config( + config = client.runtimes._create_config( mode="create", staging_bucket=os.environ["GCS_BUCKET"], display_name=agent_display_name, @@ -51,12 +51,12 @@ def test_create_config_lightweight(client): def test_create_with_labels(client): labels = {"test-label": "test-value"} - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={"labels": labels}, ) - assert agent_engine.api_resource.labels == labels + assert runtime.api_resource.labels == labels # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) def test_create_with_context_spec(client): @@ -99,7 +99,7 @@ def test_create_with_context_spec(client): **generation_trigger_config ) - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "context_spec": { "memory_bank_config": { @@ -117,8 +117,8 @@ def test_create_with_context_spec(client): "http_options": {"api_version": "v1beta1"}, }, ) - agent_engine = client.agent_engines.get(name=agent_engine.api_resource.name) - memory_bank_config = agent_engine.api_resource.context_spec.memory_bank_config + runtime = client.runtimes.get(name=runtime.api_resource.name) + memory_bank_config = runtime.api_resource.context_spec.memory_bank_config assert memory_bank_config.generation_config.model == generation_model assert ( memory_bank_config.generation_config.generation_trigger_config @@ -132,13 +132,13 @@ def test_create_with_context_spec(client): memory_bank_customization_config ] # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) def test_create_with_source_packages( client, - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): """Tests creating an agent engine with source packages.""" if sys.version_info >= (3, 13): @@ -168,10 +168,10 @@ def _update_ver(obj): except Exception: pass with ( - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "display_name": "test-agent-engine-source-packages", "source_packages": [ @@ -187,32 +187,32 @@ def _update_ver(obj): }, }, ) - assert agent_engine.api_resource.display_name == "test-agent-engine-source-packages" + assert runtime.api_resource.display_name == "test-agent-engine-source-packages" # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) def test_create_with_identity_type(client): """Tests creating an agent engine with identity type.""" - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "identity_type": types.IdentityType.AGENT_IDENTITY, "http_options": {"api_version": "v1beta1"}, }, ) assert ( - agent_engine.api_resource.spec.identity_type + runtime.api_resource.spec.identity_type == types.IdentityType.AGENT_IDENTITY ) assert _AGENT_IDENTITY_REGEX.match( - agent_engine.api_resource.spec.effective_identity + runtime.api_resource.spec.effective_identity ) # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.create", + test_method="runtimes.create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a.py index a6104ece90..015736d8e3 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a.py +++ b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a.py @@ -19,7 +19,7 @@ from unittest import mock from tests.unit.agentplatform.genai.replays import pytest_helper from agentplatform._genai import types -from agentplatform.agent_engines.templates.a2a import default_a2a_agent +from agentplatform.frameworks.a2a import default_a2a_agent import pytest @@ -45,16 +45,16 @@ def test_create_a2a_agent(client, is_replay_mode): staging_bucket = os.environ["GCS_BUCKET"] # In replay mode, GCS operations are mocked and blob.open("rb") returns a mock - # that fails when cloudpickle.load expects bytes. We mock _upload_agent_engine + # that fails when cloudpickle.load expects bytes. We mock _upload_runtime # to skip this verification step, which is not needed when replaying API calls. upload_patch = ( - mock.patch("agentplatform._genai._agent_engines_utils._upload_agent_engine") + mock.patch("agentplatform._genai._runtimes_utils._upload_runtime") if is_replay_mode else contextlib.nullcontext() ) with upload_patch: - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( agent=my_agent, config={ "staging_bucket": staging_bucket, @@ -69,11 +69,11 @@ def test_create_a2a_agent(client, is_replay_mode): ) - assert isinstance(agent_engine, types.AgentEngine) - assert agent_engine.api_resource.display_name == "test-a2a-agent" + assert isinstance(runtime, types.Runtime) + assert runtime.api_resource.display_name == "test-a2a-agent" # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a_task.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a_task.py deleted file mode 100644 index 6d6355cdf5..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a_task.py +++ /dev/null @@ -1,173 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -from google.genai import types as genai_types -import pytest - - -def test_create_simple_a2a_task(client): - # Use the autopush environment. - client._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client._api_client._http_options.api_version = "internal" - - task = client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig( - context_id="context123", - metadata={ - "key": "value", - "key2": [{"key3": "value3", "key4": "value4"}], - }, - status_details=types.TaskStatusDetails( - task_message=types.TaskMessage( - role="user", - message_id="message123", - parts=[ - genai_types.Part( - text="hello123", - ) - ], - metadata={ - "key42": "value42", - }, - ), - ), - output=types.TaskOutput( - artifacts=[ - types.TaskArtifact( - artifact_id="artifact123", - display_name="display_name123", - description="description123", - parts=[ - genai_types.Part( - text="hello456", - ) - ], - ) - ], - ), - ), - ) - - assert isinstance(task, types.A2aTask) - assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" - assert task.context_id == "context123" - assert task.state == types.State.SUBMITTED - assert task.status_details.task_message.role == "user" - assert task.status_details.task_message.message_id == "message123" - assert task.status_details.task_message.parts[0].text == "hello123" - assert task.status_details.task_message.metadata["key42"] == "value42" - assert task.output.artifacts[0].artifact_id == "artifact123" - assert task.output.artifacts[0].display_name == "display_name123" - assert task.output.artifacts[0].description == "description123" - assert task.output.artifacts[0].parts[0].text == "hello456" - assert task.metadata == { - "key": "value", - "key2": [{"key3": "value3", "key4": "value4"}], - } - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), -) - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_create_simple_a2a_task_async(client): - # Use the autopush environment. - client.aio._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client.aio._api_client._http_options.api_version = "internal" - - task = await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig( - context_id="context123", - metadata={ - "key": "value", - "key2": [{"key3": "value3", "key4": "value4"}], - }, - status_details=types.TaskStatusDetails( - task_message=types.TaskMessage( - role="user", - message_id="message123", - parts=[ - genai_types.Part( - text="hello123", - ) - ], - metadata={ - "key42": "value42", - }, - ), - ), - output=types.TaskOutput( - artifacts=[ - types.TaskArtifact( - artifact_id="artifact123", - display_name="display_name123", - description="description123", - parts=[ - genai_types.Part( - text="hello456", - ) - ], - ) - ], - ), - ), - ) - - assert isinstance(task, types.A2aTask) - assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" - assert task.context_id == "context123" - assert task.state == types.State.SUBMITTED - assert task.status_details.task_message.role == "user" - assert task.status_details.task_message.message_id == "message123" - assert task.status_details.task_message.parts[0].text == "hello123" - assert task.status_details.task_message.metadata["key42"] == "value42" - assert task.output.artifacts[0].artifact_id == "artifact123" - assert task.output.artifacts[0].display_name == "display_name123" - assert task.output.artifacts[0].description == "description123" - assert task.output.artifacts[0].parts[0].text == "hello456" - assert task.metadata == { - "key": "value", - "key2": [{"key3": "value3", "key4": "value4"}], - } - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_developer_connect.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_developer_connect.py index f26d0eb715..6ac5afc0a8 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_developer_connect.py +++ b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_developer_connect.py @@ -57,7 +57,7 @@ def _update_ver(obj): revision="main", dir="test", ) - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "display_name": "test-agent-engine-dev-connect", "developer_connect_source": developer_connect_source_config, @@ -70,25 +70,25 @@ def _update_ver(obj): }, }, ) - assert agent_engine.api_resource.display_name == "test-agent-engine-dev-connect" + assert runtime.api_resource.display_name == "test-agent-engine-dev-connect" assert ( - agent_engine.api_resource.spec.source_code_spec.developer_connect_source.config.git_repository_link + runtime.api_resource.spec.source_code_spec.developer_connect_source.config.git_repository_link == developer_connect_source_config.git_repository_link ) assert ( - agent_engine.api_resource.spec.source_code_spec.developer_connect_source.config.revision + runtime.api_resource.spec.source_code_spec.developer_connect_source.config.revision == developer_connect_source_config.revision ) assert ( - agent_engine.api_resource.spec.source_code_spec.developer_connect_source.config.dir + runtime.api_resource.spec.source_code_spec.developer_connect_source.config.dir == developer_connect_source_config.dir ) # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.create", + test_method="runtimes.create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_docker.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_docker.py index 051264f4ef..73f9692a8a 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_docker.py +++ b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_docker.py @@ -26,8 +26,8 @@ def test_create_with_docker( client, - mock_agent_engine_create_docker_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_docker_base64_encoded_tarball, + mock_runtime_create_path_exists, ): """Tests creating an agent engine with docker spec.""" if sys.version_info >= (3, 13): @@ -57,10 +57,10 @@ def _update_ver(obj): except Exception: pass with ( - mock_agent_engine_create_docker_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_docker_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "display_name": "test-agent-engine-docker", "description": "test agent engine with docker spec", @@ -74,13 +74,13 @@ def _update_ver(obj): }, }, ) - assert agent_engine.api_resource.display_name == "test-agent-engine-docker" + assert runtime.api_resource.display_name == "test-agent-engine-docker" # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.create", + test_method="runtimes.create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_memory.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_memory.py index f321a921c0..c89fffd46c 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_memory.py +++ b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_memory.py @@ -21,9 +21,9 @@ def test_create_memory_with_ttl(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) metadata = { "my_string_key": types.MemoryMetadataValue(string_value="my_string_value"), @@ -36,20 +36,20 @@ def test_create_memory_with_ttl(client): ), } - operation = client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + operation = client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, - config=types.AgentEngineMemoryConfig( + config=types.RuntimeMemoryConfig( display_name="my_memory_fact", ttl="120s", metadata=metadata, ), ) - assert isinstance(operation, types.AgentEngineMemoryOperation) + assert isinstance(operation, types.RuntimeMemoryOperation) assert operation.response.fact == "memory_fact" assert operation.response.scope == {"user_id": "123"} - assert operation.response.name.startswith(agent_engine.api_resource.name) + assert operation.response.name.startswith(runtime.api_resource.name) # Expire time is calculated by the server, so we only check that it is # within a reasonable range to avoid flakiness. assert ( @@ -59,60 +59,60 @@ def test_create_memory_with_ttl(client): ) assert operation.response.metadata == metadata # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) def test_create_memory_with_expire_time(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) expire_time = datetime.datetime( 2027, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc ) - operation = client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + operation = client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, - config=types.AgentEngineMemoryConfig( + config=types.RuntimeMemoryConfig( display_name="my_memory_fact", expire_time=expire_time ), ) - assert isinstance(operation, types.AgentEngineMemoryOperation) + assert isinstance(operation, types.RuntimeMemoryOperation) assert operation.response.fact == "memory_fact" assert operation.response.scope == {"user_id": "123"} - assert operation.response.name.startswith(agent_engine.api_resource.name) + assert operation.response.name.startswith(runtime.api_resource.name) assert operation.response.expire_time == expire_time # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) def test_create_memory_with_custom_memory_id(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) - operation = client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + operation = client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, - config=types.AgentEngineMemoryConfig( + config=types.RuntimeMemoryConfig( display_name="my_memory_fact", memory_id="my-memory-id" ), ) - assert isinstance(operation, types.AgentEngineMemoryOperation) + assert isinstance(operation, types.RuntimeMemoryOperation) assert operation.response.fact == "memory_fact" assert operation.response.scope == {"user_id": "123"} assert ( operation.response.name - == f"{agent_engine.api_resource.name}/memories/my-memory-id" + == f"{runtime.api_resource.name}/memories/my-memory-id" ) # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.create_memory", + test_method="runtimes.create_memory", ) diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_sandbox.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_sandbox.py index c8257565d7..a3a6a52124 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_sandbox.py +++ b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_sandbox.py @@ -19,33 +19,33 @@ def test_create_sandbox(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) - operation = client.agent_engines.sandboxes.create( - name=agent_engine.api_resource.name, + operation = client.sandboxes.create( + name=runtime.api_resource.name, poll_interval_seconds=1, spec={ "code_execution_environment": { "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB" } }, - config=types.CreateAgentEngineSandboxConfig( + config=types.CreateRuntimeSandboxConfig( display_name="test_sandbox", ttl="3600s" ), ) - assert isinstance(operation, types.AgentEngineSandboxOperation) + assert isinstance(operation, types.RuntimeSandboxOperation) assert operation.response.display_name == "test_sandbox" assert ( operation.response.spec.code_execution_environment.machine_config == "MACHINE_CONFIG_VCPU4_RAM4GIB" ) - assert operation.response.name.startswith(agent_engine.api_resource.name) + assert operation.response.name.startswith(runtime.api_resource.name) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.create", + test_method="sandboxes.create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_session.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_session.py index 28361394fc..36e7768510 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_session.py +++ b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_session.py @@ -21,27 +21,27 @@ def test_create_session_with_ttl(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() try: - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) - operation = client.agent_engines.create_session( - name=agent_engine.api_resource.name, + operation = client.sessions.create( + name=runtime.api_resource.name, user_id="test-user-123", - config=types.CreateAgentEngineSessionConfig( + config=types.CreateRuntimeSessionConfig( display_name="my_session", session_state={"foo": "bar"}, ttl="1200000s", labels={"label_key": "label_value"}, ), ) - assert isinstance(operation, types.AgentEngineSessionOperation) + assert isinstance(operation, types.RuntimeSessionOperation) assert operation.response.display_name == "my_session" assert operation.response.session_state == {"foo": "bar"} assert operation.response.user_id == "test-user-123" assert operation.response.labels == {"label_key": "label_value"} - assert operation.response.name.startswith(agent_engine.api_resource.name) + assert operation.response.name.startswith(runtime.api_resource.name) assert operation.done # Expire time is calculated by the server, so we only check that it is # within a reasonable range to avoid flakiness. @@ -52,70 +52,70 @@ def test_create_session_with_ttl(client): ) finally: # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) def test_create_session_with_expire_time(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() try: - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) expire_time = datetime.datetime( 2028, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc ) - operation = client.agent_engines.sessions.create( - name=agent_engine.api_resource.name, + operation = client.sessions.create( + name=runtime.api_resource.name, user_id="test-user-123", - config=types.CreateAgentEngineSessionConfig( + config=types.CreateRuntimeSessionConfig( display_name="my_session", session_state={"foo": "bar"}, expire_time=expire_time, ), ) - assert isinstance(operation, types.AgentEngineSessionOperation) + assert isinstance(operation, types.RuntimeSessionOperation) assert operation.response.display_name == "my_session" assert operation.response.session_state == {"foo": "bar"} assert operation.response.user_id == "test-user-123" - assert operation.response.name.startswith(agent_engine.api_resource.name) + assert operation.response.name.startswith(runtime.api_resource.name) assert operation.response.expire_time == expire_time assert operation.done finally: # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) def test_create_session_with_custom_session_id(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() try: - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) - operation = client.agent_engines.sessions.create( - name=agent_engine.api_resource.name, + operation = client.sessions.create( + name=runtime.api_resource.name, user_id="test-user-123", - config=types.CreateAgentEngineSessionConfig( + config=types.CreateRuntimeSessionConfig( display_name="my_session", session_state={"foo": "bar"}, session_id="my-session-id", ), ) - assert isinstance(operation, types.AgentEngineSessionOperation) + assert isinstance(operation, types.RuntimeSessionOperation) assert operation.response.display_name == "my_session" assert operation.response.session_state == {"foo": "bar"} assert operation.response.user_id == "test-user-123" assert ( operation.response.name - == f"{agent_engine.api_resource.name}/sessions/my-session-id" + == f"{runtime.api_resource.name}/sessions/my-session-id" ) assert operation.done finally: # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions.create", + test_method="sessions.create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_create_feedback_entry.py b/tests/unit/agentplatform/genai/replays/test_create_feedback_entry.py index 9dec74af48..60ce635270 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_feedback_entry.py +++ b/tests/unit/agentplatform/genai/replays/test_create_feedback_entry.py @@ -32,13 +32,13 @@ def test_create(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) try: operation = client.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_123", event_id="event_456", feedback_type=types.FeedbackType.THUMBS_UP, @@ -68,21 +68,21 @@ def test_create(client): } finally: # Clean up resources. - client.agent_engines.delete( - name=agent_engine.api_resource.name, + client.runtimes.delete( + name=runtime.api_resource.name, force=True, ) @pytest.mark.asyncio async def test_create_async(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) try: operation = await client.aio.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_123", event_id="event_456", feedback_type=types.FeedbackType.THUMBS_UP, @@ -112,7 +112,7 @@ async def test_create_async(client): } finally: # Clean up resources. - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True, ) diff --git a/tests/unit/agentplatform/genai/replays/test_delete_ae_runtime_revision.py b/tests/unit/agentplatform/genai/replays/test_delete_ae_runtime_revision.py index 9a214b4baa..e23939cd17 100644 --- a/tests/unit/agentplatform/genai/replays/test_delete_ae_runtime_revision.py +++ b/tests/unit/agentplatform/genai/replays/test_delete_ae_runtime_revision.py @@ -27,8 +27,8 @@ def test_delete_runtime_revision( client, - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): client._api_client._http_options.base_url = ( "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" @@ -36,10 +36,10 @@ def test_delete_runtime_revision( client._api_client._http_options.api_version = "v1beta1" with ( - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "display_name": "test-agent-engine-delete-runtime-revision", "source_packages": [ @@ -58,11 +58,11 @@ def test_delete_runtime_revision( # Create a second runtime revision, # since it's not possible to delete if there is only one runtime revision. with ( - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - updated_agent_engine = client.agent_engines.update( - name=agent_engine.api_resource.name, + updated_runtime = client.runtimes.update( + name=runtime.api_resource.name, config={ "display_name": "test-agent-engine-update-traffic-with-agent-after-update", "source_packages": [ @@ -79,19 +79,19 @@ def test_delete_runtime_revision( }, ) - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( - name=updated_agent_engine.api_resource.name, + runtime_revisions_iter = client.runtimes.revisions.list( + name=updated_runtime.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) assert len(runtime_revisions_list) == 2 revision_to_delete = runtime_revisions_list[1] - operation = client.agent_engines.runtimes.revisions.delete( + operation = client.runtimes.revisions.delete( name=revision_to_delete.api_resource.name, ) - assert isinstance(operation, types.DeleteAgentEngineRuntimeRevisionOperation) + assert isinstance(operation, types.DeleteRuntimeRevisionOperation) assert operation.done - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( - name=updated_agent_engine.api_resource.name, + runtime_revisions_iter = client.runtimes.revisions.list( + name=updated_runtime.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) assert len(runtime_revisions_list) == 1 @@ -99,13 +99,13 @@ def test_delete_runtime_revision( runtime_revisions_list[0].api_resource.name != revision_to_delete.api_resource.name ) - client.agent_engines.delete(name=updated_agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=updated_runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.runtimes.revisions.delete", + test_method="runtimes.revisions.delete", ) @@ -115,8 +115,8 @@ def test_delete_runtime_revision( @pytest.mark.asyncio async def test_delete_runtime_revision_async( client, - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): client._api_client._http_options.base_url = ( "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" @@ -124,10 +124,10 @@ async def test_delete_runtime_revision_async( client._api_client._http_options.api_version = "v1beta1" with ( - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "display_name": "test-agent-engine-delete-runtime-revision", "source_packages": [ @@ -146,11 +146,11 @@ async def test_delete_runtime_revision_async( # Create a second runtime revision, # since it's not possible to delete if there is only one runtime revision. with ( - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - updated_agent_engine = client.agent_engines.update( - name=agent_engine.api_resource.name, + updated_runtime = client.runtimes.update( + name=runtime.api_resource.name, config={ "display_name": "test-agent-engine-update-traffic-with-agent-after-update", "source_packages": [ @@ -167,18 +167,18 @@ async def test_delete_runtime_revision_async( }, ) - runtime_revisions_iter = client.aio.agent_engines.runtimes.revisions.list( - name=updated_agent_engine.api_resource.name, + runtime_revisions_iter = client.aio.runtimes.revisions.list( + name=updated_runtime.api_resource.name, ) runtime_revisions_list = [] async for revision in runtime_revisions_iter: runtime_revisions_list.append(revision) assert len(runtime_revisions_list) == 2 revision_to_delete = runtime_revisions_list[1] - operation = await client.aio.agent_engines.runtimes.revisions.delete( + operation = await client.aio.runtimes.revisions.delete( name=revision_to_delete.api_resource.name, ) - assert isinstance(operation, types.DeleteAgentEngineRuntimeRevisionOperation) - await client.aio.agent_engines.delete( - name=updated_agent_engine.api_resource.name, force=True + assert isinstance(operation, types.DeleteRuntimeRevisionOperation) + await client.aio.runtimes.delete( + name=updated_runtime.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine.py b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine.py index 4a159ed0f6..033c6879f7 100644 --- a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine.py +++ b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine.py @@ -24,17 +24,17 @@ def test_agent_engine_delete(client, caplog): caplog.set_level(logging.INFO) - agent_engine = client.agent_engines.create() - operation = client.agent_engines.delete(name=agent_engine.api_resource.name) - assert isinstance(operation, types.DeleteAgentEngineOperation) - assert "Deleting AgentEngine resource" in caplog.text - assert f"Started AgentEngine delete operation: {operation.name}" in caplog.text + runtime = client.runtimes.create() + operation = client.runtimes.delete(name=runtime.api_resource.name) + assert isinstance(operation, types.DeleteRuntimeOperation) + assert "Deleting Runtime resource" in caplog.text + assert f"Started Runtime delete operation: {operation.name}" in caplog.text pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.delete", + test_method="runtimes.delete", ) @@ -45,10 +45,10 @@ def test_agent_engine_delete(client, caplog): async def test_agent_engine_delete_async(client, caplog): caplog.set_level(logging.INFO) # TODO(b/431785750): use async methods for create() when available - agent_engine = client.agent_engines.create() - operation = await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name + runtime = client.runtimes.create() + operation = await client.aio.runtimes.delete( + name=runtime.api_resource.name ) - assert isinstance(operation, types.DeleteAgentEngineOperation) - assert "Deleting AgentEngine resource" in caplog.text - assert f"Started AgentEngine delete operation: {operation.name}" in caplog.text + assert isinstance(operation, types.DeleteRuntimeOperation) + assert "Deleting Runtime resource" in caplog.text + assert f"Started Runtime delete operation: {operation.name}" in caplog.text diff --git a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_a2a_task.py b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_a2a_task.py deleted file mode 100644 index b2859dea24..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_a2a_task.py +++ /dev/null @@ -1,96 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -from google.genai import errors -import pytest - - -def test_delete_a2a_task(client): - # Use the autopush environment. - client._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client._api_client._http_options.api_version = "internal" - - created_task = client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - assert isinstance(created_task, types.A2aTask) - - # Test validity check. - task = client.agent_engines.a2a_tasks.get( - name=created_task.name, - ) - assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" - - client.agent_engines.a2a_tasks.delete(name=created_task.name) - - with pytest.raises(errors.ClientError, match="404 NOT_FOUND"): - client.agent_engines.a2a_tasks.get(name=created_task.name) - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), -) - - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_delete_a2a_task_async(client): - # Use the autopush environment. - client.aio._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client.aio._api_client._http_options.api_version = "internal" - - created_task = await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - assert isinstance(created_task, types.A2aTask) - - # Test validity check. - task = await client.aio.agent_engines.a2a_tasks.get( - name=created_task.name, - ) - assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" - - await client.aio.agent_engines.a2a_tasks.delete(name=created_task.name) - - with pytest.raises(errors.ClientError, match="404 NOT_FOUND"): - await client.aio.agent_engines.a2a_tasks.get(name=created_task.name) - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_memory.py b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_memory.py index 9d3dea4c72..480db8f95f 100644 --- a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_memory.py +++ b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_memory.py @@ -22,23 +22,23 @@ def test_delete_memory(client): - agent_engine = client.agent_engines.create() - operation = client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + runtime = client.runtimes.create() + operation = client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, ) memory = operation.response - operation = client.agent_engines.memories.delete(name=memory.name) - assert isinstance(operation, types.DeleteAgentEngineMemoryOperation) + operation = client.runtimes.memories.delete(name=memory.name) + assert isinstance(operation, types.DeleteRuntimeMemoryOperation) assert operation.name.startswith(memory.name + "/operations/") - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.delete_memory", + test_method="runtimes.delete_memory", ) @@ -47,16 +47,16 @@ def test_delete_memory(client): @pytest.mark.asyncio async def test_delete_memory_async(client): - agent_engine = client.agent_engines.create() - operation = await client.aio.agent_engines.memories.create( - name=agent_engine.api_resource.name, + runtime = client.runtimes.create() + operation = await client.aio.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, ) memory = operation.response - operation = await client.aio.agent_engines.memories.delete(name=memory.name) - assert isinstance(operation, types.DeleteAgentEngineMemoryOperation) + operation = await client.aio.runtimes.memories.delete(name=memory.name) + assert isinstance(operation, types.DeleteRuntimeMemoryOperation) assert operation.name.startswith(memory.name + "/operations/") - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_sandbox.py b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_sandbox.py index 4346ad62dd..c66763aacc 100644 --- a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_sandbox.py +++ b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_sandbox.py @@ -18,29 +18,29 @@ def test_delete_sandbox(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) - operation = client.agent_engines.sandboxes.create( - name=agent_engine.api_resource.name, + operation = client.sandboxes.create( + name=runtime.api_resource.name, spec={ "code_execution_environment": { "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB" } }, - config=types.CreateAgentEngineSandboxConfig(display_name="test_sandbox"), + config=types.CreateRuntimeSandboxConfig(display_name="test_sandbox"), ) - assert isinstance(operation, types.AgentEngineSandboxOperation) - delete_operation = client.agent_engines.sandboxes.delete( + assert isinstance(operation, types.RuntimeSandboxOperation) + delete_operation = client.sandboxes.delete( name=operation.response.name, ) - assert isinstance(delete_operation, types.DeleteAgentEngineSandboxOperation) + assert isinstance(delete_operation, types.DeleteRuntimeSandboxOperation) assert "/operations/" in delete_operation.name pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.delete", + test_method="sandboxes.delete", ) diff --git a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_session.py b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_session.py deleted file mode 100644 index 17775bb39b..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_session.py +++ /dev/null @@ -1,37 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types - - -def test_delete_session(client): - agent_engine = client.agent_engines.create() - operation = client.agent_engines.create_session( - name=agent_engine.api_resource.name, - user_id="test-user-123", - ) - session = operation.response - operation = client.agent_engines.delete_session(name=session.name) - assert isinstance(operation, types.DeleteAgentEngineSessionOperation) - assert "/operations/" in operation.name - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.delete_session", -) diff --git a/tests/unit/agentplatform/genai/replays/test_delete_feedback_entry.py b/tests/unit/agentplatform/genai/replays/test_delete_feedback_entry.py index 94809cc761..1a03380bd5 100644 --- a/tests/unit/agentplatform/genai/replays/test_delete_feedback_entry.py +++ b/tests/unit/agentplatform/genai/replays/test_delete_feedback_entry.py @@ -33,13 +33,13 @@ def test_delete(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) try: create_operation = client.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_123", event_id="event_456", feedback_type=types.FeedbackType.THUMBS_UP, @@ -59,21 +59,21 @@ def test_delete(client): client.feedback_entries.get(name=create_operation.response.name) finally: # Clean up resources. - client.agent_engines.delete( - name=agent_engine.api_resource.name, + client.runtimes.delete( + name=runtime.api_resource.name, force=True, ) @pytest.mark.asyncio async def test_delete_async(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) try: create_operation = await client.aio.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_123", event_id="event_456", feedback_type=types.FeedbackType.THUMBS_UP, @@ -92,7 +92,7 @@ async def test_delete_async(client): await client.aio.feedback_entries.get(name=create_operation.response.name) finally: # Clean up resources. - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True, ) diff --git a/tests/unit/agentplatform/genai/replays/test_evaluate_instances.py b/tests/unit/agentplatform/genai/replays/test_evaluate_instances.py index 639b57a165..c43049cb89 100644 --- a/tests/unit/agentplatform/genai/replays/test_evaluate_instances.py +++ b/tests/unit/agentplatform/genai/replays/test_evaluate_instances.py @@ -180,7 +180,7 @@ def test_run_inference_with_agent(client): agent="projects/977012026409/locations/us-central1/reasoningEngines/7188347537655332864", src=test_df, ) - assert inference_result.candidate_name == "agent_engine_0" + assert inference_result.candidate_name == "runtime_0" assert inference_result.gcs_source is None @@ -193,8 +193,7 @@ def test_evaluation_with_interaction(client): ), gemini_agent_config=types.GeminiAgentConfig( gemini_agent=( - "projects/977012026409/locations/global/agents/" - "test-agent-eval" + "projects/977012026409/locations/global/agents/test-agent-eval" ), ), ) @@ -207,6 +206,7 @@ def test_evaluation_with_interaction(client): ) assert response is not None + def test_evaluate_method_with_interaction(client): eval_case = types.EvalCase( interactions_data_source=types.InteractionsDataSource( diff --git a/tests/unit/agentplatform/genai/replays/test_execute_code_agent_engine_sandbox.py b/tests/unit/agentplatform/genai/replays/test_execute_code_agent_engine_sandbox.py index 928c1f7839..399039d2a0 100644 --- a/tests/unit/agentplatform/genai/replays/test_execute_code_agent_engine_sandbox.py +++ b/tests/unit/agentplatform/genai/replays/test_execute_code_agent_engine_sandbox.py @@ -19,20 +19,20 @@ def test_execute_code_sandbox(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) - operation = client.agent_engines.sandboxes.create( - name=agent_engine.api_resource.name, + operation = client.sandboxes.create( + name=runtime.api_resource.name, spec={ "code_execution_environment": { "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB" } }, - config=types.CreateAgentEngineSandboxConfig(display_name="test_sandbox"), + config=types.CreateRuntimeSandboxConfig(display_name="test_sandbox"), ) - assert isinstance(operation, types.AgentEngineSandboxOperation) + assert isinstance(operation, types.RuntimeSandboxOperation) code = """ with open("test.txt", "r") as input: @@ -50,7 +50,7 @@ def test_execute_code_sandbox(client): } ], } - response = client.agent_engines.sandboxes.execute_code( + response = client.sandboxes.execute_code( name=operation.response.name, input_data=input_data, ) @@ -62,5 +62,5 @@ def test_execute_code_sandbox(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.execute_code", + test_method="sandboxes.execute_code", ) diff --git a/tests/unit/agentplatform/genai/replays/test_generate_agent_engine_memories.py b/tests/unit/agentplatform/genai/replays/test_generate_agent_engine_memories.py index 4942e3ac9e..62e2b9a50f 100644 --- a/tests/unit/agentplatform/genai/replays/test_generate_agent_engine_memories.py +++ b/tests/unit/agentplatform/genai/replays/test_generate_agent_engine_memories.py @@ -24,16 +24,16 @@ def test_generate_and_rollback_memories(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() assert not list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name, + client.runtimes.memories.list( + name=runtime.api_resource.name, ) ) # Generate memories using source content. This result is non-deterministic, # because an LLM is used to generate the memories. - client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, + client.runtimes.memories.generate( + name=runtime.api_resource.name, scope={"user_id": "test-user-id"}, direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource( events=[ @@ -49,20 +49,20 @@ def test_generate_and_rollback_memories(client): ) ] ), - config=types.GenerateAgentEngineMemoriesConfig( + config=types.GenerateRuntimeMemoriesConfig( revision_labels={"key": "value"} ), ) memories = list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name, + client.runtimes.memories.list( + name=runtime.api_resource.name, ) ) assert len(memories) >= 1 # Every action that modifies a memory creates a new revision. memory_revisions = list( - client.agent_engines.memories.revisions.list( + client.runtimes.memories.revisions.list( name=memories[0].name, ) ) @@ -72,29 +72,29 @@ def test_generate_and_rollback_memories(client): revision_name = memory_revisions[0].name # Update the memory. - client.agent_engines.memories._update( + client.runtimes.memories._update( name=memories[0].name, fact="This is temporary", scope={"user_id": "test-user-id"}, ) - memory = client.agent_engines.memories.get(name=memories[0].name) + memory = client.runtimes.memories.get(name=memories[0].name) assert memory.fact == "This is temporary" # Rollback to the revision with the original fact that was created by the # generation request. - client.agent_engines.memories.rollback( + client.runtimes.memories.rollback( name=memories[0].name, target_revision_id=revision_name.split("/")[-1], ) - memory = client.agent_engines.memories.get(name=memories[0].name) + memory = client.runtimes.memories.get(name=memories[0].name) assert memory.fact == memory_revisions[0].fact # Update the memory again using generation. We use the original source # content to ensure that the original memory is updated. The response should # refer to the previous revision. pre_extracted_fact = "I am a software engineer focusing in security" - response = client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, + response = client.runtimes.memories.generate( + name=runtime.api_resource.name, scope={"user_id": "test-user-id"}, direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource( direct_memories=[ @@ -107,19 +107,19 @@ def test_generate_and_rollback_memories(client): # The memory was updated, so the previous revision is set. assert response.response.generated_memories[0].previous_revision is not None memory_revisions = list( - client.agent_engines.memories.revisions.list(name=memories[0].name) + client.runtimes.memories.revisions.list(name=memories[0].name) ) # Memory Revisions are returned in descending order by revision create time. # We can't make an assertion on the actual value, since it's # generated and thus non-deterministic. assert memory_revisions[0].extracted_memories[0].fact is not None - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) def test_generate_memories_direct_memories_source(client): - agent_engine = client.agent_engines.create() - client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, + runtime = client.runtimes.create() + client.runtimes.memories.generate( + name=runtime.api_resource.name, scope={"user_id": "test-user-id"}, direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource( direct_memories=[ @@ -131,23 +131,23 @@ def test_generate_memories_direct_memories_source(client): ), ] ), - config=types.GenerateAgentEngineMemoriesConfig(wait_for_completion=True), + config=types.GenerateRuntimeMemoriesConfig(wait_for_completion=True), ) assert ( len( list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name, + client.runtimes.memories.list( + name=runtime.api_resource.name, ) ) ) >= 1 ) - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) def test_generate_memories_with_metadata(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() metadata = { "my_string_key": types.MemoryMetadataValue(string_value="my_string_value"), "my_double_key": types.MemoryMetadataValue(double_value=123.456), @@ -169,14 +169,14 @@ def test_generate_memories_with_metadata(client): ) scope = {"user_id": "test-user-id"} - operation = client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, + operation = client.runtimes.memories.generate( + name=runtime.api_resource.name, scope=scope, direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig(metadata=metadata), + config=types.GenerateRuntimeMemoriesConfig(metadata=metadata), ) assert len(operation.response.generated_memories) >= 1 - memory = client.agent_engines.memories.get( + memory = client.runtimes.memories.get( name=operation.response.generated_memories[0].memory.name ) assert memory.metadata == metadata @@ -185,11 +185,11 @@ def test_generate_memories_with_metadata(client): overwrite_metadata = { "my_string_key": types.MemoryMetadataValue(string_value="new_value"), } - operation = client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, + operation = client.runtimes.memories.generate( + name=runtime.api_resource.name, scope=scope, direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig( + config=types.GenerateRuntimeMemoriesConfig( metadata=overwrite_metadata, metadata_merge_strategy=types.MemoryMetadataMergeStrategy.OVERWRITE, ), @@ -199,7 +199,7 @@ def test_generate_memories_with_metadata(client): operation.response.generated_memories[0].action == types.GenerateMemoriesResponseGeneratedMemoryAction.UPDATED ) - memory = client.agent_engines.memories.get( + memory = client.runtimes.memories.get( name=operation.response.generated_memories[0].memory.name ) assert memory.metadata == overwrite_metadata @@ -208,11 +208,11 @@ def test_generate_memories_with_metadata(client): new_metadata = { "my_double_key": types.MemoryMetadataValue(double_value=123.456), } - operation = client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, + operation = client.runtimes.memories.generate( + name=runtime.api_resource.name, scope=scope, direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig( + config=types.GenerateRuntimeMemoriesConfig( metadata=new_metadata, metadata_merge_strategy=types.MemoryMetadataMergeStrategy.MERGE, ), @@ -222,7 +222,7 @@ def test_generate_memories_with_metadata(client): operation.response.generated_memories[0].action == types.GenerateMemoriesResponseGeneratedMemoryAction.UPDATED ) - memory = client.agent_engines.memories.get( + memory = client.runtimes.memories.get( name=operation.response.generated_memories[0].memory.name ) assert memory.metadata == {**overwrite_metadata, **new_metadata} @@ -233,11 +233,11 @@ def test_generate_memories_with_metadata(client): restricted_metadata = { "my_string_key": types.MemoryMetadataValue(string_value="new_value2"), } - operation = client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, + operation = client.runtimes.memories.generate( + name=runtime.api_resource.name, scope=scope, direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig( + config=types.GenerateRuntimeMemoriesConfig( metadata=restricted_metadata, metadata_merge_strategy="REQUIRE_EXACT_MATCH", ), @@ -248,18 +248,18 @@ def test_generate_memories_with_metadata(client): operation.response.generated_memories[0].action == types.GenerateMemoriesResponseGeneratedMemoryAction.CREATED ) - memory = client.agent_engines.memories.get( + memory = client.runtimes.memories.get( name=operation.response.generated_memories[0].memory.name ) assert memory.metadata == restricted_metadata # Send a second request where the metadata matches only one of the existing # memories. - operation = client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, + operation = client.runtimes.memories.generate( + name=runtime.api_resource.name, scope=scope, direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig( + config=types.GenerateRuntimeMemoriesConfig( metadata=restricted_metadata, metadata_merge_strategy="REQUIRE_EXACT_MATCH", ), @@ -271,13 +271,13 @@ def test_generate_memories_with_metadata(client): ) assert operation.response.generated_memories[0].memory.name == memory.name - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.generate_memories", + test_method="runtimes.generate_memories", ) @@ -286,9 +286,9 @@ def test_generate_memories_with_metadata(client): @pytest.mark.asyncio async def test_generate_and_rollback_memories_async(client): - agent_engine = client.agent_engines.create() - await client.aio.agent_engines.memories.generate( - name=agent_engine.api_resource.name, + runtime = client.runtimes.create() + await client.aio.runtimes.memories.generate( + name=runtime.api_resource.name, scope={"user_id": "test-user-id"}, direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource( direct_memories=[ @@ -300,15 +300,15 @@ async def test_generate_and_rollback_memories_async(client): ), ] ), - config=types.GenerateAgentEngineMemoriesConfig(wait_for_completion=True), + config=types.GenerateRuntimeMemoriesConfig(wait_for_completion=True), ) - memories_pager = await client.aio.agent_engines.memories.list( - name=agent_engine.api_resource.name + memories_pager = await client.aio.runtimes.memories.list( + name=runtime.api_resource.name ) memory_list = [item async for item in memories_pager] assert len(memory_list) >= 1 - revisions_pager = await client.aio.agent_engines.memories.revisions.list( + revisions_pager = await client.aio.runtimes.memories.revisions.list( name=memory_list[0].name ) memory_revisions = [item async for item in revisions_pager] @@ -316,32 +316,32 @@ async def test_generate_and_rollback_memories_async(client): revision_name = memory_revisions[0].name # Update the memory. - client.agent_engines.memories._update( + client.runtimes.memories._update( name=memory_list[0].name, fact="This is temporary", scope={"user_id": "test-user-id"}, ) - memory = await client.aio.agent_engines.memories.get(name=memory_list[0].name) + memory = await client.aio.runtimes.memories.get(name=memory_list[0].name) assert memory.fact == "This is temporary" # Rollback to the revision with the original fact that was created by the # generation request. - await client.aio.agent_engines.memories.rollback( + await client.aio.runtimes.memories.rollback( name=memory_list[0].name, target_revision_id=revision_name.split("/")[-1], ) - memory = await client.aio.agent_engines.memories.get(name=memory_list[0].name) + memory = await client.aio.runtimes.memories.get(name=memory_list[0].name) assert memory.fact == memory_revisions[0].fact - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True ) def test_generate_memories_with_allowed_topics(client): - agent_engine = client.agent_engines.create() - client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, + runtime = client.runtimes.create() + client.runtimes.memories.generate( + name=runtime.api_resource.name, scope={"user_id": "test-user-id"}, direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource( events=[ @@ -359,7 +359,7 @@ def test_generate_memories_with_allowed_topics(client): ), ] ), - config=types.GenerateAgentEngineMemoriesConfig( + config=types.GenerateRuntimeMemoriesConfig( allowed_topics=[ types.MemoryTopicId( managed_memory_topic=types.ManagedTopicEnum.USER_PREFERENCES @@ -371,11 +371,11 @@ def test_generate_memories_with_allowed_topics(client): assert ( len( list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name, + client.runtimes.memories.list( + name=runtime.api_resource.name, ) ) ) == 1 ) - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_get_ae_runtime_revision.py b/tests/unit/agentplatform/genai/replays/test_get_ae_runtime_revision.py index ef3413b5e7..e780db2bca 100644 --- a/tests/unit/agentplatform/genai/replays/test_get_ae_runtime_revision.py +++ b/tests/unit/agentplatform/genai/replays/test_get_ae_runtime_revision.py @@ -26,8 +26,8 @@ def test_get_runtime_revisions( client, - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): client._api_client._http_options.base_url = ( "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" @@ -35,10 +35,10 @@ def test_get_runtime_revisions( client._api_client._http_options.api_version = "v1beta1" with ( - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "display_name": "test-agent-engine-get-runtime-revisions", "source_packages": [ @@ -55,30 +55,30 @@ def test_get_runtime_revisions( }, ) assert ( - agent_engine.api_resource.display_name + runtime.api_resource.display_name == "test-agent-engine-get-runtime-revisions" ) - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( - name=agent_engine.api_resource.name, + runtime_revisions_iter = client.runtimes.revisions.list( + name=runtime.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) assert len(runtime_revisions_list) == 1 - assert isinstance(runtime_revisions_list[0], types.AgentEngineRuntimeRevision) + assert isinstance(runtime_revisions_list[0], types.RuntimeRevision) runtime_revision_name = runtime_revisions_list[0].api_resource.name - runtime_revision = client.agent_engines.runtimes.revisions.get( + runtime_revision = client.runtimes.revisions.get( name=runtime_revision_name, ) - assert isinstance(runtime_revision, types.AgentEngineRuntimeRevision) + assert isinstance(runtime_revision, types.RuntimeRevision) assert runtime_revision.api_resource.name == runtime_revision_name # Clean up resources. - agent_engine.delete(force=True) + runtime.delete(force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.runtimes.revisions.get", + test_method="runtimes.revisions.get", ) pytest_plugins = ("pytest_asyncio",) @@ -87,8 +87,8 @@ def test_get_runtime_revisions( @pytest.mark.asyncio async def test_async_get_runtime_revisions( client, - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): client._api_client._http_options.base_url = ( "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" @@ -96,10 +96,10 @@ async def test_async_get_runtime_revisions( client._api_client._http_options.api_version = "v1beta1" with ( - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "display_name": "test-agent-engine-get-runtime-revisions", "source_packages": [ @@ -116,24 +116,24 @@ async def test_async_get_runtime_revisions( }, ) assert ( - agent_engine.api_resource.display_name + runtime.api_resource.display_name == "test-agent-engine-get-runtime-revisions" ) - runtime_revisions_iter = client.aio.agent_engines.runtimes.revisions.list( - name=agent_engine.api_resource.name, + runtime_revisions_iter = client.aio.runtimes.revisions.list( + name=runtime.api_resource.name, ) runtime_revisions_list = [] async for revision in runtime_revisions_iter: runtime_revisions_list.append(revision) assert len(runtime_revisions_list) == 1 - assert isinstance(runtime_revisions_list[0], types.AgentEngineRuntimeRevision) + assert isinstance(runtime_revisions_list[0], types.RuntimeRevision) runtime_revision_name = runtime_revisions_list[0].api_resource.name - runtime_revision = await client.aio.agent_engines.runtimes.revisions.get( + runtime_revision = await client.aio.runtimes.revisions.get( name=runtime_revision_name, ) - assert isinstance(runtime_revision, types.AgentEngineRuntimeRevision) + assert isinstance(runtime_revision, types.RuntimeRevision) assert runtime_revision.api_resource.name == runtime_revision_name # Clean up resources. - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_a2a_task.py b/tests/unit/agentplatform/genai/replays/test_get_agent_engine_a2a_task.py deleted file mode 100644 index 6ffb4b4b50..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_a2a_task.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -import pytest - - -def test_get_a2a_task(client): - # Use the autopush environment. - client._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client._api_client._http_options.api_version = "internal" - - created_task = client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - assert isinstance(created_task, types.A2aTask) - - task = client.agent_engines.a2a_tasks.get( - name=created_task.name, - ) - assert isinstance(task, types.A2aTask) - assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" - assert task.context_id == "context123" - assert task.state == types.State.SUBMITTED - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), -) - - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_get_a2a_task_async(client): - # Use the autopush environment. - client.aio._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client.aio._api_client._http_options.api_version = "internal" - - created_task = await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - assert isinstance(created_task, types.A2aTask) - - task = await client.aio.agent_engines.a2a_tasks.get( - name=created_task.name, - ) - assert isinstance(task, types.A2aTask) - assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" - assert task.context_id == "context123" - assert task.state == types.State.SUBMITTED - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py b/tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py index 3ec3093562..61a69f9fa2 100644 --- a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py +++ b/tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py @@ -21,25 +21,25 @@ def test_get_memory(client): - agent_engine = client.agent_engines.create() - operation = client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + runtime = client.runtimes.create() + operation = client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, ) - assert isinstance(operation, types.AgentEngineMemoryOperation) - memory = client.agent_engines.memories.get( + assert isinstance(operation, types.RuntimeMemoryOperation) + memory = client.runtimes.memories.get( name=operation.response.name, ) assert isinstance(memory, types.Memory) assert memory.name == operation.response.name - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.get", + test_method="runtimes.memories.get", ) @@ -48,18 +48,18 @@ def test_get_memory(client): @pytest.mark.asyncio async def test_get_memory_async(client): - agent_engine = client.agent_engines.create() - operation = await client.aio.agent_engines.memories.create( - name=agent_engine.api_resource.name, + runtime = client.runtimes.create() + operation = await client.aio.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, ) - assert isinstance(operation, types.AgentEngineMemoryOperation) - memory = await client.aio.agent_engines.memories.get( + assert isinstance(operation, types.RuntimeMemoryOperation) + memory = await client.aio.runtimes.memories.get( name=operation.response.name, ) assert isinstance(memory, types.Memory) assert memory.name == operation.response.name - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_sandbox.py b/tests/unit/agentplatform/genai/replays/test_get_agent_engine_sandbox.py index 4c921cb5e7..984c261bfc 100644 --- a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_sandbox.py +++ b/tests/unit/agentplatform/genai/replays/test_get_agent_engine_sandbox.py @@ -19,21 +19,21 @@ def test_get_sandbox(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) - operation = client.agent_engines.sandboxes.create( - name=agent_engine.api_resource.name, + operation = client.sandboxes.create( + name=runtime.api_resource.name, spec={ "code_execution_environment": { "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB" } }, - config=types.CreateAgentEngineSandboxConfig(display_name="test_sandbox"), + config=types.CreateRuntimeSandboxConfig(display_name="test_sandbox"), ) - assert isinstance(operation, types.AgentEngineSandboxOperation) - sandbox = client.agent_engines.sandboxes.get( + assert isinstance(operation, types.RuntimeSandboxOperation) + sandbox = client.sandboxes.get( name=operation.response.name, ) assert isinstance(sandbox, types.SandboxEnvironment) @@ -44,5 +44,5 @@ def test_get_sandbox(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.get", + test_method="sandboxes.get", ) diff --git a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_session.py b/tests/unit/agentplatform/genai/replays/test_get_agent_engine_session.py deleted file mode 100644 index 7d3970f266..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_session.py +++ /dev/null @@ -1,39 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types - - -def test_get_session(client): - agent_engine = client.agent_engines.create() - operation = client.agent_engines.create_session( - name=agent_engine.api_resource.name, - user_id="test-user-123", - ) - assert isinstance(operation, types.AgentEngineSessionOperation) - session = client.agent_engines.get_session( - name=operation.response.name, - ) - assert isinstance(session, types.Session) - assert session.name == operation.response.name - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), - test_method="agent_engines.get_session", -) diff --git a/tests/unit/agentplatform/genai/replays/test_get_feedback_entry.py b/tests/unit/agentplatform/genai/replays/test_get_feedback_entry.py index d73446cf6f..a59d4145c2 100644 --- a/tests/unit/agentplatform/genai/replays/test_get_feedback_entry.py +++ b/tests/unit/agentplatform/genai/replays/test_get_feedback_entry.py @@ -32,13 +32,13 @@ def test_get(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) try: operation = client.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_123", event_id="event_456", feedback_type=types.FeedbackType.THUMBS_UP, @@ -66,18 +66,18 @@ def test_get(client): assert feedback.custom_metadata == {"key1": "val1", "key2": "val2"} finally: # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) @pytest.mark.asyncio async def test_get_async(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) try: operation = await client.aio.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_123", event_id="event_456", feedback_type=types.FeedbackType.THUMBS_UP, @@ -105,6 +105,6 @@ async def test_get_async(client): assert feedback.custom_metadata == {"key1": "val1", "key2": "val2"} finally: # Clean up resources. - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_ingest_events_memory_bank.py b/tests/unit/agentplatform/genai/replays/test_ingest_events_memory_bank.py index ff7ad46583..da88063da5 100644 --- a/tests/unit/agentplatform/genai/replays/test_ingest_events_memory_bank.py +++ b/tests/unit/agentplatform/genai/replays/test_ingest_events_memory_bank.py @@ -18,17 +18,17 @@ def test_ingest_events(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() assert not list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name, + client.runtimes.memories.list( + name=runtime.api_resource.name, ) ) scope = {"user_id": "test-user-id"} # Generate memories using source content. This result is non-deterministic, # because an LLM is used to generate the memories. - client.agent_engines.memories.ingest_events( - name=agent_engine.api_resource.name, + client.runtimes.memories.ingest_events( + name=runtime.api_resource.name, scope=scope, direct_contents_source={ "events": [ @@ -48,8 +48,8 @@ def test_ingest_events(client): }, ) memories = list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, + client.runtimes.memories.retrieve( + name=runtime.api_resource.name, scope=scope, ) ) @@ -58,8 +58,8 @@ def test_ingest_events(client): # of inactivity. assert len(memories) == 0 - client.agent_engines.memories.ingest_events( - name=agent_engine.api_resource.name, + client.runtimes.memories.ingest_events( + name=runtime.api_resource.name, scope=scope, direct_contents_source={ "events": [ @@ -84,8 +84,8 @@ def test_ingest_events(client): }, ) memories = list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, + client.runtimes.memories.retrieve( + name=runtime.api_resource.name, scope=scope, simple_retrieval_params={ "page_size": 1, @@ -101,18 +101,18 @@ def test_ingest_events(client): # The user-provided `revision_labels` are applied to the generated memory's # revision (not the Memory itself), so list the memory's revisions to verify. revisions = list( - client.agent_engines.memories.revisions.list( + client.runtimes.memories.revisions.list( name=memories[0].memory.name, ) ) assert revisions assert revisions[0].labels == {"source": "ingest-events-test"} - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.ingest_events", + test_method="runtimes.memories.ingest_events", ) diff --git a/tests/unit/agentplatform/genai/replays/test_list_ae_runtime_revisions.py b/tests/unit/agentplatform/genai/replays/test_list_ae_runtime_revisions.py index 49f9251479..88474c58d5 100644 --- a/tests/unit/agentplatform/genai/replays/test_list_ae_runtime_revisions.py +++ b/tests/unit/agentplatform/genai/replays/test_list_ae_runtime_revisions.py @@ -26,8 +26,8 @@ def test_list_runtime_revisions( client, - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): client._api_client._http_options.base_url = ( "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" @@ -35,10 +35,10 @@ def test_list_runtime_revisions( client._api_client._http_options.api_version = "v1beta1" with ( - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "display_name": "test-agent-engine-list-runtime-revisions", "source_packages": [ @@ -55,26 +55,26 @@ def test_list_runtime_revisions( }, ) assert ( - agent_engine.api_resource.display_name + runtime.api_resource.display_name == "test-agent-engine-list-runtime-revisions" ) - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( - name=agent_engine.api_resource.name, + runtime_revisions_iter = client.runtimes.revisions.list( + name=runtime.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) assert len(runtime_revisions_list) == 1 - assert isinstance(runtime_revisions_list[0], types.AgentEngineRuntimeRevision) + assert isinstance(runtime_revisions_list[0], types.RuntimeRevision) assert isinstance( runtime_revisions_list[0].api_resource, types.ReasoningEngineRuntimeRevision ) # Clean up resources. - agent_engine.delete(force=True) + runtime.delete(force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.runtimes.revisions.list", + test_method="runtimes.revisions.list", ) pytest_plugins = ("pytest_asyncio",) @@ -83,8 +83,8 @@ def test_list_runtime_revisions( @pytest.mark.asyncio async def test_async_list_runtime_revisions( client, - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): client._api_client._http_options.base_url = ( "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" @@ -92,10 +92,10 @@ async def test_async_list_runtime_revisions( client._api_client._http_options.api_version = "v1beta1" with ( - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "display_name": "test-agent-engine-list-runtime-revisions", "source_packages": [ @@ -112,21 +112,21 @@ async def test_async_list_runtime_revisions( }, ) assert ( - agent_engine.api_resource.display_name + runtime.api_resource.display_name == "test-agent-engine-list-runtime-revisions" ) - runtime_revisions_iter = client.aio.agent_engines.runtimes.revisions.list( - name=agent_engine.api_resource.name, + runtime_revisions_iter = client.aio.runtimes.revisions.list( + name=runtime.api_resource.name, ) runtime_revisions_list = [] async for revision in runtime_revisions_iter: runtime_revisions_list.append(revision) assert len(runtime_revisions_list) == 1 - assert isinstance(runtime_revisions_list[0], types.AgentEngineRuntimeRevision) + assert isinstance(runtime_revisions_list[0], types.RuntimeRevision) assert isinstance( runtime_revisions_list[0].api_resource, types.ReasoningEngineRuntimeRevision ) # Clean up resources. - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_task_events.py b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_task_events.py deleted file mode 100644 index 58b26bc3c8..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_task_events.py +++ /dev/null @@ -1,160 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -import pytest - - -def test_list_simple_a2a_task_events(client): - # Use the autopush environment. - client._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client._api_client._http_options.api_version = "internal" - task = client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task999", - config=types.CreateAgentEngineTaskConfig(context_id="context999"), - ) - assert isinstance(task, types.A2aTask) - - events = list( - client.agent_engines.a2a_tasks.events.list( - name=task.name, - ) - ) - assert len(events) == 1 - assert events[0].event_data.state_change.new_state == "SUBMITTED" - - client.agent_engines.a2a_tasks.events.append( - name=task.name, - task_events=[ - types.TaskEvent( - event_data=types.TaskEventData( - metadata_change=types.TaskMetadataChange( - new_metadata={"key1": "value1"} - ) - ), - event_sequence_number=1, - ), - types.TaskEvent( - event_data=types.TaskEventData( - metadata_change=types.TaskMetadataChange( - new_metadata={"key2": "value2"} - ) - ), - event_sequence_number=2, - ), - ], - ) - - result = list( - client.agent_engines.a2a_tasks.events.list( - name=task.name, - config=types.ListAgentEngineTaskEventsConfig( - order_by="create_time desc", - ), - ) - ) - - assert len(result) == 3 - assert result[0].event_sequence_number == 2 - assert result[0].event_data.metadata_change.new_metadata == {"key2": "value2"} - assert result[1].event_sequence_number == 1 - assert result[1].event_data.metadata_change.new_metadata == {"key1": "value1"} - assert result[2].event_data.state_change.new_state == "SUBMITTED" - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), -) - - -@pytest.mark.asyncio -async def test_list_simple_a2a_task_events_async(client): - # Use the autopush environment. - client.aio._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client.aio._api_client._http_options.api_version = "internal" - task = await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task999", - config=types.CreateAgentEngineTaskConfig(context_id="context999"), - ) - assert isinstance(task, types.A2aTask) - - events = list( - await client.aio.agent_engines.a2a_tasks.events.list( - name=task.name, - ) - ) - assert len(events) == 1 - assert events[0].event_data.state_change.new_state == "SUBMITTED" - - await client.aio.agent_engines.a2a_tasks.events.append( - name=task.name, - task_events=[ - types.TaskEvent( - event_data=types.TaskEventData( - metadata_change=types.TaskMetadataChange( - new_metadata={"key1": "value1"} - ) - ), - event_sequence_number=1, - ), - types.TaskEvent( - event_data=types.TaskEventData( - metadata_change=types.TaskMetadataChange( - new_metadata={"key2": "value2"} - ) - ), - event_sequence_number=2, - ), - ], - ) - - result = list( - await client.aio.agent_engines.a2a_tasks.events.list( - name=task.name, - config=types.ListAgentEngineTaskEventsConfig( - order_by="create_time desc", - ), - ) - ) - - assert len(result) == 3 - assert result[0].event_sequence_number == 2 - assert result[0].event_data.metadata_change.new_metadata == {"key2": "value2"} - assert result[1].event_sequence_number == 1 - assert result[1].event_data.metadata_change.new_metadata == {"key1": "value1"} - assert result[2].event_data.state_change.new_state == "SUBMITTED" - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_tasks.py b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_tasks.py deleted file mode 100644 index 1f580531b2..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_tasks.py +++ /dev/null @@ -1,117 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -import pytest - - -def test_list_a2a_tasks(client): - # Use the autopush environment. - client._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client._api_client._http_options.api_version = "internal" - - assert not list( - client.agent_engines.a2a_tasks.list( - name=agent_engine.api_resource.name, - ) - ) - client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task456", - config=types.CreateAgentEngineTaskConfig(context_id="context456"), - ) - a2a_tasks_list = client.agent_engines.a2a_tasks.list( - name=agent_engine.api_resource.name, - config=types.ListAgentEngineTasksConfig( - page_size=1, - order_by="create_time asc", - ), - ) - assert len(a2a_tasks_list) == 1 - assert isinstance(a2a_tasks_list[0], types.A2aTask) - assert a2a_tasks_list[0].name == ( - f"{agent_engine.api_resource.name}/a2aTasks/task123" - ) - assert a2a_tasks_list[0].context_id == "context123" - - # Clean up resources. - agent_engine.delete(force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), -) - - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_list_a2a_tasks_async(client): - # Use the autopush environment. - client.aio._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client.aio._api_client._http_options.api_version = "internal" - - assert not list( - await client.aio.agent_engines.a2a_tasks.list( - name=agent_engine.api_resource.name, - ) - ) - await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task456", - config=types.CreateAgentEngineTaskConfig(context_id="context456"), - ) - a2a_tasks_list = await client.aio.agent_engines.a2a_tasks.list( - name=agent_engine.api_resource.name, - config=types.ListAgentEngineTasksConfig( - page_size=1, - order_by="create_time asc", - ), - ) - assert len(a2a_tasks_list) == 1 - assert isinstance(a2a_tasks_list[0], types.A2aTask) - assert a2a_tasks_list[0].name == ( - f"{agent_engine.api_resource.name}/a2aTasks/task123" - ) - assert a2a_tasks_list[0].context_id == "context123" - - # Clean up resources. - agent_engine.delete(force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_memories.py b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_memories.py index 4373b4bf25..4e082e8491 100644 --- a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_memories.py +++ b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_memories.py @@ -21,31 +21,31 @@ def test_list_memories(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() assert not list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name, + client.runtimes.memories.list( + name=runtime.api_resource.name, ) ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, config={ "wait_for_completion": True, }, ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_2", scope={"user_id": "456"}, config={ "wait_for_completion": True, }, ) - memory_list = client.agent_engines.memories.list( - name=agent_engine.api_resource.name, - config=types.ListAgentEngineMemoryConfig( + memory_list = client.runtimes.memories.list( + name=runtime.api_resource.name, + config=types.ListRuntimeMemoryConfig( page_size=1, order_by="create_time asc", ), @@ -55,13 +55,13 @@ def test_list_memories(client): assert memory_list[0].fact == "memory_fact" assert memory_list[0].scope["user_id"] == "123" # Clean up resources. - agent_engine.delete(force=True) + runtime.delete(force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.list", + test_method="runtimes.memories.list", ) @@ -70,27 +70,27 @@ def test_list_memories(client): @pytest.mark.asyncio async def test_async_list_memories(client): - agent_engine = client.agent_engines.create() - pager = await client.aio.agent_engines.memories.list( - name=agent_engine.api_resource.name + runtime = client.runtimes.create() + pager = await client.aio.runtimes.memories.list( + name=runtime.api_resource.name ) assert not [item async for item in pager] - await client.aio.agent_engines.memories.create( - name=agent_engine.api_resource.name, + await client.aio.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_2", scope={"user_id": "456"}, config={ "wait_for_completion": True, }, ) - pager = await client.aio.agent_engines.memories.list( - name=agent_engine.api_resource.name + pager = await client.aio.runtimes.memories.list( + name=runtime.api_resource.name ) memory_list = [item async for item in pager] assert len(memory_list) == 1 assert isinstance(memory_list[0], types.Memory) - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_sandboxes.py b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_sandboxes.py index 0e141495d1..dae49d98dd 100644 --- a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_sandboxes.py +++ b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_sandboxes.py @@ -19,26 +19,26 @@ def test_list_sandboxes(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() assert not list( - client.agent_engines.sandboxes.list( - name=agent_engine.api_resource.name, + client.sandboxes.list( + name=runtime.api_resource.name, ) ) - operation = client.agent_engines.sandboxes.create( - name=agent_engine.api_resource.name, + operation = client.sandboxes.create( + name=runtime.api_resource.name, spec={ "code_execution_environment": { "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB" } }, - config=types.CreateAgentEngineSandboxConfig(display_name="test_sandbox"), + config=types.CreateRuntimeSandboxConfig(display_name="test_sandbox"), ) - assert isinstance(operation, types.AgentEngineSandboxOperation) + assert isinstance(operation, types.RuntimeSandboxOperation) - sandbox_list = client.agent_engines.sandboxes.list( - name=agent_engine.api_resource.name, + sandbox_list = client.sandboxes.list( + name=runtime.api_resource.name, ) assert len(sandbox_list) == 1 assert isinstance(sandbox_list[0], types.SandboxEnvironment) @@ -48,5 +48,5 @@ def test_list_sandboxes(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.list", + test_method="sandboxes.list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_session_events.py b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_session_events.py index 616342dd78..355a2e54cf 100644 --- a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_session_events.py +++ b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_session_events.py @@ -22,18 +22,18 @@ def test_list_session_events(client): - agent_engine = client.agent_engines.create() - operation = client.agent_engines.sessions.create( - name=agent_engine.api_resource.name, + runtime = client.runtimes.create() + operation = client.sessions.create( + name=runtime.api_resource.name, user_id="test-user-123", ) session = operation.response assert not list( - client.agent_engines.sessions.events.list( + client.sessions.events.list( name=session.name, ) ) - client.agent_engines.sessions.events.append( + client.sessions.events.append( name=session.name, author="test-user-123", invocation_id="test-invocation-id", @@ -44,7 +44,7 @@ def test_list_session_events(client): }, }, ) - session_event_list = client.agent_engines.sessions.events.list( + session_event_list = client.sessions.events.list( name=session.name, ) assert len(session_event_list) == 1 @@ -55,7 +55,7 @@ def test_list_session_events(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions.events.list", + test_method="sessions.events.list", ) @@ -64,24 +64,24 @@ def test_list_session_events(client): @pytest.mark.asyncio async def test_async_list_session_events(client): - agent_engine = client.agent_engines.create() - operation = await client.aio.agent_engines.sessions.create( - name=agent_engine.api_resource.name, + runtime = client.runtimes.create() + operation = await client.aio.sessions.create( + name=runtime.api_resource.name, user_id="test-user-123", ) session = operation.response - pager = await client.aio.agent_engines.sessions.events.list(name=session.name) + pager = await client.aio.sessions.events.list(name=session.name) assert not [item async for item in pager] - await client.aio.agent_engines.sessions.events.append( + await client.aio.sessions.events.append( name=session.name, author="test-user-123", invocation_id="test-invocation-id", timestamp=datetime.datetime.fromtimestamp(1234567890, tz=datetime.timezone.utc), ) - pager = await client.aio.agent_engines.sessions.events.list(name=session.name) + pager = await client.aio.sessions.events.list(name=session.name) session_event_list = [item async for item in pager] assert len(session_event_list) == 1 assert isinstance(session_event_list[0], types.SessionEvent) - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_sessions.py b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_sessions.py index e34cca3c60..f20e63197c 100644 --- a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_sessions.py +++ b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_sessions.py @@ -22,29 +22,29 @@ def test_list_sessions(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() assert not list( - client.agent_engines.sessions.list( - name=agent_engine.api_resource.name, + client.sessions.list( + name=runtime.api_resource.name, ) ) - client.agent_engines.sessions.create( - name=agent_engine.api_resource.name, + client.sessions.create( + name=runtime.api_resource.name, user_id="test-user-123", ) - session_list = client.agent_engines.sessions.list( - name=agent_engine.api_resource.name, + session_list = client.sessions.list( + name=runtime.api_resource.name, ) assert len(session_list) == 1 assert isinstance(session_list[0], types.Session) - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions.list", + test_method="sessions.list", ) pytest_plugins = ("pytest_asyncio",) @@ -52,21 +52,21 @@ def test_list_sessions(client): @pytest.mark.asyncio async def test_async_list_sessions(client): - agent_engine = client.agent_engines.create() - pager = await client.aio.agent_engines.sessions.list( - name=agent_engine.api_resource.name + runtime = client.runtimes.create() + pager = await client.aio.sessions.list( + name=runtime.api_resource.name ) assert not [item async for item in pager] - await client.aio.agent_engines.sessions.create( - name=agent_engine.api_resource.name, + await client.aio.sessions.create( + name=runtime.api_resource.name, user_id="test-user-123", ) - pager = await client.aio.agent_engines.sessions.list( - name=agent_engine.api_resource.name, + pager = await client.aio.sessions.list( + name=runtime.api_resource.name, ) session_list = [item async for item in pager] assert len(session_list) == 1 assert isinstance(session_list[0], types.Session) - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_list_feedback_entries.py b/tests/unit/agentplatform/genai/replays/test_list_feedback_entries.py index e46b38bd58..875d696979 100644 --- a/tests/unit/agentplatform/genai/replays/test_list_feedback_entries.py +++ b/tests/unit/agentplatform/genai/replays/test_list_feedback_entries.py @@ -32,14 +32,14 @@ def test_list(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) try: # Create THUMBS_UP feedback entry. operation_up = client.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_123", event_id="event_456", feedback_type=types.FeedbackType.THUMBS_UP, @@ -56,7 +56,7 @@ def test_list(client): # Create THUMBS_DOWN feedback entry. operation_down = client.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_abc", event_id="event_xyz", feedback_type=types.FeedbackType.THUMBS_DOWN, @@ -72,7 +72,7 @@ def test_list(client): assert operation_down.done # List and verify the feedback entries. - response = client.feedback_entries.list(parent=agent_engine.api_resource.name) + response = client.feedback_entries.list(parent=runtime.api_resource.name) feedback_entries = list(response) assert len(feedback_entries) == 2 @@ -104,22 +104,22 @@ def test_list(client): assert thumbs_down_entry.custom_metadata == {"key_a": "val_a"} finally: # Clean up resources. - client.agent_engines.delete( - name=agent_engine.api_resource.name, + client.runtimes.delete( + name=runtime.api_resource.name, force=True, ) @pytest.mark.asyncio async def test_list_async(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) try: # Create THUMBS_UP feedback entry. operation_up = await client.aio.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_123", event_id="event_456", feedback_type=types.FeedbackType.THUMBS_UP, @@ -136,7 +136,7 @@ async def test_list_async(client): # Create THUMBS_DOWN feedback entry. operation_down = await client.aio.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_abc", event_id="event_xyz", feedback_type=types.FeedbackType.THUMBS_DOWN, @@ -153,7 +153,7 @@ async def test_list_async(client): # List and verify the feedback entries. response = await client.aio.feedback_entries.list( - parent=agent_engine.api_resource.name + parent=runtime.api_resource.name ) feedback_entries = [f async for f in response] assert len(feedback_entries) == 2 @@ -186,7 +186,7 @@ async def test_list_async(client): assert thumbs_down_entry.custom_metadata == {"key_a": "val_a"} finally: # Clean up resources. - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True, ) diff --git a/tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py b/tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py index e123597953..c4633b7b9f 100644 --- a/tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py +++ b/tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py @@ -22,28 +22,28 @@ def test_purge_memories(client): """Tests purging memories.""" - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() try: - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_1", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_2", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_3", scope={"user_id": "456"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_4", scope={"user_id": "456"}, config={ @@ -51,8 +51,8 @@ def test_purge_memories(client): "metadata": {"my_key": {"string_value": "my_value"}}, }, ) - operation = client.agent_engines.memories.purge( - name=agent_engine.api_resource.name, + operation = client.runtimes.memories.purge( + name=runtime.api_resource.name, filter="scope.user_id=123", config={"wait_for_completion": True}, ) @@ -62,16 +62,16 @@ def test_purge_memories(client): assert ( len( list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name + client.runtimes.memories.list( + name=runtime.api_resource.name ) ) ) == 4 ) # Now, actually purge the memories. - operation = client.agent_engines.memories.purge( - name=agent_engine.api_resource.name, + operation = client.runtimes.memories.purge( + name=runtime.api_resource.name, filter="scope.user_id=123", force=True, config={"wait_for_completion": True}, @@ -81,16 +81,16 @@ def test_purge_memories(client): assert ( len( list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name + client.runtimes.memories.list( + name=runtime.api_resource.name ) ) ) == 2 ) # Purge memories using filter groups. - operation = client.agent_engines.memories.purge( - name=agent_engine.api_resource.name, + operation = client.runtimes.memories.purge( + name=runtime.api_resource.name, force=True, filter_groups=[ {"filters": [{"key": "my_key", "value": {"string_value": "my_value"}}]} @@ -104,21 +104,21 @@ def test_purge_memories(client): assert ( len( list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name + client.runtimes.memories.list( + name=runtime.api_resource.name ) ) ) == 1 ) finally: - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.purge", + test_method="runtimes.memories.purge", ) @@ -127,29 +127,29 @@ def test_purge_memories(client): @pytest.mark.asyncio async def test_purge_memories_async(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() try: - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_1", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_2", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_3", scope={"user_id": "456"}, config={"wait_for_completion": True}, ) - operation = await client.aio.agent_engines.memories.purge( - name=agent_engine.api_resource.name, + operation = await client.aio.runtimes.memories.purge( + name=runtime.api_resource.name, filter="scope.user_id=123", config={"wait_for_completion": True}, ) @@ -159,16 +159,16 @@ async def test_purge_memories_async(client): assert ( len( list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name + client.runtimes.memories.list( + name=runtime.api_resource.name ) ) ) == 3 ) # Now, actually purge the memories. - operation = await client.aio.agent_engines.memories.purge( - name=agent_engine.api_resource.name, + operation = await client.aio.runtimes.memories.purge( + name=runtime.api_resource.name, filter="scope.user_id=123", force=True, config={"wait_for_completion": True}, @@ -178,12 +178,12 @@ async def test_purge_memories_async(client): assert ( len( list( - client.agent_engines.memories.list( - name=agent_engine.api_resource.name + client.runtimes.memories.list( + name=runtime.api_resource.name ) ) ) == 1 ) finally: - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_retrieve_agent_engine_memories.py b/tests/unit/agentplatform/genai/replays/test_retrieve_agent_engine_memories.py index 02a3e034d0..d8c3981a2e 100644 --- a/tests/unit/agentplatform/genai/replays/test_retrieve_agent_engine_memories.py +++ b/tests/unit/agentplatform/genai/replays/test_retrieve_agent_engine_memories.py @@ -24,26 +24,26 @@ def test_retrieve_memories_with_similarity_search_params(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() assert not list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, + client.runtimes.memories.retrieve( + name=runtime.api_resource.name, scope={"user_id": "123"}, similarity_search_params=types.RetrieveMemoriesRequestSimilaritySearchParams( search_query="memory_fact_1", ), ) ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_1", scope={"user_id": "123"}, ) assert ( len( list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, + client.runtimes.memories.retrieve( + name=runtime.api_resource.name, scope={"user_id": "123"}, ) ) @@ -51,21 +51,21 @@ def test_retrieve_memories_with_similarity_search_params(client): == 1 ) assert not list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, + client.runtimes.memories.retrieve( + name=runtime.api_resource.name, scope={"user_id": "456"}, ) ) - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_2", scope={"user_id": "123"}, ) assert ( len( list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, + client.runtimes.memories.retrieve( + name=runtime.api_resource.name, scope={"user_id": "123"}, ) ) @@ -73,18 +73,18 @@ def test_retrieve_memories_with_similarity_search_params(client): == 2 ) # Clean up resources. - agent_engine.delete(force=True) + runtime.delete(force=True) def test_retrieve_memories_with_simple_retrieval_params(client): - agent_engine = client.agent_engines.create() - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + runtime = client.runtimes.create() + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_1", scope={"user_id": "123"}, ) - memories = client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, + memories = client.runtimes.memories.retrieve( + name=runtime.api_resource.name, scope={"user_id": "123"}, simple_retrieval_params=types.RetrieveMemoriesRequestSimpleRetrievalParams( page_size=1, @@ -94,18 +94,18 @@ def test_retrieve_memories_with_simple_retrieval_params(client): assert isinstance(memories.page[0], types.RetrieveMemoriesResponseRetrievedMemory) assert memories.page_size == 1 - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_2", scope={"user_id": "123"}, ) - memories = client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, scope={"user_id": "123"} + memories = client.runtimes.memories.retrieve( + name=runtime.api_resource.name, scope={"user_id": "123"} ) assert memories.page_size == 2 - memories = client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, + memories = client.runtimes.memories.retrieve( + name=runtime.api_resource.name, scope={"user_id": "123"}, config={"filter": 'fact="memory_fact_2"'}, ) @@ -113,11 +113,11 @@ def test_retrieve_memories_with_simple_retrieval_params(client): assert memories.page[0].memory.fact == "memory_fact_2" # Clean up resources. - agent_engine.delete(force=True) + runtime.delete(force=True) def test_retrieve_memories_with_metadata(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() metadata = { "my_string_key": types.MemoryMetadataValue(string_value="my_string_value"), "my_double_key": types.MemoryMetadataValue(double_value=123.456), @@ -129,21 +129,21 @@ def test_retrieve_memories_with_metadata(client): ), } scope = {"user_id": "123"} - client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_1", scope=scope, ) - operation = client.agent_engines.memories.create( - name=agent_engine.api_resource.name, + operation = client.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact_2", scope=scope, config={"metadata": metadata}, ) memory_name2 = operation.response.name - results = client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, + results = client.runtimes.memories.retrieve( + name=runtime.api_resource.name, scope=scope, config={ "filter_groups": [ @@ -162,13 +162,13 @@ def test_retrieve_memories_with_metadata(client): assert results[0].memory.name == memory_name2 # Clean up resources. - agent_engine.delete(force=True) + runtime.delete(force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.retrieve", + test_method="runtimes.memories.retrieve", ) @@ -177,20 +177,20 @@ def test_retrieve_memories_with_metadata(client): @pytest.mark.asyncio async def test_retrieve_memories_async(client): - agent_engine = client.agent_engines.create() - operation = await client.aio.agent_engines.memories.create( - name=agent_engine.api_resource.name, + runtime = client.runtimes.create() + operation = await client.aio.runtimes.memories.create( + name=runtime.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, ) - assert isinstance(operation, types.AgentEngineMemoryOperation) - pager = await client.aio.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, + assert isinstance(operation, types.RuntimeMemoryOperation) + pager = await client.aio.runtimes.memories.retrieve( + name=runtime.api_resource.name, scope={"user_id": "123"}, ) memories = [item async for item in pager] assert len(memories) == 1 assert isinstance(memories[0], types.RetrieveMemoriesResponseRetrievedMemory) - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, force=True + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_run_inference.py b/tests/unit/agentplatform/genai/replays/test_run_inference.py index 0c4afb9eaf..e29ab24fbc 100644 --- a/tests/unit/agentplatform/genai/replays/test_run_inference.py +++ b/tests/unit/agentplatform/genai/replays/test_run_inference.py @@ -97,7 +97,7 @@ def test_inference_with_eval_cases_multi_turn_agent_data(client): def test_inference_with_eval_cases_agent_engine_agent_data(client): """Tests N+1 inference with agent_data via remote Agent Engine.""" - agent_engine = client.agent_engines.get( + runtime = client.runtimes.get( name="projects/977012026409/locations/us-central1" "/reasoningEngines/7188347537655332864" ) @@ -144,7 +144,7 @@ def test_inference_with_eval_cases_agent_engine_agent_data(client): eval_dataset = types.EvaluationDataset(eval_cases=[eval_case]) inference_result = client.evals.run_inference( - agent=agent_engine, + agent=runtime, src=eval_dataset, ) assert isinstance(inference_result, types.EvaluationDataset) @@ -325,51 +325,6 @@ def test_inference_with_gemini_agent(client): assert result_df["interaction_id"].iloc[0] -def test_inference_with_gemini_agent_user_simulation(client): - """Tests multi-turn user simulation against a Gemini Agents API agent. - - An ADK LlmBackedUserSimulator generates the user turns while the - conversation is driven client-side through the Interactions API, chaining - turns via previous_interaction_id. The result is a multi-turn agent_data - trace per scenario. - """ - import pandas as pd - - prompts_df = pd.DataFrame( - { - "starting_prompt": [ - "I'm planning a weekend trip. Can you help me pick a destination?" - ], - "conversation_plan": [ - "Step 1: User asks for help planning a weekend trip. " - "Step 2: User asks for a cheaper alternative destination. " - "Step 3: User asks the assistant to summarize the final plan." - ], - } - ) - - eval_dataset = client.evals.run_inference( - src=prompts_df, - agent="projects/model-evaluation-dev/locations/global/agents/test-agent-eval", - config=types.EvalRunInferenceConfig( - user_simulator_config=types.evals.UserSimulatorConfig( - model_name="gemini-2.5-flash", - max_turn=3, - ) - ), - ) - - assert isinstance(eval_dataset, types.EvaluationDataset) - result_df = eval_dataset.eval_dataset_df - assert result_df is not None - assert set(["interaction_id", "agent_data"]).issubset(set(result_df.columns)) - assert len(result_df) == 1 - assert result_df["interaction_id"].iloc[0] - agent_data = result_df["agent_data"].iloc[0] - assert agent_data - assert len(agent_data["turns"]) >= 2 - - pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), diff --git a/tests/unit/agentplatform/genai/replays/test_structured_memories.py b/tests/unit/agentplatform/genai/replays/test_structured_memories.py index 5d1c2a39a5..7452637ad2 100644 --- a/tests/unit/agentplatform/genai/replays/test_structured_memories.py +++ b/tests/unit/agentplatform/genai/replays/test_structured_memories.py @@ -44,7 +44,7 @@ def test_generate_and_retrieve_profile(client): structured_memory_config_obj = types.StructuredMemoryConfig( **structured_memory_config ) - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "context_spec": { "memory_bank_config": { @@ -56,8 +56,8 @@ def test_generate_and_retrieve_profile(client): }, ) try: - agent_engine = client.agent_engines.get(name=agent_engine.api_resource.name) - memory_bank_config = agent_engine.api_resource.context_spec.memory_bank_config + runtime = client.runtimes.get(name=runtime.api_resource.name) + memory_bank_config = runtime.api_resource.context_spec.memory_bank_config assert memory_bank_config.customization_configs == [ memory_bank_customization_config ] @@ -66,16 +66,16 @@ def test_generate_and_retrieve_profile(client): ] scope = {"user_id": "123"} - client.agent_engines.memories.generate( - name=agent_engine.api_resource.name, + client.runtimes.memories.generate( + name=runtime.api_resource.name, scope=scope, direct_contents_source={ "events": [{"content": {"parts": [{"text": "My name is Kim."}]}}] }, ) memories = list( - client.agent_engines.memories.retrieve( - name=agent_engine.api_resource.name, + client.runtimes.memories.retrieve( + name=runtime.api_resource.name, scope=scope, config={"memory_types": ["STRUCTURED_PROFILE"]}, ) @@ -83,18 +83,18 @@ def test_generate_and_retrieve_profile(client): assert len(memories) >= 1 assert memories[0].memory.structured_content is not None - response = client.agent_engines.memories.retrieve_profiles( - name=agent_engine.api_resource.name, scope=scope + response = client.runtimes.memories.retrieve_profiles( + name=runtime.api_resource.name, scope=scope ) assert len(response.profiles) == 1 finally: # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=runtime.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.retrieve_profiles", + test_method="runtimes.retrieve_profiles", ) diff --git a/tests/unit/agentplatform/genai/replays/test_update_agent_engine.py b/tests/unit/agentplatform/genai/replays/test_update_agent_engine.py index 1829e79a02..4525e38c62 100644 --- a/tests/unit/agentplatform/genai/replays/test_update_agent_engine.py +++ b/tests/unit/agentplatform/genai/replays/test_update_agent_engine.py @@ -20,25 +20,25 @@ def test_agent_engines_update(client): - agent_engine = client.agent_engines.create() - assert agent_engine.api_resource.display_name is None + runtime = client.runtimes.create() + assert runtime.api_resource.display_name is None - updated_agent_engine = client.agent_engines.update( - name=agent_engine.api_resource.name, - config=types.AgentEngineConfig( + updated_runtime = client.runtimes.update( + name=runtime.api_resource.name, + config=types.AgentRuntimeConfig( display_name="updated_display_name", description="updated description", ), ) - assert isinstance(updated_agent_engine, types.AgentEngine) - assert updated_agent_engine.api_resource.name == agent_engine.api_resource.name + assert isinstance(updated_runtime, types.Runtime) + assert updated_runtime.api_resource.name == runtime.api_resource.name - assert updated_agent_engine.api_resource.display_name == "updated_display_name" - assert updated_agent_engine.api_resource.description == "updated description" + assert updated_runtime.api_resource.display_name == "updated_display_name" + assert updated_runtime.api_resource.description == "updated description" pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.update", + test_method="runtimes.update", ) diff --git a/tests/unit/agentplatform/genai/replays/test_update_feedback_context.py b/tests/unit/agentplatform/genai/replays/test_update_feedback_context.py index 2b19237aaa..e93eb72935 100644 --- a/tests/unit/agentplatform/genai/replays/test_update_feedback_context.py +++ b/tests/unit/agentplatform/genai/replays/test_update_feedback_context.py @@ -33,11 +33,11 @@ def test_update_and_get(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() try: operation = client.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_123", event_id="event_456", feedback_type=types.FeedbackType.THUMBS_UP, @@ -110,19 +110,19 @@ def test_update_and_get(client): finally: # Clean up resources. - client.agent_engines.delete( - name=agent_engine.api_resource.name, + client.runtimes.delete( + name=runtime.api_resource.name, force=True, ) @pytest.mark.asyncio async def test_update_and_get_async(client): - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() try: operation = await client.aio.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_123", event_id="event_456", feedback_type=types.FeedbackType.THUMBS_UP, @@ -195,7 +195,7 @@ async def test_update_and_get_async(client): finally: # Clean up resources. - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True, ) diff --git a/tests/unit/agentplatform/genai/replays/test_update_feedback_entry.py b/tests/unit/agentplatform/genai/replays/test_update_feedback_entry.py index b532b786f4..8574b19d27 100644 --- a/tests/unit/agentplatform/genai/replays/test_update_feedback_entry.py +++ b/tests/unit/agentplatform/genai/replays/test_update_feedback_entry.py @@ -32,13 +32,13 @@ def test_update(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) try: operation = client.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_123", event_id="event_456", feedback_type=types.FeedbackType.THUMBS_UP, @@ -113,21 +113,21 @@ def test_update(client): finally: # Clean up resources. - client.agent_engines.delete( - name=agent_engine.api_resource.name, + client.runtimes.delete( + name=runtime.api_resource.name, force=True, ) @pytest.mark.asyncio async def test_update_async(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) + runtime = client.runtimes.create() + assert isinstance(runtime, types.Runtime) + assert isinstance(runtime.api_resource, types.ReasoningEngine) try: operation = await client.aio.feedback_entries.create( - name=agent_engine.api_resource.name, + name=runtime.api_resource.name, session_id="session_123", event_id="event_456", feedback_type=types.FeedbackType.THUMBS_UP, @@ -202,7 +202,7 @@ async def test_update_async(client): finally: # Clean up resources. - await client.aio.agent_engines.delete( - name=agent_engine.api_resource.name, + await client.aio.runtimes.delete( + name=runtime.api_resource.name, force=True, ) diff --git a/tests/unit/agentplatform/genai/replays/test_update_traffic_agent_engine.py b/tests/unit/agentplatform/genai/replays/test_update_traffic_agent_engine.py index 059fcc557b..822b47821b 100644 --- a/tests/unit/agentplatform/genai/replays/test_update_traffic_agent_engine.py +++ b/tests/unit/agentplatform/genai/replays/test_update_traffic_agent_engine.py @@ -30,28 +30,28 @@ def test_agent_engines_update_traffic_to_always_latest(client): ) client._api_client._http_options.api_version = "v1beta1" - agent_engine = client.agent_engines.create() + runtime = client.runtimes.create() traffic_config = types.ReasoningEngineTrafficConfig( traffic_split_always_latest=types.ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest(), ) - updated_agent_engine = client.agent_engines.update( - name=agent_engine.api_resource.name, - config=types.AgentEngineConfig( + updated_runtime = client.runtimes.update( + name=runtime.api_resource.name, + config=types.AgentRuntimeConfig( traffic_config=traffic_config, ), ) - assert updated_agent_engine.api_resource.traffic_config == traffic_config + assert updated_runtime.api_resource.traffic_config == traffic_config - agent_engine.delete(force=True) + runtime.delete(force=True) def test_agent_engines_update_traffic_to_manual_split( client, - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): client._api_client._http_options.base_url = ( @@ -60,10 +60,10 @@ def test_agent_engines_update_traffic_to_manual_split( client._api_client._http_options.api_version = "v1beta1" with ( - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "display_name": "test-agent-engine-update-traffic-to-manual-split", "source_packages": [ @@ -80,8 +80,8 @@ def test_agent_engines_update_traffic_to_manual_split( }, ) - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( - name=agent_engine.api_resource.name, + runtime_revisions_iter = client.runtimes.revisions.list( + name=runtime.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) assert len(runtime_revisions_list) == 1 @@ -101,21 +101,21 @@ def test_agent_engines_update_traffic_to_manual_split( ), ) - updated_agent_engine = client.agent_engines.update( - name=agent_engine.api_resource.name, - config=types.AgentEngineConfig( + updated_runtime = client.runtimes.update( + name=runtime.api_resource.name, + config=types.AgentRuntimeConfig( traffic_config=traffic_config, ), ) - assert updated_agent_engine.api_resource.traffic_config == traffic_config - agent_engine.delete(force=True) + assert updated_runtime.api_resource.traffic_config == traffic_config + runtime.delete(force=True) def test_agent_engines_update_traffic_with_agent_update( client, - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): client._api_client._http_options.base_url = ( @@ -124,10 +124,10 @@ def test_agent_engines_update_traffic_with_agent_update( client._api_client._http_options.api_version = "v1beta1" with ( - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - agent_engine = client.agent_engines.create( + runtime = client.runtimes.create( config={ "display_name": "test-agent-engine-update-traffic-with-agent-before-update", "source_packages": [ @@ -144,12 +144,12 @@ def test_agent_engines_update_traffic_with_agent_update( }, ) assert ( - agent_engine.api_resource.display_name + runtime.api_resource.display_name == "test-agent-engine-update-traffic-with-agent-before-update" ) - assert agent_engine.api_resource.traffic_config is None - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( - name=agent_engine.api_resource.name, + assert runtime.api_resource.traffic_config is None + runtime_revisions_iter = client.runtimes.revisions.list( + name=runtime.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) assert len(runtime_revisions_list) == 1 @@ -170,11 +170,11 @@ def test_agent_engines_update_traffic_with_agent_update( ) with ( - mock_agent_engine_create_base64_encoded_tarball, - mock_agent_engine_create_path_exists, + mock_runtime_create_base64_encoded_tarball, + mock_runtime_create_path_exists, ): - updated_agent_engine = client.agent_engines.update( - name=agent_engine.api_resource.name, + updated_runtime = client.runtimes.update( + name=runtime.api_resource.name, config={ "display_name": "test-agent-engine-update-traffic-with-agent-after-update", "source_packages": [ @@ -193,12 +193,12 @@ def test_agent_engines_update_traffic_with_agent_update( ) assert ( - updated_agent_engine.api_resource.display_name + updated_runtime.api_resource.display_name == "test-agent-engine-update-traffic-with-agent-after-update" ) - assert updated_agent_engine.api_resource.traffic_config == traffic_config - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( - name=agent_engine.api_resource.name, + assert updated_runtime.api_resource.traffic_config == traffic_config + runtime_revisions_iter = client.runtimes.revisions.list( + name=runtime.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) assert len(runtime_revisions_list) == 2 @@ -213,11 +213,11 @@ def test_agent_engines_update_traffic_with_agent_update( ) # new revision assert runtime_revisions_list[1].api_resource.name == runtime_revision_name - agent_engine.delete(force=True) + runtime.delete(force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.update", + test_method="runtimes.update", ) diff --git a/tests/unit/agentplatform/genai/test_agent_engine_runtime_revisions.py b/tests/unit/agentplatform/genai/test_agent_engine_runtime_revisions.py index a2e3b5798e..fcb9add6b1 100644 --- a/tests/unit/agentplatform/genai/test_agent_engine_runtime_revisions.py +++ b/tests/unit/agentplatform/genai/test_agent_engine_runtime_revisions.py @@ -27,7 +27,7 @@ from google.cloud import aiplatform import agentplatform from google.cloud.aiplatform import initializer -from agentplatform._genai import _agent_engines_utils +from agentplatform._genai import _runtimes_utils from agentplatform._genai import runtime_revisions from agentplatform._genai import types as _genai_types from google.genai import types as genai_types @@ -298,7 +298,7 @@ def clone(self): return self def register_operations(self) -> Dict[str, List[str]]: - # Registered method `missing_method` is not a method of the AgentEngine. + # Registered method `missing_method` is not a method of the Runtime. return { _TEST_STANDARD_API_MODE: [ _TEST_DEFAULT_METHOD_NAME, @@ -316,7 +316,7 @@ def method_to_be_unregistered(self, unused_arbitrary_string_name: str) -> str: return unused_arbitrary_string_name.upper() def register_operations(self) -> Dict[str, List[str]]: - # Registered method `missing_method` is not a method of the AgentEngine. + # Registered method `missing_method` is not a method of the Runtime. return {_TEST_STANDARD_API_MODE: [_TEST_METHOD_TO_BE_UNREGISTERED_NAME]} @@ -340,30 +340,30 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_AGENT_ENGINE_DISPLAY_NAME = "Agent Engine Display Name" _TEST_AGENT_ENGINE_DESCRIPTION = "Agent Engine Description" _TEST_LIST_FILTER = f'name="{_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME}"' -_TEST_GCS_DIR_NAME = _agent_engines_utils._DEFAULT_GCS_DIR_NAME -_TEST_BLOB_FILENAME = _agent_engines_utils._BLOB_FILENAME -_TEST_REQUIREMENTS_FILE = _agent_engines_utils._REQUIREMENTS_FILE -_TEST_EXTRA_PACKAGES_FILE = _agent_engines_utils._EXTRA_PACKAGES_FILE -_TEST_STANDARD_API_MODE = _agent_engines_utils._STANDARD_API_MODE -_TEST_ASYNC_API_MODE = _agent_engines_utils._ASYNC_API_MODE -_TEST_STREAM_API_MODE = _agent_engines_utils._STREAM_API_MODE -_TEST_ASYNC_STREAM_API_MODE = _agent_engines_utils._ASYNC_STREAM_API_MODE -_TEST_BIDI_STREAM_API_MODE = _agent_engines_utils._BIDI_STREAM_API_MODE -_TEST_DEFAULT_METHOD_NAME = _agent_engines_utils._DEFAULT_METHOD_NAME -_TEST_DEFAULT_ASYNC_METHOD_NAME = _agent_engines_utils._DEFAULT_ASYNC_METHOD_NAME -_TEST_DEFAULT_STREAM_METHOD_NAME = _agent_engines_utils._DEFAULT_STREAM_METHOD_NAME +_TEST_GCS_DIR_NAME = _runtimes_utils._DEFAULT_GCS_DIR_NAME +_TEST_BLOB_FILENAME = _runtimes_utils._BLOB_FILENAME +_TEST_REQUIREMENTS_FILE = _runtimes_utils._REQUIREMENTS_FILE +_TEST_EXTRA_PACKAGES_FILE = _runtimes_utils._EXTRA_PACKAGES_FILE +_TEST_STANDARD_API_MODE = _runtimes_utils._STANDARD_API_MODE +_TEST_ASYNC_API_MODE = _runtimes_utils._ASYNC_API_MODE +_TEST_STREAM_API_MODE = _runtimes_utils._STREAM_API_MODE +_TEST_ASYNC_STREAM_API_MODE = _runtimes_utils._ASYNC_STREAM_API_MODE +_TEST_BIDI_STREAM_API_MODE = _runtimes_utils._BIDI_STREAM_API_MODE +_TEST_DEFAULT_METHOD_NAME = _runtimes_utils._DEFAULT_METHOD_NAME +_TEST_DEFAULT_ASYNC_METHOD_NAME = _runtimes_utils._DEFAULT_ASYNC_METHOD_NAME +_TEST_DEFAULT_STREAM_METHOD_NAME = _runtimes_utils._DEFAULT_STREAM_METHOD_NAME _TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME = ( - _agent_engines_utils._DEFAULT_ASYNC_STREAM_METHOD_NAME + _runtimes_utils._DEFAULT_ASYNC_STREAM_METHOD_NAME ) _TEST_DEFAULT_BIDI_STREAM_METHOD_NAME = ( - _agent_engines_utils._DEFAULT_BIDI_STREAM_METHOD_NAME + _runtimes_utils._DEFAULT_BIDI_STREAM_METHOD_NAME ) _TEST_CAPITALIZE_ENGINE_METHOD_DOCSTRING = "Runs the engine." _TEST_STREAM_METHOD_DOCSTRING = "Runs the stream engine." _TEST_ASYNC_STREAM_METHOD_DOCSTRING = "Runs the async stream engine." _TEST_BIDI_STREAM_METHOD_DOCSTRING = "Runs the bidi stream engine." -_TEST_MODE_KEY_IN_SCHEMA = _agent_engines_utils._MODE_KEY_IN_SCHEMA -_TEST_METHOD_NAME_KEY_IN_SCHEMA = _agent_engines_utils._METHOD_NAME_KEY_IN_SCHEMA +_TEST_MODE_KEY_IN_SCHEMA = _runtimes_utils._MODE_KEY_IN_SCHEMA +_TEST_METHOD_NAME_KEY_IN_SCHEMA = _runtimes_utils._METHOD_NAME_KEY_IN_SCHEMA _TEST_CUSTOM_METHOD_NAME = "custom_method" _TEST_CUSTOM_ASYNC_METHOD_NAME = "custom_async_method" _TEST_CUSTOM_STREAM_METHOD_NAME = "custom_stream_method" @@ -452,14 +452,14 @@ def register_operations(self) -> Dict[str, List[str]]: "lib", "main.py", ] -_TEST_AGENT_ENGINE_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_QUERY_SCHEMA = _runtimes_utils._generate_schema( CapitalizeEngine().query, schema_name=_TEST_DEFAULT_METHOD_NAME, ) _TEST_AGENT_ENGINE_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STANDARD_API_MODE _TEST_PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}" _TEST_PYTHON_VERSION_OVERRIDE = "3.11" -_TEST_AGENT_ENGINE_FRAMEWORK = _agent_engines_utils._DEFAULT_AGENT_FRAMEWORK +_TEST_AGENT_ENGINE_FRAMEWORK = _runtimes_utils._DEFAULT_AGENT_FRAMEWORK _TEST_AGENT_ENGINE_CLASS_METHOD_1 = { "description": "Runs the engine.", "name": "query", @@ -567,38 +567,38 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE = [{"output": "hello"}, {"output": "world"}] _TEST_AGENT_ENGINE_OPERATION_SCHEMAS = [] _TEST_AGENT_ENGINE_EXTRA_PACKAGE = "fake.py" -_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA = _runtimes_utils._generate_schema( AsyncQueryEngine().async_query, schema_name=_TEST_DEFAULT_ASYNC_METHOD_NAME, ) _TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_ASYNC_API_MODE -_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA = _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_method, schema_name=_TEST_CUSTOM_METHOD_NAME, ) _TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( _TEST_STANDARD_API_MODE ) -_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA = _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_async_method, schema_name=_TEST_CUSTOM_ASYNC_METHOD_NAME, ) _TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( _TEST_ASYNC_API_MODE ) -_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA = _runtimes_utils._generate_schema( StreamQueryEngine().stream_query, schema_name=_TEST_DEFAULT_STREAM_METHOD_NAME, ) _TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STREAM_API_MODE -_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA = _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_stream_method, schema_name=_TEST_CUSTOM_STREAM_METHOD_NAME, ) _TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( _TEST_STREAM_API_MODE ) -_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA = _runtimes_utils._generate_schema( AsyncStreamQueryEngine().async_stream_query, schema_name=_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME, ) @@ -606,7 +606,7 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_ASYNC_STREAM_API_MODE ) _TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA = ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_async_stream_method, schema_name=_TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME, ) @@ -614,7 +614,7 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( _TEST_ASYNC_STREAM_API_MODE ) -_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA = _runtimes_utils._generate_schema( OperationRegistrableEngine().bidi_stream_query, schema_name=_TEST_DEFAULT_BIDI_STREAM_METHOD_NAME, ) @@ -622,7 +622,7 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_BIDI_STREAM_API_MODE ) _TEST_AGENT_ENGINE_CUSTOM_BIDI_STREAM_QUERY_SCHEMA = ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_bidi_stream_method, schema_name=_TEST_CUSTOM_BIDI_STREAM_METHOD_NAME, ) @@ -652,7 +652,7 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_NO_OPERATION_REGISTRABLE_SCHEMAS = [ _TEST_AGENT_ENGINE_QUERY_SCHEMA, ] -_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA = _runtimes_utils._generate_schema( MethodToBeUnregisteredEngine().method_to_be_unregistered, schema_name=_TEST_METHOD_TO_BE_UNREGISTERED_NAME, ) @@ -844,7 +844,7 @@ def teardown_method(self): def test_get_delete_runtime_revision_operation(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request" + self.client.runtimes.revisions._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body=json.dumps( @@ -854,7 +854,7 @@ def test_get_delete_runtime_revision_operation(self): } ), ) - operation = self.client.agent_engines.runtimes.revisions._get_delete_runtime_revision_operation( + operation = self.client.runtimes.revisions._get_delete_runtime_revision_operation( operation_name=_TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, ) request_mock.assert_called_with( @@ -864,13 +864,13 @@ def test_get_delete_runtime_revision_operation(self): None, ) assert isinstance( - operation, _genai_types.DeleteAgentEngineRuntimeRevisionOperation + operation, _genai_types.DeleteRuntimeRevisionOperation ) assert operation.done def test_await_operation(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request" + self.client.runtimes.revisions._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body=json.dumps( @@ -880,9 +880,9 @@ def test_await_operation(self): } ), ) - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=_TEST_AGENT_ENGINE_REVISION_OPERATION_NAME, - get_operation_fn=self.client.agent_engines.runtimes.revisions._get_delete_runtime_revision_operation, + get_operation_fn=self.client.runtimes.revisions._get_delete_runtime_revision_operation, ) request_mock.assert_called_with( "get", @@ -891,13 +891,13 @@ def test_await_operation(self): None, ) assert isinstance( - operation, _genai_types.DeleteAgentEngineRuntimeRevisionOperation + operation, _genai_types.DeleteRuntimeRevisionOperation ) def test_register_api_methods(self): - agent = self.client.agent_engines.runtimes.revisions._register_api_methods( - agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( - api_client=self.client.agent_engines.runtimes.revisions._api_client, + agent = self.client.runtimes.revisions._register_api_methods( + runtime_revision=_genai_types.RuntimeRevision( + api_client=self.client.runtimes.revisions._api_client, api_resource=_genai_types.ReasoningEngineRuntimeRevision( spec=_genai_types.ReasoningEngineSpec( class_methods=[ @@ -932,10 +932,10 @@ def teardown_method(self): def test_get_runtime_revision(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request" + self.client.runtimes.revisions._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.runtimes.revisions.get( + self.client.runtimes.revisions.get( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME ) request_mock.assert_called_with( @@ -947,12 +947,12 @@ def test_get_runtime_revision(self): def test_list_runtime_revisions(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request" + self.client.runtimes.revisions._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") expected_query_params = {"filter": _TEST_LIST_FILTER} list( - self.client.agent_engines.runtimes.revisions.list( + self.client.runtimes.revisions.list( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config=expected_query_params ) ) @@ -968,7 +968,7 @@ def test_list_runtime_revisions(self): def test_delete_runtime_revision(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request" + self.client.runtimes.revisions._api_client, "request" ) as request_mock: request_mock.side_effect = [ # First call: response to delete. @@ -984,7 +984,7 @@ def test_delete_runtime_revision(self): ), ] - self.client.agent_engines.runtimes.revisions.delete( + self.client.runtimes.revisions.delete( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME ) request_mock.call_args_list[0].assert_called_with( @@ -1003,12 +1003,12 @@ def test_delete_runtime_revision(self): def test_query_runtime_revision(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request" + self.client.runtimes.revisions._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - agent = self.client.agent_engines.runtimes.revisions._register_api_methods( - agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( - api_client=self.client.agent_engines.runtimes.revisions, + agent = self.client.runtimes.revisions._register_api_methods( + runtime_revision=_genai_types.RuntimeRevision( + api_client=self.client.runtimes.revisions, api_resource=_genai_types.ReasoningEngineRuntimeRevision( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -1031,11 +1031,11 @@ def test_query_runtime_revision(self): None, ) - def test_query_agent_engine_async(self): - agent = self.client.agent_engines.runtimes.revisions._register_api_methods( - agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( + def test_query_runtime_async(self): + agent = self.client.runtimes.revisions._register_api_methods( + runtime_revision=_genai_types.RuntimeRevision( api_async_client=runtime_revisions.AsyncRuntimeRevisions( - api_client_=self.client.agent_engines.runtimes.revisions._api_client + api_client_=self.client.runtimes.revisions._api_client ), api_resource=_genai_types.ReasoningEngineRuntimeRevision( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, @@ -1048,7 +1048,7 @@ def test_query_agent_engine_async(self): ) ) with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "async_request" + self.client.runtimes.revisions._api_client, "async_request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") asyncio.run(agent.async_query(query=_TEST_QUERY_PROMPT)) @@ -1063,13 +1063,13 @@ def test_query_agent_engine_async(self): None, ) - def test_query_agent_engine_stream(self): + def test_query_runtime_stream(self): with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, "request_streamed" + self.client.runtimes.revisions._api_client, "request_streamed" ) as request_mock: - agent = self.client.agent_engines.runtimes.revisions._register_api_methods( - agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( - api_client=self.client.agent_engines.runtimes.revisions, + agent = self.client.runtimes.revisions._register_api_methods( + runtime_revision=_genai_types.RuntimeRevision( + api_client=self.client.runtimes.revisions, api_resource=_genai_types.ReasoningEngineRuntimeRevision( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -1092,18 +1092,18 @@ def test_query_agent_engine_stream(self): None, ) - def test_query_agent_engine_async_stream(self): + def test_query_runtime_async_stream(self): async def mock_async_generator(): yield genai_types.HttpResponse(body=b"") with mock.patch.object( - self.client.agent_engines.runtimes.revisions._api_client, + self.client.runtimes.revisions._api_client, "async_request_streamed", ) as request_mock: request_mock.return_value = mock_async_generator() - agent = self.client.agent_engines.runtimes.revisions._register_api_methods( - agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( - api_client=self.client.agent_engines.runtimes.revisions, + agent = self.client.runtimes.revisions._register_api_methods( + runtime_revision=_genai_types.RuntimeRevision( + api_client=self.client.runtimes.revisions, api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -1154,7 +1154,7 @@ def teardown_method(self): def test_delete_runtime_revision(self): with mock.patch.object( - self.client.aio.agent_engines.runtimes.revisions._api_client, + self.client.aio.runtimes.revisions._api_client, "async_request", ) as request_mock: request_mock.side_effect = [ @@ -1171,7 +1171,7 @@ def test_delete_runtime_revision(self): ), ] asyncio.run( - self.client.aio.agent_engines.runtimes.revisions.delete( + self.client.aio.runtimes.revisions.delete( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME ) ) diff --git a/tests/unit/agentplatform/genai/test_agent_engines.py b/tests/unit/agentplatform/genai/test_agent_engines.py index e92f121e8d..97ce0c45d5 100644 --- a/tests/unit/agentplatform/genai/test_agent_engines.py +++ b/tests/unit/agentplatform/genai/test_agent_engines.py @@ -31,9 +31,9 @@ from google.cloud import aiplatform import agentplatform from google.cloud.aiplatform import initializer -from agentplatform.agent_engines.templates import adk -from agentplatform._genai import _agent_engines_utils -from agentplatform._genai import agent_engines +from agentplatform.frameworks import adk +from agentplatform._genai import _runtimes_utils +from agentplatform._genai import runtimes from agentplatform._genai import types as _genai_types from google.genai import client as genai_client from google.genai import types as genai_types @@ -305,7 +305,7 @@ def clone(self): return self def register_operations(self) -> Dict[str, List[str]]: - # Registered method `missing_method` is not a method of the AgentEngine. + # Registered method `missing_method` is not a method of the Runtime. return { _TEST_STANDARD_API_MODE: [ _TEST_DEFAULT_METHOD_NAME, @@ -323,7 +323,7 @@ def method_to_be_unregistered(self, unused_arbitrary_string_name: str) -> str: return unused_arbitrary_string_name.upper() def register_operations(self) -> Dict[str, List[str]]: - # Registered method `missing_method` is not a method of the AgentEngine. + # Registered method `missing_method` is not a method of the Runtime. return {_TEST_STANDARD_API_MODE: [_TEST_METHOD_TO_BE_UNREGISTERED_NAME]} @@ -341,30 +341,30 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_AGENT_ENGINE_DISPLAY_NAME = "Agent Engine Display Name" _TEST_AGENT_ENGINE_DESCRIPTION = "Agent Engine Description" _TEST_AGENT_ENGINE_LIST_FILTER = f'display_name="{_TEST_AGENT_ENGINE_DISPLAY_NAME}"' -_TEST_GCS_DIR_NAME = _agent_engines_utils._DEFAULT_GCS_DIR_NAME -_TEST_BLOB_FILENAME = _agent_engines_utils._BLOB_FILENAME -_TEST_REQUIREMENTS_FILE = _agent_engines_utils._REQUIREMENTS_FILE -_TEST_EXTRA_PACKAGES_FILE = _agent_engines_utils._EXTRA_PACKAGES_FILE -_TEST_STANDARD_API_MODE = _agent_engines_utils._STANDARD_API_MODE -_TEST_ASYNC_API_MODE = _agent_engines_utils._ASYNC_API_MODE -_TEST_STREAM_API_MODE = _agent_engines_utils._STREAM_API_MODE -_TEST_ASYNC_STREAM_API_MODE = _agent_engines_utils._ASYNC_STREAM_API_MODE -_TEST_BIDI_STREAM_API_MODE = _agent_engines_utils._BIDI_STREAM_API_MODE -_TEST_DEFAULT_METHOD_NAME = _agent_engines_utils._DEFAULT_METHOD_NAME -_TEST_DEFAULT_ASYNC_METHOD_NAME = _agent_engines_utils._DEFAULT_ASYNC_METHOD_NAME -_TEST_DEFAULT_STREAM_METHOD_NAME = _agent_engines_utils._DEFAULT_STREAM_METHOD_NAME +_TEST_GCS_DIR_NAME = _runtimes_utils._DEFAULT_GCS_DIR_NAME +_TEST_BLOB_FILENAME = _runtimes_utils._BLOB_FILENAME +_TEST_REQUIREMENTS_FILE = _runtimes_utils._REQUIREMENTS_FILE +_TEST_EXTRA_PACKAGES_FILE = _runtimes_utils._EXTRA_PACKAGES_FILE +_TEST_STANDARD_API_MODE = _runtimes_utils._STANDARD_API_MODE +_TEST_ASYNC_API_MODE = _runtimes_utils._ASYNC_API_MODE +_TEST_STREAM_API_MODE = _runtimes_utils._STREAM_API_MODE +_TEST_ASYNC_STREAM_API_MODE = _runtimes_utils._ASYNC_STREAM_API_MODE +_TEST_BIDI_STREAM_API_MODE = _runtimes_utils._BIDI_STREAM_API_MODE +_TEST_DEFAULT_METHOD_NAME = _runtimes_utils._DEFAULT_METHOD_NAME +_TEST_DEFAULT_ASYNC_METHOD_NAME = _runtimes_utils._DEFAULT_ASYNC_METHOD_NAME +_TEST_DEFAULT_STREAM_METHOD_NAME = _runtimes_utils._DEFAULT_STREAM_METHOD_NAME _TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME = ( - _agent_engines_utils._DEFAULT_ASYNC_STREAM_METHOD_NAME + _runtimes_utils._DEFAULT_ASYNC_STREAM_METHOD_NAME ) _TEST_DEFAULT_BIDI_STREAM_METHOD_NAME = ( - _agent_engines_utils._DEFAULT_BIDI_STREAM_METHOD_NAME + _runtimes_utils._DEFAULT_BIDI_STREAM_METHOD_NAME ) _TEST_CAPITALIZE_ENGINE_METHOD_DOCSTRING = "Runs the engine." _TEST_STREAM_METHOD_DOCSTRING = "Runs the stream engine." _TEST_ASYNC_STREAM_METHOD_DOCSTRING = "Runs the async stream engine." _TEST_BIDI_STREAM_METHOD_DOCSTRING = "Runs the bidi stream engine." -_TEST_MODE_KEY_IN_SCHEMA = _agent_engines_utils._MODE_KEY_IN_SCHEMA -_TEST_METHOD_NAME_KEY_IN_SCHEMA = _agent_engines_utils._METHOD_NAME_KEY_IN_SCHEMA +_TEST_MODE_KEY_IN_SCHEMA = _runtimes_utils._MODE_KEY_IN_SCHEMA +_TEST_METHOD_NAME_KEY_IN_SCHEMA = _runtimes_utils._METHOD_NAME_KEY_IN_SCHEMA _TEST_CUSTOM_METHOD_NAME = "custom_method" _TEST_CUSTOM_ASYNC_METHOD_NAME = "custom_async_method" _TEST_CUSTOM_STREAM_METHOD_NAME = "custom_stream_method" @@ -454,14 +454,14 @@ def register_operations(self) -> Dict[str, List[str]]: "lib", "main.py", ] -_TEST_AGENT_ENGINE_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_QUERY_SCHEMA = _runtimes_utils._generate_schema( CapitalizeEngine().query, schema_name=_TEST_DEFAULT_METHOD_NAME, ) _TEST_AGENT_ENGINE_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STANDARD_API_MODE _TEST_PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}" _TEST_PYTHON_VERSION_OVERRIDE = "3.11" -_TEST_AGENT_ENGINE_FRAMEWORK = _agent_engines_utils._DEFAULT_AGENT_FRAMEWORK +_TEST_AGENT_ENGINE_FRAMEWORK = _runtimes_utils._DEFAULT_AGENT_FRAMEWORK _TEST_AGENT_ENGINE_CLASS_METHOD_1 = { "description": "Runs the engine.", "name": "query", @@ -590,38 +590,38 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE = [{"output": "hello"}, {"output": "world"}] _TEST_AGENT_ENGINE_OPERATION_SCHEMAS = [] _TEST_AGENT_ENGINE_EXTRA_PACKAGE = "fake.py" -_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA = _runtimes_utils._generate_schema( AsyncQueryEngine().async_query, schema_name=_TEST_DEFAULT_ASYNC_METHOD_NAME, ) _TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_ASYNC_API_MODE -_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA = _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_method, schema_name=_TEST_CUSTOM_METHOD_NAME, ) _TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( _TEST_STANDARD_API_MODE ) -_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA = _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_async_method, schema_name=_TEST_CUSTOM_ASYNC_METHOD_NAME, ) _TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( _TEST_ASYNC_API_MODE ) -_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA = _runtimes_utils._generate_schema( StreamQueryEngine().stream_query, schema_name=_TEST_DEFAULT_STREAM_METHOD_NAME, ) _TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STREAM_API_MODE -_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA = _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_stream_method, schema_name=_TEST_CUSTOM_STREAM_METHOD_NAME, ) _TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( _TEST_STREAM_API_MODE ) -_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA = _runtimes_utils._generate_schema( AsyncStreamQueryEngine().async_stream_query, schema_name=_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME, ) @@ -629,7 +629,7 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_ASYNC_STREAM_API_MODE ) _TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA = ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_async_stream_method, schema_name=_TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME, ) @@ -637,7 +637,7 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( _TEST_ASYNC_STREAM_API_MODE ) -_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA = _runtimes_utils._generate_schema( OperationRegistrableEngine().bidi_stream_query, schema_name=_TEST_DEFAULT_BIDI_STREAM_METHOD_NAME, ) @@ -645,7 +645,7 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_BIDI_STREAM_API_MODE ) _TEST_AGENT_ENGINE_CUSTOM_BIDI_STREAM_QUERY_SCHEMA = ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_bidi_stream_method, schema_name=_TEST_CUSTOM_BIDI_STREAM_METHOD_NAME, ) @@ -675,7 +675,7 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_NO_OPERATION_REGISTRABLE_SCHEMAS = [ _TEST_AGENT_ENGINE_QUERY_SCHEMA, ] -_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA = _runtimes_utils._generate_schema( MethodToBeUnregisteredEngine().method_to_be_unregistered, schema_name=_TEST_METHOD_TO_BE_UNREGISTERED_NAME, ) @@ -849,7 +849,7 @@ class FakeObject: @pytest.mark.usefixtures("google_auth_mock") -class TestAgentEngineHelpers: +class TestRuntimeHelpers: def setup_method(self): importlib.reload(initializer) importlib.reload(aiplatform) @@ -866,9 +866,9 @@ def setup_method(self): def teardown_method(self): initializer.global_pool.shutdown(wait=True) - @mock.patch.object(_agent_engines_utils, "_prepare") - def test_create_agent_engine_config_lightweight(self, mock_prepare): - config = self.client.agent_engines._create_config( + @mock.patch.object(_runtimes_utils, "_prepare") + def test_create_runtime_config_lightweight(self, mock_prepare): + config = self.client.runtimes._create_config( mode="create", staging_bucket=_TEST_STAGING_BUCKET, display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, @@ -879,7 +879,7 @@ def test_create_agent_engine_config_lightweight(self, mock_prepare): "description": _TEST_AGENT_ENGINE_DESCRIPTION, } - @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_prepare") @pytest.mark.parametrize( "env_vars,expected_env_vars", [ @@ -902,7 +902,7 @@ def test_create_agent_engine_config_lightweight(self, mock_prepare): ), ], ) - def test_agent_engine_adk_telemetry_enablement( + def test_runtime_adk_telemetry_enablement( self, mock_prepare: mock.Mock, env_vars: dict[str, str], @@ -912,7 +912,7 @@ def test_agent_engine_adk_telemetry_enablement( agent.clone = lambda: agent agent.register_operations = lambda: {} - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", agent=agent, staging_bucket=_TEST_STAGING_BUCKET, @@ -927,11 +927,11 @@ def test_agent_engine_adk_telemetry_enablement( ] @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_create_base64_encoded_tarball", return_value="test_tarball", ) - @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_prepare") @pytest.mark.parametrize( "env_vars,expected_env_vars", [ @@ -954,7 +954,7 @@ def test_agent_engine_adk_telemetry_enablement( ), ], ) - def test_agent_engine_adk_telemetry_enablement_through_source_packages( + def test_runtime_adk_telemetry_enablement_through_source_packages( self, mock_prepare: mock.Mock, mock_create_base64_encoded_tarball: mock.Mock, @@ -965,7 +965,7 @@ def test_agent_engine_adk_telemetry_enablement_through_source_packages( test_file_path = os.path.join(tmpdir, "test_file.txt") with open(test_file_path, "w") as f: f.write("test content") - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -982,9 +982,9 @@ def test_agent_engine_adk_telemetry_enablement_through_source_packages( {"name": key, "value": value} for key, value in expected_env_vars.items() ] - @mock.patch.object(_agent_engines_utils, "_prepare") - def test_create_agent_engine_config_full(self, mock_prepare): - config = self.client.agent_engines._create_config( + @mock.patch.object(_runtimes_utils, "_prepare") + def test_create_runtime_config_full(self, mock_prepare): + config = self.client.runtimes._create_config( mode="create", agent=self.test_agent, staging_bucket=_TEST_STAGING_BUCKET, @@ -1047,11 +1047,11 @@ def test_create_agent_engine_config_full(self, mock_prepare): ) @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_create_base64_encoded_tarball", return_value="test_tarball", ) - def test_create_agent_engine_config_with_source_packages( + def test_create_runtime_config_with_source_packages( self, mock_create_base64_encoded_tarball ): with tempfile.TemporaryDirectory() as tmpdir: @@ -1062,7 +1062,7 @@ def test_create_agent_engine_config_with_source_packages( with open(requirements_file_path, "w") as f: f.write("requests==2.0.0") - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1097,7 +1097,7 @@ def test_create_agent_engine_config_with_source_packages( ) assert "keep_alive_probe" not in config["spec"].get("deployment_spec", {}) - def test_create_agent_engine_config_with_developer_connect_source(self): + def test_create_runtime_config_with_developer_connect_source(self): with tempfile.TemporaryDirectory() as tmpdir: requirements_file_path = os.path.join(tmpdir, "requirements.txt") with open(requirements_file_path, "w") as f: @@ -1107,7 +1107,7 @@ def test_create_agent_engine_config_with_developer_connect_source(self): "revision": "main", "dir": "agent", } - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1140,18 +1140,18 @@ def test_create_agent_engine_config_with_developer_connect_source(self): assert "keep_alive_probe" not in config["spec"].get("deployment_spec", {}) @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_create_base64_encoded_tarball", return_value="test_tarball", ) - def test_create_agent_engine_config_with_empty_keep_alive_probe( + def test_create_runtime_config_with_empty_keep_alive_probe( self, mock_create_base64_encoded_tarball ): with tempfile.TemporaryDirectory() as tmpdir: test_file_path = os.path.join(tmpdir, "test_file.txt") with open(test_file_path, "w") as f: f.write("test content") - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", source_packages=[test_file_path], class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, @@ -1161,7 +1161,7 @@ def test_create_agent_engine_config_with_empty_keep_alive_probe( ) assert "keep_alive_probe" in config["spec"].get("deployment_spec", {}) - def test_create_agent_engine_config_with_agent_config_source_and_requirements_file( + def test_create_runtime_config_with_agent_config_source_and_requirements_file( self, ): with tempfile.TemporaryDirectory() as tmpdir: @@ -1169,7 +1169,7 @@ def test_create_agent_engine_config_with_agent_config_source_and_requirements_fi with open(requirements_file_path, "w") as f: f.write("requests==2.0.0") - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1189,12 +1189,12 @@ def test_create_agent_engine_config_with_agent_config_source_and_requirements_fi }, } - def test_create_agent_engine_config_with_agent_config_source_and_entrypoint_module_warns( + def test_create_runtime_config_with_agent_config_source_and_entrypoint_module_warns( self, caplog ): caplog.set_level(logging.WARNING, logger="vertexai_genai.agentengines") - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1219,11 +1219,11 @@ def test_create_agent_engine_config_with_agent_config_source_and_entrypoint_modu # entrypoint_module is NOT in python_spec @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_create_base64_encoded_tarball", return_value="test_tarball", ) - def test_create_agent_engine_config_with_source_packages_and_image_spec_raises( + def test_create_runtime_config_with_source_packages_and_image_spec_raises( self, mock_create_base64_encoded_tarball ): with tempfile.TemporaryDirectory() as tmpdir: @@ -1235,7 +1235,7 @@ def test_create_agent_engine_config_with_source_packages_and_image_spec_raises( f.write("requests==2.0.0") with pytest.raises(ValueError) as excinfo: - self.client.agent_engines._create_config( + self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1252,11 +1252,11 @@ def test_create_agent_engine_config_with_source_packages_and_image_spec_raises( assert "`image_spec` cannot be specified alongside" in str(excinfo.value) @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_create_base64_encoded_tarball", return_value="test_tarball", ) - def test_create_agent_engine_config_with_agent_config_source_and_image_spec_raises( + def test_create_runtime_config_with_agent_config_source_and_image_spec_raises( self, mock_create_base64_encoded_tarball ): with tempfile.TemporaryDirectory() as tmpdir: @@ -1268,7 +1268,7 @@ def test_create_agent_engine_config_with_agent_config_source_and_image_spec_rais f.write("requests==2.0.0") with pytest.raises(ValueError) as excinfo: - self.client.agent_engines._create_config( + self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1281,8 +1281,8 @@ def test_create_agent_engine_config_with_agent_config_source_and_image_spec_rais ) assert "`image_spec` cannot be specified alongside" in str(excinfo.value) - def test_create_agent_engine_config_with_agent_config_source(self): - config = self.client.agent_engines._create_config( + def test_create_runtime_config_with_agent_config_source(self): + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1306,11 +1306,11 @@ def test_create_agent_engine_config_with_agent_config_source(self): ) @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_create_base64_encoded_tarball", return_value="test_tarball", ) - def test_create_agent_engine_config_with_source_packages_and_agent_config_source( + def test_create_runtime_config_with_source_packages_and_agent_config_source( self, mock_create_base64_encoded_tarball ): with tempfile.TemporaryDirectory() as tmpdir: @@ -1321,7 +1321,7 @@ def test_create_agent_engine_config_with_source_packages_and_agent_config_source with open(requirements_file_path, "w") as f: f.write("requests==2.0.0") - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1351,9 +1351,9 @@ def test_create_agent_engine_config_with_source_packages_and_agent_config_source == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT ) - def test_create_agent_engine_config_with_container_spec(self): + def test_create_runtime_config_with_container_spec(self): container_spec = {"image_uri": "gcr.io/test-project/test-image"} - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1371,11 +1371,11 @@ def test_create_agent_engine_config_with_container_spec(self): ) assert "keep_alive_probe" not in config["spec"].get("deployment_spec", {}) - def test_create_agent_engine_config_with_container_spec_and_keep_alive_probe( + def test_create_runtime_config_with_container_spec_and_keep_alive_probe( self, ): container_spec = {"image_uri": "gcr.io/test-project/test-image"} - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1397,10 +1397,10 @@ def test_create_agent_engine_config_with_container_spec_and_keep_alive_probe( == _TEST_AGENT_ENGINE_KEEP_ALIVE_PROBE ) - def test_create_agent_engine_config_with_container_spec_and_others_raises(self): + def test_create_runtime_config_with_container_spec_and_others_raises(self): container_spec = {"image_uri": "gcr.io/test-project/test-image"} with pytest.raises(ValueError) as excinfo: - self.client.agent_engines._create_config( + self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1410,7 +1410,7 @@ def test_create_agent_engine_config_with_container_spec_and_others_raises(self): assert "please do not specify `agent`" in str(excinfo.value) with pytest.raises(ValueError) as excinfo: - self.client.agent_engines._create_config( + self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1420,12 +1420,12 @@ def test_create_agent_engine_config_with_container_spec_and_others_raises(self): assert "please do not specify `source_packages`" in str(excinfo.value) @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_create_base64_encoded_tarball", return_value="test_tarball", ) - @mock.patch.object(_agent_engines_utils, "_validate_packages_or_raise") - def test_create_agent_engine_config_with_source_packages_and_build_options( + @mock.patch.object(_runtimes_utils, "_validate_packages_or_raise") + def test_create_runtime_config_with_source_packages_and_build_options( self, mock_validate_packages, mock_create_base64_encoded_tarball ): with tempfile.TemporaryDirectory() as tmpdir: @@ -1438,7 +1438,7 @@ def test_create_agent_engine_config_with_source_packages_and_build_options( source_packages = [test_file_path, "installation_scripts/install.sh"] mock_validate_packages.return_value = source_packages - self.client.agent_engines._create_config( + self.client.runtimes._create_config( mode="create", source_packages=source_packages, entrypoint_module="main", @@ -1451,15 +1451,15 @@ def test_create_agent_engine_config_with_source_packages_and_build_options( build_options=build_options, ) - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_validate_packages_or_raise") - def test_create_agent_engine_config_with_build_options( + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_validate_packages_or_raise") + def test_create_runtime_config_with_build_options( self, mock_validate_packages, mock_prepare ): build_options = {"installation_scripts": ["install.sh"]} extra_packages = ["install.sh"] - self.client.agent_engines._create_config( + self.client.runtimes._create_config( mode="create", agent=self.test_agent, staging_bucket=_TEST_STAGING_BUCKET, @@ -1473,9 +1473,9 @@ def test_create_agent_engine_config_with_build_options( build_options=build_options, ) - @mock.patch.object(_agent_engines_utils, "_prepare") - def test_update_agent_engine_config_full(self, mock_prepare): - config = self.client.agent_engines._create_config( + @mock.patch.object(_runtimes_utils, "_prepare") + def test_update_runtime_config_full(self, mock_prepare): + config = self.client.runtimes._create_config( mode="update", agent=self.test_agent, staging_bucket=_TEST_STAGING_BUCKET, @@ -1538,9 +1538,9 @@ def test_update_agent_engine_config_full(self, mock_prepare): ] ) - @mock.patch.object(_agent_engines_utils, "_prepare") - def test_update_agent_engine_clear_service_account(self, mock_prepare): - config = self.client.agent_engines._create_config( + @mock.patch.object(_runtimes_utils, "_prepare") + def test_update_runtime_clear_service_account(self, mock_prepare): + config = self.client.runtimes._create_config( mode="update", service_account="", identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, @@ -1559,7 +1559,7 @@ def test_update_agent_engine_clear_service_account(self, mock_prepare): def test_get_agent_operation(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body=json.dumps( @@ -1570,7 +1570,7 @@ def test_get_agent_operation(self): } ), ) - operation = self.client.agent_engines._get_agent_operation( + operation = self.client.runtimes._get_agent_operation( operation_name=_TEST_AGENT_ENGINE_OPERATION_NAME, ) request_mock.assert_called_with( @@ -1579,13 +1579,13 @@ def test_get_agent_operation(self): {"_url": {"operationName": _TEST_AGENT_ENGINE_OPERATION_NAME}}, None, ) - assert isinstance(operation, _genai_types.AgentEngineOperation) + assert isinstance(operation, _genai_types.RuntimeOperation) assert operation.done assert isinstance(operation.response, _genai_types.ReasoningEngine) def test_await_operation(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body=json.dumps( @@ -1596,9 +1596,9 @@ def test_await_operation(self): } ), ) - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=_TEST_AGENT_ENGINE_OPERATION_NAME, - get_operation_fn=self.client.agent_engines._get_agent_operation, + get_operation_fn=self.client.runtimes._get_agent_operation, ) request_mock.assert_called_with( "get", @@ -1606,12 +1606,12 @@ def test_await_operation(self): {"_url": {"operationName": _TEST_AGENT_ENGINE_OPERATION_NAME}}, None, ) - assert isinstance(operation, _genai_types.AgentEngineOperation) + assert isinstance(operation, _genai_types.RuntimeOperation) def test_register_api_methods(self): - agent = self.client.agent_engines._register_api_methods( - agent_engine=_genai_types.AgentEngine( - api_client=self.client.agent_engines._api_client, + agent = self.client.runtimes._register_api_methods( + runtime=_genai_types.Runtime( + api_client=self.client.runtimes._api_client, api_resource=_genai_types.ReasoningEngine( spec=_genai_types.ReasoningEngineSpec( class_methods=[ @@ -1627,7 +1627,7 @@ def test_register_api_methods(self): @pytest.mark.usefixtures("caplog") def test_invalid_requirement_warning(self, caplog): - _agent_engines_utils._parse_constraints( + _runtimes_utils._parse_constraints( constraints=["invalid requirement line"], ) assert "Failed to parse constraint" in caplog.text @@ -1638,7 +1638,7 @@ def test_requirements_with_whl_files(self): "/content/wxPython-4.2.3-cp39-cp39-macosx_12_0_x86_64.whl", "https://wxpython.org/Phoenix/snapshot-builds/wxPython-4.2.2-cp38-cp38-macosx_12_0_x86_64.whl", ] - result = _agent_engines_utils._parse_constraints( + result = _runtimes_utils._parse_constraints( constraints=whl_files, ) assert result == { @@ -1650,7 +1650,7 @@ def test_requirements_with_whl_files(self): def test_compare_requirements_with_required_packages(self): requirements = {"requests": "2.0.0"} constraints = ["requests==1.0.0"] - result = _agent_engines_utils._compare_requirements( + result = _runtimes_utils._compare_requirements( requirements=requirements, constraints=constraints, ) @@ -1666,7 +1666,7 @@ def test_compare_requirements_with_required_packages(self): def test_scan_simple_object(self): """Test scanning an object importing a known third-party package.""" fake_obj = _create_fake_object_with_module("requests") - requirements = _agent_engines_utils._scan_requirements( + requirements = _runtimes_utils._scan_requirements( obj=fake_obj, package_distributions=_TEST_PACKAGE_DISTRIBUTIONS, ) @@ -1680,7 +1680,7 @@ def test_scan_simple_object(self): def test_scan_object_with_stdlib_module(self): """Test that stdlib modules are ignored by default.""" fake_obj_stdlib = _create_fake_object_with_module("json") - requirements = _agent_engines_utils._scan_requirements( + requirements = _runtimes_utils._scan_requirements( obj=fake_obj_stdlib, package_distributions=_TEST_PACKAGE_DISTRIBUTIONS, ) @@ -1695,13 +1695,13 @@ def test_scan_object_with_stdlib_module(self): def test_scan_with_default_ignore_modules(self, monkeypatch): """Test implicitly ignoring a module.""" fake_obj = _create_fake_object_with_module("requests") - original_base = _agent_engines_utils._BASE_MODULES + original_base = _runtimes_utils._BASE_MODULES monkeypatch.setattr( - _agent_engines_utils, + _runtimes_utils, "_BASE_MODULES", set(original_base) | {"requests"}, ) - requirements = _agent_engines_utils._scan_requirements( + requirements = _runtimes_utils._scan_requirements( obj=fake_obj, package_distributions=_TEST_PACKAGE_DISTRIBUTIONS, ) @@ -1716,7 +1716,7 @@ def test_scan_with_default_ignore_modules(self, monkeypatch): def test_scan_with_explicit_ignore_modules(self): """Test explicitly ignoring a module.""" fake_obj = _create_fake_object_with_module("requests") - requirements = _agent_engines_utils._scan_requirements( + requirements = _runtimes_utils._scan_requirements( obj=fake_obj, ignore_modules=["requests"], package_distributions=_TEST_PACKAGE_DISTRIBUTIONS, @@ -1765,7 +1765,7 @@ def test_scan_with_explicit_ignore_modules(self): ], ) def test_to_parsed_json(self, obj, expected): - for got, want in zip(_agent_engines_utils._yield_parsed_json(obj), expected): + for got, want in zip(_runtimes_utils._yield_parsed_json(obj), expected): assert got == want # pytest does not allow absl.testing.parameterized.named_parameters. @@ -1802,12 +1802,12 @@ def test_to_parsed_json(self, obj, expected): ], ) def test_yield_parsed_json_from_httpbody(self, obj, expected): - got = list(_agent_engines_utils._yield_parsed_json_from_httpbody(obj)) + got = list(_runtimes_utils._yield_parsed_json_from_httpbody(obj)) assert got == expected def test_yield_parsed_json_from_httpbody_non_json_content_type(self): body = httpbody_pb2.HttpBody(content_type="text/plain", data=b"hello") - assert list(_agent_engines_utils._yield_parsed_json_from_httpbody(body)) == [ + assert list(_runtimes_utils._yield_parsed_json_from_httpbody(body)) == [ body ] @@ -1820,7 +1820,7 @@ def test_create_base64_encoded_tarball(self): origin_dir = os.getcwd() try: os.chdir(tmpdir) - encoded_tarball = _agent_engines_utils._create_base64_encoded_tarball( + encoded_tarball = _runtimes_utils._create_base64_encoded_tarball( source_packages=["test_file.txt"] ) finally: @@ -1843,30 +1843,30 @@ def test_create_base64_encoded_tarball_outside_project_dir_raises(self): try: os.chdir(project_dir) with pytest.raises(ValueError) as excinfo: - _agent_engines_utils._create_base64_encoded_tarball( + _runtimes_utils._create_base64_encoded_tarball( source_packages=["../sibling.txt"] ) assert "is outside the project directory" in str(excinfo.value) finally: os.chdir(origin_dir) - @mock.patch.object(_agent_engines_utils, "_upload_requirements") - @mock.patch.object(_agent_engines_utils, "_upload_extra_packages") - @mock.patch.object(_agent_engines_utils, "_upload_agent_engine") - @mock.patch.object(_agent_engines_utils, "_scan_requirements") - @mock.patch.object(_agent_engines_utils, "_get_gcs_bucket") + @mock.patch.object(_runtimes_utils, "_upload_requirements") + @mock.patch.object(_runtimes_utils, "_upload_extra_packages") + @mock.patch.object(_runtimes_utils, "_upload_runtime") + @mock.patch.object(_runtimes_utils, "_scan_requirements") + @mock.patch.object(_runtimes_utils, "_get_gcs_bucket") def test_prepare_with_creds( self, mock_get_gcs_bucket, mock_scan_requirements, - mock_upload_agent_engine, + mock_upload_runtime, mock_upload_extra_packages, mock_upload_requirements, ): mock_scan_requirements.return_value = {} mock_creds = mock.Mock(spec=auth_credentials.AnonymousCredentials()) mock_creds.universe_domain = "googleapis.com" - _agent_engines_utils._prepare( + _runtimes_utils._prepare( agent=self.test_agent, project=_TEST_PROJECT, location=_TEST_LOCATION, @@ -1876,31 +1876,31 @@ def test_prepare_with_creds( requirements=[], extra_packages=[], ) - mock_upload_agent_engine.assert_called_once_with( + mock_upload_runtime.assert_called_once_with( agent=self.test_agent, gcs_bucket=mock.ANY, gcs_dir_name=_TEST_GCS_DIR_NAME, ) - @mock.patch.object(_agent_engines_utils, "_upload_requirements") - @mock.patch.object(_agent_engines_utils, "_upload_extra_packages") - @mock.patch.object(_agent_engines_utils, "_upload_agent_engine") - @mock.patch.object(_agent_engines_utils, "_scan_requirements") + @mock.patch.object(_runtimes_utils, "_upload_requirements") + @mock.patch.object(_runtimes_utils, "_upload_extra_packages") + @mock.patch.object(_runtimes_utils, "_upload_runtime") + @mock.patch.object(_runtimes_utils, "_scan_requirements") @mock.patch("google.auth.default") - @mock.patch.object(_agent_engines_utils, "_get_gcs_bucket") + @mock.patch.object(_runtimes_utils, "_get_gcs_bucket") def test_prepare_without_creds( self, mock_get_gcs_bucket, mock_auth_default, mock_scan_requirements, - mock_upload_agent_engine, + mock_upload_runtime, mock_upload_extra_packages, mock_upload_requirements, ): mock_scan_requirements.return_value = {} mock_creds = mock.Mock(spec=auth_credentials.AnonymousCredentials()) mock_auth_default.return_value = (mock_creds, _TEST_PROJECT) - _agent_engines_utils._prepare( + _runtimes_utils._prepare( agent=self.test_agent, project=_TEST_PROJECT, location=_TEST_LOCATION, @@ -1915,7 +1915,7 @@ def test_prepare_without_creds( staging_bucket=_TEST_STAGING_BUCKET, credentials=None, ) - mock_upload_agent_engine.assert_called_once_with( + mock_upload_runtime.assert_called_once_with( agent=self.test_agent, gcs_bucket=mock.ANY, gcs_dir_name=_TEST_GCS_DIR_NAME, @@ -1992,13 +1992,13 @@ def test_get_reasoning_engine_id( ): if expected_exception: with pytest.raises(expected_exception) as excinfo: - _agent_engines_utils._get_reasoning_engine_id( + _runtimes_utils._get_reasoning_engine_id( operation_name=operation_name, resource_name=resource_name ) assert expected_message in str(excinfo.value) else: assert ( - _agent_engines_utils._get_reasoning_engine_id( + _runtimes_utils._get_reasoning_engine_id( operation_name=operation_name, resource_name=resource_name ) == expected_id @@ -2006,7 +2006,7 @@ def test_get_reasoning_engine_id( @pytest.mark.usefixtures("google_auth_mock") -class TestAgentEngine: +class TestRuntime: def setup_method(self): importlib.reload(initializer) importlib.reload(aiplatform) @@ -2023,12 +2023,12 @@ def setup_method(self): def teardown_method(self): initializer.global_pool.shutdown(wait=True) - def test_get_agent_engine(self): + def test_get_runtime(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.get(name=_TEST_AGENT_ENGINE_RESOURCE_NAME) + self.client.runtimes.get(name=_TEST_AGENT_ENGINE_RESOURCE_NAME) request_mock.assert_called_with( "get", _TEST_AGENT_ENGINE_RESOURCE_NAME, @@ -2036,13 +2036,13 @@ def test_get_agent_engine(self): None, ) - def test_list_agent_engine(self): + def test_list_runtime(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") expected_query_params = {"filter": _TEST_AGENT_ENGINE_LIST_FILTER} - list(self.client.agent_engines.list(config=expected_query_params)) + list(self.client.runtimes.list(config=expected_query_params)) request_mock.assert_called_with( "get", f"reasoningEngines?{urlencode(expected_query_params)}", @@ -2051,21 +2051,21 @@ def test_list_agent_engine(self): ) @pytest.mark.usefixtures("caplog") - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) - def test_create_agent_engine( + def test_create_runtime( self, mock_get_reasoning_engine_id, mock_await_operation, mock_prepare, caplog, ): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, @@ -2073,12 +2073,12 @@ def test_create_agent_engine( ) caplog.set_level(logging.INFO, logger="vertexai_genai.agentengines") with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, @@ -2106,41 +2106,41 @@ def test_create_agent_engine( None, ) assert "View progress and logs at" in caplog.text - assert "Agent Engine created. To use it in another session:" in caplog.text + assert "Agent Runtime created. To use it in another session:" in caplog.text assert ( - f"agent_engine=client.agent_engines.get(name=" + f"runtime=client.runtimes.get(name=" f"'{_TEST_AGENT_ENGINE_RESOURCE_NAME}')" in caplog.text ) - @mock.patch.object(agent_engines.AgentEngines, "_create_config") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(runtimes.Runtimes, "_create_config") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) - def test_create_agent_engine_lightweight( + def test_create_runtime_lightweight( self, mock_get_reasoning_engine_id, mock_await_operation, mock_create_config, ): - mock_create_config.return_value = _genai_types.CreateAgentEngineConfig( + mock_create_config.return_value = _genai_types.CreateRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, ) - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( - config=_genai_types.AgentEngineConfig( + self.client.runtimes.create( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, ) @@ -2155,14 +2155,14 @@ def test_create_agent_engine_lightweight( None, ) - @mock.patch.object(agent_engines.AgentEngines, "_create_config") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(runtimes.Runtimes, "_create_config") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) - def test_create_agent_engine_with_env_vars_dict( + def test_create_runtime_with_env_vars_dict( self, mock_get_reasoning_engine_id, mock_await_operation, @@ -2181,19 +2181,19 @@ def test_create_agent_engine_with_env_vars_dict( "agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK, }, } - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], @@ -2258,14 +2258,14 @@ def test_create_agent_engine_with_env_vars_dict( None, ) - @mock.patch.object(agent_engines.AgentEngines, "_create_config") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(runtimes.Runtimes, "_create_config") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) - def test_create_agent_engine_with_custom_service_account( + def test_create_runtime_with_custom_service_account( self, mock_get_reasoning_engine_id, mock_await_operation, @@ -2286,19 +2286,19 @@ def test_create_agent_engine_with_custom_service_account( "agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK, }, } - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], @@ -2364,14 +2364,14 @@ def test_create_agent_engine_with_custom_service_account( None, ) - @mock.patch.object(agent_engines.AgentEngines, "_create_config") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(runtimes.Runtimes, "_create_config") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) - def test_create_agent_engine_with_experimental_mode( + def test_create_runtime_with_experimental_mode( self, mock_get_reasoning_engine_id, mock_await_operation, @@ -2392,19 +2392,19 @@ def test_create_agent_engine_with_experimental_mode( "class_methods": [_TEST_AGENT_ENGINE_CLASS_METHOD_1], }, } - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], @@ -2470,23 +2470,23 @@ def test_create_agent_engine_with_experimental_mode( ) @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_create_base64_encoded_tarball", return_value="test_tarball", ) - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) - def test_create_agent_engine_with_source_packages( + def test_create_runtime_with_source_packages( self, mock_get_reasoning_engine_id, mock_await_operation, mock_create_base64_encoded_tarball, ): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, @@ -2501,11 +2501,11 @@ def test_create_agent_engine_with_source_packages( f.write("requests==2.0.0") with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( - config=_genai_types.AgentEngineConfig( + self.client.runtimes.create( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, source_packages=[test_file_path], @@ -2541,14 +2541,14 @@ def test_create_agent_engine_with_source_packages( source_packages=[test_file_path] ) - @mock.patch.object(agent_engines.AgentEngines, "_create_config") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(runtimes.Runtimes, "_create_config") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) - def test_create_agent_engine_with_class_methods( + def test_create_runtime_with_class_methods( self, mock_get_reasoning_engine_id, mock_await_operation, @@ -2566,19 +2566,19 @@ def test_create_agent_engine_with_class_methods( "class_methods": _TEST_AGENT_ENGINE_CLASS_METHODS, }, } - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], @@ -2640,14 +2640,14 @@ def test_create_agent_engine_with_class_methods( None, ) - @mock.patch.object(agent_engines.AgentEngines, "_create_config") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(runtimes.Runtimes, "_create_config") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) - def test_create_agent_engine_with_agent_framework( + def test_create_runtime_with_agent_framework( self, mock_get_reasoning_engine_id, mock_await_operation, @@ -2666,19 +2666,19 @@ def test_create_agent_engine_with_agent_framework( "agent_framework": _TEST_AGENT_FRAMEWORK, }, } - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], @@ -2742,12 +2742,12 @@ def test_create_agent_engine_with_agent_framework( ) @pytest.mark.usefixtures("caplog") - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_await_operation") - def test_update_agent_engine_requirements( + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_requirements( self, mock_await_operation, mock_prepare, caplog ): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, @@ -2755,13 +2755,13 @@ def test_update_agent_engine_requirements( ) caplog.set_level(logging.INFO, logger="vertexai_genai.agentengines") with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.update( + self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( staging_bucket=_TEST_STAGING_BUCKET, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, ), @@ -2793,31 +2793,31 @@ def test_update_agent_engine_requirements( }, None, ) - assert "Agent Engine updated. To use it in another session:" in caplog.text + assert "Agent Runtime updated. To use it in another session:" in caplog.text assert ( - f"agent_engine=client.agent_engines.get(" + f"runtime=client.runtimes.get(" f"name='{_TEST_AGENT_ENGINE_RESOURCE_NAME}')" in caplog.text ) - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_await_operation") - def test_update_agent_engine_extra_packages( + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_extra_packages( self, mock_await_operation, mock_prepare ): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.update( + self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( staging_bucket=_TEST_STAGING_BUCKET, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], @@ -2853,38 +2853,38 @@ def test_update_agent_engine_extra_packages( None, ) - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_await_operation") - def test_update_agent_engine_deployment_config_without_agent_raises( + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_deployment_config_without_agent_raises( self, mock_await_operation, mock_prepare ): with pytest.raises(ValueError, match="To update `env_vars`"): - self.client.agent_engines.update( + self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT ), ) - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_await_operation") - def test_update_agent_engine_env_vars( + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_env_vars( self, mock_await_operation, mock_prepare, caplog ): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.update( + self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( staging_bucket=_TEST_STAGING_BUCKET, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT, @@ -2923,25 +2923,25 @@ def test_update_agent_engine_env_vars( None, ) - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_await_operation") - def test_update_agent_engine_with_empty_keep_alive_probe( + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_with_empty_keep_alive_probe( self, mock_await_operation, mock_prepare ): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.update( + self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( staging_bucket=_TEST_STAGING_BUCKET, keep_alive_probe={}, ), @@ -2976,11 +2976,11 @@ def test_update_agent_engine_with_empty_keep_alive_probe( None, ) - @mock.patch.object(_agent_engines_utils, "_await_operation") - def test_update_agent_engine_with_container_spec_and_keep_alive_probe( + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_with_container_spec_and_keep_alive_probe( self, mock_await_operation ): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, @@ -2988,12 +2988,12 @@ def test_update_agent_engine_with_container_spec_and_keep_alive_probe( ) container_spec = {"image_uri": "gcr.io/test-project/test-image"} with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.update( + self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( container_spec=container_spec, keep_alive_probe=_TEST_AGENT_ENGINE_KEEP_ALIVE_PROBE, class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, @@ -3026,21 +3026,21 @@ def test_update_agent_engine_with_container_spec_and_keep_alive_probe( None, ) - @mock.patch.object(_agent_engines_utils, "_await_operation") - def test_update_agent_engine_display_name(self, mock_await_operation): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_display_name(self, mock_await_operation): + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.update( + self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, ), ) @@ -3055,21 +3055,21 @@ def test_update_agent_engine_display_name(self, mock_await_operation): None, ) - @mock.patch.object(_agent_engines_utils, "_await_operation") - def test_update_agent_engine_description(self, mock_await_operation): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_description(self, mock_await_operation): + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.update( + self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( description=_TEST_AGENT_ENGINE_DESCRIPTION, ), ) @@ -3084,12 +3084,12 @@ def test_update_agent_engine_description(self, mock_await_operation): None, ) - def test_delete_agent_engine(self): + def test_delete_runtime(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.delete(name=_TEST_AGENT_ENGINE_RESOURCE_NAME) + self.client.runtimes.delete(name=_TEST_AGENT_ENGINE_RESOURCE_NAME) request_mock.assert_called_with( "delete", _TEST_AGENT_ENGINE_RESOURCE_NAME, @@ -3097,12 +3097,12 @@ def test_delete_agent_engine(self): None, ) - def test_delete_agent_engine_force(self): + def test_delete_runtime_force(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.delete( + self.client.runtimes.delete( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, force=True, ) @@ -3128,18 +3128,18 @@ def test_delete_agent_engine_force(self): ], ids=["default", "http_options_instance", "http_options_dict"], ) - def test_query_agent_engine(self, http_options_arg, expected_http_options): + def test_query_runtime(self, http_options_arg, expected_http_options): """Sync query: forwards http_options to the HTTP layer (or None by default).""" kwargs = {"query": _TEST_QUERY_PROMPT} if http_options_arg is not _UNSET_HTTP_OPTIONS: kwargs["http_options"] = http_options_arg with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - agent = self.client.agent_engines._register_api_methods( - agent_engine=_genai_types.AgentEngine( - api_client=self.client.agent_engines, + agent = self.client.runtimes._register_api_methods( + runtime=_genai_types.Runtime( + api_client=self.client.runtimes, api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -3163,11 +3163,11 @@ def test_query_agent_engine(self, http_options_arg, expected_http_options): ) @mock.patch("google.cloud.storage.Client") - @mock.patch.object(agent_engines.AgentEngines, "_get") + @mock.patch.object(runtimes.Runtimes, "_get") @mock.patch("uuid.uuid4") - def test_run_query_job_agent_engine(self, mock_uuid, get_mock, mock_storage_client): + def test_run_query_job_runtime(self, mock_uuid, get_mock, mock_storage_client): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body='{"name": "projects/123/locations/us-central1/reasoningEngines/456/operations/789"}' @@ -3194,7 +3194,7 @@ def test_run_query_job_agent_engine(self, mock_uuid, get_mock, mock_storage_clie ), ) - result = self.client.agent_engines.run_query_job( + result = self.client.runtimes.run_query_job( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config={ "query": _TEST_QUERY_PROMPT, @@ -3224,32 +3224,32 @@ def test_run_query_job_agent_engine(self, mock_uuid, get_mock, mock_storage_clie None, ) - def test_run_query_job_agent_engine_missing_query(self): + def test_run_query_job_runtime_missing_query(self): with pytest.raises( ValueError, match="`query` is required in the config object." ): - self.client.agent_engines.run_query_job( + self.client.runtimes.run_query_job( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config={"output_gcs_uri": "gs://my-input-bucket/"}, ) - def test_run_query_job_agent_engine_missing_uri(self): + def test_run_query_job_runtime_missing_uri(self): with pytest.raises( ValueError, match="`output_gcs_uri` is required in the config object." ): - self.client.agent_engines.run_query_job( + self.client.runtimes.run_query_job( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config={"query": _TEST_QUERY_PROMPT}, ) @mock.patch("google.cloud.storage.Client") - @mock.patch.object(agent_engines.AgentEngines, "_get") + @mock.patch.object(runtimes.Runtimes, "_get") @mock.patch("uuid.uuid4") - def test_run_query_job_agent_engine_bucket_creation_forbidden( + def test_run_query_job_runtime_bucket_creation_forbidden( self, mock_uuid, get_mock, mock_storage_client ): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body='{"name": "projects/123/locations/us-central1/reasoningEngines/456/operations/789"}' @@ -3279,7 +3279,7 @@ def test_run_query_job_agent_engine_bucket_creation_forbidden( with pytest.raises( ValueError, match="Permission denied to check existence of bucket" ): - self.client.agent_engines.run_query_job( + self.client.runtimes.run_query_job( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config={ "query": _TEST_QUERY_PROMPT, @@ -3288,13 +3288,13 @@ def test_run_query_job_agent_engine_bucket_creation_forbidden( ) @mock.patch("google.cloud.storage.Client") - @mock.patch.object(agent_engines.AgentEngines, "_get") + @mock.patch.object(runtimes.Runtimes, "_get") @mock.patch("uuid.uuid4") - def test_run_query_job_agent_engine_file_uri( + def test_run_query_job_runtime_file_uri( self, mock_uuid, get_mock, mock_storage_client ): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body='{"name": "projects/123/locations/us-central1/reasoningEngines/456/operations/789"}' @@ -3315,7 +3315,7 @@ def test_run_query_job_agent_engine_file_uri( ), ) - result = self.client.agent_engines.run_query_job( + result = self.client.runtimes.run_query_job( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config={ "query": _TEST_QUERY_PROMPT, @@ -3333,13 +3333,13 @@ def test_run_query_job_agent_engine_file_uri( ) @mock.patch("google.cloud.storage.Client") - @mock.patch.object(agent_engines.AgentEngines, "_get") + @mock.patch.object(runtimes.Runtimes, "_get") @mock.patch("uuid.uuid4") - def test_run_query_job_agent_engine_directory_no_slash( + def test_run_query_job_runtime_directory_no_slash( self, mock_uuid, get_mock, mock_storage_client ): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body='{"name": "projects/123/locations/us-central1/reasoningEngines/456/operations/789"}' @@ -3362,7 +3362,7 @@ def test_run_query_job_agent_engine_directory_no_slash( ), ) - result = self.client.agent_engines.run_query_job( + result = self.client.runtimes.run_query_job( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config={ "query": _TEST_QUERY_PROMPT, @@ -3395,15 +3395,15 @@ def test_run_query_job_agent_engine_directory_no_slash( ], ids=["default", "http_options_instance", "http_options_dict"], ) - def test_query_agent_engine_async(self, http_options_arg, expected_http_options): + def test_query_runtime_async(self, http_options_arg, expected_http_options): """Async query: forwards http_options to the HTTP layer (or None by default).""" kwargs = {"query": _TEST_QUERY_PROMPT} if http_options_arg is not _UNSET_HTTP_OPTIONS: kwargs["http_options"] = http_options_arg - agent = self.client.agent_engines._register_api_methods( - agent_engine=_genai_types.AgentEngine( - api_async_client=agent_engines.AsyncAgentEngines( - api_client_=self.client.agent_engines._api_client + agent = self.client.runtimes._register_api_methods( + runtime=_genai_types.Runtime( + api_async_client=runtimes.AsyncRuntimes( + api_client_=self.client.runtimes._api_client ), api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, @@ -3416,7 +3416,7 @@ def test_query_agent_engine_async(self, http_options_arg, expected_http_options) ) ) with mock.patch.object( - self.client.agent_engines._api_client, "async_request" + self.client.runtimes._api_client, "async_request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") asyncio.run(agent.async_query(**kwargs)) @@ -3431,13 +3431,13 @@ def test_query_agent_engine_async(self, http_options_arg, expected_http_options) expected_http_options, ) - def test_cancel_query_job_agent_engine(self): + def test_cancel_query_job_runtime(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="{}") - result = self.client.agent_engines.cancel_query_job( + result = self.client.runtimes.cancel_query_job( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config={"operation_name": _TEST_AGENT_ENGINE_OPERATION_NAME}, ) @@ -3453,9 +3453,9 @@ def test_cancel_query_job_agent_engine(self): None, ) - def test_check_query_job_agent_engine(self): + def test_check_query_job_runtime(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( headers={}, @@ -3473,7 +3473,7 @@ def test_check_query_job_agent_engine(self): mock_bucket.blob.return_value = mock_blob mock_storage_client.return_value.bucket.return_value = mock_bucket - result = self.client.agent_engines.check_query_job( + result = self.client.runtimes.check_query_job( name="projects/123/locations/us-central1/reasoningEngines/456/operations/789", config={"retrieve_result": True}, ) @@ -3490,9 +3490,9 @@ def test_check_query_job_agent_engine(self): {}, ) - def test_check_query_job_agent_engine_running(self): + def test_check_query_job_runtime_running(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( headers={}, @@ -3503,7 +3503,7 @@ def test_check_query_job_agent_engine_running(self): ), ) - result = self.client.agent_engines.check_query_job( + result = self.client.runtimes.check_query_job( name="projects/123/locations/us-central1/reasoningEngines/456/operations/789", config={"retrieve_result": True}, ) @@ -3515,16 +3515,16 @@ def test_check_query_job_agent_engine_running(self): result=None, ) - def test_check_query_job_agent_engine_failed(self): + def test_check_query_job_runtime_failed(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( headers={}, body='{"done": true, "error": {"message": "Job failed with errors."}}', ) - result = self.client.agent_engines.check_query_job( + result = self.client.runtimes.check_query_job( name="projects/123/locations/us-central1/reasoningEngines/456/operations/789", config={"retrieve_result": True}, ) @@ -3536,9 +3536,9 @@ def test_check_query_job_agent_engine_failed(self): result="{'message': 'Job failed with errors.'}", ) - def test_check_query_job_agent_engine_no_retrieve(self): + def test_check_query_job_runtime_no_retrieve(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( headers={}, @@ -3549,7 +3549,7 @@ def test_check_query_job_agent_engine_no_retrieve(self): ), ) - result = self.client.agent_engines.check_query_job( + result = self.client.runtimes.check_query_job( name="projects/123/locations/us-central1/reasoningEngines/456/operations/789", config={"retrieve_result": False}, ) @@ -3561,9 +3561,9 @@ def test_check_query_job_agent_engine_no_retrieve(self): result=None, ) - def test_check_query_job_agent_engine_blob_not_exists(self): + def test_check_query_job_runtime_blob_not_exists(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( headers={}, @@ -3585,7 +3585,7 @@ def test_check_query_job_agent_engine_blob_not_exists(self): ValueError, match="Failed to retrieve blob results for gs://my-output-bucket/output.json", ): - self.client.agent_engines.check_query_job( + self.client.runtimes.check_query_job( name="projects/123/locations/us-central1/reasoningEngines/456/operations/789", config={"retrieve_result": True}, ) @@ -3605,17 +3605,17 @@ def test_check_query_job_agent_engine_blob_not_exists(self): ], ids=["default", "http_options_instance", "http_options_dict"], ) - def test_query_agent_engine_stream(self, http_options_arg, expected_http_options): + def test_query_runtime_stream(self, http_options_arg, expected_http_options): """Streaming query: forwards http_options to the HTTP layer (or None by default).""" kwargs = {"query": _TEST_QUERY_PROMPT} if http_options_arg is not _UNSET_HTTP_OPTIONS: kwargs["http_options"] = http_options_arg with mock.patch.object( - self.client.agent_engines._api_client, "request_streamed" + self.client.runtimes._api_client, "request_streamed" ) as request_mock: - agent = self.client.agent_engines._register_api_methods( - agent_engine=_genai_types.AgentEngine( - api_client=self.client.agent_engines, + agent = self.client.runtimes._register_api_methods( + runtime=_genai_types.Runtime( + api_client=self.client.runtimes, api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -3653,7 +3653,7 @@ def test_query_agent_engine_stream(self, http_options_arg, expected_http_options ], ids=["default", "http_options_instance", "http_options_dict"], ) - def test_query_agent_engine_async_stream( + def test_query_runtime_async_stream( self, http_options_arg, expected_http_options ): """Async streaming query: forwards http_options to the HTTP layer (or None).""" @@ -3666,12 +3666,12 @@ async def mock_async_generator(): yield genai_types.HttpResponse(body=b"") with mock.patch.object( - self.client.agent_engines._api_client, "async_request_streamed" + self.client.runtimes._api_client, "async_request_streamed" ) as request_mock: request_mock.return_value = mock_async_generator() - agent = self.client.agent_engines._register_api_methods( - agent_engine=_genai_types.AgentEngine( - api_client=self.client.agent_engines, + agent = self.client.runtimes._register_api_methods( + runtime=_genai_types.Runtime( + api_client=self.client.runtimes, api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -3708,7 +3708,7 @@ async def consume(): _TEST_NO_OPERATION_REGISTRABLE_SCHEMAS, [ ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( CapitalizeEngine().query, schema_name=_TEST_DEFAULT_METHOD_NAME, ), @@ -3721,70 +3721,70 @@ async def consume(): _TEST_OPERATION_REGISTRABLE_SCHEMAS, [ ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().query, schema_name=_TEST_DEFAULT_METHOD_NAME, ), _TEST_STANDARD_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_method, schema_name=_TEST_CUSTOM_METHOD_NAME, ), _TEST_STANDARD_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().async_query, schema_name=_TEST_DEFAULT_ASYNC_METHOD_NAME, ), _TEST_ASYNC_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_async_method, schema_name=_TEST_CUSTOM_ASYNC_METHOD_NAME, ), _TEST_ASYNC_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().stream_query, schema_name=_TEST_DEFAULT_STREAM_METHOD_NAME, ), _TEST_STREAM_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_stream_method, schema_name=_TEST_CUSTOM_STREAM_METHOD_NAME, ), _TEST_STREAM_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().async_stream_query, schema_name=_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME, ), _TEST_ASYNC_STREAM_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_async_stream_method, schema_name=_TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME, ), _TEST_ASYNC_STREAM_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().bidi_stream_query, schema_name=_TEST_DEFAULT_BIDI_STREAM_METHOD_NAME, ), _TEST_BIDI_STREAM_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_bidi_stream_method, schema_name=_TEST_CUSTOM_BIDI_STREAM_METHOD_NAME, ), @@ -3797,7 +3797,7 @@ async def consume(): _TEST_OPERATION_NOT_REGISTERED_SCHEMAS, [ ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationNotRegisteredEngine().custom_method, schema_name=_TEST_CUSTOM_METHOD_NAME, ), @@ -3808,7 +3808,7 @@ async def consume(): ], ) @mock.patch.object(genai_client.Client, "_get_api_client") - @mock.patch.object(agent_engines.AgentEngines, "_get") + @mock.patch.object(runtimes.Runtimes, "_get") def test_operation_schemas( self, mock_get, @@ -3817,7 +3817,7 @@ def test_operation_schemas( test_class_methods_spec, want_operation_schema_api_modes, ): - test_agent_engine = _genai_types.AgentEngine( + test_runtime = _genai_types.Runtime( api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -3829,17 +3829,17 @@ def test_operation_schemas( for want_operation_schema, api_mode in want_operation_schema_api_modes: want_operation_schema[_TEST_MODE_KEY_IN_SCHEMA] = api_mode want_operation_schemas.append(want_operation_schema) - assert test_agent_engine.operation_schemas() == want_operation_schemas + assert test_runtime.operation_schemas() == want_operation_schemas - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(agent_engines.AgentEngines, "_create") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(runtimes.Runtimes, "_create") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) - def test_create_agent_engine_with_creds( + def test_create_runtime_with_creds( self, mock_get_reasoning_engine_id, mock_await_operation, @@ -3849,15 +3849,15 @@ def test_create_agent_engine_with_creds( mock_operation = mock.Mock() mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME mock_create.return_value = mock_operation - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, staging_bucket=_TEST_STAGING_BUCKET, ), @@ -3871,16 +3871,16 @@ def test_create_agent_engine_with_creds( assert mock_kwargs["credentials"] == _TEST_CREDENTIALS assert mock_kwargs["gcs_dir_name"] == "agent_engine" - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(agent_engines.AgentEngines, "_create") + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(runtimes.Runtimes, "_create") @mock.patch("google.auth.default") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) - def test_create_agent_engine_without_creds( + def test_create_runtime_without_creds( self, mock_get_reasoning_engine_id, mock_await_operation, @@ -3891,7 +3891,7 @@ def test_create_agent_engine_without_creds( mock_operation = mock.Mock() mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME mock_create.return_value = mock_operation - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, @@ -3903,9 +3903,9 @@ def test_create_agent_engine_without_creds( client = agentplatform.Client( project=_TEST_PROJECT, location=_TEST_LOCATION, credentials=mock_creds ) - client.agent_engines.create( + client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, staging_bucket=_TEST_STAGING_BUCKET, ), @@ -3919,15 +3919,15 @@ def test_create_agent_engine_without_creds( assert mock_kwargs["credentials"] == mock_creds assert mock_kwargs["gcs_dir_name"] == "agent_engine" - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(agent_engines.AgentEngines, "_create") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(runtimes.Runtimes, "_create") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) - def test_create_agent_engine_with_no_creds_in_client( + def test_create_runtime_with_no_creds_in_client( self, mock_get_reasoning_engine_id, mock_await_operation, @@ -3937,7 +3937,7 @@ def test_create_agent_engine_with_no_creds_in_client( mock_operation = mock.Mock() mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME mock_create.return_value = mock_operation - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, @@ -3946,9 +3946,9 @@ def test_create_agent_engine_with_no_creds_in_client( client = agentplatform.Client( project=_TEST_PROJECT, location=_TEST_LOCATION, credentials=None ) - client.agent_engines.create( + client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, staging_bucket=_TEST_STAGING_BUCKET, ), @@ -3964,7 +3964,7 @@ def test_create_agent_engine_with_no_creds_in_client( @pytest.mark.usefixtures("google_auth_mock") -class TestAgentEngineErrors: +class TestRuntimeErrors: def setup_method(self): importlib.reload(initializer) importlib.reload(aiplatform) @@ -3976,27 +3976,27 @@ def setup_method(self): ) self.test_agent = CapitalizeEngine() - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) - def test_create_agent_engine_error( + def test_create_runtime_error( self, mock_get_reasoning_engine_id, mock_await_operation, mock_prepare ): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( error=_TEST_AGENT_ENGINE_ERROR, ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") with pytest.raises(RuntimeError) as excinfo: - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, @@ -4008,19 +4008,19 @@ def test_create_agent_engine_error( ) assert "Failed to create agent engine" in str(excinfo.value) - @mock.patch.object(_agent_engines_utils, "_await_operation") - def test_update_agent_engine_description(self, mock_await_operation): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + @mock.patch.object(_runtimes_utils, "_await_operation") + def test_update_runtime_description(self, mock_await_operation): + mock_await_operation.return_value = _genai_types.RuntimeOperation( error=_TEST_AGENT_ENGINE_ERROR, ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") with pytest.raises(RuntimeError) as excinfo: - self.client.agent_engines.update( + self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, - config=_genai_types.AgentEngineConfig( + config=_genai_types.AgentRuntimeConfig( description=_TEST_AGENT_ENGINE_DESCRIPTION, ), ) @@ -4078,7 +4078,7 @@ def test_update_agent_engine_description(self, mock_await_operation): ], ) @pytest.mark.usefixtures("caplog") - @mock.patch.object(agent_engines.AgentEngines, "_get") + @mock.patch.object(runtimes.Runtimes, "_get") def test_invalid_operation_schema( self, mock_get, @@ -4091,7 +4091,7 @@ def test_invalid_operation_schema( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec(class_methods=test_operation_schemas), ) - self.client.agent_engines.get(name=_TEST_AGENT_ENGINE_RESOURCE_NAME) + self.client.runtimes.get(name=_TEST_AGENT_ENGINE_RESOURCE_NAME) assert want_log_output in caplog.text @pytest.mark.parametrize( @@ -4158,12 +4158,12 @@ def test_validate_resource_limits_or_raise( self, resource_limits, expected_exception, expected_message ): with pytest.raises(expected_exception) as excinfo: - _agent_engines_utils._validate_resource_limits_or_raise(resource_limits) + _runtimes_utils._validate_resource_limits_or_raise(resource_limits) assert expected_message in str(excinfo.value) @pytest.mark.usefixtures("google_auth_mock") -class TestAsyncAgentEngine: +class TestAsyncRuntime: def setup_method(self): importlib.reload(initializer) importlib.reload(aiplatform) @@ -4180,13 +4180,13 @@ def setup_method(self): def teardown_method(self): initializer.global_pool.shutdown(wait=True) - def test_delete_agent_engine(self): + def test_delete_runtime(self): with mock.patch.object( - self.client.aio.agent_engines._api_client, "async_request" + self.client.aio.runtimes._api_client, "async_request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") asyncio.run( - self.client.aio.agent_engines.delete( + self.client.aio.runtimes.delete( name=_TEST_AGENT_ENGINE_RESOURCE_NAME ) ) @@ -4197,13 +4197,13 @@ def test_delete_agent_engine(self): None, ) - def test_delete_agent_engine_force(self): + def test_delete_runtime_force(self): with mock.patch.object( - self.client.aio.agent_engines._api_client, "async_request" + self.client.aio.runtimes._api_client, "async_request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") asyncio.run( - self.client.aio.agent_engines.delete( + self.client.aio.runtimes.delete( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, force=True, ) diff --git a/tests/unit/agentplatform/genai/test_evals.py b/tests/unit/agentplatform/genai/test_evals.py index 39c5224f36..a44d0a061a 100644 --- a/tests/unit/agentplatform/genai/test_evals.py +++ b/tests/unit/agentplatform/genai/test_evals.py @@ -2422,8 +2422,8 @@ def setup_method(self): importlib.reload(_evals_metric_handlers) importlib.reload(_genai.evals) - if hasattr(_evals_common._thread_local_data, "agent_engine_instances"): - del _evals_common._thread_local_data.agent_engine_instances + if hasattr(_evals_common._thread_local_data, "runtime_instances"): + del _evals_common._thread_local_data.runtime_instances agentplatform.init( project=_TEST_PROJECT, @@ -3311,7 +3311,7 @@ def test_inference_with_multimodal_content( @mock.patch.object(_evals_utils, "EvalDatasetLoader") @mock.patch.object(_evals_common.agentplatform, "Client") - def test_run_inference_with_agent_engine_and_session_inputs_dict( + def test_run_inference_with_runtime_and_session_inputs_dict( self, mock_agentplatform_client, mock_eval_dataset_loader, @@ -3331,8 +3331,8 @@ def test_run_inference_with_agent_engine_and_session_inputs_dict( orient="records" ) - mock_agent_engine = mock.Mock() - mock_agent_engine.create_session.return_value = {"id": "session1"} + mock_runtime = mock.Mock() + mock_runtime.create_session.return_value = {"id": "session1"} stream_query_return_value = [ { "id": "1", @@ -3348,9 +3348,9 @@ def test_run_inference_with_agent_engine_and_session_inputs_dict( }, ] - mock_agent_engine.stream_query.return_value = iter(stream_query_return_value) - mock_agentplatform_client.return_value.agent_engines.get.return_value = ( - mock_agent_engine + mock_runtime.stream_query.return_value = iter(stream_query_return_value) + mock_agentplatform_client.return_value.runtimes.get.return_value = ( + mock_runtime ) inference_result = self.client.evals.run_inference( @@ -3359,13 +3359,13 @@ def test_run_inference_with_agent_engine_and_session_inputs_dict( ) mock_eval_dataset_loader.return_value.load.assert_called_once_with(mock_df) - mock_agentplatform_client.return_value.agent_engines.get.assert_called_once_with( + mock_agentplatform_client.return_value.runtimes.get.assert_called_once_with( name="projects/test-project/locations/us-central1/reasoningEngines/123" ) - mock_agent_engine.create_session.assert_called_once_with( + mock_runtime.create_session.assert_called_once_with( user_id="123", state={"a": "1"} ) - mock_agent_engine.stream_query.assert_called_once_with( + mock_runtime.stream_query.assert_called_once_with( user_id="123", session_id="session1", message="agent prompt" ) @@ -3419,12 +3419,12 @@ def test_run_inference_with_agent_engine_and_session_inputs_dict( } ), ) - assert inference_result.candidate_name == "agent_engine_0" + assert inference_result.candidate_name == "runtime_0" assert inference_result.gcs_source is None @mock.patch.object(_evals_utils, "EvalDatasetLoader") @mock.patch.object(_evals_common.agentplatform, "Client") - def test_run_inference_with_agent_engine_and_session_inputs_literal_string( + def test_run_inference_with_runtime_and_session_inputs_literal_string( self, mock_agentplatform_client, mock_eval_dataset_loader, @@ -3440,8 +3440,8 @@ def test_run_inference_with_agent_engine_and_session_inputs_literal_string( orient="records" ) - mock_agent_engine = mock.Mock() - mock_agent_engine.create_session.return_value = {"id": "session1"} + mock_runtime = mock.Mock() + mock_runtime.create_session.return_value = {"id": "session1"} stream_query_return_value = [ { "id": "1", @@ -3457,9 +3457,9 @@ def test_run_inference_with_agent_engine_and_session_inputs_literal_string( }, ] - mock_agent_engine.stream_query.return_value = iter(stream_query_return_value) - mock_agentplatform_client.return_value.agent_engines.get.return_value = ( - mock_agent_engine + mock_runtime.stream_query.return_value = iter(stream_query_return_value) + mock_agentplatform_client.return_value.runtimes.get.return_value = ( + mock_runtime ) inference_result = self.client.evals.run_inference( @@ -3468,13 +3468,13 @@ def test_run_inference_with_agent_engine_and_session_inputs_literal_string( ) mock_eval_dataset_loader.return_value.load.assert_called_once_with(mock_df) - mock_agentplatform_client.return_value.agent_engines.get.assert_called_once_with( + mock_agentplatform_client.return_value.runtimes.get.assert_called_once_with( name="projects/test-project/locations/us-central1/reasoningEngines/123" ) - mock_agent_engine.create_session.assert_called_once_with( + mock_runtime.create_session.assert_called_once_with( user_id="123", state={"a": "1"} ) - mock_agent_engine.stream_query.assert_called_once_with( + mock_runtime.stream_query.assert_called_once_with( user_id="123", session_id="session1", message="agent prompt" ) @@ -3523,12 +3523,12 @@ def test_run_inference_with_agent_engine_and_session_inputs_literal_string( } ), ) - assert inference_result.candidate_name == "agent_engine_0" + assert inference_result.candidate_name == "runtime_0" assert inference_result.gcs_source is None @mock.patch.object(_evals_utils, "EvalDatasetLoader") @mock.patch.object(_evals_common.agentplatform, "Client") - def test_run_inference_with_agent_engine_with_response_column_raises_error( + def test_run_inference_with_runtime_with_response_column_raises_error( self, mock_agentplatform_client, mock_eval_dataset_loader, @@ -3549,9 +3549,9 @@ def test_run_inference_with_agent_engine_with_response_column_raises_error( orient="records" ) - mock_agent_engine = mock.Mock() - mock_agentplatform_client.return_value.agent_engines.get.return_value = ( - mock_agent_engine + mock_runtime = mock.Mock() + mock_agentplatform_client.return_value.runtimes.get.return_value = ( + mock_runtime ) with pytest.raises(ValueError) as excinfo: @@ -3566,7 +3566,7 @@ def test_run_inference_with_agent_engine_with_response_column_raises_error( @mock.patch.object(_evals_utils, "EvalDatasetLoader") @mock.patch.object(_evals_common.agentplatform, "Client") - def test_run_inference_with_agent_engine_falls_back_to_managed_sessions_api( + def test_run_inference_with_runtime_falls_back_to_managed_sessions_api( self, mock_agentplatform_client, mock_eval_dataset_loader, @@ -3590,10 +3590,10 @@ def test_run_inference_with_agent_engine_falls_back_to_managed_sessions_api( # Create a mock agent engine WITHOUT create_session (simulates agents # deployed via Console, gcloud, or source code deployment). - mock_agent_engine = mock.Mock( + mock_runtime = mock.Mock( spec=["api_client", "api_resource", "stream_query"], ) - mock_agent_engine.api_resource.name = ( + mock_runtime.api_resource.name = ( "projects/test-project/locations/us-central1/reasoningEngines/123" ) @@ -3603,7 +3603,7 @@ def test_run_inference_with_agent_engine_falls_back_to_managed_sessions_api( "projects/test-project/locations/us-central1" "/reasoningEngines/123/sessions/managed-session-1" ) - mock_agent_engine.api_client.sessions.create.return_value = ( + mock_runtime.api_client.sessions.create.return_value = ( mock_session_operation ) @@ -3621,9 +3621,9 @@ def test_run_inference_with_agent_engine_falls_back_to_managed_sessions_api( "author": "model", }, ] - mock_agent_engine.stream_query.return_value = iter(stream_query_return_value) - mock_agentplatform_client.return_value.agent_engines.get.return_value = ( - mock_agent_engine + mock_runtime.stream_query.return_value = iter(stream_query_return_value) + mock_agentplatform_client.return_value.runtimes.get.return_value = ( + mock_runtime ) inference_result = self.client.evals.run_inference( @@ -3632,17 +3632,17 @@ def test_run_inference_with_agent_engine_falls_back_to_managed_sessions_api( ) # Verify the managed Sessions API was called as fallback. - mock_agent_engine.api_client.sessions.create.assert_called_once_with( + mock_runtime.api_client.sessions.create.assert_called_once_with( name="projects/test-project/locations/us-central1/reasoningEngines/123", user_id="123", - config=agentplatform_genai_types.CreateAgentEngineSessionConfig( + config=agentplatform_genai_types.CreateRuntimeSessionConfig( session_state={"a": "1"}, ), ) # Verify stream_query was called with the session ID extracted from # the managed session's resource name. - mock_agent_engine.stream_query.assert_called_once_with( + mock_runtime.stream_query.assert_called_once_with( user_id="123", session_id="managed-session-1", message="agent prompt", @@ -3650,7 +3650,7 @@ def test_run_inference_with_agent_engine_falls_back_to_managed_sessions_api( # Verify the inference results are correct. assert inference_result.eval_dataset_df["response"].iloc[0] == "agent response" - assert inference_result.candidate_name == "agent_engine_0" + assert inference_result.candidate_name == "runtime_0" @mock.patch.object(_evals_utils, "EvalDatasetLoader") def test_run_inference_with_local_agent( @@ -4280,7 +4280,7 @@ def test_run_inference_gemini_agent_continues_on_failure( @mock.patch.object(_evals_common, "_get_interactions_client") @mock.patch.object(_evals_utils, "EvalDatasetLoader") @mock.patch.object(_evals_common.agentplatform, "Client") - def test_run_inference_non_gemini_agent_routes_to_agent_engine( + def test_run_inference_non_gemini_agent_routes_to_runtime( self, mock_agentplatform_client, mock_eval_dataset_loader, @@ -4291,9 +4291,9 @@ def test_run_inference_non_gemini_agent_routes_to_agent_engine( orient="records" ) - mock_agent_engine = mock.Mock() - mock_agent_engine.create_session.return_value = {"id": "session1"} - mock_agent_engine.stream_query.return_value = iter( + mock_runtime = mock.Mock() + mock_runtime.create_session.return_value = {"id": "session1"} + mock_runtime.stream_query.return_value = iter( [ { "id": "1", @@ -4303,8 +4303,8 @@ def test_run_inference_non_gemini_agent_routes_to_agent_engine( } ] ) - mock_agentplatform_client.return_value.agent_engines.get.return_value = ( - mock_agent_engine + mock_agentplatform_client.return_value.runtimes.get.return_value = ( + mock_runtime ) self.client.evals.run_inference( @@ -4312,7 +4312,7 @@ def test_run_inference_non_gemini_agent_routes_to_agent_engine( agent=_TEST_AGENT_ENGINE, ) - mock_agentplatform_client.return_value.agent_engines.get.assert_called_once_with( + mock_agentplatform_client.return_value.runtimes.get.assert_called_once_with( name=_TEST_AGENT_ENGINE ) mock_get_interactions_client.assert_not_called() @@ -4717,7 +4717,7 @@ def mock_execute(*args, **kwargs): _evals_common._run_agent( api_client=mock_api_client_fixture, - agent_engine=mock.Mock(), + runtime=mock.Mock(), agent=None, prompt_dataset=prompt_dataset, user_simulator_config=user_simulator_config, @@ -4740,7 +4740,7 @@ def test_run_agent_does_not_raise_for_gemini_3_model( os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1" _evals_common._run_agent( api_client=mock_api_client_fixture, - agent_engine=mock.Mock(), + runtime=mock.Mock(), agent=None, prompt_dataset=prompt_dataset, user_simulator_config=user_simulator_config, @@ -4774,11 +4774,11 @@ def test_run_agent_internal_success(self, mock_run_agent): ] ] prompt_dataset = pd.DataFrame({"prompt": ["prompt1"]}) - mock_agent_engine = mock.Mock() + mock_runtime = mock.Mock() mock_api_client = mock.Mock() result_df = _evals_common._run_agent_internal( api_client=mock_api_client, - agent_engine=mock_agent_engine, + runtime=mock_runtime, agent=None, prompt_dataset=prompt_dataset, ) @@ -4830,11 +4830,11 @@ def test_run_agent_internal_success(self, mock_run_agent): def test_run_agent_internal_error_response(self, mock_run_agent): mock_run_agent.return_value = [{"error": "agent run failed"}] prompt_dataset = pd.DataFrame({"prompt": ["prompt1"]}) - mock_agent_engine = mock.Mock() + mock_runtime = mock.Mock() mock_api_client = mock.Mock() result_df = _evals_common._run_agent_internal( api_client=mock_api_client, - agent_engine=mock_agent_engine, + runtime=mock_runtime, agent=None, prompt_dataset=prompt_dataset, ) @@ -4853,11 +4853,11 @@ def test_run_agent_internal_multi_turn_success(self, mock_run_agent): ] ] prompt_dataset = pd.DataFrame({"prompt": ["p1"], "conversation_plan": ["plan"]}) - mock_agent_engine = mock.Mock() + mock_runtime = mock.Mock() mock_api_client = mock.Mock() result_df = _evals_common._run_agent_internal( api_client=mock_api_client, - agent_engine=mock_agent_engine, + runtime=mock_runtime, agent=None, prompt_dataset=prompt_dataset, ) @@ -4886,7 +4886,7 @@ def test_run_agent_internal_multi_turn_with_agent(self, mock_run_agent): mock_api_client = mock.Mock() result_df = _evals_common._run_agent_internal( api_client=mock_api_client, - agent_engine=None, + runtime=None, agent=mock_agent, prompt_dataset=prompt_dataset, ) @@ -5024,11 +5024,11 @@ def test_run_agent_internal_malformed_event(self, mock_run_agent): ] ] prompt_dataset = pd.DataFrame({"prompt": ["prompt1"]}) - mock_agent_engine = mock.Mock() + mock_runtime = mock.Mock() mock_api_client = mock.Mock() result_df = _evals_common._run_agent_internal( api_client=mock_api_client, - agent_engine=mock_agent_engine, + runtime=mock_runtime, agent=None, prompt_dataset=prompt_dataset, ) @@ -10684,7 +10684,7 @@ class TestIsGeminiAgentResource: def test_gemini_agent_resource_is_detected(self): assert _evals_common._is_gemini_agent_resource(_TEST_GEMINI_AGENT) is True - def test_agent_engine_resource_is_not_gemini(self): + def test_runtime_resource_is_not_gemini(self): assert _evals_common._is_gemini_agent_resource(_TEST_AGENT_ENGINE) is False def test_non_resource_string_is_not_gemini(self): @@ -10888,7 +10888,7 @@ def test_create_evaluation_run_builds_gemini_agent_config(self): assert "gemini-agent" in inference_configs assert _evals_common._DEFAULT_CANDIDATE_NAME not in inference_configs - def test_create_evaluation_run_agent_engine_does_not_set_gemini(self): + def test_create_evaluation_run_runtime_does_not_set_gemini(self): evals_module = evals.Evals(api_client_=self.mock_api_client) evals_module.create_evaluation_run( @@ -10980,7 +10980,7 @@ def test_create_evaluation_run_no_agent_no_agent_info_no_inference(self): "inferenceConfigs" ) - def test_create_evaluation_run_agent_engine_without_agent_info(self): + def test_create_evaluation_run_runtime_without_agent_info(self): """Agent Engine resource alone triggers inference_configs auto-construction.""" evals_module = evals.Evals(api_client_=self.mock_api_client) diff --git a/tests/unit/agentplatform/genai/test_live_agent_engines.py b/tests/unit/agentplatform/genai/test_live_agent_engines.py index 8af10749a3..49dceb25d0 100644 --- a/tests/unit/agentplatform/genai/test_live_agent_engines.py +++ b/tests/unit/agentplatform/genai/test_live_agent_engines.py @@ -21,7 +21,7 @@ from google.cloud import aiplatform import agentplatform from google.cloud.aiplatform import initializer as aiplatform_initializer -from agentplatform._genai import live_agent_engines +from agentplatform._genai import live_runtimes import pytest @@ -30,7 +30,7 @@ pytestmark = pytest.mark.usefixtures("google_auth_mock") -class TestLiveAgentEngines: +class TestLiveRuntimes: """Unit tests for the GenAI client.""" def setup_method(self): @@ -44,12 +44,12 @@ def setup_method(self): @pytest.mark.asyncio @pytest.mark.usefixtures("google_auth_mock") - @mock.patch.object(live_agent_engines, "ws_connect") + @mock.patch.object(live_runtimes, "ws_connect") @mock.patch.object(google.auth, "default") - async def test_async_live_agent_engines_connect( + async def test_async_live_runtimes_connect( self, mock_auth_default, mock_ws_connect ): - """Tests the AsyncLiveAgentEngines.connect method, as well as the AsyncLiveAgentEngineSession methods.""" + """Tests the AsyncLiveRuntimes.connect method, as well as the AsyncLiveRuntimeSession methods.""" # Mock credentials to avoid refresh issues mock_creds = mock.Mock(spec=google.auth.credentials.Credentials) mock_creds.token = "test-token" @@ -67,8 +67,8 @@ async def test_async_live_agent_engines_connect( json.dumps({"output": "WORLD"}).encode("utf-8"), ] - async with test_client.aio.live.agent_engines.connect( - agent_engine="test-agent-engine", + async with test_client.aio.live.runtimes.connect( + runtime="test-agent-engine", config={"class_method": "bidi_stream_query", "input": {"query": "hello"}}, ) as session: assert session is not None diff --git a/tests/unit/agentplatform/genai/test_sandbox.py b/tests/unit/agentplatform/genai/test_sandbox.py index 8f3f331d2c..32a4315c50 100644 --- a/tests/unit/agentplatform/genai/test_sandbox.py +++ b/tests/unit/agentplatform/genai/test_sandbox.py @@ -84,7 +84,7 @@ def test_send_command(self, mock_get_api_client): body=b"{}", headers={} ) - self.client.agent_engines.sandboxes.send_command( + self.client.sandboxes.send_command( http_method="GET", access_token="test_token", sandbox_environment=mock_sandbox, @@ -120,7 +120,7 @@ def test_generate_browser_ws_headers( body=b'{"endpoint": "test/endpoint"}', headers={} ) ws_url, headers = ( - self.client.agent_engines.sandboxes.generate_browser_ws_headers( + self.client.sandboxes.generate_browser_ws_headers( sandbox_environment=mock_sandbox, service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL, timeout=3600,