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 9aca69446c..815f309150 100644 --- a/agentplatform/_genai/client.py +++ b/agentplatform/_genai/client.py @@ -30,7 +30,7 @@ if TYPE_CHECKING: from agentplatform._genai import ( - agent_engines as agent_engines_module, + runtimes as runtimes_module, ) from agentplatform._genai import datasets as datasets_module from agentplatform._genai import evals as evals_module @@ -47,6 +47,11 @@ from agentplatform._genai import ( feedback_entries as feedback_entries_module, ) + from agentplatform._genai import sessions as sessions_module + from agentplatform._genai import ( + sandboxes as sandboxes_module, + ) + from agentplatform._genai import ( endpoints as endpoints_module, ) @@ -93,7 +98,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 @@ -103,6 +108,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 @@ -137,22 +144,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": @@ -343,7 +386,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 @@ -352,6 +395,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 @@ -410,22 +455,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..80f15f7e6e 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,1460 @@ logger.setLevel(logging.INFO) +def _CancelQueryJobRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["operation_name"]) is not None: + setv(parent_object, ["operationName"], getv(from_object, ["operation_name"])) + + return to_object + + +def _CancelQueryJobRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + setv( + to_object, + ["config"], + _CancelQueryJobRuntimeConfig_to_vertex( + getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _CheckQueryJobResult_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(parent_object, ["operationName"]) is not None: + setv(to_object, ["operation_name"], getv(parent_object, ["operationName"])) + + if getv(parent_object, ["outputGcsUri"]) is not None: + setv(to_object, ["output_gcs_uri"], getv(parent_object, ["outputGcsUri"])) + + if getv(parent_object, ["status"]) is not None: + setv(to_object, ["status"], getv(parent_object, ["status"])) + + if getv(parent_object, ["result"]) is not None: + setv(to_object, ["result"], getv(parent_object, ["result"])) + + return to_object + + +def _CheckQueryJobRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["retrieve_result"]) is not None: + setv(parent_object, ["retrieveResult"], getv(from_object, ["retrieve_result"])) + + return to_object + + +def _CheckQueryJobRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + setv( + to_object, + ["config"], + _CheckQueryJobRuntimeConfig_to_vertex( + getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _CreateRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["display_name"]) is not None: + setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) + + if getv(from_object, ["description"]) is not None: + setv(parent_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["spec"]) is not None: + setv(parent_object, ["spec"], getv(from_object, ["spec"])) + + if getv(from_object, ["context_spec"]) is not None: + setv( + parent_object, + ["contextSpec"], + _ReasoningEngineContextSpec_to_vertex( + getv(from_object, ["context_spec"]), to_object + ), + ) + + if getv(from_object, ["psc_interface_config"]) is not None: + setv( + parent_object, + ["pscInterfaceConfig"], + getv(from_object, ["psc_interface_config"]), + ) + + if getv(from_object, ["agent_gateway_config"]) is not None: + setv( + parent_object, + ["agentGatewayConfig"], + getv(from_object, ["agent_gateway_config"]), + ) + + if getv(from_object, ["encryption_spec"]) is not None: + setv(parent_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"])) + + if getv(from_object, ["labels"]) is not None: + setv(parent_object, ["labels"], getv(from_object, ["labels"])) + + if getv(from_object, ["source_packages"]) is not None: + setv(parent_object, ["sourcePackages"], getv(from_object, ["source_packages"])) + + if getv(from_object, ["entrypoint_module"]) is not None: + setv( + parent_object, + ["entrypointModule"], + getv(from_object, ["entrypoint_module"]), + ) + + if getv(from_object, ["entrypoint_object"]) is not None: + setv( + parent_object, + ["entrypointObject"], + getv(from_object, ["entrypoint_object"]), + ) + + if getv(from_object, ["requirements_file"]) is not None: + setv( + parent_object, + ["requirementsFile"], + getv(from_object, ["requirements_file"]), + ) + + if getv(from_object, ["agent_framework"]) is not None: + setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) + + if getv(from_object, ["python_version"]) is not None: + setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) + + return to_object + + +def _CreateRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["config"]) is not None: + _CreateRuntimeConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _DeleteRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["force"]) is not None: + setv(to_object, ["force"], getv(from_object, ["force"])) + + return to_object + + +def _GetRuntimeOperationParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["operation_name"]) is not None: + setv( + to_object, ["_url", "operationName"], getv(from_object, ["operation_name"]) + ) + + return to_object + + +def _GetRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + return to_object + + +def _ListReasoningEnginesResponse_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["sdkHttpResponse"]) is not None: + setv(to_object, ["sdk_http_response"], getv(from_object, ["sdkHttpResponse"])) + + if getv(from_object, ["nextPageToken"]) is not None: + setv(to_object, ["next_page_token"], getv(from_object, ["nextPageToken"])) + + if getv(from_object, ["reasoningEngines"]) is not None: + setv( + to_object, + ["reasoning_engines"], + [ + _ReasoningEngine_from_vertex(item, to_object) + for item in getv(from_object, ["reasoningEngines"]) + ], + ) + + return to_object + + +def _ListRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["page_size"]) is not None: + setv(parent_object, ["_query", "pageSize"], getv(from_object, ["page_size"])) + + if getv(from_object, ["page_token"]) is not None: + setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"])) + + if getv(from_object, ["filter"]) is not None: + setv(parent_object, ["_query", "filter"], getv(from_object, ["filter"])) + + return to_object + + +def _ListRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["config"]) is not None: + _ListRuntimeConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _QueryRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["class_method"]) is not None: + setv(parent_object, ["classMethod"], getv(from_object, ["class_method"])) + + if getv(from_object, ["input"]) is not None: + setv(parent_object, ["input"], getv(from_object, ["input"])) + + if getv(from_object, ["include_all_fields"]) is not None: + setv(to_object, ["includeAllFields"], getv(from_object, ["include_all_fields"])) + + return to_object + + +def _QueryRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + _QueryRuntimeConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _ReasoningEngineContextSpecMemoryBankConfig_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["customizationConfigs"]) is not None: + setv( + to_object, + ["customization_configs"], + [item for item in getv(from_object, ["customizationConfigs"])], + ) + + if getv(from_object, ["disableMemoryRevisions"]) is not None: + setv( + to_object, + ["disable_memory_revisions"], + getv(from_object, ["disableMemoryRevisions"]), + ) + + if getv(from_object, ["generationConfig"]) is not None: + setv(to_object, ["generation_config"], getv(from_object, ["generationConfig"])) + + if getv(from_object, ["similaritySearchConfig"]) is not None: + setv( + to_object, + ["similarity_search_config"], + getv(from_object, ["similaritySearchConfig"]), + ) + + if getv(from_object, ["ttlConfig"]) is not None: + setv(to_object, ["ttl_config"], getv(from_object, ["ttlConfig"])) + + if getv(from_object, ["structuredMemoryConfigs"]) is not None: + setv( + to_object, + ["structured_memory_configs"], + [ + _StructuredMemoryConfig_from_vertex(item, to_object) + for item in getv(from_object, ["structuredMemoryConfigs"]) + ], + ) + + return to_object + + +def _ReasoningEngineContextSpecMemoryBankConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["customization_configs"]) is not None: + setv( + to_object, + ["customizationConfigs"], + [item for item in getv(from_object, ["customization_configs"])], + ) + + if getv(from_object, ["disable_memory_revisions"]) is not None: + setv( + to_object, + ["disableMemoryRevisions"], + getv(from_object, ["disable_memory_revisions"]), + ) + + if getv(from_object, ["generation_config"]) is not None: + setv(to_object, ["generationConfig"], getv(from_object, ["generation_config"])) + + if getv(from_object, ["similarity_search_config"]) is not None: + setv( + to_object, + ["similaritySearchConfig"], + getv(from_object, ["similarity_search_config"]), + ) + + if getv(from_object, ["ttl_config"]) is not None: + setv(to_object, ["ttlConfig"], getv(from_object, ["ttl_config"])) + + if getv(from_object, ["structured_memory_configs"]) is not None: + setv( + to_object, + ["structuredMemoryConfigs"], + [ + _StructuredMemoryConfig_to_vertex(item, to_object) + for item in getv(from_object, ["structured_memory_configs"]) + ], + ) + + return to_object + + +def _ReasoningEngineContextSpec_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["memoryBankConfig"]) is not None: + setv( + to_object, + ["memory_bank_config"], + _ReasoningEngineContextSpecMemoryBankConfig_from_vertex( + getv(from_object, ["memoryBankConfig"]), to_object + ), + ) + + return to_object + + +def _ReasoningEngineContextSpec_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["memory_bank_config"]) is not None: + setv( + to_object, + ["memoryBankConfig"], + _ReasoningEngineContextSpecMemoryBankConfig_to_vertex( + getv(from_object, ["memory_bank_config"]), to_object + ), + ) + + return to_object + + +def _ReasoningEngine_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["encryptionSpec"]) is not None: + setv(to_object, ["encryption_spec"], getv(from_object, ["encryptionSpec"])) + + if getv(from_object, ["contextSpec"]) is not None: + setv( + to_object, + ["context_spec"], + _ReasoningEngineContextSpec_from_vertex( + getv(from_object, ["contextSpec"]), to_object + ), + ) + + if getv(from_object, ["createTime"]) is not None: + setv(to_object, ["create_time"], getv(from_object, ["createTime"])) + + if getv(from_object, ["description"]) is not None: + setv(to_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["displayName"]) is not None: + setv(to_object, ["display_name"], getv(from_object, ["displayName"])) + + if getv(from_object, ["etag"]) is not None: + setv(to_object, ["etag"], getv(from_object, ["etag"])) + + if getv(from_object, ["labels"]) is not None: + setv(to_object, ["labels"], getv(from_object, ["labels"])) + + if getv(from_object, ["name"]) is not None: + setv(to_object, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["spec"]) is not None: + setv(to_object, ["spec"], getv(from_object, ["spec"])) + + if getv(from_object, ["updateTime"]) is not None: + setv(to_object, ["update_time"], getv(from_object, ["updateTime"])) + + if getv(from_object, ["trafficConfig"]) is not None: + setv(to_object, ["traffic_config"], getv(from_object, ["trafficConfig"])) + + return to_object + + +def _RunQueryJobRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["input_gcs_uri"]) is not None: + setv(parent_object, ["inputGcsUri"], getv(from_object, ["input_gcs_uri"])) + + if getv(from_object, ["output_gcs_uri"]) is not None: + setv(parent_object, ["outputGcsUri"], getv(from_object, ["output_gcs_uri"])) + + return to_object + + +def _RunQueryJobRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + setv( + to_object, + ["config"], + _RunQueryJobRuntimeConfig_to_vertex( + getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _RuntimeOperation_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["metadata"]) is not None: + setv(to_object, ["metadata"], getv(from_object, ["metadata"])) + + if getv(from_object, ["done"]) is not None: + setv(to_object, ["done"], getv(from_object, ["done"])) + + if getv(from_object, ["error"]) is not None: + setv(to_object, ["error"], getv(from_object, ["error"])) + + if getv(from_object, ["response"]) is not None: + setv( + to_object, + ["response"], + _ReasoningEngine_from_vertex(getv(from_object, ["response"]), to_object), + ) + + return to_object + + +def _StructuredMemoryConfig_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["schemaConfigs"]) is not None: + setv( + to_object, + ["schema_configs"], + [ + _StructuredMemorySchemaConfig_from_vertex(item, to_object) + for item in getv(from_object, ["schemaConfigs"]) + ], + ) + + if getv(from_object, ["scopeKeys"]) is not None: + setv(to_object, ["scope_keys"], getv(from_object, ["scopeKeys"])) + + return to_object + + +def _StructuredMemoryConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["schema_configs"]) is not None: + setv( + to_object, + ["schemaConfigs"], + [ + _StructuredMemorySchemaConfig_to_vertex(item, to_object) + for item in getv(from_object, ["schema_configs"]) + ], + ) + + if getv(from_object, ["scope_keys"]) is not None: + setv(to_object, ["scopeKeys"], getv(from_object, ["scope_keys"])) + + return to_object + + +def _StructuredMemorySchemaConfig_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["schema"]) is not None: + setv(to_object, ["memory_schema"], getv(from_object, ["schema"])) + + if getv(from_object, ["id"]) is not None: + setv(to_object, ["id"], getv(from_object, ["id"])) + + if getv(from_object, ["memoryType"]) is not None: + setv(to_object, ["memory_type"], getv(from_object, ["memoryType"])) + + return to_object + + +def _StructuredMemorySchemaConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["memory_schema"]) is not None: + setv(to_object, ["schema"], getv(from_object, ["memory_schema"])) + + if getv(from_object, ["id"]) is not None: + setv(to_object, ["id"], getv(from_object, ["id"])) + + if getv(from_object, ["memory_type"]) is not None: + setv(to_object, ["memoryType"], getv(from_object, ["memory_type"])) + + return to_object + + +def _UpdateRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["display_name"]) is not None: + setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) + + if getv(from_object, ["description"]) is not None: + setv(parent_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["spec"]) is not None: + setv(parent_object, ["spec"], getv(from_object, ["spec"])) + + if getv(from_object, ["context_spec"]) is not None: + setv( + parent_object, + ["contextSpec"], + _ReasoningEngineContextSpec_to_vertex( + getv(from_object, ["context_spec"]), to_object + ), + ) + + if getv(from_object, ["psc_interface_config"]) is not None: + setv( + parent_object, + ["pscInterfaceConfig"], + getv(from_object, ["psc_interface_config"]), + ) + + if getv(from_object, ["agent_gateway_config"]) is not None: + setv( + parent_object, + ["agentGatewayConfig"], + getv(from_object, ["agent_gateway_config"]), + ) + + if getv(from_object, ["encryption_spec"]) is not None: + setv(parent_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"])) + + if getv(from_object, ["labels"]) is not None: + setv(parent_object, ["labels"], getv(from_object, ["labels"])) + + if getv(from_object, ["source_packages"]) is not None: + setv(parent_object, ["sourcePackages"], getv(from_object, ["source_packages"])) + + if getv(from_object, ["entrypoint_module"]) is not None: + setv( + parent_object, + ["entrypointModule"], + getv(from_object, ["entrypoint_module"]), + ) + + if getv(from_object, ["entrypoint_object"]) is not None: + setv( + parent_object, + ["entrypointObject"], + getv(from_object, ["entrypoint_object"]), + ) + + if getv(from_object, ["requirements_file"]) is not None: + setv( + parent_object, + ["requirementsFile"], + getv(from_object, ["requirements_file"]), + ) + + if getv(from_object, ["agent_framework"]) is not None: + setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) + + if getv(from_object, ["python_version"]) is not None: + setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) + + if getv(from_object, ["update_mask"]) is not None: + setv( + parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"]) + ) + + if getv(from_object, ["traffic_config"]) is not None: + setv(parent_object, ["trafficConfig"], getv(from_object, ["traffic_config"])) + + return to_object + + +def _UpdateRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + _UpdateRuntimeConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + class Runtimes(_api_module.BaseModule): + def cancel_query_job( + self, + *, + name: str, + config: Optional[types.CancelQueryJobRuntimeConfigOrDict] = None, + ) -> types.CancelQueryJobResult: + """ + Cancels a long-running query job on an Agent Runtime. + + Args: + name (str): + Required. The reasoning engine resource name. + config (CancelQueryJobRuntimeConfigOrDict): + Optional. The configuration for the cancel_query_job. + + """ + + parameter_model = types._CancelQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CancelQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:cancelAsyncQuery".format_map(request_url_dict) + else: + path = "{name}:cancelAsyncQuery" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.CancelQueryJobResult._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _check_query_job( + self, + *, + name: str, + config: Optional[types.CheckQueryJobRuntimeConfigOrDict] = None, + ) -> types.CheckQueryJobResult: + """ + Query an Agent Runtime asynchronously. + """ + + parameter_model = types._CheckQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CheckQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:checkQueryJob".format_map(request_url_dict) + else: + path = "{name}:checkQueryJob" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _CheckQueryJobResult_from_vertex(response_dict) + + return_value = types.CheckQueryJobResult._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _run_query_job( + self, + *, + name: str, + config: Optional[types._RunQueryJobRuntimeConfigOrDict] = None, + ) -> types.RuntimeOperation: + """ + Run a query job on an agent runtime. + """ + + parameter_model = types._RunQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _RunQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:asyncQuery".format_map(request_url_dict) + else: + path = "{name}:asyncQuery" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _create( + self, *, config: Optional[types.CreateRuntimeConfigOrDict] = None + ) -> types.RuntimeOperation: + """ + Creates a new Agent Runtime. + """ + + parameter_model = types._CreateRuntimeRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CreateRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _delete( + self, + *, + name: str, + force: Optional[bool] = None, + config: Optional[types.DeleteRuntimeConfigOrDict] = None, + ) -> types.DeleteRuntimeOperation: + """ + Delete an Agent Runtime resource. + + Args: + name (str): + Required. The name of the Agent Runtime to be deleted. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` + or `reasoningEngines/{resource_id}`. + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Agent Runtime has undeleted child resources. Defaults to False. + config (DeleteRuntimeConfig): + Optional. Additional configurations for deleting the Agent Runtime. + + """ + + parameter_model = types._DeleteRuntimeRequestParameters( + name=name, + force=force, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _DeleteRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("delete", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteRuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _get( + self, *, name: str, config: Optional[types.GetRuntimeConfigOrDict] = None + ) -> types.ReasoningEngine: + """ + Get an Agent Runtime instance. + """ + + parameter_model = types._GetRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ReasoningEngine_from_vertex(response_dict) + + return_value = types.ReasoningEngine._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _list( + self, *, config: Optional[types.ListRuntimeConfigOrDict] = None + ) -> types.ListReasoningEnginesResponse: + """ + Lists Agent Runtimes. + """ + + parameter_model = types._ListRuntimeRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _ListRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) + + return_value = types.ListReasoningEnginesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _get_agent_operation( + self, + *, + operation_name: str, + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeOperation: + parameter_model = types._GetRuntimeOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRuntimeOperationParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _query( + self, *, name: str, config: Optional[types.QueryRuntimeConfigOrDict] = None + ) -> types.QueryReasoningEngineResponse: + """ + Query an Agent Runtime. + """ + + parameter_model = types._QueryRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _QueryRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:query".format_map(request_url_dict) + else: + path = "{name}:query" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.QueryReasoningEngineResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _update( + self, *, name: str, config: Optional[types.UpdateRuntimeConfigOrDict] = None + ) -> types.RuntimeOperation: + """ + Updates an Agent Runtime. + """ + + parameter_model = types._UpdateRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpdateRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("patch", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + _revisions = None @property @@ -49,17 +1511,2136 @@ 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) + + if self._api_client.vertexai: + response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) + + return_value = types.ListReasoningEnginesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _get_agent_operation( + self, + *, + operation_name: str, + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeOperation: + parameter_model = types._GetRuntimeOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRuntimeOperationParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _query( + self, *, name: str, config: Optional[types.QueryRuntimeConfigOrDict] = None + ) -> types.QueryReasoningEngineResponse: + """ + Query an Agent Runtime. + """ + + parameter_model = types._QueryRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _QueryRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:query".format_map(request_url_dict) + else: + path = "{name}:query" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.QueryReasoningEngineResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _update( + self, *, name: str, config: Optional[types.UpdateRuntimeConfigOrDict] = None + ) -> types.RuntimeOperation: + """ + Updates an Agent Runtime. + """ + + parameter_model = types._UpdateRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpdateRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "patch", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + _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 +3652,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 05573aefec..4686e3d3fa 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 @@ -126,6 +114,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 @@ -135,13 +131,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 @@ -154,6 +143,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 @@ -162,8 +156,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 @@ -172,14 +166,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 @@ -188,50 +180,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 @@ -239,18 +199,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 @@ -312,26 +266,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 @@ -364,21 +318,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 @@ -421,9 +360,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 @@ -456,33 +407,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 @@ -546,12 +470,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 @@ -763,9 +711,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 @@ -902,33 +850,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 @@ -995,6 +925,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 @@ -1007,6 +940,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 @@ -1094,39 +1039,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 @@ -1211,12 +1123,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 @@ -1556,15 +1489,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 @@ -1869,7 +1802,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 @@ -1917,18 +1849,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 @@ -2152,39 +2102,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 @@ -2267,12 +2184,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 @@ -2294,12 +2205,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 @@ -2342,81 +2259,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", @@ -2864,21 +2706,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", @@ -3014,39 +2856,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", @@ -3458,30 +3324,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", @@ -3494,45 +3336,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", @@ -3578,15 +3420,15 @@ "ListSandboxEnvironmentTemplatesResponse", "ListSandboxEnvironmentTemplatesResponseDict", "ListSandboxEnvironmentTemplatesResponseOrDict", - "CreateAgentEngineSandboxSnapshotConfig", - "CreateAgentEngineSandboxSnapshotConfigDict", - "CreateAgentEngineSandboxSnapshotConfigOrDict", + "CreateRuntimeSandboxSnapshotConfig", + "CreateRuntimeSandboxSnapshotConfigDict", + "CreateRuntimeSandboxSnapshotConfigOrDict", "SandboxEnvironmentSnapshot", "SandboxEnvironmentSnapshotDict", "SandboxEnvironmentSnapshotOrDict", - "AgentEngineSandboxSnapshotOperation", - "AgentEngineSandboxSnapshotOperationDict", - "AgentEngineSandboxSnapshotOperationOrDict", + "RuntimeSandboxSnapshotOperation", + "RuntimeSandboxSnapshotOperationDict", + "RuntimeSandboxSnapshotOperationOrDict", "DeleteSandboxEnvironmentSnapshotConfig", "DeleteSandboxEnvironmentSnapshotConfigDict", "DeleteSandboxEnvironmentSnapshotConfigOrDict", @@ -3602,54 +3444,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", @@ -4310,21 +4152,6 @@ "EvalRunInferenceConfig", "EvalRunInferenceConfigDict", "EvalRunInferenceConfigOrDict", - "AgentEngine", - "AgentEngineDict", - "AgentEngineOrDict", - "AgentEngineConfig", - "AgentEngineConfigDict", - "AgentEngineConfigOrDict", - "RunQueryJobAgentEngineConfig", - "RunQueryJobAgentEngineConfigDict", - "RunQueryJobAgentEngineConfigOrDict", - "RunQueryJobResult", - "RunQueryJobResultDict", - "RunQueryJobResultOrDict", - "CheckQueryJobResponse", - "CheckQueryJobResponseDict", - "CheckQueryJobResponseOrDict", "AssembleDataset", "AssembleDatasetDict", "AssembleDatasetOrDict", @@ -4361,9 +4188,6 @@ "OptimizeJobConfig", "OptimizeJobConfigDict", "OptimizeJobConfigOrDict", - "AgentEngineRuntimeRevision", - "AgentEngineRuntimeRevisionDict", - "AgentEngineRuntimeRevisionOrDict", "ListDeployableModelsConfig", "ListDeployableModelsConfigDict", "ListDeployableModelsConfigOrDict", @@ -4385,9 +4209,24 @@ "DeployOption", "DeployOptionDict", "DeployOptionOrDict", - "A2aTaskState", - "Role", - "State", + "Runtime", + "RuntimeDict", + "RuntimeOrDict", + "AgentRuntimeConfig", + "AgentRuntimeConfigDict", + "AgentRuntimeConfigOrDict", + "RunQueryJobRuntimeConfig", + "RunQueryJobRuntimeConfigDict", + "RunQueryJobRuntimeConfigOrDict", + "RunQueryJobResult", + "RunQueryJobResultDict", + "RunQueryJobResultOrDict", + "CheckQueryJobResponse", + "CheckQueryJobResponseDict", + "CheckQueryJobResponseOrDict", + "RuntimeRevision", + "RuntimeRevisionDict", + "RuntimeRevisionOrDict", "Strategy", "AcceleratorType", "Type", @@ -4395,6 +4234,7 @@ "ManagedTopicEnum", "IdentityType", "AgentServerMode", + "State", "MemoryType", "Operator", "RagFileType", @@ -4460,12 +4300,6 @@ "MessageDict", "Importance", "ParsedResponseUnion", - "_DeleteAgentEngineTaskRequestParameters", - "_GetAgentEngineTaskRequestParameters", - "_ListAgentEngineTasksRequestParameters", - "_CreateAgentEngineTaskRequestParameters", - "_AppendAgentEngineTaskEventRequestParameters", - "_ListAgentEngineTaskEventsRequestParameters", "_CreateEvaluationExperimentParameters", "_CreateEvaluationItemParameters", "_CreateEvaluationMetricParameters", @@ -4491,16 +4325,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", "_IngestEventsRequestParameters", @@ -4536,17 +4375,12 @@ "_ImportRagFilesRequestParameters", "_GetImportFilesOperationParameters", "_UploadRagFileParameters", - "_GetAgentEngineRuntimeRevisionRequestParameters", - "_ListAgentEngineRuntimeRevisionsRequestParameters", - "_DeleteAgentEngineRuntimeRevisionRequestParameters", - "_GetDeleteAgentEngineRuntimeRevisionOperationParameters", - "_QueryAgentEngineRuntimeRevisionRequestParameters", - "_CreateAgentEngineSandboxRequestParameters", - "_DeleteAgentEngineSandboxRequestParameters", - "_ExecuteCodeAgentEngineSandboxRequestParameters", - "_GetAgentEngineSandboxRequestParameters", - "_ListAgentEngineSandboxesRequestParameters", - "_GetAgentEngineSandboxOperationParameters", + "_CreateRuntimeSandboxRequestParameters", + "_DeleteRuntimeSandboxRequestParameters", + "_ExecuteCodeRuntimeSandboxRequestParameters", + "_GetRuntimeSandboxRequestParameters", + "_ListRuntimeSandboxesRequestParameters", + "_GetRuntimeSandboxOperationParameters", "_CreateSandboxEnvironmentTemplateRequestParameters", "_DeleteSandboxEnvironmentTemplateRequestParameters", "_GetSandboxEnvironmentTemplateRequestParameters", @@ -4556,15 +4390,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 e95aa77edf..81b61d6666 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 for a particular scope.""" + + 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 for a particular scope.""" - 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 for a particular scope.""" +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 for a particular scope.""" - - enable_third_person_memories: Optional[bool] - """Optional. Indicates whether the memories will be generated in the third person (i.e. "The user generates memories with Memory Bank."). By default, the memories will be generated in the first person (i.e. "I generate memories with Memory Bank.")""" - - generate_memories_examples: Optional[ - list[MemoryBankCustomizationConfigGenerateMemoriesExampleDict] - ] - """Optional. Provides examples of how to generate memories for a particular scope.""" +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 runtimes.""" - 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 runtimes. + """, ) -class ReasoningEngineSpecContainerSpecDict(TypedDict, total=False): - """Specification for deploying from a container image.""" +class ListReasoningEnginesResponseDict(TypedDict, total=False): + """Response for listing agent runtimes.""" - 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 runtimes. + """ -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.""", + description="""The minimum number of instances to run for the Agent Runtime. + Defaults to 1. Range: [0, 10]. + """, ) - etag: Optional[str] = Field( + max_instances: Optional[int] = 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. - 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,456 +10364,455 @@ class QueryAgentEngineConfigDict(TypedDict, total=False): """""" -QueryAgentEngineConfigOrDict = Union[QueryAgentEngineConfig, QueryAgentEngineConfigDict] +QueryRuntimeRevisionConfigOrDict = Union[ + QueryRuntimeRevisionConfig, QueryRuntimeRevisionConfigDict +] -class _QueryAgentEngineRequestParameters(_common.BaseModel): - """Parameters for querying agent engines.""" +class _QueryRuntimeRevisionRequestParameters(_common.BaseModel): + """Parameters for querying agent runtime revisions.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime revision.""" + ) + config: Optional[QueryRuntimeRevisionConfig] = Field( + default=None, description="""""" ) - config: Optional[QueryAgentEngineConfig] = Field(default=None, description="""""") -class _QueryAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for querying agent engines.""" +class _QueryRuntimeRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for querying agent runtime revisions.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime revision.""" - config: Optional[QueryAgentEngineConfigDict] + config: Optional[QueryRuntimeRevisionConfigDict] """""" -_QueryAgentEngineRequestParametersOrDict = Union[ - _QueryAgentEngineRequestParameters, _QueryAgentEngineRequestParametersDict +_QueryRuntimeRevisionRequestParametersOrDict = Union[ + _QueryRuntimeRevisionRequestParameters, _QueryRuntimeRevisionRequestParametersDict ] -class QueryReasoningEngineResponse(_common.BaseModel): - """The response for querying an agent engine.""" +class CreateMemoryBankConfig(_common.BaseModel): + """Config for create memory bank.""" - output: Optional[Any] = Field( - default=None, - description="""Response provided by users in JSON object format.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class QueryReasoningEngineResponseDict(TypedDict, total=False): - """The response for querying an agent engine.""" +class CreateMemoryBankConfigDict(TypedDict, total=False): + """Config for create memory bank.""" - output: Optional[Any] - """Response provided by users in JSON object format.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -QueryReasoningEngineResponseOrDict = Union[ - QueryReasoningEngineResponse, QueryReasoningEngineResponseDict -] +CreateMemoryBankConfigOrDict = Union[CreateMemoryBankConfig, CreateMemoryBankConfigDict] -class UpdateAgentEngineConfig(_common.BaseModel): - """Config for updating agent engine.""" +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.""" - ) - display_name: Optional[str] = Field( - default=None, - description="""The user-defined name of the Agent Engine. + config: Optional[CreateMemoryBankConfig] = Field(default=None, description="""""") - 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( + +class _CreateMemoryBankRequestParametersDict(TypedDict, total=False): + """Parameters for creating memory banks.""" + + config: Optional[CreateMemoryBankConfigDict] + """""" + + +_CreateMemoryBankRequestParametersOrDict = Union[ + _CreateMemoryBankRequestParameters, _CreateMemoryBankRequestParametersDict +] + + +class MemoryBank(_common.BaseModel): + """A memory bank.""" + + name: Optional[str] = Field( default=None, - description="""Optional. The context spec to be used for the Agent Engine.""", + 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.""", ) - psc_interface_config: Optional[PscInterfaceConfig] = Field( + + +class MemoryBankDict(TypedDict, total=False): + """A 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.""" + + +MemoryBankOrDict = Union[MemoryBank, MemoryBankDict] + + +class MemoryBankOperation(_common.BaseModel): + """Operation that has an memory bank as a response.""" + + name: Optional[str] = Field( default=None, - description="""Optional. The PSC interface config for PSC-I to be used for the - Agent Engine.""", + 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}`.""", ) - min_instances: Optional[int] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""The minimum number of instances to run for the Agent Engine. - Defaults to 1. Range: [0, 10]. - """, + 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.""", ) - max_instances: Optional[int] = Field( + done: Optional[bool] = 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="""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_limits: Optional[dict[str, str]] = Field( + error: Optional[dict[str, Any]] = 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 error result of the operation in case of failure or cancellation.""", ) - 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. - """, + response: Optional[MemoryBank] = Field( + default=None, description="""The created Memory Bank.""" ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, - description="""The encryption spec to be used for the Agent Engine.""", + + +class MemoryBankOperationDict(TypedDict, total=False): + """Operation that has an memory bank 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[MemoryBankDict] + """The created Memory Bank.""" + + +MemoryBankOperationOrDict = Union[MemoryBankOperation, MemoryBankOperationDict] + + +class DeleteMemoryBankConfig(_common.BaseModel): + """Config for delete memory bank.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - labels: Optional[dict[str, str]] = Field( - default=None, description="""The labels to be used for the Agent Engine.""" + + +class DeleteMemoryBankConfigDict(TypedDict, total=False): + """Config for delete memory bank.""" + + 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.""" ) - 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. - """, + force: Optional[bool] = Field( + default=False, + description="""If set to true, any child resources will also be deleted.""", ) - 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`. - """, - ) - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig - ] = Field( + 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] + """""" + + +_DeleteMemoryBankRequestParametersOrDict = Union[ + _DeleteMemoryBankRequestParameters, _DeleteMemoryBankRequestParametersDict +] + + +class DeleteMemoryBankOperation(_common.BaseModel): + """Operation for deleting a memory bank.""" + + name: Optional[str] = Field( default=None, - description="""Agent Gateway configuration for a Reasoning Engine deployment.""", + 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}`.""", ) - keep_alive_probe: Optional[KeepAliveProbe] = Field( + metadata: Optional[dict[str, Any]] = 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="""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_mask: Optional[str] = Field( + done: Optional[bool] = Field( default=None, - description="""The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""", + 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.""", ) - traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Traffic distribution configuration for the Reasoning Engine.""", + description="""The error result of the operation in case of failure or cancellation.""", ) -class UpdateAgentEngineConfigDict(TypedDict, total=False): - """Config for updating agent engine.""" +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 Agent 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.""" - 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 Agent Engine.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - spec: Optional[ReasoningEngineSpecDict] - """Optional. Configurations of the Agent Engine.""" - context_spec: Optional[ReasoningEngineContextSpecDict] - """Optional. The context spec to be used for the Agent Engine.""" +DeleteMemoryBankOperationOrDict = Union[ + DeleteMemoryBankOperation, DeleteMemoryBankOperationDict +] - psc_interface_config: Optional[PscInterfaceConfigDict] - """Optional. The PSC interface config for PSC-I to be used for the - Agent Engine.""" - min_instances: Optional[int] - """The minimum number of instances to run for the Agent Engine. - Defaults to 1. Range: [0, 10]. - """ +class IngestionDirectContentsSourceEvent(_common.BaseModel): + """The direct contents source event for ingesting events.""" - 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]. - """ + content: Optional[genai_types.Content] = Field( + default=None, description="""Required. The content of the event.""" + ) + event_id: 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.""", + ) + event_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.""", + ) - 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 IngestionDirectContentsSourceEventDict(TypedDict, total=False): + """The direct contents source event for ingesting events.""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """The encryption spec to be used for the Agent Engine.""" + content: Optional[genai_types.Content] + """Required. The content of the event.""" - labels: Optional[dict[str, str]] - """The labels to be used for the Agent Engine.""" + 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.""" - 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. - """ + 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.""" - 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.""" +IngestionDirectContentsSourceEventOrDict = Union[ + IngestionDirectContentsSourceEvent, IngestionDirectContentsSourceEventDict +] - 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.""" +class IngestionDirectContentsSource(_common.BaseModel): + """The direct contents source for ingesting events.""" - 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. - """ + events: Optional[list[IngestionDirectContentsSourceEvent]] = Field( + default=None, description="""Required. The events to ingest.""" + ) - 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". - """ +class IngestionDirectContentsSourceDict(TypedDict, total=False): + """The direct contents source for ingesting events.""" - 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`. - """ + events: Optional[list[IngestionDirectContentsSourceEventDict]] + """Required. The events to ingest.""" - 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.""" +IngestionDirectContentsSourceOrDict = Union[ + IngestionDirectContentsSource, IngestionDirectContentsSourceDict +] - 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.""" - - -UpdateAgentEngineConfigOrDict = Union[ - UpdateAgentEngineConfig, UpdateAgentEngineConfigDict -] +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.""" + ) + string_value: Optional[str] = Field( + default=None, description="""Represents a string value.""" + ) + timestamp_value: Optional[datetime.datetime] = Field( + default=None, + description="""Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""", + ) -class _UpdateAgentEngineRequestParameters(_common.BaseModel): - """Parameters for updating agent engines.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[UpdateAgentEngineConfig] = Field(default=None, description="""""") +class MemoryMetadataValueDict(TypedDict, total=False): + """The metadata values for memories.""" + bool_value: Optional[bool] + """Represents a boolean value.""" -class _UpdateAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for updating agent engines.""" + double_value: Optional[float] + """Represents a double value.""" - name: Optional[str] - """Name of the agent engine.""" + string_value: Optional[str] + """Represents a string value.""" - config: Optional[UpdateAgentEngineConfigDict] - """""" + timestamp_value: Optional[datetime.datetime] + """Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""" -_UpdateAgentEngineRequestParametersOrDict = Union[ - _UpdateAgentEngineRequestParameters, _UpdateAgentEngineRequestParametersDict -] +MemoryMetadataValueOrDict = Union[MemoryMetadataValue, MemoryMetadataValueDict] -class CreateMemoryBankConfig(_common.BaseModel): - """Config for create memory bank.""" +class IngestEventsConfig(_common.BaseModel): + """Config for ingesting events.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) + 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="""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`.""", + ) + 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( + 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 revisions will be created for this request.""", + ) + metadata: Optional[dict[str, MemoryMetadataValue]] = Field( + default=None, + description="""Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""", + ) + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] = Field( + default=None, + description="""Optional. The strategy to use when applying metadata to existing memories.""", + ) -class CreateMemoryBankConfigDict(TypedDict, total=False): - """Config for create memory bank.""" +class IngestEventsConfigDict(TypedDict, total=False): + """Config for ingesting events.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" + wait_for_completion: Optional[bool] + """Waits for the underlying memory generation operation to complete + before returning. Defaults to false.""" -CreateMemoryBankConfigOrDict = Union[CreateMemoryBankConfig, CreateMemoryBankConfigDict] + 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`.""" + 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.""" -class _CreateMemoryBankRequestParameters(_common.BaseModel): - """Parameters for creating memory banks.""" + 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.""" - config: Optional[CreateMemoryBankConfig] = Field(default=None, description="""""") + 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 _CreateMemoryBankRequestParametersDict(TypedDict, total=False): - """Parameters for creating memory banks.""" + metadata: Optional[dict[str, MemoryMetadataValueDict]] + """Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""" - config: Optional[CreateMemoryBankConfigDict] - """""" + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] + """Optional. The strategy to use when applying metadata to existing memories.""" -_CreateMemoryBankRequestParametersOrDict = Union[ - _CreateMemoryBankRequestParameters, _CreateMemoryBankRequestParametersDict -] +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="""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.""" + ) + direct_contents_source: Optional[IngestionDirectContentsSource] = 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="""The direct memories source of the events that should be ingested.""", + ) + scope: Optional[dict[str, str]] = 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 '*'.""", + ) + generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + default=None, + description="""The configuration for the memory generation trigger.""", ) + config: Optional[IngestEventsConfig] = Field(default=None, description="""""") -class MemoryBankDict(TypedDict, total=False): - """A memory bank.""" +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.""" + stream_id: Optional[str] + """The ID of the stream to ingest events into.""" -MemoryBankOrDict = Union[MemoryBank, MemoryBankDict] + direct_contents_source: Optional[IngestionDirectContentsSourceDict] + """The direct memories source of the events that should be ingested.""" + scope: Optional[dict[str, str]] + """The scope of the memories that should be generated from the stream. -class MemoryBankOperation(_common.BaseModel): - """Operation that has an memory bank as a response.""" + Memories will be consolidated across memories with the same scope. Scope + values cannot contain the wildcard character '*'.""" + + generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] + """The configuration for the memory generation trigger.""" + + config: Optional[IngestEventsConfigDict] + """""" + + +_IngestEventsRequestParametersOrDict = Union[ + _IngestEventsRequestParameters, _IngestEventsRequestParametersDict +] + + +class MemoryBankIngestEventsOperation(_common.BaseModel): + """Operation that ingests events into a memory bank.""" name: Optional[str] = Field( default=None, @@ -11242,13 +10830,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}`.""" @@ -11262,330 +10847,363 @@ 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 GetMemoryBankOperationConfig(_common.BaseModel): http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class DeleteMemoryBankConfigDict(TypedDict, total=False): - """Config for delete memory bank.""" +class GetMemoryBankOperationConfigDict(TypedDict, total=False): http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteMemoryBankConfigOrDict = Union[DeleteMemoryBankConfig, DeleteMemoryBankConfigDict] +GetMemoryBankOperationConfigOrDict = Union[ + GetMemoryBankOperationConfig, GetMemoryBankOperationConfigDict +] -class _DeleteMemoryBankRequestParameters(_common.BaseModel): - """Parameters for deleting 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.""" ) - force: Optional[bool] = Field( - default=False, - description="""If set to true, any child resources will also be deleted.""", + config: Optional[GetMemoryBankOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) - 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.""" +class _GetMemoryBankOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a memory bank as a response.""" - force: Optional[bool] - """If set to true, any child resources will also be deleted.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - config: Optional[DeleteMemoryBankConfigDict] - """""" + config: Optional[GetMemoryBankOperationConfigDict] + """Used to override the default configuration.""" -_DeleteMemoryBankRequestParametersOrDict = Union[ - _DeleteMemoryBankRequestParameters, _DeleteMemoryBankRequestParametersDict +_GetMemoryBankOperationParametersOrDict = Union[ + _GetMemoryBankOperationParameters, _GetMemoryBankOperationParametersDict ] -class DeleteMemoryBankOperation(_common.BaseModel): - """Operation for deleting a memory bank.""" +class MemoryConfig(_common.BaseModel): + """Config for creating a Memory.""" - name: 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="""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. Input only. The TTL for this resource. + + The expiration time is computed: now + TTL.""", ) - metadata: Optional[dict[str, Any]] = Field( + expire_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="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", ) - done: Optional[bool] = Field( + revision_expire_time: Optional[datetime.datetime] = 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. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", ) - error: Optional[dict[str, Any]] = Field( + revision_ttl: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + 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 DeleteMemoryBankOperationDict(TypedDict, total=False): - """Operation for deleting 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 MemoryConfigDict(TypedDict, total=False): + """Config for creating a Memory.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + display_name: Optional[str] + """The display name of the memory.""" -DeleteMemoryBankOperationOrDict = Union[ - DeleteMemoryBankOperation, DeleteMemoryBankOperationDict -] + description: Optional[str] + """The description of the memory.""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" -class IngestionDirectContentsSourceEvent(_common.BaseModel): - """The direct contents source event for ingesting events.""" + ttl: Optional[str] + """Optional. Input only. The TTL for this resource. - content: Optional[genai_types.Content] = Field( - default=None, description="""Required. The content of the event.""" - ) - event_id: 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.""", - ) - event_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.""", - ) + 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.""" -class IngestionDirectContentsSourceEventDict(TypedDict, total=False): - """The direct contents source event for ingesting events.""" + 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.""" - content: Optional[genai_types.Content] - """Required. The content of the event.""" + revision_ttl: Optional[str] + """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" - 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.""" + disable_memory_revisions: Optional[bool] + """Optional. Input only. If true, no revision will be created for this request.""" - 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.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The topics of the memory.""" + metadata: Optional[dict[str, MemoryMetadataValueDict]] + """Optional. User-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""" -IngestionDirectContentsSourceEventOrDict = Union[ - IngestionDirectContentsSourceEvent, IngestionDirectContentsSourceEventDict -] + 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 IngestionDirectContentsSource(_common.BaseModel): - """The direct contents source for ingesting events.""" +MemoryConfigOrDict = Union[MemoryConfig, MemoryConfigDict] - events: Optional[list[IngestionDirectContentsSourceEvent]] = Field( - default=None, description="""Required. The events to ingest.""" + +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.""", + ) + fact: Optional[str] = Field( + default=None, + description="""The fact of the memory. + + This is the semantic knowledge extracted from the source content).""", + ) + scope: Optional[dict[str, str]] = Field( + default=None, + description="""The scope of the memory. + + Memories are isolated within their scope. The scope is defined when + creating or generating memories. Up to 5 key-value pairs are accepted, + and scope values cannot contain the wildcard character '*'.""", ) + config: Optional[MemoryConfig] = Field(default=None, description="""""") -class IngestionDirectContentsSourceDict(TypedDict, total=False): - """The direct contents source for ingesting events.""" +class _CreateMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for creating Memories.""" - events: Optional[list[IngestionDirectContentsSourceEventDict]] - """Required. The events to ingest.""" + name: Optional[str] + """Name of the memory bank to create the memory under.""" + + fact: Optional[str] + """The fact of the memory. + This is the semantic knowledge extracted from the source content).""" -IngestionDirectContentsSourceOrDict = Union[ - IngestionDirectContentsSource, IngestionDirectContentsSourceDict + scope: Optional[dict[str, str]] + """The scope of the memory. + + Memories are isolated within their scope. The scope is defined when + creating or generating memories. Up to 5 key-value pairs are accepted, + and scope values cannot contain the wildcard character '*'.""" + + config: Optional[MemoryConfigDict] + """""" + + +_CreateMemoryRequestParametersOrDict = Union[ + _CreateMemoryRequestParameters, _CreateMemoryRequestParametersDict ] -class MemoryMetadataValue(_common.BaseModel): - """The metadata values for memories.""" +class MemoryStructuredContent(_common.BaseModel): + """Represents the structured value of the memory.""" - bool_value: Optional[bool] = Field( - default=None, description="""Represents a boolean value.""" - ) - double_value: Optional[float] = Field( - default=None, description="""Represents a double value.""" - ) - string_value: Optional[str] = Field( - default=None, description="""Represents a string value.""" + data: Optional[dict[str, Any]] = Field( + default=None, + description="""Required. Represents the structured value of the memory.""", ) - timestamp_value: Optional[datetime.datetime] = Field( + schema_id: Optional[str] = Field( default=None, - description="""Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""", + description="""Required. Represents the schema ID for which this structured memory belongs to.""", ) -class MemoryMetadataValueDict(TypedDict, total=False): - """The metadata values for memories.""" - - bool_value: Optional[bool] - """Represents a boolean value.""" - - double_value: Optional[float] - """Represents a double value.""" +class MemoryStructuredContentDict(TypedDict, total=False): + """Represents the structured value of the memory.""" - string_value: Optional[str] - """Represents a string value.""" + data: Optional[dict[str, Any]] + """Required. Represents the structured value of the memory.""" - timestamp_value: Optional[datetime.datetime] - """Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""" + schema_id: Optional[str] + """Required. Represents the schema ID for which this structured memory belongs to.""" -MemoryMetadataValueOrDict = Union[MemoryMetadataValue, MemoryMetadataValueDict] +MemoryStructuredContentOrDict = Union[ + MemoryStructuredContent, MemoryStructuredContentDict +] -class IngestEventsConfig(_common.BaseModel): - """Config for ingesting events.""" +class Memory(_common.BaseModel): + """A memory.""" - 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. Represents the timestamp when this Memory was created.""", ) - wait_for_completion: Optional[bool] = Field( - default=False, - description="""Waits for the underlying memory generation operation to complete - before returning. Defaults to false.""", + description: Optional[str] = Field( + default=None, + description="""Optional. Represents the description of the Memory.""", ) - force_flush: Optional[bool] = Field( + disable_memory_revisions: Optional[bool] = 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="""Optional. Input only. Indicates whether no revision will be created for this request.""", ) - revision_labels: Optional[dict[str, str]] = Field( + display_name: 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 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.""" - - 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`.""" - - 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.""" - - revision_ttl: Optional[str] - """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" + description: Optional[str] + """Optional. Represents the description of the Memory.""" disable_memory_revisions: Optional[bool] - """Optional. Input only. If true, no revisions will be created for this request.""" - - 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.""" - + """Optional. Input only. Indicates whether no revision will be created for this request.""" -IngestEventsConfigOrDict = Union[IngestEventsConfig, IngestEventsConfigDict] + 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.""" -class _IngestEventsRequestParameters(_common.BaseModel): - """Parameters for ingesting events to Memory Bank.""" + fact: Optional[str] + """Optional. Represents semantic knowledge extracted from the source content.""" - 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.""" - ) - direct_contents_source: Optional[IngestionDirectContentsSource] = Field( - default=None, - description="""The direct memories source of the events that should be ingested.""", - ) - scope: Optional[dict[str, str]] = Field( - default=None, - description="""The scope of the memories that should be generated from the stream. + 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.""" - 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="""""") + 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.""" -class _IngestEventsRequestParametersDict(TypedDict, total=False): - """Parameters for ingesting events to Memory Bank.""" + 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.""" - name: Optional[str] - """Name of the Memory Bank to ingest events into.""" + revision_ttl: Optional[str] + """Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""" - stream_id: Optional[str] - """The ID of the stream to ingest events into.""" + 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 '*'.""" - direct_contents_source: Optional[IngestionDirectContentsSourceDict] - """The direct memories source of the events that should be ingested.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. Represents the Topics of the Memory.""" - scope: Optional[dict[str, str]] - """The scope of the memories that should be generated from the stream. + ttl: Optional[str] + """Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""" - Memories will be consolidated across memories with the same scope. Scope - values cannot contain the wildcard character '*'.""" + update_time: Optional[datetime.datetime] + """Output only. Represents the timestamp when this Memory was most recently updated.""" - generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] - """The configuration for the memory generation trigger.""" + 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.""" - config: Optional[IngestEventsConfigDict] - """""" + structured_content: Optional[MemoryStructuredContentDict] + """Optional. Represents the structured content of the memory.""" -_IngestEventsRequestParametersOrDict = Union[ - _IngestEventsRequestParameters, _IngestEventsRequestParametersDict -] +MemoryOrDict = Union[Memory, MemoryDict] -class MemoryBankIngestEventsOperation(_common.BaseModel): - """Operation that ingests events into a memory bank.""" +class MemoryOperation(_common.BaseModel): + """Operation that has a memory as a response.""" name: Optional[str] = Field( default=None, @@ -11603,10 +11221,11 @@ class MemoryBankIngestEventsOperation(_common.BaseModel): 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 MemoryBankIngestEventsOperationDict(TypedDict, total=False): - """Operation that ingests events into a memory bank.""" +class MemoryOperationDict(TypedDict, total=False): + """Operation that has a memory as a response.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -11620,449 +11239,451 @@ class MemoryBankIngestEventsOperationDict(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.""" -MemoryBankIngestEventsOperationOrDict = Union[ - MemoryBankIngestEventsOperation, MemoryBankIngestEventsOperationDict -] +MemoryOperationOrDict = Union[MemoryOperation, MemoryOperationDict] -class GetMemoryBankOperationConfig(_common.BaseModel): + +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 GetMemoryBankOperationConfigDict(TypedDict, total=False): +class DeleteMemoryConfigDict(TypedDict, total=False): + """Config for deleting a Memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetMemoryBankOperationConfigOrDict = Union[ - GetMemoryBankOperationConfig, GetMemoryBankOperationConfigDict -] +DeleteMemoryConfigOrDict = Union[DeleteMemoryConfig, DeleteMemoryConfigDict] -class _GetMemoryBankOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a memory bank as a response.""" +class _DeleteMemoryRequestParameters(_common.BaseModel): + """Parameters for deleting memories.""" - 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.""" + name: Optional[str] = Field( + default=None, description="""Name of the memory to delete.""" ) + config: Optional[DeleteMemoryConfig] = Field(default=None, description="""""") -class _GetMemoryBankOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a memory bank as a response.""" +class _DeleteMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for deleting memories.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + name: Optional[str] + """Name of the memory to delete.""" - config: Optional[GetMemoryBankOperationConfigDict] - """Used to override the default configuration.""" + config: Optional[DeleteMemoryConfigDict] + """""" -_GetMemoryBankOperationParametersOrDict = Union[ - _GetMemoryBankOperationParameters, _GetMemoryBankOperationParametersDict +_DeleteMemoryRequestParametersOrDict = Union[ + _DeleteMemoryRequestParameters, _DeleteMemoryRequestParametersDict ] -class MemoryConfig(_common.BaseModel): - """Config for creating a Memory.""" +class DeleteMemoryOperation(_common.BaseModel): + """Operation for deleting 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( + name: Optional[str] = Field( default=None, - description="""Optional. Input only. The TTL for this resource. - - The expiration time is computed: now + TTL.""", + 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}`.""", ) - expire_time: Optional[datetime.datetime] = Field( + metadata: Optional[dict[str, Any]] = 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="""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_expire_time: Optional[datetime.datetime] = Field( + done: Optional[bool] = 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="""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_ttl: Optional[str] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", + description="""The error result of the operation in case of failure or cancellation.""", ) - disable_memory_revisions: Optional[bool] = Field( + + +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}`.""" + + 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.""" + + +DeleteMemoryOperationOrDict = Union[DeleteMemoryOperation, DeleteMemoryOperationDict] + + +class GenerateMemoriesRequestVertexSessionSource(_common.BaseModel): + """The vertex session source for generating memories.""" + + end_time: Optional[datetime.datetime] = 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. End time (exclusive) of the time range. If not set, the end time is unbounded.""", ) - metadata: Optional[dict[str, MemoryMetadataValue]] = Field( + session: Optional[str] = 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="""Required. The resource name of the Session to generate memories for. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}`""", ) - memory_id: Optional[str] = Field( + start_time: Optional[datetime.datetime] = 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. 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 MemoryConfigDict(TypedDict, total=False): - """Config for creating a Memory.""" +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.""" - display_name: Optional[str] - """The display name of the memory.""" + session: Optional[str] + """Required. The resource name of the Session to generate memories for. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}`""" - description: Optional[str] - """The description of the memory.""" + 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.""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" - ttl: Optional[str] - """Optional. Input only. The TTL for this resource. +GenerateMemoriesRequestVertexSessionSourceOrDict = Union[ + GenerateMemoriesRequestVertexSessionSource, + GenerateMemoriesRequestVertexSessionSourceDict, +] - 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.""" +class GenerateMemoriesRequestDirectContentsSourceEvent(_common.BaseModel): - 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.""" + content: Optional[genai_types.Content] = Field( + default=None, + description="""Required. A single piece of content from which to generate memories.""", + ) - revision_ttl: Optional[str] - """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" - disable_memory_revisions: Optional[bool] - """Optional. Input only. If true, no revision will be created for this request.""" - - topics: Optional[list[MemoryTopicIdDict]] - """Optional. The topics of the memory.""" - - metadata: Optional[dict[str, MemoryMetadataValueDict]] - """Optional. User-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""" +class GenerateMemoriesRequestDirectContentsSourceEventDict(TypedDict, total=False): - 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.""" + content: Optional[genai_types.Content] + """Required. A single piece of content from which to generate memories.""" -MemoryConfigOrDict = Union[MemoryConfig, MemoryConfigDict] +GenerateMemoriesRequestDirectContentsSourceEventOrDict = Union[ + GenerateMemoriesRequestDirectContentsSourceEvent, + GenerateMemoriesRequestDirectContentsSourceEventDict, +] -class _CreateMemoryRequestParameters(_common.BaseModel): - """Parameters for creating Memories.""" +class GenerateMemoriesRequestDirectContentsSource(_common.BaseModel): + """The direct contents source for generating memories.""" - name: Optional[str] = Field( + events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEvent]] = Field( default=None, - description="""Name of the memory bank to create the memory under.""", + description="""Required. The source content (i.e. chat history) to generate memories from.""", ) - fact: Optional[str] = Field( - default=None, - description="""The fact of the memory. - This is the semantic knowledge extracted from the source content).""", - ) - scope: Optional[dict[str, str]] = Field( - default=None, - description="""The scope of the memory. - Memories are isolated within their scope. The scope is defined when - creating or generating memories. Up to 5 key-value pairs are accepted, - and scope values cannot contain the wildcard character '*'.""", - ) - config: Optional[MemoryConfig] = Field(default=None, description="""""") +class GenerateMemoriesRequestDirectContentsSourceDict(TypedDict, total=False): + """The direct contents source for generating memories.""" + events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEventDict]] + """Required. The source content (i.e. chat history) to generate memories from.""" -class _CreateMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for creating Memories.""" - name: Optional[str] - """Name of the memory bank to create the memory under.""" +GenerateMemoriesRequestDirectContentsSourceOrDict = Union[ + GenerateMemoriesRequestDirectContentsSource, + GenerateMemoriesRequestDirectContentsSourceDict, +] - fact: Optional[str] - """The fact of the memory. - This is the semantic knowledge extracted from the source content).""" +class GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(_common.BaseModel): + """A direct memory to upload to Memory Bank.""" - scope: Optional[dict[str, str]] - """The scope of the memory. + fact: Optional[str] = Field( + default=None, + description="""Required. The fact to consolidate with existing memories.""", + ) + topics: Optional[list[MemoryTopicId]] = Field( + default=None, + description="""Optional. The topics that the consolidated memories should be associated with.""", + ) - 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[MemoryConfigDict] - """""" +class GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict( + TypedDict, total=False +): + """A direct memory to upload to Memory Bank.""" + fact: Optional[str] + """Required. The fact to consolidate with existing memories.""" -_CreateMemoryRequestParametersOrDict = Union[ - _CreateMemoryRequestParameters, _CreateMemoryRequestParametersDict + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The topics that the consolidated memories should be associated with.""" + + +GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryOrDict = Union[ + GenerateMemoriesRequestDirectMemoriesSourceDirectMemory, + GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict, ] -class MemoryStructuredContent(_common.BaseModel): - """Represents the structured value of the memory.""" +class GenerateMemoriesRequestDirectMemoriesSource(_common.BaseModel): + """The direct memories source for generating memories.""" - data: Optional[dict[str, Any]] = Field( - default=None, - description="""Required. Represents the structured value of the memory.""", - ) - schema_id: Optional[str] = Field( + direct_memories: Optional[ + list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemory] + ] = Field( default=None, - description="""Required. Represents the schema ID for which this structured memory belongs to.""", + description="""Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""", ) -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.""" +class GenerateMemoriesRequestDirectMemoriesSourceDict(TypedDict, total=False): + """The direct memories source for generating memories.""" - schema_id: Optional[str] - """Required. Represents the schema ID for which this structured memory belongs to.""" + direct_memories: Optional[ + list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict] + ] + """Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""" -MemoryStructuredContentOrDict = Union[ - MemoryStructuredContent, MemoryStructuredContentDict +GenerateMemoriesRequestDirectMemoriesSourceOrDict = Union[ + GenerateMemoriesRequestDirectMemoriesSource, + GenerateMemoriesRequestDirectMemoriesSourceDict, ] -class Memory(_common.BaseModel): - """A memory.""" +class GenerateMemoriesConfig(_common.BaseModel): + """Config 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - fact: Optional[str] = Field( + disable_consolidation: Optional[bool] = Field( default=None, - description="""Optional. Represents semantic knowledge extracted from the source content.""", + description="""Whether to disable consolidation of memories. + + 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.""", ) - 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.""", + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) - name: Optional[str] = Field( + revision_labels: Optional[dict[str, str]] = Field( default=None, - description="""Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""", + 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, - 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.""", + 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. Represents the TTL for the revision. The expiration time is computed: now + TTL.""", - ) - scope: Optional[dict[str, str]] = 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 '*'.""", - ) - topics: Optional[list[MemoryTopicId]] = Field( - default=None, description="""Optional. Represents the Topics of the Memory.""" + description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", ) - ttl: Optional[str] = Field( + disable_memory_revisions: Optional[bool] = Field( default=None, - description="""Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""", + description="""Optional. Input only. If true, no revisions will be created for this request.""", ) - update_time: Optional[datetime.datetime] = Field( + metadata: Optional[dict[str, MemoryMetadataValue]] = Field( default=None, - description="""Output only. Represents the timestamp when this Memory was most recently updated.""", + description="""Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""", ) - memory_type: Optional[MemoryType] = Field( + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] = 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.""", + description="""Optional. The strategy to use when applying metadata to existing memories.""", ) - structured_content: Optional[MemoryStructuredContent] = Field( + allowed_topics: Optional[list[MemoryTopicId]] = Field( default=None, - description="""Optional. Represents the structured content of the memory.""", + description="""Optional. Restricts memory generation to a subset of memory topics.""", ) -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.""" +class GenerateMemoriesConfigDict(TypedDict, total=False): + """Config for generating memories.""" - display_name: Optional[str] - """Optional. Represents the display name of the Memory.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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.""" + disable_consolidation: Optional[bool] + """Whether to disable consolidation of memories. - fact: Optional[str] - """Optional. Represents semantic knowledge extracted from the source content.""" + 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.""" - 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.""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" - name: Optional[str] - """Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""" + 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. 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.""" + """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. Represents the TTL for the revision. The expiration time is computed: now + TTL.""" + """Optional. Input only. 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 '*'.""" + disable_memory_revisions: Optional[bool] + """Optional. Input only. If true, no revisions will be created for this request.""" - topics: Optional[list[MemoryTopicIdDict]] - """Optional. Represents the Topics of the Memory.""" + metadata: Optional[dict[str, MemoryMetadataValueDict]] + """Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""" - ttl: Optional[str] - """Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""" + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] + """Optional. The strategy to use when applying metadata to existing memories.""" - update_time: Optional[datetime.datetime] - """Output only. Represents the timestamp when this Memory was most recently updated.""" + allowed_topics: Optional[list[MemoryTopicIdDict]] + """Optional. Restricts memory generation to a subset of memory topics.""" - 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.""" +GenerateMemoriesConfigOrDict = Union[GenerateMemoriesConfig, GenerateMemoriesConfigDict] -MemoryOrDict = Union[Memory, MemoryDict] - - -class MemoryOperation(_common.BaseModel): - """Operation that has a memory as a response.""" +class _GenerateMemoriesRequestParameters(_common.BaseModel): + """Parameters for generating 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}`.""", + description="""Name of the Memory Bank to generate memories with.""", ) - metadata: Optional[dict[str, Any]] = Field( + vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSource] = 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="""The vertex session source of the memories that should be generated.""", ) - 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.""", + direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSource] = ( + Field( + default=None, + description="""The direct contents source of the memories that should be generated.""", + ) ) - error: Optional[dict[str, Any]] = Field( + 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 error result of the operation in case of failure or cancellation.""", + description="""The scope of the memories that should be generated. + + 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 '*'.""", ) - response: Optional[Memory] = Field(default=None, description="""The Memory.""") + config: Optional[GenerateMemoriesConfig] = Field(default=None, description="""""") -class MemoryOperationDict(TypedDict, total=False): - """Operation that has a memory as a response.""" +class _GenerateMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for generating 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}`.""" + """Name of the Memory Bank to generate memories with.""" - 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.""" + vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSourceDict] + """The vertex session source of the memories that should be generated.""" - 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.""" + direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSourceDict] + """The direct contents source of the memories that should be generated.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSourceDict] + """The direct memories source of the memories that should be generated.""" - response: Optional[MemoryDict] - """The Memory.""" + scope: Optional[dict[str, str]] + """The scope of the memories that should be generated. + + 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[GenerateMemoriesConfigDict] + """""" -MemoryOperationOrDict = Union[MemoryOperation, MemoryOperationDict] +_GenerateMemoriesRequestParametersOrDict = Union[ + _GenerateMemoriesRequestParameters, _GenerateMemoriesRequestParametersDict +] -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 GenerateMemoriesResponseGeneratedMemory(_common.BaseModel): + """A memory that was generated.""" + + memory: Optional[Memory] = Field( + default=None, description="""The generated memory.""" + ) + action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] = Field( + default=None, description="""The action to take.""" + ) + previous_revision: Optional[str] = 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}`""", ) -class DeleteMemoryConfigDict(TypedDict, total=False): - """Config for deleting a Memory.""" +class GenerateMemoriesResponseGeneratedMemoryDict(TypedDict, total=False): + """A memory that was generated.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + memory: Optional[MemoryDict] + """The generated memory.""" + action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] + """The action to take.""" -DeleteMemoryConfigOrDict = Union[DeleteMemoryConfig, DeleteMemoryConfigDict] + 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}`""" -class _DeleteMemoryRequestParameters(_common.BaseModel): - """Parameters for deleting memories.""" +GenerateMemoriesResponseGeneratedMemoryOrDict = Union[ + GenerateMemoriesResponseGeneratedMemory, GenerateMemoriesResponseGeneratedMemoryDict +] - name: Optional[str] = Field( - default=None, description="""Name of the memory to delete.""" - ) - config: Optional[DeleteMemoryConfig] = Field(default=None, description="""""") +class GenerateMemoriesResponse(_common.BaseModel): + """The response for generating memories.""" -class _DeleteMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for deleting memories.""" + generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemory]] = Field( + default=None, description="""The generated memories.""" + ) - name: Optional[str] - """Name of the memory to delete.""" - config: Optional[DeleteMemoryConfigDict] - """""" +class GenerateMemoriesResponseDict(TypedDict, total=False): + """The response for generating memories.""" + + generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemoryDict]] + """The generated memories.""" -_DeleteMemoryRequestParametersOrDict = Union[ - _DeleteMemoryRequestParameters, _DeleteMemoryRequestParametersDict +GenerateMemoriesResponseOrDict = Union[ + GenerateMemoriesResponse, GenerateMemoriesResponseDict ] -class DeleteMemoryOperation(_common.BaseModel): - """Operation for deleting memories.""" +class GenerateMemoriesOperation(_common.BaseModel): + """Operation that generates memories with a Memory Bank.""" name: Optional[str] = Field( default=None, @@ -12080,10 +11701,13 @@ class DeleteMemoryOperation(_common.BaseModel): default=None, description="""The error result of the operation in case of failure or cancellation.""", ) + response: Optional[GenerateMemoriesResponse] = Field( + default=None, description="""The response for generating memories.""" + ) -class DeleteMemoryOperationDict(TypedDict, total=False): - """Operation for deleting 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}`.""" @@ -12097,3667 +11721,3666 @@ class DeleteMemoryOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" + response: Optional[GenerateMemoriesResponseDict] + """The response for generating memories.""" -DeleteMemoryOperationOrDict = Union[DeleteMemoryOperation, DeleteMemoryOperationDict] +GenerateMemoriesOperationOrDict = Union[ + GenerateMemoriesOperation, GenerateMemoriesOperationDict +] -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 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 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 GetMemoryConfigDict(TypedDict, total=False): + """Config for getting a Memory.""" - session: Optional[str] - """Required. The resource name of the Session to generate memories for. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}`""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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.""" +GetMemoryConfigOrDict = Union[GetMemoryConfig, GetMemoryConfigDict] -GenerateMemoriesRequestVertexSessionSourceOrDict = Union[ - GenerateMemoriesRequestVertexSessionSource, - GenerateMemoriesRequestVertexSessionSourceDict, -] +class _GetMemoryRequestParameters(_common.BaseModel): + """Parameters for getting a Memory.""" -class GenerateMemoriesRequestDirectContentsSourceEvent(_common.BaseModel): + name: Optional[str] = Field(default=None, description="""Name of the memory.""") + config: Optional[GetMemoryConfig] = 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 _GetMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for getting a Memory.""" -class GenerateMemoriesRequestDirectContentsSourceEventDict(TypedDict, total=False): + name: Optional[str] + """Name of the memory.""" - content: Optional[genai_types.Content] - """Required. A single piece of content from which to generate memories.""" + config: Optional[GetMemoryConfigDict] + """""" -GenerateMemoriesRequestDirectContentsSourceEventOrDict = Union[ - GenerateMemoriesRequestDirectContentsSourceEvent, - GenerateMemoriesRequestDirectContentsSourceEventDict, +_GetMemoryRequestParametersOrDict = Union[ + _GetMemoryRequestParameters, _GetMemoryRequestParametersDict ] -class GenerateMemoriesRequestDirectContentsSource(_common.BaseModel): - """The direct contents source for generating memories.""" +class ListMemoriesConfig(_common.BaseModel): + """Config for listing memories.""" - events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEvent]] = Field( + 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="""Required. The source content (i.e. chat history) to generate memories from.""", + description="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""", ) + order_by: Optional[str] = Field( + default=None, + description="""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 GenerateMemoriesRequestDirectContentsSourceDict(TypedDict, total=False): - """The direct contents source for generating memories.""" - - events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEventDict]] - """Required. The source content (i.e. chat history) to generate memories from.""" +class ListMemoriesConfigDict(TypedDict, total=False): + """Config for listing memories.""" -GenerateMemoriesRequestDirectContentsSourceOrDict = Union[ - GenerateMemoriesRequestDirectContentsSource, - GenerateMemoriesRequestDirectContentsSourceDict, -] + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + page_size: Optional[int] + """""" -class GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(_common.BaseModel): - """A direct memory to upload to Memory Bank.""" + page_token: Optional[str] + """""" - fact: Optional[str] = Field( - default=None, - description="""Required. The fact to consolidate with existing memories.""", - ) - topics: Optional[list[MemoryTopicId]] = Field( - default=None, - description="""Optional. The topics that the consolidated memories should be associated with.""", + 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`""" + + +ListMemoriesConfigOrDict = Union[ListMemoriesConfig, ListMemoriesConfigDict] + + +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 GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict( - TypedDict, total=False -): - """A direct memory to upload to Memory Bank.""" +class _ListMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for listing memories.""" - fact: Optional[str] - """Required. The fact to consolidate with existing memories.""" + name: Optional[str] + """Name of the Memory Bank.""" - topics: Optional[list[MemoryTopicIdDict]] - """Optional. The topics that the consolidated memories should be associated with.""" + config: Optional[ListMemoriesConfigDict] + """""" -GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryOrDict = Union[ - GenerateMemoriesRequestDirectMemoriesSourceDirectMemory, - GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict, +_ListMemoriesRequestParametersOrDict = Union[ + _ListMemoriesRequestParameters, _ListMemoriesRequestParametersDict ] -class GenerateMemoriesRequestDirectMemoriesSource(_common.BaseModel): - """The direct memories source for generating memories.""" +class ListMemoriesResponse(_common.BaseModel): + """Response for listing 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.""", + 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 GenerateMemoriesRequestDirectMemoriesSourceDict(TypedDict, total=False): - """The direct memories source for generating memories.""" +class ListMemoriesResponseDict(TypedDict, total=False): + """Response for listing memories.""" - direct_memories: Optional[ - list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict] - ] - """Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + next_page_token: Optional[str] + """""" -GenerateMemoriesRequestDirectMemoriesSourceOrDict = Union[ - GenerateMemoriesRequestDirectMemoriesSource, - GenerateMemoriesRequestDirectMemoriesSourceDict, -] + memories: Optional[list[MemoryDict]] + """List of memories.""" -class GenerateMemoriesConfig(_common.BaseModel): - """Config for generating memories.""" +ListMemoriesResponseOrDict = Union[ListMemoriesResponse, ListMemoriesResponseDict] - 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. - 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( - 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 revisions will be created for this request.""", - ) - metadata: Optional[dict[str, MemoryMetadataValue]] = Field( - default=None, - description="""Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""", - ) - metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] = Field( - default=None, - description="""Optional. The strategy to use when applying metadata to existing memories.""", +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.""" ) - allowed_topics: Optional[list[MemoryTopicId]] = Field( - default=None, - description="""Optional. Restricts memory generation to a subset of memory topics.""", + config: Optional[GetMemoryBankOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -class GenerateMemoriesConfigDict(TypedDict, total=False): - """Config for generating memories.""" +class _GetMemoryOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a memory as a response.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - disable_consolidation: Optional[bool] - """Whether to disable consolidation of memories. + config: Optional[GetMemoryBankOperationConfigDict] + """Used to override the default configuration.""" - 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.""" +_GetMemoryOperationParametersOrDict = Union[ + _GetMemoryOperationParameters, _GetMemoryOperationParametersDict +] - 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.""" +class _GetGenerateMemoriesOperationParameters(_common.BaseModel): + """Parameters for getting an operation with generated memories as a response.""" - revision_ttl: Optional[str] - """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" + 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.""" + ) - disable_memory_revisions: Optional[bool] - """Optional. Input only. If true, no revisions will be created for this request.""" - metadata: Optional[dict[str, MemoryMetadataValueDict]] - """Optional. User-provided metadata for the generated memories. This is not generated by Memory Bank.""" +class _GetGenerateMemoriesOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with generated memories as a response.""" - metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] - """Optional. The strategy to use when applying metadata to existing memories.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - allowed_topics: Optional[list[MemoryTopicIdDict]] - """Optional. Restricts memory generation to a subset of memory topics.""" + config: Optional[GetMemoryBankOperationConfigDict] + """Used to override the default configuration.""" -GenerateMemoriesConfigOrDict = Union[GenerateMemoriesConfig, GenerateMemoriesConfigDict] +_GetGenerateMemoriesOperationParametersOrDict = Union[ + _GetGenerateMemoriesOperationParameters, _GetGenerateMemoriesOperationParametersDict +] -class _GenerateMemoriesRequestParameters(_common.BaseModel): - """Parameters for generating memories.""" +class RetrieveMemoriesRequestSimilaritySearchParams(_common.BaseModel): + """The parameters for semantic similarity search based retrieval.""" - name: Optional[str] = Field( + search_query: Optional[str] = Field( default=None, - description="""Name of the Memory Bank to generate memories with.""", + description="""Required. Query to use for similarity search retrieval. If provided, then the parent ReasoningEngine must have ReasoningEngineContextSpec.MemoryBankConfig.SimilaritySearchConfig set.""", ) - vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSource] = Field( + top_k: Optional[int] = 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.""", - ) - ) - direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSource] = ( - Field( - default=None, - description="""The direct memories source of the memories that should be generated.""", - ) + 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.""", ) - scope: Optional[dict[str, str]] = Field( - default=None, - description="""The scope of the memories that should be generated. - 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[GenerateMemoriesConfig] = Field(default=None, description="""""") +class RetrieveMemoriesRequestSimilaritySearchParamsDict(TypedDict, total=False): + """The parameters for semantic similarity search based retrieval.""" -class _GenerateMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for generating memories.""" + 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] - """Name of the Memory Bank to generate memories with.""" + 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.""" - vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSourceDict] - """The vertex session source of the memories that should be generated.""" - direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSourceDict] - """The direct contents source of the memories that should be generated.""" +RetrieveMemoriesRequestSimilaritySearchParamsOrDict = Union[ + RetrieveMemoriesRequestSimilaritySearchParams, + RetrieveMemoriesRequestSimilaritySearchParamsDict, +] - 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 memories that should be generated. +class RetrieveMemoriesRequestSimpleRetrievalParams(_common.BaseModel): + """The parameters for simple (non-similarity search) retrieval.""" - 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 '*'.""" + 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( + default=None, + description="""Optional. A page token, received from a previous `RetrieveMemories` call. Provide this to retrieve the subsequent page.""", + ) - config: Optional[GenerateMemoriesConfigDict] - """""" +class RetrieveMemoriesRequestSimpleRetrievalParamsDict(TypedDict, total=False): + """The parameters for simple (non-similarity search) retrieval.""" -_GenerateMemoriesRequestParametersOrDict = Union[ - _GenerateMemoriesRequestParameters, _GenerateMemoriesRequestParametersDict + 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 GenerateMemoriesResponseGeneratedMemory(_common.BaseModel): - """A memory that was generated.""" +class MemoryFilter(_common.BaseModel): + """Filter to apply when retrieving memories.""" - memory: Optional[Memory] = Field( - default=None, description="""The generated memory.""" + 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".""", ) - action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] = Field( - default=None, description="""The action to take.""" + negate: Optional[bool] = Field( + default=None, description="""Indicates whether the filter will be negated.""" ) - previous_revision: Optional[str] = Field( + op: Optional[Operator] = 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="""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.""" ) -class GenerateMemoriesResponseGeneratedMemoryDict(TypedDict, total=False): - """A memory that was generated.""" +class MemoryFilterDict(TypedDict, total=False): + """Filter to apply when retrieving memories.""" - memory: Optional[MemoryDict] - """The generated memory.""" + key: Optional[str] + """Represents the key of the filter. For example, "author" would apply to `metadata` entries with the key "author".""" - action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] - """The action to take.""" + negate: Optional[bool] + """Indicates whether the filter will be negated.""" - 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}`""" + op: Optional[Operator] + """Represents the operator to apply to the filter. If not set, then EQUAL will be used.""" + value: Optional[MemoryMetadataValueDict] + """Represents the value to compare to.""" -GenerateMemoriesResponseGeneratedMemoryOrDict = Union[ - GenerateMemoriesResponseGeneratedMemory, GenerateMemoriesResponseGeneratedMemoryDict -] +MemoryFilterOrDict = Union[MemoryFilter, MemoryFilterDict] -class GenerateMemoriesResponse(_common.BaseModel): - """The response for generating memories.""" - generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemory]] = Field( - default=None, description="""The generated memories.""" +class MemoryConjunctionFilter(_common.BaseModel): + """The conjunction filter for memories.""" + + filters: Optional[list[MemoryFilter]] = Field( + default=None, + description="""Represents filters that will be combined using AND logic.""", ) -class GenerateMemoriesResponseDict(TypedDict, total=False): - """The response for generating memories.""" +class MemoryConjunctionFilterDict(TypedDict, total=False): + """The conjunction filter for memories.""" - generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemoryDict]] - """The generated memories.""" + filters: Optional[list[MemoryFilterDict]] + """Represents filters that will be combined using AND logic.""" -GenerateMemoriesResponseOrDict = Union[ - GenerateMemoriesResponse, GenerateMemoriesResponseDict +MemoryConjunctionFilterOrDict = Union[ + MemoryConjunctionFilter, MemoryConjunctionFilterDict ] -class GenerateMemoriesOperation(_common.BaseModel): - """Operation that generates memories with a Memory Bank.""" +class RetrieveMemoriesConfig(_common.BaseModel): + """Config 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}`.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - metadata: Optional[dict[str, Any]] = Field( + filter: 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="""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` + """, ) - done: Optional[bool] = Field( + filter_groups: Optional[list[MemoryConjunctionFilter]] = 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="""Metadata filters that will be applied to the retrieved 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"))`. + """, ) - error: Optional[dict[str, Any]] = Field( + memory_types: Optional[list[MemoryType]] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", - ) - response: Optional[GenerateMemoriesResponse] = Field( - default=None, description="""The response for generating memories.""" + 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.""", ) -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}`.""" +class RetrieveMemoriesConfigDict(TypedDict, total=False): + """Config for retrieving 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.""" + 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.""" + 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). - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + Supported fields: + * `fact` + * `create_time` + * `update_time` + """ - response: Optional[GenerateMemoriesResponseDict] - """The response for generating memories.""" + 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). + 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}]}]` -GenerateMemoriesOperationOrDict = Union[ - GenerateMemoriesOperation, GenerateMemoriesOperationDict -] + 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]] + """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.""" -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.""" - ) +RetrieveMemoriesConfigOrDict = Union[RetrieveMemoriesConfig, RetrieveMemoriesConfigDict] -class GetMemoryConfigDict(TypedDict, total=False): - """Config for getting a Memory.""" +class _RetrieveMemoriesRequestParameters(_common.BaseModel): + """Parameters for retrieving memories.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + 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. + 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.""", + ) + ) + config: Optional[RetrieveMemoriesConfig] = Field(default=None, description="""""") -GetMemoryConfigOrDict = Union[GetMemoryConfig, GetMemoryConfigDict] +class _RetrieveMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for retrieving memories.""" -class _GetMemoryRequestParameters(_common.BaseModel): - """Parameters for getting a Memory.""" + name: Optional[str] + """Name of the Memory Bank to retrieve memories from.""" - name: Optional[str] = Field(default=None, description="""Name of the memory.""") - config: Optional[GetMemoryConfig] = Field(default=None, description="""""") + 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.""" -class _GetMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for getting a Memory.""" + similarity_search_params: Optional[ + RetrieveMemoriesRequestSimilaritySearchParamsDict + ] + """Parameters for semantic similarity search based retrieval.""" - name: Optional[str] - """Name of the memory.""" + simple_retrieval_params: Optional[RetrieveMemoriesRequestSimpleRetrievalParamsDict] + """Parameters for simple (non-similarity search) retrieval.""" - config: Optional[GetMemoryConfigDict] + config: Optional[RetrieveMemoriesConfigDict] """""" -_GetMemoryRequestParametersOrDict = Union[ - _GetMemoryRequestParameters, _GetMemoryRequestParametersDict +_RetrieveMemoriesRequestParametersOrDict = Union[ + _RetrieveMemoriesRequestParameters, _RetrieveMemoriesRequestParametersDict ] -class ListMemoriesConfig(_common.BaseModel): - """Config for listing memories.""" +class RetrieveMemoriesResponseRetrievedMemory(_common.BaseModel): + """A retrieved memory.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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.""", ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( - default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""", - ) - order_by: Optional[str] = Field( - default=None, - description="""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`""", + memory: Optional[Memory] = Field( + default=None, description="""The retrieved Memory.""" ) -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] - """""" - - 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.""" +class RetrieveMemoriesResponseRetrievedMemoryDict(TypedDict, total=False): + """A retrieved memory.""" - 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). + 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.""" - Supported fields: - * `create_time` - * `update_time`""" + memory: Optional[MemoryDict] + """The retrieved Memory.""" -ListMemoriesConfigOrDict = Union[ListMemoriesConfig, ListMemoriesConfigDict] +RetrieveMemoriesResponseRetrievedMemoryOrDict = Union[ + RetrieveMemoriesResponseRetrievedMemory, RetrieveMemoriesResponseRetrievedMemoryDict +] -class _ListMemoriesRequestParameters(_common.BaseModel): - """Parameters for listing memories.""" +class RetrieveMemoriesResponse(_common.BaseModel): + """The response for retrieving memories.""" - name: Optional[str] = Field( - default=None, description="""Name of the Memory Bank.""" + next_page_token: Optional[str] = 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.""", + ) + retrieved_memories: Optional[list[RetrieveMemoriesResponseRetrievedMemory]] = Field( + default=None, description="""The retrieved memories.""" ) - config: Optional[ListMemoriesConfig] = Field(default=None, description="""""") -class _ListMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for listing memories.""" +class RetrieveMemoriesResponseDict(TypedDict, total=False): + """The response for retrieving memories.""" - name: Optional[str] - """Name of the 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.""" - config: Optional[ListMemoriesConfigDict] - """""" + retrieved_memories: Optional[list[RetrieveMemoriesResponseRetrievedMemoryDict]] + """The retrieved memories.""" -_ListMemoriesRequestParametersOrDict = Union[ - _ListMemoriesRequestParameters, _ListMemoriesRequestParametersDict +RetrieveMemoriesResponseOrDict = Union[ + RetrieveMemoriesResponse, RetrieveMemoriesResponseDict ] -class ListMemoriesResponse(_common.BaseModel): - """Response for listing memories.""" +class RetrieveMemoryProfilesConfig(_common.BaseModel): + """Config 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.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class ListMemoriesResponseDict(TypedDict, total=False): - """Response for listing memories.""" - - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" +class RetrieveMemoryProfilesConfigDict(TypedDict, total=False): + """Config for retrieving memory profiles.""" - next_page_token: Optional[str] - """""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - memories: Optional[list[MemoryDict]] - """List of memories.""" +RetrieveMemoryProfilesConfigOrDict = Union[ + RetrieveMemoryProfilesConfig, RetrieveMemoryProfilesConfigDict +] -ListMemoriesResponseOrDict = Union[ListMemoriesResponse, ListMemoriesResponseDict] +class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): + """Parameters for retrieving memory profiles.""" -class _GetMemoryOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a memory as a response.""" + 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. - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + 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[GetMemoryBankOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + config: Optional[RetrieveMemoryProfilesConfig] = Field( + default=None, description="""""" ) -class _GetMemoryOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a memory as a response.""" +class _RetrieveMemoryProfilesRequestParametersDict(TypedDict, total=False): + """Parameters for retrieving memory profiles.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + name: Optional[str] + """Name of the Memory Bank to retrieve memory profiles from.""" - config: Optional[GetMemoryBankOperationConfigDict] - """Used to override the default configuration.""" + 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.""" + + config: Optional[RetrieveMemoryProfilesConfigDict] + """""" -_GetMemoryOperationParametersOrDict = Union[ - _GetMemoryOperationParameters, _GetMemoryOperationParametersDict +_RetrieveMemoryProfilesRequestParametersOrDict = Union[ + _RetrieveMemoryProfilesRequestParameters, + _RetrieveMemoryProfilesRequestParametersDict, ] -class _GetGenerateMemoriesOperationParameters(_common.BaseModel): - """Parameters for getting an operation with generated memories as a response.""" +class MemoryProfile(_common.BaseModel): + """A memory profile.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + 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.""", ) - config: Optional[GetMemoryBankOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + profile: Optional[dict[str, Any]] = Field( + default=None, description="""Represents the profile data.""" ) -class _GetGenerateMemoriesOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with generated memories as a response.""" +class MemoryProfileDict(TypedDict, total=False): + """A memory profile.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + 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[GetMemoryBankOperationConfigDict] - """Used to override the default configuration.""" + profile: Optional[dict[str, Any]] + """Represents the profile data.""" -_GetGenerateMemoriesOperationParametersOrDict = Union[ - _GetGenerateMemoriesOperationParameters, _GetGenerateMemoriesOperationParametersDict -] +MemoryProfileOrDict = Union[MemoryProfile, MemoryProfileDict] -class RetrieveMemoriesRequestSimilaritySearchParams(_common.BaseModel): - """The parameters for semantic similarity search based retrieval.""" +class RetrieveProfilesResponse(_common.BaseModel): + """The response for retrieving memory profiles.""" - search_query: 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.""", - ) - top_k: Optional[int] = Field( + profiles: Optional[dict[str, MemoryProfile]] = 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="""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 RetrieveMemoriesRequestSimilaritySearchParamsDict(TypedDict, total=False): - """The parameters for semantic similarity search based retrieval.""" - - search_query: Optional[str] - """Required. Query to use for similarity search retrieval. If provided, then the parent ReasoningEngine must have ReasoningEngineContextSpec.MemoryBankConfig.SimilaritySearchConfig set.""" +class RetrieveProfilesResponseDict(TypedDict, total=False): + """The response for retrieving memory profiles.""" - 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.""" + 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`.""" -RetrieveMemoriesRequestSimilaritySearchParamsOrDict = Union[ - RetrieveMemoriesRequestSimilaritySearchParams, - RetrieveMemoriesRequestSimilaritySearchParamsDict, +RetrieveProfilesResponseOrDict = Union[ + RetrieveProfilesResponse, RetrieveProfilesResponseDict ] -class RetrieveMemoriesRequestSimpleRetrievalParams(_common.BaseModel): - """The parameters for simple (non-similarity search) retrieval.""" +class RollbackMemoryConfig(_common.BaseModel): + """Config for rolling back 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - page_token: Optional[str] = Field( - default=None, - description="""Optional. A page token, received from a previous `RetrieveMemories` call. Provide this to retrieve the subsequent page.""", + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) -class RetrieveMemoriesRequestSimpleRetrievalParamsDict(TypedDict, total=False): - """The parameters for simple (non-similarity search) retrieval.""" +class RollbackMemoryConfigDict(TypedDict, total=False): + """Config for rolling back 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.""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" -RetrieveMemoriesRequestSimpleRetrievalParamsOrDict = Union[ - RetrieveMemoriesRequestSimpleRetrievalParams, - RetrieveMemoriesRequestSimpleRetrievalParamsDict, -] +RollbackMemoryConfigOrDict = Union[RollbackMemoryConfig, RollbackMemoryConfigDict] -class MemoryFilter(_common.BaseModel): - """Filter to apply when retrieving memories.""" +class _RollbackMemoryRequestParameters(_common.BaseModel): + """Parameters for generating memories.""" - 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.""", + name: Optional[str] = Field( + default=None, description="""Name of the memory to rollback.""" ) - value: Optional[MemoryMetadataValue] = Field( - default=None, description="""Represents the value to compare to.""" + 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 MemoryFilterDict(TypedDict, total=False): - """Filter to apply when retrieving memories.""" - - key: Optional[str] - """Represents the key of the filter. For example, "author" would apply to `metadata` entries with the key "author".""" +class _RollbackMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for generating memories.""" - negate: Optional[bool] - """Indicates whether the filter will be negated.""" + name: Optional[str] + """Name of the memory to rollback.""" - op: Optional[Operator] - """Represents the operator to apply to the filter. If not set, then EQUAL will be used.""" + target_revision_id: Optional[str] + """The ID of the revision to rollback to.""" - value: Optional[MemoryMetadataValueDict] - """Represents the value to compare to.""" + config: Optional[RollbackMemoryConfigDict] + """""" -MemoryFilterOrDict = Union[MemoryFilter, MemoryFilterDict] +_RollbackMemoryRequestParametersOrDict = Union[ + _RollbackMemoryRequestParameters, _RollbackMemoryRequestParametersDict +] -class MemoryConjunctionFilter(_common.BaseModel): - """The conjunction filter for memories.""" +class RollbackMemoryOperation(_common.BaseModel): + """Operation that rolls back a memory.""" - filters: Optional[list[MemoryFilter]] = Field( + name: Optional[str] = Field( default=None, - description="""Represents filters that will be combined using AND logic.""", + 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 MemoryConjunctionFilterDict(TypedDict, total=False): - """The conjunction filter for memories.""" +class RollbackMemoryOperationDict(TypedDict, total=False): + """Operation that rolls back a memory.""" - filters: Optional[list[MemoryFilterDict]] - """Represents filters that will be combined using AND logic.""" + 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.""" -MemoryConjunctionFilterOrDict = Union[ - MemoryConjunctionFilter, MemoryConjunctionFilterDict + 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.""" + + +RollbackMemoryOperationOrDict = Union[ + RollbackMemoryOperation, RollbackMemoryOperationDict ] -class RetrieveMemoriesConfig(_common.BaseModel): - """Config for retrieving memories.""" +class UpdateMemoryConfig(_common.BaseModel): + """Config for updating a memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - filter: 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="""The standard list filter that will be applied to the retrieved - memories. More detail in [AIP-160](https://google.aip.dev/160). + description="""Optional. Input only. The TTL for this resource. - Supported fields: - * `fact` - * `create_time` - * `update_time` - """, + The expiration time is computed: now + TTL.""", ) - filter_groups: Optional[list[MemoryConjunctionFilter]] = Field( + expire_time: Optional[datetime.datetime] = Field( default=None, - 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). - - 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="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", ) - memory_types: Optional[list[MemoryType]] = Field( + revision_expire_time: Optional[datetime.datetime] = 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="""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 RetrieveMemoriesConfigDict(TypedDict, total=False): - """Config for retrieving memories.""" +class UpdateMemoryConfigDict(TypedDict, total=False): + """Config for updating a memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - 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). + display_name: Optional[str] + """The display name of the memory.""" - Supported fields: - * `fact` - * `create_time` - * `update_time` - """ + description: Optional[str] + """The description of the memory.""" - 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). + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" - 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}]}]` + ttl: Optional[str] + """Optional. Input only. The TTL for this resource. - would be equivalent to the logical expression: - `(metadata.author = "agent 123" OR (metadata.label = "travel" AND - metadata.author = "agent 321"))`. - """ + The expiration time is computed: now + TTL.""" - 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.""" + 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.""" -RetrieveMemoriesConfigOrDict = Union[RetrieveMemoriesConfig, RetrieveMemoriesConfigDict] + 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 _RetrieveMemoriesRequestParameters(_common.BaseModel): - """Parameters for retrieving memories.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The topics of the memory.""" + + metadata: Optional[dict[str, MemoryMetadataValueDict]] + """Optional. User-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""" + + memory_id: Optional[str] + """Optional. The user defined ID to use for memory, which will become the final component of the memory resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""" + + update_mask: Optional[str] + """The update mask to apply. For the `FieldMask` definition, see + https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" + + +UpdateMemoryConfigOrDict = Union[UpdateMemoryConfig, UpdateMemoryConfigDict] + + +class _UpdateMemoryRequestParameters(_common.BaseModel): + """Parameters for updating memories.""" name: Optional[str] = Field( - default=None, - description="""Name of the Memory Bank to retrieve memories from.""", + default=None, description="""Name of the memory to update.""" ) - scope: Optional[dict[str, str]] = Field( + fact: Optional[str] = Field( default=None, - description="""The scope of the memories to retrieve. + description="""The updated fact of the memory. - 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.""", + This is the semantic knowledge extracted from the source content.""", ) - similarity_search_params: Optional[ - RetrieveMemoriesRequestSimilaritySearchParams - ] = Field( + scope: Optional[dict[str, str]] = 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.""", - ) + 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[RetrieveMemoriesConfig] = Field(default=None, description="""""") + config: Optional[UpdateMemoryConfig] = Field(default=None, description="""""") -class _RetrieveMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving memories.""" +class _UpdateMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for updating memories.""" name: Optional[str] - """Name of the Memory Bank to retrieve memories from.""" + """Name of the memory to update.""" - scope: Optional[dict[str, str]] - """The scope of the memories to retrieve. + fact: Optional[str] + """The updated fact of the memory. - 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.""" + This is the semantic knowledge extracted from the source content.""" - similarity_search_params: Optional[ - RetrieveMemoriesRequestSimilaritySearchParamsDict - ] - """Parameters for semantic similarity search based retrieval.""" + scope: Optional[dict[str, str]] + """The updated scope of the memory. - simple_retrieval_params: Optional[RetrieveMemoriesRequestSimpleRetrievalParamsDict] - """Parameters for simple (non-similarity search) retrieval.""" + 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[RetrieveMemoriesConfigDict] + config: Optional[UpdateMemoryConfigDict] """""" -_RetrieveMemoriesRequestParametersOrDict = Union[ - _RetrieveMemoriesRequestParameters, _RetrieveMemoriesRequestParametersDict +_UpdateMemoryRequestParametersOrDict = Union[ + _UpdateMemoryRequestParameters, _UpdateMemoryRequestParametersDict ] -class RetrieveMemoriesResponseRetrievedMemory(_common.BaseModel): - """A retrieved memory.""" +class PurgeMemoriesConfig(_common.BaseModel): + """Config for purging 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - memory: Optional[Memory] = Field( - default=None, description="""The retrieved Memory.""" + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) -class RetrieveMemoriesResponseRetrievedMemoryDict(TypedDict, total=False): - """A retrieved memory.""" +class PurgeMemoriesConfigDict(TypedDict, total=False): + """Config for purging 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.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - memory: Optional[MemoryDict] - """The retrieved Memory.""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" -RetrieveMemoriesResponseRetrievedMemoryOrDict = Union[ - RetrieveMemoriesResponseRetrievedMemory, RetrieveMemoriesResponseRetrievedMemoryDict -] +PurgeMemoriesConfigOrDict = Union[PurgeMemoriesConfig, PurgeMemoriesConfigDict] -class RetrieveMemoriesResponse(_common.BaseModel): - """The response for retrieving memories.""" +class _PurgeMemoriesRequestParameters(_common.BaseModel): + """Parameters for purging memories.""" - next_page_token: Optional[str] = Field( + name: Optional[str] = Field( + default=None, description="""Name of the Memory Bank to purge memories from.""" + ) + filter: Optional[str] = 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="""The standard list filter to determine which memories to purge. + More detail in [AIP-160](https://google.aip.dev/160).""", ) - retrieved_memories: Optional[list[RetrieveMemoriesResponseRetrievedMemory]] = Field( - default=None, description="""The retrieved memories.""" + filter_groups: Optional[list[MemoryConjunctionFilter]] = 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"))`. + """, + ) + force: Optional[bool] = Field( + default=None, + 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 RetrieveMemoriesResponseDict(TypedDict, total=False): - """The response for retrieving memories.""" +class _PurgeMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for purging memories.""" - 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] + """Name of the Memory Bank to purge memories from.""" - retrieved_memories: Optional[list[RetrieveMemoriesResponseRetrievedMemoryDict]] - """The retrieved memories.""" + filter: Optional[str] + """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 memories' + `metadata` using OR logic. Filters are defined using disjunctive normal + form (OR of ANDs). -RetrieveMemoriesResponseOrDict = Union[ - RetrieveMemoriesResponse, RetrieveMemoriesResponseDict + 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"))`. + """ + + 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] + """""" + + +_PurgeMemoriesRequestParametersOrDict = Union[ + _PurgeMemoriesRequestParameters, _PurgeMemoriesRequestParametersDict ] -class RetrieveMemoryProfilesConfig(_common.BaseModel): - """Config for retrieving memory profiles.""" +class PurgeMemoriesResponse(_common.BaseModel): + """The response for purging memories.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + purge_count: Optional[int] = Field( + default=None, description="""The number of memories that were purged.""" ) -class RetrieveMemoryProfilesConfigDict(TypedDict, total=False): - """Config for retrieving memory profiles.""" +class PurgeMemoriesResponseDict(TypedDict, total=False): + """The response for purging memories.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + purge_count: Optional[int] + """The number of memories that were purged.""" -RetrieveMemoryProfilesConfigOrDict = Union[ - RetrieveMemoryProfilesConfig, RetrieveMemoryProfilesConfigDict -] +PurgeMemoriesResponseOrDict = Union[PurgeMemoriesResponse, PurgeMemoriesResponseDict] -class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): - """Parameters for retrieving memory profiles.""" +class PurgeMemoriesOperation(_common.BaseModel): + """Operation that purges memories from a Memory Bank.""" name: Optional[str] = Field( default=None, - description="""Name of the Memory Bank to retrieve memory profiles from.""", + 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}`.""", ) - scope: Optional[dict[str, str]] = Field( + metadata: Optional[dict[str, Any]] = 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.""", + 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[RetrieveMemoryProfilesConfig] = Field( - default=None, description="""""" + 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[PurgeMemoriesResponse] = Field( + default=None, description="""The response for purging memories.""" ) -class _RetrieveMemoryProfilesRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving memory profiles.""" +class PurgeMemoriesOperationDict(TypedDict, total=False): + """Operation that purges memories from a Memory Bank.""" name: Optional[str] - """Name of the Memory Bank to retrieve memory profiles from.""" + """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}`.""" - scope: Optional[dict[str, str]] - """The scope of the memories 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.""" - 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.""" + 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.""" - config: Optional[RetrieveMemoryProfilesConfigDict] - """""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + response: Optional[PurgeMemoriesResponseDict] + """The response for purging memories.""" -_RetrieveMemoryProfilesRequestParametersOrDict = Union[ - _RetrieveMemoryProfilesRequestParameters, - _RetrieveMemoryProfilesRequestParametersDict, -] +PurgeMemoriesOperationOrDict = Union[PurgeMemoriesOperation, PurgeMemoriesOperationDict] -class MemoryProfile(_common.BaseModel): - """A memory profile.""" - 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.""" - ) +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 MemoryProfileDict(TypedDict, total=False): - """A memory profile.""" - schema_id: Optional[str] - """Represents the ID of the schema. This ID corresponds to the `schema_id` defined inside the SchemaConfig, under StructuredMemoryCustomizationConfig.""" +class GetMemoryRevisionConfigDict(TypedDict, total=False): + """Config for getting a Memory Revision.""" - profile: Optional[dict[str, Any]] - """Represents the profile data.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -MemoryProfileOrDict = Union[MemoryProfile, MemoryProfileDict] +GetMemoryRevisionConfigOrDict = Union[ + GetMemoryRevisionConfig, GetMemoryRevisionConfigDict +] -class RetrieveProfilesResponse(_common.BaseModel): - """The response for retrieving memory profiles.""" +class _GetMemoryRevisionRequestParameters(_common.BaseModel): + """Parameters for getting a memory revision.""" - 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`.""", + name: Optional[str] = Field( + default=None, description="""Name of the Memory Revision.""" ) + config: Optional[GetMemoryRevisionConfig] = Field(default=None, description="""""") -class RetrieveProfilesResponseDict(TypedDict, total=False): - """The response for retrieving memory profiles.""" +class _GetMemoryRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for getting a memory revision.""" - 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`.""" + name: Optional[str] + """Name of the Memory Revision.""" + config: Optional[GetMemoryRevisionConfigDict] + """""" -RetrieveProfilesResponseOrDict = Union[ - RetrieveProfilesResponse, RetrieveProfilesResponseDict -] - - -class RollbackMemoryConfig(_common.BaseModel): - """Config for rolling back a memory.""" - - 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 RollbackMemoryConfigDict(TypedDict, total=False): - """Config for rolling back a memory.""" - - 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.""" - - -RollbackMemoryConfigOrDict = Union[RollbackMemoryConfig, RollbackMemoryConfigDict] - - -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.""" - ) - config: Optional[RollbackMemoryConfig] = Field(default=None, description="""""") - - -class _RollbackMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for generating memories.""" - - name: Optional[str] - """Name of the memory to rollback.""" - - target_revision_id: Optional[str] - """The ID of the revision to rollback to.""" - config: Optional[RollbackMemoryConfigDict] - """""" - - -_RollbackMemoryRequestParametersOrDict = Union[ - _RollbackMemoryRequestParameters, _RollbackMemoryRequestParametersDict +_GetMemoryRevisionRequestParametersOrDict = Union[ + _GetMemoryRevisionRequestParameters, _GetMemoryRevisionRequestParametersDict ] -class RollbackMemoryOperation(_common.BaseModel): - """Operation that rolls back a memory.""" +class IntermediateExtractedMemory(_common.BaseModel): + """An extracted memory that is the intermediate result before consolidation.""" - 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( + fact: 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. Represents the fact of the extracted memory.""", ) - done: Optional[bool] = Field( + context: 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. Represents the explanation of why the information was extracted from the source content.""", ) - error: Optional[dict[str, Any]] = Field( + structured_data: Optional[dict[str, Any]] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Output only. Represents the structured value of the extracted memory.""", ) -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 IntermediateExtractedMemoryDict(TypedDict, total=False): + """An extracted memory that is the intermediate result before consolidation.""" - 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.""" + fact: Optional[str] + """Output only. Represents the fact of the extracted memory.""" - 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.""" + context: Optional[str] + """Output only. Represents the explanation of why the information was extracted from the source content.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + structured_data: Optional[dict[str, Any]] + """Output only. Represents the structured value of the extracted memory.""" -RollbackMemoryOperationOrDict = Union[ - RollbackMemoryOperation, RollbackMemoryOperationDict +IntermediateExtractedMemoryOrDict = Union[ + IntermediateExtractedMemory, IntermediateExtractedMemoryDict ] -class UpdateMemoryConfig(_common.BaseModel): - """Config for updating a memory.""" +class MemoryRevision(_common.BaseModel): + """A memory revision.""" - 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( + create_time: Optional[datetime.datetime] = Field( default=None, - description="""Optional. Input only. The TTL for this resource. - - The expiration time is computed: now + TTL.""", + description="""Output only. Represents the timestamp when this Memory Revision was created.""", ) 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.""", + description="""Output only. Represents the timestamp of when this resource is considered expired.""", ) - revision_ttl: Optional[str] = Field( + extracted_memories: Optional[list[IntermediateExtractedMemory]] = Field( default=None, - description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", + 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.""", ) - disable_memory_revisions: Optional[bool] = Field( + fact: Optional[str] = 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="""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.""", ) - metadata: Optional[dict[str, MemoryMetadataValue]] = Field( + labels: Optional[dict[str, str]] = 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="""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`.""", ) - memory_id: Optional[str] = Field( + name: 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.""", + description="""Identifier. Represents the resource name of the Memory Revision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{memory_revision}`""", ) - update_mask: Optional[str] = Field( + structured_data: Optional[dict[str, Any]] = Field( default=None, - description="""The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""", + description="""Output only. Represents the structured value of the memory at the time of revision creation.""", ) -class UpdateMemoryConfigDict(TypedDict, total=False): - """Config for updating a memory.""" +class MemoryRevisionDict(TypedDict, total=False): + """A memory revision.""" - 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 Revision was created.""" - display_name: Optional[str] - """The display name of the memory.""" + expire_time: Optional[datetime.datetime] + """Output only. Represents the timestamp of when this resource is considered expired.""" - description: Optional[str] - """The description of the memory.""" + 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.""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + 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.""" - ttl: Optional[str] - """Optional. Input only. The TTL for this resource. + 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`.""" - The expiration time is computed: now + TTL.""" + 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}`""" - 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.""" + structured_data: Optional[dict[str, Any]] + """Output only. Represents the structured value of the memory at the time of revision creation.""" - 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.""" +MemoryRevisionOrDict = Union[MemoryRevision, MemoryRevisionDict] - 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.""" +class ListMemoryRevisionsConfig(_common.BaseModel): + """Config for listing memory revisions.""" - 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.""" + 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.""", + ) - 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.""" - update_mask: Optional[str] - """The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" +class ListMemoryRevisionsConfigDict(TypedDict, total=False): + """Config for listing memory revisions.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -UpdateMemoryConfigOrDict = Union[UpdateMemoryConfig, UpdateMemoryConfigDict] + page_size: Optional[int] + """""" + page_token: Optional[str] + """""" -class _UpdateMemoryRequestParameters(_common.BaseModel): - """Parameters for updating memories.""" + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported.""" - name: Optional[str] = Field( - default=None, description="""Name of the memory to update.""" - ) - fact: Optional[str] = Field( - default=None, - 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. +ListMemoryRevisionsConfigOrDict = Union[ + ListMemoryRevisionsConfig, ListMemoryRevisionsConfigDict +] - 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 '*'.""", + +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[UpdateMemoryConfig] = Field(default=None, description="""""") -class _UpdateMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for updating memories.""" +class _ListMemoryRevisionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing memory revisions.""" name: Optional[str] - """Name of the memory to update.""" + """Name of the memory""" - fact: Optional[str] - """The updated fact of the memory. + config: Optional[ListMemoryRevisionsConfigDict] + """""" - This is the semantic knowledge extracted from the source content.""" - scope: Optional[dict[str, str]] - """The updated scope of the memory. +_ListMemoryRevisionsRequestParametersOrDict = Union[ + _ListMemoryRevisionsRequestParameters, _ListMemoryRevisionsRequestParametersDict +] - 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] +class ListMemoryRevisionsResponse(_common.BaseModel): + """Response for listing memory 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="""""") + memory_revisions: Optional[list[MemoryRevision]] = Field( + default=None, description="""List of memory revisions.""" + ) + + +class ListMemoryRevisionsResponseDict(TypedDict, total=False): + """Response for listing memory revisions.""" + + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + + next_page_token: Optional[str] """""" + memory_revisions: Optional[list[MemoryRevisionDict]] + """List of memory revisions.""" -_UpdateMemoryRequestParametersOrDict = Union[ - _UpdateMemoryRequestParameters, _UpdateMemoryRequestParametersDict + +ListMemoryRevisionsResponseOrDict = Union[ + ListMemoryRevisionsResponse, ListMemoryRevisionsResponseDict ] -class PurgeMemoriesConfig(_common.BaseModel): - """Config for purging memories.""" +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.""" ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", - ) + tools: Optional[list[genai_types.Tool]] = Field(default=None, description="""""") -class PurgeMemoriesConfigDict(TypedDict, total=False): - """Config for purging memories.""" +class AskContextsConfigDict(TypedDict, total=False): + """Config for asking RAG Contexts.""" 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.""" + tools: Optional[list[genai_types.Tool]] + """""" -PurgeMemoriesConfigOrDict = Union[PurgeMemoriesConfig, PurgeMemoriesConfigDict] +AskContextsConfigOrDict = Union[AskContextsConfig, AskContextsConfigDict] -class _PurgeMemoriesRequestParameters(_common.BaseModel): - """Parameters for purging memories.""" +class RagQueryRanking(_common.BaseModel): + """Configurations for hybrid search results ranking.""" - name: Optional[str] = Field( - default=None, description="""Name of the Memory Bank to purge memories from.""" - ) - filter: Optional[str] = Field( + alpha: 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="""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.""", ) - filter_groups: Optional[list[MemoryConjunctionFilter]] = 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"))`. - """, +class RagQueryRankingDict(TypedDict, total=False): + """Configurations for hybrid search results ranking.""" + + 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.""" + + +RagQueryRankingOrDict = Union[RagQueryRanking, RagQueryRankingDict] + + +class RagQuery(_common.BaseModel): + """A query to retrieve relevant contexts.""" + + rag_retrieval_config: Optional[genai_types.RagRetrievalConfig] = Field( + default=None, description="""Optional. The retrieval config for the query.""" ) - force: Optional[bool] = Field( + ranking: Optional[RagQueryRanking] = Field( default=None, - description="""If true, the memories will actually be purged. If false, the purge request will be validated but not executed.""", + description="""Optional. Configurations for hybrid search results ranking.""", + ) + similarity_top_k: Optional[int] = Field( + default=None, description="""Optional. The number of contexts to retrieve.""" + ) + text: Optional[str] = Field( + default=None, + description="""Optional. The query in text format to get relevant contexts.""", ) - config: Optional[PurgeMemoriesConfig] = Field(default=None, description="""""") - - -class _PurgeMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for purging memories.""" - name: Optional[str] - """Name of the Memory Bank to purge memories from.""" - filter: Optional[str] - """The standard list filter to determine which memories to purge. - More detail in [AIP-160](https://google.aip.dev/160).""" +class RagQueryDict(TypedDict, total=False): + """A query to retrieve relevant contexts.""" - 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). + rag_retrieval_config: Optional[genai_types.RagRetrievalConfigDict] + """Optional. The retrieval config for the query.""" - 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}]}]` + ranking: Optional[RagQueryRankingDict] + """Optional. Configurations for hybrid search results ranking.""" - would be equivalent to the logical expression: - `(metadata.author = "agent 123" OR (metadata.label = "travel" AND - metadata.author = "agent 321"))`. - """ + similarity_top_k: Optional[int] + """Optional. The number of contexts to retrieve.""" - force: Optional[bool] - """If true, the memories will actually be purged. If false, the purge request will be validated but not executed.""" + text: Optional[str] + """Optional. The query in text format to get relevant contexts.""" - config: Optional[PurgeMemoriesConfigDict] - """""" +RagQueryOrDict = Union[RagQuery, RagQueryDict] -_PurgeMemoriesRequestParametersOrDict = Union[ - _PurgeMemoriesRequestParameters, _PurgeMemoriesRequestParametersDict -] +class _AskContextsRequestParameters(_common.BaseModel): + """Parameters for asking RAG Contexts.""" -class PurgeMemoriesResponse(_common.BaseModel): - """The response for purging memories.""" + query: Optional[RagQuery] = Field(default=None, description="""""") + config: Optional[AskContextsConfig] = Field(default=None, description="""""") - purge_count: Optional[int] = Field( - default=None, description="""The number of memories that were purged.""" - ) +class _AskContextsRequestParametersDict(TypedDict, total=False): + """Parameters for asking RAG Contexts.""" -class PurgeMemoriesResponseDict(TypedDict, total=False): - """The response for purging memories.""" + query: Optional[RagQueryDict] + """""" - purge_count: Optional[int] - """The number of memories that were purged.""" + config: Optional[AskContextsConfigDict] + """""" -PurgeMemoriesResponseOrDict = Union[PurgeMemoriesResponse, PurgeMemoriesResponseDict] +_AskContextsRequestParametersOrDict = Union[ + _AskContextsRequestParameters, _AskContextsRequestParametersDict +] -class PurgeMemoriesOperation(_common.BaseModel): - """Operation that purges memories from a Memory Bank.""" +class RagContextsContext(_common.BaseModel): + """A context of the query.""" - 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}`.""", + chunk: Optional[genai_types.RagChunk] = Field( + default=None, description="""Context of the retrieved chunk.""" ) - metadata: Optional[dict[str, Any]] = Field( + distance: 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="""The distance between the query dense embedding vector and the context text vector.""", ) - done: Optional[bool] = Field( + score: Optional[float] = 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="""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.""", ) - error: Optional[dict[str, Any]] = Field( + source_display_name: Optional[str] = Field( + default=None, description="""The file display name.""" + ) + source_uri: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + 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.""", ) - response: Optional[PurgeMemoriesResponse] = Field( - default=None, description="""The response for purging memories.""" + 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 PurgeMemoriesOperationDict(TypedDict, total=False): - """Operation that purges memories from a Memory Bank.""" +class RagContextsContextDict(TypedDict, total=False): + """A context of the query.""" - 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}`.""" + chunk: Optional[genai_types.RagChunkDict] + """Context of the retrieved chunk.""" - 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.""" + distance: Optional[float] + """The distance between the query dense embedding vector and the context text vector.""" - 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.""" + 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.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + source_display_name: Optional[str] + """The file display name.""" - response: Optional[PurgeMemoriesResponseDict] - """The response for purging memories.""" + 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.""" + sparse_distance: Optional[float] + """The distance between the query sparse embedding vector and the context text vector.""" -PurgeMemoriesOperationOrDict = Union[PurgeMemoriesOperation, PurgeMemoriesOperationDict] + text: Optional[str] + """The text chunk.""" -class GetMemoryRevisionConfig(_common.BaseModel): - """Config for getting a Memory Revision.""" +RagContextsContextOrDict = Union[RagContextsContext, RagContextsContextDict] - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) +class RagContexts(_common.BaseModel): + """Relevant contexts for one query.""" -class GetMemoryRevisionConfigDict(TypedDict, total=False): - """Config for getting a Memory Revision.""" + contexts: Optional[list[RagContextsContext]] = Field( + default=None, description="""All its contexts.""" + ) - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class RagContextsDict(TypedDict, total=False): + """Relevant contexts for one query.""" -GetMemoryRevisionConfigOrDict = Union[ - GetMemoryRevisionConfig, GetMemoryRevisionConfigDict -] + contexts: Optional[list[RagContextsContextDict]] + """All its contexts.""" -class _GetMemoryRevisionRequestParameters(_common.BaseModel): - """Parameters for getting a memory revision.""" +RagContextsOrDict = Union[RagContexts, RagContextsDict] - name: Optional[str] = Field( - default=None, description="""Name of the Memory Revision.""" + +class AskContextsResponse(_common.BaseModel): + + contexts: Optional[RagContexts] = Field( + default=None, description="""The contexts of the query.""" + ) + response: Optional[str] = Field( + default=None, description="""The Retrieval Response.""" ) - config: Optional[GetMemoryRevisionConfig] = Field(default=None, description="""""") -class _GetMemoryRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for getting a memory revision.""" +class AskContextsResponseDict(TypedDict, total=False): - name: Optional[str] - """Name of the Memory Revision.""" + contexts: Optional[RagContextsDict] + """The contexts of the query.""" - config: Optional[GetMemoryRevisionConfigDict] - """""" + response: Optional[str] + """The Retrieval Response.""" -_GetMemoryRevisionRequestParametersOrDict = Union[ - _GetMemoryRevisionRequestParameters, _GetMemoryRevisionRequestParametersDict -] +AskContextsResponseOrDict = Union[AskContextsResponse, AskContextsResponseDict] -class IntermediateExtractedMemory(_common.BaseModel): - """An extracted memory that is the intermediate result before consolidation.""" +class CorpusStatus(_common.BaseModel): + """RagCorpus status.""" - fact: Optional[str] = Field( - default=None, - description="""Output only. Represents the fact of the extracted memory.""", - ) - context: Optional[str] = Field( + error_status: Optional[str] = Field( default=None, - description="""Output only. Represents the explanation of why the information was extracted from the source content.""", + description="""Output only. Only when the `state` field is ERROR.""", ) - structured_data: Optional[dict[str, Any]] = Field( - default=None, - description="""Output only. Represents the structured value of the extracted memory.""", + state: Optional[Literal["UNKNOWN", "INITIALIZED", "ACTIVE", "ERROR"]] = Field( + default=None, description="""Output only. RagCorpus life state.""" ) -class IntermediateExtractedMemoryDict(TypedDict, total=False): - """An extracted memory that is the intermediate result before consolidation.""" +class CorpusStatusDict(TypedDict, total=False): + """RagCorpus status.""" - fact: Optional[str] - """Output only. Represents the fact of the extracted memory.""" + error_status: Optional[str] + """Output only. Only when the `state` field is ERROR.""" - context: Optional[str] - """Output only. Represents the explanation of why the information was extracted from the source content.""" + state: Optional[Literal["UNKNOWN", "INITIALIZED", "ACTIVE", "ERROR"]] + """Output only. RagCorpus life state.""" - structured_data: Optional[dict[str, Any]] - """Output only. Represents the structured value of the extracted memory.""" +CorpusStatusOrDict = Union[CorpusStatus, CorpusStatusDict] -IntermediateExtractedMemoryOrDict = Union[ - IntermediateExtractedMemory, IntermediateExtractedMemoryDict + +class RagCorpusCorpusTypeConfigDocumentCorpus(_common.BaseModel): + """Config for the document corpus.""" + + pass + + +class RagCorpusCorpusTypeConfigDocumentCorpusDict(TypedDict, total=False): + """Config for the document corpus.""" + + pass + + +RagCorpusCorpusTypeConfigDocumentCorpusOrDict = Union[ + RagCorpusCorpusTypeConfigDocumentCorpus, RagCorpusCorpusTypeConfigDocumentCorpusDict ] -class MemoryRevision(_common.BaseModel): - """A memory revision.""" +class RagFileParsingConfigLlmParser(_common.BaseModel): + """Specifies the LLM parsing for RagFiles.""" - 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( + custom_parsing_prompt: 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.""", + description="""The prompt to use for parsing. If not specified, a default prompt will be used.""", ) - labels: Optional[dict[str, str]] = Field( + global_max_parsing_requests_per_min: Optional[int] = 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`.""", + 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.""", ) - name: Optional[str] = Field( + max_parsing_requests_per_min: Optional[int] = 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}`""", + 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.""", ) - structured_data: Optional[dict[str, Any]] = Field( + model_name: Optional[str] = Field( default=None, - description="""Output only. Represents the structured value of the memory at the time of revision creation.""", + description="""The name of a LLM model used for parsing. Format: * `projects/{project_id}/locations/{location}/publishers/{publisher}/models/{model}`""", ) -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.""" +class RagFileParsingConfigLlmParserDict(TypedDict, total=False): + """Specifies the LLM parsing for RagFiles.""" - 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.""" + custom_parsing_prompt: Optional[str] + """The prompt to use for parsing. If not specified, a default prompt will be used.""" - 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`.""" + 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.""" - 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}`""" + 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.""" - structured_data: Optional[dict[str, Any]] - """Output only. Represents the structured value of the memory at the time of revision creation.""" + model_name: Optional[str] + """The name of a LLM model used for parsing. Format: * `projects/{project_id}/locations/{location}/publishers/{publisher}/models/{model}`""" -MemoryRevisionOrDict = Union[MemoryRevision, MemoryRevisionDict] +RagFileParsingConfigLlmParserOrDict = Union[ + RagFileParsingConfigLlmParser, RagFileParsingConfigLlmParserDict +] -class ListMemoryRevisionsConfig(_common.BaseModel): - """Config for listing memory revisions.""" +class RagCorpusCorpusTypeConfigMemoryCorpus(_common.BaseModel): + """Config for the memory corpus.""" - 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.""", + llm_parser: Optional[RagFileParsingConfigLlmParser] = Field( + default=None, description="""The LLM parser to use 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] - """""" - - page_token: Optional[str] - """""" +class RagCorpusCorpusTypeConfigMemoryCorpusDict(TypedDict, total=False): + """Config for the memory corpus.""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" + llm_parser: Optional[RagFileParsingConfigLlmParserDict] + """The LLM parser to use for the memory corpus.""" -ListMemoryRevisionsConfigOrDict = Union[ - ListMemoryRevisionsConfig, ListMemoryRevisionsConfigDict +RagCorpusCorpusTypeConfigMemoryCorpusOrDict = Union[ + RagCorpusCorpusTypeConfigMemoryCorpus, RagCorpusCorpusTypeConfigMemoryCorpusDict ] -class _ListMemoryRevisionsRequestParameters(_common.BaseModel): - """Parameters for listing memory revisions.""" +class RagCorpusCorpusTypeConfig(_common.BaseModel): + """The config for the corpus type of the RagCorpus.""" - name: Optional[str] = Field(default=None, description="""Name of the memory""") - config: Optional[ListMemoryRevisionsConfig] = Field( - default=None, description="""""" + document_corpus: Optional[RagCorpusCorpusTypeConfigDocumentCorpus] = Field( + default=None, description="""Optional. Config for the document corpus.""" + ) + memory_corpus: Optional[RagCorpusCorpusTypeConfigMemoryCorpus] = Field( + default=None, description="""Optional. Config for the memory corpus.""" ) -class _ListMemoryRevisionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing memory revisions.""" +class RagCorpusCorpusTypeConfigDict(TypedDict, total=False): + """The config for the corpus type of the RagCorpus.""" - name: Optional[str] - """Name of the memory""" + document_corpus: Optional[RagCorpusCorpusTypeConfigDocumentCorpusDict] + """Optional. Config for the document corpus.""" - config: Optional[ListMemoryRevisionsConfigDict] - """""" + memory_corpus: Optional[RagCorpusCorpusTypeConfigMemoryCorpusDict] + """Optional. Config for the memory corpus.""" -_ListMemoryRevisionsRequestParametersOrDict = Union[ - _ListMemoryRevisionsRequestParameters, _ListMemoryRevisionsRequestParametersDict +RagCorpusCorpusTypeConfigOrDict = Union[ + RagCorpusCorpusTypeConfig, RagCorpusCorpusTypeConfigDict ] -class ListMemoryRevisionsResponse(_common.BaseModel): - """Response for listing memory revisions.""" +class RagEmbeddingModelConfigVertexPredictionEndpoint(_common.BaseModel): + """Config representing a model hosted on Vertex Prediction Endpoint.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + 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}`""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - memory_revisions: Optional[list[MemoryRevision]] = Field( - default=None, description="""List of memory revisions.""" + 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 ListMemoryRevisionsResponseDict(TypedDict, total=False): - """Response for listing memory revisions.""" +class RagEmbeddingModelConfigVertexPredictionEndpointDict(TypedDict, total=False): + """Config representing a model hosted on Vertex Prediction Endpoint.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + endpoint: Optional[str] + """Required. The endpoint resource name. Format: `projects/{project}/locations/{location}/publishers/{publisher}/models/{model}` or `projects/{project}/locations/{location}/endpoints/{endpoint}`""" - next_page_token: Optional[str] - """""" + 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}`""" - memory_revisions: Optional[list[MemoryRevisionDict]] - """List of memory revisions.""" + 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.""" -ListMemoryRevisionsResponseOrDict = Union[ - ListMemoryRevisionsResponse, ListMemoryRevisionsResponseDict +RagEmbeddingModelConfigVertexPredictionEndpointOrDict = Union[ + RagEmbeddingModelConfigVertexPredictionEndpoint, + RagEmbeddingModelConfigVertexPredictionEndpointDict, ] -class AskContextsConfig(_common.BaseModel): - """Config for asking RAG Contexts.""" +class RagEmbeddingModelConfigSparseEmbeddingConfigBm25(_common.BaseModel): + """Message for BM25 parameters.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - tools: Optional[list[genai_types.Tool]] = Field(default=None, description="""""") + 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.""", + ) + 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 AskContextsConfigDict(TypedDict, total=False): - """Config for asking RAG Contexts.""" +class RagEmbeddingModelConfigSparseEmbeddingConfigBm25Dict(TypedDict, total=False): + """Message for BM25 parameters.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + 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.""" - tools: Optional[list[genai_types.Tool]] - """""" + 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.""" -AskContextsConfigOrDict = Union[AskContextsConfig, AskContextsConfigDict] +RagEmbeddingModelConfigSparseEmbeddingConfigBm25OrDict = Union[ + RagEmbeddingModelConfigSparseEmbeddingConfigBm25, + RagEmbeddingModelConfigSparseEmbeddingConfigBm25Dict, +] -class RagQueryRanking(_common.BaseModel): - """Configurations for hybrid search results ranking.""" - 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.""", +class RagEmbeddingModelConfigSparseEmbeddingConfig(_common.BaseModel): + """Configuration for sparse emebdding generation.""" + + bm25: Optional[RagEmbeddingModelConfigSparseEmbeddingConfigBm25] = Field( + default=None, description="""Use BM25 scoring algorithm.""" ) -class RagQueryRankingDict(TypedDict, total=False): - """Configurations for hybrid search results ranking.""" +class RagEmbeddingModelConfigSparseEmbeddingConfigDict(TypedDict, total=False): + """Configuration for sparse emebdding generation.""" - 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.""" + bm25: Optional[RagEmbeddingModelConfigSparseEmbeddingConfigBm25Dict] + """Use BM25 scoring algorithm.""" -RagQueryRankingOrDict = Union[RagQueryRanking, RagQueryRankingDict] +RagEmbeddingModelConfigSparseEmbeddingConfigOrDict = Union[ + RagEmbeddingModelConfigSparseEmbeddingConfig, + RagEmbeddingModelConfigSparseEmbeddingConfigDict, +] -class RagQuery(_common.BaseModel): - """A query to retrieve relevant contexts.""" +class RagEmbeddingModelConfigHybridSearchConfig(_common.BaseModel): + """Config for hybrid search.""" - rag_retrieval_config: Optional[genai_types.RagRetrievalConfig] = Field( - default=None, description="""Optional. The retrieval config for the query.""" - ) - ranking: Optional[RagQueryRanking] = Field( + dense_embedding_model_prediction_endpoint: Optional[ + RagEmbeddingModelConfigVertexPredictionEndpoint + ] = 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.""" + description="""Required. The Vertex AI Prediction Endpoint that hosts the embedding model for dense embedding generations.""", ) - text: Optional[str] = Field( - default=None, - description="""Optional. The query in text format to get relevant contexts.""", + 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 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.""" +class RagEmbeddingModelConfigHybridSearchConfigDict(TypedDict, total=False): + """Config for hybrid search.""" - similarity_top_k: Optional[int] - """Optional. The number of contexts to retrieve.""" + dense_embedding_model_prediction_endpoint: Optional[ + RagEmbeddingModelConfigVertexPredictionEndpointDict + ] + """Required. The Vertex AI Prediction Endpoint that hosts the embedding model for dense embedding generations.""" - text: Optional[str] - """Optional. The query in text format to get relevant contexts.""" + 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.""" -RagQueryOrDict = Union[RagQuery, RagQueryDict] +RagEmbeddingModelConfigHybridSearchConfigOrDict = Union[ + RagEmbeddingModelConfigHybridSearchConfig, + RagEmbeddingModelConfigHybridSearchConfigDict, +] -class _AskContextsRequestParameters(_common.BaseModel): - """Parameters for asking RAG Contexts.""" +class RagEmbeddingModelConfig(_common.BaseModel): + """Config for the embedding model to use for RAG.""" - query: Optional[RagQuery] = Field(default=None, description="""""") - config: Optional[AskContextsConfig] = Field(default=None, description="""""") + hybrid_search_config: Optional[RagEmbeddingModelConfigHybridSearchConfig] = Field( + default=None, description="""Configuration for hybrid search.""" + ) + vertex_prediction_endpoint: Optional[ + RagEmbeddingModelConfigVertexPredictionEndpoint + ] = 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.""", + ) -class _AskContextsRequestParametersDict(TypedDict, total=False): - """Parameters for asking RAG Contexts.""" +class RagEmbeddingModelConfigDict(TypedDict, total=False): + """Config for the embedding model to use for RAG.""" - query: Optional[RagQueryDict] - """""" + hybrid_search_config: Optional[RagEmbeddingModelConfigHybridSearchConfigDict] + """Configuration for hybrid search.""" - config: Optional[AskContextsConfigDict] - """""" + 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.""" -_AskContextsRequestParametersOrDict = Union[ - _AskContextsRequestParameters, _AskContextsRequestParametersDict +RagEmbeddingModelConfigOrDict = Union[ + RagEmbeddingModelConfig, RagEmbeddingModelConfigDict ] -class RagContextsContext(_common.BaseModel): - """A context of the query.""" +class RagVectorDbConfigPinecone(_common.BaseModel): + """The config for the Pinecone.""" - 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( - 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.""", - ) - sparse_distance: Optional[float] = Field( + index_name: Optional[str] = Field( default=None, - description="""The distance between the query sparse embedding vector and the context text vector.""", + description="""Pinecone index name. This value cannot be changed after it's set.""", ) - text: Optional[str] = Field(default=None, description="""The text chunk.""") -class RagContextsContextDict(TypedDict, total=False): - """A context of the query.""" +class RagVectorDbConfigPineconeDict(TypedDict, total=False): + """The config for the Pinecone.""" - chunk: Optional[genai_types.RagChunkDict] - """Context of the retrieved chunk.""" + index_name: Optional[str] + """Pinecone index name. This value cannot be changed after it's set.""" - 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.""" +RagVectorDbConfigPineconeOrDict = Union[ + RagVectorDbConfigPinecone, RagVectorDbConfigPineconeDict +] - 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 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.""" - sparse_distance: Optional[float] - """The distance between the query sparse embedding vector and the context text vector.""" + 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.""", + ) - text: Optional[str] - """The text chunk.""" +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.""" -RagContextsContextOrDict = Union[RagContextsContext, RagContextsContextDict] + 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.""" + 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.""" -class RagContexts(_common.BaseModel): - """Relevant contexts for one query.""" - contexts: Optional[list[RagContextsContext]] = Field( - default=None, description="""All its contexts.""" - ) +RagVectorDbConfigRagManagedDbANNOrDict = Union[ + RagVectorDbConfigRagManagedDbANN, RagVectorDbConfigRagManagedDbANNDict +] -class RagContextsDict(TypedDict, total=False): - """Relevant contexts for one query.""" +class RagVectorDbConfigRagManagedDbKNN(_common.BaseModel): + """Config for KNN search.""" - contexts: Optional[list[RagContextsContextDict]] - """All its contexts.""" + pass -RagContextsOrDict = Union[RagContexts, RagContextsDict] +class RagVectorDbConfigRagManagedDbKNNDict(TypedDict, total=False): + """Config for KNN search.""" + pass -class AskContextsResponse(_common.BaseModel): - contexts: Optional[RagContexts] = Field( - default=None, description="""The contexts of the query.""" +RagVectorDbConfigRagManagedDbKNNOrDict = Union[ + RagVectorDbConfigRagManagedDbKNN, RagVectorDbConfigRagManagedDbKNNDict +] + + +class RagVectorDbConfigRagManagedDb(_common.BaseModel): + """The config for the default RAG-managed Vector DB.""" + + 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.""" - - error_status: Optional[str] - """Output only. Only when the `state` field is ERROR.""" +class RagVectorDbConfigRagManagedVertexVectorSearchDict(TypedDict, total=False): + """The config for the RAG-managed Vertex Vector Search 2.0.""" - state: Optional[Literal["UNKNOWN", "INITIALIZED", "ACTIVE", "ERROR"]] - """Output only. RagCorpus life state.""" + 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}`""" -CorpusStatusOrDict = Union[CorpusStatus, CorpusStatusDict] +RagVectorDbConfigRagManagedVertexVectorSearchOrDict = Union[ + RagVectorDbConfigRagManagedVertexVectorSearch, + RagVectorDbConfigRagManagedVertexVectorSearchDict, +] -class RagCorpusCorpusTypeConfigDocumentCorpus(_common.BaseModel): - """Config for the document corpus.""" +class RagVectorDbConfigVertexFeatureStore(_common.BaseModel): + """The config for the Vertex Feature Store.""" - pass + 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}`""", + ) -class RagCorpusCorpusTypeConfigDocumentCorpusDict(TypedDict, total=False): - """Config for the document corpus.""" +class RagVectorDbConfigVertexFeatureStoreDict(TypedDict, total=False): + """The config for the Vertex Feature Store.""" - pass + feature_view_resource_name: Optional[str] + """The resource name of the FeatureView. Format: `projects/{project}/locations/{location}/featureOnlineStores/{feature_online_store}/featureViews/{feature_view}`""" -RagCorpusCorpusTypeConfigDocumentCorpusOrDict = Union[ - RagCorpusCorpusTypeConfigDocumentCorpus, RagCorpusCorpusTypeConfigDocumentCorpusDict +RagVectorDbConfigVertexFeatureStoreOrDict = Union[ + RagVectorDbConfigVertexFeatureStore, RagVectorDbConfigVertexFeatureStoreDict ] -class RagFileParsingConfigLlmParser(_common.BaseModel): - """Specifies the LLM parsing for RagFiles.""" +class RagVectorDbConfigVertexVectorSearch(_common.BaseModel): + """The config for the Vertex Vector Search.""" - 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( + pinecone: Optional[RagVectorDbConfigPinecone] = Field( + default=None, description="""The config for the Pinecone.""" + ) + rag_embedding_model_config: Optional[RagEmbeddingModelConfig] = 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}`""", + description="""Optional. Immutable. The embedding model config of the Vector DB.""", ) - model_version_id: Optional[str] = Field( + 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="""Output only. Version ID of the model that is deployed on the endpoint. Present only when the endpoint is not a publisher model.""", + 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 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}`""" +class RagVectorDbConfigDict(TypedDict, total=False): + """Config for the Vector DB to use for RAG.""" - 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.""" + api_auth: Optional[genai_types.ApiAuthDict] + """Authentication config for the chosen Vector DB.""" + pinecone: Optional[RagVectorDbConfigPineconeDict] + """The config for the Pinecone.""" -RagEmbeddingModelConfigVertexPredictionEndpointOrDict = Union[ - RagEmbeddingModelConfigVertexPredictionEndpoint, - RagEmbeddingModelConfigVertexPredictionEndpointDict, -] + rag_embedding_model_config: Optional[RagEmbeddingModelConfigDict] + """Optional. Immutable. The embedding model config of the Vector DB.""" + rag_managed_db: Optional[RagVectorDbConfigRagManagedDbDict] + """The config for the RAG-managed Vector DB.""" -class RagEmbeddingModelConfigSparseEmbeddingConfigBm25(_common.BaseModel): - """Message for BM25 parameters.""" + rag_managed_vertex_vector_search: Optional[ + RagVectorDbConfigRagManagedVertexVectorSearchDict + ] + """The config for the RAG-managed Vertex Vector Search 2.0.""" - 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.""", - ) - 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.""", - ) + vertex_feature_store: Optional[RagVectorDbConfigVertexFeatureStoreDict] + """The config for the Vertex Feature Store.""" + vertex_vector_search: Optional[RagVectorDbConfigVertexVectorSearchDict] + """The config for the Vertex Vector Search.""" -class RagEmbeddingModelConfigSparseEmbeddingConfigBm25Dict(TypedDict, total=False): - """Message for BM25 parameters.""" + weaviate: Optional[RagVectorDbConfigWeaviateDict] + """The config for the Weaviate.""" - 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.""" +RagVectorDbConfigOrDict = Union[RagVectorDbConfig, RagVectorDbConfigDict] - multilingual: Optional[bool] - """Optional. Use multilingual tokenizer if set to true.""" +class VertexAiSearchConfig(_common.BaseModel): + """Config for the Vertex AI Search.""" -RagEmbeddingModelConfigSparseEmbeddingConfigBm25OrDict = Union[ - RagEmbeddingModelConfigSparseEmbeddingConfigBm25, - RagEmbeddingModelConfigSparseEmbeddingConfigBm25Dict, -] + 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}`.""", + ) -class RagEmbeddingModelConfigSparseEmbeddingConfig(_common.BaseModel): - """Configuration for sparse emebdding generation.""" +class VertexAiSearchConfigDict(TypedDict, total=False): + """Config for the Vertex AI Search.""" - bm25: Optional[RagEmbeddingModelConfigSparseEmbeddingConfigBm25] = Field( - default=None, description="""Use BM25 scoring algorithm.""" - ) + 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}`.""" -class RagEmbeddingModelConfigSparseEmbeddingConfigDict(TypedDict, total=False): - """Configuration for sparse emebdding generation.""" +VertexAiSearchConfigOrDict = Union[VertexAiSearchConfig, VertexAiSearchConfigDict] - bm25: Optional[RagEmbeddingModelConfigSparseEmbeddingConfigBm25Dict] - """Use BM25 scoring algorithm.""" +class RagCorpus(_common.BaseModel): + """A RAG Corpus.""" -RagEmbeddingModelConfigSparseEmbeddingConfigOrDict = Union[ - RagEmbeddingModelConfigSparseEmbeddingConfig, - RagEmbeddingModelConfigSparseEmbeddingConfigDict, -] - - -class RagEmbeddingModelConfigHybridSearchConfig(_common.BaseModel): - """Config for hybrid search.""" - - dense_embedding_model_prediction_endpoint: Optional[ - RagEmbeddingModelConfigVertexPredictionEndpoint - ] = Field( + corpus_status: Optional[CorpusStatus] = Field( + default=None, description="""Output only. RagCorpus state.""" + ) + corpus_type_config: Optional[RagCorpusCorpusTypeConfig] = Field( default=None, - description="""Required. The Vertex AI Prediction Endpoint that hosts the embedding model for dense embedding generations.""", + description="""Optional. The corpus type config of the RagCorpus.""", ) - 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.""", - ) + 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 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.""" - - 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.""" - - -RagEmbeddingModelConfigHybridSearchConfigOrDict = Union[ - RagEmbeddingModelConfigHybridSearchConfig, - RagEmbeddingModelConfigHybridSearchConfigDict, -] +class RagCorpusDict(TypedDict, total=False): + """A RAG Corpus.""" + corpus_status: Optional[CorpusStatusDict] + """Output only. RagCorpus state.""" -class RagEmbeddingModelConfig(_common.BaseModel): - """Config for the embedding model to use for RAG.""" + corpus_type_config: Optional[RagCorpusCorpusTypeConfigDict] + """Optional. The corpus type config of the RagCorpus.""" - hybrid_search_config: Optional[RagEmbeddingModelConfigHybridSearchConfig] = Field( - default=None, description="""Configuration for hybrid search.""" - ) - vertex_prediction_endpoint: Optional[ - RagEmbeddingModelConfigVertexPredictionEndpoint - ] = 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.""", - ) + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this RagCorpus was created.""" + description: Optional[str] + """Optional. The description of the RagCorpus.""" -class RagEmbeddingModelConfigDict(TypedDict, total=False): - """Config for the embedding model to use for RAG.""" + 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.""" - hybrid_search_config: Optional[RagEmbeddingModelConfigHybridSearchConfigDict] - """Configuration for hybrid search.""" + 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.""" - 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.""" + 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.""" -RagEmbeddingModelConfigOrDict = Union[ - RagEmbeddingModelConfig, RagEmbeddingModelConfigDict -] + 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.""" -class RagVectorDbConfigPinecone(_common.BaseModel): - """The config for the Pinecone.""" + satisfies_pzi: Optional[bool] + """Output only. Reserved for future use.""" - index_name: Optional[str] = Field( - default=None, - description="""Pinecone index name. This value cannot be changed after it's set.""", - ) + satisfies_pzs: Optional[bool] + """Output only. Reserved for future use.""" + update_time: Optional[datetime.datetime] + """Output only. Timestamp when this RagCorpus was last updated.""" -class RagVectorDbConfigPineconeDict(TypedDict, total=False): - """The config for the Pinecone.""" + vector_db_config: Optional[RagVectorDbConfigDict] + """Optional. Immutable. The config for the Vector DBs.""" - index_name: Optional[str] - """Pinecone index name. This value cannot be changed after it's set.""" + vertex_ai_search_config: Optional[VertexAiSearchConfigDict] + """Optional. Immutable. The config for the Vertex AI Search.""" -RagVectorDbConfigPineconeOrDict = Union[ - RagVectorDbConfigPinecone, RagVectorDbConfigPineconeDict -] +RagCorpusOrDict = Union[RagCorpus, RagCorpusDict] -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.""" +class CreateRagCorpusConfig(_common.BaseModel): + """Config for creating a RAG corpus.""" - 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.""", + 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.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -RagVectorDbConfigVertexVectorSearchOrDict = Union[ - RagVectorDbConfigVertexVectorSearch, RagVectorDbConfigVertexVectorSearchDict -] + response: Optional[RagCorpusDict] + """The created Corpus.""" -class RagVectorDbConfigWeaviate(_common.BaseModel): - """The config for the Weaviate.""" +CorpusOperationOrDict = Union[CorpusOperation, CorpusOperationDict] - 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.""", + +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.""" + + 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.""" + + config: Optional[GetRagCorpusConfigDict] + """""" + + name: Optional[str] + """""" + + +_GetRagCorpusRequestParametersOrDict = Union[ + _GetRagCorpusRequestParameters, _GetRagCorpusRequestParametersDict ] -class RagVectorDbConfig(_common.BaseModel): - """Config for the Vector DB to use for RAG.""" +class ListRagCorporaConfig(_common.BaseModel): + """Config for listing RagCorpora.""" - 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.""" + 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 RagVectorDbConfigDict(TypedDict, total=False): - """Config for the Vector DB to use for RAG.""" +class ListRagCorporaConfigDict(TypedDict, total=False): + """Config for listing RagCorpora.""" - api_auth: Optional[genai_types.ApiAuthDict] - """Authentication config for the chosen Vector DB.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - pinecone: Optional[RagVectorDbConfigPineconeDict] - """The config for the Pinecone.""" + page_size: Optional[int] + """""" - rag_embedding_model_config: Optional[RagEmbeddingModelConfigDict] - """Optional. Immutable. The embedding model config of the Vector DB.""" + page_token: Optional[str] + """""" - 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.""" +ListRagCorporaConfigOrDict = Union[ListRagCorporaConfig, ListRagCorporaConfigDict] - vertex_feature_store: Optional[RagVectorDbConfigVertexFeatureStoreDict] - """The config for the Vertex Feature Store.""" - vertex_vector_search: Optional[RagVectorDbConfigVertexVectorSearchDict] - """The config for the Vertex Vector Search.""" +class _ListRagCorporaRequestParameters(_common.BaseModel): + """Parameters for listing RagCorpora.""" - weaviate: Optional[RagVectorDbConfigWeaviateDict] - """The config for the Weaviate.""" + config: Optional[ListRagCorporaConfig] = Field(default=None, description="""""") -RagVectorDbConfigOrDict = Union[RagVectorDbConfig, RagVectorDbConfigDict] +class _ListRagCorporaRequestParametersDict(TypedDict, total=False): + """Parameters for listing RagCorpora.""" + config: Optional[ListRagCorporaConfigDict] + """""" -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}`.""", +_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 VertexAiSearchConfigDict(TypedDict, total=False): - """Config for the Vertex AI Search.""" +class ListRagCorporaResponseDict(TypedDict, total=False): + """Response 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}`.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + next_page_token: Optional[str] + """""" -VertexAiSearchConfigOrDict = Union[VertexAiSearchConfig, VertexAiSearchConfigDict] + rag_corpora: Optional[list[RagCorpusDict]] + """List of RagCorpus instances.""" -class RagCorpus(_common.BaseModel): - """A RAG Corpus.""" +ListRagCorporaResponseOrDict = Union[ListRagCorporaResponse, ListRagCorporaResponseDict] - 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.""", - ) - vertex_ai_search_config: Optional[VertexAiSearchConfig] = Field( - default=None, - description="""Optional. Immutable. The config for the Vertex AI Search.""", + +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 RagCorpusDict(TypedDict, total=False): - """A RAG Corpus.""" +class GetRagFileConfigDict(TypedDict, total=False): + """Config for getting a RAG corpus.""" - corpus_status: Optional[CorpusStatusDict] - """Output only. RagCorpus state.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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.""" +GetRagFileConfigOrDict = Union[GetRagFileConfig, GetRagFileConfigDict] - 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.""" +class _GetRagFileRequestParameters(_common.BaseModel): + """Parameters for getting a RAG corpus.""" - 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.""" + 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] - """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.""" +_GetRagFileRequestParametersOrDict = Union[ + _GetRagFileRequestParameters, _GetRagFileRequestParametersDict +] - 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.""" +class RagFileStatus(_common.BaseModel): + """RagFile status.""" - satisfies_pzs: Optional[bool] - """Output only. Reserved for future use.""" + state: Optional[RagFileState] = Field( + default=None, description="""The state of the RagFile.""" + ) - 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.""" +class RagFileStatusDict(TypedDict, total=False): + """RagFile status.""" - vertex_ai_search_config: Optional[VertexAiSearchConfigDict] - """Optional. Immutable. The config for the Vertex AI Search.""" + state: Optional[RagFileState] + """The state of the RagFile.""" -RagCorpusOrDict = Union[RagCorpus, RagCorpusDict] +RagFileStatusOrDict = Union[RagFileStatus, RagFileStatusDict] -class CreateRagCorpusConfig(_common.BaseModel): - """Config for creating a RAG corpus.""" +class DirectUploadSource(_common.BaseModel): + """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 -class CreateRagCorpusConfigDict(TypedDict, total=False): - """Config for creating a RAG corpus.""" +class DirectUploadSourceDict(TypedDict, total=False): + """The input content is encapsulated and uploaded in the request.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + pass -CreateRagCorpusConfigOrDict = Union[CreateRagCorpusConfig, CreateRagCorpusConfigDict] +DirectUploadSourceOrDict = Union[DirectUploadSource, DirectUploadSourceDict] -class _CreateRagCorpusRequestParameters(_common.BaseModel): - """Parameters for creating a RAG corpus.""" +class GoogleDriveSourceResourceId(_common.BaseModel): + """The type and ID of the Google Drive resource.""" - rag_corpus: Optional[RagCorpus] = Field(default=None, description="""""") - config: Optional[CreateRagCorpusConfig] = Field(default=None, description="""""") + 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 _CreateRagCorpusRequestParametersDict(TypedDict, total=False): - """Parameters for creating a RAG corpus.""" +class GoogleDriveSourceResourceIdDict(TypedDict, total=False): + """The type and ID of the Google Drive resource.""" - rag_corpus: Optional[RagCorpusDict] - """""" + resource_id: Optional[str] + """Required. The ID of the Google Drive resource.""" - config: Optional[CreateRagCorpusConfigDict] - """""" + resource_type: Optional[ResourceType] + """Required. The type of the Google Drive resource.""" -_CreateRagCorpusRequestParametersOrDict = Union[ - _CreateRagCorpusRequestParameters, _CreateRagCorpusRequestParametersDict +GoogleDriveSourceResourceIdOrDict = Union[ + GoogleDriveSourceResourceId, GoogleDriveSourceResourceIdDict ] -class CreateRagCorpusOperation(_common.BaseModel): - """Operation for creating a RAG corpus.""" +class GoogleDriveSource(_common.BaseModel): + """The Google Drive location for the input content.""" - 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}`.""", + resource_ids: Optional[list[GoogleDriveSourceResourceId]] = Field( + default=None, description="""Required. Google Drive resource IDs.""" ) - metadata: Optional[dict[str, Any]] = Field( + + +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="""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. 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/).""", ) - done: Optional[bool] = Field( + custom_queries: Optional[list[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="""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/""", ) - error: Optional[dict[str, Any]] = Field( + email: Optional[str] = Field( + default=None, description="""Required. The Jira email address.""" + ) + 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.""", + ) + server_uri: Optional[str] = Field( + default=None, description="""Required. The Jira server URI.""" ) -class CreateRagCorpusOperationDict(TypedDict, total=False): - """Operation for creating a RAG corpus.""" +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.""" + + server_uri: Optional[str] + """Required. The Jira server URI.""" -CreateRagCorpusOperationOrDict = Union[ - CreateRagCorpusOperation, CreateRagCorpusOperationDict -] +JiraSourceJiraQueriesOrDict = Union[JiraSourceJiraQueries, JiraSourceJiraQueriesDict] -class GetCorpusOperationConfig(_common.BaseModel): - """Config for getting a corpus operation.""" +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 GetCorpusOperationConfigDict(TypedDict, total=False): - """Config for getting a corpus operation.""" +class JiraSourceDict(TypedDict, total=False): + """The Jira source for the ImportRagFilesRequest.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + jira_queries: Optional[list[JiraSourceJiraQueriesDict]] + """Required. The Jira queries.""" -GetCorpusOperationConfigOrDict = Union[ - GetCorpusOperationConfig, GetCorpusOperationConfigDict -] +JiraSourceOrDict = Union[JiraSource, JiraSourceDict] -class _GetCorpusOperationParameters(_common.BaseModel): - """Parameters for getting a corpus operation.""" +class SharePointSourcesSharePointSource(_common.BaseModel): + """An individual SharePointSource.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + client_id: Optional[str] = Field( + default=None, + description="""The Application ID for the app registered in Microsoft Azure Portal. The application must also be configured with MS Graph permissions "Files.ReadAll", "Sites.ReadAll" and BrowserSiteLists.Read.All.""", ) - config: Optional[GetCorpusOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + client_secret: Optional[genai_types.ApiAuthApiKeyConfig] = Field( + default=None, + description="""The application secret for the app registered in Azure.""", ) - - -class _GetCorpusOperationParametersDict(TypedDict, total=False): - """Parameters for getting a corpus operation.""" - - operation_name: Optional[str] - """The server-assigned name for the operation.""" - - config: Optional[GetCorpusOperationConfigDict] - """Used to override the default configuration.""" - - -_GetCorpusOperationParametersOrDict = Union[ - _GetCorpusOperationParameters, _GetCorpusOperationParametersDict -] - - -class CorpusOperation(_common.BaseModel): - """Operation that has a corpus 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}`.""", + drive_id: Optional[str] = Field( + default=None, description="""The ID of the drive to download from.""" ) - metadata: Optional[dict[str, Any]] = Field( + drive_name: Optional[str] = Field( + default=None, description="""The name of the drive to download from.""" + ) + file_id: 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 SharePoint file id. Output only.""", ) - done: Optional[bool] = Field( + sharepoint_folder_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 SharePoint folder to download from.""", ) - error: Optional[dict[str, Any]] = Field( + sharepoint_folder_path: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""The path of the SharePoint folder to download from.""", ) - response: Optional[RagCorpus] = Field( - default=None, description="""The created Corpus.""" + sharepoint_site_name: Optional[str] = Field( + default=None, + description="""The name of the SharePoint site to download from. This can be the site name or the site id.""", ) - - -class CorpusOperationDict(TypedDict, total=False): - """Operation that has a corpus 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[RagCorpusDict] - """The created Corpus.""" - - -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.""" + tenant_id: Optional[str] = Field( + default=None, + description="""Unique identifier of the Azure Active Directory Instance.""", ) -class GetRagCorpusConfigDict(TypedDict, total=False): - """Config for getting a RAG corpus.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class SharePointSourcesSharePointSourceDict(TypedDict, total=False): + """An individual SharePointSource.""" + client_id: Optional[str] + """The Application ID for the app registered in Microsoft Azure Portal. The application must also be configured with MS Graph permissions "Files.ReadAll", "Sites.ReadAll" and BrowserSiteLists.Read.All.""" -GetRagCorpusConfigOrDict = Union[GetRagCorpusConfig, GetRagCorpusConfigDict] + client_secret: Optional[genai_types.ApiAuthApiKeyConfigDict] + """The application secret for the app registered in Azure.""" + drive_id: Optional[str] + """The ID of the drive to download from.""" -class _GetRagCorpusRequestParameters(_common.BaseModel): - """Parameters for getting a RAG corpus.""" + drive_name: Optional[str] + """The name of the drive to download from.""" - config: Optional[GetRagCorpusConfig] = Field(default=None, description="""""") - name: Optional[str] = Field(default=None, description="""""") + file_id: Optional[str] + """Output only. The SharePoint file id. Output only.""" + sharepoint_folder_id: Optional[str] + """The ID of the SharePoint folder to download from.""" -class _GetRagCorpusRequestParametersDict(TypedDict, total=False): - """Parameters for getting a RAG corpus.""" + sharepoint_folder_path: Optional[str] + """The path of the SharePoint folder to download from.""" - config: Optional[GetRagCorpusConfigDict] - """""" + sharepoint_site_name: Optional[str] + """The name of the SharePoint site to download from. This can be the site name or the site id.""" - name: Optional[str] - """""" + tenant_id: Optional[str] + """Unique identifier of the Azure Active Directory Instance.""" -_GetRagCorpusRequestParametersOrDict = Union[ - _GetRagCorpusRequestParameters, _GetRagCorpusRequestParametersDict +SharePointSourcesSharePointSourceOrDict = Union[ + SharePointSourcesSharePointSource, SharePointSourcesSharePointSourceDict ] -class ListRagCorporaConfig(_common.BaseModel): - """Config for listing RagCorpora.""" +class SharePointSources(_common.BaseModel): + """The SharePointSources to pass to ImportRagFiles.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + share_point_sources: Optional[list[SharePointSourcesSharePointSource]] = Field( + default=None, description="""The SharePoint sources.""" ) - 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.""" +class SharePointSourcesDict(TypedDict, total=False): + """The SharePointSources to pass to ImportRagFiles.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + share_point_sources: Optional[list[SharePointSourcesSharePointSourceDict]] + """The SharePoint sources.""" - page_size: Optional[int] - """""" - page_token: Optional[str] - """""" +SharePointSourcesOrDict = Union[SharePointSources, SharePointSourcesDict] -ListRagCorporaConfigOrDict = Union[ListRagCorporaConfig, ListRagCorporaConfigDict] +class SlackSourceSlackChannelsSlackChannel(_common.BaseModel): + """SlackChannel contains the Slack channel ID and the time range to import.""" + channel_id: Optional[str] = Field( + default=None, description="""Required. The Slack channel ID.""" + ) + end_time: Optional[datetime.datetime] = Field( + default=None, + description="""Optional. The ending timestamp for messages to import.""", + ) + start_time: Optional[datetime.datetime] = Field( + default=None, + description="""Optional. The starting timestamp for messages to import.""", + ) -class _ListRagCorporaRequestParameters(_common.BaseModel): - """Parameters for listing RagCorpora.""" - config: Optional[ListRagCorporaConfig] = Field(default=None, description="""""") +class SlackSourceSlackChannelsSlackChannelDict(TypedDict, total=False): + """SlackChannel contains the Slack channel ID and the time range to import.""" + channel_id: Optional[str] + """Required. The Slack channel ID.""" -class _ListRagCorporaRequestParametersDict(TypedDict, total=False): - """Parameters for listing RagCorpora.""" + end_time: Optional[datetime.datetime] + """Optional. The ending timestamp for messages to import.""" - config: Optional[ListRagCorporaConfigDict] - """""" + start_time: Optional[datetime.datetime] + """Optional. The starting timestamp for messages to import.""" -_ListRagCorporaRequestParametersOrDict = Union[ - _ListRagCorporaRequestParameters, _ListRagCorporaRequestParametersDict +SlackSourceSlackChannelsSlackChannelOrDict = Union[ + SlackSourceSlackChannelsSlackChannel, SlackSourceSlackChannelsSlackChannelDict ] -class ListRagCorporaResponse(_common.BaseModel): - """Response for listing RagCorpora.""" +class SlackSourceSlackChannels(_common.BaseModel): + """SlackChannels contains the Slack channels and corresponding access token.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + 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 Slack channel access token that has access to the slack channel IDs. See: https://api.slack.com/tutorials/tracks/getting-a-token.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - rag_corpora: Optional[list[RagCorpus]] = Field( - default=None, description="""List of RagCorpus instances.""" + channels: Optional[list[SlackSourceSlackChannelsSlackChannel]] = Field( + default=None, description="""Required. The Slack channel IDs.""" ) -class ListRagCorporaResponseDict(TypedDict, total=False): - """Response for listing RagCorpora.""" - - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" +class SlackSourceSlackChannelsDict(TypedDict, total=False): + """SlackChannels contains the Slack channels and corresponding access token.""" - next_page_token: Optional[str] - """""" + api_key_config: Optional[genai_types.ApiAuthApiKeyConfigDict] + """Required. The SecretManager secret version resource name (e.g. projects/{project}/secrets/{secret}/versions/{version}) storing the Slack channel access token that has access to the slack channel IDs. See: https://api.slack.com/tutorials/tracks/getting-a-token.""" - rag_corpora: Optional[list[RagCorpusDict]] - """List of RagCorpus instances.""" + channels: Optional[list[SlackSourceSlackChannelsSlackChannelDict]] + """Required. The Slack channel IDs.""" -ListRagCorporaResponseOrDict = Union[ListRagCorporaResponse, ListRagCorporaResponseDict] +SlackSourceSlackChannelsOrDict = Union[ + SlackSourceSlackChannels, SlackSourceSlackChannelsDict +] -class GetRagFileConfig(_common.BaseModel): - """Config for getting a RAG corpus.""" +class SlackSource(_common.BaseModel): + """The Slack source for the ImportRagFilesRequest.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + channels: Optional[list[SlackSourceSlackChannels]] = Field( + default=None, description="""Required. The Slack channels.""" ) -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] - """""" +class SlackSourceDict(TypedDict, total=False): + """The Slack source for the ImportRagFilesRequest.""" - name: Optional[str] - """""" + channels: Optional[list[SlackSourceSlackChannelsDict]] + """Required. The Slack channels.""" -_GetRagFileRequestParametersOrDict = Union[ - _GetRagFileRequestParameters, _GetRagFileRequestParametersDict -] +SlackSourceOrDict = Union[SlackSource, SlackSourceDict] -class RagFileStatus(_common.BaseModel): - """RagFile status.""" +class RagFile(_common.BaseModel): + """A RAG File.""" - state: Optional[RagFileState] = Field( - default=None, description="""The state of the RagFile.""" + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this RagFile was created.""", + ) + description: Optional[str] = Field( + default=None, description="""Optional. The description of the RagFile.""" + ) + direct_upload_source: Optional[DirectUploadSource] = Field( + default=None, + description="""Output only. The RagFile is encapsulated and uploaded in the UploadRagFile request.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""Required. The display name of the RagFile. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", + ) + file_status: Optional[RagFileStatus] = Field( + default=None, description="""Output only. State of the RagFile.""" + ) + gcs_source: Optional[genai_types.GcsSource] = Field( + default=None, + description="""Output only. Google Cloud Storage location of the RagFile. It does not support wildcards in the Cloud Storage uri for now.""", + ) + google_drive_source: Optional[GoogleDriveSource] = Field( + default=None, + description="""Output only. Google Drive location. Supports importing individual files as well as Google Drive folders.""", + ) + jira_source: Optional[JiraSource] = Field( + default=None, description="""The RagFile is imported from a Jira query.""" + ) + name: Optional[str] = Field( + default=None, description="""Output only. The resource name of the RagFile.""" + ) + rag_file_type: Optional[RagFileType] = Field( + default=None, description="""Output only. The type of the RagFile.""" + ) + share_point_sources: Optional[SharePointSources] = Field( + default=None, + description="""The RagFile is imported from a SharePoint source.""", + ) + size_bytes: Optional[int] = Field( + default=None, description="""Output only. The size of the RagFile in bytes.""" + ) + slack_source: Optional[SlackSource] = Field( + default=None, description="""The RagFile is imported from a Slack channel.""" + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this RagFile was last updated.""", + ) + user_metadata: Optional[str] = Field( + default=None, + description="""Output only. The metadata for metadata search. The user_metadata Needs to be in JSON format.""", ) -class RagFileStatusDict(TypedDict, total=False): - """RagFile status.""" - - state: Optional[RagFileState] - """The state of the RagFile.""" - - -RagFileStatusOrDict = Union[RagFileStatus, RagFileStatusDict] +class RagFileDict(TypedDict, total=False): + """A RAG File.""" + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this RagFile was created.""" -class DirectUploadSource(_common.BaseModel): - """The input content is encapsulated and uploaded in the request.""" + description: Optional[str] + """Optional. The description of the RagFile.""" - pass + direct_upload_source: Optional[DirectUploadSourceDict] + """Output only. The RagFile is encapsulated and uploaded in the UploadRagFile request.""" + display_name: Optional[str] + """Required. The display name of the RagFile. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" -class DirectUploadSourceDict(TypedDict, total=False): - """The input content is encapsulated and uploaded in the request.""" + file_status: Optional[RagFileStatusDict] + """Output only. State of the RagFile.""" - pass + gcs_source: Optional[genai_types.GcsSourceDict] + """Output only. Google Cloud Storage location of the RagFile. It does not support wildcards in the Cloud Storage uri for now.""" + google_drive_source: Optional[GoogleDriveSourceDict] + """Output only. Google Drive location. Supports importing individual files as well as Google Drive folders.""" -DirectUploadSourceOrDict = Union[DirectUploadSource, DirectUploadSourceDict] + jira_source: Optional[JiraSourceDict] + """The RagFile is imported from a Jira query.""" + name: Optional[str] + """Output only. The resource name of the RagFile.""" -class GoogleDriveSourceResourceId(_common.BaseModel): - """The type and ID of the Google Drive resource.""" + rag_file_type: Optional[RagFileType] + """Output only. The type of the RagFile.""" - 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.""" - ) + share_point_sources: Optional[SharePointSourcesDict] + """The RagFile is imported from a SharePoint source.""" + size_bytes: Optional[int] + """Output only. The size of the RagFile in bytes.""" -class GoogleDriveSourceResourceIdDict(TypedDict, total=False): - """The type and ID of the Google Drive resource.""" + slack_source: Optional[SlackSourceDict] + """The RagFile is imported from a Slack channel.""" - resource_id: Optional[str] - """Required. The ID of the Google Drive resource.""" + update_time: Optional[datetime.datetime] + """Output only. Timestamp when this RagFile was last updated.""" - resource_type: Optional[ResourceType] - """Required. The type of the Google Drive resource.""" + user_metadata: Optional[str] + """Output only. The metadata for metadata search. The user_metadata Needs to be in JSON format.""" -GoogleDriveSourceResourceIdOrDict = Union[ - GoogleDriveSourceResourceId, GoogleDriveSourceResourceIdDict -] +RagFileOrDict = Union[RagFile, RagFileDict] -class GoogleDriveSource(_common.BaseModel): - """The Google Drive location for the input content.""" +class ListRagFilesConfig(_common.BaseModel): + """Config for listing RagFile instances within a RagCorpus.""" - resource_ids: Optional[list[GoogleDriveSourceResourceId]] = Field( - default=None, description="""Required. Google Drive resource IDs.""" + 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 GoogleDriveSourceDict(TypedDict, total=False): - """The Google Drive location for the input content.""" - - resource_ids: Optional[list[GoogleDriveSourceResourceIdDict]] - """Required. Google Drive resource IDs.""" +class ListRagFilesConfigDict(TypedDict, total=False): + """Config for listing RagFile instances within a RagCorpus.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -GoogleDriveSourceOrDict = Union[GoogleDriveSource, GoogleDriveSourceDict] + page_size: Optional[int] + """""" + page_token: Optional[str] + """""" -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.""" - ) +ListRagFilesConfigOrDict = Union[ListRagFilesConfig, ListRagFilesConfigDict] -class JiraSourceJiraQueriesDict(TypedDict, total=False): - """JiraQueries contains the Jira queries and corresponding authentication.""" +class _ListRagFilesRequestParameters(_common.BaseModel): + """Parameters for listing RagFile instances within a RagCorpus.""" - 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/).""" + config: Optional[ListRagFilesConfig] = Field(default=None, description="""""") + name: Optional[str] = Field(default=None, description="""""") - 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.""" +class _ListRagFilesRequestParametersDict(TypedDict, total=False): + """Parameters for listing RagFile instances within a RagCorpus.""" - projects: Optional[list[str]] - """A list of Jira projects to import in their entirety.""" + config: Optional[ListRagFilesConfigDict] + """""" - server_uri: Optional[str] - """Required. The Jira server URI.""" + name: Optional[str] + """""" -JiraSourceJiraQueriesOrDict = Union[JiraSourceJiraQueries, JiraSourceJiraQueriesDict] +_ListRagFilesRequestParametersOrDict = Union[ + _ListRagFilesRequestParameters, _ListRagFilesRequestParametersDict +] -class JiraSource(_common.BaseModel): - """The Jira source for the ImportRagFilesRequest.""" +class ListRagFilesResponse(_common.BaseModel): + """Response for listing RagFile instances within a RagCorpus.""" - jira_queries: Optional[list[JiraSourceJiraQueries]] = Field( - default=None, description="""Required. The Jira queries.""" + 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_files: Optional[list[RagFile]] = Field( + default=None, description="""List of RagFile instances within a RagCorpus.""" ) -class JiraSourceDict(TypedDict, total=False): - """The Jira source for the ImportRagFilesRequest.""" +class ListRagFilesResponseDict(TypedDict, total=False): + """Response for listing RagFile instances within a RagCorpus.""" - jira_queries: Optional[list[JiraSourceJiraQueriesDict]] - """Required. The Jira queries.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + next_page_token: Optional[str] + """""" -JiraSourceOrDict = Union[JiraSource, JiraSourceDict] + rag_files: Optional[list[RagFileDict]] + """List of RagFile instances within a RagCorpus.""" -class SharePointSourcesSharePointSource(_common.BaseModel): - """An individual SharePointSource.""" +ListRagFilesResponseOrDict = Union[ListRagFilesResponse, ListRagFilesResponseDict] - client_id: Optional[str] = Field( - default=None, - description="""The Application ID for the app registered in Microsoft Azure Portal. The application must also be configured with MS Graph permissions "Files.ReadAll", "Sites.ReadAll" and BrowserSiteLists.Read.All.""", - ) - client_secret: Optional[genai_types.ApiAuthApiKeyConfig] = Field( - default=None, - description="""The application secret for the app registered in Azure.""", - ) - drive_id: Optional[str] = Field( - default=None, description="""The ID of the drive to download from.""" - ) - drive_name: Optional[str] = Field( - default=None, description="""The name of the drive to download from.""" - ) - file_id: Optional[str] = Field( - default=None, - description="""Output only. The SharePoint file id. Output only.""", - ) - sharepoint_folder_id: Optional[str] = Field( - default=None, - description="""The ID of the SharePoint folder to download from.""", - ) - sharepoint_folder_path: Optional[str] = Field( - default=None, - description="""The path of the SharePoint folder to download from.""", - ) - sharepoint_site_name: Optional[str] = Field( - default=None, - description="""The name of the SharePoint site to download from. This can be the site name or the site id.""", - ) - tenant_id: Optional[str] = Field( - default=None, - description="""Unique identifier of the Azure Active Directory Instance.""", + +class GetRagConfig(_common.BaseModel): + """Config for getting a RAG Config.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class SharePointSourcesSharePointSourceDict(TypedDict, total=False): - """An individual SharePointSource.""" +class GetRagConfigDict(TypedDict, total=False): + """Config for getting a RAG Config.""" - client_id: Optional[str] - """The Application ID for the app registered in Microsoft Azure Portal. The application must also be configured with MS Graph permissions "Files.ReadAll", "Sites.ReadAll" and BrowserSiteLists.Read.All.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - client_secret: Optional[genai_types.ApiAuthApiKeyConfigDict] - """The application secret for the app registered in Azure.""" - drive_id: Optional[str] - """The ID of the drive to download from.""" +GetRagConfigOrDict = Union[GetRagConfig, GetRagConfigDict] - drive_name: Optional[str] - """The name of the drive to download from.""" - file_id: Optional[str] - """Output only. The SharePoint file id. Output only.""" +class _GetRagConfigRequestParameters(_common.BaseModel): + """Parameters for getting a RAG Config.""" - sharepoint_folder_id: Optional[str] - """The ID of the SharePoint folder to download from.""" + config: Optional[GetRagConfig] = Field(default=None, description="""""") - sharepoint_folder_path: Optional[str] - """The path of the SharePoint folder to download from.""" - sharepoint_site_name: Optional[str] - """The name of the SharePoint site to download from. This can be the site name or the site id.""" +class _GetRagConfigRequestParametersDict(TypedDict, total=False): + """Parameters for getting a RAG Config.""" - tenant_id: Optional[str] - """Unique identifier of the Azure Active Directory Instance.""" + config: Optional[GetRagConfigDict] + """""" -SharePointSourcesSharePointSourceOrDict = Union[ - SharePointSourcesSharePointSource, SharePointSourcesSharePointSourceDict +_GetRagConfigRequestParametersOrDict = Union[ + _GetRagConfigRequestParameters, _GetRagConfigRequestParametersDict ] -class SharePointSources(_common.BaseModel): - """The SharePointSources to pass to ImportRagFiles.""" +class RagManagedDbConfigBasic(_common.BaseModel): + """Basic tier is a cost-effective and low compute tier suitable for the following cases: * Experimenting with RagManagedDb. * Small data size. * Latency insensitive workload. * Only using RAG Engine with external vector DBs. NOTE: This is the default tier under Spanner mode if not explicitly chosen.""" - share_point_sources: Optional[list[SharePointSourcesSharePointSource]] = Field( - default=None, description="""The SharePoint sources.""" - ) + pass -class SharePointSourcesDict(TypedDict, total=False): - """The SharePointSources to pass to ImportRagFiles.""" +class RagManagedDbConfigBasicDict(TypedDict, total=False): + """Basic tier is a cost-effective and low compute tier suitable for the following cases: * Experimenting with RagManagedDb. * Small data size. * Latency insensitive workload. * Only using RAG Engine with external vector DBs. NOTE: This is the default tier under Spanner mode if not explicitly chosen.""" - share_point_sources: Optional[list[SharePointSourcesSharePointSourceDict]] - """The SharePoint sources.""" + pass -SharePointSourcesOrDict = Union[SharePointSources, SharePointSourcesDict] +RagManagedDbConfigBasicOrDict = Union[ + RagManagedDbConfigBasic, RagManagedDbConfigBasicDict +] -class SlackSourceSlackChannelsSlackChannel(_common.BaseModel): - """SlackChannel contains the Slack channel ID and the time range to import.""" +class RagManagedDbConfigEnterprise(_common.BaseModel): + """Enterprise tier offers production grade performance along with autoscaling functionality. It is suitable for customers with large amounts of data or performance sensitive workloads.""" - channel_id: Optional[str] = Field( - default=None, description="""Required. The Slack channel ID.""" - ) - end_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. The ending timestamp for messages to import.""", - ) - start_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. The starting timestamp for messages to import.""", - ) + pass -class SlackSourceSlackChannelsSlackChannelDict(TypedDict, total=False): - """SlackChannel contains the Slack channel ID and the time range to import.""" +class RagManagedDbConfigEnterpriseDict(TypedDict, total=False): + """Enterprise tier offers production grade performance along with autoscaling functionality. It is suitable for customers with large amounts of data or performance sensitive workloads.""" - channel_id: Optional[str] - """Required. The Slack channel ID.""" + pass - end_time: Optional[datetime.datetime] - """Optional. The ending timestamp for messages to import.""" - start_time: Optional[datetime.datetime] - """Optional. The starting timestamp for messages to import.""" +RagManagedDbConfigEnterpriseOrDict = Union[ + RagManagedDbConfigEnterprise, RagManagedDbConfigEnterpriseDict +] -SlackSourceSlackChannelsSlackChannelOrDict = Union[ - SlackSourceSlackChannelsSlackChannel, SlackSourceSlackChannelsSlackChannelDict -] +class RagManagedDbConfigScaled(_common.BaseModel): + """Scaled tier offers production grade performance along with autoscaling functionality. It is suitable for customers with large amounts of data or performance sensitive workloads.""" + pass -class SlackSourceSlackChannels(_common.BaseModel): - """SlackChannels contains the Slack channels and corresponding access token.""" - 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 Slack channel access token that has access to the slack channel IDs. See: https://api.slack.com/tutorials/tracks/getting-a-token.""", - ) - channels: Optional[list[SlackSourceSlackChannelsSlackChannel]] = Field( - default=None, description="""Required. The Slack channel IDs.""" - ) +class RagManagedDbConfigScaledDict(TypedDict, total=False): + """Scaled tier offers production grade performance along with autoscaling functionality. It is suitable for customers with large amounts of data or performance sensitive workloads.""" + pass -class SlackSourceSlackChannelsDict(TypedDict, total=False): - """SlackChannels contains the Slack channels and corresponding access token.""" - api_key_config: Optional[genai_types.ApiAuthApiKeyConfigDict] - """Required. The SecretManager secret version resource name (e.g. projects/{project}/secrets/{secret}/versions/{version}) storing the Slack channel access token that has access to the slack channel IDs. See: https://api.slack.com/tutorials/tracks/getting-a-token.""" +RagManagedDbConfigScaledOrDict = Union[ + RagManagedDbConfigScaled, RagManagedDbConfigScaledDict +] - channels: Optional[list[SlackSourceSlackChannelsSlackChannelDict]] - """Required. The Slack channel IDs.""" +class RagManagedDbConfigServerless(_common.BaseModel): + """Message to configure the serverless mode offered by RAG Engine.""" + + pass -SlackSourceSlackChannelsOrDict = Union[ - SlackSourceSlackChannels, SlackSourceSlackChannelsDict + +class RagManagedDbConfigServerlessDict(TypedDict, total=False): + """Message to configure the serverless mode offered by RAG Engine.""" + + pass + + +RagManagedDbConfigServerlessOrDict = Union[ + RagManagedDbConfigServerless, RagManagedDbConfigServerlessDict ] -class SlackSource(_common.BaseModel): - """The Slack source for the ImportRagFilesRequest.""" +class RagManagedDbConfigUnprovisioned(_common.BaseModel): + """Disables the RAG Engine service and deletes all your data held within this service. This will halt the billing of the service. NOTE: Once deleted the data cannot be recovered. To start using RAG Engine again, you will need to update the tier by calling the UpdateRagEngineConfig API.""" - channels: Optional[list[SlackSourceSlackChannels]] = Field( - default=None, description="""Required. The Slack channels.""" - ) + pass -class SlackSourceDict(TypedDict, total=False): - """The Slack source for the ImportRagFilesRequest.""" +class RagManagedDbConfigUnprovisionedDict(TypedDict, total=False): + """Disables the RAG Engine service and deletes all your data held within this service. This will halt the billing of the service. NOTE: Once deleted the data cannot be recovered. To start using RAG Engine again, you will need to update the tier by calling the UpdateRagEngineConfig API.""" - channels: Optional[list[SlackSourceSlackChannelsDict]] - """Required. The Slack channels.""" + pass -SlackSourceOrDict = Union[SlackSource, SlackSourceDict] +RagManagedDbConfigUnprovisionedOrDict = Union[ + RagManagedDbConfigUnprovisioned, RagManagedDbConfigUnprovisionedDict +] -class RagFile(_common.BaseModel): - """A RAG File.""" +class RagManagedDbConfigSpanner(_common.BaseModel): + """Message to configure the Spanner database used by RagManagedDb.""" - create_time: Optional[datetime.datetime] = Field( + basic: Optional[RagManagedDbConfigBasic] = Field( default=None, - description="""Output only. Timestamp when this RagFile was created.""", + description="""Sets the RagManagedDb to the Basic tier. This is the default tier for Spanner mode if not explicitly chosen.""", ) - description: Optional[str] = Field( - default=None, description="""Optional. The description of the RagFile.""" + scaled: Optional[RagManagedDbConfigScaled] = Field( + default=None, description="""Sets the RagManagedDb to the Scaled tier.""" ) - direct_upload_source: Optional[DirectUploadSource] = Field( - default=None, - description="""Output only. The RagFile is encapsulated and uploaded in the UploadRagFile request.""", + unprovisioned: Optional[RagManagedDbConfigUnprovisioned] = Field( + default=None, description="""Sets the RagManagedDb to the Unprovisioned tier.""" ) - display_name: Optional[str] = Field( + + +class RagManagedDbConfigSpannerDict(TypedDict, total=False): + """Message to configure the Spanner database used by RagManagedDb.""" + + basic: Optional[RagManagedDbConfigBasicDict] + """Sets the RagManagedDb to the Basic tier. This is the default tier for Spanner mode if not explicitly chosen.""" + + scaled: Optional[RagManagedDbConfigScaledDict] + """Sets the RagManagedDb to the Scaled tier.""" + + unprovisioned: Optional[RagManagedDbConfigUnprovisionedDict] + """Sets the RagManagedDb to the Unprovisioned tier.""" + + +RagManagedDbConfigSpannerOrDict = Union[ + RagManagedDbConfigSpanner, RagManagedDbConfigSpannerDict +] + + +class RagManagedDbConfig(_common.BaseModel): + """Configuration message for RagManagedDb used by RagEngine.""" + + basic: Optional[RagManagedDbConfigBasic] = Field( default=None, - description="""Required. The display name of the RagFile. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", - ) - file_status: Optional[RagFileStatus] = Field( - default=None, description="""Output only. State of the RagFile.""" + description="""Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Basic tier.""", ) - gcs_source: Optional[genai_types.GcsSource] = Field( - default=None, - description="""Output only. Google Cloud Storage location of the RagFile. It does not support wildcards in the Cloud Storage uri for now.""", + enterprise: Optional[RagManagedDbConfigEnterprise] = Field( + default=None, description="""Sets the RagManagedDb to the Enterprise tier.""" ) - google_drive_source: Optional[GoogleDriveSource] = Field( + scaled: Optional[RagManagedDbConfigScaled] = Field( default=None, - description="""Output only. Google Drive location. Supports importing individual files as well as Google Drive folders.""", - ) - jira_source: Optional[JiraSource] = Field( - default=None, description="""The RagFile is imported from a Jira query.""" - ) - name: Optional[str] = Field( - default=None, description="""Output only. The resource name of the RagFile.""" - ) - rag_file_type: Optional[RagFileType] = Field( - default=None, description="""Output only. The type of the RagFile.""" + description="""Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Scaled tier.""", ) - share_point_sources: Optional[SharePointSources] = Field( + serverless: Optional[RagManagedDbConfigServerless] = Field( default=None, - description="""The RagFile is imported from a SharePoint source.""", - ) - size_bytes: Optional[int] = Field( - default=None, description="""Output only. The size of the RagFile in bytes.""" - ) - slack_source: Optional[SlackSource] = Field( - default=None, description="""The RagFile is imported from a Slack channel.""" + description="""Sets the backend to be the serverless mode offered by RAG Engine.""", ) - update_time: Optional[datetime.datetime] = Field( + spanner: Optional[RagManagedDbConfigSpanner] = Field( default=None, - description="""Output only. Timestamp when this RagFile was last updated.""", + description="""Sets the RAG Engine backend to be RagManagedDb, built on top of Spanner. NOTE: This is the default mode (w/ Basic Tier) if not explicitly chosen.""", ) - user_metadata: Optional[str] = Field( + unprovisioned: Optional[RagManagedDbConfigUnprovisioned] = Field( default=None, - description="""Output only. The metadata for metadata search. The user_metadata Needs to be in JSON format.""", + description="""Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Unprovisioned tier.""", ) -class RagFileDict(TypedDict, total=False): - """A RAG File.""" +class RagManagedDbConfigDict(TypedDict, total=False): + """Configuration message for RagManagedDb used by RagEngine.""" - create_time: Optional[datetime.datetime] - """Output only. Timestamp when this RagFile was created.""" + basic: Optional[RagManagedDbConfigBasicDict] + """Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Basic tier.""" - description: Optional[str] - """Optional. The description of the RagFile.""" + enterprise: Optional[RagManagedDbConfigEnterpriseDict] + """Sets the RagManagedDb to the Enterprise tier.""" - direct_upload_source: Optional[DirectUploadSourceDict] - """Output only. The RagFile is encapsulated and uploaded in the UploadRagFile request.""" + scaled: Optional[RagManagedDbConfigScaledDict] + """Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Scaled tier.""" - display_name: Optional[str] - """Required. The display name of the RagFile. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" + serverless: Optional[RagManagedDbConfigServerlessDict] + """Sets the backend to be the serverless mode offered by RAG Engine.""" - file_status: Optional[RagFileStatusDict] - """Output only. State of the RagFile.""" + spanner: Optional[RagManagedDbConfigSpannerDict] + """Sets the RAG Engine backend to be RagManagedDb, built on top of Spanner. NOTE: This is the default mode (w/ Basic Tier) if not explicitly chosen.""" - gcs_source: Optional[genai_types.GcsSourceDict] - """Output only. Google Cloud Storage location of the RagFile. It does not support wildcards in the Cloud Storage uri for now.""" + unprovisioned: Optional[RagManagedDbConfigUnprovisionedDict] + """Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Unprovisioned tier.""" - google_drive_source: Optional[GoogleDriveSourceDict] - """Output only. Google Drive location. Supports importing individual files as well as Google Drive folders.""" - jira_source: Optional[JiraSourceDict] - """The RagFile is imported from a Jira query.""" +RagManagedDbConfigOrDict = Union[RagManagedDbConfig, RagManagedDbConfigDict] - name: Optional[str] - """Output only. The resource name of the RagFile.""" - rag_file_type: Optional[RagFileType] - """Output only. The type of the RagFile.""" +class RagEngineConfig(_common.BaseModel): + """The config of the RAG Engine.""" - share_point_sources: Optional[SharePointSourcesDict] - """The RagFile is imported from a SharePoint source.""" + name: Optional[str] = Field( + default=None, + description="""Identifier. The name of the RagEngineConfig. Format: `projects/{project}/locations/{location}/ragEngineConfig`""", + ) + rag_managed_db_config: Optional[RagManagedDbConfig] = Field( + default=None, + description="""The config of the RagManagedDb used by RagEngine.""", + ) - size_bytes: Optional[int] - """Output only. The size of the RagFile in bytes.""" - slack_source: Optional[SlackSourceDict] - """The RagFile is imported from a Slack channel.""" +class RagEngineConfigDict(TypedDict, total=False): + """The config of the RAG Engine.""" - update_time: Optional[datetime.datetime] - """Output only. Timestamp when this RagFile was last updated.""" + name: Optional[str] + """Identifier. The name of the RagEngineConfig. Format: `projects/{project}/locations/{location}/ragEngineConfig`""" - user_metadata: Optional[str] - """Output only. The metadata for metadata search. The user_metadata Needs to be in JSON format.""" + rag_managed_db_config: Optional[RagManagedDbConfigDict] + """The config of the RagManagedDb used by RagEngine.""" -RagFileOrDict = Union[RagFile, RagFileDict] +RagEngineConfigOrDict = Union[RagEngineConfig, RagEngineConfigDict] -class ListRagFilesConfig(_common.BaseModel): - """Config for listing RagFile instances within a RagCorpus.""" +class UpdateRagCorpusConfig(_common.BaseModel): + """Config for updating a RAG corpus.""" 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 ListRagFilesConfigDict(TypedDict, total=False): - """Config for listing RagFile instances within a RagCorpus.""" +class UpdateRagCorpusConfigDict(TypedDict, total=False): + """Config for updating a RAG corpus.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - page_size: Optional[int] - """""" - - page_token: Optional[str] - """""" - -ListRagFilesConfigOrDict = Union[ListRagFilesConfig, ListRagFilesConfigDict] +UpdateRagCorpusConfigOrDict = Union[UpdateRagCorpusConfig, UpdateRagCorpusConfigDict] -class _ListRagFilesRequestParameters(_common.BaseModel): - """Parameters for listing RagFile instances within a RagCorpus.""" +class _UpdateRagCorpusRequestParameters(_common.BaseModel): + """Parameters for updating a RAG corpus.""" - config: Optional[ListRagFilesConfig] = Field(default=None, description="""""") + config: Optional[UpdateRagCorpusConfig] = Field(default=None, description="""""") name: Optional[str] = Field(default=None, description="""""") + rag_corpus: Optional[RagCorpus] = Field(default=None, description="""""") -class _ListRagFilesRequestParametersDict(TypedDict, total=False): - """Parameters for listing RagFile instances within a RagCorpus.""" +class _UpdateRagCorpusRequestParametersDict(TypedDict, total=False): + """Parameters for updating a RAG corpus.""" - config: Optional[ListRagFilesConfigDict] + config: Optional[UpdateRagCorpusConfigDict] """""" name: Optional[str] """""" + rag_corpus: Optional[RagCorpusDict] + """""" + -_ListRagFilesRequestParametersOrDict = Union[ - _ListRagFilesRequestParameters, _ListRagFilesRequestParametersDict +_UpdateRagCorpusRequestParametersOrDict = Union[ + _UpdateRagCorpusRequestParameters, _UpdateRagCorpusRequestParametersDict ] -class ListRagFilesResponse(_common.BaseModel): - """Response for listing RagFile instances within a RagCorpus.""" +class UpdateRagCorpusOperation(_common.BaseModel): + """Operation for updating a RAG corpus.""" - 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="""""") - rag_files: Optional[list[RagFile]] = Field( - default=None, description="""List of RagFile instances within a RagCorpus.""" + 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 ListRagFilesResponseDict(TypedDict, total=False): - """Response for listing RagFile instances within a RagCorpus.""" +class UpdateRagCorpusOperationDict(TypedDict, total=False): + """Operation for updating a RAG corpus.""" - 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.""" - rag_files: Optional[list[RagFileDict]] - """List of RagFile instances within a RagCorpus.""" + 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.""" -ListRagFilesResponseOrDict = Union[ListRagFilesResponse, ListRagFilesResponseDict] +UpdateRagCorpusOperationOrDict = Union[ + UpdateRagCorpusOperation, UpdateRagCorpusOperationDict +] -class GetRagConfig(_common.BaseModel): - """Config for getting a RAG Config.""" + +class DeleteRagCorpusConfig(_common.BaseModel): + """Config for deleting a RAG corpus.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetRagConfigDict(TypedDict, total=False): - """Config for getting a RAG Config.""" +class DeleteRagCorpusConfigDict(TypedDict, total=False): + """Config for deleting a RAG corpus.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetRagConfigOrDict = Union[GetRagConfig, GetRagConfigDict] +DeleteRagCorpusConfigOrDict = Union[DeleteRagCorpusConfig, DeleteRagCorpusConfigDict] -class _GetRagConfigRequestParameters(_common.BaseModel): - """Parameters for getting a RAG Config.""" +class _DeleteRagCorpusRequestParameters(_common.BaseModel): + """Parameters for deleting a RAG corpus.""" - config: Optional[GetRagConfig] = Field(default=None, description="""""") + config: Optional[DeleteRagCorpusConfig] = Field(default=None, description="""""") + name: Optional[str] = Field(default=None, description="""""") -class _GetRagConfigRequestParametersDict(TypedDict, total=False): - """Parameters for getting a RAG Config.""" +class _DeleteRagCorpusRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a RAG corpus.""" - config: Optional[GetRagConfigDict] + config: Optional[DeleteRagCorpusConfigDict] + """""" + + name: Optional[str] """""" -_GetRagConfigRequestParametersOrDict = Union[ - _GetRagConfigRequestParameters, _GetRagConfigRequestParametersDict +_DeleteRagCorpusRequestParametersOrDict = Union[ + _DeleteRagCorpusRequestParameters, _DeleteRagCorpusRequestParametersDict ] -class RagManagedDbConfigBasic(_common.BaseModel): - """Basic tier is a cost-effective and low compute tier suitable for the following cases: * Experimenting with RagManagedDb. * Small data size. * Latency insensitive workload. * Only using RAG Engine with external vector DBs. NOTE: This is the default tier under Spanner mode if not explicitly chosen.""" +class DeleteRagCorpusOperation(_common.BaseModel): + """Operation for deleting a RAG corpus.""" - pass + 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 RagManagedDbConfigBasicDict(TypedDict, total=False): - """Basic tier is a cost-effective and low compute tier suitable for the following cases: * Experimenting with RagManagedDb. * Small data size. * Latency insensitive workload. * Only using RAG Engine with external vector DBs. NOTE: This is the default tier under Spanner mode if not explicitly chosen.""" +class DeleteRagCorpusOperationDict(TypedDict, total=False): + """Operation for deleting a RAG corpus.""" - pass + 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.""" -RagManagedDbConfigBasicOrDict = Union[ - RagManagedDbConfigBasic, RagManagedDbConfigBasicDict -] + 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 RagManagedDbConfigEnterprise(_common.BaseModel): - """Enterprise tier offers production grade performance along with autoscaling functionality. It is suitable for customers with large amounts of data or performance sensitive workloads.""" - pass +DeleteRagCorpusOperationOrDict = Union[ + DeleteRagCorpusOperation, DeleteRagCorpusOperationDict +] -class RagManagedDbConfigEnterpriseDict(TypedDict, total=False): - """Enterprise tier offers production grade performance along with autoscaling functionality. It is suitable for customers with large amounts of data or performance sensitive workloads.""" +class DeleteRagFileConfig(_common.BaseModel): + """Config for deleting a RAG File.""" - pass + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -RagManagedDbConfigEnterpriseOrDict = Union[ - RagManagedDbConfigEnterprise, RagManagedDbConfigEnterpriseDict -] +class DeleteRagFileConfigDict(TypedDict, total=False): + """Config for deleting a RAG File.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -class RagManagedDbConfigScaled(_common.BaseModel): - """Scaled tier offers production grade performance along with autoscaling functionality. It is suitable for customers with large amounts of data or performance sensitive workloads.""" - pass +DeleteRagFileConfigOrDict = Union[DeleteRagFileConfig, DeleteRagFileConfigDict] -class RagManagedDbConfigScaledDict(TypedDict, total=False): - """Scaled tier offers production grade performance along with autoscaling functionality. It is suitable for customers with large amounts of data or performance sensitive workloads.""" +class _DeleteRagFileRequestParameters(_common.BaseModel): + """Parameters for deleting a RAG File.""" - pass + config: Optional[DeleteRagFileConfig] = Field(default=None, description="""""") + name: Optional[str] = Field(default=None, description="""""") -RagManagedDbConfigScaledOrDict = Union[ - RagManagedDbConfigScaled, RagManagedDbConfigScaledDict -] +class _DeleteRagFileRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a RAG File.""" + config: Optional[DeleteRagFileConfigDict] + """""" -class RagManagedDbConfigServerless(_common.BaseModel): - """Message to configure the serverless mode offered by RAG Engine.""" + name: Optional[str] + """""" - pass +_DeleteRagFileRequestParametersOrDict = Union[ + _DeleteRagFileRequestParameters, _DeleteRagFileRequestParametersDict +] -class RagManagedDbConfigServerlessDict(TypedDict, total=False): - """Message to configure the serverless mode offered by RAG Engine.""" - pass +class DeleteRagFileOperation(_common.BaseModel): + """Operation for deleting a RAG File.""" + 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.""", + ) -RagManagedDbConfigServerlessOrDict = Union[ - RagManagedDbConfigServerless, RagManagedDbConfigServerlessDict -] +class DeleteRagFileOperationDict(TypedDict, total=False): + """Operation for deleting a RAG File.""" -class RagManagedDbConfigUnprovisioned(_common.BaseModel): - """Disables the RAG Engine service and deletes all your data held within this service. This will halt the billing of the service. NOTE: Once deleted the data cannot be recovered. To start using RAG Engine again, you will need to update the tier by calling the UpdateRagEngineConfig API.""" - - pass + 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 RagManagedDbConfigUnprovisionedDict(TypedDict, total=False): - """Disables the RAG Engine service and deletes all your data held within this service. This will halt the billing of the service. NOTE: Once deleted the data cannot be recovered. To start using RAG Engine again, you will need to update the tier by calling the UpdateRagEngineConfig API.""" + 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.""" - pass + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -RagManagedDbConfigUnprovisionedOrDict = Union[ - RagManagedDbConfigUnprovisioned, RagManagedDbConfigUnprovisionedDict -] +DeleteRagFileOperationOrDict = Union[DeleteRagFileOperation, DeleteRagFileOperationDict] -class RagManagedDbConfigSpanner(_common.BaseModel): - """Message to configure the Spanner database used by RagManagedDb.""" +class UpdateRagConfig(_common.BaseModel): + """Config for updating a RAG Config.""" - basic: Optional[RagManagedDbConfigBasic] = Field( - default=None, - description="""Sets the RagManagedDb to the Basic tier. This is the default tier for Spanner mode if not explicitly chosen.""", - ) - scaled: Optional[RagManagedDbConfigScaled] = Field( - default=None, description="""Sets the RagManagedDb to the Scaled tier.""" - ) - unprovisioned: Optional[RagManagedDbConfigUnprovisioned] = Field( - default=None, description="""Sets the RagManagedDb to the Unprovisioned tier.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class RagManagedDbConfigSpannerDict(TypedDict, total=False): - """Message to configure the Spanner database used by RagManagedDb.""" +class UpdateRagConfigDict(TypedDict, total=False): + """Config for updating a RAG Config.""" - basic: Optional[RagManagedDbConfigBasicDict] - """Sets the RagManagedDb to the Basic tier. This is the default tier for Spanner mode if not explicitly chosen.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - scaled: Optional[RagManagedDbConfigScaledDict] - """Sets the RagManagedDb to the Scaled tier.""" - unprovisioned: Optional[RagManagedDbConfigUnprovisionedDict] - """Sets the RagManagedDb to the Unprovisioned tier.""" +UpdateRagConfigOrDict = Union[UpdateRagConfig, UpdateRagConfigDict] -RagManagedDbConfigSpannerOrDict = Union[ - RagManagedDbConfigSpanner, RagManagedDbConfigSpannerDict +class _UpdateRagConfigRequestParameters(_common.BaseModel): + """Parameters for updating a RAG Config.""" + + updated_config: Optional[RagEngineConfig] = Field(default=None, description="""""") + config: Optional[UpdateRagConfig] = Field(default=None, description="""""") + + +class _UpdateRagConfigRequestParametersDict(TypedDict, total=False): + """Parameters for updating a RAG Config.""" + + updated_config: Optional[RagEngineConfigDict] + """""" + + config: Optional[UpdateRagConfigDict] + """""" + + +_UpdateRagConfigRequestParametersOrDict = Union[ + _UpdateRagConfigRequestParameters, _UpdateRagConfigRequestParametersDict ] -class RagManagedDbConfig(_common.BaseModel): - """Configuration message for RagManagedDb used by RagEngine.""" +class UpdateRagConfigOperation(_common.BaseModel): + """Operation for updating a RAG Config.""" - basic: Optional[RagManagedDbConfigBasic] = Field( - default=None, - description="""Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Basic tier.""", - ) - enterprise: Optional[RagManagedDbConfigEnterprise] = Field( - default=None, description="""Sets the RagManagedDb to the Enterprise tier.""" - ) - scaled: Optional[RagManagedDbConfigScaled] = Field( + name: Optional[str] = Field( default=None, - description="""Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Scaled tier.""", + 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}`.""", ) - serverless: Optional[RagManagedDbConfigServerless] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Sets the backend to be the serverless mode offered by RAG Engine.""", + 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.""", ) - spanner: Optional[RagManagedDbConfigSpanner] = Field( + done: Optional[bool] = Field( default=None, - description="""Sets the RAG Engine backend to be RagManagedDb, built on top of Spanner. NOTE: This is the default mode (w/ Basic Tier) if not explicitly chosen.""", + 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.""", ) - unprovisioned: Optional[RagManagedDbConfigUnprovisioned] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Unprovisioned tier.""", + description="""The error result of the operation in case of failure or cancellation.""", ) -class RagManagedDbConfigDict(TypedDict, total=False): - """Configuration message for RagManagedDb used by RagEngine.""" +class UpdateRagConfigOperationDict(TypedDict, total=False): + """Operation for updating a RAG Config.""" - basic: Optional[RagManagedDbConfigBasicDict] - """Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Basic tier.""" + 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}`.""" - enterprise: Optional[RagManagedDbConfigEnterpriseDict] - """Sets the RagManagedDb to the Enterprise tier.""" + 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.""" - scaled: Optional[RagManagedDbConfigScaledDict] - """Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Scaled tier.""" + 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.""" - serverless: Optional[RagManagedDbConfigServerlessDict] - """Sets the backend to be the serverless mode offered by RAG Engine.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - spanner: Optional[RagManagedDbConfigSpannerDict] - """Sets the RAG Engine backend to be RagManagedDb, built on top of Spanner. NOTE: This is the default mode (w/ Basic Tier) if not explicitly chosen.""" - unprovisioned: Optional[RagManagedDbConfigUnprovisionedDict] - """Deprecated: Use `mode` instead to set the tier under Spanner. Sets the RagManagedDb to the Unprovisioned tier.""" +UpdateRagConfigOperationOrDict = Union[ + UpdateRagConfigOperation, UpdateRagConfigOperationDict +] -RagManagedDbConfigOrDict = Union[RagManagedDbConfig, RagManagedDbConfigDict] +class RetrieveContextsConfig(_common.BaseModel): + """Config for retrieving RAG Contexts.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -class RagEngineConfig(_common.BaseModel): - """The config of the RAG Engine.""" - name: Optional[str] = Field( - default=None, - description="""Identifier. The name of the RagEngineConfig. Format: `projects/{project}/locations/{location}/ragEngineConfig`""", - ) - rag_managed_db_config: Optional[RagManagedDbConfig] = Field( - default=None, - description="""The config of the RagManagedDb used by RagEngine.""", +class RetrieveContextsConfigDict(TypedDict, total=False): + """Config for retrieving RAG Contexts.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +RetrieveContextsConfigOrDict = Union[RetrieveContextsConfig, RetrieveContextsConfigDict] + + +class _RetrieveRagContextsRequestParameters(_common.BaseModel): + """Parameters for retrieving RAG Contexts.""" + + vertex_rag_store: Optional[genai_types.VertexRagStore] = Field( + default=None, description="""""" ) + query: Optional[RagQuery] = Field(default=None, description="""""") + config: Optional[RetrieveContextsConfig] = Field(default=None, description="""""") -class RagEngineConfigDict(TypedDict, total=False): - """The config of the RAG Engine.""" +class _RetrieveRagContextsRequestParametersDict(TypedDict, total=False): + """Parameters for retrieving RAG Contexts.""" - name: Optional[str] - """Identifier. The name of the RagEngineConfig. Format: `projects/{project}/locations/{location}/ragEngineConfig`""" + vertex_rag_store: Optional[genai_types.VertexRagStoreDict] + """""" - rag_managed_db_config: Optional[RagManagedDbConfigDict] - """The config of the RagManagedDb used by RagEngine.""" + query: Optional[RagQueryDict] + """""" + + config: Optional[RetrieveContextsConfigDict] + """""" -RagEngineConfigOrDict = Union[RagEngineConfig, RagEngineConfigDict] +_RetrieveRagContextsRequestParametersOrDict = Union[ + _RetrieveRagContextsRequestParameters, _RetrieveRagContextsRequestParametersDict +] -class UpdateRagCorpusConfig(_common.BaseModel): - """Config for updating a RAG corpus.""" +class RetrieveContextsResponse(_common.BaseModel): + + contexts: Optional[RagContexts] = Field( + default=None, description="""The contexts of the query.""" + ) + + +class RetrieveContextsResponseDict(TypedDict, total=False): + + contexts: Optional[RagContextsDict] + """The contexts of the query.""" + + +RetrieveContextsResponseOrDict = Union[ + RetrieveContextsResponse, RetrieveContextsResponseDict +] + + +class GetRagConfigOperationConfig(_common.BaseModel): + """Config for getting a RAG config operation.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class UpdateRagCorpusConfigDict(TypedDict, total=False): - """Config for updating a RAG corpus.""" +class GetRagConfigOperationConfigDict(TypedDict, total=False): + """Config for getting a RAG config operation.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -UpdateRagCorpusConfigOrDict = Union[UpdateRagCorpusConfig, UpdateRagCorpusConfigDict] - +GetRagConfigOperationConfigOrDict = Union[ + GetRagConfigOperationConfig, GetRagConfigOperationConfigDict +] -class _UpdateRagCorpusRequestParameters(_common.BaseModel): - """Parameters for updating a RAG corpus.""" - config: Optional[UpdateRagCorpusConfig] = Field(default=None, description="""""") - name: Optional[str] = Field(default=None, description="""""") - rag_corpus: Optional[RagCorpus] = Field(default=None, description="""""") +class _GetRagConfigOperationParameters(_common.BaseModel): + """Parameters for getting a RAG config operation.""" + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetRagConfigOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" + ) -class _UpdateRagCorpusRequestParametersDict(TypedDict, total=False): - """Parameters for updating a RAG corpus.""" - config: Optional[UpdateRagCorpusConfigDict] - """""" +class _GetRagConfigOperationParametersDict(TypedDict, total=False): + """Parameters for getting a RAG config operation.""" - name: Optional[str] - """""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - rag_corpus: Optional[RagCorpusDict] - """""" + config: Optional[GetRagConfigOperationConfigDict] + """Used to override the default configuration.""" -_UpdateRagCorpusRequestParametersOrDict = Union[ - _UpdateRagCorpusRequestParameters, _UpdateRagCorpusRequestParametersDict +_GetRagConfigOperationParametersOrDict = Union[ + _GetRagConfigOperationParameters, _GetRagConfigOperationParametersDict ] -class UpdateRagCorpusOperation(_common.BaseModel): - """Operation for updating a RAG corpus.""" +class RagEngineConfigOperation(_common.BaseModel): + """Operation for getting a RAG config.""" name: Optional[str] = Field( default=None, @@ -15777,8 +15400,8 @@ class UpdateRagCorpusOperation(_common.BaseModel): ) -class UpdateRagCorpusOperationDict(TypedDict, total=False): - """Operation for updating a RAG corpus.""" +class RagEngineConfigOperationDict(TypedDict, total=False): + """Operation for getting a RAG config.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -15793,713 +15416,614 @@ class UpdateRagCorpusOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -UpdateRagCorpusOperationOrDict = Union[ - UpdateRagCorpusOperation, UpdateRagCorpusOperationDict +RagEngineConfigOperationOrDict = Union[ + RagEngineConfigOperation, RagEngineConfigOperationDict ] -class DeleteRagCorpusConfig(_common.BaseModel): - """Config for deleting a RAG corpus.""" +class ImportRagFilesRequestConfig(_common.BaseModel): + """Config for importing RAG files.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class DeleteRagCorpusConfigDict(TypedDict, total=False): - """Config for deleting a RAG corpus.""" +class ImportRagFilesRequestConfigDict(TypedDict, total=False): + """Config for importing RAG files.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteRagCorpusConfigOrDict = Union[DeleteRagCorpusConfig, DeleteRagCorpusConfigDict] +ImportRagFilesRequestConfigOrDict = Union[ + ImportRagFilesRequestConfig, ImportRagFilesRequestConfigDict +] -class _DeleteRagCorpusRequestParameters(_common.BaseModel): - """Parameters for deleting a RAG corpus.""" - - config: Optional[DeleteRagCorpusConfig] = Field(default=None, description="""""") - name: Optional[str] = Field(default=None, description="""""") +class BigQueryDestination(_common.BaseModel): + """The BigQuery location for the output content.""" + output_uri: Optional[str] = Field( + default=None, + description="""Required. BigQuery URI to a project or table, up to 2000 characters long. When only the project is specified, the Dataset and Table is created. When the full table reference is specified, the Dataset must exist and table must not exist. Accepted forms: * BigQuery path. For example: `bq://projectId` or `bq://projectId.bqDatasetId` or `bq://projectId.bqDatasetId.bqTableId`.""", + ) -class _DeleteRagCorpusRequestParametersDict(TypedDict, total=False): - """Parameters for deleting a RAG corpus.""" - config: Optional[DeleteRagCorpusConfigDict] - """""" +class BigQueryDestinationDict(TypedDict, total=False): + """The BigQuery location for the output content.""" - name: Optional[str] - """""" + output_uri: Optional[str] + """Required. BigQuery URI to a project or table, up to 2000 characters long. When only the project is specified, the Dataset and Table is created. When the full table reference is specified, the Dataset must exist and table must not exist. Accepted forms: * BigQuery path. For example: `bq://projectId` or `bq://projectId.bqDatasetId` or `bq://projectId.bqDatasetId.bqTableId`.""" -_DeleteRagCorpusRequestParametersOrDict = Union[ - _DeleteRagCorpusRequestParameters, _DeleteRagCorpusRequestParametersDict -] +BigQueryDestinationOrDict = Union[BigQueryDestination, BigQueryDestinationDict] -class DeleteRagCorpusOperation(_common.BaseModel): - """Operation for deleting a RAG corpus.""" +class RagFileChunkingConfigFixedLengthChunking(_common.BaseModel): + """Specifies the fixed length chunking config.""" - 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.""", + chunk_overlap: Optional[int] = Field( + default=None, description="""The overlap between chunks.""" ) - error: Optional[dict[str, Any]] = Field( - default=None, - description="""The error result of the operation in case of failure or cancellation.""", + chunk_size: Optional[int] = Field( + default=None, description="""The size of the chunks.""" ) -class DeleteRagCorpusOperationDict(TypedDict, total=False): - """Operation for deleting a RAG corpus.""" - - 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 RagFileChunkingConfigFixedLengthChunkingDict(TypedDict, total=False): + """Specifies the fixed length chunking config.""" - 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.""" + chunk_overlap: Optional[int] + """The overlap between chunks.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + chunk_size: Optional[int] + """The size of the chunks.""" -DeleteRagCorpusOperationOrDict = Union[ - DeleteRagCorpusOperation, DeleteRagCorpusOperationDict +RagFileChunkingConfigFixedLengthChunkingOrDict = Union[ + RagFileChunkingConfigFixedLengthChunking, + RagFileChunkingConfigFixedLengthChunkingDict, ] -class DeleteRagFileConfig(_common.BaseModel): - """Config for deleting a RAG File.""" +class RagFileChunkingConfig(_common.BaseModel): + """Specifies the size and overlap of chunks for RagFiles.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + chunk_overlap: Optional[int] = Field( + default=None, description="""The overlap between chunks.""" + ) + chunk_size: Optional[int] = Field( + default=None, description="""The size of the chunks.""" + ) + fixed_length_chunking: Optional[RagFileChunkingConfigFixedLengthChunking] = Field( + default=None, description="""Specifies the fixed length chunking config.""" ) -class DeleteRagFileConfigDict(TypedDict, total=False): - """Config for deleting a RAG File.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -DeleteRagFileConfigOrDict = Union[DeleteRagFileConfig, DeleteRagFileConfigDict] - - -class _DeleteRagFileRequestParameters(_common.BaseModel): - """Parameters for deleting a RAG File.""" - - config: Optional[DeleteRagFileConfig] = Field(default=None, description="""""") - name: Optional[str] = Field(default=None, description="""""") - +class RagFileChunkingConfigDict(TypedDict, total=False): + """Specifies the size and overlap of chunks for RagFiles.""" -class _DeleteRagFileRequestParametersDict(TypedDict, total=False): - """Parameters for deleting a RAG File.""" + chunk_overlap: Optional[int] + """The overlap between chunks.""" - config: Optional[DeleteRagFileConfigDict] - """""" + chunk_size: Optional[int] + """The size of the chunks.""" - name: Optional[str] - """""" + fixed_length_chunking: Optional[RagFileChunkingConfigFixedLengthChunkingDict] + """Specifies the fixed length chunking config.""" -_DeleteRagFileRequestParametersOrDict = Union[ - _DeleteRagFileRequestParameters, _DeleteRagFileRequestParametersDict -] +RagFileChunkingConfigOrDict = Union[RagFileChunkingConfig, RagFileChunkingConfigDict] -class DeleteRagFileOperation(_common.BaseModel): - """Operation for deleting a RAG File.""" +class RagFileMetadataConfig(_common.BaseModel): + """Metadata config for RagFile.""" - name: Optional[str] = Field( + gcs_metadata_schema_source: Optional[genai_types.GcsSource] = 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="""Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/metadata_schema.json` - `gs://bucket_name/my_directory` If the user provides a directory, the metadata schema will be read from the files that ends with "metadata_schema.json" in the directory.""", ) - metadata: Optional[dict[str, Any]] = Field( + gcs_metadata_source: Optional[genai_types.GcsSource] = 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="""Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/metadata.json` - `gs://bucket_name/my_directory` If the user provides a directory, the metadata will be read from the files that ends with "metadata.json" in the directory.""", ) - done: Optional[bool] = Field( + google_drive_metadata_schema_source: Optional[GoogleDriveSource] = 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="""Google Drive location. Supports importing individual files as well as Google Drive folders. If the user provides a folder, the metadata schema will be read from the files that ends with "metadata_schema.json" in the directory.""", ) - error: Optional[dict[str, Any]] = Field( + google_drive_metadata_source: Optional[GoogleDriveSource] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Google Drive location. Supports importing individual files as well as Google Drive folders. If the user provides a directory, the metadata will be read from the files that ends with "metadata.json" in the directory.""", + ) + inline_metadata_schema_source: Optional[str] = Field( + default=None, + description="""Inline metadata schema source. Must be a JSON string.""", + ) + inline_metadata_source: Optional[str] = Field( + default=None, description="""Inline metadata source. Must be a JSON string.""" ) -class DeleteRagFileOperationDict(TypedDict, total=False): - """Operation for deleting a RAG File.""" - - 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.""" - - -DeleteRagFileOperationOrDict = Union[DeleteRagFileOperation, DeleteRagFileOperationDict] - +class RagFileMetadataConfigDict(TypedDict, total=False): + """Metadata config for RagFile.""" -class UpdateRagConfig(_common.BaseModel): - """Config for updating a RAG Config.""" + gcs_metadata_schema_source: Optional[genai_types.GcsSourceDict] + """Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/metadata_schema.json` - `gs://bucket_name/my_directory` If the user provides a directory, the metadata schema will be read from the files that ends with "metadata_schema.json" in the directory.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + gcs_metadata_source: Optional[genai_types.GcsSourceDict] + """Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/metadata.json` - `gs://bucket_name/my_directory` If the user provides a directory, the metadata will be read from the files that ends with "metadata.json" in the directory.""" + google_drive_metadata_schema_source: Optional[GoogleDriveSourceDict] + """Google Drive location. Supports importing individual files as well as Google Drive folders. If the user provides a folder, the metadata schema will be read from the files that ends with "metadata_schema.json" in the directory.""" -class UpdateRagConfigDict(TypedDict, total=False): - """Config for updating a RAG Config.""" + google_drive_metadata_source: Optional[GoogleDriveSourceDict] + """Google Drive location. Supports importing individual files as well as Google Drive folders. If the user provides a directory, the metadata will be read from the files that ends with "metadata.json" in the directory.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + inline_metadata_schema_source: Optional[str] + """Inline metadata schema source. Must be a JSON string.""" + inline_metadata_source: Optional[str] + """Inline metadata source. Must be a JSON string.""" -UpdateRagConfigOrDict = Union[UpdateRagConfig, UpdateRagConfigDict] +RagFileMetadataConfigOrDict = Union[RagFileMetadataConfig, RagFileMetadataConfigDict] -class _UpdateRagConfigRequestParameters(_common.BaseModel): - """Parameters for updating a RAG Config.""" - updated_config: Optional[RagEngineConfig] = Field(default=None, description="""""") - config: Optional[UpdateRagConfig] = Field(default=None, description="""""") +class RagFileParsingConfigAdvancedParser(_common.BaseModel): + """Specifies the advanced parsing for RagFiles.""" + use_advanced_pdf_parsing: Optional[bool] = Field( + default=None, description="""Whether to use advanced PDF parsing.""" + ) -class _UpdateRagConfigRequestParametersDict(TypedDict, total=False): - """Parameters for updating a RAG Config.""" - updated_config: Optional[RagEngineConfigDict] - """""" +class RagFileParsingConfigAdvancedParserDict(TypedDict, total=False): + """Specifies the advanced parsing for RagFiles.""" - config: Optional[UpdateRagConfigDict] - """""" + use_advanced_pdf_parsing: Optional[bool] + """Whether to use advanced PDF parsing.""" -_UpdateRagConfigRequestParametersOrDict = Union[ - _UpdateRagConfigRequestParameters, _UpdateRagConfigRequestParametersDict +RagFileParsingConfigAdvancedParserOrDict = Union[ + RagFileParsingConfigAdvancedParser, RagFileParsingConfigAdvancedParserDict ] -class UpdateRagConfigOperation(_common.BaseModel): - """Operation for updating a RAG Config.""" +class RagFileParsingConfigLayoutParser(_common.BaseModel): + """Document AI Layout Parser config.""" - 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( + global_max_parsing_requests_per_min: Optional[int] = 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="""The maximum number of requests the job is allowed to make to the Document AI processor per minute in this project. Consult https://cloud.google.com/document-ai/quotas and the Quota page for your project to set an appropriate value here. If this value is not specified, max_parsing_requests_per_min will be used by indexing pipeline as the global limit.""", ) - done: Optional[bool] = Field( + max_parsing_requests_per_min: Optional[int] = Field( default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + description="""The maximum number of requests the job is allowed to make to the Document AI processor per minute. Consult https://cloud.google.com/document-ai/quotas and the Quota page for your project to set an appropriate value here. If unspecified, a default value of 120 QPM would be used.""", ) - error: Optional[dict[str, Any]] = Field( + processor_name: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""The full resource name of a Document AI processor or processor version. The processor must have type `LAYOUT_PARSER_PROCESSOR`. If specified, the `additional_config.parse_as_scanned_pdf` field must be false. Format: * `projects/{project_id}/locations/{location}/processors/{processor_id}` * `projects/{project_id}/locations/{location}/processors/{processor_id}/processorVersions/{processor_version_id}`""", ) -class UpdateRagConfigOperationDict(TypedDict, total=False): - """Operation for updating a RAG Config.""" - - name: Optional[str] - """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" +class RagFileParsingConfigLayoutParserDict(TypedDict, total=False): + """Document AI Layout Parser config.""" - 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.""" + global_max_parsing_requests_per_min: Optional[int] + """The maximum number of requests the job is allowed to make to the Document AI processor per minute in this project. Consult https://cloud.google.com/document-ai/quotas and the Quota page for your project to set an appropriate value here. If this value is not specified, max_parsing_requests_per_min will be used by indexing pipeline as the global limit.""" - 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.""" + max_parsing_requests_per_min: Optional[int] + """The maximum number of requests the job is allowed to make to the Document AI processor per minute. Consult https://cloud.google.com/document-ai/quotas and the Quota page for your project to set an appropriate value here. If unspecified, a default value of 120 QPM would be used.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + processor_name: Optional[str] + """The full resource name of a Document AI processor or processor version. The processor must have type `LAYOUT_PARSER_PROCESSOR`. If specified, the `additional_config.parse_as_scanned_pdf` field must be false. Format: * `projects/{project_id}/locations/{location}/processors/{processor_id}` * `projects/{project_id}/locations/{location}/processors/{processor_id}/processorVersions/{processor_version_id}`""" -UpdateRagConfigOperationOrDict = Union[ - UpdateRagConfigOperation, UpdateRagConfigOperationDict +RagFileParsingConfigLayoutParserOrDict = Union[ + RagFileParsingConfigLayoutParser, RagFileParsingConfigLayoutParserDict ] -class RetrieveContextsConfig(_common.BaseModel): - """Config for retrieving RAG Contexts.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" +class RagFileParsingConfig(_common.BaseModel): + """Specifies the parsing config for RagFiles.""" + + advanced_parser: Optional[RagFileParsingConfigAdvancedParser] = Field( + default=None, description="""The Advanced Parser to use for RagFiles.""" + ) + layout_parser: Optional[RagFileParsingConfigLayoutParser] = Field( + default=None, description="""The Layout Parser to use for RagFiles.""" + ) + llm_parser: Optional[RagFileParsingConfigLlmParser] = Field( + default=None, description="""The LLM Parser to use for RagFiles.""" + ) + use_advanced_pdf_parsing: Optional[bool] = Field( + default=None, description="""Whether to use advanced PDF parsing.""" ) -class RetrieveContextsConfigDict(TypedDict, total=False): - """Config for retrieving RAG Contexts.""" +class RagFileParsingConfigDict(TypedDict, total=False): + """Specifies the parsing config for RagFiles.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + advanced_parser: Optional[RagFileParsingConfigAdvancedParserDict] + """The Advanced Parser to use for RagFiles.""" + layout_parser: Optional[RagFileParsingConfigLayoutParserDict] + """The Layout Parser to use for RagFiles.""" -RetrieveContextsConfigOrDict = Union[RetrieveContextsConfig, RetrieveContextsConfigDict] + llm_parser: Optional[RagFileParsingConfigLlmParserDict] + """The LLM Parser to use for RagFiles.""" + use_advanced_pdf_parsing: Optional[bool] + """Whether to use advanced PDF parsing.""" -class _RetrieveRagContextsRequestParameters(_common.BaseModel): - """Parameters for retrieving RAG Contexts.""" - vertex_rag_store: Optional[genai_types.VertexRagStore] = Field( - default=None, description="""""" - ) - query: Optional[RagQuery] = Field(default=None, description="""""") - config: Optional[RetrieveContextsConfig] = Field(default=None, description="""""") +RagFileParsingConfigOrDict = Union[RagFileParsingConfig, RagFileParsingConfigDict] -class _RetrieveRagContextsRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving RAG Contexts.""" +class RagFileTransformationConfig(_common.BaseModel): + """Specifies the transformation config for RagFiles.""" - vertex_rag_store: Optional[genai_types.VertexRagStoreDict] - """""" + rag_file_chunking_config: Optional[RagFileChunkingConfig] = Field( + default=None, description="""Specifies the chunking config for RagFiles.""" + ) - query: Optional[RagQueryDict] - """""" - config: Optional[RetrieveContextsConfigDict] - """""" +class RagFileTransformationConfigDict(TypedDict, total=False): + """Specifies the transformation config for RagFiles.""" + + rag_file_chunking_config: Optional[RagFileChunkingConfigDict] + """Specifies the chunking config for RagFiles.""" -_RetrieveRagContextsRequestParametersOrDict = Union[ - _RetrieveRagContextsRequestParameters, _RetrieveRagContextsRequestParametersDict +RagFileTransformationConfigOrDict = Union[ + RagFileTransformationConfig, RagFileTransformationConfigDict ] -class RetrieveContextsResponse(_common.BaseModel): +class ImportRagFilesConfig(_common.BaseModel): + """Config for importing RagFiles.""" - contexts: Optional[RagContexts] = Field( - default=None, description="""The contexts of the query.""" + gcs_source: Optional[genai_types.GcsSource] = Field( + default=None, + description="""Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/my_file.txt` - `gs://bucket_name/my_directory`""", + ) + global_max_embedding_requests_per_min: Optional[int] = Field( + default=None, + description="""Optional. The max number of queries per minute that the indexing pipeline job is allowed to make to the embedding model specified in the project. Please follow the quota usage guideline of the embedding model you use to set the value properly.If this value is not specified, max_embedding_requests_per_min will be used by indexing pipeline job as the global limit.""", + ) + google_drive_source: Optional[GoogleDriveSource] = Field( + default=None, + description="""Google Drive location. Supports importing individual files as well as Google Drive folders.""", + ) + import_result_bigquery_sink: Optional[BigQueryDestination] = Field( + default=None, + description="""The BigQuery destination to write import result to. It should be a bigquery table resource name (e.g. "bq://projectId.bqDatasetId.bqTableId"). The dataset must exist. If the table does not exist, it will be created with the expected schema. If the table exists, the schema will be validated and data will be added to this existing table.""", + ) + import_result_gcs_sink: Optional[genai_types.GcsDestination] = Field( + default=None, + description="""The Cloud Storage path to write import result to.""", + ) + jira_source: Optional[JiraSource] = Field( + default=None, + description="""Jira queries with their corresponding authentication.""", + ) + max_embedding_requests_per_min: Optional[int] = Field( + default=None, + description="""Optional. The max number of queries per minute that this job is allowed to make to the embedding model specified on the corpus. This value is specific to this job and not shared across other import jobs. Consult the Quotas page on the project to set an appropriate value here. If unspecified, a default value of 1,000 QPM would be used.""", + ) + partial_failure_bigquery_sink: Optional[BigQueryDestination] = Field( + default=None, + description="""The BigQuery destination to write partial failures to. It should be a bigquery table resource name (e.g. "bq://projectId.bqDatasetId.bqTableId"). The dataset must exist. If the table does not exist, it will be created with the expected schema. If the table exists, the schema will be validated and data will be added to this existing table. Deprecated. Prefer to use `import_result_bq_sink`.""", + ) + partial_failure_gcs_sink: Optional[genai_types.GcsDestination] = Field( + default=None, + description="""The Cloud Storage path to write partial failures to. Deprecated. Prefer to use `import_result_gcs_sink`.""", + ) + rag_file_chunking_config: Optional[RagFileChunkingConfig] = Field( + default=None, + description="""Specifies the size and overlap of chunks after importing RagFiles.""", + ) + rag_file_metadata_config: Optional[RagFileMetadataConfig] = Field( + default=None, + description="""Specifies the metadata config for RagFiles. Including paths for metadata schema and metadata. Deprecated: Not in use.""", + ) + rag_file_parsing_config: Optional[RagFileParsingConfig] = Field( + default=None, + description="""Optional. Specifies the parsing config for RagFiles. RAG will use the default parser if this field is not set.""", + ) + rag_file_transformation_config: Optional[RagFileTransformationConfig] = Field( + default=None, + description="""Specifies the transformation config for RagFiles.""", + ) + rebuild_ann_index: Optional[bool] = Field( + default=None, + description="""Rebuilds the ANN index to optimize for recall on the imported data. Only applicable for RagCorpora running on RagManagedDb with `retrieval_strategy` set to `ANN`. The rebuild will be performed using the existing ANN config set on the RagCorpus. To change the ANN config, please use the UpdateRagCorpus API. Default is false, i.e., index is not rebuilt.""", + ) + share_point_sources: Optional[SharePointSources] = Field( + default=None, description="""SharePoint sources.""" + ) + slack_source: Optional[SlackSource] = Field( + default=None, + description="""Slack channels with their corresponding access tokens.""", ) -class RetrieveContextsResponseDict(TypedDict, total=False): - - contexts: Optional[RagContextsDict] - """The contexts of the query.""" - +class ImportRagFilesConfigDict(TypedDict, total=False): + """Config for importing RagFiles.""" -RetrieveContextsResponseOrDict = Union[ - RetrieveContextsResponse, RetrieveContextsResponseDict -] + gcs_source: Optional[genai_types.GcsSourceDict] + """Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/my_file.txt` - `gs://bucket_name/my_directory`""" + global_max_embedding_requests_per_min: Optional[int] + """Optional. The max number of queries per minute that the indexing pipeline job is allowed to make to the embedding model specified in the project. Please follow the quota usage guideline of the embedding model you use to set the value properly.If this value is not specified, max_embedding_requests_per_min will be used by indexing pipeline job as the global limit.""" -class GetRagConfigOperationConfig(_common.BaseModel): - """Config for getting a RAG config operation.""" + google_drive_source: Optional[GoogleDriveSourceDict] + """Google Drive location. Supports importing individual files as well as Google Drive folders.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + import_result_bigquery_sink: Optional[BigQueryDestinationDict] + """The BigQuery destination to write import result to. It should be a bigquery table resource name (e.g. "bq://projectId.bqDatasetId.bqTableId"). The dataset must exist. If the table does not exist, it will be created with the expected schema. If the table exists, the schema will be validated and data will be added to this existing table.""" + import_result_gcs_sink: Optional[genai_types.GcsDestinationDict] + """The Cloud Storage path to write import result to.""" -class GetRagConfigOperationConfigDict(TypedDict, total=False): - """Config for getting a RAG config operation.""" + jira_source: Optional[JiraSourceDict] + """Jira queries with their corresponding authentication.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + max_embedding_requests_per_min: Optional[int] + """Optional. The max number of queries per minute that this job is allowed to make to the embedding model specified on the corpus. This value is specific to this job and not shared across other import jobs. Consult the Quotas page on the project to set an appropriate value here. If unspecified, a default value of 1,000 QPM would be used.""" + partial_failure_bigquery_sink: Optional[BigQueryDestinationDict] + """The BigQuery destination to write partial failures to. It should be a bigquery table resource name (e.g. "bq://projectId.bqDatasetId.bqTableId"). The dataset must exist. If the table does not exist, it will be created with the expected schema. If the table exists, the schema will be validated and data will be added to this existing table. Deprecated. Prefer to use `import_result_bq_sink`.""" -GetRagConfigOperationConfigOrDict = Union[ - GetRagConfigOperationConfig, GetRagConfigOperationConfigDict -] + partial_failure_gcs_sink: Optional[genai_types.GcsDestinationDict] + """The Cloud Storage path to write partial failures to. Deprecated. Prefer to use `import_result_gcs_sink`.""" + rag_file_chunking_config: Optional[RagFileChunkingConfigDict] + """Specifies the size and overlap of chunks after importing RagFiles.""" -class _GetRagConfigOperationParameters(_common.BaseModel): - """Parameters for getting a RAG config operation.""" + rag_file_metadata_config: Optional[RagFileMetadataConfigDict] + """Specifies the metadata config for RagFiles. Including paths for metadata schema and metadata. Deprecated: Not in use.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" - ) - config: Optional[GetRagConfigOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" - ) + rag_file_parsing_config: Optional[RagFileParsingConfigDict] + """Optional. Specifies the parsing config for RagFiles. RAG will use the default parser if this field is not set.""" + rag_file_transformation_config: Optional[RagFileTransformationConfigDict] + """Specifies the transformation config for RagFiles.""" -class _GetRagConfigOperationParametersDict(TypedDict, total=False): - """Parameters for getting a RAG config operation.""" + rebuild_ann_index: Optional[bool] + """Rebuilds the ANN index to optimize for recall on the imported data. Only applicable for RagCorpora running on RagManagedDb with `retrieval_strategy` set to `ANN`. The rebuild will be performed using the existing ANN config set on the RagCorpus. To change the ANN config, please use the UpdateRagCorpus API. Default is false, i.e., index is not rebuilt.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + share_point_sources: Optional[SharePointSourcesDict] + """SharePoint sources.""" - config: Optional[GetRagConfigOperationConfigDict] - """Used to override the default configuration.""" + slack_source: Optional[SlackSourceDict] + """Slack channels with their corresponding access tokens.""" -_GetRagConfigOperationParametersOrDict = Union[ - _GetRagConfigOperationParameters, _GetRagConfigOperationParametersDict -] +ImportRagFilesConfigOrDict = Union[ImportRagFilesConfig, ImportRagFilesConfigDict] -class RagEngineConfigOperation(_common.BaseModel): - """Operation for getting a RAG config.""" +class ImportRagFilesRequest(_common.BaseModel): + """Request message for VertexRagDataService.ImportRagFiles.""" - 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( + import_rag_files_config: Optional[ImportRagFilesConfig] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Required. The config for the RagFiles to be synced and imported into the RagCorpus. VertexRagDataService.ImportRagFiles.""", ) -class RagEngineConfigOperationDict(TypedDict, total=False): - """Operation for getting a RAG config.""" +class ImportRagFilesRequestDict(TypedDict, total=False): + """Request message for VertexRagDataService.ImportRagFiles.""" - 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}`.""" + import_rag_files_config: Optional[ImportRagFilesConfigDict] + """Required. The config for the RagFiles to be synced and imported into the RagCorpus. VertexRagDataService.ImportRagFiles.""" - 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.""" +ImportRagFilesRequestOrDict = Union[ImportRagFilesRequest, ImportRagFilesRequestDict] - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" +class _ImportRagFilesRequestParameters(_common.BaseModel): + """Parameters for importing RAG files.""" -RagEngineConfigOperationOrDict = Union[ - RagEngineConfigOperation, RagEngineConfigOperationDict -] + name: Optional[str] = Field(default=None, description="""""") + import_rag_files_request: Optional[ImportRagFilesRequest] = Field( + default=None, description="""""" + ) + config: Optional[ImportRagFilesRequestConfig] = Field( + default=None, description="""""" + ) -class ImportRagFilesRequestConfig(_common.BaseModel): - """Config for importing RAG files.""" +class _ImportRagFilesRequestParametersDict(TypedDict, total=False): + """Parameters for importing RAG files.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + name: Optional[str] + """""" + import_rag_files_request: Optional[ImportRagFilesRequestDict] + """""" -class ImportRagFilesRequestConfigDict(TypedDict, total=False): - """Config for importing RAG files.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + config: Optional[ImportRagFilesRequestConfigDict] + """""" -ImportRagFilesRequestConfigOrDict = Union[ - ImportRagFilesRequestConfig, ImportRagFilesRequestConfigDict +_ImportRagFilesRequestParametersOrDict = Union[ + _ImportRagFilesRequestParameters, _ImportRagFilesRequestParametersDict ] -class BigQueryDestination(_common.BaseModel): - """The BigQuery location for the output content.""" +class ImportRagFilesResponse(_common.BaseModel): + """Response from importing RAG files.""" - output_uri: Optional[str] = Field( + imported_rag_files_count: Optional[int] = Field( default=None, - description="""Required. BigQuery URI to a project or table, up to 2000 characters long. When only the project is specified, the Dataset and Table is created. When the full table reference is specified, the Dataset must exist and table must not exist. Accepted forms: * BigQuery path. For example: `bq://projectId` or `bq://projectId.bqDatasetId` or `bq://projectId.bqDatasetId.bqTableId`.""", - ) - - -class BigQueryDestinationDict(TypedDict, total=False): - """The BigQuery location for the output content.""" - - output_uri: Optional[str] - """Required. BigQuery URI to a project or table, up to 2000 characters long. When only the project is specified, the Dataset and Table is created. When the full table reference is specified, the Dataset must exist and table must not exist. Accepted forms: * BigQuery path. For example: `bq://projectId` or `bq://projectId.bqDatasetId` or `bq://projectId.bqDatasetId.bqTableId`.""" - - -BigQueryDestinationOrDict = Union[BigQueryDestination, BigQueryDestinationDict] - - -class RagFileChunkingConfigFixedLengthChunking(_common.BaseModel): - """Specifies the fixed length chunking config.""" - - chunk_overlap: Optional[int] = Field( - default=None, description="""The overlap between chunks.""" + description="""The number of RagFiles imported into the RagCorpus.""", ) - chunk_size: Optional[int] = Field( - default=None, description="""The size of the chunks.""" + failed_rag_files_count: Optional[int] = Field( + default=None, + description="""The number of RagFiles that failed to import into the RagCorpus.""", ) - - -class RagFileChunkingConfigFixedLengthChunkingDict(TypedDict, total=False): - """Specifies the fixed length chunking config.""" - - chunk_overlap: Optional[int] - """The overlap between chunks.""" - - chunk_size: Optional[int] - """The size of the chunks.""" - - -RagFileChunkingConfigFixedLengthChunkingOrDict = Union[ - RagFileChunkingConfigFixedLengthChunking, - RagFileChunkingConfigFixedLengthChunkingDict, -] - - -class RagFileChunkingConfig(_common.BaseModel): - """Specifies the size and overlap of chunks for RagFiles.""" - - chunk_overlap: Optional[int] = Field( - default=None, description="""The overlap between chunks.""" + skipped_rag_files_count: Optional[int] = Field( + default=None, + description="""The number of RagFiles skipped while importing into the RagCorpus.""", ) - chunk_size: Optional[int] = Field( - default=None, description="""The size of the chunks.""" + partial_failures_gcs_path: Optional[str] = Field( + default=None, + description="""The Google Cloud Storage path with the partial failures.""", ) - fixed_length_chunking: Optional[RagFileChunkingConfigFixedLengthChunking] = Field( - default=None, description="""Specifies the fixed length chunking config.""" + partial_failures_bigquery_table: Optional[str] = Field( + default=None, + description="""The BigQuery table where the partial failures were written.""", ) -class RagFileChunkingConfigDict(TypedDict, total=False): - """Specifies the size and overlap of chunks for RagFiles.""" +class ImportRagFilesResponseDict(TypedDict, total=False): + """Response from importing RAG files.""" - chunk_overlap: Optional[int] - """The overlap between chunks.""" + imported_rag_files_count: Optional[int] + """The number of RagFiles imported into the RagCorpus.""" - chunk_size: Optional[int] - """The size of the chunks.""" + failed_rag_files_count: Optional[int] + """The number of RagFiles that failed to import into the RagCorpus.""" - fixed_length_chunking: Optional[RagFileChunkingConfigFixedLengthChunkingDict] - """Specifies the fixed length chunking config.""" + skipped_rag_files_count: Optional[int] + """The number of RagFiles skipped while importing into the RagCorpus.""" + partial_failures_gcs_path: Optional[str] + """The Google Cloud Storage path with the partial failures.""" -RagFileChunkingConfigOrDict = Union[RagFileChunkingConfig, RagFileChunkingConfigDict] + partial_failures_bigquery_table: Optional[str] + """The BigQuery table where the partial failures were written.""" -class RagFileMetadataConfig(_common.BaseModel): - """Metadata config for RagFile.""" +ImportRagFilesResponseOrDict = Union[ImportRagFilesResponse, ImportRagFilesResponseDict] - gcs_metadata_schema_source: Optional[genai_types.GcsSource] = Field( - default=None, - description="""Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/metadata_schema.json` - `gs://bucket_name/my_directory` If the user provides a directory, the metadata schema will be read from the files that ends with "metadata_schema.json" in the directory.""", - ) - gcs_metadata_source: Optional[genai_types.GcsSource] = Field( + +class ImportRagFilesOperation(_common.BaseModel): + """Operation for importing RAG files.""" + + name: Optional[str] = Field( default=None, - description="""Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/metadata.json` - `gs://bucket_name/my_directory` If the user provides a directory, the metadata will be read from the files that ends with "metadata.json" in the directory.""", + 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}`.""", ) - google_drive_metadata_schema_source: Optional[GoogleDriveSource] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Google Drive location. Supports importing individual files as well as Google Drive folders. If the user provides a folder, the metadata schema will be read from the files that ends with "metadata_schema.json" in the directory.""", + 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.""", ) - google_drive_metadata_source: Optional[GoogleDriveSource] = Field( + done: Optional[bool] = Field( default=None, - description="""Google Drive location. Supports importing individual files as well as Google Drive folders. If the user provides a directory, the metadata will be read from the files that ends with "metadata.json" in the directory.""", + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", ) - inline_metadata_schema_source: Optional[str] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Inline metadata schema source. Must be a JSON string.""", + description="""The error result of the operation in case of failure or cancellation.""", ) - inline_metadata_source: Optional[str] = Field( - default=None, description="""Inline metadata source. Must be a JSON string.""" + response: Optional[ImportRagFilesResponse] = Field( + default=None, description="""The response from the import operation.""" ) -class RagFileMetadataConfigDict(TypedDict, total=False): - """Metadata config for RagFile.""" - - gcs_metadata_schema_source: Optional[genai_types.GcsSourceDict] - """Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/metadata_schema.json` - `gs://bucket_name/my_directory` If the user provides a directory, the metadata schema will be read from the files that ends with "metadata_schema.json" in the directory.""" +class ImportRagFilesOperationDict(TypedDict, total=False): + """Operation for importing RAG files.""" - gcs_metadata_source: Optional[genai_types.GcsSourceDict] - """Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/metadata.json` - `gs://bucket_name/my_directory` If the user provides a directory, the metadata will be read from the files that ends with "metadata.json" in the directory.""" + 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}`.""" - google_drive_metadata_schema_source: Optional[GoogleDriveSourceDict] - """Google Drive location. Supports importing individual files as well as Google Drive folders. If the user provides a folder, the metadata schema will be read from the files that ends with "metadata_schema.json" in the directory.""" + 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.""" - google_drive_metadata_source: Optional[GoogleDriveSourceDict] - """Google Drive location. Supports importing individual files as well as Google Drive folders. If the user provides a directory, the metadata will be read from the files that ends with "metadata.json" in the directory.""" + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - inline_metadata_schema_source: Optional[str] - """Inline metadata schema source. Must be a JSON string.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - inline_metadata_source: Optional[str] - """Inline metadata source. Must be a JSON string.""" + response: Optional[ImportRagFilesResponseDict] + """The response from the import operation.""" -RagFileMetadataConfigOrDict = Union[RagFileMetadataConfig, RagFileMetadataConfigDict] +ImportRagFilesOperationOrDict = Union[ + ImportRagFilesOperation, ImportRagFilesOperationDict +] -class RagFileParsingConfigAdvancedParser(_common.BaseModel): - """Specifies the advanced parsing for RagFiles.""" +class GetImportFilesOperationConfig(_common.BaseModel): + """Config for getting an import files operation.""" - use_advanced_pdf_parsing: Optional[bool] = Field( - default=None, description="""Whether to use advanced PDF parsing.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class RagFileParsingConfigAdvancedParserDict(TypedDict, total=False): - """Specifies the advanced parsing for RagFiles.""" +class GetImportFilesOperationConfigDict(TypedDict, total=False): + """Config for getting an import files operation.""" - use_advanced_pdf_parsing: Optional[bool] - """Whether to use advanced PDF parsing.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -RagFileParsingConfigAdvancedParserOrDict = Union[ - RagFileParsingConfigAdvancedParser, RagFileParsingConfigAdvancedParserDict +GetImportFilesOperationConfigOrDict = Union[ + GetImportFilesOperationConfig, GetImportFilesOperationConfigDict ] -class RagFileParsingConfigLayoutParser(_common.BaseModel): - """Document AI Layout Parser config.""" +class _GetImportFilesOperationParameters(_common.BaseModel): + """Parameters for getting an import files operation.""" - 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 Document AI processor per minute in this project. Consult https://cloud.google.com/document-ai/quotas and the Quota page for your project to set an appropriate value here. If this value is not specified, max_parsing_requests_per_min will be used by indexing pipeline as the global limit.""", - ) - max_parsing_requests_per_min: Optional[int] = Field( - default=None, - description="""The maximum number of requests the job is allowed to make to the Document AI processor per minute. Consult https://cloud.google.com/document-ai/quotas and the Quota page for your project to set an appropriate value here. If unspecified, a default value of 120 QPM would be used.""", + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - processor_name: Optional[str] = Field( - default=None, - description="""The full resource name of a Document AI processor or processor version. The processor must have type `LAYOUT_PARSER_PROCESSOR`. If specified, the `additional_config.parse_as_scanned_pdf` field must be false. Format: * `projects/{project_id}/locations/{location}/processors/{processor_id}` * `projects/{project_id}/locations/{location}/processors/{processor_id}/processorVersions/{processor_version_id}`""", + config: Optional[GetImportFilesOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -class RagFileParsingConfigLayoutParserDict(TypedDict, total=False): - """Document AI Layout Parser config.""" - - global_max_parsing_requests_per_min: Optional[int] - """The maximum number of requests the job is allowed to make to the Document AI processor per minute in this project. Consult https://cloud.google.com/document-ai/quotas and the Quota page for your project to set an appropriate value here. If this value is not specified, max_parsing_requests_per_min will be used by indexing pipeline as the global limit.""" +class _GetImportFilesOperationParametersDict(TypedDict, total=False): + """Parameters for getting an import files operation.""" - max_parsing_requests_per_min: Optional[int] - """The maximum number of requests the job is allowed to make to the Document AI processor per minute. Consult https://cloud.google.com/document-ai/quotas and the Quota page for your project to set an appropriate value here. If unspecified, a default value of 120 QPM would be used.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - processor_name: Optional[str] - """The full resource name of a Document AI processor or processor version. The processor must have type `LAYOUT_PARSER_PROCESSOR`. If specified, the `additional_config.parse_as_scanned_pdf` field must be false. Format: * `projects/{project_id}/locations/{location}/processors/{processor_id}` * `projects/{project_id}/locations/{location}/processors/{processor_id}/processorVersions/{processor_version_id}`""" + config: Optional[GetImportFilesOperationConfigDict] + """Used to override the default configuration.""" -RagFileParsingConfigLayoutParserOrDict = Union[ - RagFileParsingConfigLayoutParser, RagFileParsingConfigLayoutParserDict +_GetImportFilesOperationParametersOrDict = Union[ + _GetImportFilesOperationParameters, _GetImportFilesOperationParametersDict ] -class RagFileParsingConfig(_common.BaseModel): - """Specifies the parsing config for RagFiles.""" +class UploadRagFileRequestConfig(_common.BaseModel): + """Config for the request to upload a Rag File.""" - advanced_parser: Optional[RagFileParsingConfigAdvancedParser] = Field( - default=None, description="""The Advanced Parser to use for RagFiles.""" - ) - layout_parser: Optional[RagFileParsingConfigLayoutParser] = Field( - default=None, description="""The Layout Parser to use for RagFiles.""" - ) - llm_parser: Optional[RagFileParsingConfigLlmParser] = Field( - default=None, description="""The LLM Parser to use for RagFiles.""" - ) - use_advanced_pdf_parsing: Optional[bool] = Field( - default=None, description="""Whether to use advanced PDF parsing.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class RagFileParsingConfigDict(TypedDict, total=False): - """Specifies the parsing config for RagFiles.""" +class UploadRagFileRequestConfigDict(TypedDict, total=False): + """Config for the request to upload a Rag File.""" - advanced_parser: Optional[RagFileParsingConfigAdvancedParserDict] - """The Advanced Parser to use for RagFiles.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - layout_parser: Optional[RagFileParsingConfigLayoutParserDict] - """The Layout Parser to use for RagFiles.""" - llm_parser: Optional[RagFileParsingConfigLlmParserDict] - """The LLM Parser to use for RagFiles.""" +UploadRagFileRequestConfigOrDict = Union[ + UploadRagFileRequestConfig, UploadRagFileRequestConfigDict +] - use_advanced_pdf_parsing: Optional[bool] - """Whether to use advanced PDF parsing.""" +class UploadRagFileConfig(_common.BaseModel): + """Config for uploading RagFile.""" -RagFileParsingConfigOrDict = Union[RagFileParsingConfig, RagFileParsingConfigDict] - - -class RagFileTransformationConfig(_common.BaseModel): - """Specifies the transformation config for RagFiles.""" - - rag_file_chunking_config: Optional[RagFileChunkingConfig] = Field( - default=None, description="""Specifies the chunking config for RagFiles.""" - ) - - -class RagFileTransformationConfigDict(TypedDict, total=False): - """Specifies the transformation config for RagFiles.""" - - rag_file_chunking_config: Optional[RagFileChunkingConfigDict] - """Specifies the chunking config for RagFiles.""" - - -RagFileTransformationConfigOrDict = Union[ - RagFileTransformationConfig, RagFileTransformationConfigDict -] - - -class ImportRagFilesConfig(_common.BaseModel): - """Config for importing RagFiles.""" - - gcs_source: Optional[genai_types.GcsSource] = Field( - default=None, - description="""Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/my_file.txt` - `gs://bucket_name/my_directory`""", - ) - global_max_embedding_requests_per_min: Optional[int] = Field( - default=None, - description="""Optional. The max number of queries per minute that the indexing pipeline job is allowed to make to the embedding model specified in the project. Please follow the quota usage guideline of the embedding model you use to set the value properly.If this value is not specified, max_embedding_requests_per_min will be used by indexing pipeline job as the global limit.""", - ) - google_drive_source: Optional[GoogleDriveSource] = Field( - default=None, - description="""Google Drive location. Supports importing individual files as well as Google Drive folders.""", - ) - import_result_bigquery_sink: Optional[BigQueryDestination] = Field( - default=None, - description="""The BigQuery destination to write import result to. It should be a bigquery table resource name (e.g. "bq://projectId.bqDatasetId.bqTableId"). The dataset must exist. If the table does not exist, it will be created with the expected schema. If the table exists, the schema will be validated and data will be added to this existing table.""", - ) - import_result_gcs_sink: Optional[genai_types.GcsDestination] = Field( - default=None, - description="""The Cloud Storage path to write import result to.""", - ) - jira_source: Optional[JiraSource] = Field( - default=None, - description="""Jira queries with their corresponding authentication.""", - ) - max_embedding_requests_per_min: Optional[int] = Field( - default=None, - description="""Optional. The max number of queries per minute that this job is allowed to make to the embedding model specified on the corpus. This value is specific to this job and not shared across other import jobs. Consult the Quotas page on the project to set an appropriate value here. If unspecified, a default value of 1,000 QPM would be used.""", - ) - partial_failure_bigquery_sink: Optional[BigQueryDestination] = Field( - default=None, - description="""The BigQuery destination to write partial failures to. It should be a bigquery table resource name (e.g. "bq://projectId.bqDatasetId.bqTableId"). The dataset must exist. If the table does not exist, it will be created with the expected schema. If the table exists, the schema will be validated and data will be added to this existing table. Deprecated. Prefer to use `import_result_bq_sink`.""", - ) - partial_failure_gcs_sink: Optional[genai_types.GcsDestination] = Field( - default=None, - description="""The Cloud Storage path to write partial failures to. Deprecated. Prefer to use `import_result_gcs_sink`.""", - ) rag_file_chunking_config: Optional[RagFileChunkingConfig] = Field( default=None, - description="""Specifies the size and overlap of chunks after importing RagFiles.""", + description="""Specifies the size and overlap of chunks after uploading RagFile.""", ) rag_file_metadata_config: Optional[RagFileMetadataConfig] = Field( default=None, - description="""Specifies the metadata config for RagFiles. Including paths for metadata schema and metadata. Deprecated: Not in use.""", + description="""Optional. Specifies the metadata config for RagFiles. Including paths for metadata schema and metadata. Alteratively, inline metadata schema and metadata can be provided. Deprecated: Not in use.""", ) rag_file_parsing_config: Optional[RagFileParsingConfig] = Field( default=None, @@ -16509,54 +16033,16 @@ class ImportRagFilesConfig(_common.BaseModel): default=None, description="""Specifies the transformation config for RagFiles.""", ) - rebuild_ann_index: Optional[bool] = Field( - default=None, - description="""Rebuilds the ANN index to optimize for recall on the imported data. Only applicable for RagCorpora running on RagManagedDb with `retrieval_strategy` set to `ANN`. The rebuild will be performed using the existing ANN config set on the RagCorpus. To change the ANN config, please use the UpdateRagCorpus API. Default is false, i.e., index is not rebuilt.""", - ) - share_point_sources: Optional[SharePointSources] = Field( - default=None, description="""SharePoint sources.""" - ) - slack_source: Optional[SlackSource] = Field( - default=None, - description="""Slack channels with their corresponding access tokens.""", - ) - - -class ImportRagFilesConfigDict(TypedDict, total=False): - """Config for importing RagFiles.""" - - gcs_source: Optional[genai_types.GcsSourceDict] - """Google Cloud Storage location. Supports importing individual files as well as entire Google Cloud Storage directories. Sample formats: - `gs://bucket_name/my_directory/object_name/my_file.txt` - `gs://bucket_name/my_directory`""" - - global_max_embedding_requests_per_min: Optional[int] - """Optional. The max number of queries per minute that the indexing pipeline job is allowed to make to the embedding model specified in the project. Please follow the quota usage guideline of the embedding model you use to set the value properly.If this value is not specified, max_embedding_requests_per_min will be used by indexing pipeline job as the global limit.""" - - google_drive_source: Optional[GoogleDriveSourceDict] - """Google Drive location. Supports importing individual files as well as Google Drive folders.""" - - import_result_bigquery_sink: Optional[BigQueryDestinationDict] - """The BigQuery destination to write import result to. It should be a bigquery table resource name (e.g. "bq://projectId.bqDatasetId.bqTableId"). The dataset must exist. If the table does not exist, it will be created with the expected schema. If the table exists, the schema will be validated and data will be added to this existing table.""" - - import_result_gcs_sink: Optional[genai_types.GcsDestinationDict] - """The Cloud Storage path to write import result to.""" - jira_source: Optional[JiraSourceDict] - """Jira queries with their corresponding authentication.""" - - max_embedding_requests_per_min: Optional[int] - """Optional. The max number of queries per minute that this job is allowed to make to the embedding model specified on the corpus. This value is specific to this job and not shared across other import jobs. Consult the Quotas page on the project to set an appropriate value here. If unspecified, a default value of 1,000 QPM would be used.""" - - partial_failure_bigquery_sink: Optional[BigQueryDestinationDict] - """The BigQuery destination to write partial failures to. It should be a bigquery table resource name (e.g. "bq://projectId.bqDatasetId.bqTableId"). The dataset must exist. If the table does not exist, it will be created with the expected schema. If the table exists, the schema will be validated and data will be added to this existing table. Deprecated. Prefer to use `import_result_bq_sink`.""" - partial_failure_gcs_sink: Optional[genai_types.GcsDestinationDict] - """The Cloud Storage path to write partial failures to. Deprecated. Prefer to use `import_result_gcs_sink`.""" +class UploadRagFileConfigDict(TypedDict, total=False): + """Config for uploading RagFile.""" rag_file_chunking_config: Optional[RagFileChunkingConfigDict] - """Specifies the size and overlap of chunks after importing RagFiles.""" + """Specifies the size and overlap of chunks after uploading RagFile.""" rag_file_metadata_config: Optional[RagFileMetadataConfigDict] - """Specifies the metadata config for RagFiles. Including paths for metadata schema and metadata. Deprecated: Not in use.""" + """Optional. Specifies the metadata config for RagFiles. Including paths for metadata schema and metadata. Alteratively, inline metadata schema and metadata can be provided. Deprecated: Not in use.""" rag_file_parsing_config: Optional[RagFileParsingConfigDict] """Optional. Specifies the parsing config for RagFiles. RAG will use the default parser if this field is not set.""" @@ -16564,579 +16050,517 @@ class ImportRagFilesConfigDict(TypedDict, total=False): rag_file_transformation_config: Optional[RagFileTransformationConfigDict] """Specifies the transformation config for RagFiles.""" - rebuild_ann_index: Optional[bool] - """Rebuilds the ANN index to optimize for recall on the imported data. Only applicable for RagCorpora running on RagManagedDb with `retrieval_strategy` set to `ANN`. The rebuild will be performed using the existing ANN config set on the RagCorpus. To change the ANN config, please use the UpdateRagCorpus API. Default is false, i.e., index is not rebuilt.""" - share_point_sources: Optional[SharePointSourcesDict] - """SharePoint sources.""" +UploadRagFileConfigOrDict = Union[UploadRagFileConfig, UploadRagFileConfigDict] - slack_source: Optional[SlackSourceDict] - """Slack channels with their corresponding access tokens.""" +class _UploadRagFileParameters(_common.BaseModel): + """Parameters for uploading a Rag File.""" -ImportRagFilesConfigOrDict = Union[ImportRagFilesConfig, ImportRagFilesConfigDict] + name: Optional[str] = Field( + default=None, + description="""The name of the RagCorpus resource into which to upload the file.""", + ) + rag_file: Optional[RagFile] = Field( + default=None, description="""The RagFile metadata to upload.""" + ) + upload_rag_file_config: Optional[UploadRagFileConfig] = Field( + default=None, + description="""The config for the RagFiles to be uploaded into the RagCorpus.""", + ) + config: Optional[UploadRagFileRequestConfig] = Field( + default=None, description="""Used to override the default configuration.""" + ) -class ImportRagFilesRequest(_common.BaseModel): - """Request message for VertexRagDataService.ImportRagFiles.""" +class _UploadRagFileParametersDict(TypedDict, total=False): + """Parameters for uploading a Rag File.""" - import_rag_files_config: Optional[ImportRagFilesConfig] = Field( - default=None, - description="""Required. The config for the RagFiles to be synced and imported into the RagCorpus. VertexRagDataService.ImportRagFiles.""", - ) + name: Optional[str] + """The name of the RagCorpus resource into which to upload the file.""" + rag_file: Optional[RagFileDict] + """The RagFile metadata to upload.""" -class ImportRagFilesRequestDict(TypedDict, total=False): - """Request message for VertexRagDataService.ImportRagFiles.""" + upload_rag_file_config: Optional[UploadRagFileConfigDict] + """The config for the RagFiles to be uploaded into the RagCorpus.""" - import_rag_files_config: Optional[ImportRagFilesConfigDict] - """Required. The config for the RagFiles to be synced and imported into the RagCorpus. VertexRagDataService.ImportRagFiles.""" + config: Optional[UploadRagFileRequestConfigDict] + """Used to override the default configuration.""" -ImportRagFilesRequestOrDict = Union[ImportRagFilesRequest, ImportRagFilesRequestDict] +_UploadRagFileParametersOrDict = Union[ + _UploadRagFileParameters, _UploadRagFileParametersDict +] -class _ImportRagFilesRequestParameters(_common.BaseModel): - """Parameters for importing RAG files.""" +class UploadRagFileResponse(_common.BaseModel): + """Response for uploading a Rag File.""" - name: Optional[str] = Field(default=None, description="""""") - import_rag_files_request: Optional[ImportRagFilesRequest] = Field( - default=None, description="""""" + error: Optional[genai_types.GoogleRpcStatus] = Field( + default=None, + description="""The error that occurred while processing the RagFile.""", ) - config: Optional[ImportRagFilesRequestConfig] = Field( - default=None, description="""""" + rag_file: Optional[RagFile] = Field( + default=None, + description="""The RagFile that had been uploaded into the RagCorpus.""", ) -class _ImportRagFilesRequestParametersDict(TypedDict, total=False): - """Parameters for importing RAG files.""" - - name: Optional[str] - """""" +class UploadRagFileResponseDict(TypedDict, total=False): + """Response for uploading a Rag File.""" - import_rag_files_request: Optional[ImportRagFilesRequestDict] - """""" + error: Optional[genai_types.GoogleRpcStatusDict] + """The error that occurred while processing the RagFile.""" - config: Optional[ImportRagFilesRequestConfigDict] - """""" + rag_file: Optional[RagFileDict] + """The RagFile that had been uploaded into the RagCorpus.""" -_ImportRagFilesRequestParametersOrDict = Union[ - _ImportRagFilesRequestParameters, _ImportRagFilesRequestParametersDict -] +UploadRagFileResponseOrDict = Union[UploadRagFileResponse, UploadRagFileResponseDict] -class ImportRagFilesResponse(_common.BaseModel): - """Response from importing RAG files.""" +class SandboxEnvironmentSpecCodeExecutionEnvironment(_common.BaseModel): + """The code execution environment with customized settings.""" - imported_rag_files_count: Optional[int] = Field( - default=None, - description="""The number of RagFiles imported into the RagCorpus.""", - ) - failed_rag_files_count: Optional[int] = Field( - default=None, - description="""The number of RagFiles that failed to import into the RagCorpus.""", - ) - skipped_rag_files_count: Optional[int] = Field( - default=None, - description="""The number of RagFiles skipped while importing into the RagCorpus.""", - ) - partial_failures_gcs_path: Optional[str] = Field( + code_language: Optional[Language] = Field( default=None, - description="""The Google Cloud Storage path with the partial failures.""", + description="""The coding language supported in this environment.""", ) - partial_failures_bigquery_table: Optional[str] = Field( + machine_config: Optional[MachineConfig] = Field( default=None, - description="""The BigQuery table where the partial failures were written.""", + description="""The machine config of the code execution environment.""", ) -class ImportRagFilesResponseDict(TypedDict, total=False): - """Response from importing RAG files.""" +class SandboxEnvironmentSpecCodeExecutionEnvironmentDict(TypedDict, total=False): + """The code execution environment with customized settings.""" - imported_rag_files_count: Optional[int] - """The number of RagFiles imported into the RagCorpus.""" + code_language: Optional[Language] + """The coding language supported in this environment.""" - failed_rag_files_count: Optional[int] - """The number of RagFiles that failed to import into the RagCorpus.""" + machine_config: Optional[MachineConfig] + """The machine config of the code execution environment.""" - skipped_rag_files_count: Optional[int] - """The number of RagFiles skipped while importing into the RagCorpus.""" - partial_failures_gcs_path: Optional[str] - """The Google Cloud Storage path with the partial failures.""" +SandboxEnvironmentSpecCodeExecutionEnvironmentOrDict = Union[ + SandboxEnvironmentSpecCodeExecutionEnvironment, + SandboxEnvironmentSpecCodeExecutionEnvironmentDict, +] - partial_failures_bigquery_table: Optional[str] - """The BigQuery table where the partial failures were written.""" +class SandboxEnvironmentSpecComputerUseEnvironment(_common.BaseModel): + """The computer use environment with customized settings.""" -ImportRagFilesResponseOrDict = Union[ImportRagFilesResponse, ImportRagFilesResponseDict] + pass -class ImportRagFilesOperation(_common.BaseModel): - """Operation for importing RAG files.""" +class SandboxEnvironmentSpecComputerUseEnvironmentDict(TypedDict, total=False): + """The computer use environment with customized settings.""" - 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[ImportRagFilesResponse] = Field( - default=None, description="""The response from the import operation.""" - ) - - -class ImportRagFilesOperationDict(TypedDict, total=False): - """Operation for importing RAG files.""" - - 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[ImportRagFilesResponseDict] - """The response from the import operation.""" + pass -ImportRagFilesOperationOrDict = Union[ - ImportRagFilesOperation, ImportRagFilesOperationDict +SandboxEnvironmentSpecComputerUseEnvironmentOrDict = Union[ + SandboxEnvironmentSpecComputerUseEnvironment, + SandboxEnvironmentSpecComputerUseEnvironmentDict, ] -class GetImportFilesOperationConfig(_common.BaseModel): - """Config for getting an import files operation.""" +class SandboxEnvironmentSpecShellEnvironment(_common.BaseModel): + """The shell environment with customized settings.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + pass -class GetImportFilesOperationConfigDict(TypedDict, total=False): - """Config for getting an import files operation.""" +class SandboxEnvironmentSpecShellEnvironmentDict(TypedDict, total=False): + """The shell environment with customized settings.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + pass -GetImportFilesOperationConfigOrDict = Union[ - GetImportFilesOperationConfig, GetImportFilesOperationConfigDict +SandboxEnvironmentSpecShellEnvironmentOrDict = Union[ + SandboxEnvironmentSpecShellEnvironment, SandboxEnvironmentSpecShellEnvironmentDict ] -class _GetImportFilesOperationParameters(_common.BaseModel): - """Parameters for getting an import files operation.""" +class SandboxEnvironmentSpec(_common.BaseModel): + """The specification of a sandbox environment.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + 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.""") ) - config: Optional[GetImportFilesOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + shell_environment: Optional[SandboxEnvironmentSpecShellEnvironment] = Field( + default=None, description="""Optional. The shell environment.""" ) -class _GetImportFilesOperationParametersDict(TypedDict, total=False): - """Parameters for getting an import files operation.""" +class SandboxEnvironmentSpecDict(TypedDict, total=False): + """The specification of a sandbox environment.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + code_execution_environment: Optional[ + SandboxEnvironmentSpecCodeExecutionEnvironmentDict + ] + """Optional. The code execution environment.""" - config: Optional[GetImportFilesOperationConfigDict] - """Used to override the default configuration.""" + computer_use_environment: Optional[SandboxEnvironmentSpecComputerUseEnvironmentDict] + """Optional. The computer use environment.""" + shell_environment: Optional[SandboxEnvironmentSpecShellEnvironmentDict] + """Optional. The shell environment.""" -_GetImportFilesOperationParametersOrDict = Union[ - _GetImportFilesOperationParameters, _GetImportFilesOperationParametersDict -] +SandboxEnvironmentSpecOrDict = Union[SandboxEnvironmentSpec, SandboxEnvironmentSpecDict] -class UploadRagFileRequestConfig(_common.BaseModel): - """Config for the request to upload a Rag File.""" + +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.""" ) - - -class UploadRagFileRequestConfigDict(TypedDict, total=False): - """Config for the request to upload a Rag File.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -UploadRagFileRequestConfigOrDict = Union[ - UploadRagFileRequestConfig, UploadRagFileRequestConfigDict -] - - -class UploadRagFileConfig(_common.BaseModel): - """Config for uploading RagFile.""" - - rag_file_chunking_config: Optional[RagFileChunkingConfig] = Field( + 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="""Specifies the size and overlap of chunks after uploading RagFile.""", + description="""The TTL for this resource. The expiration time is computed: now + TTL.""", ) - rag_file_metadata_config: Optional[RagFileMetadataConfig] = Field( + sandbox_environment_template: Optional[str] = Field( default=None, - description="""Optional. Specifies the metadata config for RagFiles. Including paths for metadata schema and metadata. Alteratively, inline metadata schema and metadata can be provided. Deprecated: Not in use.""", + 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}""", ) - rag_file_parsing_config: Optional[RagFileParsingConfig] = Field( + sandbox_environment_snapshot: Optional[str] = Field( default=None, - description="""Optional. Specifies the parsing config for RagFiles. RAG will use the default parser if this field is not set.""", + 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}""", ) - rag_file_transformation_config: Optional[RagFileTransformationConfig] = Field( + owner: Optional[str] = Field( default=None, - description="""Specifies the transformation config for RagFiles.""", + description="""Owner information for this sandbox environment. A sandbox can only be restored from a snapshot belonging to the same owner.""", ) -class UploadRagFileConfigDict(TypedDict, total=False): - """Config for uploading RagFile.""" +class CreateRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for creating a Sandbox.""" - rag_file_chunking_config: Optional[RagFileChunkingConfigDict] - """Specifies the size and overlap of chunks after uploading RagFile.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - rag_file_metadata_config: Optional[RagFileMetadataConfigDict] - """Optional. Specifies the metadata config for RagFiles. Including paths for metadata schema and metadata. Alteratively, inline metadata schema and metadata can be provided. Deprecated: Not in use.""" + display_name: Optional[str] + """The display name of the sandbox.""" - rag_file_parsing_config: Optional[RagFileParsingConfigDict] - """Optional. Specifies the parsing config for RagFiles. RAG will use the default parser if this field is not set.""" + description: Optional[str] + """The description of the sandbox.""" - rag_file_transformation_config: Optional[RagFileTransformationConfigDict] - """Specifies the transformation config for RagFiles.""" + 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}""" -UploadRagFileConfigOrDict = Union[UploadRagFileConfig, UploadRagFileConfigDict] + 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 _UploadRagFileParameters(_common.BaseModel): - """Parameters for uploading a Rag File.""" + +CreateRuntimeSandboxConfigOrDict = Union[ + CreateRuntimeSandboxConfig, CreateRuntimeSandboxConfigDict +] + + +class _CreateRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for creating Agent Runtime Sandboxes.""" name: Optional[str] = Field( default=None, - description="""The name of the RagCorpus resource into which to upload the file.""", - ) - rag_file: Optional[RagFile] = Field( - default=None, description="""The RagFile metadata to upload.""" + description="""Name of the Agent Runtime to create the sandbox under.""", ) - upload_rag_file_config: Optional[UploadRagFileConfig] = Field( - default=None, - description="""The config for the RagFiles to be uploaded into the RagCorpus.""", + spec: Optional[SandboxEnvironmentSpec] = Field( + default=None, description="""The specification of the sandbox.""" ) - config: Optional[UploadRagFileRequestConfig] = Field( - default=None, description="""Used to override the default configuration.""" + config: Optional[CreateRuntimeSandboxConfig] = Field( + default=None, description="""""" ) -class _UploadRagFileParametersDict(TypedDict, total=False): - """Parameters for uploading a Rag File.""" +class _CreateRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for creating Agent Runtime Sandboxes.""" name: Optional[str] - """The name of the RagCorpus resource into which to upload the file.""" - - rag_file: Optional[RagFileDict] - """The RagFile metadata to upload.""" + """Name of the Agent Runtime to create the sandbox under.""" - upload_rag_file_config: Optional[UploadRagFileConfigDict] - """The config for the RagFiles to be uploaded into the RagCorpus.""" + spec: Optional[SandboxEnvironmentSpecDict] + """The specification of the sandbox.""" - config: Optional[UploadRagFileRequestConfigDict] - """Used to override the default configuration.""" + config: Optional[CreateRuntimeSandboxConfigDict] + """""" -_UploadRagFileParametersOrDict = Union[ - _UploadRagFileParameters, _UploadRagFileParametersDict +_CreateRuntimeSandboxRequestParametersOrDict = Union[ + _CreateRuntimeSandboxRequestParameters, _CreateRuntimeSandboxRequestParametersDict ] -class UploadRagFileResponse(_common.BaseModel): - """Response for uploading a Rag File.""" +class SandboxEnvironmentConnectionInfo(_common.BaseModel): + """The connection information of the SandboxEnvironment.""" - error: Optional[genai_types.GoogleRpcStatus] = Field( + 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="""The error that occurred while processing the RagFile.""", + description="""Output only. The IP address of the load balancer.""", ) - rag_file: Optional[RagFile] = Field( + sandbox_internal_ip: Optional[str] = Field( default=None, - description="""The RagFile that had been uploaded into the RagCorpus.""", + description="""Output only. The internal IP address of the SandboxEnvironment.""", + ) + routing_token: Optional[str] = Field( + default=None, + description="""Output only. The routing token for the SandboxEnvironment.""", ) -class UploadRagFileResponseDict(TypedDict, total=False): - """Response for uploading a Rag File.""" - - error: Optional[genai_types.GoogleRpcStatusDict] - """The error that occurred while processing the RagFile.""" - - rag_file: Optional[RagFileDict] - """The RagFile that had been uploaded into the RagCorpus.""" - - -UploadRagFileResponseOrDict = Union[UploadRagFileResponse, UploadRagFileResponseDict] - - -class GetAgentEngineRuntimeRevisionConfig(_common.BaseModel): - """Config for getting an Agent Engine Runtime Revision.""" +class SandboxEnvironmentConnectionInfoDict(TypedDict, total=False): + """The connection information of the SandboxEnvironment.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + load_balancer_hostname: Optional[str] + """Output only. The hostname of the load balancer.""" + load_balancer_ip: Optional[str] + """Output only. The IP address of the load balancer.""" -class GetAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Runtime Revision.""" + sandbox_internal_ip: Optional[str] + """Output only. The internal IP address of the SandboxEnvironment.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + routing_token: Optional[str] + """Output only. The routing token for the SandboxEnvironment.""" -GetAgentEngineRuntimeRevisionConfigOrDict = Union[ - GetAgentEngineRuntimeRevisionConfig, GetAgentEngineRuntimeRevisionConfigDict +SandboxEnvironmentConnectionInfoOrDict = Union[ + SandboxEnvironmentConnectionInfo, SandboxEnvironmentConnectionInfoDict ] -class _GetAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine runtime revision.""" +class SandboxEnvironment(_common.BaseModel): + """A sandbox environment.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine runtime revision.""" - ) - config: Optional[GetAgentEngineRuntimeRevisionConfig] = Field( - default=None, description="""""" + 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.""", ) - - -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.""" - create_time: Optional[datetime.datetime] = Field( default=None, - description="""Output only. Timestamp when this ReasoningEngineRuntimeRevision was created.""", + 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="""Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""", + description="""Optional. The configuration of the SandboxEnvironment.""", ) - spec: Optional[ReasoningEngineSpec] = Field( + state: Optional[SandboxState] = Field( default=None, - description="""Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""", + description="""Output only. The runtime state of the SandboxEnvironment.""", ) - state: Optional[State] = Field( - default=None, description="""Output only. The state of the revision.""" + ttl: Optional[str] = Field( + default=None, + description="""Optional. Input only. The TTL for the sandbox environment. The expiration time is computed: now + TTL.""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. The timestamp when this SandboxEnvironment was most recently updated.""", + ) + latest_sandbox_environment_snapshot: Optional[str] = Field( + default=None, + description="""Output only. The resource name of the latest snapshot taken for this SandboxEnvironment.""", + ) + owner: Optional[str] = Field( + default=None, + description="""Optional. Owner information for this sandbox environment. A Sandbox can only be restored from a snapshot that belongs to the same owner. If not set, sandbox will be created as the default owner.""", + ) + sandbox_environment_snapshot: Optional[str] = Field( + default=None, + description="""Optional. The resource name of the SandboxEnvironmentSnapshot to use for creating this SandboxEnvironment. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}`""", + ) + sandbox_environment_template: Optional[str] = Field( + default=None, + description="""Optional. The name of the SandboxEnvironmentTemplate specified in the parent Agent Engine resource that this SandboxEnvironment is created from.""", ) -class 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 SandboxEnvironmentDict(TypedDict, total=False): + """A sandbox environment.""" - spec: Optional[ReasoningEngineSpecDict] - """Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""" + expire_time: Optional[datetime.datetime] + """Expiration time of the sandbox environment. + """ - state: Optional[State] - """Output only. The state of the revision.""" + connection_info: Optional[SandboxEnvironmentConnectionInfoDict] + """Output only. The connection information of the SandboxEnvironment.""" + create_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironment was created.""" -ReasoningEngineRuntimeRevisionOrDict = Union[ - ReasoningEngineRuntimeRevision, ReasoningEngineRuntimeRevisionDict -] + display_name: Optional[str] + """Required. The display name of the SandboxEnvironment.""" + name: Optional[str] + """Identifier. The name of the SandboxEnvironment.""" -class ListAgentEngineRuntimeRevisionsConfig(_common.BaseModel): - """Config for listing reasoning engine runtime revisions.""" + spec: Optional[SandboxEnvironmentSpecDict] + """Optional. The configuration of the SandboxEnvironment.""" - 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.""", - ) + state: Optional[SandboxState] + """Output only. The runtime state of the SandboxEnvironment.""" + ttl: Optional[str] + """Optional. Input only. The TTL for the sandbox environment. The expiration time is computed: now + TTL.""" -class ListAgentEngineRuntimeRevisionsConfigDict(TypedDict, total=False): - """Config for listing reasoning engine runtime revisions.""" + update_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironment was most recently updated.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + latest_sandbox_environment_snapshot: Optional[str] + """Output only. The resource name of the latest snapshot taken for this SandboxEnvironment.""" - page_size: Optional[int] - """""" + 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.""" - page_token: Optional[str] - """""" + 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}`""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" + sandbox_environment_template: Optional[str] + """Optional. The name of the SandboxEnvironmentTemplate specified in the parent Agent Engine resource that this SandboxEnvironment is created from.""" -ListAgentEngineRuntimeRevisionsConfigOrDict = Union[ - ListAgentEngineRuntimeRevisionsConfig, ListAgentEngineRuntimeRevisionsConfigDict -] +SandboxEnvironmentOrDict = Union[SandboxEnvironment, SandboxEnvironmentDict] -class _ListAgentEngineRuntimeRevisionsRequestParameters(_common.BaseModel): - """Parameters for listing reasoning engine runtime revisions.""" +class RuntimeSandboxOperation(_common.BaseModel): + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] = Field( - default=None, description="""Name of the reasoning 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[ListAgentEngineRuntimeRevisionsConfig] = 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[SandboxEnvironment] = Field( + default=None, description="""The Agent Runtime Sandbox.""" ) -class _ListAgentEngineRuntimeRevisionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing reasoning engine runtime revisions.""" +class RuntimeSandboxOperationDict(TypedDict, total=False): + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] - """Name of the reasoning engine.""" - - config: Optional[ListAgentEngineRuntimeRevisionsConfigDict] - """""" - - -_ListAgentEngineRuntimeRevisionsRequestParametersOrDict = Union[ - _ListAgentEngineRuntimeRevisionsRequestParameters, - _ListAgentEngineRuntimeRevisionsRequestParametersDict, -] - - -class ListReasoningEnginesRuntimeRevisionsResponse(_common.BaseModel): - """Response for listing agent engine runtime revisions.""" - - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" - ) - next_page_token: Optional[str] = Field(default=None, description="""""") - reasoning_engine_runtime_revisions: Optional[ - list[ReasoningEngineRuntimeRevision] - ] = Field( - default=None, description="""List of reasoning engine runtime revisions.""" - ) - + """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 ListReasoningEnginesRuntimeRevisionsResponseDict(TypedDict, total=False): - """Response for listing agent engine runtime 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.""" - reasoning_engine_runtime_revisions: Optional[ - list[ReasoningEngineRuntimeRevisionDict] - ] - """List of reasoning engine runtime revisions.""" + response: Optional[SandboxEnvironmentDict] + """The Agent Runtime Sandbox.""" -ListReasoningEnginesRuntimeRevisionsResponseOrDict = Union[ - ListReasoningEnginesRuntimeRevisionsResponse, - ListReasoningEnginesRuntimeRevisionsResponseDict, +RuntimeSandboxOperationOrDict = Union[ + RuntimeSandboxOperation, RuntimeSandboxOperationDict ] -class DeleteAgentEngineRuntimeRevisionConfig(_common.BaseModel): - """Config for deleting an Agent Engine Runtime Revision.""" +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.""" ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", - ) -class DeleteAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Runtime Revision.""" +class DeleteRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for deleting an Agent Runtime Sandbox.""" 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.""" - -DeleteAgentEngineRuntimeRevisionConfigOrDict = Union[ - DeleteAgentEngineRuntimeRevisionConfig, DeleteAgentEngineRuntimeRevisionConfigDict +DeleteRuntimeSandboxConfigOrDict = Union[ + DeleteRuntimeSandboxConfig, DeleteRuntimeSandboxConfigDict ] -class _DeleteAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): - """Parameters for deleting 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 to delete.""", + default=None, description="""Name of the agent runtime sandbox to delete.""" ) - config: Optional[DeleteAgentEngineRuntimeRevisionConfig] = Field( + config: Optional[DeleteRuntimeSandboxConfig] = Field( default=None, description="""""" ) -class _DeleteAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engine runtime revisions.""" +class _DeleteRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtimes.""" name: Optional[str] - """Name of the agent engine runtime revision to delete.""" + """Name of the agent runtime sandbox to delete.""" - config: Optional[DeleteAgentEngineRuntimeRevisionConfigDict] + config: Optional[DeleteRuntimeSandboxConfigDict] """""" -_DeleteAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ - _DeleteAgentEngineRuntimeRevisionRequestParameters, - _DeleteAgentEngineRuntimeRevisionRequestParametersDict, +_DeleteRuntimeSandboxRequestParametersOrDict = Union[ + _DeleteRuntimeSandboxRequestParameters, _DeleteRuntimeSandboxRequestParametersDict ] -class DeleteAgentEngineRuntimeRevisionOperation(_common.BaseModel): - """Operation for deleting agent engine runtime revisions.""" +class DeleteRuntimeSandboxOperation(_common.BaseModel): + """Operation for deleting agent runtimes.""" name: Optional[str] = Field( default=None, @@ -17156,8 +16580,8 @@ class DeleteAgentEngineRuntimeRevisionOperation(_common.BaseModel): ) -class DeleteAgentEngineRuntimeRevisionOperationDict(TypedDict, total=False): - """Operation for deleting agent engine runtime revisions.""" +class DeleteRuntimeSandboxOperationDict(TypedDict, total=False): + """Operation for deleting agent runtimes.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -17172,1320 +16596,1287 @@ class DeleteAgentEngineRuntimeRevisionOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -DeleteAgentEngineRuntimeRevisionOperationOrDict = Union[ - DeleteAgentEngineRuntimeRevisionOperation, - DeleteAgentEngineRuntimeRevisionOperationDict, +DeleteRuntimeSandboxOperationOrDict = Union[ + DeleteRuntimeSandboxOperation, DeleteRuntimeSandboxOperationDict ] -class GetDeleteAgentEngineRuntimeRevisionOperationConfig(_common.BaseModel): +class Metadata(_common.BaseModel): + """Metadata for a chunk.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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 GetDeleteAgentEngineRuntimeRevisionOperationConfigDict(TypedDict, total=False): +class MetadataDict(TypedDict, total=False): + """Metadata for a chunk.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + 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.""" -GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict = Union[ - GetDeleteAgentEngineRuntimeRevisionOperationConfig, - GetDeleteAgentEngineRuntimeRevisionOperationConfigDict, -] +MetadataOrDict = Union[Metadata, MetadataDict] -class _GetDeleteAgentEngineRuntimeRevisionOperationParameters(_common.BaseModel): - """Parameters for getting an operation that deletes a agent engine runtime revision.""" +class Chunk(_common.BaseModel): + """A chunk of data.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + data: Optional[bytes] = Field( + default=None, description="""Required. The data in the chunk.""" ) - config: Optional[GetDeleteAgentEngineRuntimeRevisionOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + 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 _GetDeleteAgentEngineRuntimeRevisionOperationParametersDict( - TypedDict, total=False -): - """Parameters for getting an operation that deletes a agent engine runtime revision.""" +class ChunkDict(TypedDict, total=False): + """A chunk of data.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + data: Optional[bytes] + """Required. The data in the chunk.""" - config: Optional[GetDeleteAgentEngineRuntimeRevisionOperationConfigDict] - """Used to override the default configuration.""" + metadata: Optional[MetadataDict] + """Optional. Metadata that is associated with the data in the payload.""" + + mime_type: Optional[str] + """Required. Mime type of the chunk data. See https://www.iana.org/assignments/media-types/media-types.xhtml for the full list.""" -_GetDeleteAgentEngineRuntimeRevisionOperationParametersOrDict = Union[ - _GetDeleteAgentEngineRuntimeRevisionOperationParameters, - _GetDeleteAgentEngineRuntimeRevisionOperationParametersDict, -] +ChunkOrDict = Union[Chunk, ChunkDict] -class QueryAgentEngineRuntimeRevisionConfig(_common.BaseModel): - """Config for querying agent engine runtime revisions.""" +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_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 ExecuteCodeRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for executing code in an Agent Runtime sandbox.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - 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 +ExecuteCodeRuntimeSandboxConfigOrDict = Union[ + ExecuteCodeRuntimeSandboxConfig, ExecuteCodeRuntimeSandboxConfigDict ] -class _QueryAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): - """Parameters for querying agent engine runtime revisions.""" +class _ExecuteCodeRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for executing code in an agent runtime sandbox.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine runtime revision.""" + default=None, + description="""Name of the agent runtime sandbox to execute code in.""", + ) + inputs: Optional[list[Chunk]] = Field( + default=None, description="""Inputs to the code execution.""" ) - config: Optional[QueryAgentEngineRuntimeRevisionConfig] = Field( + config: Optional[ExecuteCodeRuntimeSandboxConfig] = Field( default=None, description="""""" ) -class _QueryAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for querying agent engine runtime revisions.""" +class _ExecuteCodeRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for executing code in an agent runtime sandbox.""" name: Optional[str] - """Name of the agent engine runtime revision.""" + """Name of the agent runtime sandbox to execute code in.""" - config: Optional[QueryAgentEngineRuntimeRevisionConfigDict] + inputs: Optional[list[ChunkDict]] + """Inputs to the code execution.""" + + config: Optional[ExecuteCodeRuntimeSandboxConfigDict] """""" -_QueryAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ - _QueryAgentEngineRuntimeRevisionRequestParameters, - _QueryAgentEngineRuntimeRevisionRequestParametersDict, +_ExecuteCodeRuntimeSandboxRequestParametersOrDict = Union[ + _ExecuteCodeRuntimeSandboxRequestParameters, + _ExecuteCodeRuntimeSandboxRequestParametersDict, ] -class SandboxEnvironmentSpecCodeExecutionEnvironment(_common.BaseModel): - """The code execution environment with customized settings.""" +class ExecuteSandboxEnvironmentResponse(_common.BaseModel): + """The response for executing a sandbox environment.""" - code_language: Optional[Language] = Field( - default=None, - description="""The coding language supported in this environment.""", - ) - machine_config: Optional[MachineConfig] = Field( - default=None, - description="""The machine config of the code execution environment.""", + outputs: Optional[list[Chunk]] = Field( + default=None, description="""The outputs from the sandbox environment.""" ) -class SandboxEnvironmentSpecCodeExecutionEnvironmentDict(TypedDict, total=False): - """The code execution environment with customized settings.""" - - code_language: Optional[Language] - """The coding language supported in this environment.""" +class ExecuteSandboxEnvironmentResponseDict(TypedDict, total=False): + """The response for executing a sandbox environment.""" - machine_config: Optional[MachineConfig] - """The machine config of the code execution environment.""" + outputs: Optional[list[ChunkDict]] + """The outputs from the sandbox environment.""" -SandboxEnvironmentSpecCodeExecutionEnvironmentOrDict = Union[ - SandboxEnvironmentSpecCodeExecutionEnvironment, - SandboxEnvironmentSpecCodeExecutionEnvironmentDict, +ExecuteSandboxEnvironmentResponseOrDict = Union[ + ExecuteSandboxEnvironmentResponse, ExecuteSandboxEnvironmentResponseDict ] -class SandboxEnvironmentSpecComputerUseEnvironment(_common.BaseModel): - """The computer use environment with customized settings.""" +class GetRuntimeSandboxConfig(_common.BaseModel): + """Config for getting an Agent Runtime Memory.""" - pass + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -class SandboxEnvironmentSpecComputerUseEnvironmentDict(TypedDict, total=False): - """The computer use environment with customized settings.""" +class GetRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for getting an Agent Runtime Memory.""" - pass + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -SandboxEnvironmentSpecComputerUseEnvironmentOrDict = Union[ - SandboxEnvironmentSpecComputerUseEnvironment, - SandboxEnvironmentSpecComputerUseEnvironmentDict, +GetRuntimeSandboxConfigOrDict = Union[ + GetRuntimeSandboxConfig, GetRuntimeSandboxConfigDict ] -class SandboxEnvironmentSpecShellEnvironment(_common.BaseModel): - """The shell environment with customized settings.""" +class _GetRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for getting an agent runtime sandbox.""" - pass + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime sandbox.""" + ) + config: Optional[GetRuntimeSandboxConfig] = Field(default=None, description="""""") -class SandboxEnvironmentSpecShellEnvironmentDict(TypedDict, total=False): - """The shell environment with customized settings.""" +class _GetRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for getting an agent runtime sandbox.""" - pass + name: Optional[str] + """Name of the agent runtime sandbox.""" + + config: Optional[GetRuntimeSandboxConfigDict] + """""" -SandboxEnvironmentSpecShellEnvironmentOrDict = Union[ - SandboxEnvironmentSpecShellEnvironment, SandboxEnvironmentSpecShellEnvironmentDict +_GetRuntimeSandboxRequestParametersOrDict = Union[ + _GetRuntimeSandboxRequestParameters, _GetRuntimeSandboxRequestParametersDict ] -class SandboxEnvironmentSpec(_common.BaseModel): - """The specification of a sandbox environment.""" +class ListRuntimeSandboxesConfig(_common.BaseModel): + """Config for listing agent runtime sandboxes.""" - 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.""") + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - shell_environment: Optional[SandboxEnvironmentSpecShellEnvironment] = Field( - default=None, description="""Optional. The shell environment.""" + 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 SandboxEnvironmentSpecDict(TypedDict, total=False): - """The specification of a sandbox environment.""" +class ListRuntimeSandboxesConfigDict(TypedDict, total=False): + """Config for listing agent runtime sandboxes.""" - code_execution_environment: Optional[ - SandboxEnvironmentSpecCodeExecutionEnvironmentDict - ] - """Optional. The code execution environment.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - computer_use_environment: Optional[SandboxEnvironmentSpecComputerUseEnvironmentDict] - """Optional. The computer use environment.""" + page_size: Optional[int] + """""" - shell_environment: Optional[SandboxEnvironmentSpecShellEnvironmentDict] - """Optional. The shell environment.""" + 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.""" -SandboxEnvironmentSpecOrDict = Union[SandboxEnvironmentSpec, SandboxEnvironmentSpecDict] +ListRuntimeSandboxesConfigOrDict = Union[ + ListRuntimeSandboxesConfig, ListRuntimeSandboxesConfigDict +] -class CreateAgentEngineSandboxConfig(_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}""", +class _ListRuntimeSandboxesRequestParameters(_common.BaseModel): + """Parameters for listing agent runtime sandboxes.""" + + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime.""" ) - 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.""", + config: Optional[ListRuntimeSandboxesConfig] = Field( + default=None, description="""""" ) -class CreateAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for creating a Sandbox.""" +class _ListRuntimeSandboxesRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtime sandboxes.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + name: Optional[str] + """Name of the agent runtime.""" - display_name: Optional[str] - """The display name of the sandbox.""" + config: Optional[ListRuntimeSandboxesConfigDict] + """""" - description: Optional[str] - """The description of the sandbox.""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" +_ListRuntimeSandboxesRequestParametersOrDict = Union[ + _ListRuntimeSandboxesRequestParameters, _ListRuntimeSandboxesRequestParametersDict +] - 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}""" +class ListRuntimeSandboxesResponse(_common.BaseModel): + """Response for listing agent runtime sandboxes.""" - 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}""" + 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.""" + ) - owner: Optional[str] - """Owner information for this sandbox environment. A sandbox can only be restored from a snapshot belonging to the same owner.""" + +class ListRuntimeSandboxesResponseDict(TypedDict, total=False): + """Response for listing agent runtime sandboxes.""" + + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + + next_page_token: Optional[str] + """""" + + sandbox_environments: Optional[list[SandboxEnvironmentDict]] + """List of agent runtime sandboxes.""" -CreateAgentEngineSandboxConfigOrDict = Union[ - CreateAgentEngineSandboxConfig, CreateAgentEngineSandboxConfigDict +ListRuntimeSandboxesResponseOrDict = Union[ + ListRuntimeSandboxesResponse, ListRuntimeSandboxesResponseDict ] -class _CreateAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for creating Agent Engine Sandboxes.""" +class _GetRuntimeSandboxOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a sandbox as a response.""" - 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.""" + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - config: Optional[CreateAgentEngineSandboxConfig] = Field( - default=None, description="""""" + config: Optional[GetRuntimeOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -class _CreateAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for creating Agent Engine Sandboxes.""" - - name: Optional[str] - """Name of the agent engine to create the sandbox under.""" +class _GetRuntimeSandboxOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a sandbox as a response.""" - spec: Optional[SandboxEnvironmentSpecDict] - """The specification of the sandbox.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - config: Optional[CreateAgentEngineSandboxConfigDict] - """""" + config: Optional[GetRuntimeOperationConfigDict] + """Used to override the default configuration.""" -_CreateAgentEngineSandboxRequestParametersOrDict = Union[ - _CreateAgentEngineSandboxRequestParameters, - _CreateAgentEngineSandboxRequestParametersDict, +_GetRuntimeSandboxOperationParametersOrDict = Union[ + _GetRuntimeSandboxOperationParameters, _GetRuntimeSandboxOperationParametersDict ] -class SandboxEnvironmentConnectionInfo(_common.BaseModel): - """The connection information of the SandboxEnvironment.""" +class SandboxEnvironmentTemplateCustomContainerSpec(_common.BaseModel): + """Specification for deploying from a custom container image.""" - load_balancer_hostname: Optional[str] = Field( - default=None, description="""Output only. The hostname of the load balancer.""" - ) - load_balancer_ip: Optional[str] = Field( + image_uri: Optional[str] = Field( default=None, - description="""Output only. The IP address of the load balancer.""", + 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.""", ) - sandbox_internal_ip: Optional[str] = Field( + + +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( default=None, - description="""Output only. The internal IP address of the SandboxEnvironment.""", + description="""Optional. Port number to expose. This must be a valid port number, between 1 and 65535.""", ) - routing_token: Optional[str] = Field( + protocol: Optional[Protocol] = Field( default=None, - description="""Output only. The routing token for the SandboxEnvironment.""", + description="""Optional. Protocol for port. Defaults to TCP if not specified.""", ) -class SandboxEnvironmentConnectionInfoDict(TypedDict, total=False): - """The connection information of the SandboxEnvironment.""" - - load_balancer_hostname: Optional[str] - """Output only. The hostname of the load balancer.""" - - load_balancer_ip: Optional[str] - """Output only. The IP address of the load balancer.""" +class SandboxEnvironmentTemplateNetworkPortDict(TypedDict, total=False): + """Represents a network port in a container.""" - sandbox_internal_ip: Optional[str] - """Output only. The internal IP address of the SandboxEnvironment.""" + port: Optional[int] + """Optional. Port number to expose. This must be a valid port number, between 1 and 65535.""" - routing_token: Optional[str] - """Output only. The routing token for the SandboxEnvironment.""" + protocol: Optional[Protocol] + """Optional. Protocol for port. Defaults to TCP if not specified.""" -SandboxEnvironmentConnectionInfoOrDict = Union[ - SandboxEnvironmentConnectionInfo, SandboxEnvironmentConnectionInfoDict +SandboxEnvironmentTemplateNetworkPortOrDict = Union[ + SandboxEnvironmentTemplateNetworkPort, SandboxEnvironmentTemplateNetworkPortDict ] -class SandboxEnvironment(_common.BaseModel): - """A sandbox environment.""" +class SandboxEnvironmentTemplateResourceRequirements(_common.BaseModel): + """Message to define resource requests and limits (mirroring Kubernetes) for each sandbox instance created from this template.""" - expire_time: Optional[datetime.datetime] = Field( + limits: Optional[dict[str, str]] = Field( default=None, - description="""Expiration time of the sandbox environment. - """, + description="""Optional. The maximum amounts of compute resources allowed. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "500m", "1Gi").""", ) - connection_info: Optional[SandboxEnvironmentConnectionInfo] = Field( + requests: Optional[dict[str, str]] = 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.""", + description="""Optional. The requested amounts of compute resources. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "250m", "512Mi").""", ) -class SandboxEnvironmentDict(TypedDict, total=False): - """A sandbox environment.""" - - expire_time: Optional[datetime.datetime] - """Expiration time of the sandbox environment. - """ +class SandboxEnvironmentTemplateResourceRequirementsDict(TypedDict, total=False): + """Message to define resource requests and limits (mirroring Kubernetes) for each sandbox instance created from this template.""" - connection_info: Optional[SandboxEnvironmentConnectionInfoDict] - """Output only. The connection information of the SandboxEnvironment.""" + 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").""" - create_time: Optional[datetime.datetime] - """Output only. The timestamp when this SandboxEnvironment was created.""" + 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").""" - display_name: Optional[str] - """Required. The display name of the SandboxEnvironment.""" - name: Optional[str] - """Identifier. The name of the SandboxEnvironment.""" +SandboxEnvironmentTemplateResourceRequirementsOrDict = Union[ + SandboxEnvironmentTemplateResourceRequirements, + SandboxEnvironmentTemplateResourceRequirementsDict, +] - spec: Optional[SandboxEnvironmentSpecDict] - """Optional. The configuration of the SandboxEnvironment.""" - state: Optional[SandboxState] - """Output only. The runtime state of the SandboxEnvironment.""" +class SandboxEnvironmentTemplateCustomContainerEnvironment(_common.BaseModel): + """The customized sandbox runtime environment for BYOC.""" - ttl: Optional[str] - """Optional. Input only. The TTL for the sandbox environment. The expiration time is computed: now + TTL.""" + 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.""" + ) - 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 SandboxEnvironmentTemplateCustomContainerEnvironmentDict(TypedDict, total=False): + """The customized sandbox runtime environment for BYOC.""" - 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.""" + custom_container_spec: Optional[SandboxEnvironmentTemplateCustomContainerSpecDict] + """The specification of the custom container environment.""" - 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}`""" + ports: Optional[list[SandboxEnvironmentTemplateNetworkPortDict]] + """Ports to expose from the container.""" - sandbox_environment_template: Optional[str] - """Optional. The name of the SandboxEnvironmentTemplate specified in the parent Agent Engine resource that this SandboxEnvironment is created from.""" + resources: Optional[SandboxEnvironmentTemplateResourceRequirementsDict] + """Resource requests and limits for the container.""" -SandboxEnvironmentOrDict = Union[SandboxEnvironment, SandboxEnvironmentDict] +SandboxEnvironmentTemplateCustomContainerEnvironmentOrDict = Union[ + SandboxEnvironmentTemplateCustomContainerEnvironment, + SandboxEnvironmentTemplateCustomContainerEnvironmentDict, +] -class AgentEngineSandboxOperation(_common.BaseModel): - """Operation that has an agent engine sandbox as a response.""" +class SandboxEnvironmentTemplateDefaultContainerEnvironment(_common.BaseModel): + """The default sandbox runtime environment for default container workloads.""" - name: Optional[str] = Field( - default=None, - description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, - description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", - ) - done: Optional[bool] = Field( + default_container_category: Optional[DefaultContainerCategory] = Field( default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + description="""Required. The category of the default container image.""", ) - error: Optional[dict[str, Any]] = Field( + resources: Optional[SandboxEnvironmentTemplateResourceRequirements] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", - ) - response: Optional[SandboxEnvironment] = Field( - default=None, description="""The Agent Engine Sandbox.""" + description="""Optional. Resource requests and limits for the default container.""", ) -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 SandboxEnvironmentTemplateDefaultContainerEnvironmentDict(TypedDict, total=False): + """The default sandbox runtime environment for default container workloads.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + default_container_category: Optional[DefaultContainerCategory] + """Required. The category of the default container image.""" - response: Optional[SandboxEnvironmentDict] - """The Agent Engine Sandbox.""" + resources: Optional[SandboxEnvironmentTemplateResourceRequirementsDict] + """Optional. Resource requests and limits for the default container.""" -AgentEngineSandboxOperationOrDict = Union[ - AgentEngineSandboxOperation, AgentEngineSandboxOperationDict +SandboxEnvironmentTemplateDefaultContainerEnvironmentOrDict = Union[ + SandboxEnvironmentTemplateDefaultContainerEnvironment, + SandboxEnvironmentTemplateDefaultContainerEnvironmentDict, ] -class DeleteAgentEngineSandboxConfig(_common.BaseModel): - """Config for deleting 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.""", ) - - -class DeleteAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Sandbox.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -DeleteAgentEngineSandboxConfigOrDict = Union[ - DeleteAgentEngineSandboxConfig, DeleteAgentEngineSandboxConfigDict -] - - -class _DeleteAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for deleting agent engines.""" - - name: Optional[str] = Field( - default=None, description="""Name of the agent engine sandbox to delete.""" + 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.""", ) - config: Optional[DeleteAgentEngineSandboxConfig] = Field( - default=None, description="""""" + 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 _DeleteAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engines.""" +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.""" - name: Optional[str] - """Name of the agent engine sandbox to delete.""" + 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[DeleteAgentEngineSandboxConfigDict] - """""" + 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.""" -_DeleteAgentEngineSandboxRequestParametersOrDict = Union[ - _DeleteAgentEngineSandboxRequestParameters, - _DeleteAgentEngineSandboxRequestParametersDict, +SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigOrDict = Union[ + SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfig, + SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigDict, ] -class DeleteAgentEngineSandboxOperation(_common.BaseModel): - """Operation for deleting agent engines.""" +class SandboxEnvironmentTemplateEgressControlConfig(_common.BaseModel): + """Configuration for egress control of sandbox instances.""" - 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}`.""", + internet_access: Optional[bool] = Field( + default=None, description="""Optional. Whether to allow internet access.""" ) - metadata: Optional[dict[str, Any]] = Field( + customer_vpc_network: 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 customer VPC network that sandbox egress is routed into.""", ) - done: Optional[bool] = Field( + dns_peering_configs: Optional[ + list[SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfig] + ] = 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. DNS peering configurations that allow sandbox egress to resolve customer-internal domains via the customer VPC.""", ) - error: Optional[dict[str, Any]] = Field( + network_attachment: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Optional. The name of the customer VPC NetworkAttachment used to draw a PSC interface IP into the customer VPC for sandbox egress.""", ) -class DeleteAgentEngineSandboxOperationDict(TypedDict, total=False): - """Operation for deleting agent engines.""" +class SandboxEnvironmentTemplateEgressControlConfigDict(TypedDict, total=False): + """Configuration for egress control of sandbox instances.""" - 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}`.""" + internet_access: Optional[bool] + """Optional. Whether to allow internet 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.""" + customer_vpc_network: Optional[str] + """Optional. The customer VPC network that sandbox egress is routed into.""" - 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.""" + dns_peering_configs: Optional[ + list[SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigDict] + ] + """Optional. DNS peering configurations that allow sandbox egress to resolve customer-internal domains via the customer VPC.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + 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.""" -DeleteAgentEngineSandboxOperationOrDict = Union[ - DeleteAgentEngineSandboxOperation, DeleteAgentEngineSandboxOperationDict +SandboxEnvironmentTemplateEgressControlConfigOrDict = Union[ + SandboxEnvironmentTemplateEgressControlConfig, + SandboxEnvironmentTemplateEgressControlConfigDict, ] -class Metadata(_common.BaseModel): - """Metadata for a chunk.""" +class CreateSandboxEnvironmentTemplateConfig(_common.BaseModel): + """Config for creating a Sandbox Template.""" - attributes: Optional[dict[str, bytes]] = Field( + 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="""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.""", + description="""The custom container environment for the sandbox template.""", ) - - -class MetadataDict(TypedDict, total=False): - """Metadata for a chunk.""" - - attributes: Optional[dict[str, bytes]] - """Optional. Attributes attached to the data. The keys have semantic conventions and the consumers of the attributes should know how to deserialize the value bytes based on the keys.""" - - -MetadataOrDict = Union[Metadata, MetadataDict] - - -class Chunk(_common.BaseModel): - """A chunk of data.""" - - data: Optional[bytes] = Field( - default=None, description="""Required. The data in the chunk.""" - ) - metadata: Optional[Metadata] = Field( + default_container_environment: Optional[ + SandboxEnvironmentTemplateDefaultContainerEnvironment + ] = Field( default=None, - description="""Optional. Metadata that is associated with the data in the payload.""", + description="""The default container environment for the sandbox template.""", ) - 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.""", + egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = ( + Field( + default=None, + description="""The egress control config for the sandbox template.""", + ) ) -class ChunkDict(TypedDict, total=False): - """A chunk of data.""" - - data: Optional[bytes] - """Required. The data in the chunk.""" - - metadata: Optional[MetadataDict] - """Optional. Metadata that is associated with the data in the payload.""" - - mime_type: Optional[str] - """Required. Mime type of the chunk data. See https://www.iana.org/assignments/media-types/media-types.xhtml for the full list.""" - - -ChunkOrDict = Union[Chunk, ChunkDict] - +class CreateSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): + """Config for creating a Sandbox Template.""" -class ExecuteCodeAgentEngineSandboxConfig(_common.BaseModel): - """Config for executing code in an Agent Engine sandbox.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" + custom_container_environment: Optional[ + SandboxEnvironmentTemplateCustomContainerEnvironmentDict + ] + """The custom container environment for the sandbox template.""" -class ExecuteCodeAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for executing code in an Agent Engine sandbox.""" + default_container_environment: Optional[ + SandboxEnvironmentTemplateDefaultContainerEnvironmentDict + ] + """The default container environment for the sandbox template.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfigDict] + """The egress control config for the sandbox template.""" -ExecuteCodeAgentEngineSandboxConfigOrDict = Union[ - ExecuteCodeAgentEngineSandboxConfig, ExecuteCodeAgentEngineSandboxConfigDict +CreateSandboxEnvironmentTemplateConfigOrDict = Union[ + CreateSandboxEnvironmentTemplateConfig, CreateSandboxEnvironmentTemplateConfigDict ] -class _ExecuteCodeAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for executing code in an agent engine sandbox.""" +class _CreateSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): + """Parameters for creating Sandbox Environment Templates.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine sandbox to execute code in.""", + description="""Name of the agent runtime to create the template under.""", ) - inputs: Optional[list[Chunk]] = Field( - default=None, description="""Inputs to the code execution.""" - ) - config: Optional[ExecuteCodeAgentEngineSandboxConfig] = Field( + config: Optional[CreateSandboxEnvironmentTemplateConfig] = Field( default=None, description="""""" ) + display_name: Optional[str] = Field( + default=None, description="""The display name of the sandbox template.""" + ) -class _ExecuteCodeAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for executing code in an agent engine sandbox.""" +class _CreateSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): + """Parameters for creating Sandbox Environment Templates.""" name: Optional[str] - """Name of the agent engine sandbox to execute code in.""" - - inputs: Optional[list[ChunkDict]] - """Inputs to the code execution.""" + """Name of the agent runtime to create the template under.""" - config: Optional[ExecuteCodeAgentEngineSandboxConfigDict] + config: Optional[CreateSandboxEnvironmentTemplateConfigDict] """""" + display_name: Optional[str] + """The display name of the sandbox template.""" + -_ExecuteCodeAgentEngineSandboxRequestParametersOrDict = Union[ - _ExecuteCodeAgentEngineSandboxRequestParameters, - _ExecuteCodeAgentEngineSandboxRequestParametersDict, +_CreateSandboxEnvironmentTemplateRequestParametersOrDict = Union[ + _CreateSandboxEnvironmentTemplateRequestParameters, + _CreateSandboxEnvironmentTemplateRequestParametersDict, ] -class ExecuteSandboxEnvironmentResponse(_common.BaseModel): - """The response for executing a sandbox environment.""" +class SandboxEnvironmentTemplate(_common.BaseModel): + """A sandbox environment template.""" - outputs: Optional[list[Chunk]] = Field( - default=None, description="""The outputs from the sandbox environment.""" + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. The timestamp when this SandboxEnvironmentTemplate was created.""", + ) + custom_container_environment: Optional[ + SandboxEnvironmentTemplateCustomContainerEnvironment + ] = Field( + default=None, + description="""The sandbox environment for custom container workloads.""", + ) + default_container_environment: Optional[ + SandboxEnvironmentTemplateDefaultContainerEnvironment + ] = Field( + default=None, + description="""The sandbox environment for default container workloads.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""Required. The display name of the SandboxEnvironmentTemplate.""", + ) + egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = ( + Field( + default=None, + description="""Optional. The configuration for egress control of this template.""", + ) + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. The resource name of the SandboxEnvironmentTemplate. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}`""", + ) + state: Optional[ + Literal[ + "UNSPECIFIED", + "PROVISIONING", + "ACTIVE", + "DEPROVISIONING", + "DELETED", + "FAILED", + ] + ] = Field( + default=None, + description="""Output only. The state of the sandbox environment template.""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. The timestamp when this SandboxEnvironmentTemplate was most recently updated.""", ) -class ExecuteSandboxEnvironmentResponseDict(TypedDict, total=False): - """The response for executing a sandbox environment.""" - - outputs: Optional[list[ChunkDict]] - """The outputs from the sandbox environment.""" +class SandboxEnvironmentTemplateDict(TypedDict, total=False): + """A sandbox environment template.""" + create_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironmentTemplate was created.""" -ExecuteSandboxEnvironmentResponseOrDict = Union[ - ExecuteSandboxEnvironmentResponse, ExecuteSandboxEnvironmentResponseDict -] + custom_container_environment: Optional[ + SandboxEnvironmentTemplateCustomContainerEnvironmentDict + ] + """The sandbox environment for custom container workloads.""" + default_container_environment: Optional[ + SandboxEnvironmentTemplateDefaultContainerEnvironmentDict + ] + """The sandbox environment for default container workloads.""" -class GetAgentEngineSandboxConfig(_common.BaseModel): - """Config for getting an Agent Engine Memory.""" + display_name: Optional[str] + """Required. The display name of the SandboxEnvironmentTemplate.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfigDict] + """Optional. The configuration for egress control of this template.""" + name: Optional[str] + """Identifier. The resource name of the SandboxEnvironmentTemplate. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}`""" -class GetAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Memory.""" + state: Optional[ + Literal[ + "UNSPECIFIED", + "PROVISIONING", + "ACTIVE", + "DEPROVISIONING", + "DELETED", + "FAILED", + ] + ] + """Output only. The state of the sandbox environment template.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + update_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironmentTemplate was most recently updated.""" -GetAgentEngineSandboxConfigOrDict = Union[ - GetAgentEngineSandboxConfig, GetAgentEngineSandboxConfigDict +SandboxEnvironmentTemplateOrDict = Union[ + SandboxEnvironmentTemplate, SandboxEnvironmentTemplateDict ] -class _GetAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine sandbox.""" +class SandboxEnvironmentTemplateOperation(_common.BaseModel): + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine sandbox.""" + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - config: Optional[GetAgentEngineSandboxConfig] = Field( - default=None, description="""""" + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + response: Optional[SandboxEnvironmentTemplate] = Field( + default=None, description="""The Agent Runtime Sandbox Template.""" ) -class _GetAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine sandbox.""" +class SandboxEnvironmentTemplateOperationDict(TypedDict, total=False): + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] - """Name of the agent engine sandbox.""" + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - config: Optional[GetAgentEngineSandboxConfigDict] - """""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + response: Optional[SandboxEnvironmentTemplateDict] + """The Agent Runtime Sandbox Template.""" -_GetAgentEngineSandboxRequestParametersOrDict = Union[ - _GetAgentEngineSandboxRequestParameters, _GetAgentEngineSandboxRequestParametersDict +SandboxEnvironmentTemplateOperationOrDict = Union[ + SandboxEnvironmentTemplateOperation, SandboxEnvironmentTemplateOperationDict ] -class ListAgentEngineSandboxesConfig(_common.BaseModel): - """Config for listing agent engine sandboxes.""" +class DeleteSandboxEnvironmentTemplateConfig(_common.BaseModel): + """Config for deleting a Sandbox Template.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( - default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""", - ) -class ListAgentEngineSandboxesConfigDict(TypedDict, total=False): - """Config for listing agent engine sandboxes.""" +class DeleteSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): + """Config for deleting a Sandbox Template.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - page_size: Optional[int] - """""" - - page_token: Optional[str] - """""" - - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" - -ListAgentEngineSandboxesConfigOrDict = Union[ - ListAgentEngineSandboxesConfig, ListAgentEngineSandboxesConfigDict +DeleteSandboxEnvironmentTemplateConfigOrDict = Union[ + DeleteSandboxEnvironmentTemplateConfig, DeleteSandboxEnvironmentTemplateConfigDict ] -class _ListAgentEngineSandboxesRequestParameters(_common.BaseModel): - """Parameters for listing agent engine sandboxes.""" +class _DeleteSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): + """Parameters for deleting sandbox templates.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the sandbox template to delete.""" ) - config: Optional[ListAgentEngineSandboxesConfig] = Field( + config: Optional[DeleteSandboxEnvironmentTemplateConfig] = Field( default=None, description="""""" ) -class _ListAgentEngineSandboxesRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engine sandboxes.""" +class _DeleteSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): + """Parameters for deleting sandbox templates.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the sandbox template to delete.""" - config: Optional[ListAgentEngineSandboxesConfigDict] + config: Optional[DeleteSandboxEnvironmentTemplateConfigDict] """""" -_ListAgentEngineSandboxesRequestParametersOrDict = Union[ - _ListAgentEngineSandboxesRequestParameters, - _ListAgentEngineSandboxesRequestParametersDict, +_DeleteSandboxEnvironmentTemplateRequestParametersOrDict = Union[ + _DeleteSandboxEnvironmentTemplateRequestParameters, + _DeleteSandboxEnvironmentTemplateRequestParametersDict, ] -class ListAgentEngineSandboxesResponse(_common.BaseModel): - """Response for listing agent engine sandboxes.""" - - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" +class DeleteSandboxEnvironmentTemplateOperation(_common.BaseModel): + """Operation for deleting sandbox templates.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - sandbox_environments: Optional[list[SandboxEnvironment]] = Field( - default=None, description="""List of agent engine sandboxes.""" + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", ) -class ListAgentEngineSandboxesResponseDict(TypedDict, total=False): - """Response for listing agent engine sandboxes.""" +class DeleteSandboxEnvironmentTemplateOperationDict(TypedDict, total=False): + """Operation for deleting sandbox templates.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - next_page_token: Optional[str] - """""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - sandbox_environments: Optional[list[SandboxEnvironmentDict]] - """List of agent engine sandboxes.""" + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" -ListAgentEngineSandboxesResponseOrDict = Union[ - ListAgentEngineSandboxesResponse, ListAgentEngineSandboxesResponseDict +DeleteSandboxEnvironmentTemplateOperationOrDict = Union[ + DeleteSandboxEnvironmentTemplateOperation, + DeleteSandboxEnvironmentTemplateOperationDict, ] -class _GetAgentEngineSandboxOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a sandbox as a response.""" +class GetSandboxEnvironmentTemplateConfig(_common.BaseModel): + """Config for getting a Sandbox Template.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" - ) - config: Optional[GetAgentEngineOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class _GetAgentEngineSandboxOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a sandbox as a response.""" - - operation_name: Optional[str] - """The server-assigned name for the operation.""" +class GetSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): + """Config for getting a Sandbox Template.""" - config: Optional[GetAgentEngineOperationConfigDict] - """Used to override the default configuration.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -_GetAgentEngineSandboxOperationParametersOrDict = Union[ - _GetAgentEngineSandboxOperationParameters, - _GetAgentEngineSandboxOperationParametersDict, +GetSandboxEnvironmentTemplateConfigOrDict = Union[ + GetSandboxEnvironmentTemplateConfig, GetSandboxEnvironmentTemplateConfigDict ] -class SandboxEnvironmentTemplateCustomContainerSpec(_common.BaseModel): - """Specification for deploying from a custom container image.""" +class _GetSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): + """Parameters for getting a sandbox template.""" - image_uri: Optional[str] = Field( - default=None, - description="""Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""", + name: Optional[str] = Field( + default=None, description="""Name of the sandbox template.""" + ) + config: Optional[GetSandboxEnvironmentTemplateConfig] = Field( + default=None, description="""""" ) -class SandboxEnvironmentTemplateCustomContainerSpecDict(TypedDict, total=False): - """Specification for deploying from a custom container image.""" +class _GetSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): + """Parameters for getting a sandbox template.""" - image_uri: Optional[str] - """Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""" + name: Optional[str] + """Name of the sandbox template.""" + config: Optional[GetSandboxEnvironmentTemplateConfigDict] + """""" -SandboxEnvironmentTemplateCustomContainerSpecOrDict = Union[ - SandboxEnvironmentTemplateCustomContainerSpec, - SandboxEnvironmentTemplateCustomContainerSpecDict, + +_GetSandboxEnvironmentTemplateRequestParametersOrDict = Union[ + _GetSandboxEnvironmentTemplateRequestParameters, + _GetSandboxEnvironmentTemplateRequestParametersDict, ] -class SandboxEnvironmentTemplateNetworkPort(_common.BaseModel): - """Represents a network port in a container.""" +class ListSandboxEnvironmentTemplatesConfig(_common.BaseModel): + """Config for listing sandbox templates.""" - port: Optional[int] = Field( - default=None, - description="""Optional. Port number to expose. This must be a valid port number, between 1 and 65535.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - protocol: Optional[Protocol] = Field( + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( default=None, - description="""Optional. Protocol for port. Defaults to TCP if not specified.""", + description="""An expression for filtering the results of the request.""", ) -class SandboxEnvironmentTemplateNetworkPortDict(TypedDict, total=False): - """Represents a network port in a container.""" +class ListSandboxEnvironmentTemplatesConfigDict(TypedDict, total=False): + """Config for listing sandbox templates.""" - port: Optional[int] - """Optional. Port number to expose. This must be a valid port number, between 1 and 65535.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - protocol: Optional[Protocol] - """Optional. Protocol for port. Defaults to TCP if not specified.""" + page_size: Optional[int] + """""" + page_token: Optional[str] + """""" -SandboxEnvironmentTemplateNetworkPortOrDict = Union[ - SandboxEnvironmentTemplateNetworkPort, SandboxEnvironmentTemplateNetworkPortDict + filter: Optional[str] + """An expression for filtering the results of the request.""" + + +ListSandboxEnvironmentTemplatesConfigOrDict = Union[ + ListSandboxEnvironmentTemplatesConfig, ListSandboxEnvironmentTemplatesConfigDict ] -class SandboxEnvironmentTemplateResourceRequirements(_common.BaseModel): - """Message to define resource requests and limits (mirroring Kubernetes) for each sandbox instance created from this template.""" +class _ListSandboxEnvironmentTemplatesRequestParameters(_common.BaseModel): + """Parameters for listing sandbox templates.""" - limits: Optional[dict[str, str]] = Field( - default=None, - description="""Optional. The maximum amounts of compute resources allowed. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "500m", "1Gi").""", + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime.""" ) - requests: Optional[dict[str, str]] = Field( - default=None, - description="""Optional. The requested amounts of compute resources. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "250m", "512Mi").""", + config: Optional[ListSandboxEnvironmentTemplatesConfig] = Field( + default=None, description="""""" ) -class SandboxEnvironmentTemplateResourceRequirementsDict(TypedDict, total=False): - """Message to define resource requests and limits (mirroring Kubernetes) for each sandbox instance created from this template.""" +class _ListSandboxEnvironmentTemplatesRequestParametersDict(TypedDict, total=False): + """Parameters for listing sandbox templates.""" - limits: Optional[dict[str, str]] - """Optional. The maximum amounts of compute resources allowed. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "500m", "1Gi").""" + name: Optional[str] + """Name of the agent runtime.""" - requests: Optional[dict[str, str]] - """Optional. The requested amounts of compute resources. Keys are resource names (e.g., "cpu", "memory"). Values are quantities (e.g., "250m", "512Mi").""" + config: Optional[ListSandboxEnvironmentTemplatesConfigDict] + """""" -SandboxEnvironmentTemplateResourceRequirementsOrDict = Union[ - SandboxEnvironmentTemplateResourceRequirements, - SandboxEnvironmentTemplateResourceRequirementsDict, +_ListSandboxEnvironmentTemplatesRequestParametersOrDict = Union[ + _ListSandboxEnvironmentTemplatesRequestParameters, + _ListSandboxEnvironmentTemplatesRequestParametersDict, ] -class SandboxEnvironmentTemplateCustomContainerEnvironment(_common.BaseModel): - """The customized sandbox runtime environment for BYOC.""" +class ListSandboxEnvironmentTemplatesResponse(_common.BaseModel): + """Response for listing sandbox templates.""" - custom_container_spec: Optional[SandboxEnvironmentTemplateCustomContainerSpec] = ( - Field( - default=None, - description="""The specification of the custom container environment.""", - ) - ) - ports: Optional[list[SandboxEnvironmentTemplateNetworkPort]] = Field( - default=None, description="""Ports to expose from the container.""" + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - resources: Optional[SandboxEnvironmentTemplateResourceRequirements] = Field( - default=None, description="""Resource requests and limits for the container.""" + next_page_token: Optional[str] = Field(default=None, description="""""") + sandbox_environment_templates: Optional[list[SandboxEnvironmentTemplate]] = Field( + default=None, description="""List of sandbox templates.""" ) -class SandboxEnvironmentTemplateCustomContainerEnvironmentDict(TypedDict, total=False): - """The customized sandbox runtime environment for BYOC.""" +class ListSandboxEnvironmentTemplatesResponseDict(TypedDict, total=False): + """Response for listing sandbox templates.""" - custom_container_spec: Optional[SandboxEnvironmentTemplateCustomContainerSpecDict] - """The specification of the custom container environment.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - ports: Optional[list[SandboxEnvironmentTemplateNetworkPortDict]] - """Ports to expose from the container.""" + next_page_token: Optional[str] + """""" - resources: Optional[SandboxEnvironmentTemplateResourceRequirementsDict] - """Resource requests and limits for the container.""" + sandbox_environment_templates: Optional[list[SandboxEnvironmentTemplateDict]] + """List of sandbox templates.""" -SandboxEnvironmentTemplateCustomContainerEnvironmentOrDict = Union[ - SandboxEnvironmentTemplateCustomContainerEnvironment, - SandboxEnvironmentTemplateCustomContainerEnvironmentDict, +ListSandboxEnvironmentTemplatesResponseOrDict = Union[ + ListSandboxEnvironmentTemplatesResponse, ListSandboxEnvironmentTemplatesResponseDict ] -class SandboxEnvironmentTemplateDefaultContainerEnvironment(_common.BaseModel): - """The default sandbox runtime environment for default container workloads.""" +class _GetSandboxEnvironmentTemplateOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a sandbox template as a response.""" - default_container_category: Optional[DefaultContainerCategory] = Field( - default=None, - description="""Required. The category of the default container image.""", + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - resources: Optional[SandboxEnvironmentTemplateResourceRequirements] = Field( - default=None, - description="""Optional. Resource requests and limits for the default container.""", + config: Optional[GetRuntimeOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -class SandboxEnvironmentTemplateDefaultContainerEnvironmentDict(TypedDict, total=False): - """The default sandbox runtime environment for default container workloads.""" +class _GetSandboxEnvironmentTemplateOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a sandbox template as a response.""" - default_container_category: Optional[DefaultContainerCategory] - """Required. The category of the default container image.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - resources: Optional[SandboxEnvironmentTemplateResourceRequirementsDict] - """Optional. Resource requests and limits for the default container.""" + config: Optional[GetRuntimeOperationConfigDict] + """Used to override the default configuration.""" -SandboxEnvironmentTemplateDefaultContainerEnvironmentOrDict = Union[ - SandboxEnvironmentTemplateDefaultContainerEnvironment, - SandboxEnvironmentTemplateDefaultContainerEnvironmentDict, +_GetSandboxEnvironmentTemplateOperationParametersOrDict = Union[ + _GetSandboxEnvironmentTemplateOperationParameters, + _GetSandboxEnvironmentTemplateOperationParametersDict, ] -class 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 CreateRuntimeSandboxSnapshotConfig(_common.BaseModel): + """Config for creating a Sandbox Environment Snapshot.""" - 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.""", + display_name: Optional[str] = Field( + default=None, description="""The display name of the sandbox snapshot.""" ) - target_project: Optional[str] = Field( + owner: Optional[str] = Field( + default=None, description="""The owner of the sandbox snapshot.""" + ) + ttl: 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.""", + description="""The TTL for this resource. The expiration time is computed: now + TTL.""", + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) -class 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.""" +class CreateRuntimeSandboxSnapshotConfigDict(TypedDict, total=False): + """Config for creating a Sandbox Environment Snapshot.""" - 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.""" + display_name: Optional[str] + """The display name of the sandbox snapshot.""" - 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.""" + owner: Optional[str] + """The owner of the sandbox snapshot.""" + ttl: Optional[str] + """The TTL for this resource. The expiration time is computed: now + TTL.""" -SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigOrDict = Union[ - SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfig, - SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigDict, + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" + + +CreateRuntimeSandboxSnapshotConfigOrDict = Union[ + CreateRuntimeSandboxSnapshotConfig, CreateRuntimeSandboxSnapshotConfigDict ] -class SandboxEnvironmentTemplateEgressControlConfig(_common.BaseModel): - """Configuration for egress control of sandbox instances.""" +class _CreateSandboxEnvironmentSnapshotRequestParameters(_common.BaseModel): + """Parameters for creating a sandbox environment snapshot.""" - internet_access: Optional[bool] = Field( - default=None, description="""Optional. Whether to allow internet access.""" + source_sandbox_environment_name: Optional[str] = Field( + default=None, description="""Name of the sandbox environment to snapshot.""" ) - customer_vpc_network: Optional[str] = Field( + config: Optional[CreateRuntimeSandboxSnapshotConfig] = Field( + default=None, description="""""" + ) + + +class _CreateSandboxEnvironmentSnapshotRequestParametersDict(TypedDict, total=False): + """Parameters for creating a sandbox environment snapshot.""" + + source_sandbox_environment_name: Optional[str] + """Name of the sandbox environment to snapshot.""" + + config: Optional[CreateRuntimeSandboxSnapshotConfigDict] + """""" + + +_CreateSandboxEnvironmentSnapshotRequestParametersOrDict = Union[ + _CreateSandboxEnvironmentSnapshotRequestParameters, + _CreateSandboxEnvironmentSnapshotRequestParametersDict, +] + + +class SandboxEnvironmentSnapshot(_common.BaseModel): + """A sandbox environment snapshot.""" + + display_name: Optional[str] = Field( default=None, - description="""Optional. The customer VPC network that sandbox egress is routed into.""", + description="""The display name of the sandbox environment snapshot.""", ) - dns_peering_configs: Optional[ - list[SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfig] - ] = Field( + expire_time: Optional[datetime.datetime] = Field( default=None, - description="""Optional. DNS peering configurations that allow sandbox egress to resolve customer-internal domains via the customer VPC.""", + description="""Expiration time of the sandbox environment snapshot. + """, ) - network_attachment: Optional[str] = Field( + create_time: Optional[datetime.datetime] = 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.""", + 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}`""", + ) + 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.""", + ) + parent_snapshot: Optional[str] = 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.""", + ) + 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.""", + ) + size_bytes: Optional[int] = Field( + default=None, + description="""Optional. Output only. Size of the snapshot data in bytes.""", + ) + source_sandbox_environment: Optional[str] = Field( + default=None, + 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 SandboxEnvironment was most recently updated.""", ) -class SandboxEnvironmentTemplateEgressControlConfigDict(TypedDict, total=False): - """Configuration for egress control of sandbox instances.""" +class SandboxEnvironmentSnapshotDict(TypedDict, total=False): + """A sandbox environment snapshot.""" - internet_access: Optional[bool] - """Optional. Whether to allow internet access.""" + display_name: Optional[str] + """The display name of the sandbox environment snapshot.""" - customer_vpc_network: Optional[str] - """Optional. The customer VPC network that sandbox egress is routed into.""" + expire_time: Optional[datetime.datetime] + """Expiration time of the sandbox environment snapshot. + """ - dns_peering_configs: Optional[ - list[SandboxEnvironmentTemplateEgressControlConfigDnsPeeringConfigDict] - ] - """Optional. DNS peering configurations that allow sandbox egress to resolve customer-internal domains via the customer VPC.""" + create_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironmentSnapshot was created.""" - 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.""" + name: Optional[str] + """Identifier. The resource name of the SandboxEnvironmentSnapshot. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentSnapshots/{sandbox_environment_snapshot}`""" + owner: Optional[str] + """Optional. Owner information for this sandbox snapshot. Different owners will have isolations on snapshot storage and identity. If not set, snapshot will be created as the default owner.""" -SandboxEnvironmentTemplateEgressControlConfigOrDict = Union[ - SandboxEnvironmentTemplateEgressControlConfig, - SandboxEnvironmentTemplateEgressControlConfigDict, + 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.""" + + ttl: Optional[str] + """Optional. Input only. The TTL for the sandbox environment snapshot. The expiration time is computed: now + TTL.""" + + update_time: Optional[datetime.datetime] + """Output only. The timestamp when this SandboxEnvironment was most recently updated.""" + + +SandboxEnvironmentSnapshotOrDict = Union[ + SandboxEnvironmentSnapshot, SandboxEnvironmentSnapshotDict ] -class CreateSandboxEnvironmentTemplateConfig(_common.BaseModel): - """Config for creating a Sandbox Template.""" +class RuntimeSandboxSnapshotOperation(_common.BaseModel): + """Operation that has an agent runtime sandbox snapshot as a response.""" - 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}`.""", ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", + 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.""", ) - custom_container_environment: Optional[ - SandboxEnvironmentTemplateCustomContainerEnvironment - ] = Field( + done: Optional[bool] = Field( default=None, - description="""The custom container environment for the sandbox template.""", + 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.""", ) - default_container_environment: Optional[ - SandboxEnvironmentTemplateDefaultContainerEnvironment - ] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""The default container environment for the sandbox template.""", + description="""The error result of the operation in case of failure or cancellation.""", ) - egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = ( - Field( - default=None, - description="""The egress control config for the sandbox template.""", - ) + response: Optional[SandboxEnvironmentSnapshot] = Field( + default=None, description="""The Agent Runtime Sandbox Snapshot.""" ) -class CreateSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): - """Config for creating a Sandbox Template.""" +class RuntimeSandboxSnapshotOperationDict(TypedDict, total=False): + """Operation that has an agent runtime sandbox snapshot as a response.""" - 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}`.""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + 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_container_environment: Optional[ - SandboxEnvironmentTemplateCustomContainerEnvironmentDict - ] - """The custom container environment for the sandbox 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.""" - default_container_environment: Optional[ - SandboxEnvironmentTemplateDefaultContainerEnvironmentDict - ] - """The default container environment for the sandbox template.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfigDict] - """The egress control config for the sandbox template.""" + response: Optional[SandboxEnvironmentSnapshotDict] + """The Agent Runtime Sandbox Snapshot.""" -CreateSandboxEnvironmentTemplateConfigOrDict = Union[ - CreateSandboxEnvironmentTemplateConfig, CreateSandboxEnvironmentTemplateConfigDict +RuntimeSandboxSnapshotOperationOrDict = Union[ + RuntimeSandboxSnapshotOperation, RuntimeSandboxSnapshotOperationDict ] -class _CreateSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): - """Parameters for creating Sandbox Environment Templates.""" +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 DeleteSandboxEnvironmentSnapshotConfigDict(TypedDict, total=False): + """Config for deleting a Sandbox Environment Snapshot.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +DeleteSandboxEnvironmentSnapshotConfigOrDict = Union[ + DeleteSandboxEnvironmentSnapshotConfig, DeleteSandboxEnvironmentSnapshotConfigDict +] + + +class _DeleteSandboxEnvironmentSnapshotRequestParameters(_common.BaseModel): + """Parameters for deleting sandbox environment snapshots.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to create the template under.""", + description="""Name of the sandbox environment snapshot to delete.""", ) - config: Optional[CreateSandboxEnvironmentTemplateConfig] = Field( + config: Optional[DeleteSandboxEnvironmentSnapshotConfig] = 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 _DeleteSandboxEnvironmentSnapshotRequestParametersDict(TypedDict, total=False): + """Parameters for deleting sandbox environment snapshots.""" name: Optional[str] - """Name of the agent engine to create the template under.""" + """Name of the sandbox environment snapshot to delete.""" - config: Optional[CreateSandboxEnvironmentTemplateConfigDict] + config: Optional[DeleteSandboxEnvironmentSnapshotConfigDict] """""" - display_name: Optional[str] - """The display name of the sandbox template.""" - -_CreateSandboxEnvironmentTemplateRequestParametersOrDict = Union[ - _CreateSandboxEnvironmentTemplateRequestParameters, - _CreateSandboxEnvironmentTemplateRequestParametersDict, +_DeleteSandboxEnvironmentSnapshotRequestParametersOrDict = Union[ + _DeleteSandboxEnvironmentSnapshotRequestParameters, + _DeleteSandboxEnvironmentSnapshotRequestParametersDict, ] -class SandboxEnvironmentTemplate(_common.BaseModel): - """A sandbox environment template.""" +class DeleteSandboxEnvironmentSnapshotOperation(_common.BaseModel): + """Operation for deleting sandbox environment snapshots.""" - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. The timestamp when this SandboxEnvironmentTemplate was created.""", - ) - custom_container_environment: Optional[ - SandboxEnvironmentTemplateCustomContainerEnvironment - ] = Field( - default=None, - description="""The sandbox environment for custom container workloads.""", - ) - default_container_environment: Optional[ - SandboxEnvironmentTemplateDefaultContainerEnvironment - ] = Field( - default=None, - description="""The sandbox environment for default container workloads.""", - ) - display_name: Optional[str] = Field( - default=None, - description="""Required. The display name of the SandboxEnvironmentTemplate.""", - ) - egress_control_config: Optional[SandboxEnvironmentTemplateEgressControlConfig] = ( - Field( - default=None, - description="""Optional. The configuration for egress control of this template.""", - ) - ) - name: Optional[str] = Field( - default=None, - description="""Identifier. The resource name of the SandboxEnvironmentTemplate. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sandboxEnvironmentTemplates/{sandbox_environment_template}`""", - ) - state: Optional[ - Literal[ - "UNSPECIFIED", - "PROVISIONING", - "ACTIVE", - "DEPROVISIONING", - "DELETED", - "FAILED", - ] - ] = Field( - default=None, - description="""Output only. The state of the sandbox environment template.""", - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. The timestamp when this SandboxEnvironmentTemplate was most recently updated.""", - ) - - -class SandboxEnvironmentTemplateDict(TypedDict, total=False): - """A sandbox environment template.""" - - create_time: Optional[datetime.datetime] - """Output only. The timestamp when this SandboxEnvironmentTemplate was created.""" - - custom_container_environment: Optional[ - SandboxEnvironmentTemplateCustomContainerEnvironmentDict - ] - """The sandbox environment for custom container workloads.""" - - default_container_environment: Optional[ - SandboxEnvironmentTemplateDefaultContainerEnvironmentDict - ] - """The sandbox environment for default container workloads.""" - - 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.""" - - -SandboxEnvironmentTemplateOrDict = Union[ - SandboxEnvironmentTemplate, SandboxEnvironmentTemplateDict -] - - -class SandboxEnvironmentTemplateOperation(_common.BaseModel): - """Operation that has an agent engine sandbox as a response.""" - - name: Optional[str] = Field( + 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}`.""", ) @@ -18501,13 +17892,10 @@ 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.""" - ) -class SandboxEnvironmentTemplateOperationDict(TypedDict, total=False): - """Operation that has an agent engine sandbox as a response.""" +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}`.""" @@ -18521,439 +17909,361 @@ 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.""" - -SandboxEnvironmentTemplateOperationOrDict = Union[ - SandboxEnvironmentTemplateOperation, SandboxEnvironmentTemplateOperationDict +DeleteSandboxEnvironmentSnapshotOperationOrDict = Union[ + DeleteSandboxEnvironmentSnapshotOperation, + DeleteSandboxEnvironmentSnapshotOperationDict, ] -class DeleteSandboxEnvironmentTemplateConfig(_common.BaseModel): - """Config for deleting 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 DeleteSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): - """Config for deleting 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.""" -DeleteSandboxEnvironmentTemplateConfigOrDict = Union[ - DeleteSandboxEnvironmentTemplateConfig, DeleteSandboxEnvironmentTemplateConfigDict +GetSandboxEnvironmentSnapshotConfigOrDict = Union[ + GetSandboxEnvironmentSnapshotConfig, GetSandboxEnvironmentSnapshotConfigDict ] -class _DeleteSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): - """Parameters for deleting sandbox templates.""" +class _GetSandboxEnvironmentSnapshotRequestParameters(_common.BaseModel): + """Parameters for getting a sandbox environment snapshot.""" name: Optional[str] = Field( - default=None, description="""Name of the sandbox template to delete.""" + default=None, description="""Name of the sandbox environment snapshot.""" ) - config: Optional[DeleteSandboxEnvironmentTemplateConfig] = Field( + config: Optional[GetSandboxEnvironmentSnapshotConfig] = Field( default=None, description="""""" ) -class _DeleteSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): - """Parameters for deleting sandbox templates.""" +class _GetSandboxEnvironmentSnapshotRequestParametersDict(TypedDict, total=False): + """Parameters for getting a sandbox environment snapshot.""" name: Optional[str] - """Name of the sandbox template to delete.""" + """Name of the sandbox environment snapshot.""" - config: Optional[DeleteSandboxEnvironmentTemplateConfigDict] + config: Optional[GetSandboxEnvironmentSnapshotConfigDict] """""" -_DeleteSandboxEnvironmentTemplateRequestParametersOrDict = Union[ - _DeleteSandboxEnvironmentTemplateRequestParameters, - _DeleteSandboxEnvironmentTemplateRequestParametersDict, +_GetSandboxEnvironmentSnapshotRequestParametersOrDict = Union[ + _GetSandboxEnvironmentSnapshotRequestParameters, + _GetSandboxEnvironmentSnapshotRequestParametersDict, ] -class DeleteSandboxEnvironmentTemplateOperation(_common.BaseModel): - """Operation for deleting sandbox templates.""" +class ListSandboxEnvironmentSnapshotsConfig(_common.BaseModel): + """Config for listing sandbox environment snapshots.""" - name: Optional[str] = Field( - default=None, - description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, - description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", - ) - done: Optional[bool] = Field( - default=None, - description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - error: Optional[dict[str, Any]] = Field( + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""An expression for filtering the results of the request.""", ) -class DeleteSandboxEnvironmentTemplateOperationDict(TypedDict, total=False): - """Operation for deleting sandbox templates.""" +class ListSandboxEnvironmentSnapshotsConfigDict(TypedDict, total=False): + """Config for listing 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}`.""" + 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.""" -DeleteSandboxEnvironmentTemplateOperationOrDict = Union[ - DeleteSandboxEnvironmentTemplateOperation, - DeleteSandboxEnvironmentTemplateOperationDict, +ListSandboxEnvironmentSnapshotsConfigOrDict = Union[ + ListSandboxEnvironmentSnapshotsConfig, ListSandboxEnvironmentSnapshotsConfigDict ] -class GetSandboxEnvironmentTemplateConfig(_common.BaseModel): - """Config for getting a Sandbox Template.""" +class _ListSandboxEnvironmentSnapshotsRequestParameters(_common.BaseModel): + """Parameters for listing sandbox environment snapshots.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + name: Optional[str] = Field( + default=None, + description="""Name of the reasoning engine to list snapshots from.""", + ) + config: Optional[ListSandboxEnvironmentSnapshotsConfig] = Field( + default=None, description="""""" ) -class GetSandboxEnvironmentTemplateConfigDict(TypedDict, total=False): - """Config for getting a Sandbox Template.""" +class _ListSandboxEnvironmentSnapshotsRequestParametersDict(TypedDict, total=False): + """Parameters for listing sandbox environment snapshots.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + name: Optional[str] + """Name of the reasoning engine to list snapshots from.""" + + config: Optional[ListSandboxEnvironmentSnapshotsConfigDict] + """""" -GetSandboxEnvironmentTemplateConfigOrDict = Union[ - GetSandboxEnvironmentTemplateConfig, GetSandboxEnvironmentTemplateConfigDict +_ListSandboxEnvironmentSnapshotsRequestParametersOrDict = Union[ + _ListSandboxEnvironmentSnapshotsRequestParameters, + _ListSandboxEnvironmentSnapshotsRequestParametersDict, ] -class _GetSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): - """Parameters for getting a sandbox template.""" +class ListSandboxEnvironmentSnapshotsResponse(_common.BaseModel): + """Response for listing sandbox environment snapshots.""" - name: Optional[str] = Field( - default=None, description="""Name of the sandbox template.""" + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - config: Optional[GetSandboxEnvironmentTemplateConfig] = Field( - default=None, description="""""" + next_page_token: Optional[str] = Field(default=None, description="""""") + sandbox_environment_snapshots: Optional[list[SandboxEnvironmentSnapshot]] = Field( + default=None, description="""List of sandbox environment snapshots.""" ) -class _GetSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=False): - """Parameters for getting a sandbox template.""" +class ListSandboxEnvironmentSnapshotsResponseDict(TypedDict, total=False): + """Response for listing sandbox environment snapshots.""" - name: Optional[str] - """Name of the sandbox template.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - config: Optional[GetSandboxEnvironmentTemplateConfigDict] + next_page_token: Optional[str] """""" + sandbox_environment_snapshots: Optional[list[SandboxEnvironmentSnapshotDict]] + """List of sandbox environment snapshots.""" + -_GetSandboxEnvironmentTemplateRequestParametersOrDict = Union[ - _GetSandboxEnvironmentTemplateRequestParameters, - _GetSandboxEnvironmentTemplateRequestParametersDict, +ListSandboxEnvironmentSnapshotsResponseOrDict = Union[ + ListSandboxEnvironmentSnapshotsResponse, ListSandboxEnvironmentSnapshotsResponseDict ] -class ListSandboxEnvironmentTemplatesConfig(_common.BaseModel): - """Config for listing sandbox templates.""" - - 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.""", - ) - - -class ListSandboxEnvironmentTemplatesConfigDict(TypedDict, total=False): - """Config for listing sandbox templates.""" - - 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.""" - - -ListSandboxEnvironmentTemplatesConfigOrDict = Union[ - ListSandboxEnvironmentTemplatesConfig, ListSandboxEnvironmentTemplatesConfigDict -] - - -class _ListSandboxEnvironmentTemplatesRequestParameters(_common.BaseModel): - """Parameters for listing sandbox templates.""" - - name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[ListSandboxEnvironmentTemplatesConfig] = Field( - default=None, description="""""" - ) - - -class _ListSandboxEnvironmentTemplatesRequestParametersDict(TypedDict, total=False): - """Parameters for listing sandbox templates.""" - - name: Optional[str] - """Name of the agent engine.""" - - config: Optional[ListSandboxEnvironmentTemplatesConfigDict] - """""" - - -_ListSandboxEnvironmentTemplatesRequestParametersOrDict = Union[ - _ListSandboxEnvironmentTemplatesRequestParameters, - _ListSandboxEnvironmentTemplatesRequestParametersDict, -] - - -class ListSandboxEnvironmentTemplatesResponse(_common.BaseModel): - """Response for listing sandbox templates.""" - - 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.""" - ) - - -class ListSandboxEnvironmentTemplatesResponseDict(TypedDict, total=False): - """Response for listing sandbox templates.""" - - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" - - next_page_token: Optional[str] - """""" - - sandbox_environment_templates: Optional[list[SandboxEnvironmentTemplateDict]] - """List of sandbox templates.""" - - -ListSandboxEnvironmentTemplatesResponseOrDict = Union[ - ListSandboxEnvironmentTemplatesResponse, ListSandboxEnvironmentTemplatesResponseDict -] - - -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.""" + session_state: Optional[dict[str, Any]] + """Optional. Session specific memory which stores key conversation points.""" - 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.""" + ttl: Optional[str] + """Optional. Input only. The TTL for this session. The minimum value is 24 hours.""" - 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.""" + update_time: Optional[datetime.datetime] + """Output only. Timestamp when the session was updated.""" - size_bytes: Optional[int] - """Optional. Output only. Size of the snapshot data in bytes.""" + user_id: Optional[str] + """Required. Immutable. String id provided by the user""" - 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.""" +SessionOrDict = Union[Session, SessionDict] - update_time: Optional[datetime.datetime] - """Output only. The timestamp when this SandboxEnvironment was most recently updated.""" - -SandboxEnvironmentSnapshotOrDict = Union[ - SandboxEnvironmentSnapshot, SandboxEnvironmentSnapshotDict -] - - -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, @@ -18971,13 +18281,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}`.""" @@ -18991,65 +18301,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, @@ -19069,8 +18377,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}`.""" @@ -19085,61 +18393,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.""" @@ -19148,12 +18452,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.""" @@ -19165,56 +18470,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.""" @@ -19222,44 +18526,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.""" @@ -19293,10 +18596,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.""" @@ -19324,305 +18635,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.""" -AgentEngineSessionOperationOrDict = Union[ - AgentEngineSessionOperation, AgentEngineSessionOperationDict -] + 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.""" + + 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.""" @@ -19636,8 +18935,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.""" @@ -19653,1293 +18952,812 @@ class ListAgentEngineSessionsConfigDict(TypedDict, total=False): For field names both snake_case and camelCase are supported.""" -ListAgentEngineSessionsConfigOrDict = Union[ - ListAgentEngineSessionsConfig, ListAgentEngineSessionsConfigDict +ListRuntimeSessionEventsConfigOrDict = Union[ + ListRuntimeSessionEventsConfig, ListRuntimeSessionEventsConfigDict ] -class _ListAgentEngineSessionsRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" +class _ListRuntimeSessionEventsRequestParameters(_common.BaseModel): + """Parameters for listing agent runtime session events.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime session.""" ) - config: Optional[ListAgentEngineSessionsConfig] = Field( + config: Optional[ListRuntimeSessionEventsConfig] = Field( default=None, description="""""" ) -class _ListAgentEngineSessionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" +class _ListRuntimeSessionEventsRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtime session events.""" name: Optional[str] - """Name of the agent engine.""" - - config: Optional[ListAgentEngineSessionsConfigDict] - """""" - - -_ListAgentEngineSessionsRequestParametersOrDict = Union[ - _ListAgentEngineSessionsRequestParameters, - _ListAgentEngineSessionsRequestParametersDict, -] - - -class ListReasoningEnginesSessionsResponse(_common.BaseModel): - """Response for listing agent engine sessions.""" - - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" - ) - next_page_token: Optional[str] = Field(default=None, description="""""") - sessions: Optional[list[Session]] = Field( - default=None, description="""List of agent engine sessions.""" - ) + """Name of the agent runtime session.""" - -class ListReasoningEnginesSessionsResponseDict(TypedDict, total=False): - """Response for listing agent engine sessions.""" - - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" - - next_page_token: Optional[str] + config: Optional[ListRuntimeSessionEventsConfigDict] """""" - sessions: Optional[list[SessionDict]] - """List of agent engine sessions.""" - - -ListReasoningEnginesSessionsResponseOrDict = Union[ - ListReasoningEnginesSessionsResponse, ListReasoningEnginesSessionsResponseDict -] - - -class _GetAgentEngineSessionOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a session as a response.""" - - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" - ) - config: Optional[GetAgentEngineOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" - ) - - -class _GetAgentEngineSessionOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a session as a response.""" - - operation_name: Optional[str] - """The server-assigned name for the operation.""" - config: Optional[GetAgentEngineOperationConfigDict] - """Used to override the default configuration.""" - - -_GetAgentEngineSessionOperationParametersOrDict = Union[ - _GetAgentEngineSessionOperationParameters, - _GetAgentEngineSessionOperationParametersDict, +_ListRuntimeSessionEventsRequestParametersOrDict = Union[ + _ListRuntimeSessionEventsRequestParameters, + _ListRuntimeSessionEventsRequestParametersDict, ] -class UpdateAgentEngineSessionConfig(_common.BaseModel): - """Config for updating agent engine session.""" +class SessionEvent(_common.BaseModel): + """A session event.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + content: Optional[genai_types.Content] = Field( + default=None, + description="""Optional. Content of the event provided by the author.""", ) - display_name: Optional[str] = Field( - default=None, description="""The display name of the session.""" + actions: Optional[EventActions] = Field( + default=None, description="""Optional. Actions executed by the agent.""" ) - session_state: Optional[dict[str, Any]] = Field( + author: Optional[str] = Field( default=None, - description="""Session state which stores key conversation points.""", - ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Waits for the operation to complete before returning.""", + description="""Required. The name of the agent that sent the event, or user.""", ) - ttl: Optional[str] = Field( + error_code: Optional[str] = Field( default=None, - description="""Optional. Input only. The TTL for this resource. - - The expiration time is computed: now + TTL.""", + description="""Optional. Error code if the response is an error. Code varies by model.""", ) - expire_time: Optional[datetime.datetime] = Field( + error_message: Optional[str] = Field( default=None, - description="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", + description="""Optional. Error message if the response is an error.""", ) - labels: Optional[dict[str, str]] = Field( + event_metadata: Optional[EventMetadata] = Field( + default=None, description="""Optional. Metadata relating to this event.""" + ) + invocation_id: Optional[str] = Field( default=None, - description="""Optional. The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", + description="""Required. The invocation id of the event, multiple events can have the same invocation id.""", ) - session_id: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Optional. The user defined ID to use for session, which will become the final component of the session resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""", + description="""Identifier. The resource name of the event. Format:`projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}/events/{event}`.""", ) - update_mask: Optional[str] = Field( + timestamp: Optional[datetime.datetime] = Field( default=None, - description="""The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""", + description="""Required. Timestamp when the event was created on client side.""", ) - user_id: Optional[str] = Field( - default=None, description="""User ID of the agent engine session to update.""" + raw_event: Optional[dict[str, Any]] = Field( + default=None, + description="""Optional. Weakly typed raw event data in proto struct format.""", ) -class UpdateAgentEngineSessionConfigDict(TypedDict, total=False): - """Config for updating agent engine session.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class SessionEventDict(TypedDict, total=False): + """A session event.""" - display_name: Optional[str] - """The display name of the session.""" + content: Optional[genai_types.Content] + """Optional. Content of the event provided by the author.""" - session_state: Optional[dict[str, Any]] - """Session state which stores key conversation points.""" + actions: Optional[EventActionsDict] + """Optional. Actions executed by the agent.""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + author: Optional[str] + """Required. The name of the agent that sent the event, or user.""" - ttl: Optional[str] - """Optional. Input only. The TTL for this resource. + error_code: Optional[str] + """Optional. Error code if the response is an error. Code varies by model.""" - The expiration time is computed: now + TTL.""" + error_message: Optional[str] + """Optional. Error message if the response is an error.""" - expire_time: Optional[datetime.datetime] - """Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""" + event_metadata: Optional[EventMetadataDict] + """Optional. Metadata relating to this event.""" - labels: Optional[dict[str, str]] - """Optional. The labels with user-defined metadata to organize your Sessions. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""" + invocation_id: Optional[str] + """Required. The invocation id of the event, multiple events can have the same invocation id.""" - session_id: Optional[str] - """Optional. The user defined ID to use for session, which will become the final component of the session resource name. If not provided, Vertex AI will generate a value for this ID. This value may be up to 63 characters, and valid characters are `[a-z0-9-]`. The first character must be a letter, and the last character must be a letter or number.""" + name: Optional[str] + """Identifier. The resource name of the event. Format:`projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}/events/{event}`.""" - update_mask: Optional[str] - """The update mask to apply. For the `FieldMask` definition, see - https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" + timestamp: Optional[datetime.datetime] + """Required. Timestamp when the event was created on client side.""" - user_id: Optional[str] - """User ID of the agent engine session to update.""" + raw_event: Optional[dict[str, Any]] + """Optional. Weakly typed raw event data in proto struct format.""" -UpdateAgentEngineSessionConfigOrDict = Union[ - UpdateAgentEngineSessionConfig, UpdateAgentEngineSessionConfigDict -] +SessionEventOrDict = Union[SessionEvent, SessionEventDict] -class _UpdateAgentEngineSessionRequestParameters(_common.BaseModel): - """Parameters for updating agent engine sessions.""" +class ListRuntimeSessionEventsResponse(_common.BaseModel): + """Response for listing agent runtime session events.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine session to update.""" + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) - config: Optional[UpdateAgentEngineSessionConfig] = Field( - default=None, description="""""" + next_page_token: Optional[str] = Field(default=None, description="""""") + session_events: Optional[list[SessionEvent]] = Field( + default=None, description="""List of session events.""" ) -class _UpdateAgentEngineSessionRequestParametersDict(TypedDict, total=False): - """Parameters for updating agent engine sessions.""" +class ListRuntimeSessionEventsResponseDict(TypedDict, total=False): + """Response for listing agent runtime session events.""" - name: Optional[str] - """Name of the agent engine session to update.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - config: Optional[UpdateAgentEngineSessionConfigDict] + next_page_token: Optional[str] """""" + session_events: Optional[list[SessionEventDict]] + """List of session events.""" + -_UpdateAgentEngineSessionRequestParametersOrDict = Union[ - _UpdateAgentEngineSessionRequestParameters, - _UpdateAgentEngineSessionRequestParametersDict, +ListRuntimeSessionEventsResponseOrDict = Union[ + ListRuntimeSessionEventsResponse, ListRuntimeSessionEventsResponseDict ] -class EventActions(_common.BaseModel): - """Actions are parts of events that are executed by the agent.""" +class GeminiExample(_common.BaseModel): + """Represents a Gemini example.""" - artifact_delta: Optional[dict[str, int]] = Field( - default=None, - description="""Optional. Indicates that the event is updating an artifact. key is the filename, value is the version.""", - ) - escalate: Optional[bool] = Field( - default=None, - description="""Optional. The agent is escalating to a higher level agent.""", - ) - requested_auth_configs: Optional[dict[str, Any]] = Field( - default=None, - description="""Optional. Will only be set by a tool response indicating tool request euc. Struct key is the function call id since one function call response (from model) could correspond to multiple function calls. Struct value is the required auth config, which can be another struct.""", - ) - skip_summarization: Optional[bool] = Field( - default=None, - description="""Optional. If true, it won't call model to summarize function response. Only used for function_response event.""", - ) - state_delta: Optional[dict[str, Any]] = Field( - default=None, - description="""Optional. Indicates that the event is updating the state with the given delta.""", + model: Optional[str] = Field( + default=None, description="""The model used to generate the Gemini example.""" ) - transfer_agent: Optional[str] = Field( - default=None, - description="""Optional. If set, the event transfers to the specified agent.""", + contents: Optional[list[genai_types.Content]] = Field( + default=None, description="""Contents of the Gemini example.""" ) - - -class EventActionsDict(TypedDict, total=False): - """Actions are parts of events that are executed by the agent.""" - - artifact_delta: Optional[dict[str, int]] - """Optional. Indicates that the event is updating an artifact. key is the filename, value is the version.""" - - escalate: Optional[bool] - """Optional. The agent is escalating to a higher level agent.""" - - requested_auth_configs: Optional[dict[str, Any]] - """Optional. Will only be set by a tool response indicating tool request euc. Struct key is the function call id since one function call response (from model) could correspond to multiple function calls. Struct value is the required auth config, which can be another struct.""" - - skip_summarization: Optional[bool] - """Optional. If true, it won't call model to summarize function response. Only used for function_response event.""" - - state_delta: Optional[dict[str, Any]] - """Optional. Indicates that the event is updating the state with the given delta.""" - - transfer_agent: Optional[str] - """Optional. If set, the event transfers to the specified agent.""" - - -EventActionsOrDict = Union[EventActions, EventActionsDict] - - -class EventMetadata(_common.BaseModel): - """Metadata relating to a LLM response event.""" - - grounding_metadata: Optional[genai_types.GroundingMetadata] = Field( - default=None, - description="""Optional. Metadata returned to client when grounding is enabled.""", + system_instruction: Optional[genai_types.Content] = Field( + default=None, description="""System instruction for the Gemini example.""" ) - branch: Optional[str] = Field( - default=None, - description="""Optional. The branch of the event. The format is like agent_1.agent_2.agent_3, where agent_1 is the parent of agent_2, and agent_2 is the parent of agent_3. Branch is used when multiple child agents shouldn't see their siblings' conversation history.""", + cached_content: Optional[str] = Field( + default=None, description="""Cached content for the Gemini example.""" ) - custom_metadata: Optional[dict[str, Any]] = Field( - default=None, description="""The custom metadata of the LlmResponse.""" + tools: Optional[list[genai_types.Tool]] = Field( + default=None, description="""Tools for the Gemini example.""" ) - interrupted: Optional[bool] = Field( - default=None, - description="""Optional. Flag indicating that LLM was interrupted when generating the content. Usually it's due to user interruption during a bidi streaming.""", + tool_config: Optional[genai_types.ToolConfig] = Field( + default=None, description="""Tools for the Gemini example.""" ) - long_running_tool_ids: Optional[list[str]] = Field( - default=None, - description="""Optional. Set of ids of the long running function calls. Agent client will know from this field about which function call is long running. Only valid for function call event.""", + safety_settings: Optional[list[genai_types.SafetySetting]] = Field( + default=None, description="""Safety settings for the Gemini example.""" ) - partial: Optional[bool] = Field( - default=None, - description="""Optional. Indicates whether the text content is part of a unfinished text stream. Only used for streaming mode and when the content is plain text.""", + generation_config: Optional[genai_types.GenerationConfig] = Field( + default=None, description="""Generation config for the Gemini example.""" ) - turn_complete: Optional[bool] = Field( + model_armor_config: Optional[genai_types.ModelArmorConfig] = Field( default=None, - description="""Optional. Indicates whether the response from the model is complete. Only used for streaming mode.""", - ) - input_transcription: Optional[genai_types.Transcription] = Field( - default=None, description="""Optional. Audio transcription of user input.""" - ) - output_transcription: Optional[genai_types.Transcription] = Field( - default=None, description="""Optional. Audio transcription of model output.""" + description="""Optional. Settings for prompt and response sanitization using the Model Armor service. If supplied, safety_settings must not be supplied.""", ) -class EventMetadataDict(TypedDict, total=False): - """Metadata relating to a LLM response event.""" +class GeminiExampleDict(TypedDict, total=False): + """Represents a Gemini example.""" - grounding_metadata: Optional[genai_types.GroundingMetadata] - """Optional. Metadata returned to client when grounding is enabled.""" + model: Optional[str] + """The model used to generate the Gemini example.""" - branch: Optional[str] - """Optional. The branch of the event. The format is like agent_1.agent_2.agent_3, where agent_1 is the parent of agent_2, and agent_2 is the parent of agent_3. Branch is used when multiple child agents shouldn't see their siblings' conversation history.""" + contents: Optional[list[genai_types.Content]] + """Contents of the Gemini example.""" - custom_metadata: Optional[dict[str, Any]] - """The custom metadata of the LlmResponse.""" + system_instruction: Optional[genai_types.Content] + """System instruction for the Gemini example.""" - interrupted: Optional[bool] - """Optional. Flag indicating that LLM was interrupted when generating the content. Usually it's due to user interruption during a bidi streaming.""" + cached_content: Optional[str] + """Cached content for the Gemini example.""" - long_running_tool_ids: Optional[list[str]] - """Optional. Set of ids of the long running function calls. Agent client will know from this field about which function call is long running. Only valid for function call event.""" + tools: Optional[list[genai_types.Tool]] + """Tools for the Gemini example.""" - partial: Optional[bool] - """Optional. Indicates whether the text content is part of a unfinished text stream. Only used for streaming mode and when the content is plain text.""" + tool_config: Optional[genai_types.ToolConfig] + """Tools for the Gemini example.""" - turn_complete: Optional[bool] - """Optional. Indicates whether the response from the model is complete. Only used for streaming mode.""" + safety_settings: Optional[list[genai_types.SafetySetting]] + """Safety settings for the Gemini example.""" - input_transcription: Optional[genai_types.Transcription] - """Optional. Audio transcription of user input.""" + generation_config: Optional[genai_types.GenerationConfig] + """Generation config for the Gemini example.""" - output_transcription: Optional[genai_types.Transcription] - """Optional. Audio transcription of model output.""" + model_armor_config: Optional[genai_types.ModelArmorConfig] + """Optional. Settings for prompt and response sanitization using the Model Armor service. If supplied, safety_settings must not be supplied.""" -EventMetadataOrDict = Union[EventMetadata, EventMetadataDict] +GeminiExampleOrDict = Union[GeminiExample, GeminiExampleDict] -class AppendAgentEngineSessionEventConfig(_common.BaseModel): - """Config for appending agent engine session event.""" +class GeminiTemplateConfig(_common.BaseModel): + """Represents a Gemini template config.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - content: Optional[genai_types.Content] = Field( - default=None, description="""The content of the session event.""" - ) - actions: Optional[EventActions] = Field( + gemini_example: Optional[GeminiExample] = Field( default=None, - description="""Actions are parts of events that are related to the session event.""", - ) - error_code: Optional[str] = Field( - default=None, description="""The error code of the session event.""" - ) - error_message: Optional[str] = Field( - default=None, description="""The error message of the session event.""" - ) - event_metadata: Optional[EventMetadata] = Field( - default=None, description="""Metadata relating to the session event.""" + description="""Required. The template that will be used for assembling the request to use for downstream applications.""", ) - raw_event: Optional[dict[str, Any]] = Field( + field_mapping: Optional[dict[str, str]] = Field( default=None, - description="""Weakly typed raw event data in proto struct format.""", + description="""Required. Map of template parameters to the columns in the dataset table.""", ) -class AppendAgentEngineSessionEventConfigDict(TypedDict, total=False): - """Config for appending agent engine session event.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - content: Optional[genai_types.Content] - """The content of the session event.""" - - actions: Optional[EventActionsDict] - """Actions are parts of events that are related to the session event.""" - - error_code: Optional[str] - """The error code of the session event.""" - - error_message: Optional[str] - """The error message of the session event.""" +class GeminiTemplateConfigDict(TypedDict, total=False): + """Represents a Gemini template config.""" - event_metadata: Optional[EventMetadataDict] - """Metadata relating to the session event.""" + gemini_example: Optional[GeminiExampleDict] + """Required. The template that will be used for assembling the request to use for downstream applications.""" - raw_event: Optional[dict[str, Any]] - """Weakly typed raw event data in proto struct format.""" + field_mapping: Optional[dict[str, str]] + """Required. Map of template parameters to the columns in the dataset table.""" -AppendAgentEngineSessionEventConfigOrDict = Union[ - AppendAgentEngineSessionEventConfig, AppendAgentEngineSessionEventConfigDict -] +GeminiTemplateConfigOrDict = Union[GeminiTemplateConfig, GeminiTemplateConfigDict] -class _AppendAgentEngineSessionEventRequestParameters(_common.BaseModel): - """Parameters for appending agent engines.""" +class GeminiRequestReadConfig(_common.BaseModel): + """Represents the config for reading Gemini requests.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine session.""" - ) - author: Optional[str] = Field( - default=None, description="""Author of the agent engine session event.""" - ) - invocation_id: Optional[str] = Field( - default=None, description="""Invocation ID of the agent engine.""" - ) - timestamp: Optional[datetime.datetime] = Field( - default=None, description="""Timestamp indicating when the event was created.""" + template_config: Optional[GeminiTemplateConfig] = Field( + default=None, description="""Gemini request template with placeholders.""" ) - config: Optional[AppendAgentEngineSessionEventConfig] = Field( - default=None, description="""""" + assembled_request_column_name: Optional[str] = Field( + default=None, + description="""Column name in the underlying BigQuery table that contains already fully assembled Gemini requests.""", ) + @classmethod + def single_turn_template( + cls, + *, + prompt: str, + response: Optional[str] = None, + system_instruction: Optional[str] = None, + model: Optional[str] = None, + cached_content: Optional[str] = None, + tools: Optional[list[Union[genai_types.Tool, dict[str, Any]]]] = None, + tool_config: Optional[Union[genai_types.ToolConfig, dict[str, Any]]] = None, + safety_settings: Optional[ + list[Union[genai_types.SafetySetting, dict[str, Any]]] + ] = None, + generation_config: Optional[ + Union[genai_types.GenerationConfig, dict[str, Any]] + ] = None, + field_mapping: Optional[dict[str, str]] = None, + ) -> "GeminiRequestReadConfig": + """Constructs a GeminiRequestReadConfig object for single-turn cases. -class _AppendAgentEngineSessionEventRequestParametersDict(TypedDict, total=False): - """Parameters for appending agent engines.""" - - name: Optional[str] - """Name of the agent engine session.""" - - author: Optional[str] - """Author of the agent engine session event.""" + 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." + ) - invocation_id: Optional[str] - """Invocation ID of the agent engine.""" + 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. - timestamp: Optional[datetime.datetime] - """Timestamp indicating when the event was created.""" - - config: Optional[AppendAgentEngineSessionEventConfigDict] - """""" - - -_AppendAgentEngineSessionEventRequestParametersOrDict = Union[ - _AppendAgentEngineSessionEventRequestParameters, - _AppendAgentEngineSessionEventRequestParametersDict, -] + Returns: + A GeminiRequestReadConfig object. + """ + contents = [] + contents.append( + genai_types.Content( + role="user", + parts=[ + genai_types.Part.from_text(text=prompt), + ], + ) + ) + if response: + contents.append( + genai_types.Content( + role="model", + parts=[ + genai_types.Part.from_text(text=response), + ], + ) + ) + system_instruction_content = None + if system_instruction: + system_instruction_content = genai_types.Content( + parts=[ + genai_types.Part.from_text(text=system_instruction), + ], + ) -class AppendAgentEngineSessionEventResponse(_common.BaseModel): - """Response for appending agent engine session event.""" + return cls( + template_config=GeminiTemplateConfig( + gemini_example=GeminiExample( + model=model, + contents=contents, + system_instruction=system_instruction_content, + cached_content=cached_content, + tools=tools, + tool_config=tool_config, + safety_settings=safety_settings, + generation_config=generation_config, + ), + field_mapping=field_mapping, + ), + ) - pass +class GeminiRequestReadConfigDict(TypedDict, total=False): + """Represents the config for reading Gemini requests.""" -class AppendAgentEngineSessionEventResponseDict(TypedDict, total=False): - """Response for appending agent engine session event.""" + template_config: Optional[GeminiTemplateConfigDict] + """Gemini request template with placeholders.""" - pass + assembled_request_column_name: Optional[str] + """Column name in the underlying BigQuery table that contains already fully assembled Gemini requests.""" -AppendAgentEngineSessionEventResponseOrDict = Union[ - AppendAgentEngineSessionEventResponse, AppendAgentEngineSessionEventResponseDict +GeminiRequestReadConfigOrDict = Union[ + GeminiRequestReadConfig, GeminiRequestReadConfigDict ] -class ListAgentEngineSessionEventsConfig(_common.BaseModel): - """Config for listing agent engine session events.""" +class AssembleDatasetConfig(_common.BaseModel): + """Config for assembling a multimodal dataset resource.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - page_size: Optional[int] = Field(default=None, description="""""") - page_token: Optional[str] = Field(default=None, description="""""") - filter: Optional[str] = Field( - default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""", + timeout: Optional[int] = Field( + default=90, + description="""The timeout for the assemble dataset request in seconds. If not + set, the default timeout is 90 seconds.""", ) -class ListAgentEngineSessionEventsConfigDict(TypedDict, total=False): - """Config for listing agent engine session events.""" +class AssembleDatasetConfigDict(TypedDict, total=False): + """Config for assembling a multimodal dataset resource.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - page_size: Optional[int] - """""" - - page_token: Optional[str] - """""" - - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" + timeout: Optional[int] + """The timeout for the assemble dataset request in seconds. If not + set, the default timeout is 90 seconds.""" -ListAgentEngineSessionEventsConfigOrDict = Union[ - ListAgentEngineSessionEventsConfig, ListAgentEngineSessionEventsConfigDict -] +AssembleDatasetConfigOrDict = Union[AssembleDatasetConfig, AssembleDatasetConfigDict] -class _ListAgentEngineSessionEventsRequestParameters(_common.BaseModel): - """Parameters for listing agent engine session events.""" +class _AssembleDatasetParameters(_common.BaseModel): + """Parameters for assembling a multimodal dataset resource.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine session.""" - ) - config: Optional[ListAgentEngineSessionEventsConfig] = Field( + name: Optional[str] = Field(default=None, description="""""") + gemini_request_read_config: Optional[GeminiRequestReadConfig] = Field( default=None, description="""""" ) + config: Optional[AssembleDatasetConfig] = Field(default=None, description="""""") -class _ListAgentEngineSessionEventsRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engine session events.""" +class _AssembleDatasetParametersDict(TypedDict, total=False): + """Parameters for assembling a multimodal dataset resource.""" name: Optional[str] - """Name of the agent engine session.""" + """""" + + gemini_request_read_config: Optional[GeminiRequestReadConfigDict] + """""" - config: Optional[ListAgentEngineSessionEventsConfigDict] + config: Optional[AssembleDatasetConfigDict] """""" -_ListAgentEngineSessionEventsRequestParametersOrDict = Union[ - _ListAgentEngineSessionEventsRequestParameters, - _ListAgentEngineSessionEventsRequestParametersDict, +_AssembleDatasetParametersOrDict = Union[ + _AssembleDatasetParameters, _AssembleDatasetParametersDict ] -class SessionEvent(_common.BaseModel): - """A session event.""" +class MultimodalDatasetOperation(_common.BaseModel): + """Represents the create dataset operation.""" - content: Optional[genai_types.Content] = Field( - default=None, - description="""Optional. Content of the event provided by the author.""", - ) - actions: Optional[EventActions] = Field( - default=None, description="""Optional. Actions executed by the agent.""" - ) - author: Optional[str] = Field( - default=None, - description="""Required. The name of the agent that sent the event, or user.""", - ) - error_code: Optional[str] = Field( - default=None, - description="""Optional. Error code if the response is an error. Code varies by model.""", - ) - error_message: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Optional. Error message if the response is an error.""", - ) - event_metadata: Optional[EventMetadata] = Field( - default=None, description="""Optional. Metadata relating to this event.""" + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", ) - invocation_id: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Required. The invocation id of the event, multiple events can have the same invocation id.""", + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", ) - name: Optional[str] = Field( + done: Optional[bool] = Field( default=None, - description="""Identifier. The resource name of the event. Format:`projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}/events/{event}`.""", + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", ) - timestamp: Optional[datetime.datetime] = Field( + error: Optional[dict[str, Any]] = Field( default=None, - description="""Required. Timestamp when the event was created on client side.""", + description="""The error result of the operation in case of failure or cancellation.""", ) - raw_event: Optional[dict[str, Any]] = Field( - default=None, - description="""Optional. Weakly typed raw event data in proto struct format.""", + response: Optional[dict[str, Any]] = Field( + default=None, description="""The result of the dataset operation.""" ) -class SessionEventDict(TypedDict, total=False): - """A session event.""" +class MultimodalDatasetOperationDict(TypedDict, total=False): + """Represents the create dataset operation.""" - content: Optional[genai_types.Content] - """Optional. Content of the event provided by the author.""" + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - actions: Optional[EventActionsDict] - """Optional. Actions executed by the agent.""" + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - author: Optional[str] - """Required. The name of the agent that sent the event, or user.""" + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - error_code: Optional[str] - """Optional. Error code if the response is an error. Code varies by model.""" + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" - error_message: Optional[str] - """Optional. Error message if the response is an error.""" + response: Optional[dict[str, Any]] + """The result of the dataset operation.""" - event_metadata: Optional[EventMetadataDict] - """Optional. Metadata relating to this event.""" - invocation_id: Optional[str] - """Required. The invocation id of the event, multiple events can have the same invocation id.""" +MultimodalDatasetOperationOrDict = Union[ + MultimodalDatasetOperation, MultimodalDatasetOperationDict +] - name: Optional[str] - """Identifier. The resource name of the event. Format:`projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}/events/{event}`.""" - timestamp: Optional[datetime.datetime] - """Required. Timestamp when the event was created on client side.""" +class TuningResourceUsageAssessmentConfig(_common.BaseModel): + """Config for tuning resource usage assessment.""" - raw_event: Optional[dict[str, Any]] - """Optional. Weakly typed raw event data in proto struct format.""" + model_name: Optional[str] = Field(default=None, description="""""") -SessionEventOrDict = Union[SessionEvent, SessionEventDict] +class TuningResourceUsageAssessmentConfigDict(TypedDict, total=False): + """Config for tuning resource usage assessment.""" + model_name: Optional[str] + """""" -class ListAgentEngineSessionEventsResponse(_common.BaseModel): - """Response for listing agent engine session events.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" - ) - next_page_token: Optional[str] = Field(default=None, description="""""") - session_events: Optional[list[SessionEvent]] = Field( - default=None, description="""List of session events.""" - ) +TuningResourceUsageAssessmentConfigOrDict = Union[ + TuningResourceUsageAssessmentConfig, TuningResourceUsageAssessmentConfigDict +] -class ListAgentEngineSessionEventsResponseDict(TypedDict, total=False): - """Response for listing agent engine session events.""" +class TuningValidationAssessmentConfig(_common.BaseModel): + """Config for tuning validation assessment.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + model_name: Optional[str] = Field(default=None, description="""""") + dataset_usage: Optional[str] = Field(default=None, description="""""") - next_page_token: Optional[str] + +class TuningValidationAssessmentConfigDict(TypedDict, total=False): + """Config for tuning validation assessment.""" + + model_name: Optional[str] """""" - session_events: Optional[list[SessionEventDict]] - """List of session events.""" + dataset_usage: Optional[str] + """""" -ListAgentEngineSessionEventsResponseOrDict = Union[ - ListAgentEngineSessionEventsResponse, ListAgentEngineSessionEventsResponseDict +TuningValidationAssessmentConfigOrDict = Union[ + TuningValidationAssessmentConfig, TuningValidationAssessmentConfigDict ] -class GeminiExample(_common.BaseModel): - """Represents a Gemini example.""" +class BatchPredictionResourceUsageAssessmentConfig(_common.BaseModel): + """Config for batch prediction resource usage assessment.""" - model: Optional[str] = Field( - default=None, description="""The model used to generate the Gemini example.""" - ) - contents: Optional[list[genai_types.Content]] = Field( - default=None, description="""Contents of the Gemini example.""" - ) - system_instruction: Optional[genai_types.Content] = Field( - default=None, description="""System instruction for the Gemini example.""" - ) - cached_content: Optional[str] = Field( - default=None, description="""Cached content for the Gemini example.""" - ) - tools: Optional[list[genai_types.Tool]] = Field( - default=None, description="""Tools for the Gemini example.""" - ) - tool_config: Optional[genai_types.ToolConfig] = Field( - default=None, description="""Tools for the Gemini example.""" - ) - safety_settings: Optional[list[genai_types.SafetySetting]] = Field( - default=None, description="""Safety settings for the Gemini example.""" - ) - generation_config: Optional[genai_types.GenerationConfig] = Field( - default=None, description="""Generation config for the Gemini example.""" - ) - model_armor_config: Optional[genai_types.ModelArmorConfig] = Field( - default=None, - description="""Optional. Settings for prompt and response sanitization using the Model Armor service. If supplied, safety_settings must not be supplied.""", - ) + model_name: Optional[str] = Field(default=None, description="""""") -class GeminiExampleDict(TypedDict, total=False): - """Represents a Gemini example.""" +class BatchPredictionResourceUsageAssessmentConfigDict(TypedDict, total=False): + """Config for batch prediction resource usage assessment.""" - model: Optional[str] - """The model used to generate the Gemini example.""" + model_name: Optional[str] + """""" - contents: Optional[list[genai_types.Content]] - """Contents of the Gemini example.""" - system_instruction: Optional[genai_types.Content] - """System instruction for the Gemini example.""" +BatchPredictionResourceUsageAssessmentConfigOrDict = Union[ + BatchPredictionResourceUsageAssessmentConfig, + BatchPredictionResourceUsageAssessmentConfigDict, +] - cached_content: Optional[str] - """Cached content for the Gemini example.""" - tools: Optional[list[genai_types.Tool]] - """Tools for the Gemini example.""" +class BatchPredictionValidationAssessmentConfig(_common.BaseModel): + """Config for batch prediction validation assessment.""" - tool_config: Optional[genai_types.ToolConfig] - """Tools for the Gemini example.""" + model_name: Optional[str] = Field(default=None, description="""""") - safety_settings: Optional[list[genai_types.SafetySetting]] - """Safety settings for the Gemini example.""" - generation_config: Optional[genai_types.GenerationConfig] - """Generation config for the Gemini example.""" +class BatchPredictionValidationAssessmentConfigDict(TypedDict, total=False): + """Config for batch prediction validation assessment.""" - 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.""" + model_name: Optional[str] + """""" -GeminiExampleOrDict = Union[GeminiExample, GeminiExampleDict] +BatchPredictionValidationAssessmentConfigOrDict = Union[ + BatchPredictionValidationAssessmentConfig, + BatchPredictionValidationAssessmentConfigDict, +] -class GeminiTemplateConfig(_common.BaseModel): - """Represents a Gemini template config.""" +class AssessDatasetConfig(_common.BaseModel): + """Config for assessing a multimodal dataset resource.""" - gemini_example: Optional[GeminiExample] = Field( - default=None, - description="""Required. The template that will be used for assembling the request to use for downstream applications.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - field_mapping: Optional[dict[str, str]] = Field( - default=None, - description="""Required. Map of template parameters to the columns in the dataset table.""", + 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 GeminiTemplateConfigDict(TypedDict, total=False): - """Represents a Gemini template config.""" +class AssessDatasetConfigDict(TypedDict, total=False): + """Config for assessing a multimodal dataset resource.""" - gemini_example: Optional[GeminiExampleDict] - """Required. The template that will be used for assembling the request to use for downstream applications.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - field_mapping: Optional[dict[str, str]] - """Required. Map of template parameters to the columns in the dataset table.""" + timeout: Optional[int] + """The timeout for the assess dataset request in seconds. If not set, + the default timeout is 90 seconds.""" -GeminiTemplateConfigOrDict = Union[GeminiTemplateConfig, GeminiTemplateConfigDict] +AssessDatasetConfigOrDict = Union[AssessDatasetConfig, AssessDatasetConfigDict] -class GeminiRequestReadConfig(_common.BaseModel): - """Represents the config for reading Gemini requests.""" +class _AssessDatasetParameters(_common.BaseModel): + """Parameters for assessing a multimodal dataset resource.""" - template_config: Optional[GeminiTemplateConfig] = Field( - default=None, description="""Gemini request template with placeholders.""" + name: Optional[str] = Field(default=None, description="""""") + gemini_request_read_config: Optional[GeminiRequestReadConfig] = Field( + default=None, description="""""" ) - assembled_request_column_name: Optional[str] = Field( - default=None, - description="""Column name in the underlying BigQuery table that contains already fully assembled Gemini requests.""", + 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="""""") - @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. - - Example: - read_config = GeminiRequestReadConfig.single_turn_template( - prompt="Which flower is this {flower_image}?", - response="This is a {label}.", - system_instruction="You are a botanical classifier." - ) - Args: - prompt: Required. User input. - response: Optional. Model response to user input. - system_instruction: Optional. System instructions for the model. - model: Optional. The model to use for the GeminiExample. - cached_content: Optional. The cached content to use for the GeminiExample. - tools: Optional. The tools to use for the GeminiExample. - tool_config: Optional. The tool config to use for the GeminiExample. - safety_settings: Optional. The safety settings to use for the GeminiExample. - generation_config: Optional. The generation config to use for the GeminiExample. - field_mapping: Optional. Mapping of placeholders to dataset columns. +class _AssessDatasetParametersDict(TypedDict, total=False): + """Parameters for assessing a multimodal dataset resource.""" - 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), - ], - ) - ) + name: Optional[str] + """""" - system_instruction_content = None - if system_instruction: - system_instruction_content = genai_types.Content( - parts=[ - genai_types.Part.from_text(text=system_instruction), - ], - ) + gemini_request_read_config: Optional[GeminiRequestReadConfigDict] + """""" - 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, - ), - ) + tuning_resource_usage_assessment_config: Optional[ + TuningResourceUsageAssessmentConfigDict + ] + """""" + tuning_validation_assessment_config: Optional[TuningValidationAssessmentConfigDict] + """""" -class GeminiRequestReadConfigDict(TypedDict, total=False): - """Represents the config for reading Gemini requests.""" + batch_prediction_resource_usage_assessment_config: Optional[ + BatchPredictionResourceUsageAssessmentConfigDict + ] + """""" - template_config: Optional[GeminiTemplateConfigDict] - """Gemini request template with placeholders.""" + batch_prediction_validation_assessment_config: Optional[ + BatchPredictionValidationAssessmentConfigDict + ] + """""" - assembled_request_column_name: Optional[str] - """Column name in the underlying BigQuery table that contains already fully assembled Gemini requests.""" + config: Optional[AssessDatasetConfigDict] + """""" -GeminiRequestReadConfigOrDict = Union[ - GeminiRequestReadConfig, GeminiRequestReadConfigDict +_AssessDatasetParametersOrDict = Union[ + _AssessDatasetParameters, _AssessDatasetParametersDict ] -class AssembleDatasetConfig(_common.BaseModel): - """Config for assembling a multimodal dataset resource.""" +class SchemaTablesDatasetMetadataBigQuerySource(_common.BaseModel): + """Represents the BigQuery source for multimodal dataset metadata.""" - 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.""", + uri: Optional[str] = Field( + default=None, + description="""The URI of the BigQuery table. This accepts the table name with or without the bq:// prefix.""", ) -class AssembleDatasetConfigDict(TypedDict, total=False): - """Config for assembling a multimodal dataset resource.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class SchemaTablesDatasetMetadataBigQuerySourceDict(TypedDict, total=False): + """Represents the BigQuery source for multimodal dataset metadata.""" - timeout: Optional[int] - """The timeout for the assemble dataset request in seconds. If not - set, the default timeout is 90 seconds.""" + uri: Optional[str] + """The URI of the BigQuery table. This accepts the table name with or without the bq:// prefix.""" -AssembleDatasetConfigOrDict = Union[AssembleDatasetConfig, AssembleDatasetConfigDict] +SchemaTablesDatasetMetadataBigQuerySourceOrDict = Union[ + SchemaTablesDatasetMetadataBigQuerySource, + SchemaTablesDatasetMetadataBigQuerySourceDict, +] -class _AssembleDatasetParameters(_common.BaseModel): - """Parameters for assembling a multimodal dataset resource.""" +class SchemaTablesDatasetMetadataInputConfig(_common.BaseModel): + """Represents the input config for multimodal dataset metadata.""" - name: Optional[str] = Field(default=None, description="""""") - gemini_request_read_config: Optional[GeminiRequestReadConfig] = Field( - default=None, description="""""" + bigquery_source: Optional[SchemaTablesDatasetMetadataBigQuerySource] = Field( + default=None, + description="""The BigQuery source for multimodal dataset metadata.""", ) - config: Optional[AssembleDatasetConfig] = Field(default=None, description="""""") -class _AssembleDatasetParametersDict(TypedDict, total=False): - """Parameters for assembling a multimodal dataset resource.""" +class SchemaTablesDatasetMetadataInputConfigDict(TypedDict, total=False): + """Represents the input config for multimodal dataset metadata.""" - name: Optional[str] - """""" + bigquery_source: Optional[SchemaTablesDatasetMetadataBigQuerySourceDict] + """The BigQuery source for multimodal dataset metadata.""" - gemini_request_read_config: Optional[GeminiRequestReadConfigDict] - """""" - config: Optional[AssembleDatasetConfigDict] - """""" - - -_AssembleDatasetParametersOrDict = Union[ - _AssembleDatasetParameters, _AssembleDatasetParametersDict +SchemaTablesDatasetMetadataInputConfigOrDict = Union[ + SchemaTablesDatasetMetadataInputConfig, SchemaTablesDatasetMetadataInputConfigDict ] -class MultimodalDatasetOperation(_common.BaseModel): - """Represents the create dataset operation.""" +class SchemaTablesDatasetMetadata(_common.BaseModel): + """Represents the metadata schema for multimodal 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( + input_config: Optional[SchemaTablesDatasetMetadataInputConfig] = 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 input config for multimodal dataset metadata.""", ) - error: Optional[dict[str, Any]] = Field( + gemini_request_read_config: Optional[GeminiRequestReadConfig] = 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="""The Gemini request read config for the multimodal dataset.""", ) -class MultimodalDatasetOperationDict(TypedDict, total=False): - """Represents the create dataset operation.""" - - name: Optional[str] - """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" - - metadata: Optional[dict[str, Any]] - """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" - - done: Optional[bool] - """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" - - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" - - response: Optional[dict[str, Any]] - """The result of the dataset operation.""" - - -MultimodalDatasetOperationOrDict = Union[ - MultimodalDatasetOperation, MultimodalDatasetOperationDict -] - - -class TuningResourceUsageAssessmentConfig(_common.BaseModel): - """Config for tuning resource usage assessment.""" - - model_name: Optional[str] = Field(default=None, description="""""") - - -class TuningResourceUsageAssessmentConfigDict(TypedDict, total=False): - """Config for tuning resource usage assessment.""" - - model_name: Optional[str] - """""" - - -TuningResourceUsageAssessmentConfigOrDict = Union[ - TuningResourceUsageAssessmentConfig, TuningResourceUsageAssessmentConfigDict -] - - -class TuningValidationAssessmentConfig(_common.BaseModel): - """Config for tuning validation assessment.""" - - model_name: Optional[str] = Field(default=None, description="""""") - dataset_usage: Optional[str] = Field(default=None, description="""""") - - -class TuningValidationAssessmentConfigDict(TypedDict, total=False): - """Config for tuning validation assessment.""" - - model_name: Optional[str] - """""" - - dataset_usage: Optional[str] - """""" - - -TuningValidationAssessmentConfigOrDict = Union[ - TuningValidationAssessmentConfig, TuningValidationAssessmentConfigDict -] - - -class BatchPredictionResourceUsageAssessmentConfig(_common.BaseModel): - """Config for batch prediction resource usage assessment.""" - - model_name: Optional[str] = Field(default=None, description="""""") - - -class BatchPredictionResourceUsageAssessmentConfigDict(TypedDict, total=False): - """Config for batch prediction resource usage assessment.""" - - model_name: Optional[str] - """""" - - -BatchPredictionResourceUsageAssessmentConfigOrDict = Union[ - BatchPredictionResourceUsageAssessmentConfig, - BatchPredictionResourceUsageAssessmentConfigDict, -] - - -class BatchPredictionValidationAssessmentConfig(_common.BaseModel): - """Config for batch prediction validation assessment.""" - - model_name: Optional[str] = Field(default=None, description="""""") - +class SchemaTablesDatasetMetadataDict(TypedDict, total=False): + """Represents the metadata schema for multimodal dataset metadata.""" -class BatchPredictionValidationAssessmentConfigDict(TypedDict, total=False): - """Config for batch prediction validation assessment.""" + input_config: Optional[SchemaTablesDatasetMetadataInputConfigDict] + """The input config for multimodal dataset metadata.""" - model_name: Optional[str] - """""" + gemini_request_read_config: Optional[GeminiRequestReadConfigDict] + """The Gemini request read config for the multimodal dataset.""" -BatchPredictionValidationAssessmentConfigOrDict = Union[ - BatchPredictionValidationAssessmentConfig, - BatchPredictionValidationAssessmentConfigDict, +SchemaTablesDatasetMetadataOrDict = Union[ + SchemaTablesDatasetMetadata, SchemaTablesDatasetMetadataDict ] -class AssessDatasetConfig(_common.BaseModel): - """Config for assessing a multimodal dataset resource.""" +class CreateMultimodalDatasetConfig(_common.BaseModel): + """Config for creating a dataset resource to store multimodal dataset.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) timeout: Optional[int] = Field( default=90, - description="""The timeout for the assess dataset request in seconds. If not set, + description="""The timeout for the create dataset request in seconds. If not set, the default timeout is 90 seconds.""", ) -class AssessDatasetConfigDict(TypedDict, total=False): - """Config for assessing a multimodal dataset resource.""" +class CreateMultimodalDatasetConfigDict(TypedDict, total=False): + """Config for creating a dataset resource to store multimodal dataset.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" timeout: Optional[int] - """The timeout for the assess dataset request in seconds. If not set, + """The timeout for the create dataset request in seconds. If not set, the default timeout is 90 seconds.""" -AssessDatasetConfigOrDict = Union[AssessDatasetConfig, AssessDatasetConfigDict] +CreateMultimodalDatasetConfigOrDict = Union[ + CreateMultimodalDatasetConfig, CreateMultimodalDatasetConfigDict +] -class _AssessDatasetParameters(_common.BaseModel): - """Parameters for assessing a multimodal dataset resource.""" +class _CreateMultimodalDatasetParameters(_common.BaseModel): + """Parameters for creating a dataset resource to store multimodal dataset.""" name: Optional[str] = Field(default=None, description="""""") - gemini_request_read_config: Optional[GeminiRequestReadConfig] = Field( + display_name: Optional[str] = Field(default=None, description="""""") + metadata_schema_uri: 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="""""" + ) + config: Optional[CreateMultimodalDatasetConfig] = Field( + default=None, description="""""" ) - batch_prediction_resource_usage_assessment_config: Optional[ - BatchPredictionResourceUsageAssessmentConfig - ] = Field(default=None, description="""""") - batch_prediction_validation_assessment_config: Optional[ - BatchPredictionValidationAssessmentConfig - ] = Field(default=None, description="""""") - config: Optional[AssessDatasetConfig] = Field(default=None, description="""""") -class _AssessDatasetParametersDict(TypedDict, total=False): - """Parameters for assessing a multimodal dataset resource.""" +class _CreateMultimodalDatasetParametersDict(TypedDict, total=False): + """Parameters for creating a dataset resource to store multimodal dataset.""" name: Optional[str] """""" - gemini_request_read_config: Optional[GeminiRequestReadConfigDict] + display_name: Optional[str] """""" - tuning_resource_usage_assessment_config: Optional[ - TuningResourceUsageAssessmentConfigDict - ] + metadata_schema_uri: 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[CreateMultimodalDatasetConfigDict] """""" -_AssessDatasetParametersOrDict = Union[ - _AssessDatasetParameters, _AssessDatasetParametersDict +_CreateMultimodalDatasetParametersOrDict = Union[ + _CreateMultimodalDatasetParameters, _CreateMultimodalDatasetParametersDict ] -class SchemaTablesDatasetMetadataBigQuerySource(_common.BaseModel): - """Represents the BigQuery source for multimodal dataset metadata.""" +class _DeleteMultimodalDatasetRequestParameters(_common.BaseModel): + """Parameters for deleting a multimodal dataset.""" - uri: Optional[str] = Field( - default=None, - description="""The URI of the BigQuery table. This accepts the table name with or without the bq:// prefix.""", + name: Optional[str] = Field( + default=None, description="""ID of the dataset to be deleted.""" ) + config: Optional[VertexBaseConfig] = Field(default=None, description="""""") -class SchemaTablesDatasetMetadataBigQuerySourceDict(TypedDict, total=False): - """Represents the BigQuery source for multimodal dataset metadata.""" +class _DeleteMultimodalDatasetRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a multimodal dataset.""" - uri: Optional[str] - """The URI of the BigQuery table. This accepts the table name with or without the bq:// prefix.""" + name: Optional[str] + """ID of the dataset to be deleted.""" + config: Optional[VertexBaseConfigDict] + """""" -SchemaTablesDatasetMetadataBigQuerySourceOrDict = Union[ - SchemaTablesDatasetMetadataBigQuerySource, - SchemaTablesDatasetMetadataBigQuerySourceDict, + +_DeleteMultimodalDatasetRequestParametersOrDict = Union[ + _DeleteMultimodalDatasetRequestParameters, + _DeleteMultimodalDatasetRequestParametersDict, ] -class SchemaTablesDatasetMetadataInputConfig(_common.BaseModel): - """Represents the input config for multimodal dataset metadata.""" +class _GetMultimodalDatasetParameters(_common.BaseModel): + """Parameters for getting a multimodal dataset resource.""" - bigquery_source: Optional[SchemaTablesDatasetMetadataBigQuerySource] = Field( - default=None, - description="""The BigQuery source for multimodal dataset metadata.""", - ) + name: Optional[str] = Field(default=None, description="""""") + config: Optional[VertexBaseConfig] = Field(default=None, description="""""") -class SchemaTablesDatasetMetadataInputConfigDict(TypedDict, total=False): - """Represents the input config for multimodal dataset metadata.""" +class _GetMultimodalDatasetParametersDict(TypedDict, total=False): + """Parameters for getting a multimodal dataset resource.""" - bigquery_source: Optional[SchemaTablesDatasetMetadataBigQuerySourceDict] - """The BigQuery source for multimodal dataset metadata.""" + name: Optional[str] + """""" + + config: Optional[VertexBaseConfigDict] + """""" -SchemaTablesDatasetMetadataInputConfigOrDict = Union[ - SchemaTablesDatasetMetadataInputConfig, SchemaTablesDatasetMetadataInputConfigDict -] - - -class SchemaTablesDatasetMetadata(_common.BaseModel): - """Represents the metadata schema for multimodal dataset metadata.""" - - 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.""", - ) - - -class SchemaTablesDatasetMetadataDict(TypedDict, total=False): - """Represents the metadata schema for multimodal dataset metadata.""" - - input_config: Optional[SchemaTablesDatasetMetadataInputConfigDict] - """The input config for multimodal dataset metadata.""" - - gemini_request_read_config: Optional[GeminiRequestReadConfigDict] - """The Gemini request read config for the multimodal dataset.""" - - -SchemaTablesDatasetMetadataOrDict = Union[ - SchemaTablesDatasetMetadata, SchemaTablesDatasetMetadataDict -] - - -class CreateMultimodalDatasetConfig(_common.BaseModel): - """Config for creating a dataset resource to store multimodal dataset.""" - - 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.""", - ) - - -class CreateMultimodalDatasetConfigDict(TypedDict, total=False): - """Config for creating a dataset resource to store multimodal dataset.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - timeout: Optional[int] - """The timeout for the create dataset request in seconds. If not set, - the default timeout is 90 seconds.""" - - -CreateMultimodalDatasetConfigOrDict = Union[ - CreateMultimodalDatasetConfig, CreateMultimodalDatasetConfigDict -] - - -class _CreateMultimodalDatasetParameters(_common.BaseModel): - """Parameters for creating a dataset resource to store multimodal dataset.""" - - 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="""""" - ) - - -class _CreateMultimodalDatasetParametersDict(TypedDict, total=False): - """Parameters for creating a dataset resource to store multimodal dataset.""" - - name: Optional[str] - """""" - - display_name: Optional[str] - """""" - - metadata_schema_uri: Optional[str] - """""" - - metadata: Optional[SchemaTablesDatasetMetadataDict] - """""" - - description: Optional[str] - """""" - - encryption_spec: Optional[genai_types.EncryptionSpec] - """""" - - config: Optional[CreateMultimodalDatasetConfigDict] - """""" - - -_CreateMultimodalDatasetParametersOrDict = Union[ - _CreateMultimodalDatasetParameters, _CreateMultimodalDatasetParametersDict -] - - -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 _DeleteMultimodalDatasetRequestParametersDict(TypedDict, total=False): - """Parameters for deleting a multimodal dataset.""" - - name: Optional[str] - """ID of the dataset to be deleted.""" - - config: Optional[VertexBaseConfigDict] - """""" - - -_DeleteMultimodalDatasetRequestParametersOrDict = Union[ - _DeleteMultimodalDatasetRequestParameters, - _DeleteMultimodalDatasetRequestParametersDict, -] - - -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 _GetMultimodalDatasetParametersDict(TypedDict, total=False): - """Parameters for getting a multimodal dataset resource.""" - - name: Optional[str] - """""" - - config: Optional[VertexBaseConfigDict] - """""" - - -_GetMultimodalDatasetParametersOrDict = Union[ - _GetMultimodalDatasetParameters, _GetMultimodalDatasetParametersDict +_GetMultimodalDatasetParametersOrDict = Union[ + _GetMultimodalDatasetParameters, _GetMultimodalDatasetParametersDict ] @@ -30784,1114 +29602,937 @@ class EvalRunInferenceConfigDict(TypedDict, total=False): EvalRunInferenceConfigOrDict = Union[EvalRunInferenceConfig, EvalRunInferenceConfigDict] -class AgentEngine(_common.BaseModel): - """An agent engine instance.""" +class AssembleDataset(_common.BaseModel): + """Represents the assembled dataset.""" - api_client: Optional[Any] = Field( - default=None, description="""The underlying API client.""" - ) - api_async_client: Optional[Any] = Field( - default=None, - description="""The underlying API client for asynchronous operations.""", - ) - api_resource: Optional[ReasoningEngine] = Field( + bigquery_destination: Optional[str] = Field( default=None, - description="""The underlying API resource (i.e. ReasoningEngine).""", + description="""The BigQuery destination of the assembled dataset.""", ) - # Allows dynamic binding of methods based on the registered operations. - model_config = ConfigDict(extra="allow") - def __repr__(self) -> str: - return ( - f"AgentEngine(api_resource.name='{self.api_resource.name}')" - if self.api_resource is not None - else "AgentEngine(api_resource.name=None)" - ) +class AssembleDatasetDict(TypedDict, total=False): + """Represents the assembled dataset.""" - def operation_schemas(self) -> Optional[list[Dict[str, Any]]]: - """Returns the schemas of all registered operations for the agent.""" - if not isinstance(self.api_resource, ReasoningEngine): - raise ValueError("api_resource is not initialized.") - if not self.api_resource.spec: - raise ValueError("api_resource.spec is not initialized.") - return self.api_resource.spec.class_methods + bigquery_destination: Optional[str] + """The BigQuery destination of the assembled dataset.""" - def delete( - self, - force: bool = False, - config: Optional[DeleteAgentEngineConfigOrDict] = None, - ) -> None: - """Deletes the agent engine. - Args: - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error when - the Agent Engine has undeleted child resources. Defaults to False. - config (DeleteAgentEngineConfig): - Optional. Additional configurations for deleting the Agent Engine. - """ - if not isinstance(self.api_resource, ReasoningEngine): - raise ValueError("api_resource is not initialized.") - self.api_client.delete(name=self.api_resource.name, force=force, config=config) # type: ignore[union-attr] +AssembleDatasetOrDict = Union[AssembleDataset, AssembleDatasetDict] -RubricContentProperty = evals_types.RubricContentProperty -RubricContentPropertyDict = evals_types.RubricContentPropertyDict -RubricContentPropertyDictOrDict = evals_types.RubricContentPropertyOrDict +class BatchPredictionResourceUsageAssessmentResult(_common.BaseModel): + """Result of batch prediction resource usage assessment.""" -RubricContent = evals_types.RubricContent -RubricContentDict = evals_types.RubricContentDict -RubricContentDictOrDict = evals_types.RubricContentOrDict + token_count: Optional[int] = Field( + default=None, description="""Number of tokens in the dataset.""" + ) + audio_token_count: Optional[int] = Field( + default=None, description="""Number of audio tokens in the dataset.""" + ) -Rubric = evals_types.Rubric -RubricDict = evals_types.RubricDict -RubricDictOrDict = evals_types.RubricOrDict -RubricVerdict = evals_types.RubricVerdict -RubricVerdictDict = evals_types.RubricVerdictDict -RubricVerdictDictOrDict = evals_types.RubricVerdictOrDict +class BatchPredictionResourceUsageAssessmentResultDict(TypedDict, total=False): + """Result of batch prediction resource usage assessment.""" -CandidateResult = evals_types.CandidateResult -CandidateResultDict = evals_types.CandidateResultDict -CandidateResultDictOrDict = evals_types.CandidateResultOrDict + token_count: Optional[int] + """Number of tokens in the dataset.""" -Event = evals_types.Event -EventDict = evals_types.EventDict -EventDictOrDict = evals_types.EventOrDict + audio_token_count: Optional[int] + """Number of audio tokens in the dataset.""" -Message = evals_types.Message -MessageDict = evals_types.MessageDict -MessageDictOrDict = evals_types.MessageOrDict -Importance = evals_types.Importance +BatchPredictionResourceUsageAssessmentResultOrDict = Union[ + BatchPredictionResourceUsageAssessmentResult, + BatchPredictionResourceUsageAssessmentResultDict, +] -class AgentEngineDict(TypedDict, total=False): - """An agent engine instance.""" +class BatchPredictionValidationAssessmentResult(_common.BaseModel): + """Result of batch prediction validation assessment.""" - api_client: Optional[Any] - """The underlying API client.""" + errors: Optional[list[str]] = Field( + default=None, description="""The list of errors found in the dataset.""" + ) - api_async_client: Optional[Any] - """The underlying API client for asynchronous operations.""" - api_resource: Optional[ReasoningEngineDict] - """The underlying API resource (i.e. ReasoningEngine).""" +class BatchPredictionValidationAssessmentResultDict(TypedDict, total=False): + """Result of batch prediction validation assessment.""" + errors: Optional[list[str]] + """The list of errors found in the dataset.""" -AgentEngineOrDict = Union[AgentEngine, AgentEngineDict] +BatchPredictionValidationAssessmentResultOrDict = Union[ + BatchPredictionValidationAssessmentResult, + BatchPredictionValidationAssessmentResultDict, +] -class AgentEngineConfig(_common.BaseModel): - """Config for agent engine methods.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - staging_bucket: Optional[str] = Field( - default=None, - description="""The GCS bucket to use for staging the artifacts needed. +class TuningResourceUsageAssessmentResult(_common.BaseModel): + """Result of tuning resource usage assessment.""" - It must be a valid GCS bucket name, e.g. "gs://bucket-name". It is - required if `agent_engine` is specified.""", + token_count: Optional[int] = Field( + default=None, description="""The number of tokens in the dataset.""" ) - requirements: Optional[Any] = Field( + billable_character_count: Optional[int] = Field( default=None, - description="""The set of PyPI dependencies needed. - - It can either be the path to a single file (requirements.txt), or an - ordered list of strings corresponding to each line of the requirements - file.""", + description="""The number of billable characters in the dataset.""", ) - display_name: Optional[str] = Field( - default=None, - description="""The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any UTF-8 - character.""", - ) - description: Optional[str] = Field( - default=None, description="""The description of the Agent Engine.""" - ) - gcs_dir_name: Optional[str] = Field( - default=None, - description="""The GCS bucket directory under `staging_bucket` to use for staging - the artifacts needed.""", - ) - extra_packages: Optional[list[str]] = Field( - default=None, - description="""The set of extra user-provided packages (if any).""", - ) - env_vars: Optional[Any] = Field( - default=None, - description="""The environment variables to be set when running the Agent Engine. - If it is a dictionary, the keys are the environment variable names, and - the values are the corresponding values.""", - ) - service_account: Optional[str] = Field( - default=None, - description="""The service account to be used for the Agent Engine. - - 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 TuningResourceUsageAssessmentResultDict(TypedDict, total=False): + """Result of tuning resource usage assessment.""" + token_count: Optional[int] + """The number of tokens in the dataset.""" -class AgentEngineConfigDict(TypedDict, total=False): - """Config for agent engine methods.""" + billable_character_count: Optional[int] + """The number of billable characters in the dataset.""" - 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. +TuningResourceUsageAssessmentResultOrDict = Union[ + TuningResourceUsageAssessmentResult, TuningResourceUsageAssessmentResultDict +] - 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. +class TuningValidationAssessmentResult(_common.BaseModel): + """The result of a tuning validation assessment.""" - 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.""" + errors: Optional[list[str]] = Field( + default=None, description="""The list of errors found in the dataset.""" + ) - 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.""" +class TuningValidationAssessmentResultDict(TypedDict, total=False): + """The result of a tuning validation assessment.""" - description: Optional[str] - """The description of the Agent Engine.""" + errors: Optional[list[str]] + """The list of errors found in the dataset.""" - 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).""" +TuningValidationAssessmentResultOrDict = Union[ + TuningValidationAssessmentResult, TuningValidationAssessmentResultDict +] - 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.""" +class Prompt(_common.BaseModel): + """Represents a prompt.""" - service_account: Optional[str] - """The service account to be used for the Agent Engine. + prompt_data: Optional["PromptData"] = Field(default=None, description="""""") + _dataset: Optional["Dataset"] = PrivateAttr(default=None) + _dataset_version: Optional["DatasetVersion"] = PrivateAttr(default=None) - If not specified, the default Reasoning Engine P6SA service agent will be used.""" + @property + def dataset(self) -> "Dataset": + return self._dataset # type: ignore[return-value] - identity_type: Optional[IdentityType] - """The identity type to use for the Agent Engine.""" + @property + def dataset_version(self) -> "DatasetVersion": + return self._dataset_version # type: ignore[return-value] - context_spec: Optional[ReasoningEngineContextSpecDict] - """The context spec to be used for the Agent Engine.""" + @property + def prompt_id(self) -> Optional[str]: + """Returns the ID associated with the prompt resource.""" + if self._dataset and self._dataset.name: + return self._dataset.name.split("/")[-1] + elif not self._dataset and ( + self._dataset_version and self._dataset_version.name + ): + return self._dataset_version.name.split("datasets/")[1].split("/")[0] + return None - psc_interface_config: Optional[PscInterfaceConfigDict] - """The PSC interface config for PSC-I to be used for the Agent Engine.""" + @property + def version_id(self) -> Optional[str]: + """Returns the ID associated with the prompt version resource.""" + if self._dataset_version and self._dataset_version.name: + return self._dataset_version.name.split("/")[-1] + return None - min_instances: Optional[int] - """The minimum number of instances to run for the Agent Engine. - Defaults to 1. Range: [0, 10]. - """ + def assemble_contents(self) -> list[genai_types.Content]: + """Transforms a Prompt object into a list with a single genai_types.Content object. - 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]. - """ + This method replaces the variables in the prompt template with the values provided in prompt.prompt_data.variables. + If no variables are provided, prompt.prompt_data.contents is returned as is. Only single-turn prompts are supported. - 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'. - """ + This can be used to call generate_content() in the Gen AI SDK. - container_concurrency: Optional[int] - """The container concurrency to be used for the Agent Engine. - Recommended value: 2 * cpu + 1. Defaults to 9. - """ + Example usage: - encryption_spec: Optional[genai_types.EncryptionSpec] - """The encryption spec to be used for the Agent Engine.""" + my_prompt = types.Prompt( + prompt_data=types.PromptData( + model="gemini-2.0-flash-001", + contents=[ + genai_types.Content( + parts=[ + genai_types.Part(text="Hello {name}!"), + ], + ), + ], + variables=[ + { + "name": genai_types.Part(text="Alice"), + }, + ], + ), + ) - labels: Optional[dict[str, str]] - """The labels to be used for the Agent Engine.""" + from google import genai - agent_server_mode: Optional[AgentServerMode] - """The agent server mode to use for deployment.""" + genai_client = genai.Client(vertexai=True, project="my-project", location="us-central1") + genai_client.models.generate_content( + model=my_prompt.prompt_data.model, + contents=my_prompt.assemble_contents(), + ) - 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. - """ + Returns: + A list with a single Content object that can be used to call + generate_content(). + """ + if not self.prompt_data or not self.prompt_data.contents: + return [] - 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) - """ + if not self.prompt_data.variables: + return self.prompt_data.contents - 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.""" + if len(self.prompt_data.contents) > 1: + raise ValueError( + "Multiple contents are not supported. Use assemble_contents() for a prompt with a single Content item." + ) - entrypoint_module: Optional[str] - """The entrypoint module to be used for the Agent Engine - This field only used when source_packages is specified.""" + parts_to_process = self.prompt_data.contents[0].parts + if parts_to_process is None: + return [] + if not isinstance(parts_to_process, list): + parts_to_process = [parts_to_process] - entrypoint_object: Optional[str] - """The entrypoint object to be used for the Agent Engine. - This field only used when source_packages is specified.""" + has_placeholders = False + variable_regex = r"\{.*?\}" + for item in parts_to_process: + part = ( + item + if isinstance(item, genai_types.Part) + else genai_types.Part(text=str(item)) + ) + if part.text and re.search(variable_regex, part.text): + has_placeholders = True + break - 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. - """ + if not has_placeholders: + return [genai_types.Content(parts=parts_to_process)] - 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".""" + all_rendered_parts: list[genai_types.Part] = [] - 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". - """ + for var_dict in self.prompt_data.variables: + for template_item in parts_to_process: + template_part = ( + template_item + if isinstance(template_item, genai_types.Part) + else genai_types.Part(text=str(template_item)) + ) + if template_part.text: + rendered_text = template_part.text - 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`. - """ + for key, value in var_dict.items(): + placeholder = f"{{{key}}}" + replacement_text = None - image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict] - """The image spec for the Agent Engine.""" + if isinstance(value, str): + replacement_text = value + elif isinstance(value, genai_types.Part): + if value.text: + replacement_text = value.text + else: + all_rendered_parts.append(value) + if ( + replacement_text is not None + and placeholder in rendered_text + ): + rendered_text = rendered_text.replace( + placeholder, replacement_text + ) + all_rendered_parts.append(genai_types.Part(text=rendered_text)) + else: + all_rendered_parts.append(template_part) + return [genai_types.Content(parts=all_rendered_parts, role="user")] - agent_config_source: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict - ] - """The agent config source for the Agent Engine.""" - container_spec: Optional[ReasoningEngineSpecContainerSpecDict] - """The container spec for the Agent Engine.""" +PromptData = SchemaPromptSpecPromptMessage +PromptDataDict = SchemaPromptSpecPromptMessageDict +PromptDataOrDict = Union[PromptData, PromptDataDict] - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict - ] - """Agent Gateway configuration for a Reasoning Engine deployment.""" +ParsedResponseUnion = Union[ + prompts_types.ParsedResponse, prompts_types.ParsedResponseFewShot +] +ParsedResponseUnionDict = Union[ + prompts_types.ParsedResponseDict, prompts_types.ParsedResponseFewShotDict +] - 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.""" +class PromptDict(TypedDict, total=False): + """Represents a prompt.""" - 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).""" + prompt_data: Optional["PromptDataDict"] + """""" -AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict] +PromptOrDict = Union[Prompt, PromptDict] -class RunQueryJobAgentEngineConfig(_common.BaseModel): - """Config for checking a query job on an agent engine.""" +class CreatePromptConfig(_common.BaseModel): + """Config for creating a prompt.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - query: Optional[str] = Field( - default=None, description="""The query to send to the agent engine.""" + prompt_display_name: Optional[str] = Field( + default=None, + description="""The display name for the prompt. If not set, a default name with a timestamp will be used.""", ) - output_gcs_uri: Optional[str] = Field( + timeout: Optional[int] = Field( + default=90, + description="""The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, - description="""The GCS URI to use for the output. - If it is a file, the system use this file to store the response. - If it represents a directory, the system automatically generate a file - for the response. - In both cases, the input query will be stored in the same directory under - the same file name prefix as the output file.""", + description="""Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""", + ) + version_display_name: Optional[str] = Field( + default=None, + description="""The display name for the prompt version. If not set, a default name with a timestamp will be used.""", + ) + max_wait_time: Optional[int] = Field( + default=60, + description="""The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""", ) -class RunQueryJobAgentEngineConfigDict(TypedDict, total=False): - """Config for checking a query job on an agent engine.""" +class CreatePromptConfigDict(TypedDict, total=False): + """Config for creating a prompt.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - query: Optional[str] - """The query to send to the agent engine.""" + prompt_display_name: Optional[str] + """The display name for the prompt. If not set, a default name with a timestamp will be used.""" - output_gcs_uri: Optional[str] - """The GCS URI to use for the output. - If it is a file, the system use this file to store the response. - If it represents a directory, the system automatically generate a file - for the response. - In both cases, the input query will be stored in the same directory under - the same file name prefix as the output file.""" + timeout: Optional[int] + """The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""" -RunQueryJobAgentEngineConfigOrDict = Union[ - RunQueryJobAgentEngineConfig, RunQueryJobAgentEngineConfigDict -] + version_display_name: Optional[str] + """The display name for the prompt version. If not set, a default name with a timestamp will be used.""" + max_wait_time: Optional[int] + """The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""" -class RunQueryJobResult(_common.BaseModel): - """Result of running a query job.""" + +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.""" ) - job_name: Optional[str] = Field( + version_display_name: Optional[str] = Field( default=None, - description="""Name of the agent engine operation to later check for status.""", + description="""The display name for the prompt version. If not set, a default name with a timestamp will be used.""", ) - input_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the input file.""" + timeout: Optional[int] = Field( + default=90, + description="""The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""", ) - output_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the output file.""" + prompt_display_name: Optional[str] = Field( + default=None, + description="""The display name for the prompt. If not set, a default name with a timestamp will be used.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""", + ) + max_wait_time: Optional[int] = Field( + default=60, + description="""The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""", ) -class RunQueryJobResultDict(TypedDict, total=False): - """Result of running a query job.""" +class CreatePromptVersionConfigDict(TypedDict, total=False): + """Config for creating a prompt version.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - job_name: Optional[str] - """Name of the agent engine operation to later check for status.""" - - 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] - + 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 CheckQueryJobResponse(_common.BaseModel): - """Response from LRO.""" + timeout: Optional[int] + """The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""" - 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.""" - ) + 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.""" -class CheckQueryJobResponseDict(TypedDict, total=False): - """Response from LRO.""" + max_wait_time: Optional[int] + """The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - output_gcs_uri: Optional[str] - """The GCS URI of the output file.""" +CreatePromptVersionConfigOrDict = Union[ + CreatePromptVersionConfig, CreatePromptVersionConfigDict +] -CheckQueryJobResponseOrDict = Union[CheckQueryJobResponse, CheckQueryJobResponseDict] +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 AssembleDataset(_common.BaseModel): - """Represents the assembled dataset.""" - bigquery_destination: Optional[str] = Field( - default=None, - description="""The BigQuery destination of the assembled dataset.""", - ) +class GetPromptConfigDict(TypedDict, total=False): + """Config for getting a prompt.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -class AssembleDatasetDict(TypedDict, total=False): - """Represents the assembled dataset.""" - bigquery_destination: Optional[str] - """The BigQuery destination of the assembled dataset.""" +GetPromptConfigOrDict = Union[GetPromptConfig, GetPromptConfigDict] -AssembleDatasetOrDict = Union[AssembleDataset, AssembleDatasetDict] +class PromptRef(_common.BaseModel): + """Reference to a prompt.""" + prompt_id: Optional[str] = Field(default=None, description="""""") + model: Optional[str] = Field(default=None, description="""""") -class BatchPredictionResourceUsageAssessmentResult(_common.BaseModel): - """Result of batch prediction resource usage assessment.""" - token_count: Optional[int] = Field( - default=None, description="""Number of tokens in the dataset.""" - ) - audio_token_count: Optional[int] = Field( - default=None, description="""Number of audio tokens in the dataset.""" - ) +class PromptRefDict(TypedDict, total=False): + """Reference to a prompt.""" + prompt_id: Optional[str] + """""" -class BatchPredictionResourceUsageAssessmentResultDict(TypedDict, total=False): - """Result of batch prediction resource usage assessment.""" + model: Optional[str] + """""" - token_count: Optional[int] - """Number of tokens in the dataset.""" - audio_token_count: Optional[int] - """Number of audio tokens in the dataset.""" +PromptRefOrDict = Union[PromptRef, PromptRefDict] -BatchPredictionResourceUsageAssessmentResultOrDict = Union[ - BatchPredictionResourceUsageAssessmentResult, - BatchPredictionResourceUsageAssessmentResultDict, -] +class PromptVersionRef(_common.BaseModel): + """Reference to a prompt version.""" + prompt_id: Optional[str] = Field(default=None, description="""""") + version_id: Optional[str] = Field(default=None, description="""""") + model: Optional[str] = Field(default=None, description="""""") -class BatchPredictionValidationAssessmentResult(_common.BaseModel): - """Result of batch prediction validation assessment.""" - errors: Optional[list[str]] = Field( - default=None, description="""The list of errors found in the dataset.""" - ) +class PromptVersionRefDict(TypedDict, total=False): + """Reference to a prompt version.""" + prompt_id: Optional[str] + """""" -class BatchPredictionValidationAssessmentResultDict(TypedDict, total=False): - """Result of batch prediction validation assessment.""" + version_id: Optional[str] + """""" - errors: Optional[list[str]] - """The list of errors found in the dataset.""" + model: Optional[str] + """""" -BatchPredictionValidationAssessmentResultOrDict = Union[ - BatchPredictionValidationAssessmentResult, - BatchPredictionValidationAssessmentResultDict, -] +PromptVersionRefOrDict = Union[PromptVersionRef, PromptVersionRefDict] -class TuningResourceUsageAssessmentResult(_common.BaseModel): - """Result of tuning resource usage assessment.""" +class OptimizeJobConfig(_common.BaseModel): + """VAPO Prompt Optimizer Config.""" - token_count: Optional[int] = Field( - default=None, description="""The number of tokens in the dataset.""" + config_path: Optional[str] = Field( + default=None, + description="""The gcs path to the config file, e.g. gs://bucket/config.json.""", ) - billable_character_count: Optional[int] = Field( + service_account: Optional[str] = Field( default=None, - description="""The number of billable characters in the dataset.""", + description="""The service account to use for the custom job. Cannot be provided at the same time as service_account_project_number.""", + ) + service_account_project_number: Optional[Union[int, str]] = Field( + default=None, + description="""The project number used to construct the default service account:{service_account_project_number}-compute@developer.gserviceaccount.comCannot be provided at the same time as "service_account".""", + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to wait for the job tocomplete. Ignored for async jobs.""", + ) + optimizer_job_display_name: Optional[str] = Field( + default=None, + description="""The display name of the optimization job. If not provided, a display name in the format of "vapo-optimizer-{timestamp}" will be used.""", ) -class TuningResourceUsageAssessmentResultDict(TypedDict, total=False): - """Result of tuning resource usage assessment.""" +class OptimizeJobConfigDict(TypedDict, total=False): + """VAPO Prompt Optimizer Config.""" - token_count: Optional[int] - """The number of tokens in the dataset.""" + config_path: Optional[str] + """The gcs path to the config file, e.g. gs://bucket/config.json.""" - billable_character_count: Optional[int] - """The number of billable characters in the dataset.""" + 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".""" -TuningResourceUsageAssessmentResultOrDict = Union[ - TuningResourceUsageAssessmentResult, TuningResourceUsageAssessmentResultDict -] + 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.""" -class TuningValidationAssessmentResult(_common.BaseModel): - """The result of a tuning validation assessment.""" - errors: Optional[list[str]] = Field( - default=None, description="""The list of errors found in the dataset.""" +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 TuningValidationAssessmentResultDict(TypedDict, total=False): - """The result of a tuning validation assessment.""" +class ListDeployableModelsConfigDict(TypedDict, total=False): + """Config for listing deployable models.""" - errors: Optional[list[str]] - """The list of errors found in the dataset.""" + include_hugging_face_models: Optional[bool] + """Whether to list Hugging Face models.""" + model_filter: Optional[str] + """Optional. A string to filter the models by.""" -TuningValidationAssessmentResultOrDict = Union[ - TuningValidationAssessmentResult, TuningValidationAssessmentResultDict -] +ListDeployableModelsConfigOrDict = Union[ + ListDeployableModelsConfig, ListDeployableModelsConfigDict +] -class Prompt(_common.BaseModel): - """Represents a prompt.""" - prompt_data: Optional["PromptData"] = Field(default=None, description="""""") - _dataset: Optional["Dataset"] = PrivateAttr(default=None) - _dataset_version: Optional["DatasetVersion"] = PrivateAttr(default=None) +class ListModelGardenModelsConfig(_common.BaseModel): + """Config for listing Model Garden models.""" - @property - def dataset(self) -> "Dataset": - return self._dataset # type: ignore[return-value] + 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.""" + ) - @property - def dataset_version(self) -> "DatasetVersion": - return self._dataset_version # type: ignore[return-value] - @property - def prompt_id(self) -> Optional[str]: - """Returns the ID associated with the prompt resource.""" - if self._dataset and self._dataset.name: - return self._dataset.name.split("/")[-1] - elif not self._dataset and ( - self._dataset_version and self._dataset_version.name - ): - return self._dataset_version.name.split("datasets/")[1].split("/")[0] - return None +class ListModelGardenModelsConfigDict(TypedDict, total=False): + """Config for listing Model Garden models.""" - @property - def version_id(self) -> Optional[str]: - """Returns the ID associated with the prompt version resource.""" - if self._dataset_version and self._dataset_version.name: - return self._dataset_version.name.split("/")[-1] - return None + include_hugging_face_models: Optional[bool] + """Whether to list Hugging Face models.""" - def assemble_contents(self) -> list[genai_types.Content]: - """Transforms a Prompt object into a list with a single genai_types.Content object. + model_filter: Optional[str] + """Optional. A string to filter the models by.""" - This method replaces the variables in the prompt template with the values provided in prompt.prompt_data.variables. - If no variables are provided, prompt.prompt_data.contents is returned as is. Only single-turn prompts are supported. - This can be used to call generate_content() in the Gen AI SDK. +ListModelGardenModelsConfigOrDict = Union[ + ListModelGardenModelsConfig, ListModelGardenModelsConfigDict +] - Example usage: - my_prompt = types.Prompt( - prompt_data=types.PromptData( - model="gemini-2.0-flash-001", - contents=[ - genai_types.Content( - parts=[ - genai_types.Part(text="Hello {name}!"), - ], - ), - ], - variables=[ - { - "name": genai_types.Part(text="Alice"), - }, - ], - ), - ) +class ListPublisherModelDeployOptionsConfig(_common.BaseModel): + """Config for listing the deploy options of a publisher model.""" - from google import genai + machine_type_filter: Optional[Union[str, list[str]]] = Field( + default=None, + description="""Optional. Case-insensitive substring filter on the machine type. + Accepts a single keyword (e.g. ``'g2'``) or a list of keywords (e.g. + ``['n1', 'g2']``); an option matches if it contains any of them.""", + ) + accelerator_type_filter: Optional[Union[str, list[str]]] = Field( + default=None, + description="""Optional. Case-insensitive substring filter on the accelerator + type. Accepts a single keyword (e.g. ``'L4'``) or a list of keywords + (e.g. ``['T4', 'L4']``); an option matches if it contains any of them.""", + ) + serving_container_image_uri_filter: Optional[Union[str, list[str]]] = Field( + default=None, + description="""Optional. Case-insensitive substring filter on the serving + container image URI. Accepts a single keyword (e.g. ``'vllm'``) or a list + of keywords (e.g. ``['vllm', 'tgi']``); an option matches if it contains + any of them.""", + ) + concise: Optional[bool] = Field( + default=None, + description="""Optional. If True, returns a human-readable string describing the + deploy options (container and machine specs) instead of a list of + ``DeployOption`` objects.""", + ) - genai_client = genai.Client(vertexai=True, project="my-project", location="us-central1") - genai_client.models.generate_content( - model=my_prompt.prompt_data.model, - contents=my_prompt.assemble_contents(), - ) - Returns: - A list with a single Content object that can be used to call - generate_content(). - """ - if not self.prompt_data or not self.prompt_data.contents: - return [] +class ListPublisherModelDeployOptionsConfigDict(TypedDict, total=False): + """Config for listing the deploy options of a publisher model.""" - if not self.prompt_data.variables: - return self.prompt_data.contents + machine_type_filter: Optional[Union[str, list[str]]] + """Optional. Case-insensitive substring filter on the machine type. + Accepts a single keyword (e.g. ``'g2'``) or a list of keywords (e.g. + ``['n1', 'g2']``); an option matches if it contains any of them.""" - if len(self.prompt_data.contents) > 1: - raise ValueError( - "Multiple contents are not supported. Use assemble_contents() for a prompt with a single Content item." - ) + 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.""" - parts_to_process = self.prompt_data.contents[0].parts - if parts_to_process is None: - return [] - if not isinstance(parts_to_process, list): - parts_to_process = [parts_to_process] + 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.""" - has_placeholders = False - variable_regex = r"\{.*?\}" - for item in parts_to_process: - part = ( - item - if isinstance(item, genai_types.Part) - else genai_types.Part(text=str(item)) - ) - if part.text and re.search(variable_regex, part.text): - has_placeholders = True - break + 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.""" - if not has_placeholders: - return [genai_types.Content(parts=parts_to_process)] - all_rendered_parts: list[genai_types.Part] = [] +ListPublisherModelDeployOptionsConfigOrDict = Union[ + ListPublisherModelDeployOptionsConfig, ListPublisherModelDeployOptionsConfigDict +] - for var_dict in self.prompt_data.variables: - for template_item in parts_to_process: - template_part = ( - template_item - if isinstance(template_item, genai_types.Part) - else genai_types.Part(text=str(template_item)) - ) - if template_part.text: - rendered_text = template_part.text - for key, value in var_dict.items(): - placeholder = f"{{{key}}}" - replacement_text = None +class ListCustomModelDeployOptionsConfig(_common.BaseModel): + """Config for listing custom model deploy options.""" - if isinstance(value, str): - replacement_text = value - elif isinstance(value, genai_types.Part): - if value.text: - replacement_text = value.text - else: - all_rendered_parts.append(value) - if ( - replacement_text is not None - and placeholder in rendered_text - ): - rendered_text = rendered_text.replace( - placeholder, replacement_text - ) - all_rendered_parts.append(genai_types.Part(text=rendered_text)) - else: - all_rendered_parts.append(template_part) - return [genai_types.Content(parts=all_rendered_parts, role="user")] + 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. -PromptData = SchemaPromptSpecPromptMessage -PromptDataDict = SchemaPromptSpecPromptMessageDict -PromptDataOrDict = Union[PromptData, PromptDataDict] + 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). + """, + ) -ParsedResponseUnion = Union[ - prompts_types.ParsedResponse, prompts_types.ParsedResponseFewShot -] -ParsedResponseUnionDict = Union[ - prompts_types.ParsedResponseDict, prompts_types.ParsedResponseFewShotDict -] +class ListCustomModelDeployOptionsConfigDict(TypedDict, total=False): + """Config for listing custom model deploy options.""" -class PromptDict(TypedDict, total=False): - """Represents a prompt.""" + filter_by_user_quota: Optional[bool] + """Whether to filter recommendations to regions with user quota. - prompt_data: Optional["PromptDataDict"] - """""" + 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] + """Whether to check per-region machine availability. -PromptOrDict = Union[Prompt, PromptDict] + When True (the default), the API returns per-region recommendations + that include the machine spec, region and user quota state. When + False, the API returns a flat list of specs without per-region or + quota information (and ``filter_by_user_quota`` has no effect). + """ -class CreatePromptConfig(_common.BaseModel): - """Config for creating a prompt.""" +ListCustomModelDeployOptionsConfigOrDict = Union[ + ListCustomModelDeployOptionsConfig, ListCustomModelDeployOptionsConfigDict +] - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - prompt_display_name: Optional[str] = Field( - default=None, - description="""The display name for the prompt. If not set, a default name with a timestamp will be used.""", - ) - timeout: Optional[int] = Field( - default=90, - description="""The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""", + +class ExportOpenModelConfig(_common.BaseModel): + """Config for export_open_model.""" + + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to block on the export long-running operation. When + ``True`` (default), returns the destination URI on completion. When + ``False``, returns the ``ExportModelOperation`` for the caller to + poll.""", ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + poll_interval_seconds: Optional[float] = Field( default=None, - description="""Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""", + description="""Seconds between LRO polls when ``wait_for_completion=True``. + Defaults to 30. Ignored when ``wait_for_completion=False``.""", ) - version_display_name: Optional[str] = Field( + timeout_seconds: Optional[float] = Field( default=None, - description="""The display name for the prompt version. If not set, a default name with a timestamp will be used.""", - ) - max_wait_time: Optional[int] = Field( - default=60, - description="""The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""", + description="""Total wall-clock seconds to wait for the export to complete + when ``wait_for_completion=True``. Defaults to 2 hours to + accommodate large model weights. Ignored when + ``wait_for_completion=False``.""", ) -class CreatePromptConfigDict(TypedDict, total=False): - """Config for creating a prompt.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - prompt_display_name: Optional[str] - """The display name for the prompt. If not set, a default name with a timestamp will be used.""" +class ExportOpenModelConfigDict(TypedDict, total=False): + """Config for export_open_model.""" - timeout: Optional[int] - """The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""" + 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.""" - 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.""" + poll_interval_seconds: Optional[float] + """Seconds between LRO polls when ``wait_for_completion=True``. + Defaults to 30. Ignored when ``wait_for_completion=False``.""" - 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_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``.""" - max_wait_time: Optional[int] - """The maximum interval between requests in seconds. If not set, the default interval is 60 seconds.""" +ExportOpenModelConfigOrDict = Union[ExportOpenModelConfig, ExportOpenModelConfigDict] -CreatePromptConfigOrDict = Union[CreatePromptConfig, CreatePromptConfigDict] +class DeployPublisherModelConfig(_common.BaseModel): + """Config for ``deploy_publisher_model``. -class CreatePromptVersionConfig(_common.BaseModel): - """Config for creating a prompt version.""" + 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. + """ - 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 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="""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 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``.""", + ) + accept_eula: Optional[bool] = Field( + default=None, + description="""Whether to accept the model's End User License Agreement.""", + ) + 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.""", + ) + machine_type: Optional[str] = Field( + default=None, + 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'``).""" ) - 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.""", + accelerator_count: Optional[int] = Field( + default=None, description="""Number of accelerators per replica.""" ) - 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.""", + 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).""", ) - prompt_display_name: Optional[str] = Field( + fast_tryout_enabled: Optional[bool] = Field( default=None, - description="""The display name for the prompt. If not set, a default name with a timestamp will be used.""", + description="""Use the fast-tryout deployment path (experimentation only, not + production). Only supported for select models and machine types.""", ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + 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="""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="""Custom serving container image URI overriding the model's + default container.""", ) - 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.""", + 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 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 DeployPublisherModelConfigDict(TypedDict, total=False): + """Config for ``deploy_publisher_model``. + 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. + """ -class GetPromptConfigDict(TypedDict, total=False): - """Config for getting a prompt.""" + 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.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + 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 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``.""" -GetPromptConfigOrDict = Union[GetPromptConfig, GetPromptConfigDict] + 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.""" -class PromptRef(_common.BaseModel): - """Reference to a prompt.""" + machine_type: Optional[str] + """Machine type (e.g. ``'g2-standard-48'``). Leave unset for + automatic resources.""" - prompt_id: Optional[str] = Field(default=None, description="""""") - model: Optional[str] = Field(default=None, description="""""") + min_replica_count: Optional[int] + """Minimum number of replicas.""" + max_replica_count: Optional[int] + """Maximum number of replicas.""" -class PromptRefDict(TypedDict, total=False): - """Reference to a prompt.""" + accelerator_type: Optional[str] + """Accelerator type (e.g. ``'NVIDIA_L4'``).""" - prompt_id: Optional[str] - """""" + accelerator_count: Optional[int] + """Number of accelerators per replica.""" - model: Optional[str] - """""" + spot: Optional[bool] + """Schedule on Spot VMs.""" + dedicated_endpoint_disabled: Optional[bool] + """Set True to serve predictions via the shared endpoint DNS + instead of the dedicated endpoint DNS (default).""" -PromptRefOrDict = Union[PromptRef, PromptRefDict] + fast_tryout_enabled: Optional[bool] + """Use the fast-tryout deployment path (experimentation only, not + production). Only supported for select models and machine types.""" + endpoint_display_name: Optional[str] + """Display name for the endpoint.""" -class PromptVersionRef(_common.BaseModel): - """Reference to a prompt version.""" + model_display_name: Optional[str] + """Display name for the deployed model.""" - prompt_id: Optional[str] = Field(default=None, description="""""") - version_id: Optional[str] = Field(default=None, description="""""") - model: Optional[str] = Field(default=None, description="""""") + serving_container_image_uri: Optional[str] + """Custom serving container image URI overriding the model's + default container.""" + container_command: Optional[list[str]] + """Serving container ENTRYPOINT override.""" -class PromptVersionRefDict(TypedDict, total=False): - """Reference to a prompt version.""" + container_args: Optional[list[str]] + """Serving container CMD override.""" - prompt_id: Optional[str] - """""" + container_variables: Optional[dict[str, str]] + """Environment variables for the serving container.""" - version_id: Optional[str] - """""" + enable_private_service_connect: Optional[bool] + """Enable Private Service Connect for the endpoint.""" - model: Optional[str] - """""" + 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.""" -PromptVersionRefOrDict = Union[PromptVersionRef, PromptVersionRefDict] +DeployPublisherModelConfigOrDict = Union[ + DeployPublisherModelConfig, DeployPublisherModelConfigDict +] -class OptimizeJobConfig(_common.BaseModel): - """VAPO Prompt Optimizer Config.""" +class DeployOption(_common.BaseModel): + """A verified deploy option for a model.""" - config_path: Optional[str] = Field( - default=None, - description="""The gcs path to the config file, e.g. gs://bucket/config.json.""", + option_name: Optional[str] = Field( + default=None, description="""The name of the deploy task.""" ) - 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.""", + serving_container_image_uri: Optional[str] = Field( + default=None, description="""The URI of the serving container.""" ) - 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".""", + machine_type: Optional[str] = Field( + default=None, description="""The machine type.""" ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Whether to wait for the job tocomplete. Ignored for async jobs.""", + accelerator_type: Optional[str] = Field( + default=None, description="""The accelerator type.""" ) - 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.""", + accelerator_count: Optional[int] = Field( + default=None, description="""The number of accelerators.""" ) -class OptimizeJobConfigDict(TypedDict, total=False): - """VAPO Prompt Optimizer Config.""" +class DeployOptionDict(TypedDict, total=False): + """A verified deploy option for a model.""" - config_path: Optional[str] - """The gcs path to the config file, e.g. gs://bucket/config.json.""" + option_name: Optional[str] + """The name of the deploy task.""" - 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.""" + serving_container_image_uri: Optional[str] + """The URI of the serving container.""" - 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".""" + machine_type: Optional[str] + """The machine type.""" - wait_for_completion: Optional[bool] - """Whether to wait for the job tocomplete. Ignored for async jobs.""" + accelerator_type: Optional[str] + """The accelerator type.""" - 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.""" + accelerator_count: Optional[int] + """The number of accelerators.""" -OptimizeJobConfigOrDict = Union[OptimizeJobConfig, OptimizeJobConfigDict] +DeployOptionOrDict = Union[DeployOption, DeployOptionDict] -class AgentEngineRuntimeRevision(_common.BaseModel): - """An agent engine runtime revision instance.""" +class Runtime(_common.BaseModel): + """An agent runtime instance.""" api_client: Optional[Any] = Field( default=None, description="""The underlying API client.""" @@ -31900,9 +30541,9 @@ class AgentEngineRuntimeRevision(_common.BaseModel): default=None, description="""The underlying API client for asynchronous operations.""", ) - api_resource: Optional[ReasoningEngineRuntimeRevision] = Field( + api_resource: Optional[ReasoningEngine] = Field( default=None, - description="""The underlying API resource (i.e. ReasoningEngineRuntimeRevision).""", + description="""The underlying API resource (i.e. ReasoningEngine).""", ) # Allows dynamic binding of methods based on the registered operations. @@ -31910,14 +30551,14 @@ class AgentEngineRuntimeRevision(_common.BaseModel): def __repr__(self) -> str: return ( - f"AgentEngineRuntimeRevision(api_resource.name='{self.api_resource.name}')" + f"Runtime(api_resource.name='{self.api_resource.name}')" if self.api_resource is not None - else "AgentEngineRuntimeRevision(api_resource.name=None)" + else "Runtime(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, ReasoningEngineRuntimeRevision): + 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.") @@ -31925,467 +30566,642 @@ def operation_schemas(self) -> Optional[list[Dict[str, Any]]]: def delete( self, - config: Optional[DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, + force: bool = False, + config: Optional[DeleteRuntimeConfigOrDict] = None, ) -> None: - """Deletes the agent engine runtime revision. + """Deletes the agent engine. Args: - config (DeleteAgentEngineRuntimeRevisionConfig): - Optional. Additional configurations for deleting the Agent Engine Runtime Revision. + 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, ReasoningEngineRuntimeRevision): + if not isinstance(self.api_resource, ReasoningEngine): raise ValueError("api_resource is not initialized.") - self.api_client.delete(name=self.api_resource.name, config=config) # type: ignore[union-attr] - - -class AgentEngineRuntimeRevisionDict(TypedDict, total=False): - """An agent engine runtime revision instance.""" - - api_client: Optional[Any] - """The underlying API client.""" - - api_async_client: Optional[Any] - """The underlying API client for asynchronous operations.""" + self.api_client.delete(name=self.api_resource.name, force=force, config=config) # type: ignore[union-attr] - api_resource: Optional[ReasoningEngineRuntimeRevisionDict] - """The underlying API resource (i.e. ReasoningEngineRuntimeRevision).""" +RubricContentProperty = evals_types.RubricContentProperty +RubricContentPropertyDict = evals_types.RubricContentPropertyDict +RubricContentPropertyDictOrDict = evals_types.RubricContentPropertyOrDict -AgentEngineRuntimeRevisionOrDict = Union[ - AgentEngineRuntimeRevision, AgentEngineRuntimeRevisionDict -] +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 ListDeployableModelsConfig(_common.BaseModel): - """Config for listing deployable models.""" +RubricVerdict = evals_types.RubricVerdict +RubricVerdictDict = evals_types.RubricVerdictDict +RubricVerdictDictOrDict = evals_types.RubricVerdictOrDict - 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.""" - ) +CandidateResult = evals_types.CandidateResult +CandidateResultDict = evals_types.CandidateResultDict +CandidateResultDictOrDict = evals_types.CandidateResultOrDict +Event = evals_types.Event +EventDict = evals_types.EventDict +EventDictOrDict = evals_types.EventOrDict -class ListDeployableModelsConfigDict(TypedDict, total=False): - """Config for listing deployable models.""" +Message = evals_types.Message +MessageDict = evals_types.MessageDict +MessageDictOrDict = evals_types.MessageOrDict - include_hugging_face_models: Optional[bool] - """Whether to list Hugging Face models.""" +Importance = evals_types.Importance - model_filter: Optional[str] - """Optional. A string to filter the models by.""" +class RuntimeDict(TypedDict, total=False): + """An agent runtime instance.""" -ListDeployableModelsConfigOrDict = Union[ - ListDeployableModelsConfig, ListDeployableModelsConfigDict -] + api_client: Optional[Any] + """The underlying API client.""" + api_async_client: Optional[Any] + """The underlying API client for asynchronous operations.""" -class ListModelGardenModelsConfig(_common.BaseModel): - """Config for listing Model Garden models.""" + api_resource: Optional[ReasoningEngineDict] + """The underlying API resource (i.e. ReasoningEngine).""" - 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.""" - ) +RuntimeOrDict = Union[Runtime, RuntimeDict] -class ListModelGardenModelsConfigDict(TypedDict, total=False): - """Config for listing Model Garden models.""" - include_hugging_face_models: Optional[bool] - """Whether to list Hugging Face models.""" +class AgentRuntimeConfig(_common.BaseModel): + """Config for agent runtime methods.""" - model_filter: Optional[str] - """Optional. A string to filter the models by.""" + 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. -ListModelGardenModelsConfigOrDict = Union[ - ListModelGardenModelsConfig, ListModelGardenModelsConfigDict -] + 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. + 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. -class ListPublisherModelDeployOptionsConfig(_common.BaseModel): - """Config for listing the deploy options of a publisher model.""" + 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. - machine_type_filter: Optional[Union[str, list[str]]] = Field( + 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="""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.""", + description="""The context spec to be used for the Agent Runtime.""", ) - accelerator_type_filter: Optional[Union[str, list[str]]] = Field( + psc_interface_config: Optional[PscInterfaceConfig] = 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.""", + description="""The PSC interface config for PSC-I to be used for the Agent Runtime.""", ) - serving_container_image_uri_filter: Optional[Union[str, list[str]]] = Field( + min_instances: Optional[int] = 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 minimum number of instances to run for the Agent Runtime. + Defaults to 1. Range: [0, 10]. + """, ) - concise: Optional[bool] = Field( + max_instances: Optional[int] = 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 maximum number of instances to run for the Agent Runtime. + Defaults to 100. Range: [1, 1000]. + If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. + """, + ) + resource_limits: Optional[dict[str, str]] = Field( + default=None, + description="""The resource limits to be applied to the Agent Runtime. + Required keys: 'cpu' and 'memory'. + Supported values for 'cpu': '1', '2', '4', '6', '8'. + Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. + """, + ) + container_concurrency: Optional[int] = Field( + default=None, + description="""The container concurrency to be used for the Agent Runtime. + Recommended value: 2 * cpu + 1. Defaults to 9. + """, + ) + keep_alive_probe: Optional[KeepAliveProbe] = Field( + default=None, + description="""Optional. Specifies the configuration for keep-alive probe. + Contains configuration on a specified endpoint that a deployment host + should use to keep the container alive based on the probe settings.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""The encryption spec to be used for the Agent Runtime.""", + ) + labels: Optional[dict[str, str]] = Field( + default=None, description="""The labels to be used for the Agent Runtime.""" + ) + agent_server_mode: Optional[AgentServerMode] = Field( + default=None, description="""The agent server mode to use for deployment.""" + ) + class_methods: Optional[list[dict[str, Any]]] = Field( + default=None, + description="""The class methods to be used for the Agent Runtime. + If specified, they'll override the class methods that are autogenerated by + default. By default, methods are generated by inspecting the agent object + and generating a corresponding method for each method defined on the + agent class. + """, + ) + source_packages: Optional[list[str]] = Field( + default=None, + description="""The user-provided paths to the source packages (if any). + If specified, the files in the source packages will be packed into a + a tarball file, uploaded to Agent Runtime's API, and deployed to the + Agent Runtime. + The following fields will be ignored: + - agent + - extra_packages + - staging_bucket + - requirements + The following fields will be used to install and use the agent from the + source packages: + - entrypoint_module (required) + - entrypoint_object (required) + - requirements_file (optional) + - class_methods (required) + """, + ) + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig + ] = Field( + default=None, + description="""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 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 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 runtime 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 Runtime. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""", + ) + python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] = Field( + default=None, + description="""The Python version to be used for the Agent Runtime. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". + """, + ) + build_options: Optional[dict[str, list[str]]] = Field( + default=None, + description="""The build options for the Agent Runtime. + The following keys are supported: + - installation_scripts: + Optional. The paths to the installation scripts to be + executed in the Docker image. + The scripts must be located in the `installation_scripts` + subdirectory and the path must be added to `extra_packages`. + """, + ) + image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpec] = Field( + default=None, description="""The image spec for the Agent Runtime.""" + ) + agent_config_source: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSource + ] = Field( + default=None, description="""The agent config source for the Agent Runtime.""" + ) + traffic_config: Optional[ReasoningEngineTrafficConfig] = Field( + default=None, description="""The traffic config for the Agent Runtime.""" + ) + container_spec: Optional[ReasoningEngineSpecContainerSpec] = Field( + default=None, description="""The container spec for the Agent Runtime.""" + ) + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + ] = Field( + default=None, + description="""Agent Gateway configuration for a Agent Runtime deployment.""", + ) + 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 ListPublisherModelDeployOptionsConfigDict(TypedDict, total=False): - """Config for listing the deploy options of a publisher model.""" - - machine_type_filter: Optional[Union[str, list[str]]] - """Optional. Case-insensitive substring filter on the machine type. - Accepts a single keyword (e.g. ``'g2'``) or a list of keywords (e.g. - ``['n1', 'g2']``); an option matches if it contains any of them.""" - - accelerator_type_filter: Optional[Union[str, list[str]]] - """Optional. Case-insensitive substring filter on the accelerator - type. Accepts a single keyword (e.g. ``'L4'``) or a list of keywords - (e.g. ``['T4', 'L4']``); an option matches if it contains any of them.""" - - serving_container_image_uri_filter: Optional[Union[str, list[str]]] - """Optional. Case-insensitive substring filter on the serving - container image URI. Accepts a single keyword (e.g. ``'vllm'``) or a list - of keywords (e.g. ``['vllm', 'tgi']``); an option matches if it contains - any of them.""" - - concise: Optional[bool] - """Optional. If True, returns a human-readable string describing the - deploy options (container and machine specs) instead of a list of - ``DeployOption`` objects.""" - +class AgentRuntimeConfigDict(TypedDict, total=False): + """Config for agent runtime methods.""" -ListPublisherModelDeployOptionsConfigOrDict = Union[ - ListPublisherModelDeployOptionsConfig, ListPublisherModelDeployOptionsConfigDict -] + 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. -class ListCustomModelDeployOptionsConfig(_common.BaseModel): - """Config for listing custom model deploy options.""" + It must be a valid GCS bucket name, e.g. "gs://bucket-name". It is + required if `agent_engine` is specified.""" - filter_by_user_quota: Optional[bool] = Field( - default=True, - description="""Whether to filter recommendations to regions with user quota. + requirements: Optional[Any] + """The set of PyPI dependencies needed. - 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. + 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.""" - 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). - """, - ) + 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.""" -class ListCustomModelDeployOptionsConfigDict(TypedDict, total=False): - """Config for listing custom model deploy options.""" + description: Optional[str] + """The description of the Agent Runtime.""" - filter_by_user_quota: Optional[bool] - """Whether to filter recommendations to regions with user quota. + gcs_dir_name: Optional[str] + """The GCS bucket directory under `staging_bucket` to use for staging + the artifacts needed.""" - 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. - """ + extra_packages: Optional[list[str]] + """The set of extra user-provided packages (if any).""" - check_machine_availability: Optional[bool] - """Whether to check per-region machine availability. + env_vars: Optional[Any] + """The environment variables to be set when running the Agent Runtime. - 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). - """ + 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. -ListCustomModelDeployOptionsConfigOrDict = Union[ - ListCustomModelDeployOptionsConfig, ListCustomModelDeployOptionsConfigDict -] + 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.""" -class ExportOpenModelConfig(_common.BaseModel): - """Config for export_open_model.""" + context_spec: Optional[ReasoningEngineContextSpecDict] + """The context spec to be used for the Agent Runtime.""" - 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``.""", - ) + 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]. + """ -class ExportOpenModelConfigDict(TypedDict, total=False): - """Config for export_open_model.""" + 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]. + """ - 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.""" + 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'. + """ - poll_interval_seconds: Optional[float] - """Seconds between LRO polls when ``wait_for_completion=True``. - Defaults to 30. Ignored when ``wait_for_completion=False``.""" + container_concurrency: Optional[int] + """The container concurrency to be used for the Agent Runtime. + Recommended value: 2 * cpu + 1. Defaults to 9. + """ - 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``.""" + 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.""" -ExportOpenModelConfigOrDict = Union[ExportOpenModelConfig, ExportOpenModelConfigDict] + 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 DeployPublisherModelConfig(_common.BaseModel): - """Config for ``deploy_publisher_model``. + 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. + """ - 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. - """ + 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) + """ - 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="""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 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``.""", - ) - accept_eula: Optional[bool] = Field( - default=None, - description="""Whether to accept the model's End User License Agreement.""", - ) - 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.""", - ) - machine_type: Optional[str] = Field( - default=None, - 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.""" + 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 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.""" + + requirements_file: Optional[str] + """The user-provided path to the requirements file (if any). + This field is only used when source_packages is specified. + If not specified, agent runtime will find and use the `requirements.txt` in + the source package. + """ + + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + ] + """The agent framework to be used for the Agent Runtime. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""" + + python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] + """The Python version to be used for the Agent Runtime. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". + """ + + 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`. + """ + + image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict] + """The image spec for the Agent Runtime.""" + + agent_config_source: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict + ] + """The agent config source for the Agent Runtime.""" + + traffic_config: Optional[ReasoningEngineTrafficConfigDict] + """The traffic config for the Agent Runtime.""" + + container_spec: Optional[ReasoningEngineSpecContainerSpecDict] + """The container spec for the Agent Runtime.""" + + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] + """Agent Gateway configuration for a Agent Runtime deployment.""" + + 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).""" + + +AgentRuntimeConfigOrDict = Union[AgentRuntimeConfig, AgentRuntimeConfigDict] + + +class RunQueryJobRuntimeConfig(_common.BaseModel): + """Config for checking a query job on an agent runtime.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - enable_private_service_connect: Optional[bool] = Field( - default=None, description="""Enable Private Service Connect for the endpoint.""" + query: Optional[str] = Field( + default=None, description="""The query to send to the agent runtime.""" ) - psc_project_allow_list: Optional[list[str]] = Field( + output_gcs_uri: Optional[str] = Field( default=None, - description="""Projects allowed to access the endpoint over Private Service - Connect. Only honored when ``enable_private_service_connect`` is True.""", + 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 DeployPublisherModelConfigDict(TypedDict, total=False): - """Config for ``deploy_publisher_model``. +class RunQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for checking a query job on an 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. - """ + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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.""" + query: Optional[str] + """The query to send to the agent runtime.""" - poll_interval_seconds: Optional[float] - """Seconds between LRO polls when ``wait_for_completion=True``. - Defaults to 30. Ignored when ``wait_for_completion=False``.""" + output_gcs_uri: Optional[str] + """The GCS URI to use for the output. + If it is a file, the system use this file to store the response. + If it represents a directory, the system automatically generate a file + for the response. + In both cases, the input query will be stored in the same directory under + the same file name prefix as the output file.""" - timeout_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``.""" - accept_eula: Optional[bool] - """Whether to accept the model's End User License Agreement.""" +RunQueryJobRuntimeConfigOrDict = Union[ + RunQueryJobRuntimeConfig, RunQueryJobRuntimeConfigDict +] - 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.""" +class RunQueryJobResult(_common.BaseModel): + """Result of running a query job.""" - min_replica_count: Optional[int] - """Minimum number of replicas.""" + 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 runtime operation to later check for status.""", + ) + input_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the input file.""" + ) + output_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the output file.""" + ) - max_replica_count: Optional[int] - """Maximum number of replicas.""" - accelerator_type: Optional[str] - """Accelerator type (e.g. ``'NVIDIA_L4'``).""" +class RunQueryJobResultDict(TypedDict, total=False): + """Result of running a query job.""" - accelerator_count: Optional[int] - """Number of accelerators per replica.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - spot: Optional[bool] - """Schedule on Spot VMs.""" + job_name: Optional[str] + """Name of the agent runtime operation to later check for status.""" - dedicated_endpoint_disabled: Optional[bool] - """Set True to serve predictions via the shared endpoint DNS - instead of the dedicated endpoint DNS (default).""" + input_gcs_uri: Optional[str] + """The GCS URI of the input file.""" - fast_tryout_enabled: Optional[bool] - """Use the fast-tryout deployment path (experimentation only, not - production). Only supported for select models and machine types.""" + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" - endpoint_display_name: Optional[str] - """Display name for the endpoint.""" - model_display_name: Optional[str] - """Display name for the deployed model.""" +RunQueryJobResultOrDict = Union[RunQueryJobResult, RunQueryJobResultDict] - serving_container_image_uri: Optional[str] - """Custom serving container image URI overriding the model's - default container.""" - container_command: Optional[list[str]] - """Serving container ENTRYPOINT override.""" +class CheckQueryJobResponse(_common.BaseModel): + """Response from LRO.""" - container_args: Optional[list[str]] - """Serving container CMD override.""" + 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.""" + ) - container_variables: Optional[dict[str, str]] - """Environment variables for the serving container.""" - enable_private_service_connect: Optional[bool] - """Enable Private Service Connect for the endpoint.""" +class CheckQueryJobResponseDict(TypedDict, total=False): + """Response from LRO.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" -DeployPublisherModelConfigOrDict = Union[ - DeployPublisherModelConfig, DeployPublisherModelConfigDict -] +CheckQueryJobResponseOrDict = Union[CheckQueryJobResponse, CheckQueryJobResponseDict] -class DeployOption(_common.BaseModel): - """A verified deploy option for a model.""" - option_name: Optional[str] = Field( - default=None, description="""The name of the deploy task.""" - ) - serving_container_image_uri: Optional[str] = Field( - default=None, description="""The URI of the serving container.""" - ) - machine_type: Optional[str] = Field( - default=None, description="""The machine type.""" +class RuntimeRevision(_common.BaseModel): + """An agent runtime revision instance.""" + + api_client: Optional[Any] = Field( + default=None, description="""The underlying API client.""" ) - accelerator_type: Optional[str] = Field( - default=None, description="""The accelerator type.""" + api_async_client: Optional[Any] = Field( + default=None, + description="""The underlying API client for asynchronous operations.""", ) - accelerator_count: Optional[int] = Field( - default=None, description="""The number of accelerators.""" + api_resource: Optional[ReasoningEngineRuntimeRevision] = Field( + default=None, + description="""The underlying API resource (i.e. ReasoningEngineRuntimeRevision).""", ) + # Allows dynamic binding of methods based on the registered operations. + model_config = ConfigDict(extra="allow") -class DeployOptionDict(TypedDict, total=False): - """A verified deploy option for a model.""" + def __repr__(self) -> str: + return ( + f"RuntimeRevision(api_resource.name='{self.api_resource.name}')" + if self.api_resource is not None + else "RuntimeRevision(api_resource.name=None)" + ) - option_name: Optional[str] - """The name of the deploy task.""" + def operation_schemas(self) -> Optional[list[Dict[str, Any]]]: + """Returns the schemas of all registered operations for the agent.""" + if not isinstance(self.api_resource, ReasoningEngineRuntimeRevision): + raise ValueError("api_resource is not initialized.") + if not self.api_resource.spec: + raise ValueError("api_resource.spec is not initialized.") + return self.api_resource.spec.class_methods - serving_container_image_uri: Optional[str] - """The URI of the serving container.""" + def delete( + self, + config: Optional[DeleteRuntimeRevisionConfigOrDict] = None, + ) -> None: + """Deletes the agent engine runtime revision. - machine_type: Optional[str] - """The machine type.""" + Args: + config (DeleteRuntimeRevisionConfig): + Optional. Additional configurations for deleting the Agent Engine Runtime Revision. + """ + if not isinstance(self.api_resource, ReasoningEngineRuntimeRevision): + raise ValueError("api_resource is not initialized.") + self.api_client.delete(name=self.api_resource.name, config=config) # type: ignore[union-attr] - accelerator_type: Optional[str] - """The accelerator type.""" - accelerator_count: Optional[int] - """The number of accelerators.""" +class RuntimeRevisionDict(TypedDict, total=False): + """An agent runtime revision instance.""" + + api_client: Optional[Any] + """The underlying API client.""" + api_async_client: Optional[Any] + """The underlying API client for asynchronous operations.""" -DeployOptionOrDict = Union[DeployOption, DeployOptionDict] + api_resource: Optional[ReasoningEngineRuntimeRevisionDict] + """The underlying API resource (i.e. ReasoningEngineRuntimeRevision).""" + + +RuntimeRevisionOrDict = Union[RuntimeRevision, RuntimeRevisionDict] diff --git a/agentplatform/agent_engines/__init__.py b/agentplatform/agent_engines/__init__.py deleted file mode 100644 index ac4802a9b2..0000000000 --- a/agentplatform/agent_engines/__init__.py +++ /dev/null @@ -1,428 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -"""Classes and functions for working with agent engines.""" - -from typing import Dict, Iterable, Optional, Sequence, Union - -from google.cloud.aiplatform import base -from google.cloud.aiplatform import initializer -from google.cloud.aiplatform import utils as aip_utils -from google.cloud.aiplatform_v1 import types as aip_types - -# We just want to re-export certain classes -# pylint: disable=g-multiple-import,g-importing-member -from agentplatform.agent_engines._agent_engines import ( - _AgentEngineInterface, - AgentEngine, - Cloneable, - ModuleAgent, - OperationRegistrable, - Queryable, - AsyncQueryable, - StreamQueryable, - AsyncStreamQueryable, -) -from agentplatform.agent_engines.templates.adk import ( - AdkApp, -) -from agentplatform.agent_engines.templates.ag2 import ( - AG2Agent, -) -from agentplatform.agent_engines.templates.langchain import ( - LangchainAgent, -) -from agentplatform.agent_engines.templates.langgraph import ( - LanggraphAgent, -) -from agentplatform.agent_engines.templates.llama_index import ( - LlamaIndexQueryPipelineAgent, -) - - -_LOGGER = base.Logger(__name__) - - -def get(resource_name: str) -> AgentEngine: - """Retrieves an Agent Engine resource. - - Args: - resource_name (str): - Required. A fully-qualified resource name or ID such as - "projects/123/locations/us-central1/reasoningEngines/456" or - "456" when project and location are initialized or passed. - """ - return AgentEngine(resource_name) - - -def create( - agent_engine: Optional[_AgentEngineInterface] = None, - *, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - build_options: Optional[Dict[str, Sequence[str]]] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[aip_types.EncryptionSpec] = None, -) -> AgentEngine: - """Creates a new Agent Engine. - - The Agent Engine will be an instance of the `agent_engine` that - was passed in, running remotely on Vertex AI. - - Sample ``src_dir`` contents (e.g. ``./user_src_dir``): - - .. code-block:: python - - user_src_dir/ - |-- main.py - |-- requirements.txt - |-- user_code/ - | |-- utils.py - | |-- ... - |-- installation_scripts/ - | |-- install_package.sh - | |-- ... - |-- ... - - To build an Agent Engine with the above files, run: - - .. code-block:: python - - remote_agent = agent_engines.create( - agent_engine=local_agent, - requirements=[ - # I.e. the PyPI dependencies listed in requirements.txt - "google-cloud-aiplatform==1.25.0", - "langchain==0.0.242", - ... - ], - extra_packages=[ - "./user_src_dir/main.py", # a single file - "./user_src_dir/user_code", # a directory - ... - ], - build_options={ - "installation": [ - "./user_src_dir/installation_scripts/install_package.sh", - ... - ], - }, - ) - - Args: - agent_engine (AgentEngineInterface): - Required. The Agent Engine to be created. - requirements (Union[str, Sequence[str]]): - Optional. The set of PyPI dependencies needed. It can either be - the path to a single file (requirements.txt), or an ordered list - of strings corresponding to each line of the requirements file. - display_name (str): - Optional. The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any - UTF-8 character. - description (str): - Optional. The description of the Agent Engine. - gcs_dir_name (str): - Optional. The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - extra_packages (Sequence[str]): - Optional. The set of extra user-provided packages (if any). - env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): - Optional. The environment variables to be set when running the - Agent Engine. If it is a list of strings, each string should be - a valid key to `os.environ`. If it is a dictionary, the keys are - the environment variable names, and the values are the - corresponding values. - build_options (Dict[str, Sequence[str]]): - Optional. The build options for the Agent Engine. This includes - options such as installation scripts. - service_account (str): - Optional. The service account to be used for the Agent Engine. If - not specified, the default reasoning engine service agent service - account will be used. - psc_interface_config (PscInterfaceConfig): - Optional. The PSC interface config for the Agent Engine. If not - specified, the default PSC interface config will be used. - min_instances (int): - Optional. The minimum number of instances to run the Agent Engine. - If not specified, the default value will be used. - max_instances (int): - Optional. The maximum number of instances to run the Agent Engine. - If not specified, the default value will be used. - resource_limits (Dict[str, str]): - Optional. The resource limits for the Agent Engine. If not - specified, the default value will be used. - container_concurrency (int): - Optional. The container concurrency for the Agent Engine. If not - specified, the default value will be used. - encryption_spec (EncryptionSpec): - Optional. The encryption spec for the Agent Engine. If not - specified, the default encryption spec will be used. - - Returns: - AgentEngine: The Agent Engine that was created. - - Raises: - ValueError: If the `project` was not set using `agentplatform.init`. - ValueError: If the `location` was not set using `agentplatform.init`. - ValueError: If the `staging_bucket` was not set using agentplatform.init. - ValueError: If the `staging_bucket` does not start with "gs://". - FileNotFoundError: If `extra_packages` includes a file or directory - that does not exist. - IOError: If requirements is a string that corresponds to a - nonexistent file. - """ - return AgentEngine.create( - agent_engine=agent_engine, - requirements=requirements, - display_name=display_name, - description=description, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - env_vars=env_vars, - build_options=build_options, - service_account=service_account, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - encryption_spec=encryption_spec, - ) - - -def list(*, filter: str = "") -> Iterable[AgentEngine]: - """List all instances of Agent Engine matching the filter. - - Example Usage: - - .. code-block:: python - import agentplatform - from agentplatform import agent_engines - - agentplatform.init(project="my_project", location="us-central1") - agent_engines.list(filter='display_name="My Custom Agent"') - - Args: - filter (str): - Optional. An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported. - - Returns: - Iterable[AgentEngine]: An iterable of Agent Engines matching the filter. - """ - api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineClientWithOverride, - ) - for agent in api_client.list_reasoning_engines( - request=aip_types.ListReasoningEnginesRequest( - parent=initializer.global_config.common_location_path(), - filter=filter, - ) - ): - yield AgentEngine(agent.name) - - -def delete( - resource_name: str, - *, - force: bool = False, - **kwargs, -) -> None: - """Delete an Agent Engine resource. - - Args: - resource_name (str): - Required. The name of the Agent Engine to be deleted. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error - when the Agent Engine has undeleted child resources. Defaults to - False. - **kwargs (dict[str, Any]): - Optional. Additional keyword arguments to pass to the - delete_reasoning_engine method. - """ - api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineClientWithOverride, - ) - _LOGGER.info(f"Deleting AgentEngine resource: {resource_name}") - operation_future = api_client.delete_reasoning_engine( - request=aip_types.DeleteReasoningEngineRequest( - name=resource_name, - force=force, - **(kwargs or {}), - ) - ) - _LOGGER.info(f"Delete AgentEngine backing LRO: {operation_future.operation.name}") - operation_future.result() - _LOGGER.info(f"AgentEngine resource deleted: {resource_name}") - - -def update( - resource_name: str, - *, - agent_engine: Optional[Union[Queryable, OperationRegistrable]] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - build_options: Optional[Dict[str, Sequence[str]]] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[aip_types.EncryptionSpec] = None, -) -> "AgentEngine": - """Updates an existing Agent Engine. - - This method updates the configuration of a deployed Agent Engine, identified - by its resource name. Unlike the `create` function which requires an - `agent_engine` object, all arguments in this method are optional. This - method allows you to modify individual aspects of the configuration by - providing any of the optional arguments. - - Args: - resource_name (str): - Required. The name of the Agent Engine to be updated. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - agent_engine (AgentEngineInterface): - Optional. The instance to be used as the updated Agent Engine. If it - is not specified, the existing instance will be used. - requirements (Union[str, Sequence[str]]): - Optional. The set of PyPI dependencies needed. It can either be - the path to a single file (requirements.txt), or an ordered list - of strings corresponding to each line of the requirements file. - If it is not specified, the existing requirements will be used. - If it is set to an empty string or list, the existing - requirements will be removed. - display_name (str): - Optional. The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any - UTF-8 character. - description (str): - Optional. The description of the Agent Engine. - gcs_dir_name (str): - Optional. The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - extra_packages (Sequence[str]): - Optional. The set of extra user-provided packages (if any). If - it is not specified, the existing extra packages will be used. - If it is set to an empty list, the existing extra packages will - be removed. - env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): - Optional. The environment variables to be set when running the - Agent Engine. If it is a list of strings, each string should be - a valid key to `os.environ`. If it is a dictionary, the keys are - the environment variable names, and the values are the - corresponding values. - build_options (Dict[str, Sequence[str]]): - Optional. The build options for the Agent Engine. This includes - options such as installation scripts. - service_account (str): - Optional. The service account to be used for the Agent Engine. If - not specified, the default reasoning engine service agent service - account will be used. - min_instances (int): - Optional. The minimum number of instances to run the Agent Engine. - If not specified, the default value will be used. - max_instances (int): - Optional. The maximum number of instances to run the Agent Engine. - If not specified, the default value will be used. - resource_limits (Dict[str, str]): - Optional. The resource limits for the Agent Engine. If not - specified, the default value will be used. - container_concurrency (int): - Optional. The container concurrency for the Agent Engine. If not - specified, the default value will be used. - encryption_spec (EncryptionSpec): - Optional. The encryption spec for the Agent Engine. If not - specified, the default encryption spec will be used. - - Returns: - AgentEngine: The Agent Engine that was updated. - - Raises: - ValueError: If the `staging_bucket` was not set using agentplatform.init. - ValueError: If the `staging_bucket` does not start with "gs://". - FileNotFoundError: If `extra_packages` includes a file or directory - that does not exist. - ValueError: if none of `display_name`, `description`, - `requirements`, `extra_packages`, `agent_engine`, or `build_options` - were specified. - IOError: If requirements is a string that corresponds to a - nonexistent file. - """ - agent = get(resource_name) - return agent.update( - agent_engine=agent_engine, - requirements=requirements, - display_name=display_name, - description=description, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - env_vars=env_vars, - build_options=build_options, - service_account=service_account, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - encryption_spec=encryption_spec, - ) - - -__all__ = ( - # Resources - "AgentEngine", - # Protocols - "Cloneable", - "OperationRegistrable", - "Queryable", - "AsyncQueryable", - "StreamQueryable", - "AsyncStreamQueryable", - # Methods - "create", - "delete", - "get", - "list", - "update", - # Templates - "AdkApp", - "ModuleAgent", - "LangchainAgent", - "LanggraphAgent", - "AG2Agent", - "LlamaIndexQueryPipelineAgent", -) diff --git a/agentplatform/agent_engines/_agent_engines.py b/agentplatform/agent_engines/_agent_engines.py deleted file mode 100644 index 35333b82a4..0000000000 --- a/agentplatform/agent_engines/_agent_engines.py +++ /dev/null @@ -1,2028 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import abc -import inspect -import io -import json -import logging -import os -import sys -import tarfile -import types -import typing -from typing import ( - Any, - AsyncIterable, - Callable, - Coroutine, - Dict, - Iterable, - List, - Optional, - Protocol, - Sequence, - Tuple, - Union, -) - -from google.api_core import exceptions -from google.cloud import storage -from google.cloud.aiplatform import base -from google.cloud.aiplatform import initializer -from google.cloud.aiplatform import utils as aip_utils -from google.cloud.aiplatform_v1 import types as aip_types -from google.cloud.aiplatform_v1.types import reasoning_engine_service -from agentplatform._genai import _agent_engines_utils -import httpx -import proto - -from google.protobuf import field_mask_pb2 - - -_LOGGER = base.Logger("agentplatform.agent_engines") - -_SUPPORTED_PYTHON_VERSIONS = ("3.10", "3.11", "3.12", "3.13", "3.14") -_DEFAULT_GCS_DIR_NAME = "agent_engine" -_BLOB_FILENAME = "agent_engine.pkl" -_REQUIREMENTS_FILE = "requirements.txt" -_EXTRA_PACKAGES_FILE = "dependencies.tar.gz" -_STANDARD_API_MODE = "" -_ASYNC_API_MODE = "async" -_STREAM_API_MODE = "stream" -_ASYNC_STREAM_API_MODE = "async_stream" -_BIDI_STREAM_API_MODE = "bidi_stream" -_A2A_EXTENSION_MODE = "a2a_extension" -_A2A_AGENT_CARD = "a2a_agent_card" -_MODE_KEY_IN_SCHEMA = "api_mode" -_METHOD_NAME_KEY_IN_SCHEMA = "name" -_DEFAULT_METHOD_NAME = "query" -_DEFAULT_ASYNC_METHOD_NAME = "async_query" -_DEFAULT_STREAM_METHOD_NAME = "stream_query" -_DEFAULT_ASYNC_STREAM_METHOD_NAME = "async_stream_query" -_DEFAULT_METHOD_RETURN_TYPE = "dict[str, Any]" -_DEFAULT_ASYNC_METHOD_RETURN_TYPE = "Coroutine[Any, Any, Any]" -_DEFAULT_STREAM_METHOD_RETURN_TYPE = "Iterable[Any]" -_DEFAULT_ASYNC_STREAM_METHOD_RETURN_TYPE = "AsyncIterable[Any]" -_DEFAULT_METHOD_DOCSTRING_TEMPLATE = """ - Runs the Agent Engine to serve the user request. - This will be based on the `.{method_name}(...)` of the python object that - was passed in when creating the Agent Engine. The method will invoke the - `{default_method_name}` API client of the python object. - Args: - **kwargs: - Optional. The arguments of the `.{method_name}(...)` method. - Returns: - {return_type}: The response from serving the user request. -""" -_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE = ( - "Failed to register API methods. Please follow the guide to " - "register the API methods: " - "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#custom-methods. " - "Error: {%s}" -) -_AGENT_FRAMEWORK_ATTR = "agent_framework" -_DEFAULT_AGENT_FRAMEWORK = "custom" -_BUILD_OPTIONS_INSTALLATION = "installation_scripts" -_DEFAULT_METHOD_NAME_MAP = { - _STANDARD_API_MODE: _DEFAULT_METHOD_NAME, - _ASYNC_API_MODE: _DEFAULT_ASYNC_METHOD_NAME, - _STREAM_API_MODE: _DEFAULT_STREAM_METHOD_NAME, - _ASYNC_STREAM_API_MODE: _DEFAULT_ASYNC_STREAM_METHOD_NAME, -} -_DEFAULT_METHOD_RETURN_TYPE_MAP = { - _STANDARD_API_MODE: _DEFAULT_METHOD_RETURN_TYPE, - _ASYNC_API_MODE: _DEFAULT_ASYNC_METHOD_RETURN_TYPE, - _STREAM_API_MODE: _DEFAULT_STREAM_METHOD_RETURN_TYPE, - _ASYNC_STREAM_API_MODE: _DEFAULT_ASYNC_STREAM_METHOD_RETURN_TYPE, -} - - -try: - from google.adk.agents import BaseAgent - - ADKAgent = BaseAgent -except (ImportError, AttributeError): - ADKAgent = None - -try: - from a2a.types import ( - AgentCard, - AgentInterface, - Message, - TaskIdParams, - TaskQueryParams, - ) - from a2a.utils.constants import TransportProtocol, PROTOCOL_VERSION_CURRENT - from a2a.client import ClientConfig, ClientFactory - - AgentCard = AgentCard - AgentInterface = AgentInterface - TransportProtocol = TransportProtocol - PROTOCOL_VERSION_CURRENT = PROTOCOL_VERSION_CURRENT - Message = Message - ClientConfig = ClientConfig - ClientFactory = ClientFactory - TaskIdParams = TaskIdParams - TaskQueryParams = TaskQueryParams -except (ImportError, AttributeError): - AgentCard = None - AgentInterface = None - TransportProtocol = None - PROTOCOL_VERSION_CURRENT = None - Message = None - ClientConfig = None - ClientFactory = None - TaskIdParams = None - TaskQueryParams = None - - -@typing.runtime_checkable -class Queryable(Protocol): - """Protocol for Agent Engines that can be queried.""" - - @abc.abstractmethod - def query(self, **kwargs) -> Any: - """Runs the Agent Engine to serve the user query.""" - - -@typing.runtime_checkable -class AsyncQueryable(Protocol): - """Protocol for Agent Engines that can be queried asynchronously.""" - - @abc.abstractmethod - def async_query(self, **kwargs) -> Coroutine[Any, Any, Any]: - """Runs the Agent Engine to serve the user query asynchronously.""" - - -@typing.runtime_checkable -class AsyncStreamQueryable(Protocol): - """Protocol for Agent Engines that can stream responses asynchronously.""" - - @abc.abstractmethod - async def async_stream_query(self, **kwargs) -> AsyncIterable[Any]: - """Asynchronously stream responses to serve the user query.""" - - -@typing.runtime_checkable -class StreamQueryable(Protocol): - """Protocol for Agent Engines that can stream responses.""" - - @abc.abstractmethod - def stream_query(self, **kwargs) -> Iterable[Any]: - """Stream responses to serve the user query.""" - - -@typing.runtime_checkable -class BidiStreamQueryable(Protocol): - """Protocol for Agent Engines that can stream requests and responses.""" - - @abc.abstractmethod - async def bidi_stream_query(self, **kwargs) -> AsyncIterable[Any]: - """Asynchronously stream requests and responses to serve the user query.""" - - -@typing.runtime_checkable -class Cloneable(Protocol): - """Protocol for Agent Engines that can be cloned.""" - - @abc.abstractmethod - def clone(self) -> Any: - """Return a clone of the object.""" - - -@typing.runtime_checkable -class OperationRegistrable(Protocol): - """Protocol for agents that have registered operations.""" - - @abc.abstractmethod - def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]: - """Register the user provided operations (modes and methods).""" - - -_AgentEngineInterface = Union[ - ADKAgent, - AsyncQueryable, - AsyncStreamQueryable, - BidiStreamQueryable, - OperationRegistrable, - Queryable, - StreamQueryable, -] - - -def _wrap_agent_operation(agent: Any, operation: str): - """Wraps an agent operation into a method (works for all API modes).""" - - def _method(self, **kwargs): - if not self._tmpl_attrs.get("agent"): - self.set_up() - return getattr(self._tmpl_attrs["agent"], operation)(**kwargs) - - _method.__name__ = operation - _method.__doc__ = getattr(agent, operation).__doc__ - return _method - - -class ModuleAgent(Cloneable, OperationRegistrable): - """Agent that is defined by a module and an agent name. - - This agent is instantiated by importing a module and instantiating an agent - from that module. It also allows to register operations that are defined in - the agent. - """ - - def __init__( - self, - *, - module_name: str, - agent_name: str, - register_operations: Dict[str, Sequence[str]], - sys_paths: Optional[Sequence[str]] = None, - agent_framework: Optional[str] = None, - ): - """Initializes a module-based agent. - - Args: - module_name (str): - Required. The name of the module to import. - agent_name (str): - Required. The name of the agent in the module to instantiate. - register_operations (Dict[str, Sequence[str]]): - Required. A dictionary of API modes to a list of method names. - sys_paths (Sequence[str]): - Optional. The system paths to search for the module. It should - be relative to the directory where the code will be running. - I.e. it should correspond to the directory being passed to - `extra_packages=...` in the create method. It will be appended - to the system path in the sequence being specified here, and - only be appended if it is not already in the system path. - """ - self.agent_framework = agent_framework - self._tmpl_attrs = { - "module_name": module_name, - "agent_name": agent_name, - "register_operations": register_operations, - "sys_paths": sys_paths, - } - - def clone(self): - """Return a clone of the agent.""" - return ModuleAgent( - module_name=self._tmpl_attrs.get("module_name"), - agent_name=self._tmpl_attrs.get("agent_name"), - register_operations=self._tmpl_attrs.get("register_operations"), - sys_paths=self._tmpl_attrs.get("sys_paths"), - agent_framework=self.agent_framework, - ) - - def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]: - return self._tmpl_attrs.get("register_operations") - - def set_up(self) -> None: - """Sets up the agent for execution of queries at runtime. - - It runs the code to import the agent from the module, and registers the - operations of the agent. - """ - if self._tmpl_attrs.get("sys_paths"): - import sys - - for sys_path in self._tmpl_attrs.get("sys_paths"): - abs_path = os.path.abspath(sys_path) - if abs_path not in sys.path: - sys.path.append(abs_path) - - import importlib - - module = importlib.import_module(self._tmpl_attrs.get("module_name")) - try: - importlib.reload(module) - except Exception as e: - _LOGGER.warning( - f"Failed to reload module {self._tmpl_attrs.get('module_name')}: {e}" - ) - agent_name = self._tmpl_attrs.get("agent_name") - try: - agent = getattr(module, agent_name) - except AttributeError as e: - raise AttributeError( - f"Agent {agent_name} not found in module " - f"{self._tmpl_attrs.get('module_name')}" - ) from e - if not self.agent_framework: - self.agent_framework = _get_agent_framework(agent) - self._tmpl_attrs["agent"] = agent - if hasattr(agent, "set_up"): - agent.set_up() - for operations in self.register_operations().values(): - for operation in operations: - op = _wrap_agent_operation(agent, operation) - setattr(self, operation, types.MethodType(op, self)) - - -class AgentEngine(base.VertexAiResourceNounWithFutureManager): - """Represents a Vertex AI Agent Engine resource.""" - - client_class = aip_utils.AgentEngineClientWithOverride - _resource_noun = "reasoning_engine" - _getter_method = "get_reasoning_engine" - _list_method = "list_reasoning_engines" - _delete_method = "delete_reasoning_engine" - _parse_resource_name_method = "parse_reasoning_engine_path" - _format_resource_name_method = "reasoning_engine_path" - - def __init__(self, resource_name: str): - """Retrieves an Agent Engine resource. - - Args: - resource_name (str): - Required. A fully-qualified resource name or ID such as - "projects/123/locations/us-central1/reasoningEngines/456" or - "456" when project and location are initialized or passed. - """ - super().__init__(resource_name=resource_name) - self.execution_api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionClientWithOverride, - ) - self.execution_async_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionAsyncClientWithOverride, - ) - self._gca_resource = self._get_gca_resource(resource_name=resource_name) - try: - _register_api_methods_or_raise(self) - except Exception as e: - _LOGGER.warning(_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e) - self._operation_schemas = None - - @property - def resource_name(self) -> str: - """Fully-qualified resource name.""" - return self._gca_resource.name - - @classmethod - def create( - cls, - agent_engine: Optional[_AgentEngineInterface] = None, - *, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - build_options: Optional[Dict[str, Sequence[str]]] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[aip_types.EncryptionSpec] = None, - ) -> "AgentEngine": - """Creates a new Agent Engine. - - The Agent Engine will be an instance of the `agent_engine` that - was passed in, running remotely on Vertex AI. - - Sample `src_dir` contents (e.g. `./user_src_dir`): - - .. code-block:: python - - user_src_dir/ - |-- main.py - |-- requirements.txt - |-- user_code/ - | |-- utils.py - | |-- ... - |-- installation_scripts/ - | |-- install_package.sh - | |-- ... - |-- ... - - To build an Agent Engine with the above files, run: - - .. code-block:: python - - remote_agent = agent_engines.create( - agent_engine=local_agent, - requirements=[ - # I.e. the PyPI dependencies listed in requirements.txt - "google-cloud-aiplatform==1.25.0", - "langchain==0.0.242", - ... - ], - extra_packages=[ - "./user_src_dir/main.py", # a single file - "./user_src_dir/user_code", # a directory - ... - ], - build_options={ - "installation_scripts": [ - "./user_src_dir/installation_scripts/install_package.sh", - ... - ], - }, - ) - - Args: - agent_engine (AgentEngineInterface): - Optional. The Agent Engine to be created. - requirements (Union[str, Sequence[str]]): - Optional. The set of PyPI dependencies needed. It can either be - the path to a single file (requirements.txt), or an ordered list - of strings corresponding to each line of the requirements file. - display_name (str): - Optional. The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any - UTF-8 character. - description (str): - Optional. The description of the Agent Engine. - gcs_dir_name (str): - Optional. The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - extra_packages (Sequence[str]): - Optional. The set of extra user-provided packages (if any). - env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): - Optional. The environment variables to be set when running the - Agent Engine. If it is a list of strings, each string should be - a valid key to `os.environ`. If it is a dictionary, the keys are - the environment variable names, and the values are the - corresponding values. - build_options (Dict[str, Sequence[str]]): - Optional. The build options for the Agent Engine. - The following keys are supported: - - installation_scripts: - Optional. The paths to the installation scripts to be - executed in the Docker image. - The scripts must be located in the `installation_scripts` - subdirectory and the path must be added to `extra_packages`. - service_account (str): - Optional. The service account to be used for the Agent Engine. - If not specified, the default reasoning engine service agent - service account will be used. - psc_interface_config (aip_types.PscInterfaceConfig): - Optional. The Private Service Connect interface config for the - Agent Engine. - min_instances (int): - Optional. The minimum number of instances to be running for the - Agent Engine. - max_instances (int): - Optional. The maximum number of instances to be running for the - Agent Engine. - resource_limits (Dict[str, str]): - Optional. The resource limits for the Agent Engine. - container_concurrency (int): - Optional. The container concurrency for the Agent Engine. - encryption_spec (aip_types.EncryptionSpec): - Optional. The Cloud KMS resource identifier of the customer - managed encryption key used to protect the model. Has the - form: - `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`. - The key needs to be in the same region as the model. - - Returns: - AgentEngine: The Agent Engine that was created. - - Raises: - ValueError: If the `project` was not set using `agentplatform.init`. - ValueError: If the `location` was not set using `agentplatform.init`. - ValueError: If the `staging_bucket` was not set using agentplatform.init. - ValueError: If the `staging_bucket` does not start with "gs://". - ValueError: If `extra_packages` is specified but `agent_engine` is None. - ValueError: If `requirements` is specified but `agent_engine` is None. - ValueError: If `env_vars` has a dictionary entry that does not - correspond to a SecretRef. - ValueError: If `env_vars` is a list which contains a string that - does not exist in `os.environ`. - TypeError: If `env_vars` is not a list of strings or a dictionary. - TypeError: If `env_vars` has a value that is not a string or SecretRef. - FileNotFoundError: If `extra_packages` includes a file or directory - that does not exist. - IOError: If requirements is a string that corresponds to a - nonexistent file. - """ - sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" - _validate_sys_version_or_raise(sys_version) - gcs_dir_name = gcs_dir_name or _DEFAULT_GCS_DIR_NAME - staging_bucket = initializer.global_config.staging_bucket - - if agent_engine is not None: - agent_engine = _validate_agent_engine_or_raise(agent_engine) - staging_bucket = _validate_staging_bucket_or_raise(staging_bucket) - if _is_adk_agent(None, agent_engine): - env_vars = _add_telemetry_enablement_env(env_vars=env_vars) - - if agent_engine is None: - if requirements is not None: - raise ValueError("requirements must be None if agent_engine is None.") - if extra_packages is not None: - raise ValueError("extra_packages must be None if agent_engine is None.") - requirements = _validate_requirements_or_raise( - agent_engine=agent_engine, - requirements=requirements, - ) - extra_packages = _validate_extra_packages_or_raise( - extra_packages=extra_packages, - build_options=build_options, - ) - - sdk_resource = cls.__new__(cls) - base.VertexAiResourceNounWithFutureManager.__init__(sdk_resource) - - # Prepares the Agent Engine for creation in Vertex AI. - # This involves packaging and uploading the artifacts for - # agent_engine, requirements and extra_packages to - # `staging_bucket/gcs_dir_name`. - _prepare( - agent_engine=agent_engine, - requirements=requirements, - project=sdk_resource.project, - location=sdk_resource.location, - staging_bucket=staging_bucket, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - ) - reasoning_engine = aip_types.ReasoningEngine( - display_name=display_name, - description=description, - encryption_spec=encryption_spec, - ) - if agent_engine is not None: - # Update the package spec. - package_spec = aip_types.ReasoningEngineSpec.PackageSpec( - python_version=sys_version, - pickle_object_gcs_uri="{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _BLOB_FILENAME, - ), - ) - if extra_packages: - package_spec.dependency_files_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _EXTRA_PACKAGES_FILE, - ) - if requirements: - package_spec.requirements_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _REQUIREMENTS_FILE, - ) - agent_engine_spec = aip_types.ReasoningEngineSpec( - package_spec=package_spec, - ) - if ( - env_vars - or psc_interface_config - or min_instances is not None - or max_instances is not None - or resource_limits - or container_concurrency is not None - ): - deployment_spec, _ = _generate_deployment_spec_or_raise( - env_vars=env_vars, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - ) - agent_engine_spec.deployment_spec = deployment_spec - class_methods_spec = _generate_class_methods_spec_or_raise( - agent_engine=agent_engine, - operations=_get_registered_operations(agent_engine), - ) - agent_engine_spec.class_methods.extend(class_methods_spec) - if service_account: - agent_engine_spec.service_account = service_account - reasoning_engine.spec = agent_engine_spec - reasoning_engine.spec.agent_framework = _get_agent_framework(agent_engine) - operation_future = sdk_resource.api_client.create_reasoning_engine( - parent=initializer.global_config.common_location_path( - project=sdk_resource.project, location=sdk_resource.location - ), - reasoning_engine=reasoning_engine, - ) - _LOGGER.log_create_with_lro(cls, operation_future) - _LOGGER.info( - f"View progress and logs at https://console.cloud.google.com/logs/query?project={sdk_resource.project}" - ) - created_resource = operation_future.result() - _LOGGER.info(f"{cls.__name__} created. Resource name: {created_resource.name}") - _LOGGER.info(f"To use this {cls.__name__} in another session:") - _LOGGER.info( - f"agent_engine = agentplatform.agent_engines.get('{created_resource.name}')" - ) - # We use `._get_gca_resource(...)` instead of `created_resource` to - # fully instantiate the attributes of the agent engine. - sdk_resource._gca_resource = sdk_resource._get_gca_resource( - resource_name=created_resource.name - ) - sdk_resource.execution_api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionClientWithOverride, - credentials=sdk_resource.credentials, - location_override=sdk_resource.location, - ) - sdk_resource.execution_async_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionAsyncClientWithOverride, - credentials=sdk_resource.credentials, - location_override=sdk_resource.location, - ) - if agent_engine is not None: - try: - _register_api_methods_or_raise(sdk_resource) - except Exception as e: - _LOGGER.warning(_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e) - sdk_resource._operation_schemas = None - return sdk_resource - - def update( - self, - *, - agent_engine: Optional[_AgentEngineInterface] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - build_options: Optional[Dict[str, Sequence[str]]] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[aip_types.EncryptionSpec] = None, - ) -> "AgentEngine": - """Updates an existing Agent Engine. - - This method updates the configuration of an existing Agent Engine - running remotely, which is identified by its resource name. - Unlike the `create` function which requires a `agent_engine` object, - all arguments in this method are optional. - This method allows you to modify individual aspects of the configuration - by providing any of the optional arguments. - - Args: - agent_engine (AgentEngineInterface): - Optional. The instance to be used as the updated Agent Engine. - If it is not specified, the existing instance will be used. - requirements (Union[str, Sequence[str]]): - Optional. The set of PyPI dependencies needed. It can either be - the path to a single file (requirements.txt), or an ordered list - of strings corresponding to each line of the requirements file. - If it is not specified, the existing requirements will be used. - If it is set to an empty string or list, the existing - requirements will be removed. - display_name (str): - Optional. The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any - UTF-8 character. - description (str): - Optional. The description of the Agent Engine. - gcs_dir_name (str): - Optional. The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - extra_packages (Sequence[str]): - Optional. The set of extra user-provided packages (if any). If - it is not specified, the existing extra packages will be used. - If it is set to an empty list, the existing extra packages will - be removed. - env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): - Optional. The environment variables to be set when running the - Agent Engine. If it is a list of strings, each string should be - a valid key to `os.environ`. If it is a dictionary, the keys are - the environment variable names, and the values are the - corresponding values. - build_options (Dict[str, Sequence[str]]): - Optional. The build options for the Agent Engine. - The following keys are supported: - - installation_scripts: - Optional. The paths to the installation scripts to be - executed in the Docker image. - The scripts must be located in the `installation_scripts` - subdirectory and the path must be added to `extra_packages`. - service_account (str): - Optional. The service account to be used for the Agent Engine. - If not specified, the default reasoning engine service agent - service account will be used. - psc_interface_config (aip_types.PscInterfaceConfig): - Optional. The Private Service Connect interface config for the - Agent Engine. - min_instances (int): - Optional. The minimum number of instances to be running for the - Agent Engine. - max_instances (int): - Optional. The maximum number of instances to be running for the - Agent Engine. - resource_limits (Dict[str, str]): - Optional. The resource limits for the Agent Engine. - container_concurrency (int): - Optional. The container concurrency for the Agent Engine. - encryption_spec (aip_types.EncryptionSpec): - Optional. The Cloud KMS resource identifier of the customer - managed encryption key used to protect the model. Has the - form: - `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`. - The key needs to be in the same region as the model. - - Returns: - AgentEngine: The Agent Engine that was updated. - - Raises: - ValueError: If the `staging_bucket` was not set using agentplatform.init. - ValueError: If the `staging_bucket` does not start with "gs://". - ValueError: If `env_vars` has a dictionary entry that does not - correspond to a SecretRef. - ValueError: If `env_vars` is a list which contains a string that - does not exist in `os.environ`. - TypeError: If `env_vars` is not a list of strings or a dictionary. - TypeError: If `env_vars` has a value that is not a string or SecretRef. - FileNotFoundError: If `extra_packages` includes a file or directory - that does not exist. - ValueError: if none of `display_name`, `description`, `requirements`, - `extra_packages`, `env_vars`, or `agent_engine` were specified. - IOError: If requirements is a string that corresponds to a - nonexistent file. - """ - staging_bucket = initializer.global_config.staging_bucket - staging_bucket = _validate_staging_bucket_or_raise(staging_bucket) - historical_operation_schemas = self.operation_schemas() - gcs_dir_name = gcs_dir_name or _DEFAULT_GCS_DIR_NAME - - # Validate the arguments. - if not any( - [ - agent_engine, - requirements, - extra_packages, - display_name, - description, - env_vars, - build_options, - service_account, - psc_interface_config, - min_instances is not None, - max_instances is not None, - resource_limits, - container_concurrency is not None, - encryption_spec, - ] - ): - raise ValueError( - "At least one of `agent_engine`, `requirements`, " - "`extra_packages`, `display_name`, `description`, " - "`env_vars`, `build_options`, `service_account`, " - "`psc_interface_config`, `min_instances`, `max_instances`, " - "`resource_limits`, `container_concurrency`, or " - "`encryption_spec` must be specified." - ) - if requirements is not None: - requirements = _validate_requirements_or_raise( - agent_engine=agent_engine, - requirements=requirements, - ) - if extra_packages is not None: - extra_packages = _validate_extra_packages_or_raise( - extra_packages=extra_packages, - build_options=build_options, - ) - if agent_engine is not None: - agent_engine = _validate_agent_engine_or_raise(agent_engine) - - if _is_adk_agent(self, agent_engine): - env_vars = _add_telemetry_enablement_env(env_vars=env_vars) - - # Prepares the Agent Engine for update in Vertex AI. This involves - # packaging and uploading the artifacts for agent_engine, requirements - # and extra_packages to `staging_bucket/gcs_dir_name`. - _prepare( - agent_engine=agent_engine, - requirements=requirements, - project=self.project, - location=self.location, - staging_bucket=staging_bucket, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - ) - update_request = _generate_update_request_or_raise( - resource_name=self.resource_name, - staging_bucket=staging_bucket, - gcs_dir_name=gcs_dir_name, - agent_engine=agent_engine, - requirements=requirements, - extra_packages=extra_packages, - display_name=display_name, - description=description, - env_vars=env_vars, - service_account=service_account, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - encryption_spec=encryption_spec, - ) - operation_future = self.api_client.update_reasoning_engine( - request=update_request - ) - _LOGGER.info( - f"Update Agent Engine backing LRO: {operation_future.operation.name}" - ) - created_resource = operation_future.result() - _LOGGER.info(f"Agent Engine updated. Resource name: {created_resource.name}") - self._operation_schemas = None - self.execution_api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionClientWithOverride, - ) - # We use `._get_gca_resource(...)` instead of `created_resource` to - # fully instantiate the attributes of the agent engine. - self._gca_resource = self._get_gca_resource(resource_name=self.resource_name) - - if ( - agent_engine is None - or historical_operation_schemas == self.operation_schemas() - ): - # The operations of the agent engine are unchanged, so we return it. - return self - - # If the agent engine has changed and the historical operation - # schemas are different from the current operation schemas, we need to - # unregister the historical operation schemas and register the current - # operation schemas. - _unregister_api_methods(self, historical_operation_schemas) - try: - _register_api_methods_or_raise(self) - except Exception as e: - _LOGGER.warning(_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e) - return self - - def delete( - self, - *, - force: bool = False, - **kwargs, - ) -> None: - """Deletes the ReasoningEngine. - - Args: - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error - when the Agent Engine has undeleted child resources. Defaults to - False. - **kwargs (dict[str, Any]): - Optional. Additional keyword arguments to pass to the - delete_reasoning_engine method. - """ - kwargs = kwargs or {} - operation_future = self.api_client.delete_reasoning_engine( - request=aip_types.DeleteReasoningEngineRequest( - name=self.resource_name, - force=force, - **kwargs, - ), - ) - _LOGGER.info( - f"Delete Agent Engine backing LRO: {operation_future.operation.name}" - ) - operation_future.result() - _LOGGER.info(f"Agent Engine deleted. Resource name: {self.resource_name}") - - def operation_schemas(self) -> Sequence[_agent_engines_utils.JsonDict]: - """Returns the (Open)API schemas for the Agent Engine.""" - spec = _agent_engines_utils._to_dict(self._gca_resource.spec) - if not hasattr(self, "_operation_schemas") or self._operation_schemas is None: - self._operation_schemas = spec.get("class_methods", []) - return self._operation_schemas - - -def _validate_sys_version_or_raise(sys_version: str) -> None: - """Tries to validate the python system version.""" - if sys_version not in _SUPPORTED_PYTHON_VERSIONS: - raise ValueError( - f"Unsupported python version: {sys_version}. AgentEngine " - f"only supports {_SUPPORTED_PYTHON_VERSIONS} at the moment." - ) - if sys_version != f"{sys.version_info.major}.{sys.version_info.minor}": - _LOGGER.warning( - f"{sys_version=} is inconsistent with {sys.version_info=}. " - "This might result in issues with deployment, and should only " - "be used as a workaround for advanced cases." - ) - - -def _validate_staging_bucket_or_raise(staging_bucket: Optional[str]) -> str: - """Tries to validate the staging bucket.""" - if not staging_bucket: - raise ValueError( - "Please provide a `staging_bucket` in `agentplatform.init(...)`" - ) - if not staging_bucket.startswith("gs://"): - raise ValueError(f"{staging_bucket=} must start with `gs://`") - return staging_bucket - - -def _validate_agent_engine_or_raise( - agent_engine: _AgentEngineInterface, - logger: base.Logger = _LOGGER, -) -> _AgentEngineInterface: - """Tries to validate the agent engine. - - The agent engine must have one of the following: - * a callable method named `query` - * a callable method named `stream_query` - * a callable method named `async_stream_query` - * a callable method named `bidi_stream_query` - * a callable method named `register_operations` - - Args: - agent_engine: The agent engine to be validated. - logger: The logger to use for logging. - - Returns: - The validated agent engine. - - Raises: - TypeError: If `agent_engine` has no callable method named `query`, - `stream_query` or `register_operations`. - ValueError: If `agent_engine` has an invalid `query`, `stream_query` or - `register_operations` signature. - """ - try: - from google.adk.agents import BaseAgent - - if isinstance(agent_engine, BaseAgent): - logger.info("Deploying google.adk.agents.Agent as an application.") - from agentplatform import agent_engines - - agent_engine = agent_engines.AdkApp(agent=agent_engine) - except Exception: - pass - is_queryable = isinstance(agent_engine, Queryable) and callable(agent_engine.query) - is_async_queryable = isinstance(agent_engine, AsyncQueryable) and callable( - agent_engine.async_query - ) - is_stream_queryable = isinstance(agent_engine, StreamQueryable) and callable( - agent_engine.stream_query - ) - is_async_stream_queryable = isinstance( - agent_engine, AsyncStreamQueryable - ) and callable(agent_engine.async_stream_query) - is_bidi_stream_queryable = isinstance( - agent_engine, BidiStreamQueryable - ) and callable(agent_engine.bidi_stream_query) - is_operation_registrable = isinstance( - agent_engine, OperationRegistrable - ) and callable(agent_engine.register_operations) - - if not ( - is_queryable - or is_async_queryable - or is_stream_queryable - or is_operation_registrable - or is_async_stream_queryable - or is_bidi_stream_queryable - ): - raise TypeError( - "agent_engine has none of the following callable methods: " - "`query`, `async_query`, `stream_query`, `async_stream_query`, " - "`bidi_stream_query` or `register_operations`." - ) - - if is_queryable: - try: - inspect.signature(getattr(agent_engine, "query")) - except ValueError as err: - raise ValueError( - "Invalid query signature. This might be due to a missing " - "`self` argument in the agent_engine.query method." - ) from err - - if is_async_queryable: - try: - inspect.signature(getattr(agent_engine, "async_query")) - except ValueError as err: - raise ValueError( - "Invalid async_query signature. This might be due to a missing " - "`self` argument in the agent_engine.async_query method." - ) from err - - if is_stream_queryable: - try: - inspect.signature(getattr(agent_engine, "stream_query")) - except ValueError as err: - raise ValueError( - "Invalid stream_query signature. This might be due to a missing" - " `self` argument in the agent_engine.stream_query method." - ) from err - - if is_async_stream_queryable: - try: - inspect.signature(getattr(agent_engine, "async_stream_query")) - except ValueError as err: - raise ValueError( - "Invalid async_stream_query signature. This might be due to a " - " missing `self` argument in the " - "agent_engine.async_stream_query method." - ) from err - - if is_bidi_stream_queryable: - try: - inspect.signature(getattr(agent_engine, "bidi_stream_query")) - except ValueError as err: - raise ValueError( - "Invalid bidi_stream_query signature. This might be due to a " - " missing `self` argument in the " - "agent_engine.bidi_stream_query method." - ) from err - - if is_operation_registrable: - try: - inspect.signature(getattr(agent_engine, "register_operations")) - except ValueError as err: - raise ValueError( - "Invalid register_operations signature. This might be due to a " - "missing `self` argument in the " - "agent_engine.register_operations method." - ) from err - - if isinstance(agent_engine, Cloneable): - # Avoid undeployable states. - agent_engine = agent_engine.clone() - return agent_engine - - -def _is_adk_agent( - agent_engine_to_update: Optional[AgentEngine], - new_agent_engine: Optional[_AgentEngineInterface], -) -> bool: - """Checks if the agent engine is an ADK agent. - - Args: - agent_engine_to_update: Existing agent engine, None if creating new one. - new_agent_engine: The new agent engine to deploy. Can be None during an update, if the Python agent implementation is not provided, and should remain unchanged. - - Returns: - True if the agent after the create/update operation, will be an ADK agent. - """ - - from agentplatform.agent_engines.templates import adk - - if new_agent_engine is not None: - return ( - getattr(new_agent_engine, "agent_framework", None) - == adk.AdkApp.agent_framework - ) - if agent_engine_to_update is not None: - return ( - agent_engine_to_update.gca_resource.spec.agent_framework - == adk.AdkApp.agent_framework - ) - return False - - -EnvVars = Optional[Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]] - - -def _add_telemetry_enablement_env(*, env_vars: EnvVars) -> EnvVars: - """Adds telemetry enablement env var to the env vars. - - This is in order to achieve default-on telemetry. - If the telemetry enablement env var is already set, we do not override it. - - Args: - env_vars: The env vars to add the telemetry enablement env var to. - - Returns: - The env vars with the telemetry enablement env var added. - """ - - GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = ( - "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" - ) - - if env_vars is None: - return {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"} - if isinstance(env_vars, dict): - return ( - env_vars - if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in env_vars - else env_vars | {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"} - ) - if isinstance(env_vars, list) or isinstance(env_vars, tuple): - if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY not in os.environ: - os.environ[GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY] = "unspecified" - - if isinstance(env_vars, list): - return env_vars + [GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY] - else: - return env_vars + (GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY,) - - raise TypeError( - f"env_vars must be a list, tuple or a dict, but got {type(env_vars)}." - ) - - -def _validate_requirements_or_raise( - *, - agent_engine: _AgentEngineInterface, - requirements: Optional[Sequence[str]] = None, - logger: logging.getLoggerClass() = _LOGGER, -) -> Sequence[str]: - """Tries to validate the requirements.""" - if requirements is None: - requirements = [] - elif isinstance(requirements, str): - try: - logger.info(f"Reading requirements from {requirements=}") - with open(requirements) as f: - requirements = f.read().splitlines() - logger.info(f"Read the following lines: {requirements}") - except IOError as err: - raise IOError(f"Failed to read requirements from {requirements=}") from err - requirements = _agent_engines_utils._validate_requirements_or_warn( - obj=agent_engine, - requirements=requirements, - ) - logger.info(f"The final list of requirements: {requirements}") - return requirements - - -def _validate_extra_packages_or_raise( - extra_packages: Optional[Sequence[str]], - build_options: Optional[Dict[str, Sequence[str]]] = None, -) -> Sequence[str]: - """Tries to validates the extra packages.""" - extra_packages = extra_packages or [] - if build_options and _BUILD_OPTIONS_INSTALLATION in build_options: - _agent_engines_utils._validate_installation_scripts_or_raise( - script_paths=build_options[_BUILD_OPTIONS_INSTALLATION], - packages=extra_packages, - ) - for extra_package in extra_packages: - if not os.path.exists(extra_package): - raise FileNotFoundError( - f"Extra package specified but not found: {extra_package=}" - ) - return extra_packages - - -def _get_gcs_bucket( - *, - project: str, - location: str, - staging_bucket: str, - logger: base.Logger = _LOGGER, -) -> storage.Bucket: - """Gets or creates the GCS bucket.""" - storage = _agent_engines_utils._import_cloud_storage_or_raise() - storage_client = storage.Client(project=project) - staging_bucket = staging_bucket.replace("gs://", "") - try: - gcs_bucket = storage_client.get_bucket(staging_bucket) - logger.info(f"Using bucket {staging_bucket}") - except exceptions.NotFound: - new_bucket = storage_client.bucket(staging_bucket) - gcs_bucket = storage_client.create_bucket(new_bucket, location=location) - logger.info(f"Creating bucket {staging_bucket} in {location=}") - return gcs_bucket - - -def _upload_agent_engine( - *, - agent_engine: _AgentEngineInterface, - gcs_bucket: storage.Bucket, - gcs_dir_name: str, - logger: base.Logger = _LOGGER, -) -> None: - """Uploads the agent engine to GCS.""" - cloudpickle = _agent_engines_utils._import_cloudpickle_or_raise() - blob = gcs_bucket.blob(f"{gcs_dir_name}/{_BLOB_FILENAME}") - with blob.open("wb") as f: - try: - cloudpickle.dump(agent_engine, f) - except Exception as e: - url = "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#deployment-considerations" - error_msg = f"Failed to serialize agent engine. Visit {url} for details." - if "google._upb._message" in str(e) or "Descriptor" in str(e): - error_msg += ( - " This is often caused by protobuf objects (like Part, AgentCard) " - "being imported at the global module level. Please move these " - "imports inside the functions or methods where they are used. " - "Alternatively, you can import the entire module: " - "`from a2a import types as a2a_types`." - ) - raise TypeError(error_msg) from e - with blob.open("rb") as f: - try: - _ = cloudpickle.load(f) - except Exception as e: - raise TypeError("Agent engine serialized to an invalid format") from e - dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" - logger.info(f"Wrote to {dir_name}/{_BLOB_FILENAME}") - - -def _upload_requirements( - *, - requirements: Sequence[str], - gcs_bucket: storage.Bucket, - gcs_dir_name: str, - logger: base.Logger = _LOGGER, -) -> None: - """Uploads the requirements file to GCS.""" - blob = gcs_bucket.blob(f"{gcs_dir_name}/{_REQUIREMENTS_FILE}") - blob.upload_from_string("\n".join(requirements)) - dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" - logger.info(f"Writing to {dir_name}/{_REQUIREMENTS_FILE}") - - -def _upload_extra_packages( - *, - extra_packages: Sequence[str], - gcs_bucket: storage.Bucket, - gcs_dir_name: str, - logger: base.Logger = _LOGGER, -) -> None: - """Uploads extra packages to GCS.""" - logger.info("Creating in-memory tarfile of extra_packages") - tar_fileobj = io.BytesIO() - with tarfile.open(fileobj=tar_fileobj, mode="w|gz") as tar: - for file in extra_packages: - tar.add(file) - tar_fileobj.seek(0) - blob = gcs_bucket.blob(f"{gcs_dir_name}/{_EXTRA_PACKAGES_FILE}") - blob.upload_from_string(tar_fileobj.read()) - dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" - logger.info(f"Writing to {dir_name}/{_EXTRA_PACKAGES_FILE}") - - -def _prepare( - agent_engine: Optional[_AgentEngineInterface], - requirements: Optional[Sequence[str]], - extra_packages: Optional[Sequence[str]], - project: str, - location: str, - staging_bucket: str, - gcs_dir_name: str, - logger: base.Logger = _LOGGER, -) -> None: - """Prepares the agent engine for creation or updates in Vertex AI. - - This involves packaging and uploading artifacts to Cloud Storage. Note that - 1. This does not actually update the Agent Engine in Vertex AI. - 2. This will only generate and upload a pickled object if specified. - 3. This will only generate and upload the dependencies.tar.gz file if - extra_packages is non-empty. - - Args: - agent_engine: The agent engine to be prepared. - requirements (Sequence[str]): The set of PyPI dependencies needed. - extra_packages (Sequence[str]): The set of extra user-provided packages. - project (str): The project for the staging bucket. - location (str): The location for the staging bucket. - staging_bucket (str): The staging bucket name in the form "gs://...". - gcs_dir_name (str): The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - """ - if agent_engine is None: - return - gcs_bucket = _get_gcs_bucket( - project=project, - location=location, - staging_bucket=staging_bucket, - logger=logger, - ) - _upload_agent_engine( - agent_engine=agent_engine, - gcs_bucket=gcs_bucket, - gcs_dir_name=gcs_dir_name, - logger=logger, - ) - if requirements is not None: - _upload_requirements( - requirements=requirements, - gcs_bucket=gcs_bucket, - gcs_dir_name=gcs_dir_name, - logger=logger, - ) - if extra_packages is not None: - _upload_extra_packages( - extra_packages=extra_packages, - gcs_bucket=gcs_bucket, - gcs_dir_name=gcs_dir_name, - logger=logger, - ) - - -def _update_deployment_spec_with_env_vars_dict_or_raise( - *, - deployment_spec: aip_types.ReasoningEngineSpec.DeploymentSpec, - env_vars: Dict[str, Union[str, aip_types.SecretRef]], -) -> None: - for key, value in env_vars.items(): - if isinstance(value, Dict): - try: - secret_ref = _agent_engines_utils._to_proto( - value, aip_types.SecretRef() - ) - except Exception as e: - raise ValueError(f"Failed to convert to secret ref: {value}") from e - deployment_spec.secret_env.append( - aip_types.SecretEnvVar(name=key, secret_ref=secret_ref) - ) - elif isinstance(value, aip_types.SecretRef): - deployment_spec.secret_env.append( - aip_types.SecretEnvVar(name=key, secret_ref=value) - ) - elif isinstance(value, str): - deployment_spec.env.append(aip_types.EnvVar(name=key, value=value)) - else: - raise TypeError( - f"Unknown value type in env_vars for {key}. " - f"Must be a str or SecretRef: {value}" - ) - - -def _update_deployment_spec_with_env_vars_list_or_raise( - *, - deployment_spec: aip_types.ReasoningEngineSpec.DeploymentSpec, - env_vars: Sequence[str], -) -> None: - for env_var in env_vars: - if env_var not in os.environ: - raise ValueError(f"Env var not found in os.environ: {env_var}.") - deployment_spec.env.append( - aip_types.EnvVar(name=env_var, value=os.environ[env_var]) - ) - - -def _generate_deployment_spec_or_raise( - *, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, -) -> Tuple[aip_types.ReasoningEngineSpec.DeploymentSpec, List[str]]: - deployment_spec = aip_types.ReasoningEngineSpec.DeploymentSpec() - update_masks = [] - if env_vars: - deployment_spec.env = [] - deployment_spec.secret_env = [] - if isinstance(env_vars, Dict): - _update_deployment_spec_with_env_vars_dict_or_raise( - deployment_spec=deployment_spec, - env_vars=env_vars, - ) - elif isinstance(env_vars, Sequence): - _update_deployment_spec_with_env_vars_list_or_raise( - deployment_spec=deployment_spec, - env_vars=env_vars, - ) - else: - raise TypeError( - f"env_vars must be a list, tuple or a dict, but got {type(env_vars)}." - ) - if deployment_spec.env: - update_masks.append("spec.deployment_spec.env") - if deployment_spec.secret_env: - update_masks.append("spec.deployment_spec.secret_env") - if psc_interface_config: - deployment_spec.psc_interface_config = psc_interface_config - update_masks.append("spec.deployment_spec.psc_interface_config") - if min_instances is not None: - deployment_spec.min_instances = min_instances - update_masks.append("spec.deployment_spec.min_instances") - if max_instances is not None: - deployment_spec.max_instances = max_instances - update_masks.append("spec.deployment_spec.max_instances") - if resource_limits: - deployment_spec.resource_limits = resource_limits - update_masks.append("spec.deployment_spec.resource_limits") - if container_concurrency is not None: - deployment_spec.container_concurrency = container_concurrency - update_masks.append("spec.deployment_spec.container_concurrency") - return deployment_spec, update_masks - - -def _get_agent_framework( - agent_engine: _AgentEngineInterface, -) -> str: - if ( - hasattr(agent_engine, _AGENT_FRAMEWORK_ATTR) - and getattr(agent_engine, _AGENT_FRAMEWORK_ATTR) is not None - and isinstance(getattr(agent_engine, _AGENT_FRAMEWORK_ATTR), str) - ): - return getattr(agent_engine, _AGENT_FRAMEWORK_ATTR) - return _DEFAULT_AGENT_FRAMEWORK - - -def _generate_update_request_or_raise( - *, - resource_name: str, - staging_bucket: str, - gcs_dir_name: str = _DEFAULT_GCS_DIR_NAME, - agent_engine: Optional[_AgentEngineInterface] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - extra_packages: Optional[Sequence[str]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[str] = None, -) -> reasoning_engine_service.UpdateReasoningEngineRequest: - """Tries to generate the update request for the agent engine.""" - is_spec_update = False - update_masks: List[str] = [] - agent_engine_spec = aip_types.ReasoningEngineSpec() - package_spec = aip_types.ReasoningEngineSpec.PackageSpec() - if requirements is not None: - is_spec_update = True - update_masks.append("spec.package_spec.requirements_gcs_uri") - package_spec.requirements_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _REQUIREMENTS_FILE, - ) - if extra_packages is not None: - is_spec_update = True - update_masks.append("spec.package_spec.dependency_files_gcs_uri") - package_spec.dependency_files_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _EXTRA_PACKAGES_FILE, - ) - if agent_engine is not None: - is_spec_update = True - update_masks.append("spec.package_spec.pickle_object_gcs_uri") - package_spec.pickle_object_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _BLOB_FILENAME, - ) - class_methods_spec = _generate_class_methods_spec_or_raise( - agent_engine=agent_engine, - operations=_get_registered_operations(agent_engine), - ) - agent_engine_spec.class_methods.extend(class_methods_spec) - update_masks.append("spec.class_methods") - agent_engine_spec.agent_framework = _get_agent_framework(agent_engine) - update_masks.append("spec.agent_framework") - if ( - env_vars is not None - or psc_interface_config - or min_instances is not None - or max_instances is not None - or resource_limits - or container_concurrency is not None - ): - is_spec_update = True - deployment_spec, deployment_update_masks = _generate_deployment_spec_or_raise( - env_vars=env_vars, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - ) - update_masks.extend(deployment_update_masks) - agent_engine_spec.deployment_spec = deployment_spec - if service_account is not None: - is_spec_update = True - update_masks.append("spec.service_account") - agent_engine_spec.service_account = service_account - - agent_engine_message = aip_types.ReasoningEngine(name=resource_name) - if is_spec_update: - if package_spec: - agent_engine_spec.package_spec = package_spec - agent_engine_message.spec = agent_engine_spec - if display_name: - agent_engine_message.display_name = display_name - update_masks.append("display_name") - if description: - agent_engine_message.description = description - update_masks.append("description") - if encryption_spec: - agent_engine_message.encryption_spec = encryption_spec - update_masks.append("encryption_spec") - if not update_masks: - raise ValueError( - "At least one of `agent_engine`, `requirements`, `extra_packages`, " - "`display_name`, `description`, `env_vars`, or " - "`encryption_spec` must be specified." - ) - return reasoning_engine_service.UpdateReasoningEngineRequest( - reasoning_engine=agent_engine_message, - update_mask=field_mask_pb2.FieldMask(paths=update_masks), - ) - - -def _wrap_query_operation( - method_name: str, -) -> Callable[..., _agent_engines_utils.JsonDict]: - """Wraps an Agent Engine method, creating a callable for `query` API. - - This function creates a callable object that executes the specified - Agent Engine method using the `query` API. It handles the creation of - the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - doc: Documentation string for the method. - - Returns: - A callable object that executes the method on the Agent Engine via - the `query` API. - """ - - def _method(self, **kwargs) -> _agent_engines_utils.JsonDict: - response = self.execution_api_client.query_reasoning_engine( - request=aip_types.QueryReasoningEngineRequest( - name=self.resource_name, - input=kwargs, - class_method=method_name, - ), - ) - output = _agent_engines_utils._to_dict(response) - return output.get("output", output) - - return _method - - -def _wrap_async_query_operation(method_name: str) -> Callable[..., Coroutine]: - """Wraps an Agent Engine method, creating an async callable for `query` API. - - This function creates a callable object that executes the specified - Agent Engine method asynchronously using the `query` API. It handles the - creation of the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - doc: Documentation string for the method. - - Returns: - A callable object that executes the method on the Agent Engine via - the `query` API. - """ - - async def _method(self, **kwargs) -> _agent_engines_utils.JsonDict: - response = await self.execution_async_client.query_reasoning_engine( - request=aip_types.QueryReasoningEngineRequest( - name=self.resource_name, - input=kwargs, - class_method=method_name, - ), - ) - output = _agent_engines_utils._to_dict(response) - return output.get("output", output) - - return _method - - -def _wrap_stream_query_operation(*, method_name: str) -> Callable[..., Iterable[Any]]: - """Wraps an Agent Engine method, creating a callable for `stream_query` API. - - This function creates a callable object that executes the specified - Agent Engine method using the `stream_query` API. It handles the - creation of the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - doc: Documentation string for the method. - - Returns: - A callable object that executes the method on the Agent Engine via - the `stream_query` API. - """ - - def _method(self, **kwargs) -> Iterable[Any]: - response = self.execution_api_client.stream_query_reasoning_engine( - request=aip_types.StreamQueryReasoningEngineRequest( - name=self.resource_name, - input=kwargs, - class_method=method_name, - ), - ) - for chunk in response: - for parsed_json in _agent_engines_utils._yield_parsed_json_from_httpbody( - chunk - ): - if parsed_json is not None: - yield parsed_json - - return _method - - -def _wrap_async_stream_query_operation( - *, method_name: str -) -> Callable[..., AsyncIterable[Any]]: - """Wraps an Agent Engine method, creating an async callable for `stream_query` API. - - This function creates a callable object that executes the specified - Agent Engine method using the `stream_query` API. It handles the - creation of the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - doc: Documentation string for the method. - - Returns: - A callable object that executes the method on the Agent Engine via - the `stream_query` API. - """ - - async def _method(self, **kwargs) -> AsyncIterable[Any]: - response = self.execution_api_client.stream_query_reasoning_engine( - request=aip_types.StreamQueryReasoningEngineRequest( - name=self.resource_name, - input=kwargs, - class_method=method_name, - ), - ) - for chunk in response: - for parsed_json in _agent_engines_utils._yield_parsed_json_from_httpbody( - chunk - ): - if parsed_json is not None: - yield parsed_json - - return _method - - -def _wrap_bidi_stream_query_operation( - *, method_name: str -) -> Callable[..., AsyncIterable[Any]]: - """Wraps an Agent Engine method, creating an async callable for `bidi_stream_query` API. - - This function creates a callable object that executes the specified - Agent Engine method using the `bidi_stream_query` API. It handles the - creation of the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - - Returns: - A callable object that executes the method on the Agent Engine via - the `bidi_stream_query` API. - """ - - async def _method(self, **kwargs) -> AsyncIterable[Any]: - # Agent Engine bidi streaming query execution should use GenAI SDK Agent - # Engine live API client directly. - raise NotImplementedError( - f"{method_name} is not implemented, please use GenAI SDK Agent " - "Enginve live API client instead." - ) - - -def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list]: - """Wraps an Agent Engine method, creating a callable for A2A API. - - Args: - method_name: The name of the Agent Engine method to call. - agent_card: The agent card to use for the A2A API call. - Example: - {'additionalInterfaces': None, - 'capabilities': {'extensions': None, - 'pushNotifications': None, - 'stateTransitionHistory': None, - 'streaming': False}, - 'defaultInputModes': ['text'], - 'defaultOutputModes': ['text'], - 'description': ( - 'A helpful assistant agent that can answer questions.' - ), - 'documentationUrl': None, - 'iconUrl': None, - 'name': 'Q&A Agent', - 'preferredTransport': 'JSONRPC', - 'protocolVersion': '0.3.0', - 'provider': None, - 'security': None, - 'securitySchemes': None, - 'signatures': None, - 'skills': [{ - 'description': ( - 'A helpful assistant agent that can answer questions.' - ), - 'examples': ['Who is leading 2025 F1 Standings?', - 'Where can i find an active volcano?'], - 'id': 'question_answer', - 'inputModes': None, - 'name': 'Q&A Agent', - 'outputModes': None, - 'security': None, - 'tags': ['Question-Answer']}], - 'supportsAuthenticatedExtendedCard': True, - 'url': 'http://localhost:8080/', - 'version': '1.0.0'} - Returns: - A callable object that executes the method on the Agent Engine via - the A2A API. - """ - - async def _method(self, **kwargs) -> Any: - """Wraps an Agent Engine method, creating a callable for A2A API.""" - a2a_agent_card = AgentCard(**json.loads(agent_card)) - - # A2A + AE integration currently only supports Rest API. - if ( - a2a_agent_card.supported_interfaces - and a2a_agent_card.supported_interfaces[0].protocol_binding - != TransportProtocol.HTTP_JSON - ): - raise ValueError( - "Only HTTP+JSON is supported for primary interface on agent card " - ) - - # Set primary interface to HTTP+JSON if not set. - if not a2a_agent_card.supported_interfaces: - a2a_agent_card.supported_interfaces = [] - a2a_agent_card.supported_interfaces.append( - AgentInterface( - protocol_binding=TransportProtocol.HTTP_JSON, - protocol_version=PROTOCOL_VERSION_CURRENT, - ) - ) - - # AE cannot support streaming yet. Turn off streaming for now. - if a2a_agent_card.capabilities and a2a_agent_card.capabilities.streaming: - raise ValueError( - "Streaming is not supported in Agent Engine, please change " - "a2a_agent_card.capabilities.streaming to False." - ) - - if not hasattr(a2a_agent_card.capabilities, "streaming"): - a2a_agent_card.capabilities.streaming = False - - # agent_card is set on the class_methods before set_up is invoked. - # Ensure that the agent_card url is set correctly before the client is created. - url = f"https://{initializer.global_config.api_endpoint}/v1beta1/{self.resource_name}/a2a" - a2a_agent_card.supported_interfaces[0].url = url - - # Using a2a client, inject the auth token from the global config. - config = ClientConfig( - supported_transports=[ - TransportProtocol.HTTP_JSON, - ], - use_client_preference=True, - httpx_client=httpx.AsyncClient( - headers={ - "Authorization": ( - f"Bearer {initializer.global_config.credentials.token}" - ) - } - ), - ) - factory = ClientFactory(config) - client = factory.create(a2a_agent_card) - - # kokoro job uses python 3.9, replaced match with if else. - if method_name == "on_message_send": - response = client.send_message(Message(**kwargs)) - chunks = [] - async for chunk in response: - chunks.append(chunk) - return chunks - elif method_name == "on_get_task": - response = await client.get_task(TaskQueryParams(**kwargs)) - elif method_name == "on_cancel_task": - response = await client.cancel_task(TaskIdParams(**kwargs)) - elif method_name == "handle_authenticated_agent_card": - response = await client.get_card() - else: - raise ValueError(f"Unknown method name: {method_name}") - - return response - - return _method - - -def _unregister_api_methods( - obj: "AgentEngine", operation_schemas: Sequence[_agent_engines_utils.JsonDict] -): - """Unregisters Agent Engine API methods based on operation schemas. - - This function iterates through operation schemas provided by the - AgentEngine object. Each schema defines an API mode and method name. - It dynamically unregisters methods on the AgentEngine object. This - should only be used when updating the object. - - Args: - obj: The AgentEngine object to augment with API methods. - operation_schemas: The operation schemas to use for method unregistration. - """ - for operation_schema in operation_schemas: - if "name" in operation_schema: - method_name = operation_schema.get("name") - if hasattr(obj, method_name): - delattr(obj, method_name) - - -def _register_api_methods_or_raise( - obj: "AgentEngine", - wrap_operation_fn: Optional[ - dict[str, Callable[[str, str], Callable[..., Any]]] - ] = None, -): - """Registers Agent Engine API methods based on operation schemas. - - This function iterates through operation schemas provided by the - AgentEngine object. Each schema defines an API mode and method name. - It dynamically creates and registers methods on the AgentEngine object - to handle API calls based on the specified API mode. - Currently, only standard API mode `` is supported. - - Args: - obj: The AgentEngine object to augment with API methods. - wrap_operation_fn: A dictionary of API modes and method wrapping - functions. - - Raises: - ValueError: If the API mode is not supported or if the operation schema - is missing any required fields (e.g. `api_mode` or `name`). - """ - for operation_schema in obj.operation_schemas(): - if _MODE_KEY_IN_SCHEMA not in operation_schema: - raise ValueError( - f"Operation schema {operation_schema} does not" - f" contain an `{_MODE_KEY_IN_SCHEMA}` field." - ) - api_mode = operation_schema.get(_MODE_KEY_IN_SCHEMA) - if _METHOD_NAME_KEY_IN_SCHEMA not in operation_schema: - raise ValueError( - f"Operation schema {operation_schema} does not" - f" contain a `{_METHOD_NAME_KEY_IN_SCHEMA}` field." - ) - method_name = operation_schema.get(_METHOD_NAME_KEY_IN_SCHEMA) - method_description = operation_schema.get( - "description", - _DEFAULT_METHOD_DOCSTRING_TEMPLATE.format( - method_name=method_name, - default_method_name=_DEFAULT_METHOD_NAME_MAP.get( - api_mode, _DEFAULT_METHOD_NAME - ), - return_type=_DEFAULT_METHOD_RETURN_TYPE_MAP.get( - api_mode, - _DEFAULT_METHOD_RETURN_TYPE, - ), - ), - ) - _wrap_operation_map = { - _STANDARD_API_MODE: _wrap_query_operation, - _ASYNC_API_MODE: _wrap_async_query_operation, - _STREAM_API_MODE: _wrap_stream_query_operation, - _ASYNC_STREAM_API_MODE: _wrap_async_stream_query_operation, - _BIDI_STREAM_API_MODE: _wrap_bidi_stream_query_operation, - _A2A_EXTENSION_MODE: _wrap_a2a_operation, - } - if isinstance(wrap_operation_fn, dict) and api_mode in wrap_operation_fn: - # Override the default function with user-specified function if it exists. - _wrap_operation = wrap_operation_fn[api_mode] - elif api_mode in _wrap_operation_map: - _wrap_operation = _wrap_operation_map[api_mode] - else: - supported_api_modes = ", ".join( - f"`{mode}`" for mode in sorted(_wrap_operation_map.keys()) - ) - raise ValueError( - f"Unsupported api mode: `{api_mode}`," - f" Supported modes are: {supported_api_modes}." - ) - - # Bind the method to the object. - if api_mode == _A2A_EXTENSION_MODE: - agent_card = operation_schema.get(_A2A_AGENT_CARD) - method = _wrap_operation(method_name=method_name, agent_card=agent_card) - else: - method = _wrap_operation(method_name=method_name) - method.__name__ = method_name - method.__doc__ = method_description - setattr(obj, method_name, types.MethodType(method, obj)) - - -def _get_registered_operations( - agent_engine: _AgentEngineInterface, -) -> Dict[str, List[str]]: - """Retrieves registered operations for a AgentEngine.""" - if isinstance(agent_engine, OperationRegistrable): - return agent_engine.register_operations() - - operations = {} - if isinstance(agent_engine, Queryable): - operations[_STANDARD_API_MODE] = [_DEFAULT_METHOD_NAME] - if isinstance(agent_engine, AsyncQueryable): - operations[_ASYNC_API_MODE] = [_DEFAULT_ASYNC_METHOD_NAME] - if isinstance(agent_engine, StreamQueryable): - operations[_STREAM_API_MODE] = [_DEFAULT_STREAM_METHOD_NAME] - if isinstance(agent_engine, AsyncStreamQueryable): - operations[_ASYNC_STREAM_API_MODE] = [_DEFAULT_ASYNC_STREAM_METHOD_NAME] - return operations - - -def _generate_class_methods_spec_or_raise( - *, - agent_engine: _AgentEngineInterface, - operations: Dict[str, List[str]], - logger: base.Logger = _LOGGER, -) -> List[proto.Message]: - """Generates a ReasoningEngineSpec based on the registered operations. - - Args: - agent_engine: The AgentEngine instance. - operations: A dictionary of API modes and method names. - - Returns: - A list of ReasoningEngineSpec.ClassMethod messages. - - Raises: - ValueError: If a method defined in `register_operations` is not found on - the AgentEngine. - """ - if isinstance(agent_engine, ModuleAgent): - # We do a dry-run of setting up the agent engine to have the operations - # needed for registration. - agent_engine = agent_engine.clone() - try: - agent_engine.set_up() - except Exception as e: - raise ValueError( - f"Failed to set up agent engine {agent_engine}: {e}" - ) from e - class_methods_spec = [] - for mode, method_names in operations.items(): - for method_name in method_names: - if mode == _BIDI_STREAM_API_MODE: - _LOGGER.warning( - "Bidi stream API mode is not supported yet in Vertex SDK, " - "please use the GenAI SDK instead. Skipping " - f"method {method_name}." - ) - continue - if not hasattr(agent_engine, method_name): - raise ValueError( - f"Method `{method_name}` defined in `register_operations`" - " not found on AgentEngine." - ) - - method = getattr(agent_engine, method_name) - try: - schema_dict = _agent_engines_utils._generate_schema( - method, schema_name=method_name - ) - except Exception as e: - logger.warning(f"failed to generate schema for {method_name}: {e}") - continue - - class_method = _agent_engines_utils._to_proto(schema_dict) - class_method[_MODE_KEY_IN_SCHEMA] = mode - # A2A agent card is a special case, when running in A2A mode, - if hasattr(agent_engine, "agent_card"): - from google.protobuf import json_format - - class_method[_A2A_AGENT_CARD] = json_format.MessageToJson( - getattr(agent_engine, "agent_card") - ) - class_methods_spec.append(class_method) - - return class_methods_spec - - -def _class_methods_to_class_methods_spec( - class_methods: List[dict[str, Any]], -) -> List[proto.Message]: - """Converts a list of class methods to a list of ReasoningEngineSpec.ClassMethod messages.""" - return [ - _agent_engines_utils._to_proto(class_method) for class_method in class_methods - ] diff --git a/agentplatform/frameworks/__init__.py b/agentplatform/frameworks/__init__.py new file mode 100644 index 0000000000..346cdc96c9 --- /dev/null +++ b/agentplatform/frameworks/__init__.py @@ -0,0 +1,39 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Classes for working with agent platform.""" + +from agentplatform.frameworks import a2a +from agentplatform.frameworks import adk +from agentplatform.frameworks import ag2 +from agentplatform.frameworks import langchain +from agentplatform.frameworks import langgraph +from agentplatform.frameworks import llama_index + + +A2aAgent = a2a.A2aAgent +AdkApp = adk.AdkApp +AG2Agent = ag2.AG2Agent +LangchainAgent = langchain.LangchainAgent +LanggraphAgent = langgraph.LanggraphAgent +LlamaIndexQueryPipelineAgent = llama_index.LlamaIndexQueryPipelineAgent + +__all__ = ( + "A2aAgent", + "AdkApp", + "AG2Agent", + "LangchainAgent", + "LanggraphAgent", + "LlamaIndexQueryPipelineAgent", +) diff --git a/agentplatform/agent_engines/templates/a2a.py b/agentplatform/frameworks/a2a.py similarity index 97% rename from agentplatform/agent_engines/templates/a2a.py rename to agentplatform/frameworks/a2a.py index 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 9be3e35b89..1ed8b57bd4 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: @@ -1001,19 +996,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 @@ -1054,7 +1046,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" @@ -1065,11 +1057,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, ) @@ -1101,7 +1095,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): @@ -1113,7 +1107,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"), ) @@ -1135,7 +1129,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): @@ -1147,7 +1141,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: @@ -1239,7 +1233,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): @@ -1299,7 +1293,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( @@ -1349,7 +1343,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): @@ -1376,7 +1370,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, @@ -1384,7 +1378,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. @@ -2123,8 +2117,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_memories_private_create.py b/tests/unit/agentplatform/genai/replays/test_memories_private_create.py index ed2904d9e4..91a4638dd3 100644 --- a/tests/unit/agentplatform/genai/replays/test_memories_private_create.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_private_create.py @@ -34,5 +34,5 @@ def test_private_create_memory(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._create", + test_method="memory_banks.memories._create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py b/tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py index ac9cc30941..6b0cb359ed 100644 --- a/tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py +++ b/tests/unit/agentplatform/genai/replays/test_memories_retrieve_profiles.py @@ -41,7 +41,7 @@ def test_generate_and_retrieve_profile(client): structured_memory_config_obj = types.StructuredMemoryConfig( **structured_memory_config ) - memory_bank = client.agent_engines.create( + memory_bank = client.runtimes.create( config={ "context_spec": { "memory_bank_config": { @@ -53,7 +53,7 @@ def test_generate_and_retrieve_profile(client): }, ) try: - memory_bank = client.agent_engines.get(name=memory_bank.api_resource.name) + memory_bank = client.runtimes.get(name=memory_bank.api_resource.name) memory_bank_config = memory_bank.api_resource.context_spec.memory_bank_config assert memory_bank_config.customization_configs == [ memory_bank_customization_config 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, )