From 48835249015ef7b9fa55a3ef7640c148c3fcd1a0 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 moved to top-level modules. BREAKING_CHANGE(agentplatform): Removed global initializer dependency. Agent frameworks now read project and location from environment variables. BREAKING_CHANGE(agentplatform): runtime templates import changed from agentplatform.agent_engines.templates to agentplatform.frameworks. BREAKING_CHANGE(agentplatform): evals.run_inference(agent=...) now accepts types.Runtime instead of types.AgentEngine. PiperOrigin-RevId: 957042628 --- agentplatform/_genai/_evals_common.py | 130 +- ...nt_engines_utils.py => _runtimes_utils.py} | 189 +- agentplatform/_genai/agent_engines.py | 3931 -- 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/runtime_revisions.py | 328 +- agentplatform/_genai/runtimes.py | 3565 +- agentplatform/_genai/sandbox_snapshots.py | 68 +- agentplatform/_genai/sandbox_templates.py | 46 +- agentplatform/_genai/sandboxes.py | 234 +- agentplatform/_genai/session_events.py | 120 +- agentplatform/_genai/sessions.py | 322 +- agentplatform/_genai/types/__init__.py | 828 +- agentplatform/_genai/types/common.py | 34148 ++++++++-------- 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 | 63 +- .../templates => frameworks}/ag2.py | 85 +- .../templates => frameworks}/langchain.py | 40 +- .../templates => frameworks}/langgraph.py | 40 +- .../templates => frameworks}/llama_index.py | 44 +- noxfile.py | 7 +- .../frameworks/test_frameworks_a2a.py | 292 +- .../frameworks/test_frameworks_adk.py | 81 +- .../frameworks/test_frameworks_ag2.py | 74 +- .../frameworks/test_frameworks_langchain.py | 50 +- .../frameworks/test_frameworks_langgraph.py | 66 +- .../frameworks/test_frameworks_llama_index.py | 82 +- .../agentplatform/genai/replays/conftest.py | 18 +- .../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 | 6 +- .../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_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_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 +- .../replays/test_get_ae_runtime_revision.py | 52 +- .../replays/test_get_agent_engine_a2a_task.py | 89 - .../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_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_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 +- .../genai/replays/test_run_inference.py | 4 +- .../genai/replays/test_update_agent_engine.py | 20 +- .../test_update_agent_engine_session.py | 50 +- .../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 | 810 +- tests/unit/agentplatform/genai/test_evals.py | 116 +- .../genai/test_live_agent_engines.py | 14 +- .../unit/agentplatform/genai/test_sandbox.py | 16 +- 105 files changed, 22664 insertions(+), 27830 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 (97%) rename agentplatform/{agent_engines/templates => frameworks}/ag2.py (90%) rename agentplatform/{agent_engines/templates => frameworks}/langchain.py (96%) rename agentplatform/{agent_engines/templates => frameworks}/langgraph.py (95%) rename agentplatform/{agent_engines/templates => frameworks}/llama_index.py (94%) 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 e64c168666..40533d383b 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( @@ -1804,7 +1804,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[ @@ -1834,7 +1834,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 @@ -1851,7 +1851,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(): @@ -1907,29 +1907,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( @@ -1944,7 +1944,7 @@ def agent_run_wrapper( # type: ignore[no-untyped-def] agent_run_wrapper, row, contents, - agent_engine, + runtime, agent, inference_fn, api_client, @@ -2559,7 +2559,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, @@ -2577,8 +2577,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. @@ -2599,10 +2599,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) @@ -2665,27 +2664,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 @@ -2701,7 +2699,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, @@ -2716,7 +2714,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" @@ -3316,7 +3314,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, @@ -3325,7 +3323,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, @@ -3367,7 +3365,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, @@ -3386,10 +3384,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, @@ -3407,12 +3405,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: @@ -3424,7 +3422,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. @@ -3435,7 +3433,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, ) @@ -3448,18 +3446,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, ), ) @@ -3481,7 +3479,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.""" @@ -3497,8 +3495,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, ) @@ -3516,19 +3514,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, ), ) @@ -3537,7 +3535,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, @@ -4159,7 +4157,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 da2445c54e..b5f0b3d541 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: @@ -256,25 +256,25 @@ def has_field(obj: Union[BaseModel, JsonDict], field_name: str) -> bool: @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] @@ -283,7 +283,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] @@ -292,7 +292,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( @@ -303,7 +303,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: @@ -332,7 +332,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, @@ -348,7 +348,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): @@ -461,20 +461,18 @@ class _RequirementsValidationResult(TypedDict): actions: _RequirementsValidationActions -AgentEngineOperationUnion = Union[genai_types.AgentEngineOperation] +RuntimeOperationUnion = Union[genai_types.RuntimeOperation] 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 @@ -506,8 +504,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}`" ) @@ -517,11 +514,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]]): @@ -545,11 +542,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]): @@ -585,7 +582,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: @@ -627,7 +624,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. @@ -644,7 +641,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: @@ -827,7 +824,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. @@ -839,8 +836,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. @@ -887,7 +884,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): @@ -1001,7 +998,7 @@ def _parse_constraints( def _prepare( *, - agent: Optional[_AgentEngineInterface], + agent: Optional[_RuntimeInterface], requirements: Optional[Sequence[str]], extra_packages: Optional[Sequence[str]], project: str, @@ -1010,16 +1007,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. @@ -1037,7 +1034,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, @@ -1058,21 +1055,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. @@ -1080,7 +1077,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: @@ -1142,15 +1139,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( @@ -1257,13 +1252,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: @@ -1271,7 +1266,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) " @@ -1285,7 +1280,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}") @@ -1428,7 +1423,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://`") @@ -1491,11 +1486,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` @@ -1506,12 +1501,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: @@ -1519,9 +1514,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) @@ -1550,7 +1545,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`." ) @@ -1629,10 +1624,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 @@ -1642,21 +1637,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: @@ -1679,10 +1674,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 @@ -1690,16 +1685,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.") @@ -1721,10 +1716,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 @@ -1732,15 +1727,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: @@ -1765,10 +1760,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 @@ -1776,15 +1771,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: @@ -1807,10 +1802,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.' ), @@ -1826,7 +1821,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. """ @@ -2050,7 +2045,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"): @@ -2063,8 +2058,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 @@ -2085,28 +2079,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 6c802d0655..0000000000 --- a/agentplatform/_genai/agent_engines.py +++ /dev/null @@ -1,3931 +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 a2a_tasks as a2a_tasks_module - from . import runtimes as runtimes_module - - _ = sessions_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 - _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 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, - build_config=config.build_config, - ) - 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, - build_config: Optional[types.ReasoningEngineSpecBuildSpecDict] = 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_card") - 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 build_config is not None: - if agent_engine_spec is None: - agent_engine_spec = {} - build_spec: dict[str, Any] = {} - if isinstance(build_config, dict): - worker_pool = build_config.get("worker_pool") - build_service_account = build_config.get("service_account") - else: - worker_pool = getattr(build_config, "worker_pool", None) - build_service_account = getattr(build_config, "service_account", None) - if worker_pool is not None: - build_spec["worker_pool"] = worker_pool - update_masks.append("spec.build_spec.worker_pool") - if build_service_account is not None: - build_spec["service_account"] = build_service_account - update_masks.append("spec.build_spec.service_account") - if build_spec: - agent_engine_spec["build_spec"] = build_spec - - 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, - build_config=config.build_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_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 - _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 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_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_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 d2061278c1..fb9238132e 100644 --- a/agentplatform/_genai/client.py +++ b/agentplatform/_genai/client.py @@ -31,7 +31,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 @@ -48,6 +48,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, + ) + from agentplatform._genai import ( endpoints as endpoints_module, ) @@ -94,7 +99,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 @@ -104,6 +109,8 @@ def __init__(self, api_client: genai_client.BaseApiClient): # type: ignore[name self._feedback_entries: Optional[ModuleType] = None self._endpoints: Optional[ModuleType] = None self._example_stores: Optional[ModuleType] = None + self._sessions: Optional[ModuleType] = None + self._sandboxes: Optional[ModuleType] = None self._memory_banks: Optional[ModuleType] = None @property @@ -138,22 +145,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": @@ -344,7 +387,7 @@ 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 @@ -353,6 +396,8 @@ def __init__( self._feedback_entries: Optional[ModuleType] = None self._endpoints: Optional[ModuleType] = None self._example_stores: Optional[ModuleType] = None + self._sessions: Optional[ModuleType] = None + self._sandboxes: Optional[ModuleType] = None self._memory_banks: Optional[ModuleType] = None @property @@ -411,22 +456,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 cc77b0ea6e..125d77c090 100644 --- a/agentplatform/_genai/evals.py +++ b/agentplatform/_genai/evals.py @@ -3031,7 +3031,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: @@ -3054,11 +3054,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, @@ -3090,7 +3090,7 @@ def run_inference( " populated." ) - agent_engine_instance = None + runtime_instance = None agent_instance = None gemini_agent_instance = None if agent: @@ -3098,15 +3098,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, @@ -3579,7 +3579,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. @@ -6219,7 +6219,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/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..8a1981bf40 100644 --- a/agentplatform/_genai/runtimes.py +++ b/agentplatform/_genai/runtimes.py @@ -13,20 +13,30 @@ # 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 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 runtime_revisions as runtime_revisions_module - _ = runtime_revisions_module + __ = runtime_revisions_module logger = logging.getLogger("agentplatform_genai.runtimes") @@ -34,8 +44,1433 @@ 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 _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) + + 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 + _revisions = None @property @@ -49,17 +1484,2133 @@ 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] + 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, + build_config=config.build_config, + ) + 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, + build_config: Optional[types.ReasoningEngineSpecBuildSpecDict] = 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_card") + 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 build_config is not None: + if runtime_spec is None: + runtime_spec = {} + build_spec: dict[str, Any] = {} + if isinstance(build_config, dict): + worker_pool = build_config.get("worker_pool") + build_service_account = build_config.get("service_account") + else: + worker_pool = getattr(build_config, "worker_pool", None) + build_service_account = getattr(build_config, "service_account", None) + if worker_pool is not None: + build_spec["worker_pool"] = worker_pool + update_masks.append("spec.build_spec.worker_pool") + if build_service_account is not None: + build_spec["service_account"] = build_service_account + update_masks.append("spec.build_spec.service_account") + if build_spec: + runtime_spec["build_spec"] = build_spec + + 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, + build_config=config.build_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 + 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) + + 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 + _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 +3622,8 @@ 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] diff --git a/agentplatform/_genai/sandbox_snapshots.py b/agentplatform/_genai/sandbox_snapshots.py index adc25c84de..70390747db 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]: @@ -66,7 +66,7 @@ def _CreateSandboxEnvironmentSnapshotRequestParameters_to_vertex( ) if getv(from_object, ["config"]) is not None: - _CreateAgentEngineSandboxSnapshotConfig_to_vertex( + _CreateRuntimeSandboxSnapshotConfig_to_vertex( getv(from_object, ["config"]), to_object ) @@ -84,7 +84,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]: @@ -149,8 +149,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. @@ -158,7 +158,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. """ @@ -203,7 +203,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=( { @@ -456,9 +456,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, ) @@ -469,7 +469,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") @@ -498,7 +498,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=( { @@ -526,34 +526,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, @@ -571,17 +571,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", @@ -596,7 +596,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 @@ -613,7 +613,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. @@ -631,8 +631,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. @@ -640,7 +640,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. """ @@ -687,7 +687,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=( { @@ -946,9 +946,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, ) @@ -959,7 +959,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") @@ -990,7 +990,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 9dfd08dcbd..9ab6a6cc32 100644 --- a/agentplatform/_genai/sandboxes.py +++ b/agentplatform/_genai/sandboxes.py @@ -35,7 +35,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") @@ -43,7 +43,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]: @@ -78,7 +78,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]: @@ -90,14 +90,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]: @@ -108,7 +106,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]: @@ -122,7 +120,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]: @@ -135,7 +133,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]: @@ -146,7 +144,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]: @@ -164,7 +162,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]: @@ -173,9 +171,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 @@ -187,13 +183,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, @@ -205,7 +201,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") @@ -234,7 +230,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=( { @@ -262,19 +258,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, ) @@ -285,7 +281,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") @@ -314,7 +310,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=( { @@ -343,13 +339,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, @@ -361,7 +357,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") @@ -415,13 +411,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 @@ -430,7 +423,7 @@ def _get( """ - parameter_model = types._GetAgentEngineSandboxRequestParameters( + parameter_model = types._GetRuntimeSandboxRequestParameters( name=name, config=config, ) @@ -441,7 +434,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") @@ -498,23 +491,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, ) @@ -525,7 +518,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") @@ -554,7 +547,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=( { @@ -582,9 +575,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, ) @@ -595,7 +588,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") @@ -624,7 +617,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=( { @@ -660,7 +653,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 @@ -675,7 +668,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 @@ -687,35 +680,35 @@ 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 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) # A sandbox environment must be provided inline via `spec` (with an # environment set), or by referencing an existing template or snapshot in # `config`. spec_has_environment = any( - _agent_engines_utils.has_field(spec, field_name) + _runtimes_utils.has_field(spec, field_name) for field_name in ( "code_execution_environment", "computer_use_environment", @@ -751,7 +744,7 @@ def create( ) for field_name, category, display_name in environments: - if not _agent_engines_utils.has_field(spec, field_name): + if not _runtimes_utils.has_field(spec, field_name): continue if ( @@ -785,7 +778,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_sandbox_operation, poll_interval_seconds=poll_interval_seconds, @@ -801,19 +794,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", @@ -827,17 +820,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: @@ -895,15 +888,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. """ @@ -913,15 +906,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) @@ -1130,13 +1123,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, @@ -1148,7 +1141,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") @@ -1179,7 +1172,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=( { @@ -1207,19 +1200,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, ) @@ -1230,7 +1223,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") @@ -1261,7 +1254,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=( { @@ -1290,13 +1283,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, @@ -1308,7 +1301,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") @@ -1364,13 +1357,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 @@ -1379,7 +1369,7 @@ async def _get( """ - parameter_model = types._GetAgentEngineSandboxRequestParameters( + parameter_model = types._GetRuntimeSandboxRequestParameters( name=name, config=config, ) @@ -1390,7 +1380,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") @@ -1449,23 +1439,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, ) @@ -1476,7 +1466,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") @@ -1507,7 +1497,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=( { @@ -1535,9 +1525,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, ) @@ -1548,7 +1538,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") @@ -1579,7 +1569,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 f9e09d3446..72877210f8 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, + config: Optional[types.UpdateRuntimeSessionConfigOrDict] = None, ) -> types.Session: """ - 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: - types.Session: The updated Agent Engine session. + types.Session: The updated 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") @@ -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 '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,24 +765,24 @@ def update( self, *, name: str, - config: Optional[types.UpdateAgentEngineSessionConfigOrDict] = None, + config: Optional[types.UpdateRuntimeSessionConfigOrDict] = None, ) -> types.Session: - """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: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (UpdateAgentEngineSessionConfig): + Required. The name of the Agent Runtime session to be updated. Format: + `projects/{project}/locations/{location}/runtimes/{resource_id}/sessions/{session_id}`. + config (UpdateRuntimeSessionConfig): Optional. The configuration for the session to update. Returns: - Session: The updated Agent Engine session. + Session: The updated Agent Runtime session. """ if config is None: - config = types.UpdateAgentEngineSessionConfig() + config = types.UpdateRuntimeSessionConfig() elif isinstance(config, dict): - config = types.UpdateAgentEngineSessionConfig.model_validate(config) + config = types.UpdateRuntimeSessionConfig.model_validate(config) return self._update( name=name, config=config, @@ -801,14 +792,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: @@ -830,24 +821,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, @@ -859,7 +850,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") @@ -890,7 +881,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=( { @@ -918,23 +909,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, ) @@ -945,7 +936,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") @@ -976,7 +967,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=( { @@ -1001,26 +992,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, ) @@ -1031,7 +1019,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") @@ -1090,23 +1078,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, ) @@ -1117,7 +1105,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") @@ -1176,9 +1164,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, ) @@ -1189,7 +1177,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") @@ -1220,7 +1208,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=( { @@ -1248,23 +1236,23 @@ async def _update( self, *, name: str, - config: Optional[types.UpdateAgentEngineSessionConfigOrDict] = None, + config: Optional[types.UpdateRuntimeSessionConfigOrDict] = None, ) -> types.Session: """ - 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: - types.Session: The updated Agent Engine session. + types.Session: The updated Agent Runtime session. """ - parameter_model = types._UpdateAgentEngineSessionRequestParameters( + parameter_model = types._UpdateRuntimeSessionRequestParameters( name=name, config=config, ) @@ -1275,7 +1263,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") @@ -1341,7 +1329,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 'sessions.events' module requires" "additional packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e @@ -1352,25 +1340,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, @@ -1378,7 +1366,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, @@ -1400,24 +1388,24 @@ async def update( self, *, name: str, - config: Optional[types.UpdateAgentEngineSessionConfigOrDict] = None, + config: Optional[types.UpdateRuntimeSessionConfigOrDict] = None, ) -> types.Session: - """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: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (UpdateAgentEngineSessionConfig): + Required. The name of the Agent Runtime session to be updated. Format: + `projects/{project}/locations/{location}/runtimes/{resource_id}/sessions/{session_id}`. + config (UpdateRuntimeSessionConfig): Optional. The configuration for the session to update. Returns: - Session: The updated Agent Engine session. + Session: The updated Agent Runtime session. """ if config is None: - config = types.UpdateAgentEngineSessionConfig() + config = types.UpdateRuntimeSessionConfig() elif isinstance(config, dict): - config = types.UpdateAgentEngineSessionConfig.model_validate(config) + config = types.UpdateRuntimeSessionConfig.model_validate(config) return await self._update( name=name, config=config, @@ -1427,14 +1415,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 9b1dec968e..45605224d3 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -22,17 +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 _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 @@ -46,16 +41,14 @@ from .common import _CreateMultimodalDatasetParameters from .common import _CreateRagCorpusRequestParameters from .common import _CreateRuntimeFeedbackEntryRequestParameters +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 _DeleteAgentEngineRequestParameters -from .common import _DeleteAgentEngineRuntimeRevisionRequestParameters -from .common import _DeleteAgentEngineSandboxRequestParameters -from .common import _DeleteAgentEngineSessionRequestParameters -from .common import _DeleteAgentEngineTaskRequestParameters from .common import _DeleteDatasetRequestParameters from .common import _DeleteEndpointRequestParameters from .common import _DeleteEvaluationExperimentParameters @@ -69,34 +62,29 @@ from .common import _DeleteRagCorpusRequestParameters from .common import _DeleteRagFileRequestParameters from .common import _DeleteRuntimeFeedbackEntryRequestParameters +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 _DeployRequestParameters from .common import _EvaluateInstancesRequestParameters -from .common import _ExecuteCodeAgentEngineSandboxRequestParameters +from .common import _ExecuteCodeRuntimeSandboxRequestParameters from .common import _ExportPublisherModelRequestParameters from .common import _FetchExamplesParameters from .common import _GenerateInstanceRubricsRequest from .common import _GenerateLossClustersParameters from .common import _GenerateMemoriesRequestParameters from .common import _GenerateUserScenariosParameters -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 _GetDeployOperationParameters from .common import _GetEndpointOperationParameters from .common import _GetEndpointParameters @@ -127,6 +115,14 @@ from .common import _GetRuntimeFeedbackOperationParameters from .common import _GetRuntimeFeedbackOperationParameters from .common import _GetRuntimeFeedbackRequestParameters +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 @@ -136,13 +132,6 @@ from .common import _ImportEvaluationSetParameters from .common import _ImportRagFilesRequestParameters from .common import _IngestEventsRequestParameters -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 @@ -156,6 +145,11 @@ from .common import _ListRagCorporaRequestParameters from .common import _ListRagFilesRequestParameters from .common import _ListRuntimeFeedbackEntriesRequestParameters +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 @@ -164,8 +158,8 @@ from .common import _OptimizeRequestParameters from .common import _PredictParameters from .common import _PurgeMemoriesRequestParameters -from .common import _QueryAgentEngineRequestParameters -from .common import _QueryAgentEngineRuntimeRevisionRequestParameters +from .common import _QueryRuntimeRequestParameters +from .common import _QueryRuntimeRevisionRequestParameters from .common import _RecommendSpecRequestParameters from .common import _RemoveExamplesParameters from .common import _RestoreVersionRequestParameters @@ -174,14 +168,12 @@ from .common import _RetrieveRagContextsRequestParameters from .common import _RetrieveSkillsRequestParameters from .common import _RollbackMemoryRequestParameters -from .common import _RunQueryJobAgentEngineConfig -from .common import _RunQueryJobAgentEngineConfigDict -from .common import _RunQueryJobAgentEngineConfigOrDict -from .common import _RunQueryJobAgentEngineRequestParameters +from .common import _RunQueryJobRuntimeConfig +from .common import _RunQueryJobRuntimeConfigDict +from .common import _RunQueryJobRuntimeConfigOrDict +from .common import _RunQueryJobRuntimeRequestParameters from .common import _SearchExamplesParameters from .common import _UndeployModelRequestParameters -from .common import _UpdateAgentEngineRequestParameters -from .common import _UpdateAgentEngineSessionRequestParameters from .common import _UpdateDatasetParameters from .common import _UpdateEvaluationExperimentParameters from .common import _UpdateMemoryRequestParameters @@ -190,50 +182,18 @@ from .common import _UpdateRagCorpusRequestParameters from .common import _UpdateRuntimeFeedbackContextRequestParameters from .common import _UpdateRuntimeFeedbackEntryRequestParameters +from .common import _UpdateRuntimeRequestParameters +from .common import _UpdateRuntimeSessionRequestParameters from .common import _UpdateSkillRequestParameters from .common import _UploadRagFileParameters from .common import _UpsertExamplesParameters -from .common import A2aPart -from .common import A2aPartDict -from .common import A2aPartOrDict -from .common import A2aTask -from .common import A2aTaskArtifact -from .common import A2aTaskArtifactDict -from .common import A2aTaskArtifactOrDict -from .common import A2aTaskDict -from .common import A2aTaskMessage -from .common import A2aTaskMessageDict -from .common import A2aTaskMessageOrDict -from .common import A2aTaskOrDict -from .common import A2aTaskState -from .common import A2aTaskStatus -from .common import A2aTaskStatusDict -from .common import A2aTaskStatusOrDict 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 AgentEngineOperation -from .common import AgentEngineOperationDict -from .common import AgentEngineOperationOrDict -from .common import AgentEngineOrDict -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 @@ -241,18 +201,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 ArrayOperator from .common import AskContextsConfig from .common import AskContextsConfigDict @@ -314,26 +268,26 @@ from .common import BlurBaselineConfig from .common import BlurBaselineConfigDict from .common import BlurBaselineConfigOrDict -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 @@ -366,21 +320,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 @@ -423,9 +362,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 @@ -458,33 +409,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 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 DeleteEndpointConfig from .common import DeleteEndpointConfigDict from .common import DeleteEndpointConfigOrDict @@ -548,12 +472,36 @@ 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 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 @@ -765,9 +713,9 @@ from .common import ExampleStoreOperationDict from .common import ExampleStoreOperationOrDict from .common import ExampleStoreOrDict -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 @@ -904,33 +852,15 @@ 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 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 GetDeployOperationConfig from .common import GetDeployOperationConfigDict from .common import GetDeployOperationConfigOrDict @@ -1000,6 +930,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 @@ -1012,6 +945,18 @@ from .common import GetRuntimeFeedbackEntryConfig from .common import GetRuntimeFeedbackEntryConfigDict from .common import GetRuntimeFeedbackEntryConfigOrDict +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 @@ -1099,39 +1044,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 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 @@ -1222,12 +1134,33 @@ 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 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 @@ -1582,15 +1515,15 @@ 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 @@ -1895,7 +1828,6 @@ from .common import RetrieveSkillsResponse from .common import RetrieveSkillsResponseDict from .common import RetrieveSkillsResponseOrDict -from .common import Role from .common import RollbackMemoryConfig from .common import RollbackMemoryConfigDict from .common import RollbackMemoryConfigOrDict @@ -1943,18 +1875,36 @@ 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 RuntimeOperation +from .common import RuntimeOperationDict +from .common import RuntimeOperationOrDict +from .common import RuntimeOrDict +from .common import RuntimeRevision +from .common import RuntimeRevisionDict +from .common import RuntimeRevisionOrDict +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 SampledShapleyAttribution from .common import SampledShapleyAttributionDict from .common import SampledShapleyAttributionOrDict @@ -2178,39 +2128,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 @@ -2293,12 +2210,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 UpdateAgentEngineSessionConfig -from .common import UpdateAgentEngineSessionConfigDict -from .common import UpdateAgentEngineSessionConfigOrDict from .common import UpdateEvaluationExperimentConfig from .common import UpdateEvaluationExperimentConfigDict from .common import UpdateEvaluationExperimentConfigOrDict @@ -2320,12 +2231,18 @@ 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 UpdateRuntimeSessionConfig +from .common import UpdateRuntimeSessionConfigDict +from .common import UpdateRuntimeSessionConfigOrDict from .common import UpdateSkillConfig from .common import UpdateSkillConfigDict from .common import UpdateSkillConfigOrDict @@ -2368,81 +2285,6 @@ from .common import XraiAttributionOrDict __all__ = [ - "DeleteAgentEngineTaskConfig", - "DeleteAgentEngineTaskConfigDict", - "DeleteAgentEngineTaskConfigOrDict", - "GetAgentEngineTaskConfig", - "GetAgentEngineTaskConfigDict", - "GetAgentEngineTaskConfigOrDict", - "TaskArtifact", - "TaskArtifactDict", - "TaskArtifactOrDict", - "TaskOutput", - "TaskOutputDict", - "TaskOutputOrDict", - "TaskMessage", - "TaskMessageDict", - "TaskMessageOrDict", - "TaskStatusDetails", - "TaskStatusDetailsDict", - "TaskStatusDetailsOrDict", - "A2aPart", - "A2aPartDict", - "A2aPartOrDict", - "A2aTaskArtifact", - "A2aTaskArtifactDict", - "A2aTaskArtifactOrDict", - "A2aTaskMessage", - "A2aTaskMessageDict", - "A2aTaskMessageOrDict", - "A2aTaskStatus", - "A2aTaskStatusDict", - "A2aTaskStatusOrDict", - "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", @@ -2890,21 +2732,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", @@ -3040,39 +2882,63 @@ "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", "CreateMemoryBankConfig", "CreateMemoryBankConfigDict", "CreateMemoryBankConfigOrDict", @@ -3505,30 +3371,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", @@ -3541,45 +3383,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", @@ -3625,15 +3467,15 @@ "ListSandboxEnvironmentTemplatesResponse", "ListSandboxEnvironmentTemplatesResponseDict", "ListSandboxEnvironmentTemplatesResponseOrDict", - "CreateAgentEngineSandboxSnapshotConfig", - "CreateAgentEngineSandboxSnapshotConfigDict", - "CreateAgentEngineSandboxSnapshotConfigOrDict", + "CreateRuntimeSandboxSnapshotConfig", + "CreateRuntimeSandboxSnapshotConfigDict", + "CreateRuntimeSandboxSnapshotConfigOrDict", "SandboxEnvironmentSnapshot", "SandboxEnvironmentSnapshotDict", "SandboxEnvironmentSnapshotOrDict", - "AgentEngineSandboxSnapshotOperation", - "AgentEngineSandboxSnapshotOperationDict", - "AgentEngineSandboxSnapshotOperationOrDict", + "RuntimeSandboxSnapshotOperation", + "RuntimeSandboxSnapshotOperationDict", + "RuntimeSandboxSnapshotOperationOrDict", "DeleteSandboxEnvironmentSnapshotConfig", "DeleteSandboxEnvironmentSnapshotConfigDict", "DeleteSandboxEnvironmentSnapshotConfigOrDict", @@ -3649,54 +3491,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", @@ -4357,21 +4199,6 @@ "EvalRunInferenceConfig", "EvalRunInferenceConfigDict", "EvalRunInferenceConfigOrDict", - "AgentEngine", - "AgentEngineDict", - "AgentEngineOrDict", - "AgentEngineConfig", - "AgentEngineConfigDict", - "AgentEngineConfigOrDict", - "RunQueryJobAgentEngineConfig", - "RunQueryJobAgentEngineConfigDict", - "RunQueryJobAgentEngineConfigOrDict", - "RunQueryJobResult", - "RunQueryJobResultDict", - "RunQueryJobResultOrDict", - "CheckQueryJobResponse", - "CheckQueryJobResponseDict", - "CheckQueryJobResponseOrDict", "AssembleDataset", "AssembleDatasetDict", "AssembleDatasetOrDict", @@ -4408,9 +4235,6 @@ "OptimizeJobConfig", "OptimizeJobConfigDict", "OptimizeJobConfigOrDict", - "AgentEngineRuntimeRevision", - "AgentEngineRuntimeRevisionDict", - "AgentEngineRuntimeRevisionOrDict", "ListDeployableModelsConfig", "ListDeployableModelsConfigDict", "ListDeployableModelsConfigOrDict", @@ -4435,9 +4259,24 @@ "ListMemoryBanksResponse", "ListMemoryBanksResponseDict", "ListMemoryBanksResponseOrDict", - "A2aTaskState", - "Role", - "State", + "Runtime", + "RuntimeDict", + "RuntimeOrDict", + "AgentRuntimeConfig", + "AgentRuntimeConfigDict", + "AgentRuntimeConfigOrDict", + "RunQueryJobRuntimeConfig", + "RunQueryJobRuntimeConfigDict", + "RunQueryJobRuntimeConfigOrDict", + "RunQueryJobResult", + "RunQueryJobResultDict", + "RunQueryJobResultOrDict", + "CheckQueryJobResponse", + "CheckQueryJobResponseDict", + "CheckQueryJobResponseOrDict", + "RuntimeRevision", + "RuntimeRevisionDict", + "RuntimeRevisionOrDict", "Strategy", "AcceleratorType", "Type", @@ -4445,6 +4284,7 @@ "ManagedTopicEnum", "IdentityType", "AgentServerMode", + "State", "MemoryType", "Operator", "RagFileType", @@ -4510,12 +4350,6 @@ "MessageDict", "Importance", "ParsedResponseUnion", - "_DeleteAgentEngineTaskRequestParameters", - "_GetAgentEngineTaskRequestParameters", - "_ListAgentEngineTasksRequestParameters", - "_CreateAgentEngineTaskRequestParameters", - "_AppendAgentEngineTaskEventRequestParameters", - "_ListAgentEngineTaskEventsRequestParameters", "_CreateEvaluationExperimentParameters", "_CreateEvaluationItemParameters", "_CreateEvaluationMetricParameters", @@ -4541,16 +4375,21 @@ "_OptimizeRequestParameters", "_CustomJobParameters", "_GetCustomJobParameters", - "_CancelQueryJobAgentEngineRequestParameters", - "_CheckQueryJobAgentEngineRequestParameters", - "_RunQueryJobAgentEngineRequestParameters", - "_CreateAgentEngineRequestParameters", - "_DeleteAgentEngineRequestParameters", - "_GetAgentEngineRequestParameters", - "_ListAgentEngineRequestParameters", - "_GetAgentEngineOperationParameters", - "_QueryAgentEngineRequestParameters", - "_UpdateAgentEngineRequestParameters", + "_CancelQueryJobRuntimeRequestParameters", + "_CheckQueryJobRuntimeRequestParameters", + "_RunQueryJobRuntimeRequestParameters", + "_CreateRuntimeRequestParameters", + "_DeleteRuntimeRequestParameters", + "_GetRuntimeRequestParameters", + "_ListRuntimeRequestParameters", + "_GetRuntimeOperationParameters", + "_QueryRuntimeRequestParameters", + "_UpdateRuntimeRequestParameters", + "_GetRuntimeRevisionRequestParameters", + "_ListRuntimeRevisionsRequestParameters", + "_DeleteRuntimeRevisionRequestParameters", + "_GetDeleteRuntimeRevisionOperationParameters", + "_QueryRuntimeRevisionRequestParameters", "_CreateMemoryBankRequestParameters", "_DeleteMemoryBankRequestParameters", "_GetMemoryBankRequestParameters", @@ -4588,17 +4427,12 @@ "_ImportRagFilesRequestParameters", "_GetImportFilesOperationParameters", "_UploadRagFileParameters", - "_GetAgentEngineRuntimeRevisionRequestParameters", - "_ListAgentEngineRuntimeRevisionsRequestParameters", - "_DeleteAgentEngineRuntimeRevisionRequestParameters", - "_GetDeleteAgentEngineRuntimeRevisionOperationParameters", - "_QueryAgentEngineRuntimeRevisionRequestParameters", - "_CreateAgentEngineSandboxRequestParameters", - "_DeleteAgentEngineSandboxRequestParameters", - "_ExecuteCodeAgentEngineSandboxRequestParameters", - "_GetAgentEngineSandboxRequestParameters", - "_ListAgentEngineSandboxesRequestParameters", - "_GetAgentEngineSandboxOperationParameters", + "_CreateRuntimeSandboxRequestParameters", + "_DeleteRuntimeSandboxRequestParameters", + "_ExecuteCodeRuntimeSandboxRequestParameters", + "_GetRuntimeSandboxRequestParameters", + "_ListRuntimeSandboxesRequestParameters", + "_GetRuntimeSandboxOperationParameters", "_CreateSandboxEnvironmentTemplateRequestParameters", "_DeleteSandboxEnvironmentTemplateRequestParameters", "_GetSandboxEnvironmentTemplateRequestParameters", @@ -4608,15 +4442,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 6d52bee882..65b0c772e2 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -91,65 +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 Role(_common.CaseInSensitiveEnum): - """The role of the sender of the message.""" - - ROLE_UNSPECIFIED = "ROLE_UNSPECIFIED" - """The role is unspecified.""" - ROLE_USER = "ROLE_USER" - """The message is from the client to the server.""" - ROLE_AGENT = "ROLE_AGENT" - """The message is from the server to the client.""" - - -class State(_common.CaseInSensitiveEnum): - """Output only. The current state of the task.""" - - STATE_UNSPECIFIED = "STATE_UNSPECIFIED" - """The task is in an unknown or indeterminate state.""" - TASK_STATE_SUBMITTED = "TASK_STATE_SUBMITTED" - """Indicates that a task has been successfully submitted and acknowledged.""" - TASK_STATE_WORKING = "TASK_STATE_WORKING" - """Indicates that a task is actively being processed by the agent.""" - TASK_STATE_COMPLETED = "TASK_STATE_COMPLETED" - """Indicates that a task has finished successfully. This is a terminal state.""" - TASK_STATE_FAILED = "TASK_STATE_FAILED" - """Indicates that a task has finished with an error. This is a terminal state.""" - TASK_STATE_CANCELED = "TASK_STATE_CANCELED" - """Indicates that a task was canceled before completion. This is a terminal state.""" - TASK_STATE_INPUT_REQUIRED = "TASK_STATE_INPUT_REQUIRED" - """Indicates that the agent requires additional user input to proceed. This is an interrupted state.""" - TASK_STATE_REJECTED = "TASK_STATE_REJECTED" - """Indicates that the agent has decided to not perform the task. This may be done during initial task creation or later once an agent has determined it can't or won't proceed. This is a terminal state.""" - TASK_STATE_AUTH_REQUIRED = "TASK_STATE_AUTH_REQUIRED" - """Indicates that authentication is required to proceed. This is an interrupted state.""" - - class Strategy(_common.CaseInSensitiveEnum): """This determines which type of scheduling strategy to use.""" @@ -289,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.""" @@ -826,2689 +778,2521 @@ 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.""" - message_id: Optional[str] - """Required. The unique identifier of the message.""" + evaluation_item_type: Optional[str] + """""" - metadata: Optional[dict[str, Any]] - """Optional. A2A message may have extension_uris or reference_task_ids. They will be stored under metadata.""" + gcs_uri: Optional[str] + """""" - parts: Optional[list[genai_types.Part]] - """The parts of the message.""" + display_name: Optional[str] + """""" - role: Optional[str] - """Required. The role of the sender of the message. e.g. "user", "agent".""" + config: Optional[CreateEvaluationItemConfigDict] + """""" -TaskMessageOrDict = Union[TaskMessage, TaskMessageDict] +_CreateEvaluationItemParametersOrDict = Union[ + _CreateEvaluationItemParameters, _CreateEvaluationItemParametersDict +] + +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 A2aPart(_common.BaseModel): - """A single part of a message or artifact. A part carries exactly one kind of content.""" +class EvaluationPrompt(_common.BaseModel): + """Represents the prompt to be evaluated.""" - data: Optional[dict[str, Any]] = Field( - default=None, description="""Optional. Structured data content.""" - ) - filename: Optional[str] = Field( + text: Optional[str] = Field(default=None, description="""Text prompt.""") + value: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. The name of the file, when the part represents a file.""", + description="""Fields and values that can be used to populate the prompt template.""", ) - media_type: Optional[str] = Field( - default=None, - description="""Optional. The IANA media type of the content, e.g. "text/plain" or "image/png".""", + prompt_template_data: Optional[PromptTemplateData] = Field( + default=None, description="""Prompt template data.""" ) - metadata: Optional[dict[str, Any]] = Field( + user_scenario: Optional[evals_types.UserScenario] = Field( default=None, - description="""Optional. Additional context or parameters related to the part. Extensions can be used to strongly type metadata values for specific use cases.""", - ) - raw: Optional[bytes] = Field( - default=None, description="""Optional. Raw binary content.""" - ) - text: Optional[str] = Field( - default=None, description="""Optional. Textual content.""" - ) - url: Optional[str] = Field( - default=None, description="""Optional. A URL pointing to the content.""" + description="""User scenario to help simulate multi-turn agent run results.""", ) -class A2aPartDict(TypedDict, total=False): - """A single part of a message or artifact. A part carries exactly one kind of content.""" - - data: Optional[dict[str, Any]] - """Optional. Structured data content.""" - - filename: Optional[str] - """Optional. The name of the file, when the part represents a file.""" - - media_type: Optional[str] - """Optional. The IANA media type of the content, e.g. "text/plain" or "image/png".""" +class EvaluationPromptDict(TypedDict, total=False): + """Represents the prompt to be evaluated.""" - metadata: Optional[dict[str, Any]] - """Optional. Additional context or parameters related to the part. Extensions can be used to strongly type metadata values for specific use cases.""" + text: Optional[str] + """Text prompt.""" - raw: Optional[bytes] - """Optional. Raw binary content.""" + value: Optional[dict[str, Any]] + """Fields and values that can be used to populate the prompt template.""" - text: Optional[str] - """Optional. Textual content.""" + prompt_template_data: Optional[PromptTemplateDataDict] + """Prompt template data.""" - url: Optional[str] - """Optional. A URL pointing to the content.""" + user_scenario: Optional[evals_types.UserScenario] + """User scenario to help simulate multi-turn agent run results.""" -A2aPartOrDict = Union[A2aPart, A2aPartDict] +EvaluationPromptOrDict = Union[EvaluationPrompt, EvaluationPromptDict] -class A2aTaskArtifact(_common.BaseModel): - """Represents a single artifact produced by a task.""" +class CandidateResponse(_common.BaseModel): + """Responses from model or agent.""" - artifact_id: Optional[str] = Field( - default=None, - description="""Required. The unique identifier of the artifact within the task.""", - ) - description: Optional[str] = Field( - default=None, - description="""Optional. A human-readable description of the artifact.""", - ) - display_name: Optional[str] = Field( + candidate: Optional[str] = Field( default=None, - description="""Optional. The human-readable name of the artifact.""", + description="""The name of the candidate that produced the response.""", ) - extensions: Optional[list[str]] = Field( + text: Optional[str] = Field(default=None, description="""The text response.""") + value: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. A2A protocol extensions associated with the artifact.""", + description="""Fields and values that can be used to populate the response template.""", ) - metadata: Optional[dict[str, Any]] = Field( + events: Optional[list[genai_types.Content]] = Field( default=None, - description="""Optional. Additional context or parameters related to the artifact. Extensions can be used to strongly type metadata values for specific use cases.""", + description="""Intermediate events (such as tool calls and responses) that led to the final response.""", ) - parts: Optional[list[A2aPart]] = Field( + agent_data: Optional[evals_types.AgentData] = Field( default=None, - description="""Required. The content parts that make up the artifact.""", + 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 A2aTaskArtifactDict(TypedDict, total=False): - """Represents a single artifact produced by a task.""" +class CandidateResponseDict(TypedDict, total=False): + """Responses from model or agent.""" - artifact_id: Optional[str] - """Required. The unique identifier of the artifact within the task.""" + candidate: Optional[str] + """The name of the candidate that produced the response.""" - description: Optional[str] - """Optional. A human-readable description of the artifact.""" + text: Optional[str] + """The text response.""" - display_name: Optional[str] - """Optional. The human-readable name of the artifact.""" + value: Optional[dict[str, Any]] + """Fields and values that can be used to populate the response template.""" - extensions: Optional[list[str]] - """Optional. A2A protocol extensions associated with the artifact.""" + events: Optional[list[genai_types.Content]] + """Intermediate events (such as tool calls and responses) that led to the final response.""" - metadata: Optional[dict[str, Any]] - """Optional. Additional context or parameters related to the artifact. Extensions can be used to strongly type metadata values for specific use cases.""" + 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.""" - parts: Optional[list[A2aPartDict]] - """Required. The content parts that make up the artifact.""" +CandidateResponseOrDict = Union[CandidateResponse, CandidateResponseDict] -A2aTaskArtifactOrDict = Union[A2aTaskArtifact, A2aTaskArtifactDict] +class RubricGroup(_common.BaseModel): + """A group of rubrics. -class A2aTaskMessage(_common.BaseModel): - """Represents a single message in a conversation, compliant with the A2A specification.""" + Used for grouping rubrics based on a metric or a version. + """ - extensions: Optional[list[str]] = Field( - default=None, - description="""Optional. A2A protocol extensions associated with the message.""", - ) - 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. Additional context or parameters related to the message. Extensions can be used to strongly type metadata values for specific use cases.""", - ) - parts: Optional[list[A2aPart]] = Field( - default=None, - description="""Required. The content parts that make up the message.""", + group_id: Optional[str] = Field( + default=None, description="""Unique identifier for the group.""" ) - reference_task_ids: Optional[list[str]] = Field( + display_name: Optional[str] = Field( default=None, - description="""Optional. The IDs of other tasks referenced by this message.""", + description="""Human-readable name for the group. This should be unique + within a given context if used for display or selection. + Example: "Instruction Following V1", "Content Quality - Summarization + Task".""", ) - role: Optional[Role] = Field( - default=None, description="""Required. The role of the sender of the message.""" + rubrics: Optional[list[evals_types.Rubric]] = Field( + default=None, description="""Rubrics that are part of this group.""" ) -class A2aTaskMessageDict(TypedDict, total=False): - """Represents a single message in a conversation, compliant with the A2A specification.""" - - extensions: Optional[list[str]] - """Optional. A2A protocol extensions associated with the message.""" - - message_id: Optional[str] - """Required. The unique identifier of the message.""" +class RubricGroupDict(TypedDict, total=False): + """A group of rubrics. - metadata: Optional[dict[str, Any]] - """Optional. Additional context or parameters related to the message. Extensions can be used to strongly type metadata values for specific use cases.""" + Used for grouping rubrics based on a metric or a version. + """ - parts: Optional[list[A2aPartDict]] - """Required. The content parts that make up the message.""" + group_id: Optional[str] + """Unique identifier for the group.""" - reference_task_ids: Optional[list[str]] - """Optional. The IDs of other tasks referenced by this message.""" + display_name: Optional[str] + """Human-readable name for the group. This should be unique + within a given context if used for display or selection. + Example: "Instruction Following V1", "Content Quality - Summarization + Task".""" - role: Optional[Role] - """Required. The role of the sender of the message.""" + rubrics: Optional[list[evals_types.Rubric]] + """Rubrics that are part of this group.""" -A2aTaskMessageOrDict = Union[A2aTaskMessage, A2aTaskMessageDict] +RubricGroupOrDict = Union[RubricGroup, RubricGroupDict] -class A2aTaskStatus(_common.BaseModel): - """Represents the status of an A2aTask.""" +class EvaluationItemRequest(_common.BaseModel): + """Single evaluation request.""" - message: Optional[A2aTaskMessage] = Field( - default=None, - description="""Output only. The status message associated with the state.""", + prompt: Optional[EvaluationPrompt] = Field( + default=None, description="""The request/prompt to evaluate.""" ) - state: Optional[State] = Field( - default=None, description="""Output only. The current state of the task.""" + golden_response: Optional[CandidateResponse] = Field( + default=None, description="""The ideal response or ground truth.""" ) - timestamp: Optional[datetime.datetime] = Field( + rubrics: 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.""", + ) + candidate_responses: Optional[list[CandidateResponse]] = Field( default=None, - description="""Output only. The time at which the state was set.""", + description="""Responses from model under test and other baseline models for comparison.""", ) -class A2aTaskStatusDict(TypedDict, total=False): - """Represents the status of an A2aTask.""" +class EvaluationItemRequestDict(TypedDict, total=False): + """Single evaluation request.""" - message: Optional[A2aTaskMessageDict] - """Output only. The status message associated with the state.""" + prompt: Optional[EvaluationPromptDict] + """The request/prompt to evaluate.""" - state: Optional[State] - """Output only. The current state of the task.""" + golden_response: Optional[CandidateResponseDict] + """The ideal response or ground truth.""" - timestamp: Optional[datetime.datetime] - """Output only. The time at which the state was set.""" + rubrics: Optional[dict[str, RubricGroupDict]] + """Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""" + candidate_responses: Optional[list[CandidateResponseDict]] + """Responses from model under test and other baseline models for comparison.""" -A2aTaskStatusOrDict = Union[A2aTaskStatus, A2aTaskStatusDict] +EvaluationItemRequestOrDict = Union[EvaluationItemRequest, EvaluationItemRequestDict] -class A2aTask(_common.BaseModel): - """A task.""" - context_id: Optional[str] = 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.""" - ) - 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}` or `projects/{project}/locations/{location}/taskStores/{task_store}/a2aTasks/{a2a_task}`""", - ) - next_event_sequence_number: Optional[int] = 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.""" - ) - state: Optional[A2aTaskState] = 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.""" - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. The last update timestamp of the task.""", - ) - expire_time: Optional[datetime.datetime] = 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""", - ) - ttl: Optional[str] = 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]""", +class EvaluationItemResult(_common.BaseModel): + """Represents the result of an evaluation item.""" + + evaluation_request: Optional[str] = Field( + default=None, description="""The request item that was evaluated.""" ) - app_id: Optional[str] = Field( + evaluation_run: Optional[str] = Field( default=None, - description="""Optional. Agent application which created the task.""", - ) - artifacts: Optional[list[A2aTaskArtifact]] = Field( - default=None, description="""Output only. The artifacts produced by the task.""" + description="""The evaluation run that was used to generate the result.""", ) - generation: Optional[int] = Field( - default=None, description="""Output only. The task generation number.""" + request: Optional[EvaluationItemRequest] = Field( + default=None, description="""The request that was evaluated.""" ) - history: Optional[list[A2aTaskMessage]] = Field( - default=None, description="""Output only. The history of the task messages.""" + metric: Optional[str] = Field( + default=None, description="""The metric that was evaluated.""" ) - status: Optional[A2aTaskStatus] = Field( - default=None, - description="""Output only. The status of the task, including the state, status message, and timestamp.""", + candidate_results: Optional[list[evals_types.CandidateResult]] = Field( + default=None, description="""The results for the metric.""" ) - user_id: Optional[str] = Field( - default=None, description="""Optional. Task owner user ID.""" + metadata: Optional[dict[str, Any]] = Field( + default=None, description="""Metadata about the evaluation result.""" ) -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).""" - - create_time: Optional[datetime.datetime] - """Output only. The creation timestamp of the task.""" - - metadata: Optional[dict[str, Any]] - """Optional. Arbitrary, user-defined metadata.""" - - name: Optional[str] - """Identifier. The resource name of the task. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/a2aTasks/{a2a_task}` or `projects/{project}/locations/{location}/taskStores/{task_store}/a2aTasks/{a2a_task}`""" - - 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.""" - - output: Optional[TaskOutputDict] - """Optional. The final output of the task.""" - - 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.""" - - 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""" +class EvaluationItemResultDict(TypedDict, total=False): + """Represents the result of an evaluation item.""" - 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]""" + evaluation_request: Optional[str] + """The request item that was evaluated.""" - app_id: Optional[str] - """Optional. Agent application which created the task.""" + evaluation_run: Optional[str] + """The evaluation run that was used to generate the result.""" - artifacts: Optional[list[A2aTaskArtifactDict]] - """Output only. The artifacts produced by the task.""" + request: Optional[EvaluationItemRequestDict] + """The request that was evaluated.""" - generation: Optional[int] - """Output only. The task generation number.""" + metric: Optional[str] + """The metric that was evaluated.""" - history: Optional[list[A2aTaskMessageDict]] - """Output only. The history of the task messages.""" + candidate_results: Optional[list[evals_types.CandidateResult]] + """The results for the metric.""" - status: Optional[A2aTaskStatusDict] - """Output only. The status of the task, including the state, status message, and timestamp.""" + metadata: Optional[dict[str, Any]] + """Metadata about the evaluation result.""" - user_id: Optional[str] - """Optional. Task owner user ID.""" +EvaluationItemResultOrDict = Union[EvaluationItemResult, EvaluationItemResultDict] -A2aTaskOrDict = Union[A2aTask, A2aTaskDict] +class EvaluationItem(_common.BaseModel): + """EvaluationItem is a single evaluation request or result. -class ListAgentEngineTasksConfig(_common.BaseModel): - """Config for listing agent engine tasks.""" + 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.""" ) - 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 EvaluationItem.""" ) - order_by: Optional[str] = Field( + 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="""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="""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.""" + ) + 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 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] - """""" + response = _camel_key_to_snake(response) + result = super()._from_response(response=response, kwargs=kwargs) + return result - 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). +class EvaluationItemDict(TypedDict, total=False): + """EvaluationItem is a single evaluation request or result. - Supported fields: - * `create_time` - * `update_time` + The content of an EvaluationItem is immutable - it cannot be updated once + created. EvaluationItems can be deleted when no longer needed. + """ - Example: `create_time desc`.""" + name: Optional[str] + """The resource name of the EvaluationItem.""" + display_name: Optional[str] + """The display name of the EvaluationItem.""" -ListAgentEngineTasksConfigOrDict = Union[ - ListAgentEngineTasksConfig, ListAgentEngineTasksConfigDict -] + metadata: Optional[dict[str, Any]] + """Metadata for the EvaluationItem.""" + labels: Optional[dict[str, str]] + """Labels for the EvaluationItem.""" -class _ListAgentEngineTasksRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" + evaluation_item_type: Optional[EvaluationItemType] + """The type of the EvaluationItem.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[ListAgentEngineTasksConfig] = Field( - default=None, description="""""" - ) + evaluation_request: Optional[EvaluationItemRequestDict] + """The request to evaluate.""" + evaluation_response: Optional[EvaluationItemResultDict] + """The response from evaluation.""" -class _ListAgentEngineTasksRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" + gcs_uri: Optional[str] + """The Cloud Storage object where the request or response is stored.""" - name: Optional[str] - """Name of the agent engine.""" + create_time: Optional[datetime.datetime] + """Timestamp when this item was created.""" - config: Optional[ListAgentEngineTasksConfigDict] - """""" + error: Optional[genai_types.GoogleRpcStatus] + """Error for the evaluation item.""" -_ListAgentEngineTasksRequestParametersOrDict = Union[ - _ListAgentEngineTasksRequestParameters, _ListAgentEngineTasksRequestParametersDict -] +EvaluationItemOrDict = Union[EvaluationItem, EvaluationItemDict] -class ListAgentEngineTasksResponse(_common.BaseModel): - """Response for listing agent engine tasks.""" +class Metric(_common.BaseModel): + """The metric used for evaluation.""" - 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 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.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - a2aTasks: Optional[list[A2aTask]] = Field( - default=None, description="""List of agent engine tasks.""" + prompt_template: Optional[str] = Field( + default=None, description="""The prompt template for the metric.""" ) - - -class ListAgentEngineTasksResponseDict(TypedDict, total=False): - """Response for listing agent engine tasks.""" - - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" - - next_page_token: Optional[str] - """""" - - a2aTasks: Optional[list[A2aTaskDict]] - """List of agent engine tasks.""" - - -ListAgentEngineTasksResponseOrDict = Union[ - ListAgentEngineTasksResponse, ListAgentEngineTasksResponseDict -] - - -class CreateAgentEngineTaskConfig(_common.BaseModel): - """Config for creating a Session.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + judge_model_system_instruction: Optional[str] = Field( + default=None, description="""The system instruction for the judge model.""" ) - context_id: Optional[str] = Field( - default=None, description="""The context id of the task to create.""" + return_raw_output: Optional[bool] = Field( + default=None, + description="""Whether to return the raw output from the judge model.""", ) - metadata: Optional[dict[str, Any]] = Field( - default=None, description="""The metadata of the task to create.""" + parse_and_reduce_fn: Optional[Callable[..., Any]] = Field( + default=None, + description="""The parse and reduce function for the judge model.""", + ) + aggregate_summary_fn: Optional[Callable[..., Any]] = Field( + default=None, + description="""The aggregate summary function for the judge model.""", ) - status_details: Optional[TaskStatusDetails] = Field( - default=None, description="""The status details of the task to create.""" + remote_custom_function: Optional[str] = Field( + default=None, + description="""The evaluation function for the custom code execution metric. This custom code is run remotely in the evaluation service.""", ) - output: Optional[TaskOutput] = Field( - default=None, description="""The output of the task to create.""" + judge_model: Optional[str] = Field( + default=None, description="""The judge model for the metric.""" ) - - -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( + judge_model_generation_config: Optional[genai_types.GenerationConfig] = Field( default=None, - description="""Name of the agent engine to create the task under.""", + description="""The generation config for the judge LLM (temperature, top_k, top_p, etc).""", ) - a2a_task_id: Optional[str] = Field( - default=None, description="""The ID of the task.""" + judge_model_sampling_count: Optional[int] = Field( + default=None, description="""The sampling count for the judge model.""" ) - config: Optional[CreateAgentEngineTaskConfig] = Field( - default=None, description="""""" + 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 _CreateAgentEngineTaskRequestParametersDict(TypedDict, total=False): - """Parameters for creating Agent Engine Tasks.""" + _is_predefined: bool = PrivateAttr(default=False) + """A boolean indicating whether the metric is predefined.""" - name: Optional[str] - """Name of the agent engine to create the task under.""" + _config_source: Optional[str] = PrivateAttr(default=None) + """An optional string indicating the source of the metric configuration.""" - a2a_task_id: Optional[str] - """The ID of the task.""" + _version: Optional[str] = PrivateAttr(default=None) + """An optional string indicating the version of the metric.""" - config: Optional[CreateAgentEngineTaskConfigDict] - """""" + @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. -_CreateAgentEngineTaskRequestParametersOrDict = Union[ - _CreateAgentEngineTaskRequestParameters, _CreateAgentEngineTaskRequestParametersDict -] + Args: + file_path: The path to the YAML file. + version: Optional version string to include in the YAML output. + 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]'." + ) -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" } }""" + 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) + } - 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""", - ) + 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 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" } }""" + with open(file_path, "w", encoding="utf-8") as f: + yaml.dump(data_to_dump, f, sort_keys=False, allow_unicode=True) - new_metadata: Optional[dict[str, Any]] - """Required. The complete state of the metadata object *after* the change.""" - 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""" +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.""", + ) -TaskMetadataChangeOrDict = Union[TaskMetadataChange, TaskMetadataChangeDict] + # 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 TaskArtifactChange(_common.BaseModel): - """Describes changes to the artifact list.""" +class LLMMetric(Metric): + """A metric that uses LLM-as-a-judge for evaluation.""" - added_artifacts: Optional[list[TaskArtifact]] = Field( - default=None, - description="""Optional. A list of brand-new artifacts created in this event.""", - ) - deleted_artifact_ids: Optional[list[str]] = Field( + rubric_group_name: Optional[str] = Field( default=None, - description="""Optional. A list of artifact IDs that were removed in this event.""", + description="""Optional. The name of the column in the EvaluationDataset containing the list of rubrics to use for this metric.""", ) - updated_artifacts: Optional[list[TaskArtifact]] = Field( + + result_parsing_function: Optional[str] = Field( default=None, - description="""Optional. A list of existing artifacts that were modified in this event.""", + 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 TaskArtifactChangeDict(TypedDict, total=False): - """Describes changes to the artifact list.""" + @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 - added_artifacts: Optional[list[TaskArtifactDict]] - """Optional. A list of brand-new artifacts created in this event.""" + @classmethod + def load(cls, config_path: str, client: Optional[Any] = None) -> "LLMMetric": + """Loads a metric configuration from a YAML or JSON file. - deleted_artifact_ids: Optional[list[str]] - """Optional. A list of artifact IDs that were removed in this 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. - updated_artifacts: Optional[list[TaskArtifactDict]] - """Optional. A list of existing artifacts that were modified in this event.""" + 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. -TaskArtifactChangeOrDict = Union[TaskArtifactChange, TaskArtifactChangeDict] + 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 TaskOutputChange(_common.BaseModel): - """An event representing a change to the task's outputs.""" + 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) - task_artifact_change: Optional[TaskArtifactChange] = Field( - default=None, - description="""Required. A granular change to the list of artifacts.""", - ) + 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 TaskOutputChangeDict(TypedDict, total=False): - """An event representing a change to the task's outputs.""" + 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) - task_artifact_change: Optional[TaskArtifactChangeDict] - """Required. A granular change to the list of artifacts.""" + if not isinstance(data, dict): + raise ValueError("Metric config content did not parse into a dictionary.") + return cls.model_validate(data) -TaskOutputChangeOrDict = Union[TaskOutputChange, TaskOutputChangeDict] +class MetricDict(TypedDict, total=False): + """The metric used for evaluation.""" -class TaskStateChange(_common.BaseModel): - """A message representing a change in a task's state.""" + name: Optional[str] + """The name of the metric.""" - new_state: Optional[State] = Field( - default=None, description="""Required. The new state of the task.""" - ) + 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 TaskStateChangeDict(TypedDict, total=False): - """A message representing a change in a task's state.""" + judge_model_system_instruction: Optional[str] + """The system instruction for the judge model.""" - new_state: Optional[State] - """Required. The new state of the task.""" + return_raw_output: Optional[bool] + """Whether to return the raw output from the judge model.""" + parse_and_reduce_fn: Optional[Callable[..., Any]] + """The parse and reduce function for the judge model.""" -TaskStateChangeOrDict = Union[TaskStateChange, TaskStateChangeDict] + 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.""" -class TaskStatusDetailsChange(_common.BaseModel): - """Represents a change to the task's status details.""" + judge_model: Optional[str] + """The judge model for the metric.""" - new_task_status: Optional[TaskStatusDetails] = Field( - default=None, - description="""Required. The complete state of the task's status *after* the change.""", - ) + 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 TaskStatusDetailsChangeDict(TypedDict, total=False): - """Represents a change to the task's status details.""" + rubric_group_name: Optional[str] + """The rubric group name for the rubric-based metric.""" - new_task_status: Optional[TaskStatusDetailsDict] - """Required. The complete state of the task's status *after* the change.""" + 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}""" -TaskStatusDetailsChangeOrDict = Union[ - TaskStatusDetailsChange, TaskStatusDetailsChangeDict -] + 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.""" -class TaskEventData(_common.BaseModel): - """Data for a TaskEvent.""" +MetricOrDict = Union[Metric, MetricDict] - 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.""", - ) +class CreateEvaluationMetricConfig(_common.BaseModel): + """Config for creating an evaluation metric.""" -class TaskEventDataDict(TypedDict, total=False): - """Data for a TaskEvent.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) - metadata_change: Optional[TaskMetadataChangeDict] - """Optional. A change to the task's metadata.""" - output_change: Optional[TaskOutputChangeDict] - """Optional. A change to the task's final outputs.""" +class CreateEvaluationMetricConfigDict(TypedDict, total=False): + """Config for creating an evaluation metric.""" - state_change: Optional[TaskStateChangeDict] - """Optional. A change in the task's state.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - status_details_change: Optional[TaskStatusDetailsChangeDict] - """Optional. A change to the framework-specific status details.""" +CreateEvaluationMetricConfigOrDict = Union[ + CreateEvaluationMetricConfig, CreateEvaluationMetricConfigDict +] -TaskEventDataOrDict = Union[TaskEventData, TaskEventDataDict] +class _CreateEvaluationMetricParameters(_common.BaseModel): + """Parameters for creating an evaluation metric.""" -class TaskEvent(_common.BaseModel): - """A task event.""" + display_name: Optional[str] = Field( + default=None, + description="""The user-defined name of the evaluation metric. - create_time: Optional[datetime.datetime] = Field( - default=None, description="""Output only. The create time of the event.""" + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """, ) - event_data: Optional[TaskEventData] = Field( - default=None, description="""Required. The delta associated with the event.""" + description: Optional[str] = Field( + default=None, description="""The description of the evaluation metric.""" ) - event_sequence_number: Optional[int] = Field( + metric: Optional[Metric] = 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="""The metric configuration of the evaluation metric.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""Customer-managed encryption key spec for this EvaluationMetric. + If set, this EvaluationMetric will be secured by this key.""", + ) + config: Optional[CreateEvaluationMetricConfig] = Field( + default=None, description="""""" ) -class TaskEventDict(TypedDict, total=False): - """A task event.""" - - create_time: Optional[datetime.datetime] - """Output only. The create time of the event.""" - - event_data: Optional[TaskEventDataDict] - """Required. The delta associated with the event.""" - - 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.""" - - -TaskEventOrDict = Union[TaskEvent, TaskEventDict] +class _CreateEvaluationMetricParametersDict(TypedDict, total=False): + """Parameters for creating an evaluation metric.""" + display_name: Optional[str] + """The user-defined name of the evaluation metric. -class AppendAgentEngineTaskEventConfig(_common.BaseModel): - """Config for appending Agent Engine task events.""" + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """ - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + description: Optional[str] + """The description of the evaluation metric.""" + metric: Optional[MetricDict] + """The metric configuration of the evaluation metric.""" -class AppendAgentEngineTaskEventConfigDict(TypedDict, total=False): - """Config for appending Agent Engine task events.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for this EvaluationMetric. + If set, this EvaluationMetric will be secured by this key.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + config: Optional[CreateEvaluationMetricConfigDict] + """""" -AppendAgentEngineTaskEventConfigOrDict = Union[ - AppendAgentEngineTaskEventConfig, AppendAgentEngineTaskEventConfigDict +_CreateEvaluationMetricParametersOrDict = Union[ + _CreateEvaluationMetricParameters, _CreateEvaluationMetricParametersDict ] -class _AppendAgentEngineTaskEventRequestParameters(_common.BaseModel): - """Parameters for appending Agent Engine task events.""" +class CustomCodeExecutionSpec(_common.BaseModel): + """Specifies a metric using remote Python function execution. - name: Optional[str] = Field( + This metric is computed by running user-defined Python functions remotely. + """ + + evaluation_function: 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.""" + 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.""", ) - config: Optional[AppendAgentEngineTaskEventConfig] = Field( - default=None, description="""""" + 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 _AppendAgentEngineTaskEventRequestParametersDict(TypedDict, total=False): - """Parameters for appending Agent Engine task events.""" - - name: Optional[str] - """Name of the Agent Engine task to append the events to.""" - - task_events: Optional[list[TaskEventDict]] - """The events to append to the task.""" - - config: Optional[AppendAgentEngineTaskEventConfigDict] - """""" - - -_AppendAgentEngineTaskEventRequestParametersOrDict = Union[ - _AppendAgentEngineTaskEventRequestParameters, - _AppendAgentEngineTaskEventRequestParametersDict, -] - - -class AppendAgentEngineTaskEventResponse(_common.BaseModel): - """Response for appending Agent Engine task events.""" - - pass +class CustomCodeExecutionSpecDict(TypedDict, total=False): + """Specifies a metric using remote Python function execution. + This metric is computed by running user-defined Python functions remotely. + """ -class AppendAgentEngineTaskEventResponseDict(TypedDict, total=False): - """Response for appending Agent Engine task events.""" + 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.""" - pass + 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].""" -AppendAgentEngineTaskEventResponseOrDict = Union[ - AppendAgentEngineTaskEventResponse, AppendAgentEngineTaskEventResponseDict +CustomCodeExecutionSpecOrDict = Union[ + CustomCodeExecutionSpec, CustomCodeExecutionSpecDict ] -class ListAgentEngineTaskEventsConfig(_common.BaseModel): - """Config for listing agent engine tasks.""" +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.""" ) - 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.""", + rouge_spec: Optional[genai_types.RougeSpec] = Field( + default=None, description="""The rouge metric spec.""" ) - 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`.""", + 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 ListAgentEngineTaskEventsConfigDict(TypedDict, total=False): - """Config for listing agent engine tasks.""" +class UnifiedMetricDict(TypedDict, total=False): + """The unified metric used for evaluation.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + bleu_spec: Optional[genai_types.BleuSpec] + """The Bleu metric spec.""" - page_size: Optional[int] - """""" + rouge_spec: Optional[genai_types.RougeSpec] + """The rouge metric spec.""" - page_token: Optional[str] - """""" + pointwise_metric_spec: Optional[genai_types.PointwiseMetricSpec] + """The pointwise metric spec.""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" + llm_based_metric_spec: Optional[genai_types.LLMBasedMetricSpec] + """The spec for an LLM based metric.""" - 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). + custom_code_execution_spec: Optional[CustomCodeExecutionSpecDict] + """The spec for a custom code execution metric.""" - Supported fields: - * `create_time` - * `update_time` + predefined_metric_spec: Optional[genai_types.PredefinedMetricSpec] + """The spec for a pre-defined metric.""" - Example: `create_time desc`.""" + computation_based_metric_spec: Optional[genai_types.ComputationBasedMetricSpec] + """The spec for a computation based metric.""" -ListAgentEngineTaskEventsConfigOrDict = Union[ - ListAgentEngineTaskEventsConfig, ListAgentEngineTaskEventsConfigDict -] +UnifiedMetricOrDict = Union[UnifiedMetric, UnifiedMetricDict] -class _ListAgentEngineTaskEventsRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" +class EvaluationMetric(_common.BaseModel): + """Represents an evaluation metric.""" name: Optional[str] = Field( - default=None, description="""Name of the Agent Engine task.""" + default=None, description="""The resource name of the evaluation metric.""" ) - config: Optional[ListAgentEngineTaskEventsConfig] = Field( - default=None, description="""""" + display_name: Optional[str] = Field( + default=None, + description="""The user-friendly display name for the EvaluationMetric.""", + ) + 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.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""Customer-managed encryption key spec for this EvaluationMetric. + If set, this EvaluationMetric will be secured by this key.""", ) -class _ListAgentEngineTaskEventsRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" +class EvaluationMetricDict(TypedDict, total=False): + """Represents an evaluation metric.""" name: Optional[str] - """Name of the Agent Engine task.""" + """The resource name of the evaluation metric.""" - config: Optional[ListAgentEngineTaskEventsConfigDict] - """""" + display_name: Optional[str] + """The user-friendly display name for the EvaluationMetric.""" + description: Optional[str] + """The description of the EvaluationMetric.""" -_ListAgentEngineTaskEventsRequestParametersOrDict = Union[ - _ListAgentEngineTaskEventsRequestParameters, - _ListAgentEngineTaskEventsRequestParametersDict, -] + metric: Optional[UnifiedMetricDict] + """The metric configuration of the evaluation metric.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for this EvaluationMetric. + If set, this EvaluationMetric will be secured by this key.""" -class ListAgentEngineTaskEventsResponse(_common.BaseModel): - """Response for listing Agent Engine tasks 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="""""") - taskEvents: Optional[list[TaskEvent]] = Field( - default=None, description="""List of Agent Engine task events.""" - ) +EvaluationMetricOrDict = Union[EvaluationMetric, EvaluationMetricDict] -class ListAgentEngineTaskEventsResponseDict(TypedDict, total=False): - """Response for listing Agent Engine tasks events.""" +class SamplingConfig(_common.BaseModel): + """Sampling config for a BigQuery request set.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + sampling_count: Optional[int] = Field(default=None, description="""""") + sampling_method: Optional[SamplingMethod] = Field(default=None, description="""""") + sampling_duration: Optional[str] = Field(default=None, description="""""") - next_page_token: Optional[str] + +class SamplingConfigDict(TypedDict, total=False): + """Sampling config for a BigQuery request set.""" + + sampling_count: Optional[int] """""" - taskEvents: Optional[list[TaskEventDict]] - """List of Agent Engine task events.""" + sampling_method: Optional[SamplingMethod] + """""" + sampling_duration: Optional[str] + """""" -ListAgentEngineTaskEventsResponseOrDict = Union[ - ListAgentEngineTaskEventsResponse, ListAgentEngineTaskEventsResponseDict -] +SamplingConfigOrDict = Union[SamplingConfig, SamplingConfigDict] -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 BigQueryRequestSet(_common.BaseModel): + """Represents a BigQuery request set.""" + + 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 CreateEvaluationExperimentConfigDict(TypedDict, total=False): - """Config to create an evaluation experiment.""" +class BigQueryRequestSetDict(TypedDict, total=False): + """Represents a BigQuery request set.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + uri: Optional[str] + """""" + prompt_column: Optional[str] + """The column name of the prompt in the BigQuery table. Used for EvaluationRun only.""" -CreateEvaluationExperimentConfigOrDict = Union[ - CreateEvaluationExperimentConfig, CreateEvaluationExperimentConfigDict -] + rubrics_column: Optional[str] + """The column name of the rubrics in the BigQuery table. Used for EvaluationRun only.""" + candidate_response_columns: Optional[dict[str, str]] + """The column name of the response candidates in the BigQuery table. Used for EvaluationRun only.""" -class _CreateEvaluationExperimentParameters(_common.BaseModel): - """Parameters for creating an evaluation experiment.""" + sampling_config: Optional[SamplingConfigDict] + """The sampling config for the BigQuery request set. Used for EvaluationRun only.""" - 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="""""" - ) - metadata: Optional[dict[str, Any]] = Field(default=None, description="""""") - config: Optional[CreateEvaluationExperimentConfig] = Field( - default=None, description="""""" - ) +BigQueryRequestSetOrDict = Union[BigQueryRequestSet, BigQueryRequestSetDict] -class _CreateEvaluationExperimentParametersDict(TypedDict, total=False): - """Parameters for creating an evaluation experiment.""" - display_name: Optional[str] - """""" +class EvaluationRunDataSource(_common.BaseModel): + """Represents an evaluation run data source.""" - labels: Optional[dict[str, str]] - """""" + evaluation_set: Optional[str] = Field(default=None, description="""""") + bigquery_request_set: Optional[BigQueryRequestSet] = Field( + default=None, description="""""" + ) - merge_strategy: Optional[EvaluationExperimentMergeStrategy] - """""" - metadata: Optional[dict[str, Any]] +class EvaluationRunDataSourceDict(TypedDict, total=False): + """Represents an evaluation run data source.""" + + evaluation_set: Optional[str] """""" - config: Optional[CreateEvaluationExperimentConfigDict] + bigquery_request_set: Optional[BigQueryRequestSetDict] """""" -_CreateEvaluationExperimentParametersOrDict = Union[ - _CreateEvaluationExperimentParameters, _CreateEvaluationExperimentParametersDict +EvaluationRunDataSourceOrDict = Union[ + EvaluationRunDataSource, EvaluationRunDataSourceDict ] -class EvaluationExperiment(_common.BaseModel): - """Represents an experiment for iterating on and visualizing evaluation runs.""" +class EvaluationRunMetric(_common.BaseModel): + """The metric used for evaluation run.""" - name: Optional[str] = Field( - default=None, - 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="""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.""" + metric: Optional[str] = Field( + default=None, description="""The name of the metric.""" ) - metadata: Optional[dict[str, Any]] = Field( + metric_resource_name: Optional[str] = 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.""" + description="""The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""", ) - update_time: Optional[datetime.datetime] = Field( - default=None, description="""Timestamp when this experiment was last updated.""" + metric_config: Optional[UnifiedMetric] = Field( + default=None, description="""The unified metric used for evaluation run.""" ) -class EvaluationExperimentDict(TypedDict, total=False): - """Represents an experiment for iterating on and visualizing evaluation runs.""" - - name: Optional[str] - """The resource name of the EvaluationExperiment. Format: - `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""" - - display_name: Optional[str] - """The display name of the evaluation experiment.""" +class EvaluationRunMetricDict(TypedDict, total=False): + """The metric used for evaluation run.""" - evaluation_runs: Optional[list[str]] - """The EvaluationRuns that are part of this experiment.""" + metric: Optional[str] + """The name of the metric.""" - labels: Optional[dict[str, str]] - """Labels for the evaluation experiment.""" + metric_resource_name: Optional[str] + """The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" - merge_strategy: Optional[EvaluationExperimentMergeStrategy] - """Merge strategy for the evaluation experiment.""" + metric_config: Optional[UnifiedMetricDict] + """The unified metric used for evaluation run.""" - metadata: Optional[dict[str, Any]] - """Metadata about the evaluation experiment, can be used by the caller - to store additional tracking information about the experiment.""" - create_time: Optional[datetime.datetime] - """Timestamp when this experiment was created.""" +EvaluationRunMetricOrDict = Union[EvaluationRunMetric, EvaluationRunMetricDict] - update_time: Optional[datetime.datetime] - """Timestamp when this experiment was last updated.""" +class EvaluationRunPromptTemplate(_common.BaseModel): + """Prompt template used for inference. -EvaluationExperimentOrDict = Union[EvaluationExperiment, EvaluationExperimentDict] + 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( + 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.""", + ) + gcs_uri: Optional[str] = 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.""", + ) -class CreateEvaluationItemConfig(_common.BaseModel): - """Config to create an evaluation item.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) +class EvaluationRunPromptTemplateDict(TypedDict, total=False): + """Prompt template used for inference. + Only one of `prompt_template` or `gcs_uri` should be set. If both are + provided, an error will be raised. + """ -class CreateEvaluationItemConfigDict(TypedDict, total=False): - """Config to create an evaluation item.""" + 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.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + 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.""" -CreateEvaluationItemConfigOrDict = Union[ - CreateEvaluationItemConfig, CreateEvaluationItemConfigDict +EvaluationRunPromptTemplateOrDict = Union[ + EvaluationRunPromptTemplate, EvaluationRunPromptTemplateDict ] -class _CreateEvaluationItemParameters(_common.BaseModel): - """Represents a job that creates an evaluation item.""" +class LossAnalysisConfig(_common.BaseModel): + """Configuration for the loss analysis job.""" - 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="""""" + metric: Optional[str] = Field( + default=None, + description="""Required. The metric to analyze (e.g., "multi_turn_tool_use_quality_v1").""", + ) + 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.""", + ) + predefined_taxonomy: Optional[str] = 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.""", + ) + 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.""", ) -class _CreateEvaluationItemParametersDict(TypedDict, total=False): - """Represents a job that creates an evaluation item.""" - - evaluation_item_type: Optional[str] - """""" +class LossAnalysisConfigDict(TypedDict, total=False): + """Configuration for the loss analysis job.""" - gcs_uri: Optional[str] - """""" + metric: Optional[str] + """Required. The metric to analyze (e.g., "multi_turn_tool_use_quality_v1").""" - display_name: Optional[str] - """""" + 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.""" - config: Optional[CreateEvaluationItemConfigDict] - """""" + 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.""" + 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.""" -_CreateEvaluationItemParametersOrDict = Union[ - _CreateEvaluationItemParameters, _CreateEvaluationItemParametersDict -] +LossAnalysisConfigOrDict = Union[LossAnalysisConfig, LossAnalysisConfigDict] -class PromptTemplateData(_common.BaseModel): - """Holds data for a prompt template. - Message to hold a prompt template and the values to populate the template. - """ +class EvaluationRunConfig(_common.BaseModel): + """The evaluation configuration used for the evaluation run.""" - values: Optional[dict[str, genai_types.Content]] = Field( - default=None, description="""The values for fields in the prompt template.""" + metrics: Optional[list[EvaluationRunMetric]] = Field( + default=None, + description="""The metrics to be calculated in the evaluation run.""", ) - - -class PromptTemplateDataDict(TypedDict, total=False): - """Holds data for a prompt template. - - Message to hold a prompt template and the values to populate the template. - """ - - values: Optional[dict[str, genai_types.Content]] - """The values for fields in the prompt template.""" - - -PromptTemplateDataOrDict = Union[PromptTemplateData, PromptTemplateDataDict] - - -class EvaluationPrompt(_common.BaseModel): - """Represents the prompt to be evaluated.""" - - text: Optional[str] = Field(default=None, description="""Text prompt.""") - value: Optional[dict[str, Any]] = Field( + 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="""Fields and values that can be used to populate the prompt template.""", + 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_data: Optional[PromptTemplateData] = Field( - default=None, description="""Prompt template data.""" + prompt_template: Optional[EvaluationRunPromptTemplate] = Field( + default=None, description="""The prompt template used for inference.""" ) - user_scenario: Optional[evals_types.UserScenario] = Field( + loss_analysis_config: Optional[list[LossAnalysisConfig]] = Field( default=None, - description="""User scenario to help simulate multi-turn agent run results.""", + 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 EvaluationPromptDict(TypedDict, total=False): - """Represents the prompt to be evaluated.""" - - text: Optional[str] - """Text prompt.""" +class EvaluationRunConfigDict(TypedDict, total=False): + """The evaluation configuration used for the evaluation run.""" - value: Optional[dict[str, Any]] - """Fields and values that can be used to populate the prompt template.""" + metrics: Optional[list[EvaluationRunMetricDict]] + """The metrics to be calculated in the evaluation run.""" - prompt_template_data: Optional[PromptTemplateDataDict] - """Prompt template data.""" + output_config: Optional[genai_types.OutputConfig] + """The output config for the evaluation run.""" - user_scenario: Optional[evals_types.UserScenario] - """User scenario to help simulate multi-turn agent run results.""" + 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.""" -EvaluationPromptOrDict = Union[EvaluationPrompt, EvaluationPromptDict] + 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 CandidateResponse(_common.BaseModel): - """Responses from model or agent.""" - 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( - default=None, - description="""Intermediate events (such as tool calls and responses) that led to the final response.""", - ) - agent_data: Optional[evals_types.AgentData] = 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.""", - ) +EvaluationRunConfigOrDict = Union[EvaluationRunConfig, EvaluationRunConfigDict] -class CandidateResponseDict(TypedDict, total=False): - """Responses from model or agent.""" +class EvaluationRunAgentConfig(_common.BaseModel): + """Agent config for an evaluation run.""" - candidate: Optional[str] - """The name of the candidate that produced the response.""" + 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.""" + ) - text: Optional[str] - """The text response.""" - value: Optional[dict[str, Any]] - """Fields and values that can be used to populate the response template.""" +class EvaluationRunAgentConfigDict(TypedDict, total=False): + """Agent config for an evaluation run.""" - events: Optional[list[genai_types.Content]] - """Intermediate events (such as tool calls and responses) that led to the final response.""" + developer_instruction: Optional[genai_types.Content] + """The developer instruction for the agent.""" - 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.""" + tools: Optional[list[genai_types.Tool]] + """The tools available to the agent.""" -CandidateResponseOrDict = Union[CandidateResponse, CandidateResponseDict] +EvaluationRunAgentConfigOrDict = Union[ + EvaluationRunAgentConfig, EvaluationRunAgentConfigDict +] -class RubricGroup(_common.BaseModel): - """A group of rubrics. +class GeminiAgentConfig(_common.BaseModel): + """Config for scraping a Gemini Agent. - Used for grouping rubrics based on a metric or a version. + A Gemini Agent is a Vertex AI Agent resource scraped via the Vertex + Interactions API. """ - group_id: Optional[str] = Field( - default=None, description="""Unique identifier for the group.""" - ) - display_name: Optional[str] = Field( + gemini_agent: Optional[str] = Field( default=None, - description="""Human-readable name for the group. This should be unique - within a given context if used for display or selection. - Example: "Instruction Following V1", "Content Quality - Summarization - Task".""", - ) - rubrics: Optional[list[evals_types.Rubric]] = Field( - default=None, description="""Rubrics that are part of this group.""" + description="""The resource name of the Gemini Agent. + Format: `projects/{project}/locations/{location}/agents/{agent}`.""", ) -class RubricGroupDict(TypedDict, total=False): - """A group of rubrics. +class GeminiAgentConfigDict(TypedDict, total=False): + """Config for scraping a Gemini Agent. - Used for grouping rubrics based on a metric or a version. + A Gemini Agent is a Vertex AI Agent resource scraped via the Vertex + Interactions API. """ - group_id: Optional[str] - """Unique identifier for the group.""" - - display_name: Optional[str] - """Human-readable name for the group. This should be unique - within a given context if used for display or selection. - Example: "Instruction Following V1", "Content Quality - Summarization - Task".""" - - rubrics: Optional[list[evals_types.Rubric]] - """Rubrics that are part of this group.""" + gemini_agent: Optional[str] + """The resource name of the Gemini Agent. + Format: `projects/{project}/locations/{location}/agents/{agent}`.""" -RubricGroupOrDict = Union[RubricGroup, RubricGroupDict] +GeminiAgentConfigOrDict = Union[GeminiAgentConfig, GeminiAgentConfigDict] -class EvaluationItemRequest(_common.BaseModel): - """Single evaluation request.""" +class AgentRunConfig(_common.BaseModel): + """Configuration for an Agent Run.""" - prompt: Optional[EvaluationPrompt] = Field( - default=None, description="""The request/prompt to evaluate.""" + session_input: Optional[evals_types.SessionInput] = Field( + default=None, description="""The session input to get agent running results.""" ) - golden_response: Optional[CandidateResponse] = Field( - default=None, description="""The ideal response or ground truth.""" + agent_engine: Optional[str] = Field( + default=None, description="""The resource name of the Agent Engine.""" ) - rubrics: Optional[dict[str, RubricGroup]] = Field( + user_simulator_config: Optional[evals_types.UserSimulatorConfig] = Field( default=None, - description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", + 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.""", ) - candidate_responses: Optional[list[CandidateResponse]] = Field( + gemini_agent_config: Optional[GeminiAgentConfig] = Field( default=None, - description="""Responses from model under test and other baseline models for comparison.""", + description="""Config for scraping a Gemini Agent (Vertex AI Agent resource). + Used to target a Gemini agent for an evaluation run.""", ) -class EvaluationItemRequestDict(TypedDict, total=False): - """Single evaluation request.""" +class AgentRunConfigDict(TypedDict, total=False): + """Configuration for an Agent Run.""" - prompt: Optional[EvaluationPromptDict] - """The request/prompt to evaluate.""" + session_input: Optional[evals_types.SessionInput] + """The session input to get agent running results.""" - golden_response: Optional[CandidateResponseDict] - """The ideal response or ground truth.""" + agent_engine: Optional[str] + """The resource name of the Agent Engine.""" - rubrics: Optional[dict[str, RubricGroupDict]] - """Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""" + 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.""" - candidate_responses: Optional[list[CandidateResponseDict]] - """Responses from model under test and other baseline models for comparison.""" + 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.""" -EvaluationItemRequestOrDict = Union[EvaluationItemRequest, EvaluationItemRequestDict] +AgentRunConfigOrDict = Union[AgentRunConfig, AgentRunConfigDict] -class EvaluationItemResult(_common.BaseModel): - """Represents the result of an evaluation item.""" +class EvaluationRunInferenceConfig(_common.BaseModel): + """Configuration that describes an agent.""" - evaluation_request: Optional[str] = Field( - default=None, description="""The request item that was evaluated.""" + agent_config: Optional[EvaluationRunAgentConfig] = Field( + default=None, description="""The agent config.""" ) - evaluation_run: Optional[str] = Field( + model: 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.""" + 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: Optional[str] = Field( - default=None, description="""The metric that was evaluated.""" + prompt_template: Optional[EvaluationRunPromptTemplate] = Field( + default=None, description="""The prompt template used for inference.""" ) - candidate_results: Optional[list[evals_types.CandidateResult]] = Field( - default=None, description="""The results for the metric.""" + agent_run_config: Optional[AgentRunConfig] = Field( + default=None, + description="""Configuration for Agent Run in evaluation management service.""", ) - metadata: Optional[dict[str, Any]] = Field( - default=None, description="""Metadata about the evaluation result.""" + agent_configs: Optional[dict[str, evals_types.AgentConfig]] = Field( + default=None, + description="""A map of agent IDs to their respective agent config.""", ) -class EvaluationItemResultDict(TypedDict, total=False): - """Represents the result of an evaluation item.""" +class EvaluationRunInferenceConfigDict(TypedDict, total=False): + """Configuration that describes an agent.""" - evaluation_request: Optional[str] - """The request item that was evaluated.""" + agent_config: Optional[EvaluationRunAgentConfigDict] + """The agent config.""" - evaluation_run: Optional[str] - """The evaluation run that was used to generate the result.""" + 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`).""" - request: Optional[EvaluationItemRequestDict] - """The request that was evaluated.""" + prompt_template: Optional[EvaluationRunPromptTemplateDict] + """The prompt template used for inference.""" - metric: Optional[str] - """The metric that was evaluated.""" - - candidate_results: Optional[list[evals_types.CandidateResult]] - """The results for the metric.""" - - metadata: Optional[dict[str, Any]] - """Metadata about the evaluation result.""" + 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.""" -EvaluationItemResultOrDict = Union[EvaluationItemResult, EvaluationItemResultDict] +EvaluationRunInferenceConfigOrDict = Union[ + EvaluationRunInferenceConfig, EvaluationRunInferenceConfigDict +] -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 VulnerableTool(_common.BaseModel): + """A tool considered high risk for prompt injection.""" - 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( + tool_name: 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.""" + description="""Optional. The name of the vulnerable function/tool (e.g., "search_flights").""", ) - error: Optional[genai_types.GoogleRpcStatus] = Field( - default=None, description="""Error for the evaluation item.""" + json_paths: Optional[list[str]] = Field( + default=None, + description="""Optional. JSON Paths within the tool's FunctionResponse where malicious content could be injected.""", ) - # 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 EvaluationItemDict(TypedDict, total=False): - """EvaluationItem is a single evaluation request or result. +class VulnerableToolDict(TypedDict, total=False): + """A tool considered high risk for prompt injection.""" - The content of an EvaluationItem is immutable - it cannot be updated once - created. EvaluationItems can be deleted when no longer needed. - """ + tool_name: Optional[str] + """Optional. The name of the vulnerable function/tool (e.g., "search_flights").""" - name: Optional[str] - """The resource name of the EvaluationItem.""" + json_paths: Optional[list[str]] + """Optional. JSON Paths within the tool's FunctionResponse where malicious content could be injected.""" - display_name: Optional[str] - """The display name of the EvaluationItem.""" - metadata: Optional[dict[str, Any]] - """Metadata for the EvaluationItem.""" +VulnerableToolOrDict = Union[VulnerableTool, VulnerableToolDict] - labels: Optional[dict[str, str]] - """Labels for the EvaluationItem.""" - evaluation_item_type: Optional[EvaluationItemType] - """The type of the EvaluationItem.""" +class RedTeamingAnalysisConfig(_common.BaseModel): + """Configuration for the automated Agent Red Teaming analysis.""" - evaluation_request: Optional[EvaluationItemRequestDict] - """The request to evaluate.""" + attack_categories: Optional[list[str]] = Field( + default=None, + description="""Optional. Specific attack categories to test against.""", + ) + vulnerable_tools: Optional[list[VulnerableTool]] = Field( + default=None, + description="""Optional. Manually defined vulnerable tools and their injection paths.""", + ) - evaluation_response: Optional[EvaluationItemResultDict] - """The response from evaluation.""" - gcs_uri: Optional[str] - """The Cloud Storage object where the request or response is stored.""" +class RedTeamingAnalysisConfigDict(TypedDict, total=False): + """Configuration for the automated Agent Red Teaming analysis.""" - create_time: Optional[datetime.datetime] - """Timestamp when this item was created.""" + attack_categories: Optional[list[str]] + """Optional. Specific attack categories to test against.""" - error: Optional[genai_types.GoogleRpcStatus] - """Error for the evaluation item.""" + vulnerable_tools: Optional[list[VulnerableToolDict]] + """Optional. Manually defined vulnerable tools and their injection paths.""" -EvaluationItemOrDict = Union[EvaluationItem, EvaluationItemDict] +RedTeamingAnalysisConfigOrDict = Union[ + RedTeamingAnalysisConfig, RedTeamingAnalysisConfigDict +] -class Metric(_common.BaseModel): - """The metric used for evaluation.""" +class AnalysisConfig(_common.BaseModel): + """Configuration for an analysis to be performed on an evaluation 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.""", + analysis_name: Optional[str] = Field( + default=None, description="""Optional. A name for this analysis.""" ) - parse_and_reduce_fn: Optional[Callable[..., Any]] = Field( + red_teaming_analysis_config: Optional[RedTeamingAnalysisConfig] = Field( default=None, - description="""The parse and reduce function for the judge model.""", + description="""Configuration for the automated Agent Red Teaming analysis.""", ) - aggregate_summary_fn: Optional[Callable[..., Any]] = Field( - default=None, - description="""The aggregate summary function for the judge model.""", + + +class AnalysisConfigDict(TypedDict, total=False): + """Configuration for an analysis to be performed on an evaluation run.""" + + analysis_name: Optional[str] + """Optional. A name for this analysis.""" + + red_teaming_analysis_config: Optional[RedTeamingAnalysisConfigDict] + """Configuration for the automated Agent Red Teaming analysis.""" + + +AnalysisConfigOrDict = Union[AnalysisConfig, AnalysisConfigDict] + + +class CreateEvaluationRunConfig(_common.BaseModel): + """Config to create an evaluation run.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - remote_custom_function: Optional[str] = Field( + allow_cross_region_model: Optional[bool] = Field( default=None, - description="""The evaluation function for the custom code execution metric. This custom code is run remotely in the evaluation service.""", + 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: Optional[str] = Field( - default=None, description="""The judge model for the metric.""" + + +class CreateEvaluationRunConfigDict(TypedDict, total=False): + """Config to create an evaluation run.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + 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.""" + + +CreateEvaluationRunConfigOrDict = Union[ + CreateEvaluationRunConfig, CreateEvaluationRunConfigDict +] + + +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="""""" ) - 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).""", + evaluation_config: Optional[EvaluationRunConfig] = Field( + default=None, description="""""" ) - judge_model_sampling_count: Optional[int] = Field( - default=None, description="""The sampling count for the judge model.""" + labels: Optional[dict[str, str]] = Field(default=None, description="""""") + inference_configs: Optional[dict[str, EvaluationRunInferenceConfig]] = Field( + default=None, description="""""" ) - rubric_group_name: Optional[str] = Field( - default=None, - description="""The rubric group name for the rubric-based metric.""", + config: Optional[CreateEvaluationRunConfig] = Field( + default=None, description="""""" ) - metric_spec_parameters: Optional[dict[str, Any]] = Field( - default=None, - description="""Optional steering instruction parameters for the automated predefined metric.""", + analysis_configs: Optional[list[AnalysisConfig]] = Field( + default=None, description="""""" ) - metric_resource_name: Optional[str] = Field( + evaluation_experiment: Optional[str] = Field( default=None, - description="""The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""", + description="""The resource name of the parent EvaluationExperiment that this run + belongs to. Format: + `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""", ) - result_parsing_function: Optional[str] = Field( + encryption_spec: Optional[genai_types.EncryptionSpec] = 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="""Customer-managed encryption key spec for this EvaluationRun. + If set, this EvaluationRun will be secured by this key.""", ) - # 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 _CreateEvaluationRunParametersDict(TypedDict, total=False): + """Represents a job that creates an evaluation run.""" - _config_source: Optional[str] = PrivateAttr(default=None) - """An optional string indicating the source of the metric configuration.""" + name: Optional[str] + """""" - _version: Optional[str] = PrivateAttr(default=None) - """An optional string indicating the version of the metric.""" + display_name: Optional[str] + """""" - @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. - - Args: - file_path: The path to the YAML file. - version: Optional version string to include in the YAML output. + data_source: Optional[EvaluationRunDataSourceDict] + """""" - 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]'." - ) + evaluation_config: Optional[EvaluationRunConfigDict] + """""" - 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) - } + labels: Optional[dict[str, str]] + """""" - data_to_dump = self.model_dump( - exclude_unset=True, - exclude_none=True, - mode="json", - exclude=fields_to_exclude if fields_to_exclude else None, - ) + inference_configs: Optional[dict[str, EvaluationRunInferenceConfigDict]] + """""" - if version: - data_to_dump["version"] = version + config: Optional[CreateEvaluationRunConfigDict] + """""" - with open(file_path, "w", encoding="utf-8") as f: - yaml.dump(data_to_dump, f, sort_keys=False, allow_unicode=True) + analysis_configs: Optional[list[AnalysisConfigDict]] + """""" + evaluation_experiment: Optional[str] + """The resource name of the parent EvaluationExperiment that this run + belongs to. Format: + `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""" -class CodeExecutionMetric(Metric): - """A metric that executes custom Python code for evaluation.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for this EvaluationRun. + If set, this EvaluationRun will be secured by this key.""" - # 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.""", - ) - # 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 +_CreateEvaluationRunParametersOrDict = Union[ + _CreateEvaluationRunParameters, _CreateEvaluationRunParametersDict +] -class LLMMetric(Metric): - """A metric that uses LLM-as-a-judge for evaluation.""" +class SummaryMetric(_common.BaseModel): + """Represents a summary metric for an evaluation run.""" - 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.""", + metrics: Optional[dict[str, Any]] = Field( + default=None, description="""Map of metric name to metric value.""" ) - - 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.""", + 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.""" ) - @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 +class SummaryMetricDict(TypedDict, total=False): + """Represents a summary metric for an evaluation run.""" - @classmethod - def load(cls, config_path: str, client: Optional[Any] = None) -> "LLMMetric": - """Loads a metric configuration from a YAML or JSON file. + metrics: Optional[dict[str, Any]] + """Map of metric name to metric value.""" - 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. + total_items: Optional[int] + """The total number of items that were evaluated.""" - 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. + failed_items: Optional[int] + """The number of items that failed to be evaluated.""" - 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" - ) +SummaryMetricOrDict = Union[SummaryMetric, SummaryMetricDict] - 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 AttackCategoryResult(_common.BaseModel): + """The red teaming outcome for a specific attack category.""" - 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 + attack_category: Optional[str] = Field( + default=None, description="""The category of the attack evaluated.""" + ) + attack_success_rate: Optional[float] = Field( + default=None, + 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.""" + ) - 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 AttackCategoryResultDict(TypedDict, total=False): + """The red teaming outcome for a specific attack category.""" - if not isinstance(data, dict): - raise ValueError("Metric config content did not parse into a dictionary.") + attack_category: Optional[str] + """The category of the attack evaluated.""" - return cls.model_validate(data) + attack_success_rate: Optional[float] + """The ratio of successful attacks given a fixed budget.""" + vulnerability_insight: Optional[str] + """Insights into why an attack succeeded or failed.""" -class MetricDict(TypedDict, total=False): - """The metric used for evaluation.""" - name: Optional[str] - """The name of the metric.""" +AttackCategoryResultOrDict = Union[AttackCategoryResult, AttackCategoryResultDict] - 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 RedTeamingAnalysisResult(_common.BaseModel): + """The top-level result for Red Teaming analysis.""" - judge_model_system_instruction: Optional[str] - """The system instruction for the judge model.""" + config: Optional[RedTeamingAnalysisConfig] = 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.""" + ) + category_results: Optional[list[AttackCategoryResult]] = Field( + default=None, description="""Detailed results by attack category.""" + ) - return_raw_output: Optional[bool] - """Whether to return the raw output from the judge model.""" - parse_and_reduce_fn: Optional[Callable[..., Any]] - """The parse and reduce function for the judge model.""" +class RedTeamingAnalysisResultDict(TypedDict, total=False): + """The top-level result for Red Teaming analysis.""" - aggregate_summary_fn: Optional[Callable[..., Any]] - """The aggregate summary function for the judge model.""" + config: Optional[RedTeamingAnalysisConfigDict] + """The configuration used to generate this analysis.""" - remote_custom_function: Optional[str] - """The evaluation function for the custom code execution metric. This custom code is run remotely in the evaluation service.""" + analysis_time: Optional[str] + """The timestamp when this analysis was performed.""" - judge_model: Optional[str] - """The judge model for the metric.""" + category_results: Optional[list[AttackCategoryResultDict]] + """Detailed results by attack category.""" - 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.""" +RedTeamingAnalysisResultOrDict = Union[ + RedTeamingAnalysisResult, RedTeamingAnalysisResultDict +] - 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 LossTaxonomyEntry(_common.BaseModel): + """A specific entry in the loss pattern taxonomy.""" - metric_resource_name: Optional[str] - """The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" + l1_category: Optional[str] = Field( + default=None, + description="""The primary category of the loss (e.g., "Hallucination", "Tool Calling").""", + ) + l2_category: Optional[str] = Field( + default=None, + description="""The secondary category of the loss (e.g., "Hallucination of Action", "Incorrect Tool Selection").""", + ) + description: Optional[str] = Field( + default=None, + description="""A detailed description of this loss pattern. Example: "The agent verbally confirms an action without executing the tool." """, + ) - 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.""" +class LossTaxonomyEntryDict(TypedDict, total=False): + """A specific entry in the loss pattern taxonomy.""" -MetricOrDict = Union[Metric, MetricDict] + 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").""" -class CreateEvaluationMetricConfig(_common.BaseModel): - """Config for creating an evaluation metric.""" + description: Optional[str] + """A detailed description of this loss pattern. Example: "The agent verbally confirms an action without executing the tool." """ - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + +LossTaxonomyEntryOrDict = Union[LossTaxonomyEntry, LossTaxonomyEntryDict] + + +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 CreateEvaluationMetricConfigDict(TypedDict, total=False): - """Config for creating an evaluation metric.""" +class FailedRubricDict(TypedDict, total=False): + """A specific failed rubric and the associated analysis.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + rubric_id: Optional[str] + """The unique ID of the rubric (if available from the metric source).""" + classification_rationale: Optional[str] + """The rationale provided by the Loss Analysis Classifier for why this failure maps to this specific Loss Cluster.""" -CreateEvaluationMetricConfigOrDict = Union[ - CreateEvaluationMetricConfig, CreateEvaluationMetricConfigDict -] +FailedRubricOrDict = Union[FailedRubric, FailedRubricDict] -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 LossExample(_common.BaseModel): + """A specific example of a loss pattern.""" - 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( + evaluation_item: Optional[str] = Field( default=None, - description="""The metric configuration of the evaluation metric.""", + description="""Reference to the persisted EvalItem resource name. Format: projects/.../locations/.../evaluationItems/{item_id}.""", ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + evaluation_result: Optional[dict[str, Any]] = Field( default=None, - description="""Customer-managed encryption key spec for this EvaluationMetric. - If set, this EvaluationMetric will be secured by this key.""", + description="""The full evaluation result object provided inline. Used when the analysis is performed on ephemeral data.""", ) - config: Optional[CreateEvaluationMetricConfig] = Field( - default=None, description="""""" + failed_rubrics: Optional[list[FailedRubric]] = 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.""", ) -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.""" - - metric: Optional[MetricDict] - """The metric configuration of the evaluation metric.""" +class LossExampleDict(TypedDict, total=False): + """A specific example of a loss pattern.""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """Customer-managed encryption key spec for this EvaluationMetric. - If set, this EvaluationMetric will be secured by this key.""" + evaluation_item: Optional[str] + """Reference to the persisted EvalItem resource name. Format: projects/.../locations/.../evaluationItems/{item_id}.""" - config: Optional[CreateEvaluationMetricConfigDict] - """""" + 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.""" -_CreateEvaluationMetricParametersOrDict = Union[ - _CreateEvaluationMetricParameters, _CreateEvaluationMetricParametersDict -] +LossExampleOrDict = Union[LossExample, LossExampleDict] -class CustomCodeExecutionSpec(_common.BaseModel): - """Specifies a metric using remote Python function execution. - This metric is computed by running user-defined Python functions remotely. - """ +class LossCluster(_common.BaseModel): + """A semantic grouping of failures (e.g., "Hallucination of Action").""" - evaluation_function: Optional[str] = Field( + cluster_id: 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.""", + description="""Unique identifier for the loss cluster within the scope of the analysis result.""", ) - remote_custom_function: Optional[str] = Field( + taxonomy_entry: Optional[LossTaxonomyEntry] = 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 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 CustomCodeExecutionSpecDict(TypedDict, total=False): - """Specifies a metric using remote Python function execution. +class LossClusterDict(TypedDict, total=False): + """A semantic grouping of failures (e.g., "Hallucination of Action").""" - This metric is computed by running user-defined Python functions remotely. - """ + cluster_id: Optional[str] + """Unique identifier for the loss cluster within the scope of the analysis result.""" - 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.""" + taxonomy_entry: Optional[LossTaxonomyEntryDict] + """The structured definition of the loss taxonomy for this cluster.""" - 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].""" + item_count: Optional[int] + """The total number of EvaluationItems falling into this cluster.""" + examples: Optional[list[LossExampleDict]] + """A list of examples that belong to this cluster. This links the cluster back to the specific EvaluationItems and Rubrics.""" -CustomCodeExecutionSpecOrDict = Union[ - CustomCodeExecutionSpec, CustomCodeExecutionSpecDict -] +LossClusterOrDict = Union[LossCluster, LossClusterDict] -class UnifiedMetric(_common.BaseModel): - """The unified metric used for evaluation.""" - 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.""" +class LossAnalysisResult(_common.BaseModel): + """The top-level result for loss analysis.""" + + config: Optional[LossAnalysisConfig] = 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.""") + 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 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.""" + _evals_visualization.display_loss_analysis_result(self) - 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 LossAnalysisResultDict(TypedDict, total=False): + """The top-level result for loss analysis.""" - custom_code_execution_spec: Optional[CustomCodeExecutionSpecDict] - """The spec for a custom code execution metric.""" + config: Optional[LossAnalysisConfigDict] + """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.""" + clusters: Optional[list[LossClusterDict]] + """The list of identified loss clusters.""" -UnifiedMetricOrDict = Union[UnifiedMetric, UnifiedMetricDict] +LossAnalysisResultOrDict = Union[LossAnalysisResult, LossAnalysisResultDict] -class EvaluationMetric(_common.BaseModel): - """Represents an evaluation metric.""" +class EvaluationRunResults(_common.BaseModel): + """Represents the results of an evaluation run.""" - name: Optional[str] = Field( - default=None, description="""The resource name of the evaluation metric.""" - ) - display_name: Optional[str] = Field( + evaluation_set: Optional[str] = Field( default=None, - description="""The user-friendly display name for the EvaluationMetric.""", + description="""The evaluation set where item level results are stored.""", ) - description: Optional[str] = Field( - default=None, description="""The description of the EvaluationMetric.""" + summary_metrics: Optional[SummaryMetric] = Field( + default=None, description="""The summary metrics for the evaluation run.""" ) - metric: Optional[UnifiedMetric] = Field( + loss_analysis_results: Optional[list[LossAnalysisResult]] = Field( default=None, - description="""The metric configuration of the evaluation metric.""", + description="""The loss analysis results for the evaluation run.""", ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, - description="""Customer-managed encryption key spec for this EvaluationMetric. - If set, this EvaluationMetric will be secured by this key.""", + red_teaming_analysis_results: Optional[list[RedTeamingAnalysisResult]] = Field( + default=None, description="""The Red Teaming analysis results.""" ) -class EvaluationMetricDict(TypedDict, total=False): - """Represents an evaluation metric.""" +class EvaluationRunResultsDict(TypedDict, total=False): + """Represents the results of an evaluation run.""" - name: Optional[str] - """The resource name of the evaluation metric.""" + evaluation_set: Optional[str] + """The evaluation set where item level results are stored.""" - display_name: Optional[str] - """The user-friendly display name for the EvaluationMetric.""" + summary_metrics: Optional[SummaryMetricDict] + """The summary metrics for the evaluation run.""" - description: Optional[str] - """The description of the EvaluationMetric.""" + loss_analysis_results: Optional[list[LossAnalysisResultDict]] + """The loss analysis results for the evaluation run.""" - metric: Optional[UnifiedMetricDict] - """The metric configuration of the evaluation metric.""" + red_teaming_analysis_results: Optional[list[RedTeamingAnalysisResultDict]] + """The Red Teaming analysis results.""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """Customer-managed encryption key spec for this EvaluationMetric. - If set, this EvaluationMetric will be secured by this key.""" +EvaluationRunResultsOrDict = Union[EvaluationRunResults, EvaluationRunResultsDict] -EvaluationMetricOrDict = Union[EvaluationMetric, EvaluationMetricDict] +class EvalCaseMetricResult(_common.BaseModel): + """Evaluation result for a single evaluation case for a single metric.""" -class SamplingConfig(_common.BaseModel): - """Sampling config for a BigQuery request set.""" + 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.""" + ) - 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 EvalCaseMetricResultDict(TypedDict, total=False): + """Evaluation result for a single evaluation case for a single metric.""" -class SamplingConfigDict(TypedDict, total=False): - """Sampling config for a BigQuery request set.""" + metric_name: Optional[str] + """Name of the metric.""" - sampling_count: Optional[int] - """""" + score: Optional[float] + """Score of the metric.""" - sampling_method: Optional[SamplingMethod] - """""" + explanation: Optional[str] + """Explanation of the metric.""" - sampling_duration: Optional[str] - """""" + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] + """The details of all the rubrics and their verdicts for rubric-based metrics.""" + raw_output: Optional[list[str]] + """Raw output of the metric.""" -SamplingConfigOrDict = Union[SamplingConfig, SamplingConfigDict] + error_message: Optional[str] + """Error message for the metric.""" -class BigQueryRequestSet(_common.BaseModel): - """Represents a BigQuery request set.""" +EvalCaseMetricResultOrDict = Union[EvalCaseMetricResult, EvalCaseMetricResultDict] - 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( + +class ResponseCandidateResult(_common.BaseModel): + """Aggregated metric results for a single response candidate.""" + + response_index: Optional[int] = Field( default=None, - description="""The column name of the response candidates in the BigQuery table. Used for EvaluationRun only.""", + description="""Index of the response candidate this result pertains to.""", ) - sampling_config: Optional[SamplingConfig] = Field( + metric_results: Optional[dict[str, EvalCaseMetricResult]] = Field( default=None, - description="""The sampling config for the BigQuery request set. Used for EvaluationRun only.""", + description="""A dictionary of metric results for this response candidate, keyed by metric name.""", ) -class BigQueryRequestSetDict(TypedDict, total=False): - """Represents a BigQuery request set.""" - - uri: Optional[str] - """""" - - prompt_column: Optional[str] - """The column name of the prompt in the BigQuery table. Used for EvaluationRun only.""" - - rubrics_column: Optional[str] - """The column name of the rubrics in the BigQuery table. Used for EvaluationRun only.""" +class ResponseCandidateResultDict(TypedDict, total=False): + """Aggregated metric results for a single response candidate.""" - candidate_response_columns: Optional[dict[str, str]] - """The column name of the response candidates in the BigQuery table. Used for EvaluationRun only.""" + response_index: Optional[int] + """Index of the response candidate this result pertains to.""" - sampling_config: Optional[SamplingConfigDict] - """The sampling config for the BigQuery request set. Used for EvaluationRun only.""" + metric_results: Optional[dict[str, EvalCaseMetricResultDict]] + """A dictionary of metric results for this response candidate, keyed by metric name.""" -BigQueryRequestSetOrDict = Union[BigQueryRequestSet, BigQueryRequestSetDict] +ResponseCandidateResultOrDict = Union[ + ResponseCandidateResult, ResponseCandidateResultDict +] -class EvaluationRunDataSource(_common.BaseModel): - """Represents an evaluation run data source.""" +class EvalCaseResult(_common.BaseModel): + """Eval result for a single evaluation case.""" - evaluation_set: Optional[str] = Field(default=None, description="""""") - bigquery_request_set: Optional[BigQueryRequestSet] = Field( - default=None, description="""""" + 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.""", ) -class EvaluationRunDataSourceDict(TypedDict, total=False): - """Represents an evaluation run data source.""" +class EvalCaseResultDict(TypedDict, total=False): + """Eval result for a single evaluation case.""" - evaluation_set: Optional[str] - """""" + eval_case_index: Optional[int] + """Index of the evaluation case.""" - bigquery_request_set: Optional[BigQueryRequestSetDict] - """""" + response_candidate_results: Optional[list[ResponseCandidateResultDict]] + """A list of results, one for each response candidate of the EvalCase.""" -EvaluationRunDataSourceOrDict = Union[ - EvaluationRunDataSource, EvaluationRunDataSourceDict -] +EvalCaseResultOrDict = Union[EvalCaseResult, EvalCaseResultDict] -class EvaluationRunMetric(_common.BaseModel): - """The metric used for evaluation run.""" +class AggregatedMetricResult(_common.BaseModel): + """Evaluation result for a single metric for an evaluation dataset.""" - metric: Optional[str] = Field( - default=None, description="""The name of the metric.""" + metric_name: Optional[str] = Field( + default=None, description="""Name of the 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}""", + num_cases_total: Optional[int] = Field( + default=None, description="""Total number of cases in the dataset.""" ) - metric_config: Optional[UnifiedMetric] = Field( - default=None, description="""The unified metric used for evaluation run.""" + 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") -class EvaluationRunMetricDict(TypedDict, total=False): - """The metric used for evaluation run.""" - metric: Optional[str] - """The name of the metric.""" +class AggregatedMetricResultDict(TypedDict, total=False): + """Evaluation result for a single metric for an evaluation dataset.""" - metric_resource_name: Optional[str] - """The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" + metric_name: Optional[str] + """Name of the metric.""" - metric_config: Optional[UnifiedMetricDict] - """The unified metric used for evaluation run.""" + num_cases_total: Optional[int] + """Total number of cases in the dataset.""" + num_cases_valid: Optional[int] + """Number of valid cases in the dataset.""" -EvaluationRunMetricOrDict = Union[EvaluationRunMetric, EvaluationRunMetricDict] + num_cases_error: Optional[int] + """Number of cases with errors in the dataset.""" + mean_score: Optional[float] + """Mean score of the metric.""" -class EvaluationRunPromptTemplate(_common.BaseModel): - """Prompt template used for inference. + stdev_score: Optional[float] + """Standard deviation of the metric.""" - Only one of `prompt_template` or `gcs_uri` should be set. If both are - provided, an error will be raised. - """ + 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.""" - prompt_template: Optional[str] = Field( + +AggregatedMetricResultOrDict = Union[AggregatedMetricResult, AggregatedMetricResultDict] + + +class WinRateStats(_common.BaseModel): + """Statistics for win rates for a single metric.""" + + win_rates: Optional[list[float]] = 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="""Win rates for the metric, one for each candidate.""", ) - gcs_uri: Optional[str] = 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.""", + tie_rate: Optional[float] = Field( + default=None, description="""Tie rate for the metric.""" ) -class EvaluationRunPromptTemplateDict(TypedDict, total=False): - """Prompt template used for inference. - - Only one of `prompt_template` or `gcs_uri` should be set. If both are - provided, an error will be raised. - """ +class WinRateStatsDict(TypedDict, total=False): + """Statistics for win rates for a single metric.""" - 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.""" + win_rates: Optional[list[float]] + """Win rates for the metric, one for each candidate.""" - 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.""" + tie_rate: Optional[float] + """Tie rate for the metric.""" -EvaluationRunPromptTemplateOrDict = Union[ - EvaluationRunPromptTemplate, EvaluationRunPromptTemplateDict -] +WinRateStatsOrDict = Union[WinRateStats, WinRateStatsDict] -class LossAnalysisConfig(_common.BaseModel): - """Configuration for the loss analysis job.""" +class ResponseCandidate(_common.BaseModel): + """A model-generated content to the prompt.""" - metric: Optional[str] = Field( - default=None, - description="""Required. The metric to analyze (e.g., "multi_turn_tool_use_quality_v1").""", - ) - candidate: Optional[str] = Field( + response: Optional[genai_types.Content] = 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.""", + description="""The final model-generated response to the `prompt`.""", ) - predefined_taxonomy: Optional[str] = Field( + + +class ResponseCandidateDict(TypedDict, total=False): + """A model-generated content to the prompt.""" + + response: Optional[genai_types.Content] + """The final model-generated response to the `prompt`.""" + + +ResponseCandidateOrDict = Union[ResponseCandidate, ResponseCandidateDict] + + +class InteractionsDataSource(_common.BaseModel): + """Source for populating agent data from an Interactions API interaction.""" + + gemini_agent_config: Optional[GeminiAgentConfig] = 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 Gemini Agent (Vertex AI Agent resource) that produced the + interaction.""", ) - max_top_cluster_count: Optional[int] = Field( + interaction: Optional[str] = Field( default=None, - description="""Optional. Limits the analysis to the top N clusters. If not specified or set to 0, all clusters are returned.""", + description="""The interaction to evaluate. Required by the backend. + Format: + `projects/{project}/locations/{location}/interactions/{interaction}`.""", ) -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").""" - - 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.""" +class InteractionsDataSourceDict(TypedDict, total=False): + """Source for populating agent data from an Interactions API interaction.""" - 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.""" + gemini_agent_config: Optional[GeminiAgentConfigDict] + """The Gemini Agent (Vertex AI Agent resource) that produced the + interaction.""" - 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.""" + interaction: Optional[str] + """The interaction to evaluate. Required by the backend. + Format: + `projects/{project}/locations/{location}/interactions/{interaction}`.""" -LossAnalysisConfigOrDict = Union[LossAnalysisConfig, LossAnalysisConfigDict] +InteractionsDataSourceOrDict = Union[InteractionsDataSource, InteractionsDataSourceDict] -class EvaluationRunConfig(_common.BaseModel): - """The evaluation configuration used for the evaluation run.""" +class EvalCase(_common.BaseModel): + """A comprehensive representation of a GenAI interaction for evaluation.""" - metrics: Optional[list[EvaluationRunMetric]] = 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="""The metrics to be calculated in the evaluation run.""", + 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.""", ) - output_config: Optional[genai_types.OutputConfig] = Field( - default=None, description="""The output config for the evaluation run.""" + reference: Optional[ResponseCandidate] = Field( + default=None, + description="""User-provided, golden reference model reply to prompt in context of chat history; Reference for last response in a conversation.""", ) - autorater_config: Optional[genai_types.AutoraterConfig] = Field( + system_instruction: Optional[genai_types.Content] = Field( + default=None, description="""System instruction for the model.""" + ) + conversation_history: Optional[list[evals_types.Message]] = 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.""", + description="""List of all prior messages in the conversation (chat history).""", ) - prompt_template: Optional[EvaluationRunPromptTemplate] = Field( - default=None, description="""The prompt template used for inference.""" + 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.""", ) - loss_analysis_config: Optional[list[LossAnalysisConfig]] = Field( + 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="""Specifications for loss analysis. Each config specifies a metric and candidate to analyze for loss patterns.""", + description="""Intermediate events of a single turn in an agent run or intermediate events of the last turn for multi-turn an agent run.""", ) - allow_cross_region_model: Optional[bool] = Field( + agent_info: Optional[evals_types.AgentInfo] = 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="""The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""", ) - - -class EvaluationRunConfigDict(TypedDict, total=False): - """The evaluation configuration used for the evaluation run.""" - - metrics: Optional[list[EvaluationRunMetricDict]] - """The metrics to be calculated in the evaluation run.""" - - output_config: Optional[genai_types.OutputConfig] - """The output config for the evaluation run.""" - - 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.""" - - 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.""" - - -EvaluationRunConfigOrDict = Union[EvaluationRunConfig, EvaluationRunConfigDict] - - -class EvaluationRunAgentConfig(_common.BaseModel): - """Agent config for an evaluation run.""" - - developer_instruction: Optional[genai_types.Content] = Field( - default=None, description="""The developer instruction for the agent.""" + agent_data: Optional[evals_types.AgentData] = Field( + default=None, description="""The agent data of the agent under evaluation.""" ) - tools: Optional[list[genai_types.Tool]] = Field( - default=None, description="""The tools available to the agent.""" + 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 EvaluationRunAgentConfigDict(TypedDict, total=False): - """Agent config for an evaluation run.""" +class EvalCaseDict(TypedDict, total=False): + """A comprehensive representation of a GenAI interaction for evaluation.""" - developer_instruction: Optional[genai_types.Content] - """The developer instruction for the agent.""" + prompt: Optional[genai_types.Content] + """The most recent user message (current input).""" - tools: Optional[list[genai_types.Tool]] - """The tools available to the agent.""" + 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.""" -EvaluationRunAgentConfigOrDict = Union[ - EvaluationRunAgentConfig, EvaluationRunAgentConfigDict -] + 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 GeminiAgentConfig(_common.BaseModel): - """Config for scraping a Gemini Agent. + 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.""" - A Gemini Agent is a Vertex AI Agent resource scraped via the Vertex - Interactions API. - """ + eval_case_id: Optional[str] + """Unique identifier for the evaluation case.""" - gemini_agent: Optional[str] = Field( - default=None, - description="""The resource name of the Gemini Agent. - Format: `projects/{project}/locations/{location}/agents/{agent}`.""", - ) + 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.""" + agent_info: Optional[evals_types.AgentInfo] + """The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""" -class GeminiAgentConfigDict(TypedDict, total=False): - """Config for scraping a Gemini Agent. + agent_data: Optional[evals_types.AgentData] + """The agent data of the agent under evaluation.""" - A Gemini Agent is a Vertex AI Agent resource scraped via the Vertex - Interactions API. - """ + user_scenario: Optional[evals_types.UserScenario] + """The user scenario for the evaluation case.""" - gemini_agent: Optional[str] - """The resource name of the Gemini Agent. - Format: `projects/{project}/locations/{location}/agents/{agent}`.""" + 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.""" -GeminiAgentConfigOrDict = Union[GeminiAgentConfig, GeminiAgentConfigDict] +EvalCaseOrDict = Union[EvalCase, EvalCaseDict] -class AgentRunConfig(_common.BaseModel): - """Configuration for an Agent Run.""" +class EvaluationDataset(_common.BaseModel): + """The dataset used for evaluation.""" - session_input: Optional[evals_types.SessionInput] = Field( - default=None, description="""The session input to get agent running results.""" + bigquery_source: Optional[genai_types.BigQuerySource] = Field( + default=None, description="""The BigQuery source for the evaluation dataset.""" ) - agent_engine: Optional[str] = Field( - default=None, description="""The resource name of the Agent Engine.""" + gcs_source: Optional[genai_types.GcsSource] = Field( + default=None, description="""The GCS source for the evaluation dataset.""" ) - user_simulator_config: Optional[evals_types.UserSimulatorConfig] = Field( + eval_cases: Optional[list[EvalCase]] = Field( + default=None, description="""The evaluation cases to be evaluated.""" + ) + eval_dataset_df: Optional[PandasDataFrame] = 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="""The evaluation dataset in the form of a Pandas DataFrame.""", ) - gemini_agent_config: Optional[GeminiAgentConfig] = Field( + candidate_name: Optional[str] = Field( default=None, - description="""Config for scraping a Gemini Agent (Vertex AI Agent resource). - Used to target a Gemini agent for an evaluation run.""", + 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 AgentRunConfigDict(TypedDict, total=False): - """Configuration for an Agent 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 - session_input: Optional[evals_types.SessionInput] - """The session input to get agent running results.""" + formats = [] + requests = [] + responses = [] + system_instructions = [] - agent_engine: Optional[str] - """The resource name of the Agent Engine.""" + for case in cases: + gcs_utils = _gcs_utils.GcsUtils( + case.api_client._api_client if case.api_client else None + ) - 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.""" + # Associate "observability" data format for given sources + formats.append("observability") - 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.""" + # Input source + request_data = gcs_utils.read_file_contents(case.input_src) + requests.append(request_data) + # Output source + response_data = gcs_utils.read_file_contents(case.output_src) + responses.append(response_data) -AgentRunConfigOrDict = Union[AgentRunConfig, AgentRunConfigDict] + # 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) + eval_dataset_df = pd.DataFrame( + { + "format": formats, + "request": requests, + "response": responses, + "system_instruction": system_instructions, + } + ) -class EvaluationRunInferenceConfig(_common.BaseModel): - """Configuration that describes an agent.""" + except ImportError as e: + raise ImportError("Pandas DataFrame library is required.") from e - 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( - default=None, - description="""A map of agent IDs to their respective agent config.""", - ) + return EvaluationDataset(eval_dataset_df=eval_dataset_df) + def show(self) -> None: + """Shows the evaluation dataset.""" + from .. import _evals_visualization -class EvaluationRunInferenceConfigDict(TypedDict, total=False): - """Configuration that describes an agent.""" + _evals_visualization.display_evaluation_dataset(self) - agent_config: Optional[EvaluationRunAgentConfigDict] - """The agent config.""" - 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`).""" +class EvaluationDatasetDict(TypedDict, total=False): + """The dataset used for evaluation.""" - prompt_template: Optional[EvaluationRunPromptTemplateDict] - """The prompt template used for inference.""" + bigquery_source: Optional[genai_types.BigQuerySource] + """The BigQuery source for the evaluation dataset.""" - agent_run_config: Optional[AgentRunConfigDict] - """Configuration for Agent Run in evaluation management service.""" + gcs_source: Optional[genai_types.GcsSource] + """The GCS source for the evaluation dataset.""" - agent_configs: Optional[dict[str, evals_types.AgentConfig]] - """A map of agent IDs to their respective agent config.""" + 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.""" -EvaluationRunInferenceConfigOrDict = Union[ - EvaluationRunInferenceConfig, EvaluationRunInferenceConfigDict -] + candidate_name: Optional[str] + """The name of the candidate model or agent for this evaluation dataset.""" -class VulnerableTool(_common.BaseModel): - """A tool considered high risk for prompt injection.""" +EvaluationDatasetOrDict = Union[EvaluationDataset, EvaluationDatasetDict] - tool_name: Optional[str] = Field( + +class EvaluationRunMetadata(_common.BaseModel): + """Metadata for an evaluation run.""" + + candidate_names: Optional[list[str]] = Field( default=None, - description="""Optional. The name of the vulnerable function/tool (e.g., "search_flights").""", + description="""Name of the candidate(s) being evaluated in the evaluation run.""", ) - json_paths: Optional[list[str]] = Field( + dataset_name: Optional[str] = Field( default=None, - description="""Optional. JSON Paths within the tool's FunctionResponse where malicious content could be injected.""", + description="""Name of the evaluation dataset used for the evaluation run.""", + ) + 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 VulnerableToolDict(TypedDict, total=False): - """A tool considered high risk for prompt injection.""" +class EvaluationRunMetadataDict(TypedDict, total=False): + """Metadata for an evaluation run.""" - tool_name: Optional[str] - """Optional. The name of the vulnerable function/tool (e.g., "search_flights").""" + candidate_names: Optional[list[str]] + """Name of the candidate(s) being evaluated in the evaluation run.""" - json_paths: Optional[list[str]] - """Optional. JSON Paths within the tool's FunctionResponse where malicious content could be injected.""" + dataset_name: Optional[str] + """Name of the evaluation dataset used for the evaluation run.""" + dataset_id: Optional[str] + """Unique identifier for the evaluation dataset used for the evaluation run.""" -VulnerableToolOrDict = Union[VulnerableTool, VulnerableToolDict] + creation_timestamp: Optional[datetime.datetime] + """Creation timestamp of the evaluation run.""" -class RedTeamingAnalysisConfig(_common.BaseModel): - """Configuration for the automated Agent Red Teaming analysis.""" +EvaluationRunMetadataOrDict = Union[EvaluationRunMetadata, EvaluationRunMetadataDict] - attack_categories: Optional[list[str]] = Field( + +class EvaluationResult(_common.BaseModel): + """Result of an evaluation run for an evaluation dataset.""" + + eval_case_results: Optional[list[EvalCaseResult]] = Field( default=None, - description="""Optional. Specific attack categories to test against.""", + description="""A list of evaluation results for each evaluation case.""", ) - vulnerable_tools: Optional[list[VulnerableTool]] = Field( + summary_metrics: Optional[list[AggregatedMetricResult]] = Field( default=None, - description="""Optional. Manually defined vulnerable tools and their injection paths.""", + description="""A list of summary-level evaluation results for each metric.""", ) - - -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.""" + 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.""", ) - red_teaming_analysis_config: Optional[RedTeamingAnalysisConfig] = Field( + evaluation_dataset: Optional[list[EvaluationDataset]] = Field( default=None, - description="""Configuration for the automated Agent Red Teaming analysis.""", + description="""The input evaluation dataset(s) for the evaluation run.""", + ) + metadata: Optional[EvaluationRunMetadata] = Field( + default=None, description="""Metadata for the evaluation 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.""", ) + def show(self, candidate_names: Optional[List[str]] = None) -> None: + """Shows the evaluation result. -class AnalysisConfigDict(TypedDict, total=False): - """Configuration for an analysis to be performed on an evaluation run.""" - - analysis_name: Optional[str] - """Optional. A name for this analysis.""" - - red_teaming_analysis_config: Optional[RedTeamingAnalysisConfigDict] - """Configuration for the automated Agent Red Teaming analysis.""" + 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) -AnalysisConfigOrDict = Union[AnalysisConfig, AnalysisConfigDict] +class EvaluationResultDict(TypedDict, total=False): + """Result of an evaluation run for an evaluation dataset.""" -class CreateEvaluationRunConfig(_common.BaseModel): - """Config to create an evaluation run.""" + eval_case_results: Optional[list[EvalCaseResultDict]] + """A list of evaluation results for each evaluation case.""" - 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.""", - ) + 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 CreateEvaluationRunConfigDict(TypedDict, total=False): - """Config to create an evaluation run.""" + evaluation_dataset: Optional[list[EvaluationDatasetDict]] + """The input evaluation dataset(s) for the evaluation run.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + metadata: Optional[EvaluationRunMetadataDict] + """Metadata for the evaluation run.""" - 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.""" + agent_info: Optional[evals_types.AgentInfo] + """The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""" -CreateEvaluationRunConfigOrDict = Union[ - CreateEvaluationRunConfig, CreateEvaluationRunConfigDict -] +EvaluationResultOrDict = Union[EvaluationResult, EvaluationResultDict] -class _CreateEvaluationRunParameters(_common.BaseModel): - """Represents a job that creates an evaluation run.""" +class EvaluationRun(_common.BaseModel): + """Represents an evaluation run.""" name: Optional[str] = Field(default=None, description="""""") display_name: Optional[str] = Field(default=None, description="""""") - data_source: Optional[EvaluationRunDataSource] = Field( + 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_config: Optional[EvaluationRunConfig] = Field( + 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="""""" ) - labels: Optional[dict[str, str]] = Field(default=None, description="""""") - inference_configs: Optional[dict[str, EvaluationRunInferenceConfig]] = Field( + data_source: Optional[EvaluationRunDataSource] = Field( default=None, description="""""" ) - config: Optional[CreateEvaluationRunConfig] = Field( - default=None, description="""""" + evaluation_run_results: Optional[EvaluationRunResults] = Field( + default=None, description="""The evaluation run formatted results.""" ) - analysis_configs: Optional[list[AnalysisConfig]] = Field( - default=None, description="""""" + evaluation_item_results: Optional[EvaluationResult] = 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_experiment: Optional[str] = Field( + 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.""" + ) + labels: Optional[dict[str, str]] = Field(default=None, description="""""") + analysis_configs: Optional[list[AnalysisConfig]] = 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 analysis configurations for the evaluation run.""", ) encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, @@ -3516,9 +3300,62 @@ class _CreateEvaluationRunParameters(_common.BaseModel): If set, this EvaluationRun will be secured by this key.""", ) + # 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") -class _CreateEvaluationRunParametersDict(TypedDict, total=False): - """Represents a job that creates an evaluation run.""" + 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 + + 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) + + +class EvaluationRunDict(TypedDict, total=False): + """Represents an evaluation run.""" name: Optional[str] """""" @@ -3526,6681 +3363,6362 @@ class _CreateEvaluationRunParametersDict(TypedDict, total=False): display_name: Optional[str] """""" - data_source: Optional[EvaluationRunDataSourceDict] + metadata: Optional[dict[str, Any]] """""" - evaluation_config: Optional[EvaluationRunConfigDict] + create_time: Optional[datetime.datetime] """""" - labels: Optional[dict[str, str]] + completion_time: Optional[datetime.datetime] """""" - inference_configs: Optional[dict[str, EvaluationRunInferenceConfigDict]] + state: Optional[EvaluationRunState] """""" - config: Optional[CreateEvaluationRunConfigDict] + evaluation_set_snapshot: Optional[str] """""" - analysis_configs: Optional[list[AnalysisConfigDict]] + error: Optional[genai_types.GoogleRpcStatus] """""" - evaluation_experiment: Optional[str] - """The resource name of the parent EvaluationExperiment that this run - belongs to. Format: - `projects/{project}/locations/{location}/evaluationExperiments/{evaluation_experiment}`.""" + data_source: Optional[EvaluationRunDataSourceDict] + """""" + + 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.""" + + evaluation_config: Optional[EvaluationRunConfigDict] + """The evaluation config for the evaluation run.""" + + inference_configs: Optional[dict[str, EvaluationRunInferenceConfigDict]] + """The inference configs for the evaluation run.""" + + labels: Optional[dict[str, str]] + """""" + + analysis_configs: Optional[list[AnalysisConfigDict]] + """The analysis configurations for the evaluation run.""" encryption_spec: Optional[genai_types.EncryptionSpec] """Customer-managed encryption key spec for this EvaluationRun. If set, this EvaluationRun will be secured by this key.""" -_CreateEvaluationRunParametersOrDict = Union[ - _CreateEvaluationRunParameters, _CreateEvaluationRunParametersDict -] +EvaluationRunOrDict = Union[EvaluationRun, EvaluationRunDict] -class SummaryMetric(_common.BaseModel): - """Represents a summary metric for an evaluation run.""" +class CreateEvaluationSetConfig(_common.BaseModel): + """Config to create an evaluation set.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class SummaryMetricDict(TypedDict, total=False): - """Represents a summary metric for an evaluation run.""" - - metrics: Optional[dict[str, Any]] - """Map of metric name to metric value.""" - - total_items: Optional[int] - """The total number of items that were evaluated.""" +class CreateEvaluationSetConfigDict(TypedDict, total=False): + """Config to create an evaluation set.""" - failed_items: Optional[int] - """The number of items that failed to be evaluated.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -SummaryMetricOrDict = Union[SummaryMetric, SummaryMetricDict] +CreateEvaluationSetConfigOrDict = Union[ + CreateEvaluationSetConfig, CreateEvaluationSetConfigDict +] -class AttackCategoryResult(_common.BaseModel): - """The red teaming outcome for a specific attack category.""" +class _CreateEvaluationSetParameters(_common.BaseModel): + """Represents a job that creates an evaluation set.""" - attack_category: Optional[str] = Field( - default=None, description="""The category of the attack evaluated.""" + evaluation_items: Optional[list[str]] = Field(default=None, description="""""") + display_name: Optional[str] = Field(default=None, description="""""") + config: Optional[CreateEvaluationSetConfig] = Field( + default=None, description="""""" ) - attack_success_rate: Optional[float] = Field( + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, - 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.""" + description="""Customer-managed encryption key spec for this EvaluationSet. + If set, this EvaluationSet will be secured by this key.""", ) -class AttackCategoryResultDict(TypedDict, total=False): - """The red teaming outcome for a specific attack category.""" +class _CreateEvaluationSetParametersDict(TypedDict, total=False): + """Represents a job that creates an evaluation set.""" - attack_category: Optional[str] - """The category of the attack evaluated.""" + evaluation_items: Optional[list[str]] + """""" - attack_success_rate: Optional[float] - """The ratio of successful attacks given a fixed budget.""" + display_name: Optional[str] + """""" - vulnerability_insight: Optional[str] - """Insights into why an attack succeeded or failed.""" + config: Optional[CreateEvaluationSetConfigDict] + """""" + + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for this EvaluationSet. + If set, this EvaluationSet will be secured by this key.""" -AttackCategoryResultOrDict = Union[AttackCategoryResult, AttackCategoryResultDict] +_CreateEvaluationSetParametersOrDict = Union[ + _CreateEvaluationSetParameters, _CreateEvaluationSetParametersDict +] -class RedTeamingAnalysisResult(_common.BaseModel): - """The top-level result for Red Teaming analysis.""" +class EvaluationSet(_common.BaseModel): + """Represents an evaluation set.""" - config: Optional[RedTeamingAnalysisConfig] = Field( + 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.""" + ) + evaluation_items: Optional[list[str]] = Field( default=None, - description="""The configuration used to generate this analysis.""", + description="""The EvaluationItems that are part of this dataset.""", ) - analysis_time: Optional[str] = Field( - default=None, description="""The timestamp when this analysis was performed.""" + create_time: Optional[datetime.datetime] = Field( + default=None, description="""The create time of the evaluation set.""" ) - category_results: Optional[list[AttackCategoryResult]] = Field( - default=None, description="""Detailed results by attack category.""" + 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.""" + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""Customer-managed encryption key spec for this EvaluationSet. + If set, this EvaluationSet will be secured by this key.""", ) -class RedTeamingAnalysisResultDict(TypedDict, total=False): - """The top-level result for Red Teaming analysis.""" +class EvaluationSetDict(TypedDict, total=False): + """Represents an evaluation set.""" - config: Optional[RedTeamingAnalysisConfigDict] - """The configuration used to generate this analysis.""" + name: Optional[str] + """The resource name of the evaluation set.""" - analysis_time: Optional[str] - """The timestamp when this analysis was performed.""" + display_name: Optional[str] + """The display name of the evaluation set.""" - category_results: Optional[list[AttackCategoryResultDict]] - """Detailed results by attack category.""" + 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.""" -RedTeamingAnalysisResultOrDict = Union[ - RedTeamingAnalysisResult, RedTeamingAnalysisResultDict -] + update_time: Optional[datetime.datetime] + """The update time of the evaluation set.""" + metadata: Optional[dict[str, Any]] + """The metadata of the evaluation set.""" -class LossTaxonomyEntry(_common.BaseModel): - """A specific entry in the loss pattern taxonomy.""" - - l1_category: Optional[str] = Field( - default=None, - description="""The primary category of the loss (e.g., "Hallucination", "Tool Calling").""", - ) - l2_category: Optional[str] = Field( - default=None, - description="""The secondary category of the loss (e.g., "Hallucination of Action", "Incorrect Tool Selection").""", - ) - description: Optional[str] = Field( - default=None, - description="""A detailed description of this loss pattern. Example: "The agent verbally confirms an action without executing the tool." """, - ) - - -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." """ + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for this EvaluationSet. + If set, this EvaluationSet will be secured by this key.""" -LossTaxonomyEntryOrDict = Union[LossTaxonomyEntry, LossTaxonomyEntryDict] +EvaluationSetOrDict = Union[EvaluationSet, EvaluationSetDict] -class FailedRubric(_common.BaseModel): - """A specific failed rubric and the associated analysis.""" +class DeleteEvaluationExperimentConfig(_common.BaseModel): + """Config for deleting an evaluation experiment.""" - 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class FailedRubricDict(TypedDict, total=False): - """A specific failed rubric and the associated analysis.""" - - rubric_id: Optional[str] - """The unique ID of the rubric (if available from the metric source).""" +class DeleteEvaluationExperimentConfigDict(TypedDict, total=False): + """Config for deleting an evaluation experiment.""" - classification_rationale: Optional[str] - """The rationale provided by the Loss Analysis Classifier for why this failure maps to this specific Loss Cluster.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -FailedRubricOrDict = Union[FailedRubric, FailedRubricDict] +DeleteEvaluationExperimentConfigOrDict = Union[ + DeleteEvaluationExperimentConfig, DeleteEvaluationExperimentConfigDict +] -class LossExample(_common.BaseModel): - """A specific example of a loss pattern.""" +class _DeleteEvaluationExperimentParameters(_common.BaseModel): + """Parameters for deleting an evaluation experiment.""" - evaluation_item: Optional[str] = Field( - default=None, - description="""Reference to the persisted EvalItem resource name. Format: projects/.../locations/.../evaluationItems/{item_id}.""", - ) - 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.""", - ) - failed_rubrics: Optional[list[FailedRubric]] = 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.""", + name: Optional[str] = Field(default=None, description="""""") + config: Optional[DeleteEvaluationExperimentConfig] = Field( + default=None, description="""""" ) -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}.""" +class _DeleteEvaluationExperimentParametersDict(TypedDict, total=False): + """Parameters for deleting an evaluation experiment.""" - evaluation_result: Optional[dict[str, Any]] - """The full evaluation result object provided inline. Used when the analysis is performed on ephemeral data.""" + name: Optional[str] + """""" - 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.""" + config: Optional[DeleteEvaluationExperimentConfigDict] + """""" -LossExampleOrDict = Union[LossExample, LossExampleDict] +_DeleteEvaluationExperimentParametersOrDict = Union[ + _DeleteEvaluationExperimentParameters, _DeleteEvaluationExperimentParametersDict +] -class LossCluster(_common.BaseModel): - """A semantic grouping of failures (e.g., "Hallucination of Action").""" +class DeleteEvaluationExperimentOperation(_common.BaseModel): + """Operation for deleting an evaluation experiment.""" - cluster_id: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Unique identifier for the loss cluster within the scope of the analysis result.""", + 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}`.""", ) - taxonomy_entry: Optional[LossTaxonomyEntry] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""The structured definition of the loss taxonomy for this cluster.""", + 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.""", ) - item_count: Optional[int] = Field( + done: Optional[bool] = Field( default=None, - description="""The total number of EvaluationItems falling into this cluster.""", + 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.""", ) - examples: Optional[list[LossExample]] = Field( + error: Optional[dict[str, Any]] = 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 error result of the operation in case of failure or cancellation.""", ) -class LossClusterDict(TypedDict, total=False): - """A semantic grouping of failures (e.g., "Hallucination of Action").""" +class DeleteEvaluationExperimentOperationDict(TypedDict, total=False): + """Operation for deleting an evaluation experiment.""" - cluster_id: Optional[str] - """Unique identifier for the loss cluster within the scope of the analysis result.""" + 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}`.""" - taxonomy_entry: Optional[LossTaxonomyEntryDict] - """The structured definition of the loss taxonomy for this cluster.""" + 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.""" - item_count: Optional[int] - """The total number of EvaluationItems falling into this cluster.""" + 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.""" - examples: Optional[list[LossExampleDict]] - """A list of examples that belong to this cluster. This links the cluster back to the specific EvaluationItems and Rubrics.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -LossClusterOrDict = Union[LossCluster, LossClusterDict] +DeleteEvaluationExperimentOperationOrDict = Union[ + DeleteEvaluationExperimentOperation, DeleteEvaluationExperimentOperationDict +] -class LossAnalysisResult(_common.BaseModel): - """The top-level result for loss analysis.""" +class DeleteEvaluationMetricConfig(_common.BaseModel): + """Config for deleting an evaluation metric.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - def show(self) -> None: - """Shows the loss analysis result with rich HTML visualization.""" - from .. import _evals_visualization - _evals_visualization.display_loss_analysis_result(self) +class DeleteEvaluationMetricConfigDict(TypedDict, total=False): + """Config for deleting an evaluation metric.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -class LossAnalysisResultDict(TypedDict, total=False): - """The top-level result for loss analysis.""" - config: Optional[LossAnalysisConfigDict] - """The configuration used to generate this analysis.""" +DeleteEvaluationMetricConfigOrDict = Union[ + DeleteEvaluationMetricConfig, DeleteEvaluationMetricConfigDict +] - analysis_time: Optional[str] - """The timestamp when this analysis was performed.""" - clusters: Optional[list[LossClusterDict]] - """The list of identified loss clusters.""" +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="""""" + ) -LossAnalysisResultOrDict = Union[LossAnalysisResult, LossAnalysisResultDict] +class _DeleteEvaluationMetricParametersDict(TypedDict, total=False): + """Parameters for deleting an evaluation metric.""" -class EvaluationRunResults(_common.BaseModel): - """Represents the results of an evaluation run.""" + metric_resource_name: Optional[str] + """""" - evaluation_set: Optional[str] = Field( + config: Optional[DeleteEvaluationMetricConfigDict] + """""" + + +_DeleteEvaluationMetricParametersOrDict = Union[ + _DeleteEvaluationMetricParameters, _DeleteEvaluationMetricParametersDict +] + + +class DeleteEvaluationMetricOperation(_common.BaseModel): + """Operation for deleting an evaluation metric.""" + + name: Optional[str] = Field( default=None, - description="""The evaluation set where item level results are stored.""", + 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}`.""", ) - summary_metrics: Optional[SummaryMetric] = Field( - default=None, description="""The summary metrics for the evaluation run.""" + 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.""", ) - loss_analysis_results: Optional[list[LossAnalysisResult]] = Field( + done: Optional[bool] = Field( default=None, - description="""The loss analysis results for the evaluation run.""", + 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.""", ) - red_teaming_analysis_results: Optional[list[RedTeamingAnalysisResult]] = Field( - default=None, description="""The Red Teaming analysis results.""" + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", ) -class EvaluationRunResultsDict(TypedDict, total=False): - """Represents the results of an evaluation run.""" +class DeleteEvaluationMetricOperationDict(TypedDict, total=False): + """Operation for deleting an evaluation metric.""" - evaluation_set: Optional[str] - """The evaluation set where item level results are stored.""" + 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}`.""" - summary_metrics: Optional[SummaryMetricDict] - """The summary metrics for the evaluation run.""" + 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.""" - loss_analysis_results: Optional[list[LossAnalysisResultDict]] - """The loss analysis results for the evaluation run.""" + 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.""" - red_teaming_analysis_results: Optional[list[RedTeamingAnalysisResultDict]] - """The Red Teaming analysis results.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -EvaluationRunResultsOrDict = Union[EvaluationRunResults, EvaluationRunResultsDict] +DeleteEvaluationMetricOperationOrDict = Union[ + DeleteEvaluationMetricOperation, DeleteEvaluationMetricOperationDict +] -class EvalCaseMetricResult(_common.BaseModel): - """Evaluation result for a single evaluation case for a single metric.""" +class DeleteEvaluationSetConfig(_common.BaseModel): + """Config for deleting an evaluation set.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class EvalCaseMetricResultDict(TypedDict, total=False): - """Evaluation result for a single evaluation case for a single metric.""" +class DeleteEvaluationSetConfigDict(TypedDict, total=False): + """Config for deleting an evaluation set.""" - metric_name: Optional[str] - """Name of the metric.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - score: Optional[float] - """Score of the metric.""" - explanation: Optional[str] - """Explanation of the metric.""" +DeleteEvaluationSetConfigOrDict = Union[ + DeleteEvaluationSetConfig, DeleteEvaluationSetConfigDict +] - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] - """The details of all the rubrics and their verdicts for rubric-based metrics.""" - raw_output: Optional[list[str]] - """Raw output of the metric.""" +class _DeleteEvaluationSetParameters(_common.BaseModel): + """Parameters for deleting an evaluation set.""" - error_message: Optional[str] - """Error message for the metric.""" + name: Optional[str] = Field(default=None, description="""""") + config: Optional[DeleteEvaluationSetConfig] = Field( + default=None, description="""""" + ) -EvalCaseMetricResultOrDict = Union[EvalCaseMetricResult, EvalCaseMetricResultDict] +class _DeleteEvaluationSetParametersDict(TypedDict, total=False): + """Parameters for deleting an evaluation set.""" + name: Optional[str] + """""" -class ResponseCandidateResult(_common.BaseModel): - """Aggregated metric results for a single response candidate.""" + config: Optional[DeleteEvaluationSetConfigDict] + """""" - response_index: Optional[int] = Field( + +_DeleteEvaluationSetParametersOrDict = Union[ + _DeleteEvaluationSetParameters, _DeleteEvaluationSetParametersDict +] + + +class DeleteEvaluationSetOperation(_common.BaseModel): + """Operation for deleting an evaluation set.""" + + name: Optional[str] = Field( default=None, - description="""Index of the response candidate this result pertains to.""", + 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}`.""", ) - metric_results: Optional[dict[str, EvalCaseMetricResult]] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""A dictionary of metric results for this response candidate, keyed by metric name.""", + 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 ResponseCandidateResultDict(TypedDict, total=False): - """Aggregated metric results for a single response candidate.""" +class DeleteEvaluationSetOperationDict(TypedDict, total=False): + """Operation for deleting an evaluation set.""" - response_index: Optional[int] - """Index of the response candidate this result pertains to.""" + 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}`.""" - metric_results: Optional[dict[str, EvalCaseMetricResultDict]] - """A dictionary of metric results for this response candidate, keyed by metric name.""" + 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.""" -ResponseCandidateResultOrDict = Union[ - ResponseCandidateResult, ResponseCandidateResultDict + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +DeleteEvaluationSetOperationOrDict = Union[ + DeleteEvaluationSetOperation, DeleteEvaluationSetOperationDict ] -class EvalCaseResult(_common.BaseModel): - """Eval result for a single evaluation case.""" +class BleuInstance(_common.BaseModel): + """Bleu instance.""" - eval_case_index: Optional[int] = Field( - default=None, description="""Index of the evaluation case.""" + prediction: Optional[str] = Field( + default=None, description="""Required. Output of the evaluated model.""" ) - response_candidate_results: Optional[list[ResponseCandidateResult]] = Field( + reference: Optional[str] = Field( default=None, - description="""A list of results, one for each response candidate of the EvalCase.""", + description="""Required. Ground truth used to compare against the prediction.""", ) -class EvalCaseResultDict(TypedDict, total=False): - """Eval result for a single evaluation case.""" +class BleuInstanceDict(TypedDict, total=False): + """Bleu instance.""" - eval_case_index: Optional[int] - """Index of the evaluation case.""" + prediction: Optional[str] + """Required. Output of the evaluated model.""" - response_candidate_results: Optional[list[ResponseCandidateResultDict]] - """A list of results, one for each response candidate of the EvalCase.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" -EvalCaseResultOrDict = Union[EvalCaseResult, EvalCaseResultDict] +BleuInstanceOrDict = Union[BleuInstance, BleuInstanceDict] -class AggregatedMetricResult(_common.BaseModel): - """Evaluation result for a single metric for an evaluation dataset.""" +class BleuInput(_common.BaseModel): - 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.""" + instances: Optional[list[BleuInstance]] = Field( + default=None, description="""Required. Repeated bleu instances.""" ) - mean_score: Optional[float] = Field( - default=None, description="""Mean score of the metric.""" + metric_spec: Optional[genai_types.BleuSpec] = Field( + default=None, description="""Required. Spec for bleu score metric.""" ) - stdev_score: Optional[float] = Field( - default=None, description="""Standard deviation of the metric.""" + + +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.""" ) - pass_rate: Optional[float] = Field( + reference: Optional[str] = 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.""", + description="""Required. Ground truth used to compare against the prediction.""", ) - # Allow extra fields to support custom aggregation stats. - model_config = ConfigDict(extra="allow") +class ExactMatchInstanceDict(TypedDict, total=False): + """Exact match instance.""" -class AggregatedMetricResultDict(TypedDict, total=False): - """Evaluation result for a single metric for an evaluation dataset.""" + prediction: Optional[str] + """Required. Output of the evaluated model.""" - metric_name: Optional[str] - """Name of the metric.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" - num_cases_total: Optional[int] - """Total number of cases in the dataset.""" - num_cases_valid: Optional[int] - """Number of valid cases in the dataset.""" +ExactMatchInstanceOrDict = Union[ExactMatchInstance, ExactMatchInstanceDict] - num_cases_error: Optional[int] - """Number of cases with errors in the dataset.""" - mean_score: Optional[float] - """Mean score of the metric.""" +class ExactMatchSpec(_common.BaseModel): + """Spec for exact match metric.""" - stdev_score: Optional[float] - """Standard deviation of the metric.""" + pass - 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.""" +class ExactMatchSpecDict(TypedDict, total=False): + """Spec for exact match metric.""" -AggregatedMetricResultOrDict = Union[AggregatedMetricResult, AggregatedMetricResultDict] + pass -class WinRateStats(_common.BaseModel): - """Statistics for win rates for a single metric.""" +ExactMatchSpecOrDict = Union[ExactMatchSpec, ExactMatchSpecDict] - win_rates: Optional[list[float]] = Field( - default=None, - description="""Win rates for the metric, one for each candidate.""", + +class ExactMatchInput(_common.BaseModel): + + instances: Optional[list[ExactMatchInstance]] = Field( + default=None, description="""Required. Repeated exact match instances.""" ) - tie_rate: Optional[float] = Field( - default=None, description="""Tie rate for the metric.""" + metric_spec: Optional[ExactMatchSpec] = Field( + default=None, description="""Required. Spec for exact match metric.""" ) -class WinRateStatsDict(TypedDict, total=False): - """Statistics for win rates for a single metric.""" +class ExactMatchInputDict(TypedDict, total=False): - win_rates: Optional[list[float]] - """Win rates for the metric, one for each candidate.""" + instances: Optional[list[ExactMatchInstanceDict]] + """Required. Repeated exact match instances.""" - tie_rate: Optional[float] - """Tie rate for the metric.""" + metric_spec: Optional[ExactMatchSpecDict] + """Required. Spec for exact match metric.""" -WinRateStatsOrDict = Union[WinRateStats, WinRateStatsDict] +ExactMatchInputOrDict = Union[ExactMatchInput, ExactMatchInputDict] -class ResponseCandidate(_common.BaseModel): - """A model-generated content to the prompt.""" +class RougeInstance(_common.BaseModel): + """Rouge instance.""" - response: Optional[genai_types.Content] = Field( + prediction: Optional[str] = Field( + default=None, description="""Required. Output of the evaluated model.""" + ) + reference: Optional[str] = Field( default=None, - description="""The final model-generated response to the `prompt`.""", + description="""Required. Ground truth used to compare against the prediction.""", ) -class ResponseCandidateDict(TypedDict, total=False): - """A model-generated content to the prompt.""" +class RougeInstanceDict(TypedDict, total=False): + """Rouge instance.""" - response: Optional[genai_types.Content] - """The final model-generated response to the `prompt`.""" + prediction: Optional[str] + """Required. Output of the evaluated model.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" -ResponseCandidateOrDict = Union[ResponseCandidate, ResponseCandidateDict] +RougeInstanceOrDict = Union[RougeInstance, RougeInstanceDict] -class InteractionsDataSource(_common.BaseModel): - """Source for populating agent data from an Interactions API interaction.""" - gemini_agent_config: Optional[GeminiAgentConfig] = Field( - default=None, - description="""The Gemini Agent (Vertex AI Agent resource) that produced the - interaction.""", +class RougeInput(_common.BaseModel): + """Rouge input.""" + + instances: Optional[list[RougeInstance]] = Field( + default=None, description="""Required. Repeated rouge instances.""" ) - interaction: Optional[str] = Field( - default=None, - description="""The interaction to evaluate. Required by the backend. - Format: - `projects/{project}/locations/{location}/interactions/{interaction}`.""", + metric_spec: Optional[genai_types.RougeSpec] = Field( + default=None, description="""Required. Spec for rouge score metric.""" ) -class InteractionsDataSourceDict(TypedDict, total=False): - """Source for populating agent data from an Interactions API interaction.""" +class RougeInputDict(TypedDict, total=False): + """Rouge input.""" - gemini_agent_config: Optional[GeminiAgentConfigDict] - """The Gemini Agent (Vertex AI Agent resource) that produced the - interaction.""" + instances: Optional[list[RougeInstanceDict]] + """Required. Repeated rouge instances.""" - interaction: Optional[str] - """The interaction to evaluate. Required by the backend. - Format: - `projects/{project}/locations/{location}/interactions/{interaction}`.""" + metric_spec: Optional[genai_types.RougeSpec] + """Required. Spec for rouge score metric.""" -InteractionsDataSourceOrDict = Union[InteractionsDataSource, InteractionsDataSourceDict] +RougeInputOrDict = Union[RougeInput, RougeInputDict] -class EvalCase(_common.BaseModel): - """A comprehensive representation of a GenAI interaction for evaluation.""" +class ContentMap(_common.BaseModel): + """Map of placeholder in metric prompt template to contents of model input.""" - prompt: Optional[genai_types.Content] = Field( - default=None, description="""The most recent user message (current input).""" - ) - responses: Optional[list[ResponseCandidate]] = 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.""", - ) - reference: Optional[ResponseCandidate] = 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.""" - ) - conversation_history: Optional[list[evals_types.Message]] = Field( - default=None, - 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.""", + values: Optional[dict[str, "ContentMapContents"]] = Field( + default=None, description="""Map of placeholder to contents.""" ) - # 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.""" - prompt: Optional[genai_types.Content] - """The most recent user message (current input).""" - 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.""" +class ContentMapDict(TypedDict, total=False): + """Map of placeholder in metric prompt template to contents of model input.""" - reference: Optional[ResponseCandidateDict] - """User-provided, golden reference model reply to prompt in context of chat history; Reference for last response in a conversation.""" + values: Optional[dict[str, "ContentMapContents"]] + """Map of placeholder to contents.""" - 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).""" +ContentMapOrDict = Union[ContentMap, ContentMapDict] - 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.""" - eval_case_id: Optional[str] - """Unique identifier for the evaluation case.""" +class PointwiseMetricInstance(_common.BaseModel): + """Pointwise metric instance.""" - 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.""" + 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.""", + ) - agent_info: Optional[evals_types.AgentInfo] - """The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""" - agent_data: Optional[evals_types.AgentData] - """The agent data of the agent under evaluation.""" +class PointwiseMetricInstanceDict(TypedDict, total=False): + """Pointwise metric instance.""" - user_scenario: Optional[evals_types.UserScenario] - """The user scenario for the evaluation case.""" + 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.""" - 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.""" + 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.""" -EvalCaseOrDict = Union[EvalCase, EvalCaseDict] +PointwiseMetricInstanceOrDict = Union[ + PointwiseMetricInstance, PointwiseMetricInstanceDict +] -class EvaluationDataset(_common.BaseModel): - """The dataset used for evaluation.""" +class PointwiseMetricInput(_common.BaseModel): + """Pointwise metric input.""" - 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.""" - ) - eval_cases: Optional[list[EvalCase]] = Field( - default=None, description="""The evaluation cases to be evaluated.""" - ) - eval_dataset_df: Optional[PandasDataFrame] = Field( - default=None, - description="""The evaluation dataset in the form of a Pandas DataFrame.""", + instance: Optional[PointwiseMetricInstance] = Field( + default=None, description="""Required. Pointwise metric instance.""" ) - candidate_name: Optional[str] = Field( - default=None, - description="""The name of the candidate model or agent for this evaluation dataset.""", + metric_spec: Optional[genai_types.PointwiseMetricSpec] = Field( + default=None, description="""Required. Spec for pointwise metric.""" ) - @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 PointwiseMetricInputDict(TypedDict, total=False): + """Pointwise metric input.""" - formats = [] - requests = [] - responses = [] - system_instructions = [] + instance: Optional[PointwiseMetricInstanceDict] + """Required. Pointwise metric instance.""" - for case in cases: - gcs_utils = _gcs_utils.GcsUtils( - case.api_client._api_client if case.api_client else None - ) + metric_spec: Optional[genai_types.PointwiseMetricSpec] + """Required. Spec for pointwise metric.""" - # 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) +PointwiseMetricInputOrDict = Union[PointwiseMetricInput, PointwiseMetricInputDict] - # 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 PairwiseMetricInstance(_common.BaseModel): + """Pairwise metric instance.""" - eval_dataset_df = pd.DataFrame( - { - "format": formats, - "request": requests, - "response": responses, - "system_instruction": system_instructions, - } - ) + 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.""", + ) - except ImportError as e: - raise ImportError("Pandas DataFrame library is required.") from e - return EvaluationDataset(eval_dataset_df=eval_dataset_df) +class PairwiseMetricInstanceDict(TypedDict, total=False): + """Pairwise metric instance.""" - def show(self) -> None: - """Shows the evaluation dataset.""" - 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 PairwiseMetricSpec.instance_prompt_template.""" - _evals_visualization.display_evaluation_dataset(self) +PairwiseMetricInstanceOrDict = Union[PairwiseMetricInstance, PairwiseMetricInstanceDict] -class EvaluationDatasetDict(TypedDict, total=False): - """The dataset used for evaluation.""" - bigquery_source: Optional[genai_types.BigQuerySource] - """The BigQuery source for the evaluation dataset.""" +class PairwiseMetricInput(_common.BaseModel): + """Pairwise metric instance.""" - gcs_source: Optional[genai_types.GcsSource] - """The GCS source for the evaluation dataset.""" + 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.""" + ) - 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 PairwiseMetricInputDict(TypedDict, total=False): + """Pairwise metric instance.""" - candidate_name: Optional[str] - """The name of the candidate model or agent for this evaluation dataset.""" + instance: Optional[PairwiseMetricInstanceDict] + """Required. Pairwise metric instance.""" + metric_spec: Optional[genai_types.PairwiseMetricSpec] + """Required. Spec for pairwise metric.""" -EvaluationDatasetOrDict = Union[EvaluationDataset, EvaluationDatasetDict] +PairwiseMetricInputOrDict = Union[PairwiseMetricInput, PairwiseMetricInputDict] -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 ToolCallValidInstance(_common.BaseModel): + """Tool call valid 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 ToolCallValidInstanceDict(TypedDict, total=False): + """Tool call valid 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.""" +ToolCallValidInstanceOrDict = Union[ToolCallValidInstance, ToolCallValidInstanceDict] -EvaluationRunMetadataOrDict = Union[EvaluationRunMetadata, EvaluationRunMetadataDict] +class ToolCallValidSpec(_common.BaseModel): + """Spec for tool call valid metric.""" -class EvaluationResult(_common.BaseModel): - """Result of an evaluation run for an evaluation dataset.""" + pass - 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.""" - ) - 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.""", - ) - def show(self, candidate_names: Optional[List[str]] = None) -> None: - """Shows the evaluation result. +class ToolCallValidSpecDict(TypedDict, total=False): + """Spec for tool call valid metric.""" - Args: - candidate_names: list of names for the evaluated candidates, used in - comparison reports. - """ - from .. import _evals_visualization + pass - _evals_visualization.display_evaluation_result(self, candidate_names) +ToolCallValidSpecOrDict = Union[ToolCallValidSpec, ToolCallValidSpecDict] -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.""" +class ToolCallValidInput(_common.BaseModel): + """Tool call valid input.""" - summary_metrics: Optional[list[AggregatedMetricResultDict]] - """A list of summary-level evaluation results for each metric.""" + 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.""" + ) - win_rates: Optional[dict[str, WinRateStatsDict]] - """A dictionary of win rates for each metric, only populated for multi-response evaluation runs.""" - evaluation_dataset: Optional[list[EvaluationDatasetDict]] - """The input evaluation dataset(s) for the evaluation run.""" +class ToolCallValidInputDict(TypedDict, total=False): + """Tool call valid input.""" - metadata: Optional[EvaluationRunMetadataDict] - """Metadata for the evaluation run.""" + instances: Optional[list[ToolCallValidInstanceDict]] + """Required. Repeated tool call valid instances.""" - agent_info: Optional[evals_types.AgentInfo] - """The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""" + metric_spec: Optional[ToolCallValidSpecDict] + """Required. Spec for tool call valid metric.""" -EvaluationResultOrDict = Union[EvaluationResult, EvaluationResultDict] +ToolCallValidInputOrDict = Union[ToolCallValidInput, ToolCallValidInputDict] -class EvaluationRun(_common.BaseModel): - """Represents an evaluation run.""" +class ToolNameMatchInstance(_common.BaseModel): + """Tool name match 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( - 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.""" - ) - labels: Optional[dict[str, str]] = Field(default=None, description="""""") - analysis_configs: Optional[list[AnalysisConfig]] = Field( - default=None, - description="""The analysis configurations for the evaluation run.""", + prediction: Optional[str] = Field( + default=None, description="""Required. Output of the evaluated model.""" ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + reference: Optional[str] = Field( default=None, - description="""Customer-managed encryption key spec for this EvaluationRun. - If set, this EvaluationRun will be secured by this key.""", + description="""Required. Ground truth used to compare against the prediction.""", ) - # 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) +class ToolNameMatchInstanceDict(TypedDict, total=False): + """Tool name match instance.""" - evaluation_run_results = response.get("evaluation_run_results") + prediction: Optional[str] + """Required. Output of the evaluated model.""" - 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 + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" - 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) +ToolNameMatchInstanceOrDict = Union[ToolNameMatchInstance, ToolNameMatchInstanceDict] -class EvaluationRunDict(TypedDict, total=False): - """Represents an evaluation run.""" +class ToolNameMatchSpec(_common.BaseModel): + """Spec for tool name match metric.""" - name: Optional[str] - """""" + pass - display_name: Optional[str] - """""" - metadata: Optional[dict[str, Any]] - """""" +class ToolNameMatchSpecDict(TypedDict, total=False): + """Spec for tool name match metric.""" - create_time: Optional[datetime.datetime] - """""" + pass - completion_time: Optional[datetime.datetime] - """""" - state: Optional[EvaluationRunState] - """""" +ToolNameMatchSpecOrDict = Union[ToolNameMatchSpec, ToolNameMatchSpecDict] - evaluation_set_snapshot: Optional[str] - """""" - error: Optional[genai_types.GoogleRpcStatus] - """""" +class ToolNameMatchInput(_common.BaseModel): + """Tool name match input.""" - data_source: Optional[EvaluationRunDataSourceDict] - """""" + instances: Optional[list[ToolNameMatchInstance]] = Field( + default=None, description="""Required. Repeated tool name match instances.""" + ) + metric_spec: Optional[ToolNameMatchSpec] = Field( + default=None, description="""Required. Spec for tool name match metric.""" + ) - 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 ToolNameMatchInputDict(TypedDict, total=False): + """Tool name match input.""" - evaluation_config: Optional[EvaluationRunConfigDict] - """The evaluation config for the evaluation run.""" + instances: Optional[list[ToolNameMatchInstanceDict]] + """Required. Repeated tool name match instances.""" - inference_configs: Optional[dict[str, EvaluationRunInferenceConfigDict]] - """The inference configs for the evaluation run.""" + metric_spec: Optional[ToolNameMatchSpecDict] + """Required. Spec for tool name match metric.""" - labels: Optional[dict[str, str]] - """""" - analysis_configs: Optional[list[AnalysisConfigDict]] - """The analysis configurations for the evaluation run.""" - - encryption_spec: Optional[genai_types.EncryptionSpec] - """Customer-managed encryption key spec for this EvaluationRun. - If set, this EvaluationRun will be secured by this key.""" - - -EvaluationRunOrDict = Union[EvaluationRun, EvaluationRunDict] +ToolNameMatchInputOrDict = Union[ToolNameMatchInput, ToolNameMatchInputDict] -class CreateEvaluationSetConfig(_common.BaseModel): - """Config to create an evaluation set.""" +class ToolParameterKeyMatchInstance(_common.BaseModel): + """Tool parameter key match instance.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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.""", ) -class CreateEvaluationSetConfigDict(TypedDict, total=False): - """Config to create an evaluation set.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - +class ToolParameterKeyMatchInstanceDict(TypedDict, total=False): + """Tool parameter key match instance.""" -CreateEvaluationSetConfigOrDict = Union[ - CreateEvaluationSetConfig, CreateEvaluationSetConfigDict -] + prediction: Optional[str] + """Required. Output of the evaluated model.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" -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="""""" - ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, - description="""Customer-managed encryption key spec for this EvaluationSet. - If set, this EvaluationSet will be secured by this key.""", - ) +ToolParameterKeyMatchInstanceOrDict = Union[ + ToolParameterKeyMatchInstance, ToolParameterKeyMatchInstanceDict +] -class _CreateEvaluationSetParametersDict(TypedDict, total=False): - """Represents a job that creates an evaluation set.""" +class ToolParameterKeyMatchSpec(_common.BaseModel): + """Spec for tool parameter key match metric.""" - evaluation_items: Optional[list[str]] - """""" + pass - display_name: Optional[str] - """""" - config: Optional[CreateEvaluationSetConfigDict] - """""" +class ToolParameterKeyMatchSpecDict(TypedDict, total=False): + """Spec for tool parameter key match metric.""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """Customer-managed encryption key spec for this EvaluationSet. - If set, this EvaluationSet will be secured by this key.""" + pass -_CreateEvaluationSetParametersOrDict = Union[ - _CreateEvaluationSetParameters, _CreateEvaluationSetParametersDict +ToolParameterKeyMatchSpecOrDict = Union[ + ToolParameterKeyMatchSpec, ToolParameterKeyMatchSpecDict ] -class EvaluationSet(_common.BaseModel): - """Represents an evaluation set.""" +class ToolParameterKeyMatchInput(_common.BaseModel): + """Tool parameter key match input.""" - 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.""" - ) - evaluation_items: Optional[list[str]] = Field( + instances: Optional[list[ToolParameterKeyMatchInstance]] = 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. Repeated tool parameter key match instances.""", ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + metric_spec: Optional[ToolParameterKeyMatchSpec] = Field( default=None, - description="""Customer-managed encryption key spec for this EvaluationSet. - If set, this EvaluationSet will be secured by this key.""", + description="""Required. Spec for tool parameter key match metric.""", ) -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.""" - - update_time: Optional[datetime.datetime] - """The update time of the evaluation set.""" +class ToolParameterKeyMatchInputDict(TypedDict, total=False): + """Tool parameter key match input.""" - metadata: Optional[dict[str, Any]] - """The metadata of the evaluation set.""" + instances: Optional[list[ToolParameterKeyMatchInstanceDict]] + """Required. Repeated tool parameter key match instances.""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """Customer-managed encryption key spec for this EvaluationSet. - If set, this EvaluationSet will be secured by this key.""" + metric_spec: Optional[ToolParameterKeyMatchSpecDict] + """Required. Spec for tool parameter key match metric.""" -EvaluationSetOrDict = Union[EvaluationSet, EvaluationSetDict] +ToolParameterKeyMatchInputOrDict = Union[ + ToolParameterKeyMatchInput, ToolParameterKeyMatchInputDict +] -class DeleteEvaluationExperimentConfig(_common.BaseModel): - """Config for deleting an evaluation experiment.""" +class ToolParameterKVMatchInstance(_common.BaseModel): + """Tool parameter kv match instance.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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.""", ) -class DeleteEvaluationExperimentConfigDict(TypedDict, total=False): - """Config for deleting an evaluation experiment.""" +class ToolParameterKVMatchInstanceDict(TypedDict, total=False): + """Tool parameter kv 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.""" -DeleteEvaluationExperimentConfigOrDict = Union[ - DeleteEvaluationExperimentConfig, DeleteEvaluationExperimentConfigDict +ToolParameterKVMatchInstanceOrDict = Union[ + ToolParameterKVMatchInstance, ToolParameterKVMatchInstanceDict ] -class _DeleteEvaluationExperimentParameters(_common.BaseModel): - """Parameters for deleting an evaluation experiment.""" +class ToolParameterKVMatchSpec(_common.BaseModel): + """Spec for tool parameter kv match metric.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[DeleteEvaluationExperimentConfig] = Field( - default=None, description="""""" + use_strict_string_match: Optional[bool] = Field( + default=None, + description="""Optional. Whether to use STRICT string match on parameter values.""", ) -class _DeleteEvaluationExperimentParametersDict(TypedDict, total=False): - """Parameters for deleting an evaluation experiment.""" - - name: Optional[str] - """""" +class ToolParameterKVMatchSpecDict(TypedDict, total=False): + """Spec for tool parameter kv match metric.""" - config: Optional[DeleteEvaluationExperimentConfigDict] - """""" + use_strict_string_match: Optional[bool] + """Optional. Whether to use STRICT string match on parameter values.""" -_DeleteEvaluationExperimentParametersOrDict = Union[ - _DeleteEvaluationExperimentParameters, _DeleteEvaluationExperimentParametersDict +ToolParameterKVMatchSpecOrDict = Union[ + ToolParameterKVMatchSpec, ToolParameterKVMatchSpecDict ] -class DeleteEvaluationExperimentOperation(_common.BaseModel): - """Operation for deleting an evaluation experiment.""" +class ToolParameterKVMatchInput(_common.BaseModel): + """Tool parameter kv 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( + instances: Optional[list[ToolParameterKVMatchInstance]] = 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. Repeated tool parameter key value match instances.""", ) - error: Optional[dict[str, Any]] = Field( + metric_spec: Optional[ToolParameterKVMatchSpec] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Required. Spec for tool parameter key value match metric.""", ) -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.""" +class ToolParameterKVMatchInputDict(TypedDict, total=False): + """Tool parameter kv 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[ToolParameterKVMatchInstanceDict]] + """Required. Repeated tool parameter key value match instances.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + metric_spec: Optional[ToolParameterKVMatchSpecDict] + """Required. Spec for tool parameter key value match metric.""" -DeleteEvaluationExperimentOperationOrDict = Union[ - DeleteEvaluationExperimentOperation, DeleteEvaluationExperimentOperationDict +ToolParameterKVMatchInputOrDict = Union[ + ToolParameterKVMatchInput, ToolParameterKVMatchInputDict ] -class DeleteEvaluationMetricConfig(_common.BaseModel): - """Config for deleting an evaluation metric.""" +class MapInstance(_common.BaseModel): + """Instance data specified as a map.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + map_instance: Optional[dict[str, evals_types.InstanceData]] = Field( + default=None, description="""Map of instance data.""" ) -class DeleteEvaluationMetricConfigDict(TypedDict, total=False): - """Config for deleting an evaluation metric.""" +class MapInstanceDict(TypedDict, total=False): + """Instance data specified as a map.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + map_instance: Optional[dict[str, evals_types.InstanceData]] + """Map of instance data.""" -DeleteEvaluationMetricConfigOrDict = Union[ - DeleteEvaluationMetricConfig, DeleteEvaluationMetricConfigDict -] +MapInstanceOrDict = Union[MapInstance, MapInstanceDict] -class _DeleteEvaluationMetricParameters(_common.BaseModel): - """Parameters for deleting an evaluation metric.""" +class EvaluationInstance(_common.BaseModel): + """A single instance to be evaluated.""" - metric_resource_name: Optional[str] = Field(default=None, description="""""") - config: Optional[DeleteEvaluationMetricConfig] = Field( - default=None, description="""""" + prompt: Optional[evals_types.InstanceData] = Field( + default=None, + description="""Data used to populate placeholder `prompt` in a metric prompt template.""", ) - - -class _DeleteEvaluationMetricParametersDict(TypedDict, total=False): - """Parameters for deleting an evaluation metric.""" - - metric_resource_name: Optional[str] - """""" - - config: Optional[DeleteEvaluationMetricConfigDict] - """""" - - -_DeleteEvaluationMetricParametersOrDict = Union[ - _DeleteEvaluationMetricParameters, _DeleteEvaluationMetricParametersDict -] - - -class DeleteEvaluationMetricOperation(_common.BaseModel): - """Operation for deleting an evaluation metric.""" - - name: Optional[str] = Field( + response: Optional[evals_types.InstanceData] = 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="""Data used to populate placeholder `response` in a metric prompt template.""", ) - metadata: Optional[dict[str, Any]] = Field( + reference: Optional[evals_types.InstanceData] = 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="""Data used to populate placeholder `reference` in a metric prompt template.""", ) - done: Optional[bool] = Field( + other_data: Optional[MapInstance] = 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="""Other data used to populate placeholders based on their key.""", ) - error: Optional[dict[str, Any]] = Field( + 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="""The error result of the operation in case of failure or cancellation.""", + 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 DeleteEvaluationMetricOperationDict(TypedDict, total=False): - """Operation for deleting an evaluation metric.""" +class EvaluationInstanceDict(TypedDict, total=False): + """A single instance to be evaluated.""" - 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}`.""" + prompt: Optional[evals_types.InstanceData] + """Data used to populate placeholder `prompt` in a metric prompt template.""" - 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.""" + response: Optional[evals_types.InstanceData] + """Data used to populate placeholder `response` in a metric prompt template.""" - 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.""" + reference: Optional[evals_types.InstanceData] + """Data used to populate placeholder `reference` in a metric prompt template.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + 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.""" -DeleteEvaluationMetricOperationOrDict = Union[ - DeleteEvaluationMetricOperation, DeleteEvaluationMetricOperationDict -] + 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.""" + 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.""" -class DeleteEvaluationSetConfig(_common.BaseModel): - """Config for deleting an evaluation set.""" + +EvaluationInstanceOrDict = Union[EvaluationInstance, EvaluationInstanceDict] + + +class EvaluateInstancesConfig(_common.BaseModel): + """Config for evaluate instances.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class DeleteEvaluationSetConfigDict(TypedDict, total=False): - """Config for deleting an evaluation set.""" +class EvaluateInstancesConfigDict(TypedDict, total=False): + """Config for evaluate instances.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteEvaluationSetConfigOrDict = Union[ - DeleteEvaluationSetConfig, DeleteEvaluationSetConfigDict +EvaluateInstancesConfigOrDict = Union[ + EvaluateInstancesConfig, EvaluateInstancesConfigDict ] -class _DeleteEvaluationSetParameters(_common.BaseModel): - """Parameters for deleting an evaluation set.""" +class RubricBasedMetricSpec(_common.BaseModel): + """Specification for a metric that is based on rubrics.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[DeleteEvaluationSetConfig] = Field( - default=None, description="""""" + 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( + default=None, + description="""Dynamically generate rubrics for evaluation using this specification.""", ) -class _DeleteEvaluationSetParametersDict(TypedDict, total=False): - """Parameters for deleting an evaluation set.""" +class RubricBasedMetricSpecDict(TypedDict, total=False): + """Specification for a metric that is based on rubrics.""" - name: Optional[str] - """""" + metric_prompt_template: Optional[str] + """Template for the prompt used by the judge model to evaluate against + rubrics.""" - config: Optional[DeleteEvaluationSetConfigDict] - """""" + 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.""" + 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`.""" -_DeleteEvaluationSetParametersOrDict = Union[ - _DeleteEvaluationSetParameters, _DeleteEvaluationSetParametersDict -] + rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] + """Dynamically generate rubrics for evaluation using this specification.""" -class DeleteEvaluationSetOperation(_common.BaseModel): - """Operation for deleting an evaluation set.""" +RubricBasedMetricSpecOrDict = Union[RubricBasedMetricSpec, RubricBasedMetricSpecDict] - name: Optional[str] = Field( + +class RubricEnhancedContents(_common.BaseModel): + """Rubric-enhanced contents for evaluation.""" + + prompt: Optional[list[genai_types.Content]] = 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="""User prompt, using the standard Content type from the Gen AI SDK.""", ) - metadata: Optional[dict[str, Any]] = Field( + rubric_groups: Optional[dict[str, "RubricGroup"]] = 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="""Named groups of rubrics associated with this prompt. + The key is a user-defined name for the rubric group.""", ) - done: Optional[bool] = Field( + response: Optional[list[genai_types.Content]] = 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="""Response, using the standard Content type from the Gen AI SDK.""", ) - error: Optional[dict[str, Any]] = Field( + other_content: Optional[ContentMap] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Other contents needed for the metric. + For example, if `reference` is needed for the metric, it can be provided + here.""", ) -class DeleteEvaluationSetOperationDict(TypedDict, total=False): - """Operation for deleting an evaluation set.""" +class RubricEnhancedContentsDict(TypedDict, total=False): + """Rubric-enhanced contents for evaluation.""" - 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}`.""" + prompt: Optional[list[genai_types.Content]] + """User prompt, using the standard Content type from the Gen AI SDK.""" - 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.""" + 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.""" - 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.""" + response: Optional[list[genai_types.Content]] + """Response, using the standard Content type from the Gen AI SDK.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + other_content: Optional[ContentMapDict] + """Other contents needed for the metric. + For example, if `reference` is needed for the metric, it can be provided + here.""" -DeleteEvaluationSetOperationOrDict = Union[ - DeleteEvaluationSetOperation, DeleteEvaluationSetOperationDict -] +RubricEnhancedContentsOrDict = Union[RubricEnhancedContents, RubricEnhancedContentsDict] -class BleuInstance(_common.BaseModel): - """Bleu instance.""" +class RubricBasedMetricInstance(_common.BaseModel): + """Defines an instance for Rubric-based metrics. - prediction: Optional[str] = Field( - default=None, description="""Required. Output of the evaluated model.""" + This class allows various input formats. + """ + + json_instance: Optional[str] = Field( + default=None, + description="""Specify evaluation fields and their string values in JSON format.""", ) - reference: Optional[str] = Field( + content_map_instance: Optional[ContentMap] = Field( default=None, - description="""Required. Ground truth used to compare against the prediction.""", + 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.""", ) -class BleuInstanceDict(TypedDict, total=False): - """Bleu instance.""" +class RubricBasedMetricInstanceDict(TypedDict, total=False): + """Defines an instance for Rubric-based metrics. - prediction: Optional[str] - """Required. Output of the evaluated model.""" + This class allows various input formats. + """ - reference: Optional[str] - """Required. Ground truth used to compare against the prediction.""" + 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.""" -BleuInstanceOrDict = Union[BleuInstance, BleuInstanceDict] + rubric_enhanced_contents: Optional[RubricEnhancedContentsDict] + """Provide input as Gemini Content along with one or more + associated rubric groups.""" -class BleuInput(_common.BaseModel): +RubricBasedMetricInstanceOrDict = Union[ + RubricBasedMetricInstance, RubricBasedMetricInstanceDict +] - instances: Optional[list[BleuInstance]] = Field( - default=None, description="""Required. Repeated bleu instances.""" + +class RubricBasedMetricInput(_common.BaseModel): + """Input for a rubric-based metrics.""" + + metric_spec: Optional[RubricBasedMetricSpec] = Field( + default=None, description="""Specification for the rubric-based metric.""" ) - metric_spec: Optional[genai_types.BleuSpec] = Field( - default=None, description="""Required. Spec for bleu score metric.""" + instance: Optional[RubricBasedMetricInstance] = Field( + default=None, description="""The instance to be evaluated.""" ) -class BleuInputDict(TypedDict, total=False): - - instances: Optional[list[BleuInstanceDict]] - """Required. Repeated bleu instances.""" +class RubricBasedMetricInputDict(TypedDict, total=False): + """Input for a rubric-based metrics.""" - metric_spec: Optional[genai_types.BleuSpec] - """Required. Spec for bleu score metric.""" + metric_spec: Optional[RubricBasedMetricSpecDict] + """Specification for the rubric-based metric.""" + instance: Optional[RubricBasedMetricInstanceDict] + """The instance to be evaluated.""" -BleuInputOrDict = Union[BleuInput, BleuInputDict] +RubricBasedMetricInputOrDict = Union[RubricBasedMetricInput, RubricBasedMetricInputDict] -class ExactMatchInstance(_common.BaseModel): - """Exact match instance.""" - prediction: Optional[str] = Field( - default=None, description="""Required. Output of the evaluated model.""" +class MetricSource(_common.BaseModel): + """The metric source used for evaluation.""" + + 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 ExactMatchInstanceDict(TypedDict, total=False): - """Exact 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}""" -ExactMatchInstanceOrDict = Union[ExactMatchInstance, ExactMatchInstanceDict] +MetricSourceOrDict = Union[MetricSource, MetricSourceDict] -class ExactMatchSpec(_common.BaseModel): - """Spec for exact 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 ExactMatchSpecDict(TypedDict, total=False): - """Spec for exact match metric.""" +class _EvaluateInstancesRequestParametersDict(TypedDict, total=False): + """Parameters for evaluating instances.""" - pass + bleu_input: Optional[BleuInputDict] + """""" + exact_match_input: Optional[ExactMatchInputDict] + """""" -ExactMatchSpecOrDict = Union[ExactMatchSpec, ExactMatchSpecDict] + rouge_input: Optional[RougeInputDict] + """""" + pointwise_metric_input: Optional[PointwiseMetricInputDict] + """""" -class ExactMatchInput(_common.BaseModel): + pairwise_metric_input: Optional[PairwiseMetricInputDict] + """""" - 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.""" - ) + tool_call_valid_input: Optional[ToolCallValidInputDict] + """""" + tool_name_match_input: Optional[ToolNameMatchInputDict] + """""" -class ExactMatchInputDict(TypedDict, total=False): + tool_parameter_key_match_input: Optional[ToolParameterKeyMatchInputDict] + """""" - instances: Optional[list[ExactMatchInstanceDict]] - """Required. Repeated exact match instances.""" + tool_parameter_kv_match_input: Optional[ToolParameterKVMatchInputDict] + """""" - metric_spec: Optional[ExactMatchSpecDict] - """Required. Spec for exact 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.""" -ExactMatchInputOrDict = Union[ExactMatchInput, ExactMatchInputDict] + instance: Optional[EvaluationInstanceDict] + """The instance to be evaluated.""" + metric_sources: Optional[list[MetricSourceDict]] + """The metrics used for evaluation.""" -class RougeInstance(_common.BaseModel): - """Rouge instance.""" + config: Optional[EvaluateInstancesConfigDict] + """""" - prediction: Optional[str] = Field( - default=None, description="""Required. Output of the evaluated model.""" + +_EvaluateInstancesRequestParametersOrDict = Union[ + _EvaluateInstancesRequestParameters, _EvaluateInstancesRequestParametersDict +] + + +class MetricResult(_common.BaseModel): + """Result for a single metric on a single instance.""" + + 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 RougeInstanceDict(TypedDict, total=False): - """Rouge 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.""" -RougeInstanceOrDict = Union[RougeInstance, RougeInstanceDict] + error: Optional[genai_types.GoogleRpcStatus] + """The error status for the metric result.""" -class RougeInput(_common.BaseModel): - """Rouge input.""" +MetricResultOrDict = Union[MetricResult, MetricResultDict] - instances: Optional[list[RougeInstance]] = Field( - default=None, description="""Required. Repeated rouge instances.""" - ) - metric_spec: Optional[genai_types.RougeSpec] = Field( - default=None, description="""Required. Spec for rouge score metric.""" - ) +class BleuResults(_common.BaseModel): + """Result of evaluating a bleu metric.""" -class RougeInputDict(TypedDict, total=False): - """Rouge input.""" + bleu_metric_values: Optional[list[genai_types.BleuMetricValue]] = Field( + default=None, description="""Output only. Bleu metric values.""" + ) - instances: Optional[list[RougeInstanceDict]] - """Required. Repeated rouge instances.""" - metric_spec: Optional[genai_types.RougeSpec] - """Required. Spec for rouge score metric.""" +class BleuResultsDict(TypedDict, total=False): + """Result of evaluating a bleu metric.""" + + bleu_metric_values: Optional[list[genai_types.BleuMetricValue]] + """Output only. Bleu metric values.""" -RougeInputOrDict = Union[RougeInput, RougeInputDict] +BleuResultsOrDict = Union[BleuResults, BleuResultsDict] -class ContentMap(_common.BaseModel): - """Map of placeholder in metric prompt template to contents of model input.""" +class ExactMatchResults(_common.BaseModel): + """Result of evaluating an exact match metric.""" - values: Optional[dict[str, "ContentMapContents"]] = Field( - default=None, description="""Map of placeholder to contents.""" + exact_match_metric_values: Optional[list[genai_types.ExactMatchMetricValue]] = ( + Field(default=None, description="""Output only. Exact match metric values.""") ) -class ContentMapDict(TypedDict, total=False): - """Map of placeholder in metric prompt template to contents of model input.""" +class ExactMatchResultsDict(TypedDict, total=False): + """Result of evaluating an exact match metric.""" - values: Optional[dict[str, "ContentMapContents"]] - """Map of placeholder to contents.""" + exact_match_metric_values: Optional[list[genai_types.ExactMatchMetricValue]] + """Output only. Exact match metric values.""" -ContentMapOrDict = Union[ContentMap, ContentMapDict] +ExactMatchResultsOrDict = Union[ExactMatchResults, ExactMatchResultsDict] -class PointwiseMetricInstance(_common.BaseModel): - """Pointwise metric instance.""" +class RougeResults(_common.BaseModel): + """Result of evaluating a rouge metric.""" - 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.""", + rouge_metric_values: Optional[list[genai_types.RougeMetricValue]] = Field( + default=None, description="""Output only. Rouge metric values.""" ) -class PointwiseMetricInstanceDict(TypedDict, total=False): - """Pointwise metric instance.""" +class RougeResultsDict(TypedDict, total=False): + """Result of evaluating a rouge metric.""" - 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.""" + rouge_metric_values: Optional[list[genai_types.RougeMetricValue]] + """Output only. Rouge metric values.""" - 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.""" +RougeResultsOrDict = Union[RougeResults, RougeResultsDict] -PointwiseMetricInstanceOrDict = Union[ - PointwiseMetricInstance, PointwiseMetricInstanceDict -] +class RubricBasedMetricResult(_common.BaseModel): + """Result for a rubric-based metric.""" -class PointwiseMetricInput(_common.BaseModel): - """Pointwise metric input.""" - - instance: Optional[PointwiseMetricInstance] = Field( - default=None, description="""Required. Pointwise metric instance.""" + score: Optional[float] = Field( + default=None, description="""Passing rate of all the rubrics.""" ) - metric_spec: Optional[genai_types.PointwiseMetricSpec] = Field( - default=None, description="""Required. Spec for pointwise metric.""" + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( + default=None, + description="""The details of all the rubrics and their verdicts.""", ) -class PointwiseMetricInputDict(TypedDict, total=False): - """Pointwise metric input.""" +class RubricBasedMetricResultDict(TypedDict, total=False): + """Result for a rubric-based metric.""" - instance: Optional[PointwiseMetricInstanceDict] - """Required. Pointwise metric instance.""" + score: Optional[float] + """Passing rate of all the rubrics.""" - metric_spec: Optional[genai_types.PointwiseMetricSpec] - """Required. Spec for pointwise metric.""" + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] + """The details of all the rubrics and their verdicts.""" -PointwiseMetricInputOrDict = Union[PointwiseMetricInput, PointwiseMetricInputDict] +RubricBasedMetricResultOrDict = Union[ + RubricBasedMetricResult, RubricBasedMetricResultDict +] -class PairwiseMetricInstance(_common.BaseModel): - """Pairwise metric instance.""" +class CometResult(_common.BaseModel): + """Spec for Comet result - calculates the comet score for the given instance using the version specified in the spec.""" - json_instance: Optional[str] = Field( + score: Optional[float] = 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="""Output only. Comet score. Range depends on version.""", ) -class PairwiseMetricInstanceDict(TypedDict, total=False): - """Pairwise metric instance.""" +class CometResultDict(TypedDict, total=False): + """Spec for Comet result - calculates the comet score for the given instance using the version specified in the spec.""" - 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.""" + score: Optional[float] + """Output only. Comet score. Range depends on version.""" -PairwiseMetricInstanceOrDict = Union[PairwiseMetricInstance, PairwiseMetricInstanceDict] +CometResultOrDict = Union[CometResult, CometResultDict] -class PairwiseMetricInput(_common.BaseModel): - """Pairwise metric instance.""" +class MetricxResult(_common.BaseModel): + """Spec for MetricX result - calculates the MetricX score for the given instance using the version specified in the spec.""" - 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.""" + score: Optional[float] = Field( + default=None, + description="""Output only. MetricX score. Range depends on version.""", ) -class PairwiseMetricInputDict(TypedDict, total=False): - """Pairwise metric instance.""" - - instance: Optional[PairwiseMetricInstanceDict] - """Required. Pairwise metric instance.""" +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_spec: Optional[genai_types.PairwiseMetricSpec] - """Required. Spec for pairwise metric.""" + score: Optional[float] + """Output only. MetricX score. Range depends on version.""" -PairwiseMetricInputOrDict = Union[PairwiseMetricInput, PairwiseMetricInputDict] +MetricxResultOrDict = Union[MetricxResult, MetricxResultDict] -class ToolCallValidInstance(_common.BaseModel): - """Tool call valid instance.""" +class ToolCallValidMetricValue(_common.BaseModel): + """Tool call valid metric value for an instance.""" - 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.""", + score: Optional[float] = Field( + default=None, description="""Output only. Tool call valid score.""" ) -class ToolCallValidInstanceDict(TypedDict, total=False): - """Tool call valid instance.""" - - prediction: Optional[str] - """Required. Output of the evaluated model.""" +class ToolCallValidMetricValueDict(TypedDict, total=False): + """Tool call valid metric value for an instance.""" - reference: Optional[str] - """Required. Ground truth used to compare against the prediction.""" + score: Optional[float] + """Output only. Tool call valid score.""" -ToolCallValidInstanceOrDict = Union[ToolCallValidInstance, ToolCallValidInstanceDict] +ToolCallValidMetricValueOrDict = Union[ + ToolCallValidMetricValue, ToolCallValidMetricValueDict +] -class ToolCallValidSpec(_common.BaseModel): - """Spec for tool call valid metric.""" +class ToolCallValidResults(_common.BaseModel): + """Results for tool call valid metric.""" - pass + tool_call_valid_metric_values: Optional[list[ToolCallValidMetricValue]] = Field( + default=None, description="""Output only. Tool call valid metric values.""" + ) -class ToolCallValidSpecDict(TypedDict, total=False): - """Spec for tool call valid metric.""" +class ToolCallValidResultsDict(TypedDict, total=False): + """Results for tool call valid metric.""" - pass + tool_call_valid_metric_values: Optional[list[ToolCallValidMetricValueDict]] + """Output only. Tool call valid metric values.""" -ToolCallValidSpecOrDict = Union[ToolCallValidSpec, ToolCallValidSpecDict] +ToolCallValidResultsOrDict = Union[ToolCallValidResults, ToolCallValidResultsDict] -class ToolCallValidInput(_common.BaseModel): - """Tool call valid input.""" +class ToolNameMatchMetricValue(_common.BaseModel): + """Tool name match metric value for an instance.""" - 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.""" + score: Optional[float] = Field( + default=None, description="""Output only. Tool name match score.""" ) -class ToolCallValidInputDict(TypedDict, total=False): - """Tool call valid input.""" - - instances: Optional[list[ToolCallValidInstanceDict]] - """Required. Repeated tool call valid instances.""" +class ToolNameMatchMetricValueDict(TypedDict, total=False): + """Tool name match metric value for an instance.""" - metric_spec: Optional[ToolCallValidSpecDict] - """Required. Spec for tool call valid metric.""" + score: Optional[float] + """Output only. Tool name match score.""" -ToolCallValidInputOrDict = Union[ToolCallValidInput, ToolCallValidInputDict] +ToolNameMatchMetricValueOrDict = Union[ + ToolNameMatchMetricValue, ToolNameMatchMetricValueDict +] -class ToolNameMatchInstance(_common.BaseModel): - """Tool name match instance.""" +class ToolNameMatchResults(_common.BaseModel): + """Results for tool name match metric.""" - 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.""", + tool_name_match_metric_values: Optional[list[ToolNameMatchMetricValue]] = Field( + default=None, description="""Output only. Tool name match metric values.""" ) -class ToolNameMatchInstanceDict(TypedDict, total=False): - """Tool name match instance.""" - - prediction: Optional[str] - """Required. Output of the evaluated model.""" +class ToolNameMatchResultsDict(TypedDict, total=False): + """Results for tool name match metric.""" - reference: Optional[str] - """Required. Ground truth used to compare against the prediction.""" + tool_name_match_metric_values: Optional[list[ToolNameMatchMetricValueDict]] + """Output only. Tool name match metric values.""" -ToolNameMatchInstanceOrDict = Union[ToolNameMatchInstance, ToolNameMatchInstanceDict] +ToolNameMatchResultsOrDict = Union[ToolNameMatchResults, ToolNameMatchResultsDict] -class ToolNameMatchSpec(_common.BaseModel): - """Spec for tool name match metric.""" +class ToolParameterKeyMatchMetricValue(_common.BaseModel): + """Tool parameter key match metric value for an instance.""" - pass + score: Optional[float] = Field( + default=None, description="""Output only. Tool parameter key match score.""" + ) -class ToolNameMatchSpecDict(TypedDict, total=False): - """Spec for tool name match metric.""" +class ToolParameterKeyMatchMetricValueDict(TypedDict, total=False): + """Tool parameter key match metric value for an instance.""" - pass + score: Optional[float] + """Output only. Tool parameter key match score.""" -ToolNameMatchSpecOrDict = Union[ToolNameMatchSpec, ToolNameMatchSpecDict] +ToolParameterKeyMatchMetricValueOrDict = Union[ + ToolParameterKeyMatchMetricValue, ToolParameterKeyMatchMetricValueDict +] -class ToolNameMatchInput(_common.BaseModel): - """Tool name match input.""" +class ToolParameterKeyMatchResults(_common.BaseModel): + """Results for tool parameter key match metric.""" - instances: Optional[list[ToolNameMatchInstance]] = Field( - default=None, description="""Required. Repeated tool name match instances.""" - ) - metric_spec: Optional[ToolNameMatchSpec] = Field( - default=None, description="""Required. Spec for tool name match metric.""" + tool_parameter_key_match_metric_values: Optional[ + list[ToolParameterKeyMatchMetricValue] + ] = Field( + default=None, + description="""Output only. Tool parameter key match metric values.""", ) -class ToolNameMatchInputDict(TypedDict, total=False): - """Tool name match input.""" - - instances: Optional[list[ToolNameMatchInstanceDict]] - """Required. Repeated tool name match instances.""" +class ToolParameterKeyMatchResultsDict(TypedDict, total=False): + """Results for tool parameter key match metric.""" - metric_spec: Optional[ToolNameMatchSpecDict] - """Required. Spec for tool name match metric.""" + tool_parameter_key_match_metric_values: Optional[ + list[ToolParameterKeyMatchMetricValueDict] + ] + """Output only. Tool parameter key match metric values.""" -ToolNameMatchInputOrDict = Union[ToolNameMatchInput, ToolNameMatchInputDict] +ToolParameterKeyMatchResultsOrDict = Union[ + ToolParameterKeyMatchResults, ToolParameterKeyMatchResultsDict +] -class ToolParameterKeyMatchInstance(_common.BaseModel): - """Tool parameter key match instance.""" +class ToolParameterKVMatchMetricValue(_common.BaseModel): + """Tool parameter key value match metric value for an instance.""" - prediction: Optional[str] = Field( - default=None, description="""Required. Output of the evaluated model.""" - ) - reference: Optional[str] = Field( + score: Optional[float] = Field( default=None, - description="""Required. Ground truth used to compare against the prediction.""", + description="""Output only. Tool parameter key value match score.""", ) -class ToolParameterKeyMatchInstanceDict(TypedDict, total=False): - """Tool parameter key match instance.""" - - prediction: Optional[str] - """Required. Output of the evaluated model.""" +class ToolParameterKVMatchMetricValueDict(TypedDict, total=False): + """Tool parameter key value match metric value for an instance.""" - reference: Optional[str] - """Required. Ground truth used to compare against the prediction.""" + score: Optional[float] + """Output only. Tool parameter key value match score.""" -ToolParameterKeyMatchInstanceOrDict = Union[ - ToolParameterKeyMatchInstance, ToolParameterKeyMatchInstanceDict +ToolParameterKVMatchMetricValueOrDict = Union[ + ToolParameterKVMatchMetricValue, ToolParameterKVMatchMetricValueDict ] -class ToolParameterKeyMatchSpec(_common.BaseModel): - """Spec for tool parameter key match metric.""" +class ToolParameterKVMatchResults(_common.BaseModel): + """Results for tool parameter key value match metric.""" - pass + tool_parameter_kv_match_metric_values: Optional[ + list[ToolParameterKVMatchMetricValue] + ] = Field( + default=None, + description="""Output only. Tool parameter key value match metric values.""", + ) -class ToolParameterKeyMatchSpecDict(TypedDict, total=False): - """Spec for tool parameter key match metric.""" +class ToolParameterKVMatchResultsDict(TypedDict, total=False): + """Results for tool parameter key value match metric.""" - pass + tool_parameter_kv_match_metric_values: Optional[ + list[ToolParameterKVMatchMetricValueDict] + ] + """Output only. Tool parameter key value match metric values.""" -ToolParameterKeyMatchSpecOrDict = Union[ - ToolParameterKeyMatchSpec, ToolParameterKeyMatchSpecDict +ToolParameterKVMatchResultsOrDict = Union[ + ToolParameterKVMatchResults, ToolParameterKVMatchResultsDict ] -class ToolParameterKeyMatchInput(_common.BaseModel): - """Tool parameter key match input.""" +class EvaluateInstancesResponse(_common.BaseModel): + """Result of evaluating an LLM metric.""" - instances: Optional[list[ToolParameterKeyMatchInstance]] = 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="""Required. Repeated tool parameter key match instances.""", + 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.""", ) - metric_spec: Optional[ToolParameterKeyMatchSpec] = 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="""Required. Spec for tool parameter key match metric.""", + 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.""", ) -class ToolParameterKeyMatchInputDict(TypedDict, total=False): - """Tool parameter key match input.""" +class EvaluateInstancesResponseDict(TypedDict, total=False): + """Result of evaluating an LLM metric.""" - instances: Optional[list[ToolParameterKeyMatchInstanceDict]] - """Required. Repeated tool parameter key match instances.""" + rubric_based_metric_result: Optional[RubricBasedMetricResultDict] + """Result for rubric based metric.""" - metric_spec: Optional[ToolParameterKeyMatchSpecDict] - """Required. Spec for tool parameter key match 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.""" -ToolParameterKeyMatchInputOrDict = Union[ - ToolParameterKeyMatchInput, ToolParameterKeyMatchInputDict -] + comet_result: Optional[CometResultDict] + """Translation metrics. Result for Comet metric.""" + exact_match_results: Optional[ExactMatchResultsDict] + """Auto metric evaluation results. Results for exact match metric.""" -class ToolParameterKVMatchInstance(_common.BaseModel): - """Tool parameter kv match instance.""" + metricx_result: Optional[MetricxResultDict] + """Result for Metricx metric.""" - 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.""", - ) + 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 ToolParameterKVMatchInstanceDict(TypedDict, total=False): - """Tool parameter kv match instance.""" + rouge_results: Optional[RougeResultsDict] + """Results for rouge metric.""" - prediction: Optional[str] - """Required. Output of the evaluated model.""" + tool_call_valid_results: Optional[ToolCallValidResultsDict] + """Tool call metrics. Results for tool call valid metric.""" - reference: Optional[str] - """Required. Ground truth used to compare against the prediction.""" + tool_name_match_results: Optional[ToolNameMatchResultsDict] + """Results for tool name match metric.""" + tool_parameter_key_match_results: Optional[ToolParameterKeyMatchResultsDict] + """Results for tool parameter key match metric.""" -ToolParameterKVMatchInstanceOrDict = Union[ - ToolParameterKVMatchInstance, ToolParameterKVMatchInstanceDict + tool_parameter_kv_match_results: Optional[ToolParameterKVMatchResultsDict] + """Results for tool parameter key value match metric.""" + + +EvaluateInstancesResponseOrDict = Union[ + EvaluateInstancesResponse, EvaluateInstancesResponseDict ] -class ToolParameterKVMatchSpec(_common.BaseModel): - """Spec for tool parameter kv match metric.""" +class GenerateUserScenariosConfig(_common.BaseModel): - use_strict_string_match: Optional[bool] = Field( - default=None, - description="""Optional. Whether to use STRICT string match on parameter values.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class ToolParameterKVMatchSpecDict(TypedDict, total=False): - """Spec for tool parameter kv match metric.""" +class GenerateUserScenariosConfigDict(TypedDict, total=False): - use_strict_string_match: Optional[bool] - """Optional. Whether to use STRICT string match on parameter values.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ToolParameterKVMatchSpecOrDict = Union[ - ToolParameterKVMatchSpec, ToolParameterKVMatchSpecDict +GenerateUserScenariosConfigOrDict = Union[ + GenerateUserScenariosConfig, GenerateUserScenariosConfigDict ] -class ToolParameterKVMatchInput(_common.BaseModel): - """Tool parameter kv match input.""" +class _GenerateUserScenariosParameters(_common.BaseModel): + """Parameters for GenerateUserScenarios.""" - instances: Optional[list[ToolParameterKVMatchInstance]] = Field( + 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="""Required. Repeated tool parameter key value match instances.""", + description="""Opt-in flag to authorize cross-region routing for LLM models.""", ) - metric_spec: Optional[ToolParameterKVMatchSpec] = Field( + gemini_agent_config: Optional[GeminiAgentConfig] = Field( default=None, - description="""Required. Spec for tool parameter key value match metric.""", + description="""If set, the server derives the agents map and root_agent_id + from the referenced Gemini Agent server-side.""", ) -class ToolParameterKVMatchInputDict(TypedDict, total=False): - """Tool parameter kv match input.""" +class _GenerateUserScenariosParametersDict(TypedDict, total=False): + """Parameters for GenerateUserScenarios.""" - instances: Optional[list[ToolParameterKVMatchInstanceDict]] - """Required. Repeated tool parameter key value match instances.""" + location: Optional[str] + """""" - metric_spec: Optional[ToolParameterKVMatchSpecDict] - """Required. Spec for tool parameter key value match metric.""" + agents: Optional[dict[str, evals_types.AgentConfig]] + """""" + root_agent_id: Optional[str] + """""" -ToolParameterKVMatchInputOrDict = Union[ - ToolParameterKVMatchInput, ToolParameterKVMatchInputDict + user_scenario_generation_config: Optional[evals_types.UserScenarioGenerationConfig] + """""" + + config: Optional[GenerateUserScenariosConfigDict] + """""" + + allow_cross_region_model: Optional[bool] + """Opt-in flag to authorize cross-region routing for LLM models.""" + + gemini_agent_config: Optional[GeminiAgentConfigDict] + """If set, the server derives the agents map and root_agent_id + from the referenced Gemini Agent server-side.""" + + +_GenerateUserScenariosParametersOrDict = Union[ + _GenerateUserScenariosParameters, _GenerateUserScenariosParametersDict ] -class MapInstance(_common.BaseModel): - """Instance data specified as a map.""" +class GenerateUserScenariosResponse(_common.BaseModel): + """Response message for DataFoundryService.GenerateUserScenarios.""" - map_instance: Optional[dict[str, evals_types.InstanceData]] = Field( - default=None, description="""Map of instance data.""" + user_scenarios: Optional[list[evals_types.UserScenario]] = Field( + default=None, description="""""" ) -class MapInstanceDict(TypedDict, total=False): - """Instance data specified as a map.""" +class GenerateUserScenariosResponseDict(TypedDict, total=False): + """Response message for DataFoundryService.GenerateUserScenarios.""" - map_instance: Optional[dict[str, evals_types.InstanceData]] - """Map of instance data.""" + user_scenarios: Optional[list[evals_types.UserScenario]] + """""" -MapInstanceOrDict = Union[MapInstance, MapInstanceDict] +GenerateUserScenariosResponseOrDict = Union[ + GenerateUserScenariosResponse, GenerateUserScenariosResponseDict +] -class EvaluationInstance(_common.BaseModel): - """A single instance to be evaluated.""" +class GenerateLossClustersConfig(_common.BaseModel): + """Config for generating loss clusters.""" - prompt: Optional[evals_types.InstanceData] = Field( + 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.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GenerateLossClustersConfigOrDict = Union[ + GenerateLossClustersConfig, GenerateLossClustersConfigDict +] + + +class _GenerateLossClustersParameters(_common.BaseModel): + """Parameters for GenerateLossClusters.""" + + location: Optional[str] = Field( default=None, - description="""Data used to populate placeholder `prompt` in a metric prompt template.""", + description="""The resource name of the Location. Format: `projects/{project}/locations/{location}`.""", ) - response: Optional[evals_types.InstanceData] = Field( + evaluation_set: Optional[str] = Field( default=None, - description="""Data used to populate placeholder `response` in a metric prompt template.""", + description="""Reference to a persisted EvaluationSet. The service will read items from this set.""", ) - reference: Optional[evals_types.InstanceData] = Field( + inline_results: Optional[list[EvaluationResult]] = Field( default=None, - description="""Data used to populate placeholder `reference` in a metric prompt template.""", + description="""Inline evaluation results. Useful for ephemeral analysis in notebooks/SDKs where data isn't persisted.""", ) - other_data: Optional[MapInstance] = Field( + configs: Optional[list[LossAnalysisConfig]] = Field( default=None, - description="""Other data used to populate placeholders based on their key.""", + description="""Configuration for the analysis algorithm. Analysis for multiple metrics and multiple candidates could be specified.""", ) - 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.""", + config: Optional[GenerateLossClustersConfig] = Field( + default=None, description="""Config for generating loss clusters.""" ) -class EvaluationInstanceDict(TypedDict, total=False): - """A single instance to be evaluated.""" +class _GenerateLossClustersParametersDict(TypedDict, total=False): + """Parameters for GenerateLossClusters.""" - prompt: Optional[evals_types.InstanceData] - """Data used to populate placeholder `prompt` in a metric prompt template.""" + location: Optional[str] + """The resource name of the Location. Format: `projects/{project}/locations/{location}`.""" - response: Optional[evals_types.InstanceData] - """Data used to populate placeholder `response` in a metric prompt template.""" + evaluation_set: Optional[str] + """Reference to a persisted EvaluationSet. The service will read items from this set.""" - reference: Optional[evals_types.InstanceData] - """Data used to populate placeholder `reference` in a metric prompt template.""" + inline_results: Optional[list[EvaluationResultDict]] + """Inline evaluation results. Useful for ephemeral analysis in notebooks/SDKs where data isn't persisted.""" - other_data: Optional[MapInstanceDict] - """Other data used to populate placeholders based on their key.""" + configs: Optional[list[LossAnalysisConfigDict]] + """Configuration for the analysis algorithm. Analysis for multiple metrics and multiple candidates could be specified.""" - agent_data: Optional[evals_types.AgentData] - """Data used for agent evaluation.""" + config: Optional[GenerateLossClustersConfigDict] + """Config for generating loss clusters.""" - 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.""" - 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.""" +_GenerateLossClustersParametersOrDict = Union[ + _GenerateLossClustersParameters, _GenerateLossClustersParametersDict +] -EvaluationInstanceOrDict = Union[EvaluationInstance, EvaluationInstanceDict] +class GenerateLossClustersResponse(_common.BaseModel): + """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" + 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.""", + ) -class EvaluateInstancesConfig(_common.BaseModel): - """Config for evaluate instances.""" + def show(self) -> None: + """Shows the loss pattern analysis report with rich HTML visualization.""" + from .. import _evals_visualization - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + _evals_visualization.display_loss_clusters_response(self) -class EvaluateInstancesConfigDict(TypedDict, total=False): - """Config for evaluate instances.""" +class GenerateLossClustersResponseDict(TypedDict, total=False): + """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + analysis_time: Optional[str] + """The timestamp when this analysis was completed.""" + results: Optional[list[LossAnalysisResultDict]] + """The analysis results, one per config provided in the request.""" -EvaluateInstancesConfigOrDict = Union[ - EvaluateInstancesConfig, EvaluateInstancesConfigDict + +GenerateLossClustersResponseOrDict = Union[ + GenerateLossClustersResponse, GenerateLossClustersResponseDict ] -class RubricBasedMetricSpec(_common.BaseModel): - """Specification for a metric that is based on rubrics.""" +class GenerateLossClustersOperation(_common.BaseModel): + """Long-running operation for generating loss clusters.""" - metric_prompt_template: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Template for the prompt used by the judge model to evaluate against - rubrics.""", + 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}`.""", ) - judge_autorater_config: Optional[genai_types.AutoraterConfig] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Optional configuration for the judge LLM (Autorater).""", + 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.""", ) - inline_rubrics: Optional[list[evals_types.Rubric]] = Field( - default=None, description="""Use rubrics provided directly in the spec.""" + 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.""", ) - rubric_group_key: Optional[str] = Field( + error: Optional[dict[str, Any]] = 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`.""", + description="""The error result of the operation in case of failure or cancellation.""", ) - rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] = Field( + response: Optional[GenerateLossClustersResponse] = Field( default=None, - description="""Dynamically generate rubrics for evaluation using this specification.""", + description="""Response message for EvaluationAnalyticsService.GenerateLossClusters.""", ) -class RubricBasedMetricSpecDict(TypedDict, total=False): - """Specification for a metric that is based on rubrics.""" +class GenerateLossClustersOperationDict(TypedDict, total=False): + """Long-running operation for generating loss clusters.""" - metric_prompt_template: Optional[str] - """Template for the prompt used by the judge model to evaluate against - rubrics.""" + 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}`.""" - judge_autorater_config: Optional[genai_types.AutoraterConfig] - """Optional configuration for the judge LLM (Autorater).""" + 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.""" - inline_rubrics: Optional[list[evals_types.Rubric]] - """Use rubrics provided directly in the spec.""" + 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.""" - 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`.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] - """Dynamically generate rubrics for evaluation using this specification.""" + response: Optional[GenerateLossClustersResponseDict] + """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" -RubricBasedMetricSpecOrDict = Union[RubricBasedMetricSpec, RubricBasedMetricSpecDict] +GenerateLossClustersOperationOrDict = Union[ + GenerateLossClustersOperation, GenerateLossClustersOperationDict +] -class RubricEnhancedContents(_common.BaseModel): - """Rubric-enhanced contents for evaluation.""" +class RubricGenerationConfig(_common.BaseModel): + """Config for generating rubrics.""" - 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -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.""" - - 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.""" - - 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 RubricGenerationConfigDict(TypedDict, total=False): + """Config for generating rubrics.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -RubricEnhancedContentsOrDict = Union[RubricEnhancedContents, RubricEnhancedContentsDict] +RubricGenerationConfigOrDict = Union[RubricGenerationConfig, RubricGenerationConfigDict] -class RubricBasedMetricInstance(_common.BaseModel): - """Defines an instance for Rubric-based metrics. - This class allows various input formats. - """ +class _GenerateInstanceRubricsRequest(_common.BaseModel): + """Parameters for generating rubrics.""" - json_instance: Optional[str] = Field( + contents: Optional[list[genai_types.Content]] = Field( default=None, - description="""Specify evaluation fields and their string values in JSON format.""", + 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.""", ) - content_map_instance: Optional[ContentMap] = Field( + 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="""Specify evaluation fields and their content values using a ContentMap.""", + description="""Specification for how the rubrics should be generated.""", ) - rubric_enhanced_contents: Optional[RubricEnhancedContents] = Field( + metric_resource_name: Optional[str] = Field( default=None, - description="""Provide input as Gemini Content along with one or more - associated rubric groups.""", + 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 RubricBasedMetricInstanceDict(TypedDict, total=False): - """Defines an instance for Rubric-based metrics. +class _GenerateInstanceRubricsRequestDict(TypedDict, total=False): + """Parameters for generating rubrics.""" - This class allows various input formats. - """ + 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.""" - json_instance: Optional[str] - """Specify evaluation fields and their string values in JSON format.""" + 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. + """ - content_map_instance: Optional[ContentMapDict] - """Specify evaluation fields and their content values using a ContentMap.""" + rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] + """Specification for how the rubrics should be generated.""" - rubric_enhanced_contents: Optional[RubricEnhancedContentsDict] - """Provide input as Gemini Content along with one or more - associated rubric groups.""" + 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.""" + + config: Optional[RubricGenerationConfigDict] + """""" -RubricBasedMetricInstanceOrDict = Union[ - RubricBasedMetricInstance, RubricBasedMetricInstanceDict +_GenerateInstanceRubricsRequestOrDict = Union[ + _GenerateInstanceRubricsRequest, _GenerateInstanceRubricsRequestDict ] -class RubricBasedMetricInput(_common.BaseModel): - """Input for a rubric-based metrics.""" +class GenerateInstanceRubricsResponse(_common.BaseModel): + """Response for generating rubrics.""" - 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.""" + generated_rubrics: Optional[list[evals_types.Rubric]] = Field( + default=None, description="""A list of generated rubrics.""" ) -class RubricBasedMetricInputDict(TypedDict, total=False): - """Input for a rubric-based metrics.""" - - metric_spec: Optional[RubricBasedMetricSpecDict] - """Specification for the rubric-based metric.""" +class GenerateInstanceRubricsResponseDict(TypedDict, total=False): + """Response for generating rubrics.""" - instance: Optional[RubricBasedMetricInstanceDict] - """The instance to be evaluated.""" + generated_rubrics: Optional[list[evals_types.Rubric]] + """A list of generated rubrics.""" -RubricBasedMetricInputOrDict = Union[RubricBasedMetricInput, RubricBasedMetricInputDict] +GenerateInstanceRubricsResponseOrDict = Union[ + GenerateInstanceRubricsResponse, GenerateInstanceRubricsResponseDict +] -class MetricSource(_common.BaseModel): - """The metric source used for evaluation.""" +class GetEvaluationExperimentConfig(_common.BaseModel): + """Config for getting an evaluation experiment.""" - 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}""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class MetricSourceDict(TypedDict, total=False): - """The metric source used for evaluation.""" - - metric: Optional[MetricDict] - """Inline metric config.""" +class GetEvaluationExperimentConfigDict(TypedDict, total=False): + """Config for getting an evaluation experiment.""" - metric_resource_name: Optional[str] - """Resource name for registered metric. Example: - projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -MetricSourceOrDict = Union[MetricSource, MetricSourceDict] +GetEvaluationExperimentConfigOrDict = Union[ + GetEvaluationExperimentConfig, GetEvaluationExperimentConfigDict +] -class _EvaluateInstancesRequestParameters(_common.BaseModel): - """Parameters for evaluating instances.""" +class _GetEvaluationExperimentParameters(_common.BaseModel): + """Parameters for getting an evaluation experiment.""" - 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( + name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationExperimentConfig] = 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 _EvaluateInstancesRequestParametersDict(TypedDict, total=False): - """Parameters for evaluating instances.""" +class _GetEvaluationExperimentParametersDict(TypedDict, total=False): + """Parameters for getting an evaluation experiment.""" - bleu_input: Optional[BleuInputDict] + name: Optional[str] """""" - exact_match_input: Optional[ExactMatchInputDict] + config: Optional[GetEvaluationExperimentConfigDict] """""" - rouge_input: Optional[RougeInputDict] - """""" - pointwise_metric_input: Optional[PointwiseMetricInputDict] - """""" +_GetEvaluationExperimentParametersOrDict = Union[ + _GetEvaluationExperimentParameters, _GetEvaluationExperimentParametersDict +] - pairwise_metric_input: Optional[PairwiseMetricInputDict] - """""" - tool_call_valid_input: Optional[ToolCallValidInputDict] - """""" +class GetEvaluationMetricConfig(_common.BaseModel): + """Config for getting an evaluation metric.""" - tool_name_match_input: Optional[ToolNameMatchInputDict] - """""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) - tool_parameter_key_match_input: Optional[ToolParameterKeyMatchInputDict] - """""" - tool_parameter_kv_match_input: Optional[ToolParameterKVMatchInputDict] - """""" +class GetEvaluationMetricConfigDict(TypedDict, total=False): + """Config for getting an evaluation metric.""" - rubric_based_metric_input: Optional[RubricBasedMetricInputDict] - """""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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.""" +GetEvaluationMetricConfigOrDict = Union[ + GetEvaluationMetricConfig, GetEvaluationMetricConfigDict +] - instance: Optional[EvaluationInstanceDict] - """The instance to be evaluated.""" - metric_sources: Optional[list[MetricSourceDict]] - """The metrics used for evaluation.""" +class _GetEvaluationMetricParameters(_common.BaseModel): + """Parameters for getting an evaluation metric.""" - config: Optional[EvaluateInstancesConfigDict] + metric_resource_name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationMetricConfig] = Field( + default=None, description="""""" + ) + + +class _GetEvaluationMetricParametersDict(TypedDict, total=False): + """Parameters for getting an evaluation metric.""" + + metric_resource_name: Optional[str] """""" + config: Optional[GetEvaluationMetricConfigDict] + """""" -_EvaluateInstancesRequestParametersOrDict = Union[ - _EvaluateInstancesRequestParameters, _EvaluateInstancesRequestParametersDict + +_GetEvaluationMetricParametersOrDict = Union[ + _GetEvaluationMetricParameters, _GetEvaluationMetricParametersDict ] -class MetricResult(_common.BaseModel): - """Result for a single metric on a single instance.""" +class GetEvaluationRunConfig(_common.BaseModel): + """Config for get evaluation run.""" - 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.""", - ) - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( - default=None, - 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.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class MetricResultDict(TypedDict, total=False): - """Result for a single metric on a single instance.""" - - score: Optional[float] - """The score for the metric. Please refer to each metric's documentation for the meaning of the score.""" - - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] - """For rubric-based metrics, the verdicts for each rubric.""" +class GetEvaluationRunConfigDict(TypedDict, total=False): + """Config for get evaluation run.""" - explanation: Optional[str] - """The explanation for the metric result.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - error: Optional[genai_types.GoogleRpcStatus] - """The error status for the metric result.""" +GetEvaluationRunConfigOrDict = Union[GetEvaluationRunConfig, GetEvaluationRunConfigDict] -MetricResultOrDict = Union[MetricResult, MetricResultDict] +class _GetEvaluationRunParameters(_common.BaseModel): + """Represents a job that runs evaluation.""" -class BleuResults(_common.BaseModel): - """Result of evaluating a bleu metric.""" + name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationRunConfig] = Field(default=None, description="""""") - bleu_metric_values: Optional[list[genai_types.BleuMetricValue]] = Field( - default=None, description="""Output only. Bleu metric values.""" - ) +class _GetEvaluationRunParametersDict(TypedDict, total=False): + """Represents a job that runs evaluation.""" -class BleuResultsDict(TypedDict, total=False): - """Result of evaluating a bleu metric.""" + name: Optional[str] + """""" - bleu_metric_values: Optional[list[genai_types.BleuMetricValue]] - """Output only. Bleu metric values.""" + config: Optional[GetEvaluationRunConfigDict] + """""" -BleuResultsOrDict = Union[BleuResults, BleuResultsDict] +_GetEvaluationRunParametersOrDict = Union[ + _GetEvaluationRunParameters, _GetEvaluationRunParametersDict +] -class ExactMatchResults(_common.BaseModel): - """Result of evaluating an exact match metric.""" +class GetEvaluationSetConfig(_common.BaseModel): + """Config for get evaluation set.""" - exact_match_metric_values: Optional[list[genai_types.ExactMatchMetricValue]] = ( - Field(default=None, description="""Output only. Exact match metric values.""") + 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 GetEvaluationSetConfigDict(TypedDict, total=False): + """Config for get evaluation set.""" - 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] +GetEvaluationSetConfigOrDict = Union[GetEvaluationSetConfig, GetEvaluationSetConfigDict] -class RougeResults(_common.BaseModel): - """Result of evaluating a rouge metric.""" +class _GetEvaluationSetParameters(_common.BaseModel): + """Represents a job that gets an evaluation set.""" - rouge_metric_values: Optional[list[genai_types.RougeMetricValue]] = Field( - default=None, description="""Output only. Rouge metric values.""" - ) + name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationSetConfig] = Field(default=None, description="""""") -class RougeResultsDict(TypedDict, total=False): - """Result of evaluating a rouge metric.""" +class _GetEvaluationSetParametersDict(TypedDict, total=False): + """Represents a job that gets an evaluation set.""" - rouge_metric_values: Optional[list[genai_types.RougeMetricValue]] - """Output only. Rouge metric values.""" + name: Optional[str] + """""" + config: Optional[GetEvaluationSetConfigDict] + """""" -RougeResultsOrDict = Union[RougeResults, RougeResultsDict] +_GetEvaluationSetParametersOrDict = Union[ + _GetEvaluationSetParameters, _GetEvaluationSetParametersDict +] -class RubricBasedMetricResult(_common.BaseModel): - """Result for a rubric-based metric.""" - 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.""", - ) +class GetEvaluationItemConfig(_common.BaseModel): + """Config for get evaluation item.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -class RubricBasedMetricResultDict(TypedDict, total=False): - """Result for a rubric-based metric.""" - score: Optional[float] - """Passing rate of all the rubrics.""" +class GetEvaluationItemConfigDict(TypedDict, total=False): + """Config for get evaluation item.""" - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] - """The details of all the rubrics and their verdicts.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -RubricBasedMetricResultOrDict = Union[ - RubricBasedMetricResult, RubricBasedMetricResultDict +GetEvaluationItemConfigOrDict = Union[ + GetEvaluationItemConfig, GetEvaluationItemConfigDict ] -class CometResult(_common.BaseModel): - """Spec for Comet result - calculates the comet score for the given instance using the version specified in the spec.""" +class _GetEvaluationItemParameters(_common.BaseModel): + """Represents a job that gets an evaluation item.""" - score: Optional[float] = Field( - default=None, - description="""Output only. Comet score. Range depends on version.""", - ) + name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationItemConfig] = 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 _GetEvaluationItemParametersDict(TypedDict, total=False): + """Represents a job that gets an evaluation item.""" - score: Optional[float] - """Output only. Comet score. Range depends on version.""" + name: Optional[str] + """""" + config: Optional[GetEvaluationItemConfigDict] + """""" -CometResultOrDict = Union[CometResult, CometResultDict] +_GetEvaluationItemParametersOrDict = Union[ + _GetEvaluationItemParameters, _GetEvaluationItemParametersDict +] -class MetricxResult(_common.BaseModel): - """Spec for MetricX result - calculates the MetricX score for the given instance using the version specified in the spec.""" - score: Optional[float] = Field( - default=None, - description="""Output only. MetricX score. Range depends on version.""", +class ImportSchemaConfig(_common.BaseModel): + """Configuration for the input data format.""" + + data_format: Optional[ImportDataFormat] = Field( + default=None, description="""The format of the input data.""" + ) + data_format_version: Optional[str] = Field( + default=None, description="""Version of the data format.""" ) -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 ImportSchemaConfigDict(TypedDict, total=False): + """Configuration for the input data format.""" - score: Optional[float] - """Output only. MetricX score. Range depends on version.""" + data_format: Optional[ImportDataFormat] + """The format of the input data.""" + data_format_version: Optional[str] + """Version of the data format.""" -MetricxResultOrDict = Union[MetricxResult, MetricxResultDict] +ImportSchemaConfigOrDict = Union[ImportSchemaConfig, ImportSchemaConfigDict] -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.""" +class EvaluationSetGcsSource(_common.BaseModel): + """Source for loading data from Cloud Storage.""" + + gcs_uri: Optional[str] = Field( + default=None, description="""The Cloud Storage location of the input data.""" + ) + import_schema_config: Optional[ImportSchemaConfig] = Field( + default=None, description="""Schema configuration for the input data.""" ) -class ToolCallValidMetricValueDict(TypedDict, total=False): - """Tool call valid metric value for an instance.""" +class EvaluationSetGcsSourceDict(TypedDict, total=False): + """Source for loading data from Cloud Storage.""" - score: Optional[float] - """Output only. Tool call valid score.""" + gcs_uri: Optional[str] + """The Cloud Storage location of the input data.""" + + import_schema_config: Optional[ImportSchemaConfigDict] + """Schema configuration for the input data.""" -ToolCallValidMetricValueOrDict = Union[ - ToolCallValidMetricValue, ToolCallValidMetricValueDict -] +EvaluationSetGcsSourceOrDict = Union[EvaluationSetGcsSource, EvaluationSetGcsSourceDict] -class ToolCallValidResults(_common.BaseModel): - """Results for tool call valid metric.""" +class EvaluationSetInlineSource(_common.BaseModel): + """Wrapper for inline data.""" - tool_call_valid_metric_values: Optional[list[ToolCallValidMetricValue]] = Field( - default=None, description="""Output only. Tool call valid metric values.""" + content: Optional[bytes] = Field( + default=None, description="""The content of the inline data.""" + ) + import_schema_config: Optional[ImportSchemaConfig] = Field( + default=None, description="""Schema configuration for the inline data.""" ) -class ToolCallValidResultsDict(TypedDict, total=False): - """Results for tool call valid metric.""" +class EvaluationSetInlineSourceDict(TypedDict, total=False): + """Wrapper for inline data.""" - tool_call_valid_metric_values: Optional[list[ToolCallValidMetricValueDict]] - """Output only. Tool call valid metric values.""" + content: Optional[bytes] + """The content of the inline data.""" + import_schema_config: Optional[ImportSchemaConfigDict] + """Schema configuration for the inline data.""" -ToolCallValidResultsOrDict = Union[ToolCallValidResults, ToolCallValidResultsDict] +EvaluationSetInlineSourceOrDict = Union[ + EvaluationSetInlineSource, EvaluationSetInlineSourceDict +] -class ToolNameMatchMetricValue(_common.BaseModel): - """Tool name match metric value for an instance.""" - score: Optional[float] = Field( - default=None, description="""Output only. Tool name match score.""" +class EvaluationSetCloudTraceSource(_common.BaseModel): + """Source for loading traces directly from Cloud Trace.""" + + project_id: Optional[str] = Field( + default=None, description="""Project ID for the Cloud Trace.""" + ) + trace_ids: Optional[list[str]] = Field( + default=None, description="""Trace IDs to import.""" + ) + session_ids: Optional[list[str]] = Field( + default=None, + description="""Session IDs to import traces for. If both trace_ids and + session_ids are specified, the union of the two will be imported.""", ) -class ToolNameMatchMetricValueDict(TypedDict, total=False): - """Tool name match metric value for an instance.""" +class EvaluationSetCloudTraceSourceDict(TypedDict, total=False): + """Source for loading traces directly from Cloud Trace.""" - score: Optional[float] - """Output only. Tool name match score.""" + project_id: Optional[str] + """Project ID for the Cloud Trace.""" + trace_ids: Optional[list[str]] + """Trace IDs to import.""" -ToolNameMatchMetricValueOrDict = Union[ - ToolNameMatchMetricValue, ToolNameMatchMetricValueDict + session_ids: Optional[list[str]] + """Session IDs to import traces for. If both trace_ids and + session_ids are specified, the union of the two will be imported.""" + + +EvaluationSetCloudTraceSourceOrDict = Union[ + EvaluationSetCloudTraceSource, EvaluationSetCloudTraceSourceDict ] -class ToolNameMatchResults(_common.BaseModel): - """Results for tool name match metric.""" +class ImportEvaluationSetConfig(_common.BaseModel): + """Config for importing an evaluation set.""" - tool_name_match_metric_values: Optional[list[ToolNameMatchMetricValue]] = Field( - default=None, description="""Output only. Tool name match metric values.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class ToolNameMatchResultsDict(TypedDict, total=False): - """Results for tool name match metric.""" +class ImportEvaluationSetConfigDict(TypedDict, total=False): + """Config for importing an evaluation set.""" - tool_name_match_metric_values: Optional[list[ToolNameMatchMetricValueDict]] - """Output only. Tool name match metric values.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ToolNameMatchResultsOrDict = Union[ToolNameMatchResults, ToolNameMatchResultsDict] +ImportEvaluationSetConfigOrDict = Union[ + ImportEvaluationSetConfig, ImportEvaluationSetConfigDict +] -class ToolParameterKeyMatchMetricValue(_common.BaseModel): - """Tool parameter key match metric value for an instance.""" +class _ImportEvaluationSetParameters(_common.BaseModel): + """Parameters for importing an evaluation set.""" - score: Optional[float] = Field( - default=None, description="""Output only. Tool parameter key match score.""" + evaluation_set: Optional[EvaluationSet] = Field( + default=None, + description="""The EvaluationSet to create. Used to specify 'display_name' and + 'metadata'. The 'evaluation_items' field is ignored and populated by the + import process.""", + ) + gcs_destination: Optional[genai_types.GcsDestination] = Field( + default=None, + description="""The Cloud Storage location where the resulting EvaluationItem + payloads will be stored.""", + ) + gcs_source: Optional[EvaluationSetGcsSource] = Field( + default=None, description="""Google Cloud Storage location.""" + ) + inline_source: Optional[EvaluationSetInlineSource] = Field( + default=None, description="""Inline source for small payloads (< 4MB).""" + ) + cloud_trace_source: Optional[EvaluationSetCloudTraceSource] = Field( + default=None, + description="""Source for loading data directly from Cloud Trace.""", + ) + config: Optional[ImportEvaluationSetConfig] = Field( + default=None, description="""""" ) -class ToolParameterKeyMatchMetricValueDict(TypedDict, total=False): - """Tool parameter key match metric value for an instance.""" +class _ImportEvaluationSetParametersDict(TypedDict, total=False): + """Parameters for importing an evaluation set.""" - score: Optional[float] - """Output only. Tool parameter key match score.""" + evaluation_set: Optional[EvaluationSetDict] + """The EvaluationSet to create. Used to specify 'display_name' and + 'metadata'. The 'evaluation_items' field is ignored and populated by the + import process.""" + gcs_destination: Optional[genai_types.GcsDestination] + """The Cloud Storage location where the resulting EvaluationItem + payloads will be stored.""" -ToolParameterKeyMatchMetricValueOrDict = Union[ - ToolParameterKeyMatchMetricValue, ToolParameterKeyMatchMetricValueDict -] + gcs_source: Optional[EvaluationSetGcsSourceDict] + """Google Cloud Storage location.""" + inline_source: Optional[EvaluationSetInlineSourceDict] + """Inline source for small payloads (< 4MB).""" -class ToolParameterKeyMatchResults(_common.BaseModel): - """Results for tool parameter key match metric.""" + cloud_trace_source: Optional[EvaluationSetCloudTraceSourceDict] + """Source for loading data directly from Cloud Trace.""" - tool_parameter_key_match_metric_values: Optional[ - list[ToolParameterKeyMatchMetricValue] - ] = Field( - default=None, - description="""Output only. Tool parameter key match metric values.""", - ) + config: Optional[ImportEvaluationSetConfigDict] + """""" -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.""" - - -ToolParameterKeyMatchResultsOrDict = Union[ - ToolParameterKeyMatchResults, ToolParameterKeyMatchResultsDict +_ImportEvaluationSetParametersOrDict = Union[ + _ImportEvaluationSetParameters, _ImportEvaluationSetParametersDict ] -class ToolParameterKVMatchMetricValue(_common.BaseModel): - """Tool parameter key value match metric value for an instance.""" +class ImportEvaluationSetOperation(_common.BaseModel): + """Operation for importing an evaluation set.""" - score: Optional[float] = Field( + name: Optional[str] = Field( default=None, - description="""Output only. Tool parameter key value match score.""", + 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.""", ) -class ToolParameterKVMatchMetricValueDict(TypedDict, total=False): - """Tool parameter key value match metric value for an instance.""" - - score: Optional[float] - """Output only. Tool parameter key value match score.""" - - -ToolParameterKVMatchMetricValueOrDict = Union[ - ToolParameterKVMatchMetricValue, ToolParameterKVMatchMetricValueDict -] - - -class ToolParameterKVMatchResults(_common.BaseModel): - """Results for tool parameter key value match metric.""" +class ImportEvaluationSetOperationDict(TypedDict, total=False): + """Operation for importing an evaluation set.""" - tool_parameter_kv_match_metric_values: Optional[ - list[ToolParameterKVMatchMetricValue] - ] = Field( - default=None, - description="""Output only. Tool parameter key value 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.""" -class ToolParameterKVMatchResultsDict(TypedDict, total=False): - """Results for tool parameter key value match metric.""" + 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.""" - tool_parameter_kv_match_metric_values: Optional[ - list[ToolParameterKVMatchMetricValueDict] - ] - """Output only. Tool parameter key value match metric values.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -ToolParameterKVMatchResultsOrDict = Union[ - ToolParameterKVMatchResults, ToolParameterKVMatchResultsDict +ImportEvaluationSetOperationOrDict = Union[ + ImportEvaluationSetOperation, ImportEvaluationSetOperationDict ] -class EvaluateInstancesResponse(_common.BaseModel): - """Result of evaluating an LLM metric.""" +class ListEvaluationExperimentsConfig(_common.BaseModel): + """Config for listing evaluation experiments.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - tool_call_valid_results: Optional[ToolCallValidResults] = Field( + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = 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.""" + 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 `_.""", ) - tool_parameter_kv_match_results: Optional[ToolParameterKVMatchResults] = Field( + order_by: Optional[str] = Field( default=None, - description="""Results for tool parameter key value match metric.""", + 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 EvaluateInstancesResponseDict(TypedDict, total=False): - """Result of evaluating an LLM metric.""" +class ListEvaluationExperimentsConfigDict(TypedDict, total=False): + """Config for listing evaluation experiments.""" - rubric_based_metric_result: Optional[RubricBasedMetricResultDict] - """Result for rubric based metric.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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.""" + page_size: Optional[int] + """""" - bleu_results: Optional[BleuResultsDict] - """Results for bleu metric.""" + page_token: Optional[str] + """""" - comet_result: Optional[CometResultDict] - """Translation metrics. Result for Comet metric.""" + 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 `_.""" - exact_match_results: Optional[ExactMatchResultsDict] - """Auto metric evaluation results. Results for exact match metric.""" + 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"``.""" - metricx_result: Optional[MetricxResultDict] - """Result for Metricx metric.""" - pairwise_metric_result: Optional[genai_types.PairwiseMetricResult] - """Result for pairwise metric.""" +ListEvaluationExperimentsConfigOrDict = Union[ + ListEvaluationExperimentsConfig, ListEvaluationExperimentsConfigDict +] - pointwise_metric_result: Optional[genai_types.PointwiseMetricResult] - """Generic metrics. Result for pointwise metric.""" - rouge_results: Optional[RougeResultsDict] - """Results for rouge metric.""" +class _ListEvaluationExperimentsParameters(_common.BaseModel): + """Parameters for listing evaluation experiments.""" - tool_call_valid_results: Optional[ToolCallValidResultsDict] - """Tool call metrics. Results for tool call valid metric.""" + config: Optional[ListEvaluationExperimentsConfig] = Field( + default=None, description="""""" + ) - tool_name_match_results: Optional[ToolNameMatchResultsDict] - """Results for tool name match metric.""" - tool_parameter_key_match_results: Optional[ToolParameterKeyMatchResultsDict] - """Results for tool parameter key match metric.""" +class _ListEvaluationExperimentsParametersDict(TypedDict, total=False): + """Parameters for listing evaluation experiments.""" - tool_parameter_kv_match_results: Optional[ToolParameterKVMatchResultsDict] - """Results for tool parameter key value match metric.""" + config: Optional[ListEvaluationExperimentsConfigDict] + """""" -EvaluateInstancesResponseOrDict = Union[ - EvaluateInstancesResponse, EvaluateInstancesResponseDict +_ListEvaluationExperimentsParametersOrDict = Union[ + _ListEvaluationExperimentsParameters, _ListEvaluationExperimentsParametersDict ] -class GenerateUserScenariosConfig(_common.BaseModel): +class ListEvaluationExperimentsResponse(_common.BaseModel): + """Response for listing evaluation experiments.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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="""List of evaluation experiments. + """, ) -class GenerateUserScenariosConfigDict(TypedDict, total=False): +class ListEvaluationExperimentsResponseDict(TypedDict, total=False): + """Response for listing evaluation experiments.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + next_page_token: Optional[str] + """""" -GenerateUserScenariosConfigOrDict = Union[ - GenerateUserScenariosConfig, GenerateUserScenariosConfigDict + evaluation_experiments: Optional[list[EvaluationExperimentDict]] + """List of evaluation experiments. + """ + + +ListEvaluationExperimentsResponseOrDict = Union[ + ListEvaluationExperimentsResponse, ListEvaluationExperimentsResponseDict ] -class _GenerateUserScenariosParameters(_common.BaseModel): - """Parameters for GenerateUserScenarios.""" +class ListEvaluationMetricsConfig(_common.BaseModel): + """Config for listing evaluation metrics.""" - 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="""""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - allow_cross_region_model: Optional[bool] = Field( + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( default=None, - description="""Opt-in flag to authorize cross-region routing for LLM models.""", + 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 `_.""", ) - gemini_agent_config: Optional[GeminiAgentConfig] = Field( + order_by: Optional[str] = Field( default=None, - description="""If set, the server derives the agents map and root_agent_id - from the referenced Gemini Agent server-side.""", + 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 _GenerateUserScenariosParametersDict(TypedDict, total=False): - """Parameters for GenerateUserScenarios.""" +class ListEvaluationMetricsConfigDict(TypedDict, total=False): + """Config for listing evaluation metrics.""" - location: Optional[str] - """""" - - agents: Optional[dict[str, evals_types.AgentConfig]] - """""" - - root_agent_id: Optional[str] - """""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - user_scenario_generation_config: Optional[evals_types.UserScenarioGenerationConfig] + page_size: Optional[int] """""" - config: Optional[GenerateUserScenariosConfigDict] + page_token: Optional[str] """""" - allow_cross_region_model: Optional[bool] - """Opt-in flag to authorize cross-region routing for LLM models.""" + 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 `_.""" - gemini_agent_config: Optional[GeminiAgentConfigDict] - """If set, the server derives the agents map and root_agent_id - from the referenced Gemini Agent server-side.""" + 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"``.""" -_GenerateUserScenariosParametersOrDict = Union[ - _GenerateUserScenariosParameters, _GenerateUserScenariosParametersDict +ListEvaluationMetricsConfigOrDict = Union[ + ListEvaluationMetricsConfig, ListEvaluationMetricsConfigDict ] -class GenerateUserScenariosResponse(_common.BaseModel): - """Response message for DataFoundryService.GenerateUserScenarios.""" +class _ListEvaluationMetricsParameters(_common.BaseModel): + """Parameters for listing evaluation metrics.""" - user_scenarios: Optional[list[evals_types.UserScenario]] = Field( + config: Optional[ListEvaluationMetricsConfig] = Field( default=None, description="""""" ) -class GenerateUserScenariosResponseDict(TypedDict, total=False): - """Response message for DataFoundryService.GenerateUserScenarios.""" +class _ListEvaluationMetricsParametersDict(TypedDict, total=False): + """Parameters for listing evaluation metrics.""" - user_scenarios: Optional[list[evals_types.UserScenario]] + config: Optional[ListEvaluationMetricsConfigDict] """""" -GenerateUserScenariosResponseOrDict = Union[ - GenerateUserScenariosResponse, GenerateUserScenariosResponseDict +_ListEvaluationMetricsParametersOrDict = Union[ + _ListEvaluationMetricsParameters, _ListEvaluationMetricsParametersDict ] -class GenerateLossClustersConfig(_common.BaseModel): - """Config for generating loss clusters.""" +class ListEvaluationMetricsResponse(_common.BaseModel): + """Response for listing evaluation metrics.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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 GenerateLossClustersConfigDict(TypedDict, total=False): - """Config for generating loss clusters.""" +class ListEvaluationMetricsResponseDict(TypedDict, total=False): + """Response for listing evaluation metrics.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + next_page_token: Optional[str] + """""" -GenerateLossClustersConfigOrDict = Union[ - GenerateLossClustersConfig, GenerateLossClustersConfigDict + evaluation_metrics: Optional[list[EvaluationMetricDict]] + """List of evaluation metrics. + """ + + +ListEvaluationMetricsResponseOrDict = Union[ + ListEvaluationMetricsResponse, ListEvaluationMetricsResponseDict ] -class _GenerateLossClustersParameters(_common.BaseModel): - """Parameters for GenerateLossClusters.""" +class ListEvaluationSetsConfig(_common.BaseModel): + """Config for listing evaluation sets.""" - 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - inline_results: Optional[list[EvaluationResult]] = Field( + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( default=None, - description="""Inline evaluation results. Useful for ephemeral analysis in notebooks/SDKs where data isn't persisted.""", + 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 `_.""", ) - configs: Optional[list[LossAnalysisConfig]] = Field( + order_by: Optional[str] = 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.""" + 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 _GenerateLossClustersParametersDict(TypedDict, total=False): - """Parameters for GenerateLossClusters.""" +class ListEvaluationSetsConfigDict(TypedDict, total=False): + """Config for listing evaluation sets.""" - location: Optional[str] - """The resource name of the Location. Format: `projects/{project}/locations/{location}`.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - evaluation_set: Optional[str] - """Reference to a persisted EvaluationSet. The service will read items from this set.""" + page_size: Optional[int] + """""" - inline_results: Optional[list[EvaluationResultDict]] - """Inline evaluation results. Useful for ephemeral analysis in notebooks/SDKs where data isn't persisted.""" + page_token: Optional[str] + """""" - configs: Optional[list[LossAnalysisConfigDict]] - """Configuration for the analysis algorithm. Analysis for multiple metrics and multiple candidates could be specified.""" + 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 `_.""" - config: Optional[GenerateLossClustersConfigDict] - """Config for generating loss clusters.""" + 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"``.""" -_GenerateLossClustersParametersOrDict = Union[ - _GenerateLossClustersParameters, _GenerateLossClustersParametersDict +ListEvaluationSetsConfigOrDict = Union[ + ListEvaluationSetsConfig, ListEvaluationSetsConfigDict ] -class GenerateLossClustersResponse(_common.BaseModel): - """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" - - 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 - - _evals_visualization.display_loss_clusters_response(self) +class _ListEvaluationSetsParameters(_common.BaseModel): + """Parameters for listing evaluation sets.""" + config: Optional[ListEvaluationSetsConfig] = Field(default=None, description="""""") -class GenerateLossClustersResponseDict(TypedDict, total=False): - """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" - analysis_time: Optional[str] - """The timestamp when this analysis was completed.""" +class _ListEvaluationSetsParametersDict(TypedDict, total=False): + """Parameters for listing evaluation sets.""" - results: Optional[list[LossAnalysisResultDict]] - """The analysis results, one per config provided in the request.""" + config: Optional[ListEvaluationSetsConfigDict] + """""" -GenerateLossClustersResponseOrDict = Union[ - GenerateLossClustersResponse, GenerateLossClustersResponseDict +_ListEvaluationSetsParametersOrDict = Union[ + _ListEvaluationSetsParameters, _ListEvaluationSetsParametersDict ] -class GenerateLossClustersOperation(_common.BaseModel): - """Long-running operation for generating loss clusters.""" +class ListEvaluationSetsResponse(_common.BaseModel): + """Response for listing evaluation sets.""" - 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.""", + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - response: Optional[GenerateLossClustersResponse] = Field( + next_page_token: Optional[str] = Field(default=None, description="""""") + evaluation_sets: Optional[list[EvaluationSet]] = Field( default=None, - description="""Response message for EvaluationAnalyticsService.GenerateLossClusters.""", + description="""List of evaluation sets. + """, ) -class GenerateLossClustersOperationDict(TypedDict, total=False): - """Long-running operation for generating loss clusters.""" - - 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 ListEvaluationSetsResponseDict(TypedDict, total=False): + """Response for listing evaluation sets.""" - 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.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + next_page_token: Optional[str] + """""" - response: Optional[GenerateLossClustersResponseDict] - """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" + evaluation_sets: Optional[list[EvaluationSetDict]] + """List of evaluation sets. + """ -GenerateLossClustersOperationOrDict = Union[ - GenerateLossClustersOperation, GenerateLossClustersOperationDict +ListEvaluationSetsResponseOrDict = Union[ + ListEvaluationSetsResponse, ListEvaluationSetsResponseDict ] -class RubricGenerationConfig(_common.BaseModel): - """Config for generating rubrics.""" +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 RubricGenerationConfigDict(TypedDict, total=False): - """Config for generating rubrics.""" +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.""" -RubricGenerationConfigOrDict = Union[RubricGenerationConfig, RubricGenerationConfigDict] - - -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="""""") - - -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. - """ + display_name: Optional[str] + """The display name of the evaluation experiment.""" - rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] - """Specification for how the rubrics should be generated.""" + labels: Optional[dict[str, str]] + """Labels for the evaluation experiment.""" - 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.""" + merge_strategy: Optional[EvaluationExperimentMergeStrategy] + """Merge strategy for the evaluation experiment.""" - config: Optional[RubricGenerationConfigDict] - """""" + metadata: Optional[dict[str, Any]] + """Metadata about the evaluation experiment.""" -_GenerateInstanceRubricsRequestOrDict = Union[ - _GenerateInstanceRubricsRequest, _GenerateInstanceRubricsRequestDict +UpdateEvaluationExperimentConfigOrDict = Union[ + UpdateEvaluationExperimentConfig, UpdateEvaluationExperimentConfigDict ] -class GenerateInstanceRubricsResponse(_common.BaseModel): - """Response for generating rubrics.""" +class _UpdateEvaluationExperimentParameters(_common.BaseModel): + """Parameters for updating an evaluation experiment.""" - generated_rubrics: Optional[list[evals_types.Rubric]] = Field( - default=None, description="""A list of generated rubrics.""" + name: Optional[str] = Field( + default=None, description="""The resource name of the EvaluationExperiment.""" + ) + config: Optional[UpdateEvaluationExperimentConfig] = Field( + default=None, description="""""" ) -class GenerateInstanceRubricsResponseDict(TypedDict, total=False): - """Response for generating rubrics.""" +class _UpdateEvaluationExperimentParametersDict(TypedDict, total=False): + """Parameters for updating an evaluation experiment.""" - generated_rubrics: Optional[list[evals_types.Rubric]] - """A list of generated rubrics.""" + name: Optional[str] + """The resource name of the EvaluationExperiment.""" + config: Optional[UpdateEvaluationExperimentConfigDict] + """""" -GenerateInstanceRubricsResponseOrDict = Union[ - GenerateInstanceRubricsResponse, GenerateInstanceRubricsResponseDict + +_UpdateEvaluationExperimentParametersOrDict = Union[ + _UpdateEvaluationExperimentParameters, _UpdateEvaluationExperimentParametersDict ] -class GetEvaluationExperimentConfig(_common.BaseModel): - """Config for getting an evaluation experiment.""" +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 GetEvaluationExperimentConfigDict(TypedDict, total=False): - """Config for getting an evaluation experiment.""" +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.""" -GetEvaluationExperimentConfigOrDict = Union[ - GetEvaluationExperimentConfig, GetEvaluationExperimentConfigDict -] + 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 _GetEvaluationExperimentParameters(_common.BaseModel): - """Parameters for getting an evaluation experiment.""" +OptimizeConfigOrDict = Union[OptimizeConfig, OptimizeConfigDict] - name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationExperimentConfig] = Field( - default=None, description="""""" - ) +class _OptimizeRequestParameters(_common.BaseModel): + """Request for the optimize_prompt method.""" -class _GetEvaluationExperimentParametersDict(TypedDict, total=False): - """Parameters for getting an evaluation experiment.""" + content: Optional[genai_types.Content] = Field(default=None, description="""""") + config: Optional[OptimizeConfig] = Field(default=None, description="""""") - name: Optional[str] + +class _OptimizeRequestParametersDict(TypedDict, total=False): + """Request for the optimize_prompt method.""" + + content: Optional[genai_types.Content] """""" - config: Optional[GetEvaluationExperimentConfigDict] + config: Optional[OptimizeConfigDict] """""" -_GetEvaluationExperimentParametersOrDict = Union[ - _GetEvaluationExperimentParameters, _GetEvaluationExperimentParametersDict +_OptimizeRequestParametersOrDict = Union[ + _OptimizeRequestParameters, _OptimizeRequestParametersDict ] -class GetEvaluationMetricConfig(_common.BaseModel): - """Config for getting an evaluation metric.""" +class OptimizeResponseEndpoint(_common.BaseModel): + """Response for the optimize_prompt method.""" - 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="""""") -class GetEvaluationMetricConfigDict(TypedDict, total=False): - """Config for getting an evaluation metric.""" +class OptimizeResponseEndpointDict(TypedDict, total=False): + """Response for the optimize_prompt method.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + content: Optional[genai_types.Content] + """""" -GetEvaluationMetricConfigOrDict = Union[ - GetEvaluationMetricConfig, GetEvaluationMetricConfigDict +OptimizeResponseEndpointOrDict = Union[ + OptimizeResponseEndpoint, OptimizeResponseEndpointDict ] -class _GetEvaluationMetricParameters(_common.BaseModel): - """Parameters for getting an evaluation metric.""" +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.""" - metric_resource_name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationMetricConfig] = Field( - default=None, description="""""" + 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 _GetEvaluationMetricParametersDict(TypedDict, total=False): - """Parameters for getting an evaluation metric.""" +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.""" - metric_resource_name: Optional[str] - """""" + 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.""" - config: Optional[GetEvaluationMetricConfigDict] - """""" + 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.""" -_GetEvaluationMetricParametersOrDict = Union[ - _GetEvaluationMetricParameters, _GetEvaluationMetricParametersDict -] +DnsPeeringConfigOrDict = Union[DnsPeeringConfig, DnsPeeringConfigDict] -class GetEvaluationRunConfig(_common.BaseModel): - """Config for get evaluation run.""" +class PscInterfaceConfig(_common.BaseModel): + """Configuration for PSC-I.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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.""", ) -class GetEvaluationRunConfigDict(TypedDict, total=False): - """Config for get evaluation run.""" +class PscInterfaceConfigDict(TypedDict, total=False): + """Configuration for PSC-I.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + 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.""" -GetEvaluationRunConfigOrDict = Union[GetEvaluationRunConfig, GetEvaluationRunConfigDict] +PscInterfaceConfigOrDict = Union[PscInterfaceConfig, PscInterfaceConfigDict] -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 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.""", + ) + max_wait_duration: 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.""", + ) + restart_job_on_worker_restart: Optional[bool] = 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.""", + ) + 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 _GetEvaluationRunParametersDict(TypedDict, total=False): - """Represents a job that runs evaluation.""" - name: Optional[str] - """""" +class SchedulingDict(TypedDict, total=False): + """All parameters related to queuing and scheduling of custom jobs.""" - config: Optional[GetEvaluationRunConfigDict] - """""" - - -_GetEvaluationRunParametersOrDict = Union[ - _GetEvaluationRunParameters, _GetEvaluationRunParametersDict -] - - -class GetEvaluationSetConfig(_common.BaseModel): - """Config for get evaluation set.""" + 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.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + 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.""" + 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.""" -class GetEvaluationSetConfigDict(TypedDict, total=False): - """Config for get evaluation set.""" + strategy: Optional[Strategy] + """Optional. This determines which type of scheduling strategy to use.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + timeout: Optional[str] + """Optional. The maximum job running time. The default is 7 days.""" -GetEvaluationSetConfigOrDict = Union[GetEvaluationSetConfig, GetEvaluationSetConfigDict] +SchedulingOrDict = Union[Scheduling, SchedulingDict] -class _GetEvaluationSetParameters(_common.BaseModel): - """Represents a job that gets an evaluation set.""" +class EnvVar(_common.BaseModel): + """Represents an environment variable present in a Container or Python Module.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationSetConfig] = 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 _GetEvaluationSetParametersDict(TypedDict, total=False): - """Represents a job that gets an evaluation set.""" +class EnvVarDict(TypedDict, total=False): + """Represents an environment variable present in a Container or Python Module.""" name: Optional[str] - """""" + """Required. Name of the environment variable. Must be a valid C identifier.""" - config: Optional[GetEvaluationSetConfigDict] - """""" + 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.""" -_GetEvaluationSetParametersOrDict = Union[ - _GetEvaluationSetParameters, _GetEvaluationSetParametersDict -] +EnvVarOrDict = Union[EnvVar, EnvVarDict] -class GetEvaluationItemConfig(_common.BaseModel): - """Config for get evaluation item.""" +class ContainerSpec(_common.BaseModel): + """The spec of a Container.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + args: Optional[list[str]] = Field( + default=None, + description="""The arguments to be passed when starting the container.""", + ) + command: Optional[list[str]] = Field( + default=None, + 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 GetEvaluationItemConfigDict(TypedDict, total=False): - """Config for get evaluation item.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -GetEvaluationItemConfigOrDict = Union[ - GetEvaluationItemConfig, GetEvaluationItemConfigDict -] - - -class _GetEvaluationItemParameters(_common.BaseModel): - """Represents a job that gets an evaluation item.""" - - name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationItemConfig] = Field(default=None, description="""""") +class ContainerSpecDict(TypedDict, total=False): + """The spec of a Container.""" + args: Optional[list[str]] + """The arguments to be passed when starting the container.""" -class _GetEvaluationItemParametersDict(TypedDict, total=False): - """Represents a job that gets an evaluation item.""" + command: Optional[list[str]] + """The command to be invoked when the container is started. It overrides the entrypoint instruction in Dockerfile when provided.""" - name: Optional[str] - """""" + env: Optional[list[EnvVarDict]] + """Environment variables to be passed to the container. Maximum limit is 100.""" - config: Optional[GetEvaluationItemConfigDict] - """""" + image_uri: Optional[str] + """Required. The URI of a container image in the Container Registry that is to be run on each worker replica.""" -_GetEvaluationItemParametersOrDict = Union[ - _GetEvaluationItemParameters, _GetEvaluationItemParametersDict -] +ContainerSpecOrDict = Union[ContainerSpec, ContainerSpecDict] -class ImportSchemaConfig(_common.BaseModel): - """Configuration for the input data format.""" +class DiskSpec(_common.BaseModel): + """Represents the spec of disk options.""" - data_format: Optional[ImportDataFormat] = Field( - default=None, description="""The format of the input data.""" + boot_disk_size_gb: Optional[int] = Field( + default=None, description="""Size in GB of the boot disk (default is 100GB).""" ) - data_format_version: Optional[str] = Field( - default=None, description="""Version of the data format.""" + 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".""", ) -class ImportSchemaConfigDict(TypedDict, total=False): - """Configuration for the input data format.""" +class DiskSpecDict(TypedDict, total=False): + """Represents the spec of disk options.""" - data_format: Optional[ImportDataFormat] - """The format of the input data.""" + boot_disk_size_gb: Optional[int] + """Size in GB of the boot disk (default is 100GB).""" - data_format_version: Optional[str] - """Version of the data format.""" + 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".""" -ImportSchemaConfigOrDict = Union[ImportSchemaConfig, ImportSchemaConfigDict] +DiskSpecOrDict = Union[DiskSpec, DiskSpecDict] -class EvaluationSetGcsSource(_common.BaseModel): - """Source for loading data from Cloud Storage.""" +class LustreMount(_common.BaseModel): + """Represents a mount configuration for Lustre file system.""" - gcs_uri: Optional[str] = Field( - default=None, description="""The Cloud Storage location of the input data.""" + filesystem: Optional[str] = Field( + default=None, description="""Required. The name of the Lustre filesystem.""" ) - import_schema_config: Optional[ImportSchemaConfig] = Field( - default=None, description="""Schema configuration for the input data.""" + instance_ip: Optional[str] = Field( + default=None, description="""Required. IP address of the Lustre instance.""" ) - - -class EvaluationSetGcsSourceDict(TypedDict, total=False): - """Source for loading data from Cloud Storage.""" - - gcs_uri: Optional[str] - """The Cloud Storage location of the input data.""" - - import_schema_config: Optional[ImportSchemaConfigDict] - """Schema configuration for the input data.""" - - -EvaluationSetGcsSourceOrDict = Union[EvaluationSetGcsSource, EvaluationSetGcsSourceDict] - - -class EvaluationSetInlineSource(_common.BaseModel): - """Wrapper for inline data.""" - - content: Optional[bytes] = Field( - default=None, description="""The content of the inline data.""" + 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/""", ) - import_schema_config: Optional[ImportSchemaConfig] = Field( - default=None, description="""Schema configuration for the inline data.""" + volume_handle: Optional[str] = Field( + default=None, + description="""Required. The unique identifier of the Lustre volume.""", ) -class EvaluationSetInlineSourceDict(TypedDict, total=False): - """Wrapper for inline data.""" +class LustreMountDict(TypedDict, total=False): + """Represents a mount configuration for Lustre file system.""" - content: Optional[bytes] - """The content of the inline data.""" + filesystem: Optional[str] + """Required. The name of the Lustre filesystem.""" - import_schema_config: Optional[ImportSchemaConfigDict] - """Schema configuration for the inline data.""" + 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/""" -EvaluationSetInlineSourceOrDict = Union[ - EvaluationSetInlineSource, EvaluationSetInlineSourceDict -] + volume_handle: Optional[str] + """Required. The unique identifier of the Lustre volume.""" -class EvaluationSetCloudTraceSource(_common.BaseModel): - """Source for loading traces directly from Cloud Trace.""" +LustreMountOrDict = Union[LustreMount, LustreMountDict] - project_id: Optional[str] = Field( - default=None, description="""Project ID for the Cloud Trace.""" + +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.""" + + 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.""", ) - trace_ids: Optional[list[str]] = Field( - default=None, description="""Trace IDs to import.""" + reservation_affinity_type: Optional[Type] = Field( + default=None, + description="""Required. Specifies the reservation affinity type.""", ) - session_ids: Optional[list[str]] = Field( + values: Optional[list[str]] = Field( default=None, - description="""Session IDs to import traces for. If both trace_ids and - session_ids are specified, the union of the two will be imported.""", + 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 EvaluationSetCloudTraceSourceDict(TypedDict, total=False): - """Source for loading traces directly from Cloud Trace.""" +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.""" - project_id: Optional[str] - """Project ID for the Cloud Trace.""" + 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.""" - trace_ids: Optional[list[str]] - """Trace IDs to import.""" + reservation_affinity_type: Optional[Type] + """Required. Specifies the reservation affinity type.""" - session_ids: Optional[list[str]] - """Session IDs to import traces for. If both trace_ids and - session_ids are specified, the union of the two will be imported.""" + 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.""" -EvaluationSetCloudTraceSourceOrDict = Union[ - EvaluationSetCloudTraceSource, EvaluationSetCloudTraceSourceDict -] +ReservationAffinityOrDict = Union[ReservationAffinity, ReservationAffinityDict] -class ImportEvaluationSetConfig(_common.BaseModel): - """Config for importing an evaluation set.""" +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.""", ) - - -class ImportEvaluationSetConfigDict(TypedDict, total=False): - """Config for importing an evaluation set.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -ImportEvaluationSetConfigOrDict = Union[ - ImportEvaluationSetConfig, ImportEvaluationSetConfigDict -] - - -class _ImportEvaluationSetParameters(_common.BaseModel): - """Parameters for importing an evaluation set.""" - - evaluation_set: Optional[EvaluationSet] = Field( + accelerator_type: Optional[AcceleratorType] = Field( default=None, - description="""The EvaluationSet to create. Used to specify 'display_name' and - 'metadata'. The 'evaluation_items' field is ignored and populated by the - import process.""", + description="""Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""", ) - gcs_destination: Optional[genai_types.GcsDestination] = Field( + gpu_partition_size: Optional[str] = Field( default=None, - description="""The Cloud Storage location where the resulting EvaluationItem - payloads will be stored.""", + 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.""", ) - gcs_source: Optional[EvaluationSetGcsSource] = Field( - default=None, description="""Google Cloud Storage location.""" + 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.""", ) - inline_source: Optional[EvaluationSetInlineSource] = Field( - default=None, description="""Inline source for small payloads (< 4MB).""" + 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.""", ) - cloud_trace_source: Optional[EvaluationSetCloudTraceSource] = Field( + multihost_gpu_node_count: Optional[int] = Field( default=None, - description="""Source for loading data directly from Cloud Trace.""", + description="""Optional. Immutable. The number of nodes per replica for multihost GPU deployments.""", ) - config: Optional[ImportEvaluationSetConfig] = Field( - default=None, description="""""" + 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 _ImportEvaluationSetParametersDict(TypedDict, total=False): - """Parameters for importing an evaluation set.""" +class MachineSpecDict(TypedDict, total=False): + """Specification of a single machine.""" - evaluation_set: Optional[EvaluationSetDict] - """The EvaluationSet to create. Used to specify 'display_name' and - 'metadata'. The 'evaluation_items' field is ignored and populated by the - import process.""" + 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.""" - gcs_destination: Optional[genai_types.GcsDestination] - """The Cloud Storage location where the resulting EvaluationItem - payloads will be stored.""" + accelerator_type: Optional[AcceleratorType] + """Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""" - gcs_source: Optional[EvaluationSetGcsSourceDict] - """Google Cloud Storage location.""" + 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.""" - inline_source: Optional[EvaluationSetInlineSourceDict] - """Inline source for small payloads (< 4MB).""" + 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.""" - cloud_trace_source: Optional[EvaluationSetCloudTraceSourceDict] - """Source for loading data directly from Cloud Trace.""" + 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.""" - config: Optional[ImportEvaluationSetConfigDict] - """""" + 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.""" + tpu_topology: Optional[str] + """Immutable. The topology of the TPUs. Corresponds to the TPU topologies available from GKE. (Example: tpu_topology: "2x2x1").""" -_ImportEvaluationSetParametersOrDict = Union[ - _ImportEvaluationSetParameters, _ImportEvaluationSetParametersDict -] +MachineSpecOrDict = Union[MachineSpec, MachineSpecDict] -class ImportEvaluationSetOperation(_common.BaseModel): - """Operation for importing an evaluation set.""" - 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( +class NfsMount(_common.BaseModel): + """Represents a mount configuration for Network File System (NFS) to mount.""" + + mount_point: 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="""Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""", ) - done: Optional[bool] = Field( + path: Optional[str] = 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. 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`""", ) - error: Optional[dict[str, Any]] = Field( - default=None, - description="""The error result of the operation in case of failure or cancellation.""", + server: Optional[str] = Field( + default=None, description="""Required. IP address of the NFS server.""" ) -class ImportEvaluationSetOperationDict(TypedDict, total=False): - """Operation for importing an evaluation set.""" - - 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 NfsMountDict(TypedDict, total=False): + """Represents a mount configuration for Network File System (NFS) to mount.""" - 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.""" + mount_point: Optional[str] + """Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""" - 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.""" + 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`""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + server: Optional[str] + """Required. IP address of the NFS server.""" -ImportEvaluationSetOperationOrDict = Union[ - ImportEvaluationSetOperation, ImportEvaluationSetOperationDict -] +NfsMountOrDict = Union[NfsMount, NfsMountDict] -class ListEvaluationExperimentsConfig(_common.BaseModel): - """Config for listing evaluation experiments.""" +class PythonPackageSpec(_common.BaseModel): + """The spec of a Python packaged code.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + args: Optional[list[str]] = Field( + default=None, + description="""Command line arguments to be passed to the Python task.""", ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( + env: Optional[list[EnvVar]] = 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="""Environment variables to be passed to the python module. Maximum limit is 100.""", ) - order_by: Optional[str] = Field( + executor_image_uri: 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="""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 ListEvaluationExperimentsConfigDict(TypedDict, total=False): - """Config for listing evaluation experiments.""" +class PythonPackageSpecDict(TypedDict, total=False): + """The spec of a Python packaged code.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + args: Optional[list[str]] + """Command line arguments to be passed to the Python task.""" - page_size: Optional[int] - """""" + env: Optional[list[EnvVarDict]] + """Environment variables to be passed to the python module. Maximum limit is 100.""" - page_token: Optional[str] - """""" + 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.""" - 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 `_.""" + 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.""" - 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"``.""" + python_module: Optional[str] + """Required. The Python module name to run after installing the packages.""" -ListEvaluationExperimentsConfigOrDict = Union[ - ListEvaluationExperimentsConfig, ListEvaluationExperimentsConfigDict -] +PythonPackageSpecOrDict = Union[PythonPackageSpec, PythonPackageSpecDict] -class _ListEvaluationExperimentsParameters(_common.BaseModel): - """Parameters for listing evaluation experiments.""" +class WorkerPoolSpec(_common.BaseModel): + """Represents the spec of a worker pool in a job.""" - config: Optional[ListEvaluationExperimentsConfig] = Field( - default=None, description="""""" + 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.""", ) -class _ListEvaluationExperimentsParametersDict(TypedDict, total=False): - """Parameters for listing evaluation experiments.""" - - config: Optional[ListEvaluationExperimentsConfigDict] - """""" +class WorkerPoolSpecDict(TypedDict, total=False): + """Represents the spec of a worker pool in a job.""" + container_spec: Optional[ContainerSpecDict] + """The custom container task.""" -_ListEvaluationExperimentsParametersOrDict = Union[ - _ListEvaluationExperimentsParameters, _ListEvaluationExperimentsParametersDict -] + disk_spec: Optional[DiskSpecDict] + """Disk spec.""" + lustre_mounts: Optional[list[LustreMountDict]] + """Optional. List of Lustre mounts.""" -class ListEvaluationExperimentsResponse(_common.BaseModel): - """Response for listing evaluation experiments.""" + machine_spec: Optional[MachineSpecDict] + """Optional. Immutable. The specification of a single machine.""" - 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="""List of evaluation experiments. - """, - ) + nfs_mounts: Optional[list[NfsMountDict]] + """Optional. List of NFS mount spec.""" + python_package_spec: Optional[PythonPackageSpecDict] + """The Python packaged task.""" -class ListEvaluationExperimentsResponseDict(TypedDict, total=False): - """Response for listing evaluation experiments.""" - - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" - - next_page_token: Optional[str] - """""" - - evaluation_experiments: Optional[list[EvaluationExperimentDict]] - """List of evaluation experiments. - """ + replica_count: Optional[int] + """Optional. The number of worker replicas to use for this worker pool.""" -ListEvaluationExperimentsResponseOrDict = Union[ - ListEvaluationExperimentsResponse, ListEvaluationExperimentsResponseDict -] +WorkerPoolSpecOrDict = Union[WorkerPoolSpec, WorkerPoolSpecDict] -class ListEvaluationMetricsConfig(_common.BaseModel): - """Config for listing evaluation metrics.""" +class CustomJobSpec(_common.BaseModel): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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/`""", ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( + enable_dashboard_access: Optional[bool] = 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. 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).""", ) - order_by: Optional[str] = Field( + enable_web_access: 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. 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.""", + ) + 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 ListEvaluationMetricsConfigDict(TypedDict, total=False): - """Config for listing evaluation metrics.""" +class CustomJobSpecDict(TypedDict, total=False): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + 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/`""" - page_size: Optional[int] - """""" + 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).""" - page_token: Optional[str] - """""" + 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).""" - 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 `_.""" + experiment: Optional[str] + """Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`""" - 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"``.""" + 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.""" -ListEvaluationMetricsConfigOrDict = Union[ - ListEvaluationMetricsConfig, ListEvaluationMetricsConfigDict -] + 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 _ListEvaluationMetricsParameters(_common.BaseModel): - """Parameters for listing evaluation metrics.""" + 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""" - config: Optional[ListEvaluationMetricsConfig] = Field( - default=None, description="""""" - ) + 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 _ListEvaluationMetricsParametersDict(TypedDict, total=False): - """Parameters for listing evaluation metrics.""" + scheduling: Optional[SchedulingDict] + """Scheduling options for a CustomJob.""" - config: Optional[ListEvaluationMetricsConfigDict] - """""" + 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}`""" -_ListEvaluationMetricsParametersOrDict = Union[ - _ListEvaluationMetricsParameters, _ListEvaluationMetricsParametersDict -] + 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.""" -class ListEvaluationMetricsResponse(_common.BaseModel): - """Response for listing evaluation metrics.""" +CustomJobSpecOrDict = Union[CustomJobSpec, CustomJobSpecDict] - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + +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.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - evaluation_metrics: Optional[list[EvaluationMetric]] = Field( + job_spec: Optional[CustomJobSpec] = Field( + default=None, description="""Required. Job spec.""" + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, - description="""List of evaluation metrics. - """, + 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.""", + ) + 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 ListEvaluationMetricsResponseDict(TypedDict, total=False): - """Response for listing evaluation metrics.""" +class CustomJobDict(TypedDict, total=False): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + 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.""" - next_page_token: Optional[str] - """""" + job_spec: Optional[CustomJobSpecDict] + """Required. Job spec.""" - evaluation_metrics: Optional[list[EvaluationMetricDict]] - """List of evaluation metrics. - """ + 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.""" -ListEvaluationMetricsResponseOrDict = Union[ - ListEvaluationMetricsResponse, ListEvaluationMetricsResponseDict -] + error: Optional[genai_types.GoogleRpcStatus] + """Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`.""" + create_time: Optional[datetime.datetime] + """Output only. Time when the CustomJob was created.""" -class ListEvaluationSetsConfig(_common.BaseModel): - """Config for listing evaluation sets.""" + 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`.""" - 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"``.""", - ) + 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.""" + name: Optional[str] + """Output only. Resource name of a CustomJob.""" -class ListEvaluationSetsConfigDict(TypedDict, total=False): - """Config for listing evaluation sets.""" + satisfies_pzi: Optional[bool] + """Output only. Reserved for future use.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + satisfies_pzs: Optional[bool] + """Output only. Reserved for future use.""" - page_size: Optional[int] - """""" + start_time: Optional[datetime.datetime] + """Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state.""" - page_token: Optional[str] - """""" + update_time: Optional[datetime.datetime] + """Output only. Time when the CustomJob was most recently updated.""" - 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 `_.""" + 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.""" - 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"``.""" +CustomJobOrDict = Union[CustomJob, CustomJobDict] -ListEvaluationSetsConfigOrDict = Union[ - ListEvaluationSetsConfig, ListEvaluationSetsConfigDict -] +class VertexBaseConfig(_common.BaseModel): + """Base config for Vertex AI.""" -class _ListEvaluationSetsParameters(_common.BaseModel): - """Parameters for listing evaluation sets.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) - config: Optional[ListEvaluationSetsConfig] = Field(default=None, description="""""") +class VertexBaseConfigDict(TypedDict, total=False): + """Base config for Vertex AI.""" -class _ListEvaluationSetsParametersDict(TypedDict, total=False): - """Parameters for listing evaluation sets.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - config: Optional[ListEvaluationSetsConfigDict] + +VertexBaseConfigOrDict = Union[VertexBaseConfig, VertexBaseConfigDict] + + +class _CustomJobParameters(_common.BaseModel): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" + + custom_job: Optional[CustomJob] = Field(default=None, description="""""") + config: Optional[VertexBaseConfig] = Field(default=None, description="""""") + + +class _CustomJobParametersDict(TypedDict, total=False): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" + + custom_job: Optional[CustomJobDict] """""" + config: Optional[VertexBaseConfigDict] + """""" -_ListEvaluationSetsParametersOrDict = Union[ - _ListEvaluationSetsParameters, _ListEvaluationSetsParametersDict -] +_CustomJobParametersOrDict = Union[_CustomJobParameters, _CustomJobParametersDict] -class ListEvaluationSetsResponse(_common.BaseModel): - """Response for listing evaluation sets.""" - 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_sets: Optional[list[EvaluationSet]] = Field( - default=None, - description="""List of evaluation sets. - """, - ) +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="""""") -class ListEvaluationSetsResponseDict(TypedDict, total=False): - """Response for listing evaluation sets.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" +class _GetCustomJobParametersDict(TypedDict, total=False): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - next_page_token: Optional[str] + name: Optional[str] """""" - evaluation_sets: Optional[list[EvaluationSetDict]] - """List of evaluation sets. - """ + config: Optional[VertexBaseConfigDict] + """""" -ListEvaluationSetsResponseOrDict = Union[ - ListEvaluationSetsResponse, ListEvaluationSetsResponseDict +_GetCustomJobParametersOrDict = Union[ + _GetCustomJobParameters, _GetCustomJobParametersDict ] -class UpdateEvaluationExperimentConfig(_common.BaseModel): - """Config for updating an evaluation experiment.""" +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.""" ) - update_mask: Optional[str] = Field( + operation_name: 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.""" + description="""Name of the longrunning operation returned from run_query_job.""", ) -class UpdateEvaluationExperimentConfigDict(TypedDict, total=False): - """Config for updating an evaluation experiment.""" +class CancelQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for canceling async querying agent runtimes.""" 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.""" - - 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.""" + operation_name: Optional[str] + """Name of the longrunning operation returned from run_query_job.""" -UpdateEvaluationExperimentConfigOrDict = Union[ - UpdateEvaluationExperimentConfig, UpdateEvaluationExperimentConfigDict +CancelQueryJobRuntimeConfigOrDict = Union[ + CancelQueryJobRuntimeConfig, CancelQueryJobRuntimeConfigDict ] -class _UpdateEvaluationExperimentParameters(_common.BaseModel): - """Parameters for updating an evaluation experiment.""" +class _CancelQueryJobRuntimeRequestParameters(_common.BaseModel): + """Parameters for canceling async querying agent runtimes.""" name: Optional[str] = Field( - default=None, description="""The resource name of the EvaluationExperiment.""" + default=None, description="""Name of the reasoning engine resource.""" ) - config: Optional[UpdateEvaluationExperimentConfig] = Field( + config: Optional[CancelQueryJobRuntimeConfig] = Field( default=None, description="""""" ) -class _UpdateEvaluationExperimentParametersDict(TypedDict, total=False): - """Parameters for updating an evaluation experiment.""" +class _CancelQueryJobRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for canceling async querying agent runtimes.""" name: Optional[str] - """The resource name of the EvaluationExperiment.""" + """Name of the reasoning engine resource.""" - config: Optional[UpdateEvaluationExperimentConfigDict] + config: Optional[CancelQueryJobRuntimeConfigDict] """""" -_UpdateEvaluationExperimentParametersOrDict = Union[ - _UpdateEvaluationExperimentParameters, _UpdateEvaluationExperimentParametersDict +_CancelQueryJobRuntimeRequestParametersOrDict = Union[ + _CancelQueryJobRuntimeRequestParameters, _CancelQueryJobRuntimeRequestParametersDict ] -class OptimizeConfig(_common.BaseModel): - """Config for Prompt Optimizer.""" +class CancelQueryJobResult(_common.BaseModel): + """Result of canceling a query job.""" 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 OptimizeConfigDict(TypedDict, total=False): - """Config for Prompt Optimizer.""" +class CancelQueryJobResultDict(TypedDict, total=False): + """Result of canceling a query job.""" 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.""" - - 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.""" - -OptimizeConfigOrDict = Union[OptimizeConfig, OptimizeConfigDict] +CancelQueryJobResultOrDict = Union[CancelQueryJobResult, CancelQueryJobResultDict] -class _OptimizeRequestParameters(_common.BaseModel): - """Request for the optimize_prompt method.""" +class CheckQueryJobRuntimeConfig(_common.BaseModel): + """Config for async querying agent runtimes.""" - content: Optional[genai_types.Content] = Field(default=None, description="""""") - config: Optional[OptimizeConfig] = Field(default=None, description="""""") + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + retrieve_result: Optional[bool] = Field( + default=None, + description="""Whether to retrieve the results of the query job.""", + ) -class _OptimizeRequestParametersDict(TypedDict, total=False): - """Request for the optimize_prompt method.""" +class CheckQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for async querying agent runtimes.""" - content: Optional[genai_types.Content] - """""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - config: Optional[OptimizeConfigDict] - """""" + retrieve_result: Optional[bool] + """Whether to retrieve the results of the query job.""" -_OptimizeRequestParametersOrDict = Union[ - _OptimizeRequestParameters, _OptimizeRequestParametersDict +CheckQueryJobRuntimeConfigOrDict = Union[ + CheckQueryJobRuntimeConfig, CheckQueryJobRuntimeConfigDict ] -class OptimizeResponseEndpoint(_common.BaseModel): - """Response for the optimize_prompt method.""" +class _CheckQueryJobRuntimeRequestParameters(_common.BaseModel): + """Parameters for async querying agent runtimes.""" - content: Optional[genai_types.Content] = Field(default=None, description="""""") + name: Optional[str] = Field(default=None, description="""Name of the query job.""") + config: Optional[CheckQueryJobRuntimeConfig] = Field( + default=None, description="""""" + ) -class OptimizeResponseEndpointDict(TypedDict, total=False): - """Response for the optimize_prompt method.""" +class _CheckQueryJobRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for async querying agent runtimes.""" - content: Optional[genai_types.Content] + name: Optional[str] + """Name of the query job.""" + + config: Optional[CheckQueryJobRuntimeConfigDict] """""" -OptimizeResponseEndpointOrDict = Union[ - OptimizeResponseEndpoint, OptimizeResponseEndpointDict +_CheckQueryJobRuntimeRequestParametersOrDict = Union[ + _CheckQueryJobRuntimeRequestParameters, _CheckQueryJobRuntimeRequestParametersDict ] -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 CheckQueryJobResult(_common.BaseModel): + """Result of checking a query job.""" - 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - 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.""", + operation_name: Optional[str] = Field( + default=None, description="""Name of the agent runtime operation.""" ) - 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.""", + output_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the output file.""" + ) + status: Optional[str] = Field( + default=None, description="""Status of the operation.""" + ) + result: Optional[str] = Field( + default=None, description="""JSON result of the operation.""" ) -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.""" +class CheckQueryJobResultDict(TypedDict, total=False): + """Result of checking a query job.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - target_network: Optional[str] - """Required. The VPC network name in the target_project where the DNS zone specified by 'domain' is visible.""" + operation_name: Optional[str] + """Name of the agent runtime operation.""" - 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( - 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.""", - ) - - -class PscInterfaceConfigDict(TypedDict, total=False): - """Configuration for PSC-I.""" + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" - 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.""" + status: Optional[str] + """Status of the operation.""" - 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.""" + result: Optional[str] + """JSON result of the operation.""" -PscInterfaceConfigOrDict = Union[PscInterfaceConfig, PscInterfaceConfigDict] +CheckQueryJobResultOrDict = Union[CheckQueryJobResult, CheckQueryJobResultDict] -class Scheduling(_common.BaseModel): - """All parameters related to queuing and scheduling of custom jobs.""" +class _RunQueryJobRuntimeConfig(_common.BaseModel): + """Config for running a query job on an agent runtime.""" - 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.""", - ) - max_wait_duration: 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.""", - ) - restart_job_on_worker_restart: Optional[bool] = 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - strategy: Optional[Strategy] = Field( - default=None, - description="""Optional. This determines which type of scheduling strategy to use.""", + input_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the input file.""" ) - timeout: Optional[str] = Field( - default=None, - description="""Optional. The maximum job running time. The default is 7 days.""", + output_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the output file.""" ) -class SchedulingDict(TypedDict, total=False): - """All parameters related to queuing and scheduling of custom jobs.""" - - 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.""" - - 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.""" +class _RunQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for running a query job on an agent runtime.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - strategy: Optional[Strategy] - """Optional. This determines which type of scheduling strategy to use.""" + input_gcs_uri: Optional[str] + """The GCS URI of the input file.""" - timeout: Optional[str] - """Optional. The maximum job running time. The default is 7 days.""" + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" -SchedulingOrDict = Union[Scheduling, SchedulingDict] +_RunQueryJobRuntimeConfigOrDict = Union[ + _RunQueryJobRuntimeConfig, _RunQueryJobRuntimeConfigDict +] -class EnvVar(_common.BaseModel): - """Represents an environment variable present in a Container or Python Module.""" +class _RunQueryJobRuntimeRequestParameters(_common.BaseModel): + """Parameters for running a query job on an agent runtime.""" name: Optional[str] = Field( - default=None, - description="""Required. Name of the environment variable. Must be a valid C identifier.""", + default=None, description="""Name of the agent runtime.""" ) - 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.""", + config: Optional[_RunQueryJobRuntimeConfig] = Field( + default=None, description="""""" ) -class EnvVarDict(TypedDict, total=False): - """Represents an environment variable present in a Container or Python Module.""" +class _RunQueryJobRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for running a query job on an agent runtime.""" name: Optional[str] - """Required. Name of the environment variable. Must be a valid C identifier.""" + """Name of the agent runtime.""" - 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.""" + config: Optional[_RunQueryJobRuntimeConfigDict] + """""" -EnvVarOrDict = Union[EnvVar, EnvVarDict] +_RunQueryJobRuntimeRequestParametersOrDict = Union[ + _RunQueryJobRuntimeRequestParameters, _RunQueryJobRuntimeRequestParametersDict +] -class ContainerSpec(_common.BaseModel): - """The spec of a Container.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent( + _common.BaseModel +): + """The conversation source event for generating memories.""" - args: Optional[list[str]] = Field( - default=None, - description="""The arguments to be passed when starting the container.""", - ) - command: Optional[list[str]] = Field( - default=None, - 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.""", + content: Optional[genai_types.Content] = Field( + default=None, description="""Required. Represents the content of the event.""" ) -class ContainerSpecDict(TypedDict, total=False): - """The spec of a Container.""" - - args: Optional[list[str]] - """The arguments to be passed when starting the container.""" - - command: Optional[list[str]] - """The command to be invoked when the container is started. It overrides the entrypoint instruction in Dockerfile when provided.""" - - env: Optional[list[EnvVarDict]] - """Environment variables to be passed to the container. Maximum limit is 100.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict( + TypedDict, total=False +): + """The conversation source event for generating memories.""" - image_uri: Optional[str] - """Required. The URI of a container image in the Container Registry that is to be run on each worker replica.""" + content: Optional[genai_types.Content] + """Required. Represents the content of the event.""" -ContainerSpecOrDict = Union[ContainerSpec, ContainerSpecDict] +MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventOrDict = ( + Union[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent, + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict, + ] +) -class DiskSpec(_common.BaseModel): - """Represents the spec of disk options.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource( + _common.BaseModel +): + """A conversation source for the example. This is similar to `DirectContentsSource`.""" - 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( + events: Optional[ + list[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent + ] + ] = 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".""", + description="""Optional. Represents the input conversation events for the example.""", ) -class DiskSpecDict(TypedDict, total=False): - """Represents the spec of disk options.""" - - boot_disk_size_gb: Optional[int] - """Size in GB of the boot disk (default is 100GB).""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict( + TypedDict, total=False +): + """A conversation source for the example. This is similar to `DirectContentsSource`.""" - 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".""" + events: Optional[ + list[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict + ] + ] + """Optional. Represents the input conversation events for the example.""" -DiskSpecOrDict = Union[DiskSpec, DiskSpecDict] +MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceOrDict = Union[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource, + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict, +] -class LustreMount(_common.BaseModel): - """Represents a mount configuration for Lustre file system.""" +class MemoryTopicId(_common.BaseModel): + """The topic ID for a memory.""" - 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( + custom_memory_topic_label: Optional[str] = Field( default=None, - description="""Required. Destination mount path. The Lustre file system will be mounted for the user under /mnt/lustre/""", + description="""Optional. Represents the custom memory topic label.""", ) - volume_handle: Optional[str] = Field( - default=None, - description="""Required. The unique identifier of the Lustre volume.""", + managed_memory_topic: Optional[ManagedTopicEnum] = Field( + default=None, description="""Optional. Represents the managed memory topic.""" ) -class LustreMountDict(TypedDict, total=False): - """Represents a mount configuration for Lustre file system.""" - - filesystem: Optional[str] - """Required. The name of the Lustre filesystem.""" - - instance_ip: Optional[str] - """Required. IP address of the Lustre instance.""" +class MemoryTopicIdDict(TypedDict, total=False): + """The topic ID for a memory.""" - mount_point: Optional[str] - """Required. Destination mount path. The Lustre file system will be mounted for the user under /mnt/lustre/""" + custom_memory_topic_label: Optional[str] + """Optional. Represents the custom memory topic label.""" - volume_handle: Optional[str] - """Required. The unique identifier of the Lustre volume.""" + managed_memory_topic: Optional[ManagedTopicEnum] + """Optional. Represents the managed memory topic.""" -LustreMountOrDict = Union[LustreMount, LustreMountDict] +MemoryTopicIdOrDict = Union[MemoryTopicId, MemoryTopicIdDict] -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 MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory( + _common.BaseModel +): + """A memory generated by the operation.""" - 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( + fact: Optional[str] = Field( default=None, - description="""Required. Specifies the reservation affinity type.""", + description="""Required. Represents the fact to generate a memory from.""", ) - values: Optional[list[str]] = Field( + topics: Optional[list[MemoryTopicId]] = 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.""", + 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 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.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict( + TypedDict, total=False +): + """A memory generated by the operation.""" - 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.""" + fact: Optional[str] + """Required. Represents the fact to generate a memory from.""" - reservation_affinity_type: Optional[Type] - """Required. Specifies the reservation affinity type.""" + 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`.""" - 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.""" - -ReservationAffinityOrDict = Union[ReservationAffinity, ReservationAffinityDict] +MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryOrDict = Union[ + MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory, + MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict, +] -class MachineSpec(_common.BaseModel): - """Specification of a single machine.""" +class MemoryBankCustomizationConfigGenerateMemoriesExample(_common.BaseModel): + """An example of how to generate memories for a particular scope.""" - 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( + conversation_source: Optional[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource + ] = Field(default=None, description="""A conversation source for the example.""") + generated_memories: Optional[ + list[MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory] + ] = Field( default=None, - description="""Optional. Immutable. The number of nodes per replica for multihost GPU deployments.""", + 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.""", ) - reservation_affinity: Optional[ReservationAffinity] = Field( - default=None, - description="""Optional. Immutable. Configuration controlling how this resource pool consumes reservation.""", + + +class MemoryBankCustomizationConfigGenerateMemoriesExampleDict(TypedDict, total=False): + """An example of how to generate memories for a particular scope.""" + + 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.""" + + +MemoryBankCustomizationConfigGenerateMemoriesExampleOrDict = Union[ + MemoryBankCustomizationConfigGenerateMemoriesExample, + MemoryBankCustomizationConfigGenerateMemoriesExampleDict, +] + + +class MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic(_common.BaseModel): + """A custom memory topic defined by the developer.""" + + label: Optional[str] = Field( + default=None, description="""Required. Represents the label of the topic.""" ) - tpu_topology: Optional[str] = Field( + description: 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="""Required. Represents the description of the memory topic. This should explain what information should be extracted for this topic.""", ) -class MachineSpecDict(TypedDict, total=False): - """Specification of a single machine.""" +class MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict( + TypedDict, total=False +): + """A custom memory topic defined by the developer.""" - 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.""" + label: Optional[str] + """Required. Represents the label of the topic.""" - accelerator_type: Optional[AcceleratorType] - """Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""" + description: Optional[str] + """Required. Represents the description of the memory topic. This should explain what information should be extracted for this topic.""" - 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.""" +MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicOrDict = Union[ + MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic, + MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict, +] - 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 MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic(_common.BaseModel): + """A managed memory topic defined by the system.""" - reservation_affinity: Optional[ReservationAffinityDict] - """Optional. Immutable. Configuration controlling how this resource pool consumes reservation.""" + managed_topic_enum: Optional[ManagedTopicEnum] = Field( + default=None, description="""Required. Represents the managed topic.""" + ) - tpu_topology: Optional[str] - """Immutable. The topology of the TPUs. Corresponds to the TPU topologies available from GKE. (Example: tpu_topology: "2x2x1").""" +class MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict( + TypedDict, total=False +): + """A managed memory topic defined by the system.""" -MachineSpecOrDict = Union[MachineSpec, MachineSpecDict] + managed_topic_enum: Optional[ManagedTopicEnum] + """Required. Represents the managed topic.""" -class NfsMount(_common.BaseModel): - """Represents a mount configuration for Network File System (NFS) to mount.""" +MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicOrDict = Union[ + MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic, + MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict, +] - 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`""", + +class MemoryBankCustomizationConfigMemoryTopic(_common.BaseModel): + """A topic of information that should be extracted from conversations and stored as memories.""" + + custom_memory_topic: Optional[ + MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic + ] = Field( + default=None, description="""A custom memory topic defined by the developer.""" ) - server: Optional[str] = Field( - default=None, description="""Required. IP address of the NFS server.""" + managed_memory_topic: Optional[ + MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic + ] = Field( + default=None, description="""A managed memory topic defined by Memory Bank.""" ) -class NfsMountDict(TypedDict, total=False): - """Represents a mount configuration for Network File System (NFS) to mount.""" +class MemoryBankCustomizationConfigMemoryTopicDict(TypedDict, total=False): + """A topic of information that should be extracted from conversations and stored as memories.""" - mount_point: Optional[str] - """Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""" + custom_memory_topic: Optional[ + MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict + ] + """A custom memory topic defined by the developer.""" - 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`""" + managed_memory_topic: Optional[ + MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict + ] + """A managed memory topic defined by Memory Bank.""" - server: Optional[str] - """Required. IP address of the NFS server.""" +MemoryBankCustomizationConfigMemoryTopicOrDict = Union[ + MemoryBankCustomizationConfigMemoryTopic, + MemoryBankCustomizationConfigMemoryTopicDict, +] -NfsMountOrDict = Union[NfsMount, NfsMountDict] +class MemoryBankCustomizationConfigConsolidationConfig(_common.BaseModel): + """Represents configuration for customizing how memories are consolidated.""" -class PythonPackageSpec(_common.BaseModel): - """The spec of a Python packaged code.""" + 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.""", + ) - args: Optional[list[str]] = Field( + +class MemoryBankCustomizationConfigConsolidationConfigDict(TypedDict, total=False): + """Represents configuration for customizing how memories are consolidated.""" + + 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.""" + + +MemoryBankCustomizationConfigConsolidationConfigOrDict = Union[ + MemoryBankCustomizationConfigConsolidationConfig, + MemoryBankCustomizationConfigConsolidationConfigDict, +] + + +class MemoryBankCustomizationConfig(_common.BaseModel): + """Represents configuration for organizing natural language memories.""" + + enable_third_person_memories: Optional[bool] = Field( default=None, - description="""Command line arguments to be passed to the Python task.""", + 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.")""", ) - env: Optional[list[EnvVar]] = Field( + generate_memories_examples: Optional[ + list[MemoryBankCustomizationConfigGenerateMemoriesExample] + ] = Field( default=None, - description="""Environment variables to be passed to the python module. Maximum limit is 100.""", + description="""Optional. Provides examples of how to generate memories for a particular scope.""", ) - executor_image_uri: Optional[str] = Field( + memory_topics: Optional[list[MemoryBankCustomizationConfigMemoryTopic]] = 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.""", + 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.""", ) - package_uris: Optional[list[str]] = Field( + scope_keys: 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.""", + 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.""", ) - python_module: Optional[str] = Field( + 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="""Required. The Python module name to run after installing the packages.""", + 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 PythonPackageSpecDict(TypedDict, total=False): - """The spec of a Python packaged code.""" +class MemoryBankCustomizationConfigDict(TypedDict, total=False): + """Represents configuration for organizing natural language memories.""" - args: Optional[list[str]] - """Command line arguments to be passed to the Python task.""" + 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.")""" - env: Optional[list[EnvVarDict]] - """Environment variables to be passed to the python module. Maximum limit is 100.""" + generate_memories_examples: Optional[ + list[MemoryBankCustomizationConfigGenerateMemoriesExampleDict] + ] + """Optional. Provides examples of how to generate memories for a particular scope.""" - 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.""" + 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.""" - 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.""" + 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.""" - python_module: Optional[str] - """Required. The Python module name to run after installing the packages.""" + 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.""" -PythonPackageSpecOrDict = Union[PythonPackageSpec, PythonPackageSpecDict] +MemoryBankCustomizationConfigOrDict = Union[ + MemoryBankCustomizationConfig, MemoryBankCustomizationConfigDict +] -class WorkerPoolSpec(_common.BaseModel): - """Represents the spec of a worker pool in a job.""" - 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( +class MemoryGenerationTriggerConfigGenerationTriggerRule(_common.BaseModel): + """Represents the active rule that determines when to flush the buffer.""" + + event_count: Optional[int] = Field( default=None, - description="""Optional. Immutable. The specification of a single machine.""", + description="""Optional. Specifies to trigger generation when the event count reaches this limit.""", ) - nfs_mounts: Optional[list[NfsMount]] = Field( - default=None, description="""Optional. List of NFS mount spec.""" + 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.""", ) - python_package_spec: Optional[PythonPackageSpec] = Field( - default=None, description="""The Python packaged task.""" + 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.""", ) - replica_count: Optional[int] = Field( + overlap_event_count: Optional[int] = Field( default=None, - description="""Optional. The number of worker replicas to use for this worker pool.""", + description="""Optional. Re-include the last N already-processed events in the next window.""", ) -class WorkerPoolSpecDict(TypedDict, total=False): - """Represents the spec of a worker pool in a job.""" +class MemoryGenerationTriggerConfigGenerationTriggerRuleDict(TypedDict, total=False): + """Represents the active rule that determines when to flush the buffer.""" - container_spec: Optional[ContainerSpecDict] - """The custom container task.""" + event_count: Optional[int] + """Optional. Specifies to trigger generation when the event count reaches this limit.""" - disk_spec: Optional[DiskSpecDict] - """Disk spec.""" + fixed_interval: Optional[str] + """Optional. Specifies to trigger generation at a fixed interval. The duration must have a minute-level granularity.""" - lustre_mounts: Optional[list[LustreMountDict]] - """Optional. List of Lustre mounts.""" + 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.""" - machine_spec: Optional[MachineSpecDict] - """Optional. Immutable. The specification of a single machine.""" + overlap_event_count: Optional[int] + """Optional. Re-include the last N already-processed events in the next window.""" - nfs_mounts: Optional[list[NfsMountDict]] - """Optional. List of NFS mount spec.""" - python_package_spec: Optional[PythonPackageSpecDict] - """The Python packaged task.""" +MemoryGenerationTriggerConfigGenerationTriggerRuleOrDict = Union[ + MemoryGenerationTriggerConfigGenerationTriggerRule, + MemoryGenerationTriggerConfigGenerationTriggerRuleDict, +] - replica_count: Optional[int] - """Optional. The number of worker replicas to use for this worker pool.""" +class MemoryGenerationTriggerConfig(_common.BaseModel): + """The configuration for triggering memory generation for ingested events.""" -WorkerPoolSpecOrDict = Union[WorkerPoolSpec, WorkerPoolSpecDict] + 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 CustomJobSpec(_common.BaseModel): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +class MemoryGenerationTriggerConfigDict(TypedDict, total=False): + """The configuration for triggering memory generation for ingested events.""" - 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.""", - ) - tensorboard: Optional[str] = Field( + 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.""" + + +MemoryGenerationTriggerConfigOrDict = Union[ + MemoryGenerationTriggerConfig, MemoryGenerationTriggerConfigDict +] + + +class ReasoningEngineContextSpecMemoryBankConfigGenerationConfig(_common.BaseModel): + """Configuration for how to generate memories.""" + + model: 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}`""", + description="""Optional. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", ) - worker_pool_specs: Optional[list[WorkerPoolSpec]] = Field( + generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = 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.""", + description="""Optional. Specifies the default trigger configuration for generating memories using `IngestEvents`.""", ) -class CustomJobSpecDict(TypedDict, total=False): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +class ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict( + TypedDict, total=False +): + """Configuration for how to generate memories.""" - 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/`""" + model: Optional[str] + """Optional. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" - 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).""" + generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] + """Optional. Specifies the default trigger configuration for generating memories using `IngestEvents`.""" - 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).""" - experiment: Optional[str] - """Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`""" +ReasoningEngineContextSpecMemoryBankConfigGenerationConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigGenerationConfig, + ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict, +] - 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.""" +class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig( + _common.BaseModel +): + """Configuration for how to perform similarity search on memories.""" - 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.""" + embedding_model: 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}`.""", + ) - 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""" +class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict( + TypedDict, total=False +): + """Configuration for how to perform similarity search on memories.""" - psc_interface_config: Optional[PscInterfaceConfigDict] - """Optional. Configuration for PSC-I for CustomJob.""" + embedding_model: Optional[str] + """Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" - 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.""" +ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig, + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict, +] - 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}`""" +class ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig( + _common.BaseModel +): + """Configuration for TTL of the memories in the Memory Bank based on the action that created or updated the memory.""" - 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.""" + create_ttl: Optional[str] = Field( + default=None, + description="""Optional. The TTL duration for memories uploaded via CreateMemory.""", + ) + generate_created_ttl: Optional[str] = Field( + default=None, + 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).""", + ) -CustomJobSpecOrDict = Union[CustomJobSpec, CustomJobSpecDict] +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.""" + generate_created_ttl: Optional[str] + """Optional. The TTL duration for memories newly generated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.CREATED).""" -class CustomJob(_common.BaseModel): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" + 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).""" - display_name: Optional[str] = Field( + +ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig, + ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict, +] + + +class ReasoningEngineContextSpecMemoryBankConfigTtlConfig(_common.BaseModel): + """Configuration for automatically setting the TTL ("time-to-live") of the memories in the Memory Bank.""" + + default_ttl: 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.""", + description="""Optional. The default TTL duration of the memories in the Memory Bank. This applies to all operations that create or update a memory.""", ) - job_spec: Optional[CustomJobSpec] = Field( - default=None, description="""Required. Job spec.""" + granular_ttl_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig + ] = Field( + default=None, + description="""Optional. The granular TTL configuration of the memories in the Memory Bank.""", ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + memory_revision_default_ttl: Optional[str] = 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.""", - ) - 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.""", + 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 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.""" +class ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict(TypedDict, total=False): + """Configuration for automatically setting the TTL ("time-to-live") of the memories in the Memory Bank.""" - job_spec: Optional[CustomJobSpecDict] - """Required. Job spec.""" + 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.""" - 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.""" + granular_ttl_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict + ] + """Optional. The granular TTL configuration of the memories in the Memory Bank.""" - state: Optional[genai_types.JobState] - """Output only. The detailed state of the job.""" + 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.""" - error: Optional[genai_types.GoogleRpcStatus] - """Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`.""" - create_time: Optional[datetime.datetime] - """Output only. Time when the CustomJob was created.""" +ReasoningEngineContextSpecMemoryBankConfigTtlConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigTtlConfig, + ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict, +] - 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.""" +class StructuredMemorySchemaConfig(_common.BaseModel): + """Represents the OpenAPI schema of the structured memories.""" - name: Optional[str] - """Output only. Resource name of a CustomJob.""" + memory_schema: Optional[genai_types.Schema] = Field( + default=None, + description="""Required. Represents the OpenAPI schema of the structured memories.""", + ) + id: Optional[str] = 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.""", + ) + 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.""", + ) - satisfies_pzi: Optional[bool] - """Output only. Reserved for future use.""" - satisfies_pzs: Optional[bool] - """Output only. Reserved for future use.""" +class StructuredMemorySchemaConfigDict(TypedDict, total=False): + """Represents the OpenAPI schema of the structured memories.""" - start_time: Optional[datetime.datetime] - """Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state.""" + memory_schema: Optional[genai_types.Schema] + """Required. Represents the OpenAPI schema of the structured memories.""" - update_time: Optional[datetime.datetime] - """Output only. Time when the CustomJob was most recently updated.""" + 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.""" - 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.""" + 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.""" -CustomJobOrDict = Union[CustomJob, CustomJobDict] +StructuredMemorySchemaConfigOrDict = Union[ + StructuredMemorySchemaConfig, StructuredMemorySchemaConfigDict +] -class VertexBaseConfig(_common.BaseModel): - """Base config for Vertex AI.""" +class StructuredMemoryConfig(_common.BaseModel): + """Configuration for organizing structured memories within a scope.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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 VertexBaseConfigDict(TypedDict, total=False): - """Base config for Vertex AI.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -VertexBaseConfigOrDict = Union[VertexBaseConfig, VertexBaseConfigDict] - +class StructuredMemoryConfigDict(TypedDict, total=False): + """Configuration for organizing structured memories within a scope.""" -class _CustomJobParameters(_common.BaseModel): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" + schema_configs: Optional[list[StructuredMemorySchemaConfigDict]] + """Optional. Represents configuration of the structured memories' schemas.""" - custom_job: Optional[CustomJob] = Field(default=None, description="""""") - config: Optional[VertexBaseConfig] = Field(default=None, description="""""") + 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 _CustomJobParametersDict(TypedDict, total=False): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +StructuredMemoryConfigOrDict = Union[StructuredMemoryConfig, StructuredMemoryConfigDict] - custom_job: Optional[CustomJobDict] - """""" - config: Optional[VertexBaseConfigDict] - """""" +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="""Optional. Configuration for how to generate memories for the Memory Bank.""", + ) + similarity_search_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig + ] = 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`.""", + ) + 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.""", + ) -_CustomJobParametersOrDict = Union[_CustomJobParameters, _CustomJobParametersDict] +class ReasoningEngineContextSpecMemoryBankConfigDict(TypedDict, total=False): + """Specification for a Memory Bank.""" -class _GetCustomJobParameters(_common.BaseModel): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" + customization_configs: Optional[list[MemoryBankCustomizationConfigDict]] + """Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[VertexBaseConfig] = Field(default=None, description="""""") + 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.""" -class _GetCustomJobParametersDict(TypedDict, total=False): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" + 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`.""" - name: Optional[str] - """""" + 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.""" - config: Optional[VertexBaseConfigDict] - """""" + structured_memory_configs: Optional[list[StructuredMemoryConfigDict]] + """Optional. Configuration for organizing structured memories for a particular scope.""" -_GetCustomJobParametersOrDict = Union[ - _GetCustomJobParameters, _GetCustomJobParametersDict +ReasoningEngineContextSpecMemoryBankConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfig, + ReasoningEngineContextSpecMemoryBankConfigDict, ] -class CancelQueryJobAgentEngineConfig(_common.BaseModel): - """Config for canceling async querying agent engines.""" +class ReasoningEngineContextSpec(_common.BaseModel): + """Configuration for how Agent Engine sub-resources should manage context.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - operation_name: Optional[str] = Field( + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfig] = Field( default=None, - description="""Name of the longrunning operation returned from run_query_job.""", + description="""Optional. Specification for a Memory Bank, which manages memories for the Agent Engine.""", ) -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 ReasoningEngineContextSpecDict(TypedDict, total=False): + """Configuration for how Agent Engine sub-resources should manage context.""" - operation_name: Optional[str] - """Name of the longrunning operation returned from run_query_job.""" + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfigDict] + """Optional. Specification for a Memory Bank, which manages memories for the Agent Engine.""" -CancelQueryJobAgentEngineConfigOrDict = Union[ - CancelQueryJobAgentEngineConfig, CancelQueryJobAgentEngineConfigDict +ReasoningEngineContextSpecOrDict = Union[ + ReasoningEngineContextSpec, ReasoningEngineContextSpecDict ] -class _CancelQueryJobAgentEngineRequestParameters(_common.BaseModel): - """Parameters for canceling async querying agent engines.""" +class SecretRef(_common.BaseModel): + """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" - name: Optional[str] = Field( - default=None, description="""Name of the reasoning engine resource.""" + secret: Optional[str] = Field( + default=None, + description="""Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""", ) - config: Optional[CancelQueryJobAgentEngineConfig] = Field( - default=None, description="""""" + version: 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.""", ) -class _CancelQueryJobAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for canceling async querying agent engines.""" +class SecretRefDict(TypedDict, total=False): + """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" - name: Optional[str] - """Name of the reasoning engine resource.""" + secret: Optional[str] + """Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""" - config: Optional[CancelQueryJobAgentEngineConfigDict] - """""" + 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.""" -_CancelQueryJobAgentEngineRequestParametersOrDict = Union[ - _CancelQueryJobAgentEngineRequestParameters, - _CancelQueryJobAgentEngineRequestParametersDict, -] +SecretRefOrDict = Union[SecretRef, SecretRefDict] -class CancelQueryJobResult(_common.BaseModel): - """Result of canceling a query job.""" +class SecretEnvVar(_common.BaseModel): + """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + name: Optional[str] = Field( + default=None, + description="""Required. Name of the secret environment variable.""", + ) + secret_ref: Optional[SecretRef] = Field( + default=None, + description="""Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""", ) -class CancelQueryJobResultDict(TypedDict, total=False): - """Result of canceling a query job.""" +class SecretEnvVarDict(TypedDict, total=False): + """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + name: Optional[str] + """Required. Name of the secret environment variable.""" + secret_ref: Optional[SecretRefDict] + """Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" -CancelQueryJobResultOrDict = Union[CancelQueryJobResult, CancelQueryJobResultDict] +SecretEnvVarOrDict = Union[SecretEnvVar, SecretEnvVarDict] -class CheckQueryJobAgentEngineConfig(_common.BaseModel): - """Config for async querying agent engines.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - retrieve_result: Optional[bool] = Field( +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig( + _common.BaseModel +): + """Configuration for traffic originating from a Reasoning Engine.""" + + agent_gateway: Optional[str] = Field( default=None, - description="""Whether to retrieve the results of the query job.""", + 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 CheckQueryJobAgentEngineConfigDict(TypedDict, total=False): - """Config for async querying agent engines.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict( + TypedDict, total=False +): + """Configuration for traffic originating from a Reasoning Engine.""" - retrieve_result: Optional[bool] - """Whether to retrieve the results of the query job.""" + 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}`""" -CheckQueryJobAgentEngineConfigOrDict = Union[ - CheckQueryJobAgentEngineConfig, CheckQueryJobAgentEngineConfigDict +ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigOrDict = Union[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig, + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict, ] -class _CheckQueryJobAgentEngineRequestParameters(_common.BaseModel): - """Parameters for async querying agent engines.""" +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig( + _common.BaseModel +): + """Configuration for traffic targeting a Reasoning Engine.""" - name: Optional[str] = Field(default=None, description="""Name of the query job.""") - config: Optional[CheckQueryJobAgentEngineConfig] = Field( - default=None, description="""""" + agent_gateway: Optional[str] = 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}`""", ) -class _CheckQueryJobAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for async querying agent engines.""" - - name: Optional[str] - """Name of the query job.""" +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict( + TypedDict, total=False +): + """Configuration for traffic targeting a Reasoning Engine.""" - config: Optional[CheckQueryJobAgentEngineConfigDict] - """""" + 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}`""" -_CheckQueryJobAgentEngineRequestParametersOrDict = Union[ - _CheckQueryJobAgentEngineRequestParameters, - _CheckQueryJobAgentEngineRequestParametersDict, +ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigOrDict = Union[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig, + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict, ] -class CheckQueryJobResult(_common.BaseModel): - """Result of checking a query job.""" +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfig(_common.BaseModel): + """Agent Gateway configuration for a Reasoning Engine deployment.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - operation_name: Optional[str] = Field( - default=None, description="""Name of the agent engine operation.""" - ) - output_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the output file.""" - ) - status: Optional[str] = Field( - default=None, description="""Status of the operation.""" + agent_to_anywhere_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig + ] = Field( + default=None, + description="""Optional. Configuration for traffic originating from the Reasoning Engine. When unset, outgoing traffic is not routed through an Agent Gateway.""", ) - result: Optional[str] = Field( - default=None, description="""JSON result of the operation.""" + client_to_agent_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig + ] = Field( + default=None, + description="""Optional. Configuration for traffic targeting the Reasoning Engine. When unset, incoming traffic is not routed through an Agent Gateway.""", ) -class CheckQueryJobResultDict(TypedDict, total=False): - """Result of checking a query job.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - operation_name: Optional[str] - """Name of the agent engine operation.""" - - output_gcs_uri: Optional[str] - """The GCS URI of the output file.""" +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict(TypedDict, total=False): + """Agent Gateway configuration for a Reasoning Engine deployment.""" - status: Optional[str] - """Status of the operation.""" + 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.""" - result: Optional[str] - """JSON result of the operation.""" + 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.""" -CheckQueryJobResultOrDict = Union[CheckQueryJobResult, CheckQueryJobResultDict] +ReasoningEngineSpecDeploymentSpecAgentGatewayConfigOrDict = Union[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig, + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict, +] -class _RunQueryJobAgentEngineConfig(_common.BaseModel): - """Config for running a query job on an agent engine.""" +class KeepAliveProbeHttpGet(_common.BaseModel): + """Specifies the HTTP GET configuration for the probe.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - input_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the input file.""" + path: Optional[str] = Field( + default=None, + description="""Required. Specifies the path of the HTTP GET request (e.g., `"/is_busy"`).""", ) - output_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the output file.""" + port: Optional[int] = Field( + default=None, + description="""Optional. Specifies the port number on the container to which the request is sent.""", ) -class _RunQueryJobAgentEngineConfigDict(TypedDict, total=False): - """Config for running a query job on an agent engine.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class KeepAliveProbeHttpGetDict(TypedDict, total=False): + """Specifies the HTTP GET configuration for the probe.""" - input_gcs_uri: Optional[str] - """The GCS URI of the input file.""" + path: Optional[str] + """Required. Specifies the path of the HTTP GET request (e.g., `"/is_busy"`).""" - output_gcs_uri: Optional[str] - """The GCS URI of the output file.""" + port: Optional[int] + """Optional. Specifies the port number on the container to which the request is sent.""" -_RunQueryJobAgentEngineConfigOrDict = Union[ - _RunQueryJobAgentEngineConfig, _RunQueryJobAgentEngineConfigDict -] +KeepAliveProbeHttpGetOrDict = Union[KeepAliveProbeHttpGet, KeepAliveProbeHttpGetDict] -class _RunQueryJobAgentEngineRequestParameters(_common.BaseModel): - """Parameters for running a query job on an agent engine.""" +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.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + http_get: Optional[KeepAliveProbeHttpGet] = Field( + default=None, + description="""Optional. Specifies the HTTP GET configuration for the probe.""", ) - config: Optional[_RunQueryJobAgentEngineConfig] = Field( - default=None, description="""""" + max_seconds: Optional[int] = 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).""", ) -class _RunQueryJobAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for running a query job on an agent engine.""" +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.""" - name: Optional[str] - """Name of the agent engine.""" + http_get: Optional[KeepAliveProbeHttpGetDict] + """Optional. Specifies the HTTP GET configuration for the probe.""" - config: Optional[_RunQueryJobAgentEngineConfigDict] - """""" + 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).""" -_RunQueryJobAgentEngineRequestParametersOrDict = Union[ - _RunQueryJobAgentEngineRequestParameters, - _RunQueryJobAgentEngineRequestParametersDict, -] +KeepAliveProbeOrDict = Union[KeepAliveProbe, KeepAliveProbeDict] -class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent( - _common.BaseModel -): - """The conversation source event for generating memories.""" +class ReasoningEngineSpecDeploymentSpec(_common.BaseModel): + """The specification of a Reasoning Engine deployment.""" - content: Optional[genai_types.Content] = Field( - default=None, description="""Required. Represents the content of the event.""" + 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( + 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].""", + ) + min_instances: Optional[int] = 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.""" + ) + resource_limits: Optional[dict[str, 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""", + ) + secret_env: Optional[list[SecretEnvVar]] = 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.""", + ) + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + ] = Field( + default=None, + description="""Optional. Agent Gateway configuration for the 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.""", ) -class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict( - TypedDict, total=False -): - """The conversation source event for generating memories.""" +class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False): + """The specification of a Reasoning Engine deployment.""" - content: Optional[genai_types.Content] - """Required. Represents the content of the event.""" + agent_server_mode: Optional[AgentServerMode] + """The agent server mode.""" + container_concurrency: Optional[int] + """Optional. Concurrency for each container and agent server. Recommended value: 2 * cpu + 1. Defaults to 9.""" -MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventOrDict = ( - Union[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent, - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict, - ] -) + 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.""" + 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].""" -class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource( - _common.BaseModel -): - """A conversation source for the example. This is similar to `DirectContentsSource`.""" + 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].""" - events: Optional[ - list[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent - ] - ] = Field( - default=None, - description="""Optional. Represents the input conversation events for the example.""", - ) + 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""" -class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict( - TypedDict, total=False -): - """A conversation source for the example. This is similar to `DirectContentsSource`.""" + 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.""" - events: Optional[ - list[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict - ] + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict ] - """Optional. Represents the input conversation events for the example.""" + """Optional. Agent Gateway configuration for the 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.""" -MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceOrDict = Union[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource, - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict, +ReasoningEngineSpecDeploymentSpecOrDict = Union[ + ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict ] -class MemoryTopicId(_common.BaseModel): - """The topic ID for a memory.""" +class ReasoningEngineSpecPackageSpec(_common.BaseModel): + """User-provided package specification, containing pickled object and package requirements.""" - custom_memory_topic_label: Optional[str] = Field( + dependency_files_gcs_uri: Optional[str] = Field( default=None, - description="""Optional. Represents the custom memory topic label.""", + description="""Optional. The Cloud Storage URI of the dependency files in tar.gz format.""", ) - managed_memory_topic: Optional[ManagedTopicEnum] = Field( - default=None, description="""Optional. Represents the managed memory topic.""" + pickle_object_gcs_uri: Optional[str] = Field( + default=None, + 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 MemoryTopicIdDict(TypedDict, total=False): - """The topic ID for a memory.""" +class ReasoningEngineSpecPackageSpecDict(TypedDict, total=False): + """User-provided package specification, containing pickled object and package requirements.""" - custom_memory_topic_label: Optional[str] - """Optional. Represents the custom memory topic label.""" + dependency_files_gcs_uri: Optional[str] + """Optional. The Cloud Storage URI of the dependency files in tar.gz format.""" - managed_memory_topic: Optional[ManagedTopicEnum] - """Optional. Represents the managed memory topic.""" + pickle_object_gcs_uri: Optional[str] + """Optional. The Cloud Storage URI of the pickled python object.""" + 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.""" -MemoryTopicIdOrDict = Union[MemoryTopicId, MemoryTopicIdDict] + requirements_gcs_uri: Optional[str] + """Optional. The Cloud Storage URI of the `requirements.txt` file""" -class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory( - _common.BaseModel -): - """A memory generated by the operation.""" +ReasoningEngineSpecPackageSpecOrDict = Union[ + ReasoningEngineSpecPackageSpec, ReasoningEngineSpecPackageSpecDict +] - fact: Optional[str] = Field( - default=None, - description="""Required. Represents the fact to generate a memory from.""", - ) - topics: Optional[list[MemoryTopicId]] = Field( + +class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig(_common.BaseModel): + """Configuration for the Agent Development Kit (ADK).""" + + json_config: Optional[dict[str, Any]] = 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="""Required. The value of the ADK config in JSON format.""", ) -class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict( +class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict( TypedDict, total=False ): - """A memory generated by the operation.""" + """Configuration for the Agent Development Kit (ADK).""" - fact: Optional[str] - """Required. Represents the fact to generate a memory from.""" + json_config: Optional[dict[str, Any]] + """Required. The value of the ADK config in JSON format.""" - 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`.""" +ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig, + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict, +] -MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryOrDict = Union[ - MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory, - MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict, + +class ReasoningEngineSpecSourceCodeSpecInlineSource(_common.BaseModel): + """Specifies source code provided as a byte stream.""" + + 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.""", + ) + + +class ReasoningEngineSpecSourceCodeSpecInlineSourceDict(TypedDict, total=False): + """Specifies source code provided as a byte stream.""" + + source_archive: Optional[bytes] + """Required. Input only. The application source code archive. It must be a compressed tarball (.tar.gz) file.""" + + +ReasoningEngineSpecSourceCodeSpecInlineSourceOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecInlineSource, + ReasoningEngineSpecSourceCodeSpecInlineSourceDict, ] -class MemoryBankCustomizationConfigGenerateMemoriesExample(_common.BaseModel): - """An example of how to generate memories for a particular scope.""" +class ReasoningEngineSpecSourceCodeSpecAgentConfigSource(_common.BaseModel): + """Specification for the deploying from agent config.""" - conversation_source: Optional[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource - ] = Field(default=None, description="""A conversation source for the example.""") - generated_memories: Optional[ - list[MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory] - ] = Field( + adk_config: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig + ] = Field(default=None, description="""Required. The ADK configuration.""") + inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSource] = 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. 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 MemoryBankCustomizationConfigGenerateMemoriesExampleDict(TypedDict, total=False): - """An example of how to generate memories for a particular scope.""" +class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict(TypedDict, total=False): + """Specification for the deploying from agent config.""" - conversation_source: Optional[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict + adk_config: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict ] - """A conversation source for the example.""" + """Required. The ADK configuration.""" - 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.""" + 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.""" -MemoryBankCustomizationConfigGenerateMemoriesExampleOrDict = Union[ - MemoryBankCustomizationConfigGenerateMemoriesExample, - MemoryBankCustomizationConfigGenerateMemoriesExampleDict, +ReasoningEngineSpecSourceCodeSpecAgentConfigSourceOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSource, + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict, ] -class MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic(_common.BaseModel): - """A custom memory topic defined by the developer.""" +class ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig(_common.BaseModel): + """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. - label: Optional[str] = Field( - default=None, description="""Required. Represents the label of the topic.""" + This includes the repository, revision, and directory to use. + """ + + 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}`.""", ) - description: Optional[str] = Field( + dir: 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. 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 MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict( +class ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict( TypedDict, total=False ): - """A custom memory topic defined by the developer.""" + """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. - label: Optional[str] - """Required. Represents the label of the topic.""" + This includes the repository, revision, and directory to use. + """ - description: Optional[str] - """Required. Represents the description of the memory topic. This should explain what information should be extracted for this topic.""" + 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.""" -MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicOrDict = Union[ - MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic, - MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict, + 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 MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic(_common.BaseModel): - """A managed memory topic defined by the system.""" +class ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource(_common.BaseModel): + """Specifies source code to be fetched from a Git repository managed through the Developer Connect service.""" - managed_topic_enum: Optional[ManagedTopicEnum] = Field( - default=None, description="""Required. Represents the managed topic.""" + config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig] = Field( + default=None, + description="""Required. The Developer Connect configuration that defines the specific repository, revision, and directory to use as the source code root.""", ) -class MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict( +class ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict( TypedDict, total=False ): - """A managed memory topic defined by the system.""" + """Specifies source code to be fetched from a Git repository managed through the Developer Connect service.""" - managed_topic_enum: Optional[ManagedTopicEnum] - """Required. Represents the managed topic.""" + config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict] + """Required. The Developer Connect configuration that defines the specific repository, revision, and directory to use as the source code root.""" -MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicOrDict = Union[ - MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic, - MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict, +ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource, + ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict, ] -class MemoryBankCustomizationConfigMemoryTopic(_common.BaseModel): - """A topic of information that should be extracted from conversations and stored as memories.""" +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.""" - custom_memory_topic: Optional[ - MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic - ] = Field( - default=None, description="""A custom memory topic defined by the developer.""" - ) - managed_memory_topic: Optional[ - MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic - ] = Field( - default=None, description="""A managed memory topic defined by Memory Bank.""" + build_args: Optional[dict[str, str]] = Field( + default=None, + description="""Optional. Build arguments to be used. They will be passed through --build-arg flags.""", ) -class MemoryBankCustomizationConfigMemoryTopicDict(TypedDict, total=False): - """A topic of information that should be extracted from conversations and stored as memories.""" +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.""" - custom_memory_topic: Optional[ - MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict - ] - """A custom memory topic defined by the developer.""" + build_args: Optional[dict[str, str]] + """Optional. Build arguments to be used. They will be passed through --build-arg flags.""" - managed_memory_topic: Optional[ - MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict - ] - """A managed memory topic defined by Memory Bank.""" - - -MemoryBankCustomizationConfigMemoryTopicOrDict = Union[ - MemoryBankCustomizationConfigMemoryTopic, - MemoryBankCustomizationConfigMemoryTopicDict, -] - - -class MemoryBankCustomizationConfigConsolidationConfig(_common.BaseModel): - """Represents configuration for customizing how memories are consolidated.""" - - 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 MemoryBankCustomizationConfigConsolidationConfigDict(TypedDict, total=False): - """Represents configuration for customizing how memories are consolidated.""" - - 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.""" - -MemoryBankCustomizationConfigConsolidationConfigOrDict = Union[ - MemoryBankCustomizationConfigConsolidationConfig, - MemoryBankCustomizationConfigConsolidationConfigDict, +ReasoningEngineSpecSourceCodeSpecImageSpecOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecImageSpec, + ReasoningEngineSpecSourceCodeSpecImageSpecDict, ] -class MemoryBankCustomizationConfig(_common.BaseModel): - """Represents configuration for organizing natural language memories.""" +class ReasoningEngineSpecSourceCodeSpecPythonSpec(_common.BaseModel): + """Specification for running a Python application from source.""" - 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( + entrypoint_module: Optional[str] = Field( default=None, - description="""Optional. Provides examples of how to generate memories for a particular scope.""", + 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`.""", ) - memory_topics: Optional[list[MemoryBankCustomizationConfigMemoryTopic]] = Field( + entrypoint_object: Optional[str] = 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.""", + 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`.""", ) - scope_keys: Optional[list[str]] = Field( + requirements_file: 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="""Optional. The path to the requirements file, relative to the source root. If not specified, defaults to "requirements.txt".""", ) - 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="""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.""", ) -class MemoryBankCustomizationConfigDict(TypedDict, total=False): - """Represents configuration for organizing natural language memories.""" - - 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.""" +class ReasoningEngineSpecSourceCodeSpecPythonSpecDict(TypedDict, total=False): + """Specification for running a Python application from source.""" - 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.""" + 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`.""" - 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.""" + 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`.""" - consolidation_config: Optional[MemoryBankCustomizationConfigConsolidationConfigDict] - """Optional. Represents configuration for customizing how memories are consolidated together.""" + requirements_file: Optional[str] + """Optional. The path to the requirements file, relative to the source root. If not specified, defaults to "requirements.txt".""" - 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] + """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.""" -MemoryBankCustomizationConfigOrDict = Union[ - MemoryBankCustomizationConfig, MemoryBankCustomizationConfigDict +ReasoningEngineSpecSourceCodeSpecPythonSpecOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecPythonSpec, + ReasoningEngineSpecSourceCodeSpecPythonSpecDict, ] -class MemoryGenerationTriggerConfigGenerationTriggerRule(_common.BaseModel): - """Represents the active rule that determines when to flush the buffer.""" +class ReasoningEngineSpecSourceCodeSpec(_common.BaseModel): + """Specification for deploying from source code.""" - event_count: Optional[int] = Field( - default=None, - description="""Optional. Specifies to trigger generation when the event count reaches this limit.""", + agent_config_source: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSource + ] = Field( + default=None, description="""Source code is generated from the agent config.""" ) - fixed_interval: Optional[str] = Field( + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource + ] = Field( default=None, - description="""Optional. Specifies to trigger generation at a fixed interval. The duration must have a minute-level granularity.""", + description="""Source code is in a Git repository managed by Developer Connect.""", ) - idle_duration: Optional[str] = Field( + image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpec] = 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="""Optional. Configuration for building an image with custom config file.""", ) - overlap_event_count: Optional[int] = Field( - default=None, - description="""Optional. Re-include the last N already-processed events in the next window.""", + 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 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 ReasoningEngineSpecSourceCodeSpecDict(TypedDict, total=False): + """Specification for deploying from source code.""" -class MemoryGenerationTriggerConfig(_common.BaseModel): - """The configuration for triggering memory generation for ingested events.""" + agent_config_source: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict + ] + """Source code is generated from the agent config.""" - 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.""", - ) - ) + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict + ] + """Source code is in a Git repository managed by Developer Connect.""" + image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict] + """Optional. Configuration for building an image with custom config file.""" -class MemoryGenerationTriggerConfigDict(TypedDict, total=False): - """The configuration for triggering memory generation for ingested events.""" + inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSourceDict] + """Source code is provided directly in the request.""" - 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.""" + python_spec: Optional[ReasoningEngineSpecSourceCodeSpecPythonSpecDict] + """Configuration for a Python application.""" -MemoryGenerationTriggerConfigOrDict = Union[ - MemoryGenerationTriggerConfig, MemoryGenerationTriggerConfigDict +ReasoningEngineSpecSourceCodeSpecOrDict = Union[ + ReasoningEngineSpecSourceCodeSpec, ReasoningEngineSpecSourceCodeSpecDict ] -class ReasoningEngineContextSpecMemoryBankConfigGenerationConfig(_common.BaseModel): - """Configuration for how to generate memories.""" +class ReasoningEngineSpecContainerSpec(_common.BaseModel): + """Specification for deploying from a container image.""" - model: Optional[str] = Field( + image_uri: Optional[str] = Field( default=None, - description="""Optional. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", + 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.""", ) - generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + port: Optional[int] = Field( default=None, - description="""Optional. Specifies the default trigger configuration for generating memories using `IngestEvents`.""", + description="""Optional. The port the container listens on. Defaults to 8080 if unset.""", ) -class ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict( - TypedDict, total=False -): - """Configuration for how to generate memories.""" +class ReasoningEngineSpecContainerSpecDict(TypedDict, total=False): + """Specification for deploying from a container image.""" - model: Optional[str] - """Optional. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + 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.""" - generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] - """Optional. Specifies the default trigger configuration for generating memories using `IngestEvents`.""" + port: Optional[int] + """Optional. The port the container listens on. Defaults to 8080 if unset.""" -ReasoningEngineContextSpecMemoryBankConfigGenerationConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigGenerationConfig, - ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict, +ReasoningEngineSpecContainerSpecOrDict = Union[ + ReasoningEngineSpecContainerSpec, ReasoningEngineSpecContainerSpecDict ] -class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig( - _common.BaseModel -): - """Configuration for how to perform similarity search on memories.""" +class ReasoningEngineSpecBuildSpec(_common.BaseModel): + """Specification for building container image.""" - embedding_model: Optional[str] = Field( + service_account: 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="""Optional. The service account that Cloud Build uses to run the build. This field is only applicable when `worker_pool` is specified (i.e., for custom worker pools). If `worker_pool` is not specified, this field is ignored and the build runs using the Google-managed service agent. Format: `projects/{project}/serviceAccounts/{service_account}` or `{service_account}@{project}.iam.gserviceaccount.com`""", + ) + worker_pool: Optional[str] = Field( + default=None, + description="""Optional. Identifier. The resource name of the Cloud Build WorkerPool to use for the build. Format: `projects/{project}/locations/{location}/workerPools/{worker_pool}`""", ) -class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict( - TypedDict, total=False -): - """Configuration for how to perform similarity search on memories.""" +class ReasoningEngineSpecBuildSpecDict(TypedDict, total=False): + """Specification for building container image.""" - embedding_model: Optional[str] - """Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + service_account: Optional[str] + """Optional. The service account that Cloud Build uses to run the build. This field is only applicable when `worker_pool` is specified (i.e., for custom worker pools). If `worker_pool` is not specified, this field is ignored and the build runs using the Google-managed service agent. Format: `projects/{project}/serviceAccounts/{service_account}` or `{service_account}@{project}.iam.gserviceaccount.com`""" + worker_pool: Optional[str] + """Optional. Identifier. The resource name of the Cloud Build WorkerPool to use for the build. Format: `projects/{project}/locations/{location}/workerPools/{worker_pool}`""" -ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig, - ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict, + +ReasoningEngineSpecBuildSpecOrDict = Union[ + ReasoningEngineSpecBuildSpec, ReasoningEngineSpecBuildSpecDict ] -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 ReasoningEngineSpec(_common.BaseModel): + """The specification of an agent runtime.""" - create_ttl: Optional[str] = Field( + agent_card: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. The TTL duration for memories uploaded via CreateMemory.""", + 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.""", ) - generate_created_ttl: Optional[str] = Field( + agent_framework: Optional[str] = Field( default=None, - description="""Optional. The TTL duration for memories newly generated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.CREATED).""", + description="""Optional. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom".""", ) - generate_updated_ttl: Optional[str] = Field( + class_methods: Optional[list[dict[str, Any]]] = 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).""", + description="""Optional. Declarations for object class methods in OpenAPI specification format.""", + ) + deployment_spec: Optional[ReasoningEngineSpecDeploymentSpec] = Field( + default=None, + 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.""", + ) + build_spec: Optional[ReasoningEngineSpecBuildSpec] = Field( + default=None, + description="""Optional. Configuration for building container image.""", ) -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.""" +class ReasoningEngineSpecDict(TypedDict, total=False): + """The specification of an agent runtime.""" - create_ttl: Optional[str] - """Optional. The TTL duration for memories uploaded via CreateMemory.""" + 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.""" - generate_created_ttl: Optional[str] - """Optional. The TTL duration for memories newly generated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.CREATED).""" + 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".""" - 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).""" + class_methods: Optional[list[dict[str, Any]]] + """Optional. Declarations for object class methods in OpenAPI specification format.""" + deployment_spec: Optional[ReasoningEngineSpecDeploymentSpecDict] + """Optional. The specification of a Reasoning Engine deployment.""" -ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig, - ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict, -] + 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)""" + 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.""" -class ReasoningEngineContextSpecMemoryBankConfigTtlConfig(_common.BaseModel): - """Configuration for automatically setting the TTL ("time-to-live") of the memories in the Memory Bank.""" + 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.""" - 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. 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.""", - ) + 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.""" + source_code_spec: Optional[ReasoningEngineSpecSourceCodeSpecDict] + """Deploy from source code files with a defined entrypoint.""" -class ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict(TypedDict, total=False): - """Configuration for automatically setting the TTL ("time-to-live") of the memories in the Memory Bank.""" + container_spec: Optional[ReasoningEngineSpecContainerSpecDict] + """Deploy from a container image with a defined entrypoint and commands.""" - 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.""" + build_spec: Optional[ReasoningEngineSpecBuildSpecDict] + """Optional. Configuration for building container image.""" - granular_ttl_config: Optional[ - ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict - ] - """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.""" +ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict] -ReasoningEngineContextSpecMemoryBankConfigTtlConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigTtlConfig, - ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict, +class ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest(_common.BaseModel): + """Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" + + pass + + +class ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict(TypedDict, total=False): + """Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" + + pass + + +ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestOrDict = Union[ + ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest, + ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict, ] -class StructuredMemorySchemaConfig(_common.BaseModel): - """Represents the OpenAPI schema of the structured memories.""" +class ReasoningEngineTrafficConfigTrafficSplitManualTarget(_common.BaseModel): + """A single target for the traffic split, specifying a Runtime Revision and the percentage of traffic to send to it.""" - memory_schema: Optional[genai_types.Schema] = Field( - default=None, - description="""Required. Represents the OpenAPI schema of the structured memories.""", - ) - id: Optional[str] = Field( + percent: Optional[int] = 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="""Required. Specifies percent of the traffic to this Runtime Revision.""", ) - memory_type: Optional[MemoryType] = Field( + runtime_revision_name: Optional[str] = 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="""Required. The Runtime Revision name to which to send this portion of traffic, if traffic allocation is by Runtime Revision.""", ) -class StructuredMemorySchemaConfigDict(TypedDict, total=False): - """Represents the OpenAPI schema of the structured memories.""" +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.""" - memory_schema: Optional[genai_types.Schema] - """Required. Represents the OpenAPI schema of the structured memories.""" + percent: Optional[int] + """Required. Specifies percent of the traffic to this Runtime Revision.""" - 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.""" + 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.""" - 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.""" + +ReasoningEngineTrafficConfigTrafficSplitManualTargetOrDict = Union[ + ReasoningEngineTrafficConfigTrafficSplitManualTarget, + ReasoningEngineTrafficConfigTrafficSplitManualTargetDict, +] -StructuredMemorySchemaConfigOrDict = Union[ - StructuredMemorySchemaConfig, StructuredMemorySchemaConfigDict +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.""" + + 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 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.""" + + targets: Optional[list[ReasoningEngineTrafficConfigTrafficSplitManualTargetDict]] + """A list of traffic targets for the Runtimes Revisions. The sum of percentages must equal to 100.""" + + +ReasoningEngineTrafficConfigTrafficSplitManualOrDict = Union[ + ReasoningEngineTrafficConfigTrafficSplitManual, + ReasoningEngineTrafficConfigTrafficSplitManualDict, ] -class StructuredMemoryConfig(_common.BaseModel): - """Configuration for organizing structured memories within a scope.""" +class ReasoningEngineTrafficConfig(_common.BaseModel): + """Traffic distribution configuration.""" - schema_configs: Optional[list[StructuredMemorySchemaConfig]] = Field( + traffic_split_always_latest: Optional[ + ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest + ] = Field( default=None, - description="""Optional. Represents configuration of the structured memories' schemas.""", + description="""Optional. Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""", ) - 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.""", + 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 StructuredMemoryConfigDict(TypedDict, total=False): - """Configuration for organizing structured memories within a scope.""" +class ReasoningEngineTrafficConfigDict(TypedDict, total=False): + """Traffic distribution configuration.""" - schema_configs: Optional[list[StructuredMemorySchemaConfigDict]] - """Optional. Represents configuration of the structured memories' schemas.""" + traffic_split_always_latest: Optional[ + ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict + ] + """Optional. Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" - 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.""" + 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.""" -StructuredMemoryConfigOrDict = Union[StructuredMemoryConfig, StructuredMemoryConfigDict] +ReasoningEngineTrafficConfigOrDict = Union[ + ReasoningEngineTrafficConfig, ReasoningEngineTrafficConfigDict +] -class ReasoningEngineContextSpecMemoryBankConfig(_common.BaseModel): - """Specification for a Memory Bank.""" +class ReasoningEngine(_common.BaseModel): + """An agent runtime.""" - customization_configs: Optional[list[MemoryBankCustomizationConfig]] = Field( + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, - description="""Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""", + 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.""", ) - disable_memory_revisions: Optional[bool] = Field( + context_spec: Optional[ReasoningEngineContextSpec] = Field( default=None, - description="""If true, no memory revisions will be created for any requests to the Memory Bank.""", + description="""Optional. Configuration for how Agent Engine sub-resources should manage context.""", ) - generation_config: Optional[ - ReasoningEngineContextSpecMemoryBankConfigGenerationConfig - ] = Field( + create_time: Optional[datetime.datetime] = Field( default=None, - description="""Optional. Configuration for how to generate memories for the Memory Bank.""", + description="""Output only. Timestamp when this ReasoningEngine was created.""", ) - similarity_search_config: Optional[ - ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig - ] = Field( + description: Optional[str] = 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. The description of the ReasoningEngine.""", ) - ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfig] = Field( + display_name: Optional[str] = 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="""Required. The display name of the ReasoningEngine.""", ) - structured_memory_configs: Optional[list[StructuredMemoryConfig]] = Field( + etag: Optional[str] = Field( default=None, - description="""Optional. Configuration for organizing structured memories for a particular scope.""", + description="""Optional. Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""", ) - - -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`.""" - - 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.""" - - structured_memory_configs: Optional[list[StructuredMemoryConfigDict]] - """Optional. Configuration for organizing structured memories for a particular scope.""" - - -ReasoningEngineContextSpecMemoryBankConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfig, - ReasoningEngineContextSpecMemoryBankConfigDict, -] - - -class ReasoningEngineContextSpec(_common.BaseModel): - """Configuration for how Agent Engine sub-resources should manage context.""" - - memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfig] = Field( + labels: Optional[dict[str, str]] = Field( + default=None, description="""Labels for the ReasoningEngine.""" + ) + name: Optional[str] = Field( default=None, - description="""Optional. Specification for a Memory Bank, which manages memories for the Agent Engine.""", + 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 ReasoningEngineContextSpecDict(TypedDict, total=False): - """Configuration for how Agent Engine sub-resources should manage context.""" +class ReasoningEngineDict(TypedDict, total=False): + """An agent runtime.""" - memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfigDict] - """Optional. Specification for a Memory Bank, which manages memories for the 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.""" -ReasoningEngineContextSpecOrDict = Union[ - ReasoningEngineContextSpec, ReasoningEngineContextSpecDict -] + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this ReasoningEngine was created.""" + description: Optional[str] + """Optional. The description of the ReasoningEngine.""" -class SecretRef(_common.BaseModel): - """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" + display_name: Optional[str] + """Required. The display name of the ReasoningEngine.""" - secret: Optional[str] = Field( - default=None, - description="""Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""", - ) - version: 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.""", - ) + 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 SecretRefDict(TypedDict, total=False): - """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" + name: Optional[str] + """Identifier. The resource name of the ReasoningEngine. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`""" - secret: Optional[str] - """Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""" + spec: Optional[ReasoningEngineSpecDict] + """Optional. Configurations of the ReasoningEngine""" - 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.""" + 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.""" -SecretRefOrDict = Union[SecretRef, SecretRefDict] +ReasoningEngineOrDict = Union[ReasoningEngine, ReasoningEngineDict] -class SecretEnvVar(_common.BaseModel): - """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" +class RuntimeOperation(_common.BaseModel): + """Operation that has an agent runtime as a response.""" name: Optional[str] = Field( default=None, - description="""Required. Name of the secret environment variable.""", + 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}`.""", ) - secret_ref: Optional[SecretRef] = Field( + metadata: 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="""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[ReasoningEngine] = Field( + default=None, description="""The created Agent Runtime.""" ) -class SecretEnvVarDict(TypedDict, total=False): - """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" +class RuntimeOperationDict(TypedDict, total=False): + """Operation that has an agent runtime as a response.""" name: Optional[str] - """Required. Name of the secret environment variable.""" - - secret_ref: Optional[SecretRefDict] - """Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" - - -SecretEnvVarOrDict = Union[SecretEnvVar, SecretEnvVarDict] - + """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 ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig( - _common.BaseModel -): - """Configuration for traffic originating from a Reasoning Engine.""" + 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.""" - agent_gateway: Optional[str] = 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}`""", - ) + 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 ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict( - TypedDict, total=False -): - """Configuration for traffic originating from a Reasoning Engine.""" + response: Optional[ReasoningEngineDict] + """The created Agent Runtime.""" - 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}`""" +RuntimeOperationOrDict = Union[RuntimeOperation, RuntimeOperationDict] -ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigOrDict = Union[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig, - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict, -] +class CreateRuntimeConfig(_common.BaseModel): + """Config for create agent runtime.""" -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig( - _common.BaseModel -): - """Configuration for traffic targeting a Reasoning Engine.""" + 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. - agent_gateway: Optional[str] = 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="""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. The context spec to be used for the Agent Runtime.""", ) - - -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict( - TypedDict, total=False -): - """Configuration for traffic targeting a Reasoning Engine.""" - - 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}`""" - - -ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigOrDict = Union[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig, - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict, -] - - -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfig(_common.BaseModel): - """Agent Gateway configuration for a Reasoning Engine deployment.""" - - agent_to_anywhere_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig - ] = Field( + psc_interface_config: Optional[PscInterfaceConfig] = 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="""Optional. The PSC interface config for PSC-I to be used for the + Agent Runtime.""", ) - client_to_agent_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig ] = Field( default=None, - description="""Optional. Configuration for traffic targeting the Reasoning Engine. When unset, incoming traffic is not routed through an Agent Gateway.""", + description="""Agent Gateway configuration for a Reasoning Engine deployment.""", ) - - -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict(TypedDict, total=False): - """Agent Gateway configuration for a Reasoning Engine deployment.""" - - 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.""" - - 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.""" - - -ReasoningEngineSpecDeploymentSpecAgentGatewayConfigOrDict = Union[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig, - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict, -] - - -class KeepAliveProbeHttpGet(_common.BaseModel): - """Specifies the HTTP GET configuration for the probe.""" - - path: Optional[str] = Field( + min_instances: Optional[int] = Field( default=None, - description="""Required. Specifies the path of the HTTP GET request (e.g., `"/is_busy"`).""", + description="""The minimum number of instances to run for the Agent Runtime. + Defaults to 1. Range: [0, 10]. + """, ) - port: Optional[int] = Field( + max_instances: Optional[int] = Field( default=None, - description="""Optional. Specifies the port number on the container to which the request is sent.""", + 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]. + """, ) - - -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"`).""" - - port: Optional[int] - """Optional. Specifies the port number on the container to which the request is sent.""" - - -KeepAliveProbeHttpGetOrDict = Union[KeepAliveProbeHttpGet, KeepAliveProbeHttpGetDict] - - -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.""" - - http_get: Optional[KeepAliveProbeHttpGet] = Field( + resource_limits: Optional[dict[str, str]] = Field( default=None, - description="""Optional. Specifies the HTTP GET configuration for the probe.""", - ) - max_seconds: Optional[int] = Field( + 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="""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="""The container concurrency to be used for the Agent Runtime. + Recommended value: 2 * cpu + 1. Defaults to 9. + """, ) - - -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.""" - - 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).""" - - -KeepAliveProbeOrDict = Union[KeepAliveProbe, KeepAliveProbeDict] - - -class ReasoningEngineSpecDeploymentSpec(_common.BaseModel): - """The specification of a Reasoning Engine deployment.""" - - agent_server_mode: Optional[AgentServerMode] = Field( - default=None, description="""The agent server mode.""" + 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.""", ) - container_concurrency: Optional[int] = Field( + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, - description="""Optional. Concurrency for each container and agent server. Recommended value: 2 * cpu + 1. Defaults to 9.""", + description="""The encryption spec to be used for the Agent Runtime.""", ) - env: Optional[list[EnvVar]] = Field( + 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. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API.""", + 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. + """, ) - max_instances: Optional[int] = Field( + source_packages: Optional[list[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="""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) + """, ) - min_instances: Optional[int] = Field( + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig + ] = 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].""", + 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.""", ) - psc_interface_config: Optional[PscInterfaceConfig] = Field( - default=None, description="""Optional. Configuration for PSC-I.""" + entrypoint_module: Optional[str] = Field( + default=None, + description="""The entrypoint module to be used for the Agent Runtime + This field only used when source_packages is specified.""", ) - resource_limits: Optional[dict[str, str]] = Field( + entrypoint_object: 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="""The entrypoint object to be used for the Agent Runtime. + This field only used when source_packages is specified.""", ) - secret_env: Optional[list[SecretEnvVar]] = Field( + requirements_file: 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="""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_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] ] = Field( default=None, - description="""Optional. Agent Gateway configuration for the Reasoning Engine deployment.""", + 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".""", ) - keep_alive_probe: Optional[KeepAliveProbe] = Field( + python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] = 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="""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 ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False): - """The specification of a Reasoning Engine deployment.""" - - agent_server_mode: Optional[AgentServerMode] - """The agent server mode.""" +class CreateRuntimeConfigDict(TypedDict, total=False): + """Config for create agent runtime.""" - container_concurrency: Optional[int] - """Optional. Concurrency for each container and agent server. Recommended value: 2 * cpu + 1. Defaults to 9.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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.""" + display_name: Optional[str] + """The user-defined name of the Agent Runtime. - 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].""" + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """ - 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].""" + description: Optional[str] + """The description of the Agent Runtime.""" - psc_interface_config: Optional[PscInterfaceConfigDict] - """Optional. Configuration for PSC-I.""" + spec: Optional[ReasoningEngineSpecDict] + """Optional. Configurations of the Agent Runtime.""" - 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""" + context_spec: Optional[ReasoningEngineContextSpecDict] + """Optional. The context spec to be used for the Agent Runtime.""" - 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.""" + psc_interface_config: Optional[PscInterfaceConfigDict] + """Optional. The PSC interface config for PSC-I to be used for the + Agent Runtime.""" agent_gateway_config: Optional[ ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict ] - """Optional. Agent Gateway configuration for the 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.""" + """Agent Gateway configuration for a Reasoning Engine deployment.""" + min_instances: Optional[int] + """The minimum number of instances to run for the Agent Runtime. + Defaults to 1. Range: [0, 10]. + """ -ReasoningEngineSpecDeploymentSpecOrDict = Union[ - ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict -] + 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]. + """ + 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'. + """ -class ReasoningEngineSpecPackageSpec(_common.BaseModel): - """User-provided package specification, containing pickled object and package requirements.""" + container_concurrency: Optional[int] + """The container concurrency to be used for the Agent Runtime. + Recommended value: 2 * cpu + 1. Defaults to 9. + """ - dependency_files_gcs_uri: Optional[str] = Field( - default=None, - description="""Optional. The Cloud Storage URI of the dependency files in tar.gz format.""", - ) - pickle_object_gcs_uri: Optional[str] = Field( - default=None, - 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""", - ) + 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 ReasoningEngineSpecPackageSpecDict(TypedDict, total=False): - """User-provided package specification, containing pickled object and package requirements.""" + labels: Optional[dict[str, str]] + """The labels to be used for the Agent Runtime.""" - dependency_files_gcs_uri: Optional[str] - """Optional. The Cloud Storage URI of the dependency files in tar.gz format.""" + 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. + """ - pickle_object_gcs_uri: Optional[str] - """Optional. The Cloud Storage URI of the pickled python object.""" + 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) + """ - 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.""" + 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.""" - requirements_gcs_uri: Optional[str] - """Optional. The Cloud Storage URI of the `requirements.txt` file""" + 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.""" -ReasoningEngineSpecPackageSpecOrDict = Union[ - ReasoningEngineSpecPackageSpec, ReasoningEngineSpecPackageSpecDict -] + 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".""" -class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig(_common.BaseModel): - """Configuration for the Agent Development Kit (ADK).""" + 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". + """ - json_config: Optional[dict[str, Any]] = Field( - default=None, - description="""Required. The value of the ADK config in JSON format.""", - ) + 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`. + """ -class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict( - TypedDict, total=False -): - """Configuration for the Agent Development Kit (ADK).""" +CreateRuntimeConfigOrDict = Union[CreateRuntimeConfig, CreateRuntimeConfigDict] - json_config: Optional[dict[str, Any]] - """Required. The value of the ADK config in JSON format.""" +class _CreateRuntimeRequestParameters(_common.BaseModel): + """Parameters for creating agent runtimes.""" -ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig, - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict, + config: Optional[CreateRuntimeConfig] = Field(default=None, description="""""") + + +class _CreateRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for creating agent runtimes.""" + + config: Optional[CreateRuntimeConfigDict] + """""" + + +_CreateRuntimeRequestParametersOrDict = Union[ + _CreateRuntimeRequestParameters, _CreateRuntimeRequestParametersDict ] -class ReasoningEngineSpecSourceCodeSpecInlineSource(_common.BaseModel): - """Specifies source code provided as a byte stream.""" +class DeleteRuntimeConfig(_common.BaseModel): + """Config for deleting agent runtime.""" - 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class ReasoningEngineSpecSourceCodeSpecInlineSourceDict(TypedDict, total=False): - """Specifies source code provided as a byte stream.""" +class DeleteRuntimeConfigDict(TypedDict, total=False): + """Config for deleting agent runtime.""" - source_archive: Optional[bytes] - """Required. Input only. The application source code archive. It must be a compressed tarball (.tar.gz) file.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ReasoningEngineSpecSourceCodeSpecInlineSourceOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecInlineSource, - ReasoningEngineSpecSourceCodeSpecInlineSourceDict, -] +DeleteRuntimeConfigOrDict = Union[DeleteRuntimeConfig, DeleteRuntimeConfigDict] -class ReasoningEngineSpecSourceCodeSpecAgentConfigSource(_common.BaseModel): - """Specification for the deploying from agent config.""" +class _DeleteRuntimeRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtimes.""" - adk_config: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig - ] = Field(default=None, description="""Required. The ADK configuration.""") - inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSource] = 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.""", + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime.""" + ) + 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 ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict(TypedDict, total=False): - """Specification for the deploying from agent config.""" +class _DeleteRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtimes.""" - adk_config: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict - ] - """Required. The ADK configuration.""" + name: Optional[str] + """Name of the agent runtime.""" - 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.""" + force: Optional[bool] + """If set to true, any child resources will also be deleted.""" + config: Optional[DeleteRuntimeConfigDict] + """""" -ReasoningEngineSpecSourceCodeSpecAgentConfigSourceOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSource, - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict, -] +_DeleteRuntimeRequestParametersOrDict = Union[ + _DeleteRuntimeRequestParameters, _DeleteRuntimeRequestParametersDict +] -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 DeleteRuntimeOperation(_common.BaseModel): + """Operation for deleting agent runtimes.""" - git_repository_link: Optional[str] = Field( + name: 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}`.""", + 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}`.""", ) - dir: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Required. Directory, relative to the source root, in which to run the build.""", + 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.""", ) - revision: Optional[str] = Field( + done: Optional[bool] = 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.""", + 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 ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict( - TypedDict, total=False -): - """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. +class DeleteRuntimeOperationDict(TypedDict, total=False): + """Operation for deleting agent runtimes.""" - This includes the repository, revision, and directory to use. - """ + 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}`.""" - 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}`.""" + 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.""" - dir: Optional[str] - """Required. Directory, relative to the source root, in which to run the build.""" + 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.""" - 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.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig, - ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict, -] +DeleteRuntimeOperationOrDict = Union[DeleteRuntimeOperation, DeleteRuntimeOperationDict] -class ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource(_common.BaseModel): - """Specifies source code to be fetched from a Git repository managed through the Developer Connect service.""" +class GetRuntimeConfig(_common.BaseModel): + """Config for create agent runtime.""" - config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig] = Field( - default=None, - description="""Required. The Developer Connect configuration that defines the specific repository, revision, and directory to use as the source code root.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict( - TypedDict, total=False -): - """Specifies source code to be fetched from a Git repository managed through the Developer Connect service.""" +class GetRuntimeConfigDict(TypedDict, total=False): + """Config for create agent runtime.""" - config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict] - """Required. The Developer Connect configuration that defines the specific repository, revision, and directory to use as the source code root.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource, - ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict, -] +GetRuntimeConfigOrDict = Union[GetRuntimeConfig, GetRuntimeConfigDict] -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 _GetRuntimeRequestParameters(_common.BaseModel): + """Parameters for getting agent runtimes.""" - build_args: Optional[dict[str, str]] = Field( - default=None, - description="""Optional. Build arguments to be used. They will be passed through --build-arg flags.""", + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime.""" ) + config: Optional[GetRuntimeConfig] = Field(default=None, description="""""") -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.""" +class _GetRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for getting agent runtimes.""" - build_args: Optional[dict[str, str]] - """Optional. Build arguments to be used. They will be passed through --build-arg flags.""" + name: Optional[str] + """Name of the agent runtime.""" + config: Optional[GetRuntimeConfigDict] + """""" -ReasoningEngineSpecSourceCodeSpecImageSpecOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecImageSpec, - ReasoningEngineSpecSourceCodeSpecImageSpecDict, + +_GetRuntimeRequestParametersOrDict = Union[ + _GetRuntimeRequestParameters, _GetRuntimeRequestParametersDict ] -class ReasoningEngineSpecSourceCodeSpecPythonSpec(_common.BaseModel): - """Specification for running a Python application from source.""" +class ListRuntimeConfig(_common.BaseModel): + """Config for listing agent runtimes.""" - entrypoint_module: Optional[str] = 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`.""", - ) - entrypoint_object: Optional[str] = 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`.""", - ) - requirements_file: 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".""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - version: 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="""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="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""", ) -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.""" +class ListRuntimeConfigDict(TypedDict, total=False): + """Config for listing agent runtimes.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ReasoningEngineSpecSourceCodeSpecPythonSpecOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecPythonSpec, - ReasoningEngineSpecSourceCodeSpecPythonSpecDict, -] + page_size: Optional[int] + """""" + page_token: Optional[str] + """""" -class ReasoningEngineSpecSourceCodeSpec(_common.BaseModel): - """Specification for deploying from source code.""" + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""" - 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="""Source code is in a Git repository managed by Developer Connect.""", - ) - image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpec] = Field( - default=None, - 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.""" - ) +ListRuntimeConfigOrDict = Union[ListRuntimeConfig, ListRuntimeConfigDict] -class ReasoningEngineSpecSourceCodeSpecDict(TypedDict, total=False): - """Specification for deploying from source code.""" - agent_config_source: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict - ] - """Source code is generated from the agent config.""" +class _ListRuntimeRequestParameters(_common.BaseModel): + """Parameters for listing agent runtimes.""" - developer_connect_source: Optional[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict - ] - """Source code is in a Git repository managed by Developer Connect.""" + config: Optional[ListRuntimeConfig] = Field(default=None, description="""""") - image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict] - """Optional. Configuration for building an image with custom config file.""" - inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSourceDict] - """Source code is provided directly in the request.""" +class _ListRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtimes.""" - python_spec: Optional[ReasoningEngineSpecSourceCodeSpecPythonSpecDict] - """Configuration for a Python application.""" + config: Optional[ListRuntimeConfigDict] + """""" -ReasoningEngineSpecSourceCodeSpecOrDict = Union[ - ReasoningEngineSpecSourceCodeSpec, ReasoningEngineSpecSourceCodeSpecDict +_ListRuntimeRequestParametersOrDict = Union[ + _ListRuntimeRequestParameters, _ListRuntimeRequestParametersDict ] -class ReasoningEngineSpecContainerSpec(_common.BaseModel): - """Specification for deploying from a container image.""" +class ListReasoningEnginesResponse(_common.BaseModel): + """Response for listing agent engines.""" - 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.""", + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - port: Optional[int] = Field( + next_page_token: Optional[str] = Field(default=None, description="""""") + reasoning_engines: Optional[list[ReasoningEngine]] = Field( default=None, - description="""Optional. The port the container listens on. Defaults to 8080 if unset.""", - ) + description="""List of agent engines. + """, + ) -class ReasoningEngineSpecContainerSpecDict(TypedDict, total=False): - """Specification for deploying from a container image.""" +class ListReasoningEnginesResponseDict(TypedDict, total=False): + """Response for listing agent engines.""" - 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.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - port: Optional[int] - """Optional. The port the container listens on. Defaults to 8080 if unset.""" + next_page_token: Optional[str] + """""" + + reasoning_engines: Optional[list[ReasoningEngineDict]] + """List of agent engines. + """ -ReasoningEngineSpecContainerSpecOrDict = Union[ - ReasoningEngineSpecContainerSpec, ReasoningEngineSpecContainerSpecDict +ListReasoningEnginesResponseOrDict = Union[ + ListReasoningEnginesResponse, ListReasoningEnginesResponseDict ] -class ReasoningEngineSpecBuildSpec(_common.BaseModel): - """Specification for building container image.""" +class GetRuntimeOperationConfig(_common.BaseModel): - service_account: Optional[str] = Field( - default=None, - description="""Optional. The service account that Cloud Build uses to run the build. This field is only applicable when `worker_pool` is specified (i.e., for custom worker pools). If `worker_pool` is not specified, this field is ignored and the build runs using the Google-managed service agent. Format: `projects/{project}/serviceAccounts/{service_account}` or `{service_account}@{project}.iam.gserviceaccount.com`""", - ) - worker_pool: Optional[str] = Field( - default=None, - description="""Optional. Identifier. The resource name of the Cloud Build WorkerPool to use for the build. Format: `projects/{project}/locations/{location}/workerPools/{worker_pool}`""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class ReasoningEngineSpecBuildSpecDict(TypedDict, total=False): - """Specification for building container image.""" - - service_account: Optional[str] - """Optional. The service account that Cloud Build uses to run the build. This field is only applicable when `worker_pool` is specified (i.e., for custom worker pools). If `worker_pool` is not specified, this field is ignored and the build runs using the Google-managed service agent. Format: `projects/{project}/serviceAccounts/{service_account}` or `{service_account}@{project}.iam.gserviceaccount.com`""" +class GetRuntimeOperationConfigDict(TypedDict, total=False): - worker_pool: Optional[str] - """Optional. Identifier. The resource name of the Cloud Build WorkerPool to use for the build. Format: `projects/{project}/locations/{location}/workerPools/{worker_pool}`""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ReasoningEngineSpecBuildSpecOrDict = Union[ - ReasoningEngineSpecBuildSpec, ReasoningEngineSpecBuildSpecDict +GetRuntimeOperationConfigOrDict = Union[ + GetRuntimeOperationConfig, GetRuntimeOperationConfigDict ] -class ReasoningEngineSpec(_common.BaseModel): - """The specification of an agent engine.""" +class _GetRuntimeOperationParameters(_common.BaseModel): + """Parameters for getting an operation with an agent runtime as a response.""" - agent_card: Optional[dict[str, Any]] = 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.""", - ) - agent_framework: Optional[str] = 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".""", - ) - class_methods: Optional[list[dict[str, Any]]] = Field( - default=None, - description="""Optional. Declarations for object class methods in OpenAPI specification format.""", - ) - deployment_spec: Optional[ReasoningEngineSpecDeploymentSpec] = Field( - default=None, - 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.""", + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - build_spec: Optional[ReasoningEngineSpecBuildSpec] = Field( - default=None, - description="""Optional. Configuration for building container image.""", + config: Optional[GetRuntimeOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -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.""" - - 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.""" - - 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)""" - - 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.""" +class _GetRuntimeOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with an agent runtime as a response.""" - 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.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - 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.""" + config: Optional[GetRuntimeOperationConfigDict] + """Used to override the default configuration.""" - source_code_spec: Optional[ReasoningEngineSpecSourceCodeSpecDict] - """Deploy from source code files with a defined entrypoint.""" - container_spec: Optional[ReasoningEngineSpecContainerSpecDict] - """Deploy from a container image with a defined entrypoint and commands.""" +_GetRuntimeOperationParametersOrDict = Union[ + _GetRuntimeOperationParameters, _GetRuntimeOperationParametersDict +] - build_spec: Optional[ReasoningEngineSpecBuildSpecDict] - """Optional. Configuration for building container image.""" +class QueryRuntimeConfig(_common.BaseModel): + """Config for querying agent runtimes.""" -ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict] + 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 ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest(_common.BaseModel): - """Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" +class QueryRuntimeConfigDict(TypedDict, total=False): + """Config for querying agent runtimes.""" - pass + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + class_method: Optional[str] + """The class method to call.""" -class ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict(TypedDict, total=False): - """Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""" + input: Optional[dict[str, Any]] + """The input to the class method.""" - pass + include_all_fields: Optional[bool] + """""" -ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestOrDict = Union[ - ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest, - ReasoningEngineTrafficConfigTrafficSplitAlwaysLatestDict, -] +QueryRuntimeConfigOrDict = Union[QueryRuntimeConfig, QueryRuntimeConfigDict] -class ReasoningEngineTrafficConfigTrafficSplitManualTarget(_common.BaseModel): - """A single target for the traffic split, specifying a Runtime Revision and the percentage of traffic to send to it.""" +class _QueryRuntimeRequestParameters(_common.BaseModel): + """Parameters for querying agent runtimes.""" - 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.""", + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime.""" ) + config: Optional[QueryRuntimeConfig] = Field(default=None, description="""""") -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.""" +class _QueryRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for querying agent runtimes.""" - percent: Optional[int] - """Required. Specifies percent of the traffic to this Runtime Revision.""" + name: Optional[str] + """Name of the agent runtime.""" - 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.""" + config: Optional[QueryRuntimeConfigDict] + """""" -ReasoningEngineTrafficConfigTrafficSplitManualTargetOrDict = Union[ - ReasoningEngineTrafficConfigTrafficSplitManualTarget, - ReasoningEngineTrafficConfigTrafficSplitManualTargetDict, +_QueryRuntimeRequestParametersOrDict = Union[ + _QueryRuntimeRequestParameters, _QueryRuntimeRequestParametersDict ] -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 QueryReasoningEngineResponse(_common.BaseModel): + """The response for querying an agent runtime.""" - 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.""", - ) + output: Optional[Any] = Field( + default=None, + description="""Response provided by users in JSON object format.""", ) -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 QueryReasoningEngineResponseDict(TypedDict, total=False): + """The response for querying an agent runtime.""" - targets: Optional[list[ReasoningEngineTrafficConfigTrafficSplitManualTargetDict]] - """A list of traffic targets for the Runtimes Revisions. The sum of percentages must equal to 100.""" + output: Optional[Any] + """Response provided by users in JSON object format.""" -ReasoningEngineTrafficConfigTrafficSplitManualOrDict = Union[ - ReasoningEngineTrafficConfigTrafficSplitManual, - ReasoningEngineTrafficConfigTrafficSplitManualDict, +QueryReasoningEngineResponseOrDict = Union[ + QueryReasoningEngineResponse, QueryReasoningEngineResponseDict ] -class ReasoningEngineTrafficConfig(_common.BaseModel): - """Traffic distribution configuration.""" +class UpdateRuntimeConfig(_common.BaseModel): + """Config for updating agent runtime.""" - traffic_split_always_latest: Optional[ - ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest - ] = Field( + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + display_name: Optional[str] = Field( default=None, - description="""Optional. Traffic distribution configuration, where all traffic is sent to the latest Runtime Revision.""", + 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. + """, ) - 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.""", - ) + description: Optional[str] = Field( + default=None, description="""The description of the Agent Runtime.""" ) - - -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.""" - - 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.""" - - -ReasoningEngineTrafficConfigOrDict = Union[ - ReasoningEngineTrafficConfig, ReasoningEngineTrafficConfigDict -] - - -class ReasoningEngine(_common.BaseModel): - """An agent engine.""" - - 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.""", + spec: Optional[ReasoningEngineSpec] = Field( + default=None, description="""Optional. Configurations of the Agent Runtime.""" ) context_spec: Optional[ReasoningEngineContextSpec] = Field( default=None, - description="""Optional. Configuration for how Agent Engine sub-resources should manage context.""", + description="""Optional. The context spec to be used for the Agent Runtime.""", ) - create_time: Optional[datetime.datetime] = Field( + psc_interface_config: Optional[PscInterfaceConfig] = Field( default=None, - description="""Output only. Timestamp when this ReasoningEngine was created.""", + description="""Optional. The PSC interface config for PSC-I to be used for the + Agent Runtime.""", ) - description: Optional[str] = Field( + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + ] = Field( default=None, - description="""Optional. The description of the ReasoningEngine.""", + description="""Agent Gateway configuration for a Reasoning Engine deployment.""", ) - display_name: Optional[str] = Field( + min_instances: Optional[int] = 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.""", - ) - traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( - default=None, - description="""Optional. Traffic distribution configuration for the Reasoning Engine.""", - ) - - -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.""" - - 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""" - - 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 AgentEngineOperation(_common.BaseModel): - """Operation that has an agent engine 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}`.""", - ) - 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[ReasoningEngine] = Field( - default=None, description="""The created Agent Engine.""" - ) - - -class AgentEngineOperationDict(TypedDict, total=False): - """Operation that has an agent engine 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.""" - - 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] - - -class CreateAgentEngineConfig(_common.BaseModel): - """Config for create agent engine.""" - - 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.""", - ) - 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. + 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 Engine. + 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 Engine. + 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'. @@ -10208,20 +9726,26 @@ class CreateAgentEngineConfig(_common.BaseModel): ) 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 @@ -10232,8 +9756,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 @@ -10255,19 +9779,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. """, ) @@ -10275,7 +9799,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". @@ -10286,14 +9810,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 @@ -10302,77 +9826,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 @@ -10382,8 +9913,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 @@ -10403,24 +9934,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". @@ -10430,13 +9961,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 @@ -10445,177 +9976,128 @@ 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( - 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( + create_time: Optional[datetime.datetime] = 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="""Output only. Timestamp when this ReasoningEngineRuntimeRevision was created.""", ) - done: Optional[bool] = Field( + name: Optional[str] = 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="""Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""", ) - error: Optional[dict[str, Any]] = Field( + spec: Optional[ReasoningEngineSpec] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""", ) - - -class DeleteAgentEngineOperationDict(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.""" - - 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.""" - - -DeleteAgentEngineOperationOrDict = Union[ - DeleteAgentEngineOperation, DeleteAgentEngineOperationDict -] - - -class GetAgentEngineConfig(_common.BaseModel): - """Config for create agent engine.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + state: Optional[State] = Field( + default=None, description="""Output only. The state of the revision.""" ) -class GetAgentEngineConfigDict(TypedDict, total=False): - """Config for create agent engine.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -GetAgentEngineConfigOrDict = Union[GetAgentEngineConfig, GetAgentEngineConfigDict] - - -class _GetAgentEngineRequestParameters(_common.BaseModel): - """Parameters for getting agent engines.""" - - name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[GetAgentEngineConfig] = Field(default=None, description="""""") - +class ReasoningEngineRuntimeRevisionDict(TypedDict, total=False): + """A runtime revision.""" -class _GetAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for getting agent engines.""" + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this ReasoningEngineRuntimeRevision was created.""" name: Optional[str] - """Name of the agent engine.""" + """Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""" - config: Optional[GetAgentEngineConfigDict] - """""" + spec: Optional[ReasoningEngineSpecDict] + """Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""" + + state: Optional[State] + """Output only. The state of the revision.""" -_GetAgentEngineRequestParametersOrDict = Union[ - _GetAgentEngineRequestParameters, _GetAgentEngineRequestParametersDict +ReasoningEngineRuntimeRevisionOrDict = Union[ + ReasoningEngineRuntimeRevision, ReasoningEngineRuntimeRevisionDict ] -class ListAgentEngineConfig(_common.BaseModel): - """Config for listing agent engines.""" +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.""" @@ -10629,8 +10111,8 @@ class ListAgentEngineConfig(_common.BaseModel): ) -class ListAgentEngineConfigDict(TypedDict, total=False): - """Config for listing agent engines.""" +class ListRuntimeRevisionsConfigDict(TypedDict, total=False): + """Config for listing agent runtime revisions.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -10646,43 +10128,51 @@ class ListAgentEngineConfigDict(TypedDict, total=False): For field names both snake_case and camelCase are supported.""" -ListAgentEngineConfigOrDict = Union[ListAgentEngineConfig, ListAgentEngineConfigDict] +ListRuntimeRevisionsConfigOrDict = Union[ + ListRuntimeRevisionsConfig, ListRuntimeRevisionsConfigDict +] + +class _ListRuntimeRevisionsRequestParameters(_common.BaseModel): + """Parameters for listing agent runtime revisions.""" -class _ListAgentEngineRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" + name: Optional[str] = Field( + default=None, description="""Name of the reasoning engine.""" + ) + config: Optional[ListRuntimeRevisionsConfig] = Field( + default=None, description="""""" + ) - config: Optional[ListAgentEngineConfig] = Field(default=None, description="""""") +class _ListRuntimeRevisionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtime revisions.""" -class _ListAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" + name: Optional[str] + """Name of the reasoning engine.""" - config: Optional[ListAgentEngineConfigDict] + config: Optional[ListRuntimeRevisionsConfigDict] """""" -_ListAgentEngineRequestParametersOrDict = Union[ - _ListAgentEngineRequestParameters, _ListAgentEngineRequestParametersDict +_ListRuntimeRevisionsRequestParametersOrDict = Union[ + _ListRuntimeRevisionsRequestParameters, _ListRuntimeRevisionsRequestParametersDict ] -class ListReasoningEnginesResponse(_common.BaseModel): - """Response for listing agent engines.""" +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_engines: Optional[list[ReasoningEngine]] = Field( - default=None, - description="""List of agent engines. - """, - ) + reasoning_engine_runtime_revisions: Optional[ + list[ReasoningEngineRuntimeRevision] + ] = Field(default=None, description="""List of agent runtime revisions.""") -class ListReasoningEnginesResponseDict(TypedDict, total=False): - """Response for listing agent engines.""" +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.""" @@ -10690,62 +10180,161 @@ class ListReasoningEnginesResponseDict(TypedDict, total=False): next_page_token: Optional[str] """""" - reasoning_engines: Optional[list[ReasoningEngineDict]] - """List of agent engines. - """ + reasoning_engine_runtime_revisions: Optional[ + list[ReasoningEngineRuntimeRevisionDict] + ] + """List of agent runtime revisions.""" -ListReasoningEnginesResponseOrDict = Union[ - ListReasoningEnginesResponse, ListReasoningEnginesResponseDict +ListReasoningEnginesRuntimeRevisionsResponseOrDict = Union[ + ListReasoningEnginesRuntimeRevisionsResponse, + ListReasoningEnginesRuntimeRevisionsResponseDict, +] + + +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.""" + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", + ) + + +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.""" + + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" + + +DeleteRuntimeRevisionConfigOrDict = Union[ + DeleteRuntimeRevisionConfig, DeleteRuntimeRevisionConfigDict +] + + +class _DeleteRuntimeRevisionRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtime runtime revisions.""" + + name: Optional[str] = Field( + default=None, + description="""Name of the agent runtime runtime revision to delete.""", + ) + config: Optional[DeleteRuntimeRevisionConfig] = Field( + default=None, description="""""" + ) + + +class _DeleteRuntimeRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtime runtime revisions.""" + + name: Optional[str] + """Name of the agent runtime runtime revision to delete.""" + + config: Optional[DeleteRuntimeRevisionConfigDict] + """""" + + +_DeleteRuntimeRevisionRequestParametersOrDict = Union[ + _DeleteRuntimeRevisionRequestParameters, _DeleteRuntimeRevisionRequestParametersDict +] + + +class DeleteRuntimeRevisionOperation(_common.BaseModel): + """Operation for deleting agent runtime revisions.""" + + 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.""", + ) + + +class DeleteRuntimeRevisionOperationDict(TypedDict, total=False): + """Operation for deleting agent runtime revisions.""" + + 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.""" + + +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.""" @@ -10759,8 +10348,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.""" @@ -10775,743 +10364,784 @@ 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 -] - - -class QueryReasoningEngineResponse(_common.BaseModel): - """The response for querying an agent engine.""" - - 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 engine.""" - - output: Optional[Any] - """Response provided by users in JSON object format.""" - - -QueryReasoningEngineResponseOrDict = Union[ - QueryReasoningEngineResponse, QueryReasoningEngineResponseDict +_QueryRuntimeRevisionRequestParametersOrDict = Union[ + _QueryRuntimeRevisionRequestParameters, _QueryRuntimeRevisionRequestParametersDict ] -class UpdateAgentEngineConfig(_common.BaseModel): - """Config for updating agent engine.""" +class CreateMemoryBankConfig(_common.BaseModel): + """Config for create memory bank.""" 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. + description="""The user-defined name of the Memory Bank. 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.""" + default=None, description="""The description of the Memory Bank.""" ) - context_spec: Optional[ReasoningEngineContextSpec] = Field( + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, - description="""Optional. The context spec to be used for the Agent Engine.""", + description="""The encryption spec to be used for the Memory Bank.""", ) - psc_interface_config: Optional[PscInterfaceConfig] = Field( - default=None, - description="""Optional. The PSC interface config for PSC-I to be used for the - Agent Engine.""", + + +class CreateMemoryBankConfigDict(TypedDict, total=False): + """Config for create memory bank.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + display_name: Optional[str] + """The user-defined name of the Memory Bank. + + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """ + + description: Optional[str] + """The description of the Memory Bank.""" + + encryption_spec: Optional[genai_types.EncryptionSpec] + """The encryption spec to be used for the Memory Bank.""" + + +CreateMemoryBankConfigOrDict = Union[CreateMemoryBankConfig, CreateMemoryBankConfigDict] + + +class _CreateMemoryBankRequestParameters(_common.BaseModel): + """Parameters for creating memory banks.""" + + config: Optional[CreateMemoryBankConfig] = Field(default=None, description="""""") + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfig] = Field( + default=None, description="""""" ) - min_instances: Optional[int] = Field( + + +class _CreateMemoryBankRequestParametersDict(TypedDict, total=False): + """Parameters for creating memory banks.""" + + config: Optional[CreateMemoryBankConfigDict] + """""" + + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfigDict] + """""" + + +_CreateMemoryBankRequestParametersOrDict = Union[ + _CreateMemoryBankRequestParameters, _CreateMemoryBankRequestParametersDict +] + + +class ManagedSemanticMemoryConfigGenerationConfig(_common.BaseModel): + """The configuration for generating memories.""" + + model: Optional[str] = Field( default=None, - description="""The minimum number of instances to run for the Agent Engine. - Defaults to 1. Range: [0, 10]. - """, + description="""The model used to generate memories. + + Format: + `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", ) - max_instances: Optional[int] = Field( + generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = 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]. - """, + description="""The configuration for triggering memory generation.""", ) - resource_limits: Optional[dict[str, str]] = Field( + + +class ManagedSemanticMemoryConfigGenerationConfigDict(TypedDict, total=False): + """The configuration for generating memories.""" + + model: Optional[str] + """The model used to generate memories. + + Format: + `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + + generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] + """The configuration for triggering memory generation.""" + + +ManagedSemanticMemoryConfigGenerationConfigOrDict = Union[ + ManagedSemanticMemoryConfigGenerationConfig, + ManagedSemanticMemoryConfigGenerationConfigDict, +] + + +class ManagedSemanticMemoryConfigSimilaritySearchConfig(_common.BaseModel): + """The configuration for similarity search.""" + + embedding_model: Optional[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'. - """, + description="""The model used to generate embeddings to look up similar memories. + Format: + `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", ) - container_concurrency: Optional[int] = Field( + + +class ManagedSemanticMemoryConfigSimilaritySearchConfigDict(TypedDict, total=False): + """The configuration for similarity search.""" + + embedding_model: Optional[str] + """The model used to generate embeddings to look up similar memories. + Format: + `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + + +ManagedSemanticMemoryConfigSimilaritySearchConfigOrDict = Union[ + ManagedSemanticMemoryConfigSimilaritySearchConfig, + ManagedSemanticMemoryConfigSimilaritySearchConfigDict, +] + + +class ManagedSemanticMemoryConfigTtlConfigGranularTtlConfig(_common.BaseModel): + """The configuration for granular TTL.""" + + create_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. The TTL duration for memories uploaded via + CreateMemory.""", ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + generate_created_ttl: Optional[str] = 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. The TTL duration for memories generated via + GenerateMemories.""", ) - class_methods: Optional[list[dict[str, Any]]] = Field( + generate_updated_ttl: Optional[str] = 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. 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).""", ) - source_packages: Optional[list[str]] = Field( + + +class ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigDict(TypedDict, total=False): + """The configuration for granular TTL.""" + + create_ttl: Optional[str] + """Optional. The TTL duration for memories uploaded via + CreateMemory.""" + + generate_created_ttl: Optional[str] + """Optional. The TTL duration for memories generated via + GenerateMemories.""" + + 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).""" + + +ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigOrDict = Union[ + ManagedSemanticMemoryConfigTtlConfigGranularTtlConfig, + ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigDict, +] + + +class ManagedSemanticMemoryConfigTtlConfig(_common.BaseModel): + """The configuration for automatic TTL ('time-to-live') of the memories.""" + + default_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="""The default TTL for 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.""", ) - developer_connect_source: Optional[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig - ] = Field( + granular_ttl_config: Optional[ + ManagedSemanticMemoryConfigTtlConfigGranularTtlConfig + ] = Field(default=None, description="""The granular TTL config for memories.""") + memory_revision_default_ttl: Optional[str] = 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="""The default TTL for memory revisions 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.""", ) - 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.""", + + +class ManagedSemanticMemoryConfigTtlConfigDict(TypedDict, total=False): + """The configuration for automatic TTL ('time-to-live') of the memories.""" + + default_ttl: Optional[str] + """The default TTL for 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.""" + + granular_ttl_config: Optional[ + ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigDict + ] + """The granular TTL config for memories.""" + + memory_revision_default_ttl: Optional[str] + """The default TTL for memory revisions 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.""" + + +ManagedSemanticMemoryConfigTtlConfigOrDict = Union[ + ManagedSemanticMemoryConfigTtlConfig, ManagedSemanticMemoryConfigTtlConfigDict +] + + +class ManagedSemanticMemoryConfig(_common.BaseModel): + """The configuration for managed semantic memory.""" + + generation_config: Optional[ManagedSemanticMemoryConfigGenerationConfig] = Field( + default=None, description="""Represents configuration for LLMs calls.""" ) - entrypoint_object: Optional[str] = Field( + ttl_config: Optional[ManagedSemanticMemoryConfigTtlConfig] = Field( default=None, - description="""The entrypoint object to be used for the Agent Engine. - This field only used when source_packages is specified.""", + description="""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.""", ) - requirements_file: Optional[str] = Field( + disable_memory_revisions: Optional[bool] = 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. - """, + description="""If true, no memory revisions will be created for any requests to + Memory Bank.""", ) - agent_framework: Optional[ - Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + similarity_search_config: Optional[ + ManagedSemanticMemoryConfigSimilaritySearchConfig ] = 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".""", + description="""Configuration for how to perform similarity search on memories.""", ) - python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] = Field( + unstructured_memory_configs: Optional[list[MemoryBankCustomizationConfig]] = 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". - """, + description="""Configuration for how to customize Memory Bank behavior for a + particular scope for unstructured memories.""", ) - build_options: Optional[dict[str, list[str]]] = Field( + structured_memory_configs: Optional[list[StructuredMemoryConfig]] = 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`. - """, + description="""Configuration for organizing structured memories for a particular + scope.""", ) - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig - ] = Field( + + +class ManagedSemanticMemoryConfigDict(TypedDict, total=False): + """The configuration for managed semantic memory.""" + + generation_config: Optional[ManagedSemanticMemoryConfigGenerationConfigDict] + """Represents configuration for LLMs calls.""" + + ttl_config: Optional[ManagedSemanticMemoryConfigTtlConfigDict] + """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.""" + + disable_memory_revisions: Optional[bool] + """If true, no memory revisions will be created for any requests to + Memory Bank.""" + + similarity_search_config: Optional[ + ManagedSemanticMemoryConfigSimilaritySearchConfigDict + ] + """Configuration for how to perform similarity search on memories.""" + + unstructured_memory_configs: Optional[list[MemoryBankCustomizationConfigDict]] + """Configuration for how to customize Memory Bank behavior for a + particular scope for unstructured memories.""" + + structured_memory_configs: Optional[list[StructuredMemoryConfigDict]] + """Configuration for organizing structured memories for a particular + scope.""" + + +ManagedSemanticMemoryConfigOrDict = Union[ + ManagedSemanticMemoryConfig, ManagedSemanticMemoryConfigDict +] + + +class MemoryBank(_common.BaseModel): + """A memory bank.""" + + name: Optional[str] = Field( default=None, - description="""Agent Gateway configuration for a Reasoning Engine deployment.""", + 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.""", ) - keep_alive_probe: Optional[KeepAliveProbe] = Field( + managed_semantic_memory_config: Optional[ManagedSemanticMemoryConfig] = 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="""Represents the configuration for managed memories in Memory Bank. If not set, then the default configuration will be used.""", ) - update_mask: Optional[str] = Field( + display_name: Optional[str] = Field( + default=None, description="""Represents the display name of the Memory Bank.""" + ) + description: Optional[str] = Field( + default=None, description="""Represents the description of the Memory Bank.""" + ) + create_time: Optional[datetime.datetime] = Field( + default=None, description="""Timestamp when this Memory Bank was created.""" + ) + update_time: 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="""Timestamp when this Memory Bank was most recently updated.""", ) - traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, - description="""Traffic distribution configuration for the Reasoning Engine.""", + description="""Customer-managed encryption key spec for a Memory Bank. If set, this Memory Bank and all sub-resources of this Memory Bank will be secured by this key.""", ) -class UpdateAgentEngineConfigDict(TypedDict, total=False): - """Config for updating agent engine.""" +class MemoryBankDict(TypedDict, total=False): + """A memory bank.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + name: 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.""" - display_name: Optional[str] - """The user-defined name of the Agent Engine. + managed_semantic_memory_config: Optional[ManagedSemanticMemoryConfigDict] + """Represents the configuration for managed memories in Memory Bank. If not set, then the default configuration will be used.""" - The display name can be up to 128 characters long and can comprise any - UTF-8 characters. - """ + display_name: Optional[str] + """Represents the display name of the Memory Bank.""" description: Optional[str] - """The description of the Agent Engine.""" + """Represents the description of the Memory Bank.""" - spec: Optional[ReasoningEngineSpecDict] - """Optional. Configurations of the Agent Engine.""" + create_time: Optional[datetime.datetime] + """Timestamp when this Memory Bank was created.""" - context_spec: Optional[ReasoningEngineContextSpecDict] - """Optional. The context spec to be used for the Agent Engine.""" + update_time: Optional[datetime.datetime] + """Timestamp when this Memory Bank was most recently updated.""" - psc_interface_config: Optional[PscInterfaceConfigDict] - """Optional. The PSC interface config for PSC-I to be used for the - Agent Engine.""" + encryption_spec: Optional[genai_types.EncryptionSpecDict] + """Customer-managed encryption key spec for a Memory Bank. If set, this Memory Bank and all sub-resources of this Memory Bank will be secured by this key.""" - 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]. - """ +MemoryBankOrDict = Union[MemoryBank, MemoryBankDict] - 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'. - """ - container_concurrency: Optional[int] - """The container concurrency to be used for the Agent Engine. - Recommended value: 2 * cpu + 1. Defaults to 9. - """ +class MemoryBankOperation(_common.BaseModel): + """Operation that has an memory bank as a response.""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """The encryption spec to be used for the Agent Engine.""" + 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[MemoryBank] = Field( + default=None, description="""The created Memory Bank.""" + ) - labels: Optional[dict[str, str]] - """The labels to be used for the Agent Engine.""" - 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. - """ +class MemoryBankOperationDict(TypedDict, total=False): + """Operation that has an memory bank as a response.""" - 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) - """ + 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[ - 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.""" + 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.""" - entrypoint_module: Optional[str] - """The entrypoint module to be used for the Agent Engine - This field only used when source_packages is specified.""" + 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.""" - entrypoint_object: Optional[str] - """The entrypoint object to be used for the Agent Engine. - This field only used when source_packages is specified.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - 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. - """ + response: Optional[MemoryBankDict] + """The created 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".""" - 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". - """ +MemoryBankOperationOrDict = Union[MemoryBankOperation, MemoryBankOperationDict] - 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`. - """ - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict - ] - """Agent Gateway configuration for a Reasoning Engine deployment.""" +class DeleteMemoryBankConfig(_common.BaseModel): + """Config for delete memory bank.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""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.""" - traffic_config: Optional[ReasoningEngineTrafficConfigDict] - """Traffic distribution configuration for the Reasoning Engine.""" +class DeleteMemoryBankConfigDict(TypedDict, total=False): + """Config for delete memory bank.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -UpdateAgentEngineConfigOrDict = Union[ - UpdateAgentEngineConfig, UpdateAgentEngineConfigDict -] +DeleteMemoryBankConfigOrDict = Union[DeleteMemoryBankConfig, DeleteMemoryBankConfigDict] -class _UpdateAgentEngineRequestParameters(_common.BaseModel): - """Parameters for updating agent engines.""" +class _DeleteMemoryBankRequestParameters(_common.BaseModel): + """Parameters for deleting a memory bank.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the memory bank.""" + ) + force: Optional[bool] = Field( + default=False, + description="""If set to true, any child resources will also be deleted.""", ) - config: Optional[UpdateAgentEngineConfig] = Field(default=None, description="""""") + config: Optional[DeleteMemoryBankConfig] = Field(default=None, description="""""") -class _UpdateAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for updating agent engines.""" +class _DeleteMemoryBankRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a memory bank.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the memory bank.""" + + force: Optional[bool] + """If set to true, any child resources will also be deleted.""" - config: Optional[UpdateAgentEngineConfigDict] + config: Optional[DeleteMemoryBankConfigDict] """""" -_UpdateAgentEngineRequestParametersOrDict = Union[ - _UpdateAgentEngineRequestParameters, _UpdateAgentEngineRequestParametersDict +_DeleteMemoryBankRequestParametersOrDict = Union[ + _DeleteMemoryBankRequestParameters, _DeleteMemoryBankRequestParametersDict ] -class CreateMemoryBankConfig(_common.BaseModel): - """Config for create memory bank.""" +class DeleteMemoryBankOperation(_common.BaseModel): + """Operation for deleting a memory bank.""" - 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}`.""", ) - display_name: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""The user-defined name of the Memory Bank. - - The display name can be up to 128 characters long and can comprise any - UTF-8 characters. - """, + 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[str] = Field( - default=None, description="""The description of the Memory Bank.""" + 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.""", ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""The encryption spec to be used for the Memory Bank.""", + description="""The error result of the operation in case of failure or cancellation.""", ) -class CreateMemoryBankConfigDict(TypedDict, total=False): - """Config for create memory bank.""" +class DeleteMemoryBankOperationDict(TypedDict, total=False): + """Operation for deleting a memory bank.""" - 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}`.""" - display_name: Optional[str] - """The user-defined name of the Memory Bank. + 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.""" - The display name can be up to 128 characters long and can comprise any - UTF-8 characters. - """ + 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.""" - description: Optional[str] - """The description of the Memory Bank.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """The encryption spec to be used for the Memory Bank.""" +DeleteMemoryBankOperationOrDict = Union[ + DeleteMemoryBankOperation, DeleteMemoryBankOperationDict +] -CreateMemoryBankConfigOrDict = Union[CreateMemoryBankConfig, CreateMemoryBankConfigDict] +class GetMemoryBankConfig(_common.BaseModel): + """Config for getting a Memory Bank.""" -class _CreateMemoryBankRequestParameters(_common.BaseModel): - """Parameters for creating memory banks.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) - config: Optional[CreateMemoryBankConfig] = Field(default=None, description="""""") - memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfig] = Field( - default=None, description="""""" + +class GetMemoryBankConfigDict(TypedDict, total=False): + """Config for getting a Memory Bank.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetMemoryBankConfigOrDict = Union[GetMemoryBankConfig, GetMemoryBankConfigDict] + + +class _GetMemoryBankRequestParameters(_common.BaseModel): + """Parameters for getting a Memory Bank.""" + + name: Optional[str] = Field( + default=None, description="""Name of the Memory Bank.""" ) + config: Optional[GetMemoryBankConfig] = Field(default=None, description="""""") -class _CreateMemoryBankRequestParametersDict(TypedDict, total=False): - """Parameters for creating memory banks.""" +class _GetMemoryBankRequestParametersDict(TypedDict, total=False): + """Parameters for getting a Memory Bank.""" - config: Optional[CreateMemoryBankConfigDict] - """""" + name: Optional[str] + """Name of the Memory Bank.""" - memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfigDict] + config: Optional[GetMemoryBankConfigDict] """""" -_CreateMemoryBankRequestParametersOrDict = Union[ - _CreateMemoryBankRequestParameters, _CreateMemoryBankRequestParametersDict +_GetMemoryBankRequestParametersOrDict = Union[ + _GetMemoryBankRequestParameters, _GetMemoryBankRequestParametersDict ] -class ManagedSemanticMemoryConfigGenerationConfig(_common.BaseModel): - """The configuration for generating memories.""" +class IngestionDirectContentsSourceEvent(_common.BaseModel): + """The direct contents source event for ingesting events.""" - model: Optional[str] = Field( + content: Optional[genai_types.Content] = Field( + default=None, description="""Required. The content of the event.""" + ) + event_id: Optional[str] = Field( default=None, - description="""The model used to generate memories. - - Format: - `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", + description="""Optional. A unique identifier for the event. If an event with the same event_id is ingested multiple times, it will be de-duplicated.""", ) - generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + event_time: Optional[datetime.datetime] = Field( default=None, - description="""The configuration for triggering memory generation.""", + description="""Optional. The time at which the event occurred. If provided, this timestamp will be used for ordering events within a stream. If not provided, the server-side ingestion time will be used.""", ) -class ManagedSemanticMemoryConfigGenerationConfigDict(TypedDict, total=False): - """The configuration for generating memories.""" +class IngestionDirectContentsSourceEventDict(TypedDict, total=False): + """The direct contents source event for ingesting events.""" - model: Optional[str] - """The model used to generate memories. + content: Optional[genai_types.Content] + """Required. The content of the event.""" - Format: - `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + event_id: Optional[str] + """Optional. A unique identifier for the event. If an event with the same event_id is ingested multiple times, it will be de-duplicated.""" - generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] - """The configuration for triggering memory generation.""" + event_time: Optional[datetime.datetime] + """Optional. The time at which the event occurred. If provided, this timestamp will be used for ordering events within a stream. If not provided, the server-side ingestion time will be used.""" -ManagedSemanticMemoryConfigGenerationConfigOrDict = Union[ - ManagedSemanticMemoryConfigGenerationConfig, - ManagedSemanticMemoryConfigGenerationConfigDict, +IngestionDirectContentsSourceEventOrDict = Union[ + IngestionDirectContentsSourceEvent, IngestionDirectContentsSourceEventDict ] -class ManagedSemanticMemoryConfigSimilaritySearchConfig(_common.BaseModel): - """The configuration for similarity search.""" +class IngestionDirectContentsSource(_common.BaseModel): + """The direct contents source for ingesting events.""" - embedding_model: Optional[str] = Field( - default=None, - description="""The model used to generate embeddings to look up similar memories. - Format: - `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", + events: Optional[list[IngestionDirectContentsSourceEvent]] = Field( + default=None, description="""Required. The events to ingest.""" ) -class ManagedSemanticMemoryConfigSimilaritySearchConfigDict(TypedDict, total=False): - """The configuration for similarity search.""" +class IngestionDirectContentsSourceDict(TypedDict, total=False): + """The direct contents source for ingesting events.""" - embedding_model: Optional[str] - """The model used to generate embeddings to look up similar memories. - Format: - `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + events: Optional[list[IngestionDirectContentsSourceEventDict]] + """Required. The events to ingest.""" -ManagedSemanticMemoryConfigSimilaritySearchConfigOrDict = Union[ - ManagedSemanticMemoryConfigSimilaritySearchConfig, - ManagedSemanticMemoryConfigSimilaritySearchConfigDict, +IngestionDirectContentsSourceOrDict = Union[ + IngestionDirectContentsSource, IngestionDirectContentsSourceDict ] -class ManagedSemanticMemoryConfigTtlConfigGranularTtlConfig(_common.BaseModel): - """The configuration for granular TTL.""" +class MemoryMetadataValue(_common.BaseModel): + """The metadata values for memories.""" - create_ttl: Optional[str] = Field( - default=None, - description="""Optional. The TTL duration for memories uploaded via - CreateMemory.""", + bool_value: Optional[bool] = Field( + default=None, description="""Represents a boolean value.""" ) - generate_created_ttl: Optional[str] = Field( - default=None, - description="""Optional. The TTL duration for memories generated via - GenerateMemories.""", + double_value: Optional[float] = Field( + default=None, description="""Represents a double value.""" ) - generate_updated_ttl: Optional[str] = Field( + string_value: Optional[str] = Field( + default=None, description="""Represents a string value.""" + ) + timestamp_value: Optional[datetime.datetime] = 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).""", + description="""Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""", ) -class ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigDict(TypedDict, total=False): - """The configuration for granular TTL.""" +class MemoryMetadataValueDict(TypedDict, total=False): + """The metadata values for memories.""" - create_ttl: Optional[str] - """Optional. The TTL duration for memories uploaded via - CreateMemory.""" + bool_value: Optional[bool] + """Represents a boolean value.""" - generate_created_ttl: Optional[str] - """Optional. The TTL duration for memories generated via - GenerateMemories.""" + double_value: Optional[float] + """Represents a double value.""" - 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).""" + 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.""" -ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigOrDict = Union[ - ManagedSemanticMemoryConfigTtlConfigGranularTtlConfig, - ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigDict, -] +MemoryMetadataValueOrDict = Union[MemoryMetadataValue, MemoryMetadataValueDict] -class ManagedSemanticMemoryConfigTtlConfig(_common.BaseModel): - """The configuration for automatic TTL ('time-to-live') of the memories.""" - default_ttl: Optional[str] = Field( - default=None, - description="""The default TTL for 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.""", +class IngestEventsConfig(_common.BaseModel): + """Config for ingesting events.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - granular_ttl_config: Optional[ - ManagedSemanticMemoryConfigTtlConfigGranularTtlConfig - ] = Field(default=None, description="""The granular TTL config for memories.""") - memory_revision_default_ttl: Optional[str] = Field( + wait_for_completion: Optional[bool] = Field( + default=False, + description="""Waits for the underlying memory generation operation to complete + before returning. Defaults to false.""", + ) + force_flush: Optional[bool] = Field( default=None, - description="""The default TTL for memory revisions 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. Forces a flush of all pending events in the stream and triggers memory generation immediately bypassing any conditions configured in the `generation_trigger_config`.""", ) - - -class ManagedSemanticMemoryConfigTtlConfigDict(TypedDict, total=False): - """The configuration for automatic TTL ('time-to-live') of the memories.""" - - default_ttl: Optional[str] - """The default TTL for 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.""" - - granular_ttl_config: Optional[ - ManagedSemanticMemoryConfigTtlConfigGranularTtlConfigDict - ] - """The granular TTL config for memories.""" - - memory_revision_default_ttl: Optional[str] - """The default TTL for memory revisions 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.""" - - -ManagedSemanticMemoryConfigTtlConfigOrDict = Union[ - ManagedSemanticMemoryConfigTtlConfig, ManagedSemanticMemoryConfigTtlConfigDict -] - - -class ManagedSemanticMemoryConfig(_common.BaseModel): - """The configuration for managed semantic memory.""" - - generation_config: Optional[ManagedSemanticMemoryConfigGenerationConfig] = Field( - default=None, description="""Represents configuration for LLMs calls.""" + revision_labels: Optional[dict[str, str]] = Field( + default=None, + description="""Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""", ) - ttl_config: Optional[ManagedSemanticMemoryConfigTtlConfig] = Field( + revision_expire_time: Optional[datetime.datetime] = Field( default=None, - description="""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. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", ) - disable_memory_revisions: Optional[bool] = Field( + revision_ttl: Optional[str] = Field( default=None, - description="""If true, no memory revisions will be created for any requests to - Memory Bank.""", + description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", ) - similarity_search_config: Optional[ - ManagedSemanticMemoryConfigSimilaritySearchConfig - ] = Field( + disable_memory_revisions: Optional[bool] = Field( default=None, - description="""Configuration for how to perform similarity search on memories.""", + description="""Optional. Input only. If true, no revisions will be created for this request.""", ) - unstructured_memory_configs: Optional[list[MemoryBankCustomizationConfig]] = Field( + metadata: Optional[dict[str, MemoryMetadataValue]] = Field( default=None, - description="""Configuration for how to customize Memory Bank behavior for a - particular scope for unstructured memories.""", + description="""Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""", ) - structured_memory_configs: Optional[list[StructuredMemoryConfig]] = Field( + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] = Field( default=None, - description="""Configuration for organizing structured memories for a particular - scope.""", + description="""Optional. The strategy to use when applying metadata to existing memories.""", ) -class ManagedSemanticMemoryConfigDict(TypedDict, total=False): - """The configuration for managed semantic memory.""" +class IngestEventsConfigDict(TypedDict, total=False): + """Config for ingesting events.""" - generation_config: Optional[ManagedSemanticMemoryConfigGenerationConfigDict] - """Represents configuration for LLMs calls.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - ttl_config: Optional[ManagedSemanticMemoryConfigTtlConfigDict] - """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.""" + wait_for_completion: Optional[bool] + """Waits for the underlying memory generation operation to complete + before returning. Defaults to false.""" - disable_memory_revisions: Optional[bool] - """If true, no memory revisions will be created for any requests to - Memory Bank.""" + force_flush: Optional[bool] + """Optional. Forces a flush of all pending events in the stream and triggers memory generation immediately bypassing any conditions configured in the `generation_trigger_config`.""" - similarity_search_config: Optional[ - ManagedSemanticMemoryConfigSimilaritySearchConfigDict - ] - """Configuration for how to perform similarity search on memories.""" + revision_labels: Optional[dict[str, str]] + """Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""" - unstructured_memory_configs: Optional[list[MemoryBankCustomizationConfigDict]] - """Configuration for how to customize Memory Bank behavior for a - particular scope for unstructured memories.""" + 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.""" - structured_memory_configs: Optional[list[StructuredMemoryConfigDict]] - """Configuration for organizing structured memories for a particular - scope.""" + 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 revisions will be created for this request.""" -ManagedSemanticMemoryConfigOrDict = Union[ - ManagedSemanticMemoryConfig, ManagedSemanticMemoryConfigDict -] + metadata: Optional[dict[str, MemoryMetadataValueDict]] + """Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""" + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] + """Optional. The strategy to use when applying metadata to existing memories.""" + + +IngestEventsConfigOrDict = Union[IngestEventsConfig, IngestEventsConfigDict] -class MemoryBank(_common.BaseModel): - """A memory bank.""" + +class _IngestEventsRequestParameters(_common.BaseModel): + """Parameters for ingesting events to Memory Bank.""" name: Optional[str] = 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.""", - ) - managed_semantic_memory_config: Optional[ManagedSemanticMemoryConfig] = Field( - default=None, - description="""Represents the configuration for managed memories in Memory Bank. If not set, then the default configuration will be used.""", - ) - display_name: Optional[str] = Field( - default=None, description="""Represents the display name of the Memory Bank.""" - ) - description: Optional[str] = Field( - default=None, description="""Represents the description of the Memory Bank.""" + default=None, description="""Name of the Memory Bank to ingest events into.""" ) - create_time: Optional[datetime.datetime] = Field( - default=None, description="""Timestamp when this Memory Bank was created.""" + stream_id: Optional[str] = Field( + default=None, description="""The ID of the stream to ingest events into.""" ) - update_time: Optional[datetime.datetime] = Field( + direct_contents_source: Optional[IngestionDirectContentsSource] = Field( default=None, - description="""Timestamp when this Memory Bank was most recently updated.""", + description="""The direct memories source of the events that should be ingested.""", ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + scope: Optional[dict[str, str]] = Field( default=None, - description="""Customer-managed encryption key spec for a Memory Bank. If set, this Memory Bank and all sub-resources of this Memory Bank will be secured by this key.""", - ) - + description="""The scope of the memories that should be generated from the stream. -class MemoryBankDict(TypedDict, total=False): - """A memory bank.""" + Memories will be consolidated across memories with the same scope. Scope + values cannot contain the wildcard character '*'.""", + ) + generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + default=None, + description="""The configuration for the memory generation trigger.""", + ) + config: Optional[IngestEventsConfig] = Field(default=None, description="""""") + + +class _IngestEventsRequestParametersDict(TypedDict, total=False): + """Parameters for ingesting events to Memory Bank.""" name: 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.""" + """Name of the Memory Bank to ingest events into.""" - managed_semantic_memory_config: Optional[ManagedSemanticMemoryConfigDict] - """Represents the configuration for managed memories in Memory Bank. If not set, then the default configuration will be used.""" + stream_id: Optional[str] + """The ID of the stream to ingest events into.""" - display_name: Optional[str] - """Represents the display name of the Memory Bank.""" + direct_contents_source: Optional[IngestionDirectContentsSourceDict] + """The direct memories source of the events that should be ingested.""" - description: Optional[str] - """Represents the description of the Memory Bank.""" + scope: Optional[dict[str, str]] + """The scope of the memories that should be generated from the stream. - create_time: Optional[datetime.datetime] - """Timestamp when this Memory Bank was created.""" + Memories will be consolidated across memories with the same scope. Scope + values cannot contain the wildcard character '*'.""" - update_time: Optional[datetime.datetime] - """Timestamp when this Memory Bank was most recently updated.""" + generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] + """The configuration for the memory generation trigger.""" - encryption_spec: Optional[genai_types.EncryptionSpecDict] - """Customer-managed encryption key spec for a Memory Bank. If set, this Memory Bank and all sub-resources of this Memory Bank will be secured by this key.""" + config: Optional[IngestEventsConfigDict] + """""" -MemoryBankOrDict = Union[MemoryBank, MemoryBankDict] +_IngestEventsRequestParametersOrDict = Union[ + _IngestEventsRequestParameters, _IngestEventsRequestParametersDict +] -class MemoryBankOperation(_common.BaseModel): - """Operation that has an memory bank as a response.""" +class MemoryBankIngestEventsOperation(_common.BaseModel): + """Operation that ingests events into a memory bank.""" name: Optional[str] = Field( default=None, @@ -11529,13 +11159,10 @@ class MemoryBankOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) - response: Optional[MemoryBank] = Field( - default=None, description="""The created Memory Bank.""" - ) -class MemoryBankOperationDict(TypedDict, total=False): - """Operation that has an memory bank as a response.""" +class MemoryBankIngestEventsOperationDict(TypedDict, total=False): + """Operation that ingests events into a memory bank.""" 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}`.""" @@ -11549,372 +11176,493 @@ class MemoryBankOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" - response: Optional[MemoryBankDict] - """The created Memory Bank.""" - -MemoryBankOperationOrDict = Union[MemoryBankOperation, MemoryBankOperationDict] +MemoryBankIngestEventsOperationOrDict = Union[ + MemoryBankIngestEventsOperation, MemoryBankIngestEventsOperationDict +] -class DeleteMemoryBankConfig(_common.BaseModel): - """Config for delete memory bank.""" +class ListMemoryBanksConfig(_common.BaseModel): + """Config for listing Memory Banks.""" 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="""""") -class DeleteMemoryBankConfigDict(TypedDict, total=False): - """Config for delete memory bank.""" +class ListMemoryBanksConfigDict(TypedDict, total=False): + """Config for listing Memory Banks.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - -DeleteMemoryBankConfigOrDict = Union[DeleteMemoryBankConfig, DeleteMemoryBankConfigDict] - - -class _DeleteMemoryBankRequestParameters(_common.BaseModel): - """Parameters for deleting a memory bank.""" - - name: Optional[str] = Field( - default=None, description="""Name of the memory bank.""" - ) - force: Optional[bool] = Field( - default=False, - description="""If set to true, any child resources will also be deleted.""", - ) - config: Optional[DeleteMemoryBankConfig] = Field(default=None, description="""""") - - -class _DeleteMemoryBankRequestParametersDict(TypedDict, total=False): - """Parameters for deleting a memory bank.""" - - name: Optional[str] - """Name of the memory bank.""" - - force: Optional[bool] - """If set to true, any child resources will also be deleted.""" - - config: Optional[DeleteMemoryBankConfigDict] + page_size: Optional[int] """""" + page_token: Optional[str] + """""" -_DeleteMemoryBankRequestParametersOrDict = Union[ - _DeleteMemoryBankRequestParameters, _DeleteMemoryBankRequestParametersDict -] - - -class DeleteMemoryBankOperation(_common.BaseModel): - """Operation for deleting a memory bank.""" - 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.""", - ) +ListMemoryBanksConfigOrDict = Union[ListMemoryBanksConfig, ListMemoryBanksConfigDict] -class DeleteMemoryBankOperationDict(TypedDict, total=False): - """Operation for deleting a memory bank.""" +class _ListMemoryBanksRequestParameters(_common.BaseModel): + """Parameters for listing Memory Banks.""" - 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}`.""" + config: Optional[ListMemoryBanksConfig] = Field(default=None, description="""""") - 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 _ListMemoryBanksRequestParametersDict(TypedDict, total=False): + """Parameters for listing Memory Banks.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + config: Optional[ListMemoryBanksConfigDict] + """""" -DeleteMemoryBankOperationOrDict = Union[ - DeleteMemoryBankOperation, DeleteMemoryBankOperationDict +_ListMemoryBanksRequestParametersOrDict = Union[ + _ListMemoryBanksRequestParameters, _ListMemoryBanksRequestParametersDict ] -class GetMemoryBankConfig(_common.BaseModel): - """Config for getting a Memory Bank.""" +class GetMemoryBankOperationConfig(_common.BaseModel): http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetMemoryBankConfigDict(TypedDict, total=False): - """Config for getting a Memory Bank.""" +class GetMemoryBankOperationConfigDict(TypedDict, total=False): http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetMemoryBankConfigOrDict = Union[GetMemoryBankConfig, GetMemoryBankConfigDict] +GetMemoryBankOperationConfigOrDict = Union[ + GetMemoryBankOperationConfig, GetMemoryBankOperationConfigDict +] -class _GetMemoryBankRequestParameters(_common.BaseModel): - """Parameters for getting a Memory Bank.""" +class _GetMemoryBankOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a memory bank as a response.""" - name: Optional[str] = Field( - default=None, description="""Name of the Memory Bank.""" + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetMemoryBankOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) - config: Optional[GetMemoryBankConfig] = Field(default=None, description="""""") -class _GetMemoryBankRequestParametersDict(TypedDict, total=False): - """Parameters for getting a Memory Bank.""" +class _GetMemoryBankOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a memory bank as a response.""" - name: Optional[str] - """Name of the Memory Bank.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - config: Optional[GetMemoryBankConfigDict] - """""" + config: Optional[GetMemoryBankOperationConfigDict] + """Used to override the default configuration.""" -_GetMemoryBankRequestParametersOrDict = Union[ - _GetMemoryBankRequestParameters, _GetMemoryBankRequestParametersDict +_GetMemoryBankOperationParametersOrDict = Union[ + _GetMemoryBankOperationParameters, _GetMemoryBankOperationParametersDict ] -class IngestionDirectContentsSourceEvent(_common.BaseModel): - """The direct contents source event for ingesting events.""" +class MemoryConfig(_common.BaseModel): + """Config for creating a Memory.""" - content: Optional[genai_types.Content] = Field( - default=None, description="""Required. The content of the event.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - event_id: Optional[str] = Field( + 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. A unique identifier for the event. If an event with the same event_id is ingested multiple times, it will be de-duplicated.""", + description="""Optional. Input only. The TTL for this resource. + + The expiration time is computed: now + TTL.""", ) - event_time: Optional[datetime.datetime] = Field( + expire_time: Optional[datetime.datetime] = Field( default=None, - description="""Optional. The time at which the event occurred. If provided, this timestamp will be used for ordering events within a stream. If not provided, the server-side ingestion time will be used.""", + 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 IngestionDirectContentsSourceEventDict(TypedDict, total=False): - """The direct contents source event for ingesting events.""" - - content: Optional[genai_types.Content] - """Required. The content of the event.""" +class MemoryConfigDict(TypedDict, total=False): + """Config for creating a Memory.""" - event_id: Optional[str] - """Optional. A unique identifier for the event. If an event with the same event_id is ingested multiple times, it will be de-duplicated.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - event_time: Optional[datetime.datetime] - """Optional. The time at which the event occurred. If provided, this timestamp will be used for ordering events within a stream. If not provided, the server-side ingestion time will be used.""" + display_name: Optional[str] + """The display name of the memory.""" + description: Optional[str] + """The description of the memory.""" -IngestionDirectContentsSourceEventOrDict = Union[ - IngestionDirectContentsSourceEvent, IngestionDirectContentsSourceEventDict -] + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" + ttl: Optional[str] + """Optional. Input only. The TTL for this resource. -class IngestionDirectContentsSource(_common.BaseModel): - """The direct contents source for ingesting events.""" + The expiration time is computed: now + TTL.""" - events: Optional[list[IngestionDirectContentsSourceEvent]] = Field( - default=None, description="""Required. The events to ingest.""" - ) + 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.""" -class IngestionDirectContentsSourceDict(TypedDict, total=False): - """The direct contents source for ingesting events.""" + revision_ttl: Optional[str] + """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" - events: Optional[list[IngestionDirectContentsSourceEventDict]] - """Required. The events to ingest.""" + 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.""" -IngestionDirectContentsSourceOrDict = Union[ - IngestionDirectContentsSource, IngestionDirectContentsSourceDict -] + 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.""" -class MemoryMetadataValue(_common.BaseModel): - """The metadata values for memories.""" - bool_value: Optional[bool] = Field( - default=None, description="""Represents a boolean value.""" - ) - double_value: Optional[float] = Field( - default=None, description="""Represents a double value.""" +MemoryConfigOrDict = Union[MemoryConfig, MemoryConfigDict] + + +class _CreateMemoryRequestParameters(_common.BaseModel): + """Parameters for creating Memories.""" + + name: Optional[str] = Field( + default=None, + description="""Name of the memory bank to create the memory under.""", ) - string_value: Optional[str] = Field( - default=None, description="""Represents a string value.""" + fact: Optional[str] = Field( + default=None, + description="""The fact of the memory. + + This is the semantic knowledge extracted from the source content).""", ) - timestamp_value: Optional[datetime.datetime] = Field( + scope: Optional[dict[str, str]] = Field( default=None, - description="""Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""", + 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[MemoryConfig] = Field(default=None, description="""""") -class MemoryMetadataValueDict(TypedDict, total=False): - """The metadata values for memories.""" +class _CreateMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for creating Memories.""" - bool_value: Optional[bool] - """Represents a boolean value.""" + name: Optional[str] + """Name of the memory bank to create the memory under.""" - double_value: Optional[float] - """Represents a double value.""" + fact: Optional[str] + """The fact of the memory. - string_value: Optional[str] - """Represents a string value.""" + This is the semantic knowledge extracted from the source content).""" - timestamp_value: Optional[datetime.datetime] - """Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""" + 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 '*'.""" -MemoryMetadataValueOrDict = Union[MemoryMetadataValue, MemoryMetadataValueDict] + config: Optional[MemoryConfigDict] + """""" -class IngestEventsConfig(_common.BaseModel): - """Config for ingesting events.""" +_CreateMemoryRequestParametersOrDict = Union[ + _CreateMemoryRequestParameters, _CreateMemoryRequestParametersDict +] - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + +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.""", ) - wait_for_completion: Optional[bool] = Field( - default=False, - description="""Waits for the underlying memory generation operation to complete - before returning. Defaults to false.""", + schema_id: Optional[str] = Field( + default=None, + description="""Required. Represents the schema ID for which this structured memory belongs to.""", ) - force_flush: Optional[bool] = Field( + + +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="""Optional. Forces a flush of all pending events in the stream and triggers memory generation immediately bypassing any conditions configured in the `generation_trigger_config`.""", + description="""Output only. Represents the timestamp when this Memory was created.""", ) - revision_labels: Optional[dict[str, str]] = Field( + description: Optional[str] = Field( default=None, - description="""Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""", + description="""Optional. Represents the description of the Memory.""", + ) + disable_memory_revisions: Optional[bool] = Field( + default=None, + description="""Optional. Input only. Indicates whether no revision will be created for this request.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""Optional. Represents the display name of the Memory.""", + ) + expire_time: Optional[datetime.datetime] = Field( + default=None, + description="""Optional. Represents the timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""", + ) + fact: Optional[str] = Field( + default=None, + description="""Optional. Represents semantic knowledge extracted from the source content.""", + ) + metadata: Optional[dict[str, MemoryMetadataValue]] = Field( + default=None, + description="""Optional. Represents user-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""", + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""", ) 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.""", + description="""Optional. Input only. Represents the timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", + ) + revision_labels: Optional[dict[str, str]] = Field( + default=None, + description="""Optional. Input only. Represents the labels to apply to the Memory Revision created as a result of this request.""", ) revision_ttl: Optional[str] = Field( default=None, - description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", + description="""Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""", ) - disable_memory_revisions: Optional[bool] = Field( + scope: Optional[dict[str, str]] = Field( default=None, - description="""Optional. Input only. If true, no revisions will be created for this request.""", + description="""Required. Immutable. Represents the scope of the Memory. Memories are isolated within their scope. The scope is defined when creating or generating memories. Scope values cannot contain the wildcard character '*'.""", ) - metadata: Optional[dict[str, MemoryMetadataValue]] = Field( + topics: Optional[list[MemoryTopicId]] = Field( + default=None, description="""Optional. Represents the Topics of the Memory.""" + ) + ttl: Optional[str] = Field( default=None, - description="""Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""", + description="""Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""", ) - metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] = Field( + update_time: Optional[datetime.datetime] = Field( default=None, - description="""Optional. The strategy to use when applying metadata to existing memories.""", + description="""Output only. Represents the timestamp when this Memory was most recently updated.""", + ) + memory_type: Optional[MemoryType] = Field( + default=None, + description="""Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""", + ) + structured_content: Optional[MemoryStructuredContent] = Field( + default=None, + description="""Optional. Represents the structured content of the memory.""", ) -class IngestEventsConfigDict(TypedDict, total=False): - """Config for ingesting events.""" +class MemoryDict(TypedDict, total=False): + """A memory.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + create_time: Optional[datetime.datetime] + """Output only. Represents the timestamp when this Memory was created.""" - wait_for_completion: Optional[bool] - """Waits for the underlying memory generation operation to complete - before returning. Defaults to false.""" + description: Optional[str] + """Optional. Represents the description of the Memory.""" - force_flush: Optional[bool] - """Optional. Forces a flush of all pending events in the stream and triggers memory generation immediately bypassing any conditions configured in the `generation_trigger_config`.""" + disable_memory_revisions: Optional[bool] + """Optional. Input only. Indicates whether no revision will be created for this request.""" - revision_labels: Optional[dict[str, str]] - """Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""" + display_name: Optional[str] + """Optional. Represents the display name of the Memory.""" + + expire_time: Optional[datetime.datetime] + """Optional. Represents the timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""" + + fact: Optional[str] + """Optional. Represents semantic knowledge extracted from the source content.""" + + metadata: Optional[dict[str, MemoryMetadataValueDict]] + """Optional. Represents user-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""" + + name: Optional[str] + """Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""" 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.""" + """Optional. Input only. Represents the timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" + + revision_labels: Optional[dict[str, str]] + """Optional. Input only. Represents the labels to apply to the Memory Revision created as a result of this request.""" revision_ttl: Optional[str] - """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" + """Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""" - disable_memory_revisions: Optional[bool] - """Optional. Input only. If true, no revisions will be created for this request.""" + scope: Optional[dict[str, str]] + """Required. Immutable. Represents the scope of the Memory. Memories are isolated within their scope. The scope is defined when creating or generating memories. Scope values cannot contain the wildcard character '*'.""" - metadata: Optional[dict[str, MemoryMetadataValueDict]] - """Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. Represents the Topics of the Memory.""" - metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] - """Optional. The strategy to use when applying metadata to existing memories.""" + ttl: Optional[str] + """Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""" + update_time: Optional[datetime.datetime] + """Output only. Represents the timestamp when this Memory was most recently updated.""" -IngestEventsConfigOrDict = Union[IngestEventsConfig, IngestEventsConfigDict] + memory_type: Optional[MemoryType] + """Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""" + structured_content: Optional[MemoryStructuredContentDict] + """Optional. Represents the structured content of the memory.""" -class _IngestEventsRequestParameters(_common.BaseModel): - """Parameters for ingesting events to Memory Bank.""" - name: Optional[str] = Field( - default=None, description="""Name of the Memory Bank to ingest events into.""" - ) - stream_id: Optional[str] = Field( - default=None, description="""The ID of the stream to ingest events into.""" +MemoryOrDict = Union[Memory, MemoryDict] + + +class MemoryOperation(_common.BaseModel): + """Operation that has a memory 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}`.""", ) - direct_contents_source: Optional[IngestionDirectContentsSource] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""The direct memories source of the events that should be ingested.""", + 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.""", ) - scope: Optional[dict[str, str]] = Field( + done: Optional[bool] = Field( default=None, - description="""The scope of the memories that should be generated from the stream. - - Memories will be consolidated across memories with the same scope. Scope - values cannot contain the wildcard character '*'.""", + 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.""", ) - generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""The configuration for the memory generation trigger.""", + description="""The error result of the operation in case of failure or cancellation.""", ) - config: Optional[IngestEventsConfig] = Field(default=None, description="""""") + response: Optional[Memory] = Field(default=None, description="""The Memory.""") -class _IngestEventsRequestParametersDict(TypedDict, total=False): - """Parameters for ingesting events to Memory Bank.""" +class MemoryOperationDict(TypedDict, total=False): + """Operation that has a memory as a response.""" name: Optional[str] - """Name of the Memory Bank to ingest events into.""" + """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}`.""" - stream_id: Optional[str] - """The ID of the stream to ingest events into.""" + 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.""" - direct_contents_source: Optional[IngestionDirectContentsSourceDict] - """The direct memories source of the events that should be ingested.""" + 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.""" - scope: Optional[dict[str, str]] - """The scope of the memories that should be generated from the stream. + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - Memories will be consolidated across memories with the same scope. Scope - values cannot contain the wildcard character '*'.""" + response: Optional[MemoryDict] + """The Memory.""" - generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] - """The configuration for the memory generation trigger.""" - config: Optional[IngestEventsConfigDict] +MemoryOperationOrDict = Union[MemoryOperation, MemoryOperationDict] + + +class DeleteMemoryConfig(_common.BaseModel): + """Config for deleting a Memory.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class DeleteMemoryConfigDict(TypedDict, total=False): + """Config for deleting a Memory.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +DeleteMemoryConfigOrDict = Union[DeleteMemoryConfig, DeleteMemoryConfigDict] + + +class _DeleteMemoryRequestParameters(_common.BaseModel): + """Parameters for deleting memories.""" + + name: Optional[str] = Field( + default=None, description="""Name of the memory to delete.""" + ) + config: Optional[DeleteMemoryConfig] = Field(default=None, description="""""") + + +class _DeleteMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for deleting memories.""" + + name: Optional[str] + """Name of the memory to delete.""" + + config: Optional[DeleteMemoryConfigDict] """""" -_IngestEventsRequestParametersOrDict = Union[ - _IngestEventsRequestParameters, _IngestEventsRequestParametersDict +_DeleteMemoryRequestParametersOrDict = Union[ + _DeleteMemoryRequestParameters, _DeleteMemoryRequestParametersDict ] -class MemoryBankIngestEventsOperation(_common.BaseModel): - """Operation that ingests events into a memory bank.""" +class DeleteMemoryOperation(_common.BaseModel): + """Operation for deleting memories.""" name: Optional[str] = Field( default=None, @@ -11934,8 +11682,8 @@ class MemoryBankIngestEventsOperation(_common.BaseModel): ) -class MemoryBankIngestEventsOperationDict(TypedDict, total=False): - """Operation that ingests events into a memory bank.""" +class DeleteMemoryOperationDict(TypedDict, total=False): + """Operation for deleting memories.""" 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}`.""" @@ -11950,124 +11698,166 @@ class MemoryBankIngestEventsOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -MemoryBankIngestEventsOperationOrDict = Union[ - MemoryBankIngestEventsOperation, MemoryBankIngestEventsOperationDict -] +DeleteMemoryOperationOrDict = Union[DeleteMemoryOperation, DeleteMemoryOperationDict] -class ListMemoryBanksConfig(_common.BaseModel): - """Config for listing Memory Banks.""" +class GenerateMemoriesRequestVertexSessionSource(_common.BaseModel): + """The vertex session source for generating memories.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + end_time: Optional[datetime.datetime] = Field( + default=None, + description="""Optional. End time (exclusive) of the time range. If not set, the end time is unbounded.""", + ) + session: Optional[str] = Field( + default=None, + description="""Required. The resource name of the Session to generate memories for. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}`""", + ) + start_time: Optional[datetime.datetime] = Field( + default=None, + description="""Optional. Time range to define which session events should be used to generate memories. Start time (inclusive) of the time range. If not set, the start time is unbounded.""", ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") -class ListMemoryBanksConfigDict(TypedDict, total=False): - """Config for listing Memory Banks.""" +class GenerateMemoriesRequestVertexSessionSourceDict(TypedDict, total=False): + """The vertex session source for generating memories.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + end_time: Optional[datetime.datetime] + """Optional. End time (exclusive) of the time range. If not set, the end time is unbounded.""" - page_size: Optional[int] - """""" + session: Optional[str] + """Required. The resource name of the Session to generate memories for. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}`""" - page_token: Optional[str] - """""" + start_time: Optional[datetime.datetime] + """Optional. Time range to define which session events should be used to generate memories. Start time (inclusive) of the time range. If not set, the start time is unbounded.""" -ListMemoryBanksConfigOrDict = Union[ListMemoryBanksConfig, ListMemoryBanksConfigDict] +GenerateMemoriesRequestVertexSessionSourceOrDict = Union[ + GenerateMemoriesRequestVertexSessionSource, + GenerateMemoriesRequestVertexSessionSourceDict, +] -class _ListMemoryBanksRequestParameters(_common.BaseModel): - """Parameters for listing Memory Banks.""" +class GenerateMemoriesRequestDirectContentsSourceEvent(_common.BaseModel): - config: Optional[ListMemoryBanksConfig] = Field(default=None, description="""""") + content: Optional[genai_types.Content] = Field( + default=None, + description="""Required. A single piece of content from which to generate memories.""", + ) -class _ListMemoryBanksRequestParametersDict(TypedDict, total=False): - """Parameters for listing Memory Banks.""" +class GenerateMemoriesRequestDirectContentsSourceEventDict(TypedDict, total=False): - config: Optional[ListMemoryBanksConfigDict] - """""" + content: Optional[genai_types.Content] + """Required. A single piece of content from which to generate memories.""" -_ListMemoryBanksRequestParametersOrDict = Union[ - _ListMemoryBanksRequestParameters, _ListMemoryBanksRequestParametersDict +GenerateMemoriesRequestDirectContentsSourceEventOrDict = Union[ + GenerateMemoriesRequestDirectContentsSourceEvent, + GenerateMemoriesRequestDirectContentsSourceEventDict, ] -class GetMemoryBankOperationConfig(_common.BaseModel): +class GenerateMemoriesRequestDirectContentsSource(_common.BaseModel): + """The direct contents source for generating memories.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEvent]] = Field( + default=None, + description="""Required. The source content (i.e. chat history) to generate memories from.""", ) -class GetMemoryBankOperationConfigDict(TypedDict, total=False): +class GenerateMemoriesRequestDirectContentsSourceDict(TypedDict, total=False): + """The direct contents source for generating memories.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEventDict]] + """Required. The source content (i.e. chat history) to generate memories from.""" -GetMemoryBankOperationConfigOrDict = Union[ - GetMemoryBankOperationConfig, GetMemoryBankOperationConfigDict +GenerateMemoriesRequestDirectContentsSourceOrDict = Union[ + GenerateMemoriesRequestDirectContentsSource, + GenerateMemoriesRequestDirectContentsSourceDict, ] -class _GetMemoryBankOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a memory bank as a response.""" +class GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(_common.BaseModel): + """A direct memory to upload to Memory Bank.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + fact: Optional[str] = Field( + default=None, + description="""Required. The fact to consolidate with existing memories.""", ) - config: Optional[GetMemoryBankOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + topics: Optional[list[MemoryTopicId]] = Field( + default=None, + description="""Optional. The topics that the consolidated memories should be associated with.""", ) -class _GetMemoryBankOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a memory bank as a response.""" +class GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict( + TypedDict, total=False +): + """A direct memory to upload to Memory Bank.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + fact: Optional[str] + """Required. The fact to consolidate with existing memories.""" - config: Optional[GetMemoryBankOperationConfigDict] - """Used to override the default configuration.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The topics that the consolidated memories should be associated with.""" -_GetMemoryBankOperationParametersOrDict = Union[ - _GetMemoryBankOperationParameters, _GetMemoryBankOperationParametersDict +GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryOrDict = Union[ + GenerateMemoriesRequestDirectMemoriesSourceDirectMemory, + GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict, ] -class MemoryConfig(_common.BaseModel): - """Config for creating a Memory.""" +class GenerateMemoriesRequestDirectMemoriesSource(_common.BaseModel): + """The direct memories source for generating memories.""" + + direct_memories: Optional[ + list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemory] + ] = Field( + default=None, + description="""Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""", + ) + + +class GenerateMemoriesRequestDirectMemoriesSourceDict(TypedDict, total=False): + """The direct memories source for generating memories.""" + + direct_memories: Optional[ + list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict] + ] + """Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""" + + +GenerateMemoriesRequestDirectMemoriesSourceOrDict = Union[ + GenerateMemoriesRequestDirectMemoriesSource, + GenerateMemoriesRequestDirectMemoriesSourceDict, +] + + +class GenerateMemoriesConfig(_common.BaseModel): + """Config for generating memories.""" 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( + disable_consolidation: Optional[bool] = Field( default=None, - description="""Optional. Input only. The TTL for this resource. + description="""Whether to disable consolidation of memories. - The expiration time is computed: now + TTL.""", + If true, generated memories will not be consolidated with existing + memories; all generated memories will be added as new memories regardless + of whether they are duplicates of or contradictory to existing memories. + By default, memory consolidation is enabled.""", ) - expire_time: Optional[datetime.datetime] = Field( + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", + ) + revision_labels: Optional[dict[str, 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="""Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""", ) revision_expire_time: Optional[datetime.datetime] = Field( default=None, @@ -12079,43 +11869,41 @@ class MemoryConfig(_common.BaseModel): ) 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.""" + description="""Optional. Input only. If true, no revisions will be created for this request.""", ) 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.""", + description="""Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""", ) - memory_id: Optional[str] = Field( + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] = 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.""", + description="""Optional. The strategy to use when applying metadata to existing memories.""", + ) + allowed_topics: Optional[list[MemoryTopicId]] = Field( + default=None, + description="""Optional. Restricts memory generation to a subset of memory topics.""", ) -class MemoryConfigDict(TypedDict, total=False): - """Config for creating a Memory.""" +class GenerateMemoriesConfigDict(TypedDict, total=False): + """Config for generating memories.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - display_name: Optional[str] - """The display name of the memory.""" + disable_consolidation: Optional[bool] + """Whether to disable consolidation of memories. - description: Optional[str] - """The description of the memory.""" + If true, generated memories will not be consolidated with existing + memories; all generated memories will be added as new memories regardless + of whether they are duplicates of or contradictory to existing memories. + By default, memory consolidation is enabled.""" 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_labels: Optional[dict[str, str]] + """Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""" 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.""" @@ -12124,254 +11912,175 @@ class MemoryConfigDict(TypedDict, total=False): """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.""" + """Optional. Input only. If true, no revisions will be created for this request.""" 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.""" + """Optional. User-provided metadata for the generated memories. This is 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.""" + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] + """Optional. The strategy to use when applying metadata to existing memories.""" + allowed_topics: Optional[list[MemoryTopicIdDict]] + """Optional. Restricts memory generation to a subset of memory topics.""" -MemoryConfigOrDict = Union[MemoryConfig, MemoryConfigDict] +GenerateMemoriesConfigOrDict = Union[GenerateMemoriesConfig, GenerateMemoriesConfigDict] -class _CreateMemoryRequestParameters(_common.BaseModel): - """Parameters for creating Memories.""" + +class _GenerateMemoriesRequestParameters(_common.BaseModel): + """Parameters for generating memories.""" name: Optional[str] = Field( default=None, - description="""Name of the memory bank to create the memory under.""", + description="""Name of the Memory Bank to generate memories with.""", ) - fact: Optional[str] = Field( + vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSource] = Field( default=None, - description="""The fact of the memory. - - This is the semantic knowledge extracted from the source content).""", + description="""The vertex session source of the memories that should be generated.""", + ) + direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSource] = ( + Field( + default=None, + description="""The direct contents source of the memories that should be generated.""", + ) + ) + direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSource] = ( + Field( + default=None, + description="""The direct memories source of the memories that should be generated.""", + ) ) scope: Optional[dict[str, str]] = Field( default=None, - description="""The scope of the memory. + description="""The scope of the memories that should be generated. - 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 '*'.""", + Memories will be consolidated across memories with the same scope. Must be + provided unless the scope is defined in the source content. If `scope` is + provided, it will override the scope defined in the source content. Scope + values cannot contain the wildcard character '*'.""", ) - config: Optional[MemoryConfig] = Field(default=None, description="""""") + config: Optional[GenerateMemoriesConfig] = Field(default=None, description="""""") -class _CreateMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for creating Memories.""" +class _GenerateMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for generating memories.""" name: Optional[str] - """Name of the memory bank to create the memory under.""" + """Name of the Memory Bank to generate memories with.""" - fact: Optional[str] - """The fact of the memory. + vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSourceDict] + """The vertex session source of the memories that should be generated.""" - This is the semantic knowledge extracted from the source content).""" + direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSourceDict] + """The direct contents source of the memories that should be generated.""" + + direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSourceDict] + """The direct memories source of the memories that should be generated.""" scope: Optional[dict[str, str]] - """The scope of the memory. + """The scope of the memories that should be generated. - 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 '*'.""" + Memories will be consolidated across memories with the same scope. Must be + provided unless the scope is defined in the source content. If `scope` is + provided, it will override the scope defined in the source content. Scope + values cannot contain the wildcard character '*'.""" - config: Optional[MemoryConfigDict] + config: Optional[GenerateMemoriesConfigDict] """""" -_CreateMemoryRequestParametersOrDict = Union[ - _CreateMemoryRequestParameters, _CreateMemoryRequestParametersDict +_GenerateMemoriesRequestParametersOrDict = Union[ + _GenerateMemoriesRequestParameters, _GenerateMemoriesRequestParametersDict ] -class MemoryStructuredContent(_common.BaseModel): - """Represents the structured value of the memory.""" +class GenerateMemoriesResponseGeneratedMemory(_common.BaseModel): + """A memory that was generated.""" - data: Optional[dict[str, Any]] = Field( - default=None, - description="""Required. Represents the structured value of the memory.""", + memory: Optional[Memory] = Field( + default=None, description="""The generated memory.""" ) - schema_id: Optional[str] = Field( + action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] = Field( + default=None, description="""The action to take.""" + ) + previous_revision: Optional[str] = Field( default=None, - description="""Required. Represents the schema ID for which this structured memory belongs to.""", + description="""The previous revision of the Memory before the action was performed. This + field is only set if the action is `UPDATED` or `DELETED`. You can use + this to rollback the Memory to the previous revision, undoing the action. + Format: + `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{revision}`""", ) -class MemoryStructuredContentDict(TypedDict, total=False): - """Represents the structured value of the memory.""" +class GenerateMemoriesResponseGeneratedMemoryDict(TypedDict, total=False): + """A memory that was generated.""" - data: Optional[dict[str, Any]] - """Required. Represents the structured value of the memory.""" + memory: Optional[MemoryDict] + """The generated memory.""" - schema_id: Optional[str] - """Required. Represents the schema ID for which this structured memory belongs to.""" + action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] + """The action to take.""" + + previous_revision: Optional[str] + """The previous revision of the Memory before the action was performed. This + field is only set if the action is `UPDATED` or `DELETED`. You can use + this to rollback the Memory to the previous revision, undoing the action. + Format: + `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{revision}`""" -MemoryStructuredContentOrDict = Union[ - MemoryStructuredContent, MemoryStructuredContentDict +GenerateMemoriesResponseGeneratedMemoryOrDict = Union[ + GenerateMemoriesResponseGeneratedMemory, GenerateMemoriesResponseGeneratedMemoryDict ] -class Memory(_common.BaseModel): - """A memory.""" +class GenerateMemoriesResponse(_common.BaseModel): + """The response for generating memories.""" - 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, - description="""Optional. Represents the description of the Memory.""", - ) - disable_memory_revisions: Optional[bool] = Field( - default=None, - description="""Optional. Input only. Indicates whether no revision will be created for this request.""", - ) - display_name: Optional[str] = Field( - default=None, - description="""Optional. Represents the display name of the Memory.""", - ) - expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. Represents the timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""", - ) - fact: Optional[str] = Field( - default=None, - description="""Optional. Represents semantic knowledge extracted from the source content.""", - ) - metadata: Optional[dict[str, MemoryMetadataValue]] = Field( - default=None, - description="""Optional. Represents user-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""", + generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemory]] = Field( + default=None, description="""The generated memories.""" ) + + +class GenerateMemoriesResponseDict(TypedDict, total=False): + """The response for generating memories.""" + + generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemoryDict]] + """The generated memories.""" + + +GenerateMemoriesResponseOrDict = Union[ + GenerateMemoriesResponse, GenerateMemoriesResponseDict +] + + +class GenerateMemoriesOperation(_common.BaseModel): + """Operation that generates memories with a Memory Bank.""" + name: Optional[str] = Field( default=None, - description="""Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""", - ) - revision_expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. Input only. Represents the timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", + 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}`.""", ) - revision_labels: Optional[dict[str, str]] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. Input only. Represents the labels to apply to the Memory Revision created as a result of this request.""", + 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.""", ) - revision_ttl: Optional[str] = Field( + done: Optional[bool] = Field( default=None, - description="""Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""", + 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.""", ) - scope: Optional[dict[str, str]] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Required. Immutable. Represents the scope of the Memory. Memories are isolated within their scope. The scope is defined when creating or generating memories. Scope values cannot contain the wildcard character '*'.""", + description="""The error result of the operation in case of failure or cancellation.""", ) - topics: Optional[list[MemoryTopicId]] = Field( - default=None, description="""Optional. Represents the Topics of the Memory.""" + response: Optional[GenerateMemoriesResponse] = Field( + default=None, description="""The response for generating memories.""" ) - ttl: Optional[str] = Field( - default=None, - description="""Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""", - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Represents the timestamp when this Memory was most recently updated.""", - ) - memory_type: Optional[MemoryType] = Field( - default=None, - description="""Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""", - ) - structured_content: Optional[MemoryStructuredContent] = Field( - default=None, - description="""Optional. Represents the structured content of the memory.""", - ) - - -class MemoryDict(TypedDict, total=False): - """A memory.""" - - create_time: Optional[datetime.datetime] - """Output only. Represents the timestamp when this Memory was created.""" - - description: Optional[str] - """Optional. Represents the description of the Memory.""" - - disable_memory_revisions: Optional[bool] - """Optional. Input only. Indicates whether no revision will be created for this request.""" - - display_name: Optional[str] - """Optional. Represents the display name of the Memory.""" - - expire_time: Optional[datetime.datetime] - """Optional. Represents the timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""" - - fact: Optional[str] - """Optional. Represents semantic knowledge extracted from the source content.""" - - metadata: Optional[dict[str, MemoryMetadataValueDict]] - """Optional. Represents user-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""" - - name: Optional[str] - """Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""" - - revision_expire_time: Optional[datetime.datetime] - """Optional. Input only. Represents the timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" - - revision_labels: Optional[dict[str, str]] - """Optional. Input only. Represents the labels to apply to the Memory Revision created as a result of this request.""" - - revision_ttl: Optional[str] - """Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""" - - scope: Optional[dict[str, str]] - """Required. Immutable. Represents the scope of the Memory. Memories are isolated within their scope. The scope is defined when creating or generating memories. Scope values cannot contain the wildcard character '*'.""" - - topics: Optional[list[MemoryTopicIdDict]] - """Optional. Represents the Topics of the Memory.""" - - ttl: Optional[str] - """Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""" - - update_time: Optional[datetime.datetime] - """Output only. Represents the timestamp when this Memory was most recently updated.""" - - memory_type: Optional[MemoryType] - """Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""" - - structured_content: Optional[MemoryStructuredContentDict] - """Optional. Represents the structured content of the memory.""" - - -MemoryOrDict = Union[Memory, MemoryDict] - - -class MemoryOperation(_common.BaseModel): - """Operation that has a memory 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}`.""", - ) - 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[Memory] = Field(default=None, description="""The Memory.""") -class MemoryOperationDict(TypedDict, total=False): - """Operation that has a memory as a response.""" +class GenerateMemoriesOperationDict(TypedDict, total=False): + """Operation that generates memories with a Memory Bank.""" 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}`.""" @@ -12385,834 +12094,937 @@ class MemoryOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" - response: Optional[MemoryDict] - """The Memory.""" + response: Optional[GenerateMemoriesResponseDict] + """The response for generating memories.""" -MemoryOperationOrDict = Union[MemoryOperation, MemoryOperationDict] +GenerateMemoriesOperationOrDict = Union[ + GenerateMemoriesOperation, GenerateMemoriesOperationDict +] -class DeleteMemoryConfig(_common.BaseModel): - """Config for deleting a Memory.""" +class GetMemoryConfig(_common.BaseModel): + """Config for getting a Memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class DeleteMemoryConfigDict(TypedDict, total=False): - """Config for deleting a Memory.""" +class GetMemoryConfigDict(TypedDict, total=False): + """Config for getting a Memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteMemoryConfigOrDict = Union[DeleteMemoryConfig, DeleteMemoryConfigDict] +GetMemoryConfigOrDict = Union[GetMemoryConfig, GetMemoryConfigDict] -class _DeleteMemoryRequestParameters(_common.BaseModel): - """Parameters for deleting memories.""" +class _GetMemoryRequestParameters(_common.BaseModel): + """Parameters for getting a Memory.""" - name: Optional[str] = Field( - default=None, description="""Name of the memory to delete.""" - ) - config: Optional[DeleteMemoryConfig] = Field(default=None, description="""""") + name: Optional[str] = Field(default=None, description="""Name of the memory.""") + config: Optional[GetMemoryConfig] = Field(default=None, description="""""") -class _DeleteMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for deleting memories.""" +class _GetMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for getting a Memory.""" name: Optional[str] - """Name of the memory to delete.""" + """Name of the memory.""" - config: Optional[DeleteMemoryConfigDict] + config: Optional[GetMemoryConfigDict] """""" -_DeleteMemoryRequestParametersOrDict = Union[ - _DeleteMemoryRequestParameters, _DeleteMemoryRequestParametersDict +_GetMemoryRequestParametersOrDict = Union[ + _GetMemoryRequestParameters, _GetMemoryRequestParametersDict ] -class DeleteMemoryOperation(_common.BaseModel): - """Operation for deleting memories.""" +class ListMemoriesConfig(_common.BaseModel): + """Config for listing memories.""" - 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - done: Optional[bool] = Field( + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = 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="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""", ) - error: Optional[dict[str, Any]] = Field( + order_by: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""The standard list order by string. If not specified, the default + order is `create_time desc`. If specified, the default sorting order of + provided fields is ascending. More detail in + [AIP-132](https://google.aip.dev/132). + + Supported fields: + * `create_time` + * `update_time`""", ) -class DeleteMemoryOperationDict(TypedDict, total=False): - """Operation for deleting memories.""" +class ListMemoriesConfigDict(TypedDict, total=False): + """Config for listing memories.""" - 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.""" + + order_by: Optional[str] + """The standard list order by string. If not specified, the default + order is `create_time desc`. If specified, the default sorting order of + provided fields is ascending. More detail in + [AIP-132](https://google.aip.dev/132). + Supported fields: + * `create_time` + * `update_time`""" -DeleteMemoryOperationOrDict = Union[DeleteMemoryOperation, DeleteMemoryOperationDict] +ListMemoriesConfigOrDict = Union[ListMemoriesConfig, ListMemoriesConfigDict] -class GenerateMemoriesRequestVertexSessionSource(_common.BaseModel): - """The vertex session source for generating memories.""" - end_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. End time (exclusive) of the time range. If not set, the end time is unbounded.""", - ) - session: Optional[str] = Field( - default=None, - description="""Required. The resource name of the Session to generate memories for. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}`""", - ) - start_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. Time range to define which session events should be used to generate memories. Start time (inclusive) of the time range. If not set, the start time is unbounded.""", - ) +class _ListMemoriesRequestParameters(_common.BaseModel): + """Parameters for listing memories.""" + name: Optional[str] = Field( + default=None, description="""Name of the Memory Bank.""" + ) + config: Optional[ListMemoriesConfig] = Field(default=None, description="""""") -class GenerateMemoriesRequestVertexSessionSourceDict(TypedDict, total=False): - """The vertex session source for generating memories.""" - end_time: Optional[datetime.datetime] - """Optional. End time (exclusive) of the time range. If not set, the end time is unbounded.""" +class _ListMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for listing memories.""" - session: Optional[str] - """Required. The resource name of the Session to generate memories for. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}`""" + name: Optional[str] + """Name of the Memory Bank.""" - start_time: Optional[datetime.datetime] - """Optional. Time range to define which session events should be used to generate memories. Start time (inclusive) of the time range. If not set, the start time is unbounded.""" + config: Optional[ListMemoriesConfigDict] + """""" -GenerateMemoriesRequestVertexSessionSourceOrDict = Union[ - GenerateMemoriesRequestVertexSessionSource, - GenerateMemoriesRequestVertexSessionSourceDict, +_ListMemoriesRequestParametersOrDict = Union[ + _ListMemoriesRequestParameters, _ListMemoriesRequestParametersDict ] -class GenerateMemoriesRequestDirectContentsSourceEvent(_common.BaseModel): +class ListMemoriesResponse(_common.BaseModel): + """Response for listing memories.""" - content: Optional[genai_types.Content] = Field( - default=None, - description="""Required. A single piece of content from which to generate 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 memories.""" ) -class GenerateMemoriesRequestDirectContentsSourceEventDict(TypedDict, total=False): +class ListMemoriesResponseDict(TypedDict, total=False): + """Response for listing memories.""" - content: Optional[genai_types.Content] - """Required. A single piece of content from which to generate memories.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + next_page_token: Optional[str] + """""" -GenerateMemoriesRequestDirectContentsSourceEventOrDict = Union[ - GenerateMemoriesRequestDirectContentsSourceEvent, - GenerateMemoriesRequestDirectContentsSourceEventDict, -] + memories: Optional[list[MemoryDict]] + """List of memories.""" -class GenerateMemoriesRequestDirectContentsSource(_common.BaseModel): - """The direct contents source for generating memories.""" +ListMemoriesResponseOrDict = Union[ListMemoriesResponse, ListMemoriesResponseDict] - events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEvent]] = Field( - default=None, - description="""Required. The source content (i.e. chat history) to generate memories from.""", + +class _GetMemoryOperationParameters(_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[GetMemoryBankOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -class GenerateMemoriesRequestDirectContentsSourceDict(TypedDict, total=False): - """The direct contents source for generating memories.""" +class _GetMemoryOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a memory as a response.""" - events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEventDict]] - """Required. The source content (i.e. chat history) to generate memories from.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" + config: Optional[GetMemoryBankOperationConfigDict] + """Used to override the default configuration.""" -GenerateMemoriesRequestDirectContentsSourceOrDict = Union[ - GenerateMemoriesRequestDirectContentsSource, - GenerateMemoriesRequestDirectContentsSourceDict, + +_GetMemoryOperationParametersOrDict = Union[ + _GetMemoryOperationParameters, _GetMemoryOperationParametersDict ] -class GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(_common.BaseModel): - """A direct memory to upload to Memory Bank.""" +class _GetGenerateMemoriesOperationParameters(_common.BaseModel): + """Parameters for getting an operation with generated memories as a response.""" - fact: Optional[str] = Field( - default=None, - description="""Required. The fact to consolidate with existing memories.""", + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - topics: Optional[list[MemoryTopicId]] = Field( - default=None, - description="""Optional. The topics that the consolidated memories should be associated with.""", + config: Optional[GetMemoryBankOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -class GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict( - TypedDict, total=False -): - """A direct memory to upload to Memory Bank.""" +class _GetGenerateMemoriesOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with generated memories as a response.""" - fact: Optional[str] - """Required. The fact to consolidate with existing memories.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - topics: Optional[list[MemoryTopicIdDict]] - """Optional. The topics that the consolidated memories should be associated with.""" + config: Optional[GetMemoryBankOperationConfigDict] + """Used to override the default configuration.""" -GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryOrDict = Union[ - GenerateMemoriesRequestDirectMemoriesSourceDirectMemory, - GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict, +_GetGenerateMemoriesOperationParametersOrDict = Union[ + _GetGenerateMemoriesOperationParameters, _GetGenerateMemoriesOperationParametersDict ] -class GenerateMemoriesRequestDirectMemoriesSource(_common.BaseModel): - """The direct memories source for generating memories.""" +class RetrieveMemoriesRequestSimilaritySearchParams(_common.BaseModel): + """The parameters for semantic similarity search based retrieval.""" - direct_memories: Optional[ - list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemory] - ] = Field( + search_query: Optional[str] = Field( default=None, - description="""Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""", + description="""Required. Query to use for similarity search retrieval. If provided, then the parent ReasoningEngine must have ReasoningEngineContextSpec.MemoryBankConfig.SimilaritySearchConfig set.""", + ) + top_k: Optional[int] = Field( + default=None, + description="""Optional. The maximum number of memories to return. The service may return fewer than this value. If unspecified, at most 3 memories will be returned. The maximum value is 100; values above 100 will be coerced to 100.""", ) -class GenerateMemoriesRequestDirectMemoriesSourceDict(TypedDict, total=False): - """The direct memories source for generating memories.""" +class RetrieveMemoriesRequestSimilaritySearchParamsDict(TypedDict, total=False): + """The parameters for semantic similarity search based retrieval.""" - direct_memories: Optional[ - list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict] - ] - """Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""" + search_query: Optional[str] + """Required. Query to use for similarity search retrieval. If provided, then the parent ReasoningEngine must have ReasoningEngineContextSpec.MemoryBankConfig.SimilaritySearchConfig set.""" + top_k: Optional[int] + """Optional. The maximum number of memories to return. The service may return fewer than this value. If unspecified, at most 3 memories will be returned. The maximum value is 100; values above 100 will be coerced to 100.""" -GenerateMemoriesRequestDirectMemoriesSourceOrDict = Union[ - GenerateMemoriesRequestDirectMemoriesSource, - GenerateMemoriesRequestDirectMemoriesSourceDict, -] +RetrieveMemoriesRequestSimilaritySearchParamsOrDict = Union[ + RetrieveMemoriesRequestSimilaritySearchParams, + RetrieveMemoriesRequestSimilaritySearchParamsDict, +] -class GenerateMemoriesConfig(_common.BaseModel): - """Config for generating memories.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - disable_consolidation: Optional[bool] = Field( - default=None, - description="""Whether to disable consolidation of memories. +class RetrieveMemoriesRequestSimpleRetrievalParams(_common.BaseModel): + """The parameters for simple (non-similarity search) retrieval.""" - If true, generated memories will not be consolidated with existing - memories; all generated memories will be added as new memories regardless - of whether they are duplicates of or contradictory to existing memories. - By default, memory consolidation is enabled.""", - ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", - ) - revision_labels: Optional[dict[str, str]] = Field( - default=None, - description="""Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""", - ) - revision_expire_time: Optional[datetime.datetime] = Field( + page_size: Optional[int] = 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.""", + description="""Optional. The maximum number of memories to return. The service may return fewer than this value. If unspecified, at most 3 memories will be returned. The maximum value is 100; values above 100 will be coerced to 100.""", ) - revision_ttl: Optional[str] = Field( + page_token: Optional[str] = Field( default=None, - description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", + description="""Optional. A page token, received from a previous `RetrieveMemories` call. Provide this to retrieve the subsequent page.""", ) - disable_memory_revisions: Optional[bool] = Field( + + +class RetrieveMemoriesRequestSimpleRetrievalParamsDict(TypedDict, total=False): + """The parameters for simple (non-similarity search) retrieval.""" + + page_size: Optional[int] + """Optional. The maximum number of memories to return. The service may return fewer than this value. If unspecified, at most 3 memories will be returned. The maximum value is 100; values above 100 will be coerced to 100.""" + + page_token: Optional[str] + """Optional. A page token, received from a previous `RetrieveMemories` call. Provide this to retrieve the subsequent page.""" + + +RetrieveMemoriesRequestSimpleRetrievalParamsOrDict = Union[ + RetrieveMemoriesRequestSimpleRetrievalParams, + RetrieveMemoriesRequestSimpleRetrievalParamsDict, +] + + +class MemoryFilter(_common.BaseModel): + """Filter to apply when retrieving memories.""" + + key: Optional[str] = Field( default=None, - description="""Optional. Input only. If true, no revisions will be created for this request.""", + description="""Represents the key of the filter. For example, "author" would apply to `metadata` entries with the key "author".""", ) - metadata: Optional[dict[str, MemoryMetadataValue]] = Field( - default=None, - description="""Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""", + negate: Optional[bool] = Field( + default=None, description="""Indicates whether the filter will be negated.""" ) - metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] = Field( + op: Optional[Operator] = Field( default=None, - description="""Optional. The strategy to use when applying metadata to existing memories.""", + description="""Represents the operator to apply to the filter. If not set, then EQUAL will be used.""", ) - allowed_topics: Optional[list[MemoryTopicId]] = Field( - default=None, - description="""Optional. Restricts memory generation to a subset of memory topics.""", + value: Optional[MemoryMetadataValue] = Field( + default=None, description="""Represents the value to compare to.""" ) -class GenerateMemoriesConfigDict(TypedDict, total=False): - """Config for generating memories.""" +class MemoryFilterDict(TypedDict, total=False): + """Filter to apply when retrieving memories.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + key: Optional[str] + """Represents the key of the filter. For example, "author" would apply to `metadata` entries with the key "author".""" - disable_consolidation: Optional[bool] - """Whether to disable consolidation of memories. + negate: Optional[bool] + """Indicates whether the filter will be negated.""" - If true, generated memories will not be consolidated with existing - memories; all generated memories will be added as new memories regardless - of whether they are duplicates of or contradictory to existing memories. - By default, memory consolidation is enabled.""" + op: Optional[Operator] + """Represents the operator to apply to the filter. If not set, then EQUAL will be used.""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + value: Optional[MemoryMetadataValueDict] + """Represents the value to compare to.""" - revision_labels: Optional[dict[str, str]] - """Labels to apply to the memory revision. For example, you can use this to label a revision with its data source.""" - 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.""" +MemoryFilterOrDict = Union[MemoryFilter, MemoryFilterDict] - 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 revisions will be created for this request.""" +class MemoryConjunctionFilter(_common.BaseModel): + """The conjunction filter for memories.""" - metadata: Optional[dict[str, MemoryMetadataValueDict]] - """Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""" + filters: Optional[list[MemoryFilter]] = Field( + default=None, + description="""Represents filters that will be combined using AND logic.""", + ) - metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] - """Optional. The strategy to use when applying metadata to existing memories.""" - allowed_topics: Optional[list[MemoryTopicIdDict]] - """Optional. Restricts memory generation to a subset of memory topics.""" +class MemoryConjunctionFilterDict(TypedDict, total=False): + """The conjunction filter for memories.""" + + filters: Optional[list[MemoryFilterDict]] + """Represents filters that will be combined using AND logic.""" -GenerateMemoriesConfigOrDict = Union[GenerateMemoriesConfig, GenerateMemoriesConfigDict] +MemoryConjunctionFilterOrDict = Union[ + MemoryConjunctionFilter, MemoryConjunctionFilterDict +] -class _GenerateMemoriesRequestParameters(_common.BaseModel): - """Parameters for generating memories.""" +class RetrieveMemoriesConfig(_common.BaseModel): + """Config for retrieving memories.""" - name: Optional[str] = Field( - default=None, - description="""Name of the Memory Bank to generate memories with.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSource] = Field( + filter: Optional[str] = Field( default=None, - description="""The vertex session source of the memories that should be generated.""", - ) - direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSource] = ( - Field( - default=None, - description="""The direct contents source of the memories that should be generated.""", - ) + description="""The standard list filter that will be applied to the retrieved + memories. More detail in [AIP-160](https://google.aip.dev/160). + + Supported fields: + * `fact` + * `create_time` + * `update_time` + """, ) - direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSource] = ( - Field( - default=None, - description="""The direct memories source of the memories that should be generated.""", - ) - ) - scope: Optional[dict[str, str]] = Field( + filter_groups: Optional[list[MemoryConjunctionFilter]] = Field( default=None, - description="""The scope of the memories that should be generated. + description="""Metadata filters that will be applied to the retrieved memories' + `metadata` using OR logic. Filters are defined using disjunctive normal + form (OR of ANDs). - Memories will be consolidated across memories with the same scope. Must be - provided unless the scope is defined in the source content. If `scope` is - provided, it will override the scope defined in the source content. Scope - values cannot contain the wildcard character '*'.""", + For example: + `filter_groups: [{filters: [{key: "author", value: {string_value: "agent + `123"}, op: EQUAL}]}, {filters: [{key: "label", value: {string_value: + "travel"}, op: EQUAL}, {key: "author", value: {string_value: "agent 321"}, + op: EQUAL}]}]` + + would be equivalent to the logical expression: + `(metadata.author = "agent 123" OR (metadata.label = "travel" AND + metadata.author = "agent 321"))`. + """, + ) + memory_types: Optional[list[MemoryType]] = Field( + default=None, + description="""Specifies the types of memories to retrieve. If this field is empty + or not provided, the request will default to retrieving only memories of + type `NATURAL_LANGUAGE_COLLECTION`. If populated, the request will + retrieve memories matching any of the specified `MemoryType` values.""", ) - config: Optional[GenerateMemoriesConfig] = Field(default=None, description="""""") -class _GenerateMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for generating memories.""" +class RetrieveMemoriesConfigDict(TypedDict, total=False): + """Config for retrieving memories.""" - name: Optional[str] - """Name of the Memory Bank to generate memories with.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSourceDict] - """The vertex session source of the memories that should be generated.""" + filter: Optional[str] + """The standard list filter that will be applied to the retrieved + memories. More detail in [AIP-160](https://google.aip.dev/160). - direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSourceDict] - """The direct contents source of the memories that should be generated.""" + Supported fields: + * `fact` + * `create_time` + * `update_time` + """ - direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSourceDict] - """The direct memories source of the memories that should be generated.""" + filter_groups: Optional[list[MemoryConjunctionFilterDict]] + """Metadata filters that will be applied to the retrieved memories' + `metadata` using OR logic. Filters are defined using disjunctive normal + form (OR of ANDs). - scope: Optional[dict[str, str]] - """The scope of the memories that should be generated. + For example: + `filter_groups: [{filters: [{key: "author", value: {string_value: "agent + `123"}, op: EQUAL}]}, {filters: [{key: "label", value: {string_value: + "travel"}, op: EQUAL}, {key: "author", value: {string_value: "agent 321"}, + op: EQUAL}]}]` - Memories will be consolidated across memories with the same scope. Must be - provided unless the scope is defined in the source content. If `scope` is - provided, it will override the scope defined in the source content. Scope - values cannot contain the wildcard character '*'.""" + would be equivalent to the logical expression: + `(metadata.author = "agent 123" OR (metadata.label = "travel" AND + metadata.author = "agent 321"))`. + """ - config: Optional[GenerateMemoriesConfigDict] - """""" + memory_types: Optional[list[MemoryType]] + """Specifies the types of memories to retrieve. If this field is empty + or not provided, the request will default to retrieving only memories of + type `NATURAL_LANGUAGE_COLLECTION`. If populated, the request will + retrieve memories matching any of the specified `MemoryType` values.""" -_GenerateMemoriesRequestParametersOrDict = Union[ - _GenerateMemoriesRequestParameters, _GenerateMemoriesRequestParametersDict -] +RetrieveMemoriesConfigOrDict = Union[RetrieveMemoriesConfig, RetrieveMemoriesConfigDict] -class GenerateMemoriesResponseGeneratedMemory(_common.BaseModel): - """A memory that was generated.""" +class _RetrieveMemoriesRequestParameters(_common.BaseModel): + """Parameters for retrieving memories.""" - memory: Optional[Memory] = Field( - default=None, description="""The generated memory.""" + name: Optional[str] = Field( + default=None, + description="""Name of the Memory Bank to retrieve memories from.""", ) - action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] = Field( - default=None, description="""The action to take.""" + scope: Optional[dict[str, str]] = Field( + default=None, + description="""The scope of the memories to retrieve. + + A memory must have exactly the same scope as the scope provided here to be + retrieved (i.e. same keys and values). Order does not matter, but it is + case-sensitive.""", ) - previous_revision: Optional[str] = Field( + similarity_search_params: Optional[ + RetrieveMemoriesRequestSimilaritySearchParams + ] = Field( default=None, - description="""The previous revision of the Memory before the action was performed. This - field is only set if the action is `UPDATED` or `DELETED`. You can use - this to rollback the Memory to the previous revision, undoing the action. - Format: - `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{revision}`""", + description="""Parameters for semantic similarity search based retrieval.""", + ) + simple_retrieval_params: Optional[RetrieveMemoriesRequestSimpleRetrievalParams] = ( + Field( + default=None, + description="""Parameters for simple (non-similarity search) retrieval.""", + ) ) + config: Optional[RetrieveMemoriesConfig] = Field(default=None, description="""""") -class GenerateMemoriesResponseGeneratedMemoryDict(TypedDict, total=False): - """A memory that was generated.""" +class _RetrieveMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for retrieving memories.""" - memory: Optional[MemoryDict] - """The generated memory.""" + name: Optional[str] + """Name of the Memory Bank to retrieve memories from.""" - action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] - """The action to take.""" + scope: Optional[dict[str, str]] + """The scope of the memories to retrieve. - previous_revision: Optional[str] - """The previous revision of the Memory before the action was performed. This - field is only set if the action is `UPDATED` or `DELETED`. You can use - this to rollback the Memory to the previous revision, undoing the action. - Format: - `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{revision}`""" + A memory must have exactly the same scope as the scope provided here to be + retrieved (i.e. same keys and values). Order does not matter, but it is + case-sensitive.""" + + similarity_search_params: Optional[ + RetrieveMemoriesRequestSimilaritySearchParamsDict + ] + """Parameters for semantic similarity search based retrieval.""" + simple_retrieval_params: Optional[RetrieveMemoriesRequestSimpleRetrievalParamsDict] + """Parameters for simple (non-similarity search) retrieval.""" -GenerateMemoriesResponseGeneratedMemoryOrDict = Union[ - GenerateMemoriesResponseGeneratedMemory, GenerateMemoriesResponseGeneratedMemoryDict + config: Optional[RetrieveMemoriesConfigDict] + """""" + + +_RetrieveMemoriesRequestParametersOrDict = Union[ + _RetrieveMemoriesRequestParameters, _RetrieveMemoriesRequestParametersDict ] -class GenerateMemoriesResponse(_common.BaseModel): - """The response for generating memories.""" +class RetrieveMemoriesResponseRetrievedMemory(_common.BaseModel): + """A retrieved memory.""" - generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemory]] = Field( - default=None, description="""The generated memories.""" + distance: Optional[float] = Field( + default=None, + description="""The distance between the query and the retrieved Memory. Smaller values indicate more similar memories. This is only set if similarity search was used for retrieval.""", + ) + memory: Optional[Memory] = Field( + default=None, description="""The retrieved Memory.""" ) -class GenerateMemoriesResponseDict(TypedDict, total=False): - """The response for generating memories.""" +class RetrieveMemoriesResponseRetrievedMemoryDict(TypedDict, total=False): + """A retrieved memory.""" - generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemoryDict]] - """The generated memories.""" + distance: Optional[float] + """The distance between the query and the retrieved Memory. Smaller values indicate more similar memories. This is only set if similarity search was used for retrieval.""" + memory: Optional[MemoryDict] + """The retrieved Memory.""" -GenerateMemoriesResponseOrDict = Union[ - GenerateMemoriesResponse, GenerateMemoriesResponseDict + +RetrieveMemoriesResponseRetrievedMemoryOrDict = Union[ + RetrieveMemoriesResponseRetrievedMemory, RetrieveMemoriesResponseRetrievedMemoryDict ] -class GenerateMemoriesOperation(_common.BaseModel): - """Operation that generates memories with a Memory Bank.""" +class RetrieveMemoriesResponse(_common.BaseModel): + """The response for retrieving memories.""" - 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( + next_page_token: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""A token that can be sent as `page_token` to retrieve the next page. If this field is omitted, there are no subsequent pages. This token is not set if similarity search was used for retrieval.""", ) - response: Optional[GenerateMemoriesResponse] = Field( - default=None, description="""The response for generating memories.""" + retrieved_memories: Optional[list[RetrieveMemoriesResponseRetrievedMemory]] = Field( + default=None, description="""The retrieved memories.""" ) -class GenerateMemoriesOperationDict(TypedDict, total=False): - """Operation that generates memories with a Memory Bank.""" - - 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 RetrieveMemoriesResponseDict(TypedDict, total=False): + """The response for retrieving memories.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + next_page_token: Optional[str] + """A token that can be sent as `page_token` to retrieve the next page. If this field is omitted, there are no subsequent pages. This token is not set if similarity search was used for retrieval.""" - response: Optional[GenerateMemoriesResponseDict] - """The response for generating memories.""" + retrieved_memories: Optional[list[RetrieveMemoriesResponseRetrievedMemoryDict]] + """The retrieved memories.""" -GenerateMemoriesOperationOrDict = Union[ - GenerateMemoriesOperation, GenerateMemoriesOperationDict +RetrieveMemoriesResponseOrDict = Union[ + RetrieveMemoriesResponse, RetrieveMemoriesResponseDict ] -class GetMemoryConfig(_common.BaseModel): - """Config for getting a Memory.""" +class RetrieveMemoryProfilesConfig(_common.BaseModel): + """Config for retrieving memory profiles.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetMemoryConfigDict(TypedDict, total=False): - """Config for getting a Memory.""" +class RetrieveMemoryProfilesConfigDict(TypedDict, total=False): + """Config for retrieving memory profiles.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetMemoryConfigOrDict = Union[GetMemoryConfig, GetMemoryConfigDict] - - -class _GetMemoryRequestParameters(_common.BaseModel): - """Parameters for getting a Memory.""" - - name: Optional[str] = Field(default=None, description="""Name of the memory.""") - config: Optional[GetMemoryConfig] = Field(default=None, description="""""") +RetrieveMemoryProfilesConfigOrDict = Union[ + RetrieveMemoryProfilesConfig, RetrieveMemoryProfilesConfigDict +] -class _GetMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for getting a Memory.""" +class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): + """Parameters for retrieving memory profiles.""" - name: Optional[str] - """Name of the memory.""" - - config: Optional[GetMemoryConfigDict] - """""" - - -_GetMemoryRequestParametersOrDict = Union[ - _GetMemoryRequestParameters, _GetMemoryRequestParametersDict -] - - -class ListMemoriesConfig(_common.BaseModel): - """Config for listing memories.""" - - 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( + name: 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.""", + description="""Name of the Memory Bank to retrieve memory profiles from.""", ) - order_by: Optional[str] = Field( + scope: Optional[dict[str, str]] = Field( default=None, - description="""The standard list order by string. If not specified, the default - order is `create_time desc`. If specified, the default sorting order of - provided fields is ascending. More detail in - [AIP-132](https://google.aip.dev/132). + description="""The scope of the memories to retrieve. - Supported fields: - * `create_time` - * `update_time`""", + A memory must have exactly the same scope as the scope provided here to be + retrieved (i.e. same keys and values). Order does not matter, but it is + case-sensitive.""", + ) + config: Optional[RetrieveMemoryProfilesConfig] = Field( + default=None, description="""""" ) -class ListMemoriesConfigDict(TypedDict, total=False): - """Config for listing memories.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - page_size: Optional[int] - """""" +class _RetrieveMemoryProfilesRequestParametersDict(TypedDict, total=False): + """Parameters for retrieving memory profiles.""" - page_token: Optional[str] - """""" + name: Optional[str] + """Name of the Memory Bank to retrieve memory profiles from.""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" + scope: Optional[dict[str, str]] + """The scope of the memories to retrieve. - order_by: Optional[str] - """The standard list order by string. If not specified, the default - order is `create_time desc`. If specified, the default sorting order of - provided fields is ascending. More detail in - [AIP-132](https://google.aip.dev/132). + A memory must have exactly the same scope as the scope provided here to be + retrieved (i.e. same keys and values). Order does not matter, but it is + case-sensitive.""" - Supported fields: - * `create_time` - * `update_time`""" + config: Optional[RetrieveMemoryProfilesConfigDict] + """""" -ListMemoriesConfigOrDict = Union[ListMemoriesConfig, ListMemoriesConfigDict] +_RetrieveMemoryProfilesRequestParametersOrDict = Union[ + _RetrieveMemoryProfilesRequestParameters, + _RetrieveMemoryProfilesRequestParametersDict, +] -class _ListMemoriesRequestParameters(_common.BaseModel): - """Parameters for listing memories.""" +class MemoryProfile(_common.BaseModel): + """A memory profile.""" - name: Optional[str] = Field( - default=None, description="""Name of the Memory Bank.""" + schema_id: Optional[str] = Field( + default=None, + description="""Represents the ID of the schema. This ID corresponds to the `schema_id` defined inside the SchemaConfig, under StructuredMemoryCustomizationConfig.""", + ) + profile: Optional[dict[str, Any]] = Field( + default=None, description="""Represents the profile data.""" ) - config: Optional[ListMemoriesConfig] = Field(default=None, description="""""") -class _ListMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for listing memories.""" +class MemoryProfileDict(TypedDict, total=False): + """A memory profile.""" - name: Optional[str] - """Name of the Memory Bank.""" + schema_id: Optional[str] + """Represents the ID of the schema. This ID corresponds to the `schema_id` defined inside the SchemaConfig, under StructuredMemoryCustomizationConfig.""" - config: Optional[ListMemoriesConfigDict] - """""" + profile: Optional[dict[str, Any]] + """Represents the profile data.""" -_ListMemoriesRequestParametersOrDict = Union[ - _ListMemoriesRequestParameters, _ListMemoriesRequestParametersDict -] +MemoryProfileOrDict = Union[MemoryProfile, MemoryProfileDict] -class ListMemoriesResponse(_common.BaseModel): - """Response for listing memories.""" +class RetrieveProfilesResponse(_common.BaseModel): + """The response for retrieving memory profiles.""" - 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 memories.""" + profiles: Optional[dict[str, MemoryProfile]] = Field( + default=None, + description="""The retrieved structured profiles, which match the schemas under the + requested scope. The key is the ID of the schema that the profile is + linked with, which corresponds to the `schema_id` defined inside the + `SchemaConfig`, under `StructuredMemoryCustomizationConfig`.""", ) -class ListMemoriesResponseDict(TypedDict, total=False): - """Response for listing memories.""" - - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" - - next_page_token: Optional[str] - """""" +class RetrieveProfilesResponseDict(TypedDict, total=False): + """The response for retrieving memory profiles.""" - memories: Optional[list[MemoryDict]] - """List of memories.""" + profiles: Optional[dict[str, MemoryProfileDict]] + """The retrieved structured profiles, which match the schemas under the + requested scope. The key is the ID of the schema that the profile is + linked with, which corresponds to the `schema_id` defined inside the + `SchemaConfig`, under `StructuredMemoryCustomizationConfig`.""" -ListMemoriesResponseOrDict = Union[ListMemoriesResponse, ListMemoriesResponseDict] +RetrieveProfilesResponseOrDict = Union[ + RetrieveProfilesResponse, RetrieveProfilesResponseDict +] -class _GetMemoryOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a memory as a response.""" +class RollbackMemoryConfig(_common.BaseModel): + """Config for rolling back a memory.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - config: Optional[GetMemoryBankOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) -class _GetMemoryOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a memory as a response.""" +class RollbackMemoryConfigDict(TypedDict, total=False): + """Config for rolling back a memory.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - config: Optional[GetMemoryBankOperationConfigDict] - """Used to override the default configuration.""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" -_GetMemoryOperationParametersOrDict = Union[ - _GetMemoryOperationParameters, _GetMemoryOperationParametersDict -] +RollbackMemoryConfigOrDict = Union[RollbackMemoryConfig, RollbackMemoryConfigDict] -class _GetGenerateMemoriesOperationParameters(_common.BaseModel): - """Parameters for getting an operation with generated memories as a response.""" +class _RollbackMemoryRequestParameters(_common.BaseModel): + """Parameters for generating memories.""" - 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 memory to rollback.""" ) - config: Optional[GetMemoryBankOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + target_revision_id: Optional[str] = Field( + default=None, description="""The ID of the revision to rollback to.""" ) + config: Optional[RollbackMemoryConfig] = Field(default=None, description="""""") -class _GetGenerateMemoriesOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with generated memories as a response.""" +class _RollbackMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for generating memories.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + name: Optional[str] + """Name of the memory to rollback.""" - config: Optional[GetMemoryBankOperationConfigDict] - """Used to override the default configuration.""" + target_revision_id: Optional[str] + """The ID of the revision to rollback to.""" + + config: Optional[RollbackMemoryConfigDict] + """""" -_GetGenerateMemoriesOperationParametersOrDict = Union[ - _GetGenerateMemoriesOperationParameters, _GetGenerateMemoriesOperationParametersDict +_RollbackMemoryRequestParametersOrDict = Union[ + _RollbackMemoryRequestParameters, _RollbackMemoryRequestParametersDict ] -class RetrieveMemoriesRequestSimilaritySearchParams(_common.BaseModel): - """The parameters for semantic similarity search based retrieval.""" +class RollbackMemoryOperation(_common.BaseModel): + """Operation that rolls back a memory.""" - search_query: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Required. Query to use for similarity search retrieval. If provided, then the parent ReasoningEngine must have ReasoningEngineContextSpec.MemoryBankConfig.SimilaritySearchConfig set.""", + 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}`.""", ) - top_k: Optional[int] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. The maximum number of memories to return. The service may return fewer than this value. If unspecified, at most 3 memories will be returned. The maximum value is 100; values above 100 will be coerced to 100.""", + 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 RetrieveMemoriesRequestSimilaritySearchParamsDict(TypedDict, total=False): - """The parameters for semantic similarity search based retrieval.""" +class RollbackMemoryOperationDict(TypedDict, total=False): + """Operation that rolls back a memory.""" - search_query: Optional[str] - """Required. Query to use for similarity search retrieval. If provided, then the parent ReasoningEngine must have ReasoningEngineContextSpec.MemoryBankConfig.SimilaritySearchConfig set.""" + 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}`.""" - top_k: Optional[int] - """Optional. The maximum number of memories to return. The service may return fewer than this value. If unspecified, at most 3 memories will be returned. The maximum value is 100; values above 100 will be coerced to 100.""" + 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.""" -RetrieveMemoriesRequestSimilaritySearchParamsOrDict = Union[ - RetrieveMemoriesRequestSimilaritySearchParams, - RetrieveMemoriesRequestSimilaritySearchParamsDict, + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +RollbackMemoryOperationOrDict = Union[ + RollbackMemoryOperation, RollbackMemoryOperationDict ] -class RetrieveMemoriesRequestSimpleRetrievalParams(_common.BaseModel): - """The parameters for simple (non-similarity search) retrieval.""" +class UpdateMemoryConfig(_common.BaseModel): + """Config for updating a memory.""" - page_size: Optional[int] = Field( - default=None, - description="""Optional. The maximum number of memories to return. The service may return fewer than this value. If unspecified, at most 3 memories will be returned. The maximum value is 100; values above 100 will be coerced to 100.""", - ) - page_token: Optional[str] = Field( + 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. A page token, received from a previous `RetrieveMemories` call. Provide this to retrieve the subsequent page.""", + 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.""", + ) + 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.""", ) -class RetrieveMemoriesRequestSimpleRetrievalParamsDict(TypedDict, total=False): - """The parameters for simple (non-similarity search) retrieval.""" +class UpdateMemoryConfigDict(TypedDict, total=False): + """Config for updating a memory.""" - page_size: Optional[int] - """Optional. The maximum number of memories to return. The service may return fewer than this value. If unspecified, at most 3 memories will be returned. The maximum value is 100; values above 100 will be coerced to 100.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - page_token: Optional[str] - """Optional. A page token, received from a previous `RetrieveMemories` call. Provide this to retrieve the subsequent page.""" + 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.""" -RetrieveMemoriesRequestSimpleRetrievalParamsOrDict = Union[ - RetrieveMemoriesRequestSimpleRetrievalParams, - RetrieveMemoriesRequestSimpleRetrievalParamsDict, -] + ttl: Optional[str] + """Optional. Input only. The TTL for this resource. + The expiration time is computed: now + TTL.""" -class MemoryFilter(_common.BaseModel): - """Filter to apply when retrieving memories.""" + 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.""" - key: Optional[str] = Field( - default=None, - description="""Represents the key of the filter. For example, "author" would apply to `metadata` entries with the key "author".""", - ) - negate: Optional[bool] = Field( - default=None, description="""Indicates whether the filter will be negated.""" - ) - op: Optional[Operator] = Field( - default=None, - description="""Represents the operator to apply to the filter. If not set, then EQUAL will be used.""", - ) - value: Optional[MemoryMetadataValue] = Field( - default=None, description="""Represents the value to compare to.""" - ) + 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.""" -class MemoryFilterDict(TypedDict, total=False): - """Filter to apply when retrieving memories.""" + disable_memory_revisions: Optional[bool] + """Optional. Input only. If true, no revision will be created for this request.""" - key: Optional[str] - """Represents the key of the filter. For example, "author" would apply to `metadata` entries with the key "author".""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The topics of the memory.""" - negate: Optional[bool] - """Indicates whether the filter will be negated.""" + 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.""" - op: Optional[Operator] - """Represents the operator to apply to the filter. If not set, then EQUAL will be used.""" + 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.""" - value: Optional[MemoryMetadataValueDict] - """Represents the value to compare to.""" + update_mask: Optional[str] + """The update mask to apply. For the `FieldMask` definition, see + https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" -MemoryFilterOrDict = Union[MemoryFilter, MemoryFilterDict] +UpdateMemoryConfigOrDict = Union[UpdateMemoryConfig, UpdateMemoryConfigDict] -class MemoryConjunctionFilter(_common.BaseModel): - """The conjunction filter for memories.""" +class _UpdateMemoryRequestParameters(_common.BaseModel): + """Parameters for updating memories.""" - filters: Optional[list[MemoryFilter]] = Field( + name: Optional[str] = Field( + default=None, description="""Name of the memory to update.""" + ) + fact: Optional[str] = Field( default=None, - description="""Represents filters that will be combined using AND logic.""", + description="""The updated fact of the memory. + + This is the semantic knowledge extracted from the source content.""", ) + scope: Optional[dict[str, str]] = Field( + default=None, + description="""The updated 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[UpdateMemoryConfig] = Field(default=None, description="""""") -class MemoryConjunctionFilterDict(TypedDict, total=False): - """The conjunction filter for memories.""" - filters: Optional[list[MemoryFilterDict]] - """Represents filters that will be combined using AND logic.""" +class _UpdateMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for updating memories.""" + name: Optional[str] + """Name of the memory to update.""" -MemoryConjunctionFilterOrDict = Union[ - MemoryConjunctionFilter, MemoryConjunctionFilterDict + fact: Optional[str] + """The updated fact of the memory. + + This is the semantic knowledge extracted from the source content.""" + + scope: Optional[dict[str, str]] + """The updated 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[UpdateMemoryConfigDict] + """""" + + +_UpdateMemoryRequestParametersOrDict = Union[ + _UpdateMemoryRequestParameters, _UpdateMemoryRequestParametersDict ] -class RetrieveMemoriesConfig(_common.BaseModel): - """Config for retrieving memories.""" +class PurgeMemoriesConfig(_common.BaseModel): + """Config for purging memories.""" 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.""", + ) + + +class PurgeMemoriesConfigDict(TypedDict, total=False): + """Config for purging memories.""" + + 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.""" + + +PurgeMemoriesConfigOrDict = Union[PurgeMemoriesConfig, PurgeMemoriesConfigDict] + + +class _PurgeMemoriesRequestParameters(_common.BaseModel): + """Parameters for purging memories.""" + + name: Optional[str] = Field( + default=None, description="""Name of the Memory Bank to purge memories from.""" + ) filter: Optional[str] = Field( default=None, - description="""The standard list filter that will be applied to the retrieved - memories. More detail in [AIP-160](https://google.aip.dev/160). - - Supported fields: - * `fact` - * `create_time` - * `update_time` - """, + description="""The standard list filter to determine which memories to purge. + More detail in [AIP-160](https://google.aip.dev/160).""", ) filter_groups: Optional[list[MemoryConjunctionFilter]] = Field( default=None, - description="""Metadata filters that will be applied to the retrieved memories' + description="""Metadata filters that will be applied to the memories' `metadata` using OR logic. Filters are defined using disjunctive normal form (OR of ANDs). @@ -13227,33 +13039,25 @@ class RetrieveMemoriesConfig(_common.BaseModel): metadata.author = "agent 321"))`. """, ) - memory_types: Optional[list[MemoryType]] = Field( + force: Optional[bool] = Field( default=None, - description="""Specifies the types of memories to retrieve. If this field is empty - or not provided, the request will default to retrieving only memories of - type `NATURAL_LANGUAGE_COLLECTION`. If populated, the request will - retrieve memories matching any of the specified `MemoryType` values.""", + description="""If true, the memories will actually be purged. If false, the purge request will be validated but not executed.""", ) + config: Optional[PurgeMemoriesConfig] = Field(default=None, description="""""") -class RetrieveMemoriesConfigDict(TypedDict, total=False): - """Config for retrieving memories.""" +class _PurgeMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for purging memories.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + name: Optional[str] + """Name of the Memory Bank to purge memories from.""" filter: Optional[str] - """The standard list filter that will be applied to the retrieved - memories. More detail in [AIP-160](https://google.aip.dev/160). - - Supported fields: - * `fact` - * `create_time` - * `update_time` - """ + """The standard list filter to determine which memories to purge. + More detail in [AIP-160](https://google.aip.dev/160).""" filter_groups: Optional[list[MemoryConjunctionFilterDict]] - """Metadata filters that will be applied to the retrieved memories' + """Metadata filters that will be applied to the memories' `metadata` using OR logic. Filters are defined using disjunctive normal form (OR of ANDs). @@ -13268,2234 +13072,1657 @@ class RetrieveMemoriesConfigDict(TypedDict, total=False): metadata.author = "agent 321"))`. """ - memory_types: Optional[list[MemoryType]] - """Specifies the types of memories to retrieve. If this field is empty - or not provided, the request will default to retrieving only memories of - type `NATURAL_LANGUAGE_COLLECTION`. If populated, the request will - retrieve memories matching any of the specified `MemoryType` values.""" + force: Optional[bool] + """If true, the memories will actually be purged. If false, the purge request will be validated but not executed.""" + config: Optional[PurgeMemoriesConfigDict] + """""" -RetrieveMemoriesConfigOrDict = Union[RetrieveMemoriesConfig, RetrieveMemoriesConfigDict] +_PurgeMemoriesRequestParametersOrDict = Union[ + _PurgeMemoriesRequestParameters, _PurgeMemoriesRequestParametersDict +] -class _RetrieveMemoriesRequestParameters(_common.BaseModel): - """Parameters for retrieving memories.""" - name: Optional[str] = Field( - default=None, - description="""Name of the Memory Bank to retrieve memories from.""", - ) - scope: Optional[dict[str, str]] = Field( - default=None, - description="""The scope of the memories to retrieve. +class PurgeMemoriesResponse(_common.BaseModel): + """The response for purging memories.""" - A memory must have exactly the same scope as the scope provided here to be - retrieved (i.e. same keys and values). Order does not matter, but it is - case-sensitive.""", - ) - similarity_search_params: Optional[ - RetrieveMemoriesRequestSimilaritySearchParams - ] = Field( - default=None, - description="""Parameters for semantic similarity search based retrieval.""", - ) - simple_retrieval_params: Optional[RetrieveMemoriesRequestSimpleRetrievalParams] = ( - Field( - default=None, - description="""Parameters for simple (non-similarity search) retrieval.""", - ) + purge_count: Optional[int] = Field( + default=None, description="""The number of memories that were purged.""" ) - config: Optional[RetrieveMemoriesConfig] = Field(default=None, description="""""") - - -class _RetrieveMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving memories.""" - - name: Optional[str] - """Name of the Memory Bank to retrieve memories from.""" - - scope: Optional[dict[str, str]] - """The scope of the memories to retrieve. - - A memory must have exactly the same scope as the scope provided here to be - retrieved (i.e. same keys and values). Order does not matter, but it is - case-sensitive.""" - similarity_search_params: Optional[ - RetrieveMemoriesRequestSimilaritySearchParamsDict - ] - """Parameters for semantic similarity search based retrieval.""" - simple_retrieval_params: Optional[RetrieveMemoriesRequestSimpleRetrievalParamsDict] - """Parameters for simple (non-similarity search) retrieval.""" +class PurgeMemoriesResponseDict(TypedDict, total=False): + """The response for purging memories.""" - config: Optional[RetrieveMemoriesConfigDict] - """""" + purge_count: Optional[int] + """The number of memories that were purged.""" -_RetrieveMemoriesRequestParametersOrDict = Union[ - _RetrieveMemoriesRequestParameters, _RetrieveMemoriesRequestParametersDict -] +PurgeMemoriesResponseOrDict = Union[PurgeMemoriesResponse, PurgeMemoriesResponseDict] -class RetrieveMemoriesResponseRetrievedMemory(_common.BaseModel): - """A retrieved memory.""" +class PurgeMemoriesOperation(_common.BaseModel): + """Operation that purges memories from a Memory Bank.""" - distance: Optional[float] = Field( + name: Optional[str] = Field( default=None, - description="""The distance between the query and the retrieved Memory. Smaller values indicate more similar memories. This is only set if similarity search was used for retrieval.""", + 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}`.""", ) - memory: Optional[Memory] = Field( - default=None, description="""The retrieved Memory.""" + 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.""", ) - - -class RetrieveMemoriesResponseRetrievedMemoryDict(TypedDict, total=False): - """A retrieved memory.""" - - distance: Optional[float] - """The distance between the query and the retrieved Memory. Smaller values indicate more similar memories. This is only set if similarity search was used for retrieval.""" - - memory: Optional[MemoryDict] - """The retrieved Memory.""" - - -RetrieveMemoriesResponseRetrievedMemoryOrDict = Union[ - RetrieveMemoriesResponseRetrievedMemory, RetrieveMemoriesResponseRetrievedMemoryDict -] - - -class RetrieveMemoriesResponse(_common.BaseModel): - """The response for retrieving memories.""" - - next_page_token: Optional[str] = Field( + done: Optional[bool] = Field( default=None, - description="""A token that can be sent as `page_token` to retrieve the next page. If this field is omitted, there are no subsequent pages. This token is not set if similarity search was used for retrieval.""", + 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.""", ) - retrieved_memories: Optional[list[RetrieveMemoriesResponseRetrievedMemory]] = Field( - default=None, description="""The retrieved memories.""" + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + response: Optional[PurgeMemoriesResponse] = Field( + default=None, description="""The response for purging memories.""" ) -class RetrieveMemoriesResponseDict(TypedDict, total=False): - """The response for retrieving memories.""" +class PurgeMemoriesOperationDict(TypedDict, total=False): + """Operation that purges memories from a Memory Bank.""" - next_page_token: Optional[str] - """A token that can be sent as `page_token` to retrieve the next page. If this field is omitted, there are no subsequent pages. This token is not set if similarity search was used for retrieval.""" + 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}`.""" - retrieved_memories: Optional[list[RetrieveMemoriesResponseRetrievedMemoryDict]] - """The retrieved memories.""" + 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.""" -RetrieveMemoriesResponseOrDict = Union[ - RetrieveMemoriesResponse, RetrieveMemoriesResponseDict -] + response: Optional[PurgeMemoriesResponseDict] + """The response for purging memories.""" -class RetrieveMemoryProfilesConfig(_common.BaseModel): - """Config for retrieving memory profiles.""" +PurgeMemoriesOperationOrDict = Union[PurgeMemoriesOperation, PurgeMemoriesOperationDict] + + +class GetMemoryRevisionConfig(_common.BaseModel): + """Config for getting a Memory Revision.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class RetrieveMemoryProfilesConfigDict(TypedDict, total=False): - """Config for retrieving memory profiles.""" +class GetMemoryRevisionConfigDict(TypedDict, total=False): + """Config for getting a Memory Revision.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -RetrieveMemoryProfilesConfigOrDict = Union[ - RetrieveMemoryProfilesConfig, RetrieveMemoryProfilesConfigDict +GetMemoryRevisionConfigOrDict = Union[ + GetMemoryRevisionConfig, GetMemoryRevisionConfigDict ] -class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): - """Parameters for retrieving memory profiles.""" +class _GetMemoryRevisionRequestParameters(_common.BaseModel): + """Parameters for getting a memory revision.""" name: Optional[str] = Field( - default=None, - description="""Name of the Memory Bank to retrieve memory profiles from.""", - ) - scope: Optional[dict[str, str]] = Field( - default=None, - description="""The scope of the memories to retrieve. - - A memory must have exactly the same scope as the scope provided here to be - retrieved (i.e. same keys and values). Order does not matter, but it is - case-sensitive.""", - ) - config: Optional[RetrieveMemoryProfilesConfig] = Field( - default=None, description="""""" + default=None, description="""Name of the Memory Revision.""" ) + config: Optional[GetMemoryRevisionConfig] = Field(default=None, description="""""") -class _RetrieveMemoryProfilesRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving memory profiles.""" +class _GetMemoryRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for getting a memory revision.""" name: Optional[str] - """Name of the Memory Bank to retrieve memory profiles from.""" - - scope: Optional[dict[str, str]] - """The scope of the memories to retrieve. - - A memory must have exactly the same scope as the scope provided here to be - retrieved (i.e. same keys and values). Order does not matter, but it is - case-sensitive.""" + """Name of the Memory Revision.""" - config: Optional[RetrieveMemoryProfilesConfigDict] + config: Optional[GetMemoryRevisionConfigDict] """""" -_RetrieveMemoryProfilesRequestParametersOrDict = Union[ - _RetrieveMemoryProfilesRequestParameters, - _RetrieveMemoryProfilesRequestParametersDict, +_GetMemoryRevisionRequestParametersOrDict = Union[ + _GetMemoryRevisionRequestParameters, _GetMemoryRevisionRequestParametersDict ] -class MemoryProfile(_common.BaseModel): - """A memory profile.""" +class IntermediateExtractedMemory(_common.BaseModel): + """An extracted memory that is the intermediate result before consolidation.""" - schema_id: Optional[str] = Field( + fact: Optional[str] = Field( default=None, - description="""Represents the ID of the schema. This ID corresponds to the `schema_id` defined inside the SchemaConfig, under StructuredMemoryCustomizationConfig.""", + description="""Output only. Represents the fact of the extracted memory.""", ) - profile: Optional[dict[str, Any]] = Field( - default=None, description="""Represents the profile data.""" + context: Optional[str] = Field( + default=None, + description="""Output only. Represents the explanation of why the information was extracted from the source content.""", + ) + structured_data: Optional[dict[str, Any]] = Field( + default=None, + description="""Output only. Represents the structured value of the extracted memory.""", ) -class MemoryProfileDict(TypedDict, total=False): - """A memory profile.""" +class IntermediateExtractedMemoryDict(TypedDict, total=False): + """An extracted memory that is the intermediate result before consolidation.""" - schema_id: Optional[str] - """Represents the ID of the schema. This ID corresponds to the `schema_id` defined inside the SchemaConfig, under StructuredMemoryCustomizationConfig.""" + fact: Optional[str] + """Output only. Represents the fact of the extracted memory.""" - profile: Optional[dict[str, Any]] - """Represents the profile data.""" + context: Optional[str] + """Output only. Represents the explanation of why the information was extracted from the source content.""" + structured_data: Optional[dict[str, Any]] + """Output only. Represents the structured value of the extracted memory.""" -MemoryProfileOrDict = Union[MemoryProfile, MemoryProfileDict] + +IntermediateExtractedMemoryOrDict = Union[ + IntermediateExtractedMemory, IntermediateExtractedMemoryDict +] -class RetrieveProfilesResponse(_common.BaseModel): - """The response for retrieving memory profiles.""" +class MemoryRevision(_common.BaseModel): + """A memory revision.""" - profiles: Optional[dict[str, MemoryProfile]] = Field( + create_time: Optional[datetime.datetime] = Field( default=None, - description="""The retrieved structured profiles, which match the schemas under the - requested scope. The key is the ID of the schema that the profile is - linked with, which corresponds to the `schema_id` defined inside the - `SchemaConfig`, under `StructuredMemoryCustomizationConfig`.""", + description="""Output only. Represents the timestamp when this Memory Revision was created.""", + ) + expire_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Represents the timestamp of when this resource is considered expired.""", + ) + extracted_memories: Optional[list[IntermediateExtractedMemory]] = Field( + default=None, + description="""Output only. Represents the extracted memories from the source content before consolidation when the memory was updated via GenerateMemories. This information was used to modify an existing Memory via Consolidation.""", + ) + fact: Optional[str] = Field( + default=None, + description="""Output only. Represents the fact of the Memory Revision. This corresponds to the `fact` field of the parent Memory at the time of revision creation.""", + ) + labels: Optional[dict[str, str]] = Field( + default=None, + description="""Output only. Represents the labels of the Memory Revision. These labels are applied to the MemoryRevision when it is created based on `GenerateMemoriesRequest.revision_labels`.""", + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. Represents the resource name of the Memory Revision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{memory_revision}`""", + ) + structured_data: Optional[dict[str, Any]] = Field( + default=None, + description="""Output only. Represents the structured value of the memory at the time of revision creation.""", ) -class RetrieveProfilesResponseDict(TypedDict, total=False): - """The response for retrieving memory profiles.""" - - profiles: Optional[dict[str, MemoryProfileDict]] - """The retrieved structured profiles, which match the schemas under the - requested scope. The key is the ID of the schema that the profile is - linked with, which corresponds to the `schema_id` defined inside the - `SchemaConfig`, under `StructuredMemoryCustomizationConfig`.""" +class MemoryRevisionDict(TypedDict, total=False): + """A memory revision.""" + create_time: Optional[datetime.datetime] + """Output only. Represents the timestamp when this Memory Revision was created.""" -RetrieveProfilesResponseOrDict = Union[ - RetrieveProfilesResponse, RetrieveProfilesResponseDict -] + expire_time: Optional[datetime.datetime] + """Output only. Represents the timestamp of when this resource is considered expired.""" + + extracted_memories: Optional[list[IntermediateExtractedMemoryDict]] + """Output only. Represents the extracted memories from the source content before consolidation when the memory was updated via GenerateMemories. This information was used to modify an existing Memory via Consolidation.""" + fact: Optional[str] + """Output only. Represents the fact of the Memory Revision. This corresponds to the `fact` field of the parent Memory at the time of revision creation.""" -class RollbackMemoryConfig(_common.BaseModel): - """Config for rolling back a memory.""" + labels: Optional[dict[str, str]] + """Output only. Represents the labels of the Memory Revision. These labels are applied to the MemoryRevision when it is created based on `GenerateMemoriesRequest.revision_labels`.""" + + name: Optional[str] + """Identifier. Represents the resource name of the Memory Revision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{memory_revision}`""" + + structured_data: Optional[dict[str, Any]] + """Output only. Represents the structured value of the memory at the time of revision creation.""" + + +MemoryRevisionOrDict = Union[MemoryRevision, MemoryRevisionDict] + + +class ListMemoryRevisionsConfig(_common.BaseModel): + """Config for listing memory revisions.""" 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.""", + 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 RollbackMemoryConfigDict(TypedDict, total=False): - """Config for rolling back a memory.""" +class ListMemoryRevisionsConfigDict(TypedDict, total=False): + """Config for listing memory revisions.""" 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.""" + 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.""" -RollbackMemoryConfigOrDict = Union[RollbackMemoryConfig, RollbackMemoryConfigDict] +ListMemoryRevisionsConfigOrDict = Union[ + ListMemoryRevisionsConfig, ListMemoryRevisionsConfigDict +] -class _RollbackMemoryRequestParameters(_common.BaseModel): - """Parameters for generating memories.""" - name: Optional[str] = Field( - default=None, description="""Name of the memory to rollback.""" - ) - target_revision_id: Optional[str] = Field( - default=None, description="""The ID of the revision to rollback to.""" +class _ListMemoryRevisionsRequestParameters(_common.BaseModel): + """Parameters for listing memory revisions.""" + + name: Optional[str] = Field(default=None, description="""Name of the memory""") + config: Optional[ListMemoryRevisionsConfig] = Field( + default=None, description="""""" ) - config: Optional[RollbackMemoryConfig] = Field(default=None, description="""""") -class _RollbackMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for generating memories.""" +class _ListMemoryRevisionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing memory revisions.""" name: Optional[str] - """Name of the memory to rollback.""" - - target_revision_id: Optional[str] - """The ID of the revision to rollback to.""" + """Name of the memory""" - config: Optional[RollbackMemoryConfigDict] + config: Optional[ListMemoryRevisionsConfigDict] """""" -_RollbackMemoryRequestParametersOrDict = Union[ - _RollbackMemoryRequestParameters, _RollbackMemoryRequestParametersDict +_ListMemoryRevisionsRequestParametersOrDict = Union[ + _ListMemoryRevisionsRequestParameters, _ListMemoryRevisionsRequestParametersDict ] -class RollbackMemoryOperation(_common.BaseModel): - """Operation that rolls back a memory.""" +class ListMemoryRevisionsResponse(_common.BaseModel): + """Response for listing memory revisions.""" - 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.""", + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - error: Optional[dict[str, Any]] = Field( - default=None, - description="""The error result of the operation in case of failure or cancellation.""", + next_page_token: Optional[str] = Field(default=None, description="""""") + memory_revisions: Optional[list[MemoryRevision]] = Field( + default=None, description="""List of memory revisions.""" ) -class RollbackMemoryOperationDict(TypedDict, total=False): - """Operation that rolls back a memory.""" - - 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 ListMemoryRevisionsResponseDict(TypedDict, total=False): + """Response for listing memory revisions.""" - 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.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - 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.""" + next_page_token: Optional[str] + """""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + memory_revisions: Optional[list[MemoryRevisionDict]] + """List of memory revisions.""" -RollbackMemoryOperationOrDict = Union[ - RollbackMemoryOperation, RollbackMemoryOperationDict +ListMemoryRevisionsResponseOrDict = Union[ + ListMemoryRevisionsResponse, ListMemoryRevisionsResponseDict ] -class UpdateMemoryConfig(_common.BaseModel): - """Config for updating a memory.""" +class AskContextsConfig(_common.BaseModel): + """Config for asking RAG Contexts.""" 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.""", - ) - 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.""", - ) + tools: Optional[list[genai_types.Tool]] = Field(default=None, description="""""") -class UpdateMemoryConfigDict(TypedDict, total=False): - """Config for updating a memory.""" +class AskContextsConfigDict(TypedDict, total=False): + """Config for asking RAG Contexts.""" 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.""" + tools: Optional[list[genai_types.Tool]] + """""" - 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.""" +AskContextsConfigOrDict = Union[AskContextsConfig, AskContextsConfigDict] - 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.""" +class RagQueryRanking(_common.BaseModel): + """Configurations for hybrid search results ranking.""" - topics: Optional[list[MemoryTopicIdDict]] - """Optional. The topics of the memory.""" + alpha: Optional[float] = Field( + default=None, + description="""Optional. Alpha value controls the weight between dense and sparse vector search results. The range is [0, 1], while 0 means sparse vector search only and 1 means dense vector search only. The default value is 0.5 which balances sparse and dense vector search equally.""", + ) - 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.""" +class RagQueryRankingDict(TypedDict, total=False): + """Configurations for hybrid search results ranking.""" - update_mask: Optional[str] - """The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" + alpha: Optional[float] + """Optional. Alpha value controls the weight between dense and sparse vector search results. The range is [0, 1], while 0 means sparse vector search only and 1 means dense vector search only. The default value is 0.5 which balances sparse and dense vector search equally.""" -UpdateMemoryConfigOrDict = Union[UpdateMemoryConfig, UpdateMemoryConfigDict] +RagQueryRankingOrDict = Union[RagQueryRanking, RagQueryRankingDict] -class _UpdateMemoryRequestParameters(_common.BaseModel): - """Parameters for updating memories.""" +class RagQuery(_common.BaseModel): + """A query to retrieve relevant contexts.""" - name: Optional[str] = Field( - default=None, description="""Name of the memory to update.""" + rag_retrieval_config: Optional[genai_types.RagRetrievalConfig] = Field( + default=None, description="""Optional. The retrieval config for the query.""" ) - fact: Optional[str] = Field( + ranking: Optional[RagQueryRanking] = Field( default=None, - description="""The updated fact of the memory. - - This is the semantic knowledge extracted from the source content.""", + description="""Optional. Configurations for hybrid search results ranking.""", ) - scope: Optional[dict[str, str]] = Field( + similarity_top_k: Optional[int] = Field( + default=None, description="""Optional. The number of contexts to retrieve.""" + ) + text: Optional[str] = Field( default=None, - description="""The updated scope of the memory. + description="""Optional. The query in text format to get relevant contexts.""", + ) - 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[UpdateMemoryConfig] = Field(default=None, description="""""") - - -class _UpdateMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for updating memories.""" - name: Optional[str] - """Name of the memory to update.""" - - fact: Optional[str] - """The updated fact of the memory. +class RagQueryDict(TypedDict, total=False): + """A query to retrieve relevant contexts.""" - This is the semantic knowledge extracted from the source content.""" + rag_retrieval_config: Optional[genai_types.RagRetrievalConfigDict] + """Optional. The retrieval config for the query.""" - scope: Optional[dict[str, str]] - """The updated scope of the memory. + ranking: Optional[RagQueryRankingDict] + """Optional. Configurations for hybrid search results ranking.""" - 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 '*'.""" + similarity_top_k: Optional[int] + """Optional. The number of contexts to retrieve.""" - config: Optional[UpdateMemoryConfigDict] - """""" + text: Optional[str] + """Optional. The query in text format to get relevant contexts.""" -_UpdateMemoryRequestParametersOrDict = Union[ - _UpdateMemoryRequestParameters, _UpdateMemoryRequestParametersDict -] +RagQueryOrDict = Union[RagQuery, RagQueryDict] -class PurgeMemoriesConfig(_common.BaseModel): - """Config for purging memories.""" +class _AskContextsRequestParameters(_common.BaseModel): + """Parameters for asking RAG Contexts.""" - 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.""", - ) + query: Optional[RagQuery] = Field(default=None, description="""""") + config: Optional[AskContextsConfig] = Field(default=None, description="""""") -class PurgeMemoriesConfigDict(TypedDict, total=False): - """Config for purging memories.""" +class _AskContextsRequestParametersDict(TypedDict, total=False): + """Parameters for asking RAG Contexts.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + query: Optional[RagQueryDict] + """""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + config: Optional[AskContextsConfigDict] + """""" -PurgeMemoriesConfigOrDict = Union[PurgeMemoriesConfig, PurgeMemoriesConfigDict] +_AskContextsRequestParametersOrDict = Union[ + _AskContextsRequestParameters, _AskContextsRequestParametersDict +] -class _PurgeMemoriesRequestParameters(_common.BaseModel): - """Parameters for purging memories.""" +class RagContextsContext(_common.BaseModel): + """A context of the query.""" - name: Optional[str] = Field( - default=None, description="""Name of the Memory Bank to purge memories from.""" + chunk: Optional[genai_types.RagChunk] = Field( + default=None, description="""Context of the retrieved chunk.""" ) - filter: Optional[str] = Field( + distance: Optional[float] = Field( default=None, - description="""The standard list filter to determine which memories to purge. - More detail in [AIP-160](https://google.aip.dev/160).""", + description="""The distance between the query dense embedding vector and the context text vector.""", ) - filter_groups: Optional[list[MemoryConjunctionFilter]] = Field( + score: Optional[float] = Field( default=None, - description="""Metadata filters that will be applied to the memories' - `metadata` using OR logic. Filters are defined using disjunctive normal - form (OR of ANDs). - - For example: - `filter_groups: [{filters: [{key: "author", value: {string_value: "agent - `123"}, op: EQUAL}]}, {filters: [{key: "label", value: {string_value: - "travel"}, op: EQUAL}, {key: "author", value: {string_value: "agent 321"}, - op: EQUAL}]}]` - - would be equivalent to the logical expression: - `(metadata.author = "agent 123" OR (metadata.label = "travel" AND - metadata.author = "agent 321"))`. - """, + description="""According to the underlying Vector DB and the selected metric type, the score can be either the distance or the similarity between the query and the context and its range depends on the metric type. For example, if the metric type is COSINE_DISTANCE, it represents the distance between the query and the context. The larger the distance, the less relevant the context is to the query. The range is [0, 2], while 0 means the most relevant and 2 means the least relevant.""", ) - force: Optional[bool] = Field( + source_display_name: Optional[str] = Field( + default=None, description="""The file display name.""" + ) + source_uri: Optional[str] = Field( default=None, - description="""If true, the memories will actually be purged. If false, the purge request will be validated but not executed.""", + description="""If the file is imported from Cloud Storage or Google Drive, source_uri will be original file URI in Cloud Storage or Google Drive; if file is uploaded, source_uri will be file display name.""", ) - config: Optional[PurgeMemoriesConfig] = Field(default=None, description="""""") + sparse_distance: Optional[float] = Field( + default=None, + description="""The distance between the query sparse embedding vector and the context text vector.""", + ) + text: Optional[str] = Field(default=None, description="""The text chunk.""") -class _PurgeMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for purging memories.""" +class RagContextsContextDict(TypedDict, total=False): + """A context of the query.""" - name: Optional[str] - """Name of the Memory Bank to purge memories from.""" + chunk: Optional[genai_types.RagChunkDict] + """Context of the retrieved chunk.""" - filter: Optional[str] - """The standard list filter to determine which memories to purge. - More detail in [AIP-160](https://google.aip.dev/160).""" + distance: Optional[float] + """The distance between the query dense embedding vector and the context text vector.""" - filter_groups: Optional[list[MemoryConjunctionFilterDict]] - """Metadata filters that will be applied to the memories' - `metadata` using OR logic. Filters are defined using disjunctive normal - form (OR of ANDs). + score: Optional[float] + """According to the underlying Vector DB and the selected metric type, the score can be either the distance or the similarity between the query and the context and its range depends on the metric type. For example, if the metric type is COSINE_DISTANCE, it represents the distance between the query and the context. The larger the distance, the less relevant the context is to the query. The range is [0, 2], while 0 means the most relevant and 2 means the least relevant.""" - For example: - `filter_groups: [{filters: [{key: "author", value: {string_value: "agent - `123"}, op: EQUAL}]}, {filters: [{key: "label", value: {string_value: - "travel"}, op: EQUAL}, {key: "author", value: {string_value: "agent 321"}, - op: EQUAL}]}]` + source_display_name: Optional[str] + """The file display name.""" - would be equivalent to the logical expression: - `(metadata.author = "agent 123" OR (metadata.label = "travel" AND - metadata.author = "agent 321"))`. - """ + source_uri: Optional[str] + """If the file is imported from Cloud Storage or Google Drive, source_uri will be original file URI in Cloud Storage or Google Drive; if file is uploaded, source_uri will be file display name.""" - force: Optional[bool] - """If true, the memories will actually be purged. If false, the purge request will be validated but not executed.""" + sparse_distance: Optional[float] + """The distance between the query sparse embedding vector and the context text vector.""" - config: Optional[PurgeMemoriesConfigDict] - """""" + text: Optional[str] + """The text chunk.""" -_PurgeMemoriesRequestParametersOrDict = Union[ - _PurgeMemoriesRequestParameters, _PurgeMemoriesRequestParametersDict -] +RagContextsContextOrDict = Union[RagContextsContext, RagContextsContextDict] -class PurgeMemoriesResponse(_common.BaseModel): - """The response for purging memories.""" +class RagContexts(_common.BaseModel): + """Relevant contexts for one query.""" - purge_count: Optional[int] = Field( - default=None, description="""The number of memories that were purged.""" + contexts: Optional[list[RagContextsContext]] = Field( + default=None, description="""All its contexts.""" ) -class PurgeMemoriesResponseDict(TypedDict, total=False): - """The response for purging memories.""" +class RagContextsDict(TypedDict, total=False): + """Relevant contexts for one query.""" - purge_count: Optional[int] - """The number of memories that were purged.""" + contexts: Optional[list[RagContextsContextDict]] + """All its contexts.""" -PurgeMemoriesResponseOrDict = Union[PurgeMemoriesResponse, PurgeMemoriesResponseDict] +RagContextsOrDict = Union[RagContexts, RagContextsDict] -class PurgeMemoriesOperation(_common.BaseModel): - """Operation that purges memories from a Memory Bank.""" +class AskContextsResponse(_common.BaseModel): - 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.""", + contexts: Optional[RagContexts] = Field( + default=None, description="""The contexts of the query.""" ) - response: Optional[PurgeMemoriesResponse] = Field( - default=None, description="""The response for purging memories.""" + response: Optional[str] = Field( + default=None, description="""The Retrieval Response.""" ) -class PurgeMemoriesOperationDict(TypedDict, total=False): - """Operation that purges memories from a Memory Bank.""" +class AskContextsResponseDict(TypedDict, total=False): - 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}`.""" + contexts: Optional[RagContextsDict] + """The contexts of the query.""" - 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.""" + response: Optional[str] + """The Retrieval Response.""" - 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.""" +AskContextsResponseOrDict = Union[AskContextsResponse, AskContextsResponseDict] - response: Optional[PurgeMemoriesResponseDict] - """The response for purging memories.""" +class CorpusStatus(_common.BaseModel): + """RagCorpus status.""" -PurgeMemoriesOperationOrDict = Union[PurgeMemoriesOperation, PurgeMemoriesOperationDict] + error_status: Optional[str] = Field( + default=None, + description="""Output only. Only when the `state` field is ERROR.""", + ) + state: Optional[Literal["UNKNOWN", "INITIALIZED", "ACTIVE", "ERROR"]] = Field( + default=None, description="""Output only. RagCorpus life state.""" + ) -class GetMemoryRevisionConfig(_common.BaseModel): - """Config for getting a Memory Revision.""" +class CorpusStatusDict(TypedDict, total=False): + """RagCorpus status.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + error_status: Optional[str] + """Output only. Only when the `state` field is ERROR.""" + state: Optional[Literal["UNKNOWN", "INITIALIZED", "ACTIVE", "ERROR"]] + """Output only. RagCorpus life state.""" -class GetMemoryRevisionConfigDict(TypedDict, total=False): - """Config for getting a Memory Revision.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +CorpusStatusOrDict = Union[CorpusStatus, CorpusStatusDict] -GetMemoryRevisionConfigOrDict = Union[ - GetMemoryRevisionConfig, GetMemoryRevisionConfigDict -] +class RagCorpusCorpusTypeConfigDocumentCorpus(_common.BaseModel): + """Config for the document corpus.""" + pass -class _GetMemoryRevisionRequestParameters(_common.BaseModel): - """Parameters for getting a memory revision.""" - name: Optional[str] = Field( - default=None, description="""Name of the Memory Revision.""" - ) - config: Optional[GetMemoryRevisionConfig] = Field(default=None, description="""""") +class RagCorpusCorpusTypeConfigDocumentCorpusDict(TypedDict, total=False): + """Config for the document corpus.""" + pass -class _GetMemoryRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for getting a memory revision.""" - name: Optional[str] - """Name of the Memory Revision.""" - - config: Optional[GetMemoryRevisionConfigDict] - """""" - - -_GetMemoryRevisionRequestParametersOrDict = Union[ - _GetMemoryRevisionRequestParameters, _GetMemoryRevisionRequestParametersDict +RagCorpusCorpusTypeConfigDocumentCorpusOrDict = Union[ + RagCorpusCorpusTypeConfigDocumentCorpus, RagCorpusCorpusTypeConfigDocumentCorpusDict ] -class IntermediateExtractedMemory(_common.BaseModel): - """An extracted memory that is the intermediate result before consolidation.""" +class RagFileParsingConfigLlmParser(_common.BaseModel): + """Specifies the LLM parsing for RagFiles.""" - fact: Optional[str] = Field( + custom_parsing_prompt: Optional[str] = Field( default=None, - description="""Output only. Represents the fact of the extracted memory.""", + description="""The prompt to use for parsing. If not specified, a default prompt will be used.""", ) - context: Optional[str] = Field( + global_max_parsing_requests_per_min: Optional[int] = Field( default=None, - description="""Output only. Represents the explanation of why the information was extracted from the source content.""", + description="""The maximum number of requests the job is allowed to make to the LLM model per minute in this project. Consult https://cloud.google.com/vertex-ai/generative-ai/docs/quotas and your document size to set an appropriate value here. If this value is not specified, max_parsing_requests_per_min will be used by indexing pipeline job as the global limit.""", ) - structured_data: Optional[dict[str, Any]] = Field( + max_parsing_requests_per_min: Optional[int] = Field( default=None, - description="""Output only. Represents the structured value of the extracted memory.""", + description="""The maximum number of requests the job is allowed to make to the LLM model per minute. Consult https://cloud.google.com/vertex-ai/generative-ai/docs/quotas and your document size to set an appropriate value here. If unspecified, a default value of 5000 QPM would be used.""", + ) + model_name: Optional[str] = Field( + default=None, + description="""The name of a LLM model used for parsing. Format: * `projects/{project_id}/locations/{location}/publishers/{publisher}/models/{model}`""", ) -class IntermediateExtractedMemoryDict(TypedDict, total=False): - """An extracted memory that is the intermediate result before consolidation.""" +class RagFileParsingConfigLlmParserDict(TypedDict, total=False): + """Specifies the LLM parsing for RagFiles.""" - fact: Optional[str] - """Output only. Represents the fact of the extracted memory.""" + custom_parsing_prompt: Optional[str] + """The prompt to use for parsing. If not specified, a default prompt will be used.""" - context: Optional[str] - """Output only. Represents the explanation of why the information was extracted from the source content.""" + global_max_parsing_requests_per_min: Optional[int] + """The maximum number of requests the job is allowed to make to the LLM model per minute in this project. Consult https://cloud.google.com/vertex-ai/generative-ai/docs/quotas and your document size to set an appropriate value here. If this value is not specified, max_parsing_requests_per_min will be used by indexing pipeline job as the global limit.""" - structured_data: Optional[dict[str, Any]] - """Output only. Represents the structured value of the extracted memory.""" + max_parsing_requests_per_min: Optional[int] + """The maximum number of requests the job is allowed to make to the LLM model per minute. Consult https://cloud.google.com/vertex-ai/generative-ai/docs/quotas and your document size to set an appropriate value here. If unspecified, a default value of 5000 QPM would be used.""" + model_name: Optional[str] + """The name of a LLM model used for parsing. Format: * `projects/{project_id}/locations/{location}/publishers/{publisher}/models/{model}`""" -IntermediateExtractedMemoryOrDict = Union[ - IntermediateExtractedMemory, IntermediateExtractedMemoryDict + +RagFileParsingConfigLlmParserOrDict = Union[ + RagFileParsingConfigLlmParser, RagFileParsingConfigLlmParserDict ] -class MemoryRevision(_common.BaseModel): - """A memory revision.""" +class RagCorpusCorpusTypeConfigMemoryCorpus(_common.BaseModel): + """Config for the memory corpus.""" - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Represents the timestamp when this Memory Revision was created.""", - ) - expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Represents the timestamp of when this resource is considered expired.""", - ) - extracted_memories: Optional[list[IntermediateExtractedMemory]] = Field( - default=None, - description="""Output only. Represents the extracted memories from the source content before consolidation when the memory was updated via GenerateMemories. This information was used to modify an existing Memory via Consolidation.""", - ) - fact: Optional[str] = Field( - default=None, - description="""Output only. Represents the fact of the Memory Revision. This corresponds to the `fact` field of the parent Memory at the time of revision creation.""", - ) - labels: Optional[dict[str, str]] = Field( - default=None, - description="""Output only. Represents the labels of the Memory Revision. These labels are applied to the MemoryRevision when it is created based on `GenerateMemoriesRequest.revision_labels`.""", - ) - name: Optional[str] = Field( - default=None, - description="""Identifier. Represents the resource name of the Memory Revision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{memory_revision}`""", - ) - structured_data: Optional[dict[str, Any]] = Field( - default=None, - description="""Output only. Represents the structured value of the memory at the time of revision creation.""", + llm_parser: Optional[RagFileParsingConfigLlmParser] = Field( + default=None, description="""The LLM parser to use for the memory corpus.""" ) -class MemoryRevisionDict(TypedDict, total=False): - """A memory revision.""" - - create_time: Optional[datetime.datetime] - """Output only. Represents the timestamp when this Memory Revision was created.""" - - expire_time: Optional[datetime.datetime] - """Output only. Represents the timestamp of when this resource is considered expired.""" - - extracted_memories: Optional[list[IntermediateExtractedMemoryDict]] - """Output only. Represents the extracted memories from the source content before consolidation when the memory was updated via GenerateMemories. This information was used to modify an existing Memory via Consolidation.""" - - fact: Optional[str] - """Output only. Represents the fact of the Memory Revision. This corresponds to the `fact` field of the parent Memory at the time of revision creation.""" - - labels: Optional[dict[str, str]] - """Output only. Represents the labels of the Memory Revision. These labels are applied to the MemoryRevision when it is created based on `GenerateMemoriesRequest.revision_labels`.""" - - name: Optional[str] - """Identifier. Represents the resource name of the Memory Revision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{memory_revision}`""" +class RagCorpusCorpusTypeConfigMemoryCorpusDict(TypedDict, total=False): + """Config for the memory corpus.""" - structured_data: Optional[dict[str, Any]] - """Output only. Represents the structured value of the memory at the time of revision creation.""" + llm_parser: Optional[RagFileParsingConfigLlmParserDict] + """The LLM parser to use for the memory corpus.""" -MemoryRevisionOrDict = Union[MemoryRevision, MemoryRevisionDict] +RagCorpusCorpusTypeConfigMemoryCorpusOrDict = Union[ + RagCorpusCorpusTypeConfigMemoryCorpus, RagCorpusCorpusTypeConfigMemoryCorpusDict +] -class ListMemoryRevisionsConfig(_common.BaseModel): - """Config for listing memory revisions.""" +class RagCorpusCorpusTypeConfig(_common.BaseModel): + """The config for the corpus type of the RagCorpus.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + document_corpus: Optional[RagCorpusCorpusTypeConfigDocumentCorpus] = Field( + default=None, description="""Optional. Config for the document corpus.""" ) - 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.""", + memory_corpus: Optional[RagCorpusCorpusTypeConfigMemoryCorpus] = Field( + default=None, description="""Optional. Config for the memory corpus.""" ) -class ListMemoryRevisionsConfigDict(TypedDict, total=False): - """Config for listing memory revisions.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - page_size: Optional[int] - """""" +class RagCorpusCorpusTypeConfigDict(TypedDict, total=False): + """The config for the corpus type of the RagCorpus.""" - page_token: Optional[str] - """""" + document_corpus: Optional[RagCorpusCorpusTypeConfigDocumentCorpusDict] + """Optional. Config for the document corpus.""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" + memory_corpus: Optional[RagCorpusCorpusTypeConfigMemoryCorpusDict] + """Optional. Config for the memory corpus.""" -ListMemoryRevisionsConfigOrDict = Union[ - ListMemoryRevisionsConfig, ListMemoryRevisionsConfigDict +RagCorpusCorpusTypeConfigOrDict = Union[ + RagCorpusCorpusTypeConfig, RagCorpusCorpusTypeConfigDict ] -class _ListMemoryRevisionsRequestParameters(_common.BaseModel): - """Parameters for listing memory revisions.""" +class RagEmbeddingModelConfigVertexPredictionEndpoint(_common.BaseModel): + """Config representing a model hosted on Vertex Prediction Endpoint.""" - name: Optional[str] = Field(default=None, description="""Name of the memory""") - config: Optional[ListMemoryRevisionsConfig] = Field( - default=None, description="""""" + endpoint: Optional[str] = Field( + default=None, + description="""Required. The endpoint resource name. Format: `projects/{project}/locations/{location}/publishers/{publisher}/models/{model}` or `projects/{project}/locations/{location}/endpoints/{endpoint}`""", + ) + model: Optional[str] = Field( + default=None, + description="""Output only. The resource name of the model that is deployed on the endpoint. Present only when the endpoint is not a publisher model. Pattern: `projects/{project}/locations/{location}/models/{model}`""", + ) + model_version_id: Optional[str] = Field( + default=None, + description="""Output only. Version ID of the model that is deployed on the endpoint. Present only when the endpoint is not a publisher model.""", ) -class _ListMemoryRevisionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing memory revisions.""" +class RagEmbeddingModelConfigVertexPredictionEndpointDict(TypedDict, total=False): + """Config representing a model hosted on Vertex Prediction Endpoint.""" - name: Optional[str] - """Name of the memory""" + endpoint: Optional[str] + """Required. The endpoint resource name. Format: `projects/{project}/locations/{location}/publishers/{publisher}/models/{model}` or `projects/{project}/locations/{location}/endpoints/{endpoint}`""" - config: Optional[ListMemoryRevisionsConfigDict] - """""" + model: Optional[str] + """Output only. The resource name of the model that is deployed on the endpoint. Present only when the endpoint is not a publisher model. Pattern: `projects/{project}/locations/{location}/models/{model}`""" + model_version_id: Optional[str] + """Output only. Version ID of the model that is deployed on the endpoint. Present only when the endpoint is not a publisher model.""" -_ListMemoryRevisionsRequestParametersOrDict = Union[ - _ListMemoryRevisionsRequestParameters, _ListMemoryRevisionsRequestParametersDict + +RagEmbeddingModelConfigVertexPredictionEndpointOrDict = Union[ + RagEmbeddingModelConfigVertexPredictionEndpoint, + RagEmbeddingModelConfigVertexPredictionEndpointDict, ] -class ListMemoryRevisionsResponse(_common.BaseModel): - """Response for listing memory revisions.""" +class RagEmbeddingModelConfigSparseEmbeddingConfigBm25(_common.BaseModel): + """Message for BM25 parameters.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + b: Optional[float] = Field( + default=None, + description="""Optional. The parameter to control document length normalization. It determines how much the document length affects the final score. b is in the range of [0, 1]. The default value is 0.75.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - memory_revisions: Optional[list[MemoryRevision]] = Field( - default=None, description="""List of memory revisions.""" + k1: Optional[float] = Field( + default=None, + description="""Optional. The parameter to control term frequency saturation. It determines the scaling between the matching term frequency and final score. k1 is in the range of [1.2, 3]. The default value is 1.2.""", + ) + multilingual: Optional[bool] = Field( + default=None, + description="""Optional. Use multilingual tokenizer if set to true.""", ) -class ListMemoryRevisionsResponseDict(TypedDict, total=False): - """Response for listing memory revisions.""" +class RagEmbeddingModelConfigSparseEmbeddingConfigBm25Dict(TypedDict, total=False): + """Message for BM25 parameters.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + b: Optional[float] + """Optional. The parameter to control document length normalization. It determines how much the document length affects the final score. b is in the range of [0, 1]. The default value is 0.75.""" - next_page_token: Optional[str] - """""" + k1: Optional[float] + """Optional. The parameter to control term frequency saturation. It determines the scaling between the matching term frequency and final score. k1 is in the range of [1.2, 3]. The default value is 1.2.""" - memory_revisions: Optional[list[MemoryRevisionDict]] - """List of memory revisions.""" + multilingual: Optional[bool] + """Optional. Use multilingual tokenizer if set to true.""" -ListMemoryRevisionsResponseOrDict = Union[ - ListMemoryRevisionsResponse, ListMemoryRevisionsResponseDict +RagEmbeddingModelConfigSparseEmbeddingConfigBm25OrDict = Union[ + RagEmbeddingModelConfigSparseEmbeddingConfigBm25, + RagEmbeddingModelConfigSparseEmbeddingConfigBm25Dict, ] -class AskContextsConfig(_common.BaseModel): - """Config for asking RAG Contexts.""" +class RagEmbeddingModelConfigSparseEmbeddingConfig(_common.BaseModel): + """Configuration for sparse emebdding generation.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + bm25: Optional[RagEmbeddingModelConfigSparseEmbeddingConfigBm25] = Field( + default=None, description="""Use BM25 scoring algorithm.""" ) - tools: Optional[list[genai_types.Tool]] = Field(default=None, description="""""") -class AskContextsConfigDict(TypedDict, total=False): - """Config for asking RAG Contexts.""" +class RagEmbeddingModelConfigSparseEmbeddingConfigDict(TypedDict, total=False): + """Configuration for sparse emebdding generation.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + bm25: Optional[RagEmbeddingModelConfigSparseEmbeddingConfigBm25Dict] + """Use BM25 scoring algorithm.""" - tools: Optional[list[genai_types.Tool]] - """""" +RagEmbeddingModelConfigSparseEmbeddingConfigOrDict = Union[ + RagEmbeddingModelConfigSparseEmbeddingConfig, + RagEmbeddingModelConfigSparseEmbeddingConfigDict, +] -AskContextsConfigOrDict = Union[AskContextsConfig, AskContextsConfigDict] +class RagEmbeddingModelConfigHybridSearchConfig(_common.BaseModel): + """Config for hybrid search.""" -class RagQueryRanking(_common.BaseModel): - """Configurations for hybrid search results ranking.""" - - alpha: Optional[float] = Field( + dense_embedding_model_prediction_endpoint: Optional[ + RagEmbeddingModelConfigVertexPredictionEndpoint + ] = Field( default=None, - description="""Optional. Alpha value controls the weight between dense and sparse vector search results. The range is [0, 1], while 0 means sparse vector search only and 1 means dense vector search only. The default value is 0.5 which balances sparse and dense vector search equally.""", + description="""Required. The Vertex AI Prediction Endpoint that hosts the embedding model for dense embedding generations.""", + ) + sparse_embedding_config: Optional[RagEmbeddingModelConfigSparseEmbeddingConfig] = ( + Field( + default=None, + description="""Optional. The configuration for sparse embedding generation. This field is optional the default behavior depends on the vector database choice on the RagCorpus.""", + ) ) -class RagQueryRankingDict(TypedDict, total=False): - """Configurations for hybrid search results ranking.""" +class RagEmbeddingModelConfigHybridSearchConfigDict(TypedDict, total=False): + """Config for hybrid search.""" - alpha: Optional[float] - """Optional. Alpha value controls the weight between dense and sparse vector search results. The range is [0, 1], while 0 means sparse vector search only and 1 means dense vector search only. The default value is 0.5 which balances sparse and dense vector search equally.""" + dense_embedding_model_prediction_endpoint: Optional[ + RagEmbeddingModelConfigVertexPredictionEndpointDict + ] + """Required. The Vertex AI Prediction Endpoint that hosts the embedding model for dense embedding generations.""" + + sparse_embedding_config: Optional[RagEmbeddingModelConfigSparseEmbeddingConfigDict] + """Optional. The configuration for sparse embedding generation. This field is optional the default behavior depends on the vector database choice on the RagCorpus.""" -RagQueryRankingOrDict = Union[RagQueryRanking, RagQueryRankingDict] +RagEmbeddingModelConfigHybridSearchConfigOrDict = Union[ + RagEmbeddingModelConfigHybridSearchConfig, + RagEmbeddingModelConfigHybridSearchConfigDict, +] -class RagQuery(_common.BaseModel): - """A query to retrieve relevant contexts.""" +class RagEmbeddingModelConfig(_common.BaseModel): + """Config for the embedding model to use for RAG.""" - rag_retrieval_config: Optional[genai_types.RagRetrievalConfig] = Field( - default=None, description="""Optional. The retrieval config for the query.""" - ) - ranking: Optional[RagQueryRanking] = Field( - default=None, - description="""Optional. Configurations for hybrid search results ranking.""", - ) - similarity_top_k: Optional[int] = Field( - default=None, description="""Optional. The number of contexts to retrieve.""" + hybrid_search_config: Optional[RagEmbeddingModelConfigHybridSearchConfig] = Field( + default=None, description="""Configuration for hybrid search.""" ) - text: Optional[str] = Field( + vertex_prediction_endpoint: Optional[ + RagEmbeddingModelConfigVertexPredictionEndpoint + ] = Field( default=None, - description="""Optional. The query in text format to get relevant contexts.""", + description="""The Vertex AI Prediction Endpoint that either refers to a publisher model or an endpoint that is hosting a 1P fine-tuned text embedding model. Endpoints hosting non-1P fine-tuned text embedding models are currently not supported. This is used for dense vector search.""", ) -class RagQueryDict(TypedDict, total=False): - """A query to retrieve relevant contexts.""" - - rag_retrieval_config: Optional[genai_types.RagRetrievalConfigDict] - """Optional. The retrieval config for the query.""" - - ranking: Optional[RagQueryRankingDict] - """Optional. Configurations for hybrid search results ranking.""" - - similarity_top_k: Optional[int] - """Optional. The number of contexts to retrieve.""" +class RagEmbeddingModelConfigDict(TypedDict, total=False): + """Config for the embedding model to use for RAG.""" - text: Optional[str] - """Optional. The query in text format to get relevant contexts.""" + hybrid_search_config: Optional[RagEmbeddingModelConfigHybridSearchConfigDict] + """Configuration for hybrid search.""" + vertex_prediction_endpoint: Optional[ + RagEmbeddingModelConfigVertexPredictionEndpointDict + ] + """The Vertex AI Prediction Endpoint that either refers to a publisher model or an endpoint that is hosting a 1P fine-tuned text embedding model. Endpoints hosting non-1P fine-tuned text embedding models are currently not supported. This is used for dense vector search.""" -RagQueryOrDict = Union[RagQuery, RagQueryDict] +RagEmbeddingModelConfigOrDict = Union[ + RagEmbeddingModelConfig, RagEmbeddingModelConfigDict +] -class _AskContextsRequestParameters(_common.BaseModel): - """Parameters for asking RAG Contexts.""" - query: Optional[RagQuery] = Field(default=None, description="""""") - config: Optional[AskContextsConfig] = Field(default=None, description="""""") +class RagVectorDbConfigPinecone(_common.BaseModel): + """The config for the Pinecone.""" + index_name: Optional[str] = Field( + default=None, + description="""Pinecone index name. This value cannot be changed after it's set.""", + ) -class _AskContextsRequestParametersDict(TypedDict, total=False): - """Parameters for asking RAG Contexts.""" - query: Optional[RagQueryDict] - """""" +class RagVectorDbConfigPineconeDict(TypedDict, total=False): + """The config for the Pinecone.""" - config: Optional[AskContextsConfigDict] - """""" + index_name: Optional[str] + """Pinecone index name. This value cannot be changed after it's set.""" -_AskContextsRequestParametersOrDict = Union[ - _AskContextsRequestParameters, _AskContextsRequestParametersDict +RagVectorDbConfigPineconeOrDict = Union[ + RagVectorDbConfigPinecone, RagVectorDbConfigPineconeDict ] -class RagContextsContext(_common.BaseModel): - """A context of the query.""" +class RagVectorDbConfigRagManagedDbANN(_common.BaseModel): + """Config for ANN search. RagManagedDb uses a tree-based structure to partition data and facilitate faster searches. As a tradeoff, it requires longer indexing time and manual triggering of index rebuild via the ImportRagFiles and UpdateRagCorpus API.""" - chunk: Optional[genai_types.RagChunk] = Field( - default=None, description="""Context of the retrieved chunk.""" - ) - distance: Optional[float] = Field( - default=None, - description="""The distance between the query dense embedding vector and the context text vector.""", - ) - score: Optional[float] = Field( - default=None, - description="""According to the underlying Vector DB and the selected metric type, the score can be either the distance or the similarity between the query and the context and its range depends on the metric type. For example, if the metric type is COSINE_DISTANCE, it represents the distance between the query and the context. The larger the distance, the less relevant the context is to the query. The range is [0, 2], while 0 means the most relevant and 2 means the least relevant.""", - ) - source_display_name: Optional[str] = Field( - default=None, description="""The file display name.""" - ) - source_uri: Optional[str] = Field( + leaf_count: Optional[int] = Field( default=None, - description="""If the file is imported from Cloud Storage or Google Drive, source_uri will be original file URI in Cloud Storage or Google Drive; if file is uploaded, source_uri will be file display name.""", + description="""Number of leaf nodes in the tree-based structure. Each leaf node contains groups of closely related vectors along with their corresponding centroid. Recommended value is 10 * sqrt(num of RagFiles in your RagCorpus). Default value is 500.""", ) - sparse_distance: Optional[float] = Field( + tree_depth: Optional[int] = Field( default=None, - description="""The distance between the query sparse embedding vector and the context text vector.""", + description="""The depth of the tree-based structure. Only depth values of 2 and 3 are supported. Recommended value is 2 if you have if you have O(10K) files in the RagCorpus and set this to 3 if more than that. Default value is 2.""", ) - text: Optional[str] = Field(default=None, description="""The text chunk.""") - - -class RagContextsContextDict(TypedDict, total=False): - """A context of the query.""" - - chunk: Optional[genai_types.RagChunkDict] - """Context of the retrieved chunk.""" - - distance: Optional[float] - """The distance between the query dense embedding vector and the context text vector.""" - - score: Optional[float] - """According to the underlying Vector DB and the selected metric type, the score can be either the distance or the similarity between the query and the context and its range depends on the metric type. For example, if the metric type is COSINE_DISTANCE, it represents the distance between the query and the context. The larger the distance, the less relevant the context is to the query. The range is [0, 2], while 0 means the most relevant and 2 means the least relevant.""" - source_display_name: Optional[str] - """The file display name.""" - source_uri: Optional[str] - """If the file is imported from Cloud Storage or Google Drive, source_uri will be original file URI in Cloud Storage or Google Drive; if file is uploaded, source_uri will be file display name.""" +class RagVectorDbConfigRagManagedDbANNDict(TypedDict, total=False): + """Config for ANN search. RagManagedDb uses a tree-based structure to partition data and facilitate faster searches. As a tradeoff, it requires longer indexing time and manual triggering of index rebuild via the ImportRagFiles and UpdateRagCorpus API.""" - sparse_distance: Optional[float] - """The distance between the query sparse embedding vector and the context text vector.""" + leaf_count: Optional[int] + """Number of leaf nodes in the tree-based structure. Each leaf node contains groups of closely related vectors along with their corresponding centroid. Recommended value is 10 * sqrt(num of RagFiles in your RagCorpus). Default value is 500.""" - text: Optional[str] - """The text chunk.""" + tree_depth: Optional[int] + """The depth of the tree-based structure. Only depth values of 2 and 3 are supported. Recommended value is 2 if you have if you have O(10K) files in the RagCorpus and set this to 3 if more than that. Default value is 2.""" -RagContextsContextOrDict = Union[RagContextsContext, RagContextsContextDict] +RagVectorDbConfigRagManagedDbANNOrDict = Union[ + RagVectorDbConfigRagManagedDbANN, RagVectorDbConfigRagManagedDbANNDict +] -class RagContexts(_common.BaseModel): - """Relevant contexts for one query.""" +class RagVectorDbConfigRagManagedDbKNN(_common.BaseModel): + """Config for KNN search.""" - contexts: Optional[list[RagContextsContext]] = Field( - default=None, description="""All its contexts.""" - ) + pass -class RagContextsDict(TypedDict, total=False): - """Relevant contexts for one query.""" +class RagVectorDbConfigRagManagedDbKNNDict(TypedDict, total=False): + """Config for KNN search.""" - contexts: Optional[list[RagContextsContextDict]] - """All its contexts.""" + pass -RagContextsOrDict = Union[RagContexts, RagContextsDict] +RagVectorDbConfigRagManagedDbKNNOrDict = Union[ + RagVectorDbConfigRagManagedDbKNN, RagVectorDbConfigRagManagedDbKNNDict +] -class AskContextsResponse(_common.BaseModel): +class RagVectorDbConfigRagManagedDb(_common.BaseModel): + """The config for the default RAG-managed Vector DB.""" - contexts: Optional[RagContexts] = Field( - default=None, description="""The contexts of the query.""" + ann: Optional[RagVectorDbConfigRagManagedDbANN] = Field( + default=None, + description="""Performs an ANN search on RagCorpus. Use this if you have a lot of files (> 10K) in your RagCorpus and want to reduce the search latency.""", ) - response: Optional[str] = Field( - default=None, description="""The Retrieval Response.""" + knn: Optional[RagVectorDbConfigRagManagedDbKNN] = Field( + default=None, + description="""Performs a KNN search on RagCorpus. Default choice if not specified.""", ) -class AskContextsResponseDict(TypedDict, total=False): +class RagVectorDbConfigRagManagedDbDict(TypedDict, total=False): + """The config for the default RAG-managed Vector DB.""" - contexts: Optional[RagContextsDict] - """The contexts of the query.""" + ann: Optional[RagVectorDbConfigRagManagedDbANNDict] + """Performs an ANN search on RagCorpus. Use this if you have a lot of files (> 10K) in your RagCorpus and want to reduce the search latency.""" - response: Optional[str] - """The Retrieval Response.""" + knn: Optional[RagVectorDbConfigRagManagedDbKNNDict] + """Performs a KNN search on RagCorpus. Default choice if not specified.""" -AskContextsResponseOrDict = Union[AskContextsResponse, AskContextsResponseDict] +RagVectorDbConfigRagManagedDbOrDict = Union[ + RagVectorDbConfigRagManagedDb, RagVectorDbConfigRagManagedDbDict +] -class CorpusStatus(_common.BaseModel): - """RagCorpus status.""" +class RagVectorDbConfigRagManagedVertexVectorSearch(_common.BaseModel): + """The config for the RAG-managed Vertex Vector Search 2.0.""" - error_status: Optional[str] = Field( + collection_name: Optional[str] = Field( default=None, - description="""Output only. Only when the `state` field is ERROR.""", - ) - state: Optional[Literal["UNKNOWN", "INITIALIZED", "ACTIVE", "ERROR"]] = Field( - default=None, description="""Output only. RagCorpus life state.""" + description="""Output only. The resource name of the Vector Search 2.0 Collection that RAG Created for the corpus. Only populated after the corpus is successfully created. Format: `projects/{project}/locations/{location}/collections/{collection_id}`""", ) -class CorpusStatusDict(TypedDict, total=False): - """RagCorpus status.""" +class RagVectorDbConfigRagManagedVertexVectorSearchDict(TypedDict, total=False): + """The config for the RAG-managed Vertex Vector Search 2.0.""" - error_status: Optional[str] - """Output only. Only when the `state` field is ERROR.""" + collection_name: Optional[str] + """Output only. The resource name of the Vector Search 2.0 Collection that RAG Created for the corpus. Only populated after the corpus is successfully created. Format: `projects/{project}/locations/{location}/collections/{collection_id}`""" - state: Optional[Literal["UNKNOWN", "INITIALIZED", "ACTIVE", "ERROR"]] - """Output only. RagCorpus life state.""" +RagVectorDbConfigRagManagedVertexVectorSearchOrDict = Union[ + RagVectorDbConfigRagManagedVertexVectorSearch, + RagVectorDbConfigRagManagedVertexVectorSearchDict, +] -CorpusStatusOrDict = Union[CorpusStatus, CorpusStatusDict] +class RagVectorDbConfigVertexFeatureStore(_common.BaseModel): + """The config for the Vertex Feature Store.""" -class RagCorpusCorpusTypeConfigDocumentCorpus(_common.BaseModel): - """Config for the document corpus.""" + feature_view_resource_name: Optional[str] = Field( + default=None, + description="""The resource name of the FeatureView. Format: `projects/{project}/locations/{location}/featureOnlineStores/{feature_online_store}/featureViews/{feature_view}`""", + ) - pass +class RagVectorDbConfigVertexFeatureStoreDict(TypedDict, total=False): + """The config for the Vertex Feature Store.""" -class RagCorpusCorpusTypeConfigDocumentCorpusDict(TypedDict, total=False): - """Config for the document corpus.""" + feature_view_resource_name: Optional[str] + """The resource name of the FeatureView. Format: `projects/{project}/locations/{location}/featureOnlineStores/{feature_online_store}/featureViews/{feature_view}`""" - pass +RagVectorDbConfigVertexFeatureStoreOrDict = Union[ + RagVectorDbConfigVertexFeatureStore, RagVectorDbConfigVertexFeatureStoreDict +] -RagCorpusCorpusTypeConfigDocumentCorpusOrDict = Union[ - RagCorpusCorpusTypeConfigDocumentCorpus, RagCorpusCorpusTypeConfigDocumentCorpusDict -] +class RagVectorDbConfigVertexVectorSearch(_common.BaseModel): + """The config for the Vertex Vector Search.""" -class RagFileParsingConfigLlmParser(_common.BaseModel): - """Specifies the LLM parsing for RagFiles.""" - - custom_parsing_prompt: Optional[str] = Field( - default=None, - description="""The prompt to use for parsing. If not specified, a default prompt will be used.""", - ) - global_max_parsing_requests_per_min: Optional[int] = Field( - default=None, - description="""The maximum number of requests the job is allowed to make to the LLM model per minute in this project. Consult https://cloud.google.com/vertex-ai/generative-ai/docs/quotas and your document size to set an appropriate value here. If this value is not specified, max_parsing_requests_per_min will be used by indexing pipeline job as the global limit.""", - ) - max_parsing_requests_per_min: Optional[int] = Field( + index: Optional[str] = Field( default=None, - description="""The maximum number of requests the job is allowed to make to the LLM model per minute. Consult https://cloud.google.com/vertex-ai/generative-ai/docs/quotas and your document size to set an appropriate value here. If unspecified, a default value of 5000 QPM would be used.""", + description="""The resource name of the Index. Format: `projects/{project}/locations/{location}/indexes/{index}`""", ) - model_name: Optional[str] = Field( + index_endpoint: Optional[str] = Field( default=None, - description="""The name of a LLM model used for parsing. Format: * `projects/{project_id}/locations/{location}/publishers/{publisher}/models/{model}`""", + description="""The resource name of the Index Endpoint. Format: `projects/{project}/locations/{location}/indexEndpoints/{index_endpoint}`""", ) -class RagFileParsingConfigLlmParserDict(TypedDict, total=False): - """Specifies the LLM parsing for RagFiles.""" - - custom_parsing_prompt: Optional[str] - """The prompt to use for parsing. If not specified, a default prompt will be used.""" - - global_max_parsing_requests_per_min: Optional[int] - """The maximum number of requests the job is allowed to make to the LLM model per minute in this project. Consult https://cloud.google.com/vertex-ai/generative-ai/docs/quotas and your document size to set an appropriate value here. If this value is not specified, max_parsing_requests_per_min will be used by indexing pipeline job as the global limit.""" - - max_parsing_requests_per_min: Optional[int] - """The maximum number of requests the job is allowed to make to the LLM model per minute. Consult https://cloud.google.com/vertex-ai/generative-ai/docs/quotas and your document size to set an appropriate value here. If unspecified, a default value of 5000 QPM would be used.""" - - model_name: Optional[str] - """The name of a LLM model used for parsing. Format: * `projects/{project_id}/locations/{location}/publishers/{publisher}/models/{model}`""" - - -RagFileParsingConfigLlmParserOrDict = Union[ - RagFileParsingConfigLlmParser, RagFileParsingConfigLlmParserDict -] - - -class RagCorpusCorpusTypeConfigMemoryCorpus(_common.BaseModel): - """Config for the memory corpus.""" - - llm_parser: Optional[RagFileParsingConfigLlmParser] = Field( - default=None, description="""The LLM parser to use for the memory corpus.""" - ) - +class RagVectorDbConfigVertexVectorSearchDict(TypedDict, total=False): + """The config for the Vertex Vector Search.""" -class RagCorpusCorpusTypeConfigMemoryCorpusDict(TypedDict, total=False): - """Config for the memory corpus.""" + index: Optional[str] + """The resource name of the Index. Format: `projects/{project}/locations/{location}/indexes/{index}`""" - llm_parser: Optional[RagFileParsingConfigLlmParserDict] - """The LLM parser to use for the memory corpus.""" + index_endpoint: Optional[str] + """The resource name of the Index Endpoint. Format: `projects/{project}/locations/{location}/indexEndpoints/{index_endpoint}`""" -RagCorpusCorpusTypeConfigMemoryCorpusOrDict = Union[ - RagCorpusCorpusTypeConfigMemoryCorpus, RagCorpusCorpusTypeConfigMemoryCorpusDict +RagVectorDbConfigVertexVectorSearchOrDict = Union[ + RagVectorDbConfigVertexVectorSearch, RagVectorDbConfigVertexVectorSearchDict ] -class RagCorpusCorpusTypeConfig(_common.BaseModel): - """The config for the corpus type of the RagCorpus.""" +class RagVectorDbConfigWeaviate(_common.BaseModel): + """The config for the Weaviate.""" - document_corpus: Optional[RagCorpusCorpusTypeConfigDocumentCorpus] = Field( - default=None, description="""Optional. Config for the document corpus.""" + collection_name: Optional[str] = Field( + default=None, + description="""The corresponding collection this corpus maps to. This value cannot be changed after it's set.""", ) - memory_corpus: Optional[RagCorpusCorpusTypeConfigMemoryCorpus] = Field( - default=None, description="""Optional. Config for the memory corpus.""" + http_endpoint: Optional[str] = Field( + default=None, + description="""Weaviate DB instance HTTP endpoint. e.g. 34.56.78.90:8080 Vertex RAG only supports HTTP connection to Weaviate. This value cannot be changed after it's set.""", ) -class RagCorpusCorpusTypeConfigDict(TypedDict, total=False): - """The config for the corpus type of the RagCorpus.""" +class RagVectorDbConfigWeaviateDict(TypedDict, total=False): + """The config for the Weaviate.""" - document_corpus: Optional[RagCorpusCorpusTypeConfigDocumentCorpusDict] - """Optional. Config for the document corpus.""" + collection_name: Optional[str] + """The corresponding collection this corpus maps to. This value cannot be changed after it's set.""" - memory_corpus: Optional[RagCorpusCorpusTypeConfigMemoryCorpusDict] - """Optional. Config for the memory corpus.""" + http_endpoint: Optional[str] + """Weaviate DB instance HTTP endpoint. e.g. 34.56.78.90:8080 Vertex RAG only supports HTTP connection to Weaviate. This value cannot be changed after it's set.""" -RagCorpusCorpusTypeConfigOrDict = Union[ - RagCorpusCorpusTypeConfig, RagCorpusCorpusTypeConfigDict +RagVectorDbConfigWeaviateOrDict = Union[ + RagVectorDbConfigWeaviate, RagVectorDbConfigWeaviateDict ] -class RagEmbeddingModelConfigVertexPredictionEndpoint(_common.BaseModel): - """Config representing a model hosted on Vertex Prediction Endpoint.""" +class RagVectorDbConfig(_common.BaseModel): + """Config for the Vector DB to use for RAG.""" - endpoint: Optional[str] = Field( - default=None, - description="""Required. The endpoint resource name. Format: `projects/{project}/locations/{location}/publishers/{publisher}/models/{model}` or `projects/{project}/locations/{location}/endpoints/{endpoint}`""", + api_auth: Optional[genai_types.ApiAuth] = Field( + default=None, description="""Authentication config for the chosen Vector DB.""" ) - model: Optional[str] = Field( - default=None, - description="""Output only. The resource name of the model that is deployed on the endpoint. Present only when the endpoint is not a publisher model. Pattern: `projects/{project}/locations/{location}/models/{model}`""", + pinecone: Optional[RagVectorDbConfigPinecone] = Field( + default=None, description="""The config for the Pinecone.""" ) - model_version_id: Optional[str] = Field( + rag_embedding_model_config: Optional[RagEmbeddingModelConfig] = Field( default=None, - description="""Output only. Version ID of the model that is deployed on the endpoint. Present only when the endpoint is not a publisher model.""", + description="""Optional. Immutable. The embedding model config of the Vector DB.""", ) - - -class RagEmbeddingModelConfigVertexPredictionEndpointDict(TypedDict, total=False): - """Config representing a model hosted on Vertex Prediction Endpoint.""" - - endpoint: Optional[str] - """Required. The endpoint resource name. Format: `projects/{project}/locations/{location}/publishers/{publisher}/models/{model}` or `projects/{project}/locations/{location}/endpoints/{endpoint}`""" - - model: Optional[str] - """Output only. The resource name of the model that is deployed on the endpoint. Present only when the endpoint is not a publisher model. Pattern: `projects/{project}/locations/{location}/models/{model}`""" - - model_version_id: Optional[str] - """Output only. Version ID of the model that is deployed on the endpoint. Present only when the endpoint is not a publisher model.""" - - -RagEmbeddingModelConfigVertexPredictionEndpointOrDict = Union[ - RagEmbeddingModelConfigVertexPredictionEndpoint, - RagEmbeddingModelConfigVertexPredictionEndpointDict, -] - - -class RagEmbeddingModelConfigSparseEmbeddingConfigBm25(_common.BaseModel): - """Message for BM25 parameters.""" - - b: Optional[float] = Field( - default=None, - description="""Optional. The parameter to control document length normalization. It determines how much the document length affects the final score. b is in the range of [0, 1]. The default value is 0.75.""", + rag_managed_db: Optional[RagVectorDbConfigRagManagedDb] = Field( + default=None, description="""The config for the RAG-managed Vector DB.""" ) - k1: Optional[float] = Field( + rag_managed_vertex_vector_search: Optional[ + RagVectorDbConfigRagManagedVertexVectorSearch + ] = Field( default=None, - description="""Optional. The parameter to control term frequency saturation. It determines the scaling between the matching term frequency and final score. k1 is in the range of [1.2, 3]. The default value is 1.2.""", + description="""The config for the RAG-managed Vertex Vector Search 2.0.""", ) - multilingual: Optional[bool] = Field( - default=None, - description="""Optional. Use multilingual tokenizer if set to true.""", + vertex_feature_store: Optional[RagVectorDbConfigVertexFeatureStore] = Field( + default=None, description="""The config for the Vertex Feature Store.""" + ) + vertex_vector_search: Optional[RagVectorDbConfigVertexVectorSearch] = Field( + default=None, description="""The config for the Vertex Vector Search.""" + ) + weaviate: Optional[RagVectorDbConfigWeaviate] = Field( + default=None, description="""The config for the Weaviate.""" ) -class RagEmbeddingModelConfigSparseEmbeddingConfigBm25Dict(TypedDict, total=False): - """Message for BM25 parameters.""" - - b: Optional[float] - """Optional. The parameter to control document length normalization. It determines how much the document length affects the final score. b is in the range of [0, 1]. The default value is 0.75.""" - - k1: Optional[float] - """Optional. The parameter to control term frequency saturation. It determines the scaling between the matching term frequency and final score. k1 is in the range of [1.2, 3]. The default value is 1.2.""" - - multilingual: Optional[bool] - """Optional. Use multilingual tokenizer if set to true.""" +class RagVectorDbConfigDict(TypedDict, total=False): + """Config for the Vector DB to use for RAG.""" + api_auth: Optional[genai_types.ApiAuthDict] + """Authentication config for the chosen Vector DB.""" -RagEmbeddingModelConfigSparseEmbeddingConfigBm25OrDict = Union[ - RagEmbeddingModelConfigSparseEmbeddingConfigBm25, - RagEmbeddingModelConfigSparseEmbeddingConfigBm25Dict, -] + pinecone: Optional[RagVectorDbConfigPineconeDict] + """The config for the Pinecone.""" + rag_embedding_model_config: Optional[RagEmbeddingModelConfigDict] + """Optional. Immutable. The embedding model config of the Vector DB.""" -class RagEmbeddingModelConfigSparseEmbeddingConfig(_common.BaseModel): - """Configuration for sparse emebdding generation.""" + rag_managed_db: Optional[RagVectorDbConfigRagManagedDbDict] + """The config for the RAG-managed Vector DB.""" - bm25: Optional[RagEmbeddingModelConfigSparseEmbeddingConfigBm25] = Field( - default=None, description="""Use BM25 scoring algorithm.""" - ) + rag_managed_vertex_vector_search: Optional[ + RagVectorDbConfigRagManagedVertexVectorSearchDict + ] + """The config for the RAG-managed Vertex Vector Search 2.0.""" + vertex_feature_store: Optional[RagVectorDbConfigVertexFeatureStoreDict] + """The config for the Vertex Feature Store.""" -class RagEmbeddingModelConfigSparseEmbeddingConfigDict(TypedDict, total=False): - """Configuration for sparse emebdding generation.""" + vertex_vector_search: Optional[RagVectorDbConfigVertexVectorSearchDict] + """The config for the Vertex Vector Search.""" - bm25: Optional[RagEmbeddingModelConfigSparseEmbeddingConfigBm25Dict] - """Use BM25 scoring algorithm.""" + weaviate: Optional[RagVectorDbConfigWeaviateDict] + """The config for the Weaviate.""" -RagEmbeddingModelConfigSparseEmbeddingConfigOrDict = Union[ - RagEmbeddingModelConfigSparseEmbeddingConfig, - RagEmbeddingModelConfigSparseEmbeddingConfigDict, -] +RagVectorDbConfigOrDict = Union[RagVectorDbConfig, RagVectorDbConfigDict] -class RagEmbeddingModelConfigHybridSearchConfig(_common.BaseModel): - """Config for hybrid search.""" +class VertexAiSearchConfig(_common.BaseModel): + """Config for the Vertex AI Search.""" - dense_embedding_model_prediction_endpoint: Optional[ - RagEmbeddingModelConfigVertexPredictionEndpoint - ] = Field( + serving_config: Optional[str] = Field( default=None, - description="""Required. The Vertex AI Prediction Endpoint that hosts the embedding model for dense embedding generations.""", - ) - sparse_embedding_config: Optional[RagEmbeddingModelConfigSparseEmbeddingConfig] = ( - Field( - default=None, - description="""Optional. The configuration for sparse embedding generation. This field is optional the default behavior depends on the vector database choice on the RagCorpus.""", - ) + description="""Vertex AI Search Serving Config resource full name. For example, `projects/{project}/locations/{location}/collections/{collection}/engines/{engine}/servingConfigs/{serving_config}` or `projects/{project}/locations/{location}/collections/{collection}/dataStores/{data_store}/servingConfigs/{serving_config}`.""", ) -class RagEmbeddingModelConfigHybridSearchConfigDict(TypedDict, total=False): - """Config for hybrid search.""" - - dense_embedding_model_prediction_endpoint: Optional[ - RagEmbeddingModelConfigVertexPredictionEndpointDict - ] - """Required. The Vertex AI Prediction Endpoint that hosts the embedding model for dense embedding generations.""" +class VertexAiSearchConfigDict(TypedDict, total=False): + """Config for the Vertex AI Search.""" - sparse_embedding_config: Optional[RagEmbeddingModelConfigSparseEmbeddingConfigDict] - """Optional. The configuration for sparse embedding generation. This field is optional the default behavior depends on the vector database choice on the RagCorpus.""" + serving_config: Optional[str] + """Vertex AI Search Serving Config resource full name. For example, `projects/{project}/locations/{location}/collections/{collection}/engines/{engine}/servingConfigs/{serving_config}` or `projects/{project}/locations/{location}/collections/{collection}/dataStores/{data_store}/servingConfigs/{serving_config}`.""" -RagEmbeddingModelConfigHybridSearchConfigOrDict = Union[ - RagEmbeddingModelConfigHybridSearchConfig, - RagEmbeddingModelConfigHybridSearchConfigDict, -] +VertexAiSearchConfigOrDict = Union[VertexAiSearchConfig, VertexAiSearchConfigDict] -class RagEmbeddingModelConfig(_common.BaseModel): - """Config for the embedding model to use for RAG.""" +class RagCorpus(_common.BaseModel): + """A RAG Corpus.""" - hybrid_search_config: Optional[RagEmbeddingModelConfigHybridSearchConfig] = Field( - default=None, description="""Configuration for hybrid search.""" + corpus_status: Optional[CorpusStatus] = Field( + default=None, description="""Output only. RagCorpus state.""" ) - vertex_prediction_endpoint: Optional[ - RagEmbeddingModelConfigVertexPredictionEndpoint - ] = Field( + corpus_type_config: Optional[RagCorpusCorpusTypeConfig] = Field( default=None, - description="""The Vertex AI Prediction Endpoint that either refers to a publisher model or an endpoint that is hosting a 1P fine-tuned text embedding model. Endpoints hosting non-1P fine-tuned text embedding models are currently not supported. This is used for dense vector search.""", + description="""Optional. The corpus type config of the RagCorpus.""", + ) + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this RagCorpus was created.""", + ) + description: Optional[str] = Field( + default=None, description="""Optional. The description of the RagCorpus.""" + ) + display_name: Optional[str] = Field( + default=None, + description="""Required. The display name of the RagCorpus. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""Optional. Immutable. The CMEK key name used to encrypt at-rest data related to this Corpus. Only applicable to RagManagedDb option for Vector DB. This field can only be set at corpus creation time, and cannot be updated or deleted.""", + ) + name: Optional[str] = Field( + default=None, description="""Output only. The resource name of the RagCorpus.""" + ) + rag_embedding_model_config: Optional[RagEmbeddingModelConfig] = Field( + default=None, + description="""Optional. Immutable. The embedding model config of the RagCorpus.""", + ) + rag_files_count: Optional[int] = Field( + default=None, + description="""Output only. Number of RagFiles in the RagCorpus. NOTE: This field is not populated in the response of VertexRagDataService.ListRagCorpora.""", + ) + rag_vector_db_config: Optional[RagVectorDbConfig] = Field( + default=None, + description="""Optional. Immutable. The Vector DB config of the RagCorpus.""", + ) + 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.""" + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this RagCorpus was last updated.""", + ) + vector_db_config: Optional[RagVectorDbConfig] = Field( + default=None, + description="""Optional. Immutable. The config for the Vector DBs.""", + ) + vertex_ai_search_config: Optional[VertexAiSearchConfig] = Field( + default=None, + description="""Optional. Immutable. The config for the Vertex AI Search.""", ) -class RagEmbeddingModelConfigDict(TypedDict, total=False): - """Config for the embedding model to use for RAG.""" +class RagCorpusDict(TypedDict, total=False): + """A RAG Corpus.""" - hybrid_search_config: Optional[RagEmbeddingModelConfigHybridSearchConfigDict] - """Configuration for hybrid search.""" + corpus_status: Optional[CorpusStatusDict] + """Output only. RagCorpus state.""" - vertex_prediction_endpoint: Optional[ - RagEmbeddingModelConfigVertexPredictionEndpointDict - ] - """The Vertex AI Prediction Endpoint that either refers to a publisher model or an endpoint that is hosting a 1P fine-tuned text embedding model. Endpoints hosting non-1P fine-tuned text embedding models are currently not supported. This is used for dense vector search.""" + corpus_type_config: Optional[RagCorpusCorpusTypeConfigDict] + """Optional. The corpus type config of the RagCorpus.""" + + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this RagCorpus was created.""" + description: Optional[str] + """Optional. The description of the RagCorpus.""" -RagEmbeddingModelConfigOrDict = Union[ - RagEmbeddingModelConfig, RagEmbeddingModelConfigDict -] + display_name: Optional[str] + """Required. The display name of the RagCorpus. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" + encryption_spec: Optional[genai_types.EncryptionSpecDict] + """Optional. Immutable. The CMEK key name used to encrypt at-rest data related to this Corpus. Only applicable to RagManagedDb option for Vector DB. This field can only be set at corpus creation time, and cannot be updated or deleted.""" -class RagVectorDbConfigPinecone(_common.BaseModel): - """The config for the Pinecone.""" + name: Optional[str] + """Output only. The resource name of the RagCorpus.""" - index_name: Optional[str] = Field( - default=None, - description="""Pinecone index name. This value cannot be changed after it's set.""", - ) + rag_embedding_model_config: Optional[RagEmbeddingModelConfigDict] + """Optional. Immutable. The embedding model config of the RagCorpus.""" + rag_files_count: Optional[int] + """Output only. Number of RagFiles in the RagCorpus. NOTE: This field is not populated in the response of VertexRagDataService.ListRagCorpora.""" -class RagVectorDbConfigPineconeDict(TypedDict, total=False): - """The config for the Pinecone.""" + rag_vector_db_config: Optional[RagVectorDbConfigDict] + """Optional. Immutable. The Vector DB config of the RagCorpus.""" - index_name: Optional[str] - """Pinecone index name. This value cannot be changed after it's set.""" + satisfies_pzi: Optional[bool] + """Output only. Reserved for future use.""" + + satisfies_pzs: Optional[bool] + """Output only. Reserved for future use.""" + update_time: Optional[datetime.datetime] + """Output only. Timestamp when this RagCorpus was last updated.""" -RagVectorDbConfigPineconeOrDict = Union[ - RagVectorDbConfigPinecone, RagVectorDbConfigPineconeDict -] + vector_db_config: Optional[RagVectorDbConfigDict] + """Optional. Immutable. The config for the Vector DBs.""" + vertex_ai_search_config: Optional[VertexAiSearchConfigDict] + """Optional. Immutable. The config for the Vertex AI Search.""" -class RagVectorDbConfigRagManagedDbANN(_common.BaseModel): - """Config for ANN search. RagManagedDb uses a tree-based structure to partition data and facilitate faster searches. As a tradeoff, it requires longer indexing time and manual triggering of index rebuild via the ImportRagFiles and UpdateRagCorpus API.""" - leaf_count: Optional[int] = Field( - default=None, - description="""Number of leaf nodes in the tree-based structure. Each leaf node contains groups of closely related vectors along with their corresponding centroid. Recommended value is 10 * sqrt(num of RagFiles in your RagCorpus). Default value is 500.""", - ) - tree_depth: Optional[int] = Field( - default=None, - description="""The depth of the tree-based structure. Only depth values of 2 and 3 are supported. Recommended value is 2 if you have if you have O(10K) files in the RagCorpus and set this to 3 if more than that. Default value is 2.""", +RagCorpusOrDict = Union[RagCorpus, RagCorpusDict] + + +class CreateRagCorpusConfig(_common.BaseModel): + """Config for creating a RAG corpus.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class RagVectorDbConfigRagManagedDbANNDict(TypedDict, total=False): - """Config for ANN search. RagManagedDb uses a tree-based structure to partition data and facilitate faster searches. As a tradeoff, it requires longer indexing time and manual triggering of index rebuild via the ImportRagFiles and UpdateRagCorpus API.""" +class CreateRagCorpusConfigDict(TypedDict, total=False): + """Config for creating a RAG corpus.""" - leaf_count: Optional[int] - """Number of leaf nodes in the tree-based structure. Each leaf node contains groups of closely related vectors along with their corresponding centroid. Recommended value is 10 * sqrt(num of RagFiles in your RagCorpus). Default value is 500.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - tree_depth: Optional[int] - """The depth of the tree-based structure. Only depth values of 2 and 3 are supported. Recommended value is 2 if you have if you have O(10K) files in the RagCorpus and set this to 3 if more than that. Default value is 2.""" +CreateRagCorpusConfigOrDict = Union[CreateRagCorpusConfig, CreateRagCorpusConfigDict] -RagVectorDbConfigRagManagedDbANNOrDict = Union[ - RagVectorDbConfigRagManagedDbANN, RagVectorDbConfigRagManagedDbANNDict -] +class _CreateRagCorpusRequestParameters(_common.BaseModel): + """Parameters for creating a RAG corpus.""" -class RagVectorDbConfigRagManagedDbKNN(_common.BaseModel): - """Config for KNN search.""" + rag_corpus: Optional[RagCorpus] = Field(default=None, description="""""") + config: Optional[CreateRagCorpusConfig] = Field(default=None, description="""""") - pass +class _CreateRagCorpusRequestParametersDict(TypedDict, total=False): + """Parameters for creating a RAG corpus.""" -class RagVectorDbConfigRagManagedDbKNNDict(TypedDict, total=False): - """Config for KNN search.""" + rag_corpus: Optional[RagCorpusDict] + """""" - pass + config: Optional[CreateRagCorpusConfigDict] + """""" -RagVectorDbConfigRagManagedDbKNNOrDict = Union[ - RagVectorDbConfigRagManagedDbKNN, RagVectorDbConfigRagManagedDbKNNDict +_CreateRagCorpusRequestParametersOrDict = Union[ + _CreateRagCorpusRequestParameters, _CreateRagCorpusRequestParametersDict ] -class RagVectorDbConfigRagManagedDb(_common.BaseModel): - """The config for the default RAG-managed Vector DB.""" +class CreateRagCorpusOperation(_common.BaseModel): + """Operation for creating a RAG corpus.""" - ann: Optional[RagVectorDbConfigRagManagedDbANN] = Field( + name: Optional[str] = Field( default=None, - description="""Performs an ANN search on RagCorpus. Use this if you have a lot of files (> 10K) in your RagCorpus and want to reduce the search latency.""", + 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}`.""", ) - knn: Optional[RagVectorDbConfigRagManagedDbKNN] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Performs a KNN search on RagCorpus. Default choice if not specified.""", + 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 RagVectorDbConfigRagManagedDbDict(TypedDict, total=False): - """The config for the default RAG-managed Vector DB.""" +class CreateRagCorpusOperationDict(TypedDict, total=False): + """Operation for creating a RAG corpus.""" - ann: Optional[RagVectorDbConfigRagManagedDbANNDict] - """Performs an ANN search on RagCorpus. Use this if you have a lot of files (> 10K) in your RagCorpus and want to reduce the search latency.""" + 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}`.""" - knn: Optional[RagVectorDbConfigRagManagedDbKNNDict] - """Performs a KNN search on RagCorpus. Default choice if not specified.""" + 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.""" -RagVectorDbConfigRagManagedDbOrDict = Union[ - RagVectorDbConfigRagManagedDb, RagVectorDbConfigRagManagedDbDict + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +CreateRagCorpusOperationOrDict = Union[ + CreateRagCorpusOperation, CreateRagCorpusOperationDict ] -class RagVectorDbConfigRagManagedVertexVectorSearch(_common.BaseModel): - """The config for the RAG-managed Vertex Vector Search 2.0.""" +class GetCorpusOperationConfig(_common.BaseModel): + """Config for getting a corpus operation.""" - collection_name: Optional[str] = Field( - default=None, - description="""Output only. The resource name of the Vector Search 2.0 Collection that RAG Created for the corpus. Only populated after the corpus is successfully created. Format: `projects/{project}/locations/{location}/collections/{collection_id}`""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class RagVectorDbConfigRagManagedVertexVectorSearchDict(TypedDict, total=False): - """The config for the RAG-managed Vertex Vector Search 2.0.""" +class GetCorpusOperationConfigDict(TypedDict, total=False): + """Config for getting a corpus operation.""" - collection_name: Optional[str] - """Output only. The resource name of the Vector Search 2.0 Collection that RAG Created for the corpus. Only populated after the corpus is successfully created. Format: `projects/{project}/locations/{location}/collections/{collection_id}`""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -RagVectorDbConfigRagManagedVertexVectorSearchOrDict = Union[ - RagVectorDbConfigRagManagedVertexVectorSearch, - RagVectorDbConfigRagManagedVertexVectorSearchDict, +GetCorpusOperationConfigOrDict = Union[ + GetCorpusOperationConfig, GetCorpusOperationConfigDict ] -class RagVectorDbConfigVertexFeatureStore(_common.BaseModel): - """The config for the Vertex Feature Store.""" +class _GetCorpusOperationParameters(_common.BaseModel): + """Parameters for getting a corpus operation.""" - feature_view_resource_name: Optional[str] = Field( - default=None, - description="""The resource name of the FeatureView. Format: `projects/{project}/locations/{location}/featureOnlineStores/{feature_online_store}/featureViews/{feature_view}`""", + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetCorpusOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -class RagVectorDbConfigVertexFeatureStoreDict(TypedDict, total=False): - """The config for the Vertex Feature Store.""" +class _GetCorpusOperationParametersDict(TypedDict, total=False): + """Parameters for getting a corpus operation.""" - feature_view_resource_name: Optional[str] - """The resource name of the FeatureView. Format: `projects/{project}/locations/{location}/featureOnlineStores/{feature_online_store}/featureViews/{feature_view}`""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" + config: Optional[GetCorpusOperationConfigDict] + """Used to override the default configuration.""" -RagVectorDbConfigVertexFeatureStoreOrDict = Union[ - RagVectorDbConfigVertexFeatureStore, RagVectorDbConfigVertexFeatureStoreDict + +_GetCorpusOperationParametersOrDict = Union[ + _GetCorpusOperationParameters, _GetCorpusOperationParametersDict ] -class RagVectorDbConfigVertexVectorSearch(_common.BaseModel): - """The config for the Vertex Vector Search.""" +class CorpusOperation(_common.BaseModel): + """Operation that has a corpus as a response.""" - index: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""The resource name of the Index. Format: `projects/{project}/locations/{location}/indexes/{index}`""", + 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}`.""", ) - index_endpoint: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""The resource name of the Index Endpoint. Format: `projects/{project}/locations/{location}/indexEndpoints/{index_endpoint}`""", + 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[RagCorpus] = Field( + default=None, description="""The created Corpus.""" ) -class RagVectorDbConfigVertexVectorSearchDict(TypedDict, total=False): - """The config for the Vertex Vector Search.""" +class CorpusOperationDict(TypedDict, total=False): + """Operation that has a corpus as a response.""" - index: Optional[str] - """The resource name of the Index. Format: `projects/{project}/locations/{location}/indexes/{index}`""" + 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}`.""" - index_endpoint: Optional[str] - """The resource name of the Index Endpoint. Format: `projects/{project}/locations/{location}/indexEndpoints/{index_endpoint}`""" + 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.""" -RagVectorDbConfigVertexVectorSearchOrDict = Union[ - RagVectorDbConfigVertexVectorSearch, RagVectorDbConfigVertexVectorSearchDict -] + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + response: Optional[RagCorpusDict] + """The created Corpus.""" -class RagVectorDbConfigWeaviate(_common.BaseModel): - """The config for the Weaviate.""" - collection_name: Optional[str] = Field( - default=None, - description="""The corresponding collection this corpus maps to. This value cannot be changed after it's set.""", - ) - http_endpoint: Optional[str] = Field( - default=None, - description="""Weaviate DB instance HTTP endpoint. e.g. 34.56.78.90:8080 Vertex RAG only supports HTTP connection to Weaviate. This value cannot be changed after it's set.""", +CorpusOperationOrDict = Union[CorpusOperation, CorpusOperationDict] + + +class GetRagCorpusConfig(_common.BaseModel): + """Config for getting a RAG corpus.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class RagVectorDbConfigWeaviateDict(TypedDict, total=False): - """The config for the Weaviate.""" +class GetRagCorpusConfigDict(TypedDict, total=False): + """Config for getting a RAG corpus.""" - collection_name: Optional[str] - """The corresponding collection this corpus maps to. This value cannot be changed after it's set.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - http_endpoint: Optional[str] - """Weaviate DB instance HTTP endpoint. e.g. 34.56.78.90:8080 Vertex RAG only supports HTTP connection to Weaviate. This value cannot be changed after it's set.""" +GetRagCorpusConfigOrDict = Union[GetRagCorpusConfig, GetRagCorpusConfigDict] -RagVectorDbConfigWeaviateOrDict = Union[ - RagVectorDbConfigWeaviate, RagVectorDbConfigWeaviateDict -] +class _GetRagCorpusRequestParameters(_common.BaseModel): + """Parameters for getting a RAG corpus.""" -class RagVectorDbConfig(_common.BaseModel): - """Config for the Vector DB to use for RAG.""" + config: Optional[GetRagCorpusConfig] = Field(default=None, description="""""") + name: Optional[str] = Field(default=None, description="""""") - api_auth: Optional[genai_types.ApiAuth] = Field( - default=None, description="""Authentication config for the chosen Vector DB.""" - ) - pinecone: Optional[RagVectorDbConfigPinecone] = Field( - default=None, description="""The config for the Pinecone.""" - ) - rag_embedding_model_config: Optional[RagEmbeddingModelConfig] = Field( - default=None, - description="""Optional. Immutable. The embedding model config of the Vector DB.""", - ) - rag_managed_db: Optional[RagVectorDbConfigRagManagedDb] = Field( - default=None, description="""The config for the RAG-managed Vector DB.""" - ) - rag_managed_vertex_vector_search: Optional[ - RagVectorDbConfigRagManagedVertexVectorSearch - ] = Field( - default=None, - description="""The config for the RAG-managed Vertex Vector Search 2.0.""", - ) - vertex_feature_store: Optional[RagVectorDbConfigVertexFeatureStore] = Field( - default=None, description="""The config for the Vertex Feature Store.""" - ) - vertex_vector_search: Optional[RagVectorDbConfigVertexVectorSearch] = Field( - default=None, description="""The config for the Vertex Vector Search.""" - ) - weaviate: Optional[RagVectorDbConfigWeaviate] = Field( - default=None, description="""The config for the Weaviate.""" - ) +class _GetRagCorpusRequestParametersDict(TypedDict, total=False): + """Parameters for getting a RAG corpus.""" -class RagVectorDbConfigDict(TypedDict, total=False): - """Config for the Vector DB to use for RAG.""" + config: Optional[GetRagCorpusConfigDict] + """""" - api_auth: Optional[genai_types.ApiAuthDict] - """Authentication config for the chosen Vector DB.""" + name: Optional[str] + """""" - pinecone: Optional[RagVectorDbConfigPineconeDict] - """The config for the Pinecone.""" - rag_embedding_model_config: Optional[RagEmbeddingModelConfigDict] - """Optional. Immutable. The embedding model config of the Vector DB.""" +_GetRagCorpusRequestParametersOrDict = Union[ + _GetRagCorpusRequestParameters, _GetRagCorpusRequestParametersDict +] - rag_managed_db: Optional[RagVectorDbConfigRagManagedDbDict] - """The config for the RAG-managed Vector DB.""" - rag_managed_vertex_vector_search: Optional[ - RagVectorDbConfigRagManagedVertexVectorSearchDict - ] - """The config for the RAG-managed Vertex Vector Search 2.0.""" +class ListRagCorporaConfig(_common.BaseModel): + """Config for listing RagCorpora.""" - vertex_feature_store: Optional[RagVectorDbConfigVertexFeatureStoreDict] - """The config for the Vertex Feature Store.""" + 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="""""") - vertex_vector_search: Optional[RagVectorDbConfigVertexVectorSearchDict] - """The config for the Vertex Vector Search.""" - weaviate: Optional[RagVectorDbConfigWeaviateDict] - """The config for the Weaviate.""" +class ListRagCorporaConfigDict(TypedDict, total=False): + """Config for listing RagCorpora.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -RagVectorDbConfigOrDict = Union[RagVectorDbConfig, RagVectorDbConfigDict] + page_size: Optional[int] + """""" + page_token: Optional[str] + """""" -class VertexAiSearchConfig(_common.BaseModel): - """Config for the Vertex AI Search.""" - serving_config: Optional[str] = Field( - default=None, - description="""Vertex AI Search Serving Config resource full name. For example, `projects/{project}/locations/{location}/collections/{collection}/engines/{engine}/servingConfigs/{serving_config}` or `projects/{project}/locations/{location}/collections/{collection}/dataStores/{data_store}/servingConfigs/{serving_config}`.""", - ) +ListRagCorporaConfigOrDict = Union[ListRagCorporaConfig, ListRagCorporaConfigDict] -class VertexAiSearchConfigDict(TypedDict, total=False): - """Config for the Vertex AI Search.""" +class _ListRagCorporaRequestParameters(_common.BaseModel): + """Parameters for listing RagCorpora.""" - serving_config: Optional[str] - """Vertex AI Search Serving Config resource full name. For example, `projects/{project}/locations/{location}/collections/{collection}/engines/{engine}/servingConfigs/{serving_config}` or `projects/{project}/locations/{location}/collections/{collection}/dataStores/{data_store}/servingConfigs/{serving_config}`.""" + config: Optional[ListRagCorporaConfig] = Field(default=None, description="""""") -VertexAiSearchConfigOrDict = Union[VertexAiSearchConfig, VertexAiSearchConfigDict] +class _ListRagCorporaRequestParametersDict(TypedDict, total=False): + """Parameters for listing RagCorpora.""" + config: Optional[ListRagCorporaConfigDict] + """""" -class RagCorpus(_common.BaseModel): - """A RAG Corpus.""" - corpus_status: Optional[CorpusStatus] = Field( - default=None, description="""Output only. RagCorpus state.""" - ) - corpus_type_config: Optional[RagCorpusCorpusTypeConfig] = Field( - default=None, - description="""Optional. The corpus type config of the RagCorpus.""", - ) - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when this RagCorpus was created.""", - ) - description: Optional[str] = Field( - default=None, description="""Optional. The description of the RagCorpus.""" - ) - display_name: Optional[str] = Field( - default=None, - description="""Required. The display name of the RagCorpus. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", - ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, - description="""Optional. Immutable. The CMEK key name used to encrypt at-rest data related to this Corpus. Only applicable to RagManagedDb option for Vector DB. This field can only be set at corpus creation time, and cannot be updated or deleted.""", - ) - name: Optional[str] = Field( - default=None, description="""Output only. The resource name of the RagCorpus.""" - ) - rag_embedding_model_config: Optional[RagEmbeddingModelConfig] = Field( - default=None, - description="""Optional. Immutable. The embedding model config of the RagCorpus.""", - ) - rag_files_count: Optional[int] = Field( - default=None, - description="""Output only. Number of RagFiles in the RagCorpus. NOTE: This field is not populated in the response of VertexRagDataService.ListRagCorpora.""", - ) - rag_vector_db_config: Optional[RagVectorDbConfig] = Field( - default=None, - description="""Optional. Immutable. The Vector DB config of the RagCorpus.""", - ) - 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.""" - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when this RagCorpus was last updated.""", - ) - vector_db_config: Optional[RagVectorDbConfig] = Field( - default=None, - description="""Optional. Immutable. The config for the Vector DBs.""", +_ListRagCorporaRequestParametersOrDict = Union[ + _ListRagCorporaRequestParameters, _ListRagCorporaRequestParametersDict +] + + +class ListRagCorporaResponse(_common.BaseModel): + """Response for listing RagCorpora.""" + + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - vertex_ai_search_config: Optional[VertexAiSearchConfig] = Field( - default=None, - description="""Optional. Immutable. The config for the Vertex AI Search.""", + next_page_token: Optional[str] = Field(default=None, description="""""") + rag_corpora: Optional[list[RagCorpus]] = Field( + default=None, description="""List of RagCorpus instances.""" ) -class RagCorpusDict(TypedDict, total=False): - """A RAG Corpus.""" +class ListRagCorporaResponseDict(TypedDict, total=False): + """Response for listing RagCorpora.""" - corpus_status: Optional[CorpusStatusDict] - """Output only. RagCorpus state.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - corpus_type_config: Optional[RagCorpusCorpusTypeConfigDict] - """Optional. The corpus type config of the RagCorpus.""" + next_page_token: Optional[str] + """""" - create_time: Optional[datetime.datetime] - """Output only. Timestamp when this RagCorpus was created.""" + rag_corpora: Optional[list[RagCorpusDict]] + """List of RagCorpus instances.""" - description: Optional[str] - """Optional. The description of the RagCorpus.""" - display_name: Optional[str] - """Required. The display name of the RagCorpus. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" +ListRagCorporaResponseOrDict = Union[ListRagCorporaResponse, ListRagCorporaResponseDict] - encryption_spec: Optional[genai_types.EncryptionSpecDict] - """Optional. Immutable. The CMEK key name used to encrypt at-rest data related to this Corpus. Only applicable to RagManagedDb option for Vector DB. This field can only be set at corpus creation time, and cannot be updated or deleted.""" - name: Optional[str] - """Output only. The resource name of the RagCorpus.""" - - rag_embedding_model_config: Optional[RagEmbeddingModelConfigDict] - """Optional. Immutable. The embedding model config of the RagCorpus.""" - - rag_files_count: Optional[int] - """Output only. Number of RagFiles in the RagCorpus. NOTE: This field is not populated in the response of VertexRagDataService.ListRagCorpora.""" - - rag_vector_db_config: Optional[RagVectorDbConfigDict] - """Optional. Immutable. The Vector DB config of the RagCorpus.""" - - satisfies_pzi: Optional[bool] - """Output only. Reserved for future use.""" - - satisfies_pzs: Optional[bool] - """Output only. Reserved for future use.""" - - update_time: Optional[datetime.datetime] - """Output only. Timestamp when this RagCorpus was last updated.""" - - vector_db_config: Optional[RagVectorDbConfigDict] - """Optional. Immutable. The config for the Vector DBs.""" - - vertex_ai_search_config: Optional[VertexAiSearchConfigDict] - """Optional. Immutable. The config for the Vertex AI Search.""" - - -RagCorpusOrDict = Union[RagCorpus, RagCorpusDict] - - -class CreateRagCorpusConfig(_common.BaseModel): - """Config for creating a RAG corpus.""" +class GetRagFileConfig(_common.BaseModel): + """Config for getting a RAG corpus.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class CreateRagCorpusConfigDict(TypedDict, total=False): - """Config for creating a RAG corpus.""" +class GetRagFileConfigDict(TypedDict, total=False): + """Config for getting a RAG corpus.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -CreateRagCorpusConfigOrDict = Union[CreateRagCorpusConfig, CreateRagCorpusConfigDict] +GetRagFileConfigOrDict = Union[GetRagFileConfig, GetRagFileConfigDict] -class _CreateRagCorpusRequestParameters(_common.BaseModel): - """Parameters for creating a RAG corpus.""" +class _GetRagFileRequestParameters(_common.BaseModel): + """Parameters for getting a RAG corpus.""" - rag_corpus: Optional[RagCorpus] = Field(default=None, description="""""") - config: Optional[CreateRagCorpusConfig] = Field(default=None, description="""""") + config: Optional[GetRagFileConfig] = Field(default=None, description="""""") + name: Optional[str] = Field(default=None, description="""""") -class _CreateRagCorpusRequestParametersDict(TypedDict, total=False): - """Parameters for creating a RAG corpus.""" +class _GetRagFileRequestParametersDict(TypedDict, total=False): + """Parameters for getting a RAG corpus.""" - rag_corpus: Optional[RagCorpusDict] + config: Optional[GetRagFileConfigDict] """""" - config: Optional[CreateRagCorpusConfigDict] + name: Optional[str] """""" -_CreateRagCorpusRequestParametersOrDict = Union[ - _CreateRagCorpusRequestParameters, _CreateRagCorpusRequestParametersDict +_GetRagFileRequestParametersOrDict = Union[ + _GetRagFileRequestParameters, _GetRagFileRequestParametersDict ] -class CreateRagCorpusOperation(_common.BaseModel): - """Operation for creating a RAG corpus.""" +class RagFileStatus(_common.BaseModel): + """RagFile status.""" - 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.""", + state: Optional[RagFileState] = Field( + default=None, description="""The state of the RagFile.""" ) -class CreateRagCorpusOperationDict(TypedDict, total=False): - """Operation for creating a RAG corpus.""" +class RagFileStatusDict(TypedDict, total=False): + """RagFile status.""" - 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}`.""" + state: Optional[RagFileState] + """The state of the RagFile.""" - 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.""" +RagFileStatusOrDict = Union[RagFileStatus, RagFileStatusDict] - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" +class DirectUploadSource(_common.BaseModel): + """The input content is encapsulated and uploaded in the request.""" -CreateRagCorpusOperationOrDict = Union[ - CreateRagCorpusOperation, CreateRagCorpusOperationDict -] + pass -class GetCorpusOperationConfig(_common.BaseModel): - """Config for getting a corpus operation.""" +class DirectUploadSourceDict(TypedDict, total=False): + """The input content is encapsulated and uploaded in the request.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + pass + + +DirectUploadSourceOrDict = Union[DirectUploadSource, DirectUploadSourceDict] + + +class GoogleDriveSourceResourceId(_common.BaseModel): + """The type and ID of the Google Drive resource.""" + + resource_id: Optional[str] = Field( + default=None, description="""Required. The ID of the Google Drive resource.""" + ) + resource_type: Optional[ResourceType] = Field( + default=None, description="""Required. The type of the Google Drive resource.""" ) -class GetCorpusOperationConfigDict(TypedDict, total=False): - """Config for getting a corpus operation.""" +class GoogleDriveSourceResourceIdDict(TypedDict, total=False): + """The type and ID of the Google Drive resource.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + resource_id: Optional[str] + """Required. The ID of the Google Drive resource.""" + + resource_type: Optional[ResourceType] + """Required. The type of the Google Drive resource.""" -GetCorpusOperationConfigOrDict = Union[ - GetCorpusOperationConfig, GetCorpusOperationConfigDict +GoogleDriveSourceResourceIdOrDict = Union[ + GoogleDriveSourceResourceId, GoogleDriveSourceResourceIdDict ] -class _GetCorpusOperationParameters(_common.BaseModel): - """Parameters for getting a corpus operation.""" +class GoogleDriveSource(_common.BaseModel): + """The Google Drive location for the input content.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" - ) - config: Optional[GetCorpusOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + resource_ids: Optional[list[GoogleDriveSourceResourceId]] = Field( + default=None, description="""Required. Google Drive resource IDs.""" ) -class _GetCorpusOperationParametersDict(TypedDict, total=False): - """Parameters for getting a corpus operation.""" - - operation_name: Optional[str] - """The server-assigned name for the operation.""" +class GoogleDriveSourceDict(TypedDict, total=False): + """The Google Drive location for the input content.""" - config: Optional[GetCorpusOperationConfigDict] - """Used to override the default configuration.""" + resource_ids: Optional[list[GoogleDriveSourceResourceIdDict]] + """Required. Google Drive resource IDs.""" -_GetCorpusOperationParametersOrDict = Union[ - _GetCorpusOperationParameters, _GetCorpusOperationParametersDict -] +GoogleDriveSourceOrDict = Union[GoogleDriveSource, GoogleDriveSourceDict] -class CorpusOperation(_common.BaseModel): - """Operation that has a corpus as a response.""" +class JiraSourceJiraQueries(_common.BaseModel): + """JiraQueries contains the Jira queries and corresponding authentication.""" - name: Optional[str] = Field( + api_key_config: Optional[genai_types.ApiAuthApiKeyConfig] = 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="""Required. The SecretManager secret version resource name (e.g. projects/{project}/secrets/{secret}/versions/{version}) storing the Jira API key. See [Manage API tokens for your Atlassian account](https://support.atlassian.com/atlassian-account/docs/manage-api-tokens-for-your-atlassian-account/).""", ) - metadata: Optional[dict[str, Any]] = Field( + custom_queries: Optional[list[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="""A list of custom Jira queries to import. For information about JQL (Jira Query Language), see https://support.atlassian.com/jira-service-management-cloud/docs/use-advanced-search-with-jira-query-language-jql/""", ) - 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.""", + email: Optional[str] = Field( + default=None, description="""Required. The Jira email address.""" ) - error: Optional[dict[str, Any]] = Field( + projects: Optional[list[str]] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""A list of Jira projects to import in their entirety.""", ) - response: Optional[RagCorpus] = Field( - default=None, description="""The created Corpus.""" + server_uri: Optional[str] = Field( + default=None, description="""Required. The Jira server URI.""" ) -class CorpusOperationDict(TypedDict, total=False): - """Operation that has a corpus as a response.""" +class JiraSourceJiraQueriesDict(TypedDict, total=False): + """JiraQueries contains the Jira queries and corresponding authentication.""" - 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}`.""" + api_key_config: Optional[genai_types.ApiAuthApiKeyConfigDict] + """Required. The SecretManager secret version resource name (e.g. projects/{project}/secrets/{secret}/versions/{version}) storing the Jira API key. See [Manage API tokens for your Atlassian account](https://support.atlassian.com/atlassian-account/docs/manage-api-tokens-for-your-atlassian-account/).""" - 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.""" + custom_queries: Optional[list[str]] + """A list of custom Jira queries to import. For information about JQL (Jira Query Language), see https://support.atlassian.com/jira-service-management-cloud/docs/use-advanced-search-with-jira-query-language-jql/""" - 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.""" + email: Optional[str] + """Required. The Jira email address.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + projects: Optional[list[str]] + """A list of Jira projects to import in their entirety.""" - response: Optional[RagCorpusDict] - """The created Corpus.""" + server_uri: Optional[str] + """Required. The Jira server URI.""" -CorpusOperationOrDict = Union[CorpusOperation, CorpusOperationDict] +JiraSourceJiraQueriesOrDict = Union[JiraSourceJiraQueries, JiraSourceJiraQueriesDict] -class GetRagCorpusConfig(_common.BaseModel): - """Config for getting a RAG corpus.""" +class JiraSource(_common.BaseModel): + """The Jira source for the ImportRagFilesRequest.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + jira_queries: Optional[list[JiraSourceJiraQueries]] = Field( + default=None, description="""Required. The Jira queries.""" ) -class GetRagCorpusConfigDict(TypedDict, total=False): - """Config for getting a RAG corpus.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -GetRagCorpusConfigOrDict = Union[GetRagCorpusConfig, GetRagCorpusConfigDict] - - -class _GetRagCorpusRequestParameters(_common.BaseModel): - """Parameters for getting a RAG corpus.""" - - config: Optional[GetRagCorpusConfig] = Field(default=None, description="""""") - name: Optional[str] = Field(default=None, description="""""") - - -class _GetRagCorpusRequestParametersDict(TypedDict, total=False): - """Parameters for getting a RAG corpus.""" +class JiraSourceDict(TypedDict, total=False): + """The Jira source for the ImportRagFilesRequest.""" - config: Optional[GetRagCorpusConfigDict] - """""" + jira_queries: Optional[list[JiraSourceJiraQueriesDict]] + """Required. The Jira queries.""" - name: Optional[str] - """""" +JiraSourceOrDict = Union[JiraSource, JiraSourceDict] -_GetRagCorpusRequestParametersOrDict = Union[ - _GetRagCorpusRequestParameters, _GetRagCorpusRequestParametersDict -] - -class ListRagCorporaConfig(_common.BaseModel): - """Config for listing RagCorpora.""" - - 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="""""") - - -class ListRagCorporaConfigDict(TypedDict, total=False): - """Config for listing RagCorpora.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - page_size: Optional[int] - """""" - - page_token: Optional[str] - """""" - - -ListRagCorporaConfigOrDict = Union[ListRagCorporaConfig, ListRagCorporaConfigDict] - - -class _ListRagCorporaRequestParameters(_common.BaseModel): - """Parameters for listing RagCorpora.""" - - config: Optional[ListRagCorporaConfig] = Field(default=None, description="""""") - - -class _ListRagCorporaRequestParametersDict(TypedDict, total=False): - """Parameters for listing RagCorpora.""" - - config: Optional[ListRagCorporaConfigDict] - """""" - - -_ListRagCorporaRequestParametersOrDict = Union[ - _ListRagCorporaRequestParameters, _ListRagCorporaRequestParametersDict -] - - -class ListRagCorporaResponse(_common.BaseModel): - """Response for listing RagCorpora.""" - - 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="""""") - rag_corpora: Optional[list[RagCorpus]] = Field( - default=None, description="""List of RagCorpus instances.""" - ) - - -class ListRagCorporaResponseDict(TypedDict, total=False): - """Response for listing RagCorpora.""" - - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" - - next_page_token: Optional[str] - """""" - - rag_corpora: Optional[list[RagCorpusDict]] - """List of RagCorpus instances.""" - - -ListRagCorporaResponseOrDict = Union[ListRagCorporaResponse, ListRagCorporaResponseDict] - - -class GetRagFileConfig(_common.BaseModel): - """Config for getting a RAG corpus.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - - -class GetRagFileConfigDict(TypedDict, total=False): - """Config for getting a RAG corpus.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -GetRagFileConfigOrDict = Union[GetRagFileConfig, GetRagFileConfigDict] - - -class _GetRagFileRequestParameters(_common.BaseModel): - """Parameters for getting a RAG corpus.""" - - config: Optional[GetRagFileConfig] = Field(default=None, description="""""") - name: Optional[str] = Field(default=None, description="""""") - - -class _GetRagFileRequestParametersDict(TypedDict, total=False): - """Parameters for getting a RAG corpus.""" - - config: Optional[GetRagFileConfigDict] - """""" - - name: Optional[str] - """""" - - -_GetRagFileRequestParametersOrDict = Union[ - _GetRagFileRequestParameters, _GetRagFileRequestParametersDict -] - - -class RagFileStatus(_common.BaseModel): - """RagFile status.""" - - state: Optional[RagFileState] = Field( - default=None, description="""The state of the RagFile.""" - ) - - -class RagFileStatusDict(TypedDict, total=False): - """RagFile status.""" - - state: Optional[RagFileState] - """The state of the RagFile.""" - - -RagFileStatusOrDict = Union[RagFileStatus, RagFileStatusDict] - - -class DirectUploadSource(_common.BaseModel): - """The input content is encapsulated and uploaded in the request.""" - - pass - - -class DirectUploadSourceDict(TypedDict, total=False): - """The input content is encapsulated and uploaded in the request.""" - - pass - - -DirectUploadSourceOrDict = Union[DirectUploadSource, DirectUploadSourceDict] - - -class GoogleDriveSourceResourceId(_common.BaseModel): - """The type and ID of the Google Drive resource.""" - - resource_id: Optional[str] = Field( - default=None, description="""Required. The ID of the Google Drive resource.""" - ) - resource_type: Optional[ResourceType] = Field( - default=None, description="""Required. The type of the Google Drive resource.""" - ) - - -class GoogleDriveSourceResourceIdDict(TypedDict, total=False): - """The type and ID of the Google Drive resource.""" - - resource_id: Optional[str] - """Required. The ID of the Google Drive resource.""" - - resource_type: Optional[ResourceType] - """Required. The type of the Google Drive resource.""" - - -GoogleDriveSourceResourceIdOrDict = Union[ - GoogleDriveSourceResourceId, GoogleDriveSourceResourceIdDict -] - - -class GoogleDriveSource(_common.BaseModel): - """The Google Drive location for the input content.""" - - resource_ids: Optional[list[GoogleDriveSourceResourceId]] = Field( - default=None, description="""Required. Google Drive resource IDs.""" - ) - - -class GoogleDriveSourceDict(TypedDict, total=False): - """The Google Drive location for the input content.""" - - resource_ids: Optional[list[GoogleDriveSourceResourceIdDict]] - """Required. Google Drive resource IDs.""" - - -GoogleDriveSourceOrDict = Union[GoogleDriveSource, GoogleDriveSourceDict] - - -class JiraSourceJiraQueries(_common.BaseModel): - """JiraQueries contains the Jira queries and corresponding authentication.""" - - api_key_config: Optional[genai_types.ApiAuthApiKeyConfig] = Field( - default=None, - description="""Required. The SecretManager secret version resource name (e.g. projects/{project}/secrets/{secret}/versions/{version}) storing the Jira API key. See [Manage API tokens for your Atlassian account](https://support.atlassian.com/atlassian-account/docs/manage-api-tokens-for-your-atlassian-account/).""", - ) - custom_queries: Optional[list[str]] = Field( - default=None, - description="""A list of custom Jira queries to import. For information about JQL (Jira Query Language), see https://support.atlassian.com/jira-service-management-cloud/docs/use-advanced-search-with-jira-query-language-jql/""", - ) - email: Optional[str] = Field( - default=None, description="""Required. The Jira email address.""" - ) - projects: Optional[list[str]] = Field( - default=None, - description="""A list of Jira projects to import in their entirety.""", - ) - server_uri: Optional[str] = Field( - default=None, description="""Required. The Jira server URI.""" - ) - - -class JiraSourceJiraQueriesDict(TypedDict, total=False): - """JiraQueries contains the Jira queries and corresponding authentication.""" - - api_key_config: Optional[genai_types.ApiAuthApiKeyConfigDict] - """Required. The SecretManager secret version resource name (e.g. projects/{project}/secrets/{secret}/versions/{version}) storing the Jira API key. See [Manage API tokens for your Atlassian account](https://support.atlassian.com/atlassian-account/docs/manage-api-tokens-for-your-atlassian-account/).""" - - custom_queries: Optional[list[str]] - """A list of custom Jira queries to import. For information about JQL (Jira Query Language), see https://support.atlassian.com/jira-service-management-cloud/docs/use-advanced-search-with-jira-query-language-jql/""" - - email: Optional[str] - """Required. The Jira email address.""" - - projects: Optional[list[str]] - """A list of Jira projects to import in their entirety.""" - - server_uri: Optional[str] - """Required. The Jira server URI.""" - - -JiraSourceJiraQueriesOrDict = Union[JiraSourceJiraQueries, JiraSourceJiraQueriesDict] - - -class JiraSource(_common.BaseModel): - """The Jira source for the ImportRagFilesRequest.""" - - jira_queries: Optional[list[JiraSourceJiraQueries]] = Field( - default=None, description="""Required. The Jira queries.""" - ) - - -class JiraSourceDict(TypedDict, total=False): - """The Jira source for the ImportRagFilesRequest.""" - - jira_queries: Optional[list[JiraSourceJiraQueriesDict]] - """Required. The Jira queries.""" - - -JiraSourceOrDict = Union[JiraSource, JiraSourceDict] - - -class SharePointSourcesSharePointSource(_common.BaseModel): - """An individual SharePointSource.""" +class SharePointSourcesSharePointSource(_common.BaseModel): + """An individual SharePointSource.""" client_id: Optional[str] = Field( default=None, @@ -17266,1597 +16493,1665 @@ class UploadRagFileResponseDict(TypedDict, total=False): UploadRagFileResponseOrDict = Union[UploadRagFileResponse, UploadRagFileResponseDict] -class GetAgentEngineRuntimeRevisionConfig(_common.BaseModel): - """Config for getting an Agent Engine Runtime Revision.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - - -class GetAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Runtime Revision.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -GetAgentEngineRuntimeRevisionConfigOrDict = Union[ - GetAgentEngineRuntimeRevisionConfig, GetAgentEngineRuntimeRevisionConfigDict -] - - -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 _GetAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine runtime revision.""" - - name: Optional[str] - """Name of the agent engine runtime revision.""" - - config: Optional[GetAgentEngineRuntimeRevisionConfigDict] - """""" - - -_GetAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ - _GetAgentEngineRuntimeRevisionRequestParameters, - _GetAgentEngineRuntimeRevisionRequestParametersDict, -] - - -class ReasoningEngineRuntimeRevision(_common.BaseModel): - """A runtime revision.""" +class SandboxEnvironmentSpecCodeExecutionEnvironment(_common.BaseModel): + """The code execution environment with customized settings.""" - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when this ReasoningEngineRuntimeRevision was created.""", - ) - name: Optional[str] = Field( + code_language: Optional[Language] = Field( default=None, - description="""Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""", + description="""The coding language supported in this environment.""", ) - spec: Optional[ReasoningEngineSpec] = Field( + machine_config: Optional[MachineConfig] = 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.""" + description="""The machine config of the code execution 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 SandboxEnvironmentSpecCodeExecutionEnvironmentDict(TypedDict, total=False): + """The code execution environment with customized settings.""" - spec: Optional[ReasoningEngineSpecDict] - """Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""" + code_language: Optional[Language] + """The coding language supported in this environment.""" - state: Optional[State] - """Output only. The state of the revision.""" + machine_config: Optional[MachineConfig] + """The machine config of the code execution environment.""" -ReasoningEngineRuntimeRevisionOrDict = Union[ - ReasoningEngineRuntimeRevision, ReasoningEngineRuntimeRevisionDict +SandboxEnvironmentSpecCodeExecutionEnvironmentOrDict = Union[ + SandboxEnvironmentSpecCodeExecutionEnvironment, + SandboxEnvironmentSpecCodeExecutionEnvironmentDict, ] -class ListAgentEngineRuntimeRevisionsConfig(_common.BaseModel): - """Config for listing reasoning engine runtime revisions.""" - - 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 ListAgentEngineRuntimeRevisionsConfigDict(TypedDict, total=False): - """Config for listing reasoning engine runtime revisions.""" +class SandboxEnvironmentSpecComputerUseEnvironment(_common.BaseModel): + """The computer use environment with customized settings.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + pass - page_size: Optional[int] - """""" - page_token: Optional[str] - """""" +class SandboxEnvironmentSpecComputerUseEnvironmentDict(TypedDict, total=False): + """The computer use environment with customized settings.""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" + pass -ListAgentEngineRuntimeRevisionsConfigOrDict = Union[ - ListAgentEngineRuntimeRevisionsConfig, ListAgentEngineRuntimeRevisionsConfigDict +SandboxEnvironmentSpecComputerUseEnvironmentOrDict = Union[ + SandboxEnvironmentSpecComputerUseEnvironment, + SandboxEnvironmentSpecComputerUseEnvironmentDict, ] -class _ListAgentEngineRuntimeRevisionsRequestParameters(_common.BaseModel): - """Parameters for listing reasoning engine runtime revisions.""" - - name: Optional[str] = Field( - default=None, description="""Name of the reasoning engine.""" - ) - config: Optional[ListAgentEngineRuntimeRevisionsConfig] = Field( - default=None, description="""""" - ) +class SandboxEnvironmentSpecShellEnvironment(_common.BaseModel): + """The shell environment with customized settings.""" + pass -class _ListAgentEngineRuntimeRevisionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing reasoning engine runtime revisions.""" - name: Optional[str] - """Name of the reasoning engine.""" +class SandboxEnvironmentSpecShellEnvironmentDict(TypedDict, total=False): + """The shell environment with customized settings.""" - config: Optional[ListAgentEngineRuntimeRevisionsConfigDict] - """""" + pass -_ListAgentEngineRuntimeRevisionsRequestParametersOrDict = Union[ - _ListAgentEngineRuntimeRevisionsRequestParameters, - _ListAgentEngineRuntimeRevisionsRequestParametersDict, +SandboxEnvironmentSpecShellEnvironmentOrDict = Union[ + SandboxEnvironmentSpecShellEnvironment, SandboxEnvironmentSpecShellEnvironmentDict ] -class ListReasoningEnginesRuntimeRevisionsResponse(_common.BaseModel): - """Response for listing agent engine runtime revisions.""" +class SandboxEnvironmentSpec(_common.BaseModel): + """The specification of a sandbox environment.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + 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.""") ) - 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.""" + shell_environment: Optional[SandboxEnvironmentSpecShellEnvironment] = Field( + default=None, description="""Optional. The shell environment.""" ) -class ListReasoningEnginesRuntimeRevisionsResponseDict(TypedDict, total=False): - """Response for listing agent engine runtime revisions.""" +class SandboxEnvironmentSpecDict(TypedDict, total=False): + """The specification of a sandbox environment.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + code_execution_environment: Optional[ + SandboxEnvironmentSpecCodeExecutionEnvironmentDict + ] + """Optional. The code execution environment.""" - next_page_token: Optional[str] - """""" + computer_use_environment: Optional[SandboxEnvironmentSpecComputerUseEnvironmentDict] + """Optional. The computer use environment.""" - reasoning_engine_runtime_revisions: Optional[ - list[ReasoningEngineRuntimeRevisionDict] - ] - """List of reasoning engine runtime revisions.""" + shell_environment: Optional[SandboxEnvironmentSpecShellEnvironmentDict] + """Optional. The shell environment.""" -ListReasoningEnginesRuntimeRevisionsResponseOrDict = Union[ - ListReasoningEnginesRuntimeRevisionsResponse, - ListReasoningEnginesRuntimeRevisionsResponseDict, -] +SandboxEnvironmentSpecOrDict = Union[SandboxEnvironmentSpec, SandboxEnvironmentSpecDict] -class DeleteAgentEngineRuntimeRevisionConfig(_common.BaseModel): - """Config for deleting an Agent Engine Runtime Revision.""" +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.""" ) + 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 DeleteAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Runtime Revision.""" +class CreateRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for creating a Sandbox.""" 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}""" + + 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.""" + -DeleteAgentEngineRuntimeRevisionConfigOrDict = Union[ - DeleteAgentEngineRuntimeRevisionConfig, DeleteAgentEngineRuntimeRevisionConfigDict +CreateRuntimeSandboxConfigOrDict = Union[ + CreateRuntimeSandboxConfig, CreateRuntimeSandboxConfigDict ] -class _DeleteAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): - """Parameters for deleting agent engine runtime revisions.""" +class _CreateRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for creating Agent Runtime Sandboxes.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine runtime revision to delete.""", + description="""Name of the Agent Runtime to create the sandbox under.""", ) - config: Optional[DeleteAgentEngineRuntimeRevisionConfig] = Field( + spec: Optional[SandboxEnvironmentSpec] = Field( + default=None, description="""The specification of the sandbox.""" + ) + config: Optional[CreateRuntimeSandboxConfig] = Field( default=None, description="""""" ) -class _DeleteAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engine runtime revisions.""" +class _CreateRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for creating Agent Runtime Sandboxes.""" name: Optional[str] - """Name of the agent engine runtime revision to delete.""" + """Name of the Agent Runtime to create the sandbox under.""" + + spec: Optional[SandboxEnvironmentSpecDict] + """The specification of the sandbox.""" - config: Optional[DeleteAgentEngineRuntimeRevisionConfigDict] + config: Optional[CreateRuntimeSandboxConfigDict] """""" -_DeleteAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ - _DeleteAgentEngineRuntimeRevisionRequestParameters, - _DeleteAgentEngineRuntimeRevisionRequestParametersDict, +_CreateRuntimeSandboxRequestParametersOrDict = Union[ + _CreateRuntimeSandboxRequestParameters, _CreateRuntimeSandboxRequestParametersDict ] -class DeleteAgentEngineRuntimeRevisionOperation(_common.BaseModel): - """Operation for deleting agent engine runtime revisions.""" +class SandboxEnvironmentConnectionInfo(_common.BaseModel): + """The connection information of the SandboxEnvironment.""" - 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}`.""", + load_balancer_hostname: Optional[str] = Field( + default=None, description="""Output only. The hostname of the load balancer.""" ) - metadata: Optional[dict[str, Any]] = Field( + load_balancer_ip: 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="""Output only. The IP address of the load balancer.""", ) - done: Optional[bool] = Field( + sandbox_internal_ip: Optional[str] = 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="""Output only. The internal IP address of the SandboxEnvironment.""", ) - error: Optional[dict[str, Any]] = Field( + routing_token: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Output only. The routing token for the SandboxEnvironment.""", ) -class DeleteAgentEngineRuntimeRevisionOperationDict(TypedDict, total=False): - """Operation for deleting agent engine runtime revisions.""" +class SandboxEnvironmentConnectionInfoDict(TypedDict, total=False): + """The connection information of the SandboxEnvironment.""" - 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}`.""" + load_balancer_hostname: Optional[str] + """Output only. The hostname of the load balancer.""" - 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.""" + load_balancer_ip: Optional[str] + """Output only. The IP address of the load balancer.""" - 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.""" + sandbox_internal_ip: Optional[str] + """Output only. The internal IP address of the SandboxEnvironment.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + routing_token: Optional[str] + """Output only. The routing token for the SandboxEnvironment.""" -DeleteAgentEngineRuntimeRevisionOperationOrDict = Union[ - DeleteAgentEngineRuntimeRevisionOperation, - DeleteAgentEngineRuntimeRevisionOperationDict, +SandboxEnvironmentConnectionInfoOrDict = Union[ + SandboxEnvironmentConnectionInfo, SandboxEnvironmentConnectionInfoDict ] -class GetDeleteAgentEngineRuntimeRevisionOperationConfig(_common.BaseModel): - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - +class SandboxEnvironment(_common.BaseModel): + """A sandbox environment.""" -class GetDeleteAgentEngineRuntimeRevisionOperationConfigDict(TypedDict, total=False): + expire_time: Optional[datetime.datetime] = Field( + default=None, + description="""Expiration time of the sandbox environment. + """, + ) + 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.""", + ) - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +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.""" -GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict = Union[ - GetDeleteAgentEngineRuntimeRevisionOperationConfig, - GetDeleteAgentEngineRuntimeRevisionOperationConfigDict, -] + create_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironment was created.""" + display_name: Optional[str] + """Required. The display name of the SandboxEnvironment.""" -class _GetDeleteAgentEngineRuntimeRevisionOperationParameters(_common.BaseModel): - """Parameters for getting an operation that deletes a agent engine runtime revision.""" + name: Optional[str] + """Identifier. The name of the SandboxEnvironment.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + spec: Optional[SandboxEnvironmentSpecDict] + """Optional. The configuration of the SandboxEnvironment.""" + + 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.""" + + update_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironment was most recently updated.""" + + latest_sandbox_environment_snapshot: Optional[str] + """Output only. The resource name of the latest snapshot taken for this SandboxEnvironment.""" + + 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.""" + + +SandboxEnvironmentOrDict = Union[SandboxEnvironment, SandboxEnvironmentDict] + + +class RuntimeSandboxOperation(_common.BaseModel): + """Operation that has an agent runtime sandbox 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}`.""", ) - config: Optional[GetDeleteAgentEngineRuntimeRevisionOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + 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[SandboxEnvironment] = Field( + default=None, description="""The Agent Runtime Sandbox.""" ) -class _GetDeleteAgentEngineRuntimeRevisionOperationParametersDict( - TypedDict, total=False -): - """Parameters for getting an operation that deletes a agent engine runtime revision.""" +class RuntimeSandboxOperationDict(TypedDict, total=False): + """Operation that has an agent runtime sandbox as a response.""" - operation_name: Optional[str] - """The server-assigned name for the 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}`.""" - config: Optional[GetDeleteAgentEngineRuntimeRevisionOperationConfigDict] - """Used to override the default configuration.""" + 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[SandboxEnvironmentDict] + """The Agent Runtime Sandbox.""" -_GetDeleteAgentEngineRuntimeRevisionOperationParametersOrDict = Union[ - _GetDeleteAgentEngineRuntimeRevisionOperationParameters, - _GetDeleteAgentEngineRuntimeRevisionOperationParametersDict, +RuntimeSandboxOperationOrDict = Union[ + RuntimeSandboxOperation, RuntimeSandboxOperationDict ] -class QueryAgentEngineRuntimeRevisionConfig(_common.BaseModel): - """Config for querying agent engine runtime revisions.""" +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_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 QueryAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): - """Config for querying agent engine runtime revisions.""" +class DeleteRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for deleting an Agent Runtime Sandbox.""" 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] - """""" - -QueryAgentEngineRuntimeRevisionConfigOrDict = Union[ - QueryAgentEngineRuntimeRevisionConfig, QueryAgentEngineRuntimeRevisionConfigDict +DeleteRuntimeSandboxConfigOrDict = Union[ + DeleteRuntimeSandboxConfig, DeleteRuntimeSandboxConfigDict ] -class _QueryAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): - """Parameters for querying agent engine runtime revisions.""" +class _DeleteRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtimes.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine runtime revision.""" + default=None, description="""Name of the agent runtime sandbox to delete.""" ) - config: Optional[QueryAgentEngineRuntimeRevisionConfig] = Field( + config: Optional[DeleteRuntimeSandboxConfig] = Field( default=None, description="""""" ) -class _QueryAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for querying agent engine runtime revisions.""" +class _DeleteRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtimes.""" name: Optional[str] - """Name of the agent engine runtime revision.""" + """Name of the agent runtime sandbox to delete.""" - config: Optional[QueryAgentEngineRuntimeRevisionConfigDict] + config: Optional[DeleteRuntimeSandboxConfigDict] """""" -_QueryAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ - _QueryAgentEngineRuntimeRevisionRequestParameters, - _QueryAgentEngineRuntimeRevisionRequestParametersDict, +_DeleteRuntimeSandboxRequestParametersOrDict = Union[ + _DeleteRuntimeSandboxRequestParameters, _DeleteRuntimeSandboxRequestParametersDict ] -class SandboxEnvironmentSpecCodeExecutionEnvironment(_common.BaseModel): - """The code execution environment with customized settings.""" +class DeleteRuntimeSandboxOperation(_common.BaseModel): + """Operation for deleting agent runtimes.""" - code_language: Optional[Language] = Field( + name: Optional[str] = Field( default=None, - description="""The coding language supported in this environment.""", + 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}`.""", ) - machine_config: Optional[MachineConfig] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""The machine config of the code execution environment.""", + 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 SandboxEnvironmentSpecCodeExecutionEnvironmentDict(TypedDict, total=False): - """The code execution environment with customized settings.""" +class DeleteRuntimeSandboxOperationDict(TypedDict, total=False): + """Operation for deleting agent runtimes.""" - code_language: Optional[Language] - """The coding language supported in this environment.""" + 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}`.""" - machine_config: Optional[MachineConfig] - """The machine config of the code execution environment.""" + 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.""" -SandboxEnvironmentSpecCodeExecutionEnvironmentOrDict = Union[ - SandboxEnvironmentSpecCodeExecutionEnvironment, - SandboxEnvironmentSpecCodeExecutionEnvironmentDict, + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +DeleteRuntimeSandboxOperationOrDict = Union[ + DeleteRuntimeSandboxOperation, DeleteRuntimeSandboxOperationDict ] -class SandboxEnvironmentSpecComputerUseEnvironment(_common.BaseModel): - """The computer use environment with customized settings.""" +class Metadata(_common.BaseModel): + """Metadata for a chunk.""" - pass + 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 SandboxEnvironmentSpecComputerUseEnvironmentDict(TypedDict, total=False): - """The computer use environment with customized settings.""" +class MetadataDict(TypedDict, total=False): + """Metadata for a chunk.""" - pass + 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.""" -SandboxEnvironmentSpecComputerUseEnvironmentOrDict = Union[ - SandboxEnvironmentSpecComputerUseEnvironment, - SandboxEnvironmentSpecComputerUseEnvironmentDict, -] +MetadataOrDict = Union[Metadata, MetadataDict] -class SandboxEnvironmentSpecShellEnvironment(_common.BaseModel): - """The shell environment with customized settings.""" +class Chunk(_common.BaseModel): + """A chunk of data.""" - pass + 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 SandboxEnvironmentSpecShellEnvironmentDict(TypedDict, total=False): - """The shell environment with customized settings.""" +class ChunkDict(TypedDict, total=False): + """A chunk of data.""" - pass + data: Optional[bytes] + """Required. The data in the chunk.""" + metadata: Optional[MetadataDict] + """Optional. Metadata that is associated with the data in the payload.""" -SandboxEnvironmentSpecShellEnvironmentOrDict = Union[ - SandboxEnvironmentSpecShellEnvironment, SandboxEnvironmentSpecShellEnvironmentDict -] + 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.""" -class SandboxEnvironmentSpec(_common.BaseModel): - """The specification of a sandbox environment.""" +ChunkOrDict = Union[Chunk, ChunkDict] - 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.""") - ) - shell_environment: Optional[SandboxEnvironmentSpecShellEnvironment] = Field( - default=None, description="""Optional. The shell environment.""" + +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 SandboxEnvironmentSpecDict(TypedDict, total=False): - """The specification of a sandbox environment.""" +class ExecuteCodeRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for executing code in an Agent Runtime sandbox.""" - code_execution_environment: Optional[ - SandboxEnvironmentSpecCodeExecutionEnvironmentDict - ] - """Optional. The code execution environment.""" - - computer_use_environment: Optional[SandboxEnvironmentSpecComputerUseEnvironmentDict] - """Optional. The computer use environment.""" - - shell_environment: Optional[SandboxEnvironmentSpecShellEnvironmentDict] - """Optional. The shell environment.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -SandboxEnvironmentSpecOrDict = Union[SandboxEnvironmentSpec, SandboxEnvironmentSpecDict] +ExecuteCodeRuntimeSandboxConfigOrDict = Union[ + ExecuteCodeRuntimeSandboxConfig, ExecuteCodeRuntimeSandboxConfigDict +] -class CreateAgentEngineSandboxConfig(_common.BaseModel): - """Config for creating a Sandbox.""" +class _ExecuteCodeRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for executing code in an agent runtime sandbox.""" - 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( + name: Optional[str] = Field( default=None, - description="""The TTL for this resource. The expiration time is computed: now + TTL.""", + description="""Name of the agent runtime sandbox to execute code in.""", ) - 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}""", + inputs: Optional[list[Chunk]] = Field( + default=None, description="""Inputs to the code execution.""" ) - 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}""", + config: Optional[ExecuteCodeRuntimeSandboxConfig] = Field( + default=None, description="""""" ) - 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 _ExecuteCodeRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for executing code in an agent runtime sandbox.""" + + name: Optional[str] + """Name of the agent runtime sandbox to execute code in.""" + + inputs: Optional[list[ChunkDict]] + """Inputs to the code execution.""" + + config: Optional[ExecuteCodeRuntimeSandboxConfigDict] + """""" + + +_ExecuteCodeRuntimeSandboxRequestParametersOrDict = Union[ + _ExecuteCodeRuntimeSandboxRequestParameters, + _ExecuteCodeRuntimeSandboxRequestParametersDict, +] + + +class ExecuteSandboxEnvironmentResponse(_common.BaseModel): + """The response for executing a sandbox environment.""" + + outputs: Optional[list[Chunk]] = Field( + default=None, description="""The outputs from the sandbox environment.""" ) -class CreateAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for creating a Sandbox.""" +class ExecuteSandboxEnvironmentResponseDict(TypedDict, total=False): + """The response for executing a sandbox environment.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + outputs: Optional[list[ChunkDict]] + """The outputs from the sandbox environment.""" - display_name: Optional[str] - """The display name of the sandbox.""" - description: Optional[str] - """The description of the sandbox.""" +ExecuteSandboxEnvironmentResponseOrDict = Union[ + ExecuteSandboxEnvironmentResponse, ExecuteSandboxEnvironmentResponseDict +] - 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.""" +class GetRuntimeSandboxConfig(_common.BaseModel): + """Config for getting an Agent Runtime Memory.""" - 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}""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) - 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.""" +class GetRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for getting an Agent Runtime Memory.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -CreateAgentEngineSandboxConfigOrDict = Union[ - CreateAgentEngineSandboxConfig, CreateAgentEngineSandboxConfigDict +GetRuntimeSandboxConfigOrDict = Union[ + GetRuntimeSandboxConfig, GetRuntimeSandboxConfigDict ] -class _CreateAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for creating Agent Engine Sandboxes.""" +class _GetRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for getting an agent runtime sandbox.""" 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.""" - ) - config: Optional[CreateAgentEngineSandboxConfig] = Field( - default=None, description="""""" + default=None, description="""Name of the agent runtime sandbox.""" ) + config: Optional[GetRuntimeSandboxConfig] = Field(default=None, description="""""") -class _CreateAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for creating Agent Engine Sandboxes.""" +class _GetRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for getting an agent runtime sandbox.""" 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 sandbox.""" - config: Optional[CreateAgentEngineSandboxConfigDict] + config: Optional[GetRuntimeSandboxConfigDict] """""" -_CreateAgentEngineSandboxRequestParametersOrDict = Union[ - _CreateAgentEngineSandboxRequestParameters, - _CreateAgentEngineSandboxRequestParametersDict, +_GetRuntimeSandboxRequestParametersOrDict = Union[ + _GetRuntimeSandboxRequestParameters, _GetRuntimeSandboxRequestParametersDict ] -class SandboxEnvironmentConnectionInfo(_common.BaseModel): - """The connection information of the SandboxEnvironment.""" +class ListRuntimeSandboxesConfig(_common.BaseModel): + """Config 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - routing_token: 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="""Output only. The routing token for the SandboxEnvironment.""", + description="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""", ) -class SandboxEnvironmentConnectionInfoDict(TypedDict, total=False): - """The connection information of the SandboxEnvironment.""" +class ListRuntimeSandboxesConfigDict(TypedDict, total=False): + """Config for listing agent runtime sandboxes.""" - load_balancer_hostname: Optional[str] - """Output only. The hostname of the load balancer.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - load_balancer_ip: Optional[str] - """Output only. The IP address of the load balancer.""" + page_size: Optional[int] + """""" - sandbox_internal_ip: Optional[str] - """Output only. The internal IP address of the SandboxEnvironment.""" + page_token: Optional[str] + """""" - routing_token: Optional[str] - """Output only. The routing token for the SandboxEnvironment.""" + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""" -SandboxEnvironmentConnectionInfoOrDict = Union[ - SandboxEnvironmentConnectionInfo, SandboxEnvironmentConnectionInfoDict +ListRuntimeSandboxesConfigOrDict = Union[ + ListRuntimeSandboxesConfig, ListRuntimeSandboxesConfigDict ] -class SandboxEnvironment(_common.BaseModel): - """A sandbox environment.""" +class _ListRuntimeSandboxesRequestParameters(_common.BaseModel): + """Parameters for listing agent runtime sandboxes.""" - expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Expiration time of the sandbox environment. - """, - ) - 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}`""", + default=None, description="""Name of the agent runtime.""" ) - 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[ListRuntimeSandboxesConfig] = Field( + default=None, description="""""" ) -class SandboxEnvironmentDict(TypedDict, total=False): - """A sandbox environment.""" - - expire_time: Optional[datetime.datetime] - """Expiration time of the sandbox environment. - """ +class _ListRuntimeSandboxesRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtime sandboxes.""" - connection_info: Optional[SandboxEnvironmentConnectionInfoDict] - """Output only. The connection information of the SandboxEnvironment.""" + name: Optional[str] + """Name of the agent runtime.""" - create_time: Optional[datetime.datetime] - """Output only. The timestamp when this SandboxEnvironment was created.""" + config: Optional[ListRuntimeSandboxesConfigDict] + """""" - display_name: Optional[str] - """Required. The display name of the SandboxEnvironment.""" - name: Optional[str] - """Identifier. The name of the SandboxEnvironment.""" +_ListRuntimeSandboxesRequestParametersOrDict = Union[ + _ListRuntimeSandboxesRequestParameters, _ListRuntimeSandboxesRequestParametersDict +] - spec: Optional[SandboxEnvironmentSpecDict] - """Optional. The configuration of the SandboxEnvironment.""" - state: Optional[SandboxState] - """Output only. The runtime state of the SandboxEnvironment.""" +class ListRuntimeSandboxesResponse(_common.BaseModel): + """Response for listing agent runtime sandboxes.""" - ttl: Optional[str] - """Optional. Input only. The TTL for the sandbox environment. The expiration time is computed: now + TTL.""" + 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_environments: Optional[list[SandboxEnvironment]] = Field( + default=None, description="""List of agent runtime sandboxes.""" + ) - update_time: Optional[datetime.datetime] - """Output only. The timestamp when this SandboxEnvironment was most recently updated.""" - latest_sandbox_environment_snapshot: Optional[str] - """Output only. The resource name of the latest snapshot taken for this SandboxEnvironment.""" +class ListRuntimeSandboxesResponseDict(TypedDict, total=False): + """Response for listing agent runtime sandboxes.""" - 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.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - 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}`""" + next_page_token: Optional[str] + """""" - sandbox_environment_template: Optional[str] - """Optional. The name of the SandboxEnvironmentTemplate specified in the parent Agent Engine resource that this SandboxEnvironment is created from.""" + sandbox_environments: Optional[list[SandboxEnvironmentDict]] + """List of agent runtime sandboxes.""" -SandboxEnvironmentOrDict = Union[SandboxEnvironment, SandboxEnvironmentDict] +ListRuntimeSandboxesResponseOrDict = Union[ + ListRuntimeSandboxesResponse, ListRuntimeSandboxesResponseDict +] -class AgentEngineSandboxOperation(_common.BaseModel): - """Operation that has an agent engine sandbox as a response.""" +class _GetRuntimeSandboxOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a sandbox 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}`.""", + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - 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.""", + config: Optional[GetRuntimeOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) - 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[SandboxEnvironment] = Field( - default=None, description="""The Agent Engine Sandbox.""" - ) - - -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 _GetRuntimeSandboxOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a sandbox as a response.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - response: Optional[SandboxEnvironmentDict] - """The Agent Engine Sandbox.""" + config: Optional[GetRuntimeOperationConfigDict] + """Used to override the default configuration.""" -AgentEngineSandboxOperationOrDict = Union[ - AgentEngineSandboxOperation, AgentEngineSandboxOperationDict +_GetRuntimeSandboxOperationParametersOrDict = Union[ + _GetRuntimeSandboxOperationParameters, _GetRuntimeSandboxOperationParametersDict ] -class DeleteAgentEngineSandboxConfig(_common.BaseModel): - """Config for deleting an Agent Engine Sandbox.""" +class SandboxEnvironmentTemplateCustomContainerSpec(_common.BaseModel): + """Specification for deploying from a custom container image.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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.""", ) -class DeleteAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Sandbox.""" +class SandboxEnvironmentTemplateCustomContainerSpecDict(TypedDict, total=False): + """Specification for deploying from a custom container image.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + 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.""" -DeleteAgentEngineSandboxConfigOrDict = Union[ - DeleteAgentEngineSandboxConfig, DeleteAgentEngineSandboxConfigDict +SandboxEnvironmentTemplateCustomContainerSpecOrDict = Union[ + SandboxEnvironmentTemplateCustomContainerSpec, + SandboxEnvironmentTemplateCustomContainerSpecDict, ] -class _DeleteAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for deleting agent engines.""" +class SandboxEnvironmentTemplateNetworkPort(_common.BaseModel): + """Represents a network port in a container.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine sandbox to delete.""" + port: Optional[int] = Field( + default=None, + description="""Optional. Port number to expose. This must be a valid port number, between 1 and 65535.""", ) - config: Optional[DeleteAgentEngineSandboxConfig] = Field( - default=None, description="""""" + protocol: Optional[Protocol] = Field( + default=None, + description="""Optional. Protocol for port. Defaults to TCP if not specified.""", ) -class _DeleteAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engines.""" +class SandboxEnvironmentTemplateNetworkPortDict(TypedDict, total=False): + """Represents a network port in a container.""" - name: Optional[str] - """Name of the agent engine sandbox to delete.""" + port: Optional[int] + """Optional. Port number to expose. This must be a valid port number, between 1 and 65535.""" - config: Optional[DeleteAgentEngineSandboxConfigDict] - """""" + protocol: Optional[Protocol] + """Optional. Protocol for port. Defaults to TCP if not specified.""" -_DeleteAgentEngineSandboxRequestParametersOrDict = Union[ - _DeleteAgentEngineSandboxRequestParameters, - _DeleteAgentEngineSandboxRequestParametersDict, +SandboxEnvironmentTemplateNetworkPortOrDict = Union[ + SandboxEnvironmentTemplateNetworkPort, SandboxEnvironmentTemplateNetworkPortDict ] -class DeleteAgentEngineSandboxOperation(_common.BaseModel): - """Operation for deleting agent engines.""" +class SandboxEnvironmentTemplateResourceRequirements(_common.BaseModel): + """Message to define resource requests and limits (mirroring Kubernetes) for each sandbox instance created from this template.""" - 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( + limits: Optional[dict[str, str]] = 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. The maximum amounts of compute resources allowed. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "500m", "1Gi").""", ) - error: Optional[dict[str, Any]] = Field( + requests: Optional[dict[str, str]] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Optional. The requested amounts of compute resources. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "250m", "512Mi").""", ) -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 SandboxEnvironmentTemplateResourceRequirementsDict(TypedDict, total=False): + """Message to define resource requests and limits (mirroring Kubernetes) for each sandbox instance created from this template.""" - 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.""" + 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").""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + 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").""" -DeleteAgentEngineSandboxOperationOrDict = Union[ - DeleteAgentEngineSandboxOperation, DeleteAgentEngineSandboxOperationDict +SandboxEnvironmentTemplateResourceRequirementsOrDict = Union[ + SandboxEnvironmentTemplateResourceRequirements, + SandboxEnvironmentTemplateResourceRequirementsDict, ] -class Metadata(_common.BaseModel): - """Metadata for a chunk.""" +class SandboxEnvironmentTemplateCustomContainerEnvironment(_common.BaseModel): + """The customized sandbox runtime environment for BYOC.""" - 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.""", + 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.""" + ) + resources: Optional[SandboxEnvironmentTemplateResourceRequirements] = Field( + default=None, description="""Resource requests and limits for the container.""" ) -class MetadataDict(TypedDict, total=False): - """Metadata for a chunk.""" +class SandboxEnvironmentTemplateCustomContainerEnvironmentDict(TypedDict, total=False): + """The customized sandbox runtime environment for BYOC.""" - 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.""" + custom_container_spec: Optional[SandboxEnvironmentTemplateCustomContainerSpecDict] + """The specification of the custom container environment.""" + ports: Optional[list[SandboxEnvironmentTemplateNetworkPortDict]] + """Ports to expose from the container.""" -MetadataOrDict = Union[Metadata, MetadataDict] + resources: Optional[SandboxEnvironmentTemplateResourceRequirementsDict] + """Resource requests and limits for the container.""" -class Chunk(_common.BaseModel): - """A chunk of data.""" +SandboxEnvironmentTemplateCustomContainerEnvironmentOrDict = Union[ + SandboxEnvironmentTemplateCustomContainerEnvironment, + SandboxEnvironmentTemplateCustomContainerEnvironmentDict, +] - data: Optional[bytes] = Field( - default=None, description="""Required. The data in the chunk.""" - ) - metadata: Optional[Metadata] = Field( + +class SandboxEnvironmentTemplateDefaultContainerEnvironment(_common.BaseModel): + """The default sandbox runtime environment for default container workloads.""" + + default_container_category: Optional[DefaultContainerCategory] = Field( default=None, - description="""Optional. Metadata that is associated with the data in the payload.""", + description="""Required. The category of the default container image.""", ) - mime_type: Optional[str] = Field( + resources: Optional[SandboxEnvironmentTemplateResourceRequirements] = 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="""Optional. Resource requests and limits for the default container.""", ) -class ChunkDict(TypedDict, total=False): - """A chunk of data.""" - - data: Optional[bytes] - """Required. The data in the chunk.""" +class SandboxEnvironmentTemplateDefaultContainerEnvironmentDict(TypedDict, total=False): + """The default sandbox runtime environment for default container workloads.""" - metadata: Optional[MetadataDict] - """Optional. Metadata that is associated with the data in the payload.""" + default_container_category: Optional[DefaultContainerCategory] + """Required. The category of the default container image.""" - 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.""" + resources: Optional[SandboxEnvironmentTemplateResourceRequirementsDict] + """Optional. Resource requests and limits for the default container.""" -ChunkOrDict = Union[Chunk, ChunkDict] +SandboxEnvironmentTemplateDefaultContainerEnvironmentOrDict = Union[ + SandboxEnvironmentTemplateDefaultContainerEnvironment, + SandboxEnvironmentTemplateDefaultContainerEnvironmentDict, +] -class ExecuteCodeAgentEngineSandboxConfig(_common.BaseModel): - """Config for executing code in an Agent Engine sandbox.""" +class SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfig(_common.BaseModel): + """Configuration for peering a customer's private DNS zone so that sandbox egress can resolve customer-internal domains via the customer VPC.""" - 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 ExecuteCodeAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for executing code in an Agent Engine sandbox.""" +class SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigDict( + TypedDict, total=False +): + """Configuration for peering a customer's private DNS zone so that sandbox egress can resolve customer-internal domains via the customer VPC.""" - 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.""" + 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.""" -ExecuteCodeAgentEngineSandboxConfigOrDict = Union[ - ExecuteCodeAgentEngineSandboxConfig, ExecuteCodeAgentEngineSandboxConfigDict + +SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigOrDict = Union[ + SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfig, + SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigDict, ] -class _ExecuteCodeAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for executing code in an agent engine sandbox.""" +class SandboxEnvironmentTemplateEgressControlConfig(_common.BaseModel): + """Configuration for egress control of sandbox instances.""" - name: Optional[str] = Field( + internet_access: Optional[bool] = Field( + default=None, description="""Optional. Whether to allow internet access.""" + ) + customer_vpc_network: Optional[str] = Field( default=None, - description="""Name of the agent engine sandbox to execute code in.""", + description="""Optional. The customer VPC network that sandbox egress is routed into.""", ) - inputs: Optional[list[Chunk]] = Field( - default=None, description="""Inputs to the code execution.""" + dns_peering_configs: Optional[ + list[SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfig] + ] = Field( + default=None, + description="""Optional. DNS peering configurations that allow sandbox egress to resolve customer-internal domains via the customer VPC.""", ) - config: Optional[ExecuteCodeAgentEngineSandboxConfig] = Field( - default=None, description="""""" + network_attachment: Optional[str] = Field( + default=None, + description="""Optional. The name of the customer VPC NetworkAttachment used to draw a PSC interface IP into the customer VPC for sandbox egress.""", ) -class _ExecuteCodeAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for executing code in an agent engine sandbox.""" +class SandboxEnvironmentTemplateEgressControlConfigDict(TypedDict, total=False): + """Configuration for egress control of sandbox instances.""" - name: Optional[str] - """Name of the agent engine sandbox to execute code in.""" + internet_access: Optional[bool] + """Optional. Whether to allow internet access.""" - inputs: Optional[list[ChunkDict]] - """Inputs to the code execution.""" + customer_vpc_network: Optional[str] + """Optional. The customer VPC network that sandbox egress is routed into.""" - config: Optional[ExecuteCodeAgentEngineSandboxConfigDict] - """""" + dns_peering_configs: Optional[ + list[SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigDict] + ] + """Optional. DNS peering configurations that allow sandbox egress to resolve customer-internal domains via the customer VPC.""" + + network_attachment: Optional[str] + """Optional. The name of the customer VPC NetworkAttachment used to draw a PSC interface IP into the customer VPC for sandbox egress.""" -_ExecuteCodeAgentEngineSandboxRequestParametersOrDict = Union[ - _ExecuteCodeAgentEngineSandboxRequestParameters, - _ExecuteCodeAgentEngineSandboxRequestParametersDict, +SandboxEnvironmentTemplateEgressControlConfigOrDict = Union[ + SandboxEnvironmentTemplateEgressControlConfig, + SandboxEnvironmentTemplateEgressControlConfigDict, ] -class ExecuteSandboxEnvironmentResponse(_common.BaseModel): - """The response for executing a sandbox environment.""" +class CreateSandboxEnvironmentTemplateConfig(_common.BaseModel): + """Config for creating a Sandbox Template.""" - outputs: Optional[list[Chunk]] = Field( - default=None, description="""The outputs from the sandbox environment.""" + 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 ExecuteSandboxEnvironmentResponseDict(TypedDict, total=False): - """The response for executing a sandbox environment.""" - - outputs: Optional[list[ChunkDict]] - """The outputs from the sandbox environment.""" - - -ExecuteSandboxEnvironmentResponseOrDict = Union[ - ExecuteSandboxEnvironmentResponse, ExecuteSandboxEnvironmentResponseDict -] - +class CreateSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): + """Config for creating a Sandbox Template.""" -class GetAgentEngineSandboxConfig(_common.BaseModel): - """Config for getting an Agent Engine Memory.""" + 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 GetAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Memory.""" + 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.""" -GetAgentEngineSandboxConfigOrDict = Union[ - GetAgentEngineSandboxConfig, GetAgentEngineSandboxConfigDict +CreateSandboxEnvironmentTemplateConfigOrDict = Union[ + CreateSandboxEnvironmentTemplateConfig, CreateSandboxEnvironmentTemplateConfigDict ] -class _GetAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for getting 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.""" + default=None, + description="""Name of the agent runtime to create the template under.""", ) - config: Optional[GetAgentEngineSandboxConfig] = Field( + config: Optional[CreateSandboxEnvironmentTemplateConfig] = Field( default=None, description="""""" ) + display_name: Optional[str] = Field( + default=None, description="""The display name of the sandbox template.""" + ) -class _GetAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine sandbox.""" +class _CreateSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): + """Parameters for creating Sandbox Environment Templates.""" name: Optional[str] - """Name of the agent engine sandbox.""" + """Name of the agent runtime to create the template under.""" - config: Optional[GetAgentEngineSandboxConfigDict] + config: Optional[CreateSandboxEnvironmentTemplateConfigDict] """""" + display_name: Optional[str] + """The display name of the sandbox template.""" + -_GetAgentEngineSandboxRequestParametersOrDict = Union[ - _GetAgentEngineSandboxRequestParameters, _GetAgentEngineSandboxRequestParametersDict +_CreateSandboxEnvironmentTemplateRequestParametersOrDict = Union[ + _CreateSandboxEnvironmentTemplateRequestParameters, + _CreateSandboxEnvironmentTemplateRequestParametersDict, ] -class ListAgentEngineSandboxesConfig(_common.BaseModel): - """Config for listing agent engine sandboxes.""" +class SandboxEnvironmentTemplate(_common.BaseModel): + """A sandbox environment template.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. The timestamp when this SandboxEnvironmentTemplate was created.""", ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( + custom_container_environment: Optional[ + SandboxEnvironmentTemplateCustomContainerEnvironment + ] = Field( default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""", + 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 ListAgentEngineSandboxesConfigDict(TypedDict, total=False): - """Config for listing agent engine sandboxes.""" +class SandboxEnvironmentTemplateDict(TypedDict, total=False): + """A sandbox environment template.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + create_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironmentTemplate was created.""" - page_size: Optional[int] - """""" + custom_container_environment: Optional[ + SandboxEnvironmentTemplateCustomContainerEnvironmentDict + ] + """The sandbox environment for custom container workloads.""" - page_token: Optional[str] - """""" + default_container_environment: Optional[ + SandboxEnvironmentTemplateDefaultContainerEnvironmentDict + ] + """The sandbox environment for default container workloads.""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" + display_name: Optional[str] + """Required. The display name of the SandboxEnvironmentTemplate.""" + + 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}`""" + + state: Optional[ + Literal[ + "UNSPECIFIED", + "PROVISIONING", + "ACTIVE", + "DEPROVISIONING", + "DELETED", + "FAILED", + ] + ] + """Output only. The state of the sandbox environment template.""" + + update_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironmentTemplate was most recently updated.""" -ListAgentEngineSandboxesConfigOrDict = Union[ - ListAgentEngineSandboxesConfig, ListAgentEngineSandboxesConfigDict +SandboxEnvironmentTemplateOrDict = Union[ + SandboxEnvironmentTemplate, SandboxEnvironmentTemplateDict ] -class _ListAgentEngineSandboxesRequestParameters(_common.BaseModel): - """Parameters for listing agent engine sandboxes.""" +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.""" + 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[ListAgentEngineSandboxesConfig] = 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 _ListAgentEngineSandboxesRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engine sandboxes.""" +class SandboxEnvironmentTemplateOperationDict(TypedDict, total=False): + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] - """Name of the agent engine.""" + """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[ListAgentEngineSandboxesConfigDict] - """""" + 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.""" -_ListAgentEngineSandboxesRequestParametersOrDict = Union[ - _ListAgentEngineSandboxesRequestParameters, - _ListAgentEngineSandboxesRequestParametersDict, -] + response: Optional[SandboxEnvironmentTemplateDict] + """The Agent Runtime Sandbox Template.""" -class ListAgentEngineSandboxesResponse(_common.BaseModel): - """Response for listing agent engine sandboxes.""" +SandboxEnvironmentTemplateOperationOrDict = Union[ + SandboxEnvironmentTemplateOperation, SandboxEnvironmentTemplateOperationDict +] - 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_environments: Optional[list[SandboxEnvironment]] = Field( - default=None, description="""List of agent engine sandboxes.""" - ) +class DeleteSandboxEnvironmentTemplateConfig(_common.BaseModel): + """Config for deleting a Sandbox Template.""" -class ListAgentEngineSandboxesResponseDict(TypedDict, total=False): - """Response for listing agent engine sandboxes.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" - next_page_token: Optional[str] - """""" +class DeleteSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): + """Config for deleting a Sandbox Template.""" - sandbox_environments: Optional[list[SandboxEnvironmentDict]] - """List of agent engine sandboxes.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ListAgentEngineSandboxesResponseOrDict = Union[ - ListAgentEngineSandboxesResponse, ListAgentEngineSandboxesResponseDict +DeleteSandboxEnvironmentTemplateConfigOrDict = Union[ + DeleteSandboxEnvironmentTemplateConfig, DeleteSandboxEnvironmentTemplateConfigDict ] -class _GetAgentEngineSandboxOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a sandbox as a response.""" +class _DeleteSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): + """Parameters for deleting sandbox templates.""" - 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 sandbox template to delete.""" ) - config: Optional[GetAgentEngineOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + config: Optional[DeleteSandboxEnvironmentTemplateConfig] = Field( + default=None, description="""""" ) -class _GetAgentEngineSandboxOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a sandbox as a response.""" +class _DeleteSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): + """Parameters for deleting sandbox templates.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + name: Optional[str] + """Name of the sandbox template to delete.""" - config: Optional[GetAgentEngineOperationConfigDict] - """Used to override the default configuration.""" + config: Optional[DeleteSandboxEnvironmentTemplateConfigDict] + """""" -_GetAgentEngineSandboxOperationParametersOrDict = Union[ - _GetAgentEngineSandboxOperationParameters, - _GetAgentEngineSandboxOperationParametersDict, +_DeleteSandboxEnvironmentTemplateRequestParametersOrDict = Union[ + _DeleteSandboxEnvironmentTemplateRequestParameters, + _DeleteSandboxEnvironmentTemplateRequestParametersDict, ] -class SandboxEnvironmentTemplateCustomContainerSpec(_common.BaseModel): - """Specification for deploying from a custom container image.""" +class DeleteSandboxEnvironmentTemplateOperation(_common.BaseModel): + """Operation for deleting sandbox templates.""" - image_uri: Optional[str] = Field( + name: 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.""", + 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 SandboxEnvironmentTemplateCustomContainerSpecDict(TypedDict, total=False): - """Specification for deploying from a custom 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.""" - - -SandboxEnvironmentTemplateCustomContainerSpecOrDict = Union[ - SandboxEnvironmentTemplateCustomContainerSpec, - SandboxEnvironmentTemplateCustomContainerSpecDict, -] - - -class SandboxEnvironmentTemplateNetworkPort(_common.BaseModel): - """Represents a network port in a container.""" - - port: Optional[int] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. Port number to expose. This must be a valid port number, between 1 and 65535.""", + 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.""", ) - protocol: Optional[Protocol] = Field( + done: Optional[bool] = Field( default=None, - description="""Optional. Protocol for port. Defaults to TCP if not specified.""", + 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 SandboxEnvironmentTemplateNetworkPortDict(TypedDict, total=False): - """Represents a network port in a container.""" +class DeleteSandboxEnvironmentTemplateOperationDict(TypedDict, total=False): + """Operation for deleting sandbox templates.""" - port: Optional[int] - """Optional. Port number to expose. This must be a valid port number, between 1 and 65535.""" + 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}`.""" - protocol: Optional[Protocol] - """Optional. Protocol for port. Defaults to TCP if not specified.""" + 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.""" -SandboxEnvironmentTemplateNetworkPortOrDict = Union[ - SandboxEnvironmentTemplateNetworkPort, SandboxEnvironmentTemplateNetworkPortDict + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +DeleteSandboxEnvironmentTemplateOperationOrDict = Union[ + DeleteSandboxEnvironmentTemplateOperation, + DeleteSandboxEnvironmentTemplateOperationDict, ] -class SandboxEnvironmentTemplateResourceRequirements(_common.BaseModel): - """Message to define resource requests and limits (mirroring Kubernetes) for each sandbox instance created from this template.""" +class GetSandboxEnvironmentTemplateConfig(_common.BaseModel): + """Config for getting a Sandbox Template.""" - 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").""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class SandboxEnvironmentTemplateResourceRequirementsDict(TypedDict, total=False): - """Message to define resource requests and limits (mirroring Kubernetes) for each sandbox instance created from this template.""" - - 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").""" +class GetSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): + """Config for getting a Sandbox Template.""" - 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").""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -SandboxEnvironmentTemplateResourceRequirementsOrDict = Union[ - SandboxEnvironmentTemplateResourceRequirements, - SandboxEnvironmentTemplateResourceRequirementsDict, +GetSandboxEnvironmentTemplateConfigOrDict = Union[ + GetSandboxEnvironmentTemplateConfig, GetSandboxEnvironmentTemplateConfigDict ] -class SandboxEnvironmentTemplateCustomContainerEnvironment(_common.BaseModel): - """The customized sandbox runtime environment for BYOC.""" +class _GetSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): + """Parameters for getting a sandbox template.""" - 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.""" + name: Optional[str] = Field( + default=None, description="""Name of the sandbox template.""" ) - resources: Optional[SandboxEnvironmentTemplateResourceRequirements] = Field( - default=None, description="""Resource requests and limits for the container.""" + config: Optional[GetSandboxEnvironmentTemplateConfig] = Field( + default=None, description="""""" ) -class SandboxEnvironmentTemplateCustomContainerEnvironmentDict(TypedDict, total=False): - """The customized sandbox runtime environment for BYOC.""" - - custom_container_spec: Optional[SandboxEnvironmentTemplateCustomContainerSpecDict] - """The specification of the custom container environment.""" +class _GetSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): + """Parameters for getting a sandbox template.""" - ports: Optional[list[SandboxEnvironmentTemplateNetworkPortDict]] - """Ports to expose from the container.""" + name: Optional[str] + """Name of the sandbox template.""" - resources: Optional[SandboxEnvironmentTemplateResourceRequirementsDict] - """Resource requests and limits for the container.""" + config: Optional[GetSandboxEnvironmentTemplateConfigDict] + """""" -SandboxEnvironmentTemplateCustomContainerEnvironmentOrDict = Union[ - SandboxEnvironmentTemplateCustomContainerEnvironment, - SandboxEnvironmentTemplateCustomContainerEnvironmentDict, +_GetSandboxEnvironmentTemplateRequestParametersOrDict = Union[ + _GetSandboxEnvironmentTemplateRequestParameters, + _GetSandboxEnvironmentTemplateRequestParametersDict, ] -class SandboxEnvironmentTemplateDefaultContainerEnvironment(_common.BaseModel): - """The default sandbox runtime environment for default container workloads.""" +class ListSandboxEnvironmentTemplatesConfig(_common.BaseModel): + """Config for listing sandbox templates.""" - default_container_category: Optional[DefaultContainerCategory] = Field( - default=None, - description="""Required. The category of the default container image.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - resources: Optional[SandboxEnvironmentTemplateResourceRequirements] = 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. Resource requests and limits for the default container.""", + description="""An expression for filtering the results of the request.""", ) -class SandboxEnvironmentTemplateDefaultContainerEnvironmentDict(TypedDict, total=False): - """The default sandbox runtime environment for default container workloads.""" +class ListSandboxEnvironmentTemplatesConfigDict(TypedDict, total=False): + """Config for listing sandbox templates.""" - default_container_category: Optional[DefaultContainerCategory] - """Required. The category of the default container image.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - resources: Optional[SandboxEnvironmentTemplateResourceRequirementsDict] - """Optional. Resource requests and limits for the default container.""" + page_size: Optional[int] + """""" + page_token: Optional[str] + """""" -SandboxEnvironmentTemplateDefaultContainerEnvironmentOrDict = Union[ - SandboxEnvironmentTemplateDefaultContainerEnvironment, - SandboxEnvironmentTemplateDefaultContainerEnvironmentDict, + filter: Optional[str] + """An expression for filtering the results of the request.""" + + +ListSandboxEnvironmentTemplatesConfigOrDict = Union[ + ListSandboxEnvironmentTemplatesConfig, ListSandboxEnvironmentTemplatesConfigDict ] -class SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfig(_common.BaseModel): - """Configuration for peering a customer's private DNS zone so that sandbox egress can resolve customer-internal domains via the customer VPC.""" +class _ListSandboxEnvironmentTemplatesRequestParameters(_common.BaseModel): + """Parameters for listing sandbox templates.""" - 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.""", + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime.""" ) - 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.""", + config: Optional[ListSandboxEnvironmentTemplatesConfig] = Field( + default=None, description="""""" ) -class SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigDict( - TypedDict, total=False -): - """Configuration for peering a customer's private DNS zone so that sandbox egress can resolve customer-internal domains via the customer VPC.""" - - 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.""" +class _ListSandboxEnvironmentTemplatesRequestParametersDict(TypedDict, total=False): + """Parameters for listing sandbox templates.""" - target_network: Optional[str] - """Required. The VPC network name in the target_project where the DNS zone specified by 'domain' is visible.""" + name: Optional[str] + """Name of the agent runtime.""" - 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.""" + config: Optional[ListSandboxEnvironmentTemplatesConfigDict] + """""" -SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigOrDict = Union[ - SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfig, - SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigDict, +_ListSandboxEnvironmentTemplatesRequestParametersOrDict = Union[ + _ListSandboxEnvironmentTemplatesRequestParameters, + _ListSandboxEnvironmentTemplatesRequestParametersDict, ] -class SandboxEnvironmentTemplateEgressControlConfig(_common.BaseModel): - """Configuration for egress control of sandbox instances.""" +class ListSandboxEnvironmentTemplatesResponse(_common.BaseModel): + """Response for listing sandbox templates.""" - internet_access: Optional[bool] = Field( - default=None, description="""Optional. Whether to allow internet access.""" - ) - customer_vpc_network: Optional[str] = Field( - default=None, - description="""Optional. The customer VPC network that sandbox egress is routed into.""", - ) - dns_peering_configs: Optional[ - list[SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfig] - ] = Field( - default=None, - description="""Optional. DNS peering configurations that allow sandbox egress to resolve customer-internal domains via the customer VPC.""", + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - network_attachment: Optional[str] = Field( - default=None, - description="""Optional. The name of the customer VPC NetworkAttachment used to draw a PSC interface IP into the customer VPC for sandbox egress.""", + next_page_token: Optional[str] = Field(default=None, description="""""") + sandbox_environment_templates: Optional[list[SandboxEnvironmentTemplate]] = Field( + default=None, description="""List of sandbox templates.""" ) -class SandboxEnvironmentTemplateEgressControlConfigDict(TypedDict, total=False): - """Configuration for egress control of sandbox instances.""" - - internet_access: Optional[bool] - """Optional. Whether to allow internet access.""" +class ListSandboxEnvironmentTemplatesResponseDict(TypedDict, total=False): + """Response for listing sandbox templates.""" - customer_vpc_network: Optional[str] - """Optional. The customer VPC network that sandbox egress is routed into.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - dns_peering_configs: Optional[ - list[SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigDict] - ] - """Optional. DNS peering configurations that allow sandbox egress to resolve customer-internal domains via the customer VPC.""" + next_page_token: Optional[str] + """""" - network_attachment: Optional[str] - """Optional. The name of the customer VPC NetworkAttachment used to draw a PSC interface IP into the customer VPC for sandbox egress.""" + sandbox_environment_templates: Optional[list[SandboxEnvironmentTemplateDict]] + """List of sandbox templates.""" -SandboxEnvironmentTemplateEgressControlConfigOrDict = Union[ - SandboxEnvironmentTemplateEgressControlConfig, - SandboxEnvironmentTemplateEgressControlConfigDict, +ListSandboxEnvironmentTemplatesResponseOrDict = Union[ + ListSandboxEnvironmentTemplatesResponse, ListSandboxEnvironmentTemplatesResponseDict ] -class CreateSandboxEnvironmentTemplateConfig(_common.BaseModel): - """Config for creating a Sandbox Template.""" +class _GetSandboxEnvironmentTemplateOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a sandbox template as a response.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", + config: Optional[GetRuntimeOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) - custom_container_environment: Optional[ - SandboxEnvironmentTemplateCustomContainerEnvironment - ] = Field( - default=None, - description="""The custom container environment for the sandbox template.""", + + +class _GetSandboxEnvironmentTemplateOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a sandbox template as a response.""" + + operation_name: Optional[str] + """The server-assigned name for the operation.""" + + config: Optional[GetRuntimeOperationConfigDict] + """Used to override the default configuration.""" + + +_GetSandboxEnvironmentTemplateOperationParametersOrDict = Union[ + _GetSandboxEnvironmentTemplateOperationParameters, + _GetSandboxEnvironmentTemplateOperationParametersDict, +] + + +class CreateRuntimeSandboxSnapshotConfig(_common.BaseModel): + """Config for creating a Sandbox Environment Snapshot.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - default_container_environment: Optional[ - SandboxEnvironmentTemplateDefaultContainerEnvironment - ] = Field( + 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 default container environment for the sandbox template.""", + description="""The TTL for this resource. The expiration time is computed: now + TTL.""", ) - egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = ( - Field( - default=None, - description="""The egress control config for the sandbox template.""", - ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) -class CreateSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): - """Config for creating a Sandbox Template.""" +class CreateRuntimeSandboxSnapshotConfigDict(TypedDict, total=False): + """Config for creating a Sandbox Environment Snapshot.""" 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.""" + display_name: Optional[str] + """The display name of the sandbox snapshot.""" - custom_container_environment: Optional[ - SandboxEnvironmentTemplateCustomContainerEnvironmentDict - ] - """The custom container environment for the sandbox template.""" + owner: Optional[str] + """The owner of the sandbox snapshot.""" - default_container_environment: Optional[ - SandboxEnvironmentTemplateDefaultContainerEnvironmentDict - ] - """The default container environment for the sandbox template.""" + ttl: Optional[str] + """The TTL for this resource. The expiration time is computed: now + TTL.""" - 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, @@ -18874,13 +18169,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}`.""" @@ -18894,64 +18189,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, @@ -18971,8 +18267,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}`.""" @@ -18987,61 +18283,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.""" @@ -19054,8 +18350,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.""" @@ -19070,52 +18366,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.""" @@ -19123,210 +18420,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}`""" + """Identifier. The resource name of the session. Format: 'projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}'.""" - 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.""" - - 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.""" - - 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.""" - - size_bytes: Optional[int] - """Optional. Output only. Size of the snapshot data in bytes.""" - - source_sandbox_environment: Optional[str] - """Required. The resource name of the source SandboxEnvironment this snapshot was taken from.""" + session_state: Optional[dict[str, Any]] + """Optional. Session specific memory which stores key conversation points.""" ttl: Optional[str] - """Optional. Input only. The TTL for the sandbox environment snapshot. The expiration time is computed: now + TTL.""" + """Optional. Input only. The TTL for this session. The minimum value is 24 hours.""" update_time: Optional[datetime.datetime] - """Output only. The timestamp when this SandboxEnvironment was most recently updated.""" + """Output only. Timestamp when the session was 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, @@ -19344,13 +18654,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}`.""" @@ -19364,65 +18674,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, @@ -19442,8 +18750,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}`.""" @@ -19458,61 +18766,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.""" @@ -19521,12 +18825,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.""" @@ -19538,56 +18843,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.""" @@ -19595,44 +18899,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.""" @@ -19666,10 +18969,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.""" @@ -19697,305 +19008,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}`.""" - - 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 -] + """Name of the agent runtime session.""" + author: Optional[str] + """Author of the agent runtime session event.""" -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.""" - ) - + 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.""" @@ -20009,8 +19308,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.""" @@ -20026,3395 +19325,3381 @@ 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.""" + """Name of the agent runtime session.""" - config: Optional[ListAgentEngineSessionsConfigDict] + config: Optional[ListRuntimeSessionEventsConfigDict] """""" -_ListAgentEngineSessionsRequestParametersOrDict = Union[ - _ListAgentEngineSessionsRequestParameters, - _ListAgentEngineSessionsRequestParametersDict, +_ListRuntimeSessionEventsRequestParametersOrDict = Union[ + _ListRuntimeSessionEventsRequestParameters, + _ListRuntimeSessionEventsRequestParametersDict, ] -class ListReasoningEnginesSessionsResponse(_common.BaseModel): - """Response for listing agent engine sessions.""" +class SessionEvent(_common.BaseModel): + """A session event.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + content: Optional[genai_types.Content] = Field( + default=None, + description="""Optional. Content of the event provided by the author.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - sessions: Optional[list[Session]] = Field( - default=None, description="""List of agent engine sessions.""" + 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( + 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.""" + ) + invocation_id: Optional[str] = Field( + default=None, + description="""Required. The invocation id of the event, multiple events can have the same invocation id.""", + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. The resource name of the event. Format:`projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}/events/{event}`.""", + ) + timestamp: Optional[datetime.datetime] = Field( + default=None, + description="""Required. Timestamp when the event was created on client side.""", + ) + raw_event: Optional[dict[str, Any]] = Field( + default=None, + description="""Optional. Weakly typed raw event data in proto struct format.""", ) -class ListReasoningEnginesSessionsResponseDict(TypedDict, total=False): - """Response for listing agent engine sessions.""" +class SessionEventDict(TypedDict, total=False): + """A session event.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + content: Optional[genai_types.Content] + """Optional. Content of the event provided by the author.""" - next_page_token: Optional[str] - """""" + actions: Optional[EventActionsDict] + """Optional. Actions executed by the agent.""" - sessions: Optional[list[SessionDict]] - """List of agent engine sessions.""" + author: Optional[str] + """Required. The name of the agent that sent the event, or user.""" + error_code: Optional[str] + """Optional. Error code if the response is an error. Code varies by model.""" -ListReasoningEnginesSessionsResponseOrDict = Union[ - ListReasoningEnginesSessionsResponse, ListReasoningEnginesSessionsResponseDict -] + error_message: Optional[str] + """Optional. Error message if the response is an error.""" + event_metadata: Optional[EventMetadataDict] + """Optional. Metadata relating to this event.""" -class _GetAgentEngineSessionOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a session as a response.""" + invocation_id: Optional[str] + """Required. The invocation id of the event, multiple events can have the same invocation id.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + 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.""" + + raw_event: Optional[dict[str, Any]] + """Optional. Weakly typed raw event data in proto struct format.""" + + +SessionEventOrDict = Union[SessionEvent, SessionEventDict] + + +class ListRuntimeSessionEventsResponse(_common.BaseModel): + """Response for listing agent runtime session events.""" + + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - config: Optional[GetAgentEngineOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + next_page_token: Optional[str] = Field(default=None, description="""""") + session_events: Optional[list[SessionEvent]] = Field( + default=None, description="""List of session events.""" ) -class _GetAgentEngineSessionOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a session as a response.""" +class ListRuntimeSessionEventsResponseDict(TypedDict, total=False): + """Response for listing agent runtime session events.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - config: Optional[GetAgentEngineOperationConfigDict] - """Used to override the default configuration.""" + next_page_token: Optional[str] + """""" + session_events: Optional[list[SessionEventDict]] + """List of session events.""" -_GetAgentEngineSessionOperationParametersOrDict = Union[ - _GetAgentEngineSessionOperationParameters, - _GetAgentEngineSessionOperationParametersDict, + +ListRuntimeSessionEventsResponseOrDict = Union[ + ListRuntimeSessionEventsResponse, ListRuntimeSessionEventsResponseDict ] -class UpdateAgentEngineSessionConfig(_common.BaseModel): - """Config for updating agent engine session.""" +class GeminiExample(_common.BaseModel): + """Represents a Gemini example.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + model: Optional[str] = Field( + default=None, description="""The model used to generate the Gemini example.""" ) - display_name: Optional[str] = Field( - default=None, description="""The display name of the session.""" + contents: Optional[list[genai_types.Content]] = Field( + default=None, description="""Contents of the Gemini example.""" ) - session_state: Optional[dict[str, Any]] = Field( - default=None, - description="""Session state which stores key conversation points.""", + system_instruction: Optional[genai_types.Content] = Field( + default=None, description="""System instruction for the Gemini example.""" ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", + cached_content: Optional[str] = Field( + default=None, description="""Cached content for the Gemini example.""" ) - ttl: Optional[str] = Field( - default=None, - description="""Optional. Input only. The TTL for this resource. - - The expiration time is computed: now + TTL.""", + tools: Optional[list[genai_types.Tool]] = Field( + default=None, description="""Tools for the Gemini example.""" ) - 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.""", + tool_config: Optional[genai_types.ToolConfig] = Field( + default=None, description="""Tools for the Gemini example.""" ) - 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.""", + safety_settings: Optional[list[genai_types.SafetySetting]] = Field( + default=None, description="""Safety settings for the Gemini example.""" ) - 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.""", + generation_config: Optional[genai_types.GenerationConfig] = Field( + default=None, description="""Generation config for the Gemini example.""" ) - update_mask: Optional[str] = Field( + model_armor_config: Optional[genai_types.ModelArmorConfig] = 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 engine session to update.""" + description="""Optional. Settings for prompt and response sanitization using the Model Armor service. If supplied, safety_settings must not be supplied.""", ) -class UpdateAgentEngineSessionConfigDict(TypedDict, total=False): - """Config for updating agent engine session.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - display_name: Optional[str] - """The display name of the session.""" +class GeminiExampleDict(TypedDict, total=False): + """Represents a Gemini example.""" - session_state: Optional[dict[str, Any]] - """Session state which stores key conversation points.""" + model: Optional[str] + """The model used to generate the Gemini example.""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + contents: Optional[list[genai_types.Content]] + """Contents of the Gemini example.""" - ttl: Optional[str] - """Optional. Input only. The TTL for this resource. + system_instruction: Optional[genai_types.Content] + """System instruction for the Gemini example.""" - The expiration time is computed: now + TTL.""" + cached_content: Optional[str] + """Cached content for the Gemini example.""" - 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.""" + tools: Optional[list[genai_types.Tool]] + """Tools for the Gemini example.""" - 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.""" + tool_config: Optional[genai_types.ToolConfig] + """Tools for the Gemini example.""" - 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.""" + safety_settings: Optional[list[genai_types.SafetySetting]] + """Safety settings for the Gemini example.""" - update_mask: Optional[str] - """The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" + generation_config: Optional[genai_types.GenerationConfig] + """Generation config for the Gemini example.""" - user_id: Optional[str] - """User ID of the agent engine session to update.""" + 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.""" -UpdateAgentEngineSessionConfigOrDict = Union[ - UpdateAgentEngineSessionConfig, UpdateAgentEngineSessionConfigDict -] +GeminiExampleOrDict = Union[GeminiExample, GeminiExampleDict] -class _UpdateAgentEngineSessionRequestParameters(_common.BaseModel): - """Parameters for updating agent engine sessions.""" +class GeminiTemplateConfig(_common.BaseModel): + """Represents a Gemini template config.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine session to update.""" + gemini_example: Optional[GeminiExample] = Field( + default=None, + description="""Required. The template that will be used for assembling the request to use for downstream applications.""", ) - config: Optional[UpdateAgentEngineSessionConfig] = Field( - default=None, description="""""" + field_mapping: Optional[dict[str, str]] = Field( + default=None, + description="""Required. Map of template parameters to the columns in the dataset table.""", ) -class _UpdateAgentEngineSessionRequestParametersDict(TypedDict, total=False): - """Parameters for updating agent engine sessions.""" +class GeminiTemplateConfigDict(TypedDict, total=False): + """Represents a Gemini template config.""" - name: Optional[str] - """Name of the agent engine session to update.""" + gemini_example: Optional[GeminiExampleDict] + """Required. The template that will be used for assembling the request to use for downstream applications.""" - config: Optional[UpdateAgentEngineSessionConfigDict] - """""" + field_mapping: Optional[dict[str, str]] + """Required. Map of template parameters to the columns in the dataset table.""" -_UpdateAgentEngineSessionRequestParametersOrDict = Union[ - _UpdateAgentEngineSessionRequestParameters, - _UpdateAgentEngineSessionRequestParametersDict, -] +GeminiTemplateConfigOrDict = Union[GeminiTemplateConfig, GeminiTemplateConfigDict] -class EventActions(_common.BaseModel): - """Actions are parts of events that are executed by the agent.""" +class GeminiRequestReadConfig(_common.BaseModel): + """Represents the config for reading Gemini requests.""" - 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.""", + template_config: Optional[GeminiTemplateConfig] = Field( + default=None, description="""Gemini request template with placeholders.""" ) - transfer_agent: Optional[str] = Field( + assembled_request_column_name: Optional[str] = Field( default=None, - description="""Optional. If set, the event transfers to the specified agent.""", + 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 EventActionsDict(TypedDict, total=False): - """Actions are parts of events that are executed by the agent.""" + 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." + ) - artifact_delta: Optional[dict[str, int]] - """Optional. Indicates that the event is updating an artifact. key is the filename, value is the version.""" + 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. - escalate: Optional[bool] - """Optional. The agent is escalating to a higher level agent.""" + 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), + ], + ) + ) - 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.""" + system_instruction_content = None + if system_instruction: + system_instruction_content = genai_types.Content( + parts=[ + genai_types.Part.from_text(text=system_instruction), + ], + ) - skip_summarization: Optional[bool] - """Optional. If true, it won't call model to summarize function response. Only used for function_response 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, + ), + ) - 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.""" +class GeminiRequestReadConfigDict(TypedDict, total=False): + """Represents the config for reading Gemini requests.""" + template_config: Optional[GeminiTemplateConfigDict] + """Gemini request template with placeholders.""" -EventActionsOrDict = Union[EventActions, EventActionsDict] + assembled_request_column_name: Optional[str] + """Column name in the underlying BigQuery table that contains already fully assembled Gemini requests.""" -class EventMetadata(_common.BaseModel): - """Metadata relating to a LLM response event.""" +GeminiRequestReadConfigOrDict = Union[ + GeminiRequestReadConfig, GeminiRequestReadConfigDict +] - grounding_metadata: Optional[genai_types.GroundingMetadata] = Field( - default=None, - description="""Optional. Metadata returned to client when grounding is enabled.""", - ) - 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.""", - ) - custom_metadata: Optional[dict[str, Any]] = Field( - default=None, description="""The custom metadata of the LlmResponse.""" - ) - 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.""", - ) - 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.""", - ) - 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 AssembleDatasetConfig(_common.BaseModel): + """Config for assembling a multimodal dataset resource.""" -class EventMetadataDict(TypedDict, total=False): - """Metadata relating to a LLM response event.""" + 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.""", + ) - grounding_metadata: Optional[genai_types.GroundingMetadata] - """Optional. Metadata returned to client when grounding is enabled.""" - 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.""" +class AssembleDatasetConfigDict(TypedDict, total=False): + """Config for assembling a multimodal dataset resource.""" - custom_metadata: Optional[dict[str, Any]] - """The custom metadata of the LlmResponse.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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.""" + timeout: Optional[int] + """The timeout for the assemble dataset request in seconds. If not + set, the default timeout is 90 seconds.""" - 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.""" +AssembleDatasetConfigOrDict = Union[AssembleDatasetConfig, AssembleDatasetConfigDict] - turn_complete: Optional[bool] - """Optional. Indicates whether the response from the model is complete. Only used for streaming mode.""" - input_transcription: Optional[genai_types.Transcription] - """Optional. Audio transcription of user input.""" +class _AssembleDatasetParameters(_common.BaseModel): + """Parameters for assembling a multimodal dataset resource.""" - output_transcription: Optional[genai_types.Transcription] - """Optional. Audio transcription of model output.""" + name: Optional[str] = Field(default=None, description="""""") + gemini_request_read_config: Optional[GeminiRequestReadConfig] = Field( + default=None, description="""""" + ) + config: Optional[AssembleDatasetConfig] = Field(default=None, description="""""") -EventMetadataOrDict = Union[EventMetadata, EventMetadataDict] +class _AssembleDatasetParametersDict(TypedDict, total=False): + """Parameters for assembling a multimodal dataset resource.""" + + name: Optional[str] + """""" + gemini_request_read_config: Optional[GeminiRequestReadConfigDict] + """""" -class AppendAgentEngineSessionEventConfig(_common.BaseModel): - """Config for appending agent engine session event.""" + config: Optional[AssembleDatasetConfigDict] + """""" - 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( + +_AssembleDatasetParametersOrDict = Union[ + _AssembleDatasetParameters, _AssembleDatasetParametersDict +] + + +class MultimodalDatasetOperation(_common.BaseModel): + """Represents the create dataset operation.""" + + name: Optional[str] = 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.""" + 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}`.""", ) - error_message: Optional[str] = Field( - default=None, description="""The error message of the session event.""" + 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.""", ) - event_metadata: Optional[EventMetadata] = Field( - default=None, description="""Metadata relating to the 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.""", ) - raw_event: Optional[dict[str, Any]] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Weakly typed raw event data in proto struct format.""", + 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 AppendAgentEngineSessionEventConfigDict(TypedDict, total=False): - """Config for appending agent engine session event.""" +class MultimodalDatasetOperationDict(TypedDict, total=False): + """Represents the create dataset operation.""" - 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}`.""" - content: Optional[genai_types.Content] - """The content of the session event.""" + 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.""" - actions: Optional[EventActionsDict] - """Actions are parts of events that are related to the session event.""" + 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] - """The error code of the session event.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - error_message: Optional[str] - """The error message of the session event.""" + response: Optional[dict[str, Any]] + """The result of the dataset operation.""" - event_metadata: Optional[EventMetadataDict] - """Metadata relating to the session event.""" - raw_event: Optional[dict[str, Any]] - """Weakly typed raw event data in proto struct format.""" +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] + """""" -AppendAgentEngineSessionEventConfigOrDict = Union[ - AppendAgentEngineSessionEventConfig, AppendAgentEngineSessionEventConfigDict +TuningResourceUsageAssessmentConfigOrDict = Union[ + TuningResourceUsageAssessmentConfig, TuningResourceUsageAssessmentConfigDict ] -class _AppendAgentEngineSessionEventRequestParameters(_common.BaseModel): - """Parameters for appending agent engines.""" +class TuningValidationAssessmentConfig(_common.BaseModel): + """Config for tuning validation assessment.""" - 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.""" - ) - config: Optional[AppendAgentEngineSessionEventConfig] = Field( - default=None, description="""""" - ) + model_name: Optional[str] = Field(default=None, description="""""") + dataset_usage: Optional[str] = Field(default=None, description="""""") -class _AppendAgentEngineSessionEventRequestParametersDict(TypedDict, total=False): - """Parameters for appending agent engines.""" +class TuningValidationAssessmentConfigDict(TypedDict, total=False): + """Config for tuning validation assessment.""" - name: Optional[str] - """Name of the agent engine session.""" + model_name: Optional[str] + """""" - author: Optional[str] - """Author of the agent engine session event.""" + dataset_usage: Optional[str] + """""" - invocation_id: Optional[str] - """Invocation ID of the agent engine.""" - timestamp: Optional[datetime.datetime] - """Timestamp indicating when the event was created.""" +TuningValidationAssessmentConfigOrDict = Union[ + TuningValidationAssessmentConfig, TuningValidationAssessmentConfigDict +] + + +class BatchPredictionResourceUsageAssessmentConfig(_common.BaseModel): + """Config for batch prediction resource usage assessment.""" - config: Optional[AppendAgentEngineSessionEventConfigDict] + model_name: Optional[str] = Field(default=None, description="""""") + + +class BatchPredictionResourceUsageAssessmentConfigDict(TypedDict, total=False): + """Config for batch prediction resource usage assessment.""" + + model_name: Optional[str] """""" -_AppendAgentEngineSessionEventRequestParametersOrDict = Union[ - _AppendAgentEngineSessionEventRequestParameters, - _AppendAgentEngineSessionEventRequestParametersDict, +BatchPredictionResourceUsageAssessmentConfigOrDict = Union[ + BatchPredictionResourceUsageAssessmentConfig, + BatchPredictionResourceUsageAssessmentConfigDict, ] -class AppendAgentEngineSessionEventResponse(_common.BaseModel): - """Response for appending agent engine session event.""" +class BatchPredictionValidationAssessmentConfig(_common.BaseModel): + """Config for batch prediction validation assessment.""" - pass + model_name: Optional[str] = Field(default=None, description="""""") -class AppendAgentEngineSessionEventResponseDict(TypedDict, total=False): - """Response for appending agent engine session event.""" +class BatchPredictionValidationAssessmentConfigDict(TypedDict, total=False): + """Config for batch prediction validation assessment.""" - pass + model_name: Optional[str] + """""" -AppendAgentEngineSessionEventResponseOrDict = Union[ - AppendAgentEngineSessionEventResponse, AppendAgentEngineSessionEventResponseDict +BatchPredictionValidationAssessmentConfigOrDict = Union[ + BatchPredictionValidationAssessmentConfig, + BatchPredictionValidationAssessmentConfigDict, ] -class ListAgentEngineSessionEventsConfig(_common.BaseModel): - """Config for listing agent engine session events.""" +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.""" ) - 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 assess 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 AssessDatasetConfigDict(TypedDict, total=False): + """Config for assessing 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 assess dataset request in seconds. If not set, + the default timeout is 90 seconds.""" -ListAgentEngineSessionEventsConfigOrDict = Union[ - ListAgentEngineSessionEventsConfig, ListAgentEngineSessionEventsConfigDict -] +AssessDatasetConfigOrDict = Union[AssessDatasetConfig, AssessDatasetConfigDict] -class _ListAgentEngineSessionEventsRequestParameters(_common.BaseModel): - """Parameters for listing agent engine session events.""" +class _AssessDatasetParameters(_common.BaseModel): + """Parameters for assessing 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="""""" ) + 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 _ListAgentEngineSessionEventsRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engine session events.""" +class _AssessDatasetParametersDict(TypedDict, total=False): + """Parameters for assessing a multimodal dataset resource.""" name: Optional[str] - """Name of the agent engine session.""" - - config: Optional[ListAgentEngineSessionEventsConfigDict] """""" + gemini_request_read_config: Optional[GeminiRequestReadConfigDict] + """""" -_ListAgentEngineSessionEventsRequestParametersOrDict = Union[ - _ListAgentEngineSessionEventsRequestParameters, - _ListAgentEngineSessionEventsRequestParametersDict, -] + tuning_resource_usage_assessment_config: Optional[ + TuningResourceUsageAssessmentConfigDict + ] + """""" + tuning_validation_assessment_config: Optional[TuningValidationAssessmentConfigDict] + """""" -class SessionEvent(_common.BaseModel): - """A session event.""" + batch_prediction_resource_usage_assessment_config: Optional[ + BatchPredictionResourceUsageAssessmentConfigDict + ] + """""" - 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( - 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.""" - ) - invocation_id: Optional[str] = Field( - default=None, - description="""Required. The invocation id of the event, multiple events can have the same invocation id.""", - ) - name: Optional[str] = Field( - default=None, - description="""Identifier. The resource name of the event. Format:`projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}/events/{event}`.""", - ) - timestamp: Optional[datetime.datetime] = Field( - default=None, - description="""Required. Timestamp when the event was created on client side.""", - ) - raw_event: Optional[dict[str, Any]] = Field( + 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="""Optional. Weakly typed raw event data in proto struct format.""", + description="""The URI of the BigQuery table. This accepts the table name with or without the bq:// prefix.""", ) -class SessionEventDict(TypedDict, total=False): - """A session event.""" - - content: Optional[genai_types.Content] - """Optional. Content of the event provided by the author.""" +class SchemaTablesDatasetMetadataBigQuerySourceDict(TypedDict, total=False): + """Represents the BigQuery source for multimodal dataset metadata.""" - actions: Optional[EventActionsDict] - """Optional. Actions executed by the agent.""" + uri: Optional[str] + """The URI of the BigQuery table. This accepts the table name with or without the bq:// prefix.""" - author: Optional[str] - """Required. The name of the agent that sent the event, or user.""" - error_code: Optional[str] - """Optional. Error code if the response is an error. Code varies by model.""" +SchemaTablesDatasetMetadataBigQuerySourceOrDict = Union[ + SchemaTablesDatasetMetadataBigQuerySource, + SchemaTablesDatasetMetadataBigQuerySourceDict, +] - error_message: Optional[str] - """Optional. Error message if the response is an error.""" - event_metadata: Optional[EventMetadataDict] - """Optional. Metadata relating to this event.""" +class SchemaTablesDatasetMetadataInputConfig(_common.BaseModel): + """Represents the input config for multimodal dataset metadata.""" - invocation_id: Optional[str] - """Required. The invocation id of the event, multiple events can have the same invocation id.""" + bigquery_source: Optional[SchemaTablesDatasetMetadataBigQuerySource] = Field( + default=None, + description="""The BigQuery source for multimodal dataset metadata.""", + ) - 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 SchemaTablesDatasetMetadataInputConfigDict(TypedDict, total=False): + """Represents the input config for multimodal dataset metadata.""" - raw_event: Optional[dict[str, Any]] - """Optional. Weakly typed raw event data in proto struct format.""" + bigquery_source: Optional[SchemaTablesDatasetMetadataBigQuerySourceDict] + """The BigQuery source for multimodal dataset metadata.""" -SessionEventOrDict = Union[SessionEvent, SessionEventDict] +SchemaTablesDatasetMetadataInputConfigOrDict = Union[ + SchemaTablesDatasetMetadataInputConfig, SchemaTablesDatasetMetadataInputConfigDict +] -class ListAgentEngineSessionEventsResponse(_common.BaseModel): - """Response for listing agent engine session events.""" +class SchemaTablesDatasetMetadata(_common.BaseModel): + """Represents the metadata schema for multimodal dataset metadata.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + input_config: Optional[SchemaTablesDatasetMetadataInputConfig] = Field( + default=None, + description="""The input config for multimodal dataset metadata.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - session_events: Optional[list[SessionEvent]] = Field( - default=None, description="""List of session events.""" + gemini_request_read_config: Optional[GeminiRequestReadConfig] = Field( + default=None, + description="""The Gemini request read config for the multimodal dataset.""", ) -class ListAgentEngineSessionEventsResponseDict(TypedDict, total=False): - """Response for listing agent engine session events.""" - - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" +class SchemaTablesDatasetMetadataDict(TypedDict, total=False): + """Represents the metadata schema for multimodal dataset metadata.""" - next_page_token: Optional[str] - """""" + input_config: Optional[SchemaTablesDatasetMetadataInputConfigDict] + """The input config for multimodal dataset metadata.""" - session_events: Optional[list[SessionEventDict]] - """List of session events.""" + gemini_request_read_config: Optional[GeminiRequestReadConfigDict] + """The Gemini request read config for the multimodal dataset.""" -ListAgentEngineSessionEventsResponseOrDict = Union[ - ListAgentEngineSessionEventsResponse, ListAgentEngineSessionEventsResponseDict +SchemaTablesDatasetMetadataOrDict = Union[ + SchemaTablesDatasetMetadata, SchemaTablesDatasetMetadataDict ] -class GeminiExample(_common.BaseModel): - """Represents a Gemini example.""" +class CreateMultimodalDatasetConfig(_common.BaseModel): + """Config for creating a dataset resource to store multimodal dataset.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - 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.""", + timeout: Optional[int] = Field( + default=90, + description="""The timeout for the create dataset request in seconds. If not set, + the default timeout is 90 seconds.""", ) -class GeminiExampleDict(TypedDict, total=False): - """Represents a Gemini example.""" +class CreateMultimodalDatasetConfigDict(TypedDict, total=False): + """Config for creating a dataset resource to store multimodal dataset.""" - model: Optional[str] - """The model used to generate the Gemini example.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - contents: Optional[list[genai_types.Content]] - """Contents of the Gemini example.""" + timeout: Optional[int] + """The timeout for the create dataset request in seconds. If not set, + the default timeout is 90 seconds.""" - system_instruction: Optional[genai_types.Content] - """System instruction for the Gemini example.""" - cached_content: Optional[str] - """Cached content for the Gemini example.""" +CreateMultimodalDatasetConfigOrDict = Union[ + CreateMultimodalDatasetConfig, CreateMultimodalDatasetConfigDict +] - tools: Optional[list[genai_types.Tool]] - """Tools for the Gemini example.""" - tool_config: Optional[genai_types.ToolConfig] - """Tools for the Gemini example.""" +class _CreateMultimodalDatasetParameters(_common.BaseModel): + """Parameters for creating a dataset resource to store multimodal dataset.""" - safety_settings: Optional[list[genai_types.SafetySetting]] - """Safety settings for the Gemini example.""" + name: Optional[str] = Field(default=None, description="""""") + display_name: Optional[str] = Field(default=None, description="""""") + metadata_schema_uri: Optional[str] = Field(default=None, description="""""") + metadata: Optional[SchemaTablesDatasetMetadata] = Field( + default=None, description="""""" + ) + description: Optional[str] = Field(default=None, description="""""") + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, description="""""" + ) + config: Optional[CreateMultimodalDatasetConfig] = Field( + default=None, description="""""" + ) - 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.""" +class _CreateMultimodalDatasetParametersDict(TypedDict, total=False): + """Parameters for creating a dataset resource to store multimodal dataset.""" + name: Optional[str] + """""" -GeminiExampleOrDict = Union[GeminiExample, GeminiExampleDict] + display_name: Optional[str] + """""" + metadata_schema_uri: Optional[str] + """""" -class GeminiTemplateConfig(_common.BaseModel): - """Represents a Gemini template config.""" + metadata: Optional[SchemaTablesDatasetMetadataDict] + """""" - 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.""", - ) + description: Optional[str] + """""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """""" -class GeminiTemplateConfigDict(TypedDict, total=False): - """Represents a Gemini template config.""" + config: Optional[CreateMultimodalDatasetConfigDict] + """""" - 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.""" +_CreateMultimodalDatasetParametersOrDict = Union[ + _CreateMultimodalDatasetParameters, _CreateMultimodalDatasetParametersDict +] -GeminiTemplateConfigOrDict = Union[GeminiTemplateConfig, GeminiTemplateConfigDict] +class _DeleteMultimodalDatasetRequestParameters(_common.BaseModel): + """Parameters for deleting a multimodal dataset.""" + name: Optional[str] = Field( + default=None, description="""ID of the dataset to be deleted.""" + ) + config: Optional[VertexBaseConfig] = Field(default=None, description="""""") -class GeminiRequestReadConfig(_common.BaseModel): - """Represents the config for reading Gemini requests.""" - 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.""", - ) +class _DeleteMultimodalDatasetRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a multimodal dataset.""" - @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. + name: Optional[str] + """ID of the dataset to be deleted.""" - 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." - ) + config: Optional[VertexBaseConfigDict] + """""" - 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), - ], - ) - ) +_DeleteMultimodalDatasetRequestParametersOrDict = Union[ + _DeleteMultimodalDatasetRequestParameters, + _DeleteMultimodalDatasetRequestParametersDict, +] - 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 _GetMultimodalDatasetParameters(_common.BaseModel): + """Parameters for getting a multimodal dataset resource.""" + name: Optional[str] = Field(default=None, description="""""") + config: Optional[VertexBaseConfig] = 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 _GetMultimodalDatasetParametersDict(TypedDict, total=False): + """Parameters for getting a multimodal dataset resource.""" - assembled_request_column_name: Optional[str] - """Column name in the underlying BigQuery table that contains already fully assembled Gemini requests.""" + name: Optional[str] + """""" + config: Optional[VertexBaseConfigDict] + """""" -GeminiRequestReadConfigOrDict = Union[ - GeminiRequestReadConfig, GeminiRequestReadConfigDict + +_GetMultimodalDatasetParametersOrDict = Union[ + _GetMultimodalDatasetParameters, _GetMultimodalDatasetParametersDict ] -class AssembleDatasetConfig(_common.BaseModel): - """Config for assembling a multimodal dataset resource.""" +class MultimodalDataset(_common.BaseModel): + """Represents a multimodal dataset.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + name: Optional[str] = Field( + default=None, description="""The ID of the multimodal dataset.""" ) - 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.""", + display_name: Optional[str] = Field( + default=None, description="""The display name of the multimodal dataset.""" + ) + metadata: Optional[SchemaTablesDatasetMetadata] = Field( + default=None, description="""The metadata of the multimodal dataset.""" + ) + description: Optional[str] = Field( + default=None, description="""The description of the multimodal dataset.""" ) + @property + def read_config(self) -> Optional[GeminiRequestReadConfig]: + """Gets the read config from the dataset metadata. Returns None if it's not set.""" + if self.metadata is None or self.metadata.gemini_request_read_config is None: + return None + return self.metadata.gemini_request_read_config -class AssembleDatasetConfigDict(TypedDict, total=False): - """Config for assembling a multimodal dataset resource.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + def set_read_config( + self, + *, + read_config: GeminiRequestReadConfigOrDict, + ) -> None: + """Sets the read config in the dataset metadata.""" + if isinstance(read_config, dict): + read_config = GeminiRequestReadConfig(**read_config) - timeout: Optional[int] - """The timeout for the assemble dataset request in seconds. If not - set, the default timeout is 90 seconds.""" + if self.metadata is None: + self.metadata = SchemaTablesDatasetMetadata() + self.metadata.gemini_request_read_config = read_config + @property + def bigquery_uri( + self, + ) -> Optional[str]: + """Gets the bigquery uri from the dataset metadata. Returns None if it's not set.""" + if ( + self.metadata is None + or self.metadata.input_config is None + or self.metadata.input_config.bigquery_source is None + ): + return None + return str(self.metadata.input_config.bigquery_source.uri) -AssembleDatasetConfigOrDict = Union[AssembleDatasetConfig, AssembleDatasetConfigDict] + def set_bigquery_uri( + self, + bigquery_uri: str, + ) -> None: + """Sets the bigquery uri in the dataset metadata. Prepends 'bq://' if it's not already present.""" + if not bigquery_uri.startswith("bq://"): + bigquery_uri = f"bq://{bigquery_uri}" + metadata = ( + SchemaTablesDatasetMetadata() if self.metadata is None else self.metadata + ) + input_config = ( + SchemaTablesDatasetMetadataInputConfig() + if metadata.input_config is None + else metadata.input_config + ) + bigquery_source = ( + SchemaTablesDatasetMetadataBigQuerySource() + if input_config.bigquery_source is None + else input_config.bigquery_source + ) + bigquery_source.uri = bigquery_uri + input_config.bigquery_source = bigquery_source + metadata.input_config = input_config + self.metadata = metadata + def to_bigframes( + self, + ) -> "bigframes.pandas.DataFrame": # type: ignore # noqa: F821 + """Converts the multimodal dataset to a BigFrames dataframe. -class _AssembleDatasetParameters(_common.BaseModel): - """Parameters for assembling a multimodal dataset resource.""" + This is the preferred method to inspect the multimodal dataset in a + notebook. - name: Optional[str] = Field(default=None, description="""""") - gemini_request_read_config: Optional[GeminiRequestReadConfig] = Field( - default=None, description="""""" - ) - config: Optional[AssembleDatasetConfig] = Field(default=None, description="""""") + Returns: + A BigFrames dataframe. + """ + from .. import _datasets_utils + bigframes = _datasets_utils._try_import_bigframes() -class _AssembleDatasetParametersDict(TypedDict, total=False): - """Parameters for assembling a multimodal dataset resource.""" + if self.bigquery_uri is None: + raise ValueError("Multimodal dataset bigquery source uri is not set.") + return bigframes.pandas.read_gbq_table(self.bigquery_uri.removeprefix("bq://")) - name: Optional[str] - """""" + def to_batch_job_source(self) -> "genai_types.BatchJobSource": + """Converts the dataset to a BatchJobSource.""" + return genai_types.BatchJobSource( + vertex_dataset_name=self.name, + ) - gemini_request_read_config: Optional[GeminiRequestReadConfigDict] - """""" + def get_batch_job_destination(self) -> "genai_types.BatchJobDestination": + """Converts the dataset to a BatchJobDestination.""" + from .. import _datasets_utils - config: Optional[AssembleDatasetConfigDict] - """""" + unique_name = _datasets_utils.get_batch_job_unique_name() + bigquery_uri = self.bigquery_uri + if bigquery_uri is None: + raise ValueError("Multimodal dataset bigquery source uri is not set.") + curr_display_name = self.display_name or "genai_batch_job" + return genai_types.BatchJobDestination( + vertex_dataset=genai_types.VertexMultimodalDatasetDestination( + display_name=f"{curr_display_name}_batch_output_{unique_name}", + bigquery_destination=f"{bigquery_uri}_batch_output_{unique_name}", + ) + ) -_AssembleDatasetParametersOrDict = Union[ - _AssembleDatasetParameters, _AssembleDatasetParametersDict -] +class MultimodalDatasetDict(TypedDict, total=False): + """Represents a multimodal dataset.""" + name: Optional[str] + """The ID of the multimodal dataset.""" -class MultimodalDatasetOperation(_common.BaseModel): - """Represents the create dataset operation.""" + display_name: Optional[str] + """The display name of the multimodal dataset.""" - 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[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.""" + metadata: Optional[SchemaTablesDatasetMetadataDict] + """The metadata of the multimodal dataset.""" - response: Optional[dict[str, Any]] - """The result of the dataset operation.""" + description: Optional[str] + """The description of the multimodal dataset.""" -MultimodalDatasetOperationOrDict = Union[ - MultimodalDatasetOperation, MultimodalDatasetOperationDict -] +MultimodalDatasetOrDict = Union[MultimodalDataset, MultimodalDatasetDict] -class TuningResourceUsageAssessmentConfig(_common.BaseModel): - """Config for tuning resource usage assessment.""" +class GetMultimodalDatasetOperationConfig(_common.BaseModel): + """Config for getting a multimodal dataset operation.""" - model_name: Optional[str] = Field(default=None, description="""""") + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -class TuningResourceUsageAssessmentConfigDict(TypedDict, total=False): - """Config for tuning resource usage assessment.""" +class GetMultimodalDatasetOperationConfigDict(TypedDict, total=False): + """Config for getting a multimodal dataset operation.""" - model_name: Optional[str] - """""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -TuningResourceUsageAssessmentConfigOrDict = Union[ - TuningResourceUsageAssessmentConfig, TuningResourceUsageAssessmentConfigDict +GetMultimodalDatasetOperationConfigOrDict = Union[ + GetMultimodalDatasetOperationConfig, GetMultimodalDatasetOperationConfigDict ] -class TuningValidationAssessmentConfig(_common.BaseModel): - """Config for tuning validation assessment.""" +class _GetMultimodalDatasetOperationParameters(_common.BaseModel): + """Parameters for getting a dataset operation.""" - model_name: Optional[str] = Field(default=None, description="""""") - dataset_usage: Optional[str] = Field(default=None, description="""""") + dataset_id: Optional[str] = Field(default=None, description="""""") + operation_id: Optional[str] = Field(default=None, description="""""") + config: Optional[GetMultimodalDatasetOperationConfig] = Field( + default=None, description="""""" + ) -class TuningValidationAssessmentConfigDict(TypedDict, total=False): - """Config for tuning validation assessment.""" +class _GetMultimodalDatasetOperationParametersDict(TypedDict, total=False): + """Parameters for getting a dataset operation.""" - model_name: Optional[str] + dataset_id: Optional[str] """""" - dataset_usage: Optional[str] + operation_id: Optional[str] """""" + config: Optional[GetMultimodalDatasetOperationConfigDict] + """""" -TuningValidationAssessmentConfigOrDict = Union[ - TuningValidationAssessmentConfig, TuningValidationAssessmentConfigDict + +_GetMultimodalDatasetOperationParametersOrDict = Union[ + _GetMultimodalDatasetOperationParameters, + _GetMultimodalDatasetOperationParametersDict, ] -class BatchPredictionResourceUsageAssessmentConfig(_common.BaseModel): - """Config for batch prediction resource usage assessment.""" +class ListMultimodalDatasetsConfig(_common.BaseModel): + """Config for listing multimodal datasets.""" - model_name: Optional[str] = Field(default=None, description="""""") + 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 BatchPredictionResourceUsageAssessmentConfigDict(TypedDict, total=False): - """Config for batch prediction resource usage assessment.""" +class ListMultimodalDatasetsConfigDict(TypedDict, total=False): + """Config for listing multimodal datasets.""" - model_name: Optional[str] + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + page_size: Optional[int] """""" + page_token: Optional[str] + """""" -BatchPredictionResourceUsageAssessmentConfigOrDict = Union[ - BatchPredictionResourceUsageAssessmentConfig, - BatchPredictionResourceUsageAssessmentConfigDict, + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""" + + +ListMultimodalDatasetsConfigOrDict = Union[ + ListMultimodalDatasetsConfig, ListMultimodalDatasetsConfigDict ] -class BatchPredictionValidationAssessmentConfig(_common.BaseModel): - """Config for batch prediction validation assessment.""" +class _ListMultimodalDatasetsRequestParameters(_common.BaseModel): + """Parameters for listing multimodal datasets.""" - model_name: Optional[str] = Field(default=None, description="""""") + config: Optional[ListMultimodalDatasetsConfig] = Field( + default=None, description="""""" + ) -class BatchPredictionValidationAssessmentConfigDict(TypedDict, total=False): - """Config for batch prediction validation assessment.""" +class _ListMultimodalDatasetsRequestParametersDict(TypedDict, total=False): + """Parameters for listing multimodal datasets.""" - model_name: Optional[str] + config: Optional[ListMultimodalDatasetsConfigDict] """""" -BatchPredictionValidationAssessmentConfigOrDict = Union[ - BatchPredictionValidationAssessmentConfig, - BatchPredictionValidationAssessmentConfigDict, +_ListMultimodalDatasetsRequestParametersOrDict = Union[ + _ListMultimodalDatasetsRequestParameters, + _ListMultimodalDatasetsRequestParametersDict, ] -class AssessDatasetConfig(_common.BaseModel): - """Config for assessing a multimodal dataset resource.""" +class ListMultimodalDatasetsResponse(_common.BaseModel): + """Response for listing multimodal datasets.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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="""""") timeout: Optional[int] = Field( default=90, - description="""The timeout for the assess dataset request in seconds. If not set, + description="""The timeout for the list datasets request in seconds. If not set, the default timeout is 90 seconds.""", ) + datasets: Optional[list[MultimodalDataset]] = Field( + default=None, + description="""List of datasets for the project. + """, + ) -class AssessDatasetConfigDict(TypedDict, total=False): - """Config for assessing a multimodal dataset resource.""" +class ListMultimodalDatasetsResponseDict(TypedDict, total=False): + """Response for listing multimodal datasets.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + + next_page_token: Optional[str] + """""" timeout: Optional[int] - """The timeout for the assess dataset request in seconds. If not set, + """The timeout for the list datasets request in seconds. If not set, the default timeout is 90 seconds.""" + datasets: Optional[list[MultimodalDatasetDict]] + """List of datasets for the project. + """ + -AssessDatasetConfigOrDict = Union[AssessDatasetConfig, AssessDatasetConfigDict] +ListMultimodalDatasetsResponseOrDict = Union[ + ListMultimodalDatasetsResponse, ListMultimodalDatasetsResponseDict +] -class _AssessDatasetParameters(_common.BaseModel): - """Parameters for assessing a multimodal dataset resource.""" +class _UpdateMultimodalDatasetParameters(_common.BaseModel): + """Parameters for updating a multimodal dataset resource.""" name: Optional[str] = Field(default=None, description="""""") - gemini_request_read_config: Optional[GeminiRequestReadConfig] = Field( + display_name: Optional[str] = Field(default=None, description="""""") + metadata: Optional[SchemaTablesDatasetMetadata] = 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="""""") + description: Optional[str] = Field(default=None, description="""""") + encryption_spec: Optional[genai_types.EncryptionSpec] = 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="""""") + config: Optional[VertexBaseConfig] = Field(default=None, description="""""") -class _AssessDatasetParametersDict(TypedDict, total=False): - """Parameters for assessing a multimodal dataset resource.""" +class _UpdateMultimodalDatasetParametersDict(TypedDict, total=False): + """Parameters for updating a multimodal dataset resource.""" name: Optional[str] """""" - gemini_request_read_config: Optional[GeminiRequestReadConfigDict] - """""" - - tuning_resource_usage_assessment_config: Optional[ - TuningResourceUsageAssessmentConfigDict - ] + display_name: Optional[str] """""" - tuning_validation_assessment_config: Optional[TuningValidationAssessmentConfigDict] + metadata: Optional[SchemaTablesDatasetMetadataDict] """""" - batch_prediction_resource_usage_assessment_config: Optional[ - BatchPredictionResourceUsageAssessmentConfigDict - ] + description: Optional[str] """""" - batch_prediction_validation_assessment_config: Optional[ - BatchPredictionValidationAssessmentConfigDict - ] + encryption_spec: Optional[genai_types.EncryptionSpec] """""" - config: Optional[AssessDatasetConfigDict] + config: Optional[VertexBaseConfigDict] """""" -_AssessDatasetParametersOrDict = Union[ - _AssessDatasetParameters, _AssessDatasetParametersDict +_UpdateMultimodalDatasetParametersOrDict = Union[ + _UpdateMultimodalDatasetParameters, _UpdateMultimodalDatasetParametersDict ] -class SchemaTablesDatasetMetadataBigQuerySource(_common.BaseModel): - """Represents the BigQuery source for multimodal dataset metadata.""" +class SchemaPredictParamsGroundingConfigSourceEntry(_common.BaseModel): + """Single source entry for the grounding checking.""" - uri: Optional[str] = Field( + enterprise_datastore: 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 Vertex AI Search data source. Deprecated. Use vertex_ai_search_datastore instead.""", + ) + inline_context: Optional[str] = Field( + default=None, + description="""The grounding text passed inline with the Predict API. It can support up to 1 million bytes.""", + ) + type: Optional[ + Literal["UNSPECIFIED", "WEB", "ENTERPRISE", "VERTEX_AI_SEARCH", "INLINE"] + ] = Field( + default=None, description="""The type of the grounding checking source.""" + ) + vertex_ai_search_datastore: Optional[str] = Field( + default=None, description="""The uri of the Vertex AI Search data source.""" ) -class SchemaTablesDatasetMetadataBigQuerySourceDict(TypedDict, total=False): - """Represents the BigQuery source for multimodal dataset metadata.""" +class SchemaPredictParamsGroundingConfigSourceEntryDict(TypedDict, total=False): + """Single source entry for the grounding checking.""" - uri: Optional[str] - """The URI of the BigQuery table. This accepts the table name with or without the bq:// prefix.""" + enterprise_datastore: Optional[str] + """The uri of the Vertex AI Search data source. Deprecated. Use vertex_ai_search_datastore instead.""" + inline_context: Optional[str] + """The grounding text passed inline with the Predict API. It can support up to 1 million bytes.""" -SchemaTablesDatasetMetadataBigQuerySourceOrDict = Union[ - SchemaTablesDatasetMetadataBigQuerySource, - SchemaTablesDatasetMetadataBigQuerySourceDict, + type: Optional[ + Literal["UNSPECIFIED", "WEB", "ENTERPRISE", "VERTEX_AI_SEARCH", "INLINE"] + ] + """The type of the grounding checking source.""" + + vertex_ai_search_datastore: Optional[str] + """The uri of the Vertex AI Search data source.""" + + +SchemaPredictParamsGroundingConfigSourceEntryOrDict = Union[ + SchemaPredictParamsGroundingConfigSourceEntry, + SchemaPredictParamsGroundingConfigSourceEntryDict, ] -class SchemaTablesDatasetMetadataInputConfig(_common.BaseModel): - """Represents the input config for multimodal dataset metadata.""" +class SchemaPredictParamsGroundingConfig(_common.BaseModel): + """The configuration for grounding checking.""" - bigquery_source: Optional[SchemaTablesDatasetMetadataBigQuerySource] = Field( + disable_attribution: Optional[bool] = Field( default=None, - description="""The BigQuery source for multimodal dataset metadata.""", + description="""If set, skip finding claim attributions (i.e not generate grounding citation).""", + ) + sources: Optional[list[SchemaPredictParamsGroundingConfigSourceEntry]] = Field( + default=None, description="""The sources for the grounding checking.""" ) -class SchemaTablesDatasetMetadataInputConfigDict(TypedDict, total=False): - """Represents the input config for multimodal dataset metadata.""" +class SchemaPredictParamsGroundingConfigDict(TypedDict, total=False): + """The configuration for grounding checking.""" - bigquery_source: Optional[SchemaTablesDatasetMetadataBigQuerySourceDict] - """The BigQuery source for multimodal dataset metadata.""" + disable_attribution: Optional[bool] + """If set, skip finding claim attributions (i.e not generate grounding citation).""" + + sources: Optional[list[SchemaPredictParamsGroundingConfigSourceEntryDict]] + """The sources for the grounding checking.""" -SchemaTablesDatasetMetadataInputConfigOrDict = Union[ - SchemaTablesDatasetMetadataInputConfig, SchemaTablesDatasetMetadataInputConfigDict +SchemaPredictParamsGroundingConfigOrDict = Union[ + SchemaPredictParamsGroundingConfig, SchemaPredictParamsGroundingConfigDict ] -class SchemaTablesDatasetMetadata(_common.BaseModel): - """Represents the metadata schema for multimodal dataset metadata.""" +class SchemaPromptSpecPartList(_common.BaseModel): + """Represents a prompt spec part list.""" - input_config: Optional[SchemaTablesDatasetMetadataInputConfig] = Field( - default=None, - description="""The input config for multimodal dataset metadata.""", - ) - gemini_request_read_config: Optional[GeminiRequestReadConfig] = Field( - default=None, - description="""The Gemini request read config for the multimodal dataset.""", + parts: Optional[list[genai_types.Part]] = Field( + default=None, description="""A list of elements that can be part of a prompt.""" ) -class SchemaTablesDatasetMetadataDict(TypedDict, total=False): - """Represents the metadata schema for multimodal dataset metadata.""" - - input_config: Optional[SchemaTablesDatasetMetadataInputConfigDict] - """The input config for multimodal dataset metadata.""" +class SchemaPromptSpecPartListDict(TypedDict, total=False): + """Represents a prompt spec part list.""" - gemini_request_read_config: Optional[GeminiRequestReadConfigDict] - """The Gemini request read config for the multimodal dataset.""" + parts: Optional[list[genai_types.Part]] + """A list of elements that can be part of a prompt.""" -SchemaTablesDatasetMetadataOrDict = Union[ - SchemaTablesDatasetMetadata, SchemaTablesDatasetMetadataDict +SchemaPromptSpecPartListOrDict = Union[ + SchemaPromptSpecPartList, SchemaPromptSpecPartListDict ] -class CreateMultimodalDatasetConfig(_common.BaseModel): - """Config for creating a dataset resource to store multimodal dataset.""" +class SchemaPromptInstanceVariableValue(_common.BaseModel): + """Represents a prompt instance variable.""" - 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 create dataset request in seconds. If not set, - the default timeout is 90 seconds.""", + part_list: Optional[SchemaPromptSpecPartList] = Field( + default=None, description="""The parts of the variable value.""" ) -class CreateMultimodalDatasetConfigDict(TypedDict, total=False): - """Config for creating a dataset resource to store multimodal dataset.""" +class SchemaPromptInstanceVariableValueDict(TypedDict, total=False): + """Represents a prompt instance variable.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + part_list: Optional[SchemaPromptSpecPartListDict] + """The parts of the variable value.""" - timeout: Optional[int] - """The timeout for the create dataset request in seconds. If not set, - the default timeout is 90 seconds.""" + +SchemaPromptInstanceVariableValueOrDict = Union[ + SchemaPromptInstanceVariableValue, SchemaPromptInstanceVariableValueDict +] -CreateMultimodalDatasetConfigOrDict = Union[ - CreateMultimodalDatasetConfig, CreateMultimodalDatasetConfigDict +class SchemaPromptInstancePromptExecution(_common.BaseModel): + """A prompt instance's parameters set that contains a set of variable values.""" + + arguments: Optional[dict[str, SchemaPromptInstanceVariableValue]] = Field( + default=None, description="""Maps variable names to their value.""" + ) + + +class SchemaPromptInstancePromptExecutionDict(TypedDict, total=False): + """A prompt instance's parameters set that contains a set of variable values.""" + + arguments: Optional[dict[str, SchemaPromptInstanceVariableValueDict]] + """Maps variable names to their value.""" + + +SchemaPromptInstancePromptExecutionOrDict = Union[ + SchemaPromptInstancePromptExecution, SchemaPromptInstancePromptExecutionDict ] -class _CreateMultimodalDatasetParameters(_common.BaseModel): - """Parameters for creating a dataset resource to store multimodal dataset.""" +class SchemaPromptSpecPromptMessage(_common.BaseModel): + """Represents a prompt message.""" - name: Optional[str] = Field(default=None, description="""""") - display_name: Optional[str] = Field(default=None, description="""""") - metadata_schema_uri: Optional[str] = Field(default=None, description="""""") - metadata: Optional[SchemaTablesDatasetMetadata] = Field( - default=None, description="""""" + generation_config: Optional[genai_types.GenerationConfig] = Field( + default=None, description="""Generation config.""" ) - description: Optional[str] = Field(default=None, description="""""") - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, description="""""" + tool_config: Optional[genai_types.FunctionCallingConfig] = Field( + default=None, + description="""Tool config. This config is shared for all tools provided in the request.""", ) - config: Optional[CreateMultimodalDatasetConfig] = Field( + tools: Optional[list[genai_types.Tool]] = Field( + default=None, + description="""A list of `Tools` the model may use to generate the next response. A `Tool` is a piece of code that enables the system to interact with external systems to perform an action, or set of actions, outside of knowledge and scope of the model.""", + ) + safety_settings: Optional[list[genai_types.SafetySetting]] = Field( + default=None, + description="""Per request settings for blocking unsafe content. Enforced on GenerateContentResponse.candidates.""", + ) + contents: Optional[list[genai_types.Content]] = Field( + default=None, + description="""The content of the current conversation with the model. For single-turn queries, this is a single instance. For multi-turn queries, this is a repeated field that contains conversation history + latest request.""", + ) + system_instruction: Optional[genai_types.Content] = Field( + default=None, + description="""The user provided system instructions for the model. Note: only text should be used in parts and content in each part will be in a separate paragraph.""", + ) + variables: Optional[list[dict[str, genai_types.Part]]] = Field( default=None, description="""""" ) + model: Optional[str] = Field(default=None, description="""The model name.""") -class _CreateMultimodalDatasetParametersDict(TypedDict, total=False): - """Parameters for creating a dataset resource to store multimodal dataset.""" +class SchemaPromptSpecPromptMessageDict(TypedDict, total=False): + """Represents a prompt message.""" - name: Optional[str] - """""" + generation_config: Optional[genai_types.GenerationConfig] + """Generation config.""" - display_name: Optional[str] - """""" + tool_config: Optional[genai_types.FunctionCallingConfig] + """Tool config. This config is shared for all tools provided in the request.""" - metadata_schema_uri: Optional[str] - """""" + tools: Optional[list[genai_types.Tool]] + """A list of `Tools` the model may use to generate the next response. A `Tool` is a piece of code that enables the system to interact with external systems to perform an action, or set of actions, outside of knowledge and scope of the model.""" - metadata: Optional[SchemaTablesDatasetMetadataDict] - """""" + safety_settings: Optional[list[genai_types.SafetySetting]] + """Per request settings for blocking unsafe content. Enforced on GenerateContentResponse.candidates.""" - description: Optional[str] - """""" + contents: Optional[list[genai_types.Content]] + """The content of the current conversation with the model. For single-turn queries, this is a single instance. For multi-turn queries, this is a repeated field that contains conversation history + latest request.""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """""" + system_instruction: Optional[genai_types.Content] + """The user provided system instructions for the model. Note: only text should be used in parts and content in each part will be in a separate paragraph.""" - config: Optional[CreateMultimodalDatasetConfigDict] + variables: Optional[list[dict[str, genai_types.Part]]] """""" + model: Optional[str] + """The model name.""" + -_CreateMultimodalDatasetParametersOrDict = Union[ - _CreateMultimodalDatasetParameters, _CreateMultimodalDatasetParametersDict +SchemaPromptSpecPromptMessageOrDict = Union[ + SchemaPromptSpecPromptMessage, SchemaPromptSpecPromptMessageDict ] -class _DeleteMultimodalDatasetRequestParameters(_common.BaseModel): - """Parameters for deleting a multimodal dataset.""" +class SchemaPromptSpecMultimodalPrompt(_common.BaseModel): + """Prompt variation that embeds preambles to prompt string.""" - name: Optional[str] = Field( - default=None, description="""ID of the dataset to be deleted.""" + prompt_message: Optional[SchemaPromptSpecPromptMessage] = Field( + default=None, description="""The prompt message.""" ) - config: Optional[VertexBaseConfig] = Field(default=None, description="""""") - -class _DeleteMultimodalDatasetRequestParametersDict(TypedDict, total=False): - """Parameters for deleting a multimodal dataset.""" - name: Optional[str] - """ID of the dataset to be deleted.""" +class SchemaPromptSpecMultimodalPromptDict(TypedDict, total=False): + """Prompt variation that embeds preambles to prompt string.""" - config: Optional[VertexBaseConfigDict] - """""" + prompt_message: Optional[SchemaPromptSpecPromptMessageDict] + """The prompt message.""" -_DeleteMultimodalDatasetRequestParametersOrDict = Union[ - _DeleteMultimodalDatasetRequestParameters, - _DeleteMultimodalDatasetRequestParametersDict, +SchemaPromptSpecMultimodalPromptOrDict = Union[ + SchemaPromptSpecMultimodalPrompt, SchemaPromptSpecMultimodalPromptDict ] -class _GetMultimodalDatasetParameters(_common.BaseModel): - """Parameters for getting a multimodal dataset resource.""" +class SchemaPromptSpecAppBuilderDataLinkedResource(_common.BaseModel): + """A linked resource attached to the application by the user.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[VertexBaseConfig] = Field(default=None, description="""""") + display_name: Optional[str] = Field( + default=None, + description="""A user-friendly name for the data source shown in the UI.""", + ) + name: Optional[str] = Field( + default=None, + description="""The unique resource name of the data source. The format is determined by the 'type' field. For type "SAVED_PROMPT": projects/{project}/locations/{location}/datasets/{dataset} For type "AI_AGENT": projects/{project}/locations/{location}/agents/{agent}""", + ) + type: Optional[str] = Field( + default=None, + description="""The type of the linked resource. e.g., "SAVED_PROMPT", "AI_AGENT" This string corresponds to the name of the LinkedResourceType enum member. See: google3/cloud/console/web/ai/platform/llm/prompts/build/services/specs_repository_service/linked_resources/linked_resource.ts""", + ) -class _GetMultimodalDatasetParametersDict(TypedDict, total=False): - """Parameters for getting a multimodal dataset resource.""" +class SchemaPromptSpecAppBuilderDataLinkedResourceDict(TypedDict, total=False): + """A linked resource attached to the application by the user.""" + + display_name: Optional[str] + """A user-friendly name for the data source shown in the UI.""" name: Optional[str] - """""" + """The unique resource name of the data source. The format is determined by the 'type' field. For type "SAVED_PROMPT": projects/{project}/locations/{location}/datasets/{dataset} For type "AI_AGENT": projects/{project}/locations/{location}/agents/{agent}""" - config: Optional[VertexBaseConfigDict] - """""" + type: Optional[str] + """The type of the linked resource. e.g., "SAVED_PROMPT", "AI_AGENT" This string corresponds to the name of the LinkedResourceType enum member. See: google3/cloud/console/web/ai/platform/llm/prompts/build/services/specs_repository_service/linked_resources/linked_resource.ts""" -_GetMultimodalDatasetParametersOrDict = Union[ - _GetMultimodalDatasetParameters, _GetMultimodalDatasetParametersDict +SchemaPromptSpecAppBuilderDataLinkedResourceOrDict = Union[ + SchemaPromptSpecAppBuilderDataLinkedResource, + SchemaPromptSpecAppBuilderDataLinkedResourceDict, ] -class MultimodalDataset(_common.BaseModel): - """Represents a multimodal dataset.""" +class SchemaPromptSpecAppBuilderData(_common.BaseModel): + """Defines data for an application builder.""" - name: Optional[str] = Field( - default=None, description="""The ID of the multimodal dataset.""" - ) - display_name: Optional[str] = Field( - default=None, description="""The display name of the multimodal dataset.""" + code_repository_state: Optional[str] = Field( + default=None, + description="""Serialized state of the code repository. This string will typically contain a JSON representation of the UI's CodeRepositoryService state (files, folders, content, and any metadata). The UI is responsible for serialization and deserialization.""", ) - metadata: Optional[SchemaTablesDatasetMetadata] = Field( - default=None, description="""The metadata of the multimodal dataset.""" + framework: Optional[Framework] = Field( + default=None, + description="""Optional. Framework used to build the application.""", ) - description: Optional[str] = Field( - default=None, description="""The description of the multimodal dataset.""" + linked_resources: Optional[list[SchemaPromptSpecAppBuilderDataLinkedResource]] = ( + Field( + default=None, + description="""Linked resources attached to the application by the user.""", + ) ) - @property - def read_config(self) -> Optional[GeminiRequestReadConfig]: - """Gets the read config from the dataset metadata. Returns None if it's not set.""" - if self.metadata is None or self.metadata.gemini_request_read_config is None: - return None - return self.metadata.gemini_request_read_config - def set_read_config( - self, - *, - read_config: GeminiRequestReadConfigOrDict, - ) -> None: - """Sets the read config in the dataset metadata.""" - if isinstance(read_config, dict): - read_config = GeminiRequestReadConfig(**read_config) +class SchemaPromptSpecAppBuilderDataDict(TypedDict, total=False): + """Defines data for an application builder.""" - if self.metadata is None: - self.metadata = SchemaTablesDatasetMetadata() - self.metadata.gemini_request_read_config = read_config + code_repository_state: Optional[str] + """Serialized state of the code repository. This string will typically contain a JSON representation of the UI's CodeRepositoryService state (files, folders, content, and any metadata). The UI is responsible for serialization and deserialization.""" - @property - def bigquery_uri( - self, - ) -> Optional[str]: - """Gets the bigquery uri from the dataset metadata. Returns None if it's not set.""" - if ( - self.metadata is None - or self.metadata.input_config is None - or self.metadata.input_config.bigquery_source is None - ): - return None - return str(self.metadata.input_config.bigquery_source.uri) + framework: Optional[Framework] + """Optional. Framework used to build the application.""" - def set_bigquery_uri( - self, - bigquery_uri: str, - ) -> None: - """Sets the bigquery uri in the dataset metadata. Prepends 'bq://' if it's not already present.""" - if not bigquery_uri.startswith("bq://"): - bigquery_uri = f"bq://{bigquery_uri}" - metadata = ( - SchemaTablesDatasetMetadata() if self.metadata is None else self.metadata - ) - input_config = ( - SchemaTablesDatasetMetadataInputConfig() - if metadata.input_config is None - else metadata.input_config - ) - bigquery_source = ( - SchemaTablesDatasetMetadataBigQuerySource() - if input_config.bigquery_source is None - else input_config.bigquery_source - ) - bigquery_source.uri = bigquery_uri - input_config.bigquery_source = bigquery_source - metadata.input_config = input_config - self.metadata = metadata + linked_resources: Optional[list[SchemaPromptSpecAppBuilderDataLinkedResourceDict]] + """Linked resources attached to the application by the user.""" - def to_bigframes( - self, - ) -> "bigframes.pandas.DataFrame": # type: ignore # noqa: F821 - """Converts the multimodal dataset to a BigFrames dataframe. - This is the preferred method to inspect the multimodal dataset in a - notebook. +SchemaPromptSpecAppBuilderDataOrDict = Union[ + SchemaPromptSpecAppBuilderData, SchemaPromptSpecAppBuilderDataDict +] - Returns: - A BigFrames dataframe. - """ - from .. import _datasets_utils - bigframes = _datasets_utils._try_import_bigframes() +class SchemaPromptSpecInteractionData(_common.BaseModel): + """Defines data for an interaction prompt.""" - if self.bigquery_uri is None: - raise ValueError("Multimodal dataset bigquery source uri is not set.") - return bigframes.pandas.read_gbq_table(self.bigquery_uri.removeprefix("bq://")) + interaction_ids: Optional[list[str]] = Field( + default=None, + description="""Optional. Lists interaction IDs associated with the prompt. This maps 1:1 to PromptMessage.contents. If InteractionData is present, every prompt message has an interaction ID.""", + ) - def to_batch_job_source(self) -> "genai_types.BatchJobSource": - """Converts the dataset to a BatchJobSource.""" - return genai_types.BatchJobSource( - vertex_dataset_name=self.name, - ) - def get_batch_job_destination(self) -> "genai_types.BatchJobDestination": - """Converts the dataset to a BatchJobDestination.""" - from .. import _datasets_utils +class SchemaPromptSpecInteractionDataDict(TypedDict, total=False): + """Defines data for an interaction prompt.""" - unique_name = _datasets_utils.get_batch_job_unique_name() - bigquery_uri = self.bigquery_uri - if bigquery_uri is None: - raise ValueError("Multimodal dataset bigquery source uri is not set.") - curr_display_name = self.display_name or "genai_batch_job" - return genai_types.BatchJobDestination( - vertex_dataset=genai_types.VertexMultimodalDatasetDestination( - display_name=f"{curr_display_name}_batch_output_{unique_name}", - bigquery_destination=f"{bigquery_uri}_batch_output_{unique_name}", - ) - ) + interaction_ids: Optional[list[str]] + """Optional. Lists interaction IDs associated with the prompt. This maps 1:1 to PromptMessage.contents. If InteractionData is present, every prompt message has an interaction ID.""" -class MultimodalDatasetDict(TypedDict, total=False): - """Represents a multimodal dataset.""" +SchemaPromptSpecInteractionDataOrDict = Union[ + SchemaPromptSpecInteractionData, SchemaPromptSpecInteractionDataDict +] - name: Optional[str] - """The ID of the multimodal dataset.""" - display_name: Optional[str] - """The display name of the multimodal dataset.""" +class SchemaPromptSpecStructuredPrompt(_common.BaseModel): + """Represents a structured prompt.""" - metadata: Optional[SchemaTablesDatasetMetadataDict] - """The metadata of the multimodal dataset.""" + context: Optional[genai_types.Content] = Field( + default=None, description="""Preamble: The context of the prompt.""" + ) + app_builder_data: Optional[SchemaPromptSpecAppBuilderData] = Field( + default=None, description="""Data for app builder use case.""" + ) + examples: Optional[list[SchemaPromptSpecPartList]] = Field( + default=None, + description="""Preamble: A set of examples for expected model response.""", + ) + infill_prefix: Optional[str] = Field( + default=None, + description="""Preamble: For infill prompt, the prefix before expected model response.""", + ) + infill_suffix: Optional[str] = Field( + default=None, + description="""Preamble: For infill prompt, the suffix after expected model response.""", + ) + input_prefixes: Optional[list[str]] = Field( + default=None, + description="""Preamble: The input prefixes before each example input.""", + ) + output_prefixes: Optional[list[str]] = Field( + default=None, + description="""Preamble: The output prefixes before each example output.""", + ) + prediction_inputs: Optional[list[SchemaPromptSpecPartList]] = Field( + default=None, + description="""Preamble: The input test data for prediction. Each PartList in this field represents one text-only input set for a single model request.""", + ) + prompt_message: Optional[SchemaPromptSpecPromptMessage] = Field( + default=None, description="""The prompt message.""" + ) + interaction_data: Optional[SchemaPromptSpecInteractionData] = Field( + default=None, description="""Data for interaction use case.""" + ) - description: Optional[str] - """The description of the multimodal dataset.""" +class SchemaPromptSpecStructuredPromptDict(TypedDict, total=False): + """Represents a structured prompt.""" -MultimodalDatasetOrDict = Union[MultimodalDataset, MultimodalDatasetDict] + context: Optional[genai_types.Content] + """Preamble: The context of the prompt.""" + app_builder_data: Optional[SchemaPromptSpecAppBuilderDataDict] + """Data for app builder use case.""" -class GetMultimodalDatasetOperationConfig(_common.BaseModel): - """Config for getting a multimodal dataset operation.""" + examples: Optional[list[SchemaPromptSpecPartListDict]] + """Preamble: A set of examples for expected model response.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + infill_prefix: Optional[str] + """Preamble: For infill prompt, the prefix before expected model response.""" + infill_suffix: Optional[str] + """Preamble: For infill prompt, the suffix after expected model response.""" -class GetMultimodalDatasetOperationConfigDict(TypedDict, total=False): - """Config for getting a multimodal dataset operation.""" + input_prefixes: Optional[list[str]] + """Preamble: The input prefixes before each example input.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + output_prefixes: Optional[list[str]] + """Preamble: The output prefixes before each example output.""" + prediction_inputs: Optional[list[SchemaPromptSpecPartListDict]] + """Preamble: The input test data for prediction. Each PartList in this field represents one text-only input set for a single model request.""" -GetMultimodalDatasetOperationConfigOrDict = Union[ - GetMultimodalDatasetOperationConfig, GetMultimodalDatasetOperationConfigDict -] + prompt_message: Optional[SchemaPromptSpecPromptMessageDict] + """The prompt message.""" + interaction_data: Optional[SchemaPromptSpecInteractionDataDict] + """Data for interaction use case.""" -class _GetMultimodalDatasetOperationParameters(_common.BaseModel): - """Parameters for getting a dataset operation.""" - dataset_id: Optional[str] = Field(default=None, description="""""") - operation_id: Optional[str] = Field(default=None, description="""""") - config: Optional[GetMultimodalDatasetOperationConfig] = Field( - default=None, description="""""" - ) +SchemaPromptSpecStructuredPromptOrDict = Union[ + SchemaPromptSpecStructuredPrompt, SchemaPromptSpecStructuredPromptDict +] -class _GetMultimodalDatasetOperationParametersDict(TypedDict, total=False): - """Parameters for getting a dataset operation.""" +class SchemaPromptSpecReferenceSentencePair(_common.BaseModel): + """A pair of sentences used as reference in source and target languages.""" - dataset_id: Optional[str] - """""" + source_sentence: Optional[str] = Field( + default=None, description="""Source sentence in the sentence pair.""" + ) + target_sentence: Optional[str] = Field( + default=None, description="""Target sentence in the sentence pair.""" + ) - operation_id: Optional[str] - """""" - config: Optional[GetMultimodalDatasetOperationConfigDict] - """""" +class SchemaPromptSpecReferenceSentencePairDict(TypedDict, total=False): + """A pair of sentences used as reference in source and target languages.""" + source_sentence: Optional[str] + """Source sentence in the sentence pair.""" -_GetMultimodalDatasetOperationParametersOrDict = Union[ - _GetMultimodalDatasetOperationParameters, - _GetMultimodalDatasetOperationParametersDict, -] + target_sentence: Optional[str] + """Target sentence in the sentence pair.""" -class ListMultimodalDatasetsConfig(_common.BaseModel): - """Config for listing multimodal datasets.""" - - 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 ListMultimodalDatasetsConfigDict(TypedDict, total=False): - """Config for listing multimodal datasets.""" - - 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.""" - - -ListMultimodalDatasetsConfigOrDict = Union[ - ListMultimodalDatasetsConfig, ListMultimodalDatasetsConfigDict +SchemaPromptSpecReferenceSentencePairOrDict = Union[ + SchemaPromptSpecReferenceSentencePair, SchemaPromptSpecReferenceSentencePairDict ] -class _ListMultimodalDatasetsRequestParameters(_common.BaseModel): - """Parameters for listing multimodal datasets.""" +class SchemaPromptSpecReferenceSentencePairList(_common.BaseModel): + """A list of reference sentence pairs.""" - config: Optional[ListMultimodalDatasetsConfig] = Field( - default=None, description="""""" + reference_sentence_pairs: Optional[list[SchemaPromptSpecReferenceSentencePair]] = ( + Field(default=None, description="""Reference sentence pairs.""") ) -class _ListMultimodalDatasetsRequestParametersDict(TypedDict, total=False): - """Parameters for listing multimodal datasets.""" +class SchemaPromptSpecReferenceSentencePairListDict(TypedDict, total=False): + """A list of reference sentence pairs.""" - config: Optional[ListMultimodalDatasetsConfigDict] - """""" + reference_sentence_pairs: Optional[list[SchemaPromptSpecReferenceSentencePairDict]] + """Reference sentence pairs.""" -_ListMultimodalDatasetsRequestParametersOrDict = Union[ - _ListMultimodalDatasetsRequestParameters, - _ListMultimodalDatasetsRequestParametersDict, +SchemaPromptSpecReferenceSentencePairListOrDict = Union[ + SchemaPromptSpecReferenceSentencePairList, + SchemaPromptSpecReferenceSentencePairListDict, ] -class ListMultimodalDatasetsResponse(_common.BaseModel): - """Response for listing multimodal datasets.""" +class SchemaPromptSpecTranslationFileInputSource(_common.BaseModel): - 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="""""") - timeout: Optional[int] = Field( - default=90, - description="""The timeout for the list datasets request in seconds. If not set, - the default timeout is 90 seconds.""", + content: Optional[str] = Field(default=None, description="""The file's contents.""") + display_name: Optional[str] = Field( + default=None, description="""The file's display name.""" ) - datasets: Optional[list[MultimodalDataset]] = Field( - default=None, - description="""List of datasets for the project. - """, + mime_type: Optional[str] = Field( + default=None, description="""The file's mime type.""" ) -class ListMultimodalDatasetsResponseDict(TypedDict, total=False): - """Response for listing multimodal datasets.""" - - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" +class SchemaPromptSpecTranslationFileInputSourceDict(TypedDict, total=False): - next_page_token: Optional[str] - """""" + content: Optional[str] + """The file's contents.""" - timeout: Optional[int] - """The timeout for the list datasets request in seconds. If not set, - the default timeout is 90 seconds.""" + display_name: Optional[str] + """The file's display name.""" - datasets: Optional[list[MultimodalDatasetDict]] - """List of datasets for the project. - """ + mime_type: Optional[str] + """The file's mime type.""" -ListMultimodalDatasetsResponseOrDict = Union[ - ListMultimodalDatasetsResponse, ListMultimodalDatasetsResponseDict +SchemaPromptSpecTranslationFileInputSourceOrDict = Union[ + SchemaPromptSpecTranslationFileInputSource, + SchemaPromptSpecTranslationFileInputSourceDict, ] -class _UpdateMultimodalDatasetParameters(_common.BaseModel): - """Parameters for updating a multimodal dataset resource.""" +class SchemaPromptSpecTranslationGcsInputSource(_common.BaseModel): - name: Optional[str] = Field(default=None, description="""""") - display_name: Optional[str] = Field(default=None, description="""""") - metadata: Optional[SchemaTablesDatasetMetadata] = Field( - default=None, description="""""" - ) - description: Optional[str] = Field(default=None, description="""""") - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, description="""""" + input_uri: Optional[str] = Field( + default=None, + description="""Source data URI. For example, `gs://my_bucket/my_object`.""", ) - config: Optional[VertexBaseConfig] = Field(default=None, description="""""") - - -class _UpdateMultimodalDatasetParametersDict(TypedDict, total=False): - """Parameters for updating a multimodal dataset resource.""" - - name: Optional[str] - """""" - - display_name: Optional[str] - """""" - metadata: Optional[SchemaTablesDatasetMetadataDict] - """""" - - description: Optional[str] - """""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """""" +class SchemaPromptSpecTranslationGcsInputSourceDict(TypedDict, total=False): - config: Optional[VertexBaseConfigDict] - """""" + input_uri: Optional[str] + """Source data URI. For example, `gs://my_bucket/my_object`.""" -_UpdateMultimodalDatasetParametersOrDict = Union[ - _UpdateMultimodalDatasetParameters, _UpdateMultimodalDatasetParametersDict +SchemaPromptSpecTranslationGcsInputSourceOrDict = Union[ + SchemaPromptSpecTranslationGcsInputSource, + SchemaPromptSpecTranslationGcsInputSourceDict, ] -class SchemaPredictParamsGroundingConfigSourceEntry(_common.BaseModel): - """Single source entry for the grounding checking.""" +class SchemaPromptSpecTranslationSentenceFileInput(_common.BaseModel): - enterprise_datastore: Optional[str] = Field( - default=None, - description="""The uri of the Vertex AI Search data source. Deprecated. Use vertex_ai_search_datastore instead.""", - ) - inline_context: Optional[str] = Field( - default=None, - description="""The grounding text passed inline with the Predict API. It can support up to 1 million bytes.""", - ) - type: Optional[ - Literal["UNSPECIFIED", "WEB", "ENTERPRISE", "VERTEX_AI_SEARCH", "INLINE"] - ] = Field( - default=None, description="""The type of the grounding checking source.""" + file_input_source: Optional[SchemaPromptSpecTranslationFileInputSource] = Field( + default=None, description="""Inlined file source.""" ) - vertex_ai_search_datastore: Optional[str] = Field( - default=None, description="""The uri of the Vertex AI Search data source.""" + gcs_input_source: Optional[SchemaPromptSpecTranslationGcsInputSource] = Field( + default=None, description="""Cloud Storage file source.""" ) -class SchemaPredictParamsGroundingConfigSourceEntryDict(TypedDict, total=False): - """Single source entry for the grounding checking.""" - - enterprise_datastore: Optional[str] - """The uri of the Vertex AI Search data source. Deprecated. Use vertex_ai_search_datastore instead.""" - - inline_context: Optional[str] - """The grounding text passed inline with the Predict API. It can support up to 1 million bytes.""" +class SchemaPromptSpecTranslationSentenceFileInputDict(TypedDict, total=False): - type: Optional[ - Literal["UNSPECIFIED", "WEB", "ENTERPRISE", "VERTEX_AI_SEARCH", "INLINE"] - ] - """The type of the grounding checking source.""" + file_input_source: Optional[SchemaPromptSpecTranslationFileInputSourceDict] + """Inlined file source.""" - vertex_ai_search_datastore: Optional[str] - """The uri of the Vertex AI Search data source.""" + gcs_input_source: Optional[SchemaPromptSpecTranslationGcsInputSourceDict] + """Cloud Storage file source.""" -SchemaPredictParamsGroundingConfigSourceEntryOrDict = Union[ - SchemaPredictParamsGroundingConfigSourceEntry, - SchemaPredictParamsGroundingConfigSourceEntryDict, +SchemaPromptSpecTranslationSentenceFileInputOrDict = Union[ + SchemaPromptSpecTranslationSentenceFileInput, + SchemaPromptSpecTranslationSentenceFileInputDict, ] -class SchemaPredictParamsGroundingConfig(_common.BaseModel): - """The configuration for grounding checking.""" +class SchemaPromptSpecTranslationExample(_common.BaseModel): + """The translation example that contains reference sentences from various sources.""" - disable_attribution: Optional[bool] = Field( - default=None, - description="""If set, skip finding claim attributions (i.e not generate grounding citation).""", - ) - sources: Optional[list[SchemaPredictParamsGroundingConfigSourceEntry]] = Field( - default=None, description="""The sources for the grounding checking.""" - ) + reference_sentence_pair_lists: Optional[ + list[SchemaPromptSpecReferenceSentencePairList] + ] = Field(default=None, description="""The reference sentences from inline text.""") + reference_sentences_file_inputs: Optional[ + list[SchemaPromptSpecTranslationSentenceFileInput] + ] = Field(default=None, description="""The reference sentences from file.""") -class SchemaPredictParamsGroundingConfigDict(TypedDict, total=False): - """The configuration for grounding checking.""" +class SchemaPromptSpecTranslationExampleDict(TypedDict, total=False): + """The translation example that contains reference sentences from various sources.""" - disable_attribution: Optional[bool] - """If set, skip finding claim attributions (i.e not generate grounding citation).""" + reference_sentence_pair_lists: Optional[ + list[SchemaPromptSpecReferenceSentencePairListDict] + ] + """The reference sentences from inline text.""" - sources: Optional[list[SchemaPredictParamsGroundingConfigSourceEntryDict]] - """The sources for the grounding checking.""" + reference_sentences_file_inputs: Optional[ + list[SchemaPromptSpecTranslationSentenceFileInputDict] + ] + """The reference sentences from file.""" -SchemaPredictParamsGroundingConfigOrDict = Union[ - SchemaPredictParamsGroundingConfig, SchemaPredictParamsGroundingConfigDict +SchemaPromptSpecTranslationExampleOrDict = Union[ + SchemaPromptSpecTranslationExample, SchemaPromptSpecTranslationExampleDict ] -class SchemaPromptSpecPartList(_common.BaseModel): - """Represents a prompt spec part list.""" +class SchemaPromptSpecTranslationOption(_common.BaseModel): + """Optional settings for translation prompt.""" - parts: Optional[list[genai_types.Part]] = Field( - default=None, description="""A list of elements that can be part of a prompt.""" + number_of_shots: Optional[int] = Field( + default=None, description="""How many shots to use.""" ) -class SchemaPromptSpecPartListDict(TypedDict, total=False): - """Represents a prompt spec part list.""" +class SchemaPromptSpecTranslationOptionDict(TypedDict, total=False): + """Optional settings for translation prompt.""" - parts: Optional[list[genai_types.Part]] - """A list of elements that can be part of a prompt.""" + number_of_shots: Optional[int] + """How many shots to use.""" -SchemaPromptSpecPartListOrDict = Union[ - SchemaPromptSpecPartList, SchemaPromptSpecPartListDict +SchemaPromptSpecTranslationOptionOrDict = Union[ + SchemaPromptSpecTranslationOption, SchemaPromptSpecTranslationOptionDict ] -class SchemaPromptInstanceVariableValue(_common.BaseModel): - """Represents a prompt instance variable.""" +class SchemaPromptSpecTranslationPrompt(_common.BaseModel): + """Prompt variation for Translation use case.""" - part_list: Optional[SchemaPromptSpecPartList] = Field( - default=None, description="""The parts of the variable value.""" + example: Optional[SchemaPromptSpecTranslationExample] = Field( + default=None, description="""The translation example.""" + ) + option: Optional[SchemaPromptSpecTranslationOption] = Field( + default=None, description="""The translation option.""" + ) + prompt_message: Optional[SchemaPromptSpecPromptMessage] = Field( + default=None, description="""The prompt message.""" + ) + source_language_code: Optional[str] = Field( + default=None, description="""The source language code.""" + ) + target_language_code: Optional[str] = Field( + default=None, description="""The target language code.""" ) -class SchemaPromptInstanceVariableValueDict(TypedDict, total=False): - """Represents a prompt instance variable.""" +class SchemaPromptSpecTranslationPromptDict(TypedDict, total=False): + """Prompt variation for Translation use case.""" - part_list: Optional[SchemaPromptSpecPartListDict] - """The parts of the variable value.""" + example: Optional[SchemaPromptSpecTranslationExampleDict] + """The translation example.""" + option: Optional[SchemaPromptSpecTranslationOptionDict] + """The translation option.""" -SchemaPromptInstanceVariableValueOrDict = Union[ - SchemaPromptInstanceVariableValue, SchemaPromptInstanceVariableValueDict + prompt_message: Optional[SchemaPromptSpecPromptMessageDict] + """The prompt message.""" + + source_language_code: Optional[str] + """The source language code.""" + + target_language_code: Optional[str] + """The target language code.""" + + +SchemaPromptSpecTranslationPromptOrDict = Union[ + SchemaPromptSpecTranslationPrompt, SchemaPromptSpecTranslationPromptDict ] -class SchemaPromptInstancePromptExecution(_common.BaseModel): - """A prompt instance's parameters set that contains a set of variable values.""" +class SchemaPromptApiSchema(_common.BaseModel): + """The A2 schema of a prompt.""" - arguments: Optional[dict[str, SchemaPromptInstanceVariableValue]] = Field( - default=None, description="""Maps variable names to their value.""" + api_schema_version: Optional[str] = Field( + default=None, + description="""The Schema version that represents changes to the API behavior.""", + ) + executions: Optional[list[SchemaPromptInstancePromptExecution]] = Field( + default=None, + description="""A list of execution instances for constructing a ready-to-use prompt.""", + ) + multimodal_prompt: Optional[SchemaPromptSpecMultimodalPrompt] = Field( + default=None, + description="""Multimodal prompt which embeds preambles to prompt string.""", + ) + structured_prompt: Optional[SchemaPromptSpecStructuredPrompt] = Field( + default=None, + description="""The prompt variation that stores preambles in separate fields.""", + ) + translation_prompt: Optional[SchemaPromptSpecTranslationPrompt] = Field( + default=None, description="""The prompt variation for Translation use case.""" ) -class SchemaPromptInstancePromptExecutionDict(TypedDict, total=False): - """A prompt instance's parameters set that contains a set of variable values.""" +class SchemaPromptApiSchemaDict(TypedDict, total=False): + """The A2 schema of a prompt.""" - arguments: Optional[dict[str, SchemaPromptInstanceVariableValueDict]] - """Maps variable names to their value.""" + api_schema_version: Optional[str] + """The Schema version that represents changes to the API behavior.""" + executions: Optional[list[SchemaPromptInstancePromptExecutionDict]] + """A list of execution instances for constructing a ready-to-use prompt.""" -SchemaPromptInstancePromptExecutionOrDict = Union[ - SchemaPromptInstancePromptExecution, SchemaPromptInstancePromptExecutionDict -] + multimodal_prompt: Optional[SchemaPromptSpecMultimodalPromptDict] + """Multimodal prompt which embeds preambles to prompt string.""" + structured_prompt: Optional[SchemaPromptSpecStructuredPromptDict] + """The prompt variation that stores preambles in separate fields.""" -class SchemaPromptSpecPromptMessage(_common.BaseModel): - """Represents a prompt message.""" + translation_prompt: Optional[SchemaPromptSpecTranslationPromptDict] + """The prompt variation for Translation use case.""" - generation_config: Optional[genai_types.GenerationConfig] = Field( - default=None, description="""Generation config.""" + +SchemaPromptApiSchemaOrDict = Union[SchemaPromptApiSchema, SchemaPromptApiSchemaDict] + + +class SchemaTextPromptDatasetMetadata(_common.BaseModel): + """Represents the text prompt dataset metadata.""" + + candidate_count: Optional[int] = Field( + default=None, description="""Number of candidates.""" ) - tool_config: Optional[genai_types.FunctionCallingConfig] = Field( + gcs_uri: Optional[str] = Field( default=None, - description="""Tool config. This config is shared for all tools provided in the request.""", + description="""The Google Cloud Storage URI that stores the prompt data.""", ) - tools: Optional[list[genai_types.Tool]] = Field( + grounding_config: Optional[SchemaPredictParamsGroundingConfig] = Field( + default=None, description="""Grounding checking configuration.""" + ) + has_prompt_variable: Optional[bool] = Field( + default=None, description="""Whether the prompt dataset has prompt variable.""" + ) + logprobs: Optional[bool] = Field( default=None, - description="""A list of `Tools` the model may use to generate the next response. A `Tool` is a piece of code that enables the system to interact with external systems to perform an action, or set of actions, outside of knowledge and scope of the model.""", + description="""Whether or not the user has enabled logit probabilities in the model parameters.""", ) - safety_settings: Optional[list[genai_types.SafetySetting]] = Field( + max_output_tokens: Optional[int] = Field( default=None, - description="""Per request settings for blocking unsafe content. Enforced on GenerateContentResponse.candidates.""", + description="""Value of the maximum number of tokens generated set when the dataset was saved.""", ) - contents: Optional[list[genai_types.Content]] = Field( + note: Optional[str] = Field( default=None, - description="""The content of the current conversation with the model. 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="""User-created prompt note. Note size limit is 2KB.""", ) - system_instruction: Optional[genai_types.Content] = Field( + prompt_api_schema: Optional[SchemaPromptApiSchema] = Field( default=None, - description="""The user provided system instructions for the model. Note: only text should be used in parts and content in each part will be in a separate paragraph.""", + description="""The API schema of the prompt to support both UI and SDK usages.""", ) - variables: Optional[list[dict[str, genai_types.Part]]] = Field( - default=None, description="""""" + prompt_type: Optional[str] = Field( + default=None, description="""Type of the prompt dataset.""" + ) + seed_enabled: Optional[bool] = Field( + default=None, + description="""Seeding enables model to return a deterministic response on a best effort basis. Determinism isn't guaranteed. This field determines whether or not seeding is enabled.""", + ) + seed_value: Optional[int] = Field( + default=None, description="""The actual value of the seed.""" + ) + stop_sequences: Optional[list[str]] = Field( + default=None, description="""Customized stop sequences.""" + ) + system_instruction: Optional[str] = Field( + default=None, + description="""The content of the prompt dataset system instruction.""", + ) + system_instruction_gcs_uri: Optional[str] = Field( + default=None, + description="""The Google Cloud Storage URI that stores the system instruction, starting with gs://.""", + ) + temperature: Optional[float] = Field( + default=None, + description="""Temperature value used for sampling set when the dataset was saved. This value is used to tune the degree of randomness.""", + ) + text: Optional[str] = Field( + default=None, description="""The content of the prompt dataset.""" + ) + top_k: Optional[int] = Field( + default=None, + description="""Top K value set when the dataset was saved. This value determines how many candidates with highest probability from the vocab would be selected for each decoding step.""", + ) + top_p: Optional[float] = Field( + default=None, + description="""Top P value set when the dataset was saved. Given topK tokens for decoding, top candidates will be selected until the sum of their probabilities is topP.""", ) - model: Optional[str] = Field(default=None, description="""The model name.""") -class SchemaPromptSpecPromptMessageDict(TypedDict, total=False): - """Represents a prompt message.""" +class SchemaTextPromptDatasetMetadataDict(TypedDict, total=False): + """Represents the text prompt dataset metadata.""" - generation_config: Optional[genai_types.GenerationConfig] - """Generation config.""" + candidate_count: Optional[int] + """Number of candidates.""" - tool_config: Optional[genai_types.FunctionCallingConfig] - """Tool config. This config is shared for all tools provided in the request.""" + gcs_uri: Optional[str] + """The Google Cloud Storage URI that stores the prompt data.""" - tools: Optional[list[genai_types.Tool]] - """A list of `Tools` the model may use to generate the next response. A `Tool` is a piece of code that enables the system to interact with external systems to perform an action, or set of actions, outside of knowledge and scope of the model.""" + grounding_config: Optional[SchemaPredictParamsGroundingConfigDict] + """Grounding checking configuration.""" - safety_settings: Optional[list[genai_types.SafetySetting]] - """Per request settings for blocking unsafe content. Enforced on GenerateContentResponse.candidates.""" + has_prompt_variable: Optional[bool] + """Whether the prompt dataset has prompt variable.""" - contents: Optional[list[genai_types.Content]] - """The content of the current conversation with the model. For single-turn queries, this is a single instance. For multi-turn queries, this is a repeated field that contains conversation history + latest request.""" + logprobs: Optional[bool] + """Whether or not the user has enabled logit probabilities in the model parameters.""" - system_instruction: Optional[genai_types.Content] - """The user provided system instructions for the model. Note: only text should be used in parts and content in each part will be in a separate paragraph.""" + max_output_tokens: Optional[int] + """Value of the maximum number of tokens generated set when the dataset was saved.""" - variables: Optional[list[dict[str, genai_types.Part]]] - """""" + note: Optional[str] + """User-created prompt note. Note size limit is 2KB.""" - model: Optional[str] - """The model name.""" + prompt_api_schema: Optional[SchemaPromptApiSchemaDict] + """The API schema of the prompt to support both UI and SDK usages.""" + prompt_type: Optional[str] + """Type of the prompt dataset.""" -SchemaPromptSpecPromptMessageOrDict = Union[ - SchemaPromptSpecPromptMessage, SchemaPromptSpecPromptMessageDict -] + seed_enabled: Optional[bool] + """Seeding enables model to return a deterministic response on a best effort basis. Determinism isn't guaranteed. This field determines whether or not seeding is enabled.""" + seed_value: Optional[int] + """The actual value of the seed.""" -class SchemaPromptSpecMultimodalPrompt(_common.BaseModel): - """Prompt variation that embeds preambles to prompt string.""" + stop_sequences: Optional[list[str]] + """Customized stop sequences.""" - prompt_message: Optional[SchemaPromptSpecPromptMessage] = Field( - default=None, description="""The prompt message.""" - ) + system_instruction: Optional[str] + """The content of the prompt dataset system instruction.""" + system_instruction_gcs_uri: Optional[str] + """The Google Cloud Storage URI that stores the system instruction, starting with gs://.""" -class SchemaPromptSpecMultimodalPromptDict(TypedDict, total=False): - """Prompt variation that embeds preambles to prompt string.""" + temperature: Optional[float] + """Temperature value used for sampling set when the dataset was saved. This value is used to tune the degree of randomness.""" - prompt_message: Optional[SchemaPromptSpecPromptMessageDict] - """The prompt message.""" + text: Optional[str] + """The content of the prompt dataset.""" + top_k: Optional[int] + """Top K value set when the dataset was saved. This value determines how many candidates with highest probability from the vocab would be selected for each decoding step.""" -SchemaPromptSpecMultimodalPromptOrDict = Union[ - SchemaPromptSpecMultimodalPrompt, SchemaPromptSpecMultimodalPromptDict + top_p: Optional[float] + """Top P value set when the dataset was saved. Given topK tokens for decoding, top candidates will be selected until the sum of their probabilities is topP.""" + + +SchemaTextPromptDatasetMetadataOrDict = Union[ + SchemaTextPromptDatasetMetadata, SchemaTextPromptDatasetMetadataDict ] -class SchemaPromptSpecAppBuilderDataLinkedResource(_common.BaseModel): - """A linked resource attached to the application by the user.""" +class CreateDatasetConfig(_common.BaseModel): + """Config for creating a dataset resource to store prompts.""" - display_name: Optional[str] = Field( - default=None, - description="""A user-friendly name for the data source shown in the UI.""", - ) - name: Optional[str] = Field( - default=None, - description="""The unique resource name of the data source. The format is determined by the 'type' field. For type "SAVED_PROMPT": projects/{project}/locations/{location}/datasets/{dataset} For type "AI_AGENT": projects/{project}/locations/{location}/agents/{agent}""", - ) - type: Optional[str] = Field( - default=None, - description="""The type of the linked resource. e.g., "SAVED_PROMPT", "AI_AGENT" This string corresponds to the name of the LinkedResourceType enum member. See: google3/cloud/console/web/ai/platform/llm/prompts/build/services/specs_repository_service/linked_resources/linked_resource.ts""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class SchemaPromptSpecAppBuilderDataLinkedResourceDict(TypedDict, total=False): - """A linked resource attached to the application by the user.""" +class CreateDatasetConfigDict(TypedDict, total=False): + """Config for creating a dataset resource to store prompts.""" - display_name: Optional[str] - """A user-friendly name for the data source shown in the UI.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - name: Optional[str] - """The unique resource name of the data source. The format is determined by the 'type' field. For type "SAVED_PROMPT": projects/{project}/locations/{location}/datasets/{dataset} For type "AI_AGENT": projects/{project}/locations/{location}/agents/{agent}""" - - type: Optional[str] - """The type of the linked resource. e.g., "SAVED_PROMPT", "AI_AGENT" This string corresponds to the name of the LinkedResourceType enum member. See: google3/cloud/console/web/ai/platform/llm/prompts/build/services/specs_repository_service/linked_resources/linked_resource.ts""" - -SchemaPromptSpecAppBuilderDataLinkedResourceOrDict = Union[ - SchemaPromptSpecAppBuilderDataLinkedResource, - SchemaPromptSpecAppBuilderDataLinkedResourceDict, -] +CreateDatasetConfigOrDict = Union[CreateDatasetConfig, CreateDatasetConfigDict] -class SchemaPromptSpecAppBuilderData(_common.BaseModel): - """Defines data for an application builder.""" +class _CreateDatasetParameters(_common.BaseModel): + """Parameters for creating a dataset resource to store prompts.""" - code_repository_state: Optional[str] = Field( - default=None, - description="""Serialized state of the code repository. This string will typically contain a JSON representation of the UI's CodeRepositoryService state (files, folders, content, and any metadata). The UI is responsible for serialization and deserialization.""", - ) - framework: Optional[Framework] = Field( - default=None, - description="""Optional. Framework used to build the application.""", + name: Optional[str] = Field(default=None, description="""""") + display_name: Optional[str] = Field(default=None, description="""""") + metadata_schema_uri: Optional[str] = Field(default=None, description="""""") + metadata: Optional[SchemaTextPromptDatasetMetadata] = Field( + default=None, description="""""" ) - linked_resources: Optional[list[SchemaPromptSpecAppBuilderDataLinkedResource]] = ( - Field( - default=None, - description="""Linked resources attached to the application by the user.""", - ) + description: Optional[str] = Field(default=None, description="""""") + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, description="""""" ) + model_reference: Optional[str] = Field(default=None, description="""""") + config: Optional[CreateDatasetConfig] = Field(default=None, description="""""") -class SchemaPromptSpecAppBuilderDataDict(TypedDict, total=False): - """Defines data for an application builder.""" - - code_repository_state: Optional[str] - """Serialized state of the code repository. This string will typically contain a JSON representation of the UI's CodeRepositoryService state (files, folders, content, and any metadata). The UI is responsible for serialization and deserialization.""" - - framework: Optional[Framework] - """Optional. Framework used to build the application.""" - - linked_resources: Optional[list[SchemaPromptSpecAppBuilderDataLinkedResourceDict]] - """Linked resources attached to the application by the user.""" +class _CreateDatasetParametersDict(TypedDict, total=False): + """Parameters for creating a dataset resource to store prompts.""" + name: Optional[str] + """""" -SchemaPromptSpecAppBuilderDataOrDict = Union[ - SchemaPromptSpecAppBuilderData, SchemaPromptSpecAppBuilderDataDict -] + display_name: Optional[str] + """""" + metadata_schema_uri: Optional[str] + """""" -class SchemaPromptSpecInteractionData(_common.BaseModel): - """Defines data for an interaction prompt.""" + metadata: Optional[SchemaTextPromptDatasetMetadataDict] + """""" - interaction_ids: Optional[list[str]] = Field( - default=None, - description="""Optional. Lists interaction IDs associated with the prompt. This maps 1:1 to PromptMessage.contents. If InteractionData is present, every prompt message has an interaction ID.""", - ) + description: Optional[str] + """""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """""" -class SchemaPromptSpecInteractionDataDict(TypedDict, total=False): - """Defines data for an interaction prompt.""" + model_reference: Optional[str] + """""" - interaction_ids: Optional[list[str]] - """Optional. Lists interaction IDs associated with the prompt. This maps 1:1 to PromptMessage.contents. If InteractionData is present, every prompt message has an interaction ID.""" + config: Optional[CreateDatasetConfigDict] + """""" -SchemaPromptSpecInteractionDataOrDict = Union[ - SchemaPromptSpecInteractionData, SchemaPromptSpecInteractionDataDict +_CreateDatasetParametersOrDict = Union[ + _CreateDatasetParameters, _CreateDatasetParametersDict ] -class SchemaPromptSpecStructuredPrompt(_common.BaseModel): - """Represents a structured prompt.""" +class DatasetOperation(_common.BaseModel): + """Represents the create dataset operation.""" - context: Optional[genai_types.Content] = Field( - default=None, description="""Preamble: The context of the prompt.""" - ) - app_builder_data: Optional[SchemaPromptSpecAppBuilderData] = Field( - default=None, description="""Data for app builder use case.""" - ) - examples: Optional[list[SchemaPromptSpecPartList]] = Field( - default=None, - description="""Preamble: A set of examples for expected model response.""", - ) - infill_prefix: Optional[str] = Field( - default=None, - description="""Preamble: For infill prompt, the prefix before expected model response.""", - ) - infill_suffix: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Preamble: For infill prompt, the suffix after expected model response.""", + 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}`.""", ) - input_prefixes: Optional[list[str]] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Preamble: The input prefixes before each example input.""", + 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.""", ) - output_prefixes: Optional[list[str]] = Field( + done: Optional[bool] = Field( default=None, - description="""Preamble: The output prefixes before each example output.""", + 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_inputs: Optional[list[SchemaPromptSpecPartList]] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Preamble: The input test data for prediction. Each PartList in this field represents one text-only input set for a single model request.""", - ) - prompt_message: Optional[SchemaPromptSpecPromptMessage] = Field( - default=None, description="""The prompt message.""" + description="""The error result of the operation in case of failure or cancellation.""", ) - interaction_data: Optional[SchemaPromptSpecInteractionData] = Field( - default=None, description="""Data for interaction use case.""" + response: Optional[dict[str, Any]] = Field( + default=None, description="""The result of the dataset operation.""" ) -class SchemaPromptSpecStructuredPromptDict(TypedDict, total=False): - """Represents a structured prompt.""" +class DatasetOperationDict(TypedDict, total=False): + """Represents the create dataset operation.""" - context: Optional[genai_types.Content] - """Preamble: The context of the prompt.""" + 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}`.""" - app_builder_data: Optional[SchemaPromptSpecAppBuilderDataDict] - """Data for app builder use case.""" + 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.""" - examples: Optional[list[SchemaPromptSpecPartListDict]] - """Preamble: A set of examples for expected model response.""" + 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.""" - infill_prefix: Optional[str] - """Preamble: For infill prompt, the prefix before expected model response.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - infill_suffix: Optional[str] - """Preamble: For infill prompt, the suffix after expected model response.""" + response: Optional[dict[str, Any]] + """The result of the dataset operation.""" - input_prefixes: Optional[list[str]] - """Preamble: The input prefixes before each example input.""" - output_prefixes: Optional[list[str]] - """Preamble: The output prefixes before each example output.""" +DatasetOperationOrDict = Union[DatasetOperation, DatasetOperationDict] - prediction_inputs: Optional[list[SchemaPromptSpecPartListDict]] - """Preamble: The input test data for prediction. Each PartList in this field represents one text-only input set for a single model request.""" - prompt_message: Optional[SchemaPromptSpecPromptMessageDict] - """The prompt message.""" +class CreateDatasetVersionConfig(_common.BaseModel): + """Config for creating a dataset version resource to store prompts.""" - interaction_data: Optional[SchemaPromptSpecInteractionDataDict] - """Data for interaction use case.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -SchemaPromptSpecStructuredPromptOrDict = Union[ - SchemaPromptSpecStructuredPrompt, SchemaPromptSpecStructuredPromptDict +class CreateDatasetVersionConfigDict(TypedDict, total=False): + """Config for creating a dataset version resource to store prompts.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +CreateDatasetVersionConfigOrDict = Union[ + CreateDatasetVersionConfig, CreateDatasetVersionConfigDict ] -class SchemaPromptSpecReferenceSentencePair(_common.BaseModel): - """A pair of sentences used as reference in source and target languages.""" +class _CreateDatasetVersionParameters(_common.BaseModel): + """Represents the create dataset version parameters.""" - source_sentence: Optional[str] = Field( - default=None, description="""Source sentence in the sentence pair.""" + dataset_name: Optional[str] = Field(default=None, description="""""") + metadata: Optional[SchemaTextPromptDatasetMetadata] = Field( + default=None, description="""""" ) - target_sentence: Optional[str] = Field( - default=None, description="""Target sentence in the sentence pair.""" + model_reference: Optional[str] = Field(default=None, description="""""") + parent: Optional[str] = Field(default=None, description="""""") + display_name: Optional[str] = Field(default=None, description="""""") + config: Optional[CreateDatasetVersionConfig] = Field( + default=None, description="""""" ) -class SchemaPromptSpecReferenceSentencePairDict(TypedDict, total=False): - """A pair of sentences used as reference in source and target languages.""" +class _CreateDatasetVersionParametersDict(TypedDict, total=False): + """Represents the create dataset version parameters.""" - source_sentence: Optional[str] - """Source sentence in the sentence pair.""" + dataset_name: Optional[str] + """""" - target_sentence: Optional[str] - """Target sentence in the sentence pair.""" + metadata: Optional[SchemaTextPromptDatasetMetadataDict] + """""" + model_reference: Optional[str] + """""" -SchemaPromptSpecReferenceSentencePairOrDict = Union[ - SchemaPromptSpecReferenceSentencePair, SchemaPromptSpecReferenceSentencePairDict -] + parent: Optional[str] + """""" + display_name: Optional[str] + """""" -class SchemaPromptSpecReferenceSentencePairList(_common.BaseModel): - """A list of reference sentence pairs.""" + config: Optional[CreateDatasetVersionConfigDict] + """""" - reference_sentence_pairs: Optional[list[SchemaPromptSpecReferenceSentencePair]] = ( - Field(default=None, description="""Reference sentence pairs.""") - ) +_CreateDatasetVersionParametersOrDict = Union[ + _CreateDatasetVersionParameters, _CreateDatasetVersionParametersDict +] -class SchemaPromptSpecReferenceSentencePairListDict(TypedDict, total=False): - """A list of reference sentence pairs.""" - reference_sentence_pairs: Optional[list[SchemaPromptSpecReferenceSentencePairDict]] - """Reference sentence pairs.""" +class _GetDatasetParameters(_common.BaseModel): + """Parameters for getting a dataset resource to store prompts.""" + name: Optional[str] = Field(default=None, description="""""") + config: Optional[VertexBaseConfig] = Field(default=None, description="""""") -SchemaPromptSpecReferenceSentencePairListOrDict = Union[ - SchemaPromptSpecReferenceSentencePairList, - SchemaPromptSpecReferenceSentencePairListDict, -] +class _GetDatasetParametersDict(TypedDict, total=False): + """Parameters for getting a dataset resource to store prompts.""" -class SchemaPromptSpecTranslationFileInputSource(_common.BaseModel): + name: Optional[str] + """""" - content: Optional[str] = Field(default=None, description="""The file's contents.""") - display_name: Optional[str] = Field( - default=None, description="""The file's display name.""" - ) - mime_type: Optional[str] = Field( - default=None, description="""The file's mime type.""" - ) + config: Optional[VertexBaseConfigDict] + """""" -class SchemaPromptSpecTranslationFileInputSourceDict(TypedDict, total=False): +_GetDatasetParametersOrDict = Union[_GetDatasetParameters, _GetDatasetParametersDict] - content: Optional[str] - """The file's contents.""" - display_name: Optional[str] - """The file's display name.""" +class SavedQuery(_common.BaseModel): + """A SavedQuery is a view of the dataset. It references a subset of annotations by problem type and filters.""" - mime_type: Optional[str] - """The file's mime type.""" + annotation_filter: Optional[str] = Field( + default=None, + description="""Output only. Filters on the Annotations in the dataset.""", + ) + annotation_spec_count: Optional[int] = Field( + default=None, + description="""Output only. Number of AnnotationSpecs in the context of the SavedQuery.""", + ) + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this SavedQuery was created.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""Required. The user-defined name of the SavedQuery. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", + ) + etag: Optional[str] = Field( + default=None, + description="""Used to perform a consistent read-modify-write update. If not set, a blind "overwrite" update happens.""", + ) + metadata: Optional[Any] = Field( + default=None, + description="""Some additional information about the SavedQuery.""", + ) + name: Optional[str] = Field( + default=None, description="""Output only. Resource name of the SavedQuery.""" + ) + problem_type: Optional[str] = Field( + default=None, + description="""Required. Problem type of the SavedQuery. Allowed values: * IMAGE_CLASSIFICATION_SINGLE_LABEL * IMAGE_CLASSIFICATION_MULTI_LABEL * IMAGE_BOUNDING_POLY * IMAGE_BOUNDING_BOX * TEXT_CLASSIFICATION_SINGLE_LABEL * TEXT_CLASSIFICATION_MULTI_LABEL * TEXT_EXTRACTION * TEXT_SENTIMENT * VIDEO_CLASSIFICATION * VIDEO_OBJECT_TRACKING""", + ) + support_automl_training: Optional[bool] = Field( + default=None, + description="""Output only. If the Annotations belonging to the SavedQuery can be used for AutoML training.""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when SavedQuery was last updated.""", + ) -SchemaPromptSpecTranslationFileInputSourceOrDict = Union[ - SchemaPromptSpecTranslationFileInputSource, - SchemaPromptSpecTranslationFileInputSourceDict, -] +class SavedQueryDict(TypedDict, total=False): + """A SavedQuery is a view of the dataset. It references a subset of annotations by problem type and filters.""" + annotation_filter: Optional[str] + """Output only. Filters on the Annotations in the dataset.""" -class SchemaPromptSpecTranslationGcsInputSource(_common.BaseModel): + annotation_spec_count: Optional[int] + """Output only. Number of AnnotationSpecs in the context of the SavedQuery.""" - input_uri: Optional[str] = Field( - default=None, - description="""Source data URI. For example, `gs://my_bucket/my_object`.""", - ) + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this SavedQuery was created.""" + display_name: Optional[str] + """Required. The user-defined name of the SavedQuery. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" -class SchemaPromptSpecTranslationGcsInputSourceDict(TypedDict, total=False): + etag: Optional[str] + """Used to perform a consistent read-modify-write update. If not set, a blind "overwrite" update happens.""" - input_uri: Optional[str] - """Source data URI. For example, `gs://my_bucket/my_object`.""" + metadata: Optional[Any] + """Some additional information about the SavedQuery.""" + name: Optional[str] + """Output only. Resource name of the SavedQuery.""" -SchemaPromptSpecTranslationGcsInputSourceOrDict = Union[ - SchemaPromptSpecTranslationGcsInputSource, - SchemaPromptSpecTranslationGcsInputSourceDict, -] + problem_type: Optional[str] + """Required. Problem type of the SavedQuery. Allowed values: * IMAGE_CLASSIFICATION_SINGLE_LABEL * IMAGE_CLASSIFICATION_MULTI_LABEL * IMAGE_BOUNDING_POLY * IMAGE_BOUNDING_BOX * TEXT_CLASSIFICATION_SINGLE_LABEL * TEXT_CLASSIFICATION_MULTI_LABEL * TEXT_EXTRACTION * TEXT_SENTIMENT * VIDEO_CLASSIFICATION * VIDEO_OBJECT_TRACKING""" + support_automl_training: Optional[bool] + """Output only. If the Annotations belonging to the SavedQuery can be used for AutoML training.""" -class SchemaPromptSpecTranslationSentenceFileInput(_common.BaseModel): + update_time: Optional[datetime.datetime] + """Output only. Timestamp when SavedQuery was last updated.""" - file_input_source: Optional[SchemaPromptSpecTranslationFileInputSource] = Field( - default=None, description="""Inlined file source.""" - ) - gcs_input_source: Optional[SchemaPromptSpecTranslationGcsInputSource] = Field( - default=None, description="""Cloud Storage file source.""" - ) +SavedQueryOrDict = Union[SavedQuery, SavedQueryDict] -class SchemaPromptSpecTranslationSentenceFileInputDict(TypedDict, total=False): - file_input_source: Optional[SchemaPromptSpecTranslationFileInputSourceDict] - """Inlined file source.""" +class Dataset(_common.BaseModel): + """Represents a dataset resource to store prompts.""" - gcs_input_source: Optional[SchemaPromptSpecTranslationGcsInputSourceDict] - """Cloud Storage file source.""" + metadata: Optional[SchemaTextPromptDatasetMetadata] = Field( + default=None, + description="""Required. Additional information about the Dataset.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""Customer-managed encryption key spec for a Dataset. If set, this Dataset and all sub-resources of this Dataset will be secured by this key.""", + ) + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this Dataset was created.""", + ) + data_item_count: Optional[int] = Field( + default=None, + description="""Output only. The number of DataItems in this Dataset. Only apply for non-structured Dataset.""", + ) + description: Optional[str] = Field( + default=None, description="""The description of the Dataset.""" + ) + display_name: Optional[str] = Field( + default=None, + description="""Required. The user-defined name of the Dataset. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", + ) + etag: Optional[str] = Field( + default=None, + description="""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="""The labels with user-defined metadata to organize your Datasets. 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. No more than 64 user labels can be associated with one Dataset (System labels are excluded). See https://goo.gl/xmQnxf for more information and examples of labels. System reserved label keys are prefixed with "aiplatform.googleapis.com/" and are immutable. Following system labels exist for each Dataset: * "aiplatform.googleapis.com/dataset_metadata_schema": output only, its value is the metadata_schema's title.""", + ) + metadata_artifact: Optional[str] = Field( + default=None, + description="""Output only. The resource name of the Artifact that was created in MetadataStore when creating the Dataset. The Artifact resource name pattern is `projects/{project}/locations/{location}/metadataStores/{metadata_store}/artifacts/{artifact}`.""", + ) + metadata_schema_uri: Optional[str] = Field( + default=None, + description="""Required. Points to a YAML file stored on Google Cloud Storage describing additional information about the Dataset. The schema is defined as an OpenAPI 3.0.2 Schema Object. The schema files that can be used here are found in gs://google-cloud-aiplatform/schema/dataset/metadata/.""", + ) + model_reference: Optional[str] = Field( + default=None, + description="""Optional. Reference to the public base model last used by the dataset. Only set for prompt datasets.""", + ) + name: Optional[str] = Field( + default=None, + description="""Output only. Identifier. The resource name of the Dataset. Format: `projects/{project}/locations/{location}/datasets/{dataset}`""", + ) + 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.""" + ) + saved_queries: Optional[list[SavedQuery]] = Field( + default=None, + description="""All SavedQueries belong to the Dataset will be returned in List/Get Dataset response. The annotation_specs field will not be populated except for UI cases which will only use annotation_spec_count. In CreateDataset request, a SavedQuery is created together if this field is set, up to one SavedQuery can be set in CreateDatasetRequest. The SavedQuery should not contain any AnnotationSpec.""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this Dataset was last updated.""", + ) + # TODO(b/448806531): Remove all the overridden _from_response methods once the + # ticket is resolved and published. + @classmethod + def _from_response( + cls: typing.Type["Dataset"], + *, + response: dict[str, object], + kwargs: dict[str, object], + ) -> "Dataset": + """Converts a dictionary response into a Dataset object.""" -SchemaPromptSpecTranslationSentenceFileInputOrDict = Union[ - SchemaPromptSpecTranslationSentenceFileInput, - SchemaPromptSpecTranslationSentenceFileInputDict, -] + response = _camel_key_to_snake(response) + result = super()._from_response(response=response, kwargs=kwargs) + return result -class SchemaPromptSpecTranslationExample(_common.BaseModel): - """The translation example that contains reference sentences from various sources.""" +class DatasetDict(TypedDict, total=False): + """Represents a dataset resource to store prompts.""" - reference_sentence_pair_lists: Optional[ - list[SchemaPromptSpecReferenceSentencePairList] - ] = Field(default=None, description="""The reference sentences from inline text.""") - reference_sentences_file_inputs: Optional[ - list[SchemaPromptSpecTranslationSentenceFileInput] - ] = Field(default=None, description="""The reference sentences from file.""") + metadata: Optional[SchemaTextPromptDatasetMetadataDict] + """Required. Additional information about the Dataset.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for a Dataset. If set, this Dataset and all sub-resources of this Dataset will be secured by this key.""" -class SchemaPromptSpecTranslationExampleDict(TypedDict, total=False): - """The translation example that contains reference sentences from various sources.""" + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this Dataset was created.""" - reference_sentence_pair_lists: Optional[ - list[SchemaPromptSpecReferenceSentencePairListDict] - ] - """The reference sentences from inline text.""" + data_item_count: Optional[int] + """Output only. The number of DataItems in this Dataset. Only apply for non-structured Dataset.""" - reference_sentences_file_inputs: Optional[ - list[SchemaPromptSpecTranslationSentenceFileInputDict] - ] - """The reference sentences from file.""" + description: Optional[str] + """The description of the Dataset.""" + display_name: Optional[str] + """Required. The user-defined name of the Dataset. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" -SchemaPromptSpecTranslationExampleOrDict = Union[ - SchemaPromptSpecTranslationExample, SchemaPromptSpecTranslationExampleDict -] + etag: Optional[str] + """Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""" + labels: Optional[dict[str, str]] + """The labels with user-defined metadata to organize your Datasets. 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. No more than 64 user labels can be associated with one Dataset (System labels are excluded). See https://goo.gl/xmQnxf for more information and examples of labels. System reserved label keys are prefixed with "aiplatform.googleapis.com/" and are immutable. Following system labels exist for each Dataset: * "aiplatform.googleapis.com/dataset_metadata_schema": output only, its value is the metadata_schema's title.""" -class SchemaPromptSpecTranslationOption(_common.BaseModel): - """Optional settings for translation prompt.""" + metadata_artifact: Optional[str] + """Output only. The resource name of the Artifact that was created in MetadataStore when creating the Dataset. The Artifact resource name pattern is `projects/{project}/locations/{location}/metadataStores/{metadata_store}/artifacts/{artifact}`.""" - number_of_shots: Optional[int] = Field( - default=None, description="""How many shots to use.""" - ) + metadata_schema_uri: Optional[str] + """Required. Points to a YAML file stored on Google Cloud Storage describing additional information about the Dataset. The schema is defined as an OpenAPI 3.0.2 Schema Object. The schema files that can be used here are found in gs://google-cloud-aiplatform/schema/dataset/metadata/.""" + model_reference: Optional[str] + """Optional. Reference to the public base model last used by the dataset. Only set for prompt datasets.""" -class SchemaPromptSpecTranslationOptionDict(TypedDict, total=False): - """Optional settings for translation prompt.""" + name: Optional[str] + """Output only. Identifier. The resource name of the Dataset. Format: `projects/{project}/locations/{location}/datasets/{dataset}`""" - number_of_shots: Optional[int] - """How many shots to use.""" + satisfies_pzi: Optional[bool] + """Output only. Reserved for future use.""" + satisfies_pzs: Optional[bool] + """Output only. Reserved for future use.""" -SchemaPromptSpecTranslationOptionOrDict = Union[ - SchemaPromptSpecTranslationOption, SchemaPromptSpecTranslationOptionDict -] + saved_queries: Optional[list[SavedQueryDict]] + """All SavedQueries belong to the Dataset will be returned in List/Get Dataset response. The annotation_specs field will not be populated except for UI cases which will only use annotation_spec_count. In CreateDataset request, a SavedQuery is created together if this field is set, up to one SavedQuery can be set in CreateDatasetRequest. The SavedQuery should not contain any AnnotationSpec.""" + update_time: Optional[datetime.datetime] + """Output only. Timestamp when this Dataset was last updated.""" -class SchemaPromptSpecTranslationPrompt(_common.BaseModel): - """Prompt variation for Translation use case.""" - example: Optional[SchemaPromptSpecTranslationExample] = Field( - default=None, description="""The translation example.""" - ) - option: Optional[SchemaPromptSpecTranslationOption] = Field( - default=None, description="""The translation option.""" - ) - prompt_message: Optional[SchemaPromptSpecPromptMessage] = Field( - default=None, description="""The prompt message.""" - ) - source_language_code: Optional[str] = Field( - default=None, description="""The source language code.""" - ) - target_language_code: Optional[str] = Field( - default=None, description="""The target language code.""" - ) +DatasetOrDict = Union[Dataset, DatasetDict] -class SchemaPromptSpecTranslationPromptDict(TypedDict, total=False): - """Prompt variation for Translation use case.""" - - example: Optional[SchemaPromptSpecTranslationExampleDict] - """The translation example.""" - - option: Optional[SchemaPromptSpecTranslationOptionDict] - """The translation option.""" - - prompt_message: Optional[SchemaPromptSpecPromptMessageDict] - """The prompt message.""" - - source_language_code: Optional[str] - """The source language code.""" - - target_language_code: Optional[str] - """The target language code.""" - - -SchemaPromptSpecTranslationPromptOrDict = Union[ - SchemaPromptSpecTranslationPrompt, SchemaPromptSpecTranslationPromptDict -] - - -class SchemaPromptApiSchema(_common.BaseModel): - """The A2 schema of a prompt.""" - - api_schema_version: Optional[str] = Field( - default=None, - description="""The Schema version that represents changes to the API behavior.""", - ) - executions: Optional[list[SchemaPromptInstancePromptExecution]] = Field( - default=None, - description="""A list of execution instances for constructing a ready-to-use prompt.""", - ) - multimodal_prompt: Optional[SchemaPromptSpecMultimodalPrompt] = Field( - default=None, - description="""Multimodal prompt which embeds preambles to prompt string.""", - ) - structured_prompt: Optional[SchemaPromptSpecStructuredPrompt] = Field( - default=None, - description="""The prompt variation that stores preambles in separate fields.""", - ) - translation_prompt: Optional[SchemaPromptSpecTranslationPrompt] = Field( - default=None, description="""The prompt variation for Translation use case.""" - ) - +class _GetDatasetVersionParameters(_common.BaseModel): + """Parameters for getting a dataset resource to store prompts.""" -class SchemaPromptApiSchemaDict(TypedDict, total=False): - """The A2 schema of a prompt.""" + dataset_id: Optional[str] = Field(default=None, description="""""") + dataset_version_id: Optional[str] = Field(default=None, description="""""") + config: Optional[VertexBaseConfig] = Field(default=None, description="""""") - api_schema_version: Optional[str] - """The Schema version that represents changes to the API behavior.""" - executions: Optional[list[SchemaPromptInstancePromptExecutionDict]] - """A list of execution instances for constructing a ready-to-use prompt.""" +class _GetDatasetVersionParametersDict(TypedDict, total=False): + """Parameters for getting a dataset resource to store prompts.""" - multimodal_prompt: Optional[SchemaPromptSpecMultimodalPromptDict] - """Multimodal prompt which embeds preambles to prompt string.""" + dataset_id: Optional[str] + """""" - structured_prompt: Optional[SchemaPromptSpecStructuredPromptDict] - """The prompt variation that stores preambles in separate fields.""" + dataset_version_id: Optional[str] + """""" - translation_prompt: Optional[SchemaPromptSpecTranslationPromptDict] - """The prompt variation for Translation use case.""" + config: Optional[VertexBaseConfigDict] + """""" -SchemaPromptApiSchemaOrDict = Union[SchemaPromptApiSchema, SchemaPromptApiSchemaDict] +_GetDatasetVersionParametersOrDict = Union[ + _GetDatasetVersionParameters, _GetDatasetVersionParametersDict +] -class SchemaTextPromptDatasetMetadata(_common.BaseModel): - """Represents the text prompt dataset metadata.""" +class DatasetVersion(_common.BaseModel): + """Represents a dataset version resource to store prompts.""" - candidate_count: Optional[int] = Field( - default=None, description="""Number of candidates.""" - ) - gcs_uri: Optional[str] = Field( - default=None, - description="""The Google Cloud Storage URI that stores the prompt data.""", - ) - grounding_config: Optional[SchemaPredictParamsGroundingConfig] = Field( - default=None, description="""Grounding checking configuration.""" - ) - has_prompt_variable: Optional[bool] = Field( - default=None, description="""Whether the prompt dataset has prompt variable.""" - ) - logprobs: Optional[bool] = Field( - default=None, - description="""Whether or not the user has enabled logit probabilities in the model parameters.""", - ) - max_output_tokens: Optional[int] = Field( + metadata: Optional[SchemaTextPromptDatasetMetadata] = Field( default=None, - description="""Value of the maximum number of tokens generated set when the dataset was saved.""", + description="""Required. Output only. Additional information about the DatasetVersion.""", ) - note: Optional[str] = Field( + big_query_dataset_name: Optional[str] = Field( default=None, - description="""User-created prompt note. Note size limit is 2KB.""", + description="""Output only. Name of the associated BigQuery dataset.""", ) - prompt_api_schema: Optional[SchemaPromptApiSchema] = Field( + create_time: Optional[datetime.datetime] = Field( default=None, - description="""The API schema of the prompt to support both UI and SDK usages.""", - ) - prompt_type: Optional[str] = Field( - default=None, description="""Type of the prompt dataset.""" + description="""Output only. Timestamp when this DatasetVersion was created.""", ) - seed_enabled: Optional[bool] = Field( + display_name: Optional[str] = Field( default=None, - description="""Seeding enables model to return a deterministic response on a best effort basis. Determinism isn't guaranteed. This field determines whether or not seeding is enabled.""", - ) - seed_value: Optional[int] = Field( - default=None, description="""The actual value of the seed.""" - ) - stop_sequences: Optional[list[str]] = Field( - default=None, description="""Customized stop sequences.""" + description="""The user-defined name of the DatasetVersion. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", ) - system_instruction: Optional[str] = Field( + etag: Optional[str] = Field( default=None, - description="""The content of the prompt dataset system instruction.""", + description="""Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""", ) - system_instruction_gcs_uri: Optional[str] = Field( + model_reference: Optional[str] = Field( default=None, - description="""The Google Cloud Storage URI that stores the system instruction, starting with gs://.""", + description="""Output only. Reference to the public base model last used by the dataset version. Only set for prompt dataset versions.""", ) - temperature: Optional[float] = Field( + name: Optional[str] = Field( default=None, - description="""Temperature value used for sampling set when the dataset was saved. This value is used to tune the degree of randomness.""", + description="""Output only. Identifier. The resource name of the DatasetVersion. Format: `projects/{project}/locations/{location}/datasets/{dataset}/datasetVersions/{dataset_version}`""", ) - text: Optional[str] = Field( - default=None, description="""The content of the prompt dataset.""" + satisfies_pzi: Optional[bool] = Field( + default=None, description="""Output only. Reserved for future use.""" ) - top_k: Optional[int] = Field( - default=None, - description="""Top K value set when the dataset was saved. This value determines how many candidates with highest probability from the vocab would be selected for each decoding step.""", + satisfies_pzs: Optional[bool] = Field( + default=None, description="""Output only. Reserved for future use.""" ) - top_p: Optional[float] = Field( + update_time: Optional[datetime.datetime] = Field( default=None, - description="""Top P value set when the dataset was saved. Given topK tokens for decoding, top candidates will be selected until the sum of their probabilities is topP.""", + description="""Output only. Timestamp when this DatasetVersion was last updated.""", ) + # TODO(b/448806531): Remove all the overridden _from_response methods once the + # ticket is resolved and published. + @classmethod + def _from_response( + cls: typing.Type["DatasetVersion"], + *, + response: dict[str, object], + kwargs: dict[str, object], + ) -> "DatasetVersion": + """Converts a dictionary response into a DatasetVersion object.""" -class SchemaTextPromptDatasetMetadataDict(TypedDict, total=False): - """Represents the text prompt dataset metadata.""" - - candidate_count: Optional[int] - """Number of candidates.""" - - gcs_uri: Optional[str] - """The Google Cloud Storage URI that stores the prompt data.""" - - grounding_config: Optional[SchemaPredictParamsGroundingConfigDict] - """Grounding checking configuration.""" - - has_prompt_variable: Optional[bool] - """Whether the prompt dataset has prompt variable.""" - - logprobs: Optional[bool] - """Whether or not the user has enabled logit probabilities in the model parameters.""" - - max_output_tokens: Optional[int] - """Value of the maximum number of tokens generated set when the dataset was saved.""" + response = _camel_key_to_snake(response) + result = super()._from_response(response=response, kwargs=kwargs) + return result - note: Optional[str] - """User-created prompt note. Note size limit is 2KB.""" - prompt_api_schema: Optional[SchemaPromptApiSchemaDict] - """The API schema of the prompt to support both UI and SDK usages.""" +class DatasetVersionDict(TypedDict, total=False): + """Represents a dataset version resource to store prompts.""" - prompt_type: Optional[str] - """Type of the prompt dataset.""" + metadata: Optional[SchemaTextPromptDatasetMetadataDict] + """Required. Output only. Additional information about the DatasetVersion.""" - seed_enabled: Optional[bool] - """Seeding enables model to return a deterministic response on a best effort basis. Determinism isn't guaranteed. This field determines whether or not seeding is enabled.""" + big_query_dataset_name: Optional[str] + """Output only. Name of the associated BigQuery dataset.""" - seed_value: Optional[int] - """The actual value of the seed.""" + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this DatasetVersion was created.""" - stop_sequences: Optional[list[str]] - """Customized stop sequences.""" + display_name: Optional[str] + """The user-defined name of the DatasetVersion. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" - system_instruction: Optional[str] - """The content of the prompt dataset system instruction.""" + etag: Optional[str] + """Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""" - system_instruction_gcs_uri: Optional[str] - """The Google Cloud Storage URI that stores the system instruction, starting with gs://.""" + model_reference: Optional[str] + """Output only. Reference to the public base model last used by the dataset version. Only set for prompt dataset versions.""" - temperature: Optional[float] - """Temperature value used for sampling set when the dataset was saved. This value is used to tune the degree of randomness.""" + name: Optional[str] + """Output only. Identifier. The resource name of the DatasetVersion. Format: `projects/{project}/locations/{location}/datasets/{dataset}/datasetVersions/{dataset_version}`""" - text: Optional[str] - """The content of the prompt dataset.""" + satisfies_pzi: Optional[bool] + """Output only. Reserved for future use.""" - top_k: Optional[int] - """Top K value set when the dataset was saved. This value determines how many candidates with highest probability from the vocab would be selected for each decoding step.""" + satisfies_pzs: Optional[bool] + """Output only. Reserved for future use.""" - top_p: Optional[float] - """Top P value set when the dataset was saved. Given topK tokens for decoding, top candidates will be selected until the sum of their probabilities is topP.""" + update_time: Optional[datetime.datetime] + """Output only. Timestamp when this DatasetVersion was last updated.""" -SchemaTextPromptDatasetMetadataOrDict = Union[ - SchemaTextPromptDatasetMetadata, SchemaTextPromptDatasetMetadataDict -] +DatasetVersionOrDict = Union[DatasetVersion, DatasetVersionDict] -class CreateDatasetConfig(_common.BaseModel): - """Config for creating a dataset resource to store prompts.""" +class GetDatasetOperationConfig(_common.BaseModel): + """Config for getting a dataset version operation.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class CreateDatasetConfigDict(TypedDict, total=False): - """Config for creating a dataset resource to store prompts.""" +class GetDatasetOperationConfigDict(TypedDict, total=False): + """Config for getting a dataset version operation.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -CreateDatasetConfigOrDict = Union[CreateDatasetConfig, CreateDatasetConfigDict] +GetDatasetOperationConfigOrDict = Union[ + GetDatasetOperationConfig, GetDatasetOperationConfigDict +] -class _CreateDatasetParameters(_common.BaseModel): - """Parameters for creating a dataset resource to store prompts.""" +class _GetDatasetOperationParameters(_common.BaseModel): + """Parameters for getting a dataset operation.""" - name: Optional[str] = Field(default=None, description="""""") - display_name: Optional[str] = Field(default=None, description="""""") - metadata_schema_uri: Optional[str] = Field(default=None, description="""""") - metadata: Optional[SchemaTextPromptDatasetMetadata] = Field( + dataset_id: Optional[str] = Field(default=None, description="""""") + operation_id: Optional[str] = Field(default=None, description="""""") + config: Optional[GetDatasetOperationConfig] = Field( default=None, description="""""" ) - description: Optional[str] = Field(default=None, description="""""") - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, description="""""" - ) - model_reference: Optional[str] = Field(default=None, description="""""") - config: Optional[CreateDatasetConfig] = Field(default=None, description="""""") - - -class _CreateDatasetParametersDict(TypedDict, total=False): - """Parameters for creating a dataset resource to store prompts.""" - - name: Optional[str] - """""" - - display_name: Optional[str] - """""" - - metadata_schema_uri: Optional[str] - """""" - metadata: Optional[SchemaTextPromptDatasetMetadataDict] - """""" - description: Optional[str] - """""" +class _GetDatasetOperationParametersDict(TypedDict, total=False): + """Parameters for getting a dataset operation.""" - encryption_spec: Optional[genai_types.EncryptionSpec] + dataset_id: Optional[str] """""" - model_reference: Optional[str] + operation_id: Optional[str] """""" - config: Optional[CreateDatasetConfigDict] + config: Optional[GetDatasetOperationConfigDict] """""" -_CreateDatasetParametersOrDict = Union[ - _CreateDatasetParameters, _CreateDatasetParametersDict +_GetDatasetOperationParametersOrDict = Union[ + _GetDatasetOperationParameters, _GetDatasetOperationParametersDict ] -class DatasetOperation(_common.BaseModel): - """Represents the create dataset operation.""" +class ListPromptsConfig(_common.BaseModel): + """Config for listing prompt datasets and dataset versions.""" - 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.""", - ) - response: Optional[dict[str, Any]] = Field( - default=None, description="""The result of the dataset operation.""" + description="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""", ) -class DatasetOperationDict(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}`.""" +class ListPromptsConfigDict(TypedDict, total=False): + """Config for listing prompt datasets and dataset versions.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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_size: Optional[int] + """""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + page_token: Optional[str] + """""" - response: Optional[dict[str, Any]] - """The result of the dataset operation.""" + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""" -DatasetOperationOrDict = Union[DatasetOperation, DatasetOperationDict] +ListPromptsConfigOrDict = Union[ListPromptsConfig, ListPromptsConfigDict] -class CreateDatasetVersionConfig(_common.BaseModel): - """Config for creating a dataset version resource to store prompts.""" +class _ListDatasetsRequestParameters(_common.BaseModel): + """Parameters for listing prompt datasets.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + config: Optional[ListPromptsConfig] = Field(default=None, description="""""") -class CreateDatasetVersionConfigDict(TypedDict, total=False): - """Config for creating a dataset version resource to store prompts.""" +class _ListDatasetsRequestParametersDict(TypedDict, total=False): + """Parameters for listing prompt datasets.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + config: Optional[ListPromptsConfigDict] + """""" -CreateDatasetVersionConfigOrDict = Union[ - CreateDatasetVersionConfig, CreateDatasetVersionConfigDict +_ListDatasetsRequestParametersOrDict = Union[ + _ListDatasetsRequestParameters, _ListDatasetsRequestParametersDict ] -class _CreateDatasetVersionParameters(_common.BaseModel): - """Represents the create dataset version parameters.""" +class ListDatasetsResponse(_common.BaseModel): + """Response for listing prompt datasets.""" - dataset_name: Optional[str] = Field(default=None, description="""""") - metadata: Optional[SchemaTextPromptDatasetMetadata] = Field( - default=None, description="""""" + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - model_reference: Optional[str] = Field(default=None, description="""""") - parent: Optional[str] = Field(default=None, description="""""") - display_name: Optional[str] = Field(default=None, description="""""") - config: Optional[CreateDatasetVersionConfig] = Field( - default=None, description="""""" + next_page_token: Optional[str] = Field(default=None, description="""""") + datasets: Optional[list[Dataset]] = Field( + default=None, + description="""List of datasets for the project. + """, ) + # TODO(b/448806531): Remove all the overridden _from_response methods once the + # ticket is resolved and published. + @classmethod + def _from_response( + cls: typing.Type["ListDatasetsResponse"], + *, + response: dict[str, object], + kwargs: dict[str, object], + ) -> "ListDatasetsResponse": + """Converts a dictionary response into a ListDatasetsResponse object.""" -class _CreateDatasetVersionParametersDict(TypedDict, total=False): - """Represents the create dataset version parameters.""" + response = _camel_key_to_snake(response) + result = super()._from_response(response=response, kwargs=kwargs) + return result - dataset_name: Optional[str] - """""" - metadata: Optional[SchemaTextPromptDatasetMetadataDict] - """""" +class ListDatasetsResponseDict(TypedDict, total=False): + """Response for listing prompt datasets.""" - model_reference: Optional[str] - """""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - parent: Optional[str] + next_page_token: Optional[str] """""" - display_name: Optional[str] - """""" + datasets: Optional[list[DatasetDict]] + """List of datasets for the project. + """ - config: Optional[CreateDatasetVersionConfigDict] - """""" +ListDatasetsResponseOrDict = Union[ListDatasetsResponse, ListDatasetsResponseDict] -_CreateDatasetVersionParametersOrDict = Union[ - _CreateDatasetVersionParameters, _CreateDatasetVersionParametersDict -] +class _ListDatasetVersionsRequestParameters(_common.BaseModel): + """Parameters for listing dataset versions.""" -class _GetDatasetParameters(_common.BaseModel): - """Parameters for getting a dataset resource to store prompts.""" + read_mask: Optional[str] = Field(default=None, description="""""") + dataset_id: Optional[str] = Field(default=None, description="""""") + config: Optional[ListPromptsConfig] = Field(default=None, description="""""") - name: Optional[str] = Field(default=None, description="""""") - config: Optional[VertexBaseConfig] = Field(default=None, description="""""") +class _ListDatasetVersionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing dataset versions.""" -class _GetDatasetParametersDict(TypedDict, total=False): - """Parameters for getting a dataset resource to store prompts.""" + read_mask: Optional[str] + """""" - name: Optional[str] + dataset_id: Optional[str] """""" - config: Optional[VertexBaseConfigDict] + config: Optional[ListPromptsConfigDict] """""" -_GetDatasetParametersOrDict = Union[_GetDatasetParameters, _GetDatasetParametersDict] +_ListDatasetVersionsRequestParametersOrDict = Union[ + _ListDatasetVersionsRequestParameters, _ListDatasetVersionsRequestParametersDict +] -class SavedQuery(_common.BaseModel): - """A SavedQuery is a view of the dataset. It references a subset of annotations by problem type and filters.""" +class ListDatasetVersionsResponse(_common.BaseModel): + """Response for listing prompt datasets.""" - annotation_filter: Optional[str] = Field( - default=None, - description="""Output only. Filters on the Annotations in the dataset.""", - ) - annotation_spec_count: Optional[int] = Field( - default=None, - description="""Output only. Number of AnnotationSpecs in the context of the SavedQuery.""", - ) - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when this SavedQuery was created.""", - ) - display_name: Optional[str] = Field( - default=None, - description="""Required. The user-defined name of the SavedQuery. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", - ) - etag: Optional[str] = Field( - default=None, - description="""Used to perform a consistent read-modify-write update. If not set, a blind "overwrite" update happens.""", - ) - metadata: Optional[Any] = Field( - default=None, - description="""Some additional information about the SavedQuery.""", - ) - name: Optional[str] = Field( - default=None, description="""Output only. Resource name of the SavedQuery.""" - ) - problem_type: Optional[str] = Field( - default=None, - description="""Required. Problem type of the SavedQuery. Allowed values: * IMAGE_CLASSIFICATION_SINGLE_LABEL * IMAGE_CLASSIFICATION_MULTI_LABEL * IMAGE_BOUNDING_POLY * IMAGE_BOUNDING_BOX * TEXT_CLASSIFICATION_SINGLE_LABEL * TEXT_CLASSIFICATION_MULTI_LABEL * TEXT_EXTRACTION * TEXT_SENTIMENT * VIDEO_CLASSIFICATION * VIDEO_OBJECT_TRACKING""", - ) - support_automl_training: Optional[bool] = Field( - default=None, - description="""Output only. If the Annotations belonging to the SavedQuery can be used for AutoML training.""", + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - update_time: Optional[datetime.datetime] = Field( + next_page_token: Optional[str] = Field(default=None, description="""""") + dataset_versions: Optional[list[DatasetVersion]] = Field( default=None, - description="""Output only. Timestamp when SavedQuery was last updated.""", + description="""List of datasets for the project. + """, ) + # TODO(b/448806531): Remove all the overridden _from_response methods once the + # ticket is resolved and published. + @classmethod + def _from_response( + cls: typing.Type["ListDatasetVersionsResponse"], + *, + response: dict[str, object], + kwargs: dict[str, object], + ) -> "ListDatasetVersionsResponse": + """Converts a dictionary response into a ListDatasetVersionsResponse object.""" -class SavedQueryDict(TypedDict, total=False): - """A SavedQuery is a view of the dataset. It references a subset of annotations by problem type and filters.""" + response = _camel_key_to_snake(response) + result = super()._from_response(response=response, kwargs=kwargs) + return result - annotation_filter: Optional[str] - """Output only. Filters on the Annotations in the dataset.""" - annotation_spec_count: Optional[int] - """Output only. Number of AnnotationSpecs in the context of the SavedQuery.""" +class ListDatasetVersionsResponseDict(TypedDict, total=False): + """Response for listing prompt datasets.""" - create_time: Optional[datetime.datetime] - """Output only. Timestamp when this SavedQuery was created.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - display_name: Optional[str] - """Required. The user-defined name of the SavedQuery. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" + next_page_token: Optional[str] + """""" - etag: Optional[str] - """Used to perform a consistent read-modify-write update. If not set, a blind "overwrite" update happens.""" + dataset_versions: Optional[list[DatasetVersionDict]] + """List of datasets for the project. + """ - metadata: Optional[Any] - """Some additional information about the SavedQuery.""" - name: Optional[str] - """Output only. Resource name of the SavedQuery.""" +ListDatasetVersionsResponseOrDict = Union[ + ListDatasetVersionsResponse, ListDatasetVersionsResponseDict +] - problem_type: Optional[str] - """Required. Problem type of the SavedQuery. Allowed values: * IMAGE_CLASSIFICATION_SINGLE_LABEL * IMAGE_CLASSIFICATION_MULTI_LABEL * IMAGE_BOUNDING_POLY * IMAGE_BOUNDING_BOX * TEXT_CLASSIFICATION_SINGLE_LABEL * TEXT_CLASSIFICATION_MULTI_LABEL * TEXT_EXTRACTION * TEXT_SENTIMENT * VIDEO_CLASSIFICATION * VIDEO_OBJECT_TRACKING""" - support_automl_training: Optional[bool] - """Output only. If the Annotations belonging to the SavedQuery can be used for AutoML training.""" +class DeletePromptConfig(_common.BaseModel): + """Config for deleting a prompt.""" - update_time: Optional[datetime.datetime] - """Output only. Timestamp when SavedQuery was last updated.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + timeout: Optional[int] = Field( + default=90, + description="""Timeout for the delete prompt operation in seconds. Defaults to 90.""", + ) + max_wait_time: Optional[int] = Field( + default=60, + description="""Maximum interval between polling requests in seconds. Defaults to 60.""", + ) -SavedQueryOrDict = Union[SavedQuery, SavedQueryDict] +class DeletePromptConfigDict(TypedDict, total=False): + """Config for deleting a prompt.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -class Dataset(_common.BaseModel): - """Represents a dataset resource to store prompts.""" + timeout: Optional[int] + """Timeout for the delete prompt operation in seconds. Defaults to 90.""" - metadata: Optional[SchemaTextPromptDatasetMetadata] = Field( - default=None, - description="""Required. Additional information about the Dataset.""", - ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, - description="""Customer-managed encryption key spec for a Dataset. If set, this Dataset and all sub-resources of this Dataset will be secured by this key.""", - ) - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when this Dataset was created.""", - ) - data_item_count: Optional[int] = Field( - default=None, - description="""Output only. The number of DataItems in this Dataset. Only apply for non-structured Dataset.""", - ) - description: Optional[str] = Field( - default=None, description="""The description of the Dataset.""" - ) - display_name: Optional[str] = Field( - default=None, - description="""Required. The user-defined name of the Dataset. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", - ) - etag: Optional[str] = Field( - default=None, - description="""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="""The labels with user-defined metadata to organize your Datasets. 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. No more than 64 user labels can be associated with one Dataset (System labels are excluded). See https://goo.gl/xmQnxf for more information and examples of labels. System reserved label keys are prefixed with "aiplatform.googleapis.com/" and are immutable. Following system labels exist for each Dataset: * "aiplatform.googleapis.com/dataset_metadata_schema": output only, its value is the metadata_schema's title.""", - ) - metadata_artifact: Optional[str] = Field( - default=None, - description="""Output only. The resource name of the Artifact that was created in MetadataStore when creating the Dataset. The Artifact resource name pattern is `projects/{project}/locations/{location}/metadataStores/{metadata_store}/artifacts/{artifact}`.""", - ) - metadata_schema_uri: Optional[str] = Field( - default=None, - description="""Required. Points to a YAML file stored on Google Cloud Storage describing additional information about the Dataset. The schema is defined as an OpenAPI 3.0.2 Schema Object. The schema files that can be used here are found in gs://google-cloud-aiplatform/schema/dataset/metadata/.""", - ) - model_reference: Optional[str] = Field( - default=None, - description="""Optional. Reference to the public base model last used by the dataset. Only set for prompt datasets.""", - ) - name: Optional[str] = Field( - default=None, - description="""Output only. Identifier. The resource name of the Dataset. Format: `projects/{project}/locations/{location}/datasets/{dataset}`""", - ) - 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.""" - ) - saved_queries: Optional[list[SavedQuery]] = Field( - default=None, - description="""All SavedQueries belong to the Dataset will be returned in List/Get Dataset response. The annotation_specs field will not be populated except for UI cases which will only use annotation_spec_count. In CreateDataset request, a SavedQuery is created together if this field is set, up to one SavedQuery can be set in CreateDatasetRequest. The SavedQuery should not contain any AnnotationSpec.""", - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when this Dataset was last updated.""", - ) + max_wait_time: Optional[int] + """Maximum interval between polling requests in seconds. Defaults to 60.""" - # TODO(b/448806531): Remove all the overridden _from_response methods once the - # ticket is resolved and published. - @classmethod - def _from_response( - cls: typing.Type["Dataset"], - *, - response: dict[str, object], - kwargs: dict[str, object], - ) -> "Dataset": - """Converts a dictionary response into a Dataset object.""" - response = _camel_key_to_snake(response) - result = super()._from_response(response=response, kwargs=kwargs) - return result +DeletePromptConfigOrDict = Union[DeletePromptConfig, DeletePromptConfigDict] -class DatasetDict(TypedDict, total=False): - """Represents a dataset resource to store prompts.""" +class _DeleteDatasetRequestParameters(_common.BaseModel): + """Parameters for deleting a prompt dataset.""" - metadata: Optional[SchemaTextPromptDatasetMetadataDict] - """Required. Additional information about the Dataset.""" + prompt_id: Optional[str] = Field( + default=None, description="""ID of the prompt dataset to be deleted.""" + ) + config: Optional[DeletePromptConfig] = Field(default=None, description="""""") - encryption_spec: Optional[genai_types.EncryptionSpec] - """Customer-managed encryption key spec for a Dataset. If set, this Dataset and all sub-resources of this Dataset will be secured by this key.""" - create_time: Optional[datetime.datetime] - """Output only. Timestamp when this Dataset was created.""" +class _DeleteDatasetRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a prompt dataset.""" - data_item_count: Optional[int] - """Output only. The number of DataItems in this Dataset. Only apply for non-structured Dataset.""" + prompt_id: Optional[str] + """ID of the prompt dataset to be deleted.""" - description: Optional[str] - """The description of the Dataset.""" + config: Optional[DeletePromptConfigDict] + """""" - display_name: Optional[str] - """Required. The user-defined name of the Dataset. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" - etag: Optional[str] - """Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""" +_DeleteDatasetRequestParametersOrDict = Union[ + _DeleteDatasetRequestParameters, _DeleteDatasetRequestParametersDict +] - labels: Optional[dict[str, str]] - """The labels with user-defined metadata to organize your Datasets. 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. No more than 64 user labels can be associated with one Dataset (System labels are excluded). See https://goo.gl/xmQnxf for more information and examples of labels. System reserved label keys are prefixed with "aiplatform.googleapis.com/" and are immutable. Following system labels exist for each Dataset: * "aiplatform.googleapis.com/dataset_metadata_schema": output only, its value is the metadata_schema's title.""" - metadata_artifact: Optional[str] - """Output only. The resource name of the Artifact that was created in MetadataStore when creating the Dataset. The Artifact resource name pattern is `projects/{project}/locations/{location}/metadataStores/{metadata_store}/artifacts/{artifact}`.""" +class DeletePromptOperation(_common.BaseModel): + """Operation for deleting prompts.""" - metadata_schema_uri: Optional[str] - """Required. Points to a YAML file stored on Google Cloud Storage describing additional information about the Dataset. The schema is defined as an OpenAPI 3.0.2 Schema Object. The schema files that can be used here are found in gs://google-cloud-aiplatform/schema/dataset/metadata/.""" + 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.""", + ) - model_reference: Optional[str] - """Optional. Reference to the public base model last used by the dataset. Only set for prompt datasets.""" - name: Optional[str] - """Output only. Identifier. The resource name of the Dataset. Format: `projects/{project}/locations/{location}/datasets/{dataset}`""" +class DeletePromptOperationDict(TypedDict, total=False): + """Operation for deleting prompts.""" - satisfies_pzi: Optional[bool] - """Output only. Reserved for future use.""" + 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}`.""" - satisfies_pzs: Optional[bool] - """Output only. Reserved for future use.""" + 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.""" - saved_queries: Optional[list[SavedQueryDict]] - """All SavedQueries belong to the Dataset will be returned in List/Get Dataset response. The annotation_specs field will not be populated except for UI cases which will only use annotation_spec_count. In CreateDataset request, a SavedQuery is created together if this field is set, up to one SavedQuery can be set in CreateDatasetRequest. The SavedQuery should not contain any AnnotationSpec.""" + 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.""" - update_time: Optional[datetime.datetime] - """Output only. Timestamp when this Dataset was last updated.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -DatasetOrDict = Union[Dataset, DatasetDict] +DeletePromptOperationOrDict = Union[DeletePromptOperation, DeletePromptOperationDict] -class _GetDatasetVersionParameters(_common.BaseModel): - """Parameters for getting a dataset resource to store prompts.""" +class _DeletePromptVersionRequestParameters(_common.BaseModel): + """Parameters for deleting a prompt version.""" - dataset_id: Optional[str] = Field(default=None, description="""""") - dataset_version_id: Optional[str] = Field(default=None, description="""""") - config: Optional[VertexBaseConfig] = Field(default=None, description="""""") + prompt_id: Optional[str] = Field( + default=None, description="""ID of the prompt to be deleted.""" + ) + version_id: Optional[str] = Field( + default=None, + description="""ID of the prompt version to be deleted within the provided prompt_id.""", + ) + config: Optional[DeletePromptConfig] = Field(default=None, description="""""") -class _GetDatasetVersionParametersDict(TypedDict, total=False): - """Parameters for getting a dataset resource to store prompts.""" +class _DeletePromptVersionRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a prompt version.""" - dataset_id: Optional[str] - """""" + prompt_id: Optional[str] + """ID of the prompt to be deleted.""" - dataset_version_id: Optional[str] - """""" + version_id: Optional[str] + """ID of the prompt version to be deleted within the provided prompt_id.""" - config: Optional[VertexBaseConfigDict] + config: Optional[DeletePromptConfigDict] """""" -_GetDatasetVersionParametersOrDict = Union[ - _GetDatasetVersionParameters, _GetDatasetVersionParametersDict +_DeletePromptVersionRequestParametersOrDict = Union[ + _DeletePromptVersionRequestParameters, _DeletePromptVersionRequestParametersDict ] -class DatasetVersion(_common.BaseModel): - """Represents a dataset version resource to store prompts.""" +class DeletePromptVersionOperation(_common.BaseModel): + """Operation for deleting prompt versions.""" - metadata: Optional[SchemaTextPromptDatasetMetadata] = Field( + name: Optional[str] = Field( default=None, - description="""Required. Output only. Additional information about the DatasetVersion.""", + 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}`.""", ) - big_query_dataset_name: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Output only. Name of the associated BigQuery dataset.""", + 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.""", ) - create_time: Optional[datetime.datetime] = Field( + done: Optional[bool] = Field( default=None, - description="""Output only. Timestamp when this DatasetVersion was created.""", + 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.""", ) - display_name: Optional[str] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""The user-defined name of the DatasetVersion. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", - ) - etag: Optional[str] = Field( - default=None, - description="""Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""", - ) - model_reference: Optional[str] = Field( - default=None, - description="""Output only. Reference to the public base model last used by the dataset version. Only set for prompt dataset versions.""", - ) - name: Optional[str] = Field( - default=None, - description="""Output only. Identifier. The resource name of the DatasetVersion. Format: `projects/{project}/locations/{location}/datasets/{dataset}/datasetVersions/{dataset_version}`""", - ) - 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.""" - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when this DatasetVersion was last updated.""", + description="""The error result of the operation in case of failure or cancellation.""", ) - # TODO(b/448806531): Remove all the overridden _from_response methods once the - # ticket is resolved and published. - @classmethod - def _from_response( - cls: typing.Type["DatasetVersion"], - *, - response: dict[str, object], - kwargs: dict[str, object], - ) -> "DatasetVersion": - """Converts a dictionary response into a DatasetVersion object.""" - response = _camel_key_to_snake(response) - result = super()._from_response(response=response, kwargs=kwargs) - return result +class DeletePromptVersionOperationDict(TypedDict, total=False): + """Operation for deleting prompt versions.""" + 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 DatasetVersionDict(TypedDict, total=False): - """Represents a dataset version resource to store prompts.""" + 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.""" - metadata: Optional[SchemaTextPromptDatasetMetadataDict] - """Required. Output only. Additional information about the DatasetVersion.""" + 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.""" - big_query_dataset_name: Optional[str] - """Output only. Name of the associated BigQuery dataset.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - create_time: Optional[datetime.datetime] - """Output only. Timestamp when this DatasetVersion was created.""" - display_name: Optional[str] - """The user-defined name of the DatasetVersion. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" +DeletePromptVersionOperationOrDict = Union[ + DeletePromptVersionOperation, DeletePromptVersionOperationDict +] - etag: Optional[str] - """Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""" - model_reference: Optional[str] - """Output only. Reference to the public base model last used by the dataset version. Only set for prompt dataset versions.""" +class RestoreVersionConfig(_common.BaseModel): + """Config for restoring a prompt version.""" - name: Optional[str] - """Output only. Identifier. The resource name of the DatasetVersion. Format: `projects/{project}/locations/{location}/datasets/{dataset}/datasetVersions/{dataset_version}`""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + timeout: Optional[int] = Field( + default=90, + description="""Timeout for the restore prompt version operation in seconds. Defaults to 90.""", + ) + max_wait_time: Optional[int] = Field( + default=60, + description="""Maximum interval between polling requests in seconds. Defaults to 60.""", + ) - satisfies_pzi: Optional[bool] - """Output only. Reserved for future use.""" - satisfies_pzs: Optional[bool] - """Output only. Reserved for future use.""" +class RestoreVersionConfigDict(TypedDict, total=False): + """Config for restoring a prompt version.""" - update_time: Optional[datetime.datetime] - """Output only. Timestamp when this DatasetVersion was last updated.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + timeout: Optional[int] + """Timeout for the restore prompt version operation in seconds. Defaults to 90.""" -DatasetVersionOrDict = Union[DatasetVersion, DatasetVersionDict] + max_wait_time: Optional[int] + """Maximum interval between polling requests in seconds. Defaults to 60.""" -class GetDatasetOperationConfig(_common.BaseModel): - """Config for getting a dataset version operation.""" +RestoreVersionConfigOrDict = Union[RestoreVersionConfig, RestoreVersionConfigDict] - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + +class _RestoreVersionRequestParameters(_common.BaseModel): + """Parameters for restoring a prompt version.""" + + dataset_id: Optional[str] = Field( + default=None, description="""ID of the prompt dataset to be restored.""" ) + version_id: Optional[str] = Field( + default=None, description="""ID of the prompt dataset version to be restored.""" + ) + config: Optional[RestoreVersionConfig] = Field(default=None, description="""""") -class GetDatasetOperationConfigDict(TypedDict, total=False): - """Config for getting a dataset version operation.""" +class _RestoreVersionRequestParametersDict(TypedDict, total=False): + """Parameters for restoring a prompt version.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + dataset_id: Optional[str] + """ID of the prompt dataset to be restored.""" + version_id: Optional[str] + """ID of the prompt dataset version to be restored.""" -GetDatasetOperationConfigOrDict = Union[ - GetDatasetOperationConfig, GetDatasetOperationConfigDict + config: Optional[RestoreVersionConfigDict] + """""" + + +_RestoreVersionRequestParametersOrDict = Union[ + _RestoreVersionRequestParameters, _RestoreVersionRequestParametersDict ] -class _GetDatasetOperationParameters(_common.BaseModel): - """Parameters for getting a dataset operation.""" +class RestoreVersionOperation(_common.BaseModel): + """Represents the restore version operation.""" - dataset_id: Optional[str] = Field(default=None, description="""""") - operation_id: Optional[str] = Field(default=None, description="""""") - config: Optional[GetDatasetOperationConfig] = Field( - default=None, description="""""" + 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.""", ) -class _GetDatasetOperationParametersDict(TypedDict, total=False): - """Parameters for getting a dataset operation.""" +class RestoreVersionOperationDict(TypedDict, total=False): + """Represents the restore version operation.""" - dataset_id: Optional[str] - """""" + 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}`.""" - operation_id: 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.""" - config: Optional[GetDatasetOperationConfigDict] - """""" + 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.""" -_GetDatasetOperationParametersOrDict = Union[ - _GetDatasetOperationParameters, _GetDatasetOperationParametersDict + +RestoreVersionOperationOrDict = Union[ + RestoreVersionOperation, RestoreVersionOperationDict ] -class ListPromptsConfig(_common.BaseModel): - """Config for listing prompt datasets and dataset versions.""" +class UpdatePromptConfig(_common.BaseModel): + """Config for creating a dataset resource to store prompts.""" 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( + prompt_display_name: Optional[str] = Field( + default=None, description="""The updated display name for the prompt.""" + ) + version_display_name: 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.""", + description="""The updated display name for the prompt version. If not set, a default name with a timestamp will be used.""", + ) + timeout: Optional[int] = Field( + default=90, + description="""The timeout for the update_dataset_resource request in seconds. If not set, the default timeout is 90 seconds.""", + ) + 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 polling requests in seconds. If not set, the default interval is 60 seconds.""", ) -class ListPromptsConfigDict(TypedDict, total=False): - """Config for listing prompt datasets and dataset versions.""" +class UpdatePromptConfigDict(TypedDict, total=False): + """Config for creating a dataset resource to store prompts.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - page_size: Optional[int] - """""" + prompt_display_name: Optional[str] + """The updated display name for the prompt.""" - page_token: Optional[str] - """""" + version_display_name: Optional[str] + """The updated display name for the prompt version. If not set, a default name with a timestamp will be used.""" - 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 update_dataset_resource 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.""" -ListPromptsConfigOrDict = Union[ListPromptsConfig, ListPromptsConfigDict] + max_wait_time: Optional[int] + """The maximum interval between polling requests in seconds. If not set, the default interval is 60 seconds.""" -class _ListDatasetsRequestParameters(_common.BaseModel): - """Parameters for listing prompt datasets.""" +UpdatePromptConfigOrDict = Union[UpdatePromptConfig, UpdatePromptConfigDict] - config: Optional[ListPromptsConfig] = Field(default=None, description="""""") +class _UpdateDatasetParameters(_common.BaseModel): + """Parameters for creating a dataset resource to store prompts.""" -class _ListDatasetsRequestParametersDict(TypedDict, total=False): - """Parameters for listing prompt datasets.""" + name: Optional[str] = Field(default=None, description="""""") + dataset_id: Optional[str] = Field(default=None, description="""""") + display_name: Optional[str] = Field(default=None, description="""""") + metadata: Optional[SchemaTextPromptDatasetMetadata] = Field( + default=None, description="""""" + ) + description: Optional[str] = Field(default=None, description="""""") + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, description="""""" + ) + model_reference: Optional[str] = Field(default=None, description="""""") + config: Optional[UpdatePromptConfig] = Field(default=None, description="""""") - config: Optional[ListPromptsConfigDict] + +class _UpdateDatasetParametersDict(TypedDict, total=False): + """Parameters for creating a dataset resource to store prompts.""" + + name: Optional[str] """""" + dataset_id: Optional[str] + """""" -_ListDatasetsRequestParametersOrDict = Union[ - _ListDatasetsRequestParameters, _ListDatasetsRequestParametersDict -] + display_name: Optional[str] + """""" + metadata: Optional[SchemaTextPromptDatasetMetadataDict] + """""" -class ListDatasetsResponse(_common.BaseModel): - """Response for listing prompt datasets.""" + description: Optional[str] + """""" - 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="""""") - datasets: Optional[list[Dataset]] = Field( - default=None, - description="""List of datasets for the project. - """, - ) + encryption_spec: Optional[genai_types.EncryptionSpec] + """""" - # TODO(b/448806531): Remove all the overridden _from_response methods once the - # ticket is resolved and published. - @classmethod - def _from_response( - cls: typing.Type["ListDatasetsResponse"], - *, - response: dict[str, object], - kwargs: dict[str, object], - ) -> "ListDatasetsResponse": - """Converts a dictionary response into a ListDatasetsResponse object.""" + model_reference: Optional[str] + """""" - response = _camel_key_to_snake(response) - result = super()._from_response(response=response, kwargs=kwargs) - return result + config: Optional[UpdatePromptConfigDict] + """""" -class ListDatasetsResponseDict(TypedDict, total=False): - """Response for listing prompt datasets.""" +_UpdateDatasetParametersOrDict = Union[ + _UpdateDatasetParameters, _UpdateDatasetParametersDict +] - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" - next_page_token: Optional[str] - """""" +class GetSkillConfig(_common.BaseModel): + """Config for getting a skill.""" - datasets: Optional[list[DatasetDict]] - """List of datasets for the project. - """ + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -ListDatasetsResponseOrDict = Union[ListDatasetsResponse, ListDatasetsResponseDict] +class GetSkillConfigDict(TypedDict, total=False): + """Config for getting a skill.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -class _ListDatasetVersionsRequestParameters(_common.BaseModel): - """Parameters for listing dataset versions.""" - read_mask: Optional[str] = Field(default=None, description="""""") - dataset_id: Optional[str] = Field(default=None, description="""""") - config: Optional[ListPromptsConfig] = Field(default=None, description="""""") +GetSkillConfigOrDict = Union[GetSkillConfig, GetSkillConfigDict] -class _ListDatasetVersionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing dataset versions.""" +class _GetSkillRequestParameters(_common.BaseModel): + """Parameters for GetSkillRequest.""" - read_mask: Optional[str] - """""" + name: Optional[str] = Field( + default=None, + description="""The resource name of the Skill to retrieve. Format: projects/{project}/locations/{location}/skills/{skill}""", + ) + config: Optional[GetSkillConfig] = Field(default=None, description="""""") - dataset_id: Optional[str] - """""" - config: Optional[ListPromptsConfigDict] +class _GetSkillRequestParametersDict(TypedDict, total=False): + """Parameters for GetSkillRequest.""" + + name: Optional[str] + """The resource name of the Skill to retrieve. Format: projects/{project}/locations/{location}/skills/{skill}""" + + config: Optional[GetSkillConfigDict] """""" -_ListDatasetVersionsRequestParametersOrDict = Union[ - _ListDatasetVersionsRequestParameters, _ListDatasetVersionsRequestParametersDict +_GetSkillRequestParametersOrDict = Union[ + _GetSkillRequestParameters, _GetSkillRequestParametersDict ] -class ListDatasetVersionsResponse(_common.BaseModel): - """Response for listing prompt datasets.""" +class Skill(_common.BaseModel): + """Represents a Skill resource. - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + Patches the type from the discovery document. + """ + + name: Optional[str] = Field( + default=None, + description="""Identifier. The resource name of the Skill. Format: `projects/{project}/locations/{location}/skills/{skill}`""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - dataset_versions: Optional[list[DatasetVersion]] = Field( + create_time: Optional[datetime.datetime] = Field( default=None, - description="""List of datasets for the project. - """, + description="""Output only. Timestamp when this Skill was created.""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this Skill was most recently updated.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""Required. Provides the display name of the Skill. This should align with `name` in the `SKILL.md` file.""", + ) + description: Optional[str] = Field( + default=None, + description="""Required. Describes the Skill. Should describe both what the skill does and when to use it. Should include specific keywords that help agents identify relevant tasks. This should align with `description` in the `SKILL.md` file.""", + ) + license: Optional[str] = Field( + default=None, + description="""Optional. Specifies the license of the Skill. This should be an SPDX license identifier (e.g., "MIT", "Apache-2.0"). See https://spdx.org/licenses/. This should align with `license` in the `SKILL.md` file.""", + ) + compatibility: Optional[str] = Field( + default=None, + description="""Optional. Specifies the compatibility of the Skill. Indicates environment requirements (intended product, system packages, network access, etc.). This should align with `compatibility` in the `SKILL.md` file.""", + ) + zipped_filesystem: Optional[str] = Field( + default=None, + description="""Required. Provides the zipped filesystem of the Skill. This should contain the `SKILL.md` file at the root of the zip and optional directories for scripts, references, and assets. Directory should align with the directory structure specified at https://agentskills.io/specification#directory-structure.""", + ) + state: Optional[SkillState] = Field( + default=None, description="""Output only. The state of the Skill.""" + ) + labels: Optional[dict[str, str]] = Field( + default=None, + description="""The labels with user-defined metadata to organize Skills.""", + ) + sha256: Optional[str] = Field( + default=None, + description="""Output only. The SHA256 checksum of the zipped filesystem.""", + ) + skill_source: Optional[SkillSource] = Field( + default=None, description="""Output only. The source of the Skill.""" ) - # TODO(b/448806531): Remove all the overridden _from_response methods once the - # ticket is resolved and published. - @classmethod - def _from_response( - cls: typing.Type["ListDatasetVersionsResponse"], - *, - response: dict[str, object], - kwargs: dict[str, object], - ) -> "ListDatasetVersionsResponse": - """Converts a dictionary response into a ListDatasetVersionsResponse object.""" - response = _camel_key_to_snake(response) - result = super()._from_response(response=response, kwargs=kwargs) - return result +class SkillDict(TypedDict, total=False): + """Represents a Skill resource. + Patches the type from the discovery document. + """ -class ListDatasetVersionsResponseDict(TypedDict, total=False): - """Response for listing prompt datasets.""" + name: Optional[str] + """Identifier. The resource name of the Skill. Format: `projects/{project}/locations/{location}/skills/{skill}`""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this Skill was created.""" - next_page_token: Optional[str] - """""" + update_time: Optional[datetime.datetime] + """Output only. Timestamp when this Skill was most recently updated.""" - dataset_versions: Optional[list[DatasetVersionDict]] - """List of datasets for the project. - """ + display_name: Optional[str] + """Required. Provides the display name of the Skill. This should align with `name` in the `SKILL.md` file.""" + description: Optional[str] + """Required. Describes the Skill. Should describe both what the skill does and when to use it. Should include specific keywords that help agents identify relevant tasks. This should align with `description` in the `SKILL.md` file.""" -ListDatasetVersionsResponseOrDict = Union[ - ListDatasetVersionsResponse, ListDatasetVersionsResponseDict -] + license: Optional[str] + """Optional. Specifies the license of the Skill. This should be an SPDX license identifier (e.g., "MIT", "Apache-2.0"). See https://spdx.org/licenses/. This should align with `license` in the `SKILL.md` file.""" + compatibility: Optional[str] + """Optional. Specifies the compatibility of the Skill. Indicates environment requirements (intended product, system packages, network access, etc.). This should align with `compatibility` in the `SKILL.md` file.""" -class DeletePromptConfig(_common.BaseModel): - """Config for deleting a prompt.""" + zipped_filesystem: Optional[str] + """Required. Provides the zipped filesystem of the Skill. This should contain the `SKILL.md` file at the root of the zip and optional directories for scripts, references, and assets. Directory should align with the directory structure specified at https://agentskills.io/specification#directory-structure.""" + + state: Optional[SkillState] + """Output only. The state of the Skill.""" + + labels: Optional[dict[str, str]] + """The labels with user-defined metadata to organize Skills.""" + + sha256: Optional[str] + """Output only. The SHA256 checksum of the zipped filesystem.""" + + skill_source: Optional[SkillSource] + """Output only. The source of the Skill.""" + + +SkillOrDict = Union[Skill, SkillDict] + + +class RetrieveSkillsConfig(_common.BaseModel): + """Config for retrieving skills.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - timeout: Optional[int] = Field( - default=90, - description="""Timeout for the delete prompt operation in seconds. Defaults to 90.""", - ) - max_wait_time: Optional[int] = Field( - default=60, - description="""Maximum interval between polling requests in seconds. Defaults to 60.""", + top_k: Optional[int] = Field( + default=None, + description="""Optional. The maximum number of skills to return. The service may + return fewer than this value. If unspecified, at most 10 skills will be + returned. The maximum value is 100. + """, ) -class DeletePromptConfigDict(TypedDict, total=False): - """Config for deleting a prompt.""" +class RetrieveSkillsConfigDict(TypedDict, total=False): + """Config for retrieving skills.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - timeout: Optional[int] - """Timeout for the delete prompt operation in seconds. Defaults to 90.""" - - max_wait_time: Optional[int] - """Maximum interval between polling requests in seconds. Defaults to 60.""" + top_k: Optional[int] + """Optional. The maximum number of skills to return. The service may + return fewer than this value. If unspecified, at most 10 skills will be + returned. The maximum value is 100. + """ -DeletePromptConfigOrDict = Union[DeletePromptConfig, DeletePromptConfigDict] +RetrieveSkillsConfigOrDict = Union[RetrieveSkillsConfig, RetrieveSkillsConfigDict] -class _DeleteDatasetRequestParameters(_common.BaseModel): - """Parameters for deleting a prompt dataset.""" +class _RetrieveSkillsRequestParameters(_common.BaseModel): + """Parameters for retrieving skills.""" - prompt_id: Optional[str] = Field( - default=None, description="""ID of the prompt dataset to be deleted.""" + query: Optional[str] = Field( + default=None, description="""Required. The query to find matching skills.""" ) - config: Optional[DeletePromptConfig] = Field(default=None, description="""""") + config: Optional[RetrieveSkillsConfig] = Field(default=None, description="""""") -class _DeleteDatasetRequestParametersDict(TypedDict, total=False): - """Parameters for deleting a prompt dataset.""" +class _RetrieveSkillsRequestParametersDict(TypedDict, total=False): + """Parameters for retrieving skills.""" - prompt_id: Optional[str] - """ID of the prompt dataset to be deleted.""" + query: Optional[str] + """Required. The query to find matching skills.""" - config: Optional[DeletePromptConfigDict] + config: Optional[RetrieveSkillsConfigDict] """""" -_DeleteDatasetRequestParametersOrDict = Union[ - _DeleteDatasetRequestParameters, _DeleteDatasetRequestParametersDict +_RetrieveSkillsRequestParametersOrDict = Union[ + _RetrieveSkillsRequestParameters, _RetrieveSkillsRequestParametersDict ] -class DeletePromptOperation(_common.BaseModel): - """Operation for deleting prompts.""" - - 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.""", - ) - - -class DeletePromptOperationDict(TypedDict, total=False): - """Operation for deleting prompts.""" - - 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.""" - - -DeletePromptOperationOrDict = Union[DeletePromptOperation, DeletePromptOperationDict] - - -class _DeletePromptVersionRequestParameters(_common.BaseModel): - """Parameters for deleting a prompt version.""" +class RetrievedSkill(_common.BaseModel): + """A retrieved skill from semantic search.""" - prompt_id: Optional[str] = Field( - default=None, description="""ID of the prompt to be deleted.""" + skill_name: Optional[str] = Field( + default=None, description="""The resource name of the skill.""" ) - version_id: Optional[str] = Field( - default=None, - description="""ID of the prompt version to be deleted within the provided prompt_id.""", + description: Optional[str] = Field( + default=None, description="""The description of the skill.""" ) - config: Optional[DeletePromptConfig] = Field(default=None, description="""""") - -class _DeletePromptVersionRequestParametersDict(TypedDict, total=False): - """Parameters for deleting a prompt version.""" - prompt_id: Optional[str] - """ID of the prompt to be deleted.""" +class RetrievedSkillDict(TypedDict, total=False): + """A retrieved skill from semantic search.""" - version_id: Optional[str] - """ID of the prompt version to be deleted within the provided prompt_id.""" + skill_name: Optional[str] + """The resource name of the skill.""" - config: Optional[DeletePromptConfigDict] - """""" + description: Optional[str] + """The description of the skill.""" -_DeletePromptVersionRequestParametersOrDict = Union[ - _DeletePromptVersionRequestParameters, _DeletePromptVersionRequestParametersDict -] +RetrievedSkillOrDict = Union[RetrievedSkill, RetrievedSkillDict] -class DeletePromptVersionOperation(_common.BaseModel): - """Operation for deleting prompt versions.""" +class RetrieveSkillsResponse(_common.BaseModel): + """Response for retrieving skills.""" - 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.""", + retrieved_skills: Optional[list[RetrievedSkill]] = Field( + default=None, description="""List of retrieved skills ranked by similarity.""" ) -class DeletePromptVersionOperationDict(TypedDict, total=False): - """Operation for deleting prompt versions.""" - - 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 RetrieveSkillsResponseDict(TypedDict, total=False): + """Response for retrieving skills.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + retrieved_skills: Optional[list[RetrievedSkillDict]] + """List of retrieved skills ranked by similarity.""" -DeletePromptVersionOperationOrDict = Union[ - DeletePromptVersionOperation, DeletePromptVersionOperationDict -] +RetrieveSkillsResponseOrDict = Union[RetrieveSkillsResponse, RetrieveSkillsResponseDict] -class RestoreVersionConfig(_common.BaseModel): - """Config for restoring a prompt version.""" +class CreateSkillConfig(_common.BaseModel): + """Config for creating a skill.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - timeout: Optional[int] = Field( - default=90, - description="""Timeout for the restore prompt version operation in seconds. Defaults to 90.""", + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to wait for the long running operation to complete.""", ) - max_wait_time: Optional[int] = Field( - default=60, - description="""Maximum interval between polling requests in seconds. Defaults to 60.""", + local_path: Optional[str] = Field( + default=None, + description="""Optional. The local path to the directory containing the Skill to + be zipped and uploaded. + """, + ) + zipped_filesystem: Optional[Any] = Field( + default=None, description="""Optional. The zipped filesystem of the Skill.""" ) -class RestoreVersionConfigDict(TypedDict, total=False): - """Config for restoring a prompt version.""" +class CreateSkillConfigDict(TypedDict, total=False): + """Config for creating a skill.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - timeout: Optional[int] - """Timeout for the restore prompt version operation in seconds. Defaults to 90.""" + wait_for_completion: Optional[bool] + """Whether to wait for the long running operation to complete.""" - max_wait_time: Optional[int] - """Maximum interval between polling requests in seconds. Defaults to 60.""" + local_path: Optional[str] + """Optional. The local path to the directory containing the Skill to + be zipped and uploaded. + """ + zipped_filesystem: Optional[Any] + """Optional. The zipped filesystem of the Skill.""" -RestoreVersionConfigOrDict = Union[RestoreVersionConfig, RestoreVersionConfigDict] +CreateSkillConfigOrDict = Union[CreateSkillConfig, CreateSkillConfigDict] -class _RestoreVersionRequestParameters(_common.BaseModel): - """Parameters for restoring a prompt version.""" - dataset_id: Optional[str] = Field( - default=None, description="""ID of the prompt dataset to be restored.""" +class _CreateSkillRequestParameters(_common.BaseModel): + """Parameters for creating a skill.""" + + display_name: Optional[str] = Field( + default=None, description="""Required. The display name of the Skill.""" ) - version_id: Optional[str] = Field( - default=None, description="""ID of the prompt dataset version to be restored.""" + description: Optional[str] = Field( + default=None, description="""Required. The description of the Skill.""" + ) + config: Optional[CreateSkillConfig] = Field(default=None, description="""""") + skill_id: Optional[str] = Field( + default=None, + description="""Required. The ID to use for the Skill, which will become the final + component of the Skill's resource name. + """, ) - config: Optional[RestoreVersionConfig] = Field(default=None, description="""""") -class _RestoreVersionRequestParametersDict(TypedDict, total=False): - """Parameters for restoring a prompt version.""" +class _CreateSkillRequestParametersDict(TypedDict, total=False): + """Parameters for creating a skill.""" - dataset_id: Optional[str] - """ID of the prompt dataset to be restored.""" + display_name: Optional[str] + """Required. The display name of the Skill.""" - version_id: Optional[str] - """ID of the prompt dataset version to be restored.""" + description: Optional[str] + """Required. The description of the Skill.""" - config: Optional[RestoreVersionConfigDict] + config: Optional[CreateSkillConfigDict] """""" + skill_id: Optional[str] + """Required. The ID to use for the Skill, which will become the final + component of the Skill's resource name. + """ + -_RestoreVersionRequestParametersOrDict = Union[ - _RestoreVersionRequestParameters, _RestoreVersionRequestParametersDict +_CreateSkillRequestParametersOrDict = Union[ + _CreateSkillRequestParameters, _CreateSkillRequestParametersDict ] -class RestoreVersionOperation(_common.BaseModel): - """Represents the restore version operation.""" +class SkillOperation(_common.BaseModel): + """Operation that has a skill as a response.""" name: Optional[str] = Field( default=None, @@ -23432,10 +22717,13 @@ class RestoreVersionOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) + response: Optional[Skill] = Field( + default=None, description="""The created Skill.""" + ) -class RestoreVersionOperationDict(TypedDict, total=False): - """Represents the restore version operation.""" +class SkillOperationDict(TypedDict, total=False): + """Operation that has a skill 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}`.""" @@ -23449,570 +22737,477 @@ class RestoreVersionOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" + response: Optional[SkillDict] + """The created Skill.""" -RestoreVersionOperationOrDict = Union[ - RestoreVersionOperation, RestoreVersionOperationDict -] +SkillOperationOrDict = Union[SkillOperation, SkillOperationDict] -class UpdatePromptConfig(_common.BaseModel): - """Config for creating a dataset resource to store prompts.""" + +class UpdateSkillConfig(_common.BaseModel): + """Config for updating a skill.""" 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 updated display name for the prompt.""" + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to wait for the long running operation to complete.""", ) - version_display_name: Optional[str] = Field( + local_path: Optional[str] = Field( default=None, - description="""The updated display name for the prompt version. If not set, a default name with a timestamp will be used.""", + description="""Optional. The local path to the directory containing the Skill to + be zipped and uploaded. + """, ) - timeout: Optional[int] = Field( - default=90, - description="""The timeout for the update_dataset_resource request in seconds. If not set, the default timeout is 90 seconds.""", + display_name: Optional[str] = Field( + default=None, description="""Optional. The display name of the Skill.""" ) - 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.""", + description: Optional[str] = Field( + default=None, description="""Optional. The description of the Skill.""" ) - max_wait_time: Optional[int] = Field( - default=60, - description="""The maximum interval between polling requests in seconds. If not set, the default interval is 60 seconds.""", + zipped_filesystem: Optional[Any] = Field( + default=None, description="""Optional. The zipped filesystem of the Skill.""" + ) + update_mask: Optional[str] = Field( + default=None, description="""Optional. The update mask to apply.""" ) -class UpdatePromptConfigDict(TypedDict, total=False): - """Config for creating a dataset resource to store prompts.""" +class UpdateSkillConfigDict(TypedDict, total=False): + """Config for updating a skill.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - prompt_display_name: Optional[str] - """The updated display name for the prompt.""" + wait_for_completion: Optional[bool] + """Whether to wait for the long running operation to complete.""" - version_display_name: Optional[str] - """The updated display name for the prompt version. If not set, a default name with a timestamp will be used.""" + local_path: Optional[str] + """Optional. The local path to the directory containing the Skill to + be zipped and uploaded. + """ - timeout: Optional[int] - """The timeout for the update_dataset_resource request in seconds. If not set, the default timeout is 90 seconds.""" + display_name: Optional[str] + """Optional. The display name of the Skill.""" - 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.""" + description: Optional[str] + """Optional. The description of the Skill.""" - max_wait_time: Optional[int] - """The maximum interval between polling requests in seconds. If not set, the default interval is 60 seconds.""" + zipped_filesystem: Optional[Any] + """Optional. The zipped filesystem of the Skill.""" + update_mask: Optional[str] + """Optional. The update mask to apply.""" -UpdatePromptConfigOrDict = Union[UpdatePromptConfig, UpdatePromptConfigDict] +UpdateSkillConfigOrDict = Union[UpdateSkillConfig, UpdateSkillConfigDict] -class _UpdateDatasetParameters(_common.BaseModel): - """Parameters for creating a dataset resource to store prompts.""" - name: Optional[str] = Field(default=None, description="""""") - dataset_id: Optional[str] = Field(default=None, description="""""") - display_name: Optional[str] = Field(default=None, description="""""") - metadata: Optional[SchemaTextPromptDatasetMetadata] = Field( - default=None, description="""""" - ) - description: Optional[str] = Field(default=None, description="""""") - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, description="""""" +class _UpdateSkillRequestParameters(_common.BaseModel): + """Parameters for updating a skill.""" + + name: Optional[str] = Field( + default=None, + description="""Required. The resource name of the Skill to update.""", ) - model_reference: Optional[str] = Field(default=None, description="""""") - config: Optional[UpdatePromptConfig] = Field(default=None, description="""""") + config: Optional[UpdateSkillConfig] = Field(default=None, description="""""") -class _UpdateDatasetParametersDict(TypedDict, total=False): - """Parameters for creating a dataset resource to store prompts.""" +class _UpdateSkillRequestParametersDict(TypedDict, total=False): + """Parameters for updating a skill.""" name: Optional[str] - """""" + """Required. The resource name of the Skill to update.""" - dataset_id: Optional[str] + config: Optional[UpdateSkillConfigDict] """""" - display_name: Optional[str] - """""" - metadata: Optional[SchemaTextPromptDatasetMetadataDict] - """""" +_UpdateSkillRequestParametersOrDict = Union[ + _UpdateSkillRequestParameters, _UpdateSkillRequestParametersDict +] - description: Optional[str] - """""" - encryption_spec: Optional[genai_types.EncryptionSpec] +class ListSkillsConfig(_common.BaseModel): + """Config for listing skills.""" + + 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="""Optional. The standard list filter.""" + ) + + +class ListSkillsConfigDict(TypedDict, total=False): + """Config for listing skills.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + page_size: Optional[int] """""" - model_reference: Optional[str] + page_token: Optional[str] """""" - config: Optional[UpdatePromptConfigDict] + filter: Optional[str] + """Optional. The standard list filter.""" + + +ListSkillsConfigOrDict = Union[ListSkillsConfig, ListSkillsConfigDict] + + +class _ListSkillsRequestParameters(_common.BaseModel): + """Parameters for listing skills.""" + + config: Optional[ListSkillsConfig] = Field(default=None, description="""""") + + +class _ListSkillsRequestParametersDict(TypedDict, total=False): + """Parameters for listing skills.""" + + config: Optional[ListSkillsConfigDict] """""" -_UpdateDatasetParametersOrDict = Union[ - _UpdateDatasetParameters, _UpdateDatasetParametersDict +_ListSkillsRequestParametersOrDict = Union[ + _ListSkillsRequestParameters, _ListSkillsRequestParametersDict ] -class GetSkillConfig(_common.BaseModel): - """Config for getting a skill.""" +class ListSkillsResponse(_common.BaseModel): + """Response for listing skills.""" + + 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="""""") + skills: Optional[list[Skill]] = Field( + default=None, description="""List of Skills.""" + ) + + +class ListSkillsResponseDict(TypedDict, total=False): + """Response for listing skills.""" + + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + + next_page_token: Optional[str] + """""" + + skills: Optional[list[SkillDict]] + """List of Skills.""" + + +ListSkillsResponseOrDict = Union[ListSkillsResponse, ListSkillsResponseDict] + + +class DeleteSkillConfig(_common.BaseModel): + """Config for deleting a skill.""" 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="""Whether to wait for the long running operation to complete.""", + ) -class GetSkillConfigDict(TypedDict, total=False): - """Config for getting a skill.""" +class DeleteSkillConfigDict(TypedDict, total=False): + """Config for deleting a skill.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" + wait_for_completion: Optional[bool] + """Whether to wait for the long running operation to complete.""" -GetSkillConfigOrDict = Union[GetSkillConfig, GetSkillConfigDict] + +DeleteSkillConfigOrDict = Union[DeleteSkillConfig, DeleteSkillConfigDict] -class _GetSkillRequestParameters(_common.BaseModel): - """Parameters for GetSkillRequest.""" +class _DeleteSkillRequestParameters(_common.BaseModel): + """Parameters for deleting a skill.""" name: Optional[str] = Field( default=None, - description="""The resource name of the Skill to retrieve. Format: projects/{project}/locations/{location}/skills/{skill}""", + description="""Required. The resource name of the Skill to delete.""", ) - config: Optional[GetSkillConfig] = Field(default=None, description="""""") + config: Optional[DeleteSkillConfig] = Field(default=None, description="""""") -class _GetSkillRequestParametersDict(TypedDict, total=False): - """Parameters for GetSkillRequest.""" +class _DeleteSkillRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a skill.""" name: Optional[str] - """The resource name of the Skill to retrieve. Format: projects/{project}/locations/{location}/skills/{skill}""" + """Required. The resource name of the Skill to delete.""" - config: Optional[GetSkillConfigDict] + config: Optional[DeleteSkillConfigDict] """""" -_GetSkillRequestParametersOrDict = Union[ - _GetSkillRequestParameters, _GetSkillRequestParametersDict +_DeleteSkillRequestParametersOrDict = Union[ + _DeleteSkillRequestParameters, _DeleteSkillRequestParametersDict ] -class Skill(_common.BaseModel): - """Represents a Skill resource. - - Patches the type from the discovery document. - """ +class DeleteSkillOperation(_common.BaseModel): + """Operation for deleting a skill.""" name: Optional[str] = Field( default=None, - description="""Identifier. The resource name of the Skill. Format: `projects/{project}/locations/{location}/skills/{skill}`""", - ) - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when this Skill was created.""", - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Timestamp when this Skill was most recently updated.""", - ) - display_name: Optional[str] = Field( - default=None, - description="""Required. Provides the display name of the Skill. This should align with `name` in the `SKILL.md` file.""", - ) - description: Optional[str] = Field( - default=None, - description="""Required. Describes the Skill. Should describe both what the skill does and when to use it. Should include specific keywords that help agents identify relevant tasks. This should align with `description` in the `SKILL.md` file.""", - ) - license: Optional[str] = Field( - default=None, - description="""Optional. Specifies the license of the Skill. This should be an SPDX license identifier (e.g., "MIT", "Apache-2.0"). See https://spdx.org/licenses/. This should align with `license` in the `SKILL.md` file.""", - ) - compatibility: Optional[str] = Field( - default=None, - description="""Optional. Specifies the compatibility of the Skill. Indicates environment requirements (intended product, system packages, network access, etc.). This should align with `compatibility` in the `SKILL.md` file.""", + 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}`.""", ) - zipped_filesystem: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Required. Provides the zipped filesystem of the Skill. This should contain the `SKILL.md` file at the root of the zip and optional directories for scripts, references, and assets. Directory should align with the directory structure specified at https://agentskills.io/specification#directory-structure.""", - ) - state: Optional[SkillState] = Field( - default=None, description="""Output only. The state of the Skill.""" + 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.""", ) - labels: Optional[dict[str, str]] = Field( + done: Optional[bool] = Field( default=None, - description="""The labels with user-defined metadata to organize Skills.""", + 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.""", ) - sha256: Optional[str] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Output only. The SHA256 checksum of the zipped filesystem.""", - ) - skill_source: Optional[SkillSource] = Field( - default=None, description="""Output only. The source of the Skill.""" + description="""The error result of the operation in case of failure or cancellation.""", ) -class SkillDict(TypedDict, total=False): - """Represents a Skill resource. - - Patches the type from the discovery document. - """ +class DeleteSkillOperationDict(TypedDict, total=False): + """Operation for deleting a skill.""" name: Optional[str] - """Identifier. The resource name of the Skill. Format: `projects/{project}/locations/{location}/skills/{skill}`""" + """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 Skill 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.""" - update_time: Optional[datetime.datetime] - """Output only. Timestamp when this Skill was most recently updated.""" + 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.""" - display_name: Optional[str] - """Required. Provides the display name of the Skill. This should align with `name` in the `SKILL.md` file.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - description: Optional[str] - """Required. Describes the Skill. Should describe both what the skill does and when to use it. Should include specific keywords that help agents identify relevant tasks. This should align with `description` in the `SKILL.md` file.""" - license: Optional[str] - """Optional. Specifies the license of the Skill. This should be an SPDX license identifier (e.g., "MIT", "Apache-2.0"). See https://spdx.org/licenses/. This should align with `license` in the `SKILL.md` file.""" +DeleteSkillOperationOrDict = Union[DeleteSkillOperation, DeleteSkillOperationDict] - compatibility: Optional[str] - """Optional. Specifies the compatibility of the Skill. Indicates environment requirements (intended product, system packages, network access, etc.). This should align with `compatibility` in the `SKILL.md` file.""" - zipped_filesystem: Optional[str] - """Required. Provides the zipped filesystem of the Skill. This should contain the `SKILL.md` file at the root of the zip and optional directories for scripts, references, and assets. Directory should align with the directory structure specified at https://agentskills.io/specification#directory-structure.""" +class GetSkillOperationConfig(_common.BaseModel): - state: Optional[SkillState] - """Output only. The state of the Skill.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) - labels: Optional[dict[str, str]] - """The labels with user-defined metadata to organize Skills.""" - - sha256: Optional[str] - """Output only. The SHA256 checksum of the zipped filesystem.""" - - skill_source: Optional[SkillSource] - """Output only. The source of the Skill.""" - - -SkillOrDict = Union[Skill, SkillDict] - - -class RetrieveSkillsConfig(_common.BaseModel): - """Config for retrieving skills.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - top_k: Optional[int] = Field( - default=None, - description="""Optional. The maximum number of skills to return. The service may - return fewer than this value. If unspecified, at most 10 skills will be - returned. The maximum value is 100. - """, - ) - -class RetrieveSkillsConfigDict(TypedDict, total=False): - """Config for retrieving skills.""" +class GetSkillOperationConfigDict(TypedDict, total=False): http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - top_k: Optional[int] - """Optional. The maximum number of skills to return. The service may - return fewer than this value. If unspecified, at most 10 skills will be - returned. The maximum value is 100. - """ - -RetrieveSkillsConfigOrDict = Union[RetrieveSkillsConfig, RetrieveSkillsConfigDict] +GetSkillOperationConfigOrDict = Union[ + GetSkillOperationConfig, GetSkillOperationConfigDict +] -class _RetrieveSkillsRequestParameters(_common.BaseModel): - """Parameters for retrieving skills.""" +class _GetSkillOperationParameters(_common.BaseModel): + """Parameters for getting an operation.""" - query: Optional[str] = Field( - default=None, description="""Required. The query to find matching skills.""" + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetSkillOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) - config: Optional[RetrieveSkillsConfig] = Field(default=None, description="""""") -class _RetrieveSkillsRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving skills.""" +class _GetSkillOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation.""" - query: Optional[str] - """Required. The query to find matching skills.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - config: Optional[RetrieveSkillsConfigDict] - """""" + config: Optional[GetSkillOperationConfigDict] + """Used to override the default configuration.""" -_RetrieveSkillsRequestParametersOrDict = Union[ - _RetrieveSkillsRequestParameters, _RetrieveSkillsRequestParametersDict +_GetSkillOperationParametersOrDict = Union[ + _GetSkillOperationParameters, _GetSkillOperationParametersDict ] -class RetrievedSkill(_common.BaseModel): - """A retrieved skill from semantic search.""" +class GetSkillRevisionConfig(_common.BaseModel): + """Configuration for getting a Skill Revision.""" - skill_name: Optional[str] = Field( - default=None, description="""The resource name of the skill.""" - ) - description: Optional[str] = Field( - default=None, description="""The description of the skill.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class RetrievedSkillDict(TypedDict, total=False): - """A retrieved skill from semantic search.""" - - skill_name: Optional[str] - """The resource name of the skill.""" +class GetSkillRevisionConfigDict(TypedDict, total=False): + """Configuration for getting a Skill Revision.""" - description: Optional[str] - """The description of the skill.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -RetrievedSkillOrDict = Union[RetrievedSkill, RetrievedSkillDict] +GetSkillRevisionConfigOrDict = Union[GetSkillRevisionConfig, GetSkillRevisionConfigDict] -class RetrieveSkillsResponse(_common.BaseModel): - """Response for retrieving skills.""" +class _GetSkillRevisionRequestParameters(_common.BaseModel): + """Parameters for getting a Skill Revision.""" - retrieved_skills: Optional[list[RetrievedSkill]] = Field( - default=None, description="""List of retrieved skills ranked by similarity.""" + name: Optional[str] = Field( + default=None, + description="""The resource name of the Skill Revision to retrieve. Format: projects/{project}/locations/{location}/skills/{skill}/revisions/{revision}""", ) + config: Optional[GetSkillRevisionConfig] = Field(default=None, description="""""") -class RetrieveSkillsResponseDict(TypedDict, total=False): - """Response for retrieving skills.""" +class _GetSkillRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for getting a Skill Revision.""" - retrieved_skills: Optional[list[RetrievedSkillDict]] - """List of retrieved skills ranked by similarity.""" + name: Optional[str] + """The resource name of the Skill Revision to retrieve. Format: projects/{project}/locations/{location}/skills/{skill}/revisions/{revision}""" + + config: Optional[GetSkillRevisionConfigDict] + """""" -RetrieveSkillsResponseOrDict = Union[RetrieveSkillsResponse, RetrieveSkillsResponseDict] +_GetSkillRevisionRequestParametersOrDict = Union[ + _GetSkillRevisionRequestParameters, _GetSkillRevisionRequestParametersDict +] -class CreateSkillConfig(_common.BaseModel): - """Config for creating a skill.""" +class SkillRevision(_common.BaseModel): + """A single revision of a Skill.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + name: Optional[str] = Field( + default=None, + description="""Identifier. The resource name of the Skill Revision. Format: `projects/{project}/locations/{location}/skills/{skill}/revisions/{revision}`""", ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Whether to wait for the long running operation to complete.""", + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this Skill Revision was created.""", ) - local_path: Optional[str] = Field( + skill: Optional[Skill] = Field( default=None, - description="""Optional. The local path to the directory containing the Skill to - be zipped and uploaded. - """, + description="""Output only. The state of the Skill at this revision.""", ) - zipped_filesystem: Optional[Any] = Field( - default=None, description="""Optional. The zipped filesystem of the Skill.""" + state: Optional[SkillRevisionState] = Field( + default=None, description="""Output only. The state of the Skill Revision.""" ) -class CreateSkillConfigDict(TypedDict, total=False): - """Config for creating a skill.""" +class SkillRevisionDict(TypedDict, total=False): + """A single revision of a Skill.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + name: Optional[str] + """Identifier. The resource name of the Skill Revision. Format: `projects/{project}/locations/{location}/skills/{skill}/revisions/{revision}`""" - wait_for_completion: Optional[bool] - """Whether to wait for the long running operation to complete.""" + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this Skill Revision was created.""" - local_path: Optional[str] - """Optional. The local path to the directory containing the Skill to - be zipped and uploaded. - """ + skill: Optional[SkillDict] + """Output only. The state of the Skill at this revision.""" - zipped_filesystem: Optional[Any] - """Optional. The zipped filesystem of the Skill.""" + state: Optional[SkillRevisionState] + """Output only. The state of the Skill Revision.""" -CreateSkillConfigOrDict = Union[CreateSkillConfig, CreateSkillConfigDict] +SkillRevisionOrDict = Union[SkillRevision, SkillRevisionDict] -class _CreateSkillRequestParameters(_common.BaseModel): - """Parameters for creating a skill.""" +class ListSkillRevisionsConfig(_common.BaseModel): + """Configuration for listing Skill Revisions.""" - display_name: Optional[str] = Field( - default=None, description="""Required. The display name of the Skill.""" - ) - description: Optional[str] = Field( - default=None, description="""Required. The description of the Skill.""" - ) - config: Optional[CreateSkillConfig] = Field(default=None, description="""""") - skill_id: Optional[str] = Field( - default=None, - description="""Required. The ID to use for the Skill, which will become the final - component of the Skill's resource name. - """, + 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="""""") -class _CreateSkillRequestParametersDict(TypedDict, total=False): - """Parameters for creating a skill.""" - - display_name: Optional[str] - """Required. The display name of the Skill.""" +class ListSkillRevisionsConfigDict(TypedDict, total=False): + """Configuration for listing Skill Revisions.""" - description: Optional[str] - """Required. The description of the Skill.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - config: Optional[CreateSkillConfigDict] + page_size: Optional[int] """""" - skill_id: Optional[str] - """Required. The ID to use for the Skill, which will become the final - component of the Skill's resource name. - """ + page_token: Optional[str] + """""" -_CreateSkillRequestParametersOrDict = Union[ - _CreateSkillRequestParameters, _CreateSkillRequestParametersDict +ListSkillRevisionsConfigOrDict = Union[ + ListSkillRevisionsConfig, ListSkillRevisionsConfigDict ] -class SkillOperation(_common.BaseModel): - """Operation that has a skill as a response.""" +class _ListSkillRevisionsRequestParameters(_common.BaseModel): + """Parameters for ListSkillRevisionsRequest.""" 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[Skill] = Field( - default=None, description="""The created Skill.""" + description="""Required. The name of the Skill to list revisions for.""", ) + config: Optional[ListSkillRevisionsConfig] = Field(default=None, description="""""") -class SkillOperationDict(TypedDict, total=False): - """Operation that has a skill as a response.""" +class _ListSkillRevisionsRequestParametersDict(TypedDict, total=False): + """Parameters for ListSkillRevisionsRequest.""" 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.""" + """Required. The name of the Skill to list revisions for.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + config: Optional[ListSkillRevisionsConfigDict] + """""" - response: Optional[SkillDict] - """The created Skill.""" +_ListSkillRevisionsRequestParametersOrDict = Union[ + _ListSkillRevisionsRequestParameters, _ListSkillRevisionsRequestParametersDict +] -SkillOperationOrDict = Union[SkillOperation, SkillOperationDict] +class ListSkillRevisionsResponse(_common.BaseModel): + """Response for listing Skill Revisions.""" -class UpdateSkillConfig(_common.BaseModel): - """Config for updating a skill.""" - - 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="""Whether to wait for the long running operation to complete.""", - ) - local_path: Optional[str] = Field( - default=None, - description="""Optional. The local path to the directory containing the Skill to - be zipped and uploaded. - """, - ) - display_name: Optional[str] = Field( - default=None, description="""Optional. The display name of the Skill.""" - ) - description: Optional[str] = Field( - default=None, description="""Optional. The description of the Skill.""" - ) - zipped_filesystem: Optional[Any] = Field( - default=None, description="""Optional. The zipped filesystem of the Skill.""" - ) - update_mask: Optional[str] = Field( - default=None, description="""Optional. The update mask to apply.""" + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - - -class UpdateSkillConfigDict(TypedDict, total=False): - """Config for updating a skill.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - wait_for_completion: Optional[bool] - """Whether to wait for the long running operation to complete.""" - - local_path: Optional[str] - """Optional. The local path to the directory containing the Skill to - be zipped and uploaded. - """ - - display_name: Optional[str] - """Optional. The display name of the Skill.""" - - description: Optional[str] - """Optional. The description of the Skill.""" - - zipped_filesystem: Optional[Any] - """Optional. The zipped filesystem of the Skill.""" - - update_mask: Optional[str] - """Optional. The update mask to apply.""" - - -UpdateSkillConfigOrDict = Union[UpdateSkillConfig, UpdateSkillConfigDict] - - -class _UpdateSkillRequestParameters(_common.BaseModel): - """Parameters for updating a skill.""" - - name: Optional[str] = Field( - default=None, - description="""Required. The resource name of the Skill to update.""", + next_page_token: Optional[str] = Field(default=None, description="""""") + skill_revisions: Optional[list[SkillRevision]] = Field( + default=None, description="""List of Skill Revisions.""" ) - config: Optional[UpdateSkillConfig] = Field(default=None, description="""""") -class _UpdateSkillRequestParametersDict(TypedDict, total=False): - """Parameters for updating a skill.""" +class ListSkillRevisionsResponseDict(TypedDict, total=False): + """Response for listing Skill Revisions.""" - name: Optional[str] - """Required. The resource name of the Skill to update.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - config: Optional[UpdateSkillConfigDict] + next_page_token: Optional[str] """""" + skill_revisions: Optional[list[SkillRevisionDict]] + """List of Skill Revisions.""" -_UpdateSkillRequestParametersOrDict = Union[ - _UpdateSkillRequestParameters, _UpdateSkillRequestParametersDict + +ListSkillRevisionsResponseOrDict = Union[ + ListSkillRevisionsResponse, ListSkillRevisionsResponseDict ] -class ListSkillsConfig(_common.BaseModel): - """Config for listing skills.""" +class ListPublisherModelsConfig(_common.BaseModel): + """Config for listing publisher models.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -24020,12 +23215,15 @@ class ListSkillsConfig(_common.BaseModel): page_size: Optional[int] = Field(default=None, description="""""") page_token: Optional[str] = Field(default=None, description="""""") filter: Optional[str] = Field( - default=None, description="""Optional. The standard list filter.""" + default=None, description="""Filter string for publisher models.""" + ) + list_all_versions: Optional[bool] = Field( + default=None, description="""Whether to list all versions.""" ) -class ListSkillsConfigDict(TypedDict, total=False): - """Config for listing skills.""" +class ListPublisherModelsConfigDict(TypedDict, total=False): + """Config for listing publisher models.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -24037,2494 +23235,2613 @@ class ListSkillsConfigDict(TypedDict, total=False): """""" filter: Optional[str] - """Optional. The standard list filter.""" + """Filter string for publisher models.""" + list_all_versions: Optional[bool] + """Whether to list all versions.""" -ListSkillsConfigOrDict = Union[ListSkillsConfig, ListSkillsConfigDict] +ListPublisherModelsConfigOrDict = Union[ + ListPublisherModelsConfig, ListPublisherModelsConfigDict +] -class _ListSkillsRequestParameters(_common.BaseModel): - """Parameters for listing skills.""" - config: Optional[ListSkillsConfig] = Field(default=None, description="""""") +class _ListPublisherModelsRequestParameters(_common.BaseModel): + """Parameters for listing publisher models.""" + parent: Optional[str] = Field(default=None, description="""""") + config: Optional[ListPublisherModelsConfig] = Field( + default=None, description="""""" + ) -class _ListSkillsRequestParametersDict(TypedDict, total=False): - """Parameters for listing skills.""" - config: Optional[ListSkillsConfigDict] +class _ListPublisherModelsRequestParametersDict(TypedDict, total=False): + """Parameters for listing publisher models.""" + + parent: Optional[str] """""" + config: Optional[ListPublisherModelsConfigDict] + """""" -_ListSkillsRequestParametersOrDict = Union[ - _ListSkillsRequestParameters, _ListSkillsRequestParametersDict + +_ListPublisherModelsRequestParametersOrDict = Union[ + _ListPublisherModelsRequestParameters, _ListPublisherModelsRequestParametersDict ] -class ListSkillsResponse(_common.BaseModel): - """Response for listing skills.""" +class PublisherModelResourceReference(_common.BaseModel): + """Reference to a resource.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + description: Optional[str] = Field( + default=None, description="""Description of the resource.""" ) - next_page_token: Optional[str] = Field(default=None, description="""""") - skills: Optional[list[Skill]] = Field( - default=None, description="""List of Skills.""" + resource_name: Optional[str] = Field( + default=None, description="""The resource name of the Google Cloud resource.""" + ) + uri: Optional[str] = Field(default=None, description="""The URI of the resource.""") + use_case: Optional[str] = Field( + default=None, description="""Use case (CUJ) of the resource.""" ) -class ListSkillsResponseDict(TypedDict, total=False): - """Response for listing skills.""" +class PublisherModelResourceReferenceDict(TypedDict, total=False): + """Reference to a resource.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + description: Optional[str] + """Description of the resource.""" - next_page_token: Optional[str] - """""" + resource_name: Optional[str] + """The resource name of the Google Cloud resource.""" - skills: Optional[list[SkillDict]] - """List of Skills.""" + uri: Optional[str] + """The URI of the resource.""" + use_case: Optional[str] + """Use case (CUJ) of the resource.""" -ListSkillsResponseOrDict = Union[ListSkillsResponse, ListSkillsResponseDict] +PublisherModelResourceReferenceOrDict = Union[ + PublisherModelResourceReference, PublisherModelResourceReferenceDict +] -class DeleteSkillConfig(_common.BaseModel): - """Config for deleting a skill.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" +class PublisherModelParent(_common.BaseModel): + """The information about the parent of a model.""" + + display_name: Optional[str] = Field( + default=None, + description="""Required. The display name of the parent. E.g., LaMDA, T5, Vision API, Natural Language API.""", ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Whether to wait for the long running operation to complete.""", + reference: Optional[PublisherModelResourceReference] = Field( + default=None, + description="""Optional. The Google Cloud resource name or the URI reference.""", ) -class DeleteSkillConfigDict(TypedDict, total=False): - """Config for deleting a skill.""" +class PublisherModelParentDict(TypedDict, total=False): + """The information about the parent of a model.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + display_name: Optional[str] + """Required. The display name of the parent. E.g., LaMDA, T5, Vision API, Natural Language API.""" - wait_for_completion: Optional[bool] - """Whether to wait for the long running operation to complete.""" + reference: Optional[PublisherModelResourceReferenceDict] + """Optional. The Google Cloud resource name or the URI reference.""" -DeleteSkillConfigOrDict = Union[DeleteSkillConfig, DeleteSkillConfigDict] +PublisherModelParentOrDict = Union[PublisherModelParent, PublisherModelParentDict] -class _DeleteSkillRequestParameters(_common.BaseModel): - """Parameters for deleting a skill.""" +class PredictSchemata(_common.BaseModel): + """Contains the schemata used in Model's predictions and explanations via PredictionService.Predict, PredictionService.Explain and BatchPredictionJob.""" - name: Optional[str] = Field( + instance_schema_uri: Optional[str] = Field( default=None, - description="""Required. The resource name of the Skill to delete.""", + description="""Immutable. Points to a YAML file stored on Google Cloud Storage describing the format of a single instance, which are used in PredictRequest.instances, ExplainRequest.instances and BatchPredictionJob.input_config. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML Models always have this field populated by Vertex AI. Note: The URI given on output will be immutable and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""", + ) + parameters_schema_uri: Optional[str] = Field( + default=None, + description="""Immutable. Points to a YAML file stored on Google Cloud Storage describing the parameters of prediction and explanation via PredictRequest.parameters, ExplainRequest.parameters and BatchPredictionJob.model_parameters. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML Models always have this field populated by Vertex AI, if no parameters are supported, then it is set to an empty string. Note: The URI given on output will be immutable and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""", + ) + prediction_schema_uri: Optional[str] = Field( + default=None, + description="""Immutable. Points to a YAML file stored on Google Cloud Storage describing the format of a single prediction produced by this Model, which are returned via PredictResponse.predictions, ExplainResponse.explanations, and BatchPredictionJob.output_config. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML Models always have this field populated by Vertex AI. Note: The URI given on output will be immutable and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""", ) - config: Optional[DeleteSkillConfig] = Field(default=None, description="""""") -class _DeleteSkillRequestParametersDict(TypedDict, total=False): - """Parameters for deleting a skill.""" +class PredictSchemataDict(TypedDict, total=False): + """Contains the schemata used in Model's predictions and explanations via PredictionService.Predict, PredictionService.Explain and BatchPredictionJob.""" - name: Optional[str] - """Required. The resource name of the Skill to delete.""" + instance_schema_uri: Optional[str] + """Immutable. Points to a YAML file stored on Google Cloud Storage describing the format of a single instance, which are used in PredictRequest.instances, ExplainRequest.instances and BatchPredictionJob.input_config. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML Models always have this field populated by Vertex AI. Note: The URI given on output will be immutable and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""" - config: Optional[DeleteSkillConfigDict] - """""" + parameters_schema_uri: Optional[str] + """Immutable. Points to a YAML file stored on Google Cloud Storage describing the parameters of prediction and explanation via PredictRequest.parameters, ExplainRequest.parameters and BatchPredictionJob.model_parameters. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML Models always have this field populated by Vertex AI, if no parameters are supported, then it is set to an empty string. Note: The URI given on output will be immutable and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""" + prediction_schema_uri: Optional[str] + """Immutable. Points to a YAML file stored on Google Cloud Storage describing the format of a single prediction produced by this Model, which are returned via PredictResponse.predictions, ExplainResponse.explanations, and BatchPredictionJob.output_config. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML Models always have this field populated by Vertex AI. Note: The URI given on output will be immutable and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""" -_DeleteSkillRequestParametersOrDict = Union[ - _DeleteSkillRequestParameters, _DeleteSkillRequestParametersDict -] +PredictSchemataOrDict = Union[PredictSchemata, PredictSchemataDict] -class DeleteSkillOperation(_common.BaseModel): - """Operation for deleting a skill.""" - name: Optional[str] = Field( +class PublisherModelCallToActionRegionalResourceReferences(_common.BaseModel): + """The regional resource name or the URI. Key is region, e.g., us-central1, europe-west2, global, etc..""" + + colab_notebook_disabled: Optional[bool] = 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. For notebook resource. When set to true, the Colab Enterprise link will be disabled in the "open notebook" dialog in UI.""", ) - 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.""", + references: Optional[dict[str, PublisherModelResourceReference]] = Field( + default=None, description="""Required.""" ) - 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.""", + resource_description: Optional[str] = Field( + default=None, description="""Optional. Description of the resource.""" ) - error: Optional[dict[str, Any]] = Field( - default=None, - description="""The error result of the operation in case of failure or cancellation.""", + resource_title: Optional[str] = Field( + default=None, description="""Optional. Title of the resource.""" + ) + resource_use_case: Optional[str] = Field( + default=None, description="""Optional. Use case (CUJ) of the resource.""" + ) + supports_workbench: Optional[bool] = Field( + default=None, + description="""Optional. For notebook resource, whether the notebook supports Workbench.""", ) + title: Optional[str] = Field(default=None, description="""Required. """) -class DeleteSkillOperationDict(TypedDict, total=False): - """Operation for deleting a skill.""" +class PublisherModelCallToActionRegionalResourceReferencesDict(TypedDict, total=False): + """The regional resource name or the URI. Key is region, e.g., us-central1, europe-west2, global, etc..""" - 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}`.""" + colab_notebook_disabled: Optional[bool] + """Optional. For notebook resource. When set to true, the Colab Enterprise link will be disabled in the "open notebook" dialog in UI.""" - 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.""" + references: Optional[dict[str, PublisherModelResourceReferenceDict]] + """Required.""" - 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.""" + resource_description: Optional[str] + """Optional. Description of the resource.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + resource_title: Optional[str] + """Optional. Title of the resource.""" + + resource_use_case: Optional[str] + """Optional. Use case (CUJ) of the resource.""" + supports_workbench: Optional[bool] + """Optional. For notebook resource, whether the notebook supports Workbench.""" -DeleteSkillOperationOrDict = Union[DeleteSkillOperation, DeleteSkillOperationDict] + title: Optional[str] + """Required. """ -class GetSkillOperationConfig(_common.BaseModel): +PublisherModelCallToActionRegionalResourceReferencesOrDict = Union[ + PublisherModelCallToActionRegionalResourceReferences, + PublisherModelCallToActionRegionalResourceReferencesDict, +] - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + +class AutomaticResources(_common.BaseModel): + """A description of resources that to large degree are decided by Agent Platform, and require only a modest additional configuration. Each Model supporting these resources documents its specific guidelines.""" + + max_replica_count: Optional[int] = Field( + default=None, + description="""Immutable. The maximum number of replicas that may be deployed on when the traffic against it increases. If the requested value is too large, the deployment will error, but if deployment succeeds then the ability to scale to that many replicas is guaranteed (barring service outages). If traffic increases beyond what its replicas at maximum may handle, a portion of the traffic will be dropped. If this value is not provided, a no upper bound for scaling under heavy traffic will be assume, though Agent Platform may be unable to scale beyond certain replica number.""", + ) + min_replica_count: Optional[int] = Field( + default=None, + description="""Immutable. The minimum number of replicas that will be always deployed on. If traffic against it increases, it may dynamically be deployed onto more replicas up to max_replica_count, and as traffic decreases, some of these extra replicas may be freed. If the requested value is too large, the deployment will error.""", ) -class GetSkillOperationConfigDict(TypedDict, total=False): +class AutomaticResourcesDict(TypedDict, total=False): + """A description of resources that to large degree are decided by Agent Platform, and require only a modest additional configuration. Each Model supporting these resources documents its specific guidelines.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + max_replica_count: Optional[int] + """Immutable. The maximum number of replicas that may be deployed on when the traffic against it increases. If the requested value is too large, the deployment will error, but if deployment succeeds then the ability to scale to that many replicas is guaranteed (barring service outages). If traffic increases beyond what its replicas at maximum may handle, a portion of the traffic will be dropped. If this value is not provided, a no upper bound for scaling under heavy traffic will be assume, though Agent Platform may be unable to scale beyond certain replica number.""" + min_replica_count: Optional[int] + """Immutable. The minimum number of replicas that will be always deployed on. If traffic against it increases, it may dynamically be deployed onto more replicas up to max_replica_count, and as traffic decreases, some of these extra replicas may be freed. If the requested value is too large, the deployment will error.""" -GetSkillOperationConfigOrDict = Union[ - GetSkillOperationConfig, GetSkillOperationConfigDict -] +AutomaticResourcesOrDict = Union[AutomaticResources, AutomaticResourcesDict] -class _GetSkillOperationParameters(_common.BaseModel): - """Parameters for getting an operation.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" - ) - config: Optional[GetSkillOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" +class Port(_common.BaseModel): + """Represents a network port in a container.""" + + container_port: Optional[int] = Field( + default=None, + description="""The number of the port to expose on the pod's IP address. Must be a valid port number, between 1 and 65535 inclusive.""", ) -class _GetSkillOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation.""" +class PortDict(TypedDict, total=False): + """Represents a network port in a container.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + container_port: Optional[int] + """The number of the port to expose on the pod's IP address. Must be a valid port number, between 1 and 65535 inclusive.""" - config: Optional[GetSkillOperationConfigDict] - """Used to override the default configuration.""" +PortOrDict = Union[Port, PortDict] -_GetSkillOperationParametersOrDict = Union[ - _GetSkillOperationParameters, _GetSkillOperationParametersDict -] +class ProbeExecAction(_common.BaseModel): + """ExecAction specifies a command to execute.""" -class GetSkillRevisionConfig(_common.BaseModel): - """Configuration for getting a Skill Revision.""" + command: Optional[list[str]] = Field( + default=None, + description="""Command is the command line to execute inside the container, the working directory for the command is root ('/') in the container's filesystem. The command is simply exec'd, it is not run inside a shell, so traditional shell instructions ('|', etc) won't work. To use a shell, you need to explicitly call out to that shell. Exit status of 0 is treated as live/healthy and non-zero is unhealthy.""", + ) - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + +class ProbeExecActionDict(TypedDict, total=False): + """ExecAction specifies a command to execute.""" + + command: Optional[list[str]] + """Command is the command line to execute inside the container, the working directory for the command is root ('/') in the container's filesystem. The command is simply exec'd, it is not run inside a shell, so traditional shell instructions ('|', etc) won't work. To use a shell, you need to explicitly call out to that shell. Exit status of 0 is treated as live/healthy and non-zero is unhealthy.""" + + +ProbeExecActionOrDict = Union[ProbeExecAction, ProbeExecActionDict] + + +class ProbeGrpcAction(_common.BaseModel): + """GrpcAction checks the health of a container using a gRPC service.""" + + port: Optional[int] = Field( + default=None, + description="""Port number of the gRPC service. Number must be in the range 1 to 65535.""", + ) + service: Optional[str] = Field( + default=None, + description="""Service is the name of the service to place in the gRPC HealthCheckRequest. See https://github.com/grpc/grpc/blob/master/doc/health-checking.md. If this is not specified, the default behavior is defined by gRPC.""", ) -class GetSkillRevisionConfigDict(TypedDict, total=False): - """Configuration for getting a Skill Revision.""" +class ProbeGrpcActionDict(TypedDict, total=False): + """GrpcAction checks the health of a container using a gRPC service.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + port: Optional[int] + """Port number of the gRPC service. Number must be in the range 1 to 65535.""" + service: Optional[str] + """Service is the name of the service to place in the gRPC HealthCheckRequest. See https://github.com/grpc/grpc/blob/master/doc/health-checking.md. If this is not specified, the default behavior is defined by gRPC.""" -GetSkillRevisionConfigOrDict = Union[GetSkillRevisionConfig, GetSkillRevisionConfigDict] + +ProbeGrpcActionOrDict = Union[ProbeGrpcAction, ProbeGrpcActionDict] -class _GetSkillRevisionRequestParameters(_common.BaseModel): - """Parameters for getting a Skill Revision.""" +class ProbeHttpHeader(_common.BaseModel): + """HttpHeader describes a custom header to be used in HTTP probes""" name: Optional[str] = Field( default=None, - description="""The resource name of the Skill Revision to retrieve. Format: projects/{project}/locations/{location}/skills/{skill}/revisions/{revision}""", + description="""The header field name. This will be canonicalized upon output, so case-variant names will be understood as the same header.""", ) - config: Optional[GetSkillRevisionConfig] = Field(default=None, description="""""") + value: Optional[str] = Field(default=None, description="""The header field value""") -class _GetSkillRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for getting a Skill Revision.""" +class ProbeHttpHeaderDict(TypedDict, total=False): + """HttpHeader describes a custom header to be used in HTTP probes""" name: Optional[str] - """The resource name of the Skill Revision to retrieve. Format: projects/{project}/locations/{location}/skills/{skill}/revisions/{revision}""" + """The header field name. This will be canonicalized upon output, so case-variant names will be understood as the same header.""" - config: Optional[GetSkillRevisionConfigDict] - """""" + value: Optional[str] + """The header field value""" -_GetSkillRevisionRequestParametersOrDict = Union[ - _GetSkillRevisionRequestParameters, _GetSkillRevisionRequestParametersDict -] +ProbeHttpHeaderOrDict = Union[ProbeHttpHeader, ProbeHttpHeaderDict] -class SkillRevision(_common.BaseModel): - """A single revision of a Skill.""" +class ProbeHttpGetAction(_common.BaseModel): + """HttpGetAction describes an action based on HTTP Get requests.""" - name: Optional[str] = Field( + host: Optional[str] = Field( default=None, - description="""Identifier. The resource name of the Skill Revision. Format: `projects/{project}/locations/{location}/skills/{skill}/revisions/{revision}`""", + description="""Host name to connect to, defaults to the model serving container's IP. You probably want to set "Host" in httpHeaders instead.""", ) - create_time: Optional[datetime.datetime] = Field( + http_headers: Optional[list[ProbeHttpHeader]] = Field( default=None, - description="""Output only. Timestamp when this Skill Revision was created.""", + description="""Custom headers to set in the request. HTTP allows repeated headers.""", ) - skill: Optional[Skill] = Field( + path: Optional[str] = Field( + default=None, description="""Path to access on the HTTP server.""" + ) + port: Optional[int] = Field( default=None, - description="""Output only. The state of the Skill at this revision.""", + description="""Number of the port to access on the container. Number must be in the range 1 to 65535.""", ) - state: Optional[SkillRevisionState] = Field( - default=None, description="""Output only. The state of the Skill Revision.""" + scheme: Optional[str] = Field( + default=None, + description="""Scheme to use for connecting to the host. Defaults to HTTP. Acceptable values are "HTTP" or "HTTPS".""", ) -class SkillRevisionDict(TypedDict, total=False): - """A single revision of a Skill.""" +class ProbeHttpGetActionDict(TypedDict, total=False): + """HttpGetAction describes an action based on HTTP Get requests.""" - name: Optional[str] - """Identifier. The resource name of the Skill Revision. Format: `projects/{project}/locations/{location}/skills/{skill}/revisions/{revision}`""" + host: Optional[str] + """Host name to connect to, defaults to the model serving container's IP. You probably want to set "Host" in httpHeaders instead.""" - create_time: Optional[datetime.datetime] - """Output only. Timestamp when this Skill Revision was created.""" + http_headers: Optional[list[ProbeHttpHeaderDict]] + """Custom headers to set in the request. HTTP allows repeated headers.""" - skill: Optional[SkillDict] - """Output only. The state of the Skill at this revision.""" + path: Optional[str] + """Path to access on the HTTP server.""" - state: Optional[SkillRevisionState] - """Output only. The state of the Skill Revision.""" + port: Optional[int] + """Number of the port to access on the container. Number must be in the range 1 to 65535.""" + + scheme: Optional[str] + """Scheme to use for connecting to the host. Defaults to HTTP. Acceptable values are "HTTP" or "HTTPS".""" -SkillRevisionOrDict = Union[SkillRevision, SkillRevisionDict] +ProbeHttpGetActionOrDict = Union[ProbeHttpGetAction, ProbeHttpGetActionDict] -class ListSkillRevisionsConfig(_common.BaseModel): - """Configuration for listing Skill Revisions.""" +class ProbeTcpSocketAction(_common.BaseModel): + """TcpSocketAction probes the health of a container by opening a TCP socket connection.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + host: Optional[str] = Field( + default=None, + description="""Optional: Host name to connect to, defaults to the model serving container's IP.""", + ) + port: Optional[int] = Field( + default=None, + description="""Number of the port to access on the container. Number must be in the range 1 to 65535.""", ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - -class ListSkillRevisionsConfigDict(TypedDict, total=False): - """Configuration for listing Skill Revisions.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class ProbeTcpSocketActionDict(TypedDict, total=False): + """TcpSocketAction probes the health of a container by opening a TCP socket connection.""" - page_size: Optional[int] - """""" + host: Optional[str] + """Optional: Host name to connect to, defaults to the model serving container's IP.""" - page_token: Optional[str] - """""" + port: Optional[int] + """Number of the port to access on the container. Number must be in the range 1 to 65535.""" -ListSkillRevisionsConfigOrDict = Union[ - ListSkillRevisionsConfig, ListSkillRevisionsConfigDict -] +ProbeTcpSocketActionOrDict = Union[ProbeTcpSocketAction, ProbeTcpSocketActionDict] -class _ListSkillRevisionsRequestParameters(_common.BaseModel): - """Parameters for ListSkillRevisionsRequest.""" +class Probe(_common.BaseModel): + """Probe describes a health check to be performed against a container to determine whether it is alive or ready to receive traffic.""" - name: Optional[str] = Field( + exec: Optional[ProbeExecAction] = Field( default=None, - description="""Required. The name of the Skill to list revisions for.""", + description="""ExecAction probes the health of a container by executing a command.""", + ) + failure_threshold: Optional[int] = Field( + default=None, + description="""Number of consecutive failures before the probe is considered failed. Defaults to 3. Minimum value is 1. Maps to Kubernetes probe argument 'failureThreshold'.""", + ) + grpc: Optional[ProbeGrpcAction] = Field( + default=None, + description="""GrpcAction probes the health of a container by sending a gRPC request.""", + ) + http_get: Optional[ProbeHttpGetAction] = Field( + default=None, + description="""HttpGetAction probes the health of a container by sending an HTTP GET request.""", + ) + initial_delay_seconds: Optional[int] = Field( + default=None, + description="""Number of seconds to wait before starting the probe. Defaults to 0. Minimum value is 0. Maps to Kubernetes probe argument 'initialDelaySeconds'.""", + ) + period_seconds: Optional[int] = Field( + default=None, + description="""How often (in seconds) to perform the probe. Default to 10 seconds. Minimum value is 1. Must be less than timeout_seconds. Maps to Kubernetes probe argument 'periodSeconds'.""", + ) + success_threshold: Optional[int] = Field( + default=None, + description="""Number of consecutive successes before the probe is considered successful. Defaults to 1. Minimum value is 1. Maps to Kubernetes probe argument 'successThreshold'.""", + ) + tcp_socket: Optional[ProbeTcpSocketAction] = Field( + default=None, + description="""TcpSocketAction probes the health of a container by opening a TCP socket connection.""", + ) + timeout_seconds: Optional[int] = Field( + default=None, + description="""Number of seconds after which the probe times out. Defaults to 1 second. Minimum value is 1. Must be greater or equal to period_seconds. Maps to Kubernetes probe argument 'timeoutSeconds'.""", ) - config: Optional[ListSkillRevisionsConfig] = Field(default=None, description="""""") - - -class _ListSkillRevisionsRequestParametersDict(TypedDict, total=False): - """Parameters for ListSkillRevisionsRequest.""" - - name: Optional[str] - """Required. The name of the Skill to list revisions for.""" - config: Optional[ListSkillRevisionsConfigDict] - """""" +class ProbeDict(TypedDict, total=False): + """Probe describes a health check to be performed against a container to determine whether it is alive or ready to receive traffic.""" -_ListSkillRevisionsRequestParametersOrDict = Union[ - _ListSkillRevisionsRequestParameters, _ListSkillRevisionsRequestParametersDict -] + exec: Optional[ProbeExecActionDict] + """ExecAction probes the health of a container by executing a command.""" + failure_threshold: Optional[int] + """Number of consecutive failures before the probe is considered failed. Defaults to 3. Minimum value is 1. Maps to Kubernetes probe argument 'failureThreshold'.""" -class ListSkillRevisionsResponse(_common.BaseModel): - """Response for listing Skill Revisions.""" + grpc: Optional[ProbeGrpcActionDict] + """GrpcAction probes the health of a container by sending a gRPC request.""" - 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="""""") - skill_revisions: Optional[list[SkillRevision]] = Field( - default=None, description="""List of Skill Revisions.""" - ) + http_get: Optional[ProbeHttpGetActionDict] + """HttpGetAction probes the health of a container by sending an HTTP GET request.""" + initial_delay_seconds: Optional[int] + """Number of seconds to wait before starting the probe. Defaults to 0. Minimum value is 0. Maps to Kubernetes probe argument 'initialDelaySeconds'.""" -class ListSkillRevisionsResponseDict(TypedDict, total=False): - """Response for listing Skill Revisions.""" + period_seconds: Optional[int] + """How often (in seconds) to perform the probe. Default to 10 seconds. Minimum value is 1. Must be less than timeout_seconds. Maps to Kubernetes probe argument 'periodSeconds'.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + success_threshold: Optional[int] + """Number of consecutive successes before the probe is considered successful. Defaults to 1. Minimum value is 1. Maps to Kubernetes probe argument 'successThreshold'.""" - next_page_token: Optional[str] - """""" + tcp_socket: Optional[ProbeTcpSocketActionDict] + """TcpSocketAction probes the health of a container by opening a TCP socket connection.""" - skill_revisions: Optional[list[SkillRevisionDict]] - """List of Skill Revisions.""" + timeout_seconds: Optional[int] + """Number of seconds after which the probe times out. Defaults to 1 second. Minimum value is 1. Must be greater or equal to period_seconds. Maps to Kubernetes probe argument 'timeoutSeconds'.""" -ListSkillRevisionsResponseOrDict = Union[ - ListSkillRevisionsResponse, ListSkillRevisionsResponseDict -] +ProbeOrDict = Union[Probe, ProbeDict] -class ListPublisherModelsConfig(_common.BaseModel): - """Config for listing publisher models.""" +class ModelContainerSpec(_common.BaseModel): + """Specification of a container for serving predictions. Some fields in this message correspond to fields in the [Kubernetes Container v1 core specification](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + args: Optional[list[str]] = Field( + default=None, + description="""Immutable. Specifies arguments for the command that runs when the container starts. This overrides the container's [`CMD`](https://docs.docker.com/engine/reference/builder/#cmd). Specify this field as an array of executable and arguments, similar to a Docker `CMD`'s "default parameters" form. If you don't specify this field but do specify the command field, then the command from the `command` field runs without any additional arguments. See the [Kubernetes documentation about how the `command` and `args` fields interact with a container's `ENTRYPOINT` and `CMD`](https://kubernetes.io/docs/tasks/inject-data-application/define-command-argument-container/#notes). If you don't specify this field and don't specify the `command` field, then the container's [`ENTRYPOINT`](https://docs.docker.com/engine/reference/builder/#cmd) and `CMD` determine what runs based on their default behavior. See the Docker documentation about [how `CMD` and `ENTRYPOINT` interact](https://docs.docker.com/engine/reference/builder/#understand-how-cmd-and-entrypoint-interact). In this field, you can reference [environment variables set by Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables) and environment variables set in the env field. You cannot reference environment variables set in the Docker image. In order for environment variables to be expanded, reference them by using the following syntax: $( VARIABLE_NAME) Note that this differs from Bash variable expansion, which does not use parentheses. If a variable cannot be resolved, the reference in the input string is used unchanged. To avoid variable expansion, you can escape this syntax with `$$`; for example: $$(VARIABLE_NAME) This field corresponds to the `args` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""", ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( - default=None, description="""Filter string for publisher models.""" + command: Optional[list[str]] = Field( + default=None, + description="""Immutable. Specifies the command that runs when the container starts. This overrides the container's [ENTRYPOINT](https://docs.docker.com/engine/reference/builder/#entrypoint). Specify this field as an array of executable and arguments, similar to a Docker `ENTRYPOINT`'s "exec" form, not its "shell" form. If you do not specify this field, then the container's `ENTRYPOINT` runs, in conjunction with the args field or the container's [`CMD`](https://docs.docker.com/engine/reference/builder/#cmd), if either exists. If this field is not specified and the container does not have an `ENTRYPOINT`, then refer to the Docker documentation about [how `CMD` and `ENTRYPOINT` interact](https://docs.docker.com/engine/reference/builder/#understand-how-cmd-and-entrypoint-interact). If you specify this field, then you can also specify the `args` field to provide additional arguments for this command. However, if you specify this field, then the container's `CMD` is ignored. See the [Kubernetes documentation about how the `command` and `args` fields interact with a container's `ENTRYPOINT` and `CMD`](https://kubernetes.io/docs/tasks/inject-data-application/define-command-argument-container/#notes). In this field, you can reference [environment variables set by Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables) and environment variables set in the env field. You cannot reference environment variables set in the Docker image. In order for environment variables to be expanded, reference them by using the following syntax: $( VARIABLE_NAME) Note that this differs from Bash variable expansion, which does not use parentheses. If a variable cannot be resolved, the reference in the input string is used unchanged. To avoid variable expansion, you can escape this syntax with `$$`; for example: $$(VARIABLE_NAME) This field corresponds to the `command` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""", ) - list_all_versions: Optional[bool] = Field( - default=None, description="""Whether to list all versions.""" + deployment_timeout: Optional[str] = Field( + default=None, + description="""Immutable. Deployment timeout. Limit for deployment timeout is 2 hours.""", + ) + env: Optional[list[EnvVar]] = Field( + default=None, + description="""Immutable. List of environment variables to set in the container. After the container starts running, code running in the container can read these environment variables. Additionally, the command and args fields can reference these variables. Later entries in this list can also reference earlier entries. For example, the following example sets the variable `VAR_2` to have the value `foo bar`: ```json [ { "name": "VAR_1", "value": "foo" }, { "name": "VAR_2", "value": "$(VAR_1) bar" } ] ``` If you switch the order of the variables in the example, then the expansion does not occur. This field corresponds to the `env` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""", + ) + grpc_ports: Optional[list[Port]] = Field( + default=None, + description="""Immutable. List of ports to expose from the container. Vertex AI sends gRPC prediction requests that it receives to the first port on this list. Vertex AI also sends liveness and health checks to this port. If you do not specify this field, gRPC requests to the container will be disabled. Vertex AI does not use ports other than the first one listed. This field corresponds to the `ports` field of the Kubernetes Containers v1 core API.""", + ) + health_probe: Optional[Probe] = Field( + default=None, + description="""Immutable. Specification for Kubernetes readiness probe.""", + ) + health_route: Optional[str] = Field( + default=None, + description="""Immutable. HTTP path on the container to send health checks to. Vertex AI intermittently sends GET requests to this path on the container's IP address and port to check that the container is healthy. Read more about [health checks](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#health). For example, if you set this field to `/bar`, then Vertex AI intermittently sends a GET request to the `/bar` path on the port of your container specified by the first value of this `ModelContainerSpec`'s ports field. If you don't specify this field, it defaults to the following value when you deploy this Model to an Endpoint: /v1/endpoints/ENDPOINT/deployedModels/ DEPLOYED_MODEL:predict The placeholders in this value are replaced as follows: * ENDPOINT: The last segment (following `endpoints/`)of the Endpoint.name][] field of the Endpoint where this Model has been deployed. (Vertex AI makes this value available to your container code as the [`AIP_ENDPOINT_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).) * DEPLOYED_MODEL: DeployedModel.id of the `DeployedModel`. (Vertex AI makes this value available to your container code as the [`AIP_DEPLOYED_MODEL_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).)""", + ) + image_uri: Optional[str] = Field( + default=None, + description="""Required. Immutable. URI of the Docker image to be used as the custom container for serving predictions. This URI must identify an image in Artifact Registry or Container Registry. Learn more about the [container publishing requirements](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#publishing), including permissions requirements for the Vertex AI Service Agent. The container image is ingested upon ModelService.UploadModel, stored internally, and this original path is afterwards not used. To learn about the requirements for the Docker image itself, see [Custom container requirements](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#). You can use the URI to one of Vertex AI's [pre-built container images for prediction](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers) in this field.""", + ) + invoke_route_prefix: Optional[str] = Field( + default=None, + description="""Immutable. Invoke route prefix for the custom container. "/*" is the only supported value right now. By setting this field, any non-root route on this model will be accessible with invoke http call eg: "/invoke/foo/bar", however the [PredictionService.Invoke] RPC is not supported yet. Only one of `predict_route` or `invoke_route_prefix` can be set, and we default to using `predict_route` if this field is not set. If this field is set, the Model can only be deployed to dedicated endpoint.""", + ) + liveness_probe: Optional[Probe] = Field( + default=None, + description="""Immutable. Specification for Kubernetes liveness probe.""", + ) + ports: Optional[list[Port]] = Field( + default=None, + description="""Immutable. List of ports to expose from the container. Vertex AI sends any prediction requests that it receives to the first port on this list. Vertex AI also sends [liveness and health checks](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#liveness) to this port. If you do not specify this field, it defaults to following value: ```json [ { "containerPort": 8080 } ] ``` Vertex AI does not use ports other than the first one listed. This field corresponds to the `ports` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""", + ) + predict_route: Optional[str] = Field( + default=None, + description="""Immutable. HTTP path on the container to send prediction requests to. Vertex AI forwards requests sent using projects.locations.endpoints.predict to this path on the container's IP address and port. Vertex AI then returns the container's response in the API response. For example, if you set this field to `/foo`, then when Vertex AI receives a prediction request, it forwards the request body in a POST request to the `/foo` path on the port of your container specified by the first value of this `ModelContainerSpec`'s ports field. If you don't specify this field, it defaults to the following value when you deploy this Model to an Endpoint: /v1/endpoints/ENDPOINT/deployedModels/DEPLOYED_MODEL:predict The placeholders in this value are replaced as follows: * ENDPOINT: The last segment (following `endpoints/`)of the Endpoint.name][] field of the Endpoint where this Model has been deployed. (Vertex AI makes this value available to your container code as the [`AIP_ENDPOINT_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).) * DEPLOYED_MODEL: DeployedModel.id of the `DeployedModel`. (Vertex AI makes this value available to your container code as the [`AIP_DEPLOYED_MODEL_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).)""", + ) + shared_memory_size_mb: Optional[int] = Field( + default=None, + description="""Immutable. The amount of the VM memory to reserve as the shared memory for the model in megabytes.""", + ) + startup_probe: Optional[Probe] = Field( + default=None, + description="""Immutable. Specification for Kubernetes startup probe.""", ) -class ListPublisherModelsConfigDict(TypedDict, total=False): - """Config for listing publisher models.""" +class ModelContainerSpecDict(TypedDict, total=False): + """Specification of a container for serving predictions. Some fields in this message correspond to fields in the [Kubernetes Container v1 core specification](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + args: Optional[list[str]] + """Immutable. Specifies arguments for the command that runs when the container starts. This overrides the container's [`CMD`](https://docs.docker.com/engine/reference/builder/#cmd). Specify this field as an array of executable and arguments, similar to a Docker `CMD`'s "default parameters" form. If you don't specify this field but do specify the command field, then the command from the `command` field runs without any additional arguments. See the [Kubernetes documentation about how the `command` and `args` fields interact with a container's `ENTRYPOINT` and `CMD`](https://kubernetes.io/docs/tasks/inject-data-application/define-command-argument-container/#notes). If you don't specify this field and don't specify the `command` field, then the container's [`ENTRYPOINT`](https://docs.docker.com/engine/reference/builder/#cmd) and `CMD` determine what runs based on their default behavior. See the Docker documentation about [how `CMD` and `ENTRYPOINT` interact](https://docs.docker.com/engine/reference/builder/#understand-how-cmd-and-entrypoint-interact). In this field, you can reference [environment variables set by Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables) and environment variables set in the env field. You cannot reference environment variables set in the Docker image. In order for environment variables to be expanded, reference them by using the following syntax: $( VARIABLE_NAME) Note that this differs from Bash variable expansion, which does not use parentheses. If a variable cannot be resolved, the reference in the input string is used unchanged. To avoid variable expansion, you can escape this syntax with `$$`; for example: $$(VARIABLE_NAME) This field corresponds to the `args` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" - page_size: Optional[int] - """""" + command: Optional[list[str]] + """Immutable. Specifies the command that runs when the container starts. This overrides the container's [ENTRYPOINT](https://docs.docker.com/engine/reference/builder/#entrypoint). Specify this field as an array of executable and arguments, similar to a Docker `ENTRYPOINT`'s "exec" form, not its "shell" form. If you do not specify this field, then the container's `ENTRYPOINT` runs, in conjunction with the args field or the container's [`CMD`](https://docs.docker.com/engine/reference/builder/#cmd), if either exists. If this field is not specified and the container does not have an `ENTRYPOINT`, then refer to the Docker documentation about [how `CMD` and `ENTRYPOINT` interact](https://docs.docker.com/engine/reference/builder/#understand-how-cmd-and-entrypoint-interact). If you specify this field, then you can also specify the `args` field to provide additional arguments for this command. However, if you specify this field, then the container's `CMD` is ignored. See the [Kubernetes documentation about how the `command` and `args` fields interact with a container's `ENTRYPOINT` and `CMD`](https://kubernetes.io/docs/tasks/inject-data-application/define-command-argument-container/#notes). In this field, you can reference [environment variables set by Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables) and environment variables set in the env field. You cannot reference environment variables set in the Docker image. In order for environment variables to be expanded, reference them by using the following syntax: $( VARIABLE_NAME) Note that this differs from Bash variable expansion, which does not use parentheses. If a variable cannot be resolved, the reference in the input string is used unchanged. To avoid variable expansion, you can escape this syntax with `$$`; for example: $$(VARIABLE_NAME) This field corresponds to the `command` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" - page_token: Optional[str] - """""" + deployment_timeout: Optional[str] + """Immutable. Deployment timeout. Limit for deployment timeout is 2 hours.""" - filter: Optional[str] - """Filter string for publisher models.""" + env: Optional[list[EnvVarDict]] + """Immutable. List of environment variables to set in the container. After the container starts running, code running in the container can read these environment variables. Additionally, the command and args fields can reference these variables. Later entries in this list can also reference earlier entries. For example, the following example sets the variable `VAR_2` to have the value `foo bar`: ```json [ { "name": "VAR_1", "value": "foo" }, { "name": "VAR_2", "value": "$(VAR_1) bar" } ] ``` If you switch the order of the variables in the example, then the expansion does not occur. This field corresponds to the `env` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" - list_all_versions: Optional[bool] - """Whether to list all versions.""" + grpc_ports: Optional[list[PortDict]] + """Immutable. List of ports to expose from the container. Vertex AI sends gRPC prediction requests that it receives to the first port on this list. Vertex AI also sends liveness and health checks to this port. If you do not specify this field, gRPC requests to the container will be disabled. Vertex AI does not use ports other than the first one listed. This field corresponds to the `ports` field of the Kubernetes Containers v1 core API.""" + health_probe: Optional[ProbeDict] + """Immutable. Specification for Kubernetes readiness probe.""" -ListPublisherModelsConfigOrDict = Union[ - ListPublisherModelsConfig, ListPublisherModelsConfigDict -] + health_route: Optional[str] + """Immutable. HTTP path on the container to send health checks to. Vertex AI intermittently sends GET requests to this path on the container's IP address and port to check that the container is healthy. Read more about [health checks](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#health). For example, if you set this field to `/bar`, then Vertex AI intermittently sends a GET request to the `/bar` path on the port of your container specified by the first value of this `ModelContainerSpec`'s ports field. If you don't specify this field, it defaults to the following value when you deploy this Model to an Endpoint: /v1/endpoints/ENDPOINT/deployedModels/ DEPLOYED_MODEL:predict The placeholders in this value are replaced as follows: * ENDPOINT: The last segment (following `endpoints/`)of the Endpoint.name][] field of the Endpoint where this Model has been deployed. (Vertex AI makes this value available to your container code as the [`AIP_ENDPOINT_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).) * DEPLOYED_MODEL: DeployedModel.id of the `DeployedModel`. (Vertex AI makes this value available to your container code as the [`AIP_DEPLOYED_MODEL_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).)""" + image_uri: Optional[str] + """Required. Immutable. URI of the Docker image to be used as the custom container for serving predictions. This URI must identify an image in Artifact Registry or Container Registry. Learn more about the [container publishing requirements](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#publishing), including permissions requirements for the Vertex AI Service Agent. The container image is ingested upon ModelService.UploadModel, stored internally, and this original path is afterwards not used. To learn about the requirements for the Docker image itself, see [Custom container requirements](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#). You can use the URI to one of Vertex AI's [pre-built container images for prediction](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers) in this field.""" -class _ListPublisherModelsRequestParameters(_common.BaseModel): - """Parameters for listing publisher models.""" + invoke_route_prefix: Optional[str] + """Immutable. Invoke route prefix for the custom container. "/*" is the only supported value right now. By setting this field, any non-root route on this model will be accessible with invoke http call eg: "/invoke/foo/bar", however the [PredictionService.Invoke] RPC is not supported yet. Only one of `predict_route` or `invoke_route_prefix` can be set, and we default to using `predict_route` if this field is not set. If this field is set, the Model can only be deployed to dedicated endpoint.""" - parent: Optional[str] = Field(default=None, description="""""") - config: Optional[ListPublisherModelsConfig] = Field( - default=None, description="""""" - ) + liveness_probe: Optional[ProbeDict] + """Immutable. Specification for Kubernetes liveness probe.""" + ports: Optional[list[PortDict]] + """Immutable. List of ports to expose from the container. Vertex AI sends any prediction requests that it receives to the first port on this list. Vertex AI also sends [liveness and health checks](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#liveness) to this port. If you do not specify this field, it defaults to following value: ```json [ { "containerPort": 8080 } ] ``` Vertex AI does not use ports other than the first one listed. This field corresponds to the `ports` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" -class _ListPublisherModelsRequestParametersDict(TypedDict, total=False): - """Parameters for listing publisher models.""" + predict_route: Optional[str] + """Immutable. HTTP path on the container to send prediction requests to. Vertex AI forwards requests sent using projects.locations.endpoints.predict to this path on the container's IP address and port. Vertex AI then returns the container's response in the API response. For example, if you set this field to `/foo`, then when Vertex AI receives a prediction request, it forwards the request body in a POST request to the `/foo` path on the port of your container specified by the first value of this `ModelContainerSpec`'s ports field. If you don't specify this field, it defaults to the following value when you deploy this Model to an Endpoint: /v1/endpoints/ENDPOINT/deployedModels/DEPLOYED_MODEL:predict The placeholders in this value are replaced as follows: * ENDPOINT: The last segment (following `endpoints/`)of the Endpoint.name][] field of the Endpoint where this Model has been deployed. (Vertex AI makes this value available to your container code as the [`AIP_ENDPOINT_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).) * DEPLOYED_MODEL: DeployedModel.id of the `DeployedModel`. (Vertex AI makes this value available to your container code as the [`AIP_DEPLOYED_MODEL_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).)""" - parent: Optional[str] - """""" + shared_memory_size_mb: Optional[int] + """Immutable. The amount of the VM memory to reserve as the shared memory for the model in megabytes.""" - config: Optional[ListPublisherModelsConfigDict] - """""" + startup_probe: Optional[ProbeDict] + """Immutable. Specification for Kubernetes startup probe.""" -_ListPublisherModelsRequestParametersOrDict = Union[ - _ListPublisherModelsRequestParameters, _ListPublisherModelsRequestParametersDict -] +ModelContainerSpecOrDict = Union[ModelContainerSpec, ModelContainerSpecDict] -class PublisherModelResourceReference(_common.BaseModel): - """Reference to a resource.""" +class AutoscalingMetricSpec(_common.BaseModel): + """The metric specification that defines the target resource utilization (CPU utilization, accelerator's duty cycle, and so on) for calculating the desired replica count.""" - description: Optional[str] = Field( - default=None, description="""Description of the resource.""" + metric_name: Optional[str] = Field( + default=None, + description="""Required. The resource metric name. Supported metrics: * For Online Prediction: * `aiplatform.googleapis.com/prediction/online/accelerator/duty_cycle` * `aiplatform.googleapis.com/prediction/online/cpu/utilization` * `aiplatform.googleapis.com/prediction/online/request_count` * `pubsub.googleapis.com/subscription/num_undelivered_messages` * `prometheus.googleapis.com/vertex_dcgm_fi_dev_gpu_util` * `prometheus.googleapis.com/vertex_vllm_gpu_cache_usage_perc` * `prometheus.googleapis.com/vertex_vllm_num_requests_waiting`""", ) - resource_name: Optional[str] = Field( - default=None, description="""The resource name of the Google Cloud resource.""" + monitored_resource_labels: Optional[dict[str, str]] = Field( + default=None, + description="""Optional. The Cloud Monitoring monitored resource labels as key value pairs used for metrics filtering. See Cloud Monitoring Labels https://cloud.google.com/monitoring/api/v3/metric-model#generic-label-info""", ) - uri: Optional[str] = Field(default=None, description="""The URI of the resource.""") - use_case: Optional[str] = Field( - default=None, description="""Use case (CUJ) of the resource.""" + target: Optional[int] = Field( + default=None, + description="""The target resource utilization in percentage (1% - 100%) for the given metric; once the real usage deviates from the target by a certain percentage, the machine replicas change. The default value is 60 (representing 60%) if not provided.""", ) -class PublisherModelResourceReferenceDict(TypedDict, total=False): - """Reference to a resource.""" +class AutoscalingMetricSpecDict(TypedDict, total=False): + """The metric specification that defines the target resource utilization (CPU utilization, accelerator's duty cycle, and so on) for calculating the desired replica count.""" - description: Optional[str] - """Description of the resource.""" + metric_name: Optional[str] + """Required. The resource metric name. Supported metrics: * For Online Prediction: * `aiplatform.googleapis.com/prediction/online/accelerator/duty_cycle` * `aiplatform.googleapis.com/prediction/online/cpu/utilization` * `aiplatform.googleapis.com/prediction/online/request_count` * `pubsub.googleapis.com/subscription/num_undelivered_messages` * `prometheus.googleapis.com/vertex_dcgm_fi_dev_gpu_util` * `prometheus.googleapis.com/vertex_vllm_gpu_cache_usage_perc` * `prometheus.googleapis.com/vertex_vllm_num_requests_waiting`""" - resource_name: Optional[str] - """The resource name of the Google Cloud resource.""" + monitored_resource_labels: Optional[dict[str, str]] + """Optional. The Cloud Monitoring monitored resource labels as key value pairs used for metrics filtering. See Cloud Monitoring Labels https://cloud.google.com/monitoring/api/v3/metric-model#generic-label-info""" - uri: Optional[str] - """The URI of the resource.""" + target: Optional[int] + """The target resource utilization in percentage (1% - 100%) for the given metric; once the real usage deviates from the target by a certain percentage, the machine replicas change. The default value is 60 (representing 60%) if not provided.""" - use_case: Optional[str] - """Use case (CUJ) of the resource.""" +AutoscalingMetricSpecOrDict = Union[AutoscalingMetricSpec, AutoscalingMetricSpecDict] -PublisherModelResourceReferenceOrDict = Union[ - PublisherModelResourceReference, PublisherModelResourceReferenceDict -] +class FlexStart(_common.BaseModel): + """FlexStart is used to schedule the deployment workload on DWS resource. It contains the max duration of the deployment.""" -class PublisherModelParent(_common.BaseModel): - """The information about the parent of a model.""" - - display_name: Optional[str] = Field( - default=None, - description="""Required. The display name of the parent. E.g., LaMDA, T5, Vision API, Natural Language API.""", - ) - reference: Optional[PublisherModelResourceReference] = Field( + max_runtime_duration: Optional[str] = Field( default=None, - description="""Optional. The Google Cloud resource name or the URI reference.""", + description="""The max duration of the deployment is max_runtime_duration. The deployment will be terminated after the duration. The max_runtime_duration can be set up to 7 days.""", ) -class PublisherModelParentDict(TypedDict, total=False): - """The information about the parent of a model.""" - - display_name: Optional[str] - """Required. The display name of the parent. E.g., LaMDA, T5, Vision API, Natural Language API.""" +class FlexStartDict(TypedDict, total=False): + """FlexStart is used to schedule the deployment workload on DWS resource. It contains the max duration of the deployment.""" - reference: Optional[PublisherModelResourceReferenceDict] - """Optional. The Google Cloud resource name or the URI reference.""" + max_runtime_duration: Optional[str] + """The max duration of the deployment is max_runtime_duration. The deployment will be terminated after the duration. The max_runtime_duration can be set up to 7 days.""" -PublisherModelParentOrDict = Union[PublisherModelParent, PublisherModelParentDict] +FlexStartOrDict = Union[FlexStart, FlexStartDict] -class PredictSchemata(_common.BaseModel): - """Contains the schemata used in Model's predictions and explanations via PredictionService.Predict, PredictionService.Explain and BatchPredictionJob.""" +class DedicatedResourcesScaleToZeroSpec(_common.BaseModel): + """Specification for scale-to-zero feature.""" - instance_schema_uri: Optional[str] = Field( - default=None, - description="""Immutable. Points to a YAML file stored on Google Cloud Storage describing the format of a single instance, which are used in PredictRequest.instances, ExplainRequest.instances and BatchPredictionJob.input_config. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML Models always have this field populated by Vertex AI. Note: The URI given on output will be immutable and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""", - ) - parameters_schema_uri: Optional[str] = Field( + idle_scaledown_period: Optional[str] = Field( default=None, - description="""Immutable. Points to a YAML file stored on Google Cloud Storage describing the parameters of prediction and explanation via PredictRequest.parameters, ExplainRequest.parameters and BatchPredictionJob.model_parameters. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML Models always have this field populated by Vertex AI, if no parameters are supported, then it is set to an empty string. Note: The URI given on output will be immutable and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""", + description="""Optional. Duration of no traffic before scaling to zero. [MinValue=300] (5 minutes) [MaxValue=28800] (8 hours)""", ) - prediction_schema_uri: Optional[str] = Field( + min_scaleup_period: Optional[str] = Field( default=None, - description="""Immutable. Points to a YAML file stored on Google Cloud Storage describing the format of a single prediction produced by this Model, which are returned via PredictResponse.predictions, ExplainResponse.explanations, and BatchPredictionJob.output_config. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML Models always have this field populated by Vertex AI. Note: The URI given on output will be immutable and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""", + description="""Optional. Minimum duration that a deployment will be scaled up before traffic is evaluated for potential scale-down. [MinValue=300] (5 minutes) [MaxValue=28800] (8 hours)""", ) -class PredictSchemataDict(TypedDict, total=False): - """Contains the schemata used in Model's predictions and explanations via PredictionService.Predict, PredictionService.Explain and BatchPredictionJob.""" - - instance_schema_uri: Optional[str] - """Immutable. Points to a YAML file stored on Google Cloud Storage describing the format of a single instance, which are used in PredictRequest.instances, ExplainRequest.instances and BatchPredictionJob.input_config. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML Models always have this field populated by Vertex AI. Note: The URI given on output will be immutable and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""" +class DedicatedResourcesScaleToZeroSpecDict(TypedDict, total=False): + """Specification for scale-to-zero feature.""" - parameters_schema_uri: Optional[str] - """Immutable. Points to a YAML file stored on Google Cloud Storage describing the parameters of prediction and explanation via PredictRequest.parameters, ExplainRequest.parameters and BatchPredictionJob.model_parameters. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML Models always have this field populated by Vertex AI, if no parameters are supported, then it is set to an empty string. Note: The URI given on output will be immutable and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""" + idle_scaledown_period: Optional[str] + """Optional. Duration of no traffic before scaling to zero. [MinValue=300] (5 minutes) [MaxValue=28800] (8 hours)""" - prediction_schema_uri: Optional[str] - """Immutable. Points to a YAML file stored on Google Cloud Storage describing the format of a single prediction produced by this Model, which are returned via PredictResponse.predictions, ExplainResponse.explanations, and BatchPredictionJob.output_config. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML Models always have this field populated by Vertex AI. Note: The URI given on output will be immutable and probably different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""" + min_scaleup_period: Optional[str] + """Optional. Minimum duration that a deployment will be scaled up before traffic is evaluated for potential scale-down. [MinValue=300] (5 minutes) [MaxValue=28800] (8 hours)""" -PredictSchemataOrDict = Union[PredictSchemata, PredictSchemataDict] +DedicatedResourcesScaleToZeroSpecOrDict = Union[ + DedicatedResourcesScaleToZeroSpec, DedicatedResourcesScaleToZeroSpecDict +] -class PublisherModelCallToActionRegionalResourceReferences(_common.BaseModel): - """The regional resource name or the URI. Key is region, e.g., us-central1, europe-west2, global, etc..""" +class DedicatedResources(_common.BaseModel): + """A description of resources that are dedicated to a DeployedModel or DeployedIndex, and that need a higher degree of manual configuration.""" - colab_notebook_disabled: Optional[bool] = Field( + autoscaling_metric_specs: Optional[list[AutoscalingMetricSpec]] = Field( default=None, - description="""Optional. For notebook resource. When set to true, the Colab Enterprise link will be disabled in the "open notebook" dialog in UI.""", - ) - references: Optional[dict[str, PublisherModelResourceReference]] = Field( - default=None, description="""Required.""" - ) - resource_description: Optional[str] = Field( - default=None, description="""Optional. Description of the resource.""" + description="""Immutable. The metric specifications that overrides a resource utilization metric (CPU utilization, accelerator's duty cycle, and so on) target value (default to 60 if not set). At most one entry is allowed per metric. If machine_spec.accelerator_count is above 0, the autoscaling will be based on both CPU utilization and accelerator's duty cycle metrics and scale up when either metrics exceeds its target value while scale down if both metrics are under their target value. The default target value is 60 for both metrics. If machine_spec.accelerator_count is 0, the autoscaling will be based on CPU utilization metric only with default target value 60 if not explicitly set. For example, in the case of Online Prediction, if you want to override target CPU utilization to 80, you should set autoscaling_metric_specs.metric_name to `aiplatform.googleapis.com/prediction/online/cpu/utilization` and autoscaling_metric_specs.target to `80`.""", ) - resource_title: Optional[str] = Field( - default=None, description="""Optional. Title of the resource.""" + flex_start: Optional[FlexStart] = Field( + default=None, + description="""Optional. Immutable. If set, use DWS resource to schedule the deployment workload. reference: (https://cloud.google.com/blog/products/compute/introducing-dynamic-workload-scheduler)""", ) - resource_use_case: Optional[str] = Field( - default=None, description="""Optional. Use case (CUJ) of the resource.""" + initial_replica_count: Optional[int] = Field( + default=None, + description="""Immutable. Number of initial replicas being deployed on when scaling the workload up from zero or when creating the workload in case min_replica_count = 0. When min_replica_count > 0 (meaning that the scale-to-zero feature is not enabled), initial_replica_count should not be set. When min_replica_count = 0 (meaning that the scale-to-zero feature is enabled), initial_replica_count should be larger than zero, but no greater than max_replica_count.""", ) - supports_workbench: Optional[bool] = Field( + machine_spec: Optional[MachineSpec] = Field( default=None, - description="""Optional. For notebook resource, whether the notebook supports Workbench.""", + description="""Required. Immutable. The specification of a single machine being used.""", ) - title: Optional[str] = Field(default=None, description="""Required. """) - - -class PublisherModelCallToActionRegionalResourceReferencesDict(TypedDict, total=False): - """The regional resource name or the URI. Key is region, e.g., us-central1, europe-west2, global, etc..""" - - colab_notebook_disabled: Optional[bool] - """Optional. For notebook resource. When set to true, the Colab Enterprise link will be disabled in the "open notebook" dialog in UI.""" - - references: Optional[dict[str, PublisherModelResourceReferenceDict]] - """Required.""" - - resource_description: Optional[str] - """Optional. Description of the resource.""" - - resource_title: Optional[str] - """Optional. Title of the resource.""" - - resource_use_case: Optional[str] - """Optional. Use case (CUJ) of the resource.""" - - supports_workbench: Optional[bool] - """Optional. For notebook resource, whether the notebook supports Workbench.""" - - title: Optional[str] - """Required. """ - - -PublisherModelCallToActionRegionalResourceReferencesOrDict = Union[ - PublisherModelCallToActionRegionalResourceReferences, - PublisherModelCallToActionRegionalResourceReferencesDict, -] - - -class AutomaticResources(_common.BaseModel): - """A description of resources that to large degree are decided by Agent Platform, and require only a modest additional configuration. Each Model supporting these resources documents its specific guidelines.""" - max_replica_count: Optional[int] = Field( default=None, - description="""Immutable. The maximum number of replicas that may be deployed on when the traffic against it increases. If the requested value is too large, the deployment will error, but if deployment succeeds then the ability to scale to that many replicas is guaranteed (barring service outages). If traffic increases beyond what its replicas at maximum may handle, a portion of the traffic will be dropped. If this value is not provided, a no upper bound for scaling under heavy traffic will be assume, though Agent Platform may be unable to scale beyond certain replica number.""", + description="""Immutable. The maximum number of replicas that may be deployed on when the traffic against it increases. If the requested value is too large, the deployment will error, but if deployment succeeds then the ability to scale to that many replicas is guaranteed (barring service outages). If traffic increases beyond what its replicas at maximum may handle, a portion of the traffic will be dropped. If this value is not provided, will use min_replica_count as the default value. The value of this field impacts the charge against Agent Platform CPU and GPU quotas. Specifically, you will be charged for (max_replica_count * number of cores in the selected machine type) and (max_replica_count * number of GPUs per replica in the selected machine type).""", ) min_replica_count: Optional[int] = Field( default=None, - description="""Immutable. The minimum number of replicas that will be always deployed on. If traffic against it increases, it may dynamically be deployed onto more replicas up to max_replica_count, and as traffic decreases, some of these extra replicas may be freed. If the requested value is too large, the deployment will error.""", + description="""Required. Immutable. The minimum number of machine replicas that will be always deployed on. This value must be greater than or equal to 1. If traffic increases, it may dynamically be deployed onto more replicas, and as traffic decreases, some of these extra replicas may be freed.""", ) - - -class AutomaticResourcesDict(TypedDict, total=False): - """A description of resources that to large degree are decided by Agent Platform, and require only a modest additional configuration. Each Model supporting these resources documents its specific guidelines.""" - - max_replica_count: Optional[int] - """Immutable. The maximum number of replicas that may be deployed on when the traffic against it increases. If the requested value is too large, the deployment will error, but if deployment succeeds then the ability to scale to that many replicas is guaranteed (barring service outages). If traffic increases beyond what its replicas at maximum may handle, a portion of the traffic will be dropped. If this value is not provided, a no upper bound for scaling under heavy traffic will be assume, though Agent Platform may be unable to scale beyond certain replica number.""" - - min_replica_count: Optional[int] - """Immutable. The minimum number of replicas that will be always deployed on. If traffic against it increases, it may dynamically be deployed onto more replicas up to max_replica_count, and as traffic decreases, some of these extra replicas may be freed. If the requested value is too large, the deployment will error.""" - - -AutomaticResourcesOrDict = Union[AutomaticResources, AutomaticResourcesDict] - - -class Port(_common.BaseModel): - """Represents a network port in a container.""" - - container_port: Optional[int] = Field( + required_replica_count: Optional[int] = Field( default=None, - description="""The number of the port to expose on the pod's IP address. Must be a valid port number, between 1 and 65535 inclusive.""", + description="""Optional. Number of required available replicas for the deployment to succeed. This field is only needed when partial deployment/mutation is desired. If set, the deploy/mutate operation will succeed once available_replica_count reaches required_replica_count, and the rest of the replicas will be retried. If not set, the default required_replica_count will be min_replica_count.""", + ) + scale_to_zero_spec: Optional[DedicatedResourcesScaleToZeroSpec] = Field( + default=None, + description="""Optional. Specification for scale-to-zero feature.""", + ) + spot: Optional[bool] = Field( + default=None, + description="""Optional. If true, schedule the deployment workload on [spot VMs](https://cloud.google.com/kubernetes-engine/docs/concepts/spot-vms).""", ) -class PortDict(TypedDict, total=False): - """Represents a network port in a container.""" +class DedicatedResourcesDict(TypedDict, total=False): + """A description of resources that are dedicated to a DeployedModel or DeployedIndex, and that need a higher degree of manual configuration.""" - container_port: Optional[int] - """The number of the port to expose on the pod's IP address. Must be a valid port number, between 1 and 65535 inclusive.""" + autoscaling_metric_specs: Optional[list[AutoscalingMetricSpecDict]] + """Immutable. The metric specifications that overrides a resource utilization metric (CPU utilization, accelerator's duty cycle, and so on) target value (default to 60 if not set). At most one entry is allowed per metric. If machine_spec.accelerator_count is above 0, the autoscaling will be based on both CPU utilization and accelerator's duty cycle metrics and scale up when either metrics exceeds its target value while scale down if both metrics are under their target value. The default target value is 60 for both metrics. If machine_spec.accelerator_count is 0, the autoscaling will be based on CPU utilization metric only with default target value 60 if not explicitly set. For example, in the case of Online Prediction, if you want to override target CPU utilization to 80, you should set autoscaling_metric_specs.metric_name to `aiplatform.googleapis.com/prediction/online/cpu/utilization` and autoscaling_metric_specs.target to `80`.""" + flex_start: Optional[FlexStartDict] + """Optional. Immutable. If set, use DWS resource to schedule the deployment workload. reference: (https://cloud.google.com/blog/products/compute/introducing-dynamic-workload-scheduler)""" -PortOrDict = Union[Port, PortDict] + initial_replica_count: Optional[int] + """Immutable. Number of initial replicas being deployed on when scaling the workload up from zero or when creating the workload in case min_replica_count = 0. When min_replica_count > 0 (meaning that the scale-to-zero feature is not enabled), initial_replica_count should not be set. When min_replica_count = 0 (meaning that the scale-to-zero feature is enabled), initial_replica_count should be larger than zero, but no greater than max_replica_count.""" + machine_spec: Optional[MachineSpecDict] + """Required. Immutable. The specification of a single machine being used.""" -class ProbeExecAction(_common.BaseModel): - """ExecAction specifies a command to execute.""" + max_replica_count: Optional[int] + """Immutable. The maximum number of replicas that may be deployed on when the traffic against it increases. If the requested value is too large, the deployment will error, but if deployment succeeds then the ability to scale to that many replicas is guaranteed (barring service outages). If traffic increases beyond what its replicas at maximum may handle, a portion of the traffic will be dropped. If this value is not provided, will use min_replica_count as the default value. The value of this field impacts the charge against Agent Platform CPU and GPU quotas. Specifically, you will be charged for (max_replica_count * number of cores in the selected machine type) and (max_replica_count * number of GPUs per replica in the selected machine type).""" - command: Optional[list[str]] = Field( - default=None, - description="""Command is the command line to execute inside the container, the working directory for the command is root ('/') in the container's filesystem. The command is simply exec'd, it is not run inside a shell, so traditional shell instructions ('|', etc) won't work. To use a shell, you need to explicitly call out to that shell. Exit status of 0 is treated as live/healthy and non-zero is unhealthy.""", - ) + min_replica_count: Optional[int] + """Required. Immutable. The minimum number of machine replicas that will be always deployed on. This value must be greater than or equal to 1. If traffic increases, it may dynamically be deployed onto more replicas, and as traffic decreases, some of these extra replicas may be freed.""" + required_replica_count: Optional[int] + """Optional. Number of required available replicas for the deployment to succeed. This field is only needed when partial deployment/mutation is desired. If set, the deploy/mutate operation will succeed once available_replica_count reaches required_replica_count, and the rest of the replicas will be retried. If not set, the default required_replica_count will be min_replica_count.""" -class ProbeExecActionDict(TypedDict, total=False): - """ExecAction specifies a command to execute.""" + scale_to_zero_spec: Optional[DedicatedResourcesScaleToZeroSpecDict] + """Optional. Specification for scale-to-zero feature.""" - command: Optional[list[str]] - """Command is the command line to execute inside the container, the working directory for the command is root ('/') in the container's filesystem. The command is simply exec'd, it is not run inside a shell, so traditional shell instructions ('|', etc) won't work. To use a shell, you need to explicitly call out to that shell. Exit status of 0 is treated as live/healthy and non-zero is unhealthy.""" + spot: Optional[bool] + """Optional. If true, schedule the deployment workload on [spot VMs](https://cloud.google.com/kubernetes-engine/docs/concepts/spot-vms).""" -ProbeExecActionOrDict = Union[ProbeExecAction, ProbeExecActionDict] +DedicatedResourcesOrDict = Union[DedicatedResources, DedicatedResourcesDict] -class ProbeGrpcAction(_common.BaseModel): - """GrpcAction checks the health of a container using a gRPC service.""" +class PublisherModelCallToActionDeployDeployMetadata(_common.BaseModel): + """Metadata information about the deployment for managing deployment config.""" - port: Optional[int] = Field( + labels: Optional[dict[str, str]] = Field( default=None, - description="""Port number of the gRPC service. Number must be in the range 1 to 65535.""", + description="""Optional. Labels for the deployment config. For managing deployment config like verifying, source of deployment config, etc.""", ) - service: Optional[str] = Field( - default=None, - description="""Service is the name of the service to place in the gRPC HealthCheckRequest. See https://github.com/grpc/grpc/blob/master/doc/health-checking.md. If this is not specified, the default behavior is defined by gRPC.""", + sample_request: Optional[str] = Field( + default=None, description="""Optional. Sample request for deployed endpoint.""" ) -class ProbeGrpcActionDict(TypedDict, total=False): - """GrpcAction checks the health of a container using a gRPC service.""" +class PublisherModelCallToActionDeployDeployMetadataDict(TypedDict, total=False): + """Metadata information about the deployment for managing deployment config.""" - port: Optional[int] - """Port number of the gRPC service. Number must be in the range 1 to 65535.""" + labels: Optional[dict[str, str]] + """Optional. Labels for the deployment config. For managing deployment config like verifying, source of deployment config, etc.""" - service: Optional[str] - """Service is the name of the service to place in the gRPC HealthCheckRequest. See https://github.com/grpc/grpc/blob/master/doc/health-checking.md. If this is not specified, the default behavior is defined by gRPC.""" + sample_request: Optional[str] + """Optional. Sample request for deployed endpoint.""" -ProbeGrpcActionOrDict = Union[ProbeGrpcAction, ProbeGrpcActionDict] +PublisherModelCallToActionDeployDeployMetadataOrDict = Union[ + PublisherModelCallToActionDeployDeployMetadata, + PublisherModelCallToActionDeployDeployMetadataDict, +] -class ProbeHttpHeader(_common.BaseModel): - """HttpHeader describes a custom header to be used in HTTP probes""" +class LargeModelReference(_common.BaseModel): + """Contains information about the Large Model.""" name: Optional[str] = Field( default=None, - description="""The header field name. This will be canonicalized upon output, so case-variant names will be understood as the same header.""", + description="""Required. The unique name of the large Foundation or pre-built model. Like "chat-bison", "text-bison". Or model name with version ID, like "chat-bison@001", "text-bison@005", etc.""", ) - value: Optional[str] = Field(default=None, description="""The header field value""") -class ProbeHttpHeaderDict(TypedDict, total=False): - """HttpHeader describes a custom header to be used in HTTP probes""" +class LargeModelReferenceDict(TypedDict, total=False): + """Contains information about the Large Model.""" name: Optional[str] - """The header field name. This will be canonicalized upon output, so case-variant names will be understood as the same header.""" - - value: Optional[str] - """The header field value""" + """Required. The unique name of the large Foundation or pre-built model. Like "chat-bison", "text-bison". Or model name with version ID, like "chat-bison@001", "text-bison@005", etc.""" -ProbeHttpHeaderOrDict = Union[ProbeHttpHeader, ProbeHttpHeaderDict] +LargeModelReferenceOrDict = Union[LargeModelReference, LargeModelReferenceDict] -class ProbeHttpGetAction(_common.BaseModel): - """HttpGetAction describes an action based on HTTP Get requests.""" +class PublisherModelCallToActionDeploy(_common.BaseModel): + """Model metadata that is needed for UploadModel or DeployModel/CreateEndpoint requests.""" - host: Optional[str] = Field( + artifact_uri: Optional[str] = Field( default=None, - description="""Host name to connect to, defaults to the model serving container's IP. You probably want to set "Host" in httpHeaders instead.""", + description="""Optional. The path to the directory containing the Model artifact and any of its supporting files.""", ) - http_headers: Optional[list[ProbeHttpHeader]] = Field( + automatic_resources: Optional[AutomaticResources] = Field( default=None, - description="""Custom headers to set in the request. HTTP allows repeated headers.""", - ) - path: Optional[str] = Field( - default=None, description="""Path to access on the HTTP server.""" + description="""A description of resources that to large degree are decided by Vertex AI, and require only a modest additional configuration.""", ) - port: Optional[int] = Field( + container_spec: Optional[ModelContainerSpec] = Field( default=None, - description="""Number of the port to access on the container. Number must be in the range 1 to 65535.""", + description="""Optional. The specification of the container that is to be used when deploying this Model in Vertex AI. Not present for Large Models.""", ) - scheme: Optional[str] = Field( + dedicated_resources: Optional[DedicatedResources] = Field( default=None, - description="""Scheme to use for connecting to the host. Defaults to HTTP. Acceptable values are "HTTP" or "HTTPS".""", + description="""A description of resources that are dedicated to the DeployedModel, and that need a higher degree of manual configuration.""", + ) + deploy_metadata: Optional[PublisherModelCallToActionDeployDeployMetadata] = Field( + default=None, + description="""Optional. Metadata information about this deployment config.""", + ) + deploy_task_name: Optional[str] = Field( + default=None, + description="""Optional. The name of the deploy task (e.g., "text to image generation").""", + ) + large_model_reference: Optional[LargeModelReference] = Field( + default=None, + description="""Optional. Large model reference. When this is set, model_artifact_spec is not needed.""", + ) + model_display_name: Optional[str] = Field( + default=None, description="""Optional. Default model display name.""" + ) + public_artifact_uri: Optional[str] = Field( + default=None, + description="""Optional. The signed URI for ephemeral Cloud Storage access to model artifact.""", + ) + shared_resources: Optional[str] = Field( + default=None, + description="""The resource name of the shared DeploymentResourcePool to deploy on. Format: `projects/{project}/locations/{location}/deploymentResourcePools/{deployment_resource_pool}`""", + ) + title: Optional[str] = Field( + default=None, + description="""Required. The title of the regional resource reference.""", ) -class ProbeHttpGetActionDict(TypedDict, total=False): - """HttpGetAction describes an action based on HTTP Get requests.""" +class PublisherModelCallToActionDeployDict(TypedDict, total=False): + """Model metadata that is needed for UploadModel or DeployModel/CreateEndpoint requests.""" - host: Optional[str] - """Host name to connect to, defaults to the model serving container's IP. You probably want to set "Host" in httpHeaders instead.""" + artifact_uri: Optional[str] + """Optional. The path to the directory containing the Model artifact and any of its supporting files.""" - http_headers: Optional[list[ProbeHttpHeaderDict]] - """Custom headers to set in the request. HTTP allows repeated headers.""" + automatic_resources: Optional[AutomaticResourcesDict] + """A description of resources that to large degree are decided by Vertex AI, and require only a modest additional configuration.""" - path: Optional[str] - """Path to access on the HTTP server.""" + container_spec: Optional[ModelContainerSpecDict] + """Optional. The specification of the container that is to be used when deploying this Model in Vertex AI. Not present for Large Models.""" - port: Optional[int] - """Number of the port to access on the container. Number must be in the range 1 to 65535.""" + dedicated_resources: Optional[DedicatedResourcesDict] + """A description of resources that are dedicated to the DeployedModel, and that need a higher degree of manual configuration.""" - scheme: Optional[str] - """Scheme to use for connecting to the host. Defaults to HTTP. Acceptable values are "HTTP" or "HTTPS".""" + deploy_metadata: Optional[PublisherModelCallToActionDeployDeployMetadataDict] + """Optional. Metadata information about this deployment config.""" + deploy_task_name: Optional[str] + """Optional. The name of the deploy task (e.g., "text to image generation").""" -ProbeHttpGetActionOrDict = Union[ProbeHttpGetAction, ProbeHttpGetActionDict] + large_model_reference: Optional[LargeModelReferenceDict] + """Optional. Large model reference. When this is set, model_artifact_spec is not needed.""" + model_display_name: Optional[str] + """Optional. Default model display name.""" -class ProbeTcpSocketAction(_common.BaseModel): - """TcpSocketAction probes the health of a container by opening a TCP socket connection.""" + public_artifact_uri: Optional[str] + """Optional. The signed URI for ephemeral Cloud Storage access to model artifact.""" - host: Optional[str] = Field( - default=None, - description="""Optional: Host name to connect to, defaults to the model serving container's IP.""", - ) - port: Optional[int] = Field( + shared_resources: Optional[str] + """The resource name of the shared DeploymentResourcePool to deploy on. Format: `projects/{project}/locations/{location}/deploymentResourcePools/{deployment_resource_pool}`""" + + title: Optional[str] + """Required. The title of the regional resource reference.""" + + +PublisherModelCallToActionDeployOrDict = Union[ + PublisherModelCallToActionDeploy, PublisherModelCallToActionDeployDict +] + + +class PublisherModelCallToActionDeployGke(_common.BaseModel): + """Configurations for PublisherModel GKE deployment""" + + gke_yaml_configs: Optional[list[str]] = Field( default=None, - description="""Number of the port to access on the container. Number must be in the range 1 to 65535.""", + description="""Optional. GKE deployment configuration in yaml format.""", ) -class ProbeTcpSocketActionDict(TypedDict, total=False): - """TcpSocketAction probes the health of a container by opening a TCP socket connection.""" - - host: Optional[str] - """Optional: Host name to connect to, defaults to the model serving container's IP.""" +class PublisherModelCallToActionDeployGkeDict(TypedDict, total=False): + """Configurations for PublisherModel GKE deployment""" - port: Optional[int] - """Number of the port to access on the container. Number must be in the range 1 to 65535.""" + gke_yaml_configs: Optional[list[str]] + """Optional. GKE deployment configuration in yaml format.""" -ProbeTcpSocketActionOrDict = Union[ProbeTcpSocketAction, ProbeTcpSocketActionDict] +PublisherModelCallToActionDeployGkeOrDict = Union[ + PublisherModelCallToActionDeployGke, PublisherModelCallToActionDeployGkeDict +] -class Probe(_common.BaseModel): - """Probe describes a health check to be performed against a container to determine whether it is alive or ready to receive traffic.""" +class PublisherModelCallToActionDeployVertex(_common.BaseModel): + """Multiple setups to deploy the PublisherModel.""" - exec: Optional[ProbeExecAction] = Field( - default=None, - description="""ExecAction probes the health of a container by executing a command.""", - ) - failure_threshold: Optional[int] = Field( - default=None, - description="""Number of consecutive failures before the probe is considered failed. Defaults to 3. Minimum value is 1. Maps to Kubernetes probe argument 'failureThreshold'.""", - ) - grpc: Optional[ProbeGrpcAction] = Field( - default=None, - description="""GrpcAction probes the health of a container by sending a gRPC request.""", - ) - http_get: Optional[ProbeHttpGetAction] = Field( - default=None, - description="""HttpGetAction probes the health of a container by sending an HTTP GET request.""", - ) - initial_delay_seconds: Optional[int] = Field( - default=None, - description="""Number of seconds to wait before starting the probe. Defaults to 0. Minimum value is 0. Maps to Kubernetes probe argument 'initialDelaySeconds'.""", - ) - period_seconds: Optional[int] = Field( - default=None, - description="""How often (in seconds) to perform the probe. Default to 10 seconds. Minimum value is 1. Must be less than timeout_seconds. Maps to Kubernetes probe argument 'periodSeconds'.""", - ) - success_threshold: Optional[int] = Field( - default=None, - description="""Number of consecutive successes before the probe is considered successful. Defaults to 1. Minimum value is 1. Maps to Kubernetes probe argument 'successThreshold'.""", - ) - tcp_socket: Optional[ProbeTcpSocketAction] = Field( - default=None, - description="""TcpSocketAction probes the health of a container by opening a TCP socket connection.""", + multi_deploy_vertex: Optional[list[PublisherModelCallToActionDeploy]] = Field( + default=None, description="""Optional. One click deployment configurations.""" ) - timeout_seconds: Optional[int] = Field( + + +class PublisherModelCallToActionDeployVertexDict(TypedDict, total=False): + """Multiple setups to deploy the PublisherModel.""" + + multi_deploy_vertex: Optional[list[PublisherModelCallToActionDeployDict]] + """Optional. One click deployment configurations.""" + + +PublisherModelCallToActionDeployVertexOrDict = Union[ + PublisherModelCallToActionDeployVertex, PublisherModelCallToActionDeployVertexDict +] + + +class PublisherModelCallToActionOpenFineTuningPipelines(_common.BaseModel): + """Open fine tuning pipelines.""" + + fine_tuning_pipelines: Optional[ + list[PublisherModelCallToActionRegionalResourceReferences] + ] = Field( default=None, - description="""Number of seconds after which the probe times out. Defaults to 1 second. Minimum value is 1. Must be greater or equal to period_seconds. Maps to Kubernetes probe argument 'timeoutSeconds'.""", + description="""Required. Regional resource references to fine tuning pipelines.""", ) -class ProbeDict(TypedDict, total=False): - """Probe describes a health check to be performed against a container to determine whether it is alive or ready to receive traffic.""" +class PublisherModelCallToActionOpenFineTuningPipelinesDict(TypedDict, total=False): + """Open fine tuning pipelines.""" - exec: Optional[ProbeExecActionDict] - """ExecAction probes the health of a container by executing a command.""" + fine_tuning_pipelines: Optional[ + list[PublisherModelCallToActionRegionalResourceReferencesDict] + ] + """Required. Regional resource references to fine tuning pipelines.""" - failure_threshold: Optional[int] - """Number of consecutive failures before the probe is considered failed. Defaults to 3. Minimum value is 1. Maps to Kubernetes probe argument 'failureThreshold'.""" - grpc: Optional[ProbeGrpcActionDict] - """GrpcAction probes the health of a container by sending a gRPC request.""" +PublisherModelCallToActionOpenFineTuningPipelinesOrDict = Union[ + PublisherModelCallToActionOpenFineTuningPipelines, + PublisherModelCallToActionOpenFineTuningPipelinesDict, +] - http_get: Optional[ProbeHttpGetActionDict] - """HttpGetAction probes the health of a container by sending an HTTP GET request.""" - initial_delay_seconds: Optional[int] - """Number of seconds to wait before starting the probe. Defaults to 0. Minimum value is 0. Maps to Kubernetes probe argument 'initialDelaySeconds'.""" +class PublisherModelCallToActionOpenNotebooks(_common.BaseModel): + """Open notebooks.""" - period_seconds: Optional[int] - """How often (in seconds) to perform the probe. Default to 10 seconds. Minimum value is 1. Must be less than timeout_seconds. Maps to Kubernetes probe argument 'periodSeconds'.""" + notebooks: Optional[list[PublisherModelCallToActionRegionalResourceReferences]] = ( + Field( + default=None, + description="""Required. Regional resource references to notebooks.""", + ) + ) - success_threshold: Optional[int] - """Number of consecutive successes before the probe is considered successful. Defaults to 1. Minimum value is 1. Maps to Kubernetes probe argument 'successThreshold'.""" - tcp_socket: Optional[ProbeTcpSocketActionDict] - """TcpSocketAction probes the health of a container by opening a TCP socket connection.""" +class PublisherModelCallToActionOpenNotebooksDict(TypedDict, total=False): + """Open notebooks.""" - timeout_seconds: Optional[int] - """Number of seconds after which the probe times out. Defaults to 1 second. Minimum value is 1. Must be greater or equal to period_seconds. Maps to Kubernetes probe argument 'timeoutSeconds'.""" + notebooks: Optional[list[PublisherModelCallToActionRegionalResourceReferencesDict]] + """Required. Regional resource references to notebooks.""" -ProbeOrDict = Union[Probe, ProbeDict] +PublisherModelCallToActionOpenNotebooksOrDict = Union[ + PublisherModelCallToActionOpenNotebooks, PublisherModelCallToActionOpenNotebooksDict +] -class ModelContainerSpec(_common.BaseModel): - """Specification of a container for serving predictions. Some fields in this message correspond to fields in the [Kubernetes Container v1 core specification](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" +class PublisherModelDocumentation(_common.BaseModel): + """A named piece of documentation.""" - args: Optional[list[str]] = Field( + content: Optional[str] = Field( default=None, - description="""Immutable. Specifies arguments for the command that runs when the container starts. This overrides the container's [`CMD`](https://docs.docker.com/engine/reference/builder/#cmd). Specify this field as an array of executable and arguments, similar to a Docker `CMD`'s "default parameters" form. If you don't specify this field but do specify the command field, then the command from the `command` field runs without any additional arguments. See the [Kubernetes documentation about how the `command` and `args` fields interact with a container's `ENTRYPOINT` and `CMD`](https://kubernetes.io/docs/tasks/inject-data-application/define-command-argument-container/#notes). If you don't specify this field and don't specify the `command` field, then the container's [`ENTRYPOINT`](https://docs.docker.com/engine/reference/builder/#cmd) and `CMD` determine what runs based on their default behavior. See the Docker documentation about [how `CMD` and `ENTRYPOINT` interact](https://docs.docker.com/engine/reference/builder/#understand-how-cmd-and-entrypoint-interact). In this field, you can reference [environment variables set by Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables) and environment variables set in the env field. You cannot reference environment variables set in the Docker image. In order for environment variables to be expanded, reference them by using the following syntax: $( VARIABLE_NAME) Note that this differs from Bash variable expansion, which does not use parentheses. If a variable cannot be resolved, the reference in the input string is used unchanged. To avoid variable expansion, you can escape this syntax with `$$`; for example: $$(VARIABLE_NAME) This field corresponds to the `args` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""", + description="""Required. Content of this piece of document (in Markdown format).""", ) - command: Optional[list[str]] = Field( + title: Optional[str] = Field( default=None, - description="""Immutable. Specifies the command that runs when the container starts. This overrides the container's [ENTRYPOINT](https://docs.docker.com/engine/reference/builder/#entrypoint). Specify this field as an array of executable and arguments, similar to a Docker `ENTRYPOINT`'s "exec" form, not its "shell" form. If you do not specify this field, then the container's `ENTRYPOINT` runs, in conjunction with the args field or the container's [`CMD`](https://docs.docker.com/engine/reference/builder/#cmd), if either exists. If this field is not specified and the container does not have an `ENTRYPOINT`, then refer to the Docker documentation about [how `CMD` and `ENTRYPOINT` interact](https://docs.docker.com/engine/reference/builder/#understand-how-cmd-and-entrypoint-interact). If you specify this field, then you can also specify the `args` field to provide additional arguments for this command. However, if you specify this field, then the container's `CMD` is ignored. See the [Kubernetes documentation about how the `command` and `args` fields interact with a container's `ENTRYPOINT` and `CMD`](https://kubernetes.io/docs/tasks/inject-data-application/define-command-argument-container/#notes). In this field, you can reference [environment variables set by Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables) and environment variables set in the env field. You cannot reference environment variables set in the Docker image. In order for environment variables to be expanded, reference them by using the following syntax: $( VARIABLE_NAME) Note that this differs from Bash variable expansion, which does not use parentheses. If a variable cannot be resolved, the reference in the input string is used unchanged. To avoid variable expansion, you can escape this syntax with `$$`; for example: $$(VARIABLE_NAME) This field corresponds to the `command` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""", + description="""Required. E.g., OVERVIEW, USE CASES, DOCUMENTATION, SDK & SAMPLES, JAVA, NODE.JS, etc..""", ) - deployment_timeout: Optional[str] = Field( - default=None, - description="""Immutable. Deployment timeout. Limit for deployment timeout is 2 hours.""", + + +class PublisherModelDocumentationDict(TypedDict, total=False): + """A named piece of documentation.""" + + content: Optional[str] + """Required. Content of this piece of document (in Markdown format).""" + + title: Optional[str] + """Required. E.g., OVERVIEW, USE CASES, DOCUMENTATION, SDK & SAMPLES, JAVA, NODE.JS, etc..""" + + +PublisherModelDocumentationOrDict = Union[ + PublisherModelDocumentation, PublisherModelDocumentationDict +] + + +class PublisherModelCallToActionViewRestApi(_common.BaseModel): + """Rest API docs.""" + + documentations: Optional[list[PublisherModelDocumentation]] = Field( + default=None, description="""Required.""" ) - env: Optional[list[EnvVar]] = Field( - default=None, - description="""Immutable. List of environment variables to set in the container. After the container starts running, code running in the container can read these environment variables. Additionally, the command and args fields can reference these variables. Later entries in this list can also reference earlier entries. For example, the following example sets the variable `VAR_2` to have the value `foo bar`: ```json [ { "name": "VAR_1", "value": "foo" }, { "name": "VAR_2", "value": "$(VAR_1) bar" } ] ``` If you switch the order of the variables in the example, then the expansion does not occur. This field corresponds to the `env` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""", + title: Optional[str] = Field( + default=None, description="""Required. The title of the view rest API.""" ) - grpc_ports: Optional[list[Port]] = Field( + + +class PublisherModelCallToActionViewRestApiDict(TypedDict, total=False): + """Rest API docs.""" + + documentations: Optional[list[PublisherModelDocumentationDict]] + """Required.""" + + title: Optional[str] + """Required. The title of the view rest API.""" + + +PublisherModelCallToActionViewRestApiOrDict = Union[ + PublisherModelCallToActionViewRestApi, PublisherModelCallToActionViewRestApiDict +] + + +class PublisherModelCallToAction(_common.BaseModel): + """Actions could take on this Publisher Model.""" + + create_application: Optional[ + PublisherModelCallToActionRegionalResourceReferences + ] = Field( default=None, - description="""Immutable. List of ports to expose from the container. Vertex AI sends gRPC prediction requests that it receives to the first port on this list. Vertex AI also sends liveness and health checks to this port. If you do not specify this field, gRPC requests to the container will be disabled. Vertex AI does not use ports other than the first one listed. This field corresponds to the `ports` field of the Kubernetes Containers v1 core API.""", + description="""Optional. Create application using the PublisherModel.""", ) - health_probe: Optional[Probe] = Field( + deploy: Optional[PublisherModelCallToActionDeploy] = Field( default=None, - description="""Immutable. Specification for Kubernetes readiness probe.""", + description="""Optional. Deploy the PublisherModel to Vertex Endpoint.""", ) - health_route: Optional[str] = Field( + deploy_gke: Optional[PublisherModelCallToActionDeployGke] = Field( default=None, - description="""Immutable. HTTP path on the container to send health checks to. Vertex AI intermittently sends GET requests to this path on the container's IP address and port to check that the container is healthy. Read more about [health checks](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#health). For example, if you set this field to `/bar`, then Vertex AI intermittently sends a GET request to the `/bar` path on the port of your container specified by the first value of this `ModelContainerSpec`'s ports field. If you don't specify this field, it defaults to the following value when you deploy this Model to an Endpoint: /v1/endpoints/ENDPOINT/deployedModels/ DEPLOYED_MODEL:predict The placeholders in this value are replaced as follows: * ENDPOINT: The last segment (following `endpoints/`)of the Endpoint.name][] field of the Endpoint where this Model has been deployed. (Vertex AI makes this value available to your container code as the [`AIP_ENDPOINT_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).) * DEPLOYED_MODEL: DeployedModel.id of the `DeployedModel`. (Vertex AI makes this value available to your container code as the [`AIP_DEPLOYED_MODEL_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).)""", + description="""Optional. Deploy PublisherModel to Google Kubernetes Engine.""", ) - image_uri: Optional[str] = Field( + multi_deploy_vertex: Optional[PublisherModelCallToActionDeployVertex] = Field( default=None, - description="""Required. Immutable. URI of the Docker image to be used as the custom container for serving predictions. This URI must identify an image in Artifact Registry or Container Registry. Learn more about the [container publishing requirements](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#publishing), including permissions requirements for the Vertex AI Service Agent. The container image is ingested upon ModelService.UploadModel, stored internally, and this original path is afterwards not used. To learn about the requirements for the Docker image itself, see [Custom container requirements](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#). You can use the URI to one of Vertex AI's [pre-built container images for prediction](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers) in this field.""", + description="""Optional. Multiple setups to deploy the PublisherModel to Vertex Endpoint.""", ) - invoke_route_prefix: Optional[str] = Field( + open_evaluation_pipeline: Optional[ + PublisherModelCallToActionRegionalResourceReferences + ] = Field( default=None, - description="""Immutable. Invoke route prefix for the custom container. "/*" is the only supported value right now. By setting this field, any non-root route on this model will be accessible with invoke http call eg: "/invoke/foo/bar", however the [PredictionService.Invoke] RPC is not supported yet. Only one of `predict_route` or `invoke_route_prefix` can be set, and we default to using `predict_route` if this field is not set. If this field is set, the Model can only be deployed to dedicated endpoint.""", + description="""Optional. Open evaluation pipeline of the PublisherModel.""", ) - liveness_probe: Optional[Probe] = Field( + open_fine_tuning_pipeline: Optional[ + PublisherModelCallToActionRegionalResourceReferences + ] = Field( default=None, - description="""Immutable. Specification for Kubernetes liveness probe.""", + description="""Optional. Open fine-tuning pipeline of the PublisherModel.""", ) - ports: Optional[list[Port]] = Field( + open_fine_tuning_pipelines: Optional[ + PublisherModelCallToActionOpenFineTuningPipelines + ] = Field( default=None, - description="""Immutable. List of ports to expose from the container. Vertex AI sends any prediction requests that it receives to the first port on this list. Vertex AI also sends [liveness and health checks](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#liveness) to this port. If you do not specify this field, it defaults to following value: ```json [ { "containerPort": 8080 } ] ``` Vertex AI does not use ports other than the first one listed. This field corresponds to the `ports` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""", + description="""Optional. Open fine-tuning pipelines of the PublisherModel.""", ) - predict_route: Optional[str] = Field( - default=None, - description="""Immutable. HTTP path on the container to send prediction requests to. Vertex AI forwards requests sent using projects.locations.endpoints.predict to this path on the container's IP address and port. Vertex AI then returns the container's response in the API response. For example, if you set this field to `/foo`, then when Vertex AI receives a prediction request, it forwards the request body in a POST request to the `/foo` path on the port of your container specified by the first value of this `ModelContainerSpec`'s ports field. If you don't specify this field, it defaults to the following value when you deploy this Model to an Endpoint: /v1/endpoints/ENDPOINT/deployedModels/DEPLOYED_MODEL:predict The placeholders in this value are replaced as follows: * ENDPOINT: The last segment (following `endpoints/`)of the Endpoint.name][] field of the Endpoint where this Model has been deployed. (Vertex AI makes this value available to your container code as the [`AIP_ENDPOINT_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).) * DEPLOYED_MODEL: DeployedModel.id of the `DeployedModel`. (Vertex AI makes this value available to your container code as the [`AIP_DEPLOYED_MODEL_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).)""", + open_generation_ai_studio: Optional[ + PublisherModelCallToActionRegionalResourceReferences + ] = Field(default=None, description="""Optional. Open in Generation AI Studio.""") + open_genie: Optional[PublisherModelCallToActionRegionalResourceReferences] = Field( + default=None, description="""Optional. Open Genie / Playground.""" ) - shared_memory_size_mb: Optional[int] = Field( - default=None, - description="""Immutable. The amount of the VM memory to reserve as the shared memory for the model in megabytes.""", + open_notebook: Optional[PublisherModelCallToActionRegionalResourceReferences] = ( + Field( + default=None, + description="""Optional. Open notebook of the PublisherModel.""", + ) ) - startup_probe: Optional[Probe] = Field( + open_notebooks: Optional[PublisherModelCallToActionOpenNotebooks] = Field( + default=None, description="""Optional. Open notebooks of the PublisherModel.""" + ) + open_prompt_tuning_pipeline: Optional[ + PublisherModelCallToActionRegionalResourceReferences + ] = Field( default=None, - description="""Immutable. Specification for Kubernetes startup probe.""", + description="""Optional. Open prompt-tuning pipeline of the PublisherModel.""", + ) + request_access: Optional[PublisherModelCallToActionRegionalResourceReferences] = ( + Field(default=None, description="""Optional. Request for access.""") + ) + view_rest_api: Optional[PublisherModelCallToActionViewRestApi] = Field( + default=None, description="""Optional. To view Rest API docs.""" ) -class ModelContainerSpecDict(TypedDict, total=False): - """Specification of a container for serving predictions. Some fields in this message correspond to fields in the [Kubernetes Container v1 core specification](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" +class PublisherModelCallToActionDict(TypedDict, total=False): + """Actions could take on this Publisher Model.""" - args: Optional[list[str]] - """Immutable. Specifies arguments for the command that runs when the container starts. This overrides the container's [`CMD`](https://docs.docker.com/engine/reference/builder/#cmd). Specify this field as an array of executable and arguments, similar to a Docker `CMD`'s "default parameters" form. If you don't specify this field but do specify the command field, then the command from the `command` field runs without any additional arguments. See the [Kubernetes documentation about how the `command` and `args` fields interact with a container's `ENTRYPOINT` and `CMD`](https://kubernetes.io/docs/tasks/inject-data-application/define-command-argument-container/#notes). If you don't specify this field and don't specify the `command` field, then the container's [`ENTRYPOINT`](https://docs.docker.com/engine/reference/builder/#cmd) and `CMD` determine what runs based on their default behavior. See the Docker documentation about [how `CMD` and `ENTRYPOINT` interact](https://docs.docker.com/engine/reference/builder/#understand-how-cmd-and-entrypoint-interact). In this field, you can reference [environment variables set by Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables) and environment variables set in the env field. You cannot reference environment variables set in the Docker image. In order for environment variables to be expanded, reference them by using the following syntax: $( VARIABLE_NAME) Note that this differs from Bash variable expansion, which does not use parentheses. If a variable cannot be resolved, the reference in the input string is used unchanged. To avoid variable expansion, you can escape this syntax with `$$`; for example: $$(VARIABLE_NAME) This field corresponds to the `args` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" + create_application: Optional[ + PublisherModelCallToActionRegionalResourceReferencesDict + ] + """Optional. Create application using the PublisherModel.""" - command: Optional[list[str]] - """Immutable. Specifies the command that runs when the container starts. This overrides the container's [ENTRYPOINT](https://docs.docker.com/engine/reference/builder/#entrypoint). Specify this field as an array of executable and arguments, similar to a Docker `ENTRYPOINT`'s "exec" form, not its "shell" form. If you do not specify this field, then the container's `ENTRYPOINT` runs, in conjunction with the args field or the container's [`CMD`](https://docs.docker.com/engine/reference/builder/#cmd), if either exists. If this field is not specified and the container does not have an `ENTRYPOINT`, then refer to the Docker documentation about [how `CMD` and `ENTRYPOINT` interact](https://docs.docker.com/engine/reference/builder/#understand-how-cmd-and-entrypoint-interact). If you specify this field, then you can also specify the `args` field to provide additional arguments for this command. However, if you specify this field, then the container's `CMD` is ignored. See the [Kubernetes documentation about how the `command` and `args` fields interact with a container's `ENTRYPOINT` and `CMD`](https://kubernetes.io/docs/tasks/inject-data-application/define-command-argument-container/#notes). In this field, you can reference [environment variables set by Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables) and environment variables set in the env field. You cannot reference environment variables set in the Docker image. In order for environment variables to be expanded, reference them by using the following syntax: $( VARIABLE_NAME) Note that this differs from Bash variable expansion, which does not use parentheses. If a variable cannot be resolved, the reference in the input string is used unchanged. To avoid variable expansion, you can escape this syntax with `$$`; for example: $$(VARIABLE_NAME) This field corresponds to the `command` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" + deploy: Optional[PublisherModelCallToActionDeployDict] + """Optional. Deploy the PublisherModel to Vertex Endpoint.""" - deployment_timeout: Optional[str] - """Immutable. Deployment timeout. Limit for deployment timeout is 2 hours.""" + deploy_gke: Optional[PublisherModelCallToActionDeployGkeDict] + """Optional. Deploy PublisherModel to Google Kubernetes Engine.""" - env: Optional[list[EnvVarDict]] - """Immutable. List of environment variables to set in the container. After the container starts running, code running in the container can read these environment variables. Additionally, the command and args fields can reference these variables. Later entries in this list can also reference earlier entries. For example, the following example sets the variable `VAR_2` to have the value `foo bar`: ```json [ { "name": "VAR_1", "value": "foo" }, { "name": "VAR_2", "value": "$(VAR_1) bar" } ] ``` If you switch the order of the variables in the example, then the expansion does not occur. This field corresponds to the `env` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" + multi_deploy_vertex: Optional[PublisherModelCallToActionDeployVertexDict] + """Optional. Multiple setups to deploy the PublisherModel to Vertex Endpoint.""" - grpc_ports: Optional[list[PortDict]] - """Immutable. List of ports to expose from the container. Vertex AI sends gRPC prediction requests that it receives to the first port on this list. Vertex AI also sends liveness and health checks to this port. If you do not specify this field, gRPC requests to the container will be disabled. Vertex AI does not use ports other than the first one listed. This field corresponds to the `ports` field of the Kubernetes Containers v1 core API.""" + open_evaluation_pipeline: Optional[ + PublisherModelCallToActionRegionalResourceReferencesDict + ] + """Optional. Open evaluation pipeline of the PublisherModel.""" - health_probe: Optional[ProbeDict] - """Immutable. Specification for Kubernetes readiness probe.""" + open_fine_tuning_pipeline: Optional[ + PublisherModelCallToActionRegionalResourceReferencesDict + ] + """Optional. Open fine-tuning pipeline of the PublisherModel.""" - health_route: Optional[str] - """Immutable. HTTP path on the container to send health checks to. Vertex AI intermittently sends GET requests to this path on the container's IP address and port to check that the container is healthy. Read more about [health checks](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#health). For example, if you set this field to `/bar`, then Vertex AI intermittently sends a GET request to the `/bar` path on the port of your container specified by the first value of this `ModelContainerSpec`'s ports field. If you don't specify this field, it defaults to the following value when you deploy this Model to an Endpoint: /v1/endpoints/ENDPOINT/deployedModels/ DEPLOYED_MODEL:predict The placeholders in this value are replaced as follows: * ENDPOINT: The last segment (following `endpoints/`)of the Endpoint.name][] field of the Endpoint where this Model has been deployed. (Vertex AI makes this value available to your container code as the [`AIP_ENDPOINT_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).) * DEPLOYED_MODEL: DeployedModel.id of the `DeployedModel`. (Vertex AI makes this value available to your container code as the [`AIP_DEPLOYED_MODEL_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).)""" + open_fine_tuning_pipelines: Optional[ + PublisherModelCallToActionOpenFineTuningPipelinesDict + ] + """Optional. Open fine-tuning pipelines of the PublisherModel.""" - image_uri: Optional[str] - """Required. Immutable. URI of the Docker image to be used as the custom container for serving predictions. This URI must identify an image in Artifact Registry or Container Registry. Learn more about the [container publishing requirements](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#publishing), including permissions requirements for the Vertex AI Service Agent. The container image is ingested upon ModelService.UploadModel, stored internally, and this original path is afterwards not used. To learn about the requirements for the Docker image itself, see [Custom container requirements](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#). You can use the URI to one of Vertex AI's [pre-built container images for prediction](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers) in this field.""" + open_generation_ai_studio: Optional[ + PublisherModelCallToActionRegionalResourceReferencesDict + ] + """Optional. Open in Generation AI Studio.""" - invoke_route_prefix: Optional[str] - """Immutable. Invoke route prefix for the custom container. "/*" is the only supported value right now. By setting this field, any non-root route on this model will be accessible with invoke http call eg: "/invoke/foo/bar", however the [PredictionService.Invoke] RPC is not supported yet. Only one of `predict_route` or `invoke_route_prefix` can be set, and we default to using `predict_route` if this field is not set. If this field is set, the Model can only be deployed to dedicated endpoint.""" + open_genie: Optional[PublisherModelCallToActionRegionalResourceReferencesDict] + """Optional. Open Genie / Playground.""" - liveness_probe: Optional[ProbeDict] - """Immutable. Specification for Kubernetes liveness probe.""" + open_notebook: Optional[PublisherModelCallToActionRegionalResourceReferencesDict] + """Optional. Open notebook of the PublisherModel.""" - ports: Optional[list[PortDict]] - """Immutable. List of ports to expose from the container. Vertex AI sends any prediction requests that it receives to the first port on this list. Vertex AI also sends [liveness and health checks](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#liveness) to this port. If you do not specify this field, it defaults to following value: ```json [ { "containerPort": 8080 } ] ``` Vertex AI does not use ports other than the first one listed. This field corresponds to the `ports` field of the Kubernetes Containers [v1 core API](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" + open_notebooks: Optional[PublisherModelCallToActionOpenNotebooksDict] + """Optional. Open notebooks of the PublisherModel.""" - predict_route: Optional[str] - """Immutable. HTTP path on the container to send prediction requests to. Vertex AI forwards requests sent using projects.locations.endpoints.predict to this path on the container's IP address and port. Vertex AI then returns the container's response in the API response. For example, if you set this field to `/foo`, then when Vertex AI receives a prediction request, it forwards the request body in a POST request to the `/foo` path on the port of your container specified by the first value of this `ModelContainerSpec`'s ports field. If you don't specify this field, it defaults to the following value when you deploy this Model to an Endpoint: /v1/endpoints/ENDPOINT/deployedModels/DEPLOYED_MODEL:predict The placeholders in this value are replaced as follows: * ENDPOINT: The last segment (following `endpoints/`)of the Endpoint.name][] field of the Endpoint where this Model has been deployed. (Vertex AI makes this value available to your container code as the [`AIP_ENDPOINT_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).) * DEPLOYED_MODEL: DeployedModel.id of the `DeployedModel`. (Vertex AI makes this value available to your container code as the [`AIP_DEPLOYED_MODEL_ID` environment variable](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements#aip-variables).)""" + open_prompt_tuning_pipeline: Optional[ + PublisherModelCallToActionRegionalResourceReferencesDict + ] + """Optional. Open prompt-tuning pipeline of the PublisherModel.""" - shared_memory_size_mb: Optional[int] - """Immutable. The amount of the VM memory to reserve as the shared memory for the model in megabytes.""" + request_access: Optional[PublisherModelCallToActionRegionalResourceReferencesDict] + """Optional. Request for access.""" - startup_probe: Optional[ProbeDict] - """Immutable. Specification for Kubernetes startup probe.""" + view_rest_api: Optional[PublisherModelCallToActionViewRestApiDict] + """Optional. To view Rest API docs.""" -ModelContainerSpecOrDict = Union[ModelContainerSpec, ModelContainerSpecDict] +PublisherModelCallToActionOrDict = Union[ + PublisherModelCallToAction, PublisherModelCallToActionDict +] -class AutoscalingMetricSpec(_common.BaseModel): - """The metric specification that defines the target resource utilization (CPU utilization, accelerator's duty cycle, and so on) for calculating the desired replica count.""" +class PublisherModel(_common.BaseModel): + """Publisher model from Model Garden.""" - metric_name: Optional[str] = Field( + frameworks: Optional[list[str]] = Field( default=None, - description="""Required. The resource metric name. Supported metrics: * For Online Prediction: * `aiplatform.googleapis.com/prediction/online/accelerator/duty_cycle` * `aiplatform.googleapis.com/prediction/online/cpu/utilization` * `aiplatform.googleapis.com/prediction/online/request_count` * `pubsub.googleapis.com/subscription/num_undelivered_messages` * `prometheus.googleapis.com/vertex_dcgm_fi_dev_gpu_util` * `prometheus.googleapis.com/vertex_vllm_gpu_cache_usage_perc` * `prometheus.googleapis.com/vertex_vllm_num_requests_waiting`""", + description="""Optional. Additional information about the model's Frameworks.""", ) - monitored_resource_labels: Optional[dict[str, str]] = Field( + launch_stage: Optional[LaunchStage] = Field( default=None, - description="""Optional. The Cloud Monitoring monitored resource labels as key value pairs used for metrics filtering. See Cloud Monitoring Labels https://cloud.google.com/monitoring/api/v3/metric-model#generic-label-info""", + description="""Optional. Indicates the launch stage of the model.""", ) - target: Optional[int] = Field( + name: Optional[str] = Field( default=None, - description="""The target resource utilization in percentage (1% - 100%) for the given metric; once the real usage deviates from the target by a certain percentage, the machine replicas change. The default value is 60 (representing 60%) if not provided.""", + description="""Output only. Identifier. The resource name of the PublisherModel.""", + ) + open_source_category: Optional[OpenSourceCategory] = Field( + default=None, + description="""Required. Indicates the open source category of the publisher model.""", + ) + parent: Optional[PublisherModelParent] = Field( + default=None, + description="""Optional. The parent that this model was customized from. E.g., Vision API, Natural Language API, LaMDA, T5, etc. Foundation models don't have parents.""", + ) + predict_schemata: Optional[PredictSchemata] = Field( + default=None, + description="""Optional. The schemata that describes formats of the PublisherModel's predictions and explanations as given and returned via PredictionService.Predict.""", + ) + publisher_model_template: Optional[str] = Field( + default=None, + description="""Optional. Output only. Immutable. Used to indicate this model has a publisher model and provide the template of the publisher model resource name.""", + ) + supported_actions: Optional[PublisherModelCallToAction] = Field( + default=None, description="""Optional. Supported call-to-action options.""" + ) + version_id: Optional[str] = Field( + default=None, + description="""Output only. Immutable. The version ID of the PublisherModel. A new version is committed when a new model version is uploaded under an existing model id. It is an auto-incrementing decimal number in string representation.""", + ) + version_state: Optional[VersionState] = Field( + default=None, + description="""Optional. Indicates the state of the model version.""", ) -class AutoscalingMetricSpecDict(TypedDict, total=False): - """The metric specification that defines the target resource utilization (CPU utilization, accelerator's duty cycle, and so on) for calculating the desired replica count.""" - - metric_name: Optional[str] - """Required. The resource metric name. Supported metrics: * For Online Prediction: * `aiplatform.googleapis.com/prediction/online/accelerator/duty_cycle` * `aiplatform.googleapis.com/prediction/online/cpu/utilization` * `aiplatform.googleapis.com/prediction/online/request_count` * `pubsub.googleapis.com/subscription/num_undelivered_messages` * `prometheus.googleapis.com/vertex_dcgm_fi_dev_gpu_util` * `prometheus.googleapis.com/vertex_vllm_gpu_cache_usage_perc` * `prometheus.googleapis.com/vertex_vllm_num_requests_waiting`""" +class PublisherModelDict(TypedDict, total=False): + """Publisher model from Model Garden.""" - monitored_resource_labels: Optional[dict[str, str]] - """Optional. The Cloud Monitoring monitored resource labels as key value pairs used for metrics filtering. See Cloud Monitoring Labels https://cloud.google.com/monitoring/api/v3/metric-model#generic-label-info""" + frameworks: Optional[list[str]] + """Optional. Additional information about the model's Frameworks.""" - target: Optional[int] - """The target resource utilization in percentage (1% - 100%) for the given metric; once the real usage deviates from the target by a certain percentage, the machine replicas change. The default value is 60 (representing 60%) if not provided.""" + launch_stage: Optional[LaunchStage] + """Optional. Indicates the launch stage of the model.""" + name: Optional[str] + """Output only. Identifier. The resource name of the PublisherModel.""" -AutoscalingMetricSpecOrDict = Union[AutoscalingMetricSpec, AutoscalingMetricSpecDict] + open_source_category: Optional[OpenSourceCategory] + """Required. Indicates the open source category of the publisher model.""" + parent: Optional[PublisherModelParentDict] + """Optional. The parent that this model was customized from. E.g., Vision API, Natural Language API, LaMDA, T5, etc. Foundation models don't have parents.""" -class FlexStart(_common.BaseModel): - """FlexStart is used to schedule the deployment workload on DWS resource. It contains the max duration of the deployment.""" + predict_schemata: Optional[PredictSchemataDict] + """Optional. The schemata that describes formats of the PublisherModel's predictions and explanations as given and returned via PredictionService.Predict.""" - max_runtime_duration: Optional[str] = Field( - default=None, - description="""The max duration of the deployment is max_runtime_duration. The deployment will be terminated after the duration. The max_runtime_duration can be set up to 7 days.""", - ) + publisher_model_template: Optional[str] + """Optional. Output only. Immutable. Used to indicate this model has a publisher model and provide the template of the publisher model resource name.""" + supported_actions: Optional[PublisherModelCallToActionDict] + """Optional. Supported call-to-action options.""" -class FlexStartDict(TypedDict, total=False): - """FlexStart is used to schedule the deployment workload on DWS resource. It contains the max duration of the deployment.""" + version_id: Optional[str] + """Output only. Immutable. The version ID of the PublisherModel. A new version is committed when a new model version is uploaded under an existing model id. It is an auto-incrementing decimal number in string representation.""" - max_runtime_duration: Optional[str] - """The max duration of the deployment is max_runtime_duration. The deployment will be terminated after the duration. The max_runtime_duration can be set up to 7 days.""" + version_state: Optional[VersionState] + """Optional. Indicates the state of the model version.""" -FlexStartOrDict = Union[FlexStart, FlexStartDict] +PublisherModelOrDict = Union[PublisherModel, PublisherModelDict] -class DedicatedResourcesScaleToZeroSpec(_common.BaseModel): - """Specification for scale-to-zero feature.""" +class ListPublisherModelsResponse(_common.BaseModel): + """Response for listing publisher models.""" - idle_scaledown_period: Optional[str] = Field( - default=None, - description="""Optional. Duration of no traffic before scaling to zero. [MinValue=300] (5 minutes) [MaxValue=28800] (8 hours)""", + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - min_scaleup_period: Optional[str] = Field( + next_page_token: Optional[str] = Field( default=None, - description="""Optional. Minimum duration that a deployment will be scaled up before traffic is evaluated for potential scale-down. [MinValue=300] (5 minutes) [MaxValue=28800] (8 hours)""", + description="""A token to retrieve next page of results. Pass to ListPublisherModels.page_token to obtain that page.""", + ) + publisher_models: Optional[list[PublisherModel]] = Field( + default=None, description="""List of PublisherModels in the requested page.""" ) -class DedicatedResourcesScaleToZeroSpecDict(TypedDict, total=False): - """Specification for scale-to-zero feature.""" +class ListPublisherModelsResponseDict(TypedDict, total=False): + """Response for listing publisher models.""" - idle_scaledown_period: Optional[str] - """Optional. Duration of no traffic before scaling to zero. [MinValue=300] (5 minutes) [MaxValue=28800] (8 hours)""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - min_scaleup_period: Optional[str] - """Optional. Minimum duration that a deployment will be scaled up before traffic is evaluated for potential scale-down. [MinValue=300] (5 minutes) [MaxValue=28800] (8 hours)""" + next_page_token: Optional[str] + """A token to retrieve next page of results. Pass to ListPublisherModels.page_token to obtain that page.""" + + publisher_models: Optional[list[PublisherModelDict]] + """List of PublisherModels in the requested page.""" -DedicatedResourcesScaleToZeroSpecOrDict = Union[ - DedicatedResourcesScaleToZeroSpec, DedicatedResourcesScaleToZeroSpecDict +ListPublisherModelsResponseOrDict = Union[ + ListPublisherModelsResponse, ListPublisherModelsResponseDict ] -class DedicatedResources(_common.BaseModel): - """A description of resources that are dedicated to a DeployedModel or DeployedIndex, and that need a higher degree of manual configuration.""" +class GetPublisherModelConfig(_common.BaseModel): + """Config for getting a publisher model.""" - autoscaling_metric_specs: Optional[list[AutoscalingMetricSpec]] = Field( - default=None, - description="""Immutable. The metric specifications that overrides a resource utilization metric (CPU utilization, accelerator's duty cycle, and so on) target value (default to 60 if not set). At most one entry is allowed per metric. If machine_spec.accelerator_count is above 0, the autoscaling will be based on both CPU utilization and accelerator's duty cycle metrics and scale up when either metrics exceeds its target value while scale down if both metrics are under their target value. The default target value is 60 for both metrics. If machine_spec.accelerator_count is 0, the autoscaling will be based on CPU utilization metric only with default target value 60 if not explicitly set. For example, in the case of Online Prediction, if you want to override target CPU utilization to 80, you should set autoscaling_metric_specs.metric_name to `aiplatform.googleapis.com/prediction/online/cpu/utilization` and autoscaling_metric_specs.target to `80`.""", - ) - flex_start: Optional[FlexStart] = Field( - default=None, - description="""Optional. Immutable. If set, use DWS resource to schedule the deployment workload. reference: (https://cloud.google.com/blog/products/compute/introducing-dynamic-workload-scheduler)""", - ) - initial_replica_count: Optional[int] = Field( - default=None, - description="""Immutable. Number of initial replicas being deployed on when scaling the workload up from zero or when creating the workload in case min_replica_count = 0. When min_replica_count > 0 (meaning that the scale-to-zero feature is not enabled), initial_replica_count should not be set. When min_replica_count = 0 (meaning that the scale-to-zero feature is enabled), initial_replica_count should be larger than zero, but no greater than max_replica_count.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - machine_spec: Optional[MachineSpec] = Field( + hugging_face_token: Optional[str] = Field( default=None, - description="""Required. Immutable. The specification of a single machine being used.""", + description="""Optional. Hugging Face access token for gated models.""", ) - max_replica_count: Optional[int] = Field( + include_equivalent_model_garden_model_deployment_configs: Optional[bool] = Field( default=None, - description="""Immutable. The maximum number of replicas that may be deployed on when the traffic against it increases. If the requested value is too large, the deployment will error, but if deployment succeeds then the ability to scale to that many replicas is guaranteed (barring service outages). If traffic increases beyond what its replicas at maximum may handle, a portion of the traffic will be dropped. If this value is not provided, will use min_replica_count as the default value. The value of this field impacts the charge against Agent Platform CPU and GPU quotas. Specifically, you will be charged for (max_replica_count * number of cores in the selected machine type) and (max_replica_count * number of GPUs per replica in the selected machine type).""", + description="""Optional. Whether to include the deploy options of equivalent + Model Garden models.""", ) - min_replica_count: Optional[int] = Field( + is_hugging_face_model: Optional[bool] = Field( default=None, - description="""Required. Immutable. The minimum number of machine replicas that will be always deployed on. This value must be greater than or equal to 1. If traffic increases, it may dynamically be deployed onto more replicas, and as traffic decreases, some of these extra replicas may be freed.""", + description="""Optional. Whether the requested model is a Hugging Face model.""", ) - required_replica_count: Optional[int] = Field( - default=None, - description="""Optional. Number of required available replicas for the deployment to succeed. This field is only needed when partial deployment/mutation is desired. If set, the deploy/mutate operation will succeed once available_replica_count reaches required_replica_count, and the rest of the replicas will be retried. If not set, the default required_replica_count will be min_replica_count.""", + + +class GetPublisherModelConfigDict(TypedDict, total=False): + """Config for getting a publisher model.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + hugging_face_token: Optional[str] + """Optional. Hugging Face access token for gated models.""" + + include_equivalent_model_garden_model_deployment_configs: Optional[bool] + """Optional. Whether to include the deploy options of equivalent + Model Garden models.""" + + is_hugging_face_model: Optional[bool] + """Optional. Whether the requested model is a Hugging Face model.""" + + +GetPublisherModelConfigOrDict = Union[ + GetPublisherModelConfig, GetPublisherModelConfigDict +] + + +class _GetPublisherModelRequestParameters(_common.BaseModel): + """Parameters for getting a publisher model.""" + + name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetPublisherModelConfig] = Field(default=None, description="""""") + + +class _GetPublisherModelRequestParametersDict(TypedDict, total=False): + """Parameters for getting a publisher model.""" + + name: Optional[str] + """""" + + config: Optional[GetPublisherModelConfigDict] + """""" + + +_GetPublisherModelRequestParametersOrDict = Union[ + _GetPublisherModelRequestParameters, _GetPublisherModelRequestParametersDict +] + + +class RecommendSpecConfig(_common.BaseModel): + """Config for recommending spec.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - scale_to_zero_spec: Optional[DedicatedResourcesScaleToZeroSpec] = Field( + check_machine_availability: Optional[bool] = Field( default=None, - description="""Optional. Specification for scale-to-zero feature.""", + description="""Whether to check per-region machine availability.""", ) - spot: Optional[bool] = Field( + check_user_quota: Optional[bool] = Field( default=None, - description="""Optional. If true, schedule the deployment workload on [spot VMs](https://cloud.google.com/kubernetes-engine/docs/concepts/spot-vms).""", + description="""Whether to filter to regions with user accelerator quota.""", ) -class DedicatedResourcesDict(TypedDict, total=False): - """A description of resources that are dedicated to a DeployedModel or DeployedIndex, and that need a higher degree of manual configuration.""" +class RecommendSpecConfigDict(TypedDict, total=False): + """Config for recommending spec.""" - autoscaling_metric_specs: Optional[list[AutoscalingMetricSpecDict]] - """Immutable. The metric specifications that overrides a resource utilization metric (CPU utilization, accelerator's duty cycle, and so on) target value (default to 60 if not set). At most one entry is allowed per metric. If machine_spec.accelerator_count is above 0, the autoscaling will be based on both CPU utilization and accelerator's duty cycle metrics and scale up when either metrics exceeds its target value while scale down if both metrics are under their target value. The default target value is 60 for both metrics. If machine_spec.accelerator_count is 0, the autoscaling will be based on CPU utilization metric only with default target value 60 if not explicitly set. For example, in the case of Online Prediction, if you want to override target CPU utilization to 80, you should set autoscaling_metric_specs.metric_name to `aiplatform.googleapis.com/prediction/online/cpu/utilization` and autoscaling_metric_specs.target to `80`.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - flex_start: Optional[FlexStartDict] - """Optional. Immutable. If set, use DWS resource to schedule the deployment workload. reference: (https://cloud.google.com/blog/products/compute/introducing-dynamic-workload-scheduler)""" + check_machine_availability: Optional[bool] + """Whether to check per-region machine availability.""" - initial_replica_count: Optional[int] - """Immutable. Number of initial replicas being deployed on when scaling the workload up from zero or when creating the workload in case min_replica_count = 0. When min_replica_count > 0 (meaning that the scale-to-zero feature is not enabled), initial_replica_count should not be set. When min_replica_count = 0 (meaning that the scale-to-zero feature is enabled), initial_replica_count should be larger than zero, but no greater than max_replica_count.""" + check_user_quota: Optional[bool] + """Whether to filter to regions with user accelerator quota.""" - machine_spec: Optional[MachineSpecDict] - """Required. Immutable. The specification of a single machine being used.""" - max_replica_count: Optional[int] - """Immutable. The maximum number of replicas that may be deployed on when the traffic against it increases. If the requested value is too large, the deployment will error, but if deployment succeeds then the ability to scale to that many replicas is guaranteed (barring service outages). If traffic increases beyond what its replicas at maximum may handle, a portion of the traffic will be dropped. If this value is not provided, will use min_replica_count as the default value. The value of this field impacts the charge against Agent Platform CPU and GPU quotas. Specifically, you will be charged for (max_replica_count * number of cores in the selected machine type) and (max_replica_count * number of GPUs per replica in the selected machine type).""" +RecommendSpecConfigOrDict = Union[RecommendSpecConfig, RecommendSpecConfigDict] - min_replica_count: Optional[int] - """Required. Immutable. The minimum number of machine replicas that will be always deployed on. This value must be greater than or equal to 1. If traffic increases, it may dynamically be deployed onto more replicas, and as traffic decreases, some of these extra replicas may be freed.""" - required_replica_count: Optional[int] - """Optional. Number of required available replicas for the deployment to succeed. This field is only needed when partial deployment/mutation is desired. If set, the deploy/mutate operation will succeed once available_replica_count reaches required_replica_count, and the rest of the replicas will be retried. If not set, the default required_replica_count will be min_replica_count.""" +class _RecommendSpecRequestParameters(_common.BaseModel): + """Parameters for recommending spec.""" - scale_to_zero_spec: Optional[DedicatedResourcesScaleToZeroSpecDict] - """Optional. Specification for scale-to-zero feature.""" + parent: Optional[str] = Field(default=None, description="""""") + gcs_uri: Optional[str] = Field(default=None, description="""""") + config: Optional[RecommendSpecConfig] = Field(default=None, description="""""") - spot: Optional[bool] - """Optional. If true, schedule the deployment workload on [spot VMs](https://cloud.google.com/kubernetes-engine/docs/concepts/spot-vms).""" +class _RecommendSpecRequestParametersDict(TypedDict, total=False): + """Parameters for recommending spec.""" -DedicatedResourcesOrDict = Union[DedicatedResources, DedicatedResourcesDict] + parent: Optional[str] + """""" + gcs_uri: Optional[str] + """""" -class PublisherModelCallToActionDeployDeployMetadata(_common.BaseModel): - """Metadata information about the deployment for managing deployment config.""" + config: Optional[RecommendSpecConfigDict] + """""" - labels: Optional[dict[str, str]] = Field( - default=None, - description="""Optional. Labels for the deployment config. For managing deployment config like verifying, source of deployment config, etc.""", + +_RecommendSpecRequestParametersOrDict = Union[ + _RecommendSpecRequestParameters, _RecommendSpecRequestParametersDict +] + + +class RecommendSpecResponseMachineAndModelContainerSpec(_common.BaseModel): + """A machine and model container spec.""" + + container_spec: Optional[ModelContainerSpec] = Field( + default=None, description="""Output only. The model container spec.""" ) - sample_request: Optional[str] = Field( - default=None, description="""Optional. Sample request for deployed endpoint.""" + machine_spec: Optional[MachineSpec] = Field( + default=None, description="""Output only. The machine spec.""" ) -class PublisherModelCallToActionDeployDeployMetadataDict(TypedDict, total=False): - """Metadata information about the deployment for managing deployment config.""" +class RecommendSpecResponseMachineAndModelContainerSpecDict(TypedDict, total=False): + """A machine and model container spec.""" - labels: Optional[dict[str, str]] - """Optional. Labels for the deployment config. For managing deployment config like verifying, source of deployment config, etc.""" + container_spec: Optional[ModelContainerSpecDict] + """Output only. The model container spec.""" - sample_request: Optional[str] - """Optional. Sample request for deployed endpoint.""" + machine_spec: Optional[MachineSpecDict] + """Output only. The machine spec.""" -PublisherModelCallToActionDeployDeployMetadataOrDict = Union[ - PublisherModelCallToActionDeployDeployMetadata, - PublisherModelCallToActionDeployDeployMetadataDict, +RecommendSpecResponseMachineAndModelContainerSpecOrDict = Union[ + RecommendSpecResponseMachineAndModelContainerSpec, + RecommendSpecResponseMachineAndModelContainerSpecDict, ] -class LargeModelReference(_common.BaseModel): - """Contains information about the Large Model.""" +class RecommendSpecResponseRecommendation(_common.BaseModel): + """Recommendation of one deployment option for the given custom weights model in one region. Contains the machine and container spec, and user accelerator quota state.""" - name: Optional[str] = Field( + region: Optional[str] = Field( + default=None, description="""The region for the deployment spec (machine).""" + ) + spec: Optional[RecommendSpecResponseMachineAndModelContainerSpec] = Field( default=None, - description="""Required. The unique name of the large Foundation or pre-built model. Like "chat-bison", "text-bison". Or model name with version ID, like "chat-bison@001", "text-bison@005", etc.""", + description="""Output only. The machine and model container specs.""", + ) + user_quota_state: Optional[QuotaState] = Field( + default=None, description="""Output only. The user accelerator quota state.""" ) -class LargeModelReferenceDict(TypedDict, total=False): - """Contains information about the Large Model.""" +class RecommendSpecResponseRecommendationDict(TypedDict, total=False): + """Recommendation of one deployment option for the given custom weights model in one region. Contains the machine and container spec, and user accelerator quota state.""" - name: Optional[str] - """Required. The unique name of the large Foundation or pre-built model. Like "chat-bison", "text-bison". Or model name with version ID, like "chat-bison@001", "text-bison@005", etc.""" + region: Optional[str] + """The region for the deployment spec (machine).""" + spec: Optional[RecommendSpecResponseMachineAndModelContainerSpecDict] + """Output only. The machine and model container specs.""" -LargeModelReferenceOrDict = Union[LargeModelReference, LargeModelReferenceDict] + user_quota_state: Optional[QuotaState] + """Output only. The user accelerator quota state.""" -class PublisherModelCallToActionDeploy(_common.BaseModel): - """Model metadata that is needed for UploadModel or DeployModel/CreateEndpoint requests.""" +RecommendSpecResponseRecommendationOrDict = Union[ + RecommendSpecResponseRecommendation, RecommendSpecResponseRecommendationDict +] - artifact_uri: Optional[str] = Field( - default=None, - description="""Optional. The path to the directory containing the Model artifact and any of its supporting files.""", - ) - automatic_resources: Optional[AutomaticResources] = Field( + +class RecommendSpecResponse(_common.BaseModel): + """Response for recommending spec.""" + + base_model: Optional[str] = Field( default=None, - description="""A description of resources that to large degree are decided by Vertex AI, and require only a modest additional configuration.""", + description="""Output only. The base model used to finetune the custom model.""", ) - container_spec: Optional[ModelContainerSpec] = Field( + recommendations: Optional[list[RecommendSpecResponseRecommendation]] = Field( default=None, - description="""Optional. The specification of the container that is to be used when deploying this Model in Vertex AI. Not present for Large Models.""", + description="""Output only. Recommendations of deployment options for the given custom weights model.""", ) - dedicated_resources: Optional[DedicatedResources] = Field( + specs: Optional[list[RecommendSpecResponseMachineAndModelContainerSpec]] = Field( default=None, - description="""A description of resources that are dedicated to the DeployedModel, and that need a higher degree of manual configuration.""", - ) - deploy_metadata: Optional[PublisherModelCallToActionDeployDeployMetadata] = Field( - default=None, - description="""Optional. Metadata information about this deployment config.""", - ) - deploy_task_name: Optional[str] = Field( - default=None, - description="""Optional. The name of the deploy task (e.g., "text to image generation").""", - ) - large_model_reference: Optional[LargeModelReference] = Field( - default=None, - description="""Optional. Large model reference. When this is set, model_artifact_spec is not needed.""", - ) - model_display_name: Optional[str] = Field( - default=None, description="""Optional. Default model display name.""" - ) - public_artifact_uri: Optional[str] = Field( - default=None, - description="""Optional. The signed URI for ephemeral Cloud Storage access to model artifact.""", - ) - shared_resources: Optional[str] = Field( - default=None, - description="""The resource name of the shared DeploymentResourcePool to deploy on. Format: `projects/{project}/locations/{location}/deploymentResourcePools/{deployment_resource_pool}`""", - ) - title: Optional[str] = Field( - default=None, - description="""Required. The title of the regional resource reference.""", + description="""Output only. The machine and model container specs.""", ) -class PublisherModelCallToActionDeployDict(TypedDict, total=False): - """Model metadata that is needed for UploadModel or DeployModel/CreateEndpoint requests.""" +class RecommendSpecResponseDict(TypedDict, total=False): + """Response for recommending spec.""" - artifact_uri: Optional[str] - """Optional. The path to the directory containing the Model artifact and any of its supporting files.""" + base_model: Optional[str] + """Output only. The base model used to finetune the custom model.""" - automatic_resources: Optional[AutomaticResourcesDict] - """A description of resources that to large degree are decided by Vertex AI, and require only a modest additional configuration.""" + recommendations: Optional[list[RecommendSpecResponseRecommendationDict]] + """Output only. Recommendations of deployment options for the given custom weights model.""" - container_spec: Optional[ModelContainerSpecDict] - """Optional. The specification of the container that is to be used when deploying this Model in Vertex AI. Not present for Large Models.""" + specs: Optional[list[RecommendSpecResponseMachineAndModelContainerSpecDict]] + """Output only. The machine and model container specs.""" - dedicated_resources: Optional[DedicatedResourcesDict] - """A description of resources that are dedicated to the DeployedModel, and that need a higher degree of manual configuration.""" - deploy_metadata: Optional[PublisherModelCallToActionDeployDeployMetadataDict] - """Optional. Metadata information about this deployment config.""" +RecommendSpecResponseOrDict = Union[RecommendSpecResponse, RecommendSpecResponseDict] - deploy_task_name: Optional[str] - """Optional. The name of the deploy task (e.g., "text to image generation").""" - large_model_reference: Optional[LargeModelReferenceDict] - """Optional. Large model reference. When this is set, model_artifact_spec is not needed.""" +class ExportPublisherModelConfig(_common.BaseModel): + """RPC-level config for ``export_publisher_model``.""" - model_display_name: Optional[str] - """Optional. Default model display name.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + destination: Optional[genai_types.GcsDestination] = Field( + default=None, description="""""" + ) - public_artifact_uri: Optional[str] - """Optional. The signed URI for ephemeral Cloud Storage access to model artifact.""" - shared_resources: Optional[str] - """The resource name of the shared DeploymentResourcePool to deploy on. Format: `projects/{project}/locations/{location}/deploymentResourcePools/{deployment_resource_pool}`""" +class ExportPublisherModelConfigDict(TypedDict, total=False): + """RPC-level config for ``export_publisher_model``.""" - title: Optional[str] - """Required. The title of the regional resource reference.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + destination: Optional[genai_types.GcsDestination] + """""" -PublisherModelCallToActionDeployOrDict = Union[ - PublisherModelCallToActionDeploy, PublisherModelCallToActionDeployDict +ExportPublisherModelConfigOrDict = Union[ + ExportPublisherModelConfig, ExportPublisherModelConfigDict ] -class PublisherModelCallToActionDeployGke(_common.BaseModel): - """Configurations for PublisherModel GKE deployment""" +class _ExportPublisherModelRequestParameters(_common.BaseModel): + """Parameters for ``export_publisher_model``.""" - gke_yaml_configs: Optional[list[str]] = Field( - default=None, - description="""Optional. GKE deployment configuration in yaml format.""", + parent: Optional[str] = Field(default=None, description="""""") + name: Optional[str] = Field(default=None, description="""""") + config: Optional[ExportPublisherModelConfig] = Field( + default=None, description="""""" ) -class PublisherModelCallToActionDeployGkeDict(TypedDict, total=False): - """Configurations for PublisherModel GKE deployment""" +class _ExportPublisherModelRequestParametersDict(TypedDict, total=False): + """Parameters for ``export_publisher_model``.""" - gke_yaml_configs: Optional[list[str]] - """Optional. GKE deployment configuration in yaml format.""" + parent: Optional[str] + """""" + name: Optional[str] + """""" -PublisherModelCallToActionDeployGkeOrDict = Union[ - PublisherModelCallToActionDeployGke, PublisherModelCallToActionDeployGkeDict + config: Optional[ExportPublisherModelConfigDict] + """""" + + +_ExportPublisherModelRequestParametersOrDict = Union[ + _ExportPublisherModelRequestParameters, _ExportPublisherModelRequestParametersDict ] -class PublisherModelCallToActionDeployVertex(_common.BaseModel): - """Multiple setups to deploy the PublisherModel.""" +class ExportPublisherModelResponse(_common.BaseModel): + """Response for the ``ExportPublisherModel`` RPC. - multi_deploy_vertex: Optional[list[PublisherModelCallToActionDeploy]] = Field( - default=None, description="""Optional. One click deployment configurations.""" + Fields are re-declared as ``SdkFieldPatch`` (both are already in the + discovery-generated class) so the SDK's dependency on ``destination_uri`` + is visible in one place and proto drift is caught at codegen time + instead of at first user call. + """ + + destination_uri: Optional[str] = Field( + default=None, + description="""Cloud Storage URI where the exported weights were written.""", + ) + publisher_model: Optional[str] = Field( + default=None, + description="""Resource name of the publisher model that was exported.""", ) -class PublisherModelCallToActionDeployVertexDict(TypedDict, total=False): - """Multiple setups to deploy the PublisherModel.""" +class ExportPublisherModelResponseDict(TypedDict, total=False): + """Response for the ``ExportPublisherModel`` RPC. - multi_deploy_vertex: Optional[list[PublisherModelCallToActionDeployDict]] - """Optional. One click deployment configurations.""" + Fields are re-declared as ``SdkFieldPatch`` (both are already in the + discovery-generated class) so the SDK's dependency on ``destination_uri`` + is visible in one place and proto drift is caught at codegen time + instead of at first user call. + """ + destination_uri: Optional[str] + """Cloud Storage URI where the exported weights were written.""" -PublisherModelCallToActionDeployVertexOrDict = Union[ - PublisherModelCallToActionDeployVertex, PublisherModelCallToActionDeployVertexDict + publisher_model: Optional[str] + """Resource name of the publisher model that was exported.""" + + +ExportPublisherModelResponseOrDict = Union[ + ExportPublisherModelResponse, ExportPublisherModelResponseDict ] -class PublisherModelCallToActionOpenFineTuningPipelines(_common.BaseModel): - """Open fine tuning pipelines.""" +class ExportModelOperation(_common.BaseModel): + """Long-running operation returned by ``ExportPublisherModel``.""" - fine_tuning_pipelines: Optional[ - list[PublisherModelCallToActionRegionalResourceReferences] - ] = Field( + name: Optional[str] = Field( default=None, - description="""Required. Regional resource references to fine tuning pipelines.""", + 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[ExportPublisherModelResponse] = Field( + default=None, description="""""" ) -class PublisherModelCallToActionOpenFineTuningPipelinesDict(TypedDict, total=False): - """Open fine tuning pipelines.""" +class ExportModelOperationDict(TypedDict, total=False): + """Long-running operation returned by ``ExportPublisherModel``.""" - fine_tuning_pipelines: Optional[ - list[PublisherModelCallToActionRegionalResourceReferencesDict] - ] - """Required. Regional resource references to fine tuning pipelines.""" + 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.""" -PublisherModelCallToActionOpenFineTuningPipelinesOrDict = Union[ - PublisherModelCallToActionOpenFineTuningPipelines, - PublisherModelCallToActionOpenFineTuningPipelinesDict, -] + 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 PublisherModelCallToActionOpenNotebooks(_common.BaseModel): - """Open notebooks.""" + response: Optional[ExportPublisherModelResponseDict] + """""" - notebooks: Optional[list[PublisherModelCallToActionRegionalResourceReferences]] = ( - Field( - default=None, - description="""Required. Regional resource references to notebooks.""", - ) + +ExportModelOperationOrDict = Union[ExportModelOperation, ExportModelOperationDict] + + +class GetExportPublisherModelOperationConfig(_common.BaseModel): + """Config for ``get_export_publisher_model_operation``.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class PublisherModelCallToActionOpenNotebooksDict(TypedDict, total=False): - """Open notebooks.""" +class GetExportPublisherModelOperationConfigDict(TypedDict, total=False): + """Config for ``get_export_publisher_model_operation``.""" - notebooks: Optional[list[PublisherModelCallToActionRegionalResourceReferencesDict]] - """Required. Regional resource references to notebooks.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -PublisherModelCallToActionOpenNotebooksOrDict = Union[ - PublisherModelCallToActionOpenNotebooks, PublisherModelCallToActionOpenNotebooksDict +GetExportPublisherModelOperationConfigOrDict = Union[ + GetExportPublisherModelOperationConfig, GetExportPublisherModelOperationConfigDict ] -class PublisherModelDocumentation(_common.BaseModel): - """A named piece of documentation.""" +class _GetExportPublisherModelOperationParameters(_common.BaseModel): + """Parameters for polling an ``export_publisher_model`` operation.""" - content: Optional[str] = Field( - default=None, - description="""Required. Content of this piece of document (in Markdown format).""", + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - title: Optional[str] = Field( - default=None, - description="""Required. E.g., OVERVIEW, USE CASES, DOCUMENTATION, SDK & SAMPLES, JAVA, NODE.JS, etc..""", + config: Optional[GetExportPublisherModelOperationConfig] = Field( + default=None, description="""""" ) -class PublisherModelDocumentationDict(TypedDict, total=False): - """A named piece of documentation.""" +class _GetExportPublisherModelOperationParametersDict(TypedDict, total=False): + """Parameters for polling an ``export_publisher_model`` operation.""" - content: Optional[str] - """Required. Content of this piece of document (in Markdown format).""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - title: Optional[str] - """Required. E.g., OVERVIEW, USE CASES, DOCUMENTATION, SDK & SAMPLES, JAVA, NODE.JS, etc..""" + config: Optional[GetExportPublisherModelOperationConfigDict] + """""" -PublisherModelDocumentationOrDict = Union[ - PublisherModelDocumentation, PublisherModelDocumentationDict +_GetExportPublisherModelOperationParametersOrDict = Union[ + _GetExportPublisherModelOperationParameters, + _GetExportPublisherModelOperationParametersDict, ] -class PublisherModelCallToActionViewRestApi(_common.BaseModel): - """Rest API docs.""" +class DeployConfig(_common.BaseModel): + """Config for deploying models.""" - documentations: Optional[list[PublisherModelDocumentation]] = Field( - default=None, description="""Required.""" - ) - title: Optional[str] = Field( - default=None, description="""Required. The title of the view rest API.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class PublisherModelCallToActionViewRestApiDict(TypedDict, total=False): - """Rest API docs.""" - - documentations: Optional[list[PublisherModelDocumentationDict]] - """Required.""" +class DeployConfigDict(TypedDict, total=False): + """Config for deploying models.""" - title: Optional[str] - """Required. The title of the view rest API.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -PublisherModelCallToActionViewRestApiOrDict = Union[ - PublisherModelCallToActionViewRestApi, PublisherModelCallToActionViewRestApiDict -] +DeployConfigOrDict = Union[DeployConfig, DeployConfigDict] -class PublisherModelCallToAction(_common.BaseModel): - """Actions could take on this Publisher Model.""" +class DeployRequestCustomModel(_common.BaseModel): + """The custom model to deploy from model weights in a Google Cloud Storage URI or Model Registry model.""" - create_application: Optional[ - PublisherModelCallToActionRegionalResourceReferences - ] = Field( + gcs_uri: Optional[str] = Field( default=None, - description="""Optional. Create application using the PublisherModel.""", + description="""Immutable. The Google Cloud Storage URI of the custom model, storing weights and config files (which can be used to infer the base model).""", ) - deploy: Optional[PublisherModelCallToActionDeploy] = Field( + model_id: Optional[str] = Field( default=None, - description="""Optional. Deploy the PublisherModel to Vertex Endpoint.""", + description="""Optional. Deprecated. Use ModelConfig.model_user_id instead.""", ) - deploy_gke: Optional[PublisherModelCallToActionDeployGke] = Field( + + +class DeployRequestCustomModelDict(TypedDict, total=False): + """The custom model to deploy from model weights in a Google Cloud Storage URI or Model Registry model.""" + + gcs_uri: Optional[str] + """Immutable. The Google Cloud Storage URI of the custom model, storing weights and config files (which can be used to infer the base model).""" + + model_id: Optional[str] + """Optional. Deprecated. Use ModelConfig.model_user_id instead.""" + + +DeployRequestCustomModelOrDict = Union[ + DeployRequestCustomModel, DeployRequestCustomModelDict +] + + +class DeployRequestModelConfig(_common.BaseModel): + """The model config to use for the deployment.""" + + accept_eula: Optional[bool] = Field( default=None, - description="""Optional. Deploy PublisherModel to Google Kubernetes Engine.""", + description="""Optional. Whether the user accepts the End User License Agreement (EULA) for the model.""", ) - multi_deploy_vertex: Optional[PublisherModelCallToActionDeployVertex] = Field( + container_spec: Optional[ModelContainerSpec] = Field( default=None, - description="""Optional. Multiple setups to deploy the PublisherModel to Vertex Endpoint.""", + description="""Optional. The specification of the container that is to be used when deploying. If not set, the default container spec will be used.""", ) - open_evaluation_pipeline: Optional[ - PublisherModelCallToActionRegionalResourceReferences - ] = Field( + hugging_face_access_token: Optional[str] = Field( default=None, - description="""Optional. Open evaluation pipeline of the PublisherModel.""", + description="""Optional. The Hugging Face read access token used to access the model artifacts of gated models.""", ) - open_fine_tuning_pipeline: Optional[ - PublisherModelCallToActionRegionalResourceReferences - ] = Field( + hugging_face_cache_enabled: Optional[bool] = Field( default=None, - description="""Optional. Open fine-tuning pipeline of the PublisherModel.""", + description="""Optional. If true, the model will deploy with a cached version instead of directly downloading the model artifacts from Hugging Face. This is suitable for VPC-SC users with limited internet access.""", ) - open_fine_tuning_pipelines: Optional[ - PublisherModelCallToActionOpenFineTuningPipelines - ] = Field( + model_display_name: Optional[str] = Field( default=None, - description="""Optional. Open fine-tuning pipelines of the PublisherModel.""", - ) - open_generation_ai_studio: Optional[ - PublisherModelCallToActionRegionalResourceReferences - ] = Field(default=None, description="""Optional. Open in Generation AI Studio.""") - open_genie: Optional[PublisherModelCallToActionRegionalResourceReferences] = Field( - default=None, description="""Optional. Open Genie / Playground.""" - ) - open_notebook: Optional[PublisherModelCallToActionRegionalResourceReferences] = ( - Field( - default=None, - description="""Optional. Open notebook of the PublisherModel.""", - ) - ) - open_notebooks: Optional[PublisherModelCallToActionOpenNotebooks] = Field( - default=None, description="""Optional. Open notebooks of the PublisherModel.""" + description="""Optional. The user-specified display name of the uploaded model. If not set, a default name will be used.""", ) - open_prompt_tuning_pipeline: Optional[ - PublisherModelCallToActionRegionalResourceReferences - ] = Field( + model_user_id: Optional[str] = Field( default=None, - description="""Optional. Open prompt-tuning pipeline of the PublisherModel.""", - ) - request_access: Optional[PublisherModelCallToActionRegionalResourceReferences] = ( - Field(default=None, description="""Optional. Request for access.""") - ) - view_rest_api: Optional[PublisherModelCallToActionViewRestApi] = Field( - default=None, description="""Optional. To view Rest API docs.""" + description="""Optional. The ID to use for the uploaded Model, which will become the final component of the model resource name. When not provided, Vertex AI will generate a value for this ID. When Model Registry model is provided, this field will be ignored. This value may be up to 63 characters, and valid characters are `[a-z0-9_-]`. The first character cannot be a number or hyphen.""", ) -class PublisherModelCallToActionDict(TypedDict, total=False): - """Actions could take on this Publisher Model.""" - - create_application: Optional[ - PublisherModelCallToActionRegionalResourceReferencesDict - ] - """Optional. Create application using the PublisherModel.""" - - deploy: Optional[PublisherModelCallToActionDeployDict] - """Optional. Deploy the PublisherModel to Vertex Endpoint.""" - - deploy_gke: Optional[PublisherModelCallToActionDeployGkeDict] - """Optional. Deploy PublisherModel to Google Kubernetes Engine.""" - - multi_deploy_vertex: Optional[PublisherModelCallToActionDeployVertexDict] - """Optional. Multiple setups to deploy the PublisherModel to Vertex Endpoint.""" - - open_evaluation_pipeline: Optional[ - PublisherModelCallToActionRegionalResourceReferencesDict - ] - """Optional. Open evaluation pipeline of the PublisherModel.""" - - open_fine_tuning_pipeline: Optional[ - PublisherModelCallToActionRegionalResourceReferencesDict - ] - """Optional. Open fine-tuning pipeline of the PublisherModel.""" - - open_fine_tuning_pipelines: Optional[ - PublisherModelCallToActionOpenFineTuningPipelinesDict - ] - """Optional. Open fine-tuning pipelines of the PublisherModel.""" - - open_generation_ai_studio: Optional[ - PublisherModelCallToActionRegionalResourceReferencesDict - ] - """Optional. Open in Generation AI Studio.""" +class DeployRequestModelConfigDict(TypedDict, total=False): + """The model config to use for the deployment.""" - open_genie: Optional[PublisherModelCallToActionRegionalResourceReferencesDict] - """Optional. Open Genie / Playground.""" + accept_eula: Optional[bool] + """Optional. Whether the user accepts the End User License Agreement (EULA) for the model.""" - open_notebook: Optional[PublisherModelCallToActionRegionalResourceReferencesDict] - """Optional. Open notebook of the PublisherModel.""" + container_spec: Optional[ModelContainerSpecDict] + """Optional. The specification of the container that is to be used when deploying. If not set, the default container spec will be used.""" - open_notebooks: Optional[PublisherModelCallToActionOpenNotebooksDict] - """Optional. Open notebooks of the PublisherModel.""" + hugging_face_access_token: Optional[str] + """Optional. The Hugging Face read access token used to access the model artifacts of gated models.""" - open_prompt_tuning_pipeline: Optional[ - PublisherModelCallToActionRegionalResourceReferencesDict - ] - """Optional. Open prompt-tuning pipeline of the PublisherModel.""" + hugging_face_cache_enabled: Optional[bool] + """Optional. If true, the model will deploy with a cached version instead of directly downloading the model artifacts from Hugging Face. This is suitable for VPC-SC users with limited internet access.""" - request_access: Optional[PublisherModelCallToActionRegionalResourceReferencesDict] - """Optional. Request for access.""" + model_display_name: Optional[str] + """Optional. The user-specified display name of the uploaded model. If not set, a default name will be used.""" - view_rest_api: Optional[PublisherModelCallToActionViewRestApiDict] - """Optional. To view Rest API docs.""" + model_user_id: Optional[str] + """Optional. The ID to use for the uploaded Model, which will become the final component of the model resource name. When not provided, Vertex AI will generate a value for this ID. When Model Registry model is provided, this field will be ignored. This value may be up to 63 characters, and valid characters are `[a-z0-9_-]`. The first character cannot be a number or hyphen.""" -PublisherModelCallToActionOrDict = Union[ - PublisherModelCallToAction, PublisherModelCallToActionDict +DeployRequestModelConfigOrDict = Union[ + DeployRequestModelConfig, DeployRequestModelConfigDict ] -class PublisherModel(_common.BaseModel): - """Publisher model from Model Garden.""" +class PSCAutomationConfig(_common.BaseModel): + """PSC config that is used to automatically create PSC endpoints in the user projects.""" - frameworks: Optional[list[str]] = Field( - default=None, - description="""Optional. Additional information about the model's Frameworks.""", - ) - launch_stage: Optional[LaunchStage] = Field( - default=None, - description="""Optional. Indicates the launch stage of the model.""", - ) - name: Optional[str] = Field( - default=None, - description="""Output only. Identifier. The resource name of the PublisherModel.""", - ) - open_source_category: Optional[OpenSourceCategory] = Field( + error_message: Optional[str] = Field( default=None, - description="""Required. Indicates the open source category of the publisher model.""", + description="""Output only. Error message if the PSC service automation failed.""", ) - parent: Optional[PublisherModelParent] = Field( + forwarding_rule: Optional[str] = Field( default=None, - description="""Optional. The parent that this model was customized from. E.g., Vision API, Natural Language API, LaMDA, T5, etc. Foundation models don't have parents.""", + description="""Output only. Forwarding rule created by the PSC service automation.""", ) - predict_schemata: Optional[PredictSchemata] = Field( + ip_address: Optional[str] = Field( default=None, - description="""Optional. The schemata that describes formats of the PublisherModel's predictions and explanations as given and returned via PredictionService.Predict.""", + description="""Output only. IP address rule created by the PSC service automation.""", ) - publisher_model_template: Optional[str] = Field( + network: Optional[str] = Field( default=None, - description="""Optional. Output only. Immutable. Used to indicate this model has a publisher model and provide the template of the publisher model resource name.""", - ) - supported_actions: Optional[PublisherModelCallToAction] = Field( - default=None, description="""Optional. Supported call-to-action options.""" + description="""Required. The full name of the Google Compute Engine [network](https://cloud.google.com/compute/docs/networks-and-firewalls#networks). [Format](https://cloud.google.com/compute/docs/reference/rest/v1/networks/get): `projects/{project}/global/networks/{network}`.""", ) - version_id: Optional[str] = Field( + project_id: Optional[str] = Field( default=None, - description="""Output only. Immutable. The version ID of the PublisherModel. A new version is committed when a new model version is uploaded under an existing model id. It is an auto-incrementing decimal number in string representation.""", + description="""Required. Project id used to create forwarding rule.""", ) - version_state: Optional[VersionState] = Field( + state: Optional[PscAutomationState] = Field( default=None, - description="""Optional. Indicates the state of the model version.""", + description="""Output only. The state of the PSC service automation.""", ) -class PublisherModelDict(TypedDict, total=False): - """Publisher model from Model Garden.""" +class PSCAutomationConfigDict(TypedDict, total=False): + """PSC config that is used to automatically create PSC endpoints in the user projects.""" - frameworks: Optional[list[str]] - """Optional. Additional information about the model's Frameworks.""" + error_message: Optional[str] + """Output only. Error message if the PSC service automation failed.""" - launch_stage: Optional[LaunchStage] - """Optional. Indicates the launch stage of the model.""" + forwarding_rule: Optional[str] + """Output only. Forwarding rule created by the PSC service automation.""" - name: Optional[str] - """Output only. Identifier. The resource name of the PublisherModel.""" + ip_address: Optional[str] + """Output only. IP address rule created by the PSC service automation.""" - open_source_category: Optional[OpenSourceCategory] - """Required. Indicates the open source category of the publisher model.""" + network: Optional[str] + """Required. The full name of the Google Compute Engine [network](https://cloud.google.com/compute/docs/networks-and-firewalls#networks). [Format](https://cloud.google.com/compute/docs/reference/rest/v1/networks/get): `projects/{project}/global/networks/{network}`.""" - parent: Optional[PublisherModelParentDict] - """Optional. The parent that this model was customized from. E.g., Vision API, Natural Language API, LaMDA, T5, etc. Foundation models don't have parents.""" + project_id: Optional[str] + """Required. Project id used to create forwarding rule.""" - predict_schemata: Optional[PredictSchemataDict] - """Optional. The schemata that describes formats of the PublisherModel's predictions and explanations as given and returned via PredictionService.Predict.""" + state: Optional[PscAutomationState] + """Output only. The state of the PSC service automation.""" - publisher_model_template: Optional[str] - """Optional. Output only. Immutable. Used to indicate this model has a publisher model and provide the template of the publisher model resource name.""" - supported_actions: Optional[PublisherModelCallToActionDict] - """Optional. Supported call-to-action options.""" +PSCAutomationConfigOrDict = Union[PSCAutomationConfig, PSCAutomationConfigDict] - version_id: Optional[str] - """Output only. Immutable. The version ID of the PublisherModel. A new version is committed when a new model version is uploaded under an existing model id. It is an auto-incrementing decimal number in string representation.""" - version_state: Optional[VersionState] - """Optional. Indicates the state of the model version.""" - - -PublisherModelOrDict = Union[PublisherModel, PublisherModelDict] - - -class ListPublisherModelsResponse(_common.BaseModel): - """Response for listing publisher models.""" +class PrivateServiceConnectConfig(_common.BaseModel): + """Represents configuration for private service connect.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + enable_private_service_connect: Optional[bool] = Field( + default=None, + description="""Required. If true, expose the IndexEndpoint via private service connect.""", ) - next_page_token: Optional[str] = Field( + enable_secure_private_service_connect: Optional[bool] = Field( default=None, - description="""A token to retrieve next page of results. Pass to ListPublisherModels.page_token to obtain that page.""", + description="""Optional. If set to true, enable secure private service connect with IAM authorization. Otherwise, private service connect will be done without authorization. Note latency will be slightly increased if authorization is enabled.""", ) - publisher_models: Optional[list[PublisherModel]] = Field( - default=None, description="""List of PublisherModels in the requested page.""" + project_allowlist: Optional[list[str]] = Field( + default=None, + description="""A list of Projects from which the forwarding rule will target the service attachment.""", + ) + psc_automation_configs: Optional[list[PSCAutomationConfig]] = Field( + default=None, + description="""Optional. List of projects and networks where the PSC endpoints will be created. This field is used by Online Inference(Prediction) only.""", + ) + service_attachment: Optional[str] = Field( + default=None, + description="""Output only. The name of the generated service attachment resource. This is only populated if the endpoint is deployed with PrivateServiceConnect.""", ) -class ListPublisherModelsResponseDict(TypedDict, total=False): - """Response for listing publisher models.""" +class PrivateServiceConnectConfigDict(TypedDict, total=False): + """Represents configuration for private service connect.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + enable_private_service_connect: Optional[bool] + """Required. If true, expose the IndexEndpoint via private service connect.""" - next_page_token: Optional[str] - """A token to retrieve next page of results. Pass to ListPublisherModels.page_token to obtain that page.""" + enable_secure_private_service_connect: Optional[bool] + """Optional. If set to true, enable secure private service connect with IAM authorization. Otherwise, private service connect will be done without authorization. Note latency will be slightly increased if authorization is enabled.""" - publisher_models: Optional[list[PublisherModelDict]] - """List of PublisherModels in the requested page.""" + project_allowlist: Optional[list[str]] + """A list of Projects from which the forwarding rule will target the service attachment.""" + psc_automation_configs: Optional[list[PSCAutomationConfigDict]] + """Optional. List of projects and networks where the PSC endpoints will be created. This field is used by Online Inference(Prediction) only.""" -ListPublisherModelsResponseOrDict = Union[ - ListPublisherModelsResponse, ListPublisherModelsResponseDict + service_attachment: Optional[str] + """Output only. The name of the generated service attachment resource. This is only populated if the endpoint is deployed with PrivateServiceConnect.""" + + +PrivateServiceConnectConfigOrDict = Union[ + PrivateServiceConnectConfig, PrivateServiceConnectConfigDict ] -class GetPublisherModelConfig(_common.BaseModel): - """Config for getting a publisher model.""" +class DeployRequestEndpointConfig(_common.BaseModel): + """The endpoint config to use for the deployment.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + dedicated_endpoint_disabled: Optional[bool] = Field( + default=None, + description="""Optional. By default, if dedicated endpoint is enabled and private service connect config is not set, the endpoint will be exposed through a dedicated DNS [Endpoint.dedicated_endpoint_dns]. If private service connect config is set, the endpoint will be exposed through private service connect. Your request to the dedicated DNS will be isolated from other users' traffic and will have better performance and reliability. Note: Once you enabled dedicated endpoint, you won't be able to send request to the shared DNS {region}-aiplatform.googleapis.com. The limitations will be removed soon. If this field is set to true, the dedicated endpoint will be disabled and the deployed model will be exposed through the shared DNS {region}-aiplatform.googleapis.com.""", ) - hugging_face_token: Optional[str] = Field( + dedicated_endpoint_enabled: Optional[bool] = Field( default=None, - description="""Optional. Hugging Face access token for gated models.""", + description="""Optional. Deprecated. Use dedicated_endpoint_disabled instead. If true, the endpoint will be exposed through a dedicated DNS [Endpoint.dedicated_endpoint_dns]. Your request to the dedicated DNS will be isolated from other users' traffic and will have better performance and reliability. Note: Once you enabled dedicated endpoint, you won't be able to send request to the shared DNS {region}-aiplatform.googleapis.com. The limitations will be removed soon.""", ) - include_equivalent_model_garden_model_deployment_configs: Optional[bool] = Field( + endpoint_display_name: Optional[str] = Field( default=None, - description="""Optional. Whether to include the deploy options of equivalent - Model Garden models.""", + description="""Optional. The user-specified display name of the endpoint. If not set, a default name will be used.""", ) - is_hugging_face_model: Optional[bool] = Field( + endpoint_user_id: Optional[str] = Field( default=None, - description="""Optional. Whether the requested model is a Hugging Face model.""", + description="""Optional. Immutable. The ID to use for endpoint, which will become the final component of the endpoint resource name. If not provided, Vertex AI will generate a value for this ID. If the first character is a letter, this value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The last character must be a letter or number. If the first character is a number, this value may be up to 9 characters, and valid characters are `[0-9]` with no leading zeros. When using HTTP/JSON, this field is populated based on a query string argument, such as `?endpoint_id=12345`. This is the fallback for fields that are not included in either the URI or the body.""", + ) + labels: Optional[dict[str, str]] = Field( + default=None, + description="""Optional. The labels with user-defined metadata to organize your Endpoints. 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.""", + ) + private_service_connect_config: Optional[PrivateServiceConnectConfig] = Field( + default=None, + description="""Optional. Configuration for private service connect. If set, the endpoint will be exposed through private service connect.""", ) -class GetPublisherModelConfigDict(TypedDict, total=False): - """Config for getting a publisher model.""" +class DeployRequestEndpointConfigDict(TypedDict, total=False): + """The endpoint config to use for the deployment.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + dedicated_endpoint_disabled: Optional[bool] + """Optional. By default, if dedicated endpoint is enabled and private service connect config is not set, the endpoint will be exposed through a dedicated DNS [Endpoint.dedicated_endpoint_dns]. If private service connect config is set, the endpoint will be exposed through private service connect. Your request to the dedicated DNS will be isolated from other users' traffic and will have better performance and reliability. Note: Once you enabled dedicated endpoint, you won't be able to send request to the shared DNS {region}-aiplatform.googleapis.com. The limitations will be removed soon. If this field is set to true, the dedicated endpoint will be disabled and the deployed model will be exposed through the shared DNS {region}-aiplatform.googleapis.com.""" - hugging_face_token: Optional[str] - """Optional. Hugging Face access token for gated models.""" + dedicated_endpoint_enabled: Optional[bool] + """Optional. Deprecated. Use dedicated_endpoint_disabled instead. If true, the endpoint will be exposed through a dedicated DNS [Endpoint.dedicated_endpoint_dns]. Your request to the dedicated DNS will be isolated from other users' traffic and will have better performance and reliability. Note: Once you enabled dedicated endpoint, you won't be able to send request to the shared DNS {region}-aiplatform.googleapis.com. The limitations will be removed soon.""" - include_equivalent_model_garden_model_deployment_configs: Optional[bool] - """Optional. Whether to include the deploy options of equivalent - Model Garden models.""" + endpoint_display_name: Optional[str] + """Optional. The user-specified display name of the endpoint. If not set, a default name will be used.""" - is_hugging_face_model: Optional[bool] - """Optional. Whether the requested model is a Hugging Face model.""" + endpoint_user_id: Optional[str] + """Optional. Immutable. The ID to use for endpoint, which will become the final component of the endpoint resource name. If not provided, Vertex AI will generate a value for this ID. If the first character is a letter, this value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The last character must be a letter or number. If the first character is a number, this value may be up to 9 characters, and valid characters are `[0-9]` with no leading zeros. When using HTTP/JSON, this field is populated based on a query string argument, such as `?endpoint_id=12345`. This is the fallback for fields that are not included in either the URI or the body.""" + labels: Optional[dict[str, str]] + """Optional. The labels with user-defined metadata to organize your Endpoints. 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.""" -GetPublisherModelConfigOrDict = Union[ - GetPublisherModelConfig, GetPublisherModelConfigDict + private_service_connect_config: Optional[PrivateServiceConnectConfigDict] + """Optional. Configuration for private service connect. If set, the endpoint will be exposed through private service connect.""" + + +DeployRequestEndpointConfigOrDict = Union[ + DeployRequestEndpointConfig, DeployRequestEndpointConfigDict ] -class _GetPublisherModelRequestParameters(_common.BaseModel): - """Parameters for getting a publisher model.""" +class DeployRequestDeployConfig(_common.BaseModel): + """The deploy config to use for the deployment.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetPublisherModelConfig] = Field(default=None, description="""""") + dedicated_resources: Optional[DedicatedResources] = Field( + default=None, + description="""Optional. The dedicated resources to use for the endpoint. If not set, the default resources will be used.""", + ) + fast_tryout_enabled: Optional[bool] = Field( + default=None, + description="""Optional. If true, enable the QMT fast tryout feature for this model if possible.""", + ) + system_labels: Optional[dict[str, str]] = Field( + default=None, + description="""Optional. System labels for Model Garden deployments. These labels are managed by Google and for tracking purposes only.""", + ) -class _GetPublisherModelRequestParametersDict(TypedDict, total=False): - """Parameters for getting a publisher model.""" +class DeployRequestDeployConfigDict(TypedDict, total=False): + """The deploy config to use for the deployment.""" - name: Optional[str] - """""" + dedicated_resources: Optional[DedicatedResourcesDict] + """Optional. The dedicated resources to use for the endpoint. If not set, the default resources will be used.""" - config: Optional[GetPublisherModelConfigDict] - """""" + fast_tryout_enabled: Optional[bool] + """Optional. If true, enable the QMT fast tryout feature for this model if possible.""" + system_labels: Optional[dict[str, str]] + """Optional. System labels for Model Garden deployments. These labels are managed by Google and for tracking purposes only.""" -_GetPublisherModelRequestParametersOrDict = Union[ - _GetPublisherModelRequestParameters, _GetPublisherModelRequestParametersDict + +DeployRequestDeployConfigOrDict = Union[ + DeployRequestDeployConfig, DeployRequestDeployConfigDict ] -class RecommendSpecConfig(_common.BaseModel): - """Config for recommending spec.""" +class _DeployRequestParameters(_common.BaseModel): + """Parameters for deployment.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + destination: Optional[str] = Field(default=None, description="""""") + publisher_model_name: Optional[str] = Field(default=None, description="""""") + hugging_face_model_id: Optional[str] = Field(default=None, description="""""") + custom_model: Optional[DeployRequestCustomModel] = Field( + default=None, description="""""" ) - check_machine_availability: Optional[bool] = Field( - default=None, - description="""Whether to check per-region machine availability.""", + model_config_val: Optional[DeployRequestModelConfig] = Field( + default=None, description="""""" ) - check_user_quota: Optional[bool] = Field( - default=None, - description="""Whether to filter to regions with user accelerator quota.""", + endpoint_config: Optional[DeployRequestEndpointConfig] = Field( + default=None, description="""""" ) + deploy_config: Optional[DeployRequestDeployConfig] = Field( + default=None, description="""""" + ) + config: Optional[DeployConfig] = Field(default=None, description="""""") -class RecommendSpecConfigDict(TypedDict, total=False): - """Config for recommending spec.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - check_machine_availability: Optional[bool] - """Whether to check per-region machine availability.""" - - check_user_quota: Optional[bool] - """Whether to filter to regions with user accelerator quota.""" - - -RecommendSpecConfigOrDict = Union[RecommendSpecConfig, RecommendSpecConfigDict] +class _DeployRequestParametersDict(TypedDict, total=False): + """Parameters for deployment.""" + destination: Optional[str] + """""" -class _RecommendSpecRequestParameters(_common.BaseModel): - """Parameters for recommending spec.""" + publisher_model_name: Optional[str] + """""" - parent: Optional[str] = Field(default=None, description="""""") - gcs_uri: Optional[str] = Field(default=None, description="""""") - config: Optional[RecommendSpecConfig] = Field(default=None, description="""""") + hugging_face_model_id: Optional[str] + """""" + custom_model: Optional[DeployRequestCustomModelDict] + """""" -class _RecommendSpecRequestParametersDict(TypedDict, total=False): - """Parameters for recommending spec.""" + model_config_val: Optional[DeployRequestModelConfigDict] + """""" - parent: Optional[str] + endpoint_config: Optional[DeployRequestEndpointConfigDict] """""" - gcs_uri: Optional[str] + deploy_config: Optional[DeployRequestDeployConfigDict] """""" - config: Optional[RecommendSpecConfigDict] + config: Optional[DeployConfigDict] """""" -_RecommendSpecRequestParametersOrDict = Union[ - _RecommendSpecRequestParameters, _RecommendSpecRequestParametersDict +_DeployRequestParametersOrDict = Union[ + _DeployRequestParameters, _DeployRequestParametersDict ] -class RecommendSpecResponseMachineAndModelContainerSpec(_common.BaseModel): - """A machine and model container spec.""" +class DeployResponse(_common.BaseModel): + """Response for deployment.""" - container_spec: Optional[ModelContainerSpec] = Field( - default=None, description="""Output only. The model container spec.""" + endpoint: Optional[str] = Field( + default=None, description="""The resource name of the deployed endpoint.""" ) - machine_spec: Optional[MachineSpec] = Field( - default=None, description="""Output only. The machine spec.""" + model: Optional[str] = Field( + default=None, description="""The resource name of the deployed model.""" ) -class RecommendSpecResponseMachineAndModelContainerSpecDict(TypedDict, total=False): - """A machine and model container spec.""" - - container_spec: Optional[ModelContainerSpecDict] - """Output only. The model container spec.""" +class DeployResponseDict(TypedDict, total=False): + """Response for deployment.""" - machine_spec: Optional[MachineSpecDict] - """Output only. The machine spec.""" + endpoint: Optional[str] + """The resource name of the deployed endpoint.""" + model: Optional[str] + """The resource name of the deployed model.""" -RecommendSpecResponseMachineAndModelContainerSpecOrDict = Union[ - RecommendSpecResponseMachineAndModelContainerSpec, - RecommendSpecResponseMachineAndModelContainerSpecDict, -] +DeployResponseOrDict = Union[DeployResponse, DeployResponseDict] -class RecommendSpecResponseRecommendation(_common.BaseModel): - """Recommendation of one deployment option for the given custom weights model in one region. Contains the machine and container spec, and user accelerator quota state.""" - region: Optional[str] = Field( - default=None, description="""The region for the deployment spec (machine).""" - ) - spec: Optional[RecommendSpecResponseMachineAndModelContainerSpec] = Field( +class DeployModelOperation(_common.BaseModel): + """Operation that has a deploy response.""" + + name: Optional[str] = Field( default=None, - description="""Output only. The machine and model container specs.""", - ) - user_quota_state: Optional[QuotaState] = Field( - default=None, description="""Output only. The user accelerator quota state.""" + 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 RecommendSpecResponseRecommendationDict(TypedDict, total=False): - """Recommendation of one deployment option for the given custom weights model in one region. Contains the machine and container spec, and user accelerator quota state.""" - - region: Optional[str] - """The region for the deployment spec (machine).""" - - spec: Optional[RecommendSpecResponseMachineAndModelContainerSpecDict] - """Output only. The machine and model container specs.""" - - user_quota_state: Optional[QuotaState] - """Output only. The user accelerator quota state.""" - - -RecommendSpecResponseRecommendationOrDict = Union[ - RecommendSpecResponseRecommendation, RecommendSpecResponseRecommendationDict -] - - -class RecommendSpecResponse(_common.BaseModel): - """Response for recommending spec.""" - - base_model: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Output only. The base model used to finetune the custom 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.""", ) - recommendations: Optional[list[RecommendSpecResponseRecommendation]] = Field( + done: Optional[bool] = Field( default=None, - description="""Output only. Recommendations of deployment options for the given custom weights model.""", + 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.""", ) - specs: Optional[list[RecommendSpecResponseMachineAndModelContainerSpec]] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Output only. The machine and model container specs.""", + description="""The error result of the operation in case of failure or cancellation.""", ) + response: Optional[DeployResponse] = Field(default=None, description="""""") -class RecommendSpecResponseDict(TypedDict, total=False): - """Response for recommending spec.""" +class DeployModelOperationDict(TypedDict, total=False): + """Operation that has a deploy response.""" - base_model: Optional[str] - """Output only. The base model used to finetune the custom model.""" + 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}`.""" - recommendations: Optional[list[RecommendSpecResponseRecommendationDict]] - """Output only. Recommendations of deployment options for the given custom weights model.""" + 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.""" - specs: Optional[list[RecommendSpecResponseMachineAndModelContainerSpecDict]] - """Output only. The machine and model container specs.""" + 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.""" -RecommendSpecResponseOrDict = Union[RecommendSpecResponse, RecommendSpecResponseDict] + response: Optional[DeployResponseDict] + """""" -class ExportPublisherModelConfig(_common.BaseModel): - """RPC-level config for ``export_publisher_model``.""" +DeployModelOperationOrDict = Union[DeployModelOperation, DeployModelOperationDict] + + +class GetDeployOperationConfig(_common.BaseModel): + """Config for ``get_deploy_publisher_model_operation``.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - destination: Optional[genai_types.GcsDestination] = Field( - default=None, description="""""" - ) -class ExportPublisherModelConfigDict(TypedDict, total=False): - """RPC-level config for ``export_publisher_model``.""" +class GetDeployOperationConfigDict(TypedDict, total=False): + """Config for ``get_deploy_publisher_model_operation``.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - destination: Optional[genai_types.GcsDestination] - """""" - -ExportPublisherModelConfigOrDict = Union[ - ExportPublisherModelConfig, ExportPublisherModelConfigDict +GetDeployOperationConfigOrDict = Union[ + GetDeployOperationConfig, GetDeployOperationConfigDict ] -class _ExportPublisherModelRequestParameters(_common.BaseModel): - """Parameters for ``export_publisher_model``.""" +class _GetDeployOperationParameters(_common.BaseModel): + """Parameters for polling a ``deploy`` operation.""" - parent: Optional[str] = Field(default=None, description="""""") - name: Optional[str] = Field(default=None, description="""""") - config: Optional[ExportPublisherModelConfig] = Field( - default=None, description="""""" + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) + config: Optional[GetDeployOperationConfig] = Field(default=None, description="""""") -class _ExportPublisherModelRequestParametersDict(TypedDict, total=False): - """Parameters for ``export_publisher_model``.""" - - parent: Optional[str] - """""" +class _GetDeployOperationParametersDict(TypedDict, total=False): + """Parameters for polling a ``deploy`` operation.""" - name: Optional[str] - """""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - config: Optional[ExportPublisherModelConfigDict] + config: Optional[GetDeployOperationConfigDict] """""" -_ExportPublisherModelRequestParametersOrDict = Union[ - _ExportPublisherModelRequestParameters, _ExportPublisherModelRequestParametersDict +_GetDeployOperationParametersOrDict = Union[ + _GetDeployOperationParameters, _GetDeployOperationParametersDict ] -class ExportPublisherModelResponse(_common.BaseModel): - """Response for the ``ExportPublisherModel`` RPC. - - Fields are re-declared as ``SdkFieldPatch`` (both are already in the - discovery-generated class) so the SDK's dependency on ``destination_uri`` - is visible in one place and proto drift is caught at codegen time - instead of at first user call. - """ +class CreateRuntimeFeedbackEntryConfig(_common.BaseModel): + """Config for creating a Feedback Entry.""" - destination_uri: Optional[str] = Field( + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + feedback_labels: Optional[list[str]] = Field( default=None, - description="""Cloud Storage URI where the exported weights were written.""", + description="""Specific labels for feedback (non-factual, offensive, etc.).""", ) - publisher_model: Optional[str] = Field( + feedback_text: Optional[str] = Field( default=None, - description="""Resource name of the publisher model that was exported.""", + description="""Qualitative free-form comments provided by the user.""", + ) + user_id: Optional[str] = Field( + default=None, description="""User provided identifier.""" + ) + source: Optional[str] = Field( + default=None, description="""Originating UI surface (e.g. 'ADK Web UI').""" + ) + custom_metadata: Optional[dict[str, str]] = Field( + default=None, + description=""" Additional key-value metadata associated with the feedback. Allows the collect data for which there is no dedicated field in the resource, ex. version, LLM temperature etc.""", + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) -class ExportPublisherModelResponseDict(TypedDict, total=False): - """Response for the ``ExportPublisherModel`` RPC. +class CreateRuntimeFeedbackEntryConfigDict(TypedDict, total=False): + """Config for creating a Feedback Entry.""" - Fields are re-declared as ``SdkFieldPatch`` (both are already in the - discovery-generated class) so the SDK's dependency on ``destination_uri`` - is visible in one place and proto drift is caught at codegen time - instead of at first user call. - """ + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - destination_uri: Optional[str] - """Cloud Storage URI where the exported weights were written.""" + feedback_labels: Optional[list[str]] + """Specific labels for feedback (non-factual, offensive, etc.).""" - publisher_model: Optional[str] - """Resource name of the publisher model that was exported.""" + feedback_text: Optional[str] + """Qualitative free-form comments provided by the user.""" + user_id: Optional[str] + """User provided identifier.""" -ExportPublisherModelResponseOrDict = Union[ - ExportPublisherModelResponse, ExportPublisherModelResponseDict + source: Optional[str] + """Originating UI surface (e.g. 'ADK Web UI').""" + + custom_metadata: Optional[dict[str, str]] + """ Additional key-value metadata associated with the feedback. Allows the collect data for which there is no dedicated field in the resource, ex. version, LLM temperature etc.""" + + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" + + +CreateRuntimeFeedbackEntryConfigOrDict = Union[ + CreateRuntimeFeedbackEntryConfig, CreateRuntimeFeedbackEntryConfigDict ] -class ExportModelOperation(_common.BaseModel): - """Long-running operation returned by ``ExportPublisherModel``.""" +class _CreateRuntimeFeedbackEntryRequestParameters(_common.BaseModel): + """Parameters for creating a Feedback Entry.""" 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}`.""", + description="""Resource name of the Runtime to create the Feedback Entry in.""", ) - 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.""", + feedback_type: Optional[FeedbackType] = Field( + default=None, description="""The type of feedback provided.""" ) - done: Optional[bool] = Field( + session_id: Optional[str] = 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="""The ID of the session to which the feedback relates to.""", ) - error: Optional[dict[str, Any]] = Field( + event_id: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""The ID of the event to which the feedback relates to.""", ) - response: Optional[ExportPublisherModelResponse] = Field( + config: Optional[CreateRuntimeFeedbackEntryConfig] = Field( default=None, description="""""" ) -class ExportModelOperationDict(TypedDict, total=False): - """Long-running operation returned by ``ExportPublisherModel``.""" +class _CreateRuntimeFeedbackEntryRequestParametersDict(TypedDict, total=False): + """Parameters for creating a Feedback Entry.""" 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}`.""" + """Resource name of the Runtime to create the Feedback Entry in.""" - 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.""" + feedback_type: Optional[FeedbackType] + """The type of feedback provided.""" - 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.""" + session_id: Optional[str] + """The ID of the session to which the feedback relates to.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + event_id: Optional[str] + """The ID of the event to which the feedback relates to.""" - response: Optional[ExportPublisherModelResponseDict] + config: Optional[CreateRuntimeFeedbackEntryConfigDict] """""" -ExportModelOperationOrDict = Union[ExportModelOperation, ExportModelOperationDict] +_CreateRuntimeFeedbackEntryRequestParametersOrDict = Union[ + _CreateRuntimeFeedbackEntryRequestParameters, + _CreateRuntimeFeedbackEntryRequestParametersDict, +] -class GetExportPublisherModelOperationConfig(_common.BaseModel): - """Config for ``get_export_publisher_model_operation``.""" +class FeedbackEntry(_common.BaseModel): + """A Feedback Entry.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. The time at which the entry was created.""", ) - - -class GetExportPublisherModelOperationConfigDict(TypedDict, total=False): - """Config for ``get_export_publisher_model_operation``.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -GetExportPublisherModelOperationConfigOrDict = Union[ - GetExportPublisherModelOperationConfig, GetExportPublisherModelOperationConfigDict -] - - -class _GetExportPublisherModelOperationParameters(_common.BaseModel): - """Parameters for polling an ``export_publisher_model`` operation.""" - - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + custom_metadata: Optional[dict[str, str]] = Field( + default=None, + description="""Optional. Additional key-value metadata associated with the feedback.""", ) - config: Optional[GetExportPublisherModelOperationConfig] = Field( - default=None, description="""""" + event_id: Optional[str] = Field( + default=None, + description="""Required. The ID of the event within the session that the feedback relates to.""", + ) + feedback_labels: Optional[list[str]] = Field( + default=None, description="""feedbackLabels""" + ) + feedback_text: Optional[str] = Field( + default=None, + description="""Optional. Qualitative free-form comments provided by the user.""", + ) + feedback_type: Optional[FeedbackType] = Field( + default=None, + description="""Required. The coarse-grained type of feedback provided by the user. Must be set to a value other than `FEEDBACK_TYPE_UNSPECIFIED`.""", + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. The resource name. Assigned by the server on create. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/feedbackEntries/{feedback_entry}`""", + ) + session_id: Optional[str] = Field( + default=None, + description="""Required. The ID of the session that the feedback relates to.""", + ) + source: Optional[str] = Field( + default=None, + description="""Optional. The surface that the feedback originated from.""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. The time at which the entry was most recently updated.""", + ) + user_id: Optional[str] = Field( + default=None, + description="""Optional. A caller-supplied identifier for the user who provided the feedback. The semantics of this field (for example whether it is an opaque token, a hashed value, or a user-visible identifier) are determined by the calling application.""", ) -class _GetExportPublisherModelOperationParametersDict(TypedDict, total=False): - """Parameters for polling an ``export_publisher_model`` operation.""" +class FeedbackEntryDict(TypedDict, total=False): + """A Feedback Entry.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + create_time: Optional[datetime.datetime] + """Output only. The time at which the entry was created.""" - config: Optional[GetExportPublisherModelOperationConfigDict] - """""" + custom_metadata: Optional[dict[str, str]] + """Optional. Additional key-value metadata associated with the feedback.""" + event_id: Optional[str] + """Required. The ID of the event within the session that the feedback relates to.""" -_GetExportPublisherModelOperationParametersOrDict = Union[ - _GetExportPublisherModelOperationParameters, - _GetExportPublisherModelOperationParametersDict, -] + feedback_labels: Optional[list[str]] + """feedbackLabels""" + feedback_text: Optional[str] + """Optional. Qualitative free-form comments provided by the user.""" -class DeployConfig(_common.BaseModel): - """Config for deploying models.""" + feedback_type: Optional[FeedbackType] + """Required. The coarse-grained type of feedback provided by the user. Must be set to a value other than `FEEDBACK_TYPE_UNSPECIFIED`.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + name: Optional[str] + """Identifier. The resource name. Assigned by the server on create. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/feedbackEntries/{feedback_entry}`""" + session_id: Optional[str] + """Required. The ID of the session that the feedback relates to.""" -class DeployConfigDict(TypedDict, total=False): - """Config for deploying models.""" + source: Optional[str] + """Optional. The surface that the feedback originated from.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + update_time: Optional[datetime.datetime] + """Output only. The time at which the entry was most recently updated.""" + user_id: Optional[str] + """Optional. A caller-supplied identifier for the user who provided the feedback. The semantics of this field (for example whether it is an opaque token, a hashed value, or a user-visible identifier) are determined by the calling application.""" -DeployConfigOrDict = Union[DeployConfig, DeployConfigDict] +FeedbackEntryOrDict = Union[FeedbackEntry, FeedbackEntryDict] -class DeployRequestCustomModel(_common.BaseModel): - """The custom model to deploy from model weights in a Google Cloud Storage URI or Model Registry model.""" - gcs_uri: Optional[str] = Field( +class RuntimeFeedbackEntryOperation(_common.BaseModel): + """Operation that has a Runtime Feedback Entry as a response.""" + + name: Optional[str] = Field( default=None, - description="""Immutable. The Google Cloud Storage URI of the custom model, storing weights and config files (which can be used to infer the base model).""", + 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}`.""", ) - model_id: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. Deprecated. Use ModelConfig.model_user_id instead.""", + 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[FeedbackEntry] = Field( + default=None, description="""The Runtime Feedback Entry.""" ) -class DeployRequestCustomModelDict(TypedDict, total=False): - """The custom model to deploy from model weights in a Google Cloud Storage URI or Model Registry model.""" +class RuntimeFeedbackEntryOperationDict(TypedDict, total=False): + """Operation that has a Runtime Feedback Entry as a response.""" - gcs_uri: Optional[str] - """Immutable. The Google Cloud Storage URI of the custom model, storing weights and config files (which can be used to infer the base model).""" + 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}`.""" - model_id: Optional[str] - """Optional. Deprecated. Use ModelConfig.model_user_id instead.""" + 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.""" -DeployRequestCustomModelOrDict = Union[ - DeployRequestCustomModel, DeployRequestCustomModelDict + response: Optional[FeedbackEntryDict] + """The Runtime Feedback Entry.""" + + +RuntimeFeedbackEntryOperationOrDict = Union[ + RuntimeFeedbackEntryOperation, RuntimeFeedbackEntryOperationDict ] -class DeployRequestModelConfig(_common.BaseModel): - """The model config to use for the deployment.""" +class DeleteRuntimeFeedbackEntryConfig(_common.BaseModel): + """Config for deleting a Feedback Entry.""" - accept_eula: Optional[bool] = Field( - default=None, - description="""Optional. Whether the user accepts the End User License Agreement (EULA) for the model.""", - ) - container_spec: Optional[ModelContainerSpec] = Field( - default=None, - description="""Optional. The specification of the container that is to be used when deploying. If not set, the default container spec will be used.""", - ) - hugging_face_access_token: Optional[str] = Field( - default=None, - description="""Optional. The Hugging Face read access token used to access the model artifacts of gated models.""", - ) - hugging_face_cache_enabled: Optional[bool] = Field( - default=None, - description="""Optional. If true, the model will deploy with a cached version instead of directly downloading the model artifacts from Hugging Face. This is suitable for VPC-SC users with limited internet access.""", - ) - model_display_name: Optional[str] = Field( - default=None, - description="""Optional. The user-specified display name of the uploaded model. If not set, a default name will be used.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - model_user_id: Optional[str] = Field( - default=None, - description="""Optional. The ID to use for the uploaded Model, which will become the final component of the model resource name. When not provided, Vertex AI will generate a value for this ID. When Model Registry model is provided, this field will be ignored. This value may be up to 63 characters, and valid characters are `[a-z0-9_-]`. The first character cannot be a number or hyphen.""", + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) -class DeployRequestModelConfigDict(TypedDict, total=False): - """The model config to use for the deployment.""" +class DeleteRuntimeFeedbackEntryConfigDict(TypedDict, total=False): + """Config for deleting a Feedback Entry.""" - accept_eula: Optional[bool] - """Optional. Whether the user accepts the End User License Agreement (EULA) for the model.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - container_spec: Optional[ModelContainerSpecDict] - """Optional. The specification of the container that is to be used when deploying. If not set, the default container spec will be used.""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" - hugging_face_access_token: Optional[str] - """Optional. The Hugging Face read access token used to access the model artifacts of gated models.""" - hugging_face_cache_enabled: Optional[bool] - """Optional. If true, the model will deploy with a cached version instead of directly downloading the model artifacts from Hugging Face. This is suitable for VPC-SC users with limited internet access.""" +DeleteRuntimeFeedbackEntryConfigOrDict = Union[ + DeleteRuntimeFeedbackEntryConfig, DeleteRuntimeFeedbackEntryConfigDict +] - model_display_name: Optional[str] - """Optional. The user-specified display name of the uploaded model. If not set, a default name will be used.""" - model_user_id: Optional[str] - """Optional. The ID to use for the uploaded Model, which will become the final component of the model resource name. When not provided, Vertex AI will generate a value for this ID. When Model Registry model is provided, this field will be ignored. This value may be up to 63 characters, and valid characters are `[a-z0-9_-]`. The first character cannot be a number or hyphen.""" +class _DeleteRuntimeFeedbackEntryRequestParameters(_common.BaseModel): + """Parameters for deleting a Feedback Entry.""" + + name: Optional[str] = Field( + default=None, description="""Name of the Feedback Entry to delete.""" + ) + config: Optional[DeleteRuntimeFeedbackEntryConfig] = Field( + default=None, description="""""" + ) -DeployRequestModelConfigOrDict = Union[ - DeployRequestModelConfig, DeployRequestModelConfigDict +class _DeleteRuntimeFeedbackEntryRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a Feedback Entry.""" + + name: Optional[str] + """Name of the Feedback Entry to delete.""" + + config: Optional[DeleteRuntimeFeedbackEntryConfigDict] + """""" + + +_DeleteRuntimeFeedbackEntryRequestParametersOrDict = Union[ + _DeleteRuntimeFeedbackEntryRequestParameters, + _DeleteRuntimeFeedbackEntryRequestParametersDict, ] -class PSCAutomationConfig(_common.BaseModel): - """PSC config that is used to automatically create PSC endpoints in the user projects.""" +class DeleteRuntimeFeedbackEntryOperation(_common.BaseModel): + """Operation for deleting a Feedback Entry.""" - error_message: Optional[str] = Field( - default=None, - description="""Output only. Error message if the PSC service automation failed.""", - ) - forwarding_rule: Optional[str] = Field( - default=None, - description="""Output only. Forwarding rule created by the PSC service automation.""", - ) - ip_address: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Output only. IP address rule created by the PSC service automation.""", + 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}`.""", ) - network: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Required. The full name of the Google Compute Engine [network](https://cloud.google.com/compute/docs/networks-and-firewalls#networks). [Format](https://cloud.google.com/compute/docs/reference/rest/v1/networks/get): `projects/{project}/global/networks/{network}`.""", + 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.""", ) - project_id: Optional[str] = Field( + done: Optional[bool] = Field( default=None, - description="""Required. Project id used to create forwarding rule.""", + 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.""", ) - state: Optional[PscAutomationState] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Output only. The state of the PSC service automation.""", + description="""The error result of the operation in case of failure or cancellation.""", ) -class PSCAutomationConfigDict(TypedDict, total=False): - """PSC config that is used to automatically create PSC endpoints in the user projects.""" +class DeleteRuntimeFeedbackEntryOperationDict(TypedDict, total=False): + """Operation for deleting a Feedback Entry.""" - error_message: Optional[str] - """Output only. Error message if the PSC service automation failed.""" + 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}`.""" - forwarding_rule: Optional[str] - """Output only. Forwarding rule created by the PSC service automation.""" + 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.""" - ip_address: Optional[str] - """Output only. IP address rule created by the PSC service automation.""" + 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.""" - network: Optional[str] - """Required. The full name of the Google Compute Engine [network](https://cloud.google.com/compute/docs/networks-and-firewalls#networks). [Format](https://cloud.google.com/compute/docs/reference/rest/v1/networks/get): `projects/{project}/global/networks/{network}`.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - project_id: Optional[str] - """Required. Project id used to create forwarding rule.""" - state: Optional[PscAutomationState] - """Output only. The state of the PSC service automation.""" +DeleteRuntimeFeedbackEntryOperationOrDict = Union[ + DeleteRuntimeFeedbackEntryOperation, DeleteRuntimeFeedbackEntryOperationDict +] -PSCAutomationConfigOrDict = Union[PSCAutomationConfig, PSCAutomationConfigDict] +class GetRuntimeFeedbackConfig(_common.BaseModel): + """Config for getting a Feedback Entry.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -class PrivateServiceConnectConfig(_common.BaseModel): - """Represents configuration for private service connect.""" - enable_private_service_connect: Optional[bool] = Field( - default=None, - description="""Required. If true, expose the IndexEndpoint via private service connect.""", - ) - enable_secure_private_service_connect: Optional[bool] = Field( - default=None, - description="""Optional. If set to true, enable secure private service connect with IAM authorization. Otherwise, private service connect will be done without authorization. Note latency will be slightly increased if authorization is enabled.""", - ) - project_allowlist: Optional[list[str]] = Field( - default=None, - description="""A list of Projects from which the forwarding rule will target the service attachment.""", - ) - psc_automation_configs: Optional[list[PSCAutomationConfig]] = Field( - default=None, - description="""Optional. List of projects and networks where the PSC endpoints will be created. This field is used by Online Inference(Prediction) only.""", - ) - service_attachment: Optional[str] = Field( - default=None, - description="""Output only. The name of the generated service attachment resource. This is only populated if the endpoint is deployed with PrivateServiceConnect.""", - ) +class GetRuntimeFeedbackConfigDict(TypedDict, total=False): + """Config for getting a Feedback Entry.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -class PrivateServiceConnectConfigDict(TypedDict, total=False): - """Represents configuration for private service connect.""" - enable_private_service_connect: Optional[bool] - """Required. If true, expose the IndexEndpoint via private service connect.""" +GetRuntimeFeedbackConfigOrDict = Union[ + GetRuntimeFeedbackConfig, GetRuntimeFeedbackConfigDict +] - enable_secure_private_service_connect: Optional[bool] - """Optional. If set to true, enable secure private service connect with IAM authorization. Otherwise, private service connect will be done without authorization. Note latency will be slightly increased if authorization is enabled.""" - project_allowlist: Optional[list[str]] - """A list of Projects from which the forwarding rule will target the service attachment.""" +class _GetRuntimeFeedbackRequestParameters(_common.BaseModel): + """Parameters for getting a Runtime Feedback Entry.""" - psc_automation_configs: Optional[list[PSCAutomationConfigDict]] - """Optional. List of projects and networks where the PSC endpoints will be created. This field is used by Online Inference(Prediction) only.""" + name: Optional[str] = Field( + default=None, description="""The name of the feedback entry to retrieve.""" + ) + config: Optional[GetRuntimeFeedbackConfig] = Field(default=None, description="""""") - service_attachment: Optional[str] - """Output only. The name of the generated service attachment resource. This is only populated if the endpoint is deployed with PrivateServiceConnect.""" +class _GetRuntimeFeedbackRequestParametersDict(TypedDict, total=False): + """Parameters for getting a Runtime Feedback Entry.""" + + name: Optional[str] + """The name of the feedback entry to retrieve.""" -PrivateServiceConnectConfigOrDict = Union[ - PrivateServiceConnectConfig, PrivateServiceConnectConfigDict + config: Optional[GetRuntimeFeedbackConfigDict] + """""" + + +_GetRuntimeFeedbackRequestParametersOrDict = Union[ + _GetRuntimeFeedbackRequestParameters, _GetRuntimeFeedbackRequestParametersDict ] -class DeployRequestEndpointConfig(_common.BaseModel): - """The endpoint config to use for the deployment.""" +class ListRuntimeFeedbackEntriesConfig(_common.BaseModel): + """Config for listing Feedback Entries.""" - dedicated_endpoint_disabled: Optional[bool] = Field( - default=None, - description="""Optional. By default, if dedicated endpoint is enabled and private service connect config is not set, the endpoint will be exposed through a dedicated DNS [Endpoint.dedicated_endpoint_dns]. If private service connect config is set, the endpoint will be exposed through private service connect. Your request to the dedicated DNS will be isolated from other users' traffic and will have better performance and reliability. Note: Once you enabled dedicated endpoint, you won't be able to send request to the shared DNS {region}-aiplatform.googleapis.com. The limitations will be removed soon. If this field is set to true, the dedicated endpoint will be disabled and the deployed model will be exposed through the shared DNS {region}-aiplatform.googleapis.com.""", - ) - dedicated_endpoint_enabled: Optional[bool] = Field( - default=None, - description="""Optional. Deprecated. Use dedicated_endpoint_disabled instead. If true, the endpoint will be exposed through a dedicated DNS [Endpoint.dedicated_endpoint_dns]. Your request to the dedicated DNS will be isolated from other users' traffic and will have better performance and reliability. Note: Once you enabled dedicated endpoint, you won't be able to send request to the shared DNS {region}-aiplatform.googleapis.com. The limitations will be removed soon.""", - ) - endpoint_display_name: Optional[str] = Field( - default=None, - description="""Optional. The user-specified display name of the endpoint. If not set, a default name will be used.""", - ) - endpoint_user_id: Optional[str] = Field( - default=None, - description="""Optional. Immutable. The ID to use for endpoint, which will become the final component of the endpoint resource name. If not provided, Vertex AI will generate a value for this ID. If the first character is a letter, this value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The last character must be a letter or number. If the first character is a number, this value may be up to 9 characters, and valid characters are `[0-9]` with no leading zeros. When using HTTP/JSON, this field is populated based on a query string argument, such as `?endpoint_id=12345`. This is the fallback for fields that are not included in either the URI or the body.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - labels: Optional[dict[str, 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="""Optional. The labels with user-defined metadata to organize your Endpoints. 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="""An expression for filtering the results of the request.""", ) - private_service_connect_config: Optional[PrivateServiceConnectConfig] = Field( - default=None, - description="""Optional. Configuration for private service connect. If set, the endpoint will be exposed through private service connect.""", + order_by: Optional[str] = Field( + default=None, description="""A comma-separated list of fields to order by.""" ) -class DeployRequestEndpointConfigDict(TypedDict, total=False): - """The endpoint config to use for the deployment.""" - - dedicated_endpoint_disabled: Optional[bool] - """Optional. By default, if dedicated endpoint is enabled and private service connect config is not set, the endpoint will be exposed through a dedicated DNS [Endpoint.dedicated_endpoint_dns]. If private service connect config is set, the endpoint will be exposed through private service connect. Your request to the dedicated DNS will be isolated from other users' traffic and will have better performance and reliability. Note: Once you enabled dedicated endpoint, you won't be able to send request to the shared DNS {region}-aiplatform.googleapis.com. The limitations will be removed soon. If this field is set to true, the dedicated endpoint will be disabled and the deployed model will be exposed through the shared DNS {region}-aiplatform.googleapis.com.""" +class ListRuntimeFeedbackEntriesConfigDict(TypedDict, total=False): + """Config for listing Feedback Entries.""" - dedicated_endpoint_enabled: Optional[bool] - """Optional. Deprecated. Use dedicated_endpoint_disabled instead. If true, the endpoint will be exposed through a dedicated DNS [Endpoint.dedicated_endpoint_dns]. Your request to the dedicated DNS will be isolated from other users' traffic and will have better performance and reliability. Note: Once you enabled dedicated endpoint, you won't be able to send request to the shared DNS {region}-aiplatform.googleapis.com. The limitations will be removed soon.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - endpoint_display_name: Optional[str] - """Optional. The user-specified display name of the endpoint. If not set, a default name will be used.""" + page_size: Optional[int] + """""" - endpoint_user_id: Optional[str] - """Optional. Immutable. The ID to use for endpoint, which will become the final component of the endpoint resource name. If not provided, Vertex AI will generate a value for this ID. If the first character is a letter, this value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The last character must be a letter or number. If the first character is a number, this value may be up to 9 characters, and valid characters are `[0-9]` with no leading zeros. When using HTTP/JSON, this field is populated based on a query string argument, such as `?endpoint_id=12345`. This is the fallback for fields that are not included in either the URI or the body.""" + page_token: Optional[str] + """""" - labels: Optional[dict[str, str]] - """Optional. The labels with user-defined metadata to organize your Endpoints. 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.""" + filter: Optional[str] + """An expression for filtering the results of the request.""" - private_service_connect_config: Optional[PrivateServiceConnectConfigDict] - """Optional. Configuration for private service connect. If set, the endpoint will be exposed through private service connect.""" + order_by: Optional[str] + """A comma-separated list of fields to order by.""" -DeployRequestEndpointConfigOrDict = Union[ - DeployRequestEndpointConfig, DeployRequestEndpointConfigDict +ListRuntimeFeedbackEntriesConfigOrDict = Union[ + ListRuntimeFeedbackEntriesConfig, ListRuntimeFeedbackEntriesConfigDict ] -class DeployRequestDeployConfig(_common.BaseModel): - """The deploy config to use for the deployment.""" +class _ListRuntimeFeedbackEntriesRequestParameters(_common.BaseModel): + """Parameters for listing Feedback Entries.""" - dedicated_resources: Optional[DedicatedResources] = Field( - default=None, - description="""Optional. The dedicated resources to use for the endpoint. If not set, the default resources will be used.""", - ) - fast_tryout_enabled: Optional[bool] = Field( + parent: Optional[str] = Field( default=None, - description="""Optional. If true, enable the QMT fast tryout feature for this model if possible.""", + description="""Resource name of the Runtime to list the Feedback Entries from.""", ) - system_labels: Optional[dict[str, str]] = Field( - default=None, - description="""Optional. System labels for Model Garden deployments. These labels are managed by Google and for tracking purposes only.""", + config: Optional[ListRuntimeFeedbackEntriesConfig] = Field( + default=None, description="""""" ) -class DeployRequestDeployConfigDict(TypedDict, total=False): - """The deploy config to use for the deployment.""" - - dedicated_resources: Optional[DedicatedResourcesDict] - """Optional. The dedicated resources to use for the endpoint. If not set, the default resources will be used.""" +class _ListRuntimeFeedbackEntriesRequestParametersDict(TypedDict, total=False): + """Parameters for listing Feedback Entries.""" - fast_tryout_enabled: Optional[bool] - """Optional. If true, enable the QMT fast tryout feature for this model if possible.""" + parent: Optional[str] + """Resource name of the Runtime to list the Feedback Entries from.""" - system_labels: Optional[dict[str, str]] - """Optional. System labels for Model Garden deployments. These labels are managed by Google and for tracking purposes only.""" + config: Optional[ListRuntimeFeedbackEntriesConfigDict] + """""" -DeployRequestDeployConfigOrDict = Union[ - DeployRequestDeployConfig, DeployRequestDeployConfigDict +_ListRuntimeFeedbackEntriesRequestParametersOrDict = Union[ + _ListRuntimeFeedbackEntriesRequestParameters, + _ListRuntimeFeedbackEntriesRequestParametersDict, ] -class _DeployRequestParameters(_common.BaseModel): - """Parameters for deployment.""" +class ListRuntimeFeedbackEntriesResponse(_common.BaseModel): + """Response for listing Feedback Entries.""" - destination: Optional[str] = Field(default=None, description="""""") - publisher_model_name: Optional[str] = Field(default=None, description="""""") - hugging_face_model_id: Optional[str] = Field(default=None, description="""""") - custom_model: Optional[DeployRequestCustomModel] = Field( - default=None, description="""""" - ) - model_config_val: Optional[DeployRequestModelConfig] = Field( - default=None, description="""""" - ) - endpoint_config: Optional[DeployRequestEndpointConfig] = Field( - default=None, description="""""" + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - deploy_config: Optional[DeployRequestDeployConfig] = Field( - default=None, description="""""" + next_page_token: Optional[str] = Field(default=None, description="""""") + feedback_entries: Optional[list[FeedbackEntry]] = Field( + default=None, description="""List of Feedback Entries.""" ) - config: Optional[DeployConfig] = Field(default=None, description="""""") - - -class _DeployRequestParametersDict(TypedDict, total=False): - """Parameters for deployment.""" - - destination: Optional[str] - """""" - - publisher_model_name: Optional[str] - """""" - hugging_face_model_id: Optional[str] - """""" - - custom_model: Optional[DeployRequestCustomModelDict] - """""" - model_config_val: Optional[DeployRequestModelConfigDict] - """""" +class ListRuntimeFeedbackEntriesResponseDict(TypedDict, total=False): + """Response for listing Feedback Entries.""" - endpoint_config: Optional[DeployRequestEndpointConfigDict] - """""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - deploy_config: Optional[DeployRequestDeployConfigDict] + next_page_token: Optional[str] """""" - config: Optional[DeployConfigDict] - """""" + feedback_entries: Optional[list[FeedbackEntryDict]] + """List of Feedback Entries.""" -_DeployRequestParametersOrDict = Union[ - _DeployRequestParameters, _DeployRequestParametersDict +ListRuntimeFeedbackEntriesResponseOrDict = Union[ + ListRuntimeFeedbackEntriesResponse, ListRuntimeFeedbackEntriesResponseDict ] -class DeployResponse(_common.BaseModel): - """Response for deployment.""" +class UpdateRuntimeFeedbackEntryConfig(_common.BaseModel): + """Config for updating a Feedback Entry.""" - endpoint: Optional[str] = Field( - default=None, description="""The resource name of the deployed endpoint.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - model: Optional[str] = Field( - default=None, description="""The resource name of the deployed model.""" + 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.""", ) - - -class DeployResponseDict(TypedDict, total=False): - """Response for deployment.""" - - endpoint: Optional[str] - """The resource name of the deployed endpoint.""" - - model: Optional[str] - """The resource name of the deployed model.""" - - -DeployResponseOrDict = Union[DeployResponse, DeployResponseDict] - - -class DeployModelOperation(_common.BaseModel): - """Operation that has a deploy 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}`.""", - ) - 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.""", + feedback_type: Optional[FeedbackType] = Field( + default=None, description="""The type of feedback provided.""" ) - done: Optional[bool] = Field( + session_id: Optional[str] = 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="""The ID of the session to which the feedback relates to.""", ) - error: Optional[dict[str, Any]] = Field( + event_id: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", - ) - response: Optional[DeployResponse] = Field(default=None, description="""""") - - -class DeployModelOperationDict(TypedDict, total=False): - """Operation that has a deploy 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.""" - - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" - - response: Optional[DeployResponseDict] - """""" - - -DeployModelOperationOrDict = Union[DeployModelOperation, DeployModelOperationDict] - - -class GetDeployOperationConfig(_common.BaseModel): - """Config for ``get_deploy_publisher_model_operation``.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - - -class GetDeployOperationConfigDict(TypedDict, total=False): - """Config for ``get_deploy_publisher_model_operation``.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -GetDeployOperationConfigOrDict = Union[ - GetDeployOperationConfig, GetDeployOperationConfigDict -] - - -class _GetDeployOperationParameters(_common.BaseModel): - """Parameters for polling a ``deploy`` operation.""" - - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" - ) - config: Optional[GetDeployOperationConfig] = Field(default=None, description="""""") - - -class _GetDeployOperationParametersDict(TypedDict, total=False): - """Parameters for polling a ``deploy`` operation.""" - - operation_name: Optional[str] - """The server-assigned name for the operation.""" - - config: Optional[GetDeployOperationConfigDict] - """""" - - -_GetDeployOperationParametersOrDict = Union[ - _GetDeployOperationParameters, _GetDeployOperationParametersDict -] - - -class CreateRuntimeFeedbackEntryConfig(_common.BaseModel): - """Config for creating a Feedback Entry.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + description="""The ID of the event to which the feedback relates to.""", ) feedback_labels: Optional[list[str]] = Field( default=None, @@ -26550,12 +25867,25 @@ class CreateRuntimeFeedbackEntryConfig(_common.BaseModel): ) -class CreateRuntimeFeedbackEntryConfigDict(TypedDict, total=False): - """Config for creating a Feedback Entry.""" +class UpdateRuntimeFeedbackEntryConfigDict(TypedDict, total=False): + """Config for updating a Feedback Entry.""" 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.""" + + feedback_type: Optional[FeedbackType] + """The type of feedback provided.""" + + session_id: Optional[str] + """The ID of the session to which the feedback relates to.""" + + event_id: Optional[str] + """The ID of the event to which the feedback relates to.""" + feedback_labels: Optional[list[str]] """Specific labels for feedback (non-factual, offensive, etc.).""" @@ -26575,251 +25905,207 @@ class CreateRuntimeFeedbackEntryConfigDict(TypedDict, total=False): """Waits for the operation to complete before returning.""" -CreateRuntimeFeedbackEntryConfigOrDict = Union[ - CreateRuntimeFeedbackEntryConfig, CreateRuntimeFeedbackEntryConfigDict +UpdateRuntimeFeedbackEntryConfigOrDict = Union[ + UpdateRuntimeFeedbackEntryConfig, UpdateRuntimeFeedbackEntryConfigDict ] -class _CreateRuntimeFeedbackEntryRequestParameters(_common.BaseModel): - """Parameters for creating a Feedback Entry.""" +class _UpdateRuntimeFeedbackEntryRequestParameters(_common.BaseModel): + """Parameters for updating a Feedback Entry.""" name: Optional[str] = Field( - default=None, - description="""Resource name of the Runtime to create the Feedback Entry in.""", - ) - feedback_type: Optional[FeedbackType] = Field( - default=None, description="""The type of feedback provided.""" - ) - session_id: Optional[str] = Field( - default=None, - description="""The ID of the session to which the feedback relates to.""", - ) - event_id: Optional[str] = Field( - default=None, - description="""The ID of the event to which the feedback relates to.""", + default=None, description="""Name of the Feedback Entry.""" ) - config: Optional[CreateRuntimeFeedbackEntryConfig] = Field( - default=None, description="""""" + config: Optional[UpdateRuntimeFeedbackEntryConfig] = Field( + default=None, description="""Config for updating a Feedback Entry.""" ) -class _CreateRuntimeFeedbackEntryRequestParametersDict(TypedDict, total=False): - """Parameters for creating a Feedback Entry.""" +class _UpdateRuntimeFeedbackEntryRequestParametersDict(TypedDict, total=False): + """Parameters for updating a Feedback Entry.""" name: Optional[str] - """Resource name of the Runtime to create the Feedback Entry in.""" - - feedback_type: Optional[FeedbackType] - """The type of feedback provided.""" - - session_id: Optional[str] - """The ID of the session to which the feedback relates to.""" - - event_id: Optional[str] - """The ID of the event to which the feedback relates to.""" + """Name of the Feedback Entry.""" - config: Optional[CreateRuntimeFeedbackEntryConfigDict] - """""" + config: Optional[UpdateRuntimeFeedbackEntryConfigDict] + """Config for updating a Feedback Entry.""" -_CreateRuntimeFeedbackEntryRequestParametersOrDict = Union[ - _CreateRuntimeFeedbackEntryRequestParameters, - _CreateRuntimeFeedbackEntryRequestParametersDict, +_UpdateRuntimeFeedbackEntryRequestParametersOrDict = Union[ + _UpdateRuntimeFeedbackEntryRequestParameters, + _UpdateRuntimeFeedbackEntryRequestParametersDict, ] -class FeedbackEntry(_common.BaseModel): - """A Feedback Entry.""" +class GetRuntimeFeedbackEntryConfig(_common.BaseModel): + """Config for getting a Feedback Entry.""" - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. The time at which the entry was created.""", - ) - custom_metadata: Optional[dict[str, str]] = Field( - default=None, - description="""Optional. Additional key-value metadata associated with the feedback.""", - ) - event_id: Optional[str] = Field( - default=None, - description="""Required. The ID of the event within the session that the feedback relates to.""", - ) - feedback_labels: Optional[list[str]] = Field( - default=None, description="""feedbackLabels""" - ) - feedback_text: Optional[str] = Field( - default=None, - description="""Optional. Qualitative free-form comments provided by the user.""", - ) - feedback_type: Optional[FeedbackType] = Field( - default=None, - description="""Required. The coarse-grained type of feedback provided by the user. Must be set to a value other than `FEEDBACK_TYPE_UNSPECIFIED`.""", - ) - name: Optional[str] = Field( - default=None, - description="""Identifier. The resource name. Assigned by the server on create. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/feedbackEntries/{feedback_entry}`""", - ) - session_id: Optional[str] = Field( - default=None, - description="""Required. The ID of the session that the feedback relates to.""", - ) - source: Optional[str] = Field( - default=None, - description="""Optional. The surface that the feedback originated from.""", - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. The time at which the entry was most recently updated.""", - ) - user_id: Optional[str] = Field( - default=None, - description="""Optional. A caller-supplied identifier for the user who provided the feedback. The semantics of this field (for example whether it is an opaque token, a hashed value, or a user-visible identifier) are determined by the calling application.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class FeedbackEntryDict(TypedDict, total=False): - """A Feedback Entry.""" +class GetRuntimeFeedbackEntryConfigDict(TypedDict, total=False): + """Config for getting a Feedback Entry.""" - create_time: Optional[datetime.datetime] - """Output only. The time at which the entry was created.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - custom_metadata: Optional[dict[str, str]] - """Optional. Additional key-value metadata associated with the feedback.""" - event_id: Optional[str] - """Required. The ID of the event within the session that the feedback relates to.""" +GetRuntimeFeedbackEntryConfigOrDict = Union[ + GetRuntimeFeedbackEntryConfig, GetRuntimeFeedbackEntryConfigDict +] - feedback_labels: Optional[list[str]] - """feedbackLabels""" - feedback_text: Optional[str] - """Optional. Qualitative free-form comments provided by the user.""" +class _GetRuntimeFeedbackOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a Feedback Entry as a response.""" - feedback_type: Optional[FeedbackType] - """Required. The coarse-grained type of feedback provided by the user. Must be set to a value other than `FEEDBACK_TYPE_UNSPECIFIED`.""" + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetRuntimeFeedbackEntryConfig] = Field( + default=None, description="""Used to override the default configuration.""" + ) - name: Optional[str] - """Identifier. The resource name. Assigned by the server on create. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/feedbackEntries/{feedback_entry}`""" - session_id: Optional[str] - """Required. The ID of the session that the feedback relates to.""" +class _GetRuntimeFeedbackOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a Feedback Entry as a response.""" - source: Optional[str] - """Optional. The surface that the feedback originated from.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - update_time: Optional[datetime.datetime] - """Output only. The time at which the entry was most recently updated.""" + config: Optional[GetRuntimeFeedbackEntryConfigDict] + """Used to override the default configuration.""" - user_id: Optional[str] - """Optional. A caller-supplied identifier for the user who provided the feedback. The semantics of this field (for example whether it is an opaque token, a hashed value, or a user-visible identifier) are determined by the calling application.""" +_GetRuntimeFeedbackOperationParametersOrDict = Union[ + _GetRuntimeFeedbackOperationParameters, _GetRuntimeFeedbackOperationParametersDict +] -FeedbackEntryOrDict = Union[FeedbackEntry, FeedbackEntryDict] +class GetRuntimeFeedbackContextConfig(_common.BaseModel): + """Config for getting a Feedback Context.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetRuntimeFeedbackContextConfigDict(TypedDict, total=False): + """Config for getting a Feedback Context.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetRuntimeFeedbackContextConfigOrDict = Union[ + GetRuntimeFeedbackContextConfig, GetRuntimeFeedbackContextConfigDict +] -class RuntimeFeedbackEntryOperation(_common.BaseModel): - """Operation that has a Runtime Feedback Entry as a response.""" + +class _GetRuntimeFeedbackContextRequestParameters(_common.BaseModel): + """Parameters for getting a Runtime Feedback Context.""" 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.""", + default=None, description="""The name of the Feedback Context to retrieve.""" ) - response: Optional[FeedbackEntry] = Field( - default=None, description="""The Runtime Feedback Entry.""" + config: Optional[GetRuntimeFeedbackContextConfig] = Field( + default=None, description="""""" ) -class RuntimeFeedbackEntryOperationDict(TypedDict, total=False): - """Operation that has a Runtime Feedback Entry as a response.""" +class _GetRuntimeFeedbackContextRequestParametersDict(TypedDict, total=False): + """Parameters for getting a Runtime Feedback Context.""" 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}`.""" + """The name of the Feedback Context to retrieve.""" - 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.""" + config: Optional[GetRuntimeFeedbackContextConfigDict] + """""" - 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.""" +_GetRuntimeFeedbackContextRequestParametersOrDict = Union[ + _GetRuntimeFeedbackContextRequestParameters, + _GetRuntimeFeedbackContextRequestParametersDict, +] - response: Optional[FeedbackEntryDict] - """The Runtime Feedback Entry.""" +class FeedbackContext(_common.BaseModel): + """A Feedback Context.""" + + context_events: Optional[list[SessionEvent]] = Field( + default=None, + description="""Optional. The session events from the originating session.""", + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. The resource name. Assigned by the server on create. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/feedbackEntries/{feedback_entry}/feedbackContext`""", + ) -RuntimeFeedbackEntryOperationOrDict = Union[ - RuntimeFeedbackEntryOperation, RuntimeFeedbackEntryOperationDict -] +class FeedbackContextDict(TypedDict, total=False): + """A Feedback Context.""" -class DeleteRuntimeFeedbackEntryConfig(_common.BaseModel): - """Config for deleting a Feedback Entry.""" + context_events: Optional[list[SessionEventDict]] + """Optional. The session events from the originating session.""" + + name: Optional[str] + """Identifier. The resource name. Assigned by the server on create. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/feedbackEntries/{feedback_entry}/feedbackContext`""" + + +FeedbackContextOrDict = Union[FeedbackContext, FeedbackContextDict] + + +class GetRuntimeFeedbackContextOperationConfig(_common.BaseModel): + """Config for getting a Feedback Context.""" 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.""", - ) -class DeleteRuntimeFeedbackEntryConfigDict(TypedDict, total=False): - """Config for deleting a Feedback Entry.""" +class GetRuntimeFeedbackContextOperationConfigDict(TypedDict, total=False): + """Config for getting a Feedback Context.""" 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.""" - -DeleteRuntimeFeedbackEntryConfigOrDict = Union[ - DeleteRuntimeFeedbackEntryConfig, DeleteRuntimeFeedbackEntryConfigDict +GetRuntimeFeedbackContextOperationConfigOrDict = Union[ + GetRuntimeFeedbackContextOperationConfig, + GetRuntimeFeedbackContextOperationConfigDict, ] -class _DeleteRuntimeFeedbackEntryRequestParameters(_common.BaseModel): - """Parameters for deleting a Feedback Entry.""" +class _GetRuntimeFeedbackContextOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a Feedback Context as a response.""" - name: Optional[str] = Field( - default=None, description="""Name of the Feedback Entry to delete.""" + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - config: Optional[DeleteRuntimeFeedbackEntryConfig] = Field( - default=None, description="""""" + config: Optional[GetRuntimeFeedbackContextOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -class _DeleteRuntimeFeedbackEntryRequestParametersDict(TypedDict, total=False): - """Parameters for deleting a Feedback Entry.""" +class _GetRuntimeFeedbackContextOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a Feedback Context as a response.""" - name: Optional[str] - """Name of the Feedback Entry to delete.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - config: Optional[DeleteRuntimeFeedbackEntryConfigDict] - """""" + config: Optional[GetRuntimeFeedbackContextOperationConfigDict] + """Used to override the default configuration.""" -_DeleteRuntimeFeedbackEntryRequestParametersOrDict = Union[ - _DeleteRuntimeFeedbackEntryRequestParameters, - _DeleteRuntimeFeedbackEntryRequestParametersDict, +_GetRuntimeFeedbackContextOperationParametersOrDict = Union[ + _GetRuntimeFeedbackContextOperationParameters, + _GetRuntimeFeedbackContextOperationParametersDict, ] -class DeleteRuntimeFeedbackEntryOperation(_common.BaseModel): - """Operation for deleting a Feedback Entry.""" +class RuntimeFeedbackContextOperation(_common.BaseModel): + """Operation that has a Runtime Feedback Context as a response.""" name: Optional[str] = Field( default=None, @@ -26837,10 +26123,13 @@ class DeleteRuntimeFeedbackEntryOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) + response: Optional[FeedbackContext] = Field( + default=None, description="""The Runtime Feedback Context.""" + ) -class DeleteRuntimeFeedbackEntryOperationDict(TypedDict, total=False): - """Operation for deleting a Feedback Entry.""" +class RuntimeFeedbackContextOperationDict(TypedDict, total=False): + """Operation that has a Runtime Feedback Context 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}`.""" @@ -26854,440 +26143,381 @@ class DeleteRuntimeFeedbackEntryOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" + response: Optional[FeedbackContextDict] + """The Runtime Feedback Context.""" -DeleteRuntimeFeedbackEntryOperationOrDict = Union[ - DeleteRuntimeFeedbackEntryOperation, DeleteRuntimeFeedbackEntryOperationDict + +RuntimeFeedbackContextOperationOrDict = Union[ + RuntimeFeedbackContextOperation, RuntimeFeedbackContextOperationDict ] -class GetRuntimeFeedbackConfig(_common.BaseModel): - """Config for getting a Feedback Entry.""" +class UpdateRuntimeFeedbackContextConfig(_common.BaseModel): + """Config for updating a Feedback Context.""" 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.""", + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", + ) -class GetRuntimeFeedbackConfigDict(TypedDict, total=False): - """Config for getting a Feedback Entry.""" +class UpdateRuntimeFeedbackContextConfigDict(TypedDict, total=False): + """Config for updating a Feedback Context.""" 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.""" + + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" -GetRuntimeFeedbackConfigOrDict = Union[ - GetRuntimeFeedbackConfig, GetRuntimeFeedbackConfigDict + +UpdateRuntimeFeedbackContextConfigOrDict = Union[ + UpdateRuntimeFeedbackContextConfig, UpdateRuntimeFeedbackContextConfigDict ] -class _GetRuntimeFeedbackRequestParameters(_common.BaseModel): - """Parameters for getting a Runtime Feedback Entry.""" +class _UpdateRuntimeFeedbackContextRequestParameters(_common.BaseModel): + """Parameters for updating a Feedback Context.""" name: Optional[str] = Field( - default=None, description="""The name of the feedback entry to retrieve.""" + default=None, + description="""Name of the Feedback Context. Format: projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/feedbackEntries/{feedback_entry}/feedbackContext""", + ) + context_events: Optional[list[SessionEvent]] = Field( + default=None, + description="""Events from the conversation relevant to the parent Feedback Entry.""", + ) + config: Optional[UpdateRuntimeFeedbackContextConfig] = Field( + default=None, description="""""" ) - config: Optional[GetRuntimeFeedbackConfig] = Field(default=None, description="""""") -class _GetRuntimeFeedbackRequestParametersDict(TypedDict, total=False): - """Parameters for getting a Runtime Feedback Entry.""" +class _UpdateRuntimeFeedbackContextRequestParametersDict(TypedDict, total=False): + """Parameters for updating a Feedback Context.""" name: Optional[str] - """The name of the feedback entry to retrieve.""" + """Name of the Feedback Context. Format: projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/feedbackEntries/{feedback_entry}/feedbackContext""" - config: Optional[GetRuntimeFeedbackConfigDict] + context_events: Optional[list[SessionEventDict]] + """Events from the conversation relevant to the parent Feedback Entry.""" + + config: Optional[UpdateRuntimeFeedbackContextConfigDict] """""" -_GetRuntimeFeedbackRequestParametersOrDict = Union[ - _GetRuntimeFeedbackRequestParameters, _GetRuntimeFeedbackRequestParametersDict +_UpdateRuntimeFeedbackContextRequestParametersOrDict = Union[ + _UpdateRuntimeFeedbackContextRequestParameters, + _UpdateRuntimeFeedbackContextRequestParametersDict, ] -class ListRuntimeFeedbackEntriesConfig(_common.BaseModel): - """Config for listing Feedback Entries.""" +class UndeployModelConfig(_common.BaseModel): + """Config for a Vertex SDK undeploy model from endpoint.""" 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( + traffic_split: Optional[dict[str, int]] = Field( default=None, - description="""An expression for filtering the results of the request.""", + description="""If this field is provided, then the Endpoint's + [traffic_split][Endpoint.traffic_split] will be overwritten with it. If + last DeployedModel is being undeployed from the Endpoint, the + [Endpoint.traffic_split] will always end up empty when this call returns. + A DeployedModel will be successfully undeployed only if it doesn't have + any traffic assigned to it when this method executes, or if this field + unassigns any traffic to it.""", ) - order_by: Optional[str] = Field( - default=None, description="""A comma-separated list of fields to order by.""" + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to wait for the long running operation to complete.""", ) -class ListRuntimeFeedbackEntriesConfigDict(TypedDict, total=False): - """Config for listing Feedback Entries.""" +class UndeployModelConfigDict(TypedDict, total=False): + """Config for a Vertex SDK undeploy model from endpoint.""" 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.""" + traffic_split: Optional[dict[str, int]] + """If this field is provided, then the Endpoint's + [traffic_split][Endpoint.traffic_split] will be overwritten with it. If + last DeployedModel is being undeployed from the Endpoint, the + [Endpoint.traffic_split] will always end up empty when this call returns. + A DeployedModel will be successfully undeployed only if it doesn't have + any traffic assigned to it when this method executes, or if this field + unassigns any traffic to it.""" - order_by: Optional[str] - """A comma-separated list of fields to order by.""" + wait_for_completion: Optional[bool] + """Whether to wait for the long running operation to complete.""" -ListRuntimeFeedbackEntriesConfigOrDict = Union[ - ListRuntimeFeedbackEntriesConfig, ListRuntimeFeedbackEntriesConfigDict -] +UndeployModelConfigOrDict = Union[UndeployModelConfig, UndeployModelConfigDict] -class _ListRuntimeFeedbackEntriesRequestParameters(_common.BaseModel): - """Parameters for listing Feedback Entries.""" +class _UndeployModelRequestParameters(_common.BaseModel): + """Parameters for undeploying a model from an endpoint.""" - parent: Optional[str] = Field( - default=None, - description="""Resource name of the Runtime to list the Feedback Entries from.""", + name: Optional[str] = Field( + default=None, description="""ID of the endpoint to undeploy the model from.""" ) - config: Optional[ListRuntimeFeedbackEntriesConfig] = Field( - default=None, description="""""" + deployed_model_id: Optional[str] = Field( + default=None, description="""ID of the deployed model to be undeployed.""" ) + config: Optional[UndeployModelConfig] = Field(default=None, description="""""") -class _ListRuntimeFeedbackEntriesRequestParametersDict(TypedDict, total=False): - """Parameters for listing Feedback Entries.""" +class _UndeployModelRequestParametersDict(TypedDict, total=False): + """Parameters for undeploying a model from an endpoint.""" - parent: Optional[str] - """Resource name of the Runtime to list the Feedback Entries from.""" + name: Optional[str] + """ID of the endpoint to undeploy the model from.""" - config: Optional[ListRuntimeFeedbackEntriesConfigDict] + deployed_model_id: Optional[str] + """ID of the deployed model to be undeployed.""" + + config: Optional[UndeployModelConfigDict] """""" -_ListRuntimeFeedbackEntriesRequestParametersOrDict = Union[ - _ListRuntimeFeedbackEntriesRequestParameters, - _ListRuntimeFeedbackEntriesRequestParametersDict, +_UndeployModelRequestParametersOrDict = Union[ + _UndeployModelRequestParameters, _UndeployModelRequestParametersDict ] -class ListRuntimeFeedbackEntriesResponse(_common.BaseModel): - """Response for listing Feedback Entries.""" +class UndeployModelOperation(_common.BaseModel): + """Operation for undeploying a model.""" - 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="""""") - feedback_entries: Optional[list[FeedbackEntry]] = Field( - default=None, description="""List of Feedback Entries.""" + 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 ListRuntimeFeedbackEntriesResponseDict(TypedDict, total=False): - """Response for listing Feedback Entries.""" +class UndeployModelOperationDict(TypedDict, total=False): + """Operation for undeploying a model.""" - 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.""" - feedback_entries: Optional[list[FeedbackEntryDict]] - """List of Feedback Entries.""" + 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.""" -ListRuntimeFeedbackEntriesResponseOrDict = Union[ - ListRuntimeFeedbackEntriesResponse, ListRuntimeFeedbackEntriesResponseDict -] +UndeployModelOperationOrDict = Union[UndeployModelOperation, UndeployModelOperationDict] -class UpdateRuntimeFeedbackEntryConfig(_common.BaseModel): - """Config for updating a Feedback Entry.""" +class PredictConfig(_common.BaseModel): + """Config for a Vertex SDK predict endpoint.""" 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.""", - ) - feedback_type: Optional[FeedbackType] = Field( - default=None, description="""The type of feedback provided.""" - ) - session_id: Optional[str] = Field( - default=None, - description="""The ID of the session to which the feedback relates to.""", - ) - event_id: Optional[str] = Field( - default=None, - description="""The ID of the event to which the feedback relates to.""", - ) - feedback_labels: Optional[list[str]] = Field( - default=None, - description="""Specific labels for feedback (non-factual, offensive, etc.).""", - ) - feedback_text: Optional[str] = Field( - default=None, - description="""Qualitative free-form comments provided by the user.""", - ) - user_id: Optional[str] = Field( - default=None, description="""User provided identifier.""" - ) - source: Optional[str] = Field( - default=None, description="""Originating UI surface (e.g. 'ADK Web UI').""" - ) - custom_metadata: Optional[dict[str, str]] = Field( + parameters: Optional[str] = Field( default=None, - description=""" Additional key-value metadata associated with the feedback. Allows the collect data for which there is no dedicated field in the resource, ex. version, LLM temperature etc.""", - ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", + description="""The parameters that govern the prediction. The schema of the + parameters may be specified via Endpoint's DeployedModels' [Model's ][ + DeployedModel.model] [PredictSchemata's][Model.predict_schemata] + [parameters_schema_uri][PredictSchemata.parameters_schema_uri]. + """, ) -class UpdateRuntimeFeedbackEntryConfigDict(TypedDict, total=False): - """Config for updating a Feedback Entry.""" +class PredictConfigDict(TypedDict, total=False): + """Config for a Vertex SDK predict endpoint.""" 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.""" + parameters: Optional[str] + """The parameters that govern the prediction. The schema of the + parameters may be specified via Endpoint's DeployedModels' [Model's ][ + DeployedModel.model] [PredictSchemata's][Model.predict_schemata] + [parameters_schema_uri][PredictSchemata.parameters_schema_uri]. + """ - feedback_type: Optional[FeedbackType] - """The type of feedback provided.""" - session_id: Optional[str] - """The ID of the session to which the feedback relates to.""" +PredictConfigOrDict = Union[PredictConfig, PredictConfigDict] - event_id: Optional[str] - """The ID of the event to which the feedback relates to.""" - feedback_labels: Optional[list[str]] - """Specific labels for feedback (non-factual, offensive, etc.).""" +class _PredictParameters(_common.BaseModel): + """Parameters for performing an online prediction.""" - feedback_text: Optional[str] - """Qualitative free-form comments provided by the user.""" + name: Optional[str] = Field( + default=None, + description="""The endpoint that serves the prediction. It could be endpoints/... + or publishers/.../models/... + """, + ) + instances: Optional[list[dict[str, Any]]] = Field( + default=None, + description="""The instances that are the input to the prediction call. The + schema of any single instance may be specified via Endpoint's DeployedModels' + [Model's ][DeployedModel.model] [PredictSchemata's][Model.predict_schemata] + [instance_schema_uri][PredictSchemata.instance_schema_uri]. + """, + ) + config: Optional[PredictConfig] = Field(default=None, description="""""") - user_id: Optional[str] - """User provided identifier.""" - source: Optional[str] - """Originating UI surface (e.g. 'ADK Web UI').""" +class _PredictParametersDict(TypedDict, total=False): + """Parameters for performing an online prediction.""" - custom_metadata: Optional[dict[str, str]] - """ Additional key-value metadata associated with the feedback. Allows the collect data for which there is no dedicated field in the resource, ex. version, LLM temperature etc.""" + name: Optional[str] + """The endpoint that serves the prediction. It could be endpoints/... + or publishers/.../models/... + """ - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + instances: Optional[list[dict[str, Any]]] + """The instances that are the input to the prediction call. The + schema of any single instance may be specified via Endpoint's DeployedModels' + [Model's ][DeployedModel.model] [PredictSchemata's][Model.predict_schemata] + [instance_schema_uri][PredictSchemata.instance_schema_uri]. + """ + config: Optional[PredictConfigDict] + """""" -UpdateRuntimeFeedbackEntryConfigOrDict = Union[ - UpdateRuntimeFeedbackEntryConfig, UpdateRuntimeFeedbackEntryConfigDict -] +_PredictParametersOrDict = Union[_PredictParameters, _PredictParametersDict] -class _UpdateRuntimeFeedbackEntryRequestParameters(_common.BaseModel): - """Parameters for updating a Feedback Entry.""" - name: Optional[str] = Field( - default=None, description="""Name of the Feedback Entry.""" +class PredictResponse(_common.BaseModel): + """Response message for PredictionService.Predict API.""" + + deployed_model_id: Optional[str] = Field( + default=None, + description="""ID of the Endpoint's DeployedModel that served this prediction.""", ) - config: Optional[UpdateRuntimeFeedbackEntryConfig] = Field( - default=None, description="""Config for updating a Feedback Entry.""" + metadata: Optional[Any] = Field( + default=None, + description="""Output only. Request-level metadata returned by the model. The metadata type will be dependent upon the model implementation.""", + ) + model: Optional[str] = Field( + default=None, + description="""Output only. The resource name of the Model which is deployed as the DeployedModel that this prediction hits.""", + ) + model_display_name: Optional[str] = Field( + default=None, + description="""Output only. The display name of the Model which is deployed as the DeployedModel that this prediction hits.""", + ) + model_version_id: Optional[str] = Field( + default=None, + description="""Output only. The version ID of the Model which is deployed as the DeployedModel that this prediction hits.""", + ) + predictions: Optional[list[Any]] = Field( + default=None, + description="""The predictions that are the output of the predictions call. The schema of any single prediction may be specified via Endpoint's DeployedModels' Model's PredictSchemata's prediction_schema_uri.""", ) -class _UpdateRuntimeFeedbackEntryRequestParametersDict(TypedDict, total=False): - """Parameters for updating a Feedback Entry.""" +class PredictResponseDict(TypedDict, total=False): + """Response message for PredictionService.Predict API.""" - name: Optional[str] - """Name of the Feedback Entry.""" + deployed_model_id: Optional[str] + """ID of the Endpoint's DeployedModel that served this prediction.""" - config: Optional[UpdateRuntimeFeedbackEntryConfigDict] - """Config for updating a Feedback Entry.""" - - -_UpdateRuntimeFeedbackEntryRequestParametersOrDict = Union[ - _UpdateRuntimeFeedbackEntryRequestParameters, - _UpdateRuntimeFeedbackEntryRequestParametersDict, -] - - -class GetRuntimeFeedbackEntryConfig(_common.BaseModel): - """Config for getting a Feedback Entry.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - - -class GetRuntimeFeedbackEntryConfigDict(TypedDict, total=False): - """Config for getting a Feedback Entry.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -GetRuntimeFeedbackEntryConfigOrDict = Union[ - GetRuntimeFeedbackEntryConfig, GetRuntimeFeedbackEntryConfigDict -] - - -class _GetRuntimeFeedbackOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a Feedback Entry as a response.""" - - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" - ) - config: Optional[GetRuntimeFeedbackEntryConfig] = Field( - default=None, description="""Used to override the default configuration.""" - ) + metadata: Optional[Any] + """Output only. Request-level metadata returned by the model. The metadata type will be dependent upon the model implementation.""" + model: Optional[str] + """Output only. The resource name of the Model which is deployed as the DeployedModel that this prediction hits.""" -class _GetRuntimeFeedbackOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a Feedback Entry as a response.""" + model_display_name: Optional[str] + """Output only. The display name of the Model which is deployed as the DeployedModel that this prediction hits.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + model_version_id: Optional[str] + """Output only. The version ID of the Model which is deployed as the DeployedModel that this prediction hits.""" - config: Optional[GetRuntimeFeedbackEntryConfigDict] - """Used to override the default configuration.""" + predictions: Optional[list[Any]] + """The predictions that are the output of the predictions call. The schema of any single prediction may be specified via Endpoint's DeployedModels' Model's PredictSchemata's prediction_schema_uri.""" -_GetRuntimeFeedbackOperationParametersOrDict = Union[ - _GetRuntimeFeedbackOperationParameters, _GetRuntimeFeedbackOperationParametersDict -] +PredictResponseOrDict = Union[PredictResponse, PredictResponseDict] -class GetRuntimeFeedbackContextConfig(_common.BaseModel): - """Config for getting a Feedback Context.""" +class DeleteEndpointConfig(_common.BaseModel): + """Config for a Vertex SDK delete endpoint.""" 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="""Whether to wait for the long running operation to complete.""", + ) -class GetRuntimeFeedbackContextConfigDict(TypedDict, total=False): - """Config for getting a Feedback Context.""" +class DeleteEndpointConfigDict(TypedDict, total=False): + """Config for a Vertex SDK delete endpoint.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - -GetRuntimeFeedbackContextConfigOrDict = Union[ - GetRuntimeFeedbackContextConfig, GetRuntimeFeedbackContextConfigDict -] - - -class _GetRuntimeFeedbackContextRequestParameters(_common.BaseModel): - """Parameters for getting a Runtime Feedback Context.""" - - name: Optional[str] = Field( - default=None, description="""The name of the Feedback Context to retrieve.""" - ) - config: Optional[GetRuntimeFeedbackContextConfig] = Field( - default=None, description="""""" - ) - - -class _GetRuntimeFeedbackContextRequestParametersDict(TypedDict, total=False): - """Parameters for getting a Runtime Feedback Context.""" - - name: Optional[str] - """The name of the Feedback Context to retrieve.""" - - config: Optional[GetRuntimeFeedbackContextConfigDict] - """""" + wait_for_completion: Optional[bool] + """Whether to wait for the long running operation to complete.""" -_GetRuntimeFeedbackContextRequestParametersOrDict = Union[ - _GetRuntimeFeedbackContextRequestParameters, - _GetRuntimeFeedbackContextRequestParametersDict, -] +DeleteEndpointConfigOrDict = Union[DeleteEndpointConfig, DeleteEndpointConfigDict] -class FeedbackContext(_common.BaseModel): - """A Feedback Context.""" +class _DeleteEndpointRequestParameters(_common.BaseModel): + """Parameters for deleting an endpoint.""" - context_events: Optional[list[SessionEvent]] = Field( - default=None, - description="""Optional. The session events from the originating session.""", - ) name: Optional[str] = Field( default=None, - description="""Identifier. The resource name. Assigned by the server on create. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/feedbackEntries/{feedback_entry}/feedbackContext`""", + description="""Required. The resource name of the endpoint to be deleted.""", ) + config: Optional[DeleteEndpointConfig] = Field(default=None, description="""""") -class FeedbackContextDict(TypedDict, total=False): - """A Feedback Context.""" - - context_events: Optional[list[SessionEventDict]] - """Optional. The session events from the originating session.""" +class _DeleteEndpointRequestParametersDict(TypedDict, total=False): + """Parameters for deleting an endpoint.""" name: Optional[str] - """Identifier. The resource name. Assigned by the server on create. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/feedbackEntries/{feedback_entry}/feedbackContext`""" - - -FeedbackContextOrDict = Union[FeedbackContext, FeedbackContextDict] - - -class GetRuntimeFeedbackContextOperationConfig(_common.BaseModel): - """Config for getting a Feedback Context.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - - -class GetRuntimeFeedbackContextOperationConfigDict(TypedDict, total=False): - """Config for getting a Feedback Context.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -GetRuntimeFeedbackContextOperationConfigOrDict = Union[ - GetRuntimeFeedbackContextOperationConfig, - GetRuntimeFeedbackContextOperationConfigDict, -] - - -class _GetRuntimeFeedbackContextOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a Feedback Context as a response.""" - - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" - ) - config: Optional[GetRuntimeFeedbackContextOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" - ) - - -class _GetRuntimeFeedbackContextOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a Feedback Context as a response.""" - - operation_name: Optional[str] - """The server-assigned name for the operation.""" + """Required. The resource name of the endpoint to be deleted.""" - config: Optional[GetRuntimeFeedbackContextOperationConfigDict] - """Used to override the default configuration.""" + config: Optional[DeleteEndpointConfigDict] + """""" -_GetRuntimeFeedbackContextOperationParametersOrDict = Union[ - _GetRuntimeFeedbackContextOperationParameters, - _GetRuntimeFeedbackContextOperationParametersDict, +_DeleteEndpointRequestParametersOrDict = Union[ + _DeleteEndpointRequestParameters, _DeleteEndpointRequestParametersDict ] -class RuntimeFeedbackContextOperation(_common.BaseModel): - """Operation that has a Runtime Feedback Context as a response.""" +class DeleteEndpointOperation(_common.BaseModel): + """Operation for deleting a endpoint.""" name: Optional[str] = Field( default=None, @@ -27305,13 +26535,10 @@ class RuntimeFeedbackContextOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) - response: Optional[FeedbackContext] = Field( - default=None, description="""The Runtime Feedback Context.""" - ) -class RuntimeFeedbackContextOperationDict(TypedDict, total=False): - """Operation that has a Runtime Feedback Context as a response.""" +class DeleteEndpointOperationDict(TypedDict, total=False): + """Operation for deleting a endpoint.""" 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}`.""" @@ -27325,991 +26552,582 @@ class RuntimeFeedbackContextOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" - response: Optional[FeedbackContextDict] - """The Runtime Feedback Context.""" - -RuntimeFeedbackContextOperationOrDict = Union[ - RuntimeFeedbackContextOperation, RuntimeFeedbackContextOperationDict +DeleteEndpointOperationOrDict = Union[ + DeleteEndpointOperation, DeleteEndpointOperationDict ] -class UpdateRuntimeFeedbackContextConfig(_common.BaseModel): - """Config for updating a Feedback Context.""" +class GetEndpointConfig(_common.BaseModel): + """Optional parameters for endpoints.get method.""" 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.""", - ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", - ) -class UpdateRuntimeFeedbackContextConfigDict(TypedDict, total=False): - """Config for updating a Feedback Context.""" +class GetEndpointConfigDict(TypedDict, total=False): + """Optional parameters for endpoints.get method.""" 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.""" - - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" - -UpdateRuntimeFeedbackContextConfigOrDict = Union[ - UpdateRuntimeFeedbackContextConfig, UpdateRuntimeFeedbackContextConfigDict -] +GetEndpointConfigOrDict = Union[GetEndpointConfig, GetEndpointConfigDict] -class _UpdateRuntimeFeedbackContextRequestParameters(_common.BaseModel): - """Parameters for updating a Feedback Context.""" +class _GetEndpointParameters(_common.BaseModel): name: Optional[str] = Field( default=None, - description="""Name of the Feedback Context. Format: projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/feedbackEntries/{feedback_entry}/feedbackContext""", - ) - context_events: Optional[list[SessionEvent]] = Field( - default=None, - description="""Events from the conversation relevant to the parent Feedback Entry.""", + description="""Required. The resource name of the Endpoint to get.""", ) - config: Optional[UpdateRuntimeFeedbackContextConfig] = Field( - default=None, description="""""" + config: Optional[GetEndpointConfig] = Field( + default=None, description="""Optional parameters for the request.""" ) -class _UpdateRuntimeFeedbackContextRequestParametersDict(TypedDict, total=False): - """Parameters for updating a Feedback Context.""" +class _GetEndpointParametersDict(TypedDict, total=False): name: Optional[str] - """Name of the Feedback Context. Format: projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/feedbackEntries/{feedback_entry}/feedbackContext""" - - context_events: Optional[list[SessionEventDict]] - """Events from the conversation relevant to the parent Feedback Entry.""" + """Required. The resource name of the Endpoint to get.""" - config: Optional[UpdateRuntimeFeedbackContextConfigDict] - """""" + config: Optional[GetEndpointConfigDict] + """Optional parameters for the request.""" -_UpdateRuntimeFeedbackContextRequestParametersOrDict = Union[ - _UpdateRuntimeFeedbackContextRequestParameters, - _UpdateRuntimeFeedbackContextRequestParametersDict, -] +_GetEndpointParametersOrDict = Union[_GetEndpointParameters, _GetEndpointParametersDict] -class UndeployModelConfig(_common.BaseModel): - """Config for a Vertex SDK undeploy model from endpoint.""" +class ClientConnectionConfig(_common.BaseModel): + """Configurations (e.g. inference timeout) that are applied on your endpoints.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + inference_timeout: Optional[str] = Field( + default=None, description="""Customizable online prediction request timeout.""" ) - traffic_split: Optional[dict[str, int]] = Field( - default=None, - description="""If this field is provided, then the Endpoint's - [traffic_split][Endpoint.traffic_split] will be overwritten with it. If - last DeployedModel is being undeployed from the Endpoint, the - [Endpoint.traffic_split] will always end up empty when this call returns. - A DeployedModel will be successfully undeployed only if it doesn't have - any traffic assigned to it when this method executes, or if this field - unassigns any traffic to it.""", - ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Whether to wait for the long running operation to complete.""", - ) - - -class UndeployModelConfigDict(TypedDict, total=False): - """Config for a Vertex SDK undeploy model from endpoint.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - traffic_split: Optional[dict[str, int]] - """If this field is provided, then the Endpoint's - [traffic_split][Endpoint.traffic_split] will be overwritten with it. If - last DeployedModel is being undeployed from the Endpoint, the - [Endpoint.traffic_split] will always end up empty when this call returns. - A DeployedModel will be successfully undeployed only if it doesn't have - any traffic assigned to it when this method executes, or if this field - unassigns any traffic to it.""" +class ClientConnectionConfigDict(TypedDict, total=False): + """Configurations (e.g. inference timeout) that are applied on your endpoints.""" - wait_for_completion: Optional[bool] - """Whether to wait for the long running operation to complete.""" + inference_timeout: Optional[str] + """Customizable online prediction request timeout.""" -UndeployModelConfigOrDict = Union[UndeployModelConfig, UndeployModelConfigDict] +ClientConnectionConfigOrDict = Union[ClientConnectionConfig, ClientConnectionConfigDict] -class _UndeployModelRequestParameters(_common.BaseModel): - """Parameters for undeploying a model from an endpoint.""" +class ExplanationMetadataInputMetadataFeatureValueDomain(_common.BaseModel): + """Domain details of the input feature value. Provides numeric information about the feature, such as its range (min, max). If the feature has been pre-processed, for example with z-scoring, then it provides information about how to recover the original feature. For example, if the input feature is an image and it has been pre-processed to obtain 0-mean and stddev = 1 values, then original_mean, and original_stddev refer to the mean and stddev of the original feature (e.g. image tensor) from which input feature (with mean = 0 and stddev = 1) was obtained.""" - name: Optional[str] = Field( - default=None, description="""ID of the endpoint to undeploy the model from.""" + max_value: Optional[float] = Field( + default=None, description="""The maximum permissible value for this feature.""" ) - deployed_model_id: Optional[str] = Field( - default=None, description="""ID of the deployed model to be undeployed.""" + min_value: Optional[float] = Field( + default=None, description="""The minimum permissible value for this feature.""" + ) + original_mean: Optional[float] = Field( + default=None, + description="""If this input feature has been normalized to a mean value of 0, the original_mean specifies the mean value of the domain prior to normalization.""", + ) + original_stddev: Optional[float] = Field( + default=None, + description="""If this input feature has been normalized to a standard deviation of 1.0, the original_stddev specifies the standard deviation of the domain prior to normalization.""", ) - config: Optional[UndeployModelConfig] = Field(default=None, description="""""") -class _UndeployModelRequestParametersDict(TypedDict, total=False): - """Parameters for undeploying a model from an endpoint.""" +class ExplanationMetadataInputMetadataFeatureValueDomainDict(TypedDict, total=False): + """Domain details of the input feature value. Provides numeric information about the feature, such as its range (min, max). If the feature has been pre-processed, for example with z-scoring, then it provides information about how to recover the original feature. For example, if the input feature is an image and it has been pre-processed to obtain 0-mean and stddev = 1 values, then original_mean, and original_stddev refer to the mean and stddev of the original feature (e.g. image tensor) from which input feature (with mean = 0 and stddev = 1) was obtained.""" - name: Optional[str] - """ID of the endpoint to undeploy the model from.""" + max_value: Optional[float] + """The maximum permissible value for this feature.""" - deployed_model_id: Optional[str] - """ID of the deployed model to be undeployed.""" + min_value: Optional[float] + """The minimum permissible value for this feature.""" - config: Optional[UndeployModelConfigDict] - """""" + original_mean: Optional[float] + """If this input feature has been normalized to a mean value of 0, the original_mean specifies the mean value of the domain prior to normalization.""" + original_stddev: Optional[float] + """If this input feature has been normalized to a standard deviation of 1.0, the original_stddev specifies the standard deviation of the domain prior to normalization.""" -_UndeployModelRequestParametersOrDict = Union[ - _UndeployModelRequestParameters, _UndeployModelRequestParametersDict + +ExplanationMetadataInputMetadataFeatureValueDomainOrDict = Union[ + ExplanationMetadataInputMetadataFeatureValueDomain, + ExplanationMetadataInputMetadataFeatureValueDomainDict, ] -class UndeployModelOperation(_common.BaseModel): - """Operation for undeploying a model.""" +class ExplanationMetadataInputMetadataVisualization(_common.BaseModel): + """Visualization configurations for image explanation.""" - name: Optional[str] = Field( + clip_percent_lowerbound: Optional[float] = 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="""Excludes attributions below the specified percentile, from the highlighted areas. Defaults to 62.""", ) - metadata: Optional[dict[str, Any]] = Field( + clip_percent_upperbound: Optional[float] = 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="""Excludes attributions above the specified percentile from the highlighted areas. Using the clip_percent_upperbound and clip_percent_lowerbound together can be useful for filtering out noise and making it easier to see areas of strong attribution. Defaults to 99.9.""", ) - done: Optional[bool] = Field( + color_map: Optional[ColorMap] = 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="""The color scheme used for the highlighted areas. Defaults to PINK_GREEN for Integrated Gradients attribution, which shows positive attributions in green and negative in pink. Defaults to VIRIDIS for XRAI attribution, which highlights the most influential regions in yellow and the least influential in blue.""", ) - error: Optional[dict[str, Any]] = Field( + overlay_type: Optional[OverlayType] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""How the original image is displayed in the visualization. Adjusting the overlay can help increase visual clarity if the original image makes it difficult to view the visualization. Defaults to NONE.""", + ) + polarity: Optional[Polarity] = Field( + default=None, + description="""Whether to only highlight pixels with positive contributions, negative or both. Defaults to POSITIVE.""", + ) + type: Optional[Type] = Field( + default=None, + description="""Type of the image visualization. Only applicable to Integrated Gradients attribution. OUTLINES shows regions of attribution, while PIXELS shows per-pixel attribution. Defaults to OUTLINES.""", ) -class UndeployModelOperationDict(TypedDict, total=False): - """Operation for undeploying a model.""" +class ExplanationMetadataInputMetadataVisualizationDict(TypedDict, total=False): + """Visualization configurations for image explanation.""" - 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}`.""" + clip_percent_lowerbound: Optional[float] + """Excludes attributions below the specified percentile, from the highlighted areas. Defaults to 62.""" - 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.""" + clip_percent_upperbound: Optional[float] + """Excludes attributions above the specified percentile from the highlighted areas. Using the clip_percent_upperbound and clip_percent_lowerbound together can be useful for filtering out noise and making it easier to see areas of strong attribution. Defaults to 99.9.""" - 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.""" + color_map: Optional[ColorMap] + """The color scheme used for the highlighted areas. Defaults to PINK_GREEN for Integrated Gradients attribution, which shows positive attributions in green and negative in pink. Defaults to VIRIDIS for XRAI attribution, which highlights the most influential regions in yellow and the least influential in blue.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + overlay_type: Optional[OverlayType] + """How the original image is displayed in the visualization. Adjusting the overlay can help increase visual clarity if the original image makes it difficult to view the visualization. Defaults to NONE.""" + polarity: Optional[Polarity] + """Whether to only highlight pixels with positive contributions, negative or both. Defaults to POSITIVE.""" -UndeployModelOperationOrDict = Union[UndeployModelOperation, UndeployModelOperationDict] + type: Optional[Type] + """Type of the image visualization. Only applicable to Integrated Gradients attribution. OUTLINES shows regions of attribution, while PIXELS shows per-pixel attribution. Defaults to OUTLINES.""" -class PredictConfig(_common.BaseModel): - """Config for a Vertex SDK predict endpoint.""" +ExplanationMetadataInputMetadataVisualizationOrDict = Union[ + ExplanationMetadataInputMetadataVisualization, + ExplanationMetadataInputMetadataVisualizationDict, +] - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + +class ExplanationMetadataInputMetadata(_common.BaseModel): + """Metadata of the input of a feature. Fields other than InputMetadata.input_baselines are applicable only for Models that are using Vertex AI-provided images for Tensorflow.""" + + dense_shape_tensor_name: Optional[str] = Field( + default=None, + description="""Specifies the shape of the values of the input if the input is a sparse representation. Refer to Tensorflow documentation for more details: https://www.tensorflow.org/api_docs/python/tf/sparse/SparseTensor.""", ) - parameters: Optional[str] = Field( + encoded_baselines: Optional[list[Any]] = Field( default=None, - description="""The parameters that govern the prediction. The schema of the - parameters may be specified via Endpoint's DeployedModels' [Model's ][ - DeployedModel.model] [PredictSchemata's][Model.predict_schemata] - [parameters_schema_uri][PredictSchemata.parameters_schema_uri]. - """, + description="""A list of baselines for the encoded tensor. The shape of each baseline should match the shape of the encoded tensor. If a scalar is provided, Vertex AI broadcasts to the same shape as the encoded tensor.""", ) - - -class PredictConfigDict(TypedDict, total=False): - """Config for a Vertex SDK predict endpoint.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - parameters: Optional[str] - """The parameters that govern the prediction. The schema of the - parameters may be specified via Endpoint's DeployedModels' [Model's ][ - DeployedModel.model] [PredictSchemata's][Model.predict_schemata] - [parameters_schema_uri][PredictSchemata.parameters_schema_uri]. - """ - - -PredictConfigOrDict = Union[PredictConfig, PredictConfigDict] - - -class _PredictParameters(_common.BaseModel): - """Parameters for performing an online prediction.""" - - name: Optional[str] = Field( + encoded_tensor_name: Optional[str] = Field( default=None, - description="""The endpoint that serves the prediction. It could be endpoints/... - or publishers/.../models/... - """, + description="""Encoded tensor is a transformation of the input tensor. Must be provided if choosing Integrated Gradients attribution or XRAI attribution and the input tensor is not differentiable. An encoded tensor is generated if the input tensor is encoded by a lookup table.""", ) - instances: Optional[list[dict[str, Any]]] = Field( + encoding: Optional[Encoding] = Field( default=None, - description="""The instances that are the input to the prediction call. The - schema of any single instance may be specified via Endpoint's DeployedModels' - [Model's ][DeployedModel.model] [PredictSchemata's][Model.predict_schemata] - [instance_schema_uri][PredictSchemata.instance_schema_uri]. - """, + description="""Defines how the feature is encoded into the input tensor. Defaults to IDENTITY.""", ) - config: Optional[PredictConfig] = Field(default=None, description="""""") - - -class _PredictParametersDict(TypedDict, total=False): - """Parameters for performing an online prediction.""" - - name: Optional[str] - """The endpoint that serves the prediction. It could be endpoints/... - or publishers/.../models/... - """ - - instances: Optional[list[dict[str, Any]]] - """The instances that are the input to the prediction call. The - schema of any single instance may be specified via Endpoint's DeployedModels' - [Model's ][DeployedModel.model] [PredictSchemata's][Model.predict_schemata] - [instance_schema_uri][PredictSchemata.instance_schema_uri]. - """ - - config: Optional[PredictConfigDict] - """""" - - -_PredictParametersOrDict = Union[_PredictParameters, _PredictParametersDict] - - -class PredictResponse(_common.BaseModel): - """Response message for PredictionService.Predict API.""" - - deployed_model_id: Optional[str] = Field( + feature_value_domain: Optional[ + ExplanationMetadataInputMetadataFeatureValueDomain + ] = Field( default=None, - description="""ID of the Endpoint's DeployedModel that served this prediction.""", + description="""The domain details of the input feature value. Like min/max, original mean or standard deviation if normalized.""", ) - metadata: Optional[Any] = Field( + group_name: Optional[str] = Field( default=None, - description="""Output only. Request-level metadata returned by the model. The metadata type will be dependent upon the model implementation.""", + description="""Name of the group that the input belongs to. Features with the same group name will be treated as one feature when computing attributions. Features grouped together can have different shapes in value. If provided, there will be one single attribution generated in Attribution.feature_attributions, keyed by the group name.""", ) - model: Optional[str] = Field( + index_feature_mapping: Optional[list[str]] = Field( default=None, - description="""Output only. The resource name of the Model which is deployed as the DeployedModel that this prediction hits.""", + description="""A list of feature names for each index in the input tensor. Required when the input InputMetadata.encoding is BAG_OF_FEATURES, BAG_OF_FEATURES_SPARSE, INDICATOR.""", ) - model_display_name: Optional[str] = Field( + indices_tensor_name: Optional[str] = Field( default=None, - description="""Output only. The display name of the Model which is deployed as the DeployedModel that this prediction hits.""", + description="""Specifies the index of the values of the input tensor. Required when the input tensor is a sparse representation. Refer to Tensorflow documentation for more details: https://www.tensorflow.org/api_docs/python/tf/sparse/SparseTensor.""", ) - model_version_id: Optional[str] = Field( + input_baselines: Optional[list[Any]] = Field( default=None, - description="""Output only. The version ID of the Model which is deployed as the DeployedModel that this prediction hits.""", + description="""Baseline inputs for this feature. If no baseline is specified, Vertex AI chooses the baseline for this feature. If multiple baselines are specified, Vertex AI returns the average attributions across them in Attribution.feature_attributions. For Vertex AI-provided Tensorflow images (both 1.x and 2.x), the shape of each baseline must match the shape of the input tensor. If a scalar is provided, we broadcast to the same shape as the input tensor. For custom images, the element of the baselines must be in the same format as the feature's input in the instance[]. The schema of any single instance may be specified via Endpoint's DeployedModels' Model's PredictSchemata's instance_schema_uri.""", ) - predictions: Optional[list[Any]] = Field( + input_tensor_name: Optional[str] = Field( default=None, - description="""The predictions that are the output of the predictions call. The schema of any single prediction may be specified via Endpoint's DeployedModels' Model's PredictSchemata's prediction_schema_uri.""", + description="""Name of the input tensor for this feature. Required and is only applicable to Vertex AI-provided images for Tensorflow.""", + ) + modality: Optional[str] = Field( + default=None, + description="""Modality of the feature. Valid values are: numeric, image. Defaults to numeric.""", + ) + visualization: Optional[ExplanationMetadataInputMetadataVisualization] = Field( + default=None, + description="""Visualization configurations for image explanation.""", ) -class PredictResponseDict(TypedDict, total=False): - """Response message for PredictionService.Predict API.""" - - deployed_model_id: Optional[str] - """ID of the Endpoint's DeployedModel that served this prediction.""" - - metadata: Optional[Any] - """Output only. Request-level metadata returned by the model. The metadata type will be dependent upon the model implementation.""" - - model: Optional[str] - """Output only. The resource name of the Model which is deployed as the DeployedModel that this prediction hits.""" +class ExplanationMetadataInputMetadataDict(TypedDict, total=False): + """Metadata of the input of a feature. Fields other than InputMetadata.input_baselines are applicable only for Models that are using Vertex AI-provided images for Tensorflow.""" - model_display_name: Optional[str] - """Output only. The display name of the Model which is deployed as the DeployedModel that this prediction hits.""" + dense_shape_tensor_name: Optional[str] + """Specifies the shape of the values of the input if the input is a sparse representation. Refer to Tensorflow documentation for more details: https://www.tensorflow.org/api_docs/python/tf/sparse/SparseTensor.""" - model_version_id: Optional[str] - """Output only. The version ID of the Model which is deployed as the DeployedModel that this prediction hits.""" + encoded_baselines: Optional[list[Any]] + """A list of baselines for the encoded tensor. The shape of each baseline should match the shape of the encoded tensor. If a scalar is provided, Vertex AI broadcasts to the same shape as the encoded tensor.""" - predictions: Optional[list[Any]] - """The predictions that are the output of the predictions call. The schema of any single prediction may be specified via Endpoint's DeployedModels' Model's PredictSchemata's prediction_schema_uri.""" + encoded_tensor_name: Optional[str] + """Encoded tensor is a transformation of the input tensor. Must be provided if choosing Integrated Gradients attribution or XRAI attribution and the input tensor is not differentiable. An encoded tensor is generated if the input tensor is encoded by a lookup table.""" + encoding: Optional[Encoding] + """Defines how the feature is encoded into the input tensor. Defaults to IDENTITY.""" -PredictResponseOrDict = Union[PredictResponse, PredictResponseDict] + feature_value_domain: Optional[ + ExplanationMetadataInputMetadataFeatureValueDomainDict + ] + """The domain details of the input feature value. Like min/max, original mean or standard deviation if normalized.""" + group_name: Optional[str] + """Name of the group that the input belongs to. Features with the same group name will be treated as one feature when computing attributions. Features grouped together can have different shapes in value. If provided, there will be one single attribution generated in Attribution.feature_attributions, keyed by the group name.""" -class DeleteEndpointConfig(_common.BaseModel): - """Config for a Vertex SDK delete endpoint.""" + index_feature_mapping: Optional[list[str]] + """A list of feature names for each index in the input tensor. Required when the input InputMetadata.encoding is BAG_OF_FEATURES, BAG_OF_FEATURES_SPARSE, INDICATOR.""" - 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="""Whether to wait for the long running operation to complete.""", - ) + indices_tensor_name: Optional[str] + """Specifies the index of the values of the input tensor. Required when the input tensor is a sparse representation. Refer to Tensorflow documentation for more details: https://www.tensorflow.org/api_docs/python/tf/sparse/SparseTensor.""" + input_baselines: Optional[list[Any]] + """Baseline inputs for this feature. If no baseline is specified, Vertex AI chooses the baseline for this feature. If multiple baselines are specified, Vertex AI returns the average attributions across them in Attribution.feature_attributions. For Vertex AI-provided Tensorflow images (both 1.x and 2.x), the shape of each baseline must match the shape of the input tensor. If a scalar is provided, we broadcast to the same shape as the input tensor. For custom images, the element of the baselines must be in the same format as the feature's input in the instance[]. The schema of any single instance may be specified via Endpoint's DeployedModels' Model's PredictSchemata's instance_schema_uri.""" -class DeleteEndpointConfigDict(TypedDict, total=False): - """Config for a Vertex SDK delete endpoint.""" + input_tensor_name: Optional[str] + """Name of the input tensor for this feature. Required and is only applicable to Vertex AI-provided images for Tensorflow.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + modality: Optional[str] + """Modality of the feature. Valid values are: numeric, image. Defaults to numeric.""" - wait_for_completion: Optional[bool] - """Whether to wait for the long running operation to complete.""" + visualization: Optional[ExplanationMetadataInputMetadataVisualizationDict] + """Visualization configurations for image explanation.""" -DeleteEndpointConfigOrDict = Union[DeleteEndpointConfig, DeleteEndpointConfigDict] +ExplanationMetadataInputMetadataOrDict = Union[ + ExplanationMetadataInputMetadata, ExplanationMetadataInputMetadataDict +] -class _DeleteEndpointRequestParameters(_common.BaseModel): - """Parameters for deleting an endpoint.""" +class ExplanationMetadataOutputMetadata(_common.BaseModel): + """Metadata of the prediction output to be explained.""" - name: Optional[str] = Field( + display_name_mapping_key: Optional[str] = Field( default=None, - description="""Required. The resource name of the endpoint to be deleted.""", + description="""Specify a field name in the prediction to look for the display name. Use this if the prediction contains the display names for the outputs. The display names in the prediction must have the same shape of the outputs, so that it can be located by Attribution.output_index for a specific output.""", + ) + index_display_name_mapping: Optional[Any] = Field( + default=None, + description="""Static mapping between the index and display name. Use this if the outputs are a deterministic n-dimensional array, e.g. a list of scores of all the classes in a pre-defined order for a multi-classification Model. It's not feasible if the outputs are non-deterministic, e.g. the Model produces top-k classes or sort the outputs by their values. The shape of the value must be an n-dimensional array of strings. The number of dimensions must match that of the outputs to be explained. The Attribution.output_display_name is populated by locating in the mapping with Attribution.output_index.""", + ) + output_tensor_name: Optional[str] = Field( + default=None, + description="""Name of the output tensor. Required and is only applicable to Vertex AI provided images for Tensorflow.""", ) - config: Optional[DeleteEndpointConfig] = Field(default=None, description="""""") -class _DeleteEndpointRequestParametersDict(TypedDict, total=False): - """Parameters for deleting an endpoint.""" +class ExplanationMetadataOutputMetadataDict(TypedDict, total=False): + """Metadata of the prediction output to be explained.""" - name: Optional[str] - """Required. The resource name of the endpoint to be deleted.""" + display_name_mapping_key: Optional[str] + """Specify a field name in the prediction to look for the display name. Use this if the prediction contains the display names for the outputs. The display names in the prediction must have the same shape of the outputs, so that it can be located by Attribution.output_index for a specific output.""" - config: Optional[DeleteEndpointConfigDict] - """""" + index_display_name_mapping: Optional[Any] + """Static mapping between the index and display name. Use this if the outputs are a deterministic n-dimensional array, e.g. a list of scores of all the classes in a pre-defined order for a multi-classification Model. It's not feasible if the outputs are non-deterministic, e.g. the Model produces top-k classes or sort the outputs by their values. The shape of the value must be an n-dimensional array of strings. The number of dimensions must match that of the outputs to be explained. The Attribution.output_display_name is populated by locating in the mapping with Attribution.output_index.""" + output_tensor_name: Optional[str] + """Name of the output tensor. Required and is only applicable to Vertex AI provided images for Tensorflow.""" -_DeleteEndpointRequestParametersOrDict = Union[ - _DeleteEndpointRequestParameters, _DeleteEndpointRequestParametersDict + +ExplanationMetadataOutputMetadataOrDict = Union[ + ExplanationMetadataOutputMetadata, ExplanationMetadataOutputMetadataDict ] -class DeleteEndpointOperation(_common.BaseModel): - """Operation for deleting a endpoint.""" +class ExplanationMetadata(_common.BaseModel): + """Metadata describing the Model's input and output for explanation.""" - name: Optional[str] = Field( + feature_attributions_schema_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}`.""", + description="""Points to a YAML file stored on Google Cloud Storage describing the format of the feature attributions. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML tabular Models always have this field populated by Vertex AI. Note: The URI given on output may be different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""", ) - metadata: Optional[dict[str, Any]] = Field( + inputs: Optional[dict[str, ExplanationMetadataInputMetadata]] = 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="""Required. Map from feature names to feature input metadata. Keys are the name of the features. Values are the specification of the feature. An empty InputMetadata is valid. It describes a text feature which has the name specified as the key in ExplanationMetadata.inputs. The baseline of the empty feature is chosen by Vertex AI. For Vertex AI-provided Tensorflow images, the key can be any friendly name of the feature. Once specified, featureAttributions are keyed by this key (if not grouped with another feature). For custom images, the key must match with the key in instance.""", ) - done: Optional[bool] = Field( + latent_space_source: Optional[str] = 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="""Name of the source to generate embeddings for example based explanations.""", ) - error: Optional[dict[str, Any]] = Field( + outputs: Optional[dict[str, ExplanationMetadataOutputMetadata]] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Required. Map from output names to output metadata. For Vertex AI-provided Tensorflow images, keys can be any user defined string that consists of any UTF-8 characters. For custom images, keys are the name of the output field in the prediction to be explained. Currently only one key is allowed.""", ) -class DeleteEndpointOperationDict(TypedDict, total=False): - """Operation for deleting a endpoint.""" +class ExplanationMetadataDict(TypedDict, total=False): + """Metadata describing the Model's input and output for explanation.""" - 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}`.""" + feature_attributions_schema_uri: Optional[str] + """Points to a YAML file stored on Google Cloud Storage describing the format of the feature attributions. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML tabular Models always have this field populated by Vertex AI. Note: The URI given on output may be different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""" - 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.""" + inputs: Optional[dict[str, ExplanationMetadataInputMetadataDict]] + """Required. Map from feature names to feature input metadata. Keys are the name of the features. Values are the specification of the feature. An empty InputMetadata is valid. It describes a text feature which has the name specified as the key in ExplanationMetadata.inputs. The baseline of the empty feature is chosen by Vertex AI. For Vertex AI-provided Tensorflow images, the key can be any friendly name of the feature. Once specified, featureAttributions are keyed by this key (if not grouped with another feature). For custom images, the key must match with the key in instance.""" - 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.""" + latent_space_source: Optional[str] + """Name of the source to generate embeddings for example based explanations.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + outputs: Optional[dict[str, ExplanationMetadataOutputMetadataDict]] + """Required. Map from output names to output metadata. For Vertex AI-provided Tensorflow images, keys can be any user defined string that consists of any UTF-8 characters. For custom images, keys are the name of the output field in the prediction to be explained. Currently only one key is allowed.""" -DeleteEndpointOperationOrDict = Union[ - DeleteEndpointOperation, DeleteEndpointOperationDict -] +ExplanationMetadataOrDict = Union[ExplanationMetadata, ExplanationMetadataDict] -class GetEndpointConfig(_common.BaseModel): - """Optional parameters for endpoints.get method.""" +class ExamplesExampleGcsSource(_common.BaseModel): + """The Cloud Storage input instances.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + data_format: Optional[DataFormat] = Field( + default=None, + description="""The format in which instances are given, if not specified, assume it's JSONL format. Currently only JSONL format is supported.""", + ) + gcs_source: Optional[genai_types.GcsSource] = Field( + default=None, + description="""The Cloud Storage location for the input instances.""", ) -class GetEndpointConfigDict(TypedDict, total=False): - """Optional parameters for endpoints.get method.""" +class ExamplesExampleGcsSourceDict(TypedDict, total=False): + """The Cloud Storage input instances.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + data_format: Optional[DataFormat] + """The format in which instances are given, if not specified, assume it's JSONL format. Currently only JSONL format is supported.""" + gcs_source: Optional[genai_types.GcsSourceDict] + """The Cloud Storage location for the input instances.""" -GetEndpointConfigOrDict = Union[GetEndpointConfig, GetEndpointConfigDict] +ExamplesExampleGcsSourceOrDict = Union[ + ExamplesExampleGcsSource, ExamplesExampleGcsSourceDict +] -class _GetEndpointParameters(_common.BaseModel): - name: Optional[str] = Field( +class Presets(_common.BaseModel): + """Preset configuration for example-based explanations""" + + modality: Optional[Modality] = Field( default=None, - description="""Required. The resource name of the Endpoint to get.""", + description="""The modality of the uploaded model, which automatically configures the distance measurement and feature normalization for the underlying example index and queries. If your model does not precisely fit one of these types, it is okay to choose the closest type.""", ) - config: Optional[GetEndpointConfig] = Field( - default=None, description="""Optional parameters for the request.""" + query: Optional[Literal["PRECISE", "FAST"]] = Field( + default=None, + description="""Preset option controlling parameters for speed-precision trade-off when querying for examples. If omitted, defaults to `PRECISE`.""", ) -class _GetEndpointParametersDict(TypedDict, total=False): +class PresetsDict(TypedDict, total=False): + """Preset configuration for example-based explanations""" - name: Optional[str] - """Required. The resource name of the Endpoint to get.""" + modality: Optional[Modality] + """The modality of the uploaded model, which automatically configures the distance measurement and feature normalization for the underlying example index and queries. If your model does not precisely fit one of these types, it is okay to choose the closest type.""" - config: Optional[GetEndpointConfigDict] - """Optional parameters for the request.""" + query: Optional[Literal["PRECISE", "FAST"]] + """Preset option controlling parameters for speed-precision trade-off when querying for examples. If omitted, defaults to `PRECISE`.""" -_GetEndpointParametersOrDict = Union[_GetEndpointParameters, _GetEndpointParametersDict] +PresetsOrDict = Union[Presets, PresetsDict] -class ClientConnectionConfig(_common.BaseModel): - """Configurations (e.g. inference timeout) that are applied on your endpoints.""" +class Examples(_common.BaseModel): + """Example-based explainability that returns the nearest neighbors from the provided dataset.""" - inference_timeout: Optional[str] = Field( - default=None, description="""Customizable online prediction request timeout.""" + example_gcs_source: Optional[ExamplesExampleGcsSource] = Field( + default=None, description="""The Cloud Storage input instances.""" ) - - -class ClientConnectionConfigDict(TypedDict, total=False): - """Configurations (e.g. inference timeout) that are applied on your endpoints.""" - - inference_timeout: Optional[str] - """Customizable online prediction request timeout.""" - - -ClientConnectionConfigOrDict = Union[ClientConnectionConfig, ClientConnectionConfigDict] - - -class ExplanationMetadataInputMetadataFeatureValueDomain(_common.BaseModel): - """Domain details of the input feature value. Provides numeric information about the feature, such as its range (min, max). If the feature has been pre-processed, for example with z-scoring, then it provides information about how to recover the original feature. For example, if the input feature is an image and it has been pre-processed to obtain 0-mean and stddev = 1 values, then original_mean, and original_stddev refer to the mean and stddev of the original feature (e.g. image tensor) from which input feature (with mean = 0 and stddev = 1) was obtained.""" - - max_value: Optional[float] = Field( - default=None, description="""The maximum permissible value for this feature.""" + gcs_source: Optional[genai_types.GcsSource] = Field( + default=None, + description="""The Cloud Storage locations that contain the instances to be indexed for approximate nearest neighbor search.""", ) - min_value: Optional[float] = Field( - default=None, description="""The minimum permissible value for this feature.""" + nearest_neighbor_search_config: Optional[Any] = Field( + default=None, + description="""The full configuration for the generated index, the semantics are the same as metadata and should match [NearestNeighborSearchConfig](https://cloud.google.com/vertex-ai/docs/explainable-ai/configuring-explanations-example-based#nearest-neighbor-search-config).""", ) - original_mean: Optional[float] = Field( + neighbor_count: Optional[int] = Field( default=None, - description="""If this input feature has been normalized to a mean value of 0, the original_mean specifies the mean value of the domain prior to normalization.""", + description="""The number of neighbors to return when querying for examples.""", ) - original_stddev: Optional[float] = Field( + presets: Optional[Presets] = Field( default=None, - description="""If this input feature has been normalized to a standard deviation of 1.0, the original_stddev specifies the standard deviation of the domain prior to normalization.""", + description="""Simplified preset configuration, which automatically sets configuration values based on the desired query speed-precision trade-off and modality.""", ) -class ExplanationMetadataInputMetadataFeatureValueDomainDict(TypedDict, total=False): - """Domain details of the input feature value. Provides numeric information about the feature, such as its range (min, max). If the feature has been pre-processed, for example with z-scoring, then it provides information about how to recover the original feature. For example, if the input feature is an image and it has been pre-processed to obtain 0-mean and stddev = 1 values, then original_mean, and original_stddev refer to the mean and stddev of the original feature (e.g. image tensor) from which input feature (with mean = 0 and stddev = 1) was obtained.""" +class ExamplesDict(TypedDict, total=False): + """Example-based explainability that returns the nearest neighbors from the provided dataset.""" - max_value: Optional[float] - """The maximum permissible value for this feature.""" + example_gcs_source: Optional[ExamplesExampleGcsSourceDict] + """The Cloud Storage input instances.""" - min_value: Optional[float] - """The minimum permissible value for this feature.""" + gcs_source: Optional[genai_types.GcsSourceDict] + """The Cloud Storage locations that contain the instances to be indexed for approximate nearest neighbor search.""" - original_mean: Optional[float] - """If this input feature has been normalized to a mean value of 0, the original_mean specifies the mean value of the domain prior to normalization.""" + nearest_neighbor_search_config: Optional[Any] + """The full configuration for the generated index, the semantics are the same as metadata and should match [NearestNeighborSearchConfig](https://cloud.google.com/vertex-ai/docs/explainable-ai/configuring-explanations-example-based#nearest-neighbor-search-config).""" - original_stddev: Optional[float] - """If this input feature has been normalized to a standard deviation of 1.0, the original_stddev specifies the standard deviation of the domain prior to normalization.""" + neighbor_count: Optional[int] + """The number of neighbors to return when querying for examples.""" + presets: Optional[PresetsDict] + """Simplified preset configuration, which automatically sets configuration values based on the desired query speed-precision trade-off and modality.""" -ExplanationMetadataInputMetadataFeatureValueDomainOrDict = Union[ - ExplanationMetadataInputMetadataFeatureValueDomain, - ExplanationMetadataInputMetadataFeatureValueDomainDict, -] +ExamplesOrDict = Union[Examples, ExamplesDict] -class ExplanationMetadataInputMetadataVisualization(_common.BaseModel): - """Visualization configurations for image explanation.""" - clip_percent_lowerbound: Optional[float] = Field( - default=None, - description="""Excludes attributions below the specified percentile, from the highlighted areas. Defaults to 62.""", - ) - clip_percent_upperbound: Optional[float] = Field( - default=None, - description="""Excludes attributions above the specified percentile from the highlighted areas. Using the clip_percent_upperbound and clip_percent_lowerbound together can be useful for filtering out noise and making it easier to see areas of strong attribution. Defaults to 99.9.""", - ) - color_map: Optional[ColorMap] = Field( - default=None, - description="""The color scheme used for the highlighted areas. Defaults to PINK_GREEN for Integrated Gradients attribution, which shows positive attributions in green and negative in pink. Defaults to VIRIDIS for XRAI attribution, which highlights the most influential regions in yellow and the least influential in blue.""", - ) - overlay_type: Optional[OverlayType] = Field( - default=None, - description="""How the original image is displayed in the visualization. Adjusting the overlay can help increase visual clarity if the original image makes it difficult to view the visualization. Defaults to NONE.""", - ) - polarity: Optional[Polarity] = Field( - default=None, - description="""Whether to only highlight pixels with positive contributions, negative or both. Defaults to POSITIVE.""", - ) - type: Optional[Type] = Field( +class BlurBaselineConfig(_common.BaseModel): + """Config for blur baseline. When enabled, a linear path from the maximally blurred image to the input image is created. Using a blurred baseline instead of zero (black image) is motivated by the BlurIG approach explained here: https://arxiv.org/abs/2004.03383""" + + max_blur_sigma: Optional[float] = Field( default=None, - description="""Type of the image visualization. Only applicable to Integrated Gradients attribution. OUTLINES shows regions of attribution, while PIXELS shows per-pixel attribution. Defaults to OUTLINES.""", + description="""The standard deviation of the blur kernel for the blurred baseline. The same blurring parameter is used for both the height and the width dimension. If not set, the method defaults to the zero (i.e. black for images) baseline.""", ) -class ExplanationMetadataInputMetadataVisualizationDict(TypedDict, total=False): - """Visualization configurations for image explanation.""" - - clip_percent_lowerbound: Optional[float] - """Excludes attributions below the specified percentile, from the highlighted areas. Defaults to 62.""" - - clip_percent_upperbound: Optional[float] - """Excludes attributions above the specified percentile from the highlighted areas. Using the clip_percent_upperbound and clip_percent_lowerbound together can be useful for filtering out noise and making it easier to see areas of strong attribution. Defaults to 99.9.""" - - color_map: Optional[ColorMap] - """The color scheme used for the highlighted areas. Defaults to PINK_GREEN for Integrated Gradients attribution, which shows positive attributions in green and negative in pink. Defaults to VIRIDIS for XRAI attribution, which highlights the most influential regions in yellow and the least influential in blue.""" - - overlay_type: Optional[OverlayType] - """How the original image is displayed in the visualization. Adjusting the overlay can help increase visual clarity if the original image makes it difficult to view the visualization. Defaults to NONE.""" - - polarity: Optional[Polarity] - """Whether to only highlight pixels with positive contributions, negative or both. Defaults to POSITIVE.""" +class BlurBaselineConfigDict(TypedDict, total=False): + """Config for blur baseline. When enabled, a linear path from the maximally blurred image to the input image is created. Using a blurred baseline instead of zero (black image) is motivated by the BlurIG approach explained here: https://arxiv.org/abs/2004.03383""" - type: Optional[Type] - """Type of the image visualization. Only applicable to Integrated Gradients attribution. OUTLINES shows regions of attribution, while PIXELS shows per-pixel attribution. Defaults to OUTLINES.""" + max_blur_sigma: Optional[float] + """The standard deviation of the blur kernel for the blurred baseline. The same blurring parameter is used for both the height and the width dimension. If not set, the method defaults to the zero (i.e. black for images) baseline.""" -ExplanationMetadataInputMetadataVisualizationOrDict = Union[ - ExplanationMetadataInputMetadataVisualization, - ExplanationMetadataInputMetadataVisualizationDict, -] +BlurBaselineConfigOrDict = Union[BlurBaselineConfig, BlurBaselineConfigDict] -class ExplanationMetadataInputMetadata(_common.BaseModel): - """Metadata of the input of a feature. Fields other than InputMetadata.input_baselines are applicable only for Models that are using Vertex AI-provided images for Tensorflow.""" +class FeatureNoiseSigmaNoiseSigmaForFeature(_common.BaseModel): + """Noise sigma for a single feature.""" - dense_shape_tensor_name: Optional[str] = Field( - default=None, - description="""Specifies the shape of the values of the input if the input is a sparse representation. Refer to Tensorflow documentation for more details: https://www.tensorflow.org/api_docs/python/tf/sparse/SparseTensor.""", - ) - encoded_baselines: Optional[list[Any]] = Field( - default=None, - description="""A list of baselines for the encoded tensor. The shape of each baseline should match the shape of the encoded tensor. If a scalar is provided, Vertex AI broadcasts to the same shape as the encoded tensor.""", - ) - encoded_tensor_name: Optional[str] = Field( - default=None, - description="""Encoded tensor is a transformation of the input tensor. Must be provided if choosing Integrated Gradients attribution or XRAI attribution and the input tensor is not differentiable. An encoded tensor is generated if the input tensor is encoded by a lookup table.""", - ) - encoding: Optional[Encoding] = Field( - default=None, - description="""Defines how the feature is encoded into the input tensor. Defaults to IDENTITY.""", - ) - feature_value_domain: Optional[ - ExplanationMetadataInputMetadataFeatureValueDomain - ] = Field( - default=None, - description="""The domain details of the input feature value. Like min/max, original mean or standard deviation if normalized.""", - ) - group_name: Optional[str] = Field( - default=None, - description="""Name of the group that the input belongs to. Features with the same group name will be treated as one feature when computing attributions. Features grouped together can have different shapes in value. If provided, there will be one single attribution generated in Attribution.feature_attributions, keyed by the group name.""", - ) - index_feature_mapping: Optional[list[str]] = Field( - default=None, - description="""A list of feature names for each index in the input tensor. Required when the input InputMetadata.encoding is BAG_OF_FEATURES, BAG_OF_FEATURES_SPARSE, INDICATOR.""", - ) - indices_tensor_name: Optional[str] = Field( - default=None, - description="""Specifies the index of the values of the input tensor. Required when the input tensor is a sparse representation. Refer to Tensorflow documentation for more details: https://www.tensorflow.org/api_docs/python/tf/sparse/SparseTensor.""", - ) - input_baselines: Optional[list[Any]] = Field( - default=None, - description="""Baseline inputs for this feature. If no baseline is specified, Vertex AI chooses the baseline for this feature. If multiple baselines are specified, Vertex AI returns the average attributions across them in Attribution.feature_attributions. For Vertex AI-provided Tensorflow images (both 1.x and 2.x), the shape of each baseline must match the shape of the input tensor. If a scalar is provided, we broadcast to the same shape as the input tensor. For custom images, the element of the baselines must be in the same format as the feature's input in the instance[]. The schema of any single instance may be specified via Endpoint's DeployedModels' Model's PredictSchemata's instance_schema_uri.""", - ) - input_tensor_name: Optional[str] = Field( - default=None, - description="""Name of the input tensor for this feature. Required and is only applicable to Vertex AI-provided images for Tensorflow.""", - ) - modality: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Modality of the feature. Valid values are: numeric, image. Defaults to numeric.""", + description="""The name of the input feature for which noise sigma is provided. The features are defined in explanation metadata inputs.""", ) - visualization: Optional[ExplanationMetadataInputMetadataVisualization] = Field( + sigma: Optional[float] = Field( default=None, - description="""Visualization configurations for image explanation.""", + description="""This represents the standard deviation of the Gaussian kernel that will be used to add noise to the feature prior to computing gradients. Similar to noise_sigma but represents the noise added to the current feature. Defaults to 0.1.""", ) -class ExplanationMetadataInputMetadataDict(TypedDict, total=False): - """Metadata of the input of a feature. Fields other than InputMetadata.input_baselines are applicable only for Models that are using Vertex AI-provided images for Tensorflow.""" - - dense_shape_tensor_name: Optional[str] - """Specifies the shape of the values of the input if the input is a sparse representation. Refer to Tensorflow documentation for more details: https://www.tensorflow.org/api_docs/python/tf/sparse/SparseTensor.""" - - encoded_baselines: Optional[list[Any]] - """A list of baselines for the encoded tensor. The shape of each baseline should match the shape of the encoded tensor. If a scalar is provided, Vertex AI broadcasts to the same shape as the encoded tensor.""" +class FeatureNoiseSigmaNoiseSigmaForFeatureDict(TypedDict, total=False): + """Noise sigma for a single feature.""" - encoded_tensor_name: Optional[str] - """Encoded tensor is a transformation of the input tensor. Must be provided if choosing Integrated Gradients attribution or XRAI attribution and the input tensor is not differentiable. An encoded tensor is generated if the input tensor is encoded by a lookup table.""" + name: Optional[str] + """The name of the input feature for which noise sigma is provided. The features are defined in explanation metadata inputs.""" - encoding: Optional[Encoding] - """Defines how the feature is encoded into the input tensor. Defaults to IDENTITY.""" + sigma: Optional[float] + """This represents the standard deviation of the Gaussian kernel that will be used to add noise to the feature prior to computing gradients. Similar to noise_sigma but represents the noise added to the current feature. Defaults to 0.1.""" - feature_value_domain: Optional[ - ExplanationMetadataInputMetadataFeatureValueDomainDict - ] - """The domain details of the input feature value. Like min/max, original mean or standard deviation if normalized.""" - group_name: Optional[str] - """Name of the group that the input belongs to. Features with the same group name will be treated as one feature when computing attributions. Features grouped together can have different shapes in value. If provided, there will be one single attribution generated in Attribution.feature_attributions, keyed by the group name.""" +FeatureNoiseSigmaNoiseSigmaForFeatureOrDict = Union[ + FeatureNoiseSigmaNoiseSigmaForFeature, FeatureNoiseSigmaNoiseSigmaForFeatureDict +] - index_feature_mapping: Optional[list[str]] - """A list of feature names for each index in the input tensor. Required when the input InputMetadata.encoding is BAG_OF_FEATURES, BAG_OF_FEATURES_SPARSE, INDICATOR.""" - indices_tensor_name: Optional[str] - """Specifies the index of the values of the input tensor. Required when the input tensor is a sparse representation. Refer to Tensorflow documentation for more details: https://www.tensorflow.org/api_docs/python/tf/sparse/SparseTensor.""" +class FeatureNoiseSigma(_common.BaseModel): + """Noise sigma by features. Noise sigma represents the standard deviation of the gaussian kernel that will be used to add noise to interpolated inputs prior to computing gradients.""" - input_baselines: Optional[list[Any]] - """Baseline inputs for this feature. If no baseline is specified, Vertex AI chooses the baseline for this feature. If multiple baselines are specified, Vertex AI returns the average attributions across them in Attribution.feature_attributions. For Vertex AI-provided Tensorflow images (both 1.x and 2.x), the shape of each baseline must match the shape of the input tensor. If a scalar is provided, we broadcast to the same shape as the input tensor. For custom images, the element of the baselines must be in the same format as the feature's input in the instance[]. The schema of any single instance may be specified via Endpoint's DeployedModels' Model's PredictSchemata's instance_schema_uri.""" + noise_sigma: Optional[list[FeatureNoiseSigmaNoiseSigmaForFeature]] = Field( + default=None, + description="""Noise sigma per feature. No noise is added to features that are not set.""", + ) - input_tensor_name: Optional[str] - """Name of the input tensor for this feature. Required and is only applicable to Vertex AI-provided images for Tensorflow.""" - modality: Optional[str] - """Modality of the feature. Valid values are: numeric, image. Defaults to numeric.""" +class FeatureNoiseSigmaDict(TypedDict, total=False): + """Noise sigma by features. Noise sigma represents the standard deviation of the gaussian kernel that will be used to add noise to interpolated inputs prior to computing gradients.""" - visualization: Optional[ExplanationMetadataInputMetadataVisualizationDict] - """Visualization configurations for image explanation.""" + noise_sigma: Optional[list[FeatureNoiseSigmaNoiseSigmaForFeatureDict]] + """Noise sigma per feature. No noise is added to features that are not set.""" -ExplanationMetadataInputMetadataOrDict = Union[ - ExplanationMetadataInputMetadata, ExplanationMetadataInputMetadataDict -] +FeatureNoiseSigmaOrDict = Union[FeatureNoiseSigma, FeatureNoiseSigmaDict] -class ExplanationMetadataOutputMetadata(_common.BaseModel): - """Metadata of the prediction output to be explained.""" +class SmoothGradConfig(_common.BaseModel): + """Config for SmoothGrad approximation of gradients. When enabled, the gradients are approximated by averaging the gradients from noisy samples in the vicinity of the inputs. Adding noise can help improve the computed gradients. Refer to this paper for more details: https://arxiv.org/pdf/1706.03825.pdf""" - display_name_mapping_key: Optional[str] = Field( + feature_noise_sigma: Optional[FeatureNoiseSigma] = Field( default=None, - description="""Specify a field name in the prediction to look for the display name. Use this if the prediction contains the display names for the outputs. The display names in the prediction must have the same shape of the outputs, so that it can be located by Attribution.output_index for a specific output.""", + description="""This is similar to noise_sigma, but provides additional flexibility. A separate noise sigma can be provided for each feature, which is useful if their distributions are different. No noise is added to features that are not set. If this field is unset, noise_sigma will be used for all features.""", ) - index_display_name_mapping: Optional[Any] = Field( + noise_sigma: Optional[float] = Field( default=None, - description="""Static mapping between the index and display name. Use this if the outputs are a deterministic n-dimensional array, e.g. a list of scores of all the classes in a pre-defined order for a multi-classification Model. It's not feasible if the outputs are non-deterministic, e.g. the Model produces top-k classes or sort the outputs by their values. The shape of the value must be an n-dimensional array of strings. The number of dimensions must match that of the outputs to be explained. The Attribution.output_display_name is populated by locating in the mapping with Attribution.output_index.""", + description="""This is a single float value and will be used to add noise to all the features. Use this field when all features are normalized to have the same distribution: scale to range [0, 1], [-1, 1] or z-scoring, where features are normalized to have 0-mean and 1-variance. Learn more about [normalization](https://developers.google.com/machine-learning/data-prep/transform/normalization). For best results the recommended value is about 10% - 20% of the standard deviation of the input feature. Refer to section 3.2 of the SmoothGrad paper: https://arxiv.org/pdf/1706.03825.pdf. Defaults to 0.1. If the distribution is different per feature, set feature_noise_sigma instead for each feature.""", ) - output_tensor_name: Optional[str] = Field( + noisy_sample_count: Optional[int] = Field( default=None, - description="""Name of the output tensor. Required and is only applicable to Vertex AI provided images for Tensorflow.""", + description="""The number of gradient samples to use for approximation. The higher this number, the more accurate the gradient is, but the runtime complexity increases by this factor as well. Valid range of its value is [1, 50]. Defaults to 3.""", ) -class ExplanationMetadataOutputMetadataDict(TypedDict, total=False): - """Metadata of the prediction output to be explained.""" +class SmoothGradConfigDict(TypedDict, total=False): + """Config for SmoothGrad approximation of gradients. When enabled, the gradients are approximated by averaging the gradients from noisy samples in the vicinity of the inputs. Adding noise can help improve the computed gradients. Refer to this paper for more details: https://arxiv.org/pdf/1706.03825.pdf""" - display_name_mapping_key: Optional[str] - """Specify a field name in the prediction to look for the display name. Use this if the prediction contains the display names for the outputs. The display names in the prediction must have the same shape of the outputs, so that it can be located by Attribution.output_index for a specific output.""" + feature_noise_sigma: Optional[FeatureNoiseSigmaDict] + """This is similar to noise_sigma, but provides additional flexibility. A separate noise sigma can be provided for each feature, which is useful if their distributions are different. No noise is added to features that are not set. If this field is unset, noise_sigma will be used for all features.""" - index_display_name_mapping: Optional[Any] - """Static mapping between the index and display name. Use this if the outputs are a deterministic n-dimensional array, e.g. a list of scores of all the classes in a pre-defined order for a multi-classification Model. It's not feasible if the outputs are non-deterministic, e.g. the Model produces top-k classes or sort the outputs by their values. The shape of the value must be an n-dimensional array of strings. The number of dimensions must match that of the outputs to be explained. The Attribution.output_display_name is populated by locating in the mapping with Attribution.output_index.""" + noise_sigma: Optional[float] + """This is a single float value and will be used to add noise to all the features. Use this field when all features are normalized to have the same distribution: scale to range [0, 1], [-1, 1] or z-scoring, where features are normalized to have 0-mean and 1-variance. Learn more about [normalization](https://developers.google.com/machine-learning/data-prep/transform/normalization). For best results the recommended value is about 10% - 20% of the standard deviation of the input feature. Refer to section 3.2 of the SmoothGrad paper: https://arxiv.org/pdf/1706.03825.pdf. Defaults to 0.1. If the distribution is different per feature, set feature_noise_sigma instead for each feature.""" - output_tensor_name: Optional[str] - """Name of the output tensor. Required and is only applicable to Vertex AI provided images for Tensorflow.""" + noisy_sample_count: Optional[int] + """The number of gradient samples to use for approximation. The higher this number, the more accurate the gradient is, but the runtime complexity increases by this factor as well. Valid range of its value is [1, 50]. Defaults to 3.""" -ExplanationMetadataOutputMetadataOrDict = Union[ - ExplanationMetadataOutputMetadata, ExplanationMetadataOutputMetadataDict -] +SmoothGradConfigOrDict = Union[SmoothGradConfig, SmoothGradConfigDict] -class ExplanationMetadata(_common.BaseModel): - """Metadata describing the Model's input and output for explanation.""" +class IntegratedGradientsAttribution(_common.BaseModel): + """An attribution method that computes the Aumann-Shapley value taking advantage of the model's fully differentiable structure. Refer to this paper for more details: https://arxiv.org/abs/1703.01365""" - feature_attributions_schema_uri: Optional[str] = Field( - default=None, - description="""Points to a YAML file stored on Google Cloud Storage describing the format of the feature attributions. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML tabular Models always have this field populated by Vertex AI. Note: The URI given on output may be different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""", - ) - inputs: Optional[dict[str, ExplanationMetadataInputMetadata]] = Field( + blur_baseline_config: Optional[BlurBaselineConfig] = Field( default=None, - description="""Required. Map from feature names to feature input metadata. Keys are the name of the features. Values are the specification of the feature. An empty InputMetadata is valid. It describes a text feature which has the name specified as the key in ExplanationMetadata.inputs. The baseline of the empty feature is chosen by Vertex AI. For Vertex AI-provided Tensorflow images, the key can be any friendly name of the feature. Once specified, featureAttributions are keyed by this key (if not grouped with another feature). For custom images, the key must match with the key in instance.""", + description="""Config for IG with blur baseline. When enabled, a linear path from the maximally blurred image to the input image is created. Using a blurred baseline instead of zero (black image) is motivated by the BlurIG approach explained here: https://arxiv.org/abs/2004.03383""", ) - latent_space_source: Optional[str] = Field( + smooth_grad_config: Optional[SmoothGradConfig] = Field( default=None, - description="""Name of the source to generate embeddings for example based explanations.""", + description="""Config for SmoothGrad approximation of gradients. When enabled, the gradients are approximated by averaging the gradients from noisy samples in the vicinity of the inputs. Adding noise can help improve the computed gradients. Refer to this paper for more details: https://arxiv.org/pdf/1706.03825.pdf""", ) - outputs: Optional[dict[str, ExplanationMetadataOutputMetadata]] = Field( + step_count: Optional[int] = Field( default=None, - description="""Required. Map from output names to output metadata. For Vertex AI-provided Tensorflow images, keys can be any user defined string that consists of any UTF-8 characters. For custom images, keys are the name of the output field in the prediction to be explained. Currently only one key is allowed.""", + description="""Required. The number of steps for approximating the path integral. A good value to start is 50 and gradually increase until the sum to diff property is within the desired error range. Valid range of its value is [1, 100], inclusively.""", ) -class ExplanationMetadataDict(TypedDict, total=False): - """Metadata describing the Model's input and output for explanation.""" - - feature_attributions_schema_uri: Optional[str] - """Points to a YAML file stored on Google Cloud Storage describing the format of the feature attributions. The schema is defined as an OpenAPI 3.0.2 [Schema Object](https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.2.md#schemaObject). AutoML tabular Models always have this field populated by Vertex AI. Note: The URI given on output may be different, including the URI scheme, than the one given on input. The output URI will point to a location where the user only has a read access.""" +class IntegratedGradientsAttributionDict(TypedDict, total=False): + """An attribution method that computes the Aumann-Shapley value taking advantage of the model's fully differentiable structure. Refer to this paper for more details: https://arxiv.org/abs/1703.01365""" - inputs: Optional[dict[str, ExplanationMetadataInputMetadataDict]] - """Required. Map from feature names to feature input metadata. Keys are the name of the features. Values are the specification of the feature. An empty InputMetadata is valid. It describes a text feature which has the name specified as the key in ExplanationMetadata.inputs. The baseline of the empty feature is chosen by Vertex AI. For Vertex AI-provided Tensorflow images, the key can be any friendly name of the feature. Once specified, featureAttributions are keyed by this key (if not grouped with another feature). For custom images, the key must match with the key in instance.""" + blur_baseline_config: Optional[BlurBaselineConfigDict] + """Config for IG with blur baseline. When enabled, a linear path from the maximally blurred image to the input image is created. Using a blurred baseline instead of zero (black image) is motivated by the BlurIG approach explained here: https://arxiv.org/abs/2004.03383""" - latent_space_source: Optional[str] - """Name of the source to generate embeddings for example based explanations.""" + smooth_grad_config: Optional[SmoothGradConfigDict] + """Config for SmoothGrad approximation of gradients. When enabled, the gradients are approximated by averaging the gradients from noisy samples in the vicinity of the inputs. Adding noise can help improve the computed gradients. Refer to this paper for more details: https://arxiv.org/pdf/1706.03825.pdf""" - outputs: Optional[dict[str, ExplanationMetadataOutputMetadataDict]] - """Required. Map from output names to output metadata. For Vertex AI-provided Tensorflow images, keys can be any user defined string that consists of any UTF-8 characters. For custom images, keys are the name of the output field in the prediction to be explained. Currently only one key is allowed.""" + step_count: Optional[int] + """Required. The number of steps for approximating the path integral. A good value to start is 50 and gradually increase until the sum to diff property is within the desired error range. Valid range of its value is [1, 100], inclusively.""" -ExplanationMetadataOrDict = Union[ExplanationMetadata, ExplanationMetadataDict] +IntegratedGradientsAttributionOrDict = Union[ + IntegratedGradientsAttribution, IntegratedGradientsAttributionDict +] -class ExamplesExampleGcsSource(_common.BaseModel): - """The Cloud Storage input instances.""" - - data_format: Optional[DataFormat] = Field( - default=None, - description="""The format in which instances are given, if not specified, assume it's JSONL format. Currently only JSONL format is supported.""", - ) - gcs_source: Optional[genai_types.GcsSource] = Field( - default=None, - description="""The Cloud Storage location for the input instances.""", - ) - - -class ExamplesExampleGcsSourceDict(TypedDict, total=False): - """The Cloud Storage input instances.""" - - data_format: Optional[DataFormat] - """The format in which instances are given, if not specified, assume it's JSONL format. Currently only JSONL format is supported.""" - - gcs_source: Optional[genai_types.GcsSourceDict] - """The Cloud Storage location for the input instances.""" - - -ExamplesExampleGcsSourceOrDict = Union[ - ExamplesExampleGcsSource, ExamplesExampleGcsSourceDict -] - - -class Presets(_common.BaseModel): - """Preset configuration for example-based explanations""" - - modality: Optional[Modality] = Field( - default=None, - description="""The modality of the uploaded model, which automatically configures the distance measurement and feature normalization for the underlying example index and queries. If your model does not precisely fit one of these types, it is okay to choose the closest type.""", - ) - query: Optional[Literal["PRECISE", "FAST"]] = Field( - default=None, - description="""Preset option controlling parameters for speed-precision trade-off when querying for examples. If omitted, defaults to `PRECISE`.""", - ) - - -class PresetsDict(TypedDict, total=False): - """Preset configuration for example-based explanations""" - - modality: Optional[Modality] - """The modality of the uploaded model, which automatically configures the distance measurement and feature normalization for the underlying example index and queries. If your model does not precisely fit one of these types, it is okay to choose the closest type.""" - - query: Optional[Literal["PRECISE", "FAST"]] - """Preset option controlling parameters for speed-precision trade-off when querying for examples. If omitted, defaults to `PRECISE`.""" - - -PresetsOrDict = Union[Presets, PresetsDict] - - -class Examples(_common.BaseModel): - """Example-based explainability that returns the nearest neighbors from the provided dataset.""" - - example_gcs_source: Optional[ExamplesExampleGcsSource] = Field( - default=None, description="""The Cloud Storage input instances.""" - ) - gcs_source: Optional[genai_types.GcsSource] = Field( - default=None, - description="""The Cloud Storage locations that contain the instances to be indexed for approximate nearest neighbor search.""", - ) - nearest_neighbor_search_config: Optional[Any] = Field( - default=None, - description="""The full configuration for the generated index, the semantics are the same as metadata and should match [NearestNeighborSearchConfig](https://cloud.google.com/vertex-ai/docs/explainable-ai/configuring-explanations-example-based#nearest-neighbor-search-config).""", - ) - neighbor_count: Optional[int] = Field( - default=None, - description="""The number of neighbors to return when querying for examples.""", - ) - presets: Optional[Presets] = Field( - default=None, - description="""Simplified preset configuration, which automatically sets configuration values based on the desired query speed-precision trade-off and modality.""", - ) - - -class ExamplesDict(TypedDict, total=False): - """Example-based explainability that returns the nearest neighbors from the provided dataset.""" - - example_gcs_source: Optional[ExamplesExampleGcsSourceDict] - """The Cloud Storage input instances.""" - - gcs_source: Optional[genai_types.GcsSourceDict] - """The Cloud Storage locations that contain the instances to be indexed for approximate nearest neighbor search.""" - - nearest_neighbor_search_config: Optional[Any] - """The full configuration for the generated index, the semantics are the same as metadata and should match [NearestNeighborSearchConfig](https://cloud.google.com/vertex-ai/docs/explainable-ai/configuring-explanations-example-based#nearest-neighbor-search-config).""" - - neighbor_count: Optional[int] - """The number of neighbors to return when querying for examples.""" - - presets: Optional[PresetsDict] - """Simplified preset configuration, which automatically sets configuration values based on the desired query speed-precision trade-off and modality.""" - - -ExamplesOrDict = Union[Examples, ExamplesDict] - - -class BlurBaselineConfig(_common.BaseModel): - """Config for blur baseline. When enabled, a linear path from the maximally blurred image to the input image is created. Using a blurred baseline instead of zero (black image) is motivated by the BlurIG approach explained here: https://arxiv.org/abs/2004.03383""" - - max_blur_sigma: Optional[float] = Field( - default=None, - description="""The standard deviation of the blur kernel for the blurred baseline. The same blurring parameter is used for both the height and the width dimension. If not set, the method defaults to the zero (i.e. black for images) baseline.""", - ) - - -class BlurBaselineConfigDict(TypedDict, total=False): - """Config for blur baseline. When enabled, a linear path from the maximally blurred image to the input image is created. Using a blurred baseline instead of zero (black image) is motivated by the BlurIG approach explained here: https://arxiv.org/abs/2004.03383""" - - max_blur_sigma: Optional[float] - """The standard deviation of the blur kernel for the blurred baseline. The same blurring parameter is used for both the height and the width dimension. If not set, the method defaults to the zero (i.e. black for images) baseline.""" - - -BlurBaselineConfigOrDict = Union[BlurBaselineConfig, BlurBaselineConfigDict] - - -class FeatureNoiseSigmaNoiseSigmaForFeature(_common.BaseModel): - """Noise sigma for a single feature.""" - - name: Optional[str] = Field( - default=None, - description="""The name of the input feature for which noise sigma is provided. The features are defined in explanation metadata inputs.""", - ) - sigma: Optional[float] = Field( - default=None, - description="""This represents the standard deviation of the Gaussian kernel that will be used to add noise to the feature prior to computing gradients. Similar to noise_sigma but represents the noise added to the current feature. Defaults to 0.1.""", - ) - - -class FeatureNoiseSigmaNoiseSigmaForFeatureDict(TypedDict, total=False): - """Noise sigma for a single feature.""" - - name: Optional[str] - """The name of the input feature for which noise sigma is provided. The features are defined in explanation metadata inputs.""" - - sigma: Optional[float] - """This represents the standard deviation of the Gaussian kernel that will be used to add noise to the feature prior to computing gradients. Similar to noise_sigma but represents the noise added to the current feature. Defaults to 0.1.""" - - -FeatureNoiseSigmaNoiseSigmaForFeatureOrDict = Union[ - FeatureNoiseSigmaNoiseSigmaForFeature, FeatureNoiseSigmaNoiseSigmaForFeatureDict -] - - -class FeatureNoiseSigma(_common.BaseModel): - """Noise sigma by features. Noise sigma represents the standard deviation of the gaussian kernel that will be used to add noise to interpolated inputs prior to computing gradients.""" - - noise_sigma: Optional[list[FeatureNoiseSigmaNoiseSigmaForFeature]] = Field( - default=None, - description="""Noise sigma per feature. No noise is added to features that are not set.""", - ) - - -class FeatureNoiseSigmaDict(TypedDict, total=False): - """Noise sigma by features. Noise sigma represents the standard deviation of the gaussian kernel that will be used to add noise to interpolated inputs prior to computing gradients.""" - - noise_sigma: Optional[list[FeatureNoiseSigmaNoiseSigmaForFeatureDict]] - """Noise sigma per feature. No noise is added to features that are not set.""" - - -FeatureNoiseSigmaOrDict = Union[FeatureNoiseSigma, FeatureNoiseSigmaDict] - - -class SmoothGradConfig(_common.BaseModel): - """Config for SmoothGrad approximation of gradients. When enabled, the gradients are approximated by averaging the gradients from noisy samples in the vicinity of the inputs. Adding noise can help improve the computed gradients. Refer to this paper for more details: https://arxiv.org/pdf/1706.03825.pdf""" - - feature_noise_sigma: Optional[FeatureNoiseSigma] = Field( - default=None, - description="""This is similar to noise_sigma, but provides additional flexibility. A separate noise sigma can be provided for each feature, which is useful if their distributions are different. No noise is added to features that are not set. If this field is unset, noise_sigma will be used for all features.""", - ) - noise_sigma: Optional[float] = Field( - default=None, - description="""This is a single float value and will be used to add noise to all the features. Use this field when all features are normalized to have the same distribution: scale to range [0, 1], [-1, 1] or z-scoring, where features are normalized to have 0-mean and 1-variance. Learn more about [normalization](https://developers.google.com/machine-learning/data-prep/transform/normalization). For best results the recommended value is about 10% - 20% of the standard deviation of the input feature. Refer to section 3.2 of the SmoothGrad paper: https://arxiv.org/pdf/1706.03825.pdf. Defaults to 0.1. If the distribution is different per feature, set feature_noise_sigma instead for each feature.""", - ) - noisy_sample_count: Optional[int] = Field( - default=None, - description="""The number of gradient samples to use for approximation. The higher this number, the more accurate the gradient is, but the runtime complexity increases by this factor as well. Valid range of its value is [1, 50]. Defaults to 3.""", - ) - - -class SmoothGradConfigDict(TypedDict, total=False): - """Config for SmoothGrad approximation of gradients. When enabled, the gradients are approximated by averaging the gradients from noisy samples in the vicinity of the inputs. Adding noise can help improve the computed gradients. Refer to this paper for more details: https://arxiv.org/pdf/1706.03825.pdf""" - - feature_noise_sigma: Optional[FeatureNoiseSigmaDict] - """This is similar to noise_sigma, but provides additional flexibility. A separate noise sigma can be provided for each feature, which is useful if their distributions are different. No noise is added to features that are not set. If this field is unset, noise_sigma will be used for all features.""" - - noise_sigma: Optional[float] - """This is a single float value and will be used to add noise to all the features. Use this field when all features are normalized to have the same distribution: scale to range [0, 1], [-1, 1] or z-scoring, where features are normalized to have 0-mean and 1-variance. Learn more about [normalization](https://developers.google.com/machine-learning/data-prep/transform/normalization). For best results the recommended value is about 10% - 20% of the standard deviation of the input feature. Refer to section 3.2 of the SmoothGrad paper: https://arxiv.org/pdf/1706.03825.pdf. Defaults to 0.1. If the distribution is different per feature, set feature_noise_sigma instead for each feature.""" - - noisy_sample_count: Optional[int] - """The number of gradient samples to use for approximation. The higher this number, the more accurate the gradient is, but the runtime complexity increases by this factor as well. Valid range of its value is [1, 50]. Defaults to 3.""" - - -SmoothGradConfigOrDict = Union[SmoothGradConfig, SmoothGradConfigDict] - - -class IntegratedGradientsAttribution(_common.BaseModel): - """An attribution method that computes the Aumann-Shapley value taking advantage of the model's fully differentiable structure. Refer to this paper for more details: https://arxiv.org/abs/1703.01365""" - - blur_baseline_config: Optional[BlurBaselineConfig] = Field( - default=None, - description="""Config for IG with blur baseline. When enabled, a linear path from the maximally blurred image to the input image is created. Using a blurred baseline instead of zero (black image) is motivated by the BlurIG approach explained here: https://arxiv.org/abs/2004.03383""", - ) - smooth_grad_config: Optional[SmoothGradConfig] = Field( - default=None, - description="""Config for SmoothGrad approximation of gradients. When enabled, the gradients are approximated by averaging the gradients from noisy samples in the vicinity of the inputs. Adding noise can help improve the computed gradients. Refer to this paper for more details: https://arxiv.org/pdf/1706.03825.pdf""", - ) - step_count: Optional[int] = Field( - default=None, - description="""Required. The number of steps for approximating the path integral. A good value to start is 50 and gradually increase until the sum to diff property is within the desired error range. Valid range of its value is [1, 100], inclusively.""", - ) - - -class IntegratedGradientsAttributionDict(TypedDict, total=False): - """An attribution method that computes the Aumann-Shapley value taking advantage of the model's fully differentiable structure. Refer to this paper for more details: https://arxiv.org/abs/1703.01365""" - - blur_baseline_config: Optional[BlurBaselineConfigDict] - """Config for IG with blur baseline. When enabled, a linear path from the maximally blurred image to the input image is created. Using a blurred baseline instead of zero (black image) is motivated by the BlurIG approach explained here: https://arxiv.org/abs/2004.03383""" - - smooth_grad_config: Optional[SmoothGradConfigDict] - """Config for SmoothGrad approximation of gradients. When enabled, the gradients are approximated by averaging the gradients from noisy samples in the vicinity of the inputs. Adding noise can help improve the computed gradients. Refer to this paper for more details: https://arxiv.org/pdf/1706.03825.pdf""" - - step_count: Optional[int] - """Required. The number of steps for approximating the path integral. A good value to start is 50 and gradually increase until the sum to diff property is within the desired error range. Valid range of its value is [1, 100], inclusively.""" - - -IntegratedGradientsAttributionOrDict = Union[ - IntegratedGradientsAttribution, IntegratedGradientsAttributionDict -] - - -class SampledShapleyAttribution(_common.BaseModel): - """An attribution method that approximates Shapley values for features that contribute to the label being predicted. A sampling strategy is used to approximate the value rather than considering all subsets of features.""" +class SampledShapleyAttribution(_common.BaseModel): + """An attribution method that approximates Shapley values for features that contribute to the label being predicted. A sampling strategy is used to approximate the value rather than considering all subsets of features.""" path_count: Optional[int] = Field( default=None, @@ -31157,629 +29975,17 @@ 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)" - ) - - 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[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] - - -RubricContentProperty = evals_types.RubricContentProperty -RubricContentPropertyDict = evals_types.RubricContentPropertyDict -RubricContentPropertyDictOrDict = evals_types.RubricContentPropertyOrDict - -RubricContent = evals_types.RubricContent -RubricContentDict = evals_types.RubricContentDict -RubricContentDictOrDict = evals_types.RubricContentOrDict - -Rubric = evals_types.Rubric -RubricDict = evals_types.RubricDict -RubricDictOrDict = evals_types.RubricOrDict - -RubricVerdict = evals_types.RubricVerdict -RubricVerdictDict = evals_types.RubricVerdictDict -RubricVerdictDictOrDict = evals_types.RubricVerdictOrDict - -CandidateResult = evals_types.CandidateResult -CandidateResultDict = evals_types.CandidateResultDict -CandidateResultDictOrDict = evals_types.CandidateResultOrDict - -Event = evals_types.Event -EventDict = evals_types.EventDict -EventDictOrDict = evals_types.EventOrDict - -Message = evals_types.Message -MessageDict = evals_types.MessageDict -MessageDictOrDict = evals_types.MessageOrDict - -Importance = evals_types.Importance - - -class AgentEngineDict(TypedDict, total=False): - """An agent engine instance.""" - - api_client: Optional[Any] - """The underlying API client.""" - - api_async_client: Optional[Any] - """The underlying API client for asynchronous operations.""" - - api_resource: Optional[ReasoningEngineDict] - """The underlying API resource (i.e. ReasoningEngine).""" - - -AgentEngineOrDict = Union[AgentEngine, AgentEngineDict] - - -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. - - 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. - - 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 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. - - 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.""" - ) - build_config: Optional[ReasoningEngineSpecBuildSpec] = Field( - default=None, - description="""The build config for the Agent Engine. Allows bringing your own Cloud Build private worker pool (BYOBP) and, optionally, a build-time service account for the container build. Supported keys: `worker_pool` (the resource name of the Cloud Build WorkerPool to use for the build) and `service_account` (the service account that Cloud Build uses to run the build; only applicable when `worker_pool` is specified).""", - ) - - -class AgentEngineConfigDict(TypedDict, total=False): - """Config for agent engine methods.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - 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.""" - - 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.""" - - display_name: Optional[str] - """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] - """The description of the Agent Engine.""" - - gcs_dir_name: Optional[str] - """The GCS bucket directory under `staging_bucket` to use for staging - the artifacts needed.""" - - extra_packages: Optional[list[str]] - """The set of extra user-provided packages (if any).""" - - env_vars: Optional[Any] - """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] - """The service account to be used for the Agent Engine. - - If not specified, the default Reasoning Engine P6SA service agent will be used.""" - - identity_type: Optional[IdentityType] - """The identity type to use for the Agent Engine.""" - - context_spec: Optional[ReasoningEngineContextSpecDict] - """The context spec to be used for the Agent Engine.""" - - psc_interface_config: Optional[PscInterfaceConfigDict] - """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'. - """ - - container_concurrency: Optional[int] - """The container concurrency to be used for the Agent Engine. - Recommended value: 2 * cpu + 1. Defaults to 9. - """ - - encryption_spec: Optional[genai_types.EncryptionSpec] - """The encryption spec to be used for the Agent Engine.""" - - labels: Optional[dict[str, str]] - """The labels to be used for the Agent Engine.""" - - 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 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]] - """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.""" - - entrypoint_module: Optional[str] - """The entrypoint module to be used for the Agent Engine - This field only used when source_packages is specified.""" - - 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. - """ - - 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".""" - - 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`. - """ - - image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict] - """The image spec for the Agent Engine.""" - - agent_config_source: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict - ] - """The agent config source for the Agent Engine.""" - - container_spec: Optional[ReasoningEngineSpecContainerSpecDict] - """The container spec for the Agent Engine.""" - - 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.""" - - traffic_config: Optional[ReasoningEngineTrafficConfigDict] - """The traffic config for the Agent Engine.""" - - build_config: Optional[ReasoningEngineSpecBuildSpecDict] - """The build config for the Agent Engine. Allows bringing your own Cloud Build private worker pool (BYOBP) and, optionally, a build-time service account for the container build. Supported keys: `worker_pool` (the resource name of the Cloud Build WorkerPool to use for the build) and `service_account` (the service account that Cloud Build uses to run the build; only applicable when `worker_pool` is specified).""" - - -AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict] - - -class RunQueryJobAgentEngineConfig(_common.BaseModel): - """Config for checking a query job on an agent engine.""" - - 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.""" - ) - output_gcs_uri: Optional[str] = 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.""", - ) - - -class RunQueryJobAgentEngineConfigDict(TypedDict, total=False): - """Config for checking a query job on an agent engine.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - query: Optional[str] - """The query to send to the agent engine.""" - - 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.""" - - -RunQueryJobAgentEngineConfigOrDict = Union[ - RunQueryJobAgentEngineConfig, RunQueryJobAgentEngineConfigDict -] - - -class RunQueryJobResult(_common.BaseModel): - """Result of running a query job.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - job_name: Optional[str] = Field( - default=None, - description="""Name of the agent engine 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 RunQueryJobResultDict(TypedDict, total=False): - """Result of running a query job.""" - - 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.""" - - input_gcs_uri: Optional[str] - """The GCS URI of the input file.""" - - output_gcs_uri: Optional[str] - """The GCS URI of the output file.""" - - -RunQueryJobResultOrDict = Union[RunQueryJobResult, RunQueryJobResultDict] - - -class CheckQueryJobResponse(_common.BaseModel): - """Response from LRO.""" - - 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.""" - - 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.""" +class AssembleDatasetDict(TypedDict, total=False): + """Represents the assembled dataset.""" bigquery_destination: Optional[str] """The BigQuery destination of the assembled dataset.""" @@ -32068,717 +30274,1327 @@ class CreatePromptConfig(_common.BaseModel): 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.""", + 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 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.""" + + 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.""" + + 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 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.""" + ) + 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.""", + ) + 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.""", + ) + 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 CreatePromptVersionConfigDict(TypedDict, total=False): + """Config for creating a prompt version.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + version_display_name: Optional[str] + """The display name for the prompt version. 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.""" + + 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.""" + + max_wait_time: Optional[int] + """The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""" + + +CreatePromptVersionConfigOrDict = Union[ + CreatePromptVersionConfig, CreatePromptVersionConfigDict +] + + +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.""" + ) + + +class GetPromptConfigDict(TypedDict, total=False): + """Config for getting a prompt.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetPromptConfigOrDict = Union[GetPromptConfig, GetPromptConfigDict] + + +class PromptRef(_common.BaseModel): + """Reference to a prompt.""" + + prompt_id: Optional[str] = Field(default=None, description="""""") + model: Optional[str] = Field(default=None, description="""""") + + +class PromptRefDict(TypedDict, total=False): + """Reference to a prompt.""" + + prompt_id: Optional[str] + """""" + + model: Optional[str] + """""" + + +PromptRefOrDict = Union[PromptRef, PromptRefDict] + + +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 PromptVersionRefDict(TypedDict, total=False): + """Reference to a prompt version.""" + + prompt_id: Optional[str] + """""" + + version_id: Optional[str] + """""" + + model: Optional[str] + """""" + + +PromptVersionRefOrDict = Union[PromptVersionRef, PromptVersionRefDict] + + +class OptimizeJobConfig(_common.BaseModel): + """VAPO Prompt Optimizer Config.""" + + config_path: Optional[str] = Field( + default=None, + description="""The gcs path to the config file, e.g. gs://bucket/config.json.""", + ) + service_account: 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.""", + ) + 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 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.""" + + 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.""" + + 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.""" + + 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.""" + + +OptimizeJobConfigOrDict = Union[OptimizeJobConfig, OptimizeJobConfigDict] + + +class ListDeployableModelsConfig(_common.BaseModel): + """Config for listing deployable models.""" + + 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.""" ) -class CreatePromptConfigDict(TypedDict, total=False): - """Config for creating a prompt.""" +class ListDeployableModelsConfigDict(TypedDict, total=False): + """Config for listing deployable models.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + include_hugging_face_models: Optional[bool] + """Whether to list Hugging Face models.""" - prompt_display_name: Optional[str] - """The display name for the prompt. If not set, a default name with a timestamp will be used.""" + model_filter: Optional[str] + """Optional. A string to filter the models by.""" - 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.""" +ListDeployableModelsConfigOrDict = Union[ + ListDeployableModelsConfig, ListDeployableModelsConfigDict +] - 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.""" +class ListModelGardenModelsConfig(_common.BaseModel): + """Config for listing Model Garden models.""" + 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.""" + ) -CreatePromptConfigOrDict = Union[CreatePromptConfig, CreatePromptConfigDict] +class ListModelGardenModelsConfigDict(TypedDict, total=False): + """Config for listing Model Garden models.""" -class CreatePromptVersionConfig(_common.BaseModel): - """Config for creating a prompt version.""" + include_hugging_face_models: Optional[bool] + """Whether to list Hugging Face models.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - version_display_name: Optional[str] = Field( + model_filter: Optional[str] + """Optional. A string to filter the models by.""" + + +ListModelGardenModelsConfigOrDict = Union[ + ListModelGardenModelsConfig, ListModelGardenModelsConfigDict +] + + +class ListPublisherModelDeployOptionsConfig(_common.BaseModel): + """Config for listing the deploy options of a publisher model.""" + + machine_type_filter: Optional[Union[str, list[str]]] = Field( default=None, - description="""The display name for the prompt version. 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.""", + 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.""", ) - prompt_display_name: Optional[str] = Field( + accelerator_type_filter: Optional[Union[str, list[str]]] = Field( default=None, - description="""The display name for the prompt. If not set, a default name with a timestamp will be used.""", + 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.""", ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + serving_container_image_uri_filter: Optional[Union[str, list[str]]] = 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="""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.""", ) - 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.""", + 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.""", ) -class CreatePromptVersionConfigDict(TypedDict, total=False): - """Config for creating a prompt version.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - version_display_name: Optional[str] - """The display name for the prompt version. If not set, a default name with a timestamp will be used.""" +class ListPublisherModelDeployOptionsConfigDict(TypedDict, total=False): + """Config for listing the deploy options of a publisher model.""" - timeout: Optional[int] - """The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""" + 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.""" - prompt_display_name: Optional[str] - """The display name for the prompt. If not set, a default name with a timestamp will be used.""" + 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.""" - 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.""" + 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.""" - max_wait_time: Optional[int] - """The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""" + 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.""" -CreatePromptVersionConfigOrDict = Union[ - CreatePromptVersionConfig, CreatePromptVersionConfigDict +ListPublisherModelDeployOptionsConfigOrDict = Union[ + ListPublisherModelDeployOptionsConfig, ListPublisherModelDeployOptionsConfigDict ] -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.""" - ) - - -class GetPromptConfigDict(TypedDict, total=False): - """Config for getting a prompt.""" +class ListCustomModelDeployOptionsConfig(_common.BaseModel): + """Config for listing custom model deploy options.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + 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. -GetPromptConfigOrDict = Union[GetPromptConfig, GetPromptConfigDict] + 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 PromptRef(_common.BaseModel): - """Reference to a prompt.""" +class ListCustomModelDeployOptionsConfigDict(TypedDict, total=False): + """Config for listing custom model deploy options.""" - prompt_id: Optional[str] = Field(default=None, description="""""") - model: Optional[str] = Field(default=None, description="""""") + 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 PromptRefDict(TypedDict, total=False): - """Reference to a prompt.""" + check_machine_availability: Optional[bool] + """Whether to check per-region machine availability. - prompt_id: Optional[str] - """""" + 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). + """ - model: Optional[str] - """""" +ListCustomModelDeployOptionsConfigOrDict = Union[ + ListCustomModelDeployOptionsConfig, ListCustomModelDeployOptionsConfigDict +] -PromptRefOrDict = Union[PromptRef, PromptRefDict] +class ExportOpenModelConfig(_common.BaseModel): + """Config for export_open_model.""" -class PromptVersionRef(_common.BaseModel): - """Reference to a prompt version.""" + 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``.""", + ) + 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``.""", + ) - prompt_id: Optional[str] = Field(default=None, description="""""") - version_id: Optional[str] = Field(default=None, description="""""") - model: Optional[str] = Field(default=None, description="""""") +class ExportOpenModelConfigDict(TypedDict, total=False): + """Config for export_open_model.""" -class PromptVersionRefDict(TypedDict, total=False): - """Reference to a prompt version.""" + 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.""" - prompt_id: Optional[str] - """""" + poll_interval_seconds: Optional[float] + """Seconds between LRO polls when ``wait_for_completion=True``. + Defaults to 30. Ignored when ``wait_for_completion=False``.""" - version_id: Optional[str] - """""" + 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``.""" - model: Optional[str] - """""" +ExportOpenModelConfigOrDict = Union[ExportOpenModelConfig, ExportOpenModelConfigDict] -PromptVersionRefOrDict = Union[PromptVersionRef, PromptVersionRefDict] +class DeployPublisherModelConfig(_common.BaseModel): + """Config for ``deploy_publisher_model``. -class OptimizeJobConfig(_common.BaseModel): - """VAPO Prompt Optimizer Config.""" + Superset of options that apply to Google open, partner and Hugging Face + publisher models. Only fields relevant to the target model are honored; + the backend rejects unsupported fields with a clear error. + """ - config_path: Optional[str] = Field( + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to block on the deployment long-running operation. When + ``True`` (default), returns the ``DeployResponse`` (deployed endpoint + and model resource names) on completion. When ``False``, returns the + ``DeployModelOperation`` for the caller to poll.""", + ) + poll_interval_seconds: Optional[float] = Field( default=None, - description="""The gcs path to the config file, e.g. gs://bucket/config.json.""", + description="""Seconds between LRO polls when ``wait_for_completion=True``. + Defaults to 30. Ignored when ``wait_for_completion=False``.""", ) - service_account: Optional[str] = Field( + timeout_seconds: Optional[float] = 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="""Total wall-clock seconds to wait for the deployment to complete + when ``wait_for_completion=True``. Defaults to 2 hours, matching the + Vertex AI Console one-click deployment timeout. Ignored when + ``wait_for_completion=False``.""", ) - service_account_project_number: Optional[Union[int, str]] = Field( + accept_eula: Optional[bool] = 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="""Whether to accept the model's End User License Agreement.""", ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Whether to wait for the job tocomplete. Ignored for async jobs.""", + hugging_face_access_token: Optional[str] = Field( + default=None, + description="""Hugging Face access token for gated HF models. See + https://huggingface.co/docs/hub/en/security-tokens.""", ) - optimizer_job_display_name: Optional[str] = Field( + machine_type: 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.""", + description="""Machine type (e.g. ``'g2-standard-48'``). Leave unset for + automatic resources.""", + ) + min_replica_count: Optional[int] = Field( + default=1, description="""Minimum number of replicas.""" + ) + max_replica_count: Optional[int] = Field( + default=1, description="""Maximum number of replicas.""" + ) + accelerator_type: Optional[str] = Field( + default=None, description="""Accelerator type (e.g. ``'NVIDIA_L4'``).""" + ) + accelerator_count: Optional[int] = Field( + default=None, description="""Number of accelerators per replica.""" + ) + spot: Optional[bool] = Field(default=None, description="""Schedule on Spot VMs.""") + dedicated_endpoint_disabled: Optional[bool] = Field( + default=None, + description="""Set True to serve predictions via the shared endpoint DNS + instead of the dedicated endpoint DNS (default).""", + ) + fast_tryout_enabled: Optional[bool] = Field( + default=None, + description="""Use the fast-tryout deployment path (experimentation only, not + production). Only supported for select models and machine types.""", + ) + endpoint_display_name: Optional[str] = Field( + default=None, description="""Display name for the endpoint.""" + ) + model_display_name: Optional[str] = Field( + default=None, description="""Display name for the deployed model.""" + ) + serving_container_image_uri: Optional[str] = Field( + default=None, + description="""Custom serving container image URI overriding the model's + default container.""", + ) + container_command: Optional[list[str]] = Field( + default=None, description="""Serving container ENTRYPOINT override.""" + ) + container_args: Optional[list[str]] = Field( + default=None, description="""Serving container CMD override.""" + ) + container_variables: Optional[dict[str, str]] = Field( + default=None, description="""Environment variables for the serving container.""" + ) + enable_private_service_connect: Optional[bool] = Field( + default=None, description="""Enable Private Service Connect for the endpoint.""" + ) + psc_project_allow_list: Optional[list[str]] = Field( + default=None, + description="""Projects allowed to access the endpoint over Private Service + Connect. Only honored when ``enable_private_service_connect`` is True.""", ) -class OptimizeJobConfigDict(TypedDict, total=False): - """VAPO Prompt Optimizer Config.""" +class DeployPublisherModelConfigDict(TypedDict, total=False): + """Config for ``deploy_publisher_model``. - config_path: Optional[str] - """The gcs path to the config file, e.g. gs://bucket/config.json.""" + Superset of options that apply to Google open, partner and Hugging Face + publisher models. Only fields relevant to the target model are honored; + the backend rejects unsupported fields with a clear error. + """ - 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.""" + wait_for_completion: Optional[bool] + """Whether to block on the deployment long-running operation. When + ``True`` (default), returns the ``DeployResponse`` (deployed endpoint + and model resource names) on completion. When ``False``, returns the + ``DeployModelOperation`` for the caller to poll.""" - 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".""" + poll_interval_seconds: Optional[float] + """Seconds between LRO polls when ``wait_for_completion=True``. + Defaults to 30. Ignored when ``wait_for_completion=False``.""" - wait_for_completion: Optional[bool] - """Whether to wait for the job tocomplete. Ignored for async jobs.""" + timeout_seconds: Optional[float] + """Total wall-clock seconds to wait for the deployment to complete + when ``wait_for_completion=True``. Defaults to 2 hours, matching the + Vertex AI Console one-click deployment timeout. Ignored when + ``wait_for_completion=False``.""" - 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.""" + accept_eula: Optional[bool] + """Whether to accept the model's End User License Agreement.""" + + hugging_face_access_token: Optional[str] + """Hugging Face access token for gated HF models. See + https://huggingface.co/docs/hub/en/security-tokens.""" + machine_type: Optional[str] + """Machine type (e.g. ``'g2-standard-48'``). Leave unset for + automatic resources.""" -OptimizeJobConfigOrDict = Union[OptimizeJobConfig, OptimizeJobConfigDict] + min_replica_count: Optional[int] + """Minimum number of replicas.""" + max_replica_count: Optional[int] + """Maximum number of replicas.""" -class AgentEngineRuntimeRevision(_common.BaseModel): - """An agent engine runtime revision instance.""" + accelerator_type: Optional[str] + """Accelerator type (e.g. ``'NVIDIA_L4'``).""" - 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).""", - ) + accelerator_count: Optional[int] + """Number of accelerators per replica.""" - # Allows dynamic binding of methods based on the registered operations. - model_config = ConfigDict(extra="allow") + spot: Optional[bool] + """Schedule on Spot VMs.""" - 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)" - ) + dedicated_endpoint_disabled: Optional[bool] + """Set True to serve predictions via the shared endpoint DNS + instead of the dedicated endpoint DNS (default).""" - 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 + fast_tryout_enabled: Optional[bool] + """Use the fast-tryout deployment path (experimentation only, not + production). Only supported for select models and machine types.""" - def delete( - self, - config: Optional[DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, - ) -> None: - """Deletes the agent engine runtime revision. + endpoint_display_name: Optional[str] + """Display name for the endpoint.""" - 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] + model_display_name: Optional[str] + """Display name for the deployed model.""" + serving_container_image_uri: Optional[str] + """Custom serving container image URI overriding the model's + default container.""" -class AgentEngineRuntimeRevisionDict(TypedDict, total=False): - """An agent engine runtime revision instance.""" + container_command: Optional[list[str]] + """Serving container ENTRYPOINT override.""" - api_client: Optional[Any] - """The underlying API client.""" + container_args: Optional[list[str]] + """Serving container CMD override.""" - api_async_client: Optional[Any] - """The underlying API client for asynchronous operations.""" + container_variables: Optional[dict[str, str]] + """Environment variables for the serving container.""" - api_resource: Optional[ReasoningEngineRuntimeRevisionDict] - """The underlying API resource (i.e. ReasoningEngineRuntimeRevision).""" + enable_private_service_connect: Optional[bool] + """Enable Private Service Connect for the endpoint.""" + + psc_project_allow_list: Optional[list[str]] + """Projects allowed to access the endpoint over Private Service + Connect. Only honored when ``enable_private_service_connect`` is True.""" -AgentEngineRuntimeRevisionOrDict = Union[ - AgentEngineRuntimeRevision, AgentEngineRuntimeRevisionDict +DeployPublisherModelConfigOrDict = Union[ + DeployPublisherModelConfig, DeployPublisherModelConfigDict ] -class ListDeployableModelsConfig(_common.BaseModel): - """Config for listing deployable models.""" +class DeployOption(_common.BaseModel): + """A verified deploy option for a model.""" - include_hugging_face_models: Optional[bool] = Field( - default=None, description="""Whether to list Hugging Face models.""" + option_name: Optional[str] = Field( + default=None, description="""The name of the deploy task.""" ) - model_filter: Optional[str] = Field( - default=None, description="""Optional. A string to filter the models by.""" + 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.""" + ) + accelerator_type: Optional[str] = Field( + default=None, description="""The accelerator type.""" + ) + accelerator_count: Optional[int] = Field( + default=None, description="""The number of accelerators.""" ) -class ListDeployableModelsConfigDict(TypedDict, total=False): - """Config for listing deployable models.""" +class DeployOptionDict(TypedDict, total=False): + """A verified deploy option for a model.""" - include_hugging_face_models: Optional[bool] - """Whether to list Hugging Face models.""" + option_name: Optional[str] + """The name of the deploy task.""" - model_filter: Optional[str] - """Optional. A string to filter the models by.""" + serving_container_image_uri: Optional[str] + """The URI of the serving container.""" + machine_type: Optional[str] + """The machine type.""" -ListDeployableModelsConfigOrDict = Union[ - ListDeployableModelsConfig, ListDeployableModelsConfigDict -] + accelerator_type: Optional[str] + """The accelerator type.""" + accelerator_count: Optional[int] + """The number of accelerators.""" -class ListModelGardenModelsConfig(_common.BaseModel): - """Config for listing Model Garden models.""" - 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.""" - ) +DeployOptionOrDict = Union[DeployOption, DeployOptionDict] + +class ListMemoryBanksResponse(_common.BaseModel): + """The response for listing Memory Banks.""" -class ListModelGardenModelsConfigDict(TypedDict, total=False): - """Config for listing Model Garden models.""" + memory_banks: Optional[list[MemoryBank]] = Field( + default=None, description="""The list of Memory Banks.""" + ) - include_hugging_face_models: Optional[bool] - """Whether to list Hugging Face models.""" - model_filter: Optional[str] - """Optional. A string to filter the models by.""" +class ListMemoryBanksResponseDict(TypedDict, total=False): + """The response for listing Memory Banks.""" + memory_banks: Optional[list[MemoryBankDict]] + """The list of Memory Banks.""" -ListModelGardenModelsConfigOrDict = Union[ - ListModelGardenModelsConfig, ListModelGardenModelsConfigDict + +ListMemoryBanksResponseOrDict = Union[ + ListMemoryBanksResponse, ListMemoryBanksResponseDict ] -class ListPublisherModelDeployOptionsConfig(_common.BaseModel): - """Config for listing the deploy options of a publisher model.""" +class Runtime(_common.BaseModel): + """An agent runtime instance.""" - 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.""", + api_client: Optional[Any] = Field( + default=None, description="""The underlying API client.""" ) - serving_container_image_uri_filter: Optional[Union[str, list[str]]] = Field( + api_async_client: Optional[Any] = 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.""", + description="""The underlying API client for asynchronous operations.""", ) - concise: Optional[bool] = Field( + api_resource: Optional[ReasoningEngine] = 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 underlying API resource (i.e. ReasoningEngine).""", ) + # Allows dynamic binding of methods based on the registered operations. + model_config = ConfigDict(extra="allow") -class ListPublisherModelDeployOptionsConfigDict(TypedDict, total=False): - """Config for listing the deploy options of a publisher model.""" + 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)" + ) - 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.""" + 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 - 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.""" + def delete( + self, + force: bool = False, + config: Optional[DeleteRuntimeConfigOrDict] = None, + ) -> None: + """Deletes the agent engine. - 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.""" + 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] - 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.""" +RubricContentProperty = evals_types.RubricContentProperty +RubricContentPropertyDict = evals_types.RubricContentPropertyDict +RubricContentPropertyDictOrDict = evals_types.RubricContentPropertyOrDict -ListPublisherModelDeployOptionsConfigOrDict = Union[ - ListPublisherModelDeployOptionsConfig, ListPublisherModelDeployOptionsConfigDict -] +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 ListCustomModelDeployOptionsConfig(_common.BaseModel): - """Config for listing custom model deploy options.""" +RubricVerdict = evals_types.RubricVerdict +RubricVerdictDict = evals_types.RubricVerdictDict +RubricVerdictDictOrDict = evals_types.RubricVerdictOrDict - filter_by_user_quota: Optional[bool] = Field( - default=True, - description="""Whether to filter recommendations to regions with user quota. +CandidateResult = evals_types.CandidateResult +CandidateResultDict = evals_types.CandidateResultDict +CandidateResultDictOrDict = evals_types.CandidateResultOrDict - 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. +Event = evals_types.Event +EventDict = evals_types.EventDict +EventDictOrDict = evals_types.EventOrDict - 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). - """, - ) +Message = evals_types.Message +MessageDict = evals_types.MessageDict +MessageDictOrDict = evals_types.MessageOrDict +Importance = evals_types.Importance -class ListCustomModelDeployOptionsConfigDict(TypedDict, total=False): - """Config for listing custom model deploy options.""" - filter_by_user_quota: Optional[bool] - """Whether to filter recommendations to regions with user quota. +class RuntimeDict(TypedDict, total=False): + """An agent runtime instance.""" - 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. - """ + api_client: Optional[Any] + """The underlying API client.""" - check_machine_availability: Optional[bool] - """Whether to check per-region machine availability. + api_async_client: Optional[Any] + """The underlying API client for asynchronous operations.""" - 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). - """ + api_resource: Optional[ReasoningEngineDict] + """The underlying API resource (i.e. ReasoningEngine).""" -ListCustomModelDeployOptionsConfigOrDict = Union[ - ListCustomModelDeployOptionsConfig, ListCustomModelDeployOptionsConfigDict -] +RuntimeOrDict = Union[Runtime, RuntimeDict] -class ExportOpenModelConfig(_common.BaseModel): - """Config for export_open_model.""" +class AgentRuntimeConfig(_common.BaseModel): + """Config for agent runtime methods.""" - 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - poll_interval_seconds: Optional[float] = Field( + staging_bucket: Optional[str] = Field( default=None, - description="""Seconds between LRO polls when ``wait_for_completion=True``. - Defaults to 30. Ignored when ``wait_for_completion=False``.""", + 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.""", ) - timeout_seconds: Optional[float] = Field( + requirements: Optional[Any] = 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``.""", - ) - - -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.""" - - poll_interval_seconds: Optional[float] - """Seconds between LRO polls when ``wait_for_completion=True``. - Defaults to 30. Ignored when ``wait_for_completion=False``.""" - - 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``.""" - - -ExportOpenModelConfigOrDict = Union[ExportOpenModelConfig, ExportOpenModelConfigDict] + 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.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""The user-defined name of the Agent Runtime. -class DeployPublisherModelConfig(_common.BaseModel): - """Config for ``deploy_publisher_model``. + 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. - Superset of options that apply to Google open, partner and Hugging Face - publisher models. Only fields relevant to the target model are honored; - the backend rejects unsupported fields with a clear error. - """ + 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. - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Whether to block on the deployment long-running operation. When - ``True`` (default), returns the ``DeployResponse`` (deployed endpoint - and model resource names) on completion. When ``False``, returns the - ``DeployModelOperation`` for the caller to poll.""", + 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.""", ) - poll_interval_seconds: Optional[float] = Field( + psc_interface_config: Optional[PscInterfaceConfig] = Field( default=None, - description="""Seconds between LRO polls when ``wait_for_completion=True``. - Defaults to 30. Ignored when ``wait_for_completion=False``.""", + description="""The PSC interface config for PSC-I to be used for the Agent Runtime.""", ) - timeout_seconds: Optional[float] = Field( + min_instances: Optional[int] = Field( default=None, - description="""Total wall-clock seconds to wait for the deployment to complete - when ``wait_for_completion=True``. Defaults to 2 hours, matching the - Vertex AI Console one-click deployment timeout. Ignored when - ``wait_for_completion=False``.""", + description="""The minimum number of instances to run for the Agent Runtime. + Defaults to 1. Range: [0, 10]. + """, ) - accept_eula: Optional[bool] = Field( + max_instances: Optional[int] = Field( default=None, - description="""Whether to accept the model's End User License Agreement.""", + 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]. + """, ) - hugging_face_access_token: Optional[str] = Field( + resource_limits: Optional[dict[str, str]] = Field( default=None, - description="""Hugging Face access token for gated HF models. See - https://huggingface.co/docs/hub/en/security-tokens.""", + 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'. + """, ) - machine_type: Optional[str] = Field( + container_concurrency: Optional[int] = Field( default=None, - description="""Machine type (e.g. ``'g2-standard-48'``). Leave unset for - automatic resources.""", + description="""The container concurrency to be used for the Agent Runtime. + Recommended value: 2 * cpu + 1. Defaults to 9. + """, ) - min_replica_count: Optional[int] = Field( - default=1, description="""Minimum number of replicas.""" + 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.""", ) - max_replica_count: Optional[int] = Field( - default=1, description="""Maximum number of replicas.""" + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""The encryption spec to be used for the Agent Runtime.""", ) - accelerator_type: Optional[str] = Field( - default=None, description="""Accelerator type (e.g. ``'NVIDIA_L4'``).""" + labels: Optional[dict[str, str]] = Field( + default=None, description="""The labels to be used for the Agent Runtime.""" ) - accelerator_count: Optional[int] = Field( - default=None, description="""Number of accelerators per replica.""" + agent_server_mode: Optional[AgentServerMode] = Field( + default=None, description="""The agent server mode to use for deployment.""" ) - spot: Optional[bool] = Field(default=None, description="""Schedule on Spot VMs.""") - dedicated_endpoint_disabled: Optional[bool] = Field( + class_methods: Optional[list[dict[str, Any]]] = Field( default=None, - description="""Set True to serve predictions via the shared endpoint DNS - instead of the dedicated endpoint DNS (default).""", + 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. + """, ) - fast_tryout_enabled: Optional[bool] = Field( + source_packages: Optional[list[str]] = Field( default=None, - description="""Use the fast-tryout deployment path (experimentation only, not - production). Only supported for select models and machine types.""", + 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) + """, ) - endpoint_display_name: Optional[str] = Field( - default=None, description="""Display name for the endpoint.""" + 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.""", ) - model_display_name: Optional[str] = Field( - default=None, description="""Display name for the deployed model.""" + entrypoint_module: Optional[str] = Field( + default=None, + description="""The entrypoint module to be used for the Agent Runtime + This field only used when source_packages is specified.""", ) - serving_container_image_uri: Optional[str] = Field( + entrypoint_object: Optional[str] = Field( default=None, - description="""Custom serving container image URI overriding the model's - default container.""", + description="""The entrypoint object to be used for the Agent Runtime. + This field only used when source_packages is specified.""", ) - container_command: Optional[list[str]] = Field( - default=None, description="""Serving container ENTRYPOINT override.""" + 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. + """, ) - container_args: Optional[list[str]] = Field( - default=None, description="""Serving container CMD override.""" + 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 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_variables: Optional[dict[str, str]] = Field( - default=None, description="""Environment variables for the serving container.""" + 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". + """, ) - enable_private_service_connect: Optional[bool] = Field( - default=None, description="""Enable Private Service Connect for the endpoint.""" + 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`. + """, ) - psc_project_allow_list: Optional[list[str]] = Field( + 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="""Projects allowed to access the endpoint over Private Service - Connect. Only honored when ``enable_private_service_connect`` is True.""", + description="""Agent Gateway configuration for a Agent Runtime deployment.""", + ) + build_config: Optional[ReasoningEngineSpecBuildSpec] = Field( + default=None, + description="""The build config for the Agent Runtime. Allows bringing your own Cloud Build private worker pool (BYOBP) and, optionally, a build-time service account for the container build. Supported keys: `worker_pool` (the resource name of the Cloud Build WorkerPool to use for the build) and `service_account` (the service account that Cloud Build uses to run the build; only applicable when `worker_pool` is specified).""", ) -class DeployPublisherModelConfigDict(TypedDict, total=False): - """Config for ``deploy_publisher_model``. +class AgentRuntimeConfigDict(TypedDict, total=False): + """Config for agent runtime methods.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + 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.""" + + 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.""" + + display_name: Optional[str] + """The user-defined name of the Agent Runtime. + + 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.""" + + gcs_dir_name: Optional[str] + """The GCS bucket directory under `staging_bucket` to use for staging + the artifacts needed.""" + + extra_packages: Optional[list[str]] + """The set of extra user-provided packages (if any).""" + + env_vars: Optional[Any] + """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] + """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] + """The identity type to use for the Agent Runtime.""" + + context_spec: Optional[ReasoningEngineContextSpecDict] + """The context spec to be used for the Agent Runtime.""" + + psc_interface_config: Optional[PscInterfaceConfigDict] + """The PSC interface config for PSC-I to be used for the Agent Runtime.""" + + 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]. + """ + + 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. + """ + + 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.""" + + labels: Optional[dict[str, str]] + """The labels to be used for the Agent Runtime.""" + + 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. + """ + + 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.""" - Superset of options that apply to Google open, partner and Hugging Face - publisher models. Only fields relevant to the target model are honored; - the backend rejects unsupported fields with a clear error. - """ + entrypoint_module: Optional[str] + """The entrypoint module to be used for the Agent Runtime + This field only used when source_packages is specified.""" - wait_for_completion: Optional[bool] - """Whether to block on the deployment long-running operation. When - ``True`` (default), returns the ``DeployResponse`` (deployed endpoint - and model resource names) on completion. When ``False``, returns the - ``DeployModelOperation`` for the caller to poll.""" + entrypoint_object: Optional[str] + """The entrypoint object to be used for the Agent Runtime. + This field only used when source_packages is specified.""" - poll_interval_seconds: Optional[float] - """Seconds between LRO polls when ``wait_for_completion=True``. - Defaults to 30. Ignored when ``wait_for_completion=False``.""" + 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. + """ - timeout_seconds: Optional[float] - """Total wall-clock seconds to wait for the deployment to complete - when ``wait_for_completion=True``. Defaults to 2 hours, matching the - Vertex AI Console one-click deployment timeout. Ignored when - ``wait_for_completion=False``.""" + 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".""" - accept_eula: Optional[bool] - """Whether to accept the model's End User License Agreement.""" + 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". + """ - hugging_face_access_token: Optional[str] - """Hugging Face access token for gated HF models. See - https://huggingface.co/docs/hub/en/security-tokens.""" + 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`. + """ - machine_type: Optional[str] - """Machine type (e.g. ``'g2-standard-48'``). Leave unset for - automatic resources.""" + image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict] + """The image spec for the Agent Runtime.""" - min_replica_count: Optional[int] - """Minimum number of replicas.""" + agent_config_source: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict + ] + """The agent config source for the Agent Runtime.""" - max_replica_count: Optional[int] - """Maximum number of replicas.""" + traffic_config: Optional[ReasoningEngineTrafficConfigDict] + """The traffic config for the Agent Runtime.""" - accelerator_type: Optional[str] - """Accelerator type (e.g. ``'NVIDIA_L4'``).""" + container_spec: Optional[ReasoningEngineSpecContainerSpecDict] + """The container spec for the Agent Runtime.""" - accelerator_count: Optional[int] - """Number of accelerators per replica.""" + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] + """Agent Gateway configuration for a Agent Runtime deployment.""" - spot: Optional[bool] - """Schedule on Spot VMs.""" + build_config: Optional[ReasoningEngineSpecBuildSpecDict] + """The build config for the Agent Runtime. Allows bringing your own Cloud Build private worker pool (BYOBP) and, optionally, a build-time service account for the container build. Supported keys: `worker_pool` (the resource name of the Cloud Build WorkerPool to use for the build) and `service_account` (the service account that Cloud Build uses to run the build; only applicable when `worker_pool` is specified).""" - dedicated_endpoint_disabled: Optional[bool] - """Set True to serve predictions via the shared endpoint DNS - instead of the dedicated endpoint DNS (default).""" - fast_tryout_enabled: Optional[bool] - """Use the fast-tryout deployment path (experimentation only, not - production). Only supported for select models and machine types.""" +AgentRuntimeConfigOrDict = Union[AgentRuntimeConfig, AgentRuntimeConfigDict] - endpoint_display_name: Optional[str] - """Display name for the endpoint.""" - model_display_name: Optional[str] - """Display name for the deployed model.""" +class RunQueryJobRuntimeConfig(_common.BaseModel): + """Config for checking a query job on an agent runtime.""" - serving_container_image_uri: Optional[str] - """Custom serving container image URI overriding the model's - default container.""" + 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 runtime.""" + ) + output_gcs_uri: Optional[str] = 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.""", + ) - container_command: Optional[list[str]] - """Serving container ENTRYPOINT override.""" - container_args: Optional[list[str]] - """Serving container CMD override.""" +class RunQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for checking a query job on an agent runtime.""" - container_variables: Optional[dict[str, str]] - """Environment variables for the serving container.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - enable_private_service_connect: Optional[bool] - """Enable Private Service Connect for the endpoint.""" + query: Optional[str] + """The query to send to the agent runtime.""" - psc_project_allow_list: Optional[list[str]] - """Projects allowed to access the endpoint over Private Service - Connect. Only honored when ``enable_private_service_connect`` is True.""" + 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.""" -DeployPublisherModelConfigOrDict = Union[ - DeployPublisherModelConfig, DeployPublisherModelConfigDict +RunQueryJobRuntimeConfigOrDict = Union[ + RunQueryJobRuntimeConfig, RunQueryJobRuntimeConfigDict ] -class DeployOption(_common.BaseModel): - """A verified deploy option for a model.""" +class RunQueryJobResult(_common.BaseModel): + """Result of running a query job.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - machine_type: Optional[str] = Field( - default=None, description="""The machine type.""" + job_name: Optional[str] = Field( + default=None, + description="""Name of the agent runtime operation to later check for status.""", ) - accelerator_type: Optional[str] = Field( - default=None, description="""The accelerator type.""" + input_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the input file.""" ) - accelerator_count: Optional[int] = Field( - default=None, description="""The number of accelerators.""" + output_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the output file.""" ) -class DeployOptionDict(TypedDict, total=False): - """A verified deploy option for a model.""" +class RunQueryJobResultDict(TypedDict, total=False): + """Result of running a query job.""" - option_name: Optional[str] - """The name of the deploy task.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - serving_container_image_uri: Optional[str] - """The URI of the serving container.""" + job_name: Optional[str] + """Name of the agent runtime operation to later check for status.""" - machine_type: Optional[str] - """The machine type.""" + input_gcs_uri: Optional[str] + """The GCS URI of the input file.""" - accelerator_type: Optional[str] - """The accelerator type.""" + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" - accelerator_count: Optional[int] - """The number of accelerators.""" +RunQueryJobResultOrDict = Union[RunQueryJobResult, RunQueryJobResultDict] -DeployOptionOrDict = Union[DeployOption, DeployOptionDict] +class CheckQueryJobResponse(_common.BaseModel): + """Response from LRO.""" -class ListMemoryBanksResponse(_common.BaseModel): - """The response for listing Memory Banks.""" + 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.""" + ) - memory_banks: Optional[list[MemoryBank]] = Field( - default=None, description="""The list of Memory Banks.""" + +class CheckQueryJobResponseDict(TypedDict, total=False): + """Response from LRO.""" + + 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 RuntimeRevision(_common.BaseModel): + """An agent runtime revision instance.""" + + 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).""", ) + # Allows dynamic binding of methods based on the registered operations. + model_config = ConfigDict(extra="allow") -class ListMemoryBanksResponseDict(TypedDict, total=False): - """The response for listing Memory Banks.""" + 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)" + ) - memory_banks: Optional[list[MemoryBankDict]] - """The list of Memory Banks.""" + 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 + def delete( + self, + config: Optional[DeleteRuntimeRevisionConfigOrDict] = None, + ) -> None: + """Deletes the agent engine runtime revision. -ListMemoryBanksResponseOrDict = Union[ - ListMemoryBanksResponse, ListMemoryBanksResponseDict -] + 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] + + +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.""" + + 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 638cba8645..dfcd4821cb 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 {}, @@ -319,14 +316,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 97% rename from agentplatform/agent_engines/templates/adk.py rename to agentplatform/frameworks/adk.py index 9bdf762756..45d755453a 100644 --- a/agentplatform/agent_engines/templates/adk.py +++ b/agentplatform/frameworks/adk.py @@ -260,13 +260,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: @@ -370,9 +370,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: @@ -485,9 +485,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. @@ -790,7 +790,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"): @@ -817,8 +816,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, @@ -829,9 +826,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: @@ -993,19 +988,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 @@ -1046,7 +1038,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" @@ -1057,11 +1049,13 @@ def set_up(self): ) if custom_instrumentor and self._tracing_enabled(): - self._tmpl_attrs["instrumentor"] = custom_instrumentor(self.project_id()) + self._tmpl_attrs["instrumentor"] = custom_instrumentor( + self._get_project_id(project) + ) if not custom_instrumentor: self._tmpl_attrs["instrumentor"] = _default_instrumentor_builder( - self.project_id(), + self._get_project_id(project), enable_tracing=self._tracing_enabled(), enable_logging=enable_logging, ) @@ -1093,7 +1087,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): @@ -1105,7 +1099,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"), ) @@ -1127,7 +1121,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): @@ -1139,7 +1133,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: @@ -1231,7 +1225,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): @@ -1291,7 +1285,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( @@ -1341,7 +1335,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): @@ -1368,7 +1362,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, @@ -1376,7 +1370,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. @@ -2120,8 +2114,7 @@ def _tracing_enabled(self) -> bool: and is_version_sufficient("1.17.0") ) - def project_id(self) -> Optional[str]: - project = self._tmpl_attrs.get("project") + def _get_project_id(self, project: str) -> Optional[str]: if project and str(project).isdigit(): try: from google.cloud.aiplatform.utils import ( diff --git a/agentplatform/agent_engines/templates/ag2.py b/agentplatform/frameworks/ag2.py similarity index 90% rename from agentplatform/agent_engines/templates/ag2.py rename to agentplatform/frameworks/ag2.py index c6d7f8c3a0..e266b2b560 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,13 +91,11 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils - 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() - ) + 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( ( openinference_autogen, @@ -158,7 +158,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(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -351,11 +351,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, @@ -367,45 +363,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. @@ -413,10 +415,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)) @@ -425,22 +427,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"), @@ -502,6 +503,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 96% rename from agentplatform/agent_engines/templates/langchain.py rename to agentplatform/frameworks/langchain.py index bf8947b230..f3cc92bc9c 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,15 +188,11 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils - 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() - ) + 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( ( openinference_langchain, @@ -263,7 +255,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(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -555,11 +547,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, @@ -600,20 +588,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 3afc19c499..8a8dfd746e 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,15 +164,11 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils - 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() - ) + 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( ( openinference_langchain, @@ -240,7 +232,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(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -481,11 +473,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, @@ -518,20 +506,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 94% rename from agentplatform/agent_engines/templates/llama_index.py rename to agentplatform/frameworks/llama_index.py index 1d0c8c12a4..7e15f99b13 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,17 +374,24 @@ 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, ) 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() ) try: import opentelemetry.exporter.otlp.proto.http.trace_exporter @@ -423,7 +421,7 @@ def set_up(self): ): credentials, _ = google.auth.default() resource = opentelemetry.sdk.resources.Resource.create( - attributes={"gcp.project_id": self._project}, + attributes={"gcp.project_id": project}, ).merge(opentelemetry.sdk.resources.OTELResourceDetector().detect()) span_exporter = _otlp_span_exporter_module.OTLPSpanExporter( session=google.auth.transport.requests.AuthorizedSession( @@ -463,9 +461,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( resource=resource ) @@ -500,8 +496,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: @@ -564,8 +560,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): @@ -575,9 +571,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/noxfile.py b/noxfile.py index b43a081d26..6ee124ecea 100644 --- a/noxfile.py +++ b/noxfile.py @@ -438,7 +438,9 @@ def unit_agentplatform_a2a(session): constraints_path = str(CURRENT_DIRECTORY / "testing" / "constraints-a2a.txt") install_unittest_dependencies(session, "-c", constraints_path) - session.install("a2a-sdk", "-c", constraints_path) + # The `http-server` extra provides `starlette` and `sse-starlette`, which + # `a2a.server.routes` imports and the agent card route tests exercise. + session.install("a2a-sdk[http-server]", "-c", constraints_path) # Run py.test against the unit tests. session.run( @@ -468,6 +470,9 @@ def unit_a2a(session): constraints_path = str(CURRENT_DIRECTORY / "testing" / "constraints-a2a.txt") install_unittest_dependencies(session, "-c", constraints_path) session.install("a2a-sdk", "-c", constraints_path) + # The `http-server` extra provides `starlette` and `sse-starlette`, which + # `a2a.server.routes` imports and the agent card route tests exercise. + session.install("a2a-sdk[http-server]", "-c", constraints_path) # Run py.test against the unit tests. session.run( diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py b/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py index dc1b5990ce..2baaa49a02 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,143 +70,71 @@ 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], - ) - - 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 - - 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, + mock_await_operation.return_value = _genai_types.RuntimeOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + ) ) - - 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") - - with pytest.raises( - TypeError, - match="Unsupported AgentCard type", - ): - 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 = _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, + ), ) + 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) + class TestA2aPublicCardRoutes: """Verifies the public agent card is served at the Agent Runtime card paths.""" def _build_agent(self): - """Builds an A2aAgent from the agentplatform template and runs set_up().""" + """Builds an A2aAgent from the agentplatform framework and runs set_up().""" # pylint: disable=g-import-not-at-top from a2a import types as a2a_types - from agentplatform.agent_engines.templates import ( - a2a as a2a_template, - ) # pylint: enable=g-import-not-at-top card = a2a_types.AgentCard( @@ -311,7 +161,7 @@ def _build_agent(self): default_input_modes=["text/plain"], default_output_modes=["text/plain"], ) - agent = a2a_template.A2aAgent(agent_card=card) + agent = A2aAgent(agent_card=card) with mock.patch.dict( os.environ, { diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py index 38b38fa12a..11932c7c1b 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 errors as genai_errors @@ -417,12 +417,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", @@ -433,8 +432,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( @@ -623,9 +620,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() @@ -1339,8 +1334,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( @@ -1366,9 +1361,7 @@ def test_project_id_skips_lookup_for_project_id( ): """Project IDs are returned as-is without a GetProject call.""" app = adk_template.AdkApp(agent=_TEST_AGENT) - app._tmpl_attrs["project"] = _TEST_PROJECT - - assert app.project_id() == _TEST_PROJECT + assert app._get_project_id(_TEST_PROJECT) == _TEST_PROJECT get_project_id_mock.assert_not_called() def test_project_id_resolves_project_number( @@ -1377,9 +1370,7 @@ def test_project_id_resolves_project_number( ): """Project numbers are resolved to project IDs via GetProject.""" app = adk_template.AdkApp(agent=_TEST_AGENT) - app._tmpl_attrs["project"] = _TEST_PROJECT_NUMBER - - assert app.project_id() == _TEST_PROJECT_ID + assert app._get_project_id(_TEST_PROJECT_NUMBER) == _TEST_PROJECT_ID get_project_id_mock.assert_called_once_with(_TEST_PROJECT_NUMBER) def test_project_id_without_project( @@ -1387,9 +1378,7 @@ def test_project_id_without_project( get_project_id_mock: mock.Mock, ): app = adk_template.AdkApp(agent=_TEST_AGENT) - app._tmpl_attrs["project"] = None - - assert app.project_id() is None + assert app._get_project_id(None) is None get_project_id_mock.assert_not_called() @pytest.mark.usefixtures("caplog") @@ -1479,7 +1468,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 @@ -1591,27 +1580,27 @@ async def _raise_server_error(*args, **kwargs): @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") @@ -1655,12 +1644,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( @@ -1669,9 +1658,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) @@ -1708,8 +1697,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, @@ -1724,13 +1713,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}, ) @@ -1763,8 +1752,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, @@ -1779,13 +1768,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 1b5fbd2eb3..85ff6783f3 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 = {} @@ -149,7 +149,7 @@ def otel_resource_detector_mock(): @pytest.fixture def is_noop_or_proxy_tracer_provider_mock(): with mock.patch.object( - _agent_engines_utils, "is_noop_or_proxy_tracer_provider" + _runtimes_utils, "is_noop_or_proxy_tracer_provider" ) as is_noop_or_proxy_tracer_provider_mock: is_noop_or_proxy_tracer_provider_mock.return_value = True yield is_noop_or_proxy_tracer_provider_mock @@ -158,7 +158,7 @@ def is_noop_or_proxy_tracer_provider_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 @@ -167,7 +167,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 @@ -177,7 +177,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() @@ -201,24 +201,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): @@ -226,7 +226,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, @@ -241,7 +241,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, @@ -251,7 +251,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, @@ -265,7 +265,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, ) @@ -292,7 +292,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, @@ -305,7 +305,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, @@ -325,7 +325,7 @@ def test_tracing_setup( is_noop_or_proxy_tracer_provider_mock, autogen_instrumentor_mock, ): - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, runnable_builder=lambda **kwargs: kwargs, @@ -371,7 +371,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], @@ -379,23 +379,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", @@ -419,13 +419,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 = [] @@ -439,7 +439,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 @@ -454,7 +454,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): @@ -463,4 +463,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 2ac9fe375c..4aff5906ab 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 @@ -135,7 +135,7 @@ def otel_resource_detector_mock(): @pytest.fixture def is_noop_or_proxy_tracer_provider_mock(): with mock.patch.object( - _agent_engines_utils, "is_noop_or_proxy_tracer_provider" + _runtimes_utils, "is_noop_or_proxy_tracer_provider" ) as is_noop_or_proxy_tracer_provider_mock: is_noop_or_proxy_tracer_provider_mock.return_value = True yield is_noop_or_proxy_tracer_provider_mock @@ -144,7 +144,7 @@ def is_noop_or_proxy_tracer_provider_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 @@ -153,7 +153,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 @@ -163,12 +163,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": ( @@ -183,13 +180,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): @@ -197,7 +197,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, @@ -211,7 +211,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, @@ -223,7 +223,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, @@ -239,7 +239,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, @@ -253,7 +253,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")) @@ -270,7 +270,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, @@ -287,7 +287,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, @@ -308,7 +308,7 @@ def test_tracing_setup( is_noop_or_proxy_tracer_provider_mock, langchain_instrumentor_mock, ): - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, prompt=self.prompt, output_parser=self.output_parser, @@ -355,7 +355,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], ) @@ -369,7 +369,7 @@ def test_raise_both_system_instruction_and_prompt_error( 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 70948abb63..1d62997858 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 @@ -146,7 +146,7 @@ def otel_resource_detector_mock(): @pytest.fixture def is_noop_or_proxy_tracer_provider_mock(): with mock.patch.object( - _agent_engines_utils, "is_noop_or_proxy_tracer_provider" + _runtimes_utils, "is_noop_or_proxy_tracer_provider" ) as is_noop_or_proxy_tracer_provider_mock: is_noop_or_proxy_tracer_provider_mock.return_value = True yield is_noop_or_proxy_tracer_provider_mock @@ -155,7 +155,7 @@ def is_noop_or_proxy_tracer_provider_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 @@ -164,7 +164,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 @@ -174,21 +174,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): @@ -196,7 +196,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, @@ -209,7 +209,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, @@ -219,7 +219,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, @@ -233,7 +233,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") @@ -248,7 +248,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")) @@ -268,7 +268,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() @@ -280,7 +280,7 @@ 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() @@ -296,7 +296,7 @@ def test_tracing_setup( is_noop_or_proxy_tracer_provider_mock, langchain_instrumentor_mock, ): - agent = agent_engines.LanggraphAgent( + agent = frameworks.LanggraphAgent( model=_TEST_MODEL, runnable_builder=lambda **kwargs: kwargs, enable_tracing=True, @@ -333,14 +333,14 @@ def test_tracing_setup( ) 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(), @@ -359,7 +359,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(), @@ -378,7 +378,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 = { @@ -388,7 +388,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 = { @@ -398,13 +398,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( @@ -412,7 +412,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( @@ -420,7 +420,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"], @@ -436,6 +436,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 670fc97bf0..8bd2ccd63a 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py @@ -14,17 +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 @@ -118,7 +116,7 @@ def otel_resource_detector_mock(): @pytest.fixture def is_noop_or_proxy_tracer_provider_mock(): with mock.patch.object( - _agent_engines_utils, "is_noop_or_proxy_tracer_provider" + _runtimes_utils, "is_noop_or_proxy_tracer_provider" ) as is_noop_or_proxy_tracer_provider_mock: is_noop_or_proxy_tracer_provider_mock.return_value = True yield is_noop_or_proxy_tracer_provider_mock @@ -127,7 +125,7 @@ def is_noop_or_proxy_tracer_provider_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 @@ -136,7 +134,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 @@ -146,7 +144,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 @@ -155,12 +153,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( @@ -175,13 +170,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): @@ -341,17 +339,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"}', ), ] @@ -362,69 +360,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_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 a2bb020b47..d234643a74 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,9 +19,9 @@ def test_private_update_session(client): - session = client.agent_engines.sessions._update( + session = 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", ), @@ -32,5 +32,5 @@ def test_private_update_session(client): 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 67ed6a6595..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.A2aTaskState.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.A2aTaskState.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_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_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_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 c715da1087..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.A2aTaskState.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.A2aTaskState.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_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_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_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_run_inference.py b/tests/unit/agentplatform/genai/replays/test_run_inference.py index 0c4afb9eaf..ba98c52be6 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) 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_agent_engine_session.py b/tests/unit/agentplatform/genai/replays/test_update_agent_engine_session.py index 51a02b8860..db4f455147 100644 --- a/tests/unit/agentplatform/genai/replays/test_update_agent_engine_session.py +++ b/tests/unit/agentplatform/genai/replays/test_update_agent_engine_session.py @@ -20,23 +20,23 @@ def test_update_session(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) - session_operation = client.agent_engines.sessions.create( - name=agent_engine.api_resource.name, + session_operation = client.sessions.create( + name=runtime.api_resource.name, user_id="test-user-123", - config=types.CreateAgentEngineSessionConfig( + config=types.CreateRuntimeSessionConfig( display_name="initial_session", ), ) - assert isinstance(session_operation, types.AgentEngineSessionOperation) + assert isinstance(session_operation, types.RuntimeSessionOperation) - updated_session = client.agent_engines.sessions.update( + updated_session = client.sessions.update( name=session_operation.response.name, - config=types.UpdateAgentEngineSessionConfig( + config=types.UpdateRuntimeSessionConfig( display_name="updated_session", user_id="test-user-123", labels={"env": "test", "tier": "dev"}, @@ -48,9 +48,9 @@ def test_update_session(client): assert updated_session.labels == {"env": "test", "tier": "dev"} # Second update: update with explicit update_mask - mask_updated_session = client.agent_engines.sessions.update( + mask_updated_session = client.sessions.update( name=session_operation.response.name, - config=types.UpdateAgentEngineSessionConfig( + config=types.UpdateRuntimeSessionConfig( display_name="session_with_mask", user_id="test-user-123", update_mask="displayName", @@ -60,22 +60,22 @@ def test_update_session(client): assert mask_updated_session.display_name == "session_with_mask" # Third update: update with ttl (duration) - ttl_updated_session = client.agent_engines.sessions.update( + ttl_updated_session = client.sessions.update( name=session_operation.response.name, - config=types.UpdateAgentEngineSessionConfig( + config=types.UpdateRuntimeSessionConfig( user_id="test-user-123", ttl="86400s", ), ) assert isinstance(ttl_updated_session, types.Session) 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.sessions.update", + test_method="sessions.update", ) pytest_plugins = ("pytest_asyncio",) @@ -83,23 +83,23 @@ def test_update_session(client): @pytest.mark.asyncio async def test_update_session_async(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) - session_operation = await client.aio.agent_engines.sessions.create( - name=agent_engine.api_resource.name, + session_operation = await client.aio.sessions.create( + name=runtime.api_resource.name, user_id="test-user-123", - config=types.CreateAgentEngineSessionConfig( + config=types.CreateRuntimeSessionConfig( display_name="initial_session", ), ) - assert isinstance(session_operation, types.AgentEngineSessionOperation) + assert isinstance(session_operation, types.RuntimeSessionOperation) - updated_session = await client.aio.agent_engines.sessions.update( + updated_session = await client.aio.sessions.update( name=session_operation.response.name, - config=types.UpdateAgentEngineSessionConfig( + config=types.UpdateRuntimeSessionConfig( display_name="updated_session", user_id="test-user-123", labels={"env": "test", "tier": "dev"}, @@ -111,4 +111,4 @@ async def test_update_session_async(client): assert updated_session.user_id == "test-user-123" assert updated_session.labels == {"env": "test", "tier": "dev"} 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_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 3686ae43bb..7e84a80ff2 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 @@ -312,7 +312,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, @@ -330,7 +330,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]} @@ -348,30 +348,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" @@ -461,14 +461,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", @@ -607,38 +607,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, ) @@ -646,7 +646,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, ) @@ -654,7 +654,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, ) @@ -662,7 +662,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, ) @@ -692,7 +692,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, ) @@ -866,7 +866,7 @@ class FakeObject: @pytest.mark.usefixtures("google_auth_mock") -class TestAgentEngineHelpers: +class TestRuntimeHelpers: def setup_method(self): importlib.reload(initializer) importlib.reload(aiplatform) @@ -883,9 +883,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, @@ -896,7 +896,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", [ @@ -919,7 +919,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], @@ -929,7 +929,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, @@ -944,11 +944,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", [ @@ -971,7 +971,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, @@ -982,7 +982,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, @@ -999,9 +999,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, @@ -1063,9 +1063,9 @@ def test_create_agent_engine_config_full(self, mock_prepare): == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT ) - @mock.patch.object(_agent_engines_utils, "_prepare") - def test_create_agent_engine_config_with_build_config(self, mock_prepare): - config = self.client.agent_engines._create_config( + @mock.patch.object(_runtimes_utils, "_prepare") + def test_create_runtime_config_with_build_config(self, mock_prepare): + config = self.client.runtimes._create_config( mode="create", agent=self.test_agent, staging_bucket=_TEST_STAGING_BUCKET, @@ -1078,11 +1078,11 @@ def test_create_agent_engine_config_with_build_config(self, mock_prepare): "service_account": _TEST_AGENT_ENGINE_BUILD_SERVICE_ACCOUNT, } - @mock.patch.object(_agent_engines_utils, "_prepare") - def test_create_agent_engine_config_with_build_config_worker_pool_only( + @mock.patch.object(_runtimes_utils, "_prepare") + def test_create_runtime_config_with_build_config_worker_pool_only( self, mock_prepare ): - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", agent=self.test_agent, staging_bucket=_TEST_STAGING_BUCKET, @@ -1094,9 +1094,9 @@ def test_create_agent_engine_config_with_build_config_worker_pool_only( "worker_pool": _TEST_AGENT_ENGINE_BUILD_WORKER_POOL, } - @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_prepare") def test_update_agent_engine_config_with_build_config(self, mock_prepare): - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="update", build_config=_TEST_AGENT_ENGINE_BUILD_CONFIG, ) @@ -1109,11 +1109,11 @@ def test_update_agent_engine_config_with_build_config(self, mock_prepare): assert "spec.build_spec.service_account" in update_mask @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: @@ -1124,7 +1124,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, @@ -1159,7 +1159,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: @@ -1169,7 +1169,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, @@ -1202,18 +1202,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, @@ -1223,7 +1223,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: @@ -1231,7 +1231,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, @@ -1251,12 +1251,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, @@ -1281,11 +1281,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: @@ -1297,7 +1297,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, @@ -1314,11 +1314,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: @@ -1330,7 +1330,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, @@ -1343,8 +1343,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, @@ -1368,11 +1368,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: @@ -1383,7 +1383,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, @@ -1413,9 +1413,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, @@ -1433,11 +1433,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, @@ -1459,10 +1459,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, @@ -1472,7 +1472,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, @@ -1482,12 +1482,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: @@ -1500,7 +1500,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", @@ -1513,15 +1513,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, @@ -1535,9 +1535,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, @@ -1600,8 +1600,8 @@ def test_update_agent_engine_config_full(self, mock_prepare): ] ) - @mock.patch.object(_agent_engines_utils, "_prepare") - def test_update_agent_engine_config_with_agent_card(self, mock_prepare): + @mock.patch.object(_runtimes_utils, "_prepare") + def test_update_runtime_config_with_agent_card(self, mock_prepare): from google.protobuf import struct_pb2 from google.protobuf import json_format @@ -1609,7 +1609,7 @@ def test_update_agent_engine_config_with_agent_card(self, mock_prepare): card["version"] = "1.3.0" agent = CapitalizeEngineWithAgentCard(agent_card=card) - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="update", agent=agent, staging_bucket=_TEST_STAGING_BUCKET, @@ -1629,11 +1629,11 @@ def test_update_agent_engine_config_with_agent_card(self, mock_prepare): "spec.agent_framework" ) - @mock.patch.object(_agent_engines_utils, "_prepare") - def test_update_agent_engine_config_without_agent_card_omits_mask( + @mock.patch.object(_runtimes_utils, "_prepare") + def test_update_runtime_config_without_agent_card_omits_mask( self, mock_prepare ): - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="update", agent=self.test_agent, staging_bucket=_TEST_STAGING_BUCKET, @@ -1643,9 +1643,9 @@ def test_update_agent_engine_config_without_agent_card_omits_mask( assert "agent_card" not in config["spec"] assert "spec.agent_card" not in config["update_mask"].split(",") - @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, @@ -1664,7 +1664,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( @@ -1675,7 +1675,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( @@ -1684,13 +1684,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( @@ -1701,9 +1701,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", @@ -1711,12 +1711,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=[ @@ -1732,7 +1732,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 @@ -1743,7 +1743,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 == { @@ -1755,7 +1755,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, ) @@ -1771,7 +1771,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, ) @@ -1785,7 +1785,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, ) @@ -1800,13 +1800,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, ) @@ -1821,7 +1821,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, @@ -1870,7 +1870,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. @@ -1907,7 +1907,7 @@ 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 # pytest does not allow absl.testing.parameterized.named_parameters. @@ -1948,7 +1948,7 @@ def test_yield_parsed_json_from_httpbody(self, obj, expected): ) def test_to_parsed_json_server_sent_events(self, obj, expected): """An SSE-framed response is parsed into the same objects as NDJSON.""" - assert list(_agent_engines_utils._yield_parsed_json(obj)) == expected + assert list(_runtimes_utils._yield_parsed_json(obj)) == expected def test_yield_parsed_json_from_httpbody_event_stream(self): """The gRPC path parses SSE instead of yielding the raw proto.""" @@ -1956,14 +1956,14 @@ def test_yield_parsed_json_from_httpbody_event_stream(self): content_type="text/event-stream", data=b'data: {"a": 1}\n\ndata: {"a": 2}\n\n', ) - assert list(_agent_engines_utils._yield_parsed_json_from_httpbody(body)) == [ + assert list(_runtimes_utils._yield_parsed_json_from_httpbody(body)) == [ {"a": 1}, {"a": 2}, ] 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 ] @@ -1976,7 +1976,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: @@ -1999,30 +1999,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, @@ -2032,31 +2032,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, @@ -2071,7 +2071,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, @@ -2148,13 +2148,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 @@ -2162,7 +2162,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) @@ -2179,12 +2179,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, @@ -2192,13 +2192,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)}", @@ -2207,21 +2207,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, @@ -2229,12 +2229,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, @@ -2262,41 +2262,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, ) @@ -2311,14 +2311,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, @@ -2337,19 +2337,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], @@ -2415,14 +2415,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, @@ -2443,19 +2443,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], @@ -2522,14 +2522,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, @@ -2550,19 +2550,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], @@ -2629,23 +2629,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, @@ -2660,11 +2660,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], @@ -2700,14 +2700,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, @@ -2725,19 +2725,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], @@ -2800,14 +2800,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, @@ -2826,19 +2826,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], @@ -2903,12 +2903,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, @@ -2916,13 +2916,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, ), @@ -2954,31 +2954,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], @@ -3014,38 +3014,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, @@ -3084,25 +3084,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={}, ), @@ -3137,11 +3137,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, @@ -3149,12 +3149,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, @@ -3187,21 +3187,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, ), ) @@ -3216,21 +3216,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, ), ) @@ -3245,12 +3245,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, @@ -3258,12 +3258,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, ) @@ -3289,18 +3289,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( @@ -3324,11 +3324,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"}' @@ -3355,7 +3355,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, @@ -3385,32 +3385,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"}' @@ -3440,7 +3440,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, @@ -3449,13 +3449,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"}' @@ -3476,7 +3476,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, @@ -3494,13 +3494,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"}' @@ -3523,7 +3523,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, @@ -3556,15 +3556,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, @@ -3577,7 +3577,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)) @@ -3592,13 +3592,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}, ) @@ -3614,9 +3614,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={}, @@ -3634,7 +3634,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}, ) @@ -3651,9 +3651,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={}, @@ -3664,7 +3664,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}, ) @@ -3676,16 +3676,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}, ) @@ -3697,9 +3697,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={}, @@ -3710,7 +3710,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}, ) @@ -3722,9 +3722,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={}, @@ -3746,7 +3746,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}, ) @@ -3766,17 +3766,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( @@ -3814,7 +3814,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).""" @@ -3827,12 +3827,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( @@ -3869,7 +3869,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, ), @@ -3882,70 +3882,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, ), @@ -3958,7 +3958,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, ), @@ -3969,7 +3969,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, @@ -3978,7 +3978,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( @@ -3990,17 +3990,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, @@ -4010,15 +4010,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, ), @@ -4032,16 +4032,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, @@ -4052,7 +4052,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, @@ -4064,9 +4064,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, ), @@ -4080,15 +4080,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, @@ -4098,7 +4098,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, @@ -4107,9 +4107,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, ), @@ -4125,7 +4125,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) @@ -4137,27 +4137,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, @@ -4169,19 +4169,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, ), ) @@ -4239,7 +4239,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, @@ -4252,7 +4252,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( @@ -4319,12 +4319,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) @@ -4341,13 +4341,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 ) ) @@ -4358,13 +4358,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 1e7670bd2f..121def9cce 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 c162ca23b8..23eba767db 100644 --- a/tests/unit/agentplatform/genai/test_sandbox.py +++ b/tests/unit/agentplatform/genai/test_sandbox.py @@ -95,7 +95,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, @@ -131,7 +131,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, @@ -148,7 +148,7 @@ def test_create_with_shell_environment_and_existing_template(self, mock_create): mock_operation = mock.Mock() mock_create.return_value = mock_operation - result = self.client.agent_engines.sandboxes.create( + result = self.client.sandboxes.create( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec={"shell_environment": {}}, config={ @@ -176,7 +176,7 @@ def test_create_with_shell_environment_creates_template_when_absent( mock_template_create.return_value = mock_template_operation mock_create.return_value = mock.Mock() - self.client.agent_engines.sandboxes.create( + self.client.sandboxes.create( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec={"shell_environment": {}}, config={"wait_for_completion": False}, @@ -207,7 +207,7 @@ def test_create_with_typed_shell_environment_creates_template_when_absent( mock_create.return_value = mock.Mock() shell_environment = agentplatform_types.SandboxEnvironmentSpecShellEnvironment() - self.client.agent_engines.sandboxes.create( + self.client.sandboxes.create( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=agentplatform_types.SandboxEnvironmentSpec( shell_environment=shell_environment, @@ -239,7 +239,7 @@ def test_create_with_computer_use_environment_creates_template_when_absent( mock_template_create.return_value = mock_template_operation mock_create.return_value = mock.Mock() - self.client.agent_engines.sandboxes.create( + self.client.sandboxes.create( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec={"computer_use_environment": {}}, config={"wait_for_completion": False}, @@ -267,7 +267,7 @@ def test_create_with_snapshot_does_not_create_template( mock_operation = mock.Mock() mock_create.return_value = mock_operation - result = self.client.agent_engines.sandboxes.create( + result = self.client.sandboxes.create( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec={"computer_use_environment": {}}, config={ @@ -289,7 +289,7 @@ def test_create_with_snapshot_does_not_create_template( def test_create_without_spec_template_or_snapshot_raises(self, mock_create): for spec in (None, {}, agentplatform_types.SandboxEnvironmentSpec()): with pytest.raises(ValueError, match="must be provided"): - self.client.agent_engines.sandboxes.create( + self.client.sandboxes.create( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=spec, )