diff --git a/agentplatform/_genai/_evals_common.py b/agentplatform/_genai/_evals_common.py index d91d352dc1..687537b89b 100644 --- a/agentplatform/_genai/_evals_common.py +++ b/agentplatform/_genai/_evals_common.py @@ -125,7 +125,7 @@ def _get_api_client_with_location( def _get_agent_engine_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 = {} @@ -135,7 +135,7 @@ def _get_agent_engine_instance( location=api_client.location, ) _thread_local_data.agent_engine_instances[agent_name] = ( - client.agent_engines.get(name=agent_name) + client.runtimes.get(name=agent_name) ) return _thread_local_data.agent_engine_instances[agent_name] @@ -1695,7 +1695,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, + agent_engine: Optional[Union[str, types.Runtime]] = None, agent: Optional["LlmAgent"] = None, # type: ignore # noqa: F821 user_simulator_config: Optional[types.evals.UserSimulatorConfig] = None, ) -> list[ @@ -2450,7 +2450,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, + agent_engine: Optional[Union[str, types.Runtime]] = None, agent: Optional["LlmAgent"] = None, # type: ignore # noqa: F821 gemini_agent: Optional[str] = None, dest: Optional[str] = None, @@ -2469,7 +2469,7 @@ def _execute_inference( 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. + 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. @@ -2569,14 +2569,14 @@ def _execute_inference( and not isinstance(agent_engine, str) and not ( hasattr(agent_engine, "api_client") - and type(agent_engine).__name__ == "AgentEngine" + and type(agent_engine).__name__ == "Runtime" ) ): raise TypeError( f"Unsupported agent_engine type: {type(agent_engine)}. 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 @@ -3204,7 +3204,7 @@ def _create_agent_results_dataframe( def _run_agent_internal( api_client: BaseApiClient, - agent_engine: Optional[Union[str, types.AgentEngine]], + agent_engine: Optional[Union[str, types.Runtime]], agent: Optional["LlmAgent"], # type: ignore # noqa: F821 prompt_dataset: pd.DataFrame, user_simulator_config: Optional[types.evals.UserSimulatorConfig] = None, @@ -3255,7 +3255,7 @@ def _run_agent_internal( def _run_agent( api_client: BaseApiClient, - agent_engine: Optional[Union[str, types.AgentEngine]], + agent_engine: Optional[Union[str, types.Runtime]], agent: Optional["LlmAgent"], # type: ignore # noqa: F821 prompt_dataset: pd.DataFrame, user_simulator_config: Optional[types.evals.UserSimulatorConfig] = None, @@ -3300,7 +3300,7 @@ def _run_agent( def _create_agent_engine_session( *, - agent_engine: types.AgentEngine, + agent_engine: types.Runtime, user_id: str, session_state: Optional[dict[str, Any]] = None, ) -> Any: @@ -3312,7 +3312,7 @@ def _create_agent_engine_session( Sessions API. Args: - agent_engine: The AgentEngine instance. + agent_engine: The Runtime instance. user_id: The user ID for the session. session_state: Optional initial state for the session. @@ -3347,7 +3347,7 @@ def _create_agent_engine_session( operation = agent_engine.api_client.sessions.create( name=agent_engine.api_resource.name, user_id=user_id, - config=types.CreateAgentEngineSessionConfig( + config=types.CreateRuntimeSessionConfig( session_state=session_state, ), ) @@ -3369,7 +3369,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, + agent_engine: types.Runtime, max_retries: int = 3, ) -> Union[list[dict[str, Any]], dict[str, Any]]: """Executes agent run over agent engine for a single prompt.""" @@ -3416,7 +3416,7 @@ def _execute_agent_run_with_retry( 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, ), ) diff --git a/agentplatform/_genai/_agent_engines_utils.py b/agentplatform/_genai/_runtimes_utils.py similarity index 93% rename from agentplatform/_genai/_agent_engines_utils.py rename to agentplatform/_genai/_runtimes_utils.py index d6b1820371..cedebbee69 100644 --- a/agentplatform/_genai/_agent_engines_utils.py +++ b/agentplatform/_genai/_runtimes_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -"""Utility functions for agent engines.""" +"""Utility functions for runtimes.""" import abc import asyncio @@ -182,9 +182,9 @@ _DEFAULT_ASYNC_STREAM_METHOD_RETURN_TYPE = "AsyncIterable[Any]" _DEFAULT_GCS_DIR_NAME = "agent_engine" _DEFAULT_METHOD_DOCSTRING_TEMPLATE = """ - Runs the Agent Engine to serve the user request. + Runs the Agent Runtime to serve the user request. This will be based on the `.{method_name}(...)` of the python object that - was passed in when creating the Agent Engine. The method will invoke the + was passed in when creating the Agent Runtime. The method will invoke the `{default_method_name}` API client of the python object. Args: **kwargs: @@ -236,25 +236,25 @@ @typing.runtime_checkable class Queryable(Protocol): - """Protocol for Agent Engines that can be queried.""" + """Protocol for Agent Runtimes that can be queried.""" @abc.abstractmethod def query(self, **kwargs): # type: ignore[no-untyped-def] - """Runs the Agent Engine to serve the user query.""" + """Runs the Agent Runtime to serve the user query.""" @typing.runtime_checkable class AsyncQueryable(Protocol): - """Protocol for Agent Engines that can be queried asynchronously.""" + """Protocol for Agent Runtimes that can be queried asynchronously.""" @abc.abstractmethod def async_query(self, **kwargs): # type: ignore[no-untyped-def] - """Runs the Agent Engine to serve the user query asynchronously.""" + """Runs the Agent Runtime to serve the user query asynchronously.""" @typing.runtime_checkable class AsyncStreamQueryable(Protocol): - """Protocol for Agent Engines that can stream responses asynchronously.""" + """Protocol for Agent Runtimes that can stream responses asynchronously.""" @abc.abstractmethod async def async_stream_query(self, **kwargs) -> AsyncIterator[Any]: # type: ignore[no-untyped-def] @@ -263,7 +263,7 @@ async def async_stream_query(self, **kwargs) -> AsyncIterator[Any]: # type: ign @typing.runtime_checkable class StreamQueryable(Protocol): - """Protocol for Agent Engines that can stream responses.""" + """Protocol for Agent Runtimes that can stream responses.""" @abc.abstractmethod def stream_query(self, **kwargs) -> Iterator[Any]: # type: ignore[no-untyped-def] @@ -272,7 +272,7 @@ def stream_query(self, **kwargs) -> Iterator[Any]: # type: ignore[no-untyped-de @typing.runtime_checkable class BidiStreamQueryable(Protocol): - """Protocol for Agent Engines that can stream requests and responses.""" + """Protocol for Agent Runtimes that can stream requests and responses.""" @abc.abstractmethod async def bidi_stream_query( @@ -283,7 +283,7 @@ async def bidi_stream_query( @typing.runtime_checkable class Cloneable(Protocol): - """Protocol for Agent Engines that can be cloned.""" + """Protocol for Agent Runtimes that can be cloned.""" @abc.abstractmethod def clone(self) -> Any: @@ -312,7 +312,7 @@ def register_operations(self, **kwargs: Any) -> dict[str, list[str]]: except (ImportError, AttributeError): ADKAgent = None # type: ignore[no-redef] -_AgentEngineInterface = Union[ +_RuntimeInterface = Union[ ADKAgent, AsyncQueryable, AsyncStreamQueryable, @@ -328,7 +328,7 @@ class _ModuleAgentAttributes(TypedDict, total=False): agent_name: str register_operations: Dict[str, list[str]] sys_paths: Optional[Sequence[str]] - agent: _AgentEngineInterface + agent: _RuntimeInterface class ModuleAgent(Cloneable, OperationRegistrable): @@ -441,24 +441,22 @@ class _RequirementsValidationResult(TypedDict): actions: _RequirementsValidationActions -AgentEngineOperationUnion = Union[ - genai_types.AgentEngineOperation, - genai_types.AgentEngineMemoryOperation, - genai_types.AgentEngineGenerateMemoriesOperation, +RuntimeOperationUnion = Union[ + genai_types.RuntimeOperation, + genai_types.RuntimeMemoryOperation, + genai_types.RuntimeGenerateMemoriesOperation, ] class GetOperationFunction(Protocol): - def __call__( - self, *, operation_name: str, **kwargs: Any - ) -> AgentEngineOperationUnion: + def __call__(self, *, operation_name: str, **kwargs: Any) -> RuntimeOperationUnion: pass class GetAsyncOperationFunction(Protocol): async def __call__( self, *, operation_name: str, **kwargs: Any - ) -> AgentEngineOperationUnion: + ) -> RuntimeOperationUnion: pass @@ -490,8 +488,7 @@ def _get_reasoning_engine_id(operation_name: str = "", resource_name: str = "") if match: return match.group(1) raise ValueError( - "Failed to parse reasoning engine ID from operation name: " - f"`{operation_name}`" + f"Failed to parse reasoning engine ID from operation name: `{operation_name}`" ) @@ -501,11 +498,11 @@ async def _await_async_operation( get_operation_fn: GetAsyncOperationFunction, poll_interval_seconds: float = 10, ) -> Any: - """Waits for the operation for creating an agent engine to complete. + """Waits for the operation for creating an agent runtime to complete. Args: operation_name (str): - Required. The name of the operation for creating the Agent Engine. + Required. The name of the operation for creating the Agent Runtime. poll_interval_seconds (float): The number of seconds to wait between each poll. get_operation_fn (Callable[[str], Awaitable[Any]]): @@ -529,11 +526,11 @@ def _await_operation( get_operation_fn: GetOperationFunction, poll_interval_seconds: float = 10, ) -> Any: - """Waits for the operation for creating an agent engine to complete. + """Waits for the operation for creating an agent runtime to complete. Args: operation_name (str): - Required. The name of the operation for creating the Agent Engine. + Required. The name of the operation for creating the Agent Runtime. poll_interval_seconds (float): The number of seconds to wait between each poll. get_operation_fn (Callable[[str], Any]): @@ -569,7 +566,7 @@ def _compare_requirements( required_packages (Iterator[str]): Optional. The set of packages that are required to be in the constraints. It defaults to the set of packages that are required - for deployment on Agent Engine. + for deployment on Agent Runtime. Returns: dict[str, dict[str, Any]]: The comparison result as a dictionary containing: @@ -611,7 +608,7 @@ def _compare_requirements( def _generate_class_methods_spec_or_raise( *, - agent: _AgentEngineInterface, + agent: _RuntimeInterface, operations: Dict[str, List[str]], ) -> List[proto.Message]: """Generates a ReasoningEngineSpec based on the registered operations. @@ -628,7 +625,7 @@ def _generate_class_methods_spec_or_raise( the AgentEngine. """ if isinstance(agent, ModuleAgent): - # We do a dry-run of setting up the agent engine to have the operations + # We do a dry-run of setting up the agent runtime to have the operations # needed for registration. agent: ModuleAgent = agent.clone() # type: ignore[no-redef] try: @@ -755,7 +752,8 @@ def _generate_schema( ) # For a bidi endpoint, it requires an asyncio.Queue as the input, but # it is not JSON serializable. We hence exclude it from the schema. - and param.annotation != asyncio.Queue and _is_pydantic_serializable(param) + and param.annotation != asyncio.Queue + and _is_pydantic_serializable(param) } parameters = pydantic.create_model(f.__name__, **fields_dict).schema() # Postprocessing @@ -811,7 +809,7 @@ def _generate_schema( def _get_agent_framework( *, agent_framework: Optional[str], - agent: _AgentEngineInterface, + agent: _RuntimeInterface, ) -> Union[str, Any]: """Gets the agent framework to use. @@ -823,8 +821,8 @@ def _get_agent_framework( Args: agent_framework (str): The agent framework provided by the user. - agent (_AgentEngineInterface): - The agent engine instance. + agent (_RuntimeInterface): + The agent runtime instance. Returns: str: The name of the agent framework to use. @@ -871,7 +869,7 @@ def _get_gcs_bucket( def _get_registered_operations( *, - agent: _AgentEngineInterface, + agent: _RuntimeInterface, ) -> dict[str, list[str]]: """Retrieves registered operations for a AgentEngine.""" if isinstance(agent, OperationRegistrable): @@ -985,7 +983,7 @@ def _parse_constraints( def _prepare( *, - agent: Optional[_AgentEngineInterface], + agent: Optional[_RuntimeInterface], requirements: Optional[Sequence[str]], extra_packages: Optional[Sequence[str]], project: str, @@ -994,16 +992,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. @@ -1042,12 +1040,12 @@ def _prepare( def _register_api_methods_or_raise( *, - agent_engine: genai_types.AgentEngine | genai_types.AgentEngineRuntimeRevision, + agent_engine: 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. @@ -1126,9 +1124,7 @@ 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 @@ -1243,11 +1239,11 @@ def _to_proto( def _upload_agent_engine( *, - agent: _AgentEngineInterface, + agent: _RuntimeInterface, gcs_bucket: _StorageBucket, gcs_dir_name: str, ) -> None: - """Uploads the agent engine to GCS.""" + """Uploads the agent runtime to GCS.""" cloudpickle = _import_cloudpickle_or_raise() blob = gcs_bucket.blob(f"{gcs_dir_name}/{_BLOB_FILENAME}") with blob.open("wb") as f: @@ -1255,7 +1251,7 @@ def _upload_agent_engine( cloudpickle.dump(agent, f) except Exception as e: url = "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#deployment-considerations" - error_msg = f"Failed to serialize agent engine. Visit {url} for details." + error_msg = f"Failed to serialize agent runtime. Visit {url} for details." if "google._upb._message" in str(e) or "Descriptor" in str(e): error_msg += ( " This is often caused by protobuf objects (like Part, AgentCard) " @@ -1269,7 +1265,7 @@ def _upload_agent_engine( try: _ = cloudpickle.load(f) except Exception as e: - raise TypeError("Agent engine serialized to an invalid format") from e + raise TypeError("agent runtime serialized to an invalid format") from e dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" logger.info(f"Wrote to {dir_name}/{_BLOB_FILENAME}") @@ -1412,7 +1408,7 @@ def _validate_staging_bucket_or_raise(*, staging_bucket: str) -> str: """Tries to validate the staging bucket.""" if not staging_bucket: raise ValueError( - "Please provide a `staging_bucket` in `client.agent_engines.create(...)`." + "Please provide a `staging_bucket` in `client.runtimes.create(...)`." ) if not staging_bucket.startswith("gs://"): raise ValueError(f"{staging_bucket=} must start with `gs://`") @@ -1475,11 +1471,11 @@ def _validate_requirements_or_raise( def _validate_agent_or_raise( *, - agent: _AgentEngineInterface, -) -> _AgentEngineInterface: - """Tries to validate the agent engine. + agent: _RuntimeInterface, +) -> _RuntimeInterface: + """Tries to validate the agent runtime. - The agent engine must have one of the following: + The agent runtime must have one of the following: * a callable method named `query` * a callable method named `stream_query` * a callable method named `async_stream_query` @@ -1490,7 +1486,7 @@ 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`, @@ -1503,9 +1499,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) @@ -1613,10 +1609,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 @@ -1632,15 +1628,15 @@ def _wrap_query_operation(*, method_name: str) -> Callable[..., Any]: ) Args: - method_name: The name of the Agent Engine method to call. + method_name: The name of the Agent Runtime method to call. doc: Documentation string for the method. Returns: - A callable object that executes the method on the Agent Engine via + A callable object that executes the method on the Agent Runtime via the `query` API. """ - def _method(self: genai_types.AgentEngine, **kwargs) -> Any: # type: ignore[no-untyped-def] + def _method(self: genai_types.Runtime, **kwargs) -> Any: # type: ignore[no-untyped-def] if not self.api_client: raise ValueError("api_client is not initialized.") if not self.api_resource: @@ -1663,10 +1659,10 @@ def _method(self: genai_types.AgentEngine, **kwargs) -> Any: # type: ignore[no- def _wrap_async_query_operation( *, method_name: str ) -> Callable[..., Coroutine[Any, Any, Any]]: - """Wraps an Agent Engine method, creating an async callable for `query` API. + """Wraps an Agent Runtime method, creating an async callable for `query` API. This function creates a callable object that executes the specified - Agent Engine method asynchronously using the `query` API. It handles the + Agent Runtime method asynchronously using the `query` API. It handles the creation of the API request and the processing of the API response. The reserved keyword argument `http_options` is consumed by this @@ -1674,16 +1670,16 @@ def _wrap_async_query_operation( `input`) and is propagated to the underlying HTTP call. Args: - method_name: The name of the Agent Engine method to call. + method_name: The name of the Agent Runtime method to call. doc: Documentation string for the method. Returns: - A callable object that executes the method on the Agent Engine via + A callable object that executes the method on the Agent Runtime via the `query` API. """ async def _method( - self: genai_types.AgentEngine, **kwargs: Any + self: genai_types.Runtime, **kwargs: Any ) -> Union[Coroutine[Any, Any, Any], Any]: if not self.api_async_client: raise ValueError("api_async_client is not initialized.") @@ -1705,10 +1701,10 @@ async def _method( def _wrap_stream_query_operation(*, method_name: str) -> Callable[..., Iterator[Any]]: - """Wraps an Agent Engine method, creating a callable for `stream_query` API. + """Wraps an Agent Runtime method, creating a callable for `stream_query` API. This function creates a callable object that executes the specified - Agent Engine method using the `stream_query` API. It handles the + Agent Runtime method using the `stream_query` API. It handles the creation of the API request and the processing of the API response. The reserved keyword argument `http_options` is consumed by this @@ -1716,15 +1712,15 @@ def _wrap_stream_query_operation(*, method_name: str) -> Callable[..., Iterator[ `input`) and is propagated to the underlying HTTP call. Args: - method_name: The name of the Agent Engine method to call. + method_name: The name of the Agent Runtime method to call. doc: Documentation string for the method. Returns: - A callable object that executes the method on the Agent Engine via + A callable object that executes the method on the Agent Runtime via the `stream_query` API. """ - def _method(self: genai_types.AgentEngine, **kwargs) -> Iterator[Any]: # type: ignore[no-untyped-def] + def _method(self: genai_types.Runtime, **kwargs) -> Iterator[Any]: # type: ignore[no-untyped-def] if not self.api_client: raise ValueError("api_client is not initialized.") if not self.api_resource: @@ -1749,10 +1745,10 @@ def _method(self: genai_types.AgentEngine, **kwargs) -> Iterator[Any]: # type: def _wrap_async_stream_query_operation( *, method_name: str ) -> Callable[..., AsyncIterator[Any]]: - """Wraps an Agent Engine method, creating an async callable for `stream_query` API. + """Wraps an Agent Runtime method, creating an async callable for `stream_query` API. This function creates a callable object that executes the specified - Agent Engine method using the `stream_query` API. It handles the + Agent Runtime method using the `stream_query` API. It handles the creation of the API request and the processing of the API response. The reserved keyword argument `http_options` is consumed by this @@ -1760,15 +1756,15 @@ def _wrap_async_stream_query_operation( `input`) and is propagated to the underlying HTTP call. Args: - method_name: The name of the Agent Engine method to call. + method_name: The name of the Agent Runtime method to call. doc: Documentation string for the method. Returns: - A callable object that executes the method on the Agent Engine via + A callable object that executes the method on the Agent Runtime via the `stream_query` API. """ - async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any]: # type: ignore[no-untyped-def] + async def _method(self: genai_types.Runtime, **kwargs) -> AsyncIterator[Any]: # type: ignore[no-untyped-def] if not self.api_client: raise ValueError("api_client is not initialized.") if not self.api_resource: @@ -1791,10 +1787,10 @@ async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any] def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list[Any]]: - """Wraps an Agent Engine method, creating a callable for A2A API. + """Wraps an Agent Runtime method, creating a callable for A2A API. Args: - method_name: The name of the Agent Engine method to call. + method_name: The name of the Agent Runtime method to call. agent_card: The agent card to use for the A2A API call. Example: { 'name': 'Sample Agent', 'description': ( 'A helpful assistant agent that can answer questions.' ), @@ -1810,7 +1806,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. """ @@ -1836,9 +1832,9 @@ async def _method(self, **kwargs) -> Any: # type: ignore[no-untyped-def] base_url = self.api_client._api_client._http_options.base_url.rstrip("/") api_version = self.api_client._api_client._http_options.api_version - a2a_agent_card.supported_interfaces[0].url = ( - f"{base_url}/{api_version}/{self.api_resource.name}/a2a" - ) + a2a_agent_card.supported_interfaces[ + 0 + ].url = f"{base_url}/{api_version}/{self.api_resource.name}/a2a" config = ClientConfig( supported_protocol_bindings=[ @@ -1991,7 +1987,7 @@ def _validate_resource_limits_or_raise(resource_limits: dict[str, str]) -> None: if cpu not in [1, 2, 4, 6, 8]: raise ValueError( - "resource_limits['cpu'] must be one of 1, 2, 4, 6, 8. Got" f" {cpu}" + f"resource_limits['cpu'] must be one of 1, 2, 4, 6, 8. Got {cpu}" ) if not isinstance(memory_str, str) or not memory_str.endswith("Gi"): @@ -2004,8 +2000,7 @@ def _validate_resource_limits_or_raise(resource_limits: dict[str, str]) -> None: memory_gb = int(memory_str[:-2]) except ValueError: raise ValueError( - f"Invalid memory value: {memory_str}. Must be an integer" - " followed by 'Gi'." + f"Invalid memory value: {memory_str}. Must be an integer followed by 'Gi'." ) # https://cloud.google.com/run/docs/configuring/memory-limits @@ -2026,28 +2021,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(agent_engine: _RuntimeInterface) -> bool: + """Checks if the agent runtime is an ADK agent. Args: - agent_engine: The agent engine to check. + agent_engine: 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(agent_engine, 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/a2a_task_events.py b/agentplatform/_genai/a2a_task_events.py index 4f3c043d9a..89bf2b9034 100644 --- a/agentplatform/_genai/a2a_task_events.py +++ b/agentplatform/_genai/a2a_task_events.py @@ -35,7 +35,7 @@ logger.setLevel(logging.INFO) -def _AppendAgentEngineTaskEventRequestParameters_to_vertex( +def _AppendRuntimeTaskEventRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -53,7 +53,7 @@ def _AppendAgentEngineTaskEventRequestParameters_to_vertex( return to_object -def _AppendAgentEngineTaskEventResponse_from_vertex( +def _AppendRuntimeTaskEventResponse_from_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -62,7 +62,7 @@ def _AppendAgentEngineTaskEventResponse_from_vertex( return to_object -def _ListAgentEngineTaskEventsConfig_to_vertex( +def _ListRuntimeTaskEventsConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -83,7 +83,7 @@ def _ListAgentEngineTaskEventsConfig_to_vertex( return to_object -def _ListAgentEngineTaskEventsRequestParameters_to_vertex( +def _ListRuntimeTaskEventsRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -92,9 +92,7 @@ def _ListAgentEngineTaskEventsRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineTaskEventsConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _ListRuntimeTaskEventsConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object @@ -106,23 +104,23 @@ def append( *, name: str, task_events: builtins.list[types.TaskEventOrDict], - config: Optional[types.AppendAgentEngineTaskEventConfigOrDict] = None, - ) -> types.AppendAgentEngineTaskEventResponse: + config: Optional[types.AppendRuntimeTaskEventConfigOrDict] = None, + ) -> types.AppendRuntimeTaskEventResponse: """ - Adds events to an Agent Engine task. + Adds events to an Agent Runtime task. Args: - name (str): Required. The name of the Agent Engine task to append the events to. Format: + name (str): Required. The name of the Agent Runtime task to append the events to. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/a2aTasks/{a2a_task_id}`. task_events (list[TaskEvent]): Required. The events to append to the task. Returns: - AppendAgentEngineTaskEventResponse: The response for appending the task events. + AppendRuntimeTaskEventResponse: The response for appending the task events. """ - parameter_model = types._AppendAgentEngineTaskEventRequestParameters( + parameter_model = types._AppendRuntimeTaskEventRequestParameters( name=name, task_events=task_events, config=config, @@ -134,7 +132,7 @@ def append( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _AppendAgentEngineTaskEventRequestParameters_to_vertex( + request_dict = _AppendRuntimeTaskEventRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -164,11 +162,9 @@ def append( response_dict = {} if not response.body else json.loads(response.body) if self._api_client.vertexai: - response_dict = _AppendAgentEngineTaskEventResponse_from_vertex( - response_dict - ) + response_dict = _AppendRuntimeTaskEventResponse_from_vertex(response_dict) - return_value = types.AppendAgentEngineTaskEventResponse._from_response( + return_value = types.AppendRuntimeTaskEventResponse._from_response( response=response_dict, kwargs=( { @@ -196,23 +192,23 @@ def _list( self, *, name: str, - config: Optional[types.ListAgentEngineTaskEventsConfigOrDict] = None, - ) -> types.ListAgentEngineTaskEventsResponse: + config: Optional[types.ListRuntimeTaskEventsConfigOrDict] = None, + ) -> types.ListRuntimeTaskEventsResponse: """ - Lists Agent Engine task events. + Lists Agent Runtime task events. Args: - name (str): Required. The name of the Agent Engine task to list events for. Format: + name (str): Required. The name of the Agent Runtime task to list events for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/a2aTasks/{a2a_task_id}`. - config (ListAgentEngineTaskEventsConfig): - Optional. Additional configurations for listing the Agent Engine tasks. + config (ListRuntimeTaskEventsConfig): + Optional. Additional configurations for listing the Agent Runtime tasks. Returns: - ListAgentEngineTaskEventsResponse: The requested Agent Engine tasks. + ListRuntimeTaskEventsResponse: The requested Agent Runtime tasks. """ - parameter_model = types._ListAgentEngineTaskEventsRequestParameters( + parameter_model = types._ListRuntimeTaskEventsRequestParameters( name=name, config=config, ) @@ -223,7 +219,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineTaskEventsRequestParameters_to_vertex( + request_dict = _ListRuntimeTaskEventsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -252,7 +248,7 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineTaskEventsResponse._from_response( + return_value = types.ListRuntimeTaskEventsResponse._from_response( response=response_dict, kwargs=( { @@ -280,13 +276,13 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineTaskEventsConfigOrDict] = None, + config: Optional[types.ListRuntimeTaskEventsConfigOrDict] = None, ) -> Iterator[types.TaskEvent]: - """Lists the A2A tasks of an Agent Engine. + """Lists the A2A tasks of a Runtime. Args: name (str): - Required. The name of the agent engine to list tasks for. + Required. The name of the agent runtime to list tasks for. config (List): Optional. The configuration for the tasks to list. @@ -309,23 +305,23 @@ async def append( *, name: str, task_events: builtins.list[types.TaskEventOrDict], - config: Optional[types.AppendAgentEngineTaskEventConfigOrDict] = None, - ) -> types.AppendAgentEngineTaskEventResponse: + config: Optional[types.AppendRuntimeTaskEventConfigOrDict] = None, + ) -> types.AppendRuntimeTaskEventResponse: """ - Adds events to an Agent Engine task. + Adds events to an Agent Runtime task. Args: - name (str): Required. The name of the Agent Engine task to append the events to. Format: + name (str): Required. The name of the Agent Runtime task to append the events to. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/a2aTasks/{a2a_task_id}`. task_events (list[TaskEvent]): Required. The events to append to the task. Returns: - AppendAgentEngineTaskEventResponse: The response for appending the task events. + AppendRuntimeTaskEventResponse: The response for appending the task events. """ - parameter_model = types._AppendAgentEngineTaskEventRequestParameters( + parameter_model = types._AppendRuntimeTaskEventRequestParameters( name=name, task_events=task_events, config=config, @@ -337,7 +333,7 @@ async def append( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _AppendAgentEngineTaskEventRequestParameters_to_vertex( + request_dict = _AppendRuntimeTaskEventRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -369,11 +365,9 @@ async def append( response_dict = {} if not response.body else json.loads(response.body) if self._api_client.vertexai: - response_dict = _AppendAgentEngineTaskEventResponse_from_vertex( - response_dict - ) + response_dict = _AppendRuntimeTaskEventResponse_from_vertex(response_dict) - return_value = types.AppendAgentEngineTaskEventResponse._from_response( + return_value = types.AppendRuntimeTaskEventResponse._from_response( response=response_dict, kwargs=( { @@ -401,23 +395,23 @@ async def _list( self, *, name: str, - config: Optional[types.ListAgentEngineTaskEventsConfigOrDict] = None, - ) -> types.ListAgentEngineTaskEventsResponse: + config: Optional[types.ListRuntimeTaskEventsConfigOrDict] = None, + ) -> types.ListRuntimeTaskEventsResponse: """ - Lists Agent Engine task events. + Lists Agent Runtime task events. Args: - name (str): Required. The name of the Agent Engine task to list events for. Format: + name (str): Required. The name of the Agent Runtime task to list events for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/a2aTasks/{a2a_task_id}`. - config (ListAgentEngineTaskEventsConfig): - Optional. Additional configurations for listing the Agent Engine tasks. + config (ListRuntimeTaskEventsConfig): + Optional. Additional configurations for listing the Agent Runtime tasks. Returns: - ListAgentEngineTaskEventsResponse: The requested Agent Engine tasks. + ListRuntimeTaskEventsResponse: The requested Agent Runtime tasks. """ - parameter_model = types._ListAgentEngineTaskEventsRequestParameters( + parameter_model = types._ListRuntimeTaskEventsRequestParameters( name=name, config=config, ) @@ -428,7 +422,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineTaskEventsRequestParameters_to_vertex( + request_dict = _ListRuntimeTaskEventsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -459,7 +453,7 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineTaskEventsResponse._from_response( + return_value = types.ListRuntimeTaskEventsResponse._from_response( response=response_dict, kwargs=( { @@ -487,13 +481,13 @@ async def list( self, *, name: str, - config: Optional[types.ListAgentEngineTaskEventsConfigOrDict] = None, + config: Optional[types.ListRuntimeTaskEventsConfigOrDict] = None, ) -> AsyncPager[types.TaskEvent]: - """Lists the A2A tasks of an Agent Engine. + """Lists the A2A tasks of an Agent Runtime. Args: name (str): - Required. The name of the agent engine to list tasks for. + Required. The name of the agent runtime to list tasks for. config (List): Optional. The configuration for the tasks to list. diff --git a/agentplatform/_genai/a2a_tasks.py b/agentplatform/_genai/a2a_tasks.py index 5cb74c1d8e..39ab9b4c47 100644 --- a/agentplatform/_genai/a2a_tasks.py +++ b/agentplatform/_genai/a2a_tasks.py @@ -42,7 +42,7 @@ logger.setLevel(logging.INFO) -def _CreateAgentEngineTaskConfig_to_vertex( +def _CreateRuntimeTaskConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -63,7 +63,7 @@ def _CreateAgentEngineTaskConfig_to_vertex( return to_object -def _CreateAgentEngineTaskRequestParameters_to_vertex( +def _CreateRuntimeTaskRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -75,12 +75,12 @@ def _CreateAgentEngineTaskRequestParameters_to_vertex( setv(to_object, ["_query", "a2a_task_id"], getv(from_object, ["a2a_task_id"])) if getv(from_object, ["config"]) is not None: - _CreateAgentEngineTaskConfig_to_vertex(getv(from_object, ["config"]), to_object) + _CreateRuntimeTaskConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _DeleteAgentEngineTaskRequestParameters_to_vertex( +def _DeleteRuntimeTaskRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -91,7 +91,7 @@ def _DeleteAgentEngineTaskRequestParameters_to_vertex( return to_object -def _GetAgentEngineTaskRequestParameters_to_vertex( +def _GetRuntimeTaskRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -102,7 +102,7 @@ def _GetAgentEngineTaskRequestParameters_to_vertex( return to_object -def _ListAgentEngineTasksConfig_to_vertex( +def _ListRuntimeTasksConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -123,7 +123,7 @@ def _ListAgentEngineTasksConfig_to_vertex( return to_object -def _ListAgentEngineTasksRequestParameters_to_vertex( +def _ListRuntimeTasksRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -132,7 +132,7 @@ def _ListAgentEngineTasksRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineTasksConfig_to_vertex(getv(from_object, ["config"]), to_object) + _ListRuntimeTasksConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object @@ -140,26 +140,23 @@ def _ListAgentEngineTasksRequestParameters_to_vertex( class A2aTasks(_api_module.BaseModule): def delete( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineTaskConfigOrDict] = None, + self, *, name: str, config: Optional[types.DeleteRuntimeTaskConfigOrDict] = None ) -> None: """ - Deletes an agent engine task. + Deletes an agent runtime task. Args: - name (str): Required. The name of the Agent Engine task to delete. Format: + name (str): Required. The name of the Agent Runtime task to delete. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/a2aTasks/{task_id}`. - config (DeleteAgentEngineTaskConfig): - Optional. Additional configurations for deleting the Agent Engine task. + config (DeleteRuntimeTaskConfig): + Optional. Additional configurations for deleting the Agent Runtime task. Returns: None """ - parameter_model = types._DeleteAgentEngineTaskRequestParameters( + parameter_model = types._DeleteRuntimeTaskRequestParameters( name=name, config=config, ) @@ -170,7 +167,7 @@ def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineTaskRequestParameters_to_vertex( + request_dict = _DeleteRuntimeTaskRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -198,26 +195,23 @@ def delete( self._api_client.request("delete", path, request_dict, http_options) def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineTaskConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeTaskConfigOrDict] = None ) -> types.A2aTask: """ - Gets an agent engine task. + Gets an agent runtime task. Args: - name (str): Required. The name of the Agent Engine task to get. Format: + name (str): Required. The name of the Agent Runtime task to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/a2aTasks/{task_id}`. - config (GetAgentEngineTaskConfig): - Optional. Additional configurations for getting the Agent Engine task. + config (GetRuntimeTaskConfig): + Optional. Additional configurations for getting the Agent Runtime task. Returns: - AgentEngineTask: The requested Agent Engine task. + RuntimeTask: The requested Agent Runtime task. """ - parameter_model = types._GetAgentEngineTaskRequestParameters( + parameter_model = types._GetRuntimeTaskRequestParameters( name=name, config=config, ) @@ -228,9 +222,7 @@ def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineTaskRequestParameters_to_vertex( - parameter_model - ) + request_dict = _GetRuntimeTaskRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -282,26 +274,23 @@ def get( return return_value def _list( - self, - *, - name: str, - config: Optional[types.ListAgentEngineTasksConfigOrDict] = None, - ) -> types.ListAgentEngineTasksResponse: + self, *, name: str, config: Optional[types.ListRuntimeTasksConfigOrDict] = None + ) -> types.ListRuntimeTasksResponse: """ - Lists Agent Engine tasks. + Lists Agent Runtime tasks. Args: - name (str): Required. The name of the Agent Engine to list tasks for. Format: + name (str): Required. The name of the Agent Runtime to list tasks for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineTasksConfig): - Optional. Additional configurations for listing the Agent Engine tasks. + config (ListRuntimeTasksConfig): + Optional. Additional configurations for listing the Agent Runtime tasks. Returns: - ListAgentEngineTasksResponse: The requested Agent Engine tasks. + ListRuntimeTasksResponse: The requested Agent Runtime tasks. """ - parameter_model = types._ListAgentEngineTasksRequestParameters( + parameter_model = types._ListRuntimeTasksRequestParameters( name=name, config=config, ) @@ -312,9 +301,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineTasksRequestParameters_to_vertex( - parameter_model - ) + request_dict = _ListRuntimeTasksRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/a2aTasks".format_map(request_url_dict) @@ -341,7 +328,7 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineTasksResponse._from_response( + return_value = types.ListRuntimeTasksResponse._from_response( response=response_dict, kwargs=( { @@ -370,25 +357,25 @@ def create( *, name: str, a2a_task_id: str, - config: Optional[types.CreateAgentEngineTaskConfigOrDict] = None, + config: Optional[types.CreateRuntimeTaskConfigOrDict] = None, ) -> types.A2aTask: """ - Creates a new task in the Agent Engine. + Creates a new task in the Agent Runtime. Args: - name (str): Required. The name of the Agent Engine to create the task under. Format: + name (str): Required. The name of the Agent Runtime to create the task under. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. a2a_task_id (str): Required. The user ID of the task. context_id (str): Required. The ID of the context to use for the task. - config (CreateAgentEngineTaskConfig): - Optional. Additional configurations for creating the Agent Engine task. + config (CreateRuntimeTaskConfig): + Optional. Additional configurations for creating the Agent Runtime task. Returns: - A2aTask: The created Agent Engine task. + A2aTask: The created Agent Runtime task. """ - parameter_model = types._CreateAgentEngineTaskRequestParameters( + parameter_model = types._CreateRuntimeTaskRequestParameters( name=name, a2a_task_id=a2a_task_id, config=config, @@ -400,7 +387,7 @@ def create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineTaskRequestParameters_to_vertex( + request_dict = _CreateRuntimeTaskRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -459,13 +446,13 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineTasksConfigOrDict] = None, + config: Optional[types.ListRuntimeTasksConfigOrDict] = None, ) -> Iterator[types.A2aTask]: - """Lists the A2A tasks of an Agent Engine. + """Lists the A2A tasks of an Agent Runtime. Args: name (str): - Required. The name of the agent engine to list tasks for. + Required. The name of the agent runtime to list tasks for. config (List): Optional. The configuration for the tasks to list. @@ -489,7 +476,7 @@ def events(self) -> "a2a_task_events_module.A2aTaskEvents": self._events = importlib.import_module(".a2a_task_events", __package__) except ImportError as e: raise ImportError( - "The 'agent_engines.a2a_tasks.events' module requires additional " + "The 'runtimes.a2a_tasks.events' module requires additional " "packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e @@ -499,26 +486,23 @@ def events(self) -> "a2a_task_events_module.A2aTaskEvents": class AsyncA2aTasks(_api_module.BaseModule): async def delete( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineTaskConfigOrDict] = None, + self, *, name: str, config: Optional[types.DeleteRuntimeTaskConfigOrDict] = None ) -> None: """ - Deletes an agent engine task. + Deletes an agent runtime task. Args: - name (str): Required. The name of the Agent Engine task to delete. Format: + name (str): Required. The name of the Agent Runtime task to delete. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/a2aTasks/{task_id}`. - config (DeleteAgentEngineTaskConfig): - Optional. Additional configurations for deleting the Agent Engine task. + config (DeleteRuntimeTaskConfig): + Optional. Additional configurations for deleting the Agent Runtime task. Returns: None """ - parameter_model = types._DeleteAgentEngineTaskRequestParameters( + parameter_model = types._DeleteRuntimeTaskRequestParameters( name=name, config=config, ) @@ -529,7 +513,7 @@ async def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineTaskRequestParameters_to_vertex( + request_dict = _DeleteRuntimeTaskRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -557,26 +541,23 @@ async def delete( await self._api_client.async_request("delete", path, request_dict, http_options) async def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineTaskConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeTaskConfigOrDict] = None ) -> types.A2aTask: """ - Gets an agent engine task. + Gets an agent runtime task. Args: - name (str): Required. The name of the Agent Engine task to get. Format: + name (str): Required. The name of the Agent Runtime task to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/a2aTasks/{task_id}`. - config (GetAgentEngineTaskConfig): - Optional. Additional configurations for getting the Agent Engine task. + config (GetRuntimeTaskConfig): + Optional. Additional configurations for getting the Agent Runtime task. Returns: - AgentEngineTask: The requested Agent Engine task. + RuntimeTask: The requested Agent Runtime task. """ - parameter_model = types._GetAgentEngineTaskRequestParameters( + parameter_model = types._GetRuntimeTaskRequestParameters( name=name, config=config, ) @@ -587,9 +568,7 @@ async def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineTaskRequestParameters_to_vertex( - parameter_model - ) + request_dict = _GetRuntimeTaskRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -643,26 +622,23 @@ async def get( return return_value async def _list( - self, - *, - name: str, - config: Optional[types.ListAgentEngineTasksConfigOrDict] = None, - ) -> types.ListAgentEngineTasksResponse: + self, *, name: str, config: Optional[types.ListRuntimeTasksConfigOrDict] = None + ) -> types.ListRuntimeTasksResponse: """ - Lists Agent Engine tasks. + Lists Agent Runtime tasks. Args: - name (str): Required. The name of the Agent Engine to list tasks for. Format: + name (str): Required. The name of the Agent Runtime to list tasks for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineTasksConfig): - Optional. Additional configurations for listing the Agent Engine tasks. + config (ListRuntimeTasksConfig): + Optional. Additional configurations for listing the Agent Runtime tasks. Returns: - ListAgentEngineTasksResponse: The requested Agent Engine tasks. + ListRuntimeTasksResponse: The requested Agent Runtime tasks. """ - parameter_model = types._ListAgentEngineTasksRequestParameters( + parameter_model = types._ListRuntimeTasksRequestParameters( name=name, config=config, ) @@ -673,9 +649,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineTasksRequestParameters_to_vertex( - parameter_model - ) + request_dict = _ListRuntimeTasksRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}/a2aTasks".format_map(request_url_dict) @@ -704,7 +678,7 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineTasksResponse._from_response( + return_value = types.ListRuntimeTasksResponse._from_response( response=response_dict, kwargs=( { @@ -733,25 +707,25 @@ async def create( *, name: str, a2a_task_id: str, - config: Optional[types.CreateAgentEngineTaskConfigOrDict] = None, + config: Optional[types.CreateRuntimeTaskConfigOrDict] = None, ) -> types.A2aTask: """ - Creates a new task in the Agent Engine. + Creates a new task in the Agent Runtime. Args: - name (str): Required. The name of the Agent Engine to create the task under. Format: + name (str): Required. The name of the Agent Runtime to create the task under. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. a2a_task_id (str): Required. The user ID of the task. context_id (str): Required. The ID of the context to use for the task. - config (CreateAgentEngineTaskConfig): - Optional. Additional configurations for creating the Agent Engine task. + config (CreateRuntimeTaskConfig): + Optional. Additional configurations for creating the Agent Runtime task. Returns: - A2aTask: The created Agent Engine task. + A2aTask: The created Agent Runtime task. """ - parameter_model = types._CreateAgentEngineTaskRequestParameters( + parameter_model = types._CreateRuntimeTaskRequestParameters( name=name, a2a_task_id=a2a_task_id, config=config, @@ -763,7 +737,7 @@ async def create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineTaskRequestParameters_to_vertex( + request_dict = _CreateRuntimeTaskRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -824,13 +798,13 @@ async def list( self, *, name: str, - config: Optional[types.ListAgentEngineTasksConfigOrDict] = None, + config: Optional[types.ListRuntimeTasksConfigOrDict] = None, ) -> AsyncPager[types.A2aTask]: - """Lists the A2A tasks of an Agent Engine. + """Lists the A2A tasks of an Agent Runtime. Args: name (str): - Required. The name of the agent engine to list tasks for. + Required. The name of the agent runtime to list tasks for. config (List): Optional. The configuration for the tasks to list. @@ -854,7 +828,7 @@ def events(self) -> "a2a_task_events_module.AsyncA2aTaskEvents": self._events = importlib.import_module(".a2a_task_events", __package__) except ImportError as e: raise ImportError( - "The 'agent_engines.a2a_tasks.events' module requires additional " + "The 'runtimes.a2a_tasks.events' module requires additional " "packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e diff --git a/agentplatform/_genai/agent_engines.py b/agentplatform/_genai/agent_engines.py deleted file mode 100644 index ef4d6f2a09..0000000000 --- a/agentplatform/_genai/agent_engines.py +++ /dev/null @@ -1,4115 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -# Code generated by the Google Gen AI SDK generator DO NOT EDIT. - -import builtins -import datetime -import importlib -import json -import logging -import typing -from typing import Any, AsyncIterator, Iterator, Optional, Sequence, Tuple, Union -from urllib.parse import urlencode -import warnings - -from google.genai import _api_module -from google.genai import _common -from google.genai import types as genai_types -from google.genai._common import get_value_by_path as getv -from google.genai._common import set_value_by_path as setv -from google.genai.pagers import Pager - -from . import _agent_engines_utils -from . import types - -if typing.TYPE_CHECKING: - from . import sessions as sessions_module - from . import memories as memories_module - from . import a2a_tasks as a2a_tasks_module - from . import runtimes as runtimes_module - - _ = sessions_module - __ = memories_module - ___ = a2a_tasks_module - ____ = runtimes_module - - -logger = logging.getLogger("agentplatform_genai.agentengines") - -logger.setLevel(logging.INFO) - - -def _AgentEngineOperation_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["name"], getv(from_object, ["name"])) - - if getv(from_object, ["metadata"]) is not None: - setv(to_object, ["metadata"], getv(from_object, ["metadata"])) - - if getv(from_object, ["done"]) is not None: - setv(to_object, ["done"], getv(from_object, ["done"])) - - if getv(from_object, ["error"]) is not None: - setv(to_object, ["error"], getv(from_object, ["error"])) - - if getv(from_object, ["response"]) is not None: - setv( - to_object, - ["response"], - _ReasoningEngine_from_vertex(getv(from_object, ["response"]), to_object), - ) - - return to_object - - -def _CancelQueryJobAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["operation_name"]) is not None: - setv(parent_object, ["operationName"], getv(from_object, ["operation_name"])) - - return to_object - - -def _CancelQueryJobAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - if getv(from_object, ["config"]) is not None: - setv( - to_object, - ["config"], - _CancelQueryJobAgentEngineConfig_to_vertex( - getv(from_object, ["config"]), to_object - ), - ) - - return to_object - - -def _CheckQueryJobAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["retrieve_result"]) is not None: - setv(parent_object, ["retrieveResult"], getv(from_object, ["retrieve_result"])) - - return to_object - - -def _CheckQueryJobAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - if getv(from_object, ["config"]) is not None: - setv( - to_object, - ["config"], - _CheckQueryJobAgentEngineConfig_to_vertex( - getv(from_object, ["config"]), to_object - ), - ) - - return to_object - - -def _CheckQueryJobResult_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(parent_object, ["operationName"]) is not None: - setv(to_object, ["operation_name"], getv(parent_object, ["operationName"])) - - if getv(parent_object, ["outputGcsUri"]) is not None: - setv(to_object, ["output_gcs_uri"], getv(parent_object, ["outputGcsUri"])) - - if getv(parent_object, ["status"]) is not None: - setv(to_object, ["status"], getv(parent_object, ["status"])) - - if getv(parent_object, ["result"]) is not None: - setv(to_object, ["result"], getv(parent_object, ["result"])) - - return to_object - - -def _CreateAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["display_name"]) is not None: - setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) - - if getv(from_object, ["description"]) is not None: - setv(parent_object, ["description"], getv(from_object, ["description"])) - - if getv(from_object, ["spec"]) is not None: - setv(parent_object, ["spec"], getv(from_object, ["spec"])) - - if getv(from_object, ["context_spec"]) is not None: - setv( - parent_object, - ["contextSpec"], - _ReasoningEngineContextSpec_to_vertex( - getv(from_object, ["context_spec"]), to_object - ), - ) - - if getv(from_object, ["psc_interface_config"]) is not None: - setv( - parent_object, - ["pscInterfaceConfig"], - getv(from_object, ["psc_interface_config"]), - ) - - if getv(from_object, ["encryption_spec"]) is not None: - setv(parent_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"])) - - if getv(from_object, ["labels"]) is not None: - setv(parent_object, ["labels"], getv(from_object, ["labels"])) - - if getv(from_object, ["source_packages"]) is not None: - setv(parent_object, ["sourcePackages"], getv(from_object, ["source_packages"])) - - if getv(from_object, ["entrypoint_module"]) is not None: - setv( - parent_object, - ["entrypointModule"], - getv(from_object, ["entrypoint_module"]), - ) - - if getv(from_object, ["entrypoint_object"]) is not None: - setv( - parent_object, - ["entrypointObject"], - getv(from_object, ["entrypoint_object"]), - ) - - if getv(from_object, ["requirements_file"]) is not None: - setv( - parent_object, - ["requirementsFile"], - getv(from_object, ["requirements_file"]), - ) - - if getv(from_object, ["agent_framework"]) is not None: - setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) - - if getv(from_object, ["python_version"]) is not None: - setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) - - if getv(from_object, ["agent_gateway_config"]) is not None: - setv( - parent_object, - ["agentGatewayConfig"], - getv(from_object, ["agent_gateway_config"]), - ) - - return to_object - - -def _CreateAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["config"]) is not None: - _CreateAgentEngineConfig_to_vertex(getv(from_object, ["config"]), to_object) - - return to_object - - -def _DeleteAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - if getv(from_object, ["force"]) is not None: - setv(to_object, ["force"], getv(from_object, ["force"])) - - return to_object - - -def _GetAgentEngineOperationParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["operation_name"]) is not None: - setv( - to_object, ["_url", "operationName"], getv(from_object, ["operation_name"]) - ) - - return to_object - - -def _GetAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - return to_object - - -def _ListAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["page_size"]) is not None: - setv(parent_object, ["_query", "pageSize"], getv(from_object, ["page_size"])) - - if getv(from_object, ["page_token"]) is not None: - setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"])) - - if getv(from_object, ["filter"]) is not None: - setv(parent_object, ["_query", "filter"], getv(from_object, ["filter"])) - - return to_object - - -def _ListAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["config"]) is not None: - _ListAgentEngineConfig_to_vertex(getv(from_object, ["config"]), to_object) - - return to_object - - -def _ListReasoningEnginesResponse_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["sdkHttpResponse"]) is not None: - setv(to_object, ["sdk_http_response"], getv(from_object, ["sdkHttpResponse"])) - - if getv(from_object, ["nextPageToken"]) is not None: - setv(to_object, ["next_page_token"], getv(from_object, ["nextPageToken"])) - - if getv(from_object, ["reasoningEngines"]) is not None: - setv( - to_object, - ["reasoning_engines"], - [ - _ReasoningEngine_from_vertex(item, to_object) - for item in getv(from_object, ["reasoningEngines"]) - ], - ) - - return to_object - - -def _QueryAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["class_method"]) is not None: - setv(parent_object, ["classMethod"], getv(from_object, ["class_method"])) - - if getv(from_object, ["input"]) is not None: - setv(parent_object, ["input"], getv(from_object, ["input"])) - - if getv(from_object, ["include_all_fields"]) is not None: - setv(to_object, ["includeAllFields"], getv(from_object, ["include_all_fields"])) - - return to_object - - -def _QueryAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - if getv(from_object, ["config"]) is not None: - _QueryAgentEngineConfig_to_vertex(getv(from_object, ["config"]), to_object) - - return to_object - - -def _ReasoningEngineContextSpecMemoryBankConfig_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["customizationConfigs"]) is not None: - setv( - to_object, - ["customization_configs"], - [item for item in getv(from_object, ["customizationConfigs"])], - ) - - if getv(from_object, ["disableMemoryRevisions"]) is not None: - setv( - to_object, - ["disable_memory_revisions"], - getv(from_object, ["disableMemoryRevisions"]), - ) - - if getv(from_object, ["generationConfig"]) is not None: - setv(to_object, ["generation_config"], getv(from_object, ["generationConfig"])) - - if getv(from_object, ["similaritySearchConfig"]) is not None: - setv( - to_object, - ["similarity_search_config"], - getv(from_object, ["similaritySearchConfig"]), - ) - - if getv(from_object, ["ttlConfig"]) is not None: - setv(to_object, ["ttl_config"], getv(from_object, ["ttlConfig"])) - - if getv(from_object, ["structuredMemoryConfigs"]) is not None: - setv( - to_object, - ["structured_memory_configs"], - [ - _StructuredMemoryConfig_from_vertex(item, to_object) - for item in getv(from_object, ["structuredMemoryConfigs"]) - ], - ) - - return to_object - - -def _ReasoningEngineContextSpecMemoryBankConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["customization_configs"]) is not None: - setv( - to_object, - ["customizationConfigs"], - [item for item in getv(from_object, ["customization_configs"])], - ) - - if getv(from_object, ["disable_memory_revisions"]) is not None: - setv( - to_object, - ["disableMemoryRevisions"], - getv(from_object, ["disable_memory_revisions"]), - ) - - if getv(from_object, ["generation_config"]) is not None: - setv(to_object, ["generationConfig"], getv(from_object, ["generation_config"])) - - if getv(from_object, ["similarity_search_config"]) is not None: - setv( - to_object, - ["similaritySearchConfig"], - getv(from_object, ["similarity_search_config"]), - ) - - if getv(from_object, ["ttl_config"]) is not None: - setv(to_object, ["ttlConfig"], getv(from_object, ["ttl_config"])) - - if getv(from_object, ["structured_memory_configs"]) is not None: - setv( - to_object, - ["structuredMemoryConfigs"], - [ - _StructuredMemoryConfig_to_vertex(item, to_object) - for item in getv(from_object, ["structured_memory_configs"]) - ], - ) - - return to_object - - -def _ReasoningEngineContextSpec_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["memoryBankConfig"]) is not None: - setv( - to_object, - ["memory_bank_config"], - _ReasoningEngineContextSpecMemoryBankConfig_from_vertex( - getv(from_object, ["memoryBankConfig"]), to_object - ), - ) - - return to_object - - -def _ReasoningEngineContextSpec_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["memory_bank_config"]) is not None: - setv( - to_object, - ["memoryBankConfig"], - _ReasoningEngineContextSpecMemoryBankConfig_to_vertex( - getv(from_object, ["memory_bank_config"]), to_object - ), - ) - - return to_object - - -def _ReasoningEngine_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["encryptionSpec"]) is not None: - setv(to_object, ["encryption_spec"], getv(from_object, ["encryptionSpec"])) - - if getv(from_object, ["contextSpec"]) is not None: - setv( - to_object, - ["context_spec"], - _ReasoningEngineContextSpec_from_vertex( - getv(from_object, ["contextSpec"]), to_object - ), - ) - - if getv(from_object, ["createTime"]) is not None: - setv(to_object, ["create_time"], getv(from_object, ["createTime"])) - - if getv(from_object, ["description"]) is not None: - setv(to_object, ["description"], getv(from_object, ["description"])) - - if getv(from_object, ["displayName"]) is not None: - setv(to_object, ["display_name"], getv(from_object, ["displayName"])) - - if getv(from_object, ["etag"]) is not None: - setv(to_object, ["etag"], getv(from_object, ["etag"])) - - if getv(from_object, ["labels"]) is not None: - setv(to_object, ["labels"], getv(from_object, ["labels"])) - - if getv(from_object, ["name"]) is not None: - setv(to_object, ["name"], getv(from_object, ["name"])) - - if getv(from_object, ["spec"]) is not None: - setv(to_object, ["spec"], getv(from_object, ["spec"])) - - if getv(from_object, ["updateTime"]) is not None: - setv(to_object, ["update_time"], getv(from_object, ["updateTime"])) - - if getv(from_object, ["trafficConfig"]) is not None: - setv(to_object, ["traffic_config"], getv(from_object, ["trafficConfig"])) - - return to_object - - -def _RunQueryJobAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["input_gcs_uri"]) is not None: - setv(parent_object, ["inputGcsUri"], getv(from_object, ["input_gcs_uri"])) - - if getv(from_object, ["output_gcs_uri"]) is not None: - setv(parent_object, ["outputGcsUri"], getv(from_object, ["output_gcs_uri"])) - - return to_object - - -def _RunQueryJobAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - if getv(from_object, ["config"]) is not None: - setv( - to_object, - ["config"], - _RunQueryJobAgentEngineConfig_to_vertex( - getv(from_object, ["config"]), to_object - ), - ) - - return to_object - - -def _StructuredMemoryConfig_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["schemaConfigs"]) is not None: - setv( - to_object, - ["schema_configs"], - [ - _StructuredMemorySchemaConfig_from_vertex(item, to_object) - for item in getv(from_object, ["schemaConfigs"]) - ], - ) - - if getv(from_object, ["scopeKeys"]) is not None: - setv(to_object, ["scope_keys"], getv(from_object, ["scopeKeys"])) - - return to_object - - -def _StructuredMemoryConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["schema_configs"]) is not None: - setv( - to_object, - ["schemaConfigs"], - [ - _StructuredMemorySchemaConfig_to_vertex(item, to_object) - for item in getv(from_object, ["schema_configs"]) - ], - ) - - if getv(from_object, ["scope_keys"]) is not None: - setv(to_object, ["scopeKeys"], getv(from_object, ["scope_keys"])) - - return to_object - - -def _StructuredMemorySchemaConfig_from_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["schema"]) is not None: - setv(to_object, ["memory_schema"], getv(from_object, ["schema"])) - - if getv(from_object, ["id"]) is not None: - setv(to_object, ["id"], getv(from_object, ["id"])) - - if getv(from_object, ["memoryType"]) is not None: - setv(to_object, ["memory_type"], getv(from_object, ["memoryType"])) - - return to_object - - -def _StructuredMemorySchemaConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["memory_schema"]) is not None: - setv(to_object, ["schema"], getv(from_object, ["memory_schema"])) - - if getv(from_object, ["id"]) is not None: - setv(to_object, ["id"], getv(from_object, ["id"])) - - if getv(from_object, ["memory_type"]) is not None: - setv(to_object, ["memoryType"], getv(from_object, ["memory_type"])) - - return to_object - - -def _UpdateAgentEngineConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["display_name"]) is not None: - setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) - - if getv(from_object, ["description"]) is not None: - setv(parent_object, ["description"], getv(from_object, ["description"])) - - if getv(from_object, ["spec"]) is not None: - setv(parent_object, ["spec"], getv(from_object, ["spec"])) - - if getv(from_object, ["context_spec"]) is not None: - setv( - parent_object, - ["contextSpec"], - _ReasoningEngineContextSpec_to_vertex( - getv(from_object, ["context_spec"]), to_object - ), - ) - - if getv(from_object, ["psc_interface_config"]) is not None: - setv( - parent_object, - ["pscInterfaceConfig"], - getv(from_object, ["psc_interface_config"]), - ) - - if getv(from_object, ["encryption_spec"]) is not None: - setv(parent_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"])) - - if getv(from_object, ["labels"]) is not None: - setv(parent_object, ["labels"], getv(from_object, ["labels"])) - - if getv(from_object, ["source_packages"]) is not None: - setv(parent_object, ["sourcePackages"], getv(from_object, ["source_packages"])) - - if getv(from_object, ["entrypoint_module"]) is not None: - setv( - parent_object, - ["entrypointModule"], - getv(from_object, ["entrypoint_module"]), - ) - - if getv(from_object, ["entrypoint_object"]) is not None: - setv( - parent_object, - ["entrypointObject"], - getv(from_object, ["entrypoint_object"]), - ) - - if getv(from_object, ["requirements_file"]) is not None: - setv( - parent_object, - ["requirementsFile"], - getv(from_object, ["requirements_file"]), - ) - - if getv(from_object, ["agent_framework"]) is not None: - setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) - - if getv(from_object, ["python_version"]) is not None: - setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) - - if getv(from_object, ["agent_gateway_config"]) is not None: - setv( - parent_object, - ["agentGatewayConfig"], - getv(from_object, ["agent_gateway_config"]), - ) - - if getv(from_object, ["update_mask"]) is not None: - setv( - parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"]) - ) - - if getv(from_object, ["traffic_config"]) is not None: - setv(parent_object, ["trafficConfig"], getv(from_object, ["traffic_config"])) - - return to_object - - -def _UpdateAgentEngineRequestParameters_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - - if getv(from_object, ["config"]) is not None: - _UpdateAgentEngineConfig_to_vertex(getv(from_object, ["config"]), to_object) - - return to_object - - -class AgentEngines(_api_module.BaseModule): - - def cancel_query_job( - self, - *, - name: str, - config: Optional[types.CancelQueryJobAgentEngineConfigOrDict] = None, - ) -> types.CancelQueryJobResult: - """ - Cancels a long-running query job on an Agent Engine. - - Args: - name (str): - Required. The reasoning engine resource name. - config (CancelQueryJobAgentEngineConfigOrDict): - Optional. The configuration for the cancel_query_job. - - """ - - parameter_model = types._CancelQueryJobAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _CancelQueryJobAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:cancelAsyncQuery".format_map(request_url_dict) - else: - path = "{name}:cancelAsyncQuery" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("post", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.CancelQueryJobResult._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _check_query_job( - self, - *, - name: str, - config: Optional[types.CheckQueryJobAgentEngineConfigOrDict] = None, - ) -> types.CheckQueryJobResult: - """ - Query an Agent Engine asynchronously. - """ - - parameter_model = types._CheckQueryJobAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _CheckQueryJobAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:checkQueryJob".format_map(request_url_dict) - else: - path = "{name}:checkQueryJob" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("post", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _CheckQueryJobResult_from_vertex(response_dict) - - return_value = types.CheckQueryJobResult._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _run_query_job( - self, - *, - name: str, - config: Optional[types._RunQueryJobAgentEngineConfigOrDict] = None, - ) -> types.AgentEngineOperation: - """ - Run a query job on an agent engine. - """ - - parameter_model = types._RunQueryJobAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _RunQueryJobAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:asyncQuery".format_map(request_url_dict) - else: - path = "{name}:asyncQuery" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("post", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _create( - self, *, config: Optional[types.CreateAgentEngineConfigOrDict] = None - ) -> types.AgentEngineOperation: - """ - Creates a new Agent Engine. - """ - - parameter_model = types._CreateAgentEngineRequestParameters( - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _CreateAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "reasoningEngines".format_map(request_url_dict) - else: - path = "reasoningEngines" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("post", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _delete( - self, - *, - name: str, - force: Optional[bool] = None, - config: Optional[types.DeleteAgentEngineConfigOrDict] = None, - ) -> types.DeleteAgentEngineOperation: - """ - Delete an Agent Engine resource. - - Args: - name (str): - Required. The name of the Agent Engine to be deleted. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` - or `reasoningEngines/{resource_id}`. - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error when - the Agent Engine has undeleted child resources. Defaults to False. - config (DeleteAgentEngineConfig): - Optional. Additional configurations for deleting the Agent Engine. - - """ - - parameter_model = types._DeleteAgentEngineRequestParameters( - name=name, - force=force, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _DeleteAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}".format_map(request_url_dict) - else: - path = "{name}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("delete", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.DeleteAgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _get( - self, *, name: str, config: Optional[types.GetAgentEngineConfigOrDict] = None - ) -> types.ReasoningEngine: - """ - Get an Agent Engine instance. - """ - - parameter_model = types._GetAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _GetAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}".format_map(request_url_dict) - else: - path = "{name}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("get", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _ReasoningEngine_from_vertex(response_dict) - - return_value = types.ReasoningEngine._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _list( - self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None - ) -> types.ListReasoningEnginesResponse: - """ - Lists Agent Engines. - """ - - parameter_model = types._ListAgentEngineRequestParameters( - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _ListAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "reasoningEngines".format_map(request_url_dict) - else: - path = "reasoningEngines" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("get", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) - - return_value = types.ListReasoningEnginesResponse._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _get_agent_operation( - self, - *, - operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineOperation: - parameter_model = types._GetAgentEngineOperationParameters( - operation_name=operation_name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _GetAgentEngineOperationParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{operationName}".format_map(request_url_dict) - else: - path = "{operationName}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("get", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _query( - self, *, name: str, config: Optional[types.QueryAgentEngineConfigOrDict] = None - ) -> types.QueryReasoningEngineResponse: - """ - Query an Agent Engine. - """ - - parameter_model = types._QueryAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _QueryAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:query".format_map(request_url_dict) - else: - path = "{name}:query" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("post", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.QueryReasoningEngineResponse._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - def _update( - self, *, name: str, config: Optional[types.UpdateAgentEngineConfigOrDict] = None - ) -> types.AgentEngineOperation: - """ - Updates an Agent Engine. - """ - - parameter_model = types._UpdateAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _UpdateAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}".format_map(request_url_dict) - else: - path = "{name}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = self._api_client.request("patch", path, request_dict, http_options) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - _a2a_tasks = None - _memories = None - _sandboxes = None - _sessions = None - _runtimes = None - - @property - def runtimes(self) -> "runtimes_module.Runtimes": - if self._runtimes is None: - try: - # We need to lazy load the runtimes module to handle the - # possibility of ImportError when dependencies are not installed. - self._runtimes = importlib.import_module(".runtimes", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.runtimes' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._runtimes.Runtimes(self._api_client) # type: ignore[no-any-return] - - @property - def a2a_tasks(self) -> "a2a_tasks_module.A2aTasks": - if self._a2a_tasks is None: - try: - # We need to lazy load the a2a_tasks module to handle the - # possibility of ImportError when dependencies are not installed. - self._a2a_tasks = importlib.import_module(".a2a_tasks", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.a2a_tasks' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._a2a_tasks.A2aTasks(self._api_client) # type: ignore[no-any-return] - - @property - def memories(self) -> "memories_module.Memories": - if self._memories is None: - try: - # We need to lazy load the memories module to handle the - # possibility of ImportError when dependencies are not installed. - self._memories = importlib.import_module(".memories", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.memories' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._memories.Memories(self._api_client) # type: ignore[no-any-return] - - @property - def sandboxes(self) -> Any: - if self._sandboxes is None: - try: - # We need to lazy load the sandboxes module to handle the - # possibility of ImportError when dependencies are not installed. - self._sandboxes = importlib.import_module(".sandboxes", __package__) - except ImportError as e: - raise ImportError( - "The agent_engines.sandboxes module requires additional packages. " - "Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._sandboxes.Sandboxes(self._api_client) - - @property - def sessions(self) -> "sessions_module.Sessions": - if self._sessions is None: - try: - # We need to lazy load the sessions module to handle the - # possibility of ImportError when dependencies are not installed. - self._sessions = importlib.import_module(".sessions", __package__) - except ImportError as e: - raise ImportError( - "The agent_engines.sessions module requires additional packages. " - "Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._sessions.Sessions(self._api_client) # type: ignore[no-any-return] - - def _list_pager( - self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None - ) -> Pager[types.ReasoningEngine]: - return Pager( - "reasoning_engines", - self._list, - self._list(config=config), - config, - ) - - def check_query_job( - self, - *, - name: str, - config: Optional[types.CheckQueryJobAgentEngineConfigOrDict] = None, - ) -> types.CheckQueryJobResult: - """Checks a query job on an agent engine and optionally returns the results. - - Args: - name (str): - Required. A fully-qualified resource name or ID. - config (CheckQueryJobAgentEngineConfigOrDict): - Optional. The configuration for the check_query_job. If not provided, - the default configuration will be used. This can be used to specify - the following fields: - - retrieve_result: Whether to retrieve the results of the query job. - """ - from google.cloud import storage # type: ignore[attr-defined] - import json - - if config is None: - config = types.CheckQueryJobAgentEngineConfig() - elif isinstance(config, dict): - config = types.CheckQueryJobAgentEngineConfig(**config) - - raw_response = self._api_client.request("get", name, {}) - if hasattr(raw_response, "body"): - operation = ( - json.loads(raw_response.body) - if isinstance(raw_response.body, str) - else raw_response.body - ) - else: - operation = raw_response - - status = "RUNNING" - if isinstance(operation, dict): - if operation.get("done"): - status = "FAILED" if operation.get("error") else "SUCCESS" - - response_dict = operation.get("response", {}) - output_gcs_uri = response_dict.get("outputGcsUri") or response_dict.get( - "output_gcs_uri" - ) - error = operation.get("error") - else: - if getattr(operation, "done", False): - status = "FAILED" if getattr(operation, "error", None) else "SUCCESS" - - response_obj = getattr(operation, "response", None) - if isinstance(response_obj, dict): - output_gcs_uri = response_obj.get("outputGcsUri") or response_obj.get( - "output_gcs_uri" - ) - else: - output_gcs_uri = ( - getattr( - response_obj, - "output_gcs_uri", - getattr(response_obj, "outputGcsUri", None), - ) - if response_obj - else None - ) - error = getattr(operation, "error", None) - - result_str = None - if status == "SUCCESS" and config.retrieve_result and output_gcs_uri: - storage_client = storage.Client( - project=self._api_client.project, - credentials=self._api_client._credentials, - ) - bucket_name = output_gcs_uri.replace("gs://", "").split("/")[0] - blob_name = output_gcs_uri.replace(f"gs://{bucket_name}/", "") - bucket = storage_client.bucket(bucket_name) - blob = bucket.blob(blob_name) - if blob.exists(): - result_str = blob.download_as_string().decode("utf-8") - else: - raise ValueError( - f"Failed to retrieve blob results for {output_gcs_uri}" - ) - - elif status == "FAILED" and error: - result_str = str(error) - - return types.CheckQueryJobResult( - operation_name=name, - output_gcs_uri=output_gcs_uri, - status=status, - result=result_str, - ) - - def _is_lightweight_creation( - self, agent: Any, config: types.AgentEngineConfig - ) -> bool: - if ( - agent - or config.source_packages - or config.developer_connect_source - or config.agent_config_source - or config.container_spec - ): - return False - return True - - def run_query_job( - self, - *, - name: str, - config: Optional[types.RunQueryJobAgentEngineConfigOrDict] = None, - ) -> types.RunQueryJobResult: - """Launches a long-running query job on an Agent Engine - - Args: - name (str): - Required. A fully-qualified resource name or ID. - config (RunQueryJobAgentEngineConfigOrDict): - Optional. The configuration for the async query. If not provided, - the default configuration will be used. This can be used to specify - the following fields: - - query: The query to send to the agent engine. - - output_gcs_uri: The GCS URI to use for the output. - """ - from google.cloud import storage # type: ignore[attr-defined] - from google.api_core import exceptions - import uuid - - if config is None: - config = types.RunQueryJobAgentEngineConfig() - elif isinstance(config, dict): - config = types.RunQueryJobAgentEngineConfig(**config) - - if not config.query: - raise ValueError("`query` is required in the config object.") - if not config.output_gcs_uri: - raise ValueError("`output_gcs_uri` is required in the config object.") - - output_gcs_uri = config.output_gcs_uri - is_file = False - last_part = "" - if not output_gcs_uri.endswith("/"): - last_part = output_gcs_uri.split("/")[-1] - if "." in last_part: - is_file = True - - if is_file: - path_parts = output_gcs_uri.split("/") - file_name = path_parts[-1] - base_uri = "/".join(path_parts[:-1]) - name_parts = file_name.rsplit(".", 1) - if len(name_parts) == 2: - name_part, ext = name_parts[0], "." + name_parts[1] - else: - name_part = name_parts[0] - ext = "" - input_gcs_uri = f"{base_uri}/{name_part}_input{ext}" - else: - job_uuid = uuid.uuid4().hex - gcs_path = output_gcs_uri.rstrip("/") - input_gcs_uri = f"{gcs_path}/{job_uuid}_input.json" - output_gcs_uri = f"{gcs_path}/{job_uuid}_output.json" - - storage_client = storage.Client( - project=self._api_client.project, credentials=self._api_client._credentials - ) - - # Handle creating the bucket if it does not exist - bucket_name = config.output_gcs_uri.replace("gs://", "").split("/")[0] - bucket = storage_client.bucket(bucket_name) - - try: - bucket_exists = bucket.exists() - except exceptions.Forbidden as e: - raise ValueError( - f"Permission denied to check existence of bucket '{bucket_name}'. " - "The service account may lack 'storage.buckets.get' permission." - ) from e - - if not bucket_exists: - try: - bucket.create() - except exceptions.Forbidden as e: - raise ValueError( - f"Permission denied to create bucket '{bucket_name}'. " - "The service account may lack 'storage.buckets.create' permission." - ) from e - - input_blob_name = input_gcs_uri.replace(f"gs://{bucket_name}/", "") - blob = bucket.blob(input_blob_name) - blob.upload_from_string(config.query) - - new_config = types._RunQueryJobAgentEngineConfig( - input_gcs_uri=input_gcs_uri, - output_gcs_uri=output_gcs_uri, - ) - - # Proceed with sending the async query via the auto-generated method - operation = self._run_query_job(name=name, config=new_config) - - return types.RunQueryJobResult( - job_name=operation.name, - input_gcs_uri=input_gcs_uri, - output_gcs_uri=output_gcs_uri, - ) - - def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineConfigOrDict] = None, - ) -> types.AgentEngine: - """Gets an agent engine. - - Args: - name (str): - Required. A fully-qualified resource name or ID such as - "projects/123/locations/us-central1/reasoningEngines/456" or - a shortened name such as "reasoningEngines/456". - """ - api_resource = self._get(name=name, config=config) - agent_engine = types.AgentEngine( - api_client=self, - api_async_client=AsyncAgentEngines(api_client_=self._api_client), - api_resource=api_resource, - ) - if api_resource.spec: - self._register_api_methods(agent_engine=agent_engine) - return agent_engine - - def delete( - self, - *, - name: str, - force: Optional[bool] = None, - config: Optional[types.DeleteAgentEngineConfigOrDict] = None, - ) -> types.DeleteAgentEngineOperation: - """ - Delete an Agent Engine resource. - - Args: - name (str): - Required. The name of the Agent Engine to be deleted. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` - or `reasoningEngines/{resource_id}`. - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error when - the Agent Engine has undeleted child resources. Defaults to False. - config (DeleteAgentEngineConfig): - Optional. Additional configurations for deleting the Agent Engine. - - """ - logger.info(f"Deleting AgentEngine resource: {name}") - operation = self._delete(name=name, force=force, config=config) - logger.info(f"Started AgentEngine delete operation: {operation.name}") - return operation - - def create( - self, - *, - agent_engine: Any = None, - agent: Any = None, - config: Optional[types.AgentEngineConfigOrDict] = None, - ) -> types.AgentEngine: - """Creates an agent engine. - - The Agent Engine will be an instance of the `agent_engine` that - was passed in, running remotely on Vertex AI. - - Sample ``src_dir`` contents (e.g. ``./user_src_dir``): - - .. code-block:: python - - user_src_dir/ - |-- main.py - |-- requirements.txt - |-- user_code/ - | |-- utils.py - | |-- ... - |-- ... - - To build an Agent Engine with the above files, run: - - .. code-block:: python - - client = agentplatform.Client( - project="your-project", - location="us-central1", - ) - remote_agent = client.agent_engines.create( - agent=local_agent, - config=dict( - requirements=[ - # I.e. the PyPI dependencies listed in requirements.txt - "google-cloud-aiplatform[agent_engines,adk]", - ... - ], - extra_packages=[ - "./user_src_dir/main.py", # a single file - "./user_src_dir/user_code", # a directory - ... - ], - ), - ) - - Args: - agent (Any): - Optional. The Agent to be created. If not specified, this will - correspond to a lightweight instance that cannot be queried - (but can be updated to future instances that can be queried). - agent_engine (Any): - Optional. This is deprecated. Please use `agent` instead. - config (AgentEngineConfig): - Optional. The configurations to use for creating the Agent Engine. - - Returns: - AgentEngine: The created Agent Engine instance. - - Raises: - ValueError: If the `project` was not set using `client.Client`. - ValueError: If the `location` was not set using `client.Client`. - ValueError: If `config.staging_bucket` was not set when `agent` - is specified. - ValueError: If `config.staging_bucket` does not start with "gs://". - ValueError: If `config.extra_packages` is specified but `agent` - is None. - ValueError: If `config.requirements` is specified but `agent` is None. - ValueError: If `config.env_vars` has a dictionary entry that does not - correspond to an environment variable value or a SecretRef. - TypeError: If `config.env_vars` is not a dictionary. - FileNotFoundError: If `config.extra_packages` includes a file or - directory that does not exist. - IOError: If ``config.requirements` is a string that corresponds to a - nonexistent file. - """ - if config is None: - config = {} - if isinstance(config, dict): - config = types.AgentEngineConfig.model_validate(config) - elif not isinstance(config, types.AgentEngineConfig): - raise TypeError( - f"config must be a dict or AgentEngineConfig, but got {type(config)}." - ) - context_spec = config.context_spec - if context_spec is not None: - # Conversion to a dict for _create_config - context_spec = json.loads(context_spec.model_dump_json()) - developer_connect_source = config.developer_connect_source - if developer_connect_source is not None: - developer_connect_source = json.loads( - developer_connect_source.model_dump_json() - ) - agent_config_source = config.agent_config_source - if agent_config_source is not None: - agent_config_source = json.loads(agent_config_source.model_dump_json()) - keep_alive_probe = config.keep_alive_probe - if keep_alive_probe is not None: - keep_alive_probe = json.loads( - keep_alive_probe.model_dump_json(exclude_none=True) - ) - if agent and agent_engine: - raise ValueError("Please specify only one of `agent` or `agent_engine`.") - elif agent_engine: - raise DeprecationWarning( - "The `agent_engine` argument is deprecated. Please use `agent` instead." - ) - agent = agent or agent_engine - api_config = self._create_config( - mode="create", - agent=agent, - identity_type=config.identity_type, - staging_bucket=config.staging_bucket, - requirements=config.requirements, - display_name=config.display_name, - description=config.description, - gcs_dir_name=config.gcs_dir_name, - extra_packages=config.extra_packages, - env_vars=config.env_vars, - service_account=config.service_account, - context_spec=context_spec, - psc_interface_config=config.psc_interface_config, - agent_gateway_config=config.agent_gateway_config, - min_instances=config.min_instances, - max_instances=config.max_instances, - resource_limits=config.resource_limits, - container_concurrency=config.container_concurrency, - encryption_spec=config.encryption_spec, - agent_server_mode=config.agent_server_mode, - labels=config.labels, - class_methods=config.class_methods, - source_packages=config.source_packages, - developer_connect_source=developer_connect_source, - entrypoint_module=config.entrypoint_module, - entrypoint_object=config.entrypoint_object, - requirements_file=config.requirements_file, - agent_framework=config.agent_framework, - python_version=config.python_version, - build_options=config.build_options, - image_spec=config.image_spec, - agent_config_source=agent_config_source, - container_spec=config.container_spec, - keep_alive_probe=keep_alive_probe, - ) - operation = self._create(config=api_config) - reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( - operation_name=operation.name - ) - logger.info( - "View progress and logs at https://console.cloud.google.com/logs/query?" - f"project={self._api_client.project}" - "&query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22%0A" - f"resource.labels.reasoning_engine_id%3D%22{reasoning_engine_id}%22." - ) - if not self._is_lightweight_creation(agent, config): - poll_interval_seconds = 10 - else: - poll_interval_seconds = 1 # Lightweight agent engine resource creation. - operation = _agent_engines_utils._await_operation( - operation_name=operation.name, - get_operation_fn=self._get_agent_operation, - poll_interval_seconds=poll_interval_seconds, - ) - - agent_engine = types.AgentEngine( - api_client=self, - api_async_client=AsyncAgentEngines(api_client_=self._api_client), - api_resource=operation.response, - ) - if agent_engine.api_resource: - logger.info("Agent Engine created. To use it in another session:") - logger.info( - f"agent_engine=client.agent_engines.get(name='{agent_engine.api_resource.name}')" - ) - elif operation.error: - raise RuntimeError(f"Failed to create Agent Engine: {operation.error}") - else: - logger.warning("The operation returned an empty response.") - if not self._is_lightweight_creation(agent, config): - # If the user did not provide an agent_engine (e.g. lightweight - # provisioning), it will not have any API methods registered. - agent_engine = self._register_api_methods(agent_engine=agent_engine) - return agent_engine # type: ignore[no-any-return] - - def _set_source_code_spec( - self, - *, - spec: types.ReasoningEngineSpecDict, - update_masks: builtins.list[str], - source_packages: Optional[Sequence[str]] = None, - developer_connect_source: Optional[ - types.ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict - ] = None, - class_methods: Optional[Sequence[dict[str, Any]]] = None, - entrypoint_module: Optional[str] = None, - entrypoint_object: Optional[str] = None, - requirements_file: Optional[str] = None, - sys_version: str, - build_options: Optional[dict[str, builtins.list[str]]] = None, - image_spec: Optional[ - types.ReasoningEngineSpecSourceCodeSpecImageSpecDict - ] = None, - agent_config_source: Optional[ - types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict - ] = None, - ) -> None: - """Sets source_code_spec for agent engine inside the `spec`.""" - source_code_spec = types.ReasoningEngineSpecSourceCodeSpecDict() - if source_packages and not agent_config_source: - source_packages = _agent_engines_utils._validate_packages_or_raise( - packages=source_packages, - build_options=build_options, - ) - update_masks.append("spec.source_code_spec.inline_source.source_archive") - source_code_spec["inline_source"] = { # type: ignore[typeddict-item] - "source_archive": _agent_engines_utils._create_base64_encoded_tarball( - source_packages=source_packages - ) - } - elif developer_connect_source: - update_masks.append("spec.source_code_spec.developer_connect_source") - source_code_spec["developer_connect_source"] = { - "config": developer_connect_source - } - elif not agent_config_source: - raise ValueError( - "Please specify one of `source_packages`, `developer_connect_source`, " - "or `agent_config_source`." - ) - if class_methods is not None: - update_masks.append("spec.class_methods") - class_methods_spec_list = ( - _agent_engines_utils._class_methods_to_class_methods_spec( - class_methods=class_methods - ) - ) - spec["class_methods"] = [ - _agent_engines_utils._to_dict(class_method_spec) - for class_method_spec in class_methods_spec_list - ] - elif image_spec is None: - raise ValueError( - "`class_methods` must be specified if `source_packages`, " - "`developer_connect_source`, or `agent_config_source` is " - "specified without a Dockerfile or `image_spec`." - ) - if image_spec is not None: - if entrypoint_module or entrypoint_object or requirements_file: - raise ValueError( - "`image_spec` cannot be specified alongside `entrypoint_module`, " - "`entrypoint_object`, or `requirements_file`, as they are " - "mutually exclusive." - ) - if agent_config_source: - raise ValueError( - "`image_spec` cannot be specified alongside `agent_config_source`, " - "as they are mutually exclusive." - ) - update_masks.append("spec.source_code_spec.image_spec") - source_code_spec["image_spec"] = image_spec - spec["source_code_spec"] = source_code_spec - return - - update_masks.append("spec.source_code_spec.python_spec.version") - python_spec: types.ReasoningEngineSpecSourceCodeSpecPythonSpecDict = { - "version": sys_version, - } - if agent_config_source is not None: - if entrypoint_module or entrypoint_object: - logger.warning( - "`entrypoint_module` and `entrypoint_object` are ignored when " - "`agent_config_source` is specified, as they are pre-defined." - ) - if source_packages: - source_packages = _agent_engines_utils._validate_packages_or_raise( - packages=source_packages, - build_options=build_options, - ) - update_masks.append( - "spec.source_code_spec.agent_config_source.inline_source.source_archive" - ) - agent_config_source["inline_source"] = { # type: ignore[typeddict-item] - "source_archive": _agent_engines_utils._create_base64_encoded_tarball( - source_packages=source_packages - ) - } - update_masks.append("spec.source_code_spec.agent_config_source") - source_code_spec["agent_config_source"] = agent_config_source - - if requirements_file is not None: - update_masks.append( - "spec.source_code_spec.python_spec.requirements_file" - ) - python_spec["requirements_file"] = requirements_file - source_code_spec["python_spec"] = python_spec - - spec["source_code_spec"] = source_code_spec - return - - if not entrypoint_module: - raise ValueError( - "`entrypoint_module` must be specified if `source_packages` or `developer_connect_source` is specified." - ) - update_masks.append("spec.source_code_spec.python_spec.entrypoint_module") - python_spec["entrypoint_module"] = entrypoint_module - if not entrypoint_object: - raise ValueError( - "`entrypoint_object` must be specified if `source_packages` or `developer_connect_source` is specified." - ) - update_masks.append("spec.source_code_spec.python_spec.entrypoint_object") - python_spec["entrypoint_object"] = entrypoint_object - if requirements_file is not None: - update_masks.append("spec.source_code_spec.python_spec.requirements_file") - python_spec["requirements_file"] = requirements_file - source_code_spec["python_spec"] = python_spec - spec["source_code_spec"] = source_code_spec - - def _set_package_spec( - self, - *, - spec: types.ReasoningEngineSpecDict, - update_masks: builtins.list[str], - agent: Any, - staging_bucket: Optional[str] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - class_methods: Optional[Sequence[dict[str, Any]]] = None, - sys_version: str, - build_options: Optional[dict[str, builtins.list[str]]] = None, - ) -> None: - """Sets package spec for agent engine.""" - project = self._api_client.project - if project is None: - raise ValueError("project must be set using `agentplatform.Client`.") - location = self._api_client.location - if location is None: - raise ValueError("location must be set using `agentplatform.Client`.") - gcs_dir_name = gcs_dir_name or _agent_engines_utils._DEFAULT_GCS_DIR_NAME - staging_bucket = _agent_engines_utils._validate_staging_bucket_or_raise( - staging_bucket=staging_bucket, - ) - requirements = _agent_engines_utils._validate_requirements_or_raise( - agent=agent, - requirements=requirements, - ) - extra_packages = _agent_engines_utils._validate_packages_or_raise( - packages=extra_packages, - build_options=build_options, - ) - # Prepares the Agent Engine for creation/update in Vertex AI. This - # involves packaging and uploading the artifacts for agent_engine, - # requirements and extra_packages to `staging_bucket/gcs_dir_name`. - _agent_engines_utils._prepare( - agent=agent, - requirements=requirements, - project=project, - location=location, - staging_bucket=staging_bucket, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - credentials=self._api_client._credentials, - ) - # Update the package spec. - update_masks.append("spec.package_spec.pickle_object_gcs_uri") - package_spec: types.ReasoningEngineSpecPackageSpecDict = { - "python_version": sys_version, - "pickle_object_gcs_uri": "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _agent_engines_utils._BLOB_FILENAME, - ), - } - if extra_packages: - update_masks.append("spec.package_spec.dependency_files_gcs_uri") - package_spec["dependency_files_gcs_uri"] = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _agent_engines_utils._EXTRA_PACKAGES_FILE, - ) - if requirements: - update_masks.append("spec.package_spec.requirements_gcs_uri") - package_spec["requirements_gcs_uri"] = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _agent_engines_utils._REQUIREMENTS_FILE, - ) - spec["package_spec"] = package_spec - - update_masks.append("spec.class_methods") - if class_methods is not None: - class_methods_spec_list = ( - _agent_engines_utils._class_methods_to_class_methods_spec( - class_methods=class_methods - ) - ) - else: - class_methods_spec_list = ( - _agent_engines_utils._generate_class_methods_spec_or_raise( - agent=agent, - operations=_agent_engines_utils._get_registered_operations( - agent=agent - ), - ) - ) - spec["class_methods"] = [ - _agent_engines_utils._to_dict(class_method_spec) - for class_method_spec in class_methods_spec_list - ] - - def _create_config( - self, - *, - mode: str, - agent: Any = None, - identity_type: Optional[types.IdentityType] = None, - staging_bucket: Optional[str] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[dict[str, Union[str, Any]]] = None, - service_account: Optional[str] = None, - context_spec: Optional[types.ReasoningEngineContextSpecDict] = None, - psc_interface_config: Optional[types.PscInterfaceConfigDict] = None, - agent_gateway_config: Optional[ - types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict - ] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[genai_types.EncryptionSpecDict] = None, - labels: Optional[dict[str, str]] = None, - agent_server_mode: Optional[types.AgentServerMode] = None, - class_methods: Optional[Sequence[dict[str, Any]]] = None, - source_packages: Optional[Sequence[str]] = None, - developer_connect_source: Optional[ - types.ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict - ] = None, - entrypoint_module: Optional[str] = None, - entrypoint_object: Optional[str] = None, - requirements_file: Optional[str] = None, - agent_framework: Optional[str] = None, - python_version: Optional[str] = None, - build_options: Optional[dict[str, builtins.list[str]]] = None, - image_spec: Optional[ - types.ReasoningEngineSpecSourceCodeSpecImageSpecDict - ] = None, - agent_config_source: Optional[ - types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict - ] = None, - container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None, - keep_alive_probe: Optional[dict[str, Any]] = None, - traffic_config: Optional[types.ReasoningEngineTrafficConfigDict] = None, - ) -> types.UpdateAgentEngineConfigDict: - import sys - - config: types.UpdateAgentEngineConfigDict = {} - update_masks = [] - if mode not in ["create", "update"]: - raise ValueError(f"Unsupported mode: {mode}") - if agent is None: - if requirements is not None: - raise ValueError("requirements must be None if agent is None.") - if extra_packages is not None: - raise ValueError("extra_packages must be None if agent is None.") - if display_name is not None: - update_masks.append("display_name") - config["display_name"] = display_name - if description is not None: - update_masks.append("description") - config["description"] = description - if context_spec is not None: - update_masks.append("context_spec") - config["context_spec"] = context_spec - if encryption_spec is not None: - update_masks.append("encryption_spec") - config["encryption_spec"] = encryption_spec - if labels is not None: - update_masks.append("labels") - config["labels"] = labels - if traffic_config is not None: - update_masks.append("traffic_config") - config["traffic_config"] = traffic_config - - if agent_framework == "google-adk": - env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) - - if python_version: - sys_version = python_version - else: - sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" - - if agent: - if source_packages: - raise ValueError( - "If you have provided `source_packages` in `config`, please " - "do not specify `agent` in `agent_engines.create()` or " - "`agent_engines.update()`." - ) - if developer_connect_source: - raise ValueError( - "If you have provided `developer_connect_source` in `config`, please " - "do not specify `agent` in `agent_engines.create()` or " - "`agent_engines.update()`." - ) - elif source_packages and developer_connect_source: - raise ValueError( - "Please specify only one of `source_packages` or `developer_connect_source` in `config`." - ) - - if container_spec: - if agent: - raise ValueError( - "If you have provided `container_spec` in `config`, please " - "do not specify `agent` in `agent_engines.create()` or " - "`agent_engines.update()`." - ) - if source_packages or developer_connect_source: - raise ValueError( - "If you have provided `container_spec` in `config`, please " - "do not specify `source_packages` or `developer_connect_source` in `config`." - ) - - agent_engine_spec: Any = None - if agent: - agent_engine_spec = {} - agent = _agent_engines_utils._validate_agent_or_raise(agent=agent) - if _agent_engines_utils._is_adk_agent(agent): - env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars) - self._set_package_spec( - spec=agent_engine_spec, - update_masks=update_masks, - agent=agent, - staging_bucket=staging_bucket, - requirements=requirements, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - class_methods=class_methods, - sys_version=sys_version, - build_options=build_options, - ) - elif ( - source_packages - or developer_connect_source - or image_spec - or agent_config_source - ): - agent_engine_spec = {} - self._set_source_code_spec( - spec=agent_engine_spec, - update_masks=update_masks, - source_packages=source_packages, - developer_connect_source=developer_connect_source, - class_methods=class_methods, - entrypoint_module=entrypoint_module, - entrypoint_object=entrypoint_object, - requirements_file=requirements_file, - sys_version=sys_version, - build_options=build_options, - image_spec=image_spec, - agent_config_source=agent_config_source, - ) - elif container_spec: - agent_engine_spec = {} - if class_methods is not None: - update_masks.append("spec.class_methods") - class_methods_spec_list = ( - _agent_engines_utils._class_methods_to_class_methods_spec( - class_methods=class_methods - ) - ) - agent_engine_spec["class_methods"] = [ - _agent_engines_utils._to_dict(class_method_spec) - for class_method_spec in class_methods_spec_list - ] - update_masks.append("spec.container_spec") - agent_engine_spec["container_spec"] = container_spec - - is_deployment_spec_updated = ( - env_vars is not None - or psc_interface_config is not None - or agent_gateway_config is not None - or min_instances is not None - or max_instances is not None - or resource_limits is not None - or container_concurrency is not None - or keep_alive_probe is not None - ) - if agent_engine_spec is None and is_deployment_spec_updated: - raise ValueError( - "To update `env_vars`, `psc_interface_config`, `min_instances`, " - "`max_instances`, `resource_limits`, `container_concurrency`, or " - "`keep_alive_probe`, you must also provide the `agent` variable or " - "the source code options (`source_packages`, " - "`developer_connect_source` or `agent_config_source`)." - ) - - if agent_engine_spec is not None: - if is_deployment_spec_updated: - ( - deployment_spec, - deployment_update_masks, - ) = self._generate_deployment_spec_or_raise( - env_vars=env_vars, - psc_interface_config=psc_interface_config, - agent_gateway_config=agent_gateway_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - keep_alive_probe=keep_alive_probe, - ) - update_masks.extend(deployment_update_masks) - agent_engine_spec["deployment_spec"] = deployment_spec - - if agent_server_mode: - if not agent_engine_spec.get("deployment_spec"): - agent_engine_spec["deployment_spec"] = ( - types.ReasoningEngineSpecDeploymentSpecDict() - ) - agent_engine_spec["deployment_spec"][ - "agent_server_mode" - ] = agent_server_mode - - agent_engine_spec["agent_framework"] = ( - _agent_engines_utils._get_agent_framework( - agent_framework=agent_framework, - agent=agent, - ) - ) - - if hasattr(agent, "agent_card"): - agent_card = getattr(agent, "agent_card") - if agent_card: - try: - from google.protobuf import json_format - - agent_engine_spec["agent_card"] = json_format.MessageToDict( - agent_card - ) - except Exception as e: - raise ValueError( - f"Failed to convert agent card to dict (serialization error): {e}" - ) from e - update_masks.append("spec.agent_framework") - - if identity_type is not None or service_account is not None: - if agent_engine_spec is None: - agent_engine_spec = {} - - if identity_type is not None: - agent_engine_spec["identity_type"] = identity_type - update_masks.append("spec.identity_type") - if service_account is not None: - # Clear the field in case of empty service_account. - if service_account: - agent_engine_spec["service_account"] = service_account - update_masks.append("spec.service_account") - - if agent_engine_spec is not None: - config["spec"] = agent_engine_spec - - if update_masks and mode == "update": - config["update_mask"] = ",".join(update_masks) - return config - - def _generate_deployment_spec_or_raise( - self, - *, - env_vars: Optional[dict[str, Union[str, Any]]] = None, - psc_interface_config: Optional[types.PscInterfaceConfigDict] = None, - agent_gateway_config: Optional[ - types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict - ] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[dict[str, str]] = None, - container_concurrency: Optional[int] = None, - keep_alive_probe: Optional[dict[str, Any]] = None, - ) -> Tuple[dict[str, Any], Sequence[str]]: - deployment_spec: dict[str, Any] = {} - update_masks = [] - if env_vars: - deployment_spec["env"] = [] - deployment_spec["secret_env"] = [] - if isinstance(env_vars, dict): - self._update_deployment_spec_with_env_vars_dict_or_raise( - deployment_spec=deployment_spec, - env_vars=env_vars, - ) - else: - raise TypeError(f"env_vars must be a dict, but got {type(env_vars)}.") - if deployment_spec.get("env"): - update_masks.append("spec.deployment_spec.env") - if deployment_spec.get("secret_env"): - update_masks.append("spec.deployment_spec.secret_env") - if psc_interface_config: - deployment_spec["psc_interface_config"] = psc_interface_config - update_masks.append("spec.deployment_spec.psc_interface_config") - if agent_gateway_config: - deployment_spec["agent_gateway_config"] = agent_gateway_config - update_masks.append("spec.deployment_spec.agent_gateway_config") - if min_instances is not None: - if not 0 <= min_instances <= 10: - raise ValueError( - f"min_instances must be between 0 and 10. Got {min_instances}" - ) - deployment_spec["min_instances"] = min_instances - update_masks.append("spec.deployment_spec.min_instances") - if max_instances is not None: - if psc_interface_config and not 1 <= max_instances <= 100: - raise ValueError( - f"max_instances must be between 1 and 100 when PSC-I is enabled. Got {max_instances}" - ) - elif not psc_interface_config and not 1 <= max_instances <= 1000: - raise ValueError( - f"max_instances must be between 1 and 1000. Got {max_instances}" - ) - deployment_spec["max_instances"] = max_instances - update_masks.append("spec.deployment_spec.max_instances") - if resource_limits: - _agent_engines_utils._validate_resource_limits_or_raise( - resource_limits=resource_limits - ) - deployment_spec["resource_limits"] = resource_limits - update_masks.append("spec.deployment_spec.resource_limits") - if container_concurrency: - deployment_spec["container_concurrency"] = container_concurrency - update_masks.append("spec.deployment_spec.container_concurrency") - if keep_alive_probe is not None: - deployment_spec["keep_alive_probe"] = keep_alive_probe - update_masks.append("spec.deployment_spec.keep_alive_probe") - return deployment_spec, update_masks - - def _update_deployment_spec_with_env_vars_dict_or_raise( - self, - *, - deployment_spec: dict[str, Any], - env_vars: dict[str, Any], - ) -> None: - for key, value in env_vars.items(): - if isinstance(value, dict): - if "secret_env" not in deployment_spec: - deployment_spec["secret_env"] = [] - deployment_spec["secret_env"].append({"name": key, "secret_ref": value}) - elif isinstance(value, str): - if "env" not in deployment_spec: - deployment_spec["env"] = [] - deployment_spec["env"].append({"name": key, "value": value}) - else: - raise TypeError( - f"Unknown value type in env_vars for {key}. " - f"Must be a str or SecretRef: {value}" - ) - - def _register_api_methods( - self, - *, - agent_engine: types.AgentEngine, - ) -> types.AgentEngine: - """Registers the API methods for the agent engine.""" - try: - _agent_engines_utils._register_api_methods_or_raise( - agent_engine=agent_engine, - wrap_operation_fn={ - "": _agent_engines_utils._wrap_query_operation, # type: ignore[dict-item] - "async": _agent_engines_utils._wrap_async_query_operation, # type: ignore[dict-item] - "stream": _agent_engines_utils._wrap_stream_query_operation, # type: ignore[dict-item] - "async_stream": _agent_engines_utils._wrap_async_stream_query_operation, # type: ignore[dict-item] - "a2a_extension": _agent_engines_utils._wrap_a2a_operation, - }, - ) - except Exception as e: - logger.warning( - _agent_engines_utils._FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e - ) - return agent_engine - - def list( - self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None - ) -> Iterator[types.AgentEngine]: - """List all instances of Agent Engine matching the filter. - - Example Usage: - - .. code-block:: python - import agentplatform - - client = agentplatform.Client(project="my_project", location="us-central1") - for agent in client.agent_engines.list( - config={"filter": "'display_name="My Custom Agent"'}, - ): - print(agent.api_resource.name) - - Args: - config (ListAgentEngineConfig): - Optional. The config (e.g. filter) for the agents to be listed. - - Returns: - Iterable[AgentEngine]: An iterable of Agent Engines matching the filter. - """ - - for reasoning_engine in self._list_pager(config=config): - yield types.AgentEngine( - api_client=self, - api_async_client=AsyncAgentEngines(api_client_=self._api_client), - api_resource=reasoning_engine, - ) - - def update( - self, - *, - name: str, - agent: Any = None, - agent_engine: Any = None, - config: types.AgentEngineConfigOrDict, - ) -> types.AgentEngine: - """Updates an existing Agent Engine. - - This method updates the configuration of an existing Agent Engine running - remotely, which is identified by its name. - - Args: - name (str): Required. A fully-qualified resource name or ID such as - "projects/123/locations/us-central1/reasoningEngines/456" or a - shortened name such as "reasoningEngines/456". - agent (Any): - Optional. The instance to be used as the updated Agent Engine. - If it is not specified, the existing instance will be used. - agent_engine (Any): - Optional. This is deprecated. Please use `agent` instead. - config (AgentEngineConfig): - Optional. The configurations to use for updating the Agent Engine. - - Returns: - AgentEngine: The updated Agent Engine. - - Raises: - ValueError: If the `project` was not set using `client.Client`. - ValueError: If the `location` was not set using `client.Client`. - ValueError: If `config.staging_bucket` was not set when `agent_engine` - is specified. - ValueError: If `config.staging_bucket` does not start with "gs://". - ValueError: If `config.extra_packages` is specified but `agent_engine` - is None. - ValueError: If `config.requirements` is specified but `agent_engine` is - None. - ValueError: If `config.env_vars` has a dictionary entry that does not - correspond to an environment variable value or a SecretRef. - TypeError: If `config.env_vars` is not a dictionary. - FileNotFoundError: If `config.extra_packages` includes a file or - directory that does not exist. - IOError: If `config.requirements` is a string that corresponds to a - nonexistent file. - """ - if isinstance(config, dict): - config = types.AgentEngineConfig.model_validate(config) - elif not isinstance(config, types.AgentEngineConfig): - raise TypeError( - f"config must be a dict or AgentEngineConfig, but got {type(config)}." - ) - context_spec = config.context_spec - if context_spec is not None: - # Conversion to a dict for _create_config - context_spec = json.loads(context_spec.model_dump_json()) - developer_connect_source = config.developer_connect_source - if developer_connect_source is not None: - developer_connect_source = json.loads( - developer_connect_source.model_dump_json() - ) - agent_config_source = config.agent_config_source - if agent_config_source is not None: - agent_config_source = json.loads(agent_config_source.model_dump_json()) - keep_alive_probe = config.keep_alive_probe - if keep_alive_probe is not None: - keep_alive_probe = json.loads( - keep_alive_probe.model_dump_json(exclude_none=True) - ) - traffic_config = config.traffic_config - if traffic_config is not None: - traffic_config = json.loads(traffic_config.model_dump_json()) - if agent and agent_engine: - raise ValueError("Please specify only one of `agent` or `agent_engine`.") - elif agent_engine: - raise DeprecationWarning( - "The `agent_engine` argument is deprecated. Please use `agent` instead." - ) - image_spec = config.image_spec - if image_spec is not None: - # Conversion to a dict for _create_config - image_spec = json.loads(image_spec.model_dump_json()) - container_spec = config.container_spec - if container_spec is not None: - # Conversion to a dict for _create_config - container_spec = json.loads(container_spec.model_dump_json()) - agent = agent or agent_engine - api_config = self._create_config( - mode="update", - agent=agent, - identity_type=config.identity_type, - staging_bucket=config.staging_bucket, - requirements=config.requirements, - display_name=config.display_name, - description=config.description, - gcs_dir_name=config.gcs_dir_name, - extra_packages=config.extra_packages, - env_vars=config.env_vars, - service_account=config.service_account, - context_spec=context_spec, - psc_interface_config=config.psc_interface_config, - agent_gateway_config=config.agent_gateway_config, - min_instances=config.min_instances, - max_instances=config.max_instances, - resource_limits=config.resource_limits, - container_concurrency=config.container_concurrency, - labels=config.labels, - class_methods=config.class_methods, - source_packages=config.source_packages, - developer_connect_source=developer_connect_source, - entrypoint_module=config.entrypoint_module, - entrypoint_object=config.entrypoint_object, - requirements_file=config.requirements_file, - agent_framework=config.agent_framework, - python_version=config.python_version, - build_options=config.build_options, - image_spec=image_spec, - agent_config_source=agent_config_source, - container_spec=container_spec, - keep_alive_probe=keep_alive_probe, - traffic_config=traffic_config, - ) - operation = self._update(name=name, config=api_config) - reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id( - resource_name=name - ) - logger.info( - "View progress and logs at https://console.cloud.google.com/logs/query?" - f"project={self._api_client.project}" - "&query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22%0A" - f"resource.labels.reasoning_engine_id%3D%22{reasoning_engine_id}%22." - ) - operation = _agent_engines_utils._await_operation( - operation_name=operation.name, - get_operation_fn=self._get_agent_operation, - ) - agent_engine = types.AgentEngine( - api_client=self, - api_async_client=AsyncAgentEngines(api_client_=self._api_client), - api_resource=operation.response, - ) - if agent_engine.api_resource: - logger.info("Agent Engine updated. To use it in another session:") - logger.info( - f"agent_engine=client.agent_engines.get(name='{agent_engine.api_resource.name}')" - ) - elif operation.error: - raise RuntimeError(f"Failed to update Agent Engine: {operation.error}") - if agent_engine.api_resource.spec: - self._register_api_methods(agent_engine=agent_engine) - return agent_engine # type: ignore[no-any-return] - - def _stream_query( - self, *, name: str, config: Optional[types.QueryAgentEngineConfigOrDict] = None - ) -> Iterator[Any]: - """Streams the response of the agent engine.""" - parameter_model = types._QueryAgentEngineRequestParameters( - name=name, - config=config, - ) - request_dict = _QueryAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) - else: - path = "{name}:streamQuery?alt=sse" - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - http_options = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - for response in self._api_client.request_streamed( - "post", path, request_dict, http_options - ): - yield response - - # TODO: b/436704146 - Replace with generated methods - # TODO: b/437129724 - Add replay test for async stream query - async def _async_stream_query( - self, - *, - name: str, - config: Optional[types.QueryAgentEngineConfigOrDict] = None, - ) -> AsyncIterator[Any]: - """Streams the response of the agent engine asynchronously.""" - parameter_model = types._QueryAgentEngineRequestParameters( - name=name, - config=config, - ) - request_dict = _QueryAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) - else: - path = "{name}:streamQuery?alt=sse" - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - http_options = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - async_iterator = await self._api_client.async_request_streamed( - "post", path, request_dict, http_options - ) - async for response in async_iterator: - yield response - - def create_memory( - self, - *, - name: str, - fact: str, - scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - """Deprecated. Use agent_engines.memories.create instead.""" - warnings.warn( - ( - "agent_engines.create_memory is deprecated. " - "Use agent_engines.memories.create instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.create( - name=name, - fact=fact, - scope=scope, - config=config, - ) - - def delete_memory( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, - ) -> types.DeleteAgentEngineMemoryOperation: - """Deprecated. Use agent_engines.memories.delete instead.""" - warnings.warn( - ( - "agent_engines.delete_memory is deprecated. " - "Use agent_engines.memories.delete instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.delete(name=name, config=config) - - def generate_memories( - self, - *, - name: str, - vertex_session_source: Optional[ - types.GenerateMemoriesRequestVertexSessionSourceOrDict - ] = None, - direct_contents_source: Optional[ - types.GenerateMemoriesRequestDirectContentsSourceOrDict - ] = None, - direct_memories_source: Optional[ - types.GenerateMemoriesRequestDirectMemoriesSourceOrDict - ] = None, - scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - """Deprecated. Use agent_engines.memories.generate instead.""" - warnings.warn( - ( - "agent_engines.generate_memories is deprecated. " - "Use agent_engines.memories.generate instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.generate( - name=name, - vertex_session_source=vertex_session_source, - direct_contents_source=direct_contents_source, - direct_memories_source=direct_memories_source, - scope=scope, - config=config, - ) - - def get_memory( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, - ) -> types.Memory: - """Deprecated. Use agent_engines.memories.get instead.""" - warnings.warn( - ( - "agent_engines.get_memory is deprecated. " - "Use agent_engines.memories.get instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.get(name=name, config=config) - - def list_memories( - self, - *, - name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, - ) -> Iterator[types.Memory]: - """Deprecated. Use agent_engines.memories.list instead.""" - warnings.warn( - ( - "agent_engines.list_memories is deprecated. " - "Use agent_engines.memories.list instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.list(name=name, config=config) - - def retrieve_memories( - self, - *, - name: str, - scope: dict[str, str], - similarity_search_params: Optional[ - types.RetrieveMemoriesRequestSimilaritySearchParamsOrDict - ] = None, - simple_retrieval_params: Optional[ - types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict - ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, - ) -> Iterator[types.RetrieveMemoriesResponseRetrievedMemory]: - """Deprecated. Use agent_engines.memories.retrieve instead.""" - warnings.warn( - ( - "agent_engines.retrieve_memories is deprecated. " - "Use agent_engines.memories.retrieve instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.memories.retrieve( - name=name, - scope=scope, - similarity_search_params=similarity_search_params, - simple_retrieval_params=simple_retrieval_params, - config=config, - ) - - def create_session( - self, - *, - name: str, - user_id: str, - config: Optional[types.CreateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: - """Deprecated. Use agent_engines.sessions.create instead.""" - warnings.warn( - ( - "agent_engines.create_session is deprecated. " - "Use agent_engines.sessions.create instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.sessions.create(name=name, user_id=user_id, config=config) - - def delete_session( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineSessionConfigOrDict] = None, - ) -> types.DeleteAgentEngineSessionOperation: - """Deprecated. Use agent_engines.sessions.delete instead.""" - warnings.warn( - ( - "agent_engines.delete_session is deprecated. " - "Use agent_engines.sessions.delete instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.sessions.delete(name=name, config=config) - - def get_session( - self, - *, - name: str, - config: Optional[types.GetAgentEngineSessionConfigOrDict] = None, - ) -> types.Session: - """Deprecated. Use agent_engines.sessions.get instead.""" - warnings.warn( - ( - "agent_engines.get_session is deprecated. " - "Use agent_engines.sessions.get instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.sessions.get(name=name, config=config) - - def list_sessions( - self, - *, - name: str, - config: Optional[types.ListAgentEngineSessionsConfigOrDict] = None, - ) -> Iterator[types.Session]: - """Deprecated. Use agent_engines.sessions.list instead.""" - warnings.warn( - ( - "agent_engines.list_sessions is deprecated. " - "Use agent_engines.sessions.list instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.sessions.list(name=name, config=config) - - def append_session_event( - self, - *, - name: str, - author: str, - invocation_id: str, - timestamp: datetime.datetime, - config: Optional[types.AppendAgentEngineSessionEventConfigOrDict] = None, - ) -> types.AppendAgentEngineSessionEventResponse: - """Deprecated. Use agent_engines.sessions.events.append instead.""" - warnings.warn( - ( - "agent_engines.append_session_event is deprecated. " - "Use agent_engines.sessions.events.append instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.sessions.events.append( - name=name, - author=author, - invocation_id=invocation_id, - timestamp=timestamp, - config=config, - ) - - def list_session_events( - self, - *, - name: str, - config: Optional[types.ListAgentEngineSessionEventsConfigOrDict] = None, - ) -> Iterator[types.SessionEvent]: - """Deprecated. Use agent_engines.sessions.events.list instead.""" - warnings.warn( - ( - "agent_engines.list_session_events is deprecated. " - "Use agent_engines.sessions.events.list instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return self.sessions.events.list(name=name, config=config) - - -class AsyncAgentEngines(_api_module.BaseModule): - - async def cancel_query_job( - self, - *, - name: str, - config: Optional[types.CancelQueryJobAgentEngineConfigOrDict] = None, - ) -> types.CancelQueryJobResult: - """ - Cancels a long-running query job on an Agent Engine. - - Args: - name (str): - Required. The reasoning engine resource name. - config (CancelQueryJobAgentEngineConfigOrDict): - Optional. The configuration for the cancel_query_job. - - """ - - parameter_model = types._CancelQueryJobAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _CancelQueryJobAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:cancelAsyncQuery".format_map(request_url_dict) - else: - path = "{name}:cancelAsyncQuery" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "post", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.CancelQueryJobResult._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _check_query_job( - self, - *, - name: str, - config: Optional[types.CheckQueryJobAgentEngineConfigOrDict] = None, - ) -> types.CheckQueryJobResult: - """ - Query an Agent Engine asynchronously. - """ - - parameter_model = types._CheckQueryJobAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _CheckQueryJobAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:checkQueryJob".format_map(request_url_dict) - else: - path = "{name}:checkQueryJob" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "post", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _CheckQueryJobResult_from_vertex(response_dict) - - return_value = types.CheckQueryJobResult._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _run_query_job( - self, - *, - name: str, - config: Optional[types._RunQueryJobAgentEngineConfigOrDict] = None, - ) -> types.AgentEngineOperation: - """ - Run a query job on an agent engine. - """ - - parameter_model = types._RunQueryJobAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _RunQueryJobAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:asyncQuery".format_map(request_url_dict) - else: - path = "{name}:asyncQuery" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "post", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _create( - self, *, config: Optional[types.CreateAgentEngineConfigOrDict] = None - ) -> types.AgentEngineOperation: - """ - Creates a new Agent Engine. - """ - - parameter_model = types._CreateAgentEngineRequestParameters( - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _CreateAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "reasoningEngines".format_map(request_url_dict) - else: - path = "reasoningEngines" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "post", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _delete( - self, - *, - name: str, - force: Optional[bool] = None, - config: Optional[types.DeleteAgentEngineConfigOrDict] = None, - ) -> types.DeleteAgentEngineOperation: - """ - Delete an Agent Engine resource. - - Args: - name (str): - Required. The name of the Agent Engine to be deleted. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` - or `reasoningEngines/{resource_id}`. - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error when - the Agent Engine has undeleted child resources. Defaults to False. - config (DeleteAgentEngineConfig): - Optional. Additional configurations for deleting the Agent Engine. - - """ - - parameter_model = types._DeleteAgentEngineRequestParameters( - name=name, - force=force, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _DeleteAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}".format_map(request_url_dict) - else: - path = "{name}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "delete", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.DeleteAgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _get( - self, *, name: str, config: Optional[types.GetAgentEngineConfigOrDict] = None - ) -> types.ReasoningEngine: - """ - Get an Agent Engine instance. - """ - - parameter_model = types._GetAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _GetAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}".format_map(request_url_dict) - else: - path = "{name}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "get", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _ReasoningEngine_from_vertex(response_dict) - - return_value = types.ReasoningEngine._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _list( - self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None - ) -> types.ListReasoningEnginesResponse: - """ - Lists Agent Engines. - """ - - parameter_model = types._ListAgentEngineRequestParameters( - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _ListAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "reasoningEngines".format_map(request_url_dict) - else: - path = "reasoningEngines" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "get", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) - - return_value = types.ListReasoningEnginesResponse._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _get_agent_operation( - self, - *, - operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineOperation: - parameter_model = types._GetAgentEngineOperationParameters( - operation_name=operation_name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _GetAgentEngineOperationParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{operationName}".format_map(request_url_dict) - else: - path = "{operationName}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "get", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _query( - self, *, name: str, config: Optional[types.QueryAgentEngineConfigOrDict] = None - ) -> types.QueryReasoningEngineResponse: - """ - Query an Agent Engine. - """ - - parameter_model = types._QueryAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _QueryAgentEngineRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}:query".format_map(request_url_dict) - else: - path = "{name}:query" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "post", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - return_value = types.QueryReasoningEngineResponse._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - async def _update( - self, *, name: str, config: Optional[types.UpdateAgentEngineConfigOrDict] = None - ) -> types.AgentEngineOperation: - """ - Updates an Agent Engine. - """ - - parameter_model = types._UpdateAgentEngineRequestParameters( - name=name, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( - "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." - ) - else: - request_dict = _UpdateAgentEngineRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{name}".format_map(request_url_dict) - else: - path = "{name}" - - query_params = request_dict.get("_query") - if query_params: - path = f"{path}?{urlencode(query_params)}" - # TODO: remove the hack that pops config. - request_dict.pop("config", None) - - http_options: Optional[types.HttpOptions] = None - if ( - parameter_model.config is not None - and parameter_model.config.http_options is not None - ): - http_options = parameter_model.config.http_options - - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) - - response = await self._api_client.async_request( - "patch", path, request_dict, http_options - ) - - response_dict = {} if not response.body else json.loads(response.body) - - if self._api_client.vertexai: - response_dict = _AgentEngineOperation_from_vertex(response_dict) - - return_value = types.AgentEngineOperation._from_response( - response=response_dict, - kwargs=( - { - "config": { - "response_schema": getattr( - parameter_model.config, "response_schema", None - ), - "response_json_schema": getattr( - parameter_model.config, "response_json_schema", None - ), - "include_all_fields": getattr( - parameter_model.config, "include_all_fields", None - ), - } - } - if getattr(parameter_model, "config", None) - else {} - ), - ) - - self._api_client._verify_response(return_value) - return return_value - - _a2a_tasks = None - _memories = None - _sessions = None - _runtimes = None - - async def delete( - self, - *, - name: str, - force: Optional[bool] = None, - config: Optional[types.DeleteAgentEngineConfigOrDict] = None, - ) -> types.DeleteAgentEngineOperation: - """ - Delete an Agent Engine resource. - - Args: - name (str): - Required. The name of the Agent Engine to be deleted. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` - or `reasoningEngines/{resource_id}`. - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error when - the Agent Engine has undeleted child resources. Defaults to False. - config (DeleteAgentEngineConfig): - Optional. Additional configurations for deleting the Agent Engine. - - """ - logger.info(f"Deleting AgentEngine resource: {name}") - operation = await self._delete(name=name, force=force, config=config) - logger.info(f"Started AgentEngine delete operation: {operation.name}") - return operation - - @property - def runtimes(self) -> "runtimes_module.AsyncRuntimes": - if self._runtimes is None: - try: - # We need to lazy load the runtimes module to handle the - # possibility of ImportError when dependencies are not installed. - self._runtimes = importlib.import_module(".runtimes", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.runtimes' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._runtimes.AsyncRuntimes(self._api_client) # type: ignore[no-any-return] - - @property - def a2a_tasks(self) -> "a2a_tasks_module.AsyncA2aTasks": - if self._a2a_tasks is None: - try: - # We need to lazy load the a2a_tasks module to handle the - # possibility of ImportError when dependencies are not installed. - self._a2a_tasks = importlib.import_module(".a2a_tasks", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.a2a_tasks' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._a2a_tasks.AsyncA2aTasks(self._api_client) # type: ignore[no-any-return] - - @property - def memories(self) -> "memories_module.AsyncMemories": - if self._memories is None: - try: - # We need to lazy load the memories module to handle the - # possibility of ImportError when dependencies are not installed. - self._memories = importlib.import_module(".memories", __package__) - except ImportError as e: - raise ImportError( - "The 'agent_engines.memories' module requires additional " - "packages. Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._memories.AsyncMemories(self._api_client) # type: ignore[no-any-return] - - @property - def sessions(self) -> "sessions_module.AsyncSessions": - if self._sessions is None: - try: - # We need to lazy load the sessions module to handle the - # possibility of ImportError when dependencies are not installed. - self._sessions = importlib.import_module(".sessions", __package__) - except ImportError as e: - raise ImportError( - "The agent_engines.sessions module requires additional packages. " - "Please install them using pip install " - "google-cloud-aiplatform[agent_engines]" - ) from e - return self._sessions.AsyncSessions(self._api_client) # type: ignore[no-any-return] - - async def append_session_event( - self, - *, - name: str, - author: str, - invocation_id: str, - timestamp: datetime.datetime, - config: Optional[types.AppendAgentEngineSessionEventConfigOrDict] = None, - ) -> types.AppendAgentEngineSessionEventResponse: - """Deprecated. Use agent_engines.sessions.events.append instead.""" - warnings.warn( - ( - "agent_engines.append_session_event is deprecated. " - "Use agent_engines.sessions.events.append instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return await self.sessions.events.append( - name=name, - author=author, - invocation_id=invocation_id, - timestamp=timestamp, - config=config, - ) - - async def delete_memory( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, - ) -> types.DeleteAgentEngineMemoryOperation: - """Deprecated. Use agent_engines.memories.delete instead.""" - warnings.warn( - ( - "agent_engines.delete_memory is deprecated. " - "Use agent_engines.memories.delete instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return await self.memories.delete(name=name, config=config) - - async def delete_session( - self, - *, - name: str, - config: Optional[types.DeleteAgentEngineSessionConfigOrDict] = None, - ) -> types.DeleteAgentEngineSessionOperation: - """Deprecated. Use agent_engines.sessions.delete instead.""" - warnings.warn( - ( - "agent_engines.delete_session is deprecated. " - "Use agent_engines.sessions.delete instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return await self.sessions.delete(name=name, config=config) - - async def get_memory( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, - ) -> types.Memory: - """Deprecated. Use agent_engines.memories.get instead.""" - warnings.warn( - ( - "agent_engines.get_memory is deprecated. " - "Use agent_engines.memories.get instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return await self.memories.get(name=name, config=config) - - async def get_session( - self, - *, - name: str, - config: Optional[types.GetAgentEngineSessionConfigOrDict] = None, - ) -> types.Session: - """Deprecated. Use agent_engines.sessions.get instead.""" - warnings.warn( - ( - "agent_engines.get_session is deprecated. " - "Use agent_engines.sessions.get instead." - ), - DeprecationWarning, - stacklevel=2, - ) - return await self.sessions.get(name=name, config=config) diff --git a/agentplatform/_genai/client.py b/agentplatform/_genai/client.py index 0dcb0d445e..a94cb1b85d 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 @@ -44,10 +44,15 @@ ) from agentplatform._genai import live as live_module from agentplatform._genai import rag as rag_module + from agentplatform._genai import sessions as sessions_module + from agentplatform._genai import ( + sandboxes as sandboxes_module, + ) from agentplatform._genai import ( feedback_entries as feedback_entries_module, ) + _GENAI_MODULES_TELEMETRY_HEADER = "vertex-genai-modules" @@ -84,7 +89,7 @@ def __init__(self, api_client: genai_client.BaseApiClient): # type: ignore[name self._api_client = api_client self._live = live.AsyncLive(self._api_client) self._evals: Optional[ModuleType] = None - self._agent_engines: Optional[ModuleType] = None + self._runtimes: Optional[ModuleType] = None self._prompt_optimizer: Optional[ModuleType] = None self._prompts: Optional[ModuleType] = None self._datasets: Optional[ModuleType] = None @@ -92,6 +97,8 @@ def __init__(self, api_client: genai_client.BaseApiClient): # type: ignore[name self._rag: Optional[ModuleType] = None self._model_garden: Optional[ModuleType] = None self._feedback_entries: Optional[ModuleType] = None + self._sessions: Optional[ModuleType] = None + self._sandboxes: Optional[ModuleType] = None @property @_common.experimental_warning( @@ -125,22 +132,58 @@ def prompt_optimizer(self) -> "prompt_optimizer_module.AsyncPromptOptimizer": return self._prompt_optimizer.AsyncPromptOptimizer(self._api_client) # type: ignore[no-any-return] @property - def agent_engines(self) -> "agent_engines_module.AsyncAgentEngines": - if self._agent_engines is None: + def runtimes(self) -> "runtimes_module.AsyncRuntimes": + if self._runtimes is None: + try: + # We need to lazy load the runtimes module to handle the + # possibility of ImportError when dependencies are not installed. + self._runtimes = importlib.import_module( + ".runtimes", + __package__, + ) + except ImportError as e: + raise ImportError( + "The 'runtimes' module requires 'additional packages'. " + "Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._runtimes.AsyncRuntimes(self._api_client) # type: ignore[no-any-return] + + @property + def sessions(self) -> "sessions_module.AsyncSessions": + if self._sessions is None: try: - # We need to lazy load the agent_engines module to handle the + # We need to lazy load the sessions module to handle the # possibility of ImportError when dependencies are not installed. - self._agent_engines = importlib.import_module( - ".agent_engines", + self._sessions = importlib.import_module( + ".sessions", __package__, ) except ImportError as e: raise ImportError( - "The 'agent_engines' module requires 'additional packages'. " + "The 'sessions' module requires 'additional packages'. " "Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e - return self._agent_engines.AsyncAgentEngines(self._api_client) # type: ignore[no-any-return] + return self._sessions.AsyncSessions(self._api_client) # type: ignore[no-any-return] + + @property + def sandboxes(self) -> "sandboxes_module.AsyncSandboxes": + if self._sandboxes is None: + try: + # We need to lazy load the sandboxes module to handle the + # possibility of ImportError when dependencies are not installed. + self._sandboxes = importlib.import_module( + ".sandboxes", + __package__, + ) + except ImportError as e: + raise ImportError( + "The 'sandboxes' module requires 'additional packages'. " + "Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._sandboxes.AsyncSandboxes(self._api_client) # type: ignore[no-any-return] @property def prompts(self) -> "prompts_module.AsyncPrompts": @@ -307,13 +350,15 @@ def __init__( self._aio = AsyncClient(self._api_client) self._evals: Optional[ModuleType] = None self._prompt_optimizer: Optional[ModuleType] = None - self._agent_engines: Optional[ModuleType] = None + self._runtimes: Optional[ModuleType] = None self._prompts: Optional[ModuleType] = None self._datasets: Optional[ModuleType] = None self._skills: Optional[ModuleType] = None self._rag: Optional[ModuleType] = None self._model_garden: Optional[ModuleType] = None self._feedback_entries: Optional[ModuleType] = None + self._sessions: Optional[ModuleType] = None + self._sandboxes: Optional[ModuleType] = None @property def evals(self) -> "evals_module.Evals": @@ -371,22 +416,58 @@ def _get_api_client( return None @property - def agent_engines(self) -> "agent_engines_module.AgentEngines": - if self._agent_engines is None: + def runtimes(self) -> "runtimes_module.Runtimes": + if self._runtimes is None: + try: + # We need to lazy load the runtimes module to handle the + # possibility of ImportError when dependencies are not installed. + self._runtimes = importlib.import_module( + ".runtimes", + __package__, + ) + except ImportError as e: + raise ImportError( + "The 'runtimes' module requires 'additional packages'. " + "Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._runtimes.Runtimes(self._api_client) # type: ignore[no-any-return] + + @property + def sessions(self) -> "sessions_module.Sessions": + if self._sessions is None: + try: + # We need to lazy load the sessions module to handle the + # possibility of ImportError when dependencies are not installed. + self._sessions = importlib.import_module( + ".sessions", + __package__, + ) + except ImportError as e: + raise ImportError( + "The 'sessions' module requires 'additional packages'. " + "Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._sessions.Sessions(self._api_client) # type: ignore[no-any-return] + + @property + def sandboxes(self) -> "sandboxes_module.Sandboxes": + if self._sandboxes is None: try: - # We need to lazy load the agent_engines module to handle the + # We need to lazy load the sandboxes module to handle the # possibility of ImportError when dependencies are not installed. - self._agent_engines = importlib.import_module( - ".agent_engines", + self._sandboxes = importlib.import_module( + ".sandboxes", __package__, ) except ImportError as e: raise ImportError( - "The 'agent_engines' module requires 'additional packages'. " + "The 'sandboxes' module requires 'additional packages'. " "Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e - return self._agent_engines.AgentEngines(self._api_client) # type: ignore[no-any-return] + return self._sandboxes.Sandboxes(self._api_client) # type: ignore[no-any-return] @property def prompts(self) -> "prompts_module.Prompts": diff --git a/agentplatform/_genai/evals.py b/agentplatform/_genai/evals.py index 1f4aa7e1de..6c37b8f048 100644 --- a/agentplatform/_genai/evals.py +++ b/agentplatform/_genai/evals.py @@ -2131,7 +2131,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: @@ -2154,11 +2154,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, @@ -2198,7 +2198,7 @@ def run_inference( agent ): gemini_agent_instance = agent - elif isinstance(agent, str) or isinstance(agent, types.AgentEngine): + elif isinstance(agent, str) or isinstance(agent, types.Runtime): agent_engine_instance = agent else: agent_instance = agent @@ -2672,7 +2672,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. @@ -4460,7 +4460,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 87% rename from agentplatform/_genai/live_agent_engines.py rename to agentplatform/_genai/live_runtimes.py index ed79ed5d0c..19f880414f 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,7 +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: " + "Failed to parse response to JSON in live_runtimes: " f"{str(response)}" ) from exc @@ -79,8 +79,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,16 +91,16 @@ 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) + agent_engine = client.runtimes.create(agent) - async with client.aio.live.agent_engines.connect( + async with client.aio.live.runtimes.connect( agent_engine=agent_engine.api_resource.name, setup={"class_method": "bidi_stream_query"}, ) as session: @@ -116,8 +116,8 @@ async def connect( self, *, agent_engine: str, - config: Optional[QueryAgentEngineConfigOrDict] = None, - ) -> AsyncIterator[AsyncLiveAgentEngineSession]: + config: Optional[QueryRuntimeConfigOrDict] = None, + ) -> AsyncIterator[AsyncLiveRuntimeSession]: """Connect to the agent deployed to Agent Engine in a live (bidirectional streaming) session. Args: @@ -129,10 +129,10 @@ 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/"): @@ -176,4 +176,4 @@ async def connect( uri, additional_headers=headers, **self._api_client._websocket_ssl_ctx ) as ws: await ws.send(request) - yield AsyncLiveAgentEngineSession(websocket=ws) + yield AsyncLiveRuntimeSession(websocket=ws) diff --git a/agentplatform/_genai/memories.py b/agentplatform/_genai/memories.py index c83cd240e1..dae85e6cde 100644 --- a/agentplatform/_genai/memories.py +++ b/agentplatform/_genai/memories.py @@ -30,7 +30,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import AsyncPager, Pager -from . import _agent_engines_utils +from . import _runtimes_utils from . import types if typing.TYPE_CHECKING: @@ -44,60 +44,7 @@ logger.setLevel(logging.INFO) -def _AgentEngineMemoryConfig_to_vertex( - from_object: Union[dict[str, Any], object], - parent_object: Optional[dict[str, Any]] = None, -) -> dict[str, Any]: - to_object: dict[str, Any] = {} - - if getv(from_object, ["display_name"]) is not None: - setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) - - if getv(from_object, ["description"]) is not None: - setv(parent_object, ["description"], getv(from_object, ["description"])) - - if getv(from_object, ["ttl"]) is not None: - setv(parent_object, ["ttl"], getv(from_object, ["ttl"])) - - if getv(from_object, ["expire_time"]) is not None: - setv(parent_object, ["expireTime"], getv(from_object, ["expire_time"])) - - if getv(from_object, ["revision_expire_time"]) is not None: - setv( - parent_object, - ["revisionExpireTime"], - getv(from_object, ["revision_expire_time"]), - ) - - if getv(from_object, ["revision_ttl"]) is not None: - setv(parent_object, ["revisionTtl"], getv(from_object, ["revision_ttl"])) - - if getv(from_object, ["disable_memory_revisions"]) is not None: - setv( - parent_object, - ["disableMemoryRevisions"], - getv(from_object, ["disable_memory_revisions"]), - ) - - if getv(from_object, ["topics"]) is not None: - setv( - parent_object, ["topics"], [item for item in getv(from_object, ["topics"])] - ) - - if getv(from_object, ["metadata"]) is not None: - setv( - parent_object, - ["metadata"], - {k: v for k, v in getv(from_object, ["metadata"]).items()}, - ) - - if getv(from_object, ["memory_id"]) is not None: - setv(parent_object, ["_query", "memoryId"], getv(from_object, ["memory_id"])) - - return to_object - - -def _CreateAgentEngineMemoryRequestParameters_to_vertex( +def _CreateRuntimeMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -112,12 +59,12 @@ def _CreateAgentEngineMemoryRequestParameters_to_vertex( setv(to_object, ["scope"], getv(from_object, ["scope"])) if getv(from_object, ["config"]) is not None: - _AgentEngineMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) + _RuntimeMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _DeleteAgentEngineMemoryRequestParameters_to_vertex( +def _DeleteRuntimeMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -128,7 +75,7 @@ def _DeleteAgentEngineMemoryRequestParameters_to_vertex( return to_object -def _GenerateAgentEngineMemoriesConfig_to_vertex( +def _GenerateRuntimeMemoriesConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -185,7 +132,7 @@ def _GenerateAgentEngineMemoriesConfig_to_vertex( return to_object -def _GenerateAgentEngineMemoriesRequestParameters_to_vertex( +def _GenerateRuntimeMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -218,14 +165,14 @@ def _GenerateAgentEngineMemoriesRequestParameters_to_vertex( setv(to_object, ["scope"], getv(from_object, ["scope"])) if getv(from_object, ["config"]) is not None: - _GenerateAgentEngineMemoriesConfig_to_vertex( + _GenerateRuntimeMemoriesConfig_to_vertex( getv(from_object, ["config"]), to_object ) return to_object -def _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( +def _GetRuntimeGenerateMemoriesOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -238,7 +185,7 @@ def _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( return to_object -def _GetAgentEngineMemoryOperationParameters_to_vertex( +def _GetRuntimeMemoryOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -251,7 +198,7 @@ def _GetAgentEngineMemoryOperationParameters_to_vertex( return to_object -def _GetAgentEngineMemoryRequestParameters_to_vertex( +def _GetRuntimeMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -342,7 +289,7 @@ def _IngestEventsRequestParameters_to_vertex( return to_object -def _ListAgentEngineMemoryConfig_to_vertex( +def _ListRuntimeMemoryConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -363,7 +310,7 @@ def _ListAgentEngineMemoryConfig_to_vertex( return to_object -def _ListAgentEngineMemoryRequestParameters_to_vertex( +def _ListRuntimeMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -372,12 +319,12 @@ def _ListAgentEngineMemoryRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) + _ListRuntimeMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _PurgeAgentEngineMemoriesRequestParameters_to_vertex( +def _PurgeRuntimeMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -401,7 +348,21 @@ def _PurgeAgentEngineMemoriesRequestParameters_to_vertex( return to_object -def _RetrieveAgentEngineMemoriesConfig_to_vertex( +def _RetrieveMemoryProfilesRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["scope"]) is not None: + setv(to_object, ["scope"], getv(from_object, ["scope"])) + + return to_object + + +def _RetrieveRuntimeMemoriesConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -423,7 +384,7 @@ def _RetrieveAgentEngineMemoriesConfig_to_vertex( return to_object -def _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( +def _RetrieveRuntimeMemoriesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -449,14 +410,14 @@ def _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( ) if getv(from_object, ["config"]) is not None: - _RetrieveAgentEngineMemoriesConfig_to_vertex( + _RetrieveRuntimeMemoriesConfig_to_vertex( getv(from_object, ["config"]), to_object ) return to_object -def _RetrieveMemoryProfilesRequestParameters_to_vertex( +def _RollbackRuntimeMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -464,27 +425,66 @@ def _RetrieveMemoryProfilesRequestParameters_to_vertex( if getv(from_object, ["name"]) is not None: setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - if getv(from_object, ["scope"]) is not None: - setv(to_object, ["scope"], getv(from_object, ["scope"])) + if getv(from_object, ["target_revision_id"]) is not None: + setv(to_object, ["targetRevisionId"], getv(from_object, ["target_revision_id"])) return to_object -def _RollbackAgentEngineMemoryRequestParameters_to_vertex( +def _RuntimeMemoryConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: to_object: dict[str, Any] = {} - if getv(from_object, ["name"]) is not None: - setv(to_object, ["_url", "name"], getv(from_object, ["name"])) - if getv(from_object, ["target_revision_id"]) is not None: - setv(to_object, ["targetRevisionId"], getv(from_object, ["target_revision_id"])) + if getv(from_object, ["display_name"]) is not None: + setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) + + if getv(from_object, ["description"]) is not None: + setv(parent_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["ttl"]) is not None: + setv(parent_object, ["ttl"], getv(from_object, ["ttl"])) + + if getv(from_object, ["expire_time"]) is not None: + setv(parent_object, ["expireTime"], getv(from_object, ["expire_time"])) + + if getv(from_object, ["revision_expire_time"]) is not None: + setv( + parent_object, + ["revisionExpireTime"], + getv(from_object, ["revision_expire_time"]), + ) + + if getv(from_object, ["revision_ttl"]) is not None: + setv(parent_object, ["revisionTtl"], getv(from_object, ["revision_ttl"])) + + if getv(from_object, ["disable_memory_revisions"]) is not None: + setv( + parent_object, + ["disableMemoryRevisions"], + getv(from_object, ["disable_memory_revisions"]), + ) + + if getv(from_object, ["topics"]) is not None: + setv( + parent_object, ["topics"], [item for item in getv(from_object, ["topics"])] + ) + + if getv(from_object, ["metadata"]) is not None: + setv( + parent_object, + ["metadata"], + {k: v for k, v in getv(from_object, ["metadata"]).items()}, + ) + + if getv(from_object, ["memory_id"]) is not None: + setv(parent_object, ["_query", "memoryId"], getv(from_object, ["memory_id"])) return to_object -def _UpdateAgentEngineMemoryConfig_to_vertex( +def _UpdateRuntimeMemoryConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -542,7 +542,7 @@ def _UpdateAgentEngineMemoryConfig_to_vertex( return to_object -def _UpdateAgentEngineMemoryRequestParameters_to_vertex( +def _UpdateRuntimeMemoryRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -557,9 +557,7 @@ def _UpdateAgentEngineMemoryRequestParameters_to_vertex( setv(to_object, ["scope"], getv(from_object, ["scope"])) if getv(from_object, ["config"]) is not None: - _UpdateAgentEngineMemoryConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _UpdateRuntimeMemoryConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object @@ -572,13 +570,13 @@ def _create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: + config: Optional[types.RuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: """ - Creates a new memory in the Agent Engine. + Creates a new memory in the Agent Runtime. """ - parameter_model = types._CreateAgentEngineMemoryRequestParameters( + parameter_model = types._CreateRuntimeMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -591,7 +589,7 @@ def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _CreateRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -620,7 +618,7 @@ def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.RuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -648,21 +646,21 @@ def delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, - ) -> types.DeleteAgentEngineMemoryOperation: + config: Optional[types.DeleteRuntimeMemoryConfigOrDict] = None, + ) -> types.DeleteRuntimeMemoryOperation: """ - Delete an Agent Engine memory. + Delete an Agent Runtime memory. Args: name (str): - Required. The name of the Agent Engine memory to be deleted. Format: + Required. The name of the Agent Runtime memory to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory}`. - config (DeleteAgentEngineMemoryConfig): - Optional. Additional configurations for deleting the Agent Engine. + config (DeleteRuntimeMemoryConfig): + Optional. Additional configurations for deleting the Agent Runtime. """ - parameter_model = types._DeleteAgentEngineMemoryRequestParameters( + parameter_model = types._DeleteRuntimeMemoryRequestParameters( name=name, config=config, ) @@ -673,7 +671,7 @@ def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _DeleteRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -702,7 +700,7 @@ def delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineMemoryOperation._from_response( + return_value = types.DeleteRuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -740,13 +738,13 @@ def _generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: + config: Optional[types.GenerateRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: """ - Generates memories for an Agent Engine. + Generates memories for an Agent Runtime. """ - parameter_model = types._GenerateAgentEngineMemoriesRequestParameters( + parameter_model = types._GenerateRuntimeMemoriesRequestParameters( name=name, vertex_session_source=vertex_session_source, direct_contents_source=direct_contents_source, @@ -761,7 +759,7 @@ def _generate( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GenerateAgentEngineMemoriesRequestParameters_to_vertex( + request_dict = _GenerateRuntimeMemoriesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -790,7 +788,7 @@ def _generate( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineGenerateMemoriesOperation._from_response( + return_value = types.RuntimeGenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -815,13 +813,10 @@ def _generate( return return_value def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeMemoryConfigOrDict] = None ) -> types.Memory: """ - Gets an agent engine memory. + Gets an agent runtime memory. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -830,7 +825,7 @@ def get( """ - parameter_model = types._GetAgentEngineMemoryRequestParameters( + parameter_model = types._GetRuntimeMemoryRequestParameters( name=name, config=config, ) @@ -841,9 +836,7 @@ def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _GetRuntimeMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -979,16 +972,13 @@ def _ingest_events( return return_value def _list( - self, - *, - name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, + self, *, name: str, config: Optional[types.ListRuntimeMemoryConfigOrDict] = None ) -> types.ListReasoningEnginesMemoriesResponse: """ - Lists Agent Engine memories. + Lists Agent Runtime memories. """ - parameter_model = types._ListAgentEngineMemoryRequestParameters( + parameter_model = types._ListRuntimeMemoryRequestParameters( name=name, config=config, ) @@ -999,7 +989,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _ListRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1056,9 +1046,9 @@ def _get_memory_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - parameter_model = types._GetAgentEngineMemoryOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: + parameter_model = types._GetRuntimeMemoryOperationParameters( operation_name=operation_name, config=config, ) @@ -1069,7 +1059,7 @@ def _get_memory_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryOperationParameters_to_vertex( + request_dict = _GetRuntimeMemoryOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1098,7 +1088,7 @@ def _get_memory_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.RuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1126,9 +1116,9 @@ def _get_generate_memories_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - parameter_model = types._GetAgentEngineGenerateMemoriesOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: + parameter_model = types._GetRuntimeGenerateMemoriesOperationParameters( operation_name=operation_name, config=config, ) @@ -1139,7 +1129,7 @@ def _get_generate_memories_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( + request_dict = _GetRuntimeGenerateMemoriesOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1168,7 +1158,7 @@ def _get_generate_memories_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineGenerateMemoriesOperation._from_response( + return_value = types.RuntimeGenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -1203,13 +1193,13 @@ def _retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveRuntimeMemoriesConfigOrDict] = None, ) -> types.RetrieveMemoriesResponse: """ - Retrieves memories for an Agent Engine. + Retrieves memories for an Agent Runtime. """ - parameter_model = types._RetrieveAgentEngineMemoriesRequestParameters( + parameter_model = types._RetrieveRuntimeMemoriesRequestParameters( name=name, scope=scope, similarity_search_params=similarity_search_params, @@ -1223,7 +1213,7 @@ def _retrieve( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( + request_dict = _RetrieveRuntimeMemoriesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1284,13 +1274,13 @@ def retrieve_profiles( config: Optional[types.RetrieveMemoryProfilesConfigOrDict] = None, ) -> types.RetrieveProfilesResponse: """ - Retrieves memory profiles for an Agent Engine. + Retrieves memory profiles for an Agent Runtime. For example, you can use the following code to retrieve all memory profiles for scope `{'user_id': '123'}`: ```python - result = client.agent_engines.memories.retrieve_profiles( + result = client.runtimes.memories.retrieve_profiles( name="projects/123/locations/us-central1/reasoningEngines/456", scope={"user_id": "123"} ) @@ -1383,13 +1373,13 @@ def _rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineRollbackMemoryOperation: + config: Optional[types.RollbackRuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeRollbackMemoryOperation: """ Rollback a memory to a previous revision. """ - parameter_model = types._RollbackAgentEngineMemoryRequestParameters( + parameter_model = types._RollbackRuntimeMemoryRequestParameters( name=name, target_revision_id=target_revision_id, config=config, @@ -1401,7 +1391,7 @@ def _rollback( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RollbackAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _RollbackRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1430,7 +1420,7 @@ def _rollback( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineRollbackMemoryOperation._from_response( + return_value = types.RuntimeRollbackMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1460,13 +1450,13 @@ def _update( name: str, fact: Optional[str] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.UpdateAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: + config: Optional[types.UpdateRuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: """ - Updates an Agent Engine memory. + Updates an Agent Runtime memory. """ - parameter_model = types._UpdateAgentEngineMemoryRequestParameters( + parameter_model = types._UpdateRuntimeMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -1479,7 +1469,7 @@ def _update( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _UpdateAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _UpdateRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1508,7 +1498,7 @@ def _update( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.RuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -1541,13 +1531,13 @@ def _purge( builtins.list[types.MemoryConjunctionFilterOrDict] ] = None, force: Optional[bool] = None, - config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEnginePurgeMemoriesOperation: + config: Optional[types.PurgeRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimePurgeMemoriesOperation: """ - Purges memories from an Agent Engine. + Purges memories from an Agent Runtime. """ - parameter_model = types._PurgeAgentEngineMemoriesRequestParameters( + parameter_model = types._PurgeRuntimeMemoriesRequestParameters( name=name, filter=filter, filter_groups=filter_groups, @@ -1561,7 +1551,7 @@ def _purge( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _PurgeAgentEngineMemoriesRequestParameters_to_vertex( + request_dict = _PurgeRuntimeMemoriesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1590,7 +1580,7 @@ def _purge( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEnginePurgeMemoriesOperation._from_response( + return_value = types.RuntimePurgeMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -1627,7 +1617,7 @@ def revisions(self) -> "memory_revisions_module.MemoryRevisions": ) except ImportError as e: raise ImportError( - "The 'agent_engines.memories.revisions' module requires " + "The 'runtimes.memories.revisions' module requires " "additional packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e @@ -1639,9 +1629,9 @@ def create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - """Creates a new memory in the Agent Engine. + config: Optional[types.RuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: + """Creates a new memory in the Agent Runtime. Args: name (str): @@ -1650,16 +1640,16 @@ def create( Required. The fact to be stored in the memory. scope (dict[str, str]): Required. The scope of the memory. For example, {"user_id": "123"}. - config (AgentEngineMemoryConfigOrDict): + config (RuntimeMemoryConfigOrDict): Optional. The configuration for the memory. Returns: - AgentEngineMemoryOperation: The operation for creating the memory. + RuntimeMemoryOperation: The operation for creating the memory. """ if config is None: - config = types.AgentEngineMemoryConfig() + config = types.RuntimeMemoryConfig() elif isinstance(config, dict): - config = types.AgentEngineMemoryConfig.model_validate(config) + config = types.RuntimeMemoryConfig.model_validate(config) operation = self._create( name=name, fact=fact, @@ -1668,7 +1658,7 @@ def create( ) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -1697,13 +1687,13 @@ def generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - """Generates memories for the agent engine. + config: Optional[types.GenerateRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: + """Generates memories for the agent runtime. Args: name (str): - Required. The name of the agent engine to generate memories for. + Required. The name of the agent runtime to generate memories for. vertex_session_source (GenerateMemoriesRequestVertexSessionSource): Optional. The vertex session source to use for generating memories. Only one of vertex_session_source, @@ -1724,13 +1714,13 @@ def generate( Optional. The configuration for the memories to generate. Returns: - AgentEngineGenerateMemoriesOperation: + RuntimeGenerateMemoriesOperation: The operation for generating the memories. """ if config is None: - config = types.GenerateAgentEngineMemoriesConfig() + config = types.GenerateRuntimeMemoriesConfig() elif isinstance(config, dict): - config = types.GenerateAgentEngineMemoriesConfig.model_validate(config) + config = types.GenerateRuntimeMemoriesConfig.model_validate(config) operation = self._generate( name=name, vertex_session_source=vertex_session_source, @@ -1740,7 +1730,7 @@ def generate( config=config, ) if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_generate_memories_operation, poll_interval_seconds=0.5, @@ -1753,14 +1743,14 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, + config: Optional[types.ListRuntimeMemoryConfigOrDict] = None, ) -> Iterator[types.Memory]: - """Lists Agent Engine memories. + """Lists Agent Runtime memories. Args: name (str): - Required. The name of the agent engine to list memories for. - config (ListAgentEngineMemoryConfig): + Required. The name of the agent runtime to list memories for. + config (ListRuntimeMemoryConfig): Optional. The configuration for the memories to list. Returns: @@ -1785,13 +1775,13 @@ def retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveRuntimeMemoriesConfigOrDict] = None, ) -> Iterator[types.RetrieveMemoriesResponseRetrievedMemory]: """Retrieves memories for the agent. Args: name (str): - Required. The name of the agent engine to retrieve memories for. + Required. The name of the agent runtime to retrieve memories for. scope (dict[str, str]): Required. The scope of the memories to retrieve. For example, {"user_id": "123"}. @@ -1801,7 +1791,7 @@ def retrieve( simple_retrieval_params (RetrieveMemoriesRequestSimpleRetrievalParams): Optional. The simple retrieval parameters to use for retrieving memories. - config (RetrieveAgentEngineMemoriesConfig): + config (RetrieveRuntimeMemoriesConfig): Optional. The configuration for the memories to retrieve. Returns: @@ -1832,8 +1822,8 @@ def rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineRollbackMemoryOperation: + config: Optional[types.RollbackRuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeRollbackMemoryOperation: """Rolls back a memory to a previous revision. Args: @@ -1841,24 +1831,24 @@ def rollback( Required. The name of the memory to rollback. target_revision_id (str): Required. The revision ID to roll back to - config (RollbackAgentEngineMemoryConfig): + config (RollbackRuntimeMemoryConfig): Optional. The configuration for the rollback. Returns: - AgentEngineRollbackMemoryOperation: + RuntimeRollbackMemoryOperation: The operation for rolling back the memory. """ if config is None: - config = types.RollbackAgentEngineMemoryConfig() + config = types.RollbackRuntimeMemoryConfig() elif isinstance(config, dict): - config = types.RollbackAgentEngineMemoryConfig.model_validate(config) + config = types.RollbackRuntimeMemoryConfig.model_validate(config) operation = self._rollback( name=name, target_revision_id=target_revision_id, config=config, ) if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -1874,13 +1864,13 @@ def purge( filter: Optional[str] = None, filter_groups: Optional[List[types.MemoryConjunctionFilter]] = None, force: bool = False, - config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEnginePurgeMemoriesOperation: - """Purges memories from an Agent Engine. + config: Optional[types.PurgeRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimePurgeMemoriesOperation: + """Purges memories from an Agent Runtime. Args: name (str): - Required. The name of the Agent Engine to purge memories from. + Required. The name of the Agent Runtime to purge memories from. filter (str): Optional. The standard list filter to determine which memories to purge. filter_groups (list[MemoryConjunctionFilter]): @@ -1890,17 +1880,17 @@ def purge( force (bool): Optional. Whether to force the purge operation. If false, the operation will be staged but not executed. - config (PurgeAgentEngineMemoriesConfig): + config (PurgeRuntimeMemoriesConfig): Optional. The configuration for the purge operation. Returns: - AgentEnginePurgeMemoriesOperation: + RuntimePurgeMemoriesOperation: The operation for purging the memories. """ if config is None: - config = types.PurgeAgentEngineMemoriesConfig() + config = types.PurgeRuntimeMemoriesConfig() elif isinstance(config, dict): - config = types.PurgeAgentEngineMemoriesConfig.model_validate(config) + config = types.PurgeRuntimeMemoriesConfig.model_validate(config) operation = self._purge( name=name, filter=filter, @@ -1909,7 +1899,7 @@ def purge( config=config, ) if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -1932,11 +1922,11 @@ def ingest_events( ] = None, config: Optional[types.IngestEventsConfigOrDict] = None, ) -> types.MemoryBankIngestEventsOperation: - """Ingests events into an Agent Engine. + """Ingests events into an Agent Runtime. Example usage: ``` - client.agent_engines.memories.ingest_events( + client.runtimes.memories.ingest_events( name="projects/test-project/locations/us-central1/reasoningEngines/test-agent-engine", scope={"user_id": "test-user-id"}, direct_contents_source={ @@ -1961,7 +1951,7 @@ def ingest_events( Args: name (str): - Required. The name of the Agent Engine to ingest events into. + Required. The name of the Agent Runtime to ingest events into. scope (dict[str, str]): Required. The scope of the events to ingest. For example, {"user_id": "123"}. @@ -1976,7 +1966,7 @@ def ingest_events( Optional. The configuration for the ingest events operation. Returns: - AgentEngineIngestEventsOperation: + RuntimeIngestEventsOperation: The operation for ingesting the events. """ if config is None: @@ -1992,7 +1982,7 @@ def ingest_events( config=config, ) if config.wait_for_completion and not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -2010,13 +2000,13 @@ async def _create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: + config: Optional[types.RuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: """ - Creates a new memory in the Agent Engine. + Creates a new memory in the Agent Runtime. """ - parameter_model = types._CreateAgentEngineMemoryRequestParameters( + parameter_model = types._CreateRuntimeMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -2029,7 +2019,7 @@ async def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _CreateRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2060,7 +2050,7 @@ async def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.RuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2088,21 +2078,21 @@ async def delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineMemoryConfigOrDict] = None, - ) -> types.DeleteAgentEngineMemoryOperation: + config: Optional[types.DeleteRuntimeMemoryConfigOrDict] = None, + ) -> types.DeleteRuntimeMemoryOperation: """ - Delete an Agent Engine memory. + Delete an Agent Runtime memory. Args: name (str): - Required. The name of the Agent Engine memory to be deleted. Format: + Required. The name of the Agent Runtime memory to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory}`. - config (DeleteAgentEngineMemoryConfig): - Optional. Additional configurations for deleting the Agent Engine. + config (DeleteRuntimeMemoryConfig): + Optional. Additional configurations for deleting the Agent Runtime. """ - parameter_model = types._DeleteAgentEngineMemoryRequestParameters( + parameter_model = types._DeleteRuntimeMemoryRequestParameters( name=name, config=config, ) @@ -2113,7 +2103,7 @@ async def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _DeleteRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2144,7 +2134,7 @@ async def delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineMemoryOperation._from_response( + return_value = types.DeleteRuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2182,13 +2172,13 @@ async def _generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: + config: Optional[types.GenerateRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: """ - Generates memories for an Agent Engine. + Generates memories for an Agent Runtime. """ - parameter_model = types._GenerateAgentEngineMemoriesRequestParameters( + parameter_model = types._GenerateRuntimeMemoriesRequestParameters( name=name, vertex_session_source=vertex_session_source, direct_contents_source=direct_contents_source, @@ -2203,7 +2193,7 @@ async def _generate( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GenerateAgentEngineMemoriesRequestParameters_to_vertex( + request_dict = _GenerateRuntimeMemoriesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2234,7 +2224,7 @@ async def _generate( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineGenerateMemoriesOperation._from_response( + return_value = types.RuntimeGenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -2259,13 +2249,10 @@ async def _generate( return return_value async def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineMemoryConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeMemoryConfigOrDict] = None ) -> types.Memory: """ - Gets an agent engine memory. + Gets an agent runtime memory. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -2274,7 +2261,7 @@ async def get( """ - parameter_model = types._GetAgentEngineMemoryRequestParameters( + parameter_model = types._GetRuntimeMemoryRequestParameters( name=name, config=config, ) @@ -2285,9 +2272,7 @@ async def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryRequestParameters_to_vertex( - parameter_model - ) + request_dict = _GetRuntimeMemoryRequestParameters_to_vertex(parameter_model) request_url_dict = request_dict.get("_url") if request_url_dict: path = "{name}".format_map(request_url_dict) @@ -2427,16 +2412,13 @@ async def _ingest_events( return return_value async def _list( - self, - *, - name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, + self, *, name: str, config: Optional[types.ListRuntimeMemoryConfigOrDict] = None ) -> types.ListReasoningEnginesMemoriesResponse: """ - Lists Agent Engine memories. + Lists Agent Runtime memories. """ - parameter_model = types._ListAgentEngineMemoryRequestParameters( + parameter_model = types._ListRuntimeMemoryRequestParameters( name=name, config=config, ) @@ -2447,7 +2429,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _ListRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2506,9 +2488,9 @@ async def _get_memory_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - parameter_model = types._GetAgentEngineMemoryOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: + parameter_model = types._GetRuntimeMemoryOperationParameters( operation_name=operation_name, config=config, ) @@ -2519,7 +2501,7 @@ async def _get_memory_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryOperationParameters_to_vertex( + request_dict = _GetRuntimeMemoryOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2550,7 +2532,7 @@ async def _get_memory_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.RuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2578,9 +2560,9 @@ async def _get_generate_memories_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - parameter_model = types._GetAgentEngineGenerateMemoriesOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: + parameter_model = types._GetRuntimeGenerateMemoriesOperationParameters( operation_name=operation_name, config=config, ) @@ -2591,7 +2573,7 @@ async def _get_generate_memories_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineGenerateMemoriesOperationParameters_to_vertex( + request_dict = _GetRuntimeGenerateMemoriesOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2622,7 +2604,7 @@ async def _get_generate_memories_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineGenerateMemoriesOperation._from_response( + return_value = types.RuntimeGenerateMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -2657,13 +2639,13 @@ async def _retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveRuntimeMemoriesConfigOrDict] = None, ) -> types.RetrieveMemoriesResponse: """ - Retrieves memories for an Agent Engine. + Retrieves memories for an Agent Runtime. """ - parameter_model = types._RetrieveAgentEngineMemoriesRequestParameters( + parameter_model = types._RetrieveRuntimeMemoriesRequestParameters( name=name, scope=scope, similarity_search_params=similarity_search_params, @@ -2677,7 +2659,7 @@ async def _retrieve( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RetrieveAgentEngineMemoriesRequestParameters_to_vertex( + request_dict = _RetrieveRuntimeMemoriesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2740,13 +2722,13 @@ async def retrieve_profiles( config: Optional[types.RetrieveMemoryProfilesConfigOrDict] = None, ) -> types.RetrieveProfilesResponse: """ - Retrieves memory profiles for an Agent Engine. + Retrieves memory profiles for an Agent Runtime. For example, you can use the following code to retrieve all memory profiles for scope `{'user_id': '123'}`: ```python - result = client.agent_engines.memories.retrieve_profiles( + result = client.runtimes.memories.retrieve_profiles( name="projects/123/locations/us-central1/reasoningEngines/456", scope={"user_id": "123"} ) @@ -2841,13 +2823,13 @@ async def _rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineRollbackMemoryOperation: + config: Optional[types.RollbackRuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeRollbackMemoryOperation: """ Rollback a memory to a previous revision. """ - parameter_model = types._RollbackAgentEngineMemoryRequestParameters( + parameter_model = types._RollbackRuntimeMemoryRequestParameters( name=name, target_revision_id=target_revision_id, config=config, @@ -2859,7 +2841,7 @@ async def _rollback( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _RollbackAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _RollbackRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2890,7 +2872,7 @@ async def _rollback( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineRollbackMemoryOperation._from_response( + return_value = types.RuntimeRollbackMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -2920,13 +2902,13 @@ async def _update( name: str, fact: Optional[str] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.UpdateAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: + config: Optional[types.UpdateRuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: """ - Updates an Agent Engine memory. + Updates an Agent Runtime memory. """ - parameter_model = types._UpdateAgentEngineMemoryRequestParameters( + parameter_model = types._UpdateRuntimeMemoryRequestParameters( name=name, fact=fact, scope=scope, @@ -2939,7 +2921,7 @@ async def _update( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _UpdateAgentEngineMemoryRequestParameters_to_vertex( + request_dict = _UpdateRuntimeMemoryRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -2970,7 +2952,7 @@ async def _update( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineMemoryOperation._from_response( + return_value = types.RuntimeMemoryOperation._from_response( response=response_dict, kwargs=( { @@ -3003,13 +2985,13 @@ async def _purge( builtins.list[types.MemoryConjunctionFilterOrDict] ] = None, force: Optional[bool] = None, - config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEnginePurgeMemoriesOperation: + config: Optional[types.PurgeRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimePurgeMemoriesOperation: """ - Purges memories from an Agent Engine. + Purges memories from an Agent Runtime. """ - parameter_model = types._PurgeAgentEngineMemoriesRequestParameters( + parameter_model = types._PurgeRuntimeMemoriesRequestParameters( name=name, filter=filter, filter_groups=filter_groups, @@ -3023,7 +3005,7 @@ async def _purge( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _PurgeAgentEngineMemoriesRequestParameters_to_vertex( + request_dict = _PurgeRuntimeMemoriesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -3054,7 +3036,7 @@ async def _purge( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEnginePurgeMemoriesOperation._from_response( + return_value = types.RuntimePurgeMemoriesOperation._from_response( response=response_dict, kwargs=( { @@ -3091,7 +3073,7 @@ def revisions(self) -> "memory_revisions_module.AsyncMemoryRevisions": ) except ImportError as e: raise ImportError( - "The 'agent_engines.memories.revisions' module requires " + "The 'runtimes.memories.revisions' module requires " "additional packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e @@ -3103,9 +3085,9 @@ async def create( name: str, fact: str, scope: dict[str, str], - config: Optional[types.AgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineMemoryOperation: - """Creates a new memory in the Agent Engine. + config: Optional[types.RuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: + """Creates a new memory in the Agent Runtime. Args: name (str): @@ -3114,16 +3096,16 @@ async def create( Required. The fact to be stored in the memory. scope (dict[str, str]): Required. The scope of the memory. For example, {"user_id": "123"}. - config (AgentEngineMemoryConfigOrDict): + config (RuntimeMemoryConfigOrDict): Optional. The configuration for the memory. Returns: - AgentEngineMemoryOperation: The operation for creating the memory. + RuntimeMemoryOperation: The operation for creating the memory. """ if config is None: - config = types.AgentEngineMemoryConfig() + config = types.RuntimeMemoryConfig() elif isinstance(config, dict): - config = types.AgentEngineMemoryConfig.model_validate(config) + config = types.RuntimeMemoryConfig.model_validate(config) operation = await self._create( name=name, fact=fact, @@ -3132,7 +3114,7 @@ async def create( ) if config.wait_for_completion: if not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -3161,13 +3143,13 @@ async def generate( types.GenerateMemoriesRequestDirectMemoriesSourceOrDict ] = None, scope: Optional[dict[str, str]] = None, - config: Optional[types.GenerateAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEngineGenerateMemoriesOperation: - """Generates memories for the agent engine. + config: Optional[types.GenerateRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: + """Generates memories for the agent runtime. Args: name (str): - Required. The name of the agent engine to generate memories for. + Required. The name of the agent runtime to generate memories for. vertex_session_source (GenerateMemoriesRequestVertexSessionSource): Optional. The vertex session source to use for generating memories. Only one of vertex_session_source, @@ -3188,13 +3170,13 @@ async def generate( Optional. The configuration for the memories to generate. Returns: - AgentEngineGenerateMemoriesOperation: + RuntimeGenerateMemoriesOperation: The operation for generating the memories. """ if config is None: - config = types.GenerateAgentEngineMemoriesConfig() + config = types.GenerateRuntimeMemoriesConfig() elif isinstance(config, dict): - config = types.GenerateAgentEngineMemoriesConfig.model_validate(config) + config = types.GenerateRuntimeMemoriesConfig.model_validate(config) operation = await self._generate( name=name, vertex_session_source=vertex_session_source, @@ -3204,7 +3186,7 @@ async def generate( config=config, ) if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_generate_memories_operation, poll_interval_seconds=0.5, @@ -3217,14 +3199,14 @@ async def list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryConfigOrDict] = None, + config: Optional[types.ListRuntimeMemoryConfigOrDict] = None, ) -> AsyncPager[types.Memory]: - """Lists Agent Engine memories. + """Lists Agent Runtime memories. Args: name (str): - Required. The name of the agent engine to list memories for. - config (ListAgentEngineMemoryConfig): + Required. The name of the agent runtime to list memories for. + config (ListRuntimeMemoryConfig): Optional. The configuration for the memories to list. Returns: @@ -3249,13 +3231,13 @@ async def retrieve( simple_retrieval_params: Optional[ types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict ] = None, - config: Optional[types.RetrieveAgentEngineMemoriesConfigOrDict] = None, + config: Optional[types.RetrieveRuntimeMemoriesConfigOrDict] = None, ) -> AsyncPager[types.RetrieveMemoriesResponseRetrievedMemory]: """Retrieves memories for the agent. Args: name (str): - Required. The name of the agent engine to retrieve memories for. + Required. The name of the agent runtime to retrieve memories for. scope (dict[str, str]): Required. The scope of the memories to retrieve. For example, {"user_id": "123"}. @@ -3265,7 +3247,7 @@ async def retrieve( simple_retrieval_params (RetrieveMemoriesRequestSimpleRetrievalParams): Optional. The simple retrieval parameters to use for retrieving memories. - config (RetrieveAgentEngineMemoriesConfig): + config (RetrieveRuntimeMemoriesConfig): Optional. The configuration for the memories to retrieve. Returns: @@ -3296,8 +3278,8 @@ async def rollback( *, name: str, target_revision_id: str, - config: Optional[types.RollbackAgentEngineMemoryConfigOrDict] = None, - ) -> types.AgentEngineRollbackMemoryOperation: + config: Optional[types.RollbackRuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeRollbackMemoryOperation: """Rolls back a memory to a previous revision. Args: @@ -3305,24 +3287,24 @@ async def rollback( Required. The name of the memory to rollback. target_revision_id (str): Required. The revision ID to roll back to - config (RollbackAgentEngineMemoryConfig): + config (RollbackRuntimeMemoryConfig): Optional. The configuration for the rollback. Returns: - AgentEngineRollbackMemoryOperation: + RuntimeRollbackMemoryOperation: The operation for rolling back the memory. """ if config is None: - config = types.RollbackAgentEngineMemoryConfig() + config = types.RollbackRuntimeMemoryConfig() elif isinstance(config, dict): - config = types.RollbackAgentEngineMemoryConfig.model_validate(config) + config = types.RollbackRuntimeMemoryConfig.model_validate(config) operation = await self._rollback( name=name, target_revision_id=target_revision_id, config=config, ) if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -3338,13 +3320,13 @@ async def purge( filter: Optional[str] = None, filter_groups: Optional[List[types.MemoryConjunctionFilter]] = None, force: bool = False, - config: Optional[types.PurgeAgentEngineMemoriesConfigOrDict] = None, - ) -> types.AgentEnginePurgeMemoriesOperation: - """Purges memories from an Agent Engine. + config: Optional[types.PurgeRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimePurgeMemoriesOperation: + """Purges memories from an Agent Runtime. Args: name (str): - Required. The name of the Agent Engine to purge memories from. + Required. The name of the Agent Runtime to purge memories from. filter (str): Optional. The standard list filter to determine which memories to purge. filter_groups (list[MemoryConjunctionFilter]): @@ -3354,17 +3336,17 @@ async def purge( force (bool): Optional. Whether to force the purge operation. If false, the operation will be staged but not executed. - config (PurgeAgentEngineMemoriesConfig): + config (PurgeRuntimeMemoriesConfig): Optional. The configuration for the purge operation. Returns: - AgentEnginePurgeMemoriesOperation: + RuntimePurgeMemoriesOperation: The operation for purging the memories. """ if config is None: - config = types.PurgeAgentEngineMemoriesConfig() + config = types.PurgeRuntimeMemoriesConfig() elif isinstance(config, dict): - config = types.PurgeAgentEngineMemoriesConfig.model_validate(config) + config = types.PurgeRuntimeMemoriesConfig.model_validate(config) operation = await self._purge( name=name, filter=filter, @@ -3373,7 +3355,7 @@ async def purge( config=config, ) if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, @@ -3396,11 +3378,11 @@ async def ingest_events( ] = None, config: Optional[types.IngestEventsConfigOrDict] = None, ) -> types.MemoryBankIngestEventsOperation: - """Ingests events into an Agent Engine. + """Ingests events into an Agent Runtime. Example usage: ``` - await client.aio.agent_engines.memories.ingest_events( + await client.aio.runtimes.memories.ingest_events( name="projects/test-project/locations/us-central1/reasoningEngines/test-agent-engine", scope={"user_id": "test-user-id"}, direct_contents_source={ @@ -3425,7 +3407,7 @@ async def ingest_events( Args: name (str): - Required. The name of the Agent Engine to ingest events into. + Required. The name of the Agent Runtime to ingest events into. scope (dict[str, str]): Required. The scope of the events to ingest. For example, {"user_id": "123"}. @@ -3440,7 +3422,7 @@ async def ingest_events( Optional. The configuration for the ingest events operation. Returns: - AgentEngineIngestEventsOperation: + RuntimeIngestEventsOperation: The operation for ingesting the events. """ if config is None: @@ -3456,7 +3438,7 @@ async def ingest_events( config=config, ) if config.wait_for_completion and not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_memory_operation, poll_interval_seconds=0.5, diff --git a/agentplatform/_genai/memory_revisions.py b/agentplatform/_genai/memory_revisions.py index c1a1e307b1..910b44ae53 100644 --- a/agentplatform/_genai/memory_revisions.py +++ b/agentplatform/_genai/memory_revisions.py @@ -34,7 +34,7 @@ logger.setLevel(logging.INFO) -def _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( +def _GetRuntimeMemoryRevisionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -45,7 +45,7 @@ def _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( return to_object -def _ListAgentEngineMemoryRevisionsConfig_to_vertex( +def _ListRuntimeMemoryRevisionsConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -63,7 +63,7 @@ def _ListAgentEngineMemoryRevisionsConfig_to_vertex( return to_object -def _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( +def _ListRuntimeMemoryRevisionsRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -72,7 +72,7 @@ def _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineMemoryRevisionsConfig_to_vertex( + _ListRuntimeMemoryRevisionsConfig_to_vertex( getv(from_object, ["config"]), to_object ) @@ -85,23 +85,23 @@ def get( self, *, name: str, - config: Optional[types.GetAgentEngineMemoryRevisionConfigOrDict] = None, + config: Optional[types.GetRuntimeMemoryRevisionConfigOrDict] = None, ) -> types.MemoryRevision: """ - Gets an agent engine memory revision. + Gets an agent runtime memory revision. Args: - name (str): Required. The name of the Agent Engine memory revision to get. Format: + name (str): Required. The name of the Agent Runtime memory revision to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}/revisions/{revision_id}`. - config (GetAgentEngineMemoryRevisionConfig): - Optional. Additional configurations for getting the Agent Engine memory revision. + config (GetRuntimeMemoryRevisionConfig): + Optional. Additional configurations for getting the Agent Runtime memory revision. Returns: - AgentEngineMemoryRevision: The requested Agent Engine memory revision. + RuntimeMemoryRevision: The requested Agent Runtime memory revision. """ - parameter_model = types._GetAgentEngineMemoryRevisionRequestParameters( + parameter_model = types._GetRuntimeMemoryRevisionRequestParameters( name=name, config=config, ) @@ -112,7 +112,7 @@ def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( + request_dict = _GetRuntimeMemoryRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -169,23 +169,23 @@ def _list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, - ) -> types.ListAgentEngineMemoryRevisionsResponse: + config: Optional[types.ListRuntimeMemoryRevisionsConfigOrDict] = None, + ) -> types.ListRuntimeMemoryRevisionsResponse: """ - Lists Agent Engine memory revisions. + Lists Agent Runtime memory revisions. Args: - name (str): Required. The name of the Agent Engine memory to list revisions for. Format: + name (str): Required. The name of the Agent Runtime memory to list revisions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}`. - config (ListAgentEngineMemoryRevisionsConfig): - Optional. Additional configurations for listing the Agent Engine memory revisions. + config (ListRuntimeMemoryRevisionsConfig): + Optional. Additional configurations for listing the Agent Runtime memory revisions. Returns: - ListAgentEngineMemoryRevisionsResponse: The requested Agent Engine memory revisions. + ListRuntimeMemoryRevisionsResponse: The requested Agent Runtime memory revisions. """ - parameter_model = types._ListAgentEngineMemoryRevisionsRequestParameters( + parameter_model = types._ListRuntimeMemoryRevisionsRequestParameters( name=name, config=config, ) @@ -196,7 +196,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( + request_dict = _ListRuntimeMemoryRevisionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -225,7 +225,7 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineMemoryRevisionsResponse._from_response( + return_value = types.ListRuntimeMemoryRevisionsResponse._from_response( response=response_dict, kwargs=( { @@ -253,14 +253,14 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, + config: Optional[types.ListRuntimeMemoryRevisionsConfigOrDict] = None, ) -> Iterator[types.MemoryRevision]: - """Lists Agent Engine memory revisions. + """Lists Agent Runtime memory revisions. Args: name (str): Required. The name of the Memory to list revisions for. - config (ListAgentEngineMemoryRevisionsConfigOrDict): + config (ListRuntimeMemoryRevisionsConfigOrDict): Optional. The configuration for the memories to list revisions. Returns: @@ -281,23 +281,23 @@ async def get( self, *, name: str, - config: Optional[types.GetAgentEngineMemoryRevisionConfigOrDict] = None, + config: Optional[types.GetRuntimeMemoryRevisionConfigOrDict] = None, ) -> types.MemoryRevision: """ - Gets an agent engine memory revision. + Gets an agent runtime memory revision. Args: - name (str): Required. The name of the Agent Engine memory revision to get. Format: + name (str): Required. The name of the Agent Runtime memory revision to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}/revisions/{revision_id}`. - config (GetAgentEngineMemoryRevisionConfig): - Optional. Additional configurations for getting the Agent Engine memory revision. + config (GetRuntimeMemoryRevisionConfig): + Optional. Additional configurations for getting the Agent Runtime memory revision. Returns: - AgentEngineMemoryRevision: The requested Agent Engine memory revision. + RuntimeMemoryRevision: The requested Agent Runtime memory revision. """ - parameter_model = types._GetAgentEngineMemoryRevisionRequestParameters( + parameter_model = types._GetRuntimeMemoryRevisionRequestParameters( name=name, config=config, ) @@ -308,7 +308,7 @@ async def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineMemoryRevisionRequestParameters_to_vertex( + request_dict = _GetRuntimeMemoryRevisionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -367,23 +367,23 @@ async def _list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, - ) -> types.ListAgentEngineMemoryRevisionsResponse: + config: Optional[types.ListRuntimeMemoryRevisionsConfigOrDict] = None, + ) -> types.ListRuntimeMemoryRevisionsResponse: """ - Lists Agent Engine memory revisions. + Lists Agent Runtime memory revisions. Args: - name (str): Required. The name of the Agent Engine memory to list revisions for. Format: + name (str): Required. The name of the Agent Runtime memory to list revisions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/memories/{memory_id}`. - config (ListAgentEngineMemoryRevisionsConfig): - Optional. Additional configurations for listing the Agent Engine memory revisions. + config (ListRuntimeMemoryRevisionsConfig): + Optional. Additional configurations for listing the Agent Runtime memory revisions. Returns: - ListAgentEngineMemoryRevisionsResponse: The requested Agent Engine memory revisions. + ListRuntimeMemoryRevisionsResponse: The requested Agent Runtime memory revisions. """ - parameter_model = types._ListAgentEngineMemoryRevisionsRequestParameters( + parameter_model = types._ListRuntimeMemoryRevisionsRequestParameters( name=name, config=config, ) @@ -394,7 +394,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineMemoryRevisionsRequestParameters_to_vertex( + request_dict = _ListRuntimeMemoryRevisionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -425,7 +425,7 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineMemoryRevisionsResponse._from_response( + return_value = types.ListRuntimeMemoryRevisionsResponse._from_response( response=response_dict, kwargs=( { @@ -453,14 +453,14 @@ async def list( self, *, name: str, - config: Optional[types.ListAgentEngineMemoryRevisionsConfigOrDict] = None, + config: Optional[types.ListRuntimeMemoryRevisionsConfigOrDict] = None, ) -> AsyncPager[types.MemoryRevision]: - """Lists Agent Engine memory revisions. + """Lists Agent Runtime memory revisions. Args: name (str): Required. The name of the Memory to list revisions for. - config (ListAgentEngineMemoryRevisionsConfigOrDict): + config (ListRuntimeMemoryRevisionsConfigOrDict): Optional. The configuration for the memories to list revisions. Returns: diff --git a/agentplatform/_genai/prompt_optimizer.py b/agentplatform/_genai/prompt_optimizer.py index 3aab5f5666..b01b65187f 100644 --- a/agentplatform/_genai/prompt_optimizer.py +++ b/agentplatform/_genai/prompt_optimizer.py @@ -100,6 +100,22 @@ def _CustomJob_from_vertex( if getv(from_object, ["webAccessUris"]) is not None: setv(to_object, ["web_access_uris"], getv(from_object, ["webAccessUris"])) + if ( + getv(from_object, ["placeholderToMakeTrialInfoOneofWorkInGeneratedProto"]) + is not None + ): + setv( + to_object, + ["placeholder_to_make_trial_info_oneof_work_in_generated_proto"], + getv(from_object, ["placeholderToMakeTrialInfoOneofWorkInGeneratedProto"]), + ) + + if getv(from_object, ["description"]) is not None: + setv(to_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["pubsubTopic"]) is not None: + setv(to_object, ["pubsub_topic"], getv(from_object, ["pubsubTopic"])) + return to_object @@ -150,6 +166,28 @@ def _CustomJob_to_vertex( if getv(from_object, ["web_access_uris"]) is not None: setv(to_object, ["webAccessUris"], getv(from_object, ["web_access_uris"])) + if ( + getv( + from_object, + ["placeholder_to_make_trial_info_oneof_work_in_generated_proto"], + ) + is not None + ): + setv( + to_object, + ["placeholderToMakeTrialInfoOneofWorkInGeneratedProto"], + getv( + from_object, + ["placeholder_to_make_trial_info_oneof_work_in_generated_proto"], + ), + ) + + if getv(from_object, ["description"]) is not None: + setv(to_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["pubsub_topic"]) is not None: + setv(to_object, ["pubsubTopic"], getv(from_object, ["pubsub_topic"])) + return to_object diff --git a/agentplatform/_genai/prompts.py b/agentplatform/_genai/prompts.py index 86bd483728..0d48721081 100644 --- a/agentplatform/_genai/prompts.py +++ b/agentplatform/_genai/prompts.py @@ -164,6 +164,22 @@ def _CustomJob_from_vertex( if getv(from_object, ["webAccessUris"]) is not None: setv(to_object, ["web_access_uris"], getv(from_object, ["webAccessUris"])) + if ( + getv(from_object, ["placeholderToMakeTrialInfoOneofWorkInGeneratedProto"]) + is not None + ): + setv( + to_object, + ["placeholder_to_make_trial_info_oneof_work_in_generated_proto"], + getv(from_object, ["placeholderToMakeTrialInfoOneofWorkInGeneratedProto"]), + ) + + if getv(from_object, ["description"]) is not None: + setv(to_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["pubsubTopic"]) is not None: + setv(to_object, ["pubsub_topic"], getv(from_object, ["pubsubTopic"])) + return to_object @@ -214,6 +230,28 @@ def _CustomJob_to_vertex( if getv(from_object, ["web_access_uris"]) is not None: setv(to_object, ["webAccessUris"], getv(from_object, ["web_access_uris"])) + if ( + getv( + from_object, + ["placeholder_to_make_trial_info_oneof_work_in_generated_proto"], + ) + is not None + ): + setv( + to_object, + ["placeholderToMakeTrialInfoOneofWorkInGeneratedProto"], + getv( + from_object, + ["placeholder_to_make_trial_info_oneof_work_in_generated_proto"], + ), + ) + + if getv(from_object, ["description"]) is not None: + setv(to_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["pubsub_topic"]) is not None: + setv(to_object, ["pubsubTopic"], getv(from_object, ["pubsub_topic"])) + return to_object diff --git a/agentplatform/_genai/rag.py b/agentplatform/_genai/rag.py index 5841da80f1..e3902f0396 100644 --- a/agentplatform/_genai/rag.py +++ b/agentplatform/_genai/rag.py @@ -962,6 +962,9 @@ def _RagFileParsingConfigLlmParser_from_vertex( if getv(from_object, ["modelName"]) is not None: setv(to_object, ["model_name"], getv(from_object, ["modelName"])) + if getv(from_object, ["memoryMode"]) is not None: + setv(to_object, ["memory_mode"], getv(from_object, ["memoryMode"])) + return to_object @@ -994,6 +997,9 @@ def _RagFileParsingConfigLlmParser_to_vertex( if getv(from_object, ["model_name"]) is not None: setv(to_object, ["modelName"], getv(from_object, ["model_name"])) + if getv(from_object, ["memory_mode"]) is not None: + setv(to_object, ["memoryMode"], getv(from_object, ["memory_mode"])) + return to_object @@ -1594,6 +1600,23 @@ def _UploadRagFileConfig_to_vertex( getv(from_object, ["rag_file_transformation_config"]), ) + if getv(from_object, ["inline_metadata"]) is not None: + setv(to_object, ["inlineMetadata"], getv(from_object, ["inline_metadata"])) + + if getv(from_object, ["max_embedding_requests_per_min"]) is not None: + setv( + to_object, + ["maxEmbeddingRequestsPerMin"], + getv(from_object, ["max_embedding_requests_per_min"]), + ) + + if getv(from_object, ["global_max_embedding_requests_per_min"]) is not None: + setv( + to_object, + ["globalMaxEmbeddingRequestsPerMin"], + getv(from_object, ["global_max_embedding_requests_per_min"]), + ) + return to_object diff --git a/agentplatform/_genai/runtime_revisions.py b/agentplatform/_genai/runtime_revisions.py index e4a8d49fd1..e0098a2322 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,21 +528,21 @@ 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( + agent_engine_runtime_revision = types.RuntimeRevision( api_client=self, api_async_client=AsyncRuntimeRevisions(api_client_=self._api_client), api_resource=api_resource, @@ -565,18 +557,18 @@ 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 +578,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 +590,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,23 +626,23 @@ 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.""" + agent_engine_runtime_revision: types.RuntimeRevision, + ) -> types.RuntimeRevision: + """Registers the API methods for the agent runtime revision.""" try: - _agent_engines_utils._register_api_methods_or_raise( + _runtimes_utils._register_api_methods_or_raise( agent_engine=agent_engine_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 @@ -658,16 +650,14 @@ 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 +686,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 +726,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 +743,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 +802,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 +810,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 +818,7 @@ async def _list( """ - parameter_model = types._ListAgentEngineRuntimeRevisionsRequestParameters( + parameter_model = types._ListRuntimeRevisionsRequestParameters( name=name, config=config, ) @@ -841,7 +829,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 +890,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 +917,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 +948,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 +976,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 +989,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 +1020,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 +1048,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 +1065,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,21 +1124,21 @@ 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( + agent_engine_runtime_revision = types.RuntimeRevision( api_client=self, api_async_client=AsyncRuntimeRevisions(api_client_=self._api_client), api_resource=api_resource, @@ -1169,18 +1153,18 @@ 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 +1174,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 +1184,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 +1220,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.""" + agent_engine_runtime_revision: types.RuntimeRevision, + ) -> types.RuntimeRevision: + """Registers the API methods for the agent runtime revision.""" try: - _agent_engines_utils._register_api_methods_or_raise( + _runtimes_utils._register_api_methods_or_raise( agent_engine=agent_engine_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 diff --git a/agentplatform/_genai/runtimes.py b/agentplatform/_genai/runtimes.py index 9f41816bf5..cc67bef82a 100644 --- a/agentplatform/_genai/runtimes.py +++ b/agentplatform/_genai/runtimes.py @@ -13,20 +13,33 @@ # limitations under the License. # -# Handwritten placeholder code for the runtimes.py file. -# Should be replaced by the generated file. +# Code generated by the Google Gen AI SDK generator DO NOT EDIT. +import builtins import importlib +import json import logging import typing +from typing import Any, AsyncIterator, Iterator, Optional, Sequence, Tuple, Union +from urllib.parse import urlencode +import warnings from google.genai import _api_module +from google.genai import _common +from google.genai import types as genai_types +from google.genai._common import get_value_by_path as getv +from google.genai._common import set_value_by_path as setv +from google.genai.pagers import Pager +from . import _runtimes_utils +from . import types if typing.TYPE_CHECKING: + from . import memories as memories_module from . import runtime_revisions as runtime_revisions_module - _ = runtime_revisions_module + _ = memories_module + __ = runtime_revisions_module logger = logging.getLogger("agentplatform_genai.runtimes") @@ -34,8 +47,1488 @@ 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 + ), + ) + + if getv(from_object, ["exampleStoreConfig"]) is not None: + setv( + to_object, + ["example_store_config"], + getv(from_object, ["exampleStoreConfig"]), + ) + + 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 + ), + ) + + if getv(from_object, ["example_store_config"]) is not None: + setv( + to_object, + ["exampleStoreConfig"], + getv(from_object, ["example_store_config"]), + ) + + 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"])) + + if getv(from_object, ["experimentConfig"]) is not None: + setv(to_object, ["experiment_config"], getv(from_object, ["experimentConfig"])) + + if getv(from_object, ["revisionGarbageCollectionStrategy"]) is not None: + setv( + to_object, + ["revision_garbage_collection_strategy"], + getv(from_object, ["revisionGarbageCollectionStrategy"]), + ) + + if getv(from_object, ["url"]) is not None: + setv(to_object, ["url"], getv(from_object, ["url"])) + + return to_object + + +def _RunQueryJobRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["input_gcs_uri"]) is not None: + setv(parent_object, ["inputGcsUri"], getv(from_object, ["input_gcs_uri"])) + + if getv(from_object, ["output_gcs_uri"]) is not None: + setv(parent_object, ["outputGcsUri"], getv(from_object, ["output_gcs_uri"])) + + return to_object + + +def _RunQueryJobRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + setv( + to_object, + ["config"], + _RunQueryJobRuntimeConfig_to_vertex( + getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _RuntimeOperation_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["metadata"]) is not None: + setv(to_object, ["metadata"], getv(from_object, ["metadata"])) + + if getv(from_object, ["done"]) is not None: + setv(to_object, ["done"], getv(from_object, ["done"])) + + if getv(from_object, ["error"]) is not None: + setv(to_object, ["error"], getv(from_object, ["error"])) + + if getv(from_object, ["response"]) is not None: + setv( + to_object, + ["response"], + _ReasoningEngine_from_vertex(getv(from_object, ["response"]), to_object), + ) + + return to_object + + +def _StructuredMemoryConfig_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["schemaConfigs"]) is not None: + setv( + to_object, + ["schema_configs"], + [ + _StructuredMemorySchemaConfig_from_vertex(item, to_object) + for item in getv(from_object, ["schemaConfigs"]) + ], + ) + + if getv(from_object, ["scopeKeys"]) is not None: + setv(to_object, ["scope_keys"], getv(from_object, ["scopeKeys"])) + + return to_object + + +def _StructuredMemoryConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["schema_configs"]) is not None: + setv( + to_object, + ["schemaConfigs"], + [ + _StructuredMemorySchemaConfig_to_vertex(item, to_object) + for item in getv(from_object, ["schema_configs"]) + ], + ) + + if getv(from_object, ["scope_keys"]) is not None: + setv(to_object, ["scopeKeys"], getv(from_object, ["scope_keys"])) + + return to_object + + +def _StructuredMemorySchemaConfig_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["schema"]) is not None: + setv(to_object, ["memory_schema"], getv(from_object, ["schema"])) + + if getv(from_object, ["id"]) is not None: + setv(to_object, ["id"], getv(from_object, ["id"])) + + if getv(from_object, ["memoryType"]) is not None: + setv(to_object, ["memory_type"], getv(from_object, ["memoryType"])) + + return to_object + + +def _StructuredMemorySchemaConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["memory_schema"]) is not None: + setv(to_object, ["schema"], getv(from_object, ["memory_schema"])) + + if getv(from_object, ["id"]) is not None: + setv(to_object, ["id"], getv(from_object, ["id"])) + + if getv(from_object, ["memory_type"]) is not None: + setv(to_object, ["memoryType"], getv(from_object, ["memory_type"])) + + return to_object + + +def _UpdateRuntimeConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["display_name"]) is not None: + setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) + + if getv(from_object, ["description"]) is not None: + setv(parent_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["spec"]) is not None: + setv(parent_object, ["spec"], getv(from_object, ["spec"])) + + if getv(from_object, ["context_spec"]) is not None: + setv( + parent_object, + ["contextSpec"], + _ReasoningEngineContextSpec_to_vertex( + getv(from_object, ["context_spec"]), to_object + ), + ) + + if getv(from_object, ["psc_interface_config"]) is not None: + setv( + parent_object, + ["pscInterfaceConfig"], + getv(from_object, ["psc_interface_config"]), + ) + + if getv(from_object, ["agent_gateway_config"]) is not None: + setv( + parent_object, + ["agentGatewayConfig"], + getv(from_object, ["agent_gateway_config"]), + ) + + if getv(from_object, ["encryption_spec"]) is not None: + setv(parent_object, ["encryptionSpec"], getv(from_object, ["encryption_spec"])) + + if getv(from_object, ["labels"]) is not None: + setv(parent_object, ["labels"], getv(from_object, ["labels"])) + + if getv(from_object, ["source_packages"]) is not None: + setv(parent_object, ["sourcePackages"], getv(from_object, ["source_packages"])) + + if getv(from_object, ["entrypoint_module"]) is not None: + setv( + parent_object, + ["entrypointModule"], + getv(from_object, ["entrypoint_module"]), + ) + + if getv(from_object, ["entrypoint_object"]) is not None: + setv( + parent_object, + ["entrypointObject"], + getv(from_object, ["entrypoint_object"]), + ) + + if getv(from_object, ["requirements_file"]) is not None: + setv( + parent_object, + ["requirementsFile"], + getv(from_object, ["requirements_file"]), + ) + + if getv(from_object, ["agent_framework"]) is not None: + setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"])) + + if getv(from_object, ["python_version"]) is not None: + setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"])) + + if getv(from_object, ["update_mask"]) is not None: + setv( + parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"]) + ) + + if getv(from_object, ["traffic_config"]) is not None: + setv(parent_object, ["trafficConfig"], getv(from_object, ["traffic_config"])) + + return to_object + + +def _UpdateRuntimeRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + _UpdateRuntimeConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + class Runtimes(_api_module.BaseModule): + def cancel_query_job( + self, + *, + name: str, + config: Optional[types.CancelQueryJobRuntimeConfigOrDict] = None, + ) -> types.CancelQueryJobResult: + """ + Cancels a long-running query job on an Agent Runtime. + + Args: + name (str): + Required. The reasoning engine resource name. + config (CancelQueryJobRuntimeConfigOrDict): + Optional. The configuration for the cancel_query_job. + + """ + + parameter_model = types._CancelQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CancelQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:cancelAsyncQuery".format_map(request_url_dict) + else: + path = "{name}:cancelAsyncQuery" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.CancelQueryJobResult._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _check_query_job( + self, + *, + name: str, + config: Optional[types.CheckQueryJobRuntimeConfigOrDict] = None, + ) -> types.CheckQueryJobResult: + """ + Query an Agent Runtime asynchronously. + """ + + parameter_model = types._CheckQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CheckQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:checkQueryJob".format_map(request_url_dict) + else: + path = "{name}:checkQueryJob" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _CheckQueryJobResult_from_vertex(response_dict) + + return_value = types.CheckQueryJobResult._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _run_query_job( + self, + *, + name: str, + config: Optional[types._RunQueryJobRuntimeConfigOrDict] = None, + ) -> types.RuntimeOperation: + """ + Run a query job on an agent runtime. + """ + + parameter_model = types._RunQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _RunQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:asyncQuery".format_map(request_url_dict) + else: + path = "{name}:asyncQuery" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _create( + self, *, config: Optional[types.CreateRuntimeConfigOrDict] = None + ) -> types.RuntimeOperation: + """ + Creates a new Agent Runtime. + """ + + parameter_model = types._CreateRuntimeRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CreateRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _delete( + self, + *, + name: str, + force: Optional[bool] = None, + config: Optional[types.DeleteRuntimeConfigOrDict] = None, + ) -> types.DeleteRuntimeOperation: + """ + Delete an Agent Runtime resource. + + Args: + name (str): + Required. The name of the Agent Runtime to be deleted. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` + or `reasoningEngines/{resource_id}`. + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Agent Runtime has undeleted child resources. Defaults to False. + config (DeleteRuntimeConfig): + Optional. Additional configurations for deleting the Agent Runtime. + + """ + + parameter_model = types._DeleteRuntimeRequestParameters( + name=name, + force=force, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _DeleteRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("delete", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteRuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _get( + self, *, name: str, config: Optional[types.GetRuntimeConfigOrDict] = None + ) -> types.ReasoningEngine: + """ + Get an Agent Runtime instance. + """ + + parameter_model = types._GetRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ReasoningEngine_from_vertex(response_dict) + + return_value = types.ReasoningEngine._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _list( + self, *, config: Optional[types.ListRuntimeConfigOrDict] = None + ) -> types.ListReasoningEnginesResponse: + """ + Lists Agent Runtimes. + """ + + parameter_model = types._ListRuntimeRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _ListRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) + + return_value = types.ListReasoningEnginesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _get_agent_operation( + self, + *, + operation_name: str, + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeOperation: + parameter_model = types._GetRuntimeOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRuntimeOperationParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _query( + self, *, name: str, config: Optional[types.QueryRuntimeConfigOrDict] = None + ) -> types.QueryReasoningEngineResponse: + """ + Query an Agent Runtime. + """ + + parameter_model = types._QueryRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _QueryRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:query".format_map(request_url_dict) + else: + path = "{name}:query" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.QueryReasoningEngineResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _update( + self, *, name: str, config: Optional[types.UpdateRuntimeConfigOrDict] = None + ) -> types.RuntimeOperation: + """ + Updates an Agent Runtime. + """ + + parameter_model = types._UpdateRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpdateRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("patch", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + _memories = None _revisions = None @property @@ -49,17 +1542,2268 @@ def revisions(self) -> "runtime_revisions_module.RuntimeRevisions": ) except ImportError as e: raise ImportError( - "The 'agent_engines.runtimes.revisions' module requires " - "additional packages. Please install them using pip install " + "The 'runtimes.revisions' module requires additional " + "packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e return self._revisions.RuntimeRevisions(self._api_client) # type: ignore[no-any-return] + @property + def memories(self) -> "memories_module.Memories": + if self._memories is None: + try: + # We need to lazy load the memories module to handle the + # possibility of ImportError when dependencies are not installed. + self._memories = importlib.import_module(".memories", __package__) + except ImportError as e: + raise ImportError( + "The 'runtimes.memories' module requires additional " + "packages. Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._memories.Memories(self._api_client) # type: ignore[no-any-return] + + def _list_pager( + self, *, config: Optional[types.ListRuntimeConfigOrDict] = None + ) -> Pager[types.ReasoningEngine]: + return Pager( + "reasoning_engines", + self._list, + self._list(config=config), + config, + ) + + def check_query_job( + self, + *, + name: str, + config: Optional[types.CheckQueryJobRuntimeConfigOrDict] = None, + ) -> types.CheckQueryJobResult: + """Checks a query job on an agent runtime and optionally returns the results. + + Args: + name (str): + Required. A fully-qualified resource name or ID. + config (CheckQueryJobRuntimeConfigOrDict): + Optional. The configuration for the check_query_job. If not provided, + the default configuration will be used. This can be used to specify + the following fields: + - retrieve_result: Whether to retrieve the results of the query job. + """ + from google.cloud import storage # type: ignore[attr-defined] + import json + + if config is None: + config = types.CheckQueryJobRuntimeConfig() + elif isinstance(config, dict): + config = types.CheckQueryJobRuntimeConfig(**config) + + raw_response = self._api_client.request("get", name, {}) + if hasattr(raw_response, "body"): + operation = ( + json.loads(raw_response.body) + if isinstance(raw_response.body, str) + else raw_response.body + ) + else: + operation = raw_response + + status = "RUNNING" + if isinstance(operation, dict): + if operation.get("done"): + status = "FAILED" if operation.get("error") else "SUCCESS" + + response_dict = operation.get("response", {}) + output_gcs_uri = response_dict.get("outputGcsUri") or response_dict.get( + "output_gcs_uri" + ) + error = operation.get("error") + else: + if getattr(operation, "done", False): + status = "FAILED" if getattr(operation, "error", None) else "SUCCESS" + + response_obj = getattr(operation, "response", None) + if isinstance(response_obj, dict): + output_gcs_uri = response_obj.get("outputGcsUri") or response_obj.get( + "output_gcs_uri" + ) + else: + output_gcs_uri = ( + getattr( + response_obj, + "output_gcs_uri", + getattr(response_obj, "outputGcsUri", None), + ) + if response_obj + else None + ) + error = getattr(operation, "error", None) + + result_str = None + if status == "SUCCESS" and config.retrieve_result and output_gcs_uri: + storage_client = storage.Client( + project=self._api_client.project, + credentials=self._api_client._credentials, + ) + bucket_name = output_gcs_uri.replace("gs://", "").split("/")[0] + blob_name = output_gcs_uri.replace(f"gs://{bucket_name}/", "") + bucket = storage_client.bucket(bucket_name) + blob = bucket.blob(blob_name) + if blob.exists(): + result_str = blob.download_as_string().decode("utf-8") + else: + raise ValueError( + f"Failed to retrieve blob results for {output_gcs_uri}" + ) + + elif status == "FAILED" and error: + result_str = str(error) + + return types.CheckQueryJobResult( + operation_name=name, + output_gcs_uri=output_gcs_uri, + status=status, + result=result_str, + ) + + def _is_lightweight_creation(self, agent: Any, config: types.RuntimeConfig) -> 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) + agent_engine = 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(agent_engine=agent_engine) + return agent_engine + + 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, + *, + agent_engine: Any = None, + agent: Any = None, + config: Optional[types.RuntimeConfigOrDict] = None, + ) -> types.Runtime: + """Creates an agent runtime. + + The Agent Runtime 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 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). + agent_engine (Any): + Optional. This is deprecated. Please use `agent` instead. + config (RuntimeConfig): + 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.RuntimeConfig.model_validate(config) + elif not isinstance(config, types.RuntimeConfig): + raise TypeError( + f"config must be a dict or RuntimeConfig, but got {type(config)}." + ) + context_spec = config.context_spec + if context_spec is not None: + # Conversion to a dict for _create_config + context_spec = json.loads(context_spec.model_dump_json()) + developer_connect_source = config.developer_connect_source + if developer_connect_source is not None: + developer_connect_source = json.loads( + developer_connect_source.model_dump_json() + ) + agent_config_source = config.agent_config_source + if agent_config_source is not None: + agent_config_source = json.loads(agent_config_source.model_dump_json()) + keep_alive_probe = config.keep_alive_probe + if keep_alive_probe is not None: + keep_alive_probe = json.loads( + keep_alive_probe.model_dump_json(exclude_none=True) + ) + if agent and agent_engine: + raise ValueError("Please specify only one of `agent` or `agent_engine`.") + elif agent_engine: + raise DeprecationWarning( + "The `agent_engine` argument is deprecated. Please use `agent` instead." + ) + agent = agent or agent_engine + api_config = self._create_config( + mode="create", + agent=agent, + identity_type=config.identity_type, + staging_bucket=config.staging_bucket, + requirements=config.requirements, + display_name=config.display_name, + description=config.description, + gcs_dir_name=config.gcs_dir_name, + extra_packages=config.extra_packages, + env_vars=config.env_vars, + service_account=config.service_account, + context_spec=context_spec, + psc_interface_config=config.psc_interface_config, + agent_gateway_config=config.agent_gateway_config, + min_instances=config.min_instances, + max_instances=config.max_instances, + resource_limits=config.resource_limits, + container_concurrency=config.container_concurrency, + encryption_spec=config.encryption_spec, + agent_server_mode=config.agent_server_mode, + labels=config.labels, + class_methods=config.class_methods, + source_packages=config.source_packages, + developer_connect_source=developer_connect_source, + entrypoint_module=config.entrypoint_module, + entrypoint_object=config.entrypoint_object, + requirements_file=config.requirements_file, + agent_framework=config.agent_framework, + python_version=config.python_version, + build_options=config.build_options, + image_spec=config.image_spec, + agent_config_source=agent_config_source, + container_spec=config.container_spec, + keep_alive_probe=keep_alive_probe, + ) + operation = self._create(config=api_config) + reasoning_engine_id = _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, + ) + + agent_engine = types.Runtime( + api_client=self, + api_async_client=AsyncRuntimes(api_client_=self._api_client), + api_resource=operation.response, + ) + if agent_engine.api_resource: + logger.info("Agent Runtime created. To use it in another session:") + logger.info( + f"agent_engine=client.runtimes.get(name='{agent_engine.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 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 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 agent_engine, + # requirements and extra_packages to `staging_bucket/gcs_dir_name`. + _runtimes_utils._prepare( + agent=agent, + requirements=requirements, + project=project, + location=location, + staging_bucket=staging_bucket, + gcs_dir_name=gcs_dir_name, + extra_packages=extra_packages, + credentials=self._api_client._credentials, + ) + # Update the package spec. + update_masks.append("spec.package_spec.pickle_object_gcs_uri") + package_spec: types.ReasoningEngineSpecPackageSpecDict = { + "python_version": sys_version, + "pickle_object_gcs_uri": "{}/{}/{}".format( + staging_bucket, + gcs_dir_name, + _runtimes_utils._BLOB_FILENAME, + ), + } + if extra_packages: + update_masks.append("spec.package_spec.dependency_files_gcs_uri") + package_spec["dependency_files_gcs_uri"] = "{}/{}/{}".format( + staging_bucket, + gcs_dir_name, + _runtimes_utils._EXTRA_PACKAGES_FILE, + ) + if requirements: + update_masks.append("spec.package_spec.requirements_gcs_uri") + package_spec["requirements_gcs_uri"] = "{}/{}/{}".format( + staging_bucket, + gcs_dir_name, + _runtimes_utils._REQUIREMENTS_FILE, + ) + spec["package_spec"] = package_spec + + update_masks.append("spec.class_methods") + if class_methods is not None: + class_methods_spec_list = ( + _runtimes_utils._class_methods_to_class_methods_spec( + class_methods=class_methods + ) + ) + else: + class_methods_spec_list = ( + _runtimes_utils._generate_class_methods_spec_or_raise( + agent=agent, + operations=_runtimes_utils._get_registered_operations(agent=agent), + ) + ) + spec["class_methods"] = [ + _runtimes_utils._to_dict(class_method_spec) + for class_method_spec in class_methods_spec_list + ] + + def _create_config( + self, + *, + mode: str, + agent: Any = None, + identity_type: Optional[types.IdentityType] = None, + staging_bucket: Optional[str] = None, + requirements: Optional[Union[str, Sequence[str]]] = None, + display_name: Optional[str] = None, + description: Optional[str] = None, + gcs_dir_name: Optional[str] = None, + extra_packages: Optional[Sequence[str]] = None, + env_vars: Optional[dict[str, Union[str, Any]]] = None, + service_account: Optional[str] = None, + context_spec: Optional[types.ReasoningEngineContextSpecDict] = None, + psc_interface_config: Optional[types.PscInterfaceConfigDict] = None, + agent_gateway_config: Optional[ + types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] = None, + min_instances: Optional[int] = None, + max_instances: Optional[int] = None, + resource_limits: Optional[dict[str, str]] = None, + container_concurrency: Optional[int] = None, + encryption_spec: Optional[genai_types.EncryptionSpecDict] = None, + labels: Optional[dict[str, str]] = None, + agent_server_mode: Optional[types.AgentServerMode] = None, + class_methods: Optional[Sequence[dict[str, Any]]] = None, + source_packages: Optional[Sequence[str]] = None, + developer_connect_source: Optional[ + types.ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict + ] = None, + entrypoint_module: Optional[str] = None, + entrypoint_object: Optional[str] = None, + requirements_file: Optional[str] = None, + agent_framework: Optional[str] = None, + python_version: Optional[str] = None, + build_options: Optional[dict[str, builtins.list[str]]] = None, + image_spec: Optional[ + types.ReasoningEngineSpecSourceCodeSpecImageSpecDict + ] = None, + agent_config_source: Optional[ + types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict + ] = None, + container_spec: Optional[types.ReasoningEngineSpecContainerSpecDict] = None, + keep_alive_probe: Optional[dict[str, Any]] = None, + traffic_config: Optional[types.ReasoningEngineTrafficConfigDict] = None, + ) -> types.UpdateRuntimeConfigDict: + import sys + + config: types.UpdateRuntimeConfigDict = {} + update_masks = [] + if mode not in ["create", "update"]: + raise ValueError(f"Unsupported mode: {mode}") + if agent is None: + if requirements is not None: + raise ValueError("requirements must be None if agent is None.") + if extra_packages is not None: + raise ValueError("extra_packages must be None if agent is None.") + if display_name is not None: + update_masks.append("display_name") + config["display_name"] = display_name + if description is not None: + update_masks.append("description") + config["description"] = description + if context_spec is not None: + update_masks.append("context_spec") + config["context_spec"] = context_spec + if encryption_spec is not None: + update_masks.append("encryption_spec") + config["encryption_spec"] = encryption_spec + if labels is not None: + update_masks.append("labels") + config["labels"] = labels + if traffic_config is not None: + update_masks.append("traffic_config") + config["traffic_config"] = traffic_config + + if agent_framework == "google-adk": + env_vars = _runtimes_utils._add_telemetry_enablement_env(env_vars) + + if python_version: + sys_version = python_version + else: + sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" + + if agent: + if source_packages: + raise ValueError( + "If you have provided `source_packages` in `config`, please " + "do not specify `agent` in `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 = _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=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 = ( + _runtimes_utils._class_methods_to_class_methods_spec( + class_methods=class_methods + ) + ) + agent_engine_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") + 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"] = _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 + + agent_engine_spec["agent_card"] = json_format.MessageToDict( + agent_card + ) + except Exception as e: + raise ValueError( + f"Failed to convert agent card to dict (serialization error): {e}" + ) from e + update_masks.append("spec.agent_framework") + + if identity_type is not None or service_account is not None: + if agent_engine_spec is None: + agent_engine_spec = {} + + if identity_type is not None: + agent_engine_spec["identity_type"] = identity_type + update_masks.append("spec.identity_type") + if service_account is not None: + # Clear the field in case of empty service_account. + if service_account: + agent_engine_spec["service_account"] = service_account + update_masks.append("spec.service_account") + + if agent_engine_spec is not None: + config["spec"] = agent_engine_spec + + if update_masks and mode == "update": + config["update_mask"] = ",".join(update_masks) + return config + + def _generate_deployment_spec_or_raise( + self, + *, + env_vars: Optional[dict[str, Union[str, Any]]] = None, + psc_interface_config: Optional[types.PscInterfaceConfigDict] = None, + agent_gateway_config: Optional[ + types.ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] = None, + min_instances: Optional[int] = None, + max_instances: Optional[int] = None, + resource_limits: Optional[dict[str, str]] = None, + container_concurrency: Optional[int] = None, + keep_alive_probe: Optional[dict[str, Any]] = None, + ) -> Tuple[dict[str, Any], Sequence[str]]: + deployment_spec: dict[str, Any] = {} + update_masks = [] + if env_vars: + deployment_spec["env"] = [] + deployment_spec["secret_env"] = [] + if isinstance(env_vars, dict): + self._update_deployment_spec_with_env_vars_dict_or_raise( + deployment_spec=deployment_spec, + env_vars=env_vars, + ) + else: + raise TypeError(f"env_vars must be a dict, but got {type(env_vars)}.") + if deployment_spec.get("env"): + update_masks.append("spec.deployment_spec.env") + if deployment_spec.get("secret_env"): + update_masks.append("spec.deployment_spec.secret_env") + if psc_interface_config: + deployment_spec["psc_interface_config"] = psc_interface_config + update_masks.append("spec.deployment_spec.psc_interface_config") + if agent_gateway_config: + deployment_spec["agent_gateway_config"] = agent_gateway_config + update_masks.append("spec.deployment_spec.agent_gateway_config") + if min_instances is not None: + if not 0 <= min_instances <= 10: + raise ValueError( + f"min_instances must be between 0 and 10. Got {min_instances}" + ) + deployment_spec["min_instances"] = min_instances + update_masks.append("spec.deployment_spec.min_instances") + if max_instances is not None: + if psc_interface_config and not 1 <= max_instances <= 100: + raise ValueError( + f"max_instances must be between 1 and 100 when PSC-I is enabled. Got {max_instances}" + ) + elif not psc_interface_config and not 1 <= max_instances <= 1000: + raise ValueError( + f"max_instances must be between 1 and 1000. Got {max_instances}" + ) + deployment_spec["max_instances"] = max_instances + update_masks.append("spec.deployment_spec.max_instances") + if resource_limits: + _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, + *, + agent_engine: types.Runtime, + ) -> types.Runtime: + """Registers the API methods for the agent runtime.""" + try: + _runtimes_utils._register_api_methods_or_raise( + agent_engine=agent_engine, + 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 agent_engine + + 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, + agent_engine: Any = None, + config: types.RuntimeConfigOrDict, + ) -> 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. + agent_engine (Any): + Optional. This is deprecated. Please use `agent` instead. + config (RuntimeConfig): + 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 `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.RuntimeConfig.model_validate(config) + elif not isinstance(config, types.RuntimeConfig): + raise TypeError( + f"config must be a dict or RuntimeConfig, but got {type(config)}." + ) + context_spec = config.context_spec + if context_spec is not None: + # Conversion to a dict for _create_config + context_spec = json.loads(context_spec.model_dump_json()) + developer_connect_source = config.developer_connect_source + if developer_connect_source is not None: + developer_connect_source = json.loads( + developer_connect_source.model_dump_json() + ) + agent_config_source = config.agent_config_source + if agent_config_source is not None: + agent_config_source = json.loads(agent_config_source.model_dump_json()) + keep_alive_probe = config.keep_alive_probe + if keep_alive_probe is not None: + keep_alive_probe = json.loads( + keep_alive_probe.model_dump_json(exclude_none=True) + ) + traffic_config = config.traffic_config + if traffic_config is not None: + traffic_config = json.loads(traffic_config.model_dump_json()) + if agent and agent_engine: + raise ValueError("Please specify only one of `agent` or `agent_engine`.") + elif agent_engine: + raise DeprecationWarning( + "The `agent_engine` argument is deprecated. Please use `agent` instead." + ) + image_spec = config.image_spec + if image_spec is not None: + # Conversion to a dict for _create_config + image_spec = json.loads(image_spec.model_dump_json()) + container_spec = config.container_spec + if container_spec is not None: + # Conversion to a dict for _create_config + container_spec = json.loads(container_spec.model_dump_json()) + agent = agent or agent_engine + api_config = self._create_config( + mode="update", + agent=agent, + identity_type=config.identity_type, + staging_bucket=config.staging_bucket, + requirements=config.requirements, + display_name=config.display_name, + description=config.description, + gcs_dir_name=config.gcs_dir_name, + extra_packages=config.extra_packages, + env_vars=config.env_vars, + service_account=config.service_account, + context_spec=context_spec, + psc_interface_config=config.psc_interface_config, + agent_gateway_config=config.agent_gateway_config, + min_instances=config.min_instances, + max_instances=config.max_instances, + resource_limits=config.resource_limits, + container_concurrency=config.container_concurrency, + labels=config.labels, + class_methods=config.class_methods, + source_packages=config.source_packages, + developer_connect_source=developer_connect_source, + entrypoint_module=config.entrypoint_module, + entrypoint_object=config.entrypoint_object, + requirements_file=config.requirements_file, + agent_framework=config.agent_framework, + python_version=config.python_version, + build_options=config.build_options, + image_spec=image_spec, + agent_config_source=agent_config_source, + container_spec=container_spec, + keep_alive_probe=keep_alive_probe, + traffic_config=traffic_config, + ) + operation = self._update(name=name, config=api_config) + reasoning_engine_id = _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, + ) + agent_engine = types.Runtime( + api_client=self, + api_async_client=AsyncRuntimes(api_client_=self._api_client), + api_resource=operation.response, + ) + if agent_engine.api_resource: + logger.info("Agent Runtime updated. To use it in another session:") + logger.info( + f"agent_engine=client.runtimes.get(name='{agent_engine.api_resource.name}')" + ) + elif operation.error: + raise RuntimeError(f"Failed to update Agent Runtime: {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.QueryRuntimeConfigOrDict] = None + ) -> Iterator[Any]: + """Streams the response of the agent runtime.""" + parameter_model = types._QueryRuntimeRequestParameters( + name=name, + config=config, + ) + request_dict = _QueryRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) + else: + path = "{name}:streamQuery?alt=sse" + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + http_options = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + for response in self._api_client.request_streamed( + "post", path, request_dict, http_options + ): + yield response + + # TODO: b/436704146 - Replace with generated methods + # TODO: b/437129724 - Add replay test for async stream query + async def _async_stream_query( + self, + *, + name: str, + config: Optional[types.QueryRuntimeConfigOrDict] = None, + ) -> AsyncIterator[Any]: + """Streams the response of the agent runtime asynchronously.""" + parameter_model = types._QueryRuntimeRequestParameters( + name=name, + config=config, + ) + request_dict = _QueryRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) + else: + path = "{name}:streamQuery?alt=sse" + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + http_options = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + async_iterator = await self._api_client.async_request_streamed( + "post", path, request_dict, http_options + ) + async for response in async_iterator: + yield response + + def create_memory( + self, + *, + name: str, + fact: str, + scope: dict[str, str], + config: Optional[types.RuntimeMemoryConfigOrDict] = None, + ) -> types.RuntimeMemoryOperation: + """Deprecated. Use runtimes.memories.create instead.""" + warnings.warn( + ( + "runtimes.create_memory is deprecated. " + "Use runtimes.memories.create instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return self.memories.create( + name=name, + fact=fact, + scope=scope, + config=config, + ) + + def delete_memory( + self, + *, + name: str, + config: Optional[types.DeleteRuntimeMemoryConfigOrDict] = None, + ) -> types.DeleteRuntimeMemoryOperation: + """Deprecated. Use runtimes.memories.delete instead.""" + warnings.warn( + ( + "runtimes.delete_memory is deprecated. " + "Use runtimes.memories.delete instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return self.memories.delete(name=name, config=config) + + def generate_memories( + self, + *, + name: str, + vertex_session_source: Optional[ + types.GenerateMemoriesRequestVertexSessionSourceOrDict + ] = None, + direct_contents_source: Optional[ + types.GenerateMemoriesRequestDirectContentsSourceOrDict + ] = None, + direct_memories_source: Optional[ + types.GenerateMemoriesRequestDirectMemoriesSourceOrDict + ] = None, + scope: Optional[dict[str, str]] = None, + config: Optional[types.GenerateRuntimeMemoriesConfigOrDict] = None, + ) -> types.RuntimeGenerateMemoriesOperation: + """Deprecated. Use runtimes.memories.generate instead.""" + warnings.warn( + ( + "runtimes.generate_memories is deprecated. " + "Use runtimes.memories.generate instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return self.memories.generate( + name=name, + vertex_session_source=vertex_session_source, + direct_contents_source=direct_contents_source, + direct_memories_source=direct_memories_source, + scope=scope, + config=config, + ) + + def get_memory( + self, + *, + name: str, + config: Optional[types.GetRuntimeMemoryConfigOrDict] = None, + ) -> types.Memory: + """Deprecated. Use runtimes.memories.get instead.""" + warnings.warn( + ( + "runtimes.get_memory is deprecated. " + "Use runtimes.memories.get instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return self.memories.get(name=name, config=config) + + def list_memories( + self, + *, + name: str, + config: Optional[types.ListRuntimeMemoryConfigOrDict] = None, + ) -> Iterator[types.Memory]: + """Deprecated. Use runtimes.memories.list instead.""" + warnings.warn( + ( + "runtimes.list_memories is deprecated. " + "Use runtimes.memories.list instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return self.memories.list(name=name, config=config) + + def retrieve_memories( + self, + *, + name: str, + scope: dict[str, str], + similarity_search_params: Optional[ + types.RetrieveMemoriesRequestSimilaritySearchParamsOrDict + ] = None, + simple_retrieval_params: Optional[ + types.RetrieveMemoriesRequestSimpleRetrievalParamsOrDict + ] = None, + config: Optional[types.RetrieveRuntimeMemoriesConfigOrDict] = None, + ) -> Iterator[types.RetrieveMemoriesResponseRetrievedMemory]: + """Deprecated. Use runtimes.memories.retrieve instead.""" + warnings.warn( + ( + "runtimes.retrieve_memories is deprecated. " + "Use runtimes.memories.retrieve instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return self.memories.retrieve( + name=name, + scope=scope, + similarity_search_params=similarity_search_params, + simple_retrieval_params=simple_retrieval_params, + config=config, + ) + class AsyncRuntimes(_api_module.BaseModule): + async def cancel_query_job( + self, + *, + name: str, + config: Optional[types.CancelQueryJobRuntimeConfigOrDict] = None, + ) -> types.CancelQueryJobResult: + """ + Cancels a long-running query job on an Agent Runtime. + + Args: + name (str): + Required. The reasoning engine resource name. + config (CancelQueryJobRuntimeConfigOrDict): + Optional. The configuration for the cancel_query_job. + + """ + + parameter_model = types._CancelQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CancelQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:cancelAsyncQuery".format_map(request_url_dict) + else: + path = "{name}:cancelAsyncQuery" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.CancelQueryJobResult._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _check_query_job( + self, + *, + name: str, + config: Optional[types.CheckQueryJobRuntimeConfigOrDict] = None, + ) -> types.CheckQueryJobResult: + """ + Query an Agent Runtime asynchronously. + """ + + parameter_model = types._CheckQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CheckQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:checkQueryJob".format_map(request_url_dict) + else: + path = "{name}:checkQueryJob" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _CheckQueryJobResult_from_vertex(response_dict) + + return_value = types.CheckQueryJobResult._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _run_query_job( + self, + *, + name: str, + config: Optional[types._RunQueryJobRuntimeConfigOrDict] = None, + ) -> types.RuntimeOperation: + """ + Run a query job on an agent runtime. + """ + + parameter_model = types._RunQueryJobRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _RunQueryJobRuntimeRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:asyncQuery".format_map(request_url_dict) + else: + path = "{name}:asyncQuery" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _create( + self, *, config: Optional[types.CreateRuntimeConfigOrDict] = None + ) -> types.RuntimeOperation: + """ + Creates a new Agent Runtime. + """ + + parameter_model = types._CreateRuntimeRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CreateRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _delete( + self, + *, + name: str, + force: Optional[bool] = None, + config: Optional[types.DeleteRuntimeConfigOrDict] = None, + ) -> types.DeleteRuntimeOperation: + """ + Delete an Agent Runtime resource. + + Args: + name (str): + Required. The name of the Agent Runtime to be deleted. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` + or `reasoningEngines/{resource_id}`. + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Agent Runtime has undeleted child resources. Defaults to False. + config (DeleteRuntimeConfig): + Optional. Additional configurations for deleting the Agent Runtime. + + """ + + parameter_model = types._DeleteRuntimeRequestParameters( + name=name, + force=force, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _DeleteRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "delete", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteRuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _get( + self, *, name: str, config: Optional[types.GetRuntimeConfigOrDict] = None + ) -> types.ReasoningEngine: + """ + Get an Agent Runtime instance. + """ + + parameter_model = types._GetRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ReasoningEngine_from_vertex(response_dict) + + return_value = types.ReasoningEngine._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _list( + self, *, config: Optional[types.ListRuntimeConfigOrDict] = None + ) -> types.ListReasoningEnginesResponse: + """ + Lists Agent Runtimes. + """ + + parameter_model = types._ListRuntimeRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _ListRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ListReasoningEnginesResponse_from_vertex(response_dict) + + return_value = types.ListReasoningEnginesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _get_agent_operation( + self, + *, + operation_name: str, + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeOperation: + parameter_model = types._GetRuntimeOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetRuntimeOperationParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _query( + self, *, name: str, config: Optional[types.QueryRuntimeConfigOrDict] = None + ) -> types.QueryReasoningEngineResponse: + """ + Query an Agent Runtime. + """ + + parameter_model = types._QueryRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _QueryRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:query".format_map(request_url_dict) + else: + path = "{name}:query" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.QueryReasoningEngineResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _update( + self, *, name: str, config: Optional[types.UpdateRuntimeConfigOrDict] = None + ) -> types.RuntimeOperation: + """ + Updates an Agent Runtime. + """ + + parameter_model = types._UpdateRuntimeRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpdateRuntimeRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "patch", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _RuntimeOperation_from_vertex(response_dict) + + return_value = types.RuntimeOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + _memories = None _revisions = None + async def delete( + self, + *, + name: str, + force: Optional[bool] = None, + config: Optional[types.DeleteRuntimeConfigOrDict] = None, + ) -> types.DeleteRuntimeOperation: + """ + Delete an Agent Runtime resource. + + Args: + name (str): + Required. The name of the Agent Runtime to be deleted. Format: + `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` + or `reasoningEngines/{resource_id}`. + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Agent Runtime has undeleted child resources. Defaults to False. + config (DeleteRuntimeConfig): + Optional. Additional configurations for deleting the Agent Runtime. + + """ + logger.info(f"Deleting Runtime resource: {name}") + operation = await self._delete(name=name, force=force, config=config) + logger.info(f"Started Runtime delete operation: {operation.name}") + return operation + @property def revisions(self) -> "runtime_revisions_module.AsyncRuntimeRevisions": if self._revisions is None: @@ -71,8 +3815,57 @@ def revisions(self) -> "runtime_revisions_module.AsyncRuntimeRevisions": ) except ImportError as e: raise ImportError( - "The 'agent_engines.runtimes.revisions' module requires " - "additional packages. Please install them using pip install " + "The 'runtimes.revisions' module requires additional " + "packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e return self._revisions.AsyncRuntimeRevisions(self._api_client) # type: ignore[no-any-return] + + @property + def memories(self) -> "memories_module.AsyncMemories": + if self._memories is None: + try: + # We need to lazy load the memories module to handle the + # possibility of ImportError when dependencies are not installed. + self._memories = importlib.import_module(".memories", __package__) + except ImportError as e: + raise ImportError( + "The 'runtimes.memories' module requires additional " + "packages. Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._memories.AsyncMemories(self._api_client) # type: ignore[no-any-return] + + async def delete_memory( + self, + *, + name: str, + config: Optional[types.DeleteRuntimeMemoryConfigOrDict] = None, + ) -> types.DeleteRuntimeMemoryOperation: + """Deprecated. Use runtimes.memories.delete instead.""" + warnings.warn( + ( + "runtimes.delete_memory is deprecated. " + "Use runtimes.memories.delete instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return await self.memories.delete(name=name, config=config) + + async def get_memory( + self, + *, + name: str, + config: Optional[types.GetRuntimeMemoryConfigOrDict] = None, + ) -> types.Memory: + """Deprecated. Use runtimes.memories.get instead.""" + warnings.warn( + ( + "runtimes.get_memory is deprecated. " + "Use runtimes.memories.get instead." + ), + DeprecationWarning, + stacklevel=2, + ) + return await self.memories.get(name=name, config=config) diff --git a/agentplatform/_genai/sandbox_snapshots.py b/agentplatform/_genai/sandbox_snapshots.py index aaeb01da1e..6e4731d482 100644 --- a/agentplatform/_genai/sandbox_snapshots.py +++ b/agentplatform/_genai/sandbox_snapshots.py @@ -27,7 +27,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import Pager -from . import _agent_engines_utils +from . import _runtimes_utils from . import types logger = logging.getLogger("agentplatform_genai.sandboxsnapshots") @@ -35,7 +35,7 @@ logger.setLevel(logging.INFO) -def _CreateAgentEngineSandboxSnapshotConfig_to_vertex( +def _CreateRuntimeSandboxSnapshotConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -69,7 +69,7 @@ def _CreateSandboxEnvironmentSnapshotRequestParameters_to_vertex( setv( to_object, ["config"], - _CreateAgentEngineSandboxSnapshotConfig_to_vertex( + _CreateRuntimeSandboxSnapshotConfig_to_vertex( getv(from_object, ["config"]), to_object ), ) @@ -88,7 +88,7 @@ def _DeleteSandboxEnvironmentSnapshotRequestParameters_to_vertex( return to_object -def _GetAgentEngineSandboxSnapshotOperationParameters_to_vertex( +def _GetRuntimeSandboxSnapshotOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -153,8 +153,8 @@ def _create( self, *, source_sandbox_environment_name: str, - config: Optional[types.CreateAgentEngineSandboxSnapshotConfigOrDict] = None, - ) -> types.AgentEngineSandboxSnapshotOperation: + config: Optional[types.CreateRuntimeSandboxSnapshotConfigOrDict] = None, + ) -> types.RuntimeSandboxSnapshotOperation: """ Snapshots an existing sandbox environment. @@ -162,7 +162,7 @@ def _create( source_sandbox_environment_name (str): Required. The name of the sandbox environment to snapshot. projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_environment_id} - config (CreateAgentEngineSandboxSnapshotConfig): + config (CreateRuntimeSandboxSnapshotConfig): Optional. The configuration for the sandbox snapshot. """ @@ -207,7 +207,7 @@ def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxSnapshotOperation._from_response( + return_value = types.RuntimeSandboxSnapshotOperation._from_response( response=response_dict, kwargs=( { @@ -460,9 +460,9 @@ def get_sandbox_snapshot_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineSandboxSnapshotOperation: - parameter_model = types._GetAgentEngineSandboxSnapshotOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeSandboxSnapshotOperation: + parameter_model = types._GetRuntimeSandboxSnapshotOperationParameters( operation_name=operation_name, config=config, ) @@ -473,7 +473,7 @@ def get_sandbox_snapshot_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSandboxSnapshotOperationParameters_to_vertex( + request_dict = _GetRuntimeSandboxSnapshotOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -502,7 +502,7 @@ def get_sandbox_snapshot_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxSnapshotOperation._from_response( + return_value = types.RuntimeSandboxSnapshotOperation._from_response( response=response_dict, kwargs=( { @@ -530,34 +530,34 @@ def create( self, *, source_sandbox_environment_name: str, - config: Optional[types.CreateAgentEngineSandboxSnapshotConfigOrDict] = None, + config: Optional[types.CreateRuntimeSandboxSnapshotConfigOrDict] = None, poll_interval_seconds: float = 0.1, - ) -> types.AgentEngineSandboxSnapshotOperation: + ) -> types.RuntimeSandboxSnapshotOperation: """Snapshots an existing sandbox environment. Args: source_sandbox_environment_name (str): Required. The name of the sandbox environment to snapshot. projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_environment_id} - config (CreateAgentEngineSandboxSnapshotConfig): + config (CreateRuntimeSandboxSnapshotConfig): Optional. The configuration for the sandbox snapshot. poll_interval_seconds (int): Optional. Seconds to wait between polling for operation status. Defaults to 0.1. Returns: - AgentEngineSandboxSnapshotOperation: The operation for creating the sandbox snapshot. + RuntimeSandboxSnapshotOperation: The operation for creating the sandbox snapshot. """ operation = self._create( source_sandbox_environment_name=source_sandbox_environment_name, config=config, ) if config is None: - config = types.CreateAgentEngineSandboxSnapshotConfig() + config = types.CreateRuntimeSandboxSnapshotConfig() elif isinstance(config, dict): - config = types.CreateAgentEngineSandboxSnapshotConfig.model_validate(config) + config = types.CreateRuntimeSandboxSnapshotConfig.model_validate(config) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self.get_sandbox_snapshot_operation, poll_interval_seconds=poll_interval_seconds, @@ -575,17 +575,17 @@ def list( name: str, config: Optional[types.ListSandboxEnvironmentSnapshotsConfigOrDict] = None, ) -> Iterator[types.SandboxEnvironmentSnapshot]: - """Lists Agent Engine sandbox snapshots. + """Lists Agent Runtime sandbox snapshots. Args: name (str): - Required. The name of the agent engine to list sandbox snapshots for. + Required. The name of the agent runtime to list sandbox snapshots for. projects/{project}/locations/{location}/reasoningEngines/{resource_id} config (ListSandboxEnvironmentSnapshotsConfig): Optional. The configuration for the sandbox snapshots to list. Returns: - Iterable[SandboxEnvironmentSnapshot]: An iterable of agent engine sandbox snapshots. + Iterable[SandboxEnvironmentSnapshot]: An iterable of agent runtime sandbox snapshots. """ return Pager( "sandbox_environment_snapshots", @@ -600,7 +600,7 @@ def get( name: str, config: Optional[types.GetSandboxEnvironmentSnapshotConfigOrDict] = None, ) -> types.SandboxEnvironmentSnapshot: - """Gets a sandbox snapshot in the Agent Engine. + """Gets a sandbox snapshot in the Agent Runtime. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -617,7 +617,7 @@ def delete( name: str, config: Optional[types.DeleteSandboxEnvironmentSnapshotConfigOrDict] = None, ) -> types.DeleteSandboxEnvironmentSnapshotOperation: - """Deletes a sandbox snapshot in the Agent Engine. + """Deletes a sandbox snapshot in the Agent Runtime. Args: name (str): Required. The name of the sandbox snapshot to delete. @@ -635,8 +635,8 @@ async def _create( self, *, source_sandbox_environment_name: str, - config: Optional[types.CreateAgentEngineSandboxSnapshotConfigOrDict] = None, - ) -> types.AgentEngineSandboxSnapshotOperation: + config: Optional[types.CreateRuntimeSandboxSnapshotConfigOrDict] = None, + ) -> types.RuntimeSandboxSnapshotOperation: """ Snapshots an existing sandbox environment. @@ -644,7 +644,7 @@ async def _create( source_sandbox_environment_name (str): Required. The name of the sandbox environment to snapshot. projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_environment_id} - config (CreateAgentEngineSandboxSnapshotConfig): + config (CreateRuntimeSandboxSnapshotConfig): Optional. The configuration for the sandbox snapshot. """ @@ -691,7 +691,7 @@ async def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxSnapshotOperation._from_response( + return_value = types.RuntimeSandboxSnapshotOperation._from_response( response=response_dict, kwargs=( { @@ -950,9 +950,9 @@ async def get_sandbox_snapshot_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineSandboxSnapshotOperation: - parameter_model = types._GetAgentEngineSandboxSnapshotOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeSandboxSnapshotOperation: + parameter_model = types._GetRuntimeSandboxSnapshotOperationParameters( operation_name=operation_name, config=config, ) @@ -963,7 +963,7 @@ async def get_sandbox_snapshot_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSandboxSnapshotOperationParameters_to_vertex( + request_dict = _GetRuntimeSandboxSnapshotOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -994,7 +994,7 @@ async def get_sandbox_snapshot_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxSnapshotOperation._from_response( + return_value = types.RuntimeSandboxSnapshotOperation._from_response( response=response_dict, kwargs=( { diff --git a/agentplatform/_genai/sandbox_templates.py b/agentplatform/_genai/sandbox_templates.py index 557cd7411e..24f51623fc 100644 --- a/agentplatform/_genai/sandbox_templates.py +++ b/agentplatform/_genai/sandbox_templates.py @@ -27,7 +27,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import Pager -from . import _agent_engines_utils +from . import _runtimes_utils from . import types logger = logging.getLogger("agentplatform_genai.sandboxtemplates") @@ -164,11 +164,11 @@ def _create( display_name: str, ) -> types.SandboxEnvironmentTemplateOperation: """ - Creates a new sandbox template in the Agent Engine. + Creates a new sandbox template in the Agent Runtime. Args: name (str): - Required. The name of the agent engine to create the template under. + Required. The name of the agent runtime to create the template under. Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id} display_name (str): Required. The display name of the sandbox template. @@ -249,7 +249,7 @@ def _delete( config: Optional[types.DeleteSandboxEnvironmentTemplateConfigOrDict] = None, ) -> types.DeleteSandboxEnvironmentTemplateOperation: """ - Delete an Agent Engine sandbox template. + Delete an Agent Runtime sandbox template. Args: name (str): @@ -331,7 +331,7 @@ def _get( config: Optional[types.GetSandboxEnvironmentTemplateConfigOrDict] = None, ) -> types.SandboxEnvironmentTemplate: """ - Gets an agent engine sandbox template. + Gets an agent runtime sandbox template. Args: name (str): The resource name of the SandboxEnvironmentTemplate. @@ -417,10 +417,10 @@ def _list( config: Optional[types.ListSandboxEnvironmentTemplatesConfigOrDict] = None, ) -> types.ListSandboxEnvironmentTemplatesResponse: """ - Lists Agent Engine sandbox templates. + Lists Agent Runtime sandbox templates. Args: - name (str): Name of the agent engine. Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id} + name (str): Name of the agent runtime. Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id} config (ListSandboxEnvironmentTemplatesConfig): Configuration for listing sandbox templates. Returns: @@ -496,7 +496,7 @@ def get_sandbox_environment_template_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, ) -> types.SandboxEnvironmentTemplateOperation: parameter_model = types._GetSandboxEnvironmentTemplateOperationParameters( operation_name=operation_name, @@ -570,11 +570,11 @@ def create( config: Optional[types.CreateSandboxEnvironmentTemplateConfigOrDict] = None, poll_interval_seconds: float = 0.1, ) -> types.SandboxEnvironmentTemplateOperation: - """Creates a new sandbox template in the Agent Engine. + """Creates a new sandbox template in the Agent Runtime. Args: name (str): - Required. The name of the agent engine to create sandbox template for. + Required. The name of the agent runtime to create sandbox template for. projects/{project}/locations/{location}/reasoningEngines/{resource_id} display_name (str): Required. The display name of the sandbox template. @@ -597,7 +597,7 @@ def create( config = types.CreateSandboxEnvironmentTemplateConfig.model_validate(config) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self.get_sandbox_environment_template_operation, poll_interval_seconds=poll_interval_seconds, @@ -615,17 +615,17 @@ def list( name: str, config: Optional[types.ListSandboxEnvironmentTemplatesConfigOrDict] = None, ) -> Iterator[types.SandboxEnvironmentTemplate]: - """Lists Agent Engine sandbox templates. + """Lists Agent Runtime sandbox templates. Args: name (str): - Required. The name of the agent engine to list sandbox templates for. + Required. The name of the agent runtime to list sandbox templates for. projects/{project}/locations/{location}/reasoningEngines/{resource_id} config (ListSandboxEnvironmentTemplatesConfig): Optional. The configuration for the sandbox templates to list. Returns: - Iterable[SandboxEnvironmentTemplate]: An iterable of agent engine sandbox templates. + Iterable[SandboxEnvironmentTemplate]: An iterable of agent runtime sandbox templates. """ return Pager( "sandbox_environment_templates", @@ -640,7 +640,7 @@ def get( name: str, config: Optional[types.GetSandboxEnvironmentTemplateConfigOrDict] = None, ) -> types.SandboxEnvironmentTemplate: - """Gets a sandbox template in the Agent Engine. + """Gets a sandbox template in the Agent Runtime. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -657,7 +657,7 @@ def delete( name: str, config: Optional[types.DeleteSandboxEnvironmentTemplateConfigOrDict] = None, ) -> types.DeleteSandboxEnvironmentTemplateOperation: - """Deletes a sandbox template in the Agent Engine. + """Deletes a sandbox template in the Agent Runtime. Args: name (str): Required. The name of the sandbox template to delete. @@ -679,11 +679,11 @@ async def _create( display_name: str, ) -> types.SandboxEnvironmentTemplateOperation: """ - Creates a new sandbox template in the Agent Engine. + Creates a new sandbox template in the Agent Runtime. Args: name (str): - Required. The name of the agent engine to create the template under. + Required. The name of the agent runtime to create the template under. Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id} display_name (str): Required. The display name of the sandbox template. @@ -766,7 +766,7 @@ async def _delete( config: Optional[types.DeleteSandboxEnvironmentTemplateConfigOrDict] = None, ) -> types.DeleteSandboxEnvironmentTemplateOperation: """ - Delete an Agent Engine sandbox template. + Delete an Agent Runtime sandbox template. Args: name (str): @@ -850,7 +850,7 @@ async def _get( config: Optional[types.GetSandboxEnvironmentTemplateConfigOrDict] = None, ) -> types.SandboxEnvironmentTemplate: """ - Gets an agent engine sandbox template. + Gets an agent runtime sandbox template. Args: name (str): The resource name of the SandboxEnvironmentTemplate. @@ -938,10 +938,10 @@ async def _list( config: Optional[types.ListSandboxEnvironmentTemplatesConfigOrDict] = None, ) -> types.ListSandboxEnvironmentTemplatesResponse: """ - Lists Agent Engine sandbox templates. + Lists Agent Runtime sandbox templates. Args: - name (str): Name of the agent engine. Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id} + name (str): Name of the agent runtime. Format: projects/{project}/locations/{location}/reasoningEngines/{resource_id} config (ListSandboxEnvironmentTemplatesConfig): Configuration for listing sandbox templates. Returns: @@ -1019,7 +1019,7 @@ async def get_sandbox_environment_template_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, ) -> types.SandboxEnvironmentTemplateOperation: parameter_model = types._GetSandboxEnvironmentTemplateOperationParameters( operation_name=operation_name, diff --git a/agentplatform/_genai/sandboxes.py b/agentplatform/_genai/sandboxes.py index 033f033359..248dbdfdb7 100644 --- a/agentplatform/_genai/sandboxes.py +++ b/agentplatform/_genai/sandboxes.py @@ -34,7 +34,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import Pager -from . import _agent_engines_utils +from . import _runtimes_utils from . import types logger = logging.getLogger("agentplatform_genai.sandboxes") @@ -42,7 +42,7 @@ logger.setLevel(logging.INFO) -def _CreateAgentEngineSandboxConfig_to_vertex( +def _CreateRuntimeSandboxConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -77,7 +77,7 @@ def _CreateAgentEngineSandboxConfig_to_vertex( return to_object -def _CreateAgentEngineSandboxRequestParameters_to_vertex( +def _CreateRuntimeSandboxRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -89,14 +89,12 @@ def _CreateAgentEngineSandboxRequestParameters_to_vertex( setv(to_object, ["spec"], getv(from_object, ["spec"])) if getv(from_object, ["config"]) is not None: - _CreateAgentEngineSandboxConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _CreateRuntimeSandboxConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _DeleteAgentEngineSandboxRequestParameters_to_vertex( +def _DeleteRuntimeSandboxRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -107,7 +105,7 @@ def _DeleteAgentEngineSandboxRequestParameters_to_vertex( return to_object -def _ExecuteCodeAgentEngineSandboxRequestParameters_to_vertex( +def _ExecuteCodeRuntimeSandboxRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -121,7 +119,7 @@ def _ExecuteCodeAgentEngineSandboxRequestParameters_to_vertex( return to_object -def _GetAgentEngineSandboxOperationParameters_to_vertex( +def _GetRuntimeSandboxOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -134,7 +132,7 @@ def _GetAgentEngineSandboxOperationParameters_to_vertex( return to_object -def _GetAgentEngineSandboxRequestParameters_to_vertex( +def _GetRuntimeSandboxRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -145,7 +143,7 @@ def _GetAgentEngineSandboxRequestParameters_to_vertex( return to_object -def _ListAgentEngineSandboxesConfig_to_vertex( +def _ListRuntimeSandboxesConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -163,7 +161,7 @@ def _ListAgentEngineSandboxesConfig_to_vertex( return to_object -def _ListAgentEngineSandboxesRequestParameters_to_vertex( +def _ListRuntimeSandboxesRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -172,9 +170,7 @@ def _ListAgentEngineSandboxesRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineSandboxesConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _ListRuntimeSandboxesConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object @@ -186,13 +182,13 @@ def _create( *, name: str, spec: Optional[types.SandboxEnvironmentSpecOrDict] = None, - config: Optional[types.CreateAgentEngineSandboxConfigOrDict] = None, - ) -> types.AgentEngineSandboxOperation: + config: Optional[types.CreateRuntimeSandboxConfigOrDict] = None, + ) -> types.RuntimeSandboxOperation: """ - Creates a new sandbox in the Agent Engine. + Creates a new sandbox in the Agent Runtime. """ - parameter_model = types._CreateAgentEngineSandboxRequestParameters( + parameter_model = types._CreateRuntimeSandboxRequestParameters( name=name, spec=spec, config=config, @@ -204,7 +200,7 @@ def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _CreateRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -233,7 +229,7 @@ def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxOperation._from_response( + return_value = types.RuntimeSandboxOperation._from_response( response=response_dict, kwargs=( { @@ -261,19 +257,19 @@ def _delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineSandboxConfigOrDict] = None, - ) -> types.DeleteAgentEngineSandboxOperation: + config: Optional[types.DeleteRuntimeSandboxConfigOrDict] = None, + ) -> types.DeleteRuntimeSandboxOperation: """ - Delete an Agent Engine sandbox. + Delete an Agent Runtime sandbox. Args: name (str): - Required. The name of the Agent Engine sandbox to be deleted. Format: + Required. The name of the Agent Runtime sandbox to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox}`. """ - parameter_model = types._DeleteAgentEngineSandboxRequestParameters( + parameter_model = types._DeleteRuntimeSandboxRequestParameters( name=name, config=config, ) @@ -284,7 +280,7 @@ def _delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _DeleteRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -313,7 +309,7 @@ def _delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineSandboxOperation._from_response( + return_value = types.DeleteRuntimeSandboxOperation._from_response( response=response_dict, kwargs=( { @@ -342,13 +338,13 @@ def _execute_code( *, name: str, inputs: Optional[builtins.list[types.ChunkOrDict]] = None, - config: Optional[types.ExecuteCodeAgentEngineSandboxConfigOrDict] = None, + config: Optional[types.ExecuteCodeRuntimeSandboxConfigOrDict] = None, ) -> types.ExecuteSandboxEnvironmentResponse: """ - Execute code in an Agent Engine sandbox. + Execute code in an Agent Runtime sandbox. """ - parameter_model = types._ExecuteCodeAgentEngineSandboxRequestParameters( + parameter_model = types._ExecuteCodeRuntimeSandboxRequestParameters( name=name, inputs=inputs, config=config, @@ -360,7 +356,7 @@ def _execute_code( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ExecuteCodeAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _ExecuteCodeRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -414,13 +410,10 @@ def _execute_code( return return_value def _get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineSandboxConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeSandboxConfigOrDict] = None ) -> types.SandboxEnvironment: """ - Gets an agent engine sandbox. + Gets an agent runtime sandbox. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -429,7 +422,7 @@ def _get( """ - parameter_model = types._GetAgentEngineSandboxRequestParameters( + parameter_model = types._GetRuntimeSandboxRequestParameters( name=name, config=config, ) @@ -440,7 +433,7 @@ def _get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _GetRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -497,23 +490,23 @@ def _list( self, *, name: str, - config: Optional[types.ListAgentEngineSandboxesConfigOrDict] = None, - ) -> types.ListAgentEngineSandboxesResponse: + config: Optional[types.ListRuntimeSandboxesConfigOrDict] = None, + ) -> types.ListRuntimeSandboxesResponse: """ - Lists Agent Engine sandboxes. + Lists Agent Runtime sandboxes. Args: - name (str): Required. The name of the Agent Engine to list sessions for. Format: + name (str): Required. The name of the Agent Runtime to list sessions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineSandboxesConfig): - Optional. Additional configurations for listing the Agent Engine sandboxes. + config (ListRuntimeSandboxesConfig): + Optional. Additional configurations for listing the Agent Runtime sandboxes. Returns: - ListReasoningEnginesSandboxesResponse: The requested Agent Engine sandboxes. + ListReasoningEnginesSandboxesResponse: The requested Agent Runtime sandboxes. """ - parameter_model = types._ListAgentEngineSandboxesRequestParameters( + parameter_model = types._ListRuntimeSandboxesRequestParameters( name=name, config=config, ) @@ -524,7 +517,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineSandboxesRequestParameters_to_vertex( + request_dict = _ListRuntimeSandboxesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -553,7 +546,7 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineSandboxesResponse._from_response( + return_value = types.ListRuntimeSandboxesResponse._from_response( response=response_dict, kwargs=( { @@ -581,9 +574,9 @@ def _get_sandbox_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineSandboxOperation: - parameter_model = types._GetAgentEngineSandboxOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeSandboxOperation: + parameter_model = types._GetRuntimeSandboxOperationParameters( operation_name=operation_name, config=config, ) @@ -594,7 +587,7 @@ def _get_sandbox_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSandboxOperationParameters_to_vertex( + request_dict = _GetRuntimeSandboxOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -623,7 +616,7 @@ def _get_sandbox_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxOperation._from_response( + return_value = types.RuntimeSandboxOperation._from_response( response=response_dict, kwargs=( { @@ -659,7 +652,7 @@ def templates(self) -> Any: ) except ImportError as e: raise ImportError( - "The 'agent_engines.sandboxes.templates' module requires " + "The 'runtimes.sandboxes.templates' module requires " "additional packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e @@ -674,7 +667,7 @@ def snapshots(self) -> Any: ) except ImportError as e: raise ImportError( - "The 'agent_engines.sandboxes.snapshots' module requires " + "The 'runtimes.sandboxes.snapshots' module requires " "additional packages. Please install them using pip install " "google-cloud-aiplatform[sandbox_snapshots]" ) from e @@ -686,24 +679,24 @@ def create( name: str, poll_interval_seconds: float = 0.1, spec: Optional[types.SandboxEnvironmentSpecOrDict] = None, - config: Optional[types.CreateAgentEngineSandboxConfigOrDict] = None, - ) -> types.AgentEngineSandboxOperation: - """Creates a new sandbox in the Agent Engine. + config: Optional[types.CreateRuntimeSandboxConfigOrDict] = None, + ) -> types.RuntimeSandboxOperation: + """Creates a new sandbox in the Agent Runtime. Args: name (str): - Required. The name of the agent engine to create sandbox for. + Required. The name of the agent runtime to create sandbox for. projects/{project}/locations/{location}/reasoningEngines/{resource_id} poll_interval_seconds (float): Optional. The interval in seconds to poll for sandbox creation completion. spec (SandboxEnvironmentSpec): Optional. The specification for the sandbox to create. - config (CreateAgentEngineSandboxConfigOrDict): + config (CreateRuntimeSandboxConfigOrDict): Optional. The configuration for the sandbox. Returns: - AgentEngineSandboxOperation: The operation for creating the sandbox. + RuntimeSandboxOperation: The operation for creating the sandbox. """ if spec: computer_use = False @@ -722,12 +715,12 @@ def create( config=config, ) if config is None: - config = types.CreateAgentEngineSandboxConfig() + config = types.CreateRuntimeSandboxConfig() elif isinstance(config, dict): - config = types.CreateAgentEngineSandboxConfig.model_validate(config) + config = types.CreateRuntimeSandboxConfig.model_validate(config) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_sandbox_operation, poll_interval_seconds=poll_interval_seconds, @@ -743,19 +736,19 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineSandboxesConfigOrDict] = None, + config: Optional[types.ListRuntimeSandboxesConfigOrDict] = None, ) -> Iterator[types.SandboxEnvironment]: - """Lists Agent Engine sandboxes. + """Lists Agent Runtime sandboxes. Args: name (str): - Required. The name of the agent engine to list sandboxes for. + Required. The name of the agent runtime to list sandboxes for. projects/{project}/locations/{location}/reasoningEngines/{resource_id} - config (ListAgentEngineSandboxConfig): + config (ListRuntimeSandboxConfig): Optional. The configuration for the sandboxes to list. Returns: - Iterable[SandboxEnvironment]: An iterable of agent engine sandboxes. + Iterable[SandboxEnvironment]: An iterable of agent runtime sandboxes. """ return Pager( "sandbox_environments", @@ -769,17 +762,17 @@ def execute_code( *, name: str, input_data: dict[str, Any], - config: Optional[types.ExecuteCodeAgentEngineSandboxConfigOrDict] = None, + config: Optional[types.ExecuteCodeRuntimeSandboxConfigOrDict] = None, ) -> types.ExecuteSandboxEnvironmentResponse: - """Executes code in the Agent Engine sandbox. + """Executes code in the Agent Runtime sandbox. Args: name (str): - Required. The name of the agent engine sandbox to run code in. + Required. The name of the agent runtime sandbox to run code in. projects/{project}/locations/{location}/reasoningEngines/{resource_id}/SandboxEnvironments/{sandbox_id} input_data (dict[str, Any]): Required. The input to the code to execute. - config (ExecuteCodeAgentEngineSandboxConfigOrDict): + config (ExecuteCodeRuntimeSandboxConfigOrDict): Optional. The configuration for the sandboxes to run code in. Returns: @@ -837,15 +830,15 @@ def get( self, *, name: str, - config: Optional[types.GetAgentEngineSandboxConfigOrDict] = None, + config: Optional[types.GetRuntimeSandboxConfigOrDict] = None, ) -> types.SandboxEnvironment: - """Gets an agent engine sandbox. + """Gets an agent runtime sandbox. Args: name (str): Required. A fully-qualified resource name or ID such as projects/{project}/locations/{location}/reasoningEngines/{resource_id}/SandboxEnvironments/{sandbox_id} or a shortened name such as "reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_id}". - config (GetAgentEngineSandboxConfigOrDict): + config (GetRuntimeSandboxConfigOrDict): Optional. The configuration for the sandbox to get. """ @@ -855,15 +848,15 @@ def delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineSandboxConfigOrDict] = None, - ) -> types.DeleteAgentEngineSandboxOperation: - """Deletes an agent engine sandbox. + config: Optional[types.DeleteRuntimeSandboxConfigOrDict] = None, + ) -> types.DeleteRuntimeSandboxOperation: + """Deletes an agent runtime sandbox. Args: name (str): Required. A fully-qualified resource name or ID such as projects/{project}/locations/{location}/reasoningEngines/{resource_id}/SandboxEnvironments/{sandbox_id} or a shortened name such as "reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox_id}". - config (DeleteAgentEngineSandboxConfigOrDict): + config (DeleteRuntimeSandboxConfigOrDict): Optional. The configuration for the sandbox to delete. """ return self._delete(name=name, config=config) @@ -1036,13 +1029,13 @@ async def _create( *, name: str, spec: Optional[types.SandboxEnvironmentSpecOrDict] = None, - config: Optional[types.CreateAgentEngineSandboxConfigOrDict] = None, - ) -> types.AgentEngineSandboxOperation: + config: Optional[types.CreateRuntimeSandboxConfigOrDict] = None, + ) -> types.RuntimeSandboxOperation: """ - Creates a new sandbox in the Agent Engine. + Creates a new sandbox in the Agent Runtime. """ - parameter_model = types._CreateAgentEngineSandboxRequestParameters( + parameter_model = types._CreateRuntimeSandboxRequestParameters( name=name, spec=spec, config=config, @@ -1054,7 +1047,7 @@ async def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _CreateRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1085,7 +1078,7 @@ async def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxOperation._from_response( + return_value = types.RuntimeSandboxOperation._from_response( response=response_dict, kwargs=( { @@ -1113,19 +1106,19 @@ async def _delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineSandboxConfigOrDict] = None, - ) -> types.DeleteAgentEngineSandboxOperation: + config: Optional[types.DeleteRuntimeSandboxConfigOrDict] = None, + ) -> types.DeleteRuntimeSandboxOperation: """ - Delete an Agent Engine sandbox. + Delete an Agent Runtime sandbox. Args: name (str): - Required. The name of the Agent Engine sandbox to be deleted. Format: + Required. The name of the Agent Runtime sandbox to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sandboxEnvironments/{sandbox}`. """ - parameter_model = types._DeleteAgentEngineSandboxRequestParameters( + parameter_model = types._DeleteRuntimeSandboxRequestParameters( name=name, config=config, ) @@ -1136,7 +1129,7 @@ async def _delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _DeleteRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1167,7 +1160,7 @@ async def _delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineSandboxOperation._from_response( + return_value = types.DeleteRuntimeSandboxOperation._from_response( response=response_dict, kwargs=( { @@ -1196,13 +1189,13 @@ async def _execute_code( *, name: str, inputs: Optional[builtins.list[types.ChunkOrDict]] = None, - config: Optional[types.ExecuteCodeAgentEngineSandboxConfigOrDict] = None, + config: Optional[types.ExecuteCodeRuntimeSandboxConfigOrDict] = None, ) -> types.ExecuteSandboxEnvironmentResponse: """ - Execute code in an Agent Engine sandbox. + Execute code in an Agent Runtime sandbox. """ - parameter_model = types._ExecuteCodeAgentEngineSandboxRequestParameters( + parameter_model = types._ExecuteCodeRuntimeSandboxRequestParameters( name=name, inputs=inputs, config=config, @@ -1214,7 +1207,7 @@ async def _execute_code( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ExecuteCodeAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _ExecuteCodeRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1270,13 +1263,10 @@ async def _execute_code( return return_value async def _get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineSandboxConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeSandboxConfigOrDict] = None ) -> types.SandboxEnvironment: """ - Gets an agent engine sandbox. + Gets an agent runtime sandbox. Args: name (str): Required. A fully-qualified resource name or ID such as @@ -1285,7 +1275,7 @@ async def _get( """ - parameter_model = types._GetAgentEngineSandboxRequestParameters( + parameter_model = types._GetRuntimeSandboxRequestParameters( name=name, config=config, ) @@ -1296,7 +1286,7 @@ async def _get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSandboxRequestParameters_to_vertex( + request_dict = _GetRuntimeSandboxRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1355,23 +1345,23 @@ async def _list( self, *, name: str, - config: Optional[types.ListAgentEngineSandboxesConfigOrDict] = None, - ) -> types.ListAgentEngineSandboxesResponse: + config: Optional[types.ListRuntimeSandboxesConfigOrDict] = None, + ) -> types.ListRuntimeSandboxesResponse: """ - Lists Agent Engine sandboxes. + Lists Agent Runtime sandboxes. Args: - name (str): Required. The name of the Agent Engine to list sessions for. Format: + name (str): Required. The name of the Agent Runtime to list sessions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineSandboxesConfig): - Optional. Additional configurations for listing the Agent Engine sandboxes. + config (ListRuntimeSandboxesConfig): + Optional. Additional configurations for listing the Agent Runtime sandboxes. Returns: - ListReasoningEnginesSandboxesResponse: The requested Agent Engine sandboxes. + ListReasoningEnginesSandboxesResponse: The requested Agent Runtime sandboxes. """ - parameter_model = types._ListAgentEngineSandboxesRequestParameters( + parameter_model = types._ListRuntimeSandboxesRequestParameters( name=name, config=config, ) @@ -1382,7 +1372,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineSandboxesRequestParameters_to_vertex( + request_dict = _ListRuntimeSandboxesRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1413,7 +1403,7 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineSandboxesResponse._from_response( + return_value = types.ListRuntimeSandboxesResponse._from_response( response=response_dict, kwargs=( { @@ -1441,9 +1431,9 @@ async def _get_sandbox_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineSandboxOperation: - parameter_model = types._GetAgentEngineSandboxOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeSandboxOperation: + parameter_model = types._GetRuntimeSandboxOperationParameters( operation_name=operation_name, config=config, ) @@ -1454,7 +1444,7 @@ async def _get_sandbox_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSandboxOperationParameters_to_vertex( + request_dict = _GetRuntimeSandboxOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1485,7 +1475,7 @@ async def _get_sandbox_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSandboxOperation._from_response( + return_value = types.RuntimeSandboxOperation._from_response( response=response_dict, kwargs=( { diff --git a/agentplatform/_genai/session_events.py b/agentplatform/_genai/session_events.py index 9333f29549..a4f46f5d49 100644 --- a/agentplatform/_genai/session_events.py +++ b/agentplatform/_genai/session_events.py @@ -35,7 +35,7 @@ logger.setLevel(logging.INFO) -def _AppendAgentEngineSessionEventConfig_to_vertex( +def _AppendRuntimeSessionEventConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -62,7 +62,7 @@ def _AppendAgentEngineSessionEventConfig_to_vertex( return to_object -def _AppendAgentEngineSessionEventRequestParameters_to_vertex( +def _AppendRuntimeSessionEventRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -80,14 +80,14 @@ def _AppendAgentEngineSessionEventRequestParameters_to_vertex( setv(to_object, ["timestamp"], getv(from_object, ["timestamp"])) if getv(from_object, ["config"]) is not None: - _AppendAgentEngineSessionEventConfig_to_vertex( + _AppendRuntimeSessionEventConfig_to_vertex( getv(from_object, ["config"]), to_object ) return to_object -def _ListAgentEngineSessionEventsConfig_to_vertex( +def _ListRuntimeSessionEventsConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -105,7 +105,7 @@ def _ListAgentEngineSessionEventsConfig_to_vertex( return to_object -def _ListAgentEngineSessionEventsRequestParameters_to_vertex( +def _ListRuntimeSessionEventsRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -114,7 +114,7 @@ def _ListAgentEngineSessionEventsRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineSessionEventsConfig_to_vertex( + _ListRuntimeSessionEventsConfig_to_vertex( getv(from_object, ["config"]), to_object ) @@ -130,26 +130,26 @@ def append( author: str, invocation_id: str, timestamp: datetime.datetime, - config: Optional[types.AppendAgentEngineSessionEventConfigOrDict] = None, - ) -> types.AppendAgentEngineSessionEventResponse: + config: Optional[types.AppendRuntimeSessionEventConfigOrDict] = None, + ) -> types.AppendRuntimeSessionEventResponse: """ - Appends Agent Engine session event. + Appends Agent Runtime session event. Args: - name (str): Required. The name of the Agent Engine session to append the event to. Format: + name (str): Required. The name of the Agent Runtime session to append the event to. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - author (str): Required. The author of the Agent Engine session event. - invocation_id (str): Required. The invocation ID of the Agent Engine session event. - timestamp (datetime.datetime): Required. The timestamp of the Agent Engine session event. - config (AppendAgentEngineSessionEventConfig): - Optional. Additional configurations for appending the Agent Engine session event. + author (str): Required. The author of the Agent Runtime session event. + invocation_id (str): Required. The invocation ID of the Agent Runtime session event. + timestamp (datetime.datetime): Required. The timestamp of the Agent Runtime session event. + config (AppendRuntimeSessionEventConfig): + Optional. Additional configurations for appending the Agent Runtime session event. Returns: - AppendAgentEngineSessionEventResponse: The requested Agent Engine session event. + AppendRuntimeSessionEventResponse: The requested Agent Runtime session event. """ - parameter_model = types._AppendAgentEngineSessionEventRequestParameters( + parameter_model = types._AppendRuntimeSessionEventRequestParameters( name=name, author=author, invocation_id=invocation_id, @@ -163,7 +163,7 @@ def append( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _AppendAgentEngineSessionEventRequestParameters_to_vertex( + request_dict = _AppendRuntimeSessionEventRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -192,7 +192,7 @@ def append( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AppendAgentEngineSessionEventResponse._from_response( + return_value = types.AppendRuntimeSessionEventResponse._from_response( response=response_dict, kwargs=( { @@ -220,23 +220,23 @@ def _list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionEventsConfigOrDict] = None, - ) -> types.ListAgentEngineSessionEventsResponse: + config: Optional[types.ListRuntimeSessionEventsConfigOrDict] = None, + ) -> types.ListRuntimeSessionEventsResponse: """ - Lists Agent Engine session events. + Lists Agent Runtime session events. Args: - name (str): Required. The name of the Agent Engine session to list events for. Format: + name (str): Required. The name of the Agent Runtime session to list events for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (ListAgentEngineSessionEventsConfig): - Optional. Additional configurations for listing the Agent Engine session events. + config (ListRuntimeSessionEventsConfig): + Optional. Additional configurations for listing the Agent Runtime session events. Returns: - ListAgentEngineSessionEventsResponse: The requested Agent Engine session events. + ListRuntimeSessionEventsResponse: The requested Agent Runtime session events. """ - parameter_model = types._ListAgentEngineSessionEventsRequestParameters( + parameter_model = types._ListRuntimeSessionEventsRequestParameters( name=name, config=config, ) @@ -247,7 +247,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineSessionEventsRequestParameters_to_vertex( + request_dict = _ListRuntimeSessionEventsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -276,7 +276,7 @@ def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineSessionEventsResponse._from_response( + return_value = types.ListRuntimeSessionEventsResponse._from_response( response=response_dict, kwargs=( { @@ -304,14 +304,14 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionEventsConfigOrDict] = None, + config: Optional[types.ListRuntimeSessionEventsConfigOrDict] = None, ) -> Iterator[types.SessionEvent]: - """Lists Agent Engine session events. + """Lists Agent Runtime session events. Args: - name (str): Required. The name of the agent engine to list session + name (str): Required. The name of the agent runtime to list session events for. - config (ListAgentEngineSessionEventsConfig): Optional. The configuration + config (ListRuntimeSessionEventsConfig): Optional. The configuration for the session events to list. Currently, the `filter` field in `config` only supports filtering by `timestamp`. The timestamp value must be enclosed in double quotes and include the time zone @@ -339,26 +339,26 @@ async def append( author: str, invocation_id: str, timestamp: datetime.datetime, - config: Optional[types.AppendAgentEngineSessionEventConfigOrDict] = None, - ) -> types.AppendAgentEngineSessionEventResponse: + config: Optional[types.AppendRuntimeSessionEventConfigOrDict] = None, + ) -> types.AppendRuntimeSessionEventResponse: """ - Appends Agent Engine session event. + Appends Agent Runtime session event. Args: - name (str): Required. The name of the Agent Engine session to append the event to. Format: + name (str): Required. The name of the Agent Runtime session to append the event to. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - author (str): Required. The author of the Agent Engine session event. - invocation_id (str): Required. The invocation ID of the Agent Engine session event. - timestamp (datetime.datetime): Required. The timestamp of the Agent Engine session event. - config (AppendAgentEngineSessionEventConfig): - Optional. Additional configurations for appending the Agent Engine session event. + author (str): Required. The author of the Agent Runtime session event. + invocation_id (str): Required. The invocation ID of the Agent Runtime session event. + timestamp (datetime.datetime): Required. The timestamp of the Agent Runtime session event. + config (AppendRuntimeSessionEventConfig): + Optional. Additional configurations for appending the Agent Runtime session event. Returns: - AppendAgentEngineSessionEventResponse: The requested Agent Engine session event. + AppendRuntimeSessionEventResponse: The requested Agent Runtime session event. """ - parameter_model = types._AppendAgentEngineSessionEventRequestParameters( + parameter_model = types._AppendRuntimeSessionEventRequestParameters( name=name, author=author, invocation_id=invocation_id, @@ -372,7 +372,7 @@ async def append( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _AppendAgentEngineSessionEventRequestParameters_to_vertex( + request_dict = _AppendRuntimeSessionEventRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -403,7 +403,7 @@ async def append( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AppendAgentEngineSessionEventResponse._from_response( + return_value = types.AppendRuntimeSessionEventResponse._from_response( response=response_dict, kwargs=( { @@ -431,23 +431,23 @@ async def _list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionEventsConfigOrDict] = None, - ) -> types.ListAgentEngineSessionEventsResponse: + config: Optional[types.ListRuntimeSessionEventsConfigOrDict] = None, + ) -> types.ListRuntimeSessionEventsResponse: """ - Lists Agent Engine session events. + Lists Agent Runtime session events. Args: - name (str): Required. The name of the Agent Engine session to list events for. Format: + name (str): Required. The name of the Agent Runtime session to list events for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (ListAgentEngineSessionEventsConfig): - Optional. Additional configurations for listing the Agent Engine session events. + config (ListRuntimeSessionEventsConfig): + Optional. Additional configurations for listing the Agent Runtime session events. Returns: - ListAgentEngineSessionEventsResponse: The requested Agent Engine session events. + ListRuntimeSessionEventsResponse: The requested Agent Runtime session events. """ - parameter_model = types._ListAgentEngineSessionEventsRequestParameters( + parameter_model = types._ListRuntimeSessionEventsRequestParameters( name=name, config=config, ) @@ -458,7 +458,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineSessionEventsRequestParameters_to_vertex( + request_dict = _ListRuntimeSessionEventsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -489,7 +489,7 @@ async def _list( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListAgentEngineSessionEventsResponse._from_response( + return_value = types.ListRuntimeSessionEventsResponse._from_response( response=response_dict, kwargs=( { @@ -517,14 +517,14 @@ async def list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionEventsConfigOrDict] = None, + config: Optional[types.ListRuntimeSessionEventsConfigOrDict] = None, ) -> AsyncPager[types.SessionEvent]: - """Lists Agent Engine session events. + """Lists Agent Runtime session events. Args: - name (str): Required. The name of the agent engine to list session + name (str): Required. The name of the agent runtime to list session events for. - config (ListAgentEngineSessionEventsConfig): Optional. The configuration + config (ListRuntimeSessionEventsConfig): Optional. The configuration for the session events to list. Currently, the `filter` field in `config` only supports filtering by `timestamp`. The timestamp value must be enclosed in double quotes and include the time zone diff --git a/agentplatform/_genai/sessions.py b/agentplatform/_genai/sessions.py index bcd4862494..61ade3fa4d 100644 --- a/agentplatform/_genai/sessions.py +++ b/agentplatform/_genai/sessions.py @@ -29,7 +29,7 @@ from google.genai._common import set_value_by_path as setv from google.genai.pagers import AsyncPager, Pager -from . import _agent_engines_utils +from . import _runtimes_utils from . import types if typing.TYPE_CHECKING: @@ -43,7 +43,7 @@ logger.setLevel(logging.INFO) -def _CreateAgentEngineSessionConfig_to_vertex( +def _CreateRuntimeSessionConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -70,7 +70,7 @@ def _CreateAgentEngineSessionConfig_to_vertex( return to_object -def _CreateAgentEngineSessionRequestParameters_to_vertex( +def _CreateRuntimeSessionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -82,14 +82,12 @@ def _CreateAgentEngineSessionRequestParameters_to_vertex( setv(to_object, ["userId"], getv(from_object, ["user_id"])) if getv(from_object, ["config"]) is not None: - _CreateAgentEngineSessionConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _CreateRuntimeSessionConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _DeleteAgentEngineSessionRequestParameters_to_vertex( +def _DeleteRuntimeSessionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -100,7 +98,7 @@ def _DeleteAgentEngineSessionRequestParameters_to_vertex( return to_object -def _GetAgentEngineSessionOperationParameters_to_vertex( +def _GetRuntimeSessionOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -113,7 +111,7 @@ def _GetAgentEngineSessionOperationParameters_to_vertex( return to_object -def _GetAgentEngineSessionRequestParameters_to_vertex( +def _GetRuntimeSessionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -124,7 +122,7 @@ def _GetAgentEngineSessionRequestParameters_to_vertex( return to_object -def _ListAgentEngineSessionsConfig_to_vertex( +def _ListRuntimeSessionsConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -142,7 +140,7 @@ def _ListAgentEngineSessionsConfig_to_vertex( return to_object -def _ListAgentEngineSessionsRequestParameters_to_vertex( +def _ListRuntimeSessionsRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -151,14 +149,12 @@ def _ListAgentEngineSessionsRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _ListAgentEngineSessionsConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _ListRuntimeSessionsConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object -def _UpdateAgentEngineSessionConfig_to_vertex( +def _UpdateRuntimeSessionConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -193,7 +189,7 @@ def _UpdateAgentEngineSessionConfig_to_vertex( return to_object -def _UpdateAgentEngineSessionRequestParameters_to_vertex( +def _UpdateRuntimeSessionRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, ) -> dict[str, Any]: @@ -202,9 +198,7 @@ def _UpdateAgentEngineSessionRequestParameters_to_vertex( setv(to_object, ["_url", "name"], getv(from_object, ["name"])) if getv(from_object, ["config"]) is not None: - _UpdateAgentEngineSessionConfig_to_vertex( - getv(from_object, ["config"]), to_object - ) + _UpdateRuntimeSessionConfig_to_vertex(getv(from_object, ["config"]), to_object) return to_object @@ -216,24 +210,24 @@ def _create( *, name: str, user_id: str, - config: Optional[types.CreateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: + config: Optional[types.CreateRuntimeSessionConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: """ - Creates a new session in the Agent Engine. + Creates a new session in the Agent Runtime. Args: - name (str): Required. The name of the Agent Engine to create the session under. Format: + name (str): Required. The name of the Agent Runtime to create the session under. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. user_id (str): Required. The user ID of the session. - config (CreateAgentEngineSessionConfig): - Optional. Additional configurations for creating the Agent Engine session. + config (CreateRuntimeSessionConfig): + Optional. Additional configurations for creating the Agent Runtime session. Returns: - AgentEngineSessionOperation: The operation for creating the Agent Engine session. + RuntimeSessionOperation: The operation for creating the Agent Runtime session. """ - parameter_model = types._CreateAgentEngineSessionRequestParameters( + parameter_model = types._CreateRuntimeSessionRequestParameters( name=name, user_id=user_id, config=config, @@ -245,7 +239,7 @@ def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineSessionRequestParameters_to_vertex( + request_dict = _CreateRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -274,7 +268,7 @@ def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSessionOperation._from_response( + return_value = types.RuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -302,23 +296,23 @@ def delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineSessionConfigOrDict] = None, - ) -> types.DeleteAgentEngineSessionOperation: + config: Optional[types.DeleteRuntimeSessionConfigOrDict] = None, + ) -> types.DeleteRuntimeSessionOperation: """ - Delete an Agent Engine session. + Delete an Agent Runtime session. Args: - name (str): Required. The name of the Agent Engine session to be deleted. Format: + name (str): Required. The name of the Agent Runtime session to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (DeleteAgentEngineSessionConfig): - Optional. Additional configurations for deleting the Agent Engine session. + config (DeleteRuntimeSessionConfig): + Optional. Additional configurations for deleting the Agent Runtime session. Returns: - DeleteAgentEngineSessionOperation: The operation for deleting the Agent Engine session. + DeleteRuntimeSessionOperation: The operation for deleting the Agent Runtime session. """ - parameter_model = types._DeleteAgentEngineSessionRequestParameters( + parameter_model = types._DeleteRuntimeSessionRequestParameters( name=name, config=config, ) @@ -329,7 +323,7 @@ def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineSessionRequestParameters_to_vertex( + request_dict = _DeleteRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -358,7 +352,7 @@ def delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineSessionOperation._from_response( + return_value = types.DeleteRuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -383,26 +377,23 @@ def delete( return return_value def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineSessionConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeSessionConfigOrDict] = None ) -> types.Session: """ - Gets an agent engine session. + Gets an agent runtime session. Args: - name (str): Required. The name of the Agent Engine session to get. Format: + name (str): Required. The name of the Agent Runtime session to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (GetAgentEngineSessionConfig): - Optional. Additional configurations for getting the Agent Engine session. + config (GetRuntimeSessionConfig): + Optional. Additional configurations for getting the Agent Runtime session. Returns: - AgentEngineSession: The requested Agent Engine session. + RuntimeSession: The requested Agent Runtime session. """ - parameter_model = types._GetAgentEngineSessionRequestParameters( + parameter_model = types._GetRuntimeSessionRequestParameters( name=name, config=config, ) @@ -413,7 +404,7 @@ def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSessionRequestParameters_to_vertex( + request_dict = _GetRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -470,23 +461,23 @@ def _list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionsConfigOrDict] = None, + config: Optional[types.ListRuntimeSessionsConfigOrDict] = None, ) -> types.ListReasoningEnginesSessionsResponse: """ - Lists Agent Engine sessions. + Lists Agent Runtime sessions. Args: - name (str): Required. The name of the Agent Engine to list sessions for. Format: + name (str): Required. The name of the Agent Runtime to list sessions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineSessionsConfig): - Optional. Additional configurations for listing the Agent Engine sessions. + config (ListRuntimeSessionsConfig): + Optional. Additional configurations for listing the Agent Runtime sessions. Returns: - ListReasoningEnginesSessionsResponse: The requested Agent Engine sessions. + ListReasoningEnginesSessionsResponse: The requested Agent Runtime sessions. """ - parameter_model = types._ListAgentEngineSessionsRequestParameters( + parameter_model = types._ListRuntimeSessionsRequestParameters( name=name, config=config, ) @@ -497,7 +488,7 @@ def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineSessionsRequestParameters_to_vertex( + request_dict = _ListRuntimeSessionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -554,9 +545,9 @@ def _get_session_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: - parameter_model = types._GetAgentEngineSessionOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: + parameter_model = types._GetRuntimeSessionOperationParameters( operation_name=operation_name, config=config, ) @@ -567,7 +558,7 @@ def _get_session_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSessionOperationParameters_to_vertex( + request_dict = _GetRuntimeSessionOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -596,7 +587,7 @@ def _get_session_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSessionOperation._from_response( + return_value = types.RuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -624,23 +615,23 @@ def _update( self, *, name: str, - config: Optional[types.UpdateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: + config: Optional[types.UpdateRuntimeSessionConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: """ - Updates an Agent Engine session. + Updates an Agent Runtime session. Args: - name (str): Required. The name of the Agent Engine session to be updated. Format: + name (str): Required. The name of the Agent Runtime session to be updated. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (UpdateAgentEngineSessionConfig): - Optional. Additional configurations for updating the Agent Engine session. + config (UpdateRuntimeSessionConfig): + Optional. Additional configurations for updating the Agent Runtime session. Returns: - AgentEngineSessionOperation: The operation for updating the Agent Engine session. + RuntimeSessionOperation: The operation for updating the Agent Runtime session. """ - parameter_model = types._UpdateAgentEngineSessionRequestParameters( + parameter_model = types._UpdateRuntimeSessionRequestParameters( name=name, config=config, ) @@ -651,7 +642,7 @@ def _update( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _UpdateAgentEngineSessionRequestParameters_to_vertex( + request_dict = _UpdateRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -680,7 +671,7 @@ def _update( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSessionOperation._from_response( + return_value = types.RuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -715,7 +706,7 @@ def events(self) -> "session_events_module.SessionEvents": self._events = importlib.import_module(".session_events", __package__) except ImportError as e: raise ImportError( - "The 'agent_engines.sessions.events' module requires" + "The 'runtimes.sessions.events' module requires" "additional packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e @@ -726,25 +717,25 @@ def create( *, name: str, user_id: str, - config: Optional[types.CreateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: - """Creates a new session in the Agent Engine. + config: Optional[types.CreateRuntimeSessionConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: + """Creates a new session in the Agent Runtime. Args: name (str): - Required. The name of the agent engine to create the session for. + Required. The name of the agent runtime to create the session for. user_id (str): Required. The user ID of the session. - config (CreateAgentEngineSessionConfig): + config (CreateRuntimeSessionConfig): Optional. The configuration for the session to create. Returns: - AgentEngineSessionOperation: The operation for creating the session. + RuntimeSessionOperation: The operation for creating the session. """ if config is None: - config = types.CreateAgentEngineSessionConfig() + config = types.CreateRuntimeSessionConfig() elif isinstance(config, dict): - config = types.CreateAgentEngineSessionConfig.model_validate(config) + config = types.CreateRuntimeSessionConfig.model_validate(config) operation = self._create( name=name, user_id=user_id, @@ -752,7 +743,7 @@ def create( ) if config.wait_for_completion: if not operation.done: - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=operation.name, get_operation_fn=self._get_session_operation, poll_interval_seconds=0.5, @@ -774,14 +765,14 @@ def list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionsConfigOrDict] = None, + config: Optional[types.ListRuntimeSessionsConfigOrDict] = None, ) -> Iterator[types.Session]: - """Lists Agent Engine sessions. + """Lists Agent Runtime sessions. Args: - name (str): Required. The name of the agent engine to list sessions + name (str): Required. The name of the agent runtime to list sessions for. - config (ListAgentEngineSessionConfig): Optional. The configuration + config (ListRuntimeSessionConfig): Optional. The configuration for the sessions to list. Returns: @@ -803,24 +794,24 @@ async def _create( *, name: str, user_id: str, - config: Optional[types.CreateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: + config: Optional[types.CreateRuntimeSessionConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: """ - Creates a new session in the Agent Engine. + Creates a new session in the Agent Runtime. Args: - name (str): Required. The name of the Agent Engine to create the session under. Format: + name (str): Required. The name of the Agent Runtime to create the session under. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. user_id (str): Required. The user ID of the session. - config (CreateAgentEngineSessionConfig): - Optional. Additional configurations for creating the Agent Engine session. + config (CreateRuntimeSessionConfig): + Optional. Additional configurations for creating the Agent Runtime session. Returns: - AgentEngineSessionOperation: The operation for creating the Agent Engine session. + RuntimeSessionOperation: The operation for creating the Agent Runtime session. """ - parameter_model = types._CreateAgentEngineSessionRequestParameters( + parameter_model = types._CreateRuntimeSessionRequestParameters( name=name, user_id=user_id, config=config, @@ -832,7 +823,7 @@ async def _create( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _CreateAgentEngineSessionRequestParameters_to_vertex( + request_dict = _CreateRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -863,7 +854,7 @@ async def _create( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSessionOperation._from_response( + return_value = types.RuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -891,23 +882,23 @@ async def delete( self, *, name: str, - config: Optional[types.DeleteAgentEngineSessionConfigOrDict] = None, - ) -> types.DeleteAgentEngineSessionOperation: + config: Optional[types.DeleteRuntimeSessionConfigOrDict] = None, + ) -> types.DeleteRuntimeSessionOperation: """ - Delete an Agent Engine session. + Delete an Agent Runtime session. Args: - name (str): Required. The name of the Agent Engine session to be deleted. Format: + name (str): Required. The name of the Agent Runtime session to be deleted. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (DeleteAgentEngineSessionConfig): - Optional. Additional configurations for deleting the Agent Engine session. + config (DeleteRuntimeSessionConfig): + Optional. Additional configurations for deleting the Agent Runtime session. Returns: - DeleteAgentEngineSessionOperation: The operation for deleting the Agent Engine session. + DeleteRuntimeSessionOperation: The operation for deleting the Agent Runtime session. """ - parameter_model = types._DeleteAgentEngineSessionRequestParameters( + parameter_model = types._DeleteRuntimeSessionRequestParameters( name=name, config=config, ) @@ -918,7 +909,7 @@ async def delete( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _DeleteAgentEngineSessionRequestParameters_to_vertex( + request_dict = _DeleteRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -949,7 +940,7 @@ async def delete( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.DeleteAgentEngineSessionOperation._from_response( + return_value = types.DeleteRuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -974,26 +965,23 @@ async def delete( return return_value async def get( - self, - *, - name: str, - config: Optional[types.GetAgentEngineSessionConfigOrDict] = None, + self, *, name: str, config: Optional[types.GetRuntimeSessionConfigOrDict] = None ) -> types.Session: """ - Gets an agent engine session. + Gets an agent runtime session. Args: - name (str): Required. The name of the Agent Engine session to get. Format: + name (str): Required. The name of the Agent Runtime session to get. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (GetAgentEngineSessionConfig): - Optional. Additional configurations for getting the Agent Engine session. + config (GetRuntimeSessionConfig): + Optional. Additional configurations for getting the Agent Runtime session. Returns: - AgentEngineSession: The requested Agent Engine session. + RuntimeSession: The requested Agent Runtime session. """ - parameter_model = types._GetAgentEngineSessionRequestParameters( + parameter_model = types._GetRuntimeSessionRequestParameters( name=name, config=config, ) @@ -1004,7 +992,7 @@ async def get( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSessionRequestParameters_to_vertex( + request_dict = _GetRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1063,23 +1051,23 @@ async def _list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionsConfigOrDict] = None, + config: Optional[types.ListRuntimeSessionsConfigOrDict] = None, ) -> types.ListReasoningEnginesSessionsResponse: """ - Lists Agent Engine sessions. + Lists Agent Runtime sessions. Args: - name (str): Required. The name of the Agent Engine to list sessions for. Format: + name (str): Required. The name of the Agent Runtime to list sessions for. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - config (ListAgentEngineSessionsConfig): - Optional. Additional configurations for listing the Agent Engine sessions. + config (ListRuntimeSessionsConfig): + Optional. Additional configurations for listing the Agent Runtime sessions. Returns: - ListReasoningEnginesSessionsResponse: The requested Agent Engine sessions. + ListReasoningEnginesSessionsResponse: The requested Agent Runtime sessions. """ - parameter_model = types._ListAgentEngineSessionsRequestParameters( + parameter_model = types._ListRuntimeSessionsRequestParameters( name=name, config=config, ) @@ -1090,7 +1078,7 @@ async def _list( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _ListAgentEngineSessionsRequestParameters_to_vertex( + request_dict = _ListRuntimeSessionsRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1149,9 +1137,9 @@ async def _get_session_operation( self, *, operation_name: str, - config: Optional[types.GetAgentEngineOperationConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: - parameter_model = types._GetAgentEngineSessionOperationParameters( + config: Optional[types.GetRuntimeOperationConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: + parameter_model = types._GetRuntimeSessionOperationParameters( operation_name=operation_name, config=config, ) @@ -1162,7 +1150,7 @@ async def _get_session_operation( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _GetAgentEngineSessionOperationParameters_to_vertex( + request_dict = _GetRuntimeSessionOperationParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1193,7 +1181,7 @@ async def _get_session_operation( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSessionOperation._from_response( + return_value = types.RuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -1221,23 +1209,23 @@ async def _update( self, *, name: str, - config: Optional[types.UpdateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: + config: Optional[types.UpdateRuntimeSessionConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: """ - Updates an Agent Engine session. + Updates an Agent Runtime session. Args: - name (str): Required. The name of the Agent Engine session to be updated. Format: + name (str): Required. The name of the Agent Runtime session to be updated. Format: `projects/{project}/locations/{location}/reasoningEngines/{resource_id}/sessions/{session_id}`. - config (UpdateAgentEngineSessionConfig): - Optional. Additional configurations for updating the Agent Engine session. + config (UpdateRuntimeSessionConfig): + Optional. Additional configurations for updating the Agent Runtime session. Returns: - AgentEngineSessionOperation: The operation for updating the Agent Engine session. + RuntimeSessionOperation: The operation for updating the Agent Runtime session. """ - parameter_model = types._UpdateAgentEngineSessionRequestParameters( + parameter_model = types._UpdateRuntimeSessionRequestParameters( name=name, config=config, ) @@ -1248,7 +1236,7 @@ async def _update( "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." ) else: - request_dict = _UpdateAgentEngineSessionRequestParameters_to_vertex( + request_dict = _UpdateRuntimeSessionRequestParameters_to_vertex( parameter_model ) request_url_dict = request_dict.get("_url") @@ -1279,7 +1267,7 @@ async def _update( response_dict = {} if not response.body else json.loads(response.body) - return_value = types.AgentEngineSessionOperation._from_response( + return_value = types.RuntimeSessionOperation._from_response( response=response_dict, kwargs=( { @@ -1314,7 +1302,7 @@ def events(self) -> "session_events_module.AsyncSessionEvents": self._events = importlib.import_module(".session_events", __package__) except ImportError as e: raise ImportError( - "The 'agent_engines.sessions.events' module requires" + "The 'runtimes.sessions.events' module requires" "additional packages. Please install them using pip install " "google-cloud-aiplatform[agent_engines]" ) from e @@ -1325,25 +1313,25 @@ async def create( *, name: str, user_id: str, - config: Optional[types.CreateAgentEngineSessionConfigOrDict] = None, - ) -> types.AgentEngineSessionOperation: - """Creates a new session in the Agent Engine. + config: Optional[types.CreateRuntimeSessionConfigOrDict] = None, + ) -> types.RuntimeSessionOperation: + """Creates a new session in the Agent Runtime. Args: name (str): - Required. The name of the agent engine to create the session for. + Required. The name of the agent runtime to create the session for. user_id (str): Required. The user ID of the session. - config (CreateAgentEngineSessionConfig): + config (CreateRuntimeSessionConfig): Optional. The configuration for the session to create. Returns: - AgentEngineSessionOperation: The operation for creating the session. + RuntimeSessionOperation: The operation for creating the session. """ if config is None: - config = types.CreateAgentEngineSessionConfig() + config = types.CreateRuntimeSessionConfig() elif isinstance(config, dict): - config = types.CreateAgentEngineSessionConfig.model_validate(config) + config = types.CreateRuntimeSessionConfig.model_validate(config) operation = await self._create( name=name, user_id=user_id, @@ -1351,7 +1339,7 @@ async def create( ) if config.wait_for_completion: if not operation.done: - operation = await _agent_engines_utils._await_async_operation( + operation = await _runtimes_utils._await_async_operation( operation_name=operation.name, get_operation_fn=self._get_session_operation, poll_interval_seconds=0.5, @@ -1373,14 +1361,14 @@ async def list( self, *, name: str, - config: Optional[types.ListAgentEngineSessionsConfigOrDict] = None, + config: Optional[types.ListRuntimeSessionsConfigOrDict] = None, ) -> AsyncPager[types.Session]: - """Lists Agent Engine sessions. + """Lists Agent Runtime sessions. Args: - name (str): Required. The name of the agent engine to list sessions + name (str): Required. The name of the agent runtime to list sessions for. - config (ListAgentEngineSessionConfig): Optional. The configuration + config (ListRuntimeSessionConfig): Optional. The configuration for the sessions to list. Returns: diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index a731e74a0c..17cbe0a7f5 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -22,18 +22,12 @@ from . import agent_engines from . import evals from . import prompts -from .common import _AppendAgentEngineSessionEventRequestParameters -from .common import _AppendAgentEngineTaskEventRequestParameters +from .common import _AppendRuntimeSessionEventRequestParameters from .common import _AskContextsRequestParameters from .common import _AssembleDatasetParameters from .common import _AssessDatasetParameters -from .common import _CancelQueryJobAgentEngineRequestParameters -from .common import _CheckQueryJobAgentEngineRequestParameters -from .common import _CreateAgentEngineMemoryRequestParameters -from .common import _CreateAgentEngineRequestParameters -from .common import _CreateAgentEngineSandboxRequestParameters -from .common import _CreateAgentEngineSessionRequestParameters -from .common import _CreateAgentEngineTaskRequestParameters +from .common import _CancelQueryJobRuntimeRequestParameters +from .common import _CheckQueryJobRuntimeRequestParameters from .common import _CreateDatasetParameters from .common import _CreateDatasetVersionParameters from .common import _CreateEvaluationItemParameters @@ -43,17 +37,15 @@ from .common import _CreateMultimodalDatasetParameters from .common import _CreateRagCorpusRequestParameters from .common import _CreateRuntimeFeedbackEntryRequestParameters +from .common import _CreateRuntimeMemoryRequestParameters +from .common import _CreateRuntimeRequestParameters +from .common import _CreateRuntimeSandboxRequestParameters +from .common import _CreateRuntimeSessionRequestParameters from .common import _CreateSandboxEnvironmentSnapshotRequestParameters from .common import _CreateSandboxEnvironmentTemplateRequestParameters from .common import _CreateSkillRequestParameters from .common import _CustomJobParameters from .common import _CustomJobParameters -from .common import _DeleteAgentEngineMemoryRequestParameters -from .common import _DeleteAgentEngineRequestParameters -from .common import _DeleteAgentEngineRuntimeRevisionRequestParameters -from .common import _DeleteAgentEngineSandboxRequestParameters -from .common import _DeleteAgentEngineSessionRequestParameters -from .common import _DeleteAgentEngineTaskRequestParameters from .common import _DeleteDatasetRequestParameters from .common import _DeleteEvaluationMetricParameters from .common import _DeleteMultimodalDatasetRequestParameters @@ -61,35 +53,27 @@ from .common import _DeleteRagCorpusRequestParameters from .common import _DeleteRagFileRequestParameters from .common import _DeleteRuntimeFeedbackEntryRequestParameters +from .common import _DeleteRuntimeMemoryRequestParameters +from .common import _DeleteRuntimeRequestParameters +from .common import _DeleteRuntimeRevisionRequestParameters +from .common import _DeleteRuntimeSandboxRequestParameters +from .common import _DeleteRuntimeSessionRequestParameters from .common import _DeleteSandboxEnvironmentSnapshotRequestParameters from .common import _DeleteSandboxEnvironmentTemplateRequestParameters from .common import _DeleteSkillRequestParameters from .common import _EvaluateInstancesRequestParameters -from .common import _ExecuteCodeAgentEngineSandboxRequestParameters -from .common import _GenerateAgentEngineMemoriesRequestParameters +from .common import _ExecuteCodeRuntimeSandboxRequestParameters from .common import _GenerateInstanceRubricsRequest from .common import _GenerateLossClustersParameters +from .common import _GenerateRuntimeMemoriesRequestParameters from .common import _GenerateUserScenariosParameters -from .common import _GetAgentEngineGenerateMemoriesOperationParameters -from .common import _GetAgentEngineMemoryOperationParameters -from .common import _GetAgentEngineMemoryRequestParameters -from .common import _GetAgentEngineMemoryRevisionRequestParameters -from .common import _GetAgentEngineOperationParameters -from .common import _GetAgentEngineRequestParameters -from .common import _GetAgentEngineRuntimeRevisionRequestParameters -from .common import _GetAgentEngineSandboxOperationParameters -from .common import _GetAgentEngineSandboxRequestParameters -from .common import _GetAgentEngineSandboxSnapshotOperationParameters -from .common import _GetAgentEngineSessionOperationParameters -from .common import _GetAgentEngineSessionRequestParameters -from .common import _GetAgentEngineTaskRequestParameters from .common import _GetCorpusOperationParameters from .common import _GetCustomJobParameters from .common import _GetCustomJobParameters from .common import _GetDatasetOperationParameters from .common import _GetDatasetParameters from .common import _GetDatasetVersionParameters -from .common import _GetDeleteAgentEngineRuntimeRevisionOperationParameters +from .common import _GetDeleteRuntimeRevisionOperationParameters from .common import _GetEvaluationItemParameters from .common import _GetEvaluationMetricParameters from .common import _GetEvaluationRunParameters @@ -107,6 +91,18 @@ from .common import _GetRuntimeFeedbackOperationParameters from .common import _GetRuntimeFeedbackOperationParameters from .common import _GetRuntimeFeedbackRequestParameters +from .common import _GetRuntimeGenerateMemoriesOperationParameters +from .common import _GetRuntimeMemoryOperationParameters +from .common import _GetRuntimeMemoryRequestParameters +from .common import _GetRuntimeMemoryRevisionRequestParameters +from .common import _GetRuntimeOperationParameters +from .common import _GetRuntimeRequestParameters +from .common import _GetRuntimeRevisionRequestParameters +from .common import _GetRuntimeSandboxOperationParameters +from .common import _GetRuntimeSandboxRequestParameters +from .common import _GetRuntimeSandboxSnapshotOperationParameters +from .common import _GetRuntimeSessionOperationParameters +from .common import _GetRuntimeSessionRequestParameters from .common import _GetSandboxEnvironmentSnapshotRequestParameters from .common import _GetSandboxEnvironmentTemplateOperationParameters from .common import _GetSandboxEnvironmentTemplateRequestParameters @@ -115,15 +111,6 @@ from .common import _GetSkillRevisionRequestParameters from .common import _ImportRagFilesRequestParameters from .common import _IngestEventsRequestParameters -from .common import _ListAgentEngineMemoryRequestParameters -from .common import _ListAgentEngineMemoryRevisionsRequestParameters -from .common import _ListAgentEngineRequestParameters -from .common import _ListAgentEngineRuntimeRevisionsRequestParameters -from .common import _ListAgentEngineSandboxesRequestParameters -from .common import _ListAgentEngineSessionEventsRequestParameters -from .common import _ListAgentEngineSessionsRequestParameters -from .common import _ListAgentEngineTaskEventsRequestParameters -from .common import _ListAgentEngineTasksRequestParameters from .common import _ListDatasetsRequestParameters from .common import _ListDatasetVersionsRequestParameters from .common import _ListEvaluationMetricsParameters @@ -132,78 +119,48 @@ from .common import _ListRagCorporaRequestParameters from .common import _ListRagFilesRequestParameters from .common import _ListRuntimeFeedbackEntriesRequestParameters +from .common import _ListRuntimeMemoryRequestParameters +from .common import _ListRuntimeMemoryRevisionsRequestParameters +from .common import _ListRuntimeRequestParameters +from .common import _ListRuntimeRevisionsRequestParameters +from .common import _ListRuntimeSandboxesRequestParameters +from .common import _ListRuntimeSessionEventsRequestParameters +from .common import _ListRuntimeSessionsRequestParameters from .common import _ListSandboxEnvironmentSnapshotsRequestParameters from .common import _ListSandboxEnvironmentTemplatesRequestParameters from .common import _ListSkillRevisionsRequestParameters from .common import _ListSkillsRequestParameters from .common import _OptimizeRequestParameters from .common import _OptimizeRequestParameters -from .common import _PurgeAgentEngineMemoriesRequestParameters -from .common import _QueryAgentEngineRequestParameters -from .common import _QueryAgentEngineRuntimeRevisionRequestParameters +from .common import _PurgeRuntimeMemoriesRequestParameters +from .common import _QueryRuntimeRequestParameters +from .common import _QueryRuntimeRevisionRequestParameters from .common import _RecommendSpecRequestParameters from .common import _RestoreVersionRequestParameters -from .common import _RetrieveAgentEngineMemoriesRequestParameters from .common import _RetrieveMemoryProfilesRequestParameters from .common import _RetrieveRagContextsRequestParameters +from .common import _RetrieveRuntimeMemoriesRequestParameters from .common import _RetrieveSkillsRequestParameters -from .common import _RollbackAgentEngineMemoryRequestParameters -from .common import _RunQueryJobAgentEngineConfig -from .common import _RunQueryJobAgentEngineConfigDict -from .common import _RunQueryJobAgentEngineConfigOrDict -from .common import _RunQueryJobAgentEngineRequestParameters -from .common import _UpdateAgentEngineMemoryRequestParameters -from .common import _UpdateAgentEngineRequestParameters -from .common import _UpdateAgentEngineSessionRequestParameters +from .common import _RollbackRuntimeMemoryRequestParameters +from .common import _RunQueryJobRuntimeConfig +from .common import _RunQueryJobRuntimeConfigDict +from .common import _RunQueryJobRuntimeConfigOrDict +from .common import _RunQueryJobRuntimeRequestParameters from .common import _UpdateDatasetParameters from .common import _UpdateMultimodalDatasetParameters from .common import _UpdateRagConfigRequestParameters from .common import _UpdateRagCorpusRequestParameters from .common import _UpdateRuntimeFeedbackContextRequestParameters from .common import _UpdateRuntimeFeedbackEntryRequestParameters +from .common import _UpdateRuntimeMemoryRequestParameters +from .common import _UpdateRuntimeRequestParameters +from .common import _UpdateRuntimeSessionRequestParameters from .common import _UpdateSkillRequestParameters from .common import _UploadRagFileParameters -from .common import A2aTask -from .common import A2aTaskDict -from .common import A2aTaskOrDict -from .common import A2aTaskState +from .common import AcceleratorRequirement +from .common import AcceleratorRequirementDict +from .common import AcceleratorRequirementOrDict from .common import AcceleratorType -from .common import AgentEngine -from .common import AgentEngineConfig -from .common import AgentEngineConfigDict -from .common import AgentEngineConfigOrDict -from .common import AgentEngineDict -from .common import AgentEngineGenerateMemoriesOperation -from .common import AgentEngineGenerateMemoriesOperationDict -from .common import AgentEngineGenerateMemoriesOperationOrDict -from .common import AgentEngineMemoryConfig -from .common import AgentEngineMemoryConfigDict -from .common import AgentEngineMemoryConfigOrDict -from .common import AgentEngineMemoryOperation -from .common import AgentEngineMemoryOperationDict -from .common import AgentEngineMemoryOperationOrDict -from .common import AgentEngineOperation -from .common import AgentEngineOperationDict -from .common import AgentEngineOperationOrDict -from .common import AgentEngineOrDict -from .common import AgentEnginePurgeMemoriesOperation -from .common import AgentEnginePurgeMemoriesOperationDict -from .common import AgentEnginePurgeMemoriesOperationOrDict -from .common import AgentEngineRollbackMemoryOperation -from .common import AgentEngineRollbackMemoryOperationDict -from .common import AgentEngineRollbackMemoryOperationOrDict -from .common import AgentEngineRuntimeRevision -from .common import AgentEngineRuntimeRevisionDict -from .common import AgentEngineRuntimeRevisionOrDict -from .common import AgentEngineSandboxOperation -from .common import AgentEngineSandboxOperationDict -from .common import AgentEngineSandboxOperationOrDict -from .common import AgentEngineSandboxSnapshotOperation -from .common import AgentEngineSandboxSnapshotOperationDict -from .common import AgentEngineSandboxSnapshotOperationOrDict -from .common import AgentEngineSessionOperation -from .common import AgentEngineSessionOperationDict -from .common import AgentEngineSessionOperationOrDict from .common import AgentRunConfig from .common import AgentRunConfigDict from .common import AgentRunConfigOrDict @@ -214,18 +171,15 @@ 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 ApiservingMediaResponseInfo +from .common import ApiservingMediaResponseInfoDict +from .common import ApiservingMediaResponseInfoOrDict +from .common import AppendRuntimeSessionEventConfig +from .common import AppendRuntimeSessionEventConfigDict +from .common import AppendRuntimeSessionEventConfigOrDict +from .common import AppendRuntimeSessionEventResponse +from .common import AppendRuntimeSessionEventResponseDict +from .common import AppendRuntimeSessionEventResponseOrDict from .common import AskContextsConfig from .common import AskContextsConfigDict from .common import AskContextsConfigOrDict @@ -277,33 +231,43 @@ from .common import BleuResults from .common import BleuResultsDict from .common import BleuResultsOrDict -from .common import CancelQueryJobAgentEngineConfig -from .common import CancelQueryJobAgentEngineConfigDict -from .common import CancelQueryJobAgentEngineConfigOrDict from .common import CancelQueryJobResult from .common import CancelQueryJobResultDict from .common import CancelQueryJobResultOrDict +from .common import CancelQueryJobRuntimeConfig +from .common import CancelQueryJobRuntimeConfigDict +from .common import CancelQueryJobRuntimeConfigOrDict from .common import CandidateResponse from .common import CandidateResponseDict from .common import CandidateResponseOrDict from .common import CandidateResult from .common import CandidateResultDict -from .common import CheckQueryJobAgentEngineConfig -from .common import CheckQueryJobAgentEngineConfigDict -from .common import CheckQueryJobAgentEngineConfigOrDict from .common import CheckQueryJobResponse from .common import CheckQueryJobResponseDict from .common import CheckQueryJobResponseOrDict from .common import CheckQueryJobResult from .common import CheckQueryJobResultDict from .common import CheckQueryJobResultOrDict +from .common import CheckQueryJobRuntimeConfig +from .common import CheckQueryJobRuntimeConfigDict +from .common import CheckQueryJobRuntimeConfigOrDict from .common import Chunk from .common import ChunkDict from .common import ChunkOrDict +from .common import CloudLoggingResult +from .common import CloudLoggingResultDict +from .common import CloudLoggingResultOrDict from .common import CodeExecutionMetric +from .common import CodeExecutionResults +from .common import CodeExecutionResultsDict +from .common import CodeExecutionResultsOrDict from .common import CometResult from .common import CometResultDict from .common import CometResultOrDict +from .common import ConfidentialInstanceConfig +from .common import ConfidentialInstanceConfigDict +from .common import ConfidentialInstanceConfigOrDict +from .common import ConfidentialInstanceType from .common import ContainerSpec from .common import ContainerSpecDict from .common import ContainerSpecOrDict @@ -319,21 +283,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 @@ -367,9 +316,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 @@ -401,39 +362,6 @@ from .common import DedicatedResourcesScaleToZeroSpecDict from .common import DedicatedResourcesScaleToZeroSpecOrDict from .common import DefaultContainerCategory -from .common import DeleteAgentEngineConfig -from .common import DeleteAgentEngineConfigDict -from .common import DeleteAgentEngineConfigOrDict -from .common import DeleteAgentEngineMemoryConfig -from .common import DeleteAgentEngineMemoryConfigDict -from .common import DeleteAgentEngineMemoryConfigOrDict -from .common import DeleteAgentEngineMemoryOperation -from .common import DeleteAgentEngineMemoryOperationDict -from .common import DeleteAgentEngineMemoryOperationOrDict -from .common import DeleteAgentEngineOperation -from .common import DeleteAgentEngineOperationDict -from .common import DeleteAgentEngineOperationOrDict -from .common import DeleteAgentEngineRuntimeRevisionConfig -from .common import DeleteAgentEngineRuntimeRevisionConfigDict -from .common import DeleteAgentEngineRuntimeRevisionConfigOrDict -from .common import DeleteAgentEngineRuntimeRevisionOperation -from .common import DeleteAgentEngineRuntimeRevisionOperationDict -from .common import DeleteAgentEngineRuntimeRevisionOperationOrDict -from .common import DeleteAgentEngineSandboxConfig -from .common import DeleteAgentEngineSandboxConfigDict -from .common import DeleteAgentEngineSandboxConfigOrDict -from .common import DeleteAgentEngineSandboxOperation -from .common import DeleteAgentEngineSandboxOperationDict -from .common import DeleteAgentEngineSandboxOperationOrDict -from .common import DeleteAgentEngineSessionConfig -from .common import DeleteAgentEngineSessionConfigDict -from .common import DeleteAgentEngineSessionConfigOrDict -from .common import DeleteAgentEngineSessionOperation -from .common import DeleteAgentEngineSessionOperationDict -from .common import DeleteAgentEngineSessionOperationOrDict -from .common import DeleteAgentEngineTaskConfig -from .common import DeleteAgentEngineTaskConfigDict -from .common import DeleteAgentEngineTaskConfigOrDict from .common import DeleteEvaluationMetricConfig from .common import DeleteEvaluationMetricConfigDict from .common import DeleteEvaluationMetricConfigOrDict @@ -461,12 +389,42 @@ from .common import DeleteRagFileOperation from .common import DeleteRagFileOperationDict from .common import DeleteRagFileOperationOrDict +from .common import DeleteRuntimeConfig +from .common import DeleteRuntimeConfigDict +from .common import DeleteRuntimeConfigOrDict from .common import DeleteRuntimeFeedbackEntryConfig from .common import DeleteRuntimeFeedbackEntryConfigDict from .common import DeleteRuntimeFeedbackEntryConfigOrDict from .common import DeleteRuntimeFeedbackEntryOperation from .common import DeleteRuntimeFeedbackEntryOperationDict from .common import DeleteRuntimeFeedbackEntryOperationOrDict +from .common import DeleteRuntimeMemoryConfig +from .common import DeleteRuntimeMemoryConfigDict +from .common import DeleteRuntimeMemoryConfigOrDict +from .common import DeleteRuntimeMemoryOperation +from .common import DeleteRuntimeMemoryOperationDict +from .common import DeleteRuntimeMemoryOperationOrDict +from .common import DeleteRuntimeOperation +from .common import DeleteRuntimeOperationDict +from .common import DeleteRuntimeOperationOrDict +from .common import DeleteRuntimeRevisionConfig +from .common import DeleteRuntimeRevisionConfigDict +from .common import DeleteRuntimeRevisionConfigOrDict +from .common import DeleteRuntimeRevisionOperation +from .common import DeleteRuntimeRevisionOperationDict +from .common import DeleteRuntimeRevisionOperationOrDict +from .common import DeleteRuntimeSandboxConfig +from .common import DeleteRuntimeSandboxConfigDict +from .common import DeleteRuntimeSandboxConfigOrDict +from .common import DeleteRuntimeSandboxOperation +from .common import DeleteRuntimeSandboxOperationDict +from .common import DeleteRuntimeSandboxOperationOrDict +from .common import DeleteRuntimeSessionConfig +from .common import DeleteRuntimeSessionConfigDict +from .common import DeleteRuntimeSessionConfigOrDict +from .common import DeleteRuntimeSessionOperation +from .common import DeleteRuntimeSessionOperationDict +from .common import DeleteRuntimeSessionOperationOrDict from .common import DeleteSandboxEnvironmentSnapshotConfig from .common import DeleteSandboxEnvironmentSnapshotConfigDict from .common import DeleteSandboxEnvironmentSnapshotConfigOrDict @@ -610,12 +568,15 @@ from .common import ExactMatchSpec from .common import ExactMatchSpecDict from .common import ExactMatchSpecOrDict -from .common import ExecuteCodeAgentEngineSandboxConfig -from .common import ExecuteCodeAgentEngineSandboxConfigDict -from .common import ExecuteCodeAgentEngineSandboxConfigOrDict +from .common import ExecuteCodeRuntimeSandboxConfig +from .common import ExecuteCodeRuntimeSandboxConfigDict +from .common import ExecuteCodeRuntimeSandboxConfigOrDict from .common import ExecuteSandboxEnvironmentResponse from .common import ExecuteSandboxEnvironmentResponseDict from .common import ExecuteSandboxEnvironmentResponseOrDict +from .common import ExperimentConfig +from .common import ExperimentConfigDict +from .common import ExperimentConfigOrDict from .common import FailedRubric from .common import FailedRubricDict from .common import FailedRubricOrDict @@ -626,10 +587,49 @@ from .common import FeedbackEntryDict from .common import FeedbackEntryOrDict from .common import FeedbackType +from .common import FleschKincaidReadabilityMetricValue +from .common import FleschKincaidReadabilityMetricValueDict +from .common import FleschKincaidReadabilityMetricValueOrDict +from .common import FleschKincaidReadabilityResults +from .common import FleschKincaidReadabilityResultsDict +from .common import FleschKincaidReadabilityResultsOrDict from .common import FlexStart from .common import FlexStartDict from .common import FlexStartOrDict from .common import Framework +from .common import GdataBlobstore2Info +from .common import GdataBlobstore2InfoDict +from .common import GdataBlobstore2InfoOrDict +from .common import GdataCompositeMedia +from .common import GdataCompositeMediaDict +from .common import GdataCompositeMediaOrDict +from .common import GdataContentTypeInfo +from .common import GdataContentTypeInfoDict +from .common import GdataContentTypeInfoOrDict +from .common import GdataDiffChecksumsResponse +from .common import GdataDiffChecksumsResponseDict +from .common import GdataDiffChecksumsResponseOrDict +from .common import GdataDiffDownloadResponse +from .common import GdataDiffDownloadResponseDict +from .common import GdataDiffDownloadResponseOrDict +from .common import GdataDiffUploadRequest +from .common import GdataDiffUploadRequestDict +from .common import GdataDiffUploadRequestOrDict +from .common import GdataDiffUploadResponse +from .common import GdataDiffUploadResponseDict +from .common import GdataDiffUploadResponseOrDict +from .common import GdataDiffVersionResponse +from .common import GdataDiffVersionResponseDict +from .common import GdataDiffVersionResponseOrDict +from .common import GdataDownloadParameters +from .common import GdataDownloadParametersDict +from .common import GdataDownloadParametersOrDict +from .common import GdataMedia +from .common import GdataMediaDict +from .common import GdataMediaOrDict +from .common import GdataObjectId +from .common import GdataObjectIdDict +from .common import GdataObjectIdOrDict from .common import GeminiAgentConfig from .common import GeminiAgentConfigDict from .common import GeminiAgentConfigOrDict @@ -642,9 +642,6 @@ from .common import GeminiTemplateConfig from .common import GeminiTemplateConfigDict from .common import GeminiTemplateConfigOrDict -from .common import GenerateAgentEngineMemoriesConfig -from .common import GenerateAgentEngineMemoriesConfigDict -from .common import GenerateAgentEngineMemoriesConfigOrDict from .common import GenerateInstanceRubricsResponse from .common import GenerateInstanceRubricsResponseDict from .common import GenerateInstanceRubricsResponseOrDict @@ -679,45 +676,24 @@ from .common import GenerateMemoriesResponseGeneratedMemoryDict from .common import GenerateMemoriesResponseGeneratedMemoryOrDict from .common import GenerateMemoriesResponseOrDict +from .common import GenerateRuntimeMemoriesConfig +from .common import GenerateRuntimeMemoriesConfigDict +from .common import GenerateRuntimeMemoriesConfigOrDict from .common import GenerateUserScenariosConfig from .common import GenerateUserScenariosConfigDict from .common import GenerateUserScenariosConfigOrDict from .common import GenerateUserScenariosResponse from .common import GenerateUserScenariosResponseDict from .common import GenerateUserScenariosResponseOrDict -from .common import GetAgentEngineConfig -from .common import GetAgentEngineConfigDict -from .common import GetAgentEngineConfigOrDict -from .common import GetAgentEngineMemoryConfig -from .common import GetAgentEngineMemoryConfigDict -from .common import GetAgentEngineMemoryConfigOrDict -from .common import GetAgentEngineMemoryRevisionConfig -from .common import GetAgentEngineMemoryRevisionConfigDict -from .common import GetAgentEngineMemoryRevisionConfigOrDict -from .common import GetAgentEngineOperationConfig -from .common import GetAgentEngineOperationConfigDict -from .common import GetAgentEngineOperationConfigOrDict -from .common import GetAgentEngineRuntimeRevisionConfig -from .common import GetAgentEngineRuntimeRevisionConfigDict -from .common import GetAgentEngineRuntimeRevisionConfigOrDict -from .common import GetAgentEngineSandboxConfig -from .common import GetAgentEngineSandboxConfigDict -from .common import GetAgentEngineSandboxConfigOrDict -from .common import GetAgentEngineSessionConfig -from .common import GetAgentEngineSessionConfigDict -from .common import GetAgentEngineSessionConfigOrDict -from .common import GetAgentEngineTaskConfig -from .common import GetAgentEngineTaskConfigDict -from .common import GetAgentEngineTaskConfigOrDict from .common import GetCorpusOperationConfig from .common import GetCorpusOperationConfigDict from .common import GetCorpusOperationConfigOrDict from .common import GetDatasetOperationConfig from .common import GetDatasetOperationConfigDict from .common import GetDatasetOperationConfigOrDict -from .common import GetDeleteAgentEngineRuntimeRevisionOperationConfig -from .common import GetDeleteAgentEngineRuntimeRevisionOperationConfigDict -from .common import GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict +from .common import GetDeleteRuntimeRevisionOperationConfig +from .common import GetDeleteRuntimeRevisionOperationConfigDict +from .common import GetDeleteRuntimeRevisionOperationConfigOrDict from .common import GetEvaluationItemConfig from .common import GetEvaluationItemConfigDict from .common import GetEvaluationItemConfigOrDict @@ -754,6 +730,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 @@ -766,6 +745,24 @@ from .common import GetRuntimeFeedbackEntryConfig from .common import GetRuntimeFeedbackEntryConfigDict from .common import GetRuntimeFeedbackEntryConfigOrDict +from .common import GetRuntimeMemoryConfig +from .common import GetRuntimeMemoryConfigDict +from .common import GetRuntimeMemoryConfigOrDict +from .common import GetRuntimeMemoryRevisionConfig +from .common import GetRuntimeMemoryRevisionConfigDict +from .common import GetRuntimeMemoryRevisionConfigOrDict +from .common import GetRuntimeOperationConfig +from .common import GetRuntimeOperationConfigDict +from .common import GetRuntimeOperationConfigOrDict +from .common import GetRuntimeRevisionConfig +from .common import GetRuntimeRevisionConfigDict +from .common import GetRuntimeRevisionConfigOrDict +from .common import GetRuntimeSandboxConfig +from .common import GetRuntimeSandboxConfigDict +from .common import GetRuntimeSandboxConfigOrDict +from .common import GetRuntimeSessionConfig +from .common import GetRuntimeSessionConfigDict +from .common import GetRuntimeSessionConfigOrDict from .common import GetSandboxEnvironmentSnapshotConfig from .common import GetSandboxEnvironmentSnapshotConfigDict from .common import GetSandboxEnvironmentSnapshotConfigOrDict @@ -837,48 +834,6 @@ from .common import LargeModelReferenceDict from .common import LargeModelReferenceOrDict from .common import LaunchStage -from .common import ListAgentEngineConfig -from .common import ListAgentEngineConfigDict -from .common import ListAgentEngineConfigOrDict -from .common import ListAgentEngineMemoryConfig -from .common import ListAgentEngineMemoryConfigDict -from .common import ListAgentEngineMemoryConfigOrDict -from .common import ListAgentEngineMemoryRevisionsConfig -from .common import ListAgentEngineMemoryRevisionsConfigDict -from .common import ListAgentEngineMemoryRevisionsConfigOrDict -from .common import ListAgentEngineMemoryRevisionsResponse -from .common import ListAgentEngineMemoryRevisionsResponseDict -from .common import ListAgentEngineMemoryRevisionsResponseOrDict -from .common import ListAgentEngineRuntimeRevisionsConfig -from .common import ListAgentEngineRuntimeRevisionsConfigDict -from .common import ListAgentEngineRuntimeRevisionsConfigOrDict -from .common import ListAgentEngineSandboxesConfig -from .common import ListAgentEngineSandboxesConfigDict -from .common import ListAgentEngineSandboxesConfigOrDict -from .common import ListAgentEngineSandboxesResponse -from .common import ListAgentEngineSandboxesResponseDict -from .common import ListAgentEngineSandboxesResponseOrDict -from .common import ListAgentEngineSessionEventsConfig -from .common import ListAgentEngineSessionEventsConfigDict -from .common import ListAgentEngineSessionEventsConfigOrDict -from .common import ListAgentEngineSessionEventsResponse -from .common import ListAgentEngineSessionEventsResponseDict -from .common import ListAgentEngineSessionEventsResponseOrDict -from .common import ListAgentEngineSessionsConfig -from .common import ListAgentEngineSessionsConfigDict -from .common import ListAgentEngineSessionsConfigOrDict -from .common import ListAgentEngineTaskEventsConfig -from .common import ListAgentEngineTaskEventsConfigDict -from .common import ListAgentEngineTaskEventsConfigOrDict -from .common import ListAgentEngineTaskEventsResponse -from .common import ListAgentEngineTaskEventsResponseDict -from .common import ListAgentEngineTaskEventsResponseOrDict -from .common import ListAgentEngineTasksConfig -from .common import ListAgentEngineTasksConfigDict -from .common import ListAgentEngineTasksConfigOrDict -from .common import ListAgentEngineTasksResponse -from .common import ListAgentEngineTasksResponseDict -from .common import ListAgentEngineTasksResponseOrDict from .common import ListCustomModelDeployOptionsConfig from .common import ListCustomModelDeployOptionsConfigDict from .common import ListCustomModelDeployOptionsConfigOrDict @@ -942,12 +897,42 @@ from .common import ListReasoningEnginesSessionsResponse from .common import ListReasoningEnginesSessionsResponseDict from .common import ListReasoningEnginesSessionsResponseOrDict +from .common import ListRuntimeConfig +from .common import ListRuntimeConfigDict +from .common import ListRuntimeConfigOrDict from .common import ListRuntimeFeedbackEntriesConfig from .common import ListRuntimeFeedbackEntriesConfigDict from .common import ListRuntimeFeedbackEntriesConfigOrDict from .common import ListRuntimeFeedbackEntriesResponse from .common import ListRuntimeFeedbackEntriesResponseDict from .common import ListRuntimeFeedbackEntriesResponseOrDict +from .common import ListRuntimeMemoryConfig +from .common import ListRuntimeMemoryConfigDict +from .common import ListRuntimeMemoryConfigOrDict +from .common import ListRuntimeMemoryRevisionsConfig +from .common import ListRuntimeMemoryRevisionsConfigDict +from .common import ListRuntimeMemoryRevisionsConfigOrDict +from .common import ListRuntimeMemoryRevisionsResponse +from .common import ListRuntimeMemoryRevisionsResponseDict +from .common import ListRuntimeMemoryRevisionsResponseOrDict +from .common import ListRuntimeRevisionsConfig +from .common import ListRuntimeRevisionsConfigDict +from .common import ListRuntimeRevisionsConfigOrDict +from .common import ListRuntimeSandboxesConfig +from .common import ListRuntimeSandboxesConfigDict +from .common import ListRuntimeSandboxesConfigOrDict +from .common import ListRuntimeSandboxesResponse +from .common import ListRuntimeSandboxesResponseDict +from .common import ListRuntimeSandboxesResponseOrDict +from .common import ListRuntimeSessionEventsConfig +from .common import ListRuntimeSessionEventsConfigDict +from .common import ListRuntimeSessionEventsConfigOrDict +from .common import ListRuntimeSessionEventsResponse +from .common import ListRuntimeSessionEventsResponseDict +from .common import ListRuntimeSessionEventsResponseOrDict +from .common import ListRuntimeSessionsConfig +from .common import ListRuntimeSessionsConfigDict +from .common import ListRuntimeSessionsConfigOrDict from .common import ListSandboxEnvironmentSnapshotsConfig from .common import ListSandboxEnvironmentSnapshotsConfigDict from .common import ListSandboxEnvironmentSnapshotsConfigOrDict @@ -972,6 +957,12 @@ from .common import ListSkillsResponse from .common import ListSkillsResponseDict from .common import ListSkillsResponseOrDict +from .common import LLMBasedClassificationMetricResult +from .common import LLMBasedClassificationMetricResultDict +from .common import LLMBasedClassificationMetricResultOrDict +from .common import LLMBasedSimilarityMetricResult +from .common import LLMBasedSimilarityMetricResultDict +from .common import LLMBasedSimilarityMetricResultOrDict from .common import LLMMetric from .common import LossAnalysisConfig from .common import LossAnalysisConfigDict @@ -999,6 +990,7 @@ from .common import MapInstance from .common import MapInstanceDict from .common import MapInstanceOrDict +from .common import MediaModality from .common import Memory from .common import MemoryBankCustomizationConfig from .common import MemoryBankCustomizationConfigConsolidationConfig @@ -1107,6 +1099,9 @@ from .common import NfsMount from .common import NfsMountDict from .common import NfsMountOrDict +from .common import NodeRecycling +from .common import NodeRecyclingDict +from .common import NodeRecyclingOrDict from .common import ObservabilityEvalCase from .common import ObservabilityEvalCaseDict from .common import ObservabilityEvalCaseOrDict @@ -1126,6 +1121,7 @@ from .common import OptimizeResponseEndpointOrDict from .common import OptimizeResponseOrDict from .common import OptimizeTarget +from .common import Outcome from .common import PairwiseMetricInput from .common import PairwiseMetricInputDict from .common import PairwiseMetricInputOrDict @@ -1139,6 +1135,9 @@ from .common import PointwiseMetricInstance from .common import PointwiseMetricInstanceDict from .common import PointwiseMetricInstanceOrDict +from .common import PointwiseSafetyMetricResult +from .common import PointwiseSafetyMetricResultDict +from .common import PointwiseSafetyMetricResultOrDict from .common import Port from .common import PortDict from .common import PortOrDict @@ -1146,6 +1145,12 @@ from .common import PredictSchemata from .common import PredictSchemataDict from .common import PredictSchemataOrDict +from .common import PrefixMatchMetricValue +from .common import PrefixMatchMetricValueDict +from .common import PrefixMatchMetricValueOrDict +from .common import PrefixMatchResults +from .common import PrefixMatchResultsDict +from .common import PrefixMatchResultsOrDict from .common import Probe from .common import ProbeDict from .common import ProbeExecAction @@ -1187,6 +1192,17 @@ from .common import PromptVersionRefDict from .common import PromptVersionRefOrDict from .common import Protocol +from .common import ProvisioningConfig +from .common import ProvisioningConfigDict +from .common import ProvisioningConfigOrDict +from .common import ProvisioningModelPriority +from .common import ProvisioningModelPriorityDict +from .common import ProvisioningModelPriorityOrDict +from .common import ProvisioningPriority +from .common import ProvisioningPriorityDict +from .common import ProvisioningPriorityOrDict +from .common import ProvisioningStrategyType +from .common import ProvisionMode from .common import PscInterfaceConfig from .common import PscInterfaceConfigDict from .common import PscInterfaceConfigOrDict @@ -1200,6 +1216,56 @@ from .common import PublisherModelCallToActionDeployGke from .common import PublisherModelCallToActionDeployGkeDict from .common import PublisherModelCallToActionDeployGkeOrDict +from .common import PublisherModelCallToActionDeployMonetizedModel +from .common import PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfig +from .common import ( + PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfigDict, +) +from .common import ( + PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfigOrDict, +) +from .common import PublisherModelCallToActionDeployMonetizedModelDict +from .common import PublisherModelCallToActionDeployMonetizedModelMarketplaceListing +from .common import PublisherModelCallToActionDeployMonetizedModelMarketplaceListingDict +from .common import ( + PublisherModelCallToActionDeployMonetizedModelMarketplaceListingOrDict, +) +from .common import ( + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnit, +) +from .common import ( + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModel, +) +from .common import ( + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModelDict, +) +from .common import ( + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModelOrDict, +) +from .common import ( + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitDict, +) +from .common import ( + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitOrDict, +) +from .common import ( + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResources, +) +from .common import ( + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResourcesDict, +) +from .common import ( + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResourcesOrDict, +) +from .common import PublisherModelCallToActionDeployMonetizedModelMonetizationConfig +from .common import PublisherModelCallToActionDeployMonetizedModelMonetizationConfigDict +from .common import ( + PublisherModelCallToActionDeployMonetizedModelMonetizationConfigOrDict, +) +from .common import PublisherModelCallToActionDeployMonetizedModelOrDict +from .common import PublisherModelCallToActionDeployMonetizedModelSingleTenant +from .common import PublisherModelCallToActionDeployMonetizedModelSingleTenantDict +from .common import PublisherModelCallToActionDeployMonetizedModelSingleTenantOrDict from .common import PublisherModelCallToActionDeployOrDict from .common import PublisherModelCallToActionDeployVertex from .common import PublisherModelCallToActionDeployVertexDict @@ -1229,25 +1295,28 @@ from .common import PublisherModelResourceReference from .common import PublisherModelResourceReferenceDict from .common import PublisherModelResourceReferenceOrDict -from .common import PurgeAgentEngineMemoriesConfig -from .common import PurgeAgentEngineMemoriesConfigDict -from .common import PurgeAgentEngineMemoriesConfigOrDict from .common import PurgeMemoriesResponse from .common import PurgeMemoriesResponseDict from .common import PurgeMemoriesResponseOrDict +from .common import PurgeRuntimeMemoriesConfig +from .common import PurgeRuntimeMemoriesConfigDict +from .common import PurgeRuntimeMemoriesConfigOrDict from .common import PythonPackageSpec from .common import PythonPackageSpecDict from .common import PythonPackageSpecOrDict -from .common import QueryAgentEngineConfig -from .common import QueryAgentEngineConfigDict -from .common import QueryAgentEngineConfigOrDict -from .common import QueryAgentEngineRuntimeRevisionConfig -from .common import QueryAgentEngineRuntimeRevisionConfigDict -from .common import QueryAgentEngineRuntimeRevisionConfigOrDict from .common import QueryReasoningEngineResponse from .common import QueryReasoningEngineResponseDict from .common import QueryReasoningEngineResponseOrDict +from .common import QueryRuntimeConfig +from .common import QueryRuntimeConfigDict +from .common import QueryRuntimeConfigOrDict +from .common import QueryRuntimeRevisionConfig +from .common import QueryRuntimeRevisionConfigDict +from .common import QueryRuntimeRevisionConfigOrDict from .common import QuotaState +from .common import RagContextRecallResult +from .common import RagContextRecallResultDict +from .common import RagContextRecallResultOrDict from .common import RagContexts from .common import RagContextsContext from .common import RagContextsContextDict @@ -1376,6 +1445,16 @@ from .common import ReasoningEngine from .common import ReasoningEngineContextSpec from .common import ReasoningEngineContextSpecDict +from .common import ReasoningEngineContextSpecExampleStoreConfig +from .common import ReasoningEngineContextSpecExampleStoreConfigDict +from .common import ReasoningEngineContextSpecExampleStoreConfigOrDict +from .common import ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfig +from .common import ( + ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfigDict, +) +from .common import ( + ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfigOrDict, +) from .common import ReasoningEngineContextSpecMemoryBankConfig from .common import ReasoningEngineContextSpecMemoryBankConfigDict from .common import ReasoningEngineContextSpecMemoryBankConfigGenerationConfig @@ -1400,10 +1479,26 @@ from .common import ReasoningEngineContextSpecOrDict from .common import ReasoningEngineDict from .common import ReasoningEngineOrDict +from .common import ReasoningEngineRevisionGarbageCollectionStrategy +from .common import ReasoningEngineRevisionGarbageCollectionStrategyDict +from .common import ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatest +from .common import ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatestDict +from .common import ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatestOrDict +from .common import ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollection +from .common import ( + ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollectionDict, +) +from .common import ( + ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollectionOrDict, +) +from .common import ReasoningEngineRevisionGarbageCollectionStrategyOrDict from .common import ReasoningEngineRuntimeRevision from .common import ReasoningEngineRuntimeRevisionDict from .common import ReasoningEngineRuntimeRevisionOrDict from .common import ReasoningEngineSpec +from .common import ReasoningEngineSpecBuildSpec +from .common import ReasoningEngineSpecBuildSpecDict +from .common import ReasoningEngineSpecBuildSpecOrDict from .common import ReasoningEngineSpecContainerSpec from .common import ReasoningEngineSpecContainerSpecDict from .common import ReasoningEngineSpecContainerSpecOrDict @@ -1500,15 +1595,15 @@ from .common import ResponseCandidateResult from .common import ResponseCandidateResultDict from .common import ResponseCandidateResultOrDict +from .common import ResponseRecallResult +from .common import ResponseRecallResultDict +from .common import ResponseRecallResultOrDict from .common import RestoreVersionConfig from .common import RestoreVersionConfigDict from .common import RestoreVersionConfigOrDict from .common import RestoreVersionOperation from .common import RestoreVersionOperationDict from .common import RestoreVersionOperationOrDict -from .common import RetrieveAgentEngineMemoriesConfig -from .common import RetrieveAgentEngineMemoriesConfigDict -from .common import RetrieveAgentEngineMemoriesConfigOrDict from .common import RetrieveContextsConfig from .common import RetrieveContextsConfigDict from .common import RetrieveContextsConfigOrDict @@ -1536,15 +1631,18 @@ from .common import RetrieveProfilesResponse from .common import RetrieveProfilesResponseDict from .common import RetrieveProfilesResponseOrDict +from .common import RetrieveRuntimeMemoriesConfig +from .common import RetrieveRuntimeMemoriesConfigDict +from .common import RetrieveRuntimeMemoriesConfigOrDict from .common import RetrieveSkillsConfig from .common import RetrieveSkillsConfigDict from .common import RetrieveSkillsConfigOrDict from .common import RetrieveSkillsResponse from .common import RetrieveSkillsResponseDict from .common import RetrieveSkillsResponseOrDict -from .common import RollbackAgentEngineMemoryConfig -from .common import RollbackAgentEngineMemoryConfigDict -from .common import RollbackAgentEngineMemoryConfigOrDict +from .common import RollbackRuntimeMemoryConfig +from .common import RollbackRuntimeMemoryConfigDict +from .common import RollbackRuntimeMemoryConfigOrDict from .common import RougeInput from .common import RougeInputDict from .common import RougeInputOrDict @@ -1583,18 +1681,54 @@ 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 RuntimeConfig +from .common import RuntimeConfigDict +from .common import RuntimeConfigOrDict +from .common import RuntimeDict from .common import RuntimeFeedbackContextOperation from .common import RuntimeFeedbackContextOperationDict from .common import RuntimeFeedbackContextOperationOrDict from .common import RuntimeFeedbackEntryOperation from .common import RuntimeFeedbackEntryOperationDict from .common import RuntimeFeedbackEntryOperationOrDict +from .common import RuntimeGenerateMemoriesOperation +from .common import RuntimeGenerateMemoriesOperationDict +from .common import RuntimeGenerateMemoriesOperationOrDict +from .common import RuntimeMemoryConfig +from .common import RuntimeMemoryConfigDict +from .common import RuntimeMemoryConfigOrDict +from .common import RuntimeMemoryOperation +from .common import RuntimeMemoryOperationDict +from .common import RuntimeMemoryOperationOrDict +from .common import RuntimeOperation +from .common import RuntimeOperationDict +from .common import RuntimeOperationOrDict +from .common import RuntimeOrDict +from .common import RuntimePurgeMemoriesOperation +from .common import RuntimePurgeMemoriesOperationDict +from .common import RuntimePurgeMemoriesOperationOrDict +from .common import RuntimeRevision +from .common import RuntimeRevisionDict +from .common import RuntimeRevisionOrDict +from .common import RuntimeRollbackMemoryOperation +from .common import RuntimeRollbackMemoryOperationDict +from .common import RuntimeRollbackMemoryOperationOrDict +from .common import RuntimeSandboxOperation +from .common import RuntimeSandboxOperationDict +from .common import RuntimeSandboxOperationOrDict +from .common import RuntimeSandboxSnapshotOperation +from .common import RuntimeSandboxSnapshotOperationDict +from .common import RuntimeSandboxSnapshotOperationOrDict +from .common import RuntimeSessionOperation +from .common import RuntimeSessionOperationDict +from .common import RuntimeSessionOperationOrDict from .common import SamplingConfig from .common import SamplingConfigDict from .common import SamplingConfigOrDict @@ -1761,6 +1895,9 @@ from .common import SlackSourceSlackChannelsSlackChannelOrDict from .common import State from .common import Strategy +from .common import StringMatchResult +from .common import StringMatchResultDict +from .common import StringMatchResultOrDict from .common import StructuredMemoryConfig from .common import StructuredMemoryConfigDict from .common import StructuredMemoryConfigOrDict @@ -1770,39 +1907,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 @@ -1863,6 +1967,7 @@ from .common import ToolParameterKVMatchSpec from .common import ToolParameterKVMatchSpecDict from .common import ToolParameterKVMatchSpecOrDict +from .common import TrafficType from .common import TuningResourceUsageAssessmentConfig from .common import TuningResourceUsageAssessmentConfigDict from .common import TuningResourceUsageAssessmentConfigOrDict @@ -1879,15 +1984,7 @@ from .common import UnifiedMetric from .common import UnifiedMetricDict from .common import UnifiedMetricOrDict -from .common import UpdateAgentEngineConfig -from .common import UpdateAgentEngineConfigDict -from .common import UpdateAgentEngineConfigOrDict -from .common import UpdateAgentEngineMemoryConfig -from .common import UpdateAgentEngineMemoryConfigDict -from .common import UpdateAgentEngineMemoryConfigOrDict -from .common import UpdateAgentEngineSessionConfig -from .common import UpdateAgentEngineSessionConfigDict -from .common import UpdateAgentEngineSessionConfigOrDict +from .common import UnknownR from .common import UpdatePromptConfig from .common import UpdatePromptConfigDict from .common import UpdatePromptConfigOrDict @@ -1903,12 +2000,21 @@ from .common import UpdateRagCorpusOperation from .common import UpdateRagCorpusOperationDict from .common import UpdateRagCorpusOperationOrDict +from .common import UpdateRuntimeConfig +from .common import UpdateRuntimeConfigDict +from .common import UpdateRuntimeConfigOrDict from .common import UpdateRuntimeFeedbackContextConfig from .common import UpdateRuntimeFeedbackContextConfigDict from .common import UpdateRuntimeFeedbackContextConfigOrDict from .common import UpdateRuntimeFeedbackEntryConfig from .common import UpdateRuntimeFeedbackEntryConfigDict from .common import UpdateRuntimeFeedbackEntryConfigOrDict +from .common import UpdateRuntimeMemoryConfig +from .common import UpdateRuntimeMemoryConfigDict +from .common import UpdateRuntimeMemoryConfigOrDict +from .common import UpdateRuntimeSessionConfig +from .common import UpdateRuntimeSessionConfigDict +from .common import UpdateRuntimeSessionConfigOrDict from .common import UpdateSkillConfig from .common import UpdateSkillConfigDict from .common import UpdateSkillConfigOrDict @@ -1934,74 +2040,17 @@ from .common import WinRateStats from .common import WinRateStatsDict from .common import WinRateStatsOrDict +from .common import WordCountMatchMetricValue +from .common import WordCountMatchMetricValueDict +from .common import WordCountMatchMetricValueOrDict +from .common import WordCountMatchResults +from .common import WordCountMatchResultsDict +from .common import WordCountMatchResultsOrDict from .common import WorkerPoolSpec from .common import WorkerPoolSpecDict from .common import WorkerPoolSpecOrDict __all__ = [ - "DeleteAgentEngineTaskConfig", - "DeleteAgentEngineTaskConfigDict", - "DeleteAgentEngineTaskConfigOrDict", - "GetAgentEngineTaskConfig", - "GetAgentEngineTaskConfigDict", - "GetAgentEngineTaskConfigOrDict", - "TaskArtifact", - "TaskArtifactDict", - "TaskArtifactOrDict", - "TaskOutput", - "TaskOutputDict", - "TaskOutputOrDict", - "TaskMessage", - "TaskMessageDict", - "TaskMessageOrDict", - "TaskStatusDetails", - "TaskStatusDetailsDict", - "TaskStatusDetailsOrDict", - "A2aTask", - "A2aTaskDict", - "A2aTaskOrDict", - "ListAgentEngineTasksConfig", - "ListAgentEngineTasksConfigDict", - "ListAgentEngineTasksConfigOrDict", - "ListAgentEngineTasksResponse", - "ListAgentEngineTasksResponseDict", - "ListAgentEngineTasksResponseOrDict", - "CreateAgentEngineTaskConfig", - "CreateAgentEngineTaskConfigDict", - "CreateAgentEngineTaskConfigOrDict", - "TaskMetadataChange", - "TaskMetadataChangeDict", - "TaskMetadataChangeOrDict", - "TaskArtifactChange", - "TaskArtifactChangeDict", - "TaskArtifactChangeOrDict", - "TaskOutputChange", - "TaskOutputChangeDict", - "TaskOutputChangeOrDict", - "TaskStateChange", - "TaskStateChangeDict", - "TaskStateChangeOrDict", - "TaskStatusDetailsChange", - "TaskStatusDetailsChangeDict", - "TaskStatusDetailsChangeOrDict", - "TaskEventData", - "TaskEventDataDict", - "TaskEventDataOrDict", - "TaskEvent", - "TaskEventDict", - "TaskEventOrDict", - "AppendAgentEngineTaskEventConfig", - "AppendAgentEngineTaskEventConfigDict", - "AppendAgentEngineTaskEventConfigOrDict", - "AppendAgentEngineTaskEventResponse", - "AppendAgentEngineTaskEventResponseDict", - "AppendAgentEngineTaskEventResponseOrDict", - "ListAgentEngineTaskEventsConfig", - "ListAgentEngineTaskEventsConfigDict", - "ListAgentEngineTaskEventsConfigOrDict", - "ListAgentEngineTaskEventsResponse", - "ListAgentEngineTaskEventsResponseDict", - "ListAgentEngineTaskEventsResponseOrDict", "CreateEvaluationItemConfig", "CreateEvaluationItemConfigDict", "CreateEvaluationItemConfigOrDict", @@ -2299,6 +2348,48 @@ "ToolParameterKVMatchResults", "ToolParameterKVMatchResultsDict", "ToolParameterKVMatchResultsOrDict", + "PrefixMatchMetricValue", + "PrefixMatchMetricValueDict", + "PrefixMatchMetricValueOrDict", + "PrefixMatchResults", + "PrefixMatchResultsDict", + "PrefixMatchResultsOrDict", + "WordCountMatchMetricValue", + "WordCountMatchMetricValueDict", + "WordCountMatchMetricValueOrDict", + "WordCountMatchResults", + "WordCountMatchResultsDict", + "WordCountMatchResultsOrDict", + "FleschKincaidReadabilityMetricValue", + "FleschKincaidReadabilityMetricValueDict", + "FleschKincaidReadabilityMetricValueOrDict", + "FleschKincaidReadabilityResults", + "FleschKincaidReadabilityResultsDict", + "FleschKincaidReadabilityResultsOrDict", + "StringMatchResult", + "StringMatchResultDict", + "StringMatchResultOrDict", + "CodeExecutionResults", + "CodeExecutionResultsDict", + "CodeExecutionResultsOrDict", + "ResponseRecallResult", + "ResponseRecallResultDict", + "ResponseRecallResultOrDict", + "RagContextRecallResult", + "RagContextRecallResultDict", + "RagContextRecallResultOrDict", + "PointwiseSafetyMetricResult", + "PointwiseSafetyMetricResultDict", + "PointwiseSafetyMetricResultOrDict", + "LLMBasedSimilarityMetricResult", + "LLMBasedSimilarityMetricResultDict", + "LLMBasedSimilarityMetricResultOrDict", + "LLMBasedClassificationMetricResult", + "LLMBasedClassificationMetricResultDict", + "LLMBasedClassificationMetricResultOrDict", + "CloudLoggingResult", + "CloudLoggingResultDict", + "CloudLoggingResultOrDict", "EvaluateInstancesResponse", "EvaluateInstancesResponseDict", "EvaluateInstancesResponseOrDict", @@ -2371,6 +2462,9 @@ "ReservationAffinity", "ReservationAffinityDict", "ReservationAffinityOrDict", + "ConfidentialInstanceConfig", + "ConfidentialInstanceConfigDict", + "ConfidentialInstanceConfigOrDict", "MachineSpec", "MachineSpecDict", "MachineSpecOrDict", @@ -2392,21 +2486,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", @@ -2464,6 +2558,12 @@ "ReasoningEngineContextSpecMemoryBankConfig", "ReasoningEngineContextSpecMemoryBankConfigDict", "ReasoningEngineContextSpecMemoryBankConfigOrDict", + "ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfig", + "ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfigDict", + "ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfigOrDict", + "ReasoningEngineContextSpecExampleStoreConfig", + "ReasoningEngineContextSpecExampleStoreConfigDict", + "ReasoningEngineContextSpecExampleStoreConfigOrDict", "ReasoningEngineContextSpec", "ReasoningEngineContextSpecDict", "ReasoningEngineContextSpecOrDict", @@ -2521,6 +2621,9 @@ "ReasoningEngineSpecContainerSpec", "ReasoningEngineSpecContainerSpecDict", "ReasoningEngineSpecContainerSpecOrDict", + "ReasoningEngineSpecBuildSpec", + "ReasoningEngineSpecBuildSpecDict", + "ReasoningEngineSpecBuildSpecOrDict", "ReasoningEngineSpec", "ReasoningEngineSpecDict", "ReasoningEngineSpecOrDict", @@ -2536,63 +2639,99 @@ "ReasoningEngineTrafficConfig", "ReasoningEngineTrafficConfigDict", "ReasoningEngineTrafficConfigOrDict", + "ExperimentConfig", + "ExperimentConfigDict", + "ExperimentConfigOrDict", + "ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollection", + "ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollectionDict", + "ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollectionOrDict", + "ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatest", + "ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatestDict", + "ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatestOrDict", + "ReasoningEngineRevisionGarbageCollectionStrategy", + "ReasoningEngineRevisionGarbageCollectionStrategyDict", + "ReasoningEngineRevisionGarbageCollectionStrategyOrDict", "ReasoningEngine", "ReasoningEngineDict", "ReasoningEngineOrDict", - "AgentEngineOperation", - "AgentEngineOperationDict", - "AgentEngineOperationOrDict", - "CreateAgentEngineConfig", - "CreateAgentEngineConfigDict", - "CreateAgentEngineConfigOrDict", - "DeleteAgentEngineConfig", - "DeleteAgentEngineConfigDict", - "DeleteAgentEngineConfigOrDict", - "DeleteAgentEngineOperation", - "DeleteAgentEngineOperationDict", - "DeleteAgentEngineOperationOrDict", - "GetAgentEngineConfig", - "GetAgentEngineConfigDict", - "GetAgentEngineConfigOrDict", - "ListAgentEngineConfig", - "ListAgentEngineConfigDict", - "ListAgentEngineConfigOrDict", + "RuntimeOperation", + "RuntimeOperationDict", + "RuntimeOperationOrDict", + "CreateRuntimeConfig", + "CreateRuntimeConfigDict", + "CreateRuntimeConfigOrDict", + "DeleteRuntimeConfig", + "DeleteRuntimeConfigDict", + "DeleteRuntimeConfigOrDict", + "DeleteRuntimeOperation", + "DeleteRuntimeOperationDict", + "DeleteRuntimeOperationOrDict", + "GetRuntimeConfig", + "GetRuntimeConfigDict", + "GetRuntimeConfigOrDict", + "ListRuntimeConfig", + "ListRuntimeConfigDict", + "ListRuntimeConfigOrDict", "ListReasoningEnginesResponse", "ListReasoningEnginesResponseDict", "ListReasoningEnginesResponseOrDict", - "GetAgentEngineOperationConfig", - "GetAgentEngineOperationConfigDict", - "GetAgentEngineOperationConfigOrDict", - "QueryAgentEngineConfig", - "QueryAgentEngineConfigDict", - "QueryAgentEngineConfigOrDict", + "GetRuntimeOperationConfig", + "GetRuntimeOperationConfigDict", + "GetRuntimeOperationConfigOrDict", + "QueryRuntimeConfig", + "QueryRuntimeConfigDict", + "QueryRuntimeConfigOrDict", "QueryReasoningEngineResponse", "QueryReasoningEngineResponseDict", "QueryReasoningEngineResponseOrDict", - "UpdateAgentEngineConfig", - "UpdateAgentEngineConfigDict", - "UpdateAgentEngineConfigOrDict", + "UpdateRuntimeConfig", + "UpdateRuntimeConfigDict", + "UpdateRuntimeConfigOrDict", + "GetRuntimeRevisionConfig", + "GetRuntimeRevisionConfigDict", + "GetRuntimeRevisionConfigOrDict", + "ReasoningEngineRuntimeRevision", + "ReasoningEngineRuntimeRevisionDict", + "ReasoningEngineRuntimeRevisionOrDict", + "ListRuntimeRevisionsConfig", + "ListRuntimeRevisionsConfigDict", + "ListRuntimeRevisionsConfigOrDict", + "ListReasoningEnginesRuntimeRevisionsResponse", + "ListReasoningEnginesRuntimeRevisionsResponseDict", + "ListReasoningEnginesRuntimeRevisionsResponseOrDict", + "DeleteRuntimeRevisionConfig", + "DeleteRuntimeRevisionConfigDict", + "DeleteRuntimeRevisionConfigOrDict", + "DeleteRuntimeRevisionOperation", + "DeleteRuntimeRevisionOperationDict", + "DeleteRuntimeRevisionOperationOrDict", + "GetDeleteRuntimeRevisionOperationConfig", + "GetDeleteRuntimeRevisionOperationConfigDict", + "GetDeleteRuntimeRevisionOperationConfigOrDict", + "QueryRuntimeRevisionConfig", + "QueryRuntimeRevisionConfigDict", + "QueryRuntimeRevisionConfigOrDict", "MemoryMetadataValue", "MemoryMetadataValueDict", "MemoryMetadataValueOrDict", - "AgentEngineMemoryConfig", - "AgentEngineMemoryConfigDict", - "AgentEngineMemoryConfigOrDict", + "RuntimeMemoryConfig", + "RuntimeMemoryConfigDict", + "RuntimeMemoryConfigOrDict", "MemoryStructuredContent", "MemoryStructuredContentDict", "MemoryStructuredContentOrDict", "Memory", "MemoryDict", "MemoryOrDict", - "AgentEngineMemoryOperation", - "AgentEngineMemoryOperationDict", - "AgentEngineMemoryOperationOrDict", - "DeleteAgentEngineMemoryConfig", - "DeleteAgentEngineMemoryConfigDict", - "DeleteAgentEngineMemoryConfigOrDict", - "DeleteAgentEngineMemoryOperation", - "DeleteAgentEngineMemoryOperationDict", - "DeleteAgentEngineMemoryOperationOrDict", + "RuntimeMemoryOperation", + "RuntimeMemoryOperationDict", + "RuntimeMemoryOperationOrDict", + "DeleteRuntimeMemoryConfig", + "DeleteRuntimeMemoryConfigDict", + "DeleteRuntimeMemoryConfigOrDict", + "DeleteRuntimeMemoryOperation", + "DeleteRuntimeMemoryOperationDict", + "DeleteRuntimeMemoryOperationOrDict", "GenerateMemoriesRequestVertexSessionSource", "GenerateMemoriesRequestVertexSessionSourceDict", "GenerateMemoriesRequestVertexSessionSourceOrDict", @@ -2608,21 +2747,21 @@ "GenerateMemoriesRequestDirectMemoriesSource", "GenerateMemoriesRequestDirectMemoriesSourceDict", "GenerateMemoriesRequestDirectMemoriesSourceOrDict", - "GenerateAgentEngineMemoriesConfig", - "GenerateAgentEngineMemoriesConfigDict", - "GenerateAgentEngineMemoriesConfigOrDict", + "GenerateRuntimeMemoriesConfig", + "GenerateRuntimeMemoriesConfigDict", + "GenerateRuntimeMemoriesConfigOrDict", "GenerateMemoriesResponseGeneratedMemory", "GenerateMemoriesResponseGeneratedMemoryDict", "GenerateMemoriesResponseGeneratedMemoryOrDict", "GenerateMemoriesResponse", "GenerateMemoriesResponseDict", "GenerateMemoriesResponseOrDict", - "AgentEngineGenerateMemoriesOperation", - "AgentEngineGenerateMemoriesOperationDict", - "AgentEngineGenerateMemoriesOperationOrDict", - "GetAgentEngineMemoryConfig", - "GetAgentEngineMemoryConfigDict", - "GetAgentEngineMemoryConfigOrDict", + "RuntimeGenerateMemoriesOperation", + "RuntimeGenerateMemoriesOperationDict", + "RuntimeGenerateMemoriesOperationOrDict", + "GetRuntimeMemoryConfig", + "GetRuntimeMemoryConfigDict", + "GetRuntimeMemoryConfigOrDict", "IngestionDirectContentsSourceEvent", "IngestionDirectContentsSourceEventDict", "IngestionDirectContentsSourceEventOrDict", @@ -2635,9 +2774,9 @@ "MemoryBankIngestEventsOperation", "MemoryBankIngestEventsOperationDict", "MemoryBankIngestEventsOperationOrDict", - "ListAgentEngineMemoryConfig", - "ListAgentEngineMemoryConfigDict", - "ListAgentEngineMemoryConfigOrDict", + "ListRuntimeMemoryConfig", + "ListRuntimeMemoryConfigDict", + "ListRuntimeMemoryConfigOrDict", "ListReasoningEnginesMemoriesResponse", "ListReasoningEnginesMemoriesResponseDict", "ListReasoningEnginesMemoriesResponseOrDict", @@ -2653,9 +2792,9 @@ "MemoryConjunctionFilter", "MemoryConjunctionFilterDict", "MemoryConjunctionFilterOrDict", - "RetrieveAgentEngineMemoriesConfig", - "RetrieveAgentEngineMemoriesConfigDict", - "RetrieveAgentEngineMemoriesConfigOrDict", + "RetrieveRuntimeMemoriesConfig", + "RetrieveRuntimeMemoriesConfigDict", + "RetrieveRuntimeMemoriesConfigOrDict", "RetrieveMemoriesResponseRetrievedMemory", "RetrieveMemoriesResponseRetrievedMemoryDict", "RetrieveMemoriesResponseRetrievedMemoryOrDict", @@ -2671,39 +2810,39 @@ "RetrieveProfilesResponse", "RetrieveProfilesResponseDict", "RetrieveProfilesResponseOrDict", - "RollbackAgentEngineMemoryConfig", - "RollbackAgentEngineMemoryConfigDict", - "RollbackAgentEngineMemoryConfigOrDict", - "AgentEngineRollbackMemoryOperation", - "AgentEngineRollbackMemoryOperationDict", - "AgentEngineRollbackMemoryOperationOrDict", - "UpdateAgentEngineMemoryConfig", - "UpdateAgentEngineMemoryConfigDict", - "UpdateAgentEngineMemoryConfigOrDict", - "PurgeAgentEngineMemoriesConfig", - "PurgeAgentEngineMemoriesConfigDict", - "PurgeAgentEngineMemoriesConfigOrDict", + "RollbackRuntimeMemoryConfig", + "RollbackRuntimeMemoryConfigDict", + "RollbackRuntimeMemoryConfigOrDict", + "RuntimeRollbackMemoryOperation", + "RuntimeRollbackMemoryOperationDict", + "RuntimeRollbackMemoryOperationOrDict", + "UpdateRuntimeMemoryConfig", + "UpdateRuntimeMemoryConfigDict", + "UpdateRuntimeMemoryConfigOrDict", + "PurgeRuntimeMemoriesConfig", + "PurgeRuntimeMemoriesConfigDict", + "PurgeRuntimeMemoriesConfigOrDict", "PurgeMemoriesResponse", "PurgeMemoriesResponseDict", "PurgeMemoriesResponseOrDict", - "AgentEnginePurgeMemoriesOperation", - "AgentEnginePurgeMemoriesOperationDict", - "AgentEnginePurgeMemoriesOperationOrDict", - "GetAgentEngineMemoryRevisionConfig", - "GetAgentEngineMemoryRevisionConfigDict", - "GetAgentEngineMemoryRevisionConfigOrDict", + "RuntimePurgeMemoriesOperation", + "RuntimePurgeMemoriesOperationDict", + "RuntimePurgeMemoriesOperationOrDict", + "GetRuntimeMemoryRevisionConfig", + "GetRuntimeMemoryRevisionConfigDict", + "GetRuntimeMemoryRevisionConfigOrDict", "IntermediateExtractedMemory", "IntermediateExtractedMemoryDict", "IntermediateExtractedMemoryOrDict", "MemoryRevision", "MemoryRevisionDict", "MemoryRevisionOrDict", - "ListAgentEngineMemoryRevisionsConfig", - "ListAgentEngineMemoryRevisionsConfigDict", - "ListAgentEngineMemoryRevisionsConfigOrDict", - "ListAgentEngineMemoryRevisionsResponse", - "ListAgentEngineMemoryRevisionsResponseDict", - "ListAgentEngineMemoryRevisionsResponseOrDict", + "ListRuntimeMemoryRevisionsConfig", + "ListRuntimeMemoryRevisionsConfigDict", + "ListRuntimeMemoryRevisionsConfigOrDict", + "ListRuntimeMemoryRevisionsResponse", + "ListRuntimeMemoryRevisionsResponseDict", + "ListRuntimeMemoryRevisionsResponseOrDict", "AskContextsConfig", "AskContextsConfigDict", "AskContextsConfigOrDict", @@ -2962,33 +3101,45 @@ "UploadRagFileConfig", "UploadRagFileConfigDict", "UploadRagFileConfigOrDict", + "GdataObjectId", + "GdataObjectIdDict", + "GdataObjectIdOrDict", + "GdataBlobstore2Info", + "GdataBlobstore2InfoDict", + "GdataBlobstore2InfoOrDict", + "GdataCompositeMedia", + "GdataCompositeMediaDict", + "GdataCompositeMediaOrDict", + "GdataDiffVersionResponse", + "GdataDiffVersionResponseDict", + "GdataDiffVersionResponseOrDict", + "GdataDiffChecksumsResponse", + "GdataDiffChecksumsResponseDict", + "GdataDiffChecksumsResponseOrDict", + "GdataDiffDownloadResponse", + "GdataDiffDownloadResponseDict", + "GdataDiffDownloadResponseOrDict", + "GdataDiffUploadRequest", + "GdataDiffUploadRequestDict", + "GdataDiffUploadRequestOrDict", + "GdataDiffUploadResponse", + "GdataDiffUploadResponseDict", + "GdataDiffUploadResponseOrDict", + "GdataContentTypeInfo", + "GdataContentTypeInfoDict", + "GdataContentTypeInfoOrDict", + "GdataDownloadParameters", + "GdataDownloadParametersDict", + "GdataDownloadParametersOrDict", + "GdataMedia", + "GdataMediaDict", + "GdataMediaOrDict", + "ApiservingMediaResponseInfo", + "ApiservingMediaResponseInfoDict", + "ApiservingMediaResponseInfoOrDict", "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", @@ -2998,45 +3149,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", @@ -3079,15 +3230,15 @@ "ListSandboxEnvironmentTemplatesResponse", "ListSandboxEnvironmentTemplatesResponseDict", "ListSandboxEnvironmentTemplatesResponseOrDict", - "CreateAgentEngineSandboxSnapshotConfig", - "CreateAgentEngineSandboxSnapshotConfigDict", - "CreateAgentEngineSandboxSnapshotConfigOrDict", + "CreateRuntimeSandboxSnapshotConfig", + "CreateRuntimeSandboxSnapshotConfigDict", + "CreateRuntimeSandboxSnapshotConfigOrDict", "SandboxEnvironmentSnapshot", "SandboxEnvironmentSnapshotDict", "SandboxEnvironmentSnapshotOrDict", - "AgentEngineSandboxSnapshotOperation", - "AgentEngineSandboxSnapshotOperationDict", - "AgentEngineSandboxSnapshotOperationOrDict", + "RuntimeSandboxSnapshotOperation", + "RuntimeSandboxSnapshotOperationDict", + "RuntimeSandboxSnapshotOperationOrDict", "DeleteSandboxEnvironmentSnapshotConfig", "DeleteSandboxEnvironmentSnapshotConfigDict", "DeleteSandboxEnvironmentSnapshotConfigOrDict", @@ -3103,54 +3254,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", @@ -3403,18 +3554,33 @@ "Probe", "ProbeDict", "ProbeOrDict", + "AcceleratorRequirement", + "AcceleratorRequirementDict", + "AcceleratorRequirementOrDict", "ModelContainerSpec", "ModelContainerSpecDict", "ModelContainerSpecOrDict", "AutoscalingMetricSpec", "AutoscalingMetricSpecDict", "AutoscalingMetricSpecOrDict", + "NodeRecycling", + "NodeRecyclingDict", + "NodeRecyclingOrDict", "FlexStart", "FlexStartDict", "FlexStartOrDict", "DedicatedResourcesScaleToZeroSpec", "DedicatedResourcesScaleToZeroSpecDict", "DedicatedResourcesScaleToZeroSpecOrDict", + "ProvisioningPriority", + "ProvisioningPriorityDict", + "ProvisioningPriorityOrDict", + "ProvisioningModelPriority", + "ProvisioningModelPriorityDict", + "ProvisioningModelPriorityOrDict", + "ProvisioningConfig", + "ProvisioningConfigDict", + "ProvisioningConfigOrDict", "DedicatedResources", "DedicatedResourcesDict", "DedicatedResourcesOrDict", @@ -3445,6 +3611,30 @@ "PublisherModelCallToActionViewRestApi", "PublisherModelCallToActionViewRestApiDict", "PublisherModelCallToActionViewRestApiOrDict", + "PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResources", + "PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResourcesDict", + "PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResourcesOrDict", + "PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModel", + "PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModelDict", + "PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModelOrDict", + "PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnit", + "PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitDict", + "PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitOrDict", + "PublisherModelCallToActionDeployMonetizedModelMarketplaceListing", + "PublisherModelCallToActionDeployMonetizedModelMarketplaceListingDict", + "PublisherModelCallToActionDeployMonetizedModelMarketplaceListingOrDict", + "PublisherModelCallToActionDeployMonetizedModelSingleTenant", + "PublisherModelCallToActionDeployMonetizedModelSingleTenantDict", + "PublisherModelCallToActionDeployMonetizedModelSingleTenantOrDict", + "PublisherModelCallToActionDeployMonetizedModelMonetizationConfig", + "PublisherModelCallToActionDeployMonetizedModelMonetizationConfigDict", + "PublisherModelCallToActionDeployMonetizedModelMonetizationConfigOrDict", + "PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfig", + "PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfigDict", + "PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfigOrDict", + "PublisherModelCallToActionDeployMonetizedModel", + "PublisherModelCallToActionDeployMonetizedModelDict", + "PublisherModelCallToActionDeployMonetizedModelOrDict", "PublisherModelCallToAction", "PublisherModelCallToActionDict", "PublisherModelCallToActionOrDict", @@ -3550,21 +3740,6 @@ "EvalRunInferenceConfig", "EvalRunInferenceConfigDict", "EvalRunInferenceConfigOrDict", - "AgentEngine", - "AgentEngineDict", - "AgentEngineOrDict", - "AgentEngineConfig", - "AgentEngineConfigDict", - "AgentEngineConfigOrDict", - "RunQueryJobAgentEngineConfig", - "RunQueryJobAgentEngineConfigDict", - "RunQueryJobAgentEngineConfigOrDict", - "RunQueryJobResult", - "RunQueryJobResultDict", - "RunQueryJobResultOrDict", - "CheckQueryJobResponse", - "CheckQueryJobResponseDict", - "CheckQueryJobResponseOrDict", "AssembleDataset", "AssembleDatasetDict", "AssembleDatasetOrDict", @@ -3604,9 +3779,6 @@ "OptimizeJobConfig", "OptimizeJobConfigDict", "OptimizeJobConfigOrDict", - "AgentEngineRuntimeRevision", - "AgentEngineRuntimeRevisionDict", - "AgentEngineRuntimeRevisionOrDict", "ListDeployableModelsConfig", "ListDeployableModelsConfigDict", "ListDeployableModelsConfigOrDict", @@ -3622,31 +3794,55 @@ "DeployOption", "DeployOptionDict", "DeployOptionOrDict", - "A2aTaskState", - "State", + "Runtime", + "RuntimeDict", + "RuntimeOrDict", + "RuntimeConfig", + "RuntimeConfigDict", + "RuntimeConfigOrDict", + "RunQueryJobRuntimeConfig", + "RunQueryJobRuntimeConfigDict", + "RunQueryJobRuntimeConfigOrDict", + "RunQueryJobResult", + "RunQueryJobResultDict", + "RunQueryJobResultOrDict", + "CheckQueryJobResponse", + "CheckQueryJobResponseDict", + "CheckQueryJobResponseOrDict", + "RuntimeRevision", + "RuntimeRevisionDict", + "RuntimeRevisionOrDict", + "Outcome", "Strategy", "AcceleratorType", "Type", + "ConfidentialInstanceType", "JobState", "ManagedTopicEnum", "IdentityType", "AgentServerMode", + "State", "MemoryType", "Operator", "RagFileType", "ResourceType", - "Language", + "UnknownR", "MachineConfig", + "Language", "SandboxState", "Protocol", "DefaultContainerCategory", "PostSnapshotAction", + "TrafficType", + "MediaModality", "Framework", - "SkillSource", "SkillState", - "LaunchStage", + "SkillSource", "OpenSourceCategory", + "LaunchStage", "VersionState", + "ProvisioningStrategyType", + "ProvisionMode", "QuotaState", "FeedbackType", "EvaluationItemType", @@ -3682,12 +3878,6 @@ "MessageDict", "Importance", "ParsedResponseUnion", - "_DeleteAgentEngineTaskRequestParameters", - "_GetAgentEngineTaskRequestParameters", - "_ListAgentEngineTasksRequestParameters", - "_CreateAgentEngineTaskRequestParameters", - "_AppendAgentEngineTaskEventRequestParameters", - "_ListAgentEngineTaskEventsRequestParameters", "_CreateEvaluationItemParameters", "_CreateEvaluationMetricParameters", "_CreateEvaluationRunParameters", @@ -3705,31 +3895,36 @@ "_OptimizeRequestParameters", "_CustomJobParameters", "_GetCustomJobParameters", - "_CancelQueryJobAgentEngineRequestParameters", - "_CheckQueryJobAgentEngineRequestParameters", - "_RunQueryJobAgentEngineRequestParameters", - "_CreateAgentEngineRequestParameters", - "_DeleteAgentEngineRequestParameters", - "_GetAgentEngineRequestParameters", - "_ListAgentEngineRequestParameters", - "_GetAgentEngineOperationParameters", - "_QueryAgentEngineRequestParameters", - "_UpdateAgentEngineRequestParameters", - "_CreateAgentEngineMemoryRequestParameters", - "_DeleteAgentEngineMemoryRequestParameters", - "_GenerateAgentEngineMemoriesRequestParameters", - "_GetAgentEngineMemoryRequestParameters", + "_CancelQueryJobRuntimeRequestParameters", + "_CheckQueryJobRuntimeRequestParameters", + "_RunQueryJobRuntimeRequestParameters", + "_CreateRuntimeRequestParameters", + "_DeleteRuntimeRequestParameters", + "_GetRuntimeRequestParameters", + "_ListRuntimeRequestParameters", + "_GetRuntimeOperationParameters", + "_QueryRuntimeRequestParameters", + "_UpdateRuntimeRequestParameters", + "_GetRuntimeRevisionRequestParameters", + "_ListRuntimeRevisionsRequestParameters", + "_DeleteRuntimeRevisionRequestParameters", + "_GetDeleteRuntimeRevisionOperationParameters", + "_QueryRuntimeRevisionRequestParameters", + "_CreateRuntimeMemoryRequestParameters", + "_DeleteRuntimeMemoryRequestParameters", + "_GenerateRuntimeMemoriesRequestParameters", + "_GetRuntimeMemoryRequestParameters", "_IngestEventsRequestParameters", - "_ListAgentEngineMemoryRequestParameters", - "_GetAgentEngineMemoryOperationParameters", - "_GetAgentEngineGenerateMemoriesOperationParameters", - "_RetrieveAgentEngineMemoriesRequestParameters", + "_ListRuntimeMemoryRequestParameters", + "_GetRuntimeMemoryOperationParameters", + "_GetRuntimeGenerateMemoriesOperationParameters", + "_RetrieveRuntimeMemoriesRequestParameters", "_RetrieveMemoryProfilesRequestParameters", - "_RollbackAgentEngineMemoryRequestParameters", - "_UpdateAgentEngineMemoryRequestParameters", - "_PurgeAgentEngineMemoriesRequestParameters", - "_GetAgentEngineMemoryRevisionRequestParameters", - "_ListAgentEngineMemoryRevisionsRequestParameters", + "_RollbackRuntimeMemoryRequestParameters", + "_UpdateRuntimeMemoryRequestParameters", + "_PurgeRuntimeMemoriesRequestParameters", + "_GetRuntimeMemoryRevisionRequestParameters", + "_ListRuntimeMemoryRevisionsRequestParameters", "_AskContextsRequestParameters", "_CreateRagCorpusRequestParameters", "_GetCorpusOperationParameters", @@ -3747,17 +3942,12 @@ "_ImportRagFilesRequestParameters", "_GetImportFilesOperationParameters", "_UploadRagFileParameters", - "_GetAgentEngineRuntimeRevisionRequestParameters", - "_ListAgentEngineRuntimeRevisionsRequestParameters", - "_DeleteAgentEngineRuntimeRevisionRequestParameters", - "_GetDeleteAgentEngineRuntimeRevisionOperationParameters", - "_QueryAgentEngineRuntimeRevisionRequestParameters", - "_CreateAgentEngineSandboxRequestParameters", - "_DeleteAgentEngineSandboxRequestParameters", - "_ExecuteCodeAgentEngineSandboxRequestParameters", - "_GetAgentEngineSandboxRequestParameters", - "_ListAgentEngineSandboxesRequestParameters", - "_GetAgentEngineSandboxOperationParameters", + "_CreateRuntimeSandboxRequestParameters", + "_DeleteRuntimeSandboxRequestParameters", + "_ExecuteCodeRuntimeSandboxRequestParameters", + "_GetRuntimeSandboxRequestParameters", + "_ListRuntimeSandboxesRequestParameters", + "_GetRuntimeSandboxOperationParameters", "_CreateSandboxEnvironmentTemplateRequestParameters", "_DeleteSandboxEnvironmentTemplateRequestParameters", "_GetSandboxEnvironmentTemplateRequestParameters", @@ -3767,15 +3957,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 dd57c62b1e..6565bfcd31 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -91,54 +91,17 @@ 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.""" +class Outcome(_common.CaseInSensitiveEnum): + """Outcome of the code execution.""" - STATE_UNSPECIFIED = "STATE_UNSPECIFIED" - """Task state unspecified. Default value if not set.""" - SUBMITTED = "SUBMITTED" - """Task is submitted and waiting to be processed.""" - WORKING = "WORKING" - """Task is actively being processed.""" - COMPLETED = "COMPLETED" - """Task is finished.""" - CANCELLED = "CANCELLED" - """Task is cancelled.""" - FAILED = "FAILED" - """Task has failed.""" - REJECTED = "REJECTED" - """Task is rejected by the system.""" - INPUT_REQUIRED = "INPUT_REQUIRED" - """Task requires input from the user.""" - AUTH_REQUIRED = "AUTH_REQUIRED" - """Task requires auth (e.g. OAuth) from the user.""" - PAUSED = "PAUSED" - """Task is paused.""" - - -class State(_common.CaseInSensitiveEnum): - """The new state of the task.""" - - STATE_UNSPECIFIED = "STATE_UNSPECIFIED" - """Task state unspecified. Default value if not set.""" - SUBMITTED = "SUBMITTED" - """Task is submitted and waiting to be processed.""" - WORKING = "WORKING" - """Task is actively being processed.""" - COMPLETED = "COMPLETED" - """Task is finished.""" - CANCELLED = "CANCELLED" - """Task is cancelled.""" - FAILED = "FAILED" - """Task has failed.""" - REJECTED = "REJECTED" - """Task is rejected by the system.""" - INPUT_REQUIRED = "INPUT_REQUIRED" - """Task requires input from the user.""" - AUTH_REQUIRED = "AUTH_REQUIRED" - """Task requires auth (e.g. OAuth) from the user.""" - PAUSED = "PAUSED" - """Task is paused.""" + OUTCOME_UNSPECIFIED = "OUTCOME_UNSPECIFIED" + """Unspecified status. This value should not be used.""" + OUTCOME_OK = "OUTCOME_OK" + """Code execution completed successfully.""" + OUTCOME_FAILED = "OUTCOME_FAILED" + """Code execution finished but with a failure. `stderr` should contain the reason.""" + OUTCOME_DEADLINE_EXCEEDED = "OUTCOME_DEADLINE_EXCEEDED" + """Code execution ran for too long, and was cancelled. There may or may not be a partial output present.""" class Strategy(_common.CaseInSensitiveEnum): @@ -154,6 +117,8 @@ class Strategy(_common.CaseInSensitiveEnum): """Standard provisioning strategy uses regular on-demand resources.""" SPOT = "SPOT" """Spot provisioning strategy uses spot resources.""" + DWS_FLEX_START = "DWS_FLEX_START" + """Deprecated. DWS Flex start strategy uses DWS to queue for resources.""" FLEX_START = "FLEX_START" """Flex Start strategy uses DWS to queue for resources.""" @@ -212,6 +177,19 @@ class Type(_common.CaseInSensitiveEnum): """Consume any reservation available, falling back to on-demand.""" SPECIFIC_RESERVATION = "SPECIFIC_RESERVATION" """Consume from a specific reservation. When chosen, the reservation must be identified via the `key` and `values` fields.""" + SPECIFIC_THEN_ANY_RESERVATION = "SPECIFIC_THEN_ANY_RESERVATION" + """Prefer to consume from a specific reservation, but still consume any reservation available if the specified reservation is not available or exhausted. Must specify key value fields for specifying the reservations. Hidden until GKE AP supports these types.""" + SPECIFIC_THEN_NO_RESERVATION = "SPECIFIC_THEN_NO_RESERVATION" + """Prefer to consume from a specific reservation, but still consume from the on-demand pool if the specified reservation is exhausted. Must specify key value fields for specifying the reservations. Hidden until GKE AP supports these types.""" + + +class ConfidentialInstanceType(_common.CaseInSensitiveEnum): + """Defines the type of technology used by the confidential instance.""" + + CONFIDENTIAL_INSTANCE_TYPE_UNSPECIFIED = "CONFIDENTIAL_INSTANCE_TYPE_UNSPECIFIED" + """No type specified.""" + SEV = "SEV" + """AMD Secure Encrypted Virtualization.""" class JobState(_common.CaseInSensitiveEnum): @@ -280,6 +258,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.""" @@ -326,15 +315,17 @@ class ResourceType(_common.CaseInSensitiveEnum): """Folder resource type.""" -class Language(_common.CaseInSensitiveEnum): - """The coding language supported in this environment.""" +class UnknownR(_common.CaseInSensitiveEnum): + """Request class to use for all Blobstore operations for this request.""" - LANGUAGE_UNSPECIFIED = "LANGUAGE_UNSPECIFIED" - """The default value. This value is unused.""" - LANGUAGE_PYTHON = "LANGUAGE_PYTHON" - """The coding language is Python.""" - LANGUAGE_JAVASCRIPT = "LANGUAGE_JAVASCRIPT" - """The coding language is JavaScript.""" + UNKNOWN_REQUEST_CLASS = "UNKNOWN_REQUEST_CLASS" + """Unpopulated request_class in log files will be taken as 0 in dremel query. GoogleSQL will try to cast it to enum by default. An unused 0 value is added to avoid GoogleSQL casting error. Please refer to b/69677280.""" + LATENCY_SENSITIVE = "LATENCY_SENSITIVE" + """A latency-sensitive request.""" + PRODUCTION_BATCH = "PRODUCTION_BATCH" + """A request generated by a batch process.""" + BEST_EFFORT = "BEST_EFFORT" + """A best-effort request.""" class MachineConfig(_common.CaseInSensitiveEnum): @@ -346,6 +337,17 @@ class MachineConfig(_common.CaseInSensitiveEnum): """The default value: milligcu 4000, memory 4 Gib""" +class Language(_common.CaseInSensitiveEnum): + """The coding language supported in this environment.""" + + LANGUAGE_UNSPECIFIED = "LANGUAGE_UNSPECIFIED" + """The default value. This value is unused.""" + LANGUAGE_PYTHON = "LANGUAGE_PYTHON" + """The coding language is Python.""" + LANGUAGE_JAVASCRIPT = "LANGUAGE_JAVASCRIPT" + """The coding language is JavaScript.""" + + class SandboxState(_common.CaseInSensitiveEnum): """Output only. The runtime state of the SandboxEnvironment.""" @@ -361,6 +363,12 @@ class SandboxState(_common.CaseInSensitiveEnum): """Sandbox has terminated with underlying runtime failure.""" STATE_DELETED = "STATE_DELETED" """Sandbox runtime has been deleted.""" + STATE_PAUSED = "STATE_PAUSED" + """Sandbox runtime is paused.""" + STATE_PAUSING = "STATE_PAUSING" + """Sandbox runtime is pausing.""" + STATE_RESUMING = "STATE_RESUMING" + """Sandbox runtime is resuming.""" class Protocol(_common.CaseInSensitiveEnum): @@ -394,6 +402,38 @@ class PostSnapshotAction(_common.CaseInSensitiveEnum): """Sandbox environment will be paused after snapshot is taken.""" +class TrafficType(_common.CaseInSensitiveEnum): + """Output only. The traffic type for this request.""" + + TRAFFIC_TYPE_UNSPECIFIED = "TRAFFIC_TYPE_UNSPECIFIED" + """Unspecified request traffic type.""" + ON_DEMAND = "ON_DEMAND" + """Type for Pay-As-You-Go traffic.""" + ON_DEMAND_PRIORITY = "ON_DEMAND_PRIORITY" + """Type for Priority Pay-As-You-Go traffic.""" + ON_DEMAND_FLEX = "ON_DEMAND_FLEX" + """Type for Flex traffic.""" + PROVISIONED_THROUGHPUT = "PROVISIONED_THROUGHPUT" + """Type for Provisioned Throughput traffic.""" + + +class MediaModality(_common.CaseInSensitiveEnum): + """The modality that this token count applies to.""" + + MODALITY_UNSPECIFIED = "MODALITY_UNSPECIFIED" + """When a modality is not specified, it is treated as `TEXT`.""" + TEXT = "TEXT" + """The `Part` contains plain text.""" + IMAGE = "IMAGE" + """The `Part` contains an image.""" + VIDEO = "VIDEO" + """The `Part` contains a video.""" + AUDIO = "AUDIO" + """The `Part` contains audio.""" + DOCUMENT = "DOCUMENT" + """The `Part` contains a document, such as a PDF.""" + + class Framework(_common.CaseInSensitiveEnum): """Framework used to build the application.""" @@ -405,17 +445,6 @@ class Framework(_common.CaseInSensitiveEnum): """Angular framework.""" -class SkillSource(_common.CaseInSensitiveEnum): - """Output only. The source of the Skill.""" - - SKILL_SOURCE_UNSPECIFIED = "SKILL_SOURCE_UNSPECIFIED" - """The skill source is unspecified.""" - USER = "USER" - """The skill is created by a user.""" - SYSTEM = "SYSTEM" - """The skill is a system skill.""" - - class SkillState(_common.CaseInSensitiveEnum): """State of the Skill.""" @@ -431,19 +460,15 @@ class SkillState(_common.CaseInSensitiveEnum): """The Skill is being deleted.""" -class LaunchStage(_common.CaseInSensitiveEnum): - """Indicates the launch stage of the model.""" +class SkillSource(_common.CaseInSensitiveEnum): + """Output only. The source of the Skill.""" - LAUNCH_STAGE_UNSPECIFIED = "LAUNCH_STAGE_UNSPECIFIED" - """The model launch stage is unspecified.""" - EXPERIMENTAL = "EXPERIMENTAL" - """Used to indicate the PublisherModel is at Experimental launch stage, available to a small set of customers.""" - PRIVATE_PREVIEW = "PRIVATE_PREVIEW" - """Used to indicate the PublisherModel is at Private Preview launch stage, only available to a small set of customers, although a larger set of customers than an Experimental launch. Previews are the first launch stage used to get feedback from customers.""" - PUBLIC_PREVIEW = "PUBLIC_PREVIEW" - """Used to indicate the PublisherModel is at Public Preview launch stage, available to all customers, although not supported for production workloads.""" - GA = "GA" - """Used to indicate the PublisherModel is at GA launch stage, available to all customers and ready for production workload.""" + SKILL_SOURCE_UNSPECIFIED = "SKILL_SOURCE_UNSPECIFIED" + """The skill source is unspecified.""" + USER = "USER" + """The skill is created by a user.""" + SYSTEM = "SYSTEM" + """The skill is a system skill.""" class OpenSourceCategory(_common.CaseInSensitiveEnum): @@ -463,6 +488,25 @@ class OpenSourceCategory(_common.CaseInSensitiveEnum): """Used to indicate the PublisherModel is a Google-owned pure open source model.""" THIRD_PARTY_OWNED_OSS = "THIRD_PARTY_OWNED_OSS" """Used to indicate the PublisherModel is a 3p-owned pure open source model.""" + THIRD_PARTY_IP_PROTECTED = "THIRD_PARTY_IP_PROTECTED" + """Used to indicate the PublisherModel is a 3p-owned IP protected model.""" + + +class LaunchStage(_common.CaseInSensitiveEnum): + """Indicates the launch stage of the model.""" + + LAUNCH_STAGE_UNSPECIFIED = "LAUNCH_STAGE_UNSPECIFIED" + """The model launch stage is unspecified.""" + DOGFOOD = "DOGFOOD" + """Used to indicate the PublisherModel is at Google internal testing launch stage.""" + EXPERIMENTAL = "EXPERIMENTAL" + """Used to indicate the PublisherModel is at Experimental launch stage, available to a small set of customers.""" + PRIVATE_PREVIEW = "PRIVATE_PREVIEW" + """Used to indicate the PublisherModel is at Private Preview launch stage, only available to a small set of customers, although a larger set of customers than an Experimental launch. Previews are the first launch stage used to get feedback from customers.""" + PUBLIC_PREVIEW = "PUBLIC_PREVIEW" + """Used to indicate the PublisherModel is at Public Preview launch stage, available to all customers, although not supported for production workloads.""" + GA = "GA" + """Used to indicate the PublisherModel is at GA launch stage, available to all customers and ready for production workload.""" class VersionState(_common.CaseInSensitiveEnum): @@ -476,6 +520,26 @@ class VersionState(_common.CaseInSensitiveEnum): """Used to indicate the version is unstable.""" +class ProvisioningStrategyType(_common.CaseInSensitiveEnum): + """A predefined, simple strategy. If this is set, the provisioning_model_priorities will be ignored.""" + + PROVISIONING_STRATEGY_TYPE_UNSPECIFIED = "PROVISIONING_STRATEGY_TYPE_UNSPECIFIED" + """Unspecified provisioning strategy.""" + CUSTOM_LIST = "CUSTOM_LIST" + """A custom list of priorities, defined in priority_levels.""" + + +class ProvisionMode(_common.CaseInSensitiveEnum): + """The predefined provisioning mode.""" + + PROVISION_MODE_UNSPECIFIED = "PROVISION_MODE_UNSPECIFIED" + """Unspecified provisioning mode.""" + SPECIFIC_RESERVATION = "SPECIFIC_RESERVATION" + """Provision from a specific reservation.""" + STANDARD = "STANDARD" + """Provision from STANDARD capacity.""" + + class QuotaState(_common.CaseInSensitiveEnum): """Output only. The user accelerator quota state.""" @@ -633,7777 +697,7493 @@ 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 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 DeleteAgentEngineTaskConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Task.""" +class CreateEvaluationItemConfigDict(TypedDict, total=False): + """Config to create an evaluation item.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteAgentEngineTaskConfigOrDict = Union[ - DeleteAgentEngineTaskConfig, DeleteAgentEngineTaskConfigDict +CreateEvaluationItemConfigOrDict = Union[ + CreateEvaluationItemConfig, CreateEvaluationItemConfigDict ] -class _DeleteAgentEngineTaskRequestParameters(_common.BaseModel): - """Parameters for deleting an agent engine task.""" +class _CreateEvaluationItemParameters(_common.BaseModel): + """Represents a job that creates an evaluation item.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine task.""" - ) - config: Optional[DeleteAgentEngineTaskConfig] = Field( + 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 _DeleteAgentEngineTaskRequestParametersDict(TypedDict, total=False): - """Parameters for deleting an agent engine task.""" +class _CreateEvaluationItemParametersDict(TypedDict, total=False): + """Represents a job that creates an evaluation item.""" + + evaluation_item_type: Optional[str] + """""" + + gcs_uri: Optional[str] + """""" - name: Optional[str] - """Name of the agent engine task.""" + display_name: Optional[str] + """""" - config: Optional[DeleteAgentEngineTaskConfigDict] + config: Optional[CreateEvaluationItemConfigDict] """""" -_DeleteAgentEngineTaskRequestParametersOrDict = Union[ - _DeleteAgentEngineTaskRequestParameters, _DeleteAgentEngineTaskRequestParametersDict +_CreateEvaluationItemParametersOrDict = Union[ + _CreateEvaluationItemParameters, _CreateEvaluationItemParametersDict ] -class GetAgentEngineTaskConfig(_common.BaseModel): - """Config for getting an Agent Engine Task.""" +class PromptTemplateData(_common.BaseModel): + """Holds data for a prompt template. - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + Message to hold a prompt template and the values to populate the template. + """ + + values: Optional[dict[str, genai_types.Content]] = Field( + default=None, description="""The values for fields in the prompt template.""" ) -class GetAgentEngineTaskConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Task.""" +class PromptTemplateDataDict(TypedDict, total=False): + """Holds data for a prompt template. + + Message to hold a prompt template and the values to populate the template. + """ - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + values: Optional[dict[str, genai_types.Content]] + """The values for fields in the prompt template.""" -GetAgentEngineTaskConfigOrDict = Union[ - GetAgentEngineTaskConfig, GetAgentEngineTaskConfigDict -] +PromptTemplateDataOrDict = Union[PromptTemplateData, PromptTemplateDataDict] -class _GetAgentEngineTaskRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine task.""" +class EvaluationPrompt(_common.BaseModel): + """Represents the prompt to be evaluated.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine task.""" + text: Optional[str] = Field(default=None, description="""Text prompt.""") + value: Optional[dict[str, Any]] = Field( + default=None, + description="""Fields and values that can be used to populate the prompt template.""", ) - config: Optional[GetAgentEngineTaskConfig] = Field(default=None, description="""""") + prompt_template_data: Optional[PromptTemplateData] = Field( + default=None, description="""Prompt template data.""" + ) + user_scenario: Optional[evals_types.UserScenario] = Field( + default=None, + description="""User scenario to help simulate multi-turn agent run results.""", + ) + +class EvaluationPromptDict(TypedDict, total=False): + """Represents the prompt to be evaluated.""" -class _GetAgentEngineTaskRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine task.""" + text: Optional[str] + """Text prompt.""" - name: Optional[str] - """Name of the agent engine task.""" + value: Optional[dict[str, Any]] + """Fields and values that can be used to populate the prompt template.""" - config: Optional[GetAgentEngineTaskConfigDict] - """""" + prompt_template_data: Optional[PromptTemplateDataDict] + """Prompt template data.""" + + user_scenario: Optional[evals_types.UserScenario] + """User scenario to help simulate multi-turn agent run results.""" -_GetAgentEngineTaskRequestParametersOrDict = Union[ - _GetAgentEngineTaskRequestParameters, _GetAgentEngineTaskRequestParametersDict -] +EvaluationPromptOrDict = Union[EvaluationPrompt, EvaluationPromptDict] -class TaskArtifact(_common.BaseModel): - """The artifact of the task event.""" +class CandidateResponse(_common.BaseModel): + """Responses from model or agent.""" - artifact_id: Optional[str] = Field( + candidate: 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="""The name of the candidate that produced the response.""", ) - description: Optional[str] = Field( + text: Optional[str] = Field(default=None, description="""The text response.""") + value: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. A human readable description of the artifact.""", + description="""Fields and values that can be used to populate the response template.""", ) - display_name: Optional[str] = Field( + events: Optional[list[genai_types.Content]] = Field( default=None, - description="""Optional. The human-readable name of the artifact provided by the creator.""", + description="""Intermediate events (such as tool calls and responses) that led to the final response.""", ) - metadata: Optional[dict[str, Any]] = Field( + agent_data: Optional[evals_types.AgentData] = 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.""", - ) - parts: Optional[list[genai_types.Part]] = Field( - default=None, description="""The parts of 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 TaskArtifactDict(TypedDict, total=False): - """The artifact of the task event.""" +class CandidateResponseDict(TypedDict, total=False): + """Responses from model or agent.""" - artifact_id: Optional[str] - """Required. The unique identifier of the artifact within the task. This id is provided by the creator of the artifact.""" + 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 provided by the creator.""" + value: Optional[dict[str, Any]] + """Fields and values that can be used to populate the response template.""" - 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.""" + events: Optional[list[genai_types.Content]] + """Intermediate events (such as tool calls and responses) that led to the final response.""" - parts: Optional[list[genai_types.Part]] - """The parts of the artifact.""" + 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.""" -TaskArtifactOrDict = Union[TaskArtifact, TaskArtifactDict] +CandidateResponseOrDict = Union[CandidateResponse, CandidateResponseDict] -class TaskOutput(_common.BaseModel): - """The output of the task event.""" +class EvaluationItemRequest(_common.BaseModel): + """Single evaluation request.""" - artifacts: Optional[list[TaskArtifact]] = Field( - default=None, description="""The artifacts of the task event.""" + prompt: Optional[EvaluationPrompt] = Field( + default=None, description="""The request/prompt to evaluate.""" ) - - -class TaskOutputDict(TypedDict, total=False): - """The output of the task event.""" - - artifacts: Optional[list[TaskArtifactDict]] - """The artifacts of the task event.""" - - -TaskOutputOrDict = Union[TaskOutput, TaskOutputDict] - - -class TaskMessage(_common.BaseModel): - """The message of the task event.""" - - message_id: Optional[str] = Field( - default=None, description="""Required. The unique identifier of the message.""" + golden_response: Optional[CandidateResponse] = Field( + default=None, description="""The ideal response or ground truth.""" ) - metadata: Optional[dict[str, Any]] = Field( + rubrics: Optional[dict[str, "RubricGroup"]] = 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.""" + description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", ) - role: Optional[str] = Field( + candidate_responses: Optional[list[CandidateResponse]] = Field( default=None, - description="""Required. The role of the sender of the message. e.g. "user", "agent".""", + description="""Responses from model under test and other baseline models for comparison.""", ) -class TaskMessageDict(TypedDict, total=False): - """The message of the task event.""" +class EvaluationItemRequestDict(TypedDict, total=False): + """Single evaluation request.""" - message_id: Optional[str] - """Required. The unique identifier of the message.""" + prompt: Optional[EvaluationPromptDict] + """The request/prompt to evaluate.""" - metadata: Optional[dict[str, Any]] - """Optional. A2A message may have extension_uris or reference_task_ids. They will be stored under metadata.""" + golden_response: Optional[CandidateResponseDict] + """The ideal response or ground truth.""" - parts: Optional[list[genai_types.Part]] - """The parts of the message.""" + rubrics: Optional[dict[str, "RubricGroupDict"]] + """Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""" - role: Optional[str] - """Required. The role of the sender of the message. e.g. "user", "agent".""" + candidate_responses: Optional[list[CandidateResponseDict]] + """Responses from model under test and other baseline models for comparison.""" -TaskMessageOrDict = Union[TaskMessage, TaskMessageDict] +EvaluationItemRequestOrDict = Union[EvaluationItemRequest, EvaluationItemRequestDict] -class TaskStatusDetails(_common.BaseModel): - """The status details of the task event.""" +class EvaluationItemResult(_common.BaseModel): + """Represents the result of an evaluation item.""" - task_message: Optional[TaskMessage] = Field( - default=None, description="""The status of the task event.""" + evaluation_request: Optional[str] = Field( + default=None, description="""The request item that was evaluated.""" ) - - -class TaskStatusDetailsDict(TypedDict, total=False): - """The status details of the task event.""" - - task_message: Optional[TaskMessageDict] - """The status of the task event.""" - - -TaskStatusDetailsOrDict = Union[TaskStatusDetails, TaskStatusDetailsDict] - - -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}`""", - ) - 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( + evaluation_run: Optional[str] = Field( default=None, - description="""Output only. The state of the task. The state of a new task is SUBMITTED by default. The state of a task can only be updated via AppendA2aTaskEvents API.""", + description="""The evaluation run that was used to generate the result.""", ) - status_details: Optional[TaskStatusDetails] = Field( - default=None, description="""Optional. The status details of the task.""" + request: Optional[EvaluationItemRequest] = Field( + default=None, description="""The request that was evaluated.""" ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. The last update timestamp of the task.""", + metric: Optional[str] = Field( + default=None, description="""The metric that was evaluated.""" ) - 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""", + candidate_results: Optional[list[evals_types.CandidateResult]] = Field( + default=None, description="""The results for the metric.""" ) - 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]""", + 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}`""" +class EvaluationItemResultDict(TypedDict, total=False): + """Represents the result of an evaluation item.""" - 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.""" + evaluation_request: Optional[str] + """The request item that was evaluated.""" - output: Optional[TaskOutputDict] - """Optional. The final output of the task.""" + evaluation_run: Optional[str] + """The evaluation run that was used to generate the result.""" - 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.""" + request: Optional[EvaluationItemRequestDict] + """The request that was evaluated.""" - status_details: Optional[TaskStatusDetailsDict] - """Optional. The status details of the task.""" + metric: Optional[str] + """The metric that was evaluated.""" - update_time: Optional[datetime.datetime] - """Output only. The last update timestamp of the task.""" + candidate_results: Optional[list[evals_types.CandidateResult]] + """The results for the metric.""" - 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""" + metadata: Optional[dict[str, Any]] + """Metadata about the evaluation result.""" - 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]""" +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( - 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`.""", + metadata: Optional[dict[str, Any]] = Field( + default=None, description="""Metadata for the EvaluationItem.""" ) - - -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] - """""" - - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" - - order_by: Optional[str] - """A comma-separated list of fields to order by, sorted in ascending order. - Use "desc" after a field name for descending. - If this field is omitted, the default ordering is `create_time` descending. - More detail in [AIP-132](https://google.aip.dev/132). - - Supported fields: - * `create_time` - * `update_time` - - Example: `create_time desc`.""" - - -ListAgentEngineTasksConfigOrDict = Union[ - ListAgentEngineTasksConfig, ListAgentEngineTasksConfigDict -] - - -class _ListAgentEngineTasksRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" - - name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + labels: Optional[dict[str, str]] = Field( + default=None, description="""Labels for the EvaluationItem.""" ) - config: Optional[ListAgentEngineTasksConfig] = Field( - default=None, description="""""" + evaluation_item_type: Optional[EvaluationItemType] = Field( + default=None, description="""The type of the EvaluationItem.""" ) - - -class _ListAgentEngineTasksRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" - - name: Optional[str] - """Name of the agent engine.""" - - config: Optional[ListAgentEngineTasksConfigDict] - """""" - - -_ListAgentEngineTasksRequestParametersOrDict = Union[ - _ListAgentEngineTasksRequestParameters, _ListAgentEngineTasksRequestParametersDict -] - - -class ListAgentEngineTasksResponse(_common.BaseModel): - """Response for listing agent engine tasks.""" - - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + evaluation_request: Optional[EvaluationItemRequest] = Field( + default=None, description="""The request to evaluate.""" ) - next_page_token: Optional[str] = Field(default=None, description="""""") - a2aTasks: Optional[list[A2aTask]] = Field( - default=None, description="""List of agent engine tasks.""" + evaluation_response: Optional[EvaluationItemResult] = Field( + default=None, description="""The response from evaluation.""" + ) + gcs_uri: Optional[str] = Field( + default=None, + description="""The Cloud Storage object where the request or response is stored.""", + ) + create_time: Optional[datetime.datetime] = Field( + default=None, description="""Timestamp when this item was created.""" + ) + 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 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] - """""" + response = _camel_key_to_snake(response) + result = super()._from_response(response=response, kwargs=kwargs) + return result - a2aTasks: Optional[list[A2aTaskDict]] - """List of agent engine tasks.""" +class EvaluationItemDict(TypedDict, total=False): + """EvaluationItem is a single evaluation request or result. -ListAgentEngineTasksResponseOrDict = Union[ - ListAgentEngineTasksResponse, ListAgentEngineTasksResponseDict -] + The content of an EvaluationItem is immutable - it cannot be updated once + created. EvaluationItems can be deleted when no longer needed. + """ + name: Optional[str] + """The resource name of the EvaluationItem.""" -class CreateAgentEngineTaskConfig(_common.BaseModel): - """Config for creating a Session.""" + display_name: Optional[str] + """The display name of the EvaluationItem.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - context_id: Optional[str] = Field( - default=None, description="""The context id of the task to create.""" - ) - metadata: Optional[dict[str, Any]] = Field( - default=None, description="""The metadata of the task to create.""" - ) - status_details: Optional[TaskStatusDetails] = Field( - default=None, description="""The status details of the task to create.""" - ) - output: Optional[TaskOutput] = Field( - default=None, description="""The output of the task to create.""" - ) + metadata: Optional[dict[str, Any]] + """Metadata for the EvaluationItem.""" + labels: Optional[dict[str, str]] + """Labels for the EvaluationItem.""" -class CreateAgentEngineTaskConfigDict(TypedDict, total=False): - """Config for creating a Session.""" + evaluation_item_type: Optional[EvaluationItemType] + """The type of the EvaluationItem.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + evaluation_request: Optional[EvaluationItemRequestDict] + """The request to evaluate.""" - context_id: Optional[str] - """The context id of the task to create.""" + evaluation_response: Optional[EvaluationItemResultDict] + """The response from evaluation.""" - metadata: Optional[dict[str, Any]] - """The metadata of the task to create.""" + gcs_uri: Optional[str] + """The Cloud Storage object where the request or response is stored.""" - status_details: Optional[TaskStatusDetailsDict] - """The status details of the task to create.""" + create_time: Optional[datetime.datetime] + """Timestamp when this item was created.""" - output: Optional[TaskOutputDict] - """The output of the task to create.""" + error: Optional[genai_types.GoogleRpcStatus] + """Error for the evaluation item.""" -CreateAgentEngineTaskConfigOrDict = Union[ - CreateAgentEngineTaskConfig, CreateAgentEngineTaskConfigDict -] +EvaluationItemOrDict = Union[EvaluationItem, EvaluationItemDict] -class _CreateAgentEngineTaskRequestParameters(_common.BaseModel): - """Parameters for creating Agent Engine Tasks.""" +class Metric(_common.BaseModel): + """The metric used for evaluation.""" - name: Optional[str] = Field( + name: Optional[str] = Field(default=None, description="""The name of the metric.""") + custom_function: Optional[Union[str, Callable[..., Any]]] = Field( default=None, - description="""Name of the agent engine to create the task under.""", + description="""The custom function that defines the end-to-end logic for metric computation.""", ) - a2a_task_id: Optional[str] = Field( - default=None, description="""The ID of the task.""" + prompt_template: Optional[str] = Field( + default=None, description="""The prompt template for the metric.""" ) - config: Optional[CreateAgentEngineTaskConfig] = Field( - default=None, description="""""" + judge_model_system_instruction: Optional[str] = Field( + default=None, description="""The system instruction for the judge model.""" ) - - -class _CreateAgentEngineTaskRequestParametersDict(TypedDict, total=False): - """Parameters for creating Agent Engine Tasks.""" - - name: Optional[str] - """Name of the agent engine to create the task under.""" - - a2a_task_id: Optional[str] - """The ID of the task.""" - - config: Optional[CreateAgentEngineTaskConfigDict] - """""" - - -_CreateAgentEngineTaskRequestParametersOrDict = Union[ - _CreateAgentEngineTaskRequestParameters, _CreateAgentEngineTaskRequestParametersDict -] - - -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" } }""" - - new_metadata: Optional[dict[str, Any]] = Field( + return_raw_output: Optional[bool] = Field( default=None, - description="""Required. The complete state of the metadata object *after* the change.""", + description="""Whether to return the raw output from the judge model.""", ) - update_mask: Optional[str] = Field( + parse_and_reduce_fn: Optional[Callable[..., Any]] = 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""", + description="""The parse and reduce function for the judge model.""", ) - - -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" } }""" - - 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""" - - -TaskMetadataChangeOrDict = Union[TaskMetadataChange, TaskMetadataChangeDict] - - -class TaskArtifactChange(_common.BaseModel): - """Describes changes to the artifact list.""" - - added_artifacts: Optional[list[TaskArtifact]] = Field( + aggregate_summary_fn: Optional[Callable[..., Any]] = Field( default=None, - description="""Optional. A list of brand-new artifacts created in this event.""", + description="""The aggregate summary function for the judge model.""", ) - deleted_artifact_ids: Optional[list[str]] = Field( + remote_custom_function: Optional[str] = Field( default=None, - description="""Optional. A list of artifact IDs that were removed in this event.""", + description="""The evaluation function for the custom code execution metric. This custom code is run remotely in the evaluation service.""", + ) + judge_model: Optional[str] = Field( + default=None, description="""The judge model for the metric.""" ) - updated_artifacts: Optional[list[TaskArtifact]] = Field( + judge_model_generation_config: Optional[genai_types.GenerationConfig] = Field( default=None, - description="""Optional. A list of existing artifacts that were modified in this event.""", + description="""The generation config for the judge LLM (temperature, top_k, top_p, etc).""", ) + judge_model_sampling_count: Optional[int] = Field( + default=None, description="""The sampling count for the judge model.""" + ) + rubric_group_name: Optional[str] = Field( + default=None, + description="""The rubric group name for the rubric-based metric.""", + ) + metric_spec_parameters: Optional[dict[str, Any]] = Field( + default=None, + description="""Optional steering instruction parameters for the automated predefined metric.""", + ) + metric_resource_name: Optional[str] = Field( + default=None, + description="""The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""", + ) + result_parsing_function: Optional[str] = Field( + default=None, + description="""Optional. A Python function string used to parse the raw output of the LLM judge model. The function must be named `parse_results` and accept a list of model response strings. It should return a dictionary with `score` (float) and `explanation` (str) keys.""", + ) + + # Allow extra fields to support metric-specific config fields. + model_config = ConfigDict(extra="allow") + + _is_predefined: bool = PrivateAttr(default=False) + """A boolean indicating whether the metric is predefined.""" + + _config_source: Optional[str] = PrivateAttr(default=None) + """An optional string indicating the source of the metric configuration.""" + _version: Optional[str] = PrivateAttr(default=None) + """An optional string indicating the version of the metric.""" -class TaskArtifactChangeDict(TypedDict, total=False): - """Describes changes to the artifact list.""" - - added_artifacts: Optional[list[TaskArtifactDict]] - """Optional. A list of brand-new artifacts created in this event.""" - - deleted_artifact_ids: Optional[list[str]] - """Optional. A list of artifact IDs that were removed in this event.""" + @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 - updated_artifacts: Optional[list[TaskArtifactDict]] - """Optional. A list of existing artifacts that were modified in this event.""" + 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. -TaskArtifactChangeOrDict = Union[TaskArtifactChange, TaskArtifactChangeDict] + Raises: + ImportError: If the pyyaml library is not installed. + """ + if yaml is None: + raise ImportError( + "YAML serialization requires the pyyaml library. Please install" + " it using 'pip install google-cloud-aiplatform[evaluation]'." + ) + fields_to_exclude = { + field_name + for field_name, field_info in self.model_fields.items() + if self.__getattribute__(field_name) is not None + and isinstance(self.__getattribute__(field_name), Callable) + } -class TaskOutputChange(_common.BaseModel): - """An event representing a change to the task's outputs.""" + data_to_dump = self.model_dump( + exclude_unset=True, + exclude_none=True, + mode="json", + exclude=fields_to_exclude if fields_to_exclude else None, + ) - task_artifact_change: Optional[TaskArtifactChange] = Field( - default=None, - description="""Required. A granular change to the list of artifacts.""", - ) + if version: + data_to_dump["version"] = version + with open(file_path, "w", encoding="utf-8") as f: + yaml.dump(data_to_dump, f, sort_keys=False, allow_unicode=True) -class TaskOutputChangeDict(TypedDict, total=False): - """An event representing a change to the task's outputs.""" - task_artifact_change: Optional[TaskArtifactChangeDict] - """Required. A granular change to the list of artifacts.""" +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.""", + ) -TaskOutputChangeOrDict = Union[TaskOutputChange, TaskOutputChangeDict] + # 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 TaskStateChange(_common.BaseModel): - """A message representing a change in a task's state.""" +class LLMMetric(Metric): + """A metric that uses LLM-as-a-judge for evaluation.""" - new_state: Optional[State] = Field( - default=None, description="""Required. The new state of the task.""" + 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.""", ) + 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.""", + ) -class TaskStateChangeDict(TypedDict, total=False): - """A message representing a change in a task's state.""" + @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 - new_state: Optional[State] - """Required. The new state of the task.""" + @field_validator("judge_model_sampling_count") + @classmethod + def validate_judge_model_sampling_count(cls, value: Optional[int]) -> Optional[int]: + """Validates judge_model_sampling_count to be between 1 and 32.""" + if value is not None and (value < 1 or value > 32): + raise ValueError("judge_model_sampling_count must be between 1 and 32.") + return value + @classmethod + def load(cls, config_path: str, client: Optional[Any] = None) -> "LLMMetric": + """Loads a metric configuration from a YAML or JSON file. -TaskStateChangeOrDict = Union[TaskStateChange, TaskStateChangeDict] + 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. + 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. -class TaskStatusDetailsChange(_common.BaseModel): - """Represents a change to the task's status details.""" + Returns: + An instance of LLMMetric configured with the loaded data. - new_task_status: Optional[TaskStatusDetails] = Field( - default=None, - description="""Required. The complete state of the task's status *after* the change.""", - ) + 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 TaskStatusDetailsChangeDict(TypedDict, total=False): - """Represents a change to the task's status details.""" + 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) - new_task_status: Optional[TaskStatusDetailsDict] - """Required. The complete state of the task's status *after* the change.""" + 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] -TaskStatusDetailsChangeOrDict = Union[ - TaskStatusDetailsChange, TaskStatusDetailsChangeDict -] + 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) + if not isinstance(data, dict): + raise ValueError("Metric config content did not parse into a dictionary.") -class TaskEventData(_common.BaseModel): - """Data for a TaskEvent.""" + return cls.model_validate(data) - 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 MetricDict(TypedDict, total=False): + """The metric used for evaluation.""" -class TaskEventDataDict(TypedDict, total=False): - """Data for a TaskEvent.""" + name: Optional[str] + """The name of the metric.""" - metadata_change: Optional[TaskMetadataChangeDict] - """Optional. A change to the task's metadata.""" + custom_function: Optional[Union[str, Callable[..., Any]]] + """The custom function that defines the end-to-end logic for metric computation.""" - output_change: Optional[TaskOutputChangeDict] - """Optional. A change to the task's final outputs.""" + prompt_template: Optional[str] + """The prompt template for the metric.""" - state_change: Optional[TaskStateChangeDict] - """Optional. A change in the task's state.""" + judge_model_system_instruction: Optional[str] + """The system instruction for the judge model.""" - status_details_change: Optional[TaskStatusDetailsChangeDict] - """Optional. A change to the framework-specific status details.""" + 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.""" -TaskEventDataOrDict = Union[TaskEventData, TaskEventDataDict] + 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 TaskEvent(_common.BaseModel): - """A task event.""" + judge_model: Optional[str] + """The judge model for the metric.""" - create_time: Optional[datetime.datetime] = Field( - default=None, description="""Output only. The create time of the event.""" - ) - event_data: Optional[TaskEventData] = Field( - default=None, description="""Required. The delta associated with the event.""" - ) - event_sequence_number: Optional[int] = 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.""", - ) + 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 TaskEventDict(TypedDict, total=False): - """A task event.""" + rubric_group_name: Optional[str] + """The rubric group name for the rubric-based metric.""" - create_time: Optional[datetime.datetime] - """Output only. The create time of the event.""" + metric_spec_parameters: Optional[dict[str, Any]] + """Optional steering instruction parameters for the automated predefined metric.""" - event_data: Optional[TaskEventDataDict] - """Required. The delta associated with the event.""" + metric_resource_name: Optional[str] + """The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" - 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.""" + 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.""" -TaskEventOrDict = Union[TaskEvent, TaskEventDict] +MetricOrDict = Union[Metric, MetricDict] -class AppendAgentEngineTaskEventConfig(_common.BaseModel): - """Config for appending Agent Engine task events.""" +class CreateEvaluationMetricConfig(_common.BaseModel): + """Config for creating an evaluation metric.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class AppendAgentEngineTaskEventConfigDict(TypedDict, total=False): - """Config for appending Agent Engine task events.""" +class CreateEvaluationMetricConfigDict(TypedDict, total=False): + """Config for creating an evaluation metric.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -AppendAgentEngineTaskEventConfigOrDict = Union[ - AppendAgentEngineTaskEventConfig, AppendAgentEngineTaskEventConfigDict +CreateEvaluationMetricConfigOrDict = Union[ + CreateEvaluationMetricConfig, CreateEvaluationMetricConfigDict ] -class _AppendAgentEngineTaskEventRequestParameters(_common.BaseModel): - """Parameters for appending Agent Engine task events.""" +class _CreateEvaluationMetricParameters(_common.BaseModel): + """Parameters for creating an evaluation metric.""" - name: Optional[str] = Field( + display_name: Optional[str] = Field( default=None, - description="""Name of the Agent Engine task to append the events to.""", + description="""The user-defined name of the evaluation metric. + + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """, ) - task_events: Optional[list[TaskEvent]] = Field( - default=None, description="""The events to append to the task.""" + description: Optional[str] = Field( + default=None, description="""The description of the evaluation metric.""" + ) + metric: Optional[Metric] = Field( + default=None, + description="""The metric configuration of the evaluation metric.""", ) - config: Optional[AppendAgentEngineTaskEventConfig] = Field( + config: Optional[CreateEvaluationMetricConfig] = Field( default=None, description="""""" ) -class _AppendAgentEngineTaskEventRequestParametersDict(TypedDict, total=False): - """Parameters for appending Agent Engine task events.""" +class _CreateEvaluationMetricParametersDict(TypedDict, total=False): + """Parameters for creating an evaluation metric.""" - name: Optional[str] - """Name of the Agent Engine task to append the events to.""" + display_name: Optional[str] + """The user-defined name of the evaluation metric. - task_events: Optional[list[TaskEventDict]] - """The events to append to the task.""" + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """ - config: Optional[AppendAgentEngineTaskEventConfigDict] - """""" + description: Optional[str] + """The description of the evaluation metric.""" + metric: Optional[MetricDict] + """The metric configuration of the evaluation metric.""" -_AppendAgentEngineTaskEventRequestParametersOrDict = Union[ - _AppendAgentEngineTaskEventRequestParameters, - _AppendAgentEngineTaskEventRequestParametersDict, + config: Optional[CreateEvaluationMetricConfigDict] + """""" + + +_CreateEvaluationMetricParametersOrDict = Union[ + _CreateEvaluationMetricParameters, _CreateEvaluationMetricParametersDict ] -class AppendAgentEngineTaskEventResponse(_common.BaseModel): - """Response for appending Agent Engine task events.""" +class CustomCodeExecutionSpec(_common.BaseModel): + """Specifies a metric using remote Python function execution. + + This metric is computed by running user-defined Python functions remotely. + """ - pass + evaluation_function: Optional[str] = Field( + default=None, + description="""Required. Python function. Expected user to define the following function, e.g.: def evaluate(instance: dict[str, Any]) -> float: Please include this function signature in the code snippet. Instance is the evaluation instance, any fields populated in the instance are available to the function as instance[field_name]. Example: Example input: ``` instance= EvaluationInstance( response=EvaluationInstance.InstanceData(text="The answer is 4."), reference=EvaluationInstance.InstanceData(text="4") ) ``` Example converted input: ``` { 'response': {'text': 'The answer is 4.'}, 'reference': {'text': '4'} } ``` Example python function: ``` def evaluate(instance: dict[str, Any]) -> float: if instance'response' == instance'reference': return 1.0 return 0.0 ``` CustomCodeExecutionSpec is also supported in Batch Evaluation (EvalDataset RPC) and Tuning Evaluation. Each line in the input jsonl file will be converted to dict[str, Any] and passed to the evaluation function.""", + ) + 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 AppendAgentEngineTaskEventResponseDict(TypedDict, total=False): - """Response for appending Agent Engine task events.""" +class CustomCodeExecutionSpecDict(TypedDict, total=False): + """Specifies a metric using remote Python function execution. - pass + This metric is computed by running user-defined Python functions remotely. + """ + + 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.""" + + 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.""", ) -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.""" -class ListAgentEngineTaskEventsResponse(_common.BaseModel): - """Response for listing Agent Engine tasks events.""" +EvaluationMetricOrDict = Union[EvaluationMetric, EvaluationMetricDict] - 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.""" - ) +class SamplingConfig(_common.BaseModel): + """Sampling config for a BigQuery request set.""" + + sampling_count: Optional[int] = Field(default=None, description="""""") + sampling_method: Optional[SamplingMethod] = Field(default=None, description="""""") + sampling_duration: Optional[str] = Field(default=None, description="""""") -class ListAgentEngineTaskEventsResponseDict(TypedDict, total=False): - """Response for listing Agent Engine tasks events.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" +class SamplingConfigDict(TypedDict, total=False): + """Sampling config for a BigQuery request set.""" - next_page_token: Optional[str] + 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 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 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 CreateEvaluationItemConfigDict(TypedDict, total=False): - """Config to create an evaluation item.""" +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.""" -CreateEvaluationItemConfigOrDict = Union[ - CreateEvaluationItemConfig, CreateEvaluationItemConfigDict -] + 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 _CreateEvaluationItemParameters(_common.BaseModel): - """Represents a job that creates an evaluation item.""" + sampling_config: Optional[SamplingConfigDict] + """The sampling config for the BigQuery request set. Used for EvaluationRun only.""" - 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="""""" - ) +BigQueryRequestSetOrDict = Union[BigQueryRequestSet, BigQueryRequestSetDict] -class _CreateEvaluationItemParametersDict(TypedDict, total=False): - """Represents a job that creates an evaluation item.""" - evaluation_item_type: Optional[str] - """""" +class EvaluationRunDataSource(_common.BaseModel): + """Represents an evaluation run data source.""" - gcs_uri: Optional[str] - """""" + evaluation_set: Optional[str] = Field(default=None, description="""""") + bigquery_request_set: Optional[BigQueryRequestSet] = Field( + default=None, description="""""" + ) - display_name: Optional[str] + +class EvaluationRunDataSourceDict(TypedDict, total=False): + """Represents an evaluation run data source.""" + + evaluation_set: Optional[str] """""" - config: Optional[CreateEvaluationItemConfigDict] + bigquery_request_set: Optional[BigQueryRequestSetDict] """""" -_CreateEvaluationItemParametersOrDict = Union[ - _CreateEvaluationItemParameters, _CreateEvaluationItemParametersDict +EvaluationRunDataSourceOrDict = Union[ + EvaluationRunDataSource, EvaluationRunDataSourceDict ] -class PromptTemplateData(_common.BaseModel): - """Holds data for a prompt template. - - Message to hold a prompt template and the values to populate the template. - """ +class EvaluationRunMetric(_common.BaseModel): + """The metric used for evaluation run.""" - values: Optional[dict[str, genai_types.Content]] = Field( - default=None, description="""The values for fields in the prompt template.""" + metric: Optional[str] = Field( + default=None, description="""The 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}""", + ) + metric_config: Optional[UnifiedMetric] = Field( + default=None, description="""The unified metric used for evaluation run.""" ) -class PromptTemplateDataDict(TypedDict, total=False): - """Holds data for a prompt template. +class EvaluationRunMetricDict(TypedDict, total=False): + """The metric used for evaluation run.""" - Message to hold a prompt template and the values to populate the template. - """ + metric: Optional[str] + """The name of the metric.""" - values: Optional[dict[str, genai_types.Content]] - """The values for fields in the prompt template.""" + metric_resource_name: Optional[str] + """The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" + metric_config: Optional[UnifiedMetricDict] + """The unified metric used for evaluation run.""" -PromptTemplateDataOrDict = Union[PromptTemplateData, PromptTemplateDataDict] +EvaluationRunMetricOrDict = Union[EvaluationRunMetric, EvaluationRunMetricDict] -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( +class EvaluationRunPromptTemplate(_common.BaseModel): + """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. + """ + + prompt_template: Optional[str] = Field( default=None, - description="""Fields and values that can be used to populate the prompt template.""", - ) - prompt_template_data: Optional[PromptTemplateData] = Field( - default=None, description="""Prompt template data.""" + description="""Inline prompt template. Template variables should be in the format + "{var_name}". Only one of `prompt_template` or `gcs_uri` should be set.""", ) - user_scenario: Optional[evals_types.UserScenario] = Field( + gcs_uri: Optional[str] = Field( default=None, - description="""User scenario to help simulate multi-turn agent run results.""", + 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 EvaluationPromptDict(TypedDict, total=False): - """Represents the prompt to be evaluated.""" - - text: Optional[str] - """Text prompt.""" +class EvaluationRunPromptTemplateDict(TypedDict, total=False): + """Prompt template used for inference. - value: Optional[dict[str, Any]] - """Fields and values that can be used to populate the prompt template.""" + Only one of `prompt_template` or `gcs_uri` should be set. If both are + provided, an error will be raised. + """ - prompt_template_data: Optional[PromptTemplateDataDict] - """Prompt template data.""" + 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.""" - user_scenario: Optional[evals_types.UserScenario] - """User scenario to help simulate multi-turn agent run results.""" + 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.""" -EvaluationPromptOrDict = Union[EvaluationPrompt, EvaluationPromptDict] +EvaluationRunPromptTemplateOrDict = Union[ + EvaluationRunPromptTemplate, EvaluationRunPromptTemplateDict +] -class CandidateResponse(_common.BaseModel): - """Responses from model or agent.""" +class LossAnalysisConfig(_common.BaseModel): + """Configuration for the loss analysis job.""" - candidate: Optional[str] = Field( + metric: Optional[str] = Field( default=None, - description="""The name of the candidate that produced the response.""", + description="""Required. The metric to analyze (e.g., "multi_turn_tool_use_quality_v1").""", ) - text: Optional[str] = Field(default=None, description="""The text response.""") - value: Optional[dict[str, Any]] = Field( + candidate: Optional[str] = Field( default=None, - description="""Fields and values that can be used to populate the response template.""", + description="""Required. The candidate model/agent to analyze (e.g., "gemini-3.1-pro-preview"). This targets the specific CandidateResult within the EvaluationResult.""", ) - events: Optional[list[genai_types.Content]] = Field( + predefined_taxonomy: Optional[str] = Field( default=None, - description="""Intermediate events (such as tool calls and responses) that led to the final response.""", + 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.""", ) - agent_data: Optional[evals_types.AgentData] = Field( + max_top_cluster_count: Optional[int] = Field( default=None, - description="""Represents the complete execution trace of an agent conversation, - which can involve single or multiple agents. This field is used to - provide the full output of an agent's run, including all turns and - events, for direct evaluation.""", + description="""Optional. Limits the analysis to the top N clusters. If not specified or set to 0, all clusters are returned.""", ) -class CandidateResponseDict(TypedDict, total=False): - """Responses from model or agent.""" - - candidate: Optional[str] - """The name of the candidate that produced the response.""" +class LossAnalysisConfigDict(TypedDict, total=False): + """Configuration for the loss analysis job.""" - text: Optional[str] - """The text response.""" + metric: Optional[str] + """Required. The metric to analyze (e.g., "multi_turn_tool_use_quality_v1").""" - value: Optional[dict[str, Any]] - """Fields and values that can be used to populate the response template.""" + 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.""" - events: Optional[list[genai_types.Content]] - """Intermediate events (such as tool calls and responses) that led to the final response.""" + 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.""" - 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.""" + 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.""" -CandidateResponseOrDict = Union[CandidateResponse, CandidateResponseDict] +LossAnalysisConfigOrDict = Union[LossAnalysisConfig, LossAnalysisConfigDict] -class EvaluationItemRequest(_common.BaseModel): - """Single evaluation request.""" +class EvaluationRunConfig(_common.BaseModel): + """The evaluation configuration used for the evaluation run.""" - prompt: Optional[EvaluationPrompt] = Field( - default=None, description="""The request/prompt to evaluate.""" + metrics: Optional[list[EvaluationRunMetric]] = Field( + default=None, + description="""The metrics to be calculated in the evaluation run.""", ) - golden_response: Optional[CandidateResponse] = Field( - default=None, description="""The ideal response or ground truth.""" + output_config: Optional[genai_types.OutputConfig] = Field( + default=None, description="""The output config for the evaluation run.""" ) - rubrics: Optional[dict[str, "RubricGroup"]] = Field( + autorater_config: Optional[genai_types.AutoraterConfig] = Field( default=None, - description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", + description="""The autorater config for the evaluation run. Not applicable for predefined metrics (PredefinedMetricSpec); the server uses its own model configuration for predefined metrics and this field is ignored.""", ) - candidate_responses: Optional[list[CandidateResponse]] = Field( + prompt_template: Optional[EvaluationRunPromptTemplate] = Field( + default=None, description="""The prompt template used for inference.""" + ) + loss_analysis_config: Optional[list[LossAnalysisConfig]] = Field( default=None, - description="""Responses from model under test and other baseline models for comparison.""", + 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 EvaluationItemRequestDict(TypedDict, total=False): - """Single evaluation request.""" +class EvaluationRunConfigDict(TypedDict, total=False): + """The evaluation configuration used for the evaluation run.""" - prompt: Optional[EvaluationPromptDict] - """The request/prompt to evaluate.""" + metrics: Optional[list[EvaluationRunMetricDict]] + """The metrics to be calculated in the evaluation run.""" - golden_response: Optional[CandidateResponseDict] - """The ideal response or ground truth.""" + output_config: Optional[genai_types.OutputConfig] + """The output config for the evaluation run.""" - rubrics: Optional[dict[str, "RubricGroupDict"]] - """Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""" + autorater_config: Optional[genai_types.AutoraterConfig] + """The autorater config for the evaluation run. Not applicable for predefined metrics (PredefinedMetricSpec); the server uses its own model configuration for predefined metrics and this field is ignored.""" - candidate_responses: Optional[list[CandidateResponseDict]] - """Responses from model under test and other baseline models for comparison.""" + 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.""" -EvaluationItemRequestOrDict = Union[EvaluationItemRequest, EvaluationItemRequestDict] + 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 EvaluationItemResult(_common.BaseModel): - """Represents the result of an evaluation item.""" +EvaluationRunConfigOrDict = Union[EvaluationRunConfig, EvaluationRunConfigDict] - evaluation_request: Optional[str] = Field( - default=None, description="""The request item that was evaluated.""" - ) - evaluation_run: Optional[str] = Field( - default=None, - description="""The evaluation run that was used to generate the result.""", - ) - request: Optional[EvaluationItemRequest] = Field( - default=None, description="""The request that was evaluated.""" - ) - metric: Optional[str] = Field( - default=None, description="""The metric that was evaluated.""" - ) - candidate_results: Optional[list[evals_types.CandidateResult]] = Field( - default=None, description="""The results for the metric.""" + +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.""" ) - metadata: Optional[dict[str, Any]] = Field( - default=None, description="""Metadata about the evaluation result.""" + tools: Optional[list[genai_types.Tool]] = Field( + default=None, description="""The tools available to the agent.""" ) -class EvaluationItemResultDict(TypedDict, total=False): - """Represents the result of an evaluation item.""" +class EvaluationRunAgentConfigDict(TypedDict, total=False): + """Agent config for an evaluation run.""" - evaluation_request: Optional[str] - """The request item that was evaluated.""" + developer_instruction: Optional[genai_types.Content] + """The developer instruction for the agent.""" - evaluation_run: Optional[str] - """The evaluation run that was used to generate the result.""" + tools: Optional[list[genai_types.Tool]] + """The tools available to the agent.""" - request: Optional[EvaluationItemRequestDict] - """The request that was evaluated.""" - metric: Optional[str] - """The metric that was evaluated.""" +EvaluationRunAgentConfigOrDict = Union[ + EvaluationRunAgentConfig, EvaluationRunAgentConfigDict +] - candidate_results: Optional[list[evals_types.CandidateResult]] - """The results for the metric.""" - metadata: Optional[dict[str, Any]] - """Metadata about the evaluation result.""" +class GeminiAgentConfig(_common.BaseModel): + """Config for scraping a Gemini Agent. + A Gemini Agent is a Vertex AI Agent resource scraped via the Vertex + Interactions API. + """ -EvaluationItemResultOrDict = Union[EvaluationItemResult, EvaluationItemResultDict] + gemini_agent: Optional[str] = Field( + default=None, + description="""The resource name of the Gemini Agent. + Format: `projects/{project}/locations/{location}/agents/{agent}`.""", + ) -class EvaluationItem(_common.BaseModel): - """EvaluationItem is a single evaluation request or result. +class GeminiAgentConfigDict(TypedDict, total=False): + """Config for scraping a Gemini Agent. - The content of an EvaluationItem is immutable - it cannot be updated once - created. EvaluationItems can be deleted when no longer needed. + A Gemini Agent is a Vertex AI Agent resource scraped via the Vertex + Interactions API. """ - 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.""" + gemini_agent: Optional[str] + """The resource name of the Gemini Agent. + Format: `projects/{project}/locations/{location}/agents/{agent}`.""" + + +GeminiAgentConfigOrDict = Union[GeminiAgentConfig, GeminiAgentConfigDict] + + +class AgentRunConfig(_common.BaseModel): + """Configuration for an Agent Run.""" + + session_input: Optional[evals_types.SessionInput] = Field( + default=None, description="""The session input to get agent running results.""" ) - evaluation_response: Optional[EvaluationItemResult] = Field( - default=None, description="""The response from evaluation.""" + agent_engine: Optional[str] = Field( + default=None, description="""The resource name of the Agent Engine.""" ) - gcs_uri: Optional[str] = Field( + user_simulator_config: Optional[evals_types.UserSimulatorConfig] = 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="""Used for multi-turn agent run. + Contains configuration for a user simulator that + uses an LLM to generate messages on behalf of the user.""", ) - error: Optional[genai_types.GoogleRpcStatus] = Field( - default=None, description="""Error for the evaluation item.""" + gemini_agent_config: Optional[GeminiAgentConfig] = Field( + default=None, + description="""Config for scraping a Gemini Agent (Vertex AI Agent resource). + Used to target a Gemini agent for an evaluation run.""", ) - # 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 AgentRunConfigDict(TypedDict, total=False): + """Configuration for an Agent Run.""" + session_input: Optional[evals_types.SessionInput] + """The session input to get agent running results.""" -class EvaluationItemDict(TypedDict, total=False): - """EvaluationItem is a single evaluation request or result. + agent_engine: Optional[str] + """The resource name of the Agent Engine.""" - The content of an EvaluationItem is immutable - it cannot be updated once - created. EvaluationItems can be deleted when no longer needed. - """ + 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.""" - name: Optional[str] - """The resource name of the EvaluationItem.""" + 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.""" - display_name: Optional[str] - """The display name of the EvaluationItem.""" - metadata: Optional[dict[str, Any]] - """Metadata for the EvaluationItem.""" +AgentRunConfigOrDict = Union[AgentRunConfig, AgentRunConfigDict] - labels: Optional[dict[str, str]] - """Labels for the EvaluationItem.""" - evaluation_item_type: Optional[EvaluationItemType] - """The type of the EvaluationItem.""" +class EvaluationRunInferenceConfig(_common.BaseModel): + """Configuration that describes an agent.""" - evaluation_request: Optional[EvaluationItemRequestDict] - """The request to evaluate.""" + 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.""", + ) - evaluation_response: Optional[EvaluationItemResultDict] - """The response from evaluation.""" - gcs_uri: Optional[str] - """The Cloud Storage object where the request or response is stored.""" +class EvaluationRunInferenceConfigDict(TypedDict, total=False): + """Configuration that describes an agent.""" - create_time: Optional[datetime.datetime] - """Timestamp when this item was created.""" + agent_config: Optional[EvaluationRunAgentConfigDict] + """The agent config.""" - error: Optional[genai_types.GoogleRpcStatus] - """Error for the evaluation item.""" + model: Optional[str] + """The model to use for inference. Accepts a short Gemini model name (e.g. `gemini-2.5-flash`), which is automatically expanded to a fully-qualified resource name using the client's project and location, or an already fully-qualified publisher-model or endpoint resource name (e.g. `projects/{project}/locations/{location}/publishers/google/models/gemini-2.5-flash`).""" + prompt_template: Optional[EvaluationRunPromptTemplateDict] + """The prompt template used for inference.""" -EvaluationItemOrDict = Union[EvaluationItem, EvaluationItemDict] + 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.""" -class Metric(_common.BaseModel): - """The metric used for evaluation.""" - name: Optional[str] = Field(default=None, description="""The name of the metric.""") - custom_function: Optional[Union[str, Callable[..., Any]]] = Field( - default=None, - description="""The custom function that defines the end-to-end logic for metric computation.""", - ) - prompt_template: Optional[str] = Field( - default=None, description="""The prompt template for the metric.""" - ) - judge_model_system_instruction: Optional[str] = Field( - default=None, description="""The system instruction for the judge model.""" - ) - return_raw_output: Optional[bool] = Field( - default=None, - description="""Whether to return the raw output from the judge model.""", - ) - parse_and_reduce_fn: Optional[Callable[..., Any]] = Field( - default=None, - description="""The parse and reduce function for the judge model.""", - ) - aggregate_summary_fn: Optional[Callable[..., Any]] = Field( - default=None, - description="""The aggregate summary function for the judge model.""", - ) - remote_custom_function: Optional[str] = Field( +EvaluationRunInferenceConfigOrDict = Union[ + EvaluationRunInferenceConfig, EvaluationRunInferenceConfigDict +] + + +class VulnerableTool(_common.BaseModel): + """A tool considered high risk for prompt injection.""" + + tool_name: 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.""", - ) - judge_model: Optional[str] = Field( - default=None, description="""The judge model for the metric.""" + description="""Optional. The name of the vulnerable function/tool (e.g., "search_flights").""", ) - judge_model_generation_config: Optional[genai_types.GenerationConfig] = Field( + json_paths: Optional[list[str]] = Field( default=None, - description="""The generation config for the judge LLM (temperature, top_k, top_p, etc).""", - ) - judge_model_sampling_count: Optional[int] = Field( - default=None, description="""The sampling count for the judge model.""" + description="""Optional. JSON Paths within the tool's FunctionResponse where malicious content could be injected.""", ) - rubric_group_name: Optional[str] = Field( + + +class VulnerableToolDict(TypedDict, total=False): + """A tool considered high risk for prompt injection.""" + + tool_name: Optional[str] + """Optional. The name of the vulnerable function/tool (e.g., "search_flights").""" + + json_paths: Optional[list[str]] + """Optional. JSON Paths within the tool's FunctionResponse where malicious content could be injected.""" + + +VulnerableToolOrDict = Union[VulnerableTool, VulnerableToolDict] + + +class RedTeamingAnalysisConfig(_common.BaseModel): + """Configuration for the automated Agent Red Teaming analysis.""" + + attack_categories: Optional[list[str]] = Field( default=None, - description="""The rubric group name for the rubric-based metric.""", + description="""Optional. Specific attack categories to test against.""", ) - metric_spec_parameters: Optional[dict[str, Any]] = Field( + vulnerable_tools: Optional[list[VulnerableTool]] = Field( default=None, - description="""Optional steering instruction parameters for the automated predefined metric.""", + description="""Optional. Manually defined vulnerable tools and their injection paths.""", ) - 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}""", + + +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.""" ) - result_parsing_function: Optional[str] = Field( + red_teaming_analysis_config: Optional[RedTeamingAnalysisConfig] = 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="""Configuration for the automated Agent Red Teaming analysis.""", ) - # 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 AnalysisConfigDict(TypedDict, total=False): + """Configuration for an analysis to be performed on an evaluation run.""" - _config_source: Optional[str] = PrivateAttr(default=None) - """An optional string indicating the source of the metric configuration.""" + analysis_name: Optional[str] + """Optional. A name for this analysis.""" - _version: Optional[str] = PrivateAttr(default=None) - """An optional string indicating the version of the metric.""" + red_teaming_analysis_config: Optional[RedTeamingAnalysisConfigDict] + """Configuration for the automated Agent Red Teaming analysis.""" - @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. +AnalysisConfigOrDict = Union[AnalysisConfig, AnalysisConfigDict] - 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]'." - ) - - fields_to_exclude = { - field_name - for field_name, field_info in self.model_fields.items() - if self.__getattribute__(field_name) is not None - and isinstance(self.__getattribute__(field_name), Callable) - } +class CreateEvaluationRunConfig(_common.BaseModel): + """Config to create an evaluation run.""" - data_to_dump = self.model_dump( - exclude_unset=True, - exclude_none=True, - mode="json", - exclude=fields_to_exclude if fields_to_exclude else None, - ) + 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.""", + ) - if version: - data_to_dump["version"] = version - with open(file_path, "w", encoding="utf-8") as f: - yaml.dump(data_to_dump, f, sort_keys=False, allow_unicode=True) +class CreateEvaluationRunConfigDict(TypedDict, total=False): + """Config to create an evaluation run.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -class CodeExecutionMetric(Metric): - """A metric that executes custom Python code for evaluation.""" + allow_cross_region_model: Optional[bool] + """Allows the evaluation run to use cross region models. When this + flag is set, the service may route traffic to other regions if a model is + unavailable in the current region (e.g., to a `global`endpoint). If a + fully-qualified model endpoint resource name with a different region than + the run location is provided elsewhere in the run config, this flag must + be set to true or the request will fail.""" - # 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 +CreateEvaluationRunConfigOrDict = Union[ + CreateEvaluationRunConfig, CreateEvaluationRunConfigDict +] -class LLMMetric(Metric): - """A metric that uses LLM-as-a-judge for evaluation.""" +class _CreateEvaluationRunParameters(_common.BaseModel): + """Represents a job that creates 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.""", + name: Optional[str] = Field(default=None, description="""""") + display_name: Optional[str] = Field(default=None, description="""""") + data_source: Optional[EvaluationRunDataSource] = Field( + default=None, description="""""" ) - - 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.""", + evaluation_config: Optional[EvaluationRunConfig] = Field( + default=None, description="""""" + ) + labels: Optional[dict[str, str]] = Field(default=None, description="""""") + inference_configs: Optional[dict[str, EvaluationRunInferenceConfig]] = Field( + default=None, description="""""" + ) + config: Optional[CreateEvaluationRunConfig] = Field( + default=None, description="""""" + ) + analysis_configs: Optional[list[AnalysisConfig]] = Field( + default=None, description="""""" ) - @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 _CreateEvaluationRunParametersDict(TypedDict, total=False): + """Represents a job that creates 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. + name: Optional[str] + """""" - 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. + display_name: Optional[str] + """""" - 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. + data_source: Optional[EvaluationRunDataSourceDict] + """""" - Returns: - An instance of LLMMetric configured with the loaded data. + evaluation_config: Optional[EvaluationRunConfigDict] + """""" - 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" - ) + labels: Optional[dict[str, str]] + """""" - content_str: str - if config_path.startswith("gs://"): - try: - from google.cloud import storage # type: ignore[attr-defined] + inference_configs: Optional[dict[str, EvaluationRunInferenceConfigDict]] + """""" - 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) + config: Optional[CreateEvaluationRunConfigDict] + """""" - bucket = storage_client.bucket(bucket_name) - blob = bucket.blob(blob_path) - content_str = blob.download_as_bytes().decode("utf-8") - except ImportError as e: - raise ImportError( - "Reading from GCS requires the 'google-cloud-storage' library. Please install it with 'pip install google-cloud-aiplatform[evaluation]'." - ) from e - except Exception as e: - raise IOError(f"Failed to read from GCS path {config_path}: {e}") from e - else: - try: - with open(config_path, "r", encoding="utf-8") as f: - content_str = f.read() - except FileNotFoundError: - raise FileNotFoundError( - f"Local configuration file not found at: {config_path}" - ) - except Exception as e: - raise IOError(f"Failed to read local file {config_path}: {e}") from e + analysis_configs: Optional[list[AnalysisConfigDict]] + """""" - 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) +_CreateEvaluationRunParametersOrDict = Union[ + _CreateEvaluationRunParameters, _CreateEvaluationRunParametersDict +] - if not isinstance(data, dict): - raise ValueError("Metric config content did not parse into a dictionary.") - return cls.model_validate(data) +class SummaryMetric(_common.BaseModel): + """Represents a summary metric for an evaluation run.""" + metrics: Optional[dict[str, Any]] = Field( + default=None, description="""Map of metric name to metric value.""" + ) + total_items: Optional[int] = Field( + default=None, description="""The total number of items that were evaluated.""" + ) + failed_items: Optional[int] = Field( + default=None, description="""The number of items that failed to be evaluated.""" + ) -class MetricDict(TypedDict, total=False): - """The metric used for evaluation.""" - name: Optional[str] - """The name of the metric.""" +class SummaryMetricDict(TypedDict, total=False): + """Represents a summary metric for an evaluation run.""" - custom_function: Optional[Union[str, Callable[..., Any]]] - """The custom function that defines the end-to-end logic for metric computation.""" + metrics: Optional[dict[str, Any]] + """Map of metric name to metric value.""" - prompt_template: Optional[str] - """The prompt template for the metric.""" + total_items: Optional[int] + """The total number of items that were evaluated.""" - judge_model_system_instruction: Optional[str] - """The system instruction for the judge model.""" + failed_items: Optional[int] + """The number of items that failed to be evaluated.""" - 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.""" +SummaryMetricOrDict = Union[SummaryMetric, SummaryMetricDict] - 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 AttackCategoryResult(_common.BaseModel): + """The red teaming outcome for a specific attack category.""" - judge_model: Optional[str] - """The judge model for the metric.""" + 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.""" + ) - 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 AttackCategoryResultDict(TypedDict, total=False): + """The red teaming outcome for a specific attack category.""" - 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.""" + attack_category: Optional[str] + """The category of the attack evaluated.""" - metric_resource_name: Optional[str] - """The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" + attack_success_rate: Optional[float] + """The ratio of successful attacks given a fixed budget.""" - 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.""" + vulnerability_insight: Optional[str] + """Insights into why an attack succeeded or failed.""" -MetricOrDict = Union[Metric, MetricDict] +AttackCategoryResultOrDict = Union[AttackCategoryResult, AttackCategoryResultDict] -class CreateEvaluationMetricConfig(_common.BaseModel): - """Config for creating an evaluation metric.""" +class RedTeamingAnalysisResult(_common.BaseModel): + """The top-level result for Red Teaming analysis.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + 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.""" ) -class CreateEvaluationMetricConfigDict(TypedDict, total=False): - """Config for creating an evaluation metric.""" +class RedTeamingAnalysisResultDict(TypedDict, total=False): + """The top-level result for Red Teaming analysis.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + config: Optional[RedTeamingAnalysisConfigDict] + """The configuration used to generate this analysis.""" + + analysis_time: Optional[str] + """The timestamp when this analysis was performed.""" + category_results: Optional[list[AttackCategoryResultDict]] + """Detailed results by attack category.""" -CreateEvaluationMetricConfigOrDict = Union[ - CreateEvaluationMetricConfig, CreateEvaluationMetricConfigDict + +RedTeamingAnalysisResultOrDict = Union[ + RedTeamingAnalysisResult, RedTeamingAnalysisResultDict ] -class _CreateEvaluationMetricParameters(_common.BaseModel): - """Parameters for creating an evaluation metric.""" +class LossTaxonomyEntry(_common.BaseModel): + """A specific entry in the loss pattern taxonomy.""" - display_name: Optional[str] = Field( + l1_category: Optional[str] = Field( default=None, - description="""The user-defined name of the evaluation metric. - - The display name can be up to 128 characters long and can comprise any - UTF-8 characters. - """, - ) - description: Optional[str] = Field( - default=None, description="""The description of the evaluation metric.""" + description="""The primary category of the loss (e.g., "Hallucination", "Tool Calling").""", ) - metric: Optional[Metric] = Field( + l2_category: Optional[str] = Field( default=None, - description="""The metric configuration of the evaluation metric.""", + description="""The secondary category of the loss (e.g., "Hallucination of Action", "Incorrect Tool Selection").""", ) - config: Optional[CreateEvaluationMetricConfig] = Field( - default=None, description="""""" + 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 _CreateEvaluationMetricParametersDict(TypedDict, total=False): - """Parameters for creating an evaluation metric.""" +class LossTaxonomyEntryDict(TypedDict, total=False): + """A specific entry in the loss pattern taxonomy.""" - display_name: Optional[str] - """The user-defined name of the evaluation metric. + l1_category: Optional[str] + """The primary category of the loss (e.g., "Hallucination", "Tool Calling").""" - The display name can be up to 128 characters long and can comprise any - UTF-8 characters. - """ + l2_category: Optional[str] + """The secondary category of the loss (e.g., "Hallucination of Action", "Incorrect Tool Selection").""" description: Optional[str] - """The description of the evaluation metric.""" - - metric: Optional[MetricDict] - """The metric configuration of the evaluation metric.""" - - config: Optional[CreateEvaluationMetricConfigDict] - """""" - + """A detailed description of this loss pattern. Example: "The agent verbally confirms an action without executing the tool." """ -_CreateEvaluationMetricParametersOrDict = Union[ - _CreateEvaluationMetricParameters, _CreateEvaluationMetricParametersDict -] +LossTaxonomyEntryOrDict = Union[LossTaxonomyEntry, LossTaxonomyEntryDict] -class CustomCodeExecutionSpec(_common.BaseModel): - """Specifies a metric using remote Python function execution. - This metric is computed by running user-defined Python functions remotely. - """ +class FailedRubric(_common.BaseModel): + """A specific failed rubric and the associated analysis.""" - evaluation_function: Optional[str] = Field( + rubric_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="""The unique ID of the rubric (if available from the metric source).""", ) - remote_custom_function: Optional[str] = Field( + classification_rationale: 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].""", + description="""The rationale provided by the Loss Analysis Classifier for why this failure maps to this specific Loss Cluster.""", ) -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 FailedRubricDict(TypedDict, total=False): + """A specific failed rubric and the associated analysis.""" - 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.""" + rubric_id: Optional[str] + """The unique ID of the rubric (if available from the metric source).""" - 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].""" + classification_rationale: Optional[str] + """The rationale provided by the Loss Analysis Classifier for why this failure maps to this specific Loss Cluster.""" -CustomCodeExecutionSpecOrDict = Union[ - CustomCodeExecutionSpec, CustomCodeExecutionSpecDict -] +FailedRubricOrDict = Union[FailedRubric, FailedRubricDict] -class UnifiedMetric(_common.BaseModel): - """The unified metric used for evaluation.""" +class LossExample(_common.BaseModel): + """A specific example of a loss pattern.""" - 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.""" + evaluation_item: Optional[str] = Field( + default=None, + description="""Reference to the persisted EvalItem resource name. Format: projects/.../locations/.../evaluationItems/{item_id}.""", ) - predefined_metric_spec: Optional[genai_types.PredefinedMetricSpec] = Field( - default=None, description="""The spec for a pre-defined metric.""" + 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.""", ) - computation_based_metric_spec: Optional[genai_types.ComputationBasedMetricSpec] = ( - Field(default=None, description="""The spec for a computation based metric.""") + 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 UnifiedMetricDict(TypedDict, total=False): - """The unified metric used for evaluation.""" - - bleu_spec: Optional[genai_types.BleuSpec] - """The Bleu metric spec.""" - - rouge_spec: Optional[genai_types.RougeSpec] - """The rouge metric spec.""" - - pointwise_metric_spec: Optional[genai_types.PointwiseMetricSpec] - """The pointwise metric spec.""" - - llm_based_metric_spec: Optional[genai_types.LLMBasedMetricSpec] - """The spec for an LLM based metric.""" +class LossExampleDict(TypedDict, total=False): + """A specific example of a loss pattern.""" - custom_code_execution_spec: Optional[CustomCodeExecutionSpecDict] - """The spec for a custom code execution metric.""" + evaluation_item: Optional[str] + """Reference to the persisted EvalItem resource name. Format: projects/.../locations/.../evaluationItems/{item_id}.""" - predefined_metric_spec: Optional[genai_types.PredefinedMetricSpec] - """The spec for a pre-defined metric.""" + evaluation_result: Optional[dict[str, Any]] + """The full evaluation result object provided inline. Used when the analysis is performed on ephemeral data.""" - computation_based_metric_spec: Optional[genai_types.ComputationBasedMetricSpec] - """The spec for a computation based metric.""" + 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.""" -UnifiedMetricOrDict = Union[UnifiedMetric, UnifiedMetricDict] +LossExampleOrDict = Union[LossExample, LossExampleDict] -class EvaluationMetric(_common.BaseModel): - """Represents an evaluation metric.""" +class LossCluster(_common.BaseModel): + """A semantic grouping of failures (e.g., "Hallucination of Action").""" - name: Optional[str] = Field( - default=None, description="""The resource name of the evaluation metric.""" + cluster_id: Optional[str] = Field( + default=None, + description="""Unique identifier for the loss cluster within the scope of the analysis result.""", ) - display_name: Optional[str] = Field( + taxonomy_entry: Optional[LossTaxonomyEntry] = Field( default=None, - description="""The user-friendly display name for the EvaluationMetric.""", + description="""The structured definition of the loss taxonomy for this cluster.""", ) - description: Optional[str] = Field( - default=None, description="""The description of the EvaluationMetric.""" + item_count: Optional[int] = Field( + default=None, + description="""The total number of EvaluationItems falling into this cluster.""", ) - metric: Optional[UnifiedMetric] = Field( + examples: Optional[list[LossExample]] = Field( default=None, - description="""The metric configuration of the evaluation metric.""", + description="""A list of examples that belong to this cluster. This links the cluster back to the specific EvaluationItems and Rubrics.""", ) -class EvaluationMetricDict(TypedDict, total=False): - """Represents an evaluation metric.""" - - name: Optional[str] - """The resource name of the evaluation metric.""" - - display_name: Optional[str] - """The user-friendly display name for the EvaluationMetric.""" +class LossClusterDict(TypedDict, total=False): + """A semantic grouping of failures (e.g., "Hallucination of Action").""" - description: Optional[str] - """The description of the EvaluationMetric.""" + cluster_id: Optional[str] + """Unique identifier for the loss cluster within the scope of the analysis result.""" - metric: Optional[UnifiedMetricDict] - """The metric configuration of the evaluation metric.""" + taxonomy_entry: Optional[LossTaxonomyEntryDict] + """The structured definition of the loss taxonomy for this cluster.""" + item_count: Optional[int] + """The total number of EvaluationItems falling into this cluster.""" -EvaluationMetricOrDict = Union[EvaluationMetric, EvaluationMetricDict] + examples: Optional[list[LossExampleDict]] + """A list of examples that belong to this cluster. This links the cluster back to the specific EvaluationItems and Rubrics.""" -class SamplingConfig(_common.BaseModel): - """Sampling config for a BigQuery request set.""" +LossClusterOrDict = Union[LossCluster, LossClusterDict] - 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 LossAnalysisResult(_common.BaseModel): + """The top-level result for loss analysis.""" -class SamplingConfigDict(TypedDict, total=False): - """Sampling config for a BigQuery request set.""" - - sampling_count: Optional[int] - """""" - - sampling_method: Optional[SamplingMethod] - """""" - - sampling_duration: Optional[str] - """""" - - -SamplingConfigOrDict = Union[SamplingConfig, SamplingConfigDict] - - -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( + config: Optional[LossAnalysisConfig] = Field( default=None, - description="""The column name of the rubrics in the BigQuery table. Used for EvaluationRun only.""", + description="""The configuration used to generate this analysis.""", ) - 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.""", + analysis_time: Optional[str] = Field( + default=None, description="""The timestamp when this analysis was performed.""" ) - sampling_config: Optional[SamplingConfig] = Field( - default=None, - description="""The sampling config for the BigQuery request set. Used for EvaluationRun only.""", + 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 BigQueryRequestSetDict(TypedDict, total=False): - """Represents a BigQuery request set.""" + _evals_visualization.display_loss_analysis_result(self) - uri: Optional[str] - """""" - prompt_column: Optional[str] - """The column name of the prompt in the BigQuery table. Used for EvaluationRun only.""" +class LossAnalysisResultDict(TypedDict, total=False): + """The top-level result for loss analysis.""" - rubrics_column: Optional[str] - """The column name of the rubrics in the BigQuery table. Used for EvaluationRun only.""" + config: Optional[LossAnalysisConfigDict] + """The configuration used to generate this analysis.""" - candidate_response_columns: Optional[dict[str, str]] - """The column name of the response candidates in the BigQuery table. Used for EvaluationRun only.""" + analysis_time: Optional[str] + """The timestamp when this analysis was performed.""" - sampling_config: Optional[SamplingConfigDict] - """The sampling config for the BigQuery request set. Used for EvaluationRun only.""" + clusters: Optional[list[LossClusterDict]] + """The list of identified loss clusters.""" -BigQueryRequestSetOrDict = Union[BigQueryRequestSet, BigQueryRequestSetDict] +LossAnalysisResultOrDict = Union[LossAnalysisResult, LossAnalysisResultDict] -class EvaluationRunDataSource(_common.BaseModel): - """Represents an evaluation run data source.""" +class EvaluationRunResults(_common.BaseModel): + """Represents the results of an evaluation run.""" - evaluation_set: Optional[str] = Field(default=None, description="""""") - bigquery_request_set: Optional[BigQueryRequestSet] = Field( - default=None, description="""""" + evaluation_set: Optional[str] = Field( + default=None, + description="""The evaluation set where item level results are stored.""", + ) + summary_metrics: Optional[SummaryMetric] = Field( + default=None, description="""The summary metrics for the evaluation run.""" + ) + loss_analysis_results: Optional[list[LossAnalysisResult]] = Field( + default=None, + description="""The loss analysis results for the evaluation run.""", + ) + red_teaming_analysis_results: Optional[list[RedTeamingAnalysisResult]] = Field( + default=None, description="""The Red Teaming analysis results.""" ) -class EvaluationRunDataSourceDict(TypedDict, total=False): - """Represents an evaluation run data source.""" +class EvaluationRunResultsDict(TypedDict, total=False): + """Represents the results of an evaluation run.""" evaluation_set: Optional[str] - """""" + """The evaluation set where item level results are stored.""" - bigquery_request_set: Optional[BigQueryRequestSetDict] - """""" + summary_metrics: Optional[SummaryMetricDict] + """The summary metrics for the evaluation run.""" + loss_analysis_results: Optional[list[LossAnalysisResultDict]] + """The loss analysis results for the evaluation run.""" -EvaluationRunDataSourceOrDict = Union[ - EvaluationRunDataSource, EvaluationRunDataSourceDict -] + red_teaming_analysis_results: Optional[list[RedTeamingAnalysisResultDict]] + """The Red Teaming analysis results.""" -class EvaluationRunMetric(_common.BaseModel): - """The metric used for evaluation run.""" +EvaluationRunResultsOrDict = Union[EvaluationRunResults, EvaluationRunResultsDict] - metric: Optional[str] = Field( - default=None, description="""The name of the metric.""" + +class EvalCaseMetricResult(_common.BaseModel): + """Evaluation result for a single evaluation case for a single metric.""" + + metric_name: Optional[str] = Field( + default=None, description="""Name of the metric.""" ) - metric_resource_name: Optional[str] = Field( + 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 resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""", + description="""The details of all the rubrics and their verdicts for rubric-based metrics.""", ) - metric_config: Optional[UnifiedMetric] = Field( - default=None, description="""The unified metric used for evaluation run.""" + raw_output: Optional[list[str]] = Field( + default=None, description="""Raw output of the metric.""" + ) + error_message: Optional[str] = Field( + default=None, description="""Error message for the metric.""" ) -class EvaluationRunMetricDict(TypedDict, total=False): - """The metric used for evaluation run.""" +class EvalCaseMetricResultDict(TypedDict, total=False): + """Evaluation result for a single evaluation case for a single metric.""" - metric: Optional[str] - """The name of the metric.""" + metric_name: Optional[str] + """Name of the metric.""" - metric_resource_name: Optional[str] - """The resource name of the metric definition. Example: projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""" + score: Optional[float] + """Score of the metric.""" - metric_config: Optional[UnifiedMetricDict] - """The unified metric used for evaluation run.""" + explanation: Optional[str] + """Explanation of the metric.""" + + rubric_verdicts: Optional[list[evals_types.RubricVerdict]] + """The details of all the rubrics and their verdicts for rubric-based metrics.""" + raw_output: Optional[list[str]] + """Raw output of the metric.""" -EvaluationRunMetricOrDict = Union[EvaluationRunMetric, EvaluationRunMetricDict] + error_message: Optional[str] + """Error message for the metric.""" -class EvaluationRunPromptTemplate(_common.BaseModel): - """Prompt template used for inference. +EvalCaseMetricResultOrDict = Union[EvalCaseMetricResult, EvalCaseMetricResultDict] - Only one of `prompt_template` or `gcs_uri` should be set. If both are - provided, an error will be raised. - """ - prompt_template: Optional[str] = Field( +class ResponseCandidateResult(_common.BaseModel): + """Aggregated metric results for a single response candidate.""" + + response_index: Optional[int] = 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="""Index of the response candidate this result pertains to.""", ) - gcs_uri: Optional[str] = Field( + metric_results: Optional[dict[str, EvalCaseMetricResult]] = Field( default=None, - description="""Prompt template stored in Cloud Storage. Format: - "gs://my-bucket/file-name.txt". Only one of `prompt_template` or `gcs_uri` - should be set.""", + description="""A dictionary of metric results for this response candidate, keyed by metric name.""", ) -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 ResponseCandidateResultDict(TypedDict, total=False): + """Aggregated metric results for a single response candidate.""" - 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.""" + response_index: Optional[int] + """Index of the response candidate this result pertains to.""" - 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.""" + metric_results: Optional[dict[str, EvalCaseMetricResultDict]] + """A dictionary of metric results for this response candidate, keyed by metric name.""" -EvaluationRunPromptTemplateOrDict = Union[ - EvaluationRunPromptTemplate, EvaluationRunPromptTemplateDict +ResponseCandidateResultOrDict = Union[ + ResponseCandidateResult, ResponseCandidateResultDict ] -class LossAnalysisConfig(_common.BaseModel): - """Configuration for the loss analysis job.""" +class EvalCaseResult(_common.BaseModel): + """Eval result for a single evaluation case.""" - 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.""", + eval_case_index: Optional[int] = Field( + default=None, description="""Index of the evaluation case.""" ) - max_top_cluster_count: Optional[int] = Field( + response_candidate_results: Optional[list[ResponseCandidateResult]] = 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="""A list of results, one for each response candidate of the EvalCase.""", ) -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 EvalCaseResultDict(TypedDict, total=False): + """Eval result for a single evaluation case.""" - 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.""" + eval_case_index: Optional[int] + """Index of the evaluation case.""" - 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.""" + response_candidate_results: Optional[list[ResponseCandidateResultDict]] + """A list of results, one for each response candidate of the EvalCase.""" -LossAnalysisConfigOrDict = Union[LossAnalysisConfig, LossAnalysisConfigDict] +EvalCaseResultOrDict = Union[EvalCaseResult, EvalCaseResultDict] -class EvaluationRunConfig(_common.BaseModel): - """The evaluation configuration used for the evaluation run.""" +class AggregatedMetricResult(_common.BaseModel): + """Evaluation result for a single metric for an evaluation dataset.""" - metrics: Optional[list[EvaluationRunMetric]] = Field( - default=None, - description="""The metrics to be calculated in the evaluation run.""", + metric_name: Optional[str] = Field( + default=None, description="""Name of the metric.""" ) - output_config: Optional[genai_types.OutputConfig] = Field( - default=None, description="""The output config for the evaluation run.""" + num_cases_total: Optional[int] = Field( + default=None, description="""Total number of cases in the dataset.""" ) - autorater_config: Optional[genai_types.AutoraterConfig] = Field( - default=None, - description="""The autorater config for the evaluation run. Not applicable for predefined metrics (PredefinedMetricSpec); the server uses its own model configuration for predefined metrics and this field is ignored.""", + num_cases_valid: Optional[int] = Field( + default=None, description="""Number of valid cases in the dataset.""" ) - prompt_template: Optional[EvaluationRunPromptTemplate] = Field( - default=None, description="""The prompt template used for inference.""" + num_cases_error: Optional[int] = Field( + default=None, description="""Number of cases with errors in the dataset.""" ) - loss_analysis_config: Optional[list[LossAnalysisConfig]] = Field( - default=None, - description="""Specifications for loss analysis. Each config specifies a metric and candidate to analyze for loss patterns.""", + mean_score: Optional[float] = Field( + default=None, description="""Mean score of the metric.""" ) - allow_cross_region_model: Optional[bool] = Field( + stdev_score: Optional[float] = Field( + default=None, description="""Standard deviation of the metric.""" + ) + pass_rate: Optional[float] = 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="""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 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.""" +class AggregatedMetricResultDict(TypedDict, total=False): + """Evaluation result for a single metric for an evaluation dataset.""" - output_config: Optional[genai_types.OutputConfig] - """The output config for the evaluation run.""" + metric_name: Optional[str] + """Name of the metric.""" - 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.""" + num_cases_total: Optional[int] + """Total number of cases in the dataset.""" - prompt_template: Optional[EvaluationRunPromptTemplateDict] - """The prompt template used for inference.""" + num_cases_valid: Optional[int] + """Number of valid cases in the dataset.""" - loss_analysis_config: Optional[list[LossAnalysisConfigDict]] - """Specifications for loss analysis. Each config specifies a metric and candidate to analyze for loss patterns.""" + num_cases_error: Optional[int] + """Number of cases with errors in the dataset.""" - 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.""" + mean_score: Optional[float] + """Mean score of the metric.""" + stdev_score: Optional[float] + """Standard deviation of the metric.""" -EvaluationRunConfigOrDict = Union[EvaluationRunConfig, EvaluationRunConfigDict] + 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 EvaluationRunAgentConfig(_common.BaseModel): - """Agent config for an evaluation run.""" +AggregatedMetricResultOrDict = Union[AggregatedMetricResult, AggregatedMetricResultDict] - developer_instruction: Optional[genai_types.Content] = Field( - default=None, description="""The developer instruction for the agent.""" + +class WinRateStats(_common.BaseModel): + """Statistics for win rates for a single metric.""" + + win_rates: Optional[list[float]] = Field( + default=None, + description="""Win rates for the metric, one for each candidate.""", ) - tools: Optional[list[genai_types.Tool]] = Field( - default=None, description="""The tools available to the agent.""" + tie_rate: Optional[float] = Field( + default=None, description="""Tie rate for the metric.""" ) -class EvaluationRunAgentConfigDict(TypedDict, total=False): - """Agent config for an evaluation run.""" - - developer_instruction: Optional[genai_types.Content] - """The developer instruction for the agent.""" +class WinRateStatsDict(TypedDict, total=False): + """Statistics for win rates for a single metric.""" - tools: Optional[list[genai_types.Tool]] - """The tools available to the agent.""" + win_rates: Optional[list[float]] + """Win rates for the metric, one for each candidate.""" + tie_rate: Optional[float] + """Tie rate for the metric.""" -EvaluationRunAgentConfigOrDict = Union[ - EvaluationRunAgentConfig, EvaluationRunAgentConfigDict -] +WinRateStatsOrDict = Union[WinRateStats, WinRateStatsDict] -class GeminiAgentConfig(_common.BaseModel): - """Config for scraping a Gemini Agent. - A Gemini Agent is a Vertex AI Agent resource scraped via the Vertex - Interactions API. - """ +class ResponseCandidate(_common.BaseModel): + """A model-generated content to the prompt.""" - gemini_agent: Optional[str] = Field( + response: Optional[genai_types.Content] = Field( default=None, - description="""The resource name of the Gemini Agent. - Format: `projects/{project}/locations/{location}/agents/{agent}`.""", + description="""The final model-generated response to the `prompt`.""", ) -class GeminiAgentConfigDict(TypedDict, total=False): - """Config for scraping a Gemini Agent. - - A Gemini Agent is a Vertex AI Agent resource scraped via the Vertex - Interactions API. - """ +class ResponseCandidateDict(TypedDict, total=False): + """A model-generated content to the prompt.""" - gemini_agent: Optional[str] - """The resource name of the Gemini Agent. - Format: `projects/{project}/locations/{location}/agents/{agent}`.""" + response: Optional[genai_types.Content] + """The final model-generated response to the `prompt`.""" -GeminiAgentConfigOrDict = Union[GeminiAgentConfig, GeminiAgentConfigDict] +ResponseCandidateOrDict = Union[ResponseCandidate, ResponseCandidateDict] -class AgentRunConfig(_common.BaseModel): - """Configuration for an Agent Run.""" +class InteractionsDataSource(_common.BaseModel): + """Source for populating agent data from an Interactions API interaction.""" - session_input: Optional[evals_types.SessionInput] = Field( - default=None, description="""The session input to get agent running results.""" - ) - agent_engine: Optional[str] = Field( - default=None, description="""The resource name of the Agent Engine.""" - ) - user_simulator_config: Optional[evals_types.UserSimulatorConfig] = Field( + gemini_agent_config: Optional[GeminiAgentConfig] = 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 Gemini Agent (Vertex AI Agent resource) that produced the + interaction.""", ) - gemini_agent_config: Optional[GeminiAgentConfig] = Field( + interaction: 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 interaction to evaluate. Required by the backend. + Format: + `projects/{project}/locations/{location}/interactions/{interaction}`.""", ) -class AgentRunConfigDict(TypedDict, total=False): - """Configuration for an Agent Run.""" - - session_input: Optional[evals_types.SessionInput] - """The session input to get agent running results.""" - - agent_engine: Optional[str] - """The resource name of the Agent Engine.""" - - user_simulator_config: Optional[evals_types.UserSimulatorConfig] - """Used for multi-turn agent run. - Contains configuration for a user simulator that - uses an LLM to generate messages on behalf of the user.""" +class InteractionsDataSourceDict(TypedDict, total=False): + """Source for populating agent data from an Interactions API interaction.""" 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.""" + """The Gemini Agent (Vertex AI Agent resource) that produced the + interaction.""" + + interaction: Optional[str] + """The interaction to evaluate. Required by the backend. + Format: + `projects/{project}/locations/{location}/interactions/{interaction}`.""" -AgentRunConfigOrDict = Union[AgentRunConfig, AgentRunConfigDict] +InteractionsDataSourceOrDict = Union[InteractionsDataSource, InteractionsDataSourceDict] -class EvaluationRunInferenceConfig(_common.BaseModel): - """Configuration that describes an agent.""" +class EvalCase(_common.BaseModel): + """A comprehensive representation of a GenAI interaction for evaluation.""" - agent_config: Optional[EvaluationRunAgentConfig] = Field( - default=None, description="""The agent config.""" + prompt: Optional[genai_types.Content] = Field( + default=None, description="""The most recent user message (current input).""" ) - model: Optional[str] = Field( + responses: Optional[list[ResponseCandidate]] = 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`).""", + 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.""", ) - prompt_template: Optional[EvaluationRunPromptTemplate] = Field( - default=None, description="""The prompt template used for inference.""" + 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.""", ) - agent_run_config: Optional[AgentRunConfig] = 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="""Configuration for Agent Run in evaluation management service.""", + description="""List of all prior messages in the conversation (chat history).""", ) - agent_configs: Optional[dict[str, evals_types.AgentConfig]] = Field( + rubric_groups: Optional[dict[str, "RubricGroup"]] = Field( default=None, - description="""A map of agent IDs to their respective agent config.""", + description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", + ) + eval_case_id: Optional[str] = Field( + default=None, description="""Unique identifier for the evaluation case.""" + ) + intermediate_events: Optional[list[evals_types.Event]] = Field( + default=None, + description="""Intermediate events of a single turn in an agent run or intermediate events of the last turn for multi-turn an agent run.""", + ) + agent_info: Optional[evals_types.AgentInfo] = Field( + default=None, + description="""The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""", + ) + agent_data: Optional[evals_types.AgentData] = Field( + default=None, description="""The agent data of the agent under evaluation.""" + ) + user_scenario: Optional[evals_types.UserScenario] = Field( + default=None, description="""The user scenario for the evaluation case.""" + ) + interactions_data_source: Optional[InteractionsDataSource] = Field( + default=None, + description="""Source for populating agent data from an Interactions API interaction. When set, the backend fetches the interaction (and the Gemini Agent config) and parses it into agent data for evaluation; agent_data must not also be set.""", ) + # Allow extra fields to support custom metric prompts and stay backward compatible. + model_config = ConfigDict(frozen=True, extra="allow") -class EvaluationRunInferenceConfigDict(TypedDict, total=False): - """Configuration that describes an agent.""" +class EvalCaseDict(TypedDict, total=False): + """A comprehensive representation of a GenAI interaction for evaluation.""" - agent_config: Optional[EvaluationRunAgentConfigDict] - """The agent config.""" + prompt: Optional[genai_types.Content] + """The most recent user message (current input).""" - 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`).""" + 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.""" - prompt_template: Optional[EvaluationRunPromptTemplateDict] - """The prompt template used for inference.""" + reference: Optional[ResponseCandidateDict] + """User-provided, golden reference model reply to prompt in context of chat history; Reference for last response in a conversation.""" - agent_run_config: Optional[AgentRunConfigDict] - """Configuration for Agent Run in evaluation management service.""" + system_instruction: Optional[genai_types.Content] + """System instruction for the model.""" - agent_configs: Optional[dict[str, evals_types.AgentConfig]] - """A map of agent IDs to their respective agent config.""" + conversation_history: Optional[list[evals_types.Message]] + """List of all prior messages in the conversation (chat history).""" + rubric_groups: Optional[dict[str, "RubricGroupDict"]] + """Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""" -EvaluationRunInferenceConfigOrDict = Union[ - EvaluationRunInferenceConfig, EvaluationRunInferenceConfigDict -] + eval_case_id: Optional[str] + """Unique identifier for the evaluation case.""" + intermediate_events: Optional[list[evals_types.Event]] + """Intermediate events of a single turn in an agent run or intermediate events of the last turn for multi-turn an agent run.""" -class VulnerableTool(_common.BaseModel): - """A tool considered high risk for prompt injection.""" + agent_info: Optional[evals_types.AgentInfo] + """The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""" - tool_name: Optional[str] = Field( - default=None, - description="""Optional. The name of the vulnerable function/tool (e.g., "search_flights").""", - ) - json_paths: Optional[list[str]] = Field( - default=None, - description="""Optional. JSON Paths within the tool's FunctionResponse where malicious content could be injected.""", - ) - - -class VulnerableToolDict(TypedDict, total=False): - """A tool considered high risk for prompt injection.""" + agent_data: Optional[evals_types.AgentData] + """The agent data of the agent under evaluation.""" - tool_name: Optional[str] - """Optional. The name of the vulnerable function/tool (e.g., "search_flights").""" + user_scenario: Optional[evals_types.UserScenario] + """The user scenario for the evaluation case.""" - json_paths: Optional[list[str]] - """Optional. JSON Paths within the tool's FunctionResponse where malicious content could be injected.""" + 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.""" -VulnerableToolOrDict = Union[VulnerableTool, VulnerableToolDict] +EvalCaseOrDict = Union[EvalCase, EvalCaseDict] -class RedTeamingAnalysisConfig(_common.BaseModel): - """Configuration for the automated Agent Red Teaming analysis.""" +class EvaluationDataset(_common.BaseModel): + """The dataset used for evaluation.""" - attack_categories: Optional[list[str]] = Field( + 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="""Optional. Specific attack categories to test against.""", + description="""The evaluation dataset in the form of a Pandas DataFrame.""", ) - vulnerable_tools: Optional[list[VulnerableTool]] = Field( + candidate_name: Optional[str] = Field( default=None, - description="""Optional. Manually defined vulnerable tools and their injection paths.""", + description="""The name of the candidate model or agent for this evaluation dataset.""", + ) + evaluation_set: Optional[str] = Field( + default=None, + description="""Optional. The full path of the EvaluationSet resource created, which contains the evaluation results. Format: `projects/{project}/locations/{location}/evaluationSets/{evaluation_set_id}`""", + ) + is_unlimited_size: Optional[bool] = Field( + default=None, + description="""Optional. Flag to validate evaluation dataset size.""", ) + @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 RedTeamingAnalysisConfigDict(TypedDict, total=False): - """Configuration for the automated Agent Red Teaming analysis.""" + @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 - attack_categories: Optional[list[str]] - """Optional. Specific attack categories to test against.""" + formats = [] + requests = [] + responses = [] + system_instructions = [] - vulnerable_tools: Optional[list[VulnerableToolDict]] - """Optional. Manually defined vulnerable tools and their injection paths.""" + for case in cases: + gcs_utils = _gcs_utils.GcsUtils( + case.api_client._api_client if case.api_client else None + ) + # Associate "observability" data format for given sources + formats.append("observability") -RedTeamingAnalysisConfigOrDict = Union[ - RedTeamingAnalysisConfig, RedTeamingAnalysisConfigDict -] + # 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) -class AnalysisConfig(_common.BaseModel): - """Configuration for an analysis to be performed on an evaluation run.""" + # System instruction source + system_instruction_data = "" + if case.system_instruction_src is not None: + system_instruction_data = gcs_utils.read_file_contents( + case.system_instruction_src + ) + system_instructions.append(system_instruction_data) - analysis_name: Optional[str] = Field( - default=None, description="""Optional. A name for this analysis.""" - ) - red_teaming_analysis_config: Optional[RedTeamingAnalysisConfig] = Field( - default=None, - description="""Configuration for the automated Agent Red Teaming analysis.""", - ) + eval_dataset_df = pd.DataFrame( + { + "format": formats, + "request": requests, + "response": responses, + "system_instruction": system_instructions, + } + ) + except ImportError as e: + raise ImportError("Pandas DataFrame library is required.") from e -class AnalysisConfigDict(TypedDict, total=False): - """Configuration for an analysis to be performed on an evaluation run.""" + return EvaluationDataset(eval_dataset_df=eval_dataset_df) - analysis_name: Optional[str] - """Optional. A name for this analysis.""" + def show(self) -> None: + """Shows the evaluation dataset.""" + from .. import _evals_visualization - red_teaming_analysis_config: Optional[RedTeamingAnalysisConfigDict] - """Configuration for the automated Agent Red Teaming analysis.""" + _evals_visualization.display_evaluation_dataset(self) -AnalysisConfigOrDict = Union[AnalysisConfig, AnalysisConfigDict] +class EvaluationDatasetDict(TypedDict, total=False): + """The dataset used for evaluation.""" + bigquery_source: Optional[genai_types.BigQuerySource] + """The BigQuery source for the evaluation dataset.""" -class CreateEvaluationRunConfig(_common.BaseModel): - """Config to create an evaluation run.""" + gcs_source: Optional[genai_types.GcsSource] + """The GCS source for the evaluation dataset.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - allow_cross_region_model: Optional[bool] = Field( - default=None, - description="""Allows the evaluation run to use cross region models. When this - flag is set, the service may route traffic to other regions if a model is - unavailable in the current region (e.g., to a `global`endpoint). If a - fully-qualified model endpoint resource name with a different region than - the run location is provided elsewhere in the run config, this flag must - be set to true or the request will fail.""", - ) + eval_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 CreateEvaluationRunConfigDict(TypedDict, total=False): - """Config to create an evaluation run.""" + candidate_name: Optional[str] + """The name of the candidate model or agent for this evaluation dataset.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + evaluation_set: Optional[str] + """Optional. The full path of the EvaluationSet resource created, which contains the evaluation results. Format: `projects/{project}/locations/{location}/evaluationSets/{evaluation_set_id}`""" - 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.""" + is_unlimited_size: Optional[bool] + """Optional. Flag to validate evaluation dataset size.""" -CreateEvaluationRunConfigOrDict = Union[ - CreateEvaluationRunConfig, CreateEvaluationRunConfigDict -] +EvaluationDatasetOrDict = Union[EvaluationDataset, EvaluationDatasetDict] -class _CreateEvaluationRunParameters(_common.BaseModel): - """Represents a job that creates an evaluation run.""" +class EvaluationRunMetadata(_common.BaseModel): + """Metadata for an evaluation run.""" - name: Optional[str] = Field(default=None, description="""""") - display_name: Optional[str] = Field(default=None, description="""""") - data_source: Optional[EvaluationRunDataSource] = Field( - default=None, description="""""" - ) - evaluation_config: Optional[EvaluationRunConfig] = Field( - default=None, description="""""" + candidate_names: Optional[list[str]] = Field( + default=None, + description="""Name of the candidate(s) being evaluated in the evaluation run.""", ) - labels: Optional[dict[str, str]] = Field(default=None, description="""""") - inference_configs: Optional[dict[str, EvaluationRunInferenceConfig]] = Field( - default=None, description="""""" + dataset_name: Optional[str] = Field( + default=None, + description="""Name of the evaluation dataset used for the evaluation run.""", ) - config: Optional[CreateEvaluationRunConfig] = Field( - default=None, description="""""" + dataset_id: Optional[str] = Field( + default=None, + description="""Unique identifier for the evaluation dataset used for the evaluation run.""", ) - analysis_configs: Optional[list[AnalysisConfig]] = Field( - default=None, description="""""" + creation_timestamp: Optional[datetime.datetime] = Field( + default=None, description="""Creation timestamp of the evaluation run.""" ) -class _CreateEvaluationRunParametersDict(TypedDict, total=False): - """Represents a job that creates an evaluation run.""" - - name: Optional[str] - """""" - - display_name: Optional[str] - """""" - - data_source: Optional[EvaluationRunDataSourceDict] - """""" - - evaluation_config: Optional[EvaluationRunConfigDict] - """""" +class EvaluationRunMetadataDict(TypedDict, total=False): + """Metadata for an evaluation run.""" - labels: Optional[dict[str, str]] - """""" + candidate_names: Optional[list[str]] + """Name of the candidate(s) being evaluated in the evaluation run.""" - inference_configs: Optional[dict[str, EvaluationRunInferenceConfigDict]] - """""" + dataset_name: Optional[str] + """Name of the evaluation dataset used for the evaluation run.""" - config: Optional[CreateEvaluationRunConfigDict] - """""" + dataset_id: Optional[str] + """Unique identifier for the evaluation dataset used for the evaluation run.""" - analysis_configs: Optional[list[AnalysisConfigDict]] - """""" + creation_timestamp: Optional[datetime.datetime] + """Creation timestamp of the evaluation run.""" -_CreateEvaluationRunParametersOrDict = Union[ - _CreateEvaluationRunParameters, _CreateEvaluationRunParametersDict -] +EvaluationRunMetadataOrDict = Union[EvaluationRunMetadata, EvaluationRunMetadataDict] -class SummaryMetric(_common.BaseModel): - """Represents a summary metric for an evaluation run.""" +class EvaluationResult(_common.BaseModel): + """Result of an evaluation run for an evaluation dataset.""" - metrics: Optional[dict[str, Any]] = Field( - default=None, description="""Map of metric name to metric value.""" + eval_case_results: Optional[list[EvalCaseResult]] = Field( + default=None, + description="""A list of evaluation results for each evaluation case.""", ) - total_items: Optional[int] = Field( - default=None, description="""The total number of items that were evaluated.""" + summary_metrics: Optional[list[AggregatedMetricResult]] = Field( + default=None, + description="""A list of summary-level evaluation results for each metric.""", ) - failed_items: Optional[int] = Field( - default=None, description="""The number of items that failed to be evaluated.""" + 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 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.""" - - failed_items: Optional[int] - """The number of items that failed to be evaluated.""" - + Args: + candidate_names: list of names for the evaluated candidates, used in + comparison reports. + """ + from .. import _evals_visualization -SummaryMetricOrDict = Union[SummaryMetric, SummaryMetricDict] + _evals_visualization.display_evaluation_result(self, candidate_names) -class AttackCategoryResult(_common.BaseModel): - """The red teaming outcome for a specific attack category.""" +class EvaluationResultDict(TypedDict, total=False): + """Result of an evaluation run for an evaluation dataset.""" - 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.""" - ) + eval_case_results: Optional[list[EvalCaseResultDict]] + """A list of evaluation results for each evaluation case.""" + summary_metrics: Optional[list[AggregatedMetricResultDict]] + """A list of summary-level evaluation results for each metric.""" -class AttackCategoryResultDict(TypedDict, total=False): - """The red teaming outcome for a specific attack category.""" + win_rates: Optional[dict[str, WinRateStatsDict]] + """A dictionary of win rates for each metric, only populated for multi-response evaluation runs.""" - attack_category: Optional[str] - """The category of the attack evaluated.""" + evaluation_dataset: Optional[list[EvaluationDatasetDict]] + """The input evaluation dataset(s) for the evaluation run.""" - attack_success_rate: Optional[float] - """The ratio of successful attacks given a fixed budget.""" + metadata: Optional[EvaluationRunMetadataDict] + """Metadata for the evaluation run.""" - vulnerability_insight: Optional[str] - """Insights into why an attack succeeded or failed.""" + agent_info: Optional[evals_types.AgentInfo] + """The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""" -AttackCategoryResultOrDict = Union[AttackCategoryResult, AttackCategoryResultDict] +EvaluationResultOrDict = Union[EvaluationResult, EvaluationResultDict] -class RedTeamingAnalysisResult(_common.BaseModel): - """The top-level result for Red Teaming analysis.""" +class EvaluationRun(_common.BaseModel): + """Represents an evaluation run.""" - config: Optional[RedTeamingAnalysisConfig] = Field( + 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 configuration used to generate this analysis.""", + description="""The parsed EvaluationItem results for the evaluation run. This is only populated when include_evaluation_items is set to True.""", ) - analysis_time: Optional[str] = Field( - default=None, description="""The timestamp when this analysis was performed.""" + evaluation_config: Optional[EvaluationRunConfig] = Field( + default=None, description="""The evaluation config for the evaluation run.""" ) - category_results: Optional[list[AttackCategoryResult]] = Field( - default=None, description="""Detailed results by attack category.""" + 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.""", ) + # TODO(b/448806531): Remove all the overridden _from_response methods once the + # ticket is resolved and published. + @classmethod + def _from_response( + cls: typing.Type["EvaluationRun"], + *, + response: dict[str, object], + kwargs: dict[str, object], + ) -> "EvaluationRun": + """Converts a dictionary response into a EvaluationRun object.""" -class RedTeamingAnalysisResultDict(TypedDict, total=False): - """The top-level result for Red Teaming analysis.""" - - config: Optional[RedTeamingAnalysisConfigDict] - """The configuration used to generate this analysis.""" + snaked_response = _camel_key_to_snake(response) - analysis_time: Optional[str] - """The timestamp when this analysis was performed.""" + evaluation_run_results = response.get("evaluation_run_results") - category_results: Optional[list[AttackCategoryResultDict]] - """Detailed results by attack category.""" + 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 -RedTeamingAnalysisResultOrDict = Union[ - RedTeamingAnalysisResult, RedTeamingAnalysisResultDict -] + 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 LossTaxonomyEntry(_common.BaseModel): - """A specific entry in the loss pattern taxonomy.""" +class EvaluationRunDict(TypedDict, total=False): + """Represents an evaluation run.""" - 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." """, - ) + name: Optional[str] + """""" + display_name: Optional[str] + """""" -class LossTaxonomyEntryDict(TypedDict, total=False): - """A specific entry in the loss pattern taxonomy.""" + metadata: Optional[dict[str, Any]] + """""" - l1_category: Optional[str] - """The primary category of the loss (e.g., "Hallucination", "Tool Calling").""" + create_time: Optional[datetime.datetime] + """""" - l2_category: Optional[str] - """The secondary category of the loss (e.g., "Hallucination of Action", "Incorrect Tool Selection").""" + completion_time: Optional[datetime.datetime] + """""" - description: Optional[str] - """A detailed description of this loss pattern. Example: "The agent verbally confirms an action without executing the tool." """ + state: Optional[EvaluationRunState] + """""" + evaluation_set_snapshot: Optional[str] + """""" -LossTaxonomyEntryOrDict = Union[LossTaxonomyEntry, LossTaxonomyEntryDict] + error: Optional[genai_types.GoogleRpcStatus] + """""" + data_source: Optional[EvaluationRunDataSourceDict] + """""" -class FailedRubric(_common.BaseModel): - """A specific failed rubric and the associated analysis.""" + evaluation_run_results: Optional[EvaluationRunResultsDict] + """The evaluation run formatted results.""" - 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.""", - ) + evaluation_item_results: Optional[EvaluationResultDict] + """The parsed EvaluationItem results for the evaluation run. This is only populated when include_evaluation_items is set to True.""" + evaluation_config: Optional[EvaluationRunConfigDict] + """The evaluation config for the evaluation run.""" -class FailedRubricDict(TypedDict, total=False): - """A specific failed rubric and the associated analysis.""" + inference_configs: Optional[dict[str, EvaluationRunInferenceConfigDict]] + """The inference configs for the evaluation run.""" - rubric_id: Optional[str] - """The unique ID of the rubric (if available from the metric source).""" + labels: Optional[dict[str, str]] + """""" - classification_rationale: Optional[str] - """The rationale provided by the Loss Analysis Classifier for why this failure maps to this specific Loss Cluster.""" + analysis_configs: Optional[list[AnalysisConfigDict]] + """The analysis configurations for the evaluation run.""" -FailedRubricOrDict = Union[FailedRubric, FailedRubricDict] +EvaluationRunOrDict = Union[EvaluationRun, EvaluationRunDict] -class LossExample(_common.BaseModel): - """A specific example of a loss pattern.""" +class CreateEvaluationSetConfig(_common.BaseModel): + """Config to create an evaluation set.""" - 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class LossExampleDict(TypedDict, total=False): - """A specific example of a loss pattern.""" +class CreateEvaluationSetConfigDict(TypedDict, total=False): + """Config to create an evaluation set.""" - evaluation_item: Optional[str] - """Reference to the persisted EvalItem resource name. Format: projects/.../locations/.../evaluationItems/{item_id}.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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.""" +CreateEvaluationSetConfigOrDict = Union[ + CreateEvaluationSetConfig, CreateEvaluationSetConfigDict +] -LossExampleOrDict = Union[LossExample, LossExampleDict] +class _CreateEvaluationSetParameters(_common.BaseModel): + """Represents a job that creates an evaluation set.""" + evaluation_items: Optional[list[str]] = Field(default=None, description="""""") + display_name: Optional[str] = Field(default=None, description="""""") + config: Optional[CreateEvaluationSetConfig] = Field( + default=None, description="""""" + ) -class LossCluster(_common.BaseModel): - """A semantic grouping of failures (e.g., "Hallucination of Action").""" - cluster_id: Optional[str] = Field( - default=None, - description="""Unique identifier for the loss cluster within the scope of the analysis result.""", +class _CreateEvaluationSetParametersDict(TypedDict, total=False): + """Represents a job that creates an evaluation set.""" + + evaluation_items: Optional[list[str]] + """""" + + display_name: Optional[str] + """""" + + config: Optional[CreateEvaluationSetConfigDict] + """""" + + +_CreateEvaluationSetParametersOrDict = Union[ + _CreateEvaluationSetParameters, _CreateEvaluationSetParametersDict +] + + +class EvaluationSet(_common.BaseModel): + """Represents an evaluation set.""" + + name: Optional[str] = Field( + default=None, description="""The resource name of the evaluation set.""" ) - taxonomy_entry: Optional[LossTaxonomyEntry] = Field( - default=None, - description="""The structured definition of the loss taxonomy for this cluster.""", + display_name: Optional[str] = Field( + default=None, description="""The display name of the evaluation set.""" ) - item_count: Optional[int] = Field( + evaluation_items: Optional[list[str]] = Field( default=None, - description="""The total number of EvaluationItems falling into this cluster.""", + description="""The EvaluationItems that are part of this dataset.""", ) - 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.""", + create_time: Optional[datetime.datetime] = Field( + default=None, description="""The create time of the evaluation set.""" + ) + update_time: Optional[datetime.datetime] = Field( + default=None, description="""The update time of the evaluation set.""" + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, description="""The metadata of the evaluation set.""" ) -class LossClusterDict(TypedDict, total=False): - """A semantic grouping of failures (e.g., "Hallucination of Action").""" +class EvaluationSetDict(TypedDict, total=False): + """Represents an evaluation set.""" - cluster_id: Optional[str] - """Unique identifier for the loss cluster within the scope of the analysis result.""" + name: Optional[str] + """The resource name of the evaluation set.""" - taxonomy_entry: Optional[LossTaxonomyEntryDict] - """The structured definition of the loss taxonomy for this cluster.""" + display_name: Optional[str] + """The display name of the evaluation set.""" - item_count: Optional[int] - """The total number of EvaluationItems falling into this cluster.""" + evaluation_items: Optional[list[str]] + """The EvaluationItems that are part of this dataset.""" - examples: Optional[list[LossExampleDict]] - """A list of examples that belong to this cluster. This links the cluster back to the specific EvaluationItems and Rubrics.""" + create_time: Optional[datetime.datetime] + """The create time of the evaluation set.""" + update_time: Optional[datetime.datetime] + """The update time of the evaluation set.""" -LossClusterOrDict = Union[LossCluster, LossClusterDict] + metadata: Optional[dict[str, Any]] + """The metadata of the evaluation set.""" -class LossAnalysisResult(_common.BaseModel): - """The top-level result for loss analysis.""" +EvaluationSetOrDict = Union[EvaluationSet, EvaluationSetDict] - 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.""" + +class DeleteEvaluationMetricConfig(_common.BaseModel): + """Config for deleting an evaluation metric.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - 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 BleuInstance(_common.BaseModel): + """Bleu instance.""" - 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.""" + prediction: Optional[str] = Field( + default=None, description="""Required. Output of the evaluated model.""" ) - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( + reference: Optional[str] = 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.""" + description="""Required. Ground truth used to compare against the prediction.""", ) -class EvalCaseMetricResultDict(TypedDict, total=False): - """Evaluation result for a single evaluation case for a single metric.""" +class BleuInstanceDict(TypedDict, total=False): + """Bleu instance.""" - metric_name: Optional[str] - """Name of the metric.""" + prediction: Optional[str] + """Required. Output of the evaluated model.""" - score: Optional[float] - """Score of the metric.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" - explanation: Optional[str] - """Explanation of the metric.""" - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] - """The details of all the rubrics and their verdicts for rubric-based metrics.""" +BleuInstanceOrDict = Union[BleuInstance, BleuInstanceDict] - raw_output: Optional[list[str]] - """Raw output of the metric.""" - error_message: Optional[str] - """Error message for the metric.""" +class BleuInput(_common.BaseModel): + instances: Optional[list[BleuInstance]] = Field( + default=None, description="""Required. Repeated bleu instances.""" + ) + metric_spec: Optional[genai_types.BleuSpec] = Field( + default=None, description="""Required. Spec for bleu score metric.""" + ) -EvalCaseMetricResultOrDict = Union[EvalCaseMetricResult, EvalCaseMetricResultDict] +class BleuInputDict(TypedDict, total=False): -class ResponseCandidateResult(_common.BaseModel): - """Aggregated metric results for a single response candidate.""" + instances: Optional[list[BleuInstanceDict]] + """Required. Repeated bleu instances.""" - response_index: Optional[int] = Field( - default=None, - description="""Index of the response candidate this result pertains to.""", + 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.""" ) - metric_results: Optional[dict[str, EvalCaseMetricResult]] = Field( + reference: Optional[str] = Field( default=None, - description="""A dictionary of metric results for this response candidate, keyed by metric name.""", + description="""Required. Ground truth used to compare against the prediction.""", ) -class ResponseCandidateResultDict(TypedDict, total=False): - """Aggregated metric results for a single response candidate.""" - - response_index: Optional[int] - """Index of the response candidate this result pertains to.""" +class ExactMatchInstanceDict(TypedDict, total=False): + """Exact match instance.""" - metric_results: Optional[dict[str, EvalCaseMetricResultDict]] - """A dictionary of metric results for this response candidate, keyed by metric name.""" + prediction: Optional[str] + """Required. Output of the evaluated model.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" -ResponseCandidateResultOrDict = Union[ - ResponseCandidateResult, ResponseCandidateResultDict -] +ExactMatchInstanceOrDict = Union[ExactMatchInstance, ExactMatchInstanceDict] -class EvalCaseResult(_common.BaseModel): - """Eval result for a single evaluation case.""" - eval_case_index: Optional[int] = Field( - default=None, description="""Index of the evaluation case.""" - ) - response_candidate_results: Optional[list[ResponseCandidateResult]] = Field( - default=None, - description="""A list of results, one for each response candidate of the EvalCase.""", - ) +class ExactMatchSpec(_common.BaseModel): + """Spec for exact match metric.""" + pass -class EvalCaseResultDict(TypedDict, total=False): - """Eval result for a single evaluation case.""" - eval_case_index: Optional[int] - """Index of the evaluation case.""" +class ExactMatchSpecDict(TypedDict, total=False): + """Spec for exact match metric.""" - response_candidate_results: Optional[list[ResponseCandidateResultDict]] - """A list of results, one for each response candidate of the EvalCase.""" + pass -EvalCaseResultOrDict = Union[EvalCaseResult, EvalCaseResultDict] +ExactMatchSpecOrDict = Union[ExactMatchSpec, ExactMatchSpecDict] -class AggregatedMetricResult(_common.BaseModel): - """Evaluation result for a single metric for an evaluation dataset.""" +class ExactMatchInput(_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.""" - ) - 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.""" + instances: Optional[list[ExactMatchInstance]] = Field( + default=None, description="""Required. Repeated exact match instances.""" ) - 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.""", + metric_spec: Optional[ExactMatchSpec] = Field( + default=None, description="""Required. Spec for exact match metric.""" ) - # Allow extra fields to support custom aggregation stats. - model_config = ConfigDict(extra="allow") +class ExactMatchInputDict(TypedDict, total=False): -class AggregatedMetricResultDict(TypedDict, total=False): - """Evaluation result for a single metric for an evaluation dataset.""" + instances: Optional[list[ExactMatchInstanceDict]] + """Required. Repeated exact match instances.""" - metric_name: Optional[str] - """Name of the metric.""" + metric_spec: Optional[ExactMatchSpecDict] + """Required. Spec for exact match metric.""" - num_cases_total: Optional[int] - """Total number of cases in the dataset.""" - num_cases_valid: Optional[int] - """Number of valid cases in the dataset.""" +ExactMatchInputOrDict = Union[ExactMatchInput, ExactMatchInputDict] - num_cases_error: Optional[int] - """Number of cases with errors in the dataset.""" - mean_score: Optional[float] - """Mean score of the metric.""" +class RougeInstance(_common.BaseModel): + """Rouge instance.""" - stdev_score: Optional[float] - """Standard deviation of the 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.""", + ) - 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 RougeInstanceDict(TypedDict, total=False): + """Rouge instance.""" -AggregatedMetricResultOrDict = Union[AggregatedMetricResult, AggregatedMetricResultDict] + prediction: Optional[str] + """Required. Output of the evaluated model.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" -class WinRateStats(_common.BaseModel): - """Statistics for win rates for a single metric.""" - win_rates: Optional[list[float]] = Field( - default=None, - description="""Win rates for the metric, one for each candidate.""", +RougeInstanceOrDict = Union[RougeInstance, RougeInstanceDict] + + +class RougeInput(_common.BaseModel): + """Rouge input.""" + + instances: Optional[list[RougeInstance]] = Field( + default=None, description="""Required. Repeated rouge instances.""" ) - tie_rate: Optional[float] = Field( - default=None, description="""Tie rate for the metric.""" + metric_spec: Optional[genai_types.RougeSpec] = Field( + default=None, description="""Required. Spec for rouge score metric.""" ) -class WinRateStatsDict(TypedDict, total=False): - """Statistics for win rates for a single metric.""" +class RougeInputDict(TypedDict, total=False): + """Rouge input.""" - win_rates: Optional[list[float]] - """Win rates for the metric, one for each candidate.""" + instances: Optional[list[RougeInstanceDict]] + """Required. Repeated rouge instances.""" - tie_rate: Optional[float] - """Tie rate for the metric.""" + metric_spec: Optional[genai_types.RougeSpec] + """Required. Spec for rouge score metric.""" -WinRateStatsOrDict = Union[WinRateStats, WinRateStatsDict] +RougeInputOrDict = Union[RougeInput, RougeInputDict] -class ResponseCandidate(_common.BaseModel): - """A model-generated content to the prompt.""" +class ContentMap(_common.BaseModel): + """Map of placeholder in metric prompt template to contents of model input.""" - response: Optional[genai_types.Content] = Field( - default=None, - description="""The final model-generated response to the `prompt`.""", + values: Optional[dict[str, "ContentMapContents"]] = Field( + default=None, description="""Map of placeholder to contents.""" ) -class ResponseCandidateDict(TypedDict, total=False): - """A model-generated content to the prompt.""" +class ContentMapDict(TypedDict, total=False): + """Map of placeholder in metric prompt template to contents of model input.""" - response: Optional[genai_types.Content] - """The final model-generated response to the `prompt`.""" + values: Optional[dict[str, "ContentMapContents"]] + """Map of placeholder to contents.""" -ResponseCandidateOrDict = Union[ResponseCandidate, ResponseCandidateDict] +ContentMapOrDict = Union[ContentMap, ContentMapDict] -class InteractionsDataSource(_common.BaseModel): - """Source for populating agent data from an Interactions API interaction.""" +class PointwiseMetricInstance(_common.BaseModel): + """Pointwise metric instance.""" - gemini_agent_config: Optional[GeminiAgentConfig] = Field( + json_instance: Optional[str] = Field( default=None, - description="""The Gemini Agent (Vertex AI Agent resource) that produced the - interaction.""", + description="""Instance specified as a json string. String key-value pairs are expected in the json_instance to render PointwiseMetricSpec.instance_prompt_template.""", ) - interaction: Optional[str] = Field( + content_map_instance: Optional[ContentMap] = Field( default=None, - description="""The interaction to evaluate. Required by the backend. - Format: - `projects/{project}/locations/{location}/interactions/{interaction}`.""", + 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.""", ) -class InteractionsDataSourceDict(TypedDict, total=False): - """Source for populating agent data from an Interactions API interaction.""" +class PointwiseMetricInstanceDict(TypedDict, total=False): + """Pointwise metric instance.""" - gemini_agent_config: Optional[GeminiAgentConfigDict] - """The Gemini Agent (Vertex AI Agent resource) that produced the - interaction.""" + 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.""" - interaction: Optional[str] - """The interaction to evaluate. Required by the backend. - Format: - `projects/{project}/locations/{location}/interactions/{interaction}`.""" + 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.""" -InteractionsDataSourceOrDict = Union[InteractionsDataSource, InteractionsDataSourceDict] +PointwiseMetricInstanceOrDict = Union[ + PointwiseMetricInstance, PointwiseMetricInstanceDict +] -class EvalCase(_common.BaseModel): - """A comprehensive representation of a GenAI interaction for evaluation.""" +class PointwiseMetricInput(_common.BaseModel): + """Pointwise metric 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.""" + instance: Optional[PointwiseMetricInstance] = Field( + default=None, description="""Required. Pointwise metric instance.""" ) - conversation_history: Optional[list[evals_types.Message]] = Field( - default=None, - description="""List of all prior messages in the conversation (chat history).""", + metric_spec: Optional[genai_types.PointwiseMetricSpec] = Field( + default=None, description="""Required. Spec for pointwise metric.""" ) - rubric_groups: Optional[dict[str, "RubricGroup"]] = Field( + + +class PointwiseMetricInputDict(TypedDict, total=False): + """Pointwise metric input.""" + + instance: Optional[PointwiseMetricInstanceDict] + """Required. Pointwise metric instance.""" + + metric_spec: Optional[genai_types.PointwiseMetricSpec] + """Required. Spec for pointwise metric.""" + + +PointwiseMetricInputOrDict = Union[PointwiseMetricInput, PointwiseMetricInputDict] + + +class PairwiseMetricInstance(_common.BaseModel): + """Pairwise metric instance.""" + + json_instance: Optional[str] = Field( default=None, - description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", + description="""Instance specified as a json string. String key-value pairs are expected in the json_instance to render PairwiseMetricSpec.instance_prompt_template.""", ) - eval_case_id: Optional[str] = Field( - default=None, description="""Unique identifier for the evaluation case.""" - ) - intermediate_events: Optional[list[evals_types.Event]] = Field( - default=None, - description="""Intermediate events of a single turn in an agent run or intermediate events of the last turn for multi-turn an agent run.""", - ) - agent_info: Optional[evals_types.AgentInfo] = Field( - default=None, - description="""The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""", - ) - agent_data: Optional[evals_types.AgentData] = Field( - default=None, description="""The agent data of the agent under evaluation.""" - ) - user_scenario: Optional[evals_types.UserScenario] = Field( - default=None, description="""The user scenario for the evaluation case.""" - ) - interactions_data_source: Optional[InteractionsDataSource] = Field( - default=None, - description="""Source for populating agent data from an Interactions API interaction. When set, the backend fetches the interaction (and the Gemini Agent config) and parses it into agent data for evaluation; agent_data must not also be set.""", - ) - # Allow extra fields to support custom metric prompts and stay backward compatible. - model_config = ConfigDict(frozen=True, extra="allow") - - -class 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 PairwiseMetricInstanceDict(TypedDict, total=False): + """Pairwise metric instance.""" - reference: Optional[ResponseCandidateDict] - """User-provided, golden reference model reply to prompt in context of chat history; Reference for last response in a conversation.""" + 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.""" - 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).""" +PairwiseMetricInstanceOrDict = Union[PairwiseMetricInstance, PairwiseMetricInstanceDict] - 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 PairwiseMetricInput(_common.BaseModel): + """Pairwise 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.""" + 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.""" + ) - 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 PairwiseMetricInputDict(TypedDict, total=False): + """Pairwise metric instance.""" - user_scenario: Optional[evals_types.UserScenario] - """The user scenario for the evaluation case.""" + instance: Optional[PairwiseMetricInstanceDict] + """Required. Pairwise metric instance.""" - interactions_data_source: Optional[InteractionsDataSourceDict] - """Source for populating agent data from an Interactions API interaction. When set, the backend fetches the interaction (and the Gemini Agent config) and parses it into agent data for evaluation; agent_data must not also be set.""" + metric_spec: Optional[genai_types.PairwiseMetricSpec] + """Required. Spec for pairwise metric.""" -EvalCaseOrDict = Union[EvalCase, EvalCaseDict] +PairwiseMetricInputOrDict = Union[PairwiseMetricInput, PairwiseMetricInputDict] -class EvaluationDataset(_common.BaseModel): - """The dataset used for evaluation.""" +class ToolCallValidInstance(_common.BaseModel): + """Tool call valid instance.""" - 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.""", + prediction: Optional[str] = Field( + default=None, description="""Required. Output of the evaluated model.""" ) - candidate_name: Optional[str] = Field( + reference: Optional[str] = Field( default=None, - description="""The name of the candidate model or agent for this evaluation dataset.""", + description="""Required. Ground truth used to compare against the prediction.""", ) - @model_validator(mode="before") - @classmethod - def _check_pandas_installed(cls, data: Any) -> Any: - if isinstance(data, dict) and data.get("eval_dataset_df") is not None: - if pd is None: - logger.warning( - "Pandas is not installed, some evals features are not available." - " Please install it with `pip install" - " google-cloud-aiplatform[evaluation]`." - ) - return data - @classmethod - def load_from_observability_eval_cases( - cls, cases: list["ObservabilityEvalCase"] - ) -> "EvaluationDataset": - """Fetches GenAI Observability data from GCS and parses into a DataFrame.""" - try: - import pandas as pd - from .. import _gcs_utils +class ToolCallValidInstanceDict(TypedDict, total=False): + """Tool call valid instance.""" - formats = [] - requests = [] - responses = [] - system_instructions = [] + prediction: Optional[str] + """Required. Output of the evaluated model.""" - for case in cases: - gcs_utils = _gcs_utils.GcsUtils( - case.api_client._api_client if case.api_client else None - ) + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" - # Associate "observability" data format for given sources - formats.append("observability") - # Input source - request_data = gcs_utils.read_file_contents(case.input_src) - requests.append(request_data) +ToolCallValidInstanceOrDict = Union[ToolCallValidInstance, ToolCallValidInstanceDict] - # 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 ToolCallValidSpec(_common.BaseModel): + """Spec for tool call valid metric.""" - eval_dataset_df = pd.DataFrame( - { - "format": formats, - "request": requests, - "response": responses, - "system_instruction": system_instructions, - } - ) + pass - except ImportError as e: - raise ImportError("Pandas DataFrame library is required.") from e - return EvaluationDataset(eval_dataset_df=eval_dataset_df) +class ToolCallValidSpecDict(TypedDict, total=False): + """Spec for tool call valid metric.""" - def show(self) -> None: - """Shows the evaluation dataset.""" - from .. import _evals_visualization + pass - _evals_visualization.display_evaluation_dataset(self) +ToolCallValidSpecOrDict = Union[ToolCallValidSpec, ToolCallValidSpecDict] -class EvaluationDatasetDict(TypedDict, total=False): - """The dataset used for evaluation.""" - bigquery_source: Optional[genai_types.BigQuerySource] - """The BigQuery source for the evaluation dataset.""" +class ToolCallValidInput(_common.BaseModel): + """Tool call valid input.""" - gcs_source: Optional[genai_types.GcsSource] - """The GCS source for the evaluation dataset.""" + 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.""" + ) - 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 ToolCallValidInputDict(TypedDict, total=False): + """Tool call valid input.""" - candidate_name: Optional[str] - """The name of the candidate model or agent for this evaluation dataset.""" + instances: Optional[list[ToolCallValidInstanceDict]] + """Required. Repeated tool call valid instances.""" + metric_spec: Optional[ToolCallValidSpecDict] + """Required. Spec for tool call valid metric.""" -EvaluationDatasetOrDict = Union[EvaluationDataset, EvaluationDatasetDict] + +ToolCallValidInputOrDict = Union[ToolCallValidInput, ToolCallValidInputDict] -class EvaluationRunMetadata(_common.BaseModel): - """Metadata for an evaluation run.""" +class ToolNameMatchInstance(_common.BaseModel): + """Tool name match instance.""" - 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.""", + 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.""" +class ToolNameMatchInstanceDict(TypedDict, total=False): + """Tool name match instance.""" - candidate_names: Optional[list[str]] - """Name of the candidate(s) being evaluated in the evaluation run.""" + prediction: Optional[str] + """Required. Output of the evaluated model.""" - dataset_name: Optional[str] - """Name of the evaluation dataset used for the evaluation run.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" - dataset_id: Optional[str] - """Unique identifier for the evaluation dataset used for the evaluation run.""" - creation_timestamp: Optional[datetime.datetime] - """Creation timestamp of the evaluation run.""" +ToolNameMatchInstanceOrDict = Union[ToolNameMatchInstance, ToolNameMatchInstanceDict] -EvaluationRunMetadataOrDict = Union[EvaluationRunMetadata, EvaluationRunMetadataDict] +class ToolNameMatchSpec(_common.BaseModel): + """Spec for tool name match metric.""" + pass -class EvaluationResult(_common.BaseModel): - """Result of an evaluation run for an evaluation dataset.""" - eval_case_results: Optional[list[EvalCaseResult]] = Field( - default=None, - description="""A list of evaluation results for each evaluation case.""", +class ToolNameMatchSpecDict(TypedDict, total=False): + """Spec for tool name match metric.""" + + pass + + +ToolNameMatchSpecOrDict = Union[ToolNameMatchSpec, ToolNameMatchSpecDict] + + +class ToolNameMatchInput(_common.BaseModel): + """Tool name match input.""" + + instances: Optional[list[ToolNameMatchInstance]] = Field( + default=None, description="""Required. Repeated tool name match instances.""" ) - 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.""", + metric_spec: Optional[ToolNameMatchSpec] = Field( + default=None, description="""Required. Spec for tool name match metric.""" ) - def show(self, candidate_names: Optional[List[str]] = None) -> None: - """Shows the evaluation result. - Args: - candidate_names: list of names for the evaluated candidates, used in - comparison reports. - """ - from .. import _evals_visualization +class ToolNameMatchInputDict(TypedDict, total=False): + """Tool name match input.""" - _evals_visualization.display_evaluation_result(self, candidate_names) + instances: Optional[list[ToolNameMatchInstanceDict]] + """Required. Repeated tool name match instances.""" + metric_spec: Optional[ToolNameMatchSpecDict] + """Required. Spec for tool name match metric.""" -class EvaluationResultDict(TypedDict, total=False): - """Result of an evaluation run for an evaluation dataset.""" - eval_case_results: Optional[list[EvalCaseResultDict]] - """A list of evaluation results for each evaluation case.""" +ToolNameMatchInputOrDict = Union[ToolNameMatchInput, ToolNameMatchInputDict] - 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 ToolParameterKeyMatchInstance(_common.BaseModel): + """Tool parameter key match instance.""" - evaluation_dataset: Optional[list[EvaluationDatasetDict]] - """The input evaluation dataset(s) for the evaluation run.""" + 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.""", + ) - metadata: Optional[EvaluationRunMetadataDict] - """Metadata for the evaluation run.""" - agent_info: Optional[evals_types.AgentInfo] - """The agent info of the agent under evaluation. This can be extended for multi-agent evaluation.""" +class ToolParameterKeyMatchInstanceDict(TypedDict, total=False): + """Tool parameter key match instance.""" + prediction: Optional[str] + """Required. Output of the evaluated model.""" -EvaluationResultOrDict = Union[EvaluationResult, EvaluationResultDict] + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" -class EvaluationRun(_common.BaseModel): - """Represents an evaluation run.""" +ToolParameterKeyMatchInstanceOrDict = Union[ + ToolParameterKeyMatchInstance, ToolParameterKeyMatchInstanceDict +] - 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.""", - ) - # TODO(b/448806531): Remove all the overridden _from_response methods once the - # ticket is resolved and published. - @classmethod - def _from_response( - cls: typing.Type["EvaluationRun"], - *, - response: dict[str, object], - kwargs: dict[str, object], - ) -> "EvaluationRun": - """Converts a dictionary response into a EvaluationRun object.""" +class ToolParameterKeyMatchSpec(_common.BaseModel): + """Spec for tool parameter key match metric.""" - snaked_response = _camel_key_to_snake(response) + pass - evaluation_run_results = response.get("evaluation_run_results") - if ( - isinstance(evaluation_run_results, dict) - and "summaryMetrics" in evaluation_run_results - ): - snaked_response["evaluation_run_results"]["summary_metrics"] = ( - evaluation_run_results["summaryMetrics"] - ) - result = super()._from_response(response=snaked_response, kwargs=kwargs) - return result +class ToolParameterKeyMatchSpecDict(TypedDict, total=False): + """Spec for tool parameter key match metric.""" - def show(self) -> None: - """Shows the evaluation result.""" - from .. import _evals_visualization + pass - 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) +ToolParameterKeyMatchSpecOrDict = Union[ + ToolParameterKeyMatchSpec, ToolParameterKeyMatchSpecDict +] -class EvaluationRunDict(TypedDict, total=False): - """Represents an evaluation run.""" - name: Optional[str] - """""" +class ToolParameterKeyMatchInput(_common.BaseModel): + """Tool parameter key match input.""" - display_name: Optional[str] - """""" + instances: Optional[list[ToolParameterKeyMatchInstance]] = Field( + default=None, + description="""Required. Repeated tool parameter key match instances.""", + ) + metric_spec: Optional[ToolParameterKeyMatchSpec] = Field( + default=None, + description="""Required. Spec for tool parameter key match metric.""", + ) - metadata: Optional[dict[str, Any]] - """""" - create_time: Optional[datetime.datetime] - """""" +class ToolParameterKeyMatchInputDict(TypedDict, total=False): + """Tool parameter key match input.""" - completion_time: Optional[datetime.datetime] - """""" + instances: Optional[list[ToolParameterKeyMatchInstanceDict]] + """Required. Repeated tool parameter key match instances.""" - state: Optional[EvaluationRunState] - """""" + metric_spec: Optional[ToolParameterKeyMatchSpecDict] + """Required. Spec for tool parameter key match metric.""" - evaluation_set_snapshot: Optional[str] - """""" - error: Optional[genai_types.GoogleRpcStatus] - """""" +ToolParameterKeyMatchInputOrDict = Union[ + ToolParameterKeyMatchInput, ToolParameterKeyMatchInputDict +] - data_source: Optional[EvaluationRunDataSourceDict] - """""" - evaluation_run_results: Optional[EvaluationRunResultsDict] - """The evaluation run formatted results.""" +class ToolParameterKVMatchInstance(_common.BaseModel): + """Tool parameter kv match instance.""" - 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.""" + 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.""", + ) - evaluation_config: Optional[EvaluationRunConfigDict] - """The evaluation config for the evaluation run.""" - inference_configs: Optional[dict[str, EvaluationRunInferenceConfigDict]] - """The inference configs for the evaluation run.""" +class ToolParameterKVMatchInstanceDict(TypedDict, total=False): + """Tool parameter kv match instance.""" - labels: Optional[dict[str, str]] - """""" + prediction: Optional[str] + """Required. Output of the evaluated model.""" - analysis_configs: Optional[list[AnalysisConfigDict]] - """The analysis configurations for the evaluation run.""" + reference: Optional[str] + """Required. Ground truth used to compare against the prediction.""" -EvaluationRunOrDict = Union[EvaluationRun, EvaluationRunDict] +ToolParameterKVMatchInstanceOrDict = Union[ + ToolParameterKVMatchInstance, ToolParameterKVMatchInstanceDict +] -class CreateEvaluationSetConfig(_common.BaseModel): - """Config to create an evaluation set.""" +class ToolParameterKVMatchSpec(_common.BaseModel): + """Spec for tool parameter kv match metric.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + use_strict_string_match: Optional[bool] = Field( + default=None, + description="""Optional. Whether to use STRICT string match on parameter values.""", ) -class CreateEvaluationSetConfigDict(TypedDict, total=False): - """Config to create an evaluation set.""" +class ToolParameterKVMatchSpecDict(TypedDict, total=False): + """Spec for tool parameter kv match metric.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + use_strict_string_match: Optional[bool] + """Optional. Whether to use STRICT string match on parameter values.""" -CreateEvaluationSetConfigOrDict = Union[ - CreateEvaluationSetConfig, CreateEvaluationSetConfigDict +ToolParameterKVMatchSpecOrDict = Union[ + ToolParameterKVMatchSpec, ToolParameterKVMatchSpecDict ] -class _CreateEvaluationSetParameters(_common.BaseModel): - """Represents a job that creates an evaluation set.""" - - evaluation_items: Optional[list[str]] = Field(default=None, description="""""") - display_name: Optional[str] = Field(default=None, description="""""") - config: Optional[CreateEvaluationSetConfig] = Field( - default=None, description="""""" - ) +class ToolParameterKVMatchInput(_common.BaseModel): + """Tool parameter kv match input.""" + instances: Optional[list[ToolParameterKVMatchInstance]] = Field( + default=None, + description="""Required. Repeated tool parameter key value match instances.""", + ) + metric_spec: Optional[ToolParameterKVMatchSpec] = Field( + default=None, + description="""Required. Spec for tool parameter key value match metric.""", + ) -class _CreateEvaluationSetParametersDict(TypedDict, total=False): - """Represents a job that creates an evaluation set.""" - evaluation_items: Optional[list[str]] - """""" +class ToolParameterKVMatchInputDict(TypedDict, total=False): + """Tool parameter kv match input.""" - display_name: Optional[str] - """""" + instances: Optional[list[ToolParameterKVMatchInstanceDict]] + """Required. Repeated tool parameter key value match instances.""" - config: Optional[CreateEvaluationSetConfigDict] - """""" + metric_spec: Optional[ToolParameterKVMatchSpecDict] + """Required. Spec for tool parameter key value match metric.""" -_CreateEvaluationSetParametersOrDict = Union[ - _CreateEvaluationSetParameters, _CreateEvaluationSetParametersDict +ToolParameterKVMatchInputOrDict = Union[ + ToolParameterKVMatchInput, ToolParameterKVMatchInputDict ] -class EvaluationSet(_common.BaseModel): - """Represents an evaluation set.""" +class MapInstance(_common.BaseModel): + """Instance data specified as a map.""" - name: Optional[str] = Field( - default=None, description="""The resource name of the evaluation set.""" + map_instance: Optional[dict[str, evals_types.InstanceData]] = Field( + default=None, description="""Map of instance data.""" ) - display_name: Optional[str] = Field( - default=None, description="""The display name of the evaluation set.""" + + +class MapInstanceDict(TypedDict, total=False): + """Instance data specified as a map.""" + + map_instance: Optional[dict[str, evals_types.InstanceData]] + """Map of instance data.""" + + +MapInstanceOrDict = Union[MapInstance, MapInstanceDict] + + +class EvaluationInstance(_common.BaseModel): + """A single instance to be evaluated.""" + + prompt: Optional[evals_types.InstanceData] = Field( + default=None, + description="""Data used to populate placeholder `prompt` in a metric prompt template.""", ) - evaluation_items: Optional[list[str]] = Field( + response: Optional[evals_types.InstanceData] = Field( default=None, - description="""The EvaluationItems that are part of this dataset.""", + description="""Data used to populate placeholder `response` in a metric prompt template.""", ) - create_time: Optional[datetime.datetime] = Field( - default=None, description="""The create time of the evaluation set.""" + reference: Optional[evals_types.InstanceData] = Field( + default=None, + description="""Data used to populate placeholder `reference` in a metric prompt template.""", ) - update_time: Optional[datetime.datetime] = Field( - default=None, description="""The update time of the evaluation set.""" + other_data: Optional[MapInstance] = Field( + default=None, + description="""Other data used to populate placeholders based on their key.""", ) - metadata: Optional[dict[str, Any]] = Field( - default=None, description="""The metadata of the evaluation set.""" + agent_data: Optional[evals_types.AgentData] = Field( + default=None, description="""Data used for agent evaluation.""" + ) + rubric_groups: Optional[dict[str, "RubricGroup"]] = Field( + default=None, + description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", + ) + interactions_data_source: Optional[InteractionsDataSource] = Field( + default=None, + description="""Source for populating agent data from an Interactions API + interaction. If set, no other agent data source may be set. The backend + fetches the interaction (and the agent that produced it) and parses it + into agent data for grading.""", ) -class EvaluationSetDict(TypedDict, total=False): - """Represents an evaluation set.""" +class EvaluationInstanceDict(TypedDict, total=False): + """A single instance to be evaluated.""" - name: Optional[str] - """The resource name of the evaluation set.""" + prompt: Optional[evals_types.InstanceData] + """Data used to populate placeholder `prompt` in a metric prompt template.""" - display_name: Optional[str] - """The display name of the evaluation set.""" + response: Optional[evals_types.InstanceData] + """Data used to populate placeholder `response` in a metric prompt template.""" - evaluation_items: Optional[list[str]] - """The EvaluationItems that are part of this dataset.""" + reference: Optional[evals_types.InstanceData] + """Data used to populate placeholder `reference` in a metric prompt template.""" - create_time: Optional[datetime.datetime] - """The create time of the evaluation set.""" + other_data: Optional[MapInstanceDict] + """Other data used to populate placeholders based on their key.""" - update_time: Optional[datetime.datetime] - """The update time of the evaluation set.""" + agent_data: Optional[evals_types.AgentData] + """Data used for agent evaluation.""" - metadata: Optional[dict[str, Any]] - """The metadata of the evaluation set.""" + 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.""" -EvaluationSetOrDict = Union[EvaluationSet, EvaluationSetDict] +EvaluationInstanceOrDict = Union[EvaluationInstance, EvaluationInstanceDict] -class DeleteEvaluationMetricConfig(_common.BaseModel): - """Config for deleting an evaluation metric.""" + +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 DeleteEvaluationMetricConfigDict(TypedDict, total=False): - """Config for deleting an evaluation metric.""" +class EvaluateInstancesConfigDict(TypedDict, total=False): + """Config for evaluate instances.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteEvaluationMetricConfigOrDict = Union[ - DeleteEvaluationMetricConfig, DeleteEvaluationMetricConfigDict +EvaluateInstancesConfigOrDict = Union[ + EvaluateInstancesConfig, EvaluateInstancesConfigDict ] -class _DeleteEvaluationMetricParameters(_common.BaseModel): - """Parameters for deleting an evaluation metric.""" +class RubricBasedMetricSpec(_common.BaseModel): + """Specification for a metric that is based on rubrics.""" - metric_resource_name: Optional[str] = Field(default=None, description="""""") - config: Optional[DeleteEvaluationMetricConfig] = 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 _DeleteEvaluationMetricParametersDict(TypedDict, total=False): - """Parameters for deleting an evaluation metric.""" +class RubricBasedMetricSpecDict(TypedDict, total=False): + """Specification for a metric that is based on rubrics.""" - metric_resource_name: Optional[str] - """""" + metric_prompt_template: Optional[str] + """Template for the prompt used by the judge model to evaluate against + rubrics.""" - config: Optional[DeleteEvaluationMetricConfigDict] - """""" + 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.""" -_DeleteEvaluationMetricParametersOrDict = Union[ - _DeleteEvaluationMetricParameters, _DeleteEvaluationMetricParametersDict -] + rubric_group_key: Optional[str] + """Use a pre-defined group of rubrics associated with the input content. + This refers to a key in the `rubric_groups` map of + `RubricEnhancedContents`.""" + + rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] + """Dynamically generate rubrics for evaluation using this specification.""" -class DeleteEvaluationMetricOperation(_common.BaseModel): - """Operation for deleting an evaluation metric.""" +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 DeleteEvaluationMetricOperationDict(TypedDict, total=False): - """Operation for deleting an evaluation metric.""" +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.""" -DeleteEvaluationMetricOperationOrDict = Union[ - DeleteEvaluationMetricOperation, DeleteEvaluationMetricOperationDict -] +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.""" - - prediction: Optional[str] - """Required. Output of the evaluated model.""" - - reference: Optional[str] - """Required. Ground truth used to compare against the prediction.""" +class RubricBasedMetricInstanceDict(TypedDict, total=False): + """Defines an instance for Rubric-based metrics. + This class allows various input formats. + """ -BleuInstanceOrDict = Union[BleuInstance, BleuInstanceDict] + json_instance: Optional[str] + """Specify evaluation fields and their string values in JSON format.""" + content_map_instance: Optional[ContentMapDict] + """Specify evaluation fields and their content values using a ContentMap.""" -class BleuInput(_common.BaseModel): + rubric_enhanced_contents: Optional[RubricEnhancedContentsDict] + """Provide input as Gemini Content along with one or more + associated rubric groups.""" - instances: Optional[list[BleuInstance]] = Field( - default=None, description="""Required. Repeated bleu instances.""" + +RubricBasedMetricInstanceOrDict = Union[ + RubricBasedMetricInstance, RubricBasedMetricInstanceDict +] + + +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): +class RubricBasedMetricInputDict(TypedDict, total=False): + """Input for a rubric-based metrics.""" - instances: Optional[list[BleuInstanceDict]] - """Required. Repeated bleu instances.""" + metric_spec: Optional[RubricBasedMetricSpecDict] + """Specification for the rubric-based metric.""" - metric_spec: Optional[genai_types.BleuSpec] - """Required. Spec for bleu score metric.""" + instance: Optional[RubricBasedMetricInstanceDict] + """The instance to be evaluated.""" -BleuInputOrDict = Union[BleuInput, BleuInputDict] +RubricBasedMetricInputOrDict = Union[RubricBasedMetricInput, RubricBasedMetricInputDict] -class ExactMatchInstance(_common.BaseModel): - """Exact match instance.""" +class MetricSource(_common.BaseModel): + """The metric source used for evaluation.""" - prediction: Optional[str] = Field( - default=None, description="""Required. Output of the evaluated model.""" + metric: Optional[Metric] = Field( + default=None, description="""Inline metric config.""" ) - reference: Optional[str] = Field( + metric_resource_name: Optional[str] = Field( default=None, - description="""Required. Ground truth used to compare against the prediction.""", + description="""Resource name for registered metric. Example: + projects/{project}/locations/{location}/evaluationMetrics/{evaluation_metric_id}""", ) -class 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.""" -ExactMatchInputOrDict = Union[ExactMatchInput, ExactMatchInputDict] + metrics: Optional[list[MetricDict]] + """The metrics used for evaluation. + Currently, we only support evaluating a single metric. If multiple metrics + are provided, only the first one will be evaluated.""" + + instance: Optional[EvaluationInstanceDict] + """The instance to be evaluated.""" + 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.""" - values: Optional[dict[str, "ContentMapContents"]] = Field( - default=None, description="""Map of placeholder to contents.""" - ) +class ExactMatchResults(_common.BaseModel): + """Result of evaluating an exact match metric.""" + exact_match_metric_values: Optional[list[genai_types.ExactMatchMetricValue]] = ( + Field(default=None, description="""Output only. Exact match metric values.""") + ) -class ContentMapDict(TypedDict, total=False): - """Map of placeholder in metric prompt template to contents of model input.""" - values: Optional[dict[str, "ContentMapContents"]] - """Map of placeholder to contents.""" +class ExactMatchResultsDict(TypedDict, total=False): + """Result of evaluating an exact match metric.""" + 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.""" - 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.""", - ) +class RougeResults(_common.BaseModel): + """Result of evaluating a rouge metric.""" + 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.""" - json_instance: Optional[str] - """Instance specified as a json string. String key-value pairs are expected in the json_instance to render PointwiseMetricSpec.instance_prompt_template.""" +class RougeResultsDict(TypedDict, total=False): + """Result of evaluating a rouge metric.""" - 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.""" + rouge_metric_values: Optional[list[genai_types.RougeMetricValue]] + """Output only. Rouge metric values.""" -PointwiseMetricInstanceOrDict = Union[ - PointwiseMetricInstance, PointwiseMetricInstanceDict -] +RougeResultsOrDict = Union[RougeResults, RougeResultsDict] -class PointwiseMetricInput(_common.BaseModel): - """Pointwise metric input.""" +class RubricBasedMetricResult(_common.BaseModel): + """Result for a rubric-based metric.""" - 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.""" +class ToolParameterKeyMatchResultsDict(TypedDict, total=False): + """Results for tool parameter key match metric.""" - instances: Optional[list[ToolNameMatchInstanceDict]] - """Required. Repeated tool name match instances.""" + tool_parameter_key_match_metric_values: Optional[ + list[ToolParameterKeyMatchMetricValueDict] + ] + """Output only. Tool parameter key match metric values.""" - metric_spec: Optional[ToolNameMatchSpecDict] - """Required. Spec for tool name match metric.""" - -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 PrefixMatchMetricValue(_common.BaseModel): + """Prefix match metric value for an instance.""" - instances: Optional[list[ToolParameterKeyMatchInstance]] = Field( - default=None, - description="""Required. Repeated tool parameter key match instances.""", + score: Optional[float] = Field( + default=None, description="""Output only. Prefix match score.""" ) - metric_spec: Optional[ToolParameterKeyMatchSpec] = Field( - default=None, - description="""Required. Spec for tool parameter key match metric.""", + + +class PrefixMatchMetricValueDict(TypedDict, total=False): + """Prefix match metric value for an instance.""" + + score: Optional[float] + """Output only. Prefix match score.""" + + +PrefixMatchMetricValueOrDict = Union[PrefixMatchMetricValue, PrefixMatchMetricValueDict] + + +class PrefixMatchResults(_common.BaseModel): + """Results for prefix match metric.""" + + prefix_match_metric_values: Optional[list[PrefixMatchMetricValue]] = Field( + default=None, description="""Output only. Prefix match metric values.""" ) -class ToolParameterKeyMatchInputDict(TypedDict, total=False): - """Tool parameter key match input.""" +class PrefixMatchResultsDict(TypedDict, total=False): + """Results for prefix match metric.""" - instances: Optional[list[ToolParameterKeyMatchInstanceDict]] - """Required. Repeated tool parameter key match instances.""" + prefix_match_metric_values: Optional[list[PrefixMatchMetricValueDict]] + """Output only. Prefix match metric values.""" - metric_spec: Optional[ToolParameterKeyMatchSpecDict] - """Required. Spec for tool parameter key match metric.""" +PrefixMatchResultsOrDict = Union[PrefixMatchResults, PrefixMatchResultsDict] -ToolParameterKeyMatchInputOrDict = Union[ - ToolParameterKeyMatchInput, ToolParameterKeyMatchInputDict + +class WordCountMatchMetricValue(_common.BaseModel): + """Word count match metric value for an instance.""" + + score: Optional[float] = Field( + default=None, description="""Output only. Word count match score.""" + ) + + +class WordCountMatchMetricValueDict(TypedDict, total=False): + """Word count match metric value for an instance.""" + + score: Optional[float] + """Output only. Word count match score.""" + + +WordCountMatchMetricValueOrDict = Union[ + WordCountMatchMetricValue, WordCountMatchMetricValueDict ] -class ToolParameterKVMatchInstance(_common.BaseModel): - """Tool parameter kv match instance.""" +class WordCountMatchResults(_common.BaseModel): + """Results for word count match metric.""" - prediction: Optional[str] = Field( - default=None, description="""Required. Output of the evaluated model.""" + word_count_match_metric_values: Optional[list[WordCountMatchMetricValue]] = Field( + default=None, description="""Output only. Word count match metric values.""" ) - reference: Optional[str] = Field( - default=None, - description="""Required. Ground truth used to compare against the prediction.""", + + +class WordCountMatchResultsDict(TypedDict, total=False): + """Results for word count match metric.""" + + word_count_match_metric_values: Optional[list[WordCountMatchMetricValueDict]] + """Output only. Word count match metric values.""" + + +WordCountMatchResultsOrDict = Union[WordCountMatchResults, WordCountMatchResultsDict] + + +class FleschKincaidReadabilityMetricValue(_common.BaseModel): + """Flesch kincaid readability metric value for an instance.""" + + score: Optional[float] = Field( + default=None, description="""Output only. Readability score.""" + ) + grade_level: Optional[float] = Field( + default=None, description="""Output only. Readability grade level.""" ) -class ToolParameterKVMatchInstanceDict(TypedDict, total=False): - """Tool parameter kv match instance.""" +class FleschKincaidReadabilityMetricValueDict(TypedDict, total=False): + """Flesch kincaid readability metric value for an instance.""" - prediction: Optional[str] - """Required. Output of the evaluated model.""" + score: Optional[float] + """Output only. Readability score.""" - reference: Optional[str] - """Required. Ground truth used to compare against the prediction.""" + grade_level: Optional[float] + """Output only. Readability grade level.""" -ToolParameterKVMatchInstanceOrDict = Union[ - ToolParameterKVMatchInstance, ToolParameterKVMatchInstanceDict +FleschKincaidReadabilityMetricValueOrDict = Union[ + FleschKincaidReadabilityMetricValue, FleschKincaidReadabilityMetricValueDict ] -class ToolParameterKVMatchSpec(_common.BaseModel): - """Spec for tool parameter kv match metric.""" +class FleschKincaidReadabilityResults(_common.BaseModel): + """Results for flesch kincaid readability metric.""" - use_strict_string_match: Optional[bool] = Field( + flesch_kincaid_readability_metric_values: Optional[ + list[FleschKincaidReadabilityMetricValue] + ] = Field( default=None, - description="""Optional. Whether to use STRICT string match on parameter values.""", + description="""Output only. Flesch kincaid readability metric values.""", ) -class ToolParameterKVMatchSpecDict(TypedDict, total=False): - """Spec for tool parameter kv match metric.""" +class FleschKincaidReadabilityResultsDict(TypedDict, total=False): + """Results for flesch kincaid readability metric.""" - use_strict_string_match: Optional[bool] - """Optional. Whether to use STRICT string match on parameter values.""" + flesch_kincaid_readability_metric_values: Optional[ + list[FleschKincaidReadabilityMetricValueDict] + ] + """Output only. Flesch kincaid readability metric values.""" -ToolParameterKVMatchSpecOrDict = Union[ - ToolParameterKVMatchSpec, ToolParameterKVMatchSpecDict +FleschKincaidReadabilityResultsOrDict = Union[ + FleschKincaidReadabilityResults, FleschKincaidReadabilityResultsDict ] -class ToolParameterKVMatchInput(_common.BaseModel): - """Tool parameter kv match input.""" +class StringMatchResult(_common.BaseModel): + """Result for StringMatch metric.""" - instances: Optional[list[ToolParameterKVMatchInstance]] = Field( - default=None, - description="""Required. Repeated tool parameter key value match instances.""", - ) - metric_spec: Optional[ToolParameterKVMatchSpec] = Field( + score: Optional[float] = Field( default=None, - description="""Required. Spec for tool parameter key value match metric.""", + description="""StringMatch metric score. If matched, the score is 1.0, otherwise 0.0.""", ) -class ToolParameterKVMatchInputDict(TypedDict, total=False): - """Tool parameter kv match input.""" - - instances: Optional[list[ToolParameterKVMatchInstanceDict]] - """Required. Repeated tool parameter key value match instances.""" +class StringMatchResultDict(TypedDict, total=False): + """Result for StringMatch metric.""" - metric_spec: Optional[ToolParameterKVMatchSpecDict] - """Required. Spec for tool parameter key value match metric.""" + score: Optional[float] + """StringMatch metric score. If matched, the score is 1.0, otherwise 0.0.""" -ToolParameterKVMatchInputOrDict = Union[ - ToolParameterKVMatchInput, ToolParameterKVMatchInputDict -] +StringMatchResultOrDict = Union[StringMatchResult, StringMatchResultDict] -class MapInstance(_common.BaseModel): - """Instance data specified as a map.""" +class CodeExecutionResults(_common.BaseModel): + """Result for CodeExecution metric.""" - map_instance: Optional[dict[str, evals_types.InstanceData]] = Field( - default=None, description="""Map of instance data.""" + outcome: Optional[Outcome] = Field( + default=None, description="""Required. Outcome of the code execution.""" + ) + score: Optional[float] = Field( + default=None, + description="""Score returned by the running code_snippet provided in the CodeExecutionInput.""", ) -class MapInstanceDict(TypedDict, total=False): - """Instance data specified as a map.""" +class CodeExecutionResultsDict(TypedDict, total=False): + """Result for CodeExecution metric.""" - map_instance: Optional[dict[str, evals_types.InstanceData]] - """Map of instance data.""" + outcome: Optional[Outcome] + """Required. Outcome of the code execution.""" + score: Optional[float] + """Score returned by the running code_snippet provided in the CodeExecutionInput.""" -MapInstanceOrDict = Union[MapInstance, MapInstanceDict] +CodeExecutionResultsOrDict = Union[CodeExecutionResults, CodeExecutionResultsDict] -class EvaluationInstance(_common.BaseModel): - """A single instance to be evaluated.""" - prompt: Optional[evals_types.InstanceData] = Field( - default=None, - description="""Data used to populate placeholder `prompt` in a metric prompt template.""", - ) - response: Optional[evals_types.InstanceData] = Field( - default=None, - description="""Data used to populate placeholder `response` in a metric prompt template.""", +class ResponseRecallResult(_common.BaseModel): + """Spec for response recall result.""" + + score: Optional[float] = Field( + default=None, description="""Output only. ResponseRecall score.""" ) - reference: Optional[evals_types.InstanceData] = Field( + explanation: Optional[str] = Field( default=None, - description="""Data used to populate placeholder `reference` in a metric prompt template.""", + description="""Output only. Explanation for response recall score.""", ) - other_data: Optional[MapInstance] = Field( - default=None, - description="""Other data used to populate placeholders based on their key.""", + confidence: Optional[float] = Field( + default=None, description="""Output only. Confidence for fulfillment score.""" ) - agent_data: Optional[evals_types.AgentData] = Field( - default=None, description="""Data used for agent evaluation.""" + + +class ResponseRecallResultDict(TypedDict, total=False): + """Spec for response recall result.""" + + score: Optional[float] + """Output only. ResponseRecall score.""" + + explanation: Optional[str] + """Output only. Explanation for response recall score.""" + + confidence: Optional[float] + """Output only. Confidence for fulfillment score.""" + + +ResponseRecallResultOrDict = Union[ResponseRecallResult, ResponseRecallResultDict] + + +class RagContextRecallResult(_common.BaseModel): + """Spec for rag context recall result.""" + + score: Optional[float] = Field( + default=None, description="""Output only. RagContextRecall score.""" ) - rubric_groups: Optional[dict[str, "RubricGroup"]] = Field( + explanation: Optional[str] = Field( default=None, - description="""Named groups of rubrics associated with this prompt. The key is a user-defined name for the rubric group.""", + description="""Output only. Explanation for rag context recall score.""", ) - interactions_data_source: Optional[InteractionsDataSource] = Field( + confidence: Optional[float] = 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.""", + description="""Output only. Confidence for rag context recall score.""", ) -class EvaluationInstanceDict(TypedDict, total=False): - """A single instance to be evaluated.""" +class RagContextRecallResultDict(TypedDict, total=False): + """Spec for rag context recall result.""" - prompt: Optional[evals_types.InstanceData] - """Data used to populate placeholder `prompt` in a metric prompt template.""" + score: Optional[float] + """Output only. RagContextRecall score.""" - response: Optional[evals_types.InstanceData] - """Data used to populate placeholder `response` in a metric prompt template.""" + explanation: Optional[str] + """Output only. Explanation for rag context recall score.""" - reference: Optional[evals_types.InstanceData] - """Data used to populate placeholder `reference` in a metric prompt template.""" + confidence: Optional[float] + """Output only. Confidence for rag context recall score.""" - 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.""" +RagContextRecallResultOrDict = Union[RagContextRecallResult, RagContextRecallResultDict] - 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 PointwiseSafetyMetricResult(_common.BaseModel): + """Spec for pointwise safety metric result.""" + + is_unsafe: Optional[bool] = Field( + default=None, description="""Output only. Pointwise safety determination.""" + ) + is_unsafe_probability: Optional[float] = Field( + default=None, description="""Output only. Unsafe probability.""" + ) + violated_policies: Optional[list[str]] = Field( + default=None, description="""Output only. Violated policies.""" + ) -EvaluationInstanceOrDict = Union[EvaluationInstance, EvaluationInstanceDict] +class PointwiseSafetyMetricResultDict(TypedDict, total=False): + """Spec for pointwise safety metric result.""" + is_unsafe: Optional[bool] + """Output only. Pointwise safety determination.""" -class EvaluateInstancesConfig(_common.BaseModel): - """Config for evaluate instances.""" + is_unsafe_probability: Optional[float] + """Output only. Unsafe probability.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + violated_policies: Optional[list[str]] + """Output only. Violated policies.""" + + +PointwiseSafetyMetricResultOrDict = Union[ + PointwiseSafetyMetricResult, PointwiseSafetyMetricResultDict +] + + +class LLMBasedSimilarityMetricResult(_common.BaseModel): + """Result for LLM-Based Similarity Metric.""" + + score: Optional[float] = Field( + default=None, description="""Output only. The score for the instance.""" ) -class EvaluateInstancesConfigDict(TypedDict, total=False): - """Config for evaluate instances.""" +class LLMBasedSimilarityMetricResultDict(TypedDict, total=False): + """Result for LLM-Based Similarity Metric.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + score: Optional[float] + """Output only. The score for the instance.""" -EvaluateInstancesConfigOrDict = Union[ - EvaluateInstancesConfig, EvaluateInstancesConfigDict +LLMBasedSimilarityMetricResultOrDict = Union[ + LLMBasedSimilarityMetricResult, LLMBasedSimilarityMetricResultDict ] -class RubricBasedMetricSpec(_common.BaseModel): - """Specification for a metric that is based on rubrics.""" +class LLMBasedClassificationMetricResult(_common.BaseModel): + """Result for LLM-Based Classification Metric.""" - 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( + score: Optional[float] = Field( default=None, - description="""Optional configuration for the judge LLM (Autorater).""", + description="""Output only. The score for each instance. The score is 1.0 if the classification is one of the passing labels, and 0.0 otherwise.""", ) - inline_rubrics: Optional[list[evals_types.Rubric]] = Field( - default=None, description="""Use rubrics provided directly in the spec.""" - ) - rubric_group_key: Optional[str] = Field( + + +class LLMBasedClassificationMetricResultDict(TypedDict, total=False): + """Result for LLM-Based Classification Metric.""" + + score: Optional[float] + """Output only. The score for each instance. The score is 1.0 if the classification is one of the passing labels, and 0.0 otherwise.""" + + +LLMBasedClassificationMetricResultOrDict = Union[ + LLMBasedClassificationMetricResult, LLMBasedClassificationMetricResultDict +] + + +class CloudLoggingResult(_common.BaseModel): + """Configuration for exporting evaluation results to Cloud Logging.""" + + insert_id: 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`.""", + description="""Output only. ID of inserted log entry. If present, export to Cloud Logging was successful.""", ) - rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] = Field( + error: Optional[genai_types.GoogleRpcStatus] = Field( default=None, - description="""Dynamically generate rubrics for evaluation using this specification.""", + description="""Output only. Error if export to Cloud Logging failed.""", ) -class RubricBasedMetricSpecDict(TypedDict, total=False): - """Specification for a metric that is based on rubrics.""" - - metric_prompt_template: Optional[str] - """Template for the prompt used by the judge model to evaluate against - rubrics.""" - - judge_autorater_config: Optional[genai_types.AutoraterConfig] - """Optional configuration for the judge LLM (Autorater).""" - - inline_rubrics: Optional[list[evals_types.Rubric]] - """Use rubrics provided directly in the spec.""" +class CloudLoggingResultDict(TypedDict, total=False): + """Configuration for exporting evaluation results to Cloud Logging.""" - 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`.""" + insert_id: Optional[str] + """Output only. ID of inserted log entry. If present, export to Cloud Logging was successful.""" - rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] - """Dynamically generate rubrics for evaluation using this specification.""" + error: Optional[genai_types.GoogleRpcStatusDict] + """Output only. Error if export to Cloud Logging failed.""" -RubricBasedMetricSpecOrDict = Union[RubricBasedMetricSpec, RubricBasedMetricSpecDict] +CloudLoggingResultOrDict = Union[CloudLoggingResult, CloudLoggingResultDict] -class RubricEnhancedContents(_common.BaseModel): - """Rubric-enhanced contents for evaluation.""" +class EvaluateInstancesResponse(_common.BaseModel): + """Result of evaluating an LLM metric.""" - prompt: Optional[list[genai_types.Content]] = 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="""User prompt, using the standard Content type from the Gen AI SDK.""", + description="""A list of metric results for each evaluation case. The order of the metric results is guaranteed to be the same as the order of the instances in the request.""", ) - rubric_groups: Optional[dict[str, "RubricGroup"]] = 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="""Named groups of rubrics associated with this prompt. - The key is a user-defined name for the rubric group.""", + description="""Auto metric evaluation results. Results for exact match metric.""", ) - response: Optional[list[genai_types.Content]] = Field( + 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="""Response, using the standard Content type from the Gen AI SDK.""", + description="""Tool call metrics. Results for tool call valid metric.""", ) - other_content: Optional[ContentMap] = Field( + 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="""Other contents needed for the metric. - For example, if `reference` is needed for the metric, it can be provided - here.""", + description="""Results for tool parameter key value match metric.""", + ) + prefix_match_results: Optional[PrefixMatchResults] = Field( + default=None, description="""Results for prefix match metric.""" + ) + word_count_match_results: Optional[WordCountMatchResults] = Field( + default=None, description="""Results for word count match metric.""" + ) + flesch_kincaid_readability_results: Optional[FleschKincaidReadabilityResults] = ( + Field( + default=None, + description="""Results for flesch kincaid readability metric.""", + ) + ) + string_match_result: Optional[StringMatchResult] = Field( + default=None, description="""Results for string match metric.""" + ) + code_execution_result: Optional[CodeExecutionResults] = Field( + default=None, description="""Results for code execution metric.""" + ) + response_recall_result: Optional[ResponseRecallResult] = Field( + default=None, description="""Result for response recall metric.""" + ) + rag_context_recall_result: Optional[RagContextRecallResult] = Field( + default=None, + description="""RAG only metrics. Result for context recall metric.""", + ) + pointwise_safety_metric_result: Optional[PointwiseSafetyMetricResult] = Field( + default=None, description="""Result for pointwise safety metric.""" + ) + llm_based_similarity_metric_result: Optional[LLMBasedSimilarityMetricResult] = ( + Field(default=None, description="""Result for LLM-based similarity metric.""") + ) + llm_based_classification_metric_result: Optional[ + LLMBasedClassificationMetricResult + ] = Field( + default=None, description="""Result for LLM-based classification metric.""" + ) + cloud_logging_result: Optional[CloudLoggingResult] = Field( + default=None, + description="""Optional. Result from exporting evaluation results to Cloud Logging.""", ) -class RubricEnhancedContentsDict(TypedDict, total=False): - """Rubric-enhanced contents for evaluation.""" - - prompt: Optional[list[genai_types.Content]] - """User prompt, using the standard Content type from the Gen AI SDK.""" +class EvaluateInstancesResponseDict(TypedDict, total=False): + """Result of evaluating an LLM metric.""" - rubric_groups: Optional[dict[str, "RubricGroup"]] - """Named groups of rubrics associated with this prompt. - The key is a user-defined name for the rubric group.""" + rubric_based_metric_result: Optional[RubricBasedMetricResultDict] + """Result for rubric based metric.""" - response: Optional[list[genai_types.Content]] - """Response, using the standard Content type from the Gen AI SDK.""" + 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.""" - other_content: Optional[ContentMapDict] - """Other contents needed for the metric. - For example, if `reference` is needed for the metric, it can be provided - here.""" + bleu_results: Optional[BleuResultsDict] + """Results for bleu metric.""" + comet_result: Optional[CometResultDict] + """Translation metrics. Result for Comet metric.""" -RubricEnhancedContentsOrDict = Union[RubricEnhancedContents, RubricEnhancedContentsDict] + exact_match_results: Optional[ExactMatchResultsDict] + """Auto metric evaluation results. Results for exact match metric.""" + metricx_result: Optional[MetricxResultDict] + """Result for Metricx metric.""" -class RubricBasedMetricInstance(_common.BaseModel): - """Defines an instance for Rubric-based metrics. + pairwise_metric_result: Optional[genai_types.PairwiseMetricResult] + """Result for pairwise metric.""" - This class allows various input formats. - """ + pointwise_metric_result: Optional[genai_types.PointwiseMetricResult] + """Generic metrics. Result for pointwise metric.""" - json_instance: Optional[str] = Field( - default=None, - description="""Specify evaluation fields and their string values in JSON format.""", - ) - content_map_instance: Optional[ContentMap] = Field( - default=None, - description="""Specify evaluation fields and their content values using a ContentMap.""", - ) - rubric_enhanced_contents: Optional[RubricEnhancedContents] = Field( - default=None, - description="""Provide input as Gemini Content along with one or more - associated rubric groups.""", - ) + rouge_results: Optional[RougeResultsDict] + """Results for rouge metric.""" + tool_call_valid_results: Optional[ToolCallValidResultsDict] + """Tool call metrics. Results for tool call valid metric.""" -class RubricBasedMetricInstanceDict(TypedDict, total=False): - """Defines an instance for Rubric-based metrics. + tool_name_match_results: Optional[ToolNameMatchResultsDict] + """Results for tool name match metric.""" - This class allows various input formats. - """ + tool_parameter_key_match_results: Optional[ToolParameterKeyMatchResultsDict] + """Results for tool parameter key match metric.""" - json_instance: Optional[str] - """Specify evaluation fields and their string values in JSON format.""" + tool_parameter_kv_match_results: Optional[ToolParameterKVMatchResultsDict] + """Results for tool parameter key value match metric.""" - content_map_instance: Optional[ContentMapDict] - """Specify evaluation fields and their content values using a ContentMap.""" + prefix_match_results: Optional[PrefixMatchResultsDict] + """Results for prefix match metric.""" - rubric_enhanced_contents: Optional[RubricEnhancedContentsDict] - """Provide input as Gemini Content along with one or more - associated rubric groups.""" + word_count_match_results: Optional[WordCountMatchResultsDict] + """Results for word count match metric.""" + flesch_kincaid_readability_results: Optional[FleschKincaidReadabilityResultsDict] + """Results for flesch kincaid readability metric.""" -RubricBasedMetricInstanceOrDict = Union[ - RubricBasedMetricInstance, RubricBasedMetricInstanceDict -] + string_match_result: Optional[StringMatchResultDict] + """Results for string match metric.""" + code_execution_result: Optional[CodeExecutionResultsDict] + """Results for code execution metric.""" -class RubricBasedMetricInput(_common.BaseModel): - """Input for a rubric-based metrics.""" + response_recall_result: Optional[ResponseRecallResultDict] + """Result for response recall metric.""" - metric_spec: Optional[RubricBasedMetricSpec] = Field( - default=None, description="""Specification for the rubric-based metric.""" - ) - instance: Optional[RubricBasedMetricInstance] = Field( - default=None, description="""The instance to be evaluated.""" - ) + rag_context_recall_result: Optional[RagContextRecallResultDict] + """RAG only metrics. Result for context recall metric.""" + pointwise_safety_metric_result: Optional[PointwiseSafetyMetricResultDict] + """Result for pointwise safety metric.""" -class RubricBasedMetricInputDict(TypedDict, total=False): - """Input for a rubric-based metrics.""" + llm_based_similarity_metric_result: Optional[LLMBasedSimilarityMetricResultDict] + """Result for LLM-based similarity metric.""" - metric_spec: Optional[RubricBasedMetricSpecDict] - """Specification for the rubric-based metric.""" + llm_based_classification_metric_result: Optional[ + LLMBasedClassificationMetricResultDict + ] + """Result for LLM-based classification metric.""" - instance: Optional[RubricBasedMetricInstanceDict] - """The instance to be evaluated.""" + cloud_logging_result: Optional[CloudLoggingResultDict] + """Optional. Result from exporting evaluation results to Cloud Logging.""" -RubricBasedMetricInputOrDict = Union[RubricBasedMetricInput, RubricBasedMetricInputDict] +EvaluateInstancesResponseOrDict = Union[ + EvaluateInstancesResponse, EvaluateInstancesResponseDict +] -class MetricSource(_common.BaseModel): - """The metric source used for evaluation.""" +class GenerateUserScenariosConfig(_common.BaseModel): - 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 GenerateUserScenariosConfigDict(TypedDict, total=False): - 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] +GenerateUserScenariosConfigOrDict = Union[ + GenerateUserScenariosConfig, GenerateUserScenariosConfigDict +] -class _EvaluateInstancesRequestParameters(_common.BaseModel): - """Parameters for evaluating instances.""" +class _GenerateUserScenariosParameters(_common.BaseModel): + """Parameters for GenerateUserScenarios.""" - 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( + location: Optional[str] = Field(default=None, description="""""") + agents: Optional[dict[str, evals_types.AgentConfig]] = Field( default=None, description="""""" ) - pairwise_metric_input: Optional[PairwiseMetricInput] = Field( + 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="""""" ) - tool_call_valid_input: Optional[ToolCallValidInput] = Field( - default=None, description="""""" + allow_cross_region_model: Optional[bool] = Field( + default=None, + description="""Opt-in flag to authorize cross-region routing for LLM models.""", ) - tool_name_match_input: Optional[ToolNameMatchInput] = Field( - default=None, description="""""" + gemini_agent_config: Optional[GeminiAgentConfig] = Field( + default=None, + description="""If set, the server derives the agents map and root_agent_id + from the referenced Gemini Agent server-side.""", ) - 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 _EvaluateInstancesRequestParametersDict(TypedDict, total=False): - """Parameters for evaluating instances.""" +class _GenerateUserScenariosParametersDict(TypedDict, total=False): + """Parameters for GenerateUserScenarios.""" - bleu_input: Optional[BleuInputDict] + location: Optional[str] """""" - exact_match_input: Optional[ExactMatchInputDict] + agents: Optional[dict[str, evals_types.AgentConfig]] """""" - rouge_input: Optional[RougeInputDict] + root_agent_id: Optional[str] """""" - pointwise_metric_input: Optional[PointwiseMetricInputDict] + user_scenario_generation_config: Optional[evals_types.UserScenarioGenerationConfig] """""" - pairwise_metric_input: Optional[PairwiseMetricInputDict] + config: Optional[GenerateUserScenariosConfigDict] """""" - tool_call_valid_input: Optional[ToolCallValidInputDict] - """""" + allow_cross_region_model: Optional[bool] + """Opt-in flag to authorize cross-region routing for LLM models.""" - tool_name_match_input: Optional[ToolNameMatchInputDict] - """""" + gemini_agent_config: Optional[GeminiAgentConfigDict] + """If set, the server derives the agents map and root_agent_id + from the referenced Gemini Agent server-side.""" - tool_parameter_key_match_input: Optional[ToolParameterKeyMatchInputDict] - """""" - tool_parameter_kv_match_input: Optional[ToolParameterKVMatchInputDict] - """""" +_GenerateUserScenariosParametersOrDict = Union[ + _GenerateUserScenariosParameters, _GenerateUserScenariosParametersDict +] - 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.""" +class GenerateUserScenariosResponse(_common.BaseModel): + """Response message for DataFoundryService.GenerateUserScenarios.""" - 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.""" + user_scenarios: Optional[list[evals_types.UserScenario]] = Field( + default=None, description="""""" + ) - instance: Optional[EvaluationInstanceDict] - """The instance to be evaluated.""" - metric_sources: Optional[list[MetricSourceDict]] - """The metrics used for evaluation.""" +class GenerateUserScenariosResponseDict(TypedDict, total=False): + """Response message for DataFoundryService.GenerateUserScenarios.""" - config: Optional[EvaluateInstancesConfigDict] + user_scenarios: Optional[list[evals_types.UserScenario]] """""" -_EvaluateInstancesRequestParametersOrDict = Union[ - _EvaluateInstancesRequestParameters, _EvaluateInstancesRequestParametersDict +GenerateUserScenariosResponseOrDict = Union[ + GenerateUserScenariosResponse, GenerateUserScenariosResponseDict ] -class MetricResult(_common.BaseModel): - """Result for a single metric on a single instance.""" +class GenerateLossClustersConfig(_common.BaseModel): + """Config for generating loss clusters.""" - 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.""" - - explanation: Optional[str] - """The explanation for the metric result.""" +class GenerateLossClustersConfigDict(TypedDict, total=False): + """Config for generating loss clusters.""" - error: Optional[genai_types.GoogleRpcStatus] - """The error status for the metric result.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -MetricResultOrDict = Union[MetricResult, MetricResultDict] +GenerateLossClustersConfigOrDict = Union[ + GenerateLossClustersConfig, GenerateLossClustersConfigDict +] -class BleuResults(_common.BaseModel): - """Result of evaluating a bleu metric.""" +class _GenerateLossClustersParameters(_common.BaseModel): + """Parameters for GenerateLossClusters.""" - bleu_metric_values: Optional[list[genai_types.BleuMetricValue]] = Field( - default=None, description="""Output only. Bleu metric values.""" + location: Optional[str] = Field( + default=None, + description="""The resource name of the Location. Format: `projects/{project}/locations/{location}`.""", + ) + evaluation_set: Optional[str] = Field( + default=None, + description="""Reference to a persisted EvaluationSet. The service will read items from this set.""", + ) + inline_results: Optional[list[EvaluationResult]] = Field( + default=None, + description="""Inline evaluation results. Useful for ephemeral analysis in notebooks/SDKs where data isn't persisted.""", + ) + configs: Optional[list[LossAnalysisConfig]] = Field( + default=None, + description="""Configuration for the analysis algorithm. Analysis for multiple metrics and multiple candidates could be specified.""", + ) + config: Optional[GenerateLossClustersConfig] = Field( + default=None, description="""Config for generating loss clusters.""" ) -class BleuResultsDict(TypedDict, total=False): - """Result of evaluating a bleu metric.""" - - bleu_metric_values: Optional[list[genai_types.BleuMetricValue]] - """Output only. Bleu metric values.""" +class _GenerateLossClustersParametersDict(TypedDict, total=False): + """Parameters for GenerateLossClusters.""" + location: Optional[str] + """The resource name of the Location. Format: `projects/{project}/locations/{location}`.""" -BleuResultsOrDict = Union[BleuResults, BleuResultsDict] + evaluation_set: Optional[str] + """Reference to a persisted EvaluationSet. The service will read items from this set.""" + inline_results: Optional[list[EvaluationResultDict]] + """Inline evaluation results. Useful for ephemeral analysis in notebooks/SDKs where data isn't persisted.""" -class ExactMatchResults(_common.BaseModel): - """Result of evaluating an exact match metric.""" + configs: Optional[list[LossAnalysisConfigDict]] + """Configuration for the analysis algorithm. Analysis for multiple metrics and multiple candidates could be specified.""" - exact_match_metric_values: Optional[list[genai_types.ExactMatchMetricValue]] = ( - Field(default=None, description="""Output only. Exact match metric values.""") - ) + config: Optional[GenerateLossClustersConfigDict] + """Config for generating loss clusters.""" -class ExactMatchResultsDict(TypedDict, total=False): - """Result of evaluating an exact match metric.""" +_GenerateLossClustersParametersOrDict = Union[ + _GenerateLossClustersParameters, _GenerateLossClustersParametersDict +] - exact_match_metric_values: Optional[list[genai_types.ExactMatchMetricValue]] - """Output only. Exact match metric values.""" +class GenerateLossClustersResponse(_common.BaseModel): + """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" -ExactMatchResultsOrDict = Union[ExactMatchResults, ExactMatchResultsDict] + analysis_time: Optional[str] = Field( + default=None, description="""The timestamp when this analysis was completed.""" + ) + results: Optional[list[LossAnalysisResult]] = Field( + default=None, + description="""The analysis results, one per config provided in the request.""", + ) + def show(self) -> None: + """Shows the loss pattern analysis report with rich HTML visualization.""" + from .. import _evals_visualization -class RougeResults(_common.BaseModel): - """Result of evaluating a rouge metric.""" + _evals_visualization.display_loss_clusters_response(self) - rouge_metric_values: Optional[list[genai_types.RougeMetricValue]] = Field( - default=None, description="""Output only. Rouge metric values.""" - ) +class GenerateLossClustersResponseDict(TypedDict, total=False): + """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" -class RougeResultsDict(TypedDict, total=False): - """Result of evaluating a rouge metric.""" + analysis_time: Optional[str] + """The timestamp when this analysis was completed.""" - rouge_metric_values: Optional[list[genai_types.RougeMetricValue]] - """Output only. Rouge metric values.""" + results: Optional[list[LossAnalysisResultDict]] + """The analysis results, one per config provided in the request.""" -RougeResultsOrDict = Union[RougeResults, RougeResultsDict] +GenerateLossClustersResponseOrDict = Union[ + GenerateLossClustersResponse, GenerateLossClustersResponseDict +] -class RubricBasedMetricResult(_common.BaseModel): - """Result for a rubric-based metric.""" +class GenerateLossClustersOperation(_common.BaseModel): + """Long-running operation for generating loss clusters.""" - score: Optional[float] = Field( - default=None, description="""Passing rate of all the rubrics.""" + 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}`.""", ) - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""The details of all the rubrics and their verdicts.""", + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + response: Optional[GenerateLossClustersResponse] = Field( + default=None, + description="""Response message for EvaluationAnalyticsService.GenerateLossClusters.""", ) -class RubricBasedMetricResultDict(TypedDict, total=False): - """Result for a rubric-based metric.""" +class GenerateLossClustersOperationDict(TypedDict, total=False): + """Long-running operation for generating loss clusters.""" - score: Optional[float] - """Passing rate of all the 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}`.""" - rubric_verdicts: Optional[list[evals_types.RubricVerdict]] - """The details of all the rubrics and their verdicts.""" + 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.""" -RubricBasedMetricResultOrDict = Union[ - RubricBasedMetricResult, RubricBasedMetricResultDict -] + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + response: Optional[GenerateLossClustersResponseDict] + """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" -class CometResult(_common.BaseModel): - """Spec for Comet result - calculates the comet score for the given instance using the version specified in the spec.""" - score: Optional[float] = Field( - default=None, - description="""Output only. Comet score. Range depends on version.""", +GenerateLossClustersOperationOrDict = Union[ + GenerateLossClustersOperation, GenerateLossClustersOperationDict +] + + +class RubricGenerationConfig(_common.BaseModel): + """Config for generating rubrics.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -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 RubricGenerationConfigDict(TypedDict, total=False): + """Config for generating rubrics.""" - score: Optional[float] - """Output only. Comet score. Range depends on version.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -CometResultOrDict = Union[CometResult, CometResultDict] +RubricGenerationConfigOrDict = Union[RubricGenerationConfig, RubricGenerationConfigDict] -class MetricxResult(_common.BaseModel): - """Spec for MetricX result - calculates the MetricX score for the given instance using the version specified in the spec.""" +class _GenerateInstanceRubricsRequest(_common.BaseModel): + """Parameters for generating rubrics.""" - score: Optional[float] = Field( + contents: Optional[list[genai_types.Content]] = Field( default=None, - description="""Output only. MetricX score. Range depends on version.""", + 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 MetricxResultDict(TypedDict, total=False): - """Spec for MetricX result - calculates the MetricX score for the given instance using the version specified in the spec.""" +class _GenerateInstanceRubricsRequestDict(TypedDict, total=False): + """Parameters for generating rubrics.""" - score: Optional[float] - """Output only. MetricX score. Range depends on version.""" + contents: Optional[list[genai_types.Content]] + """The prompt to generate rubrics from. For single-turn queries, this is a single instance. For multi-turn queries, this is a repeated field that contains conversation history + latest request.""" + + predefined_rubric_generation_spec: Optional[genai_types.PredefinedMetricSpec] + """Specification for using the rubric generation configs of a pre-defined + metric, e.g. "generic_quality_v1" and "instruction_following_v1". + Some of the configs may be only used in rubric generation and not + supporting evaluation, e.g. "fully_customized_generic_quality_v1". + If this field is set, the `rubric_generation_spec` field will be ignored. + """ + rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] + """Specification for how the rubrics should be generated.""" -MetricxResultOrDict = Union[MetricxResult, MetricxResultDict] + 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] + """""" -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.""" +_GenerateInstanceRubricsRequestOrDict = Union[ + _GenerateInstanceRubricsRequest, _GenerateInstanceRubricsRequestDict +] + + +class GenerateInstanceRubricsResponse(_common.BaseModel): + """Response for generating rubrics.""" + + generated_rubrics: Optional[list[evals_types.Rubric]] = Field( + default=None, description="""A list of generated rubrics.""" ) -class ToolCallValidMetricValueDict(TypedDict, total=False): - """Tool call valid metric value for an instance.""" +class GenerateInstanceRubricsResponseDict(TypedDict, total=False): + """Response for generating rubrics.""" - score: Optional[float] - """Output only. Tool call valid score.""" + generated_rubrics: Optional[list[evals_types.Rubric]] + """A list of generated rubrics.""" -ToolCallValidMetricValueOrDict = Union[ - ToolCallValidMetricValue, ToolCallValidMetricValueDict +GenerateInstanceRubricsResponseOrDict = Union[ + GenerateInstanceRubricsResponse, GenerateInstanceRubricsResponseDict ] -class ToolCallValidResults(_common.BaseModel): - """Results for tool call valid metric.""" +class GetEvaluationMetricConfig(_common.BaseModel): + """Config for getting an evaluation metric.""" - tool_call_valid_metric_values: Optional[list[ToolCallValidMetricValue]] = Field( - default=None, description="""Output only. Tool call valid metric values.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class ToolCallValidResultsDict(TypedDict, total=False): - """Results for tool call valid metric.""" +class GetEvaluationMetricConfigDict(TypedDict, total=False): + """Config for getting an evaluation metric.""" - tool_call_valid_metric_values: Optional[list[ToolCallValidMetricValueDict]] - """Output only. Tool call valid metric values.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ToolCallValidResultsOrDict = Union[ToolCallValidResults, ToolCallValidResultsDict] +GetEvaluationMetricConfigOrDict = Union[ + GetEvaluationMetricConfig, GetEvaluationMetricConfigDict +] -class ToolNameMatchMetricValue(_common.BaseModel): - """Tool name match metric value for an instance.""" +class _GetEvaluationMetricParameters(_common.BaseModel): + """Parameters for getting an evaluation metric.""" - score: Optional[float] = Field( - default=None, description="""Output only. Tool name match score.""" + metric_resource_name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationMetricConfig] = Field( + default=None, description="""""" ) -class ToolNameMatchMetricValueDict(TypedDict, total=False): - """Tool name match metric value for an instance.""" +class _GetEvaluationMetricParametersDict(TypedDict, total=False): + """Parameters for getting an evaluation metric.""" - score: Optional[float] - """Output only. Tool name match score.""" + metric_resource_name: Optional[str] + """""" + config: Optional[GetEvaluationMetricConfigDict] + """""" -ToolNameMatchMetricValueOrDict = Union[ - ToolNameMatchMetricValue, ToolNameMatchMetricValueDict + +_GetEvaluationMetricParametersOrDict = Union[ + _GetEvaluationMetricParameters, _GetEvaluationMetricParametersDict ] -class ToolNameMatchResults(_common.BaseModel): - """Results for tool name match metric.""" +class GetEvaluationRunConfig(_common.BaseModel): + """Config for get evaluation run.""" - 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 GetEvaluationRunConfigDict(TypedDict, total=False): + """Config for get evaluation run.""" - 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] +GetEvaluationRunConfigOrDict = Union[GetEvaluationRunConfig, GetEvaluationRunConfigDict] -class ToolParameterKeyMatchMetricValue(_common.BaseModel): - """Tool parameter key match metric value for an instance.""" +class _GetEvaluationRunParameters(_common.BaseModel): + """Represents a job that runs evaluation.""" - score: Optional[float] = Field( - default=None, description="""Output only. Tool parameter key match score.""" - ) + name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationRunConfig] = Field(default=None, description="""""") -class ToolParameterKeyMatchMetricValueDict(TypedDict, total=False): - """Tool parameter key match metric value for an instance.""" +class _GetEvaluationRunParametersDict(TypedDict, total=False): + """Represents a job that runs evaluation.""" - score: Optional[float] - """Output only. Tool parameter key match score.""" + name: Optional[str] + """""" + + config: Optional[GetEvaluationRunConfigDict] + """""" -ToolParameterKeyMatchMetricValueOrDict = Union[ - ToolParameterKeyMatchMetricValue, ToolParameterKeyMatchMetricValueDict +_GetEvaluationRunParametersOrDict = Union[ + _GetEvaluationRunParameters, _GetEvaluationRunParametersDict ] -class ToolParameterKeyMatchResults(_common.BaseModel): - """Results for tool parameter key match metric.""" +class GetEvaluationSetConfig(_common.BaseModel): + """Config for get evaluation set.""" - tool_parameter_key_match_metric_values: Optional[ - list[ToolParameterKeyMatchMetricValue] - ] = Field( - default=None, - description="""Output only. Tool parameter key match metric values.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class ToolParameterKeyMatchResultsDict(TypedDict, total=False): - """Results for tool parameter key match metric.""" +class GetEvaluationSetConfigDict(TypedDict, total=False): + """Config for get evaluation set.""" - tool_parameter_key_match_metric_values: Optional[ - list[ToolParameterKeyMatchMetricValueDict] - ] - """Output only. Tool parameter key match metric values.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ToolParameterKeyMatchResultsOrDict = Union[ - ToolParameterKeyMatchResults, ToolParameterKeyMatchResultsDict -] +GetEvaluationSetConfigOrDict = Union[GetEvaluationSetConfig, GetEvaluationSetConfigDict] -class ToolParameterKVMatchMetricValue(_common.BaseModel): - """Tool parameter key value match metric value for an instance.""" +class _GetEvaluationSetParameters(_common.BaseModel): + """Represents a job that gets an evaluation set.""" - score: Optional[float] = Field( - default=None, - description="""Output only. Tool parameter key value match score.""", - ) + name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationSetConfig] = Field(default=None, description="""""") -class ToolParameterKVMatchMetricValueDict(TypedDict, total=False): - """Tool parameter key value match metric value for an instance.""" +class _GetEvaluationSetParametersDict(TypedDict, total=False): + """Represents a job that gets an evaluation set.""" - score: Optional[float] - """Output only. Tool parameter key value match score.""" + name: Optional[str] + """""" + config: Optional[GetEvaluationSetConfigDict] + """""" -ToolParameterKVMatchMetricValueOrDict = Union[ - ToolParameterKVMatchMetricValue, ToolParameterKVMatchMetricValueDict + +_GetEvaluationSetParametersOrDict = Union[ + _GetEvaluationSetParameters, _GetEvaluationSetParametersDict ] -class ToolParameterKVMatchResults(_common.BaseModel): - """Results for tool parameter key value match metric.""" +class GetEvaluationItemConfig(_common.BaseModel): + """Config for get evaluation item.""" - tool_parameter_kv_match_metric_values: Optional[ - list[ToolParameterKVMatchMetricValue] - ] = Field( - default=None, - description="""Output only. Tool parameter key value match metric values.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class ToolParameterKVMatchResultsDict(TypedDict, total=False): - """Results for tool parameter key value match metric.""" - - tool_parameter_kv_match_metric_values: Optional[ - list[ToolParameterKVMatchMetricValueDict] - ] - """Output only. Tool parameter key value match metric values.""" +class GetEvaluationItemConfigDict(TypedDict, total=False): + """Config for get evaluation item.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -ToolParameterKVMatchResultsOrDict = Union[ - ToolParameterKVMatchResults, ToolParameterKVMatchResultsDict -] - - -class EvaluateInstancesResponse(_common.BaseModel): - """Result of evaluating an LLM metric.""" - - rubric_based_metric_result: Optional[RubricBasedMetricResult] = Field( - default=None, description="""Result for rubric based metric.""" - ) - metric_results: Optional[list[MetricResult]] = Field( - default=None, - description="""A list of metric results for each evaluation case. The order of the metric results is guaranteed to be the same as the order of the instances in the request.""", - ) - bleu_results: Optional[BleuResults] = Field( - default=None, description="""Results for bleu metric.""" - ) - comet_result: Optional[CometResult] = Field( - default=None, description="""Translation metrics. Result for Comet metric.""" - ) - exact_match_results: Optional[ExactMatchResults] = Field( - default=None, - description="""Auto metric evaluation results. Results for exact match metric.""", - ) - metricx_result: Optional[MetricxResult] = Field( - default=None, description="""Result for Metricx metric.""" - ) - pairwise_metric_result: Optional[genai_types.PairwiseMetricResult] = Field( - default=None, description="""Result for pairwise metric.""" - ) - pointwise_metric_result: Optional[genai_types.PointwiseMetricResult] = Field( - default=None, description="""Generic metrics. Result for pointwise metric.""" - ) - rouge_results: Optional[RougeResults] = Field( - default=None, description="""Results for rouge metric.""" - ) - tool_call_valid_results: Optional[ToolCallValidResults] = Field( - default=None, - description="""Tool call metrics. Results for tool call valid metric.""", - ) - tool_name_match_results: Optional[ToolNameMatchResults] = Field( - default=None, description="""Results for tool name match metric.""" - ) - tool_parameter_key_match_results: Optional[ToolParameterKeyMatchResults] = Field( - default=None, description="""Results for tool parameter key match metric.""" - ) - tool_parameter_kv_match_results: Optional[ToolParameterKVMatchResults] = Field( - default=None, - description="""Results for tool parameter key value match metric.""", - ) - - -class EvaluateInstancesResponseDict(TypedDict, total=False): - """Result of evaluating an LLM metric.""" - - rubric_based_metric_result: Optional[RubricBasedMetricResultDict] - """Result for rubric based metric.""" - - metric_results: Optional[list[MetricResultDict]] - """A list of metric results for each evaluation case. The order of the metric results is guaranteed to be the same as the order of the instances in the request.""" - - bleu_results: Optional[BleuResultsDict] - """Results for bleu metric.""" - - comet_result: Optional[CometResultDict] - """Translation metrics. Result for Comet metric.""" - - exact_match_results: Optional[ExactMatchResultsDict] - """Auto metric evaluation results. Results for exact match metric.""" - metricx_result: Optional[MetricxResultDict] - """Result for Metricx metric.""" +GetEvaluationItemConfigOrDict = Union[ + GetEvaluationItemConfig, GetEvaluationItemConfigDict +] - 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 _GetEvaluationItemParameters(_common.BaseModel): + """Represents a job that gets an evaluation item.""" - rouge_results: Optional[RougeResultsDict] - """Results for rouge metric.""" + name: Optional[str] = Field(default=None, description="""""") + config: Optional[GetEvaluationItemConfig] = Field(default=None, description="""""") - tool_call_valid_results: Optional[ToolCallValidResultsDict] - """Tool call metrics. Results for tool call valid metric.""" - tool_name_match_results: Optional[ToolNameMatchResultsDict] - """Results for tool name match metric.""" +class _GetEvaluationItemParametersDict(TypedDict, total=False): + """Represents a job that gets an evaluation item.""" - tool_parameter_key_match_results: Optional[ToolParameterKeyMatchResultsDict] - """Results for tool parameter key match metric.""" + name: Optional[str] + """""" - tool_parameter_kv_match_results: Optional[ToolParameterKVMatchResultsDict] - """Results for tool parameter key value match metric.""" + config: Optional[GetEvaluationItemConfigDict] + """""" -EvaluateInstancesResponseOrDict = Union[ - EvaluateInstancesResponse, EvaluateInstancesResponseDict +_GetEvaluationItemParametersOrDict = Union[ + _GetEvaluationItemParameters, _GetEvaluationItemParametersDict ] -class GenerateUserScenariosConfig(_common.BaseModel): +class ListEvaluationMetricsConfig(_common.BaseModel): + """Config for listing evaluation metrics.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + filter: Optional[str] = Field( + default=None, + description="""An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported. + For more information about filter syntax, see + `AIP-160 `_.""", + ) + order_by: Optional[str] = Field( + default=None, + description="""A comma-separated list of fields to order by, sorted in ascending + order by default. Use ``desc`` after a field name for descending. + Example: ``"create_time desc"``.""", + ) -class GenerateUserScenariosConfigDict(TypedDict, total=False): +class ListEvaluationMetricsConfigDict(TypedDict, total=False): + """Config for listing evaluation metrics.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" + page_size: Optional[int] + """""" -GenerateUserScenariosConfigOrDict = Union[ - GenerateUserScenariosConfig, GenerateUserScenariosConfigDict -] + page_token: Optional[str] + """""" + filter: Optional[str] + """An expression for filtering the results of the request. + For field names both snake_case and camelCase are supported. + For more information about filter syntax, see + `AIP-160 `_.""" -class _GenerateUserScenariosParameters(_common.BaseModel): - """Parameters for GenerateUserScenarios.""" + 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"``.""" - location: Optional[str] = Field(default=None, description="""""") - agents: Optional[dict[str, evals_types.AgentConfig]] = Field( - default=None, description="""""" - ) - root_agent_id: Optional[str] = Field(default=None, description="""""") - user_scenario_generation_config: Optional[ - evals_types.UserScenarioGenerationConfig - ] = Field(default=None, description="""""") - config: Optional[GenerateUserScenariosConfig] = Field( - default=None, description="""""" - ) - allow_cross_region_model: Optional[bool] = Field( - default=None, - description="""Opt-in flag to authorize cross-region routing for LLM models.""", - ) - gemini_agent_config: Optional[GeminiAgentConfig] = Field( - default=None, - description="""If set, the server derives the agents map and root_agent_id - from the referenced Gemini Agent server-side.""", - ) +ListEvaluationMetricsConfigOrDict = Union[ + ListEvaluationMetricsConfig, ListEvaluationMetricsConfigDict +] -class _GenerateUserScenariosParametersDict(TypedDict, total=False): - """Parameters for GenerateUserScenarios.""" - location: Optional[str] - """""" +class _ListEvaluationMetricsParameters(_common.BaseModel): + """Parameters for listing evaluation metrics.""" - agents: Optional[dict[str, evals_types.AgentConfig]] - """""" + config: Optional[ListEvaluationMetricsConfig] = Field( + default=None, description="""""" + ) - root_agent_id: Optional[str] - """""" - user_scenario_generation_config: Optional[evals_types.UserScenarioGenerationConfig] - """""" +class _ListEvaluationMetricsParametersDict(TypedDict, total=False): + """Parameters for listing evaluation metrics.""" - config: Optional[GenerateUserScenariosConfigDict] + config: Optional[ListEvaluationMetricsConfigDict] """""" - 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 +_ListEvaluationMetricsParametersOrDict = Union[ + _ListEvaluationMetricsParameters, _ListEvaluationMetricsParametersDict ] -class GenerateUserScenariosResponse(_common.BaseModel): - """Response message for DataFoundryService.GenerateUserScenarios.""" +class ListEvaluationMetricsResponse(_common.BaseModel): + """Response for listing evaluation metrics.""" - user_scenarios: Optional[list[evals_types.UserScenario]] = Field( - default=None, description="""""" + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" + ) + next_page_token: Optional[str] = Field(default=None, description="""""") + evaluation_metrics: Optional[list[EvaluationMetric]] = Field( + default=None, + description="""List of evaluation metrics. + """, ) -class GenerateUserScenariosResponseDict(TypedDict, total=False): - """Response message for DataFoundryService.GenerateUserScenarios.""" +class ListEvaluationMetricsResponseDict(TypedDict, total=False): + """Response for listing evaluation metrics.""" - user_scenarios: Optional[list[evals_types.UserScenario]] + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + + next_page_token: Optional[str] """""" + evaluation_metrics: Optional[list[EvaluationMetricDict]] + """List of evaluation metrics. + """ -GenerateUserScenariosResponseOrDict = Union[ - GenerateUserScenariosResponse, GenerateUserScenariosResponseDict + +ListEvaluationMetricsResponseOrDict = Union[ + ListEvaluationMetricsResponse, ListEvaluationMetricsResponseDict ] -class GenerateLossClustersConfig(_common.BaseModel): - """Config for generating loss clusters.""" +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 GenerateLossClustersConfigDict(TypedDict, total=False): - """Config for generating loss clusters.""" +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.""" -GenerateLossClustersConfigOrDict = Union[ - GenerateLossClustersConfig, GenerateLossClustersConfigDict -] + 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 _GenerateLossClustersParameters(_common.BaseModel): - """Parameters for GenerateLossClusters.""" - - location: Optional[str] = Field( - default=None, - description="""The resource name of the Location. Format: `projects/{project}/locations/{location}`.""", - ) - evaluation_set: Optional[str] = Field( - default=None, - description="""Reference to a persisted EvaluationSet. The service will read items from this set.""", - ) - inline_results: Optional[list[EvaluationResult]] = Field( - default=None, - description="""Inline evaluation results. Useful for ephemeral analysis in notebooks/SDKs where data isn't persisted.""", - ) - configs: Optional[list[LossAnalysisConfig]] = Field( - default=None, - description="""Configuration for the analysis algorithm. Analysis for multiple metrics and multiple candidates could be specified.""", - ) - config: Optional[GenerateLossClustersConfig] = Field( - default=None, description="""Config for generating loss clusters.""" - ) +OptimizeConfigOrDict = Union[OptimizeConfig, OptimizeConfigDict] -class _GenerateLossClustersParametersDict(TypedDict, total=False): - """Parameters for GenerateLossClusters.""" +class _OptimizeRequestParameters(_common.BaseModel): + """Request for the optimize_prompt method.""" - location: Optional[str] - """The resource name of the Location. Format: `projects/{project}/locations/{location}`.""" + content: Optional[genai_types.Content] = Field(default=None, description="""""") + config: Optional[OptimizeConfig] = Field(default=None, description="""""") - evaluation_set: Optional[str] - """Reference to a persisted EvaluationSet. The service will read items from this set.""" - inline_results: Optional[list[EvaluationResultDict]] - """Inline evaluation results. Useful for ephemeral analysis in notebooks/SDKs where data isn't persisted.""" +class _OptimizeRequestParametersDict(TypedDict, total=False): + """Request for the optimize_prompt method.""" - configs: Optional[list[LossAnalysisConfigDict]] - """Configuration for the analysis algorithm. Analysis for multiple metrics and multiple candidates could be specified.""" + content: Optional[genai_types.Content] + """""" - config: Optional[GenerateLossClustersConfigDict] - """Config for generating loss clusters.""" + config: Optional[OptimizeConfigDict] + """""" -_GenerateLossClustersParametersOrDict = Union[ - _GenerateLossClustersParameters, _GenerateLossClustersParametersDict +_OptimizeRequestParametersOrDict = Union[ + _OptimizeRequestParameters, _OptimizeRequestParametersDict ] -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 OptimizeResponseEndpoint(_common.BaseModel): + """Response for the optimize_prompt method.""" + content: Optional[genai_types.Content] = 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 OptimizeResponseEndpointDict(TypedDict, total=False): + """Response for the optimize_prompt method.""" - results: Optional[list[LossAnalysisResultDict]] - """The analysis results, one per config provided in the request.""" + content: Optional[genai_types.Content] + """""" -GenerateLossClustersResponseOrDict = Union[ - GenerateLossClustersResponse, GenerateLossClustersResponseDict +OptimizeResponseEndpointOrDict = Union[ + OptimizeResponseEndpoint, OptimizeResponseEndpointDict ] -class GenerateLossClustersOperation(_common.BaseModel): - """Long-running operation for generating loss clusters.""" +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.""" - 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( + domain: 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. The DNS name suffix of the zone being peered to, e.g., "my-internal-domain.corp.". Must end with a dot.""", ) - error: Optional[dict[str, Any]] = Field( + target_network: Optional[str] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""Required. The VPC network name in the target_project where the DNS zone specified by 'domain' is visible.""", ) - response: Optional[GenerateLossClustersResponse] = Field( + target_project: Optional[str] = Field( default=None, - description="""Response message for EvaluationAnalyticsService.GenerateLossClusters.""", + 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 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 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.""" - 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.""" + 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.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + target_network: Optional[str] + """Required. The VPC network name in the target_project where the DNS zone specified by 'domain' is visible.""" - response: Optional[GenerateLossClustersResponseDict] - """Response message for EvaluationAnalyticsService.GenerateLossClusters.""" + 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.""" -GenerateLossClustersOperationOrDict = Union[ - GenerateLossClustersOperation, GenerateLossClustersOperationDict -] +DnsPeeringConfigOrDict = Union[DnsPeeringConfig, DnsPeeringConfigDict] -class RubricGenerationConfig(_common.BaseModel): - """Config for generating rubrics.""" +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 RubricGenerationConfigDict(TypedDict, total=False): - """Config for generating rubrics.""" +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.""" -RubricGenerationConfigOrDict = Union[RubricGenerationConfig, RubricGenerationConfigDict] + +PscInterfaceConfigOrDict = Union[PscInterfaceConfig, PscInterfaceConfigDict] -class _GenerateInstanceRubricsRequest(_common.BaseModel): - """Parameters for generating rubrics.""" +class Scheduling(_common.BaseModel): + """All parameters related to queuing and scheduling of custom jobs.""" - contents: Optional[list[genai_types.Content]] = Field( + disable_retries: Optional[bool] = Field( default=None, - description="""The prompt to generate rubrics from. For single-turn queries, this is a single instance. For multi-turn queries, this is a repeated field that contains conversation history + latest request.""", + description="""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.""", ) - 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. - """, - ) + 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.""", ) - rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] = Field( + restart_job_on_worker_restart: Optional[bool] = Field( default=None, - description="""Specification for how the rubrics should be generated.""", + 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.""", ) - metric_resource_name: Optional[str] = Field( + strategy: Optional[Strategy] = 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.""", + 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.""", ) - config: Optional[RubricGenerationConfig] = Field(default=None, description="""""") -class _GenerateInstanceRubricsRequestDict(TypedDict, total=False): - """Parameters for generating rubrics.""" +class SchedulingDict(TypedDict, total=False): + """All parameters related to queuing and scheduling of custom jobs.""" - 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.""" + 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.""" - 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. - """ + 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.""" - rubric_generation_spec: Optional[genai_types.RubricGenerationSpec] - """Specification for how the rubrics should be generated.""" + 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.""" - 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.""" + strategy: Optional[Strategy] + """Optional. This determines which type of scheduling strategy to use.""" - config: Optional[RubricGenerationConfigDict] - """""" + timeout: Optional[str] + """Optional. The maximum job running time. The default is 7 days.""" -_GenerateInstanceRubricsRequestOrDict = Union[ - _GenerateInstanceRubricsRequest, _GenerateInstanceRubricsRequestDict -] +SchedulingOrDict = Union[Scheduling, SchedulingDict] -class GenerateInstanceRubricsResponse(_common.BaseModel): - """Response for generating rubrics.""" +class EnvVar(_common.BaseModel): + """Represents an environment variable present in a Container or Python Module.""" - generated_rubrics: Optional[list[evals_types.Rubric]] = Field( - default=None, description="""A list of generated rubrics.""" + 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 GenerateInstanceRubricsResponseDict(TypedDict, total=False): - """Response for generating rubrics.""" - - generated_rubrics: Optional[list[evals_types.Rubric]] - """A list of generated rubrics.""" +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.""" -GenerateInstanceRubricsResponseOrDict = Union[ - GenerateInstanceRubricsResponse, GenerateInstanceRubricsResponseDict -] + 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.""" -class GetEvaluationMetricConfig(_common.BaseModel): - """Config for getting an evaluation metric.""" +EnvVarOrDict = Union[EnvVar, EnvVarDict] - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) +class ContainerSpec(_common.BaseModel): + """The spec of a Container.""" -class GetEvaluationMetricConfigDict(TypedDict, total=False): - """Config for getting an evaluation metric.""" + 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.""", + ) - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class ContainerSpecDict(TypedDict, total=False): + """The spec of a Container.""" -GetEvaluationMetricConfigOrDict = Union[ - GetEvaluationMetricConfig, GetEvaluationMetricConfigDict -] + 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.""" -class _GetEvaluationMetricParameters(_common.BaseModel): - """Parameters for getting an evaluation metric.""" + env: Optional[list[EnvVarDict]] + """Environment variables to be passed to the container. Maximum limit is 100.""" - metric_resource_name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationMetricConfig] = Field( - default=None, description="""""" - ) + image_uri: Optional[str] + """Required. The URI of a container image in the Container Registry that is to be run on each worker replica.""" -class _GetEvaluationMetricParametersDict(TypedDict, total=False): - """Parameters for getting an evaluation metric.""" +ContainerSpecOrDict = Union[ContainerSpec, ContainerSpecDict] - metric_resource_name: Optional[str] - """""" - config: Optional[GetEvaluationMetricConfigDict] - """""" +class DiskSpec(_common.BaseModel): + """Represents the spec of disk options.""" + boot_disk_size_gb: Optional[int] = Field( + default=None, description="""Size in GB of the boot disk (default is 100GB).""" + ) + 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".""", + ) + local_ssd_count: Optional[int] = Field( + default=None, + description="""Number of attached local SSDs, from 0 to 8 (default is 0).""", + ) -_GetEvaluationMetricParametersOrDict = Union[ - _GetEvaluationMetricParameters, _GetEvaluationMetricParametersDict -] +class DiskSpecDict(TypedDict, total=False): + """Represents the spec of disk options.""" -class GetEvaluationRunConfig(_common.BaseModel): - """Config for get evaluation run.""" + boot_disk_size_gb: Optional[int] + """Size in GB of the boot disk (default is 100GB).""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) + 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".""" + local_ssd_count: Optional[int] + """Number of attached local SSDs, from 0 to 8 (default is 0).""" -class GetEvaluationRunConfigDict(TypedDict, total=False): - """Config for get evaluation run.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +DiskSpecOrDict = Union[DiskSpec, DiskSpecDict] -GetEvaluationRunConfigOrDict = Union[GetEvaluationRunConfig, GetEvaluationRunConfigDict] +class LustreMount(_common.BaseModel): + """Represents a mount configuration for Lustre file system.""" + filesystem: Optional[str] = Field( + default=None, description="""Required. The name of the Lustre filesystem.""" + ) + instance_ip: Optional[str] = Field( + default=None, description="""Required. IP address of the Lustre instance.""" + ) + mount_point: Optional[str] = Field( + default=None, + description="""Required. Destination mount path. The Lustre file system will be mounted for the user under /mnt/lustre/""", + ) + volume_handle: Optional[str] = Field( + default=None, + description="""Required. The unique identifier of the Lustre volume.""", + ) -class _GetEvaluationRunParameters(_common.BaseModel): - """Represents a job that runs evaluation.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationRunConfig] = Field(default=None, description="""""") +class LustreMountDict(TypedDict, total=False): + """Represents a mount configuration for Lustre file system.""" + filesystem: Optional[str] + """Required. The name of the Lustre filesystem.""" -class _GetEvaluationRunParametersDict(TypedDict, total=False): - """Represents a job that runs evaluation.""" + instance_ip: Optional[str] + """Required. IP address of the Lustre instance.""" - name: Optional[str] - """""" + mount_point: Optional[str] + """Required. Destination mount path. The Lustre file system will be mounted for the user under /mnt/lustre/""" - config: Optional[GetEvaluationRunConfigDict] - """""" + volume_handle: Optional[str] + """Required. The unique identifier of the Lustre volume.""" -_GetEvaluationRunParametersOrDict = Union[ - _GetEvaluationRunParameters, _GetEvaluationRunParametersDict -] +LustreMountOrDict = Union[LustreMount, LustreMountDict] -class GetEvaluationSetConfig(_common.BaseModel): - """Config for get evaluation set.""" +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.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + key: Optional[str] = Field( + default=None, + description="""Optional. Corresponds to the label key of a reservation resource. To target a SPECIFIC_RESERVATION by name, use `compute.googleapis.com/reservation-name` as the key and specify the name of your reservation as its value.""", + ) + reservation_affinity_type: Optional[Type] = Field( + default=None, + description="""Required. Specifies the reservation affinity type.""", + ) + values: Optional[list[str]] = Field( + default=None, + description="""Optional. Corresponds to the label values of a reservation resource. This must be the full resource name of the reservation or reservation block.""", + ) + use_reservation_pool: Optional[bool] = Field( + default=None, + description="""Optional. When set to true, resources will be drawn from go/cloud-ai-gcp-pool.""", ) -class GetEvaluationSetConfigDict(TypedDict, total=False): - """Config for get evaluation set.""" +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.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + 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.""" + reservation_affinity_type: Optional[Type] + """Required. Specifies the reservation affinity type.""" -GetEvaluationSetConfigOrDict = Union[GetEvaluationSetConfig, GetEvaluationSetConfigDict] + 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.""" + use_reservation_pool: Optional[bool] + """Optional. When set to true, resources will be drawn from go/cloud-ai-gcp-pool.""" -class _GetEvaluationSetParameters(_common.BaseModel): - """Represents a job that gets an evaluation set.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationSetConfig] = Field(default=None, description="""""") +ReservationAffinityOrDict = Union[ReservationAffinity, ReservationAffinityDict] -class _GetEvaluationSetParametersDict(TypedDict, total=False): - """Represents a job that gets an evaluation set.""" +class ConfidentialInstanceConfig(_common.BaseModel): + """A set of Confidential Instance options.""" - name: Optional[str] - """""" + confidential_instance_type: Optional[ConfidentialInstanceType] = Field( + default=None, + description="""Defines the type of technology used by the confidential instance.""", + ) - config: Optional[GetEvaluationSetConfigDict] - """""" +class ConfidentialInstanceConfigDict(TypedDict, total=False): + """A set of Confidential Instance options.""" -_GetEvaluationSetParametersOrDict = Union[ - _GetEvaluationSetParameters, _GetEvaluationSetParametersDict + confidential_instance_type: Optional[ConfidentialInstanceType] + """Defines the type of technology used by the confidential instance.""" + + +ConfidentialInstanceConfigOrDict = Union[ + ConfidentialInstanceConfig, ConfidentialInstanceConfigDict ] -class GetEvaluationItemConfig(_common.BaseModel): - """Config for get evaluation item.""" +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.""", + ) + accelerator_type: Optional[AcceleratorType] = Field( + default=None, + description="""Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""", + ) + gpu_partition_size: Optional[str] = Field( + default=None, + description="""Optional. Immutable. The Nvidia GPU partition size. When specified, the requested accelerators will be partitioned into smaller GPU partitions. For example, if the request is for 8 units of NVIDIA A100 GPUs, and gpu_partition_size="1g.10gb", the service will create 8 * 7 = 56 partitioned MIG instances. The partition size must be a value supported by the requested accelerator. Refer to [Nvidia GPU Partitioning](https://cloud.google.com/kubernetes-engine/docs/how-to/gpus-multi#multi-instance_gpu_partitions) for the available partition sizes. If set, the accelerator_count should be set to 1.""", + ) + machine_type: Optional[str] = Field( + default=None, + description="""Immutable. The type of the machine. See the [list of machine types supported for prediction](https://cloud.google.com/gemini-enterprise-agent-platform/machine-learning/predictions/configure-compute#machine-types) See the [list of machine types supported for custom training](https://cloud.google.com/gemini-enterprise-agent-platform/machine-learning/training/configure-compute#machine-types). For DeployedModel this field is optional, and the default value is `n1-standard-2`. For BatchPredictionJob or as part of WorkerPoolSpec this field is required.""", + ) + min_gpu_driver_version: Optional[str] = Field( + default=None, + description="""Optional. Immutable. The minimum GPU driver version that this machine requires. For example, "535.104.06". If not specified, the default GPU driver version will be used by the underlying infrastructure.""", + ) + multihost_gpu_node_count: Optional[int] = Field( + default=None, + description="""Optional. Immutable. The number of nodes per replica for multihost GPU deployments.""", + ) + reservation_affinity: Optional[ReservationAffinity] = Field( + default=None, + description="""Optional. Immutable. Configuration controlling how this resource pool consumes reservation.""", + ) + 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").""", + ) + confidential_instance_config: Optional[ConfidentialInstanceConfig] = Field( + default=None, + description="""Optional. Immutable. The confidential instance config for resource pool""", ) -class GetEvaluationItemConfigDict(TypedDict, total=False): - """Config for get evaluation item.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class MachineSpecDict(TypedDict, total=False): + """Specification of a single machine.""" + 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.""" -GetEvaluationItemConfigOrDict = Union[ - GetEvaluationItemConfig, GetEvaluationItemConfigDict -] + accelerator_type: Optional[AcceleratorType] + """Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""" + 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.""" -class _GetEvaluationItemParameters(_common.BaseModel): - """Represents a job that gets an evaluation item.""" + 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.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[GetEvaluationItemConfig] = Field(default=None, description="""""") + 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 _GetEvaluationItemParametersDict(TypedDict, total=False): - """Represents a job that gets an evaluation item.""" + reservation_affinity: Optional[ReservationAffinityDict] + """Optional. Immutable. Configuration controlling how this resource pool consumes reservation.""" - name: Optional[str] - """""" + tpu_topology: Optional[str] + """Immutable. The topology of the TPUs. Corresponds to the TPU topologies available from GKE. (Example: tpu_topology: "2x2x1").""" - config: Optional[GetEvaluationItemConfigDict] - """""" + confidential_instance_config: Optional[ConfidentialInstanceConfigDict] + """Optional. Immutable. The confidential instance config for resource pool""" -_GetEvaluationItemParametersOrDict = Union[ - _GetEvaluationItemParameters, _GetEvaluationItemParametersDict -] +MachineSpecOrDict = Union[MachineSpec, MachineSpecDict] -class ListEvaluationMetricsConfig(_common.BaseModel): - """Config for listing evaluation metrics.""" +class NfsMount(_common.BaseModel): + """Represents a mount configuration for Network File System (NFS) to mount.""" - 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( + mount_point: Optional[str] = Field( default=None, - description="""An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported. - For more information about filter syntax, see - `AIP-160 `_.""", + description="""Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""", ) - order_by: Optional[str] = Field( + path: 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. Source path exported from NFS server. Has to start with '/', and combined with the ip address, it indicates the source mount path in the form of `server:path`""", + ) + server: Optional[str] = Field( + default=None, description="""Required. IP address of the NFS server.""" ) -class ListEvaluationMetricsConfigDict(TypedDict, total=False): - """Config for listing evaluation metrics.""" - - 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. - For more information about filter syntax, see - `AIP-160 `_.""" - - order_by: Optional[str] - """A comma-separated list of fields to order by, sorted in ascending - order by default. Use ``desc`` after a field name for descending. - Example: ``"create_time desc"``.""" - - -ListEvaluationMetricsConfigOrDict = Union[ - ListEvaluationMetricsConfig, ListEvaluationMetricsConfigDict -] - - -class _ListEvaluationMetricsParameters(_common.BaseModel): - """Parameters for listing evaluation metrics.""" - - config: Optional[ListEvaluationMetricsConfig] = Field( - default=None, description="""""" - ) +class NfsMountDict(TypedDict, total=False): + """Represents a mount configuration for Network File System (NFS) to mount.""" + mount_point: Optional[str] + """Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""" -class _ListEvaluationMetricsParametersDict(TypedDict, total=False): - """Parameters for listing evaluation metrics.""" + 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`""" - config: Optional[ListEvaluationMetricsConfigDict] - """""" + server: Optional[str] + """Required. IP address of the NFS server.""" -_ListEvaluationMetricsParametersOrDict = Union[ - _ListEvaluationMetricsParameters, _ListEvaluationMetricsParametersDict -] +NfsMountOrDict = Union[NfsMount, NfsMountDict] -class ListEvaluationMetricsResponse(_common.BaseModel): - """Response for listing evaluation metrics.""" +class PythonPackageSpec(_common.BaseModel): + """The spec of a Python packaged code.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + args: Optional[list[str]] = Field( + default=None, + description="""Command line arguments to be passed to the Python task.""", ) - next_page_token: Optional[str] = Field(default=None, description="""""") - evaluation_metrics: Optional[list[EvaluationMetric]] = Field( + env: Optional[list[EnvVar]] = Field( default=None, - description="""List of evaluation metrics. - """, + description="""Environment variables to be passed to the python module. Maximum limit is 100.""", ) - - -class ListEvaluationMetricsResponseDict(TypedDict, total=False): - """Response for listing evaluation metrics.""" - - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" - - next_page_token: Optional[str] - """""" - - evaluation_metrics: Optional[list[EvaluationMetricDict]] - """List of evaluation metrics. - """ - - -ListEvaluationMetricsResponseOrDict = Union[ - ListEvaluationMetricsResponse, ListEvaluationMetricsResponseDict -] - - -class OptimizeConfig(_common.BaseModel): - """Config for Prompt Optimizer.""" - - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + executor_image_uri: Optional[str] = Field( + default=None, + description="""Required. The URI of a container image in Artifact Registry that will run the provided Python package. Vertex AI provides a wide range of executor images with pre-installed packages to meet users' various use cases. See the list of [pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers). You must use an image from this list.""", ) - optimization_target: Optional[OptimizeTarget] = Field( + package_uris: Optional[list[str]] = Field( default=None, - description="""The optimization target for the prompt optimizer. It must be one of the OptimizeTarget enum values: OPTIMIZATION_TARGET_GEMINI_NANO for the prompts from Android core API, OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS for the few-shot prompt optimizer with rubrics, OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE for the few-shot prompt optimizer with target responses.""", + description="""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.""", ) - examples_dataframe: Optional[PandasDataFrame] = Field( + python_module: Optional[str] = Field( default=None, - description="""The examples dataframe for the few-shot prompt optimizer. It must contain "prompt" and "model_response" columns. Depending on which optimization target is used, it also needs to contain "rubrics" and "rubrics_evaluations" or "target_response" columns.""", + description="""Required. The Python module name to run after installing the packages.""", ) -class OptimizeConfigDict(TypedDict, total=False): - """Config for Prompt Optimizer.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class PythonPackageSpecDict(TypedDict, total=False): + """The spec of a Python packaged code.""" - optimization_target: Optional[OptimizeTarget] - """The optimization target for the prompt optimizer. It must be one of the OptimizeTarget enum values: OPTIMIZATION_TARGET_GEMINI_NANO for the prompts from Android core API, OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS for the few-shot prompt optimizer with rubrics, OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE for the few-shot prompt optimizer with target responses.""" + args: Optional[list[str]] + """Command line arguments to be passed to the Python task.""" - examples_dataframe: Optional[PandasDataFrame] - """The examples dataframe for the few-shot prompt optimizer. It must contain "prompt" and "model_response" columns. Depending on which optimization target is used, it also needs to contain "rubrics" and "rubrics_evaluations" or "target_response" columns.""" + env: Optional[list[EnvVarDict]] + """Environment variables to be passed to the python module. Maximum limit is 100.""" + executor_image_uri: Optional[str] + """Required. The URI of a container image in Artifact Registry that will run the provided Python package. Vertex AI provides a wide range of executor images with pre-installed packages to meet users' various use cases. See the list of [pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers). You must use an image from this list.""" -OptimizeConfigOrDict = Union[OptimizeConfig, OptimizeConfigDict] + package_uris: Optional[list[str]] + """Required. The Google Cloud Storage location of the Python package files which are the training program and its dependent packages. The maximum number of package URIs is 100.""" + python_module: Optional[str] + """Required. The Python module name to run after installing the packages.""" -class _OptimizeRequestParameters(_common.BaseModel): - """Request for the optimize_prompt method.""" - content: Optional[genai_types.Content] = Field(default=None, description="""""") - config: Optional[OptimizeConfig] = Field(default=None, description="""""") +PythonPackageSpecOrDict = Union[PythonPackageSpec, PythonPackageSpecDict] -class _OptimizeRequestParametersDict(TypedDict, total=False): - """Request for the optimize_prompt method.""" +class WorkerPoolSpec(_common.BaseModel): + """Represents the spec of a worker pool in a job.""" - content: Optional[genai_types.Content] - """""" + container_spec: Optional[ContainerSpec] = Field( + default=None, description="""The custom container task.""" + ) + disk_spec: Optional[DiskSpec] = Field(default=None, description="""Disk spec.""") + lustre_mounts: Optional[list[LustreMount]] = Field( + default=None, description="""Optional. List of Lustre mounts.""" + ) + machine_spec: Optional[MachineSpec] = Field( + default=None, + description="""Optional. Immutable. The specification of a single machine.""", + ) + nfs_mounts: Optional[list[NfsMount]] = Field( + default=None, description="""Optional. List of NFS mount spec.""" + ) + python_package_spec: Optional[PythonPackageSpec] = Field( + default=None, description="""The Python packaged task.""" + ) + replica_count: Optional[int] = Field( + default=None, + description="""Optional. The number of worker replicas to use for this worker pool.""", + ) - config: Optional[OptimizeConfigDict] - """""" +class WorkerPoolSpecDict(TypedDict, total=False): + """Represents the spec of a worker pool in a job.""" -_OptimizeRequestParametersOrDict = Union[ - _OptimizeRequestParameters, _OptimizeRequestParametersDict -] + container_spec: Optional[ContainerSpecDict] + """The custom container task.""" + disk_spec: Optional[DiskSpecDict] + """Disk spec.""" -class OptimizeResponseEndpoint(_common.BaseModel): - """Response for the optimize_prompt method.""" + lustre_mounts: Optional[list[LustreMountDict]] + """Optional. List of Lustre mounts.""" - content: Optional[genai_types.Content] = Field(default=None, description="""""") + machine_spec: Optional[MachineSpecDict] + """Optional. Immutable. The specification of a single machine.""" + nfs_mounts: Optional[list[NfsMountDict]] + """Optional. List of NFS mount spec.""" -class OptimizeResponseEndpointDict(TypedDict, total=False): - """Response for the optimize_prompt method.""" + python_package_spec: Optional[PythonPackageSpecDict] + """The Python packaged task.""" - content: Optional[genai_types.Content] - """""" + replica_count: Optional[int] + """Optional. The number of worker replicas to use for this worker pool.""" -OptimizeResponseEndpointOrDict = Union[ - OptimizeResponseEndpoint, OptimizeResponseEndpointDict -] +WorkerPoolSpecOrDict = Union[WorkerPoolSpec, WorkerPoolSpecDict] -class DnsPeeringConfig(_common.BaseModel): - """DNS peering configuration. These configurations are used to create DNS peering zones in the Vertex tenant project VPC, enabling resolution of records within the specified domain hosted in the target network's Cloud DNS.""" +class CustomJobSpec(_common.BaseModel): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - domain: Optional[str] = Field( + base_output_directory: Optional[genai_types.GcsDestination] = Field( default=None, - description="""Required. The DNS name suffix of the zone being peered to, e.g., "my-internal-domain.corp.". Must end with a dot.""", + description="""The Cloud Storage location to store the output of this CustomJob or HyperparameterTuningJob. For HyperparameterTuningJob, the baseOutputDirectory of each child CustomJob backing a Trial is set to a subdirectory of name id under its parent HyperparameterTuningJob's baseOutputDirectory. The following Vertex AI environment variables will be passed to containers or python modules when this field is set: For CustomJob: * AIP_MODEL_DIR = `/model/` * AIP_CHECKPOINT_DIR = `/checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `/logs/` For CustomJob backing a Trial of HyperparameterTuningJob: * AIP_MODEL_DIR = `//model/` * AIP_CHECKPOINT_DIR = `//checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `//logs/`""", ) - target_network: Optional[str] = Field( + enable_dashboard_access: Optional[bool] = Field( default=None, - description="""Required. The VPC network name in the target_project where the DNS zone specified by 'domain' is visible.""", + description="""Optional. Whether you want Vertex AI to enable access to the customized dashboard in training chief container. If set to `true`, you can access the dashboard at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""", ) - target_project: Optional[str] = Field( + enable_web_access: Optional[bool] = Field( default=None, - description="""Required. The project ID hosting the Cloud DNS managed zone that contains the 'domain'. The Vertex AI Service Agent requires the dns.peer role on this project.""", + description="""Optional. Whether you want Vertex AI to enable [interactive shell access](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) to training containers. If set to `true`, you can access interactive shells at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""", ) - - -class DnsPeeringConfigDict(TypedDict, total=False): - """DNS peering configuration. These configurations are used to create DNS peering zones in the Vertex tenant project VPC, enabling resolution of records within the specified domain hosted in the target network's Cloud DNS.""" - - domain: Optional[str] - """Required. The DNS name suffix of the zone being peered to, e.g., "my-internal-domain.corp.". Must end with a dot.""" - - target_network: Optional[str] - """Required. The VPC network name in the target_project where the DNS zone specified by 'domain' is visible.""" - - target_project: Optional[str] - """Required. The project ID hosting the Cloud DNS managed zone that contains the 'domain'. The Vertex AI Service Agent requires the dns.peer role on this project.""" - - -DnsPeeringConfigOrDict = Union[DnsPeeringConfig, DnsPeeringConfigDict] - - -class PscInterfaceConfig(_common.BaseModel): - """Configuration for PSC-I.""" - - dns_peering_configs: Optional[list[DnsPeeringConfig]] = Field( + experiment: Optional[str] = Field( default=None, - description="""Optional. DNS peering configurations. When specified, Vertex AI will attempt to configure DNS peering zones in the tenant project VPC to resolve the specified domains using the target network's Cloud DNS. The user must grant the dns.peer role to the Vertex AI Service Agent on the target project.""", + description="""Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`""", ) - network_attachment: Optional[str] = Field( + experiment_run: Optional[str] = Field( default=None, - description="""Optional. The name of the Compute Engine [network attachment](https://cloud.google.com/vpc/docs/about-network-attachments) to attach to the resource within the region and user project. To specify this field, you must have already [created a network attachment] (https://cloud.google.com/vpc/docs/create-manage-network-attachments#create-network-attachments). This field is only used for resources using PSC-I.""", + description="""Optional. The Experiment Run associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}-{experiment-run-name}`""", ) - - -class PscInterfaceConfigDict(TypedDict, total=False): - """Configuration for PSC-I.""" - - dns_peering_configs: Optional[list[DnsPeeringConfigDict]] - """Optional. DNS peering configurations. When specified, Vertex AI will attempt to configure DNS peering zones in the tenant project VPC to resolve the specified domains using the target network's Cloud DNS. The user must grant the dns.peer role to the Vertex AI Service Agent on the target project.""" - - network_attachment: Optional[str] - """Optional. The name of the Compute Engine [network attachment](https://cloud.google.com/vpc/docs/about-network-attachments) to attach to the resource within the region and user project. To specify this field, you must have already [created a network attachment] (https://cloud.google.com/vpc/docs/create-manage-network-attachments#create-network-attachments). This field is only used for resources using PSC-I.""" - - -PscInterfaceConfigOrDict = Union[PscInterfaceConfig, PscInterfaceConfigDict] - - -class Scheduling(_common.BaseModel): - """All parameters related to queuing and scheduling of custom jobs.""" - - disable_retries: Optional[bool] = Field( + models: Optional[list[str]] = Field( default=None, - description="""Optional. Indicates if the job should retry for internal errors after the job starts running. If true, overrides `Scheduling.restart_job_on_worker_restart` to false.""", + description="""Optional. The name of the Model resources for which to generate a mapping to artifact URIs. Applicable only to some of the Google-provided custom jobs. Format: `projects/{project}/locations/{location}/models/{model}` In order to retrieve a specific version of the model, also provide the version ID or version alias. Example: `projects/{project}/locations/{location}/models/{model}@2` or `projects/{project}/locations/{location}/models/{model}@golden` If no version ID or alias is specified, the "default" version will be returned. The "default" version alias is created for the first version of the model, and can be moved to other versions later on. There will be exactly one default version.""", ) - max_wait_duration: Optional[str] = Field( + network: Optional[str] = Field( default=None, - description="""Optional. This is the maximum duration that a job will wait for the requested resources to be provisioned if the scheduling strategy is set to [Strategy.DWS_FLEX_START]. If set to 0, the job will wait indefinitely. The default is 24 hours.""", + description="""Optional. The full name of the Compute Engine [network](/compute/docs/networks-and-firewalls#networks) to which the Job should be peered. For example, `projects/12345/global/networks/myVPC`. [Format](/compute/docs/reference/rest/v1/networks/insert) is of the form `projects/{project}/global/networks/{network}`. Where {project} is a project number, as in `12345`, and {network} is a network name. To specify this field, you must have already [configured VPC Network Peering for Vertex AI](https://cloud.google.com/vertex-ai/docs/general/vpc-peering). If this field is left unspecified, the job is not peered with any network.""", ) - restart_job_on_worker_restart: Optional[bool] = Field( + persistent_resource_id: Optional[str] = Field( default=None, - description="""Optional. Restarts the entire CustomJob if a worker gets restarted. This feature can be used by distributed training jobs that are not resilient to workers leaving and joining a job.""", + description="""Optional. The ID of the PersistentResource in the same Project and Location which to run If this is specified, the job will be run on existing machines held by the PersistentResource instead of on-demand short-live machines. The network and CMEK configs on the job should be consistent with those on the PersistentResource, otherwise, the job will be rejected.""", ) - strategy: Optional[Strategy] = Field( + protected_artifact_location_id: Optional[str] = Field( default=None, - description="""Optional. This determines which type of scheduling strategy to use.""", + description="""The ID of the location to store protected artifacts. e.g. us-central1. Populate only when the location is different than CustomJob location. List of supported locations: https://cloud.google.com/vertex-ai/docs/general/locations""", ) - timeout: Optional[str] = Field( + psc_interface_config: Optional[PscInterfaceConfig] = Field( + default=None, description="""Optional. Configuration for PSC-I for CustomJob.""" + ) + reserved_ip_ranges: Optional[list[str]] = Field( default=None, - description="""Optional. The maximum job running time. The default is 7 days.""", + description="""Optional. A list of names for the reserved ip ranges under the VPC network that can be used for this job. If set, we will deploy the job within the provided ip ranges. Otherwise, the job will be deployed to any ip ranges under the provided VPC network. Example: ['vertex-ai-ip-range'].""", + ) + scheduling: Optional[Scheduling] = Field( + default=None, description="""Scheduling options for a CustomJob.""" + ) + service_account: Optional[str] = Field( + default=None, + description="""Specifies the service account for workload run-as account. Users submitting jobs must have act-as permission on this run-as account. If unspecified, the [Vertex AI Custom Code Service Agent](https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents) for the CustomJob's project is used.""", + ) + tensorboard: Optional[str] = Field( + default=None, + description="""Optional. The name of a Vertex AI Tensorboard resource to which this CustomJob will upload Tensorboard logs. Format: `projects/{project}/locations/{location}/tensorboards/{tensorboard}`""", + ) + worker_pool_specs: Optional[list[WorkerPoolSpec]] = Field( + default=None, + description="""Required. The spec of the worker pools including machine type and Docker image. All worker pools except the first one are optional and can be skipped by providing an empty value.""", ) -class SchedulingDict(TypedDict, total=False): - """All parameters related to queuing and scheduling of custom jobs.""" +class CustomJobSpecDict(TypedDict, total=False): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - disable_retries: Optional[bool] - """Optional. Indicates if the job should retry for internal errors after the job starts running. If true, overrides `Scheduling.restart_job_on_worker_restart` to false.""" + base_output_directory: Optional[genai_types.GcsDestination] + """The Cloud Storage location to store the output of this CustomJob or HyperparameterTuningJob. For HyperparameterTuningJob, the baseOutputDirectory of each child CustomJob backing a Trial is set to a subdirectory of name id under its parent HyperparameterTuningJob's baseOutputDirectory. The following Vertex AI environment variables will be passed to containers or python modules when this field is set: For CustomJob: * AIP_MODEL_DIR = `/model/` * AIP_CHECKPOINT_DIR = `/checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `/logs/` For CustomJob backing a Trial of HyperparameterTuningJob: * AIP_MODEL_DIR = `//model/` * AIP_CHECKPOINT_DIR = `//checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `//logs/`""" - max_wait_duration: Optional[str] - """Optional. This is the maximum duration that a job will wait for the requested resources to be provisioned if the scheduling strategy is set to [Strategy.DWS_FLEX_START]. If set to 0, the job will wait indefinitely. The default is 24 hours.""" + enable_dashboard_access: Optional[bool] + """Optional. Whether you want Vertex AI to enable access to the customized dashboard in training chief container. If set to `true`, you can access the dashboard at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""" - restart_job_on_worker_restart: Optional[bool] - """Optional. Restarts the entire CustomJob if a worker gets restarted. This feature can be used by distributed training jobs that are not resilient to workers leaving and joining a job.""" + enable_web_access: Optional[bool] + """Optional. Whether you want Vertex AI to enable [interactive shell access](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) to training containers. If set to `true`, you can access interactive shells at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""" - strategy: Optional[Strategy] - """Optional. This determines which type of scheduling strategy to use.""" + experiment: Optional[str] + """Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`""" - timeout: Optional[str] - """Optional. The maximum job running time. The default is 7 days.""" + experiment_run: Optional[str] + """Optional. The Experiment Run associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}-{experiment-run-name}`""" + models: Optional[list[str]] + """Optional. The name of the Model resources for which to generate a mapping to artifact URIs. Applicable only to some of the Google-provided custom jobs. Format: `projects/{project}/locations/{location}/models/{model}` In order to retrieve a specific version of the model, also provide the version ID or version alias. Example: `projects/{project}/locations/{location}/models/{model}@2` or `projects/{project}/locations/{location}/models/{model}@golden` If no version ID or alias is specified, the "default" version will be returned. The "default" version alias is created for the first version of the model, and can be moved to other versions later on. There will be exactly one default version.""" -SchedulingOrDict = Union[Scheduling, SchedulingDict] + network: Optional[str] + """Optional. The full name of the Compute Engine [network](/compute/docs/networks-and-firewalls#networks) to which the Job should be peered. For example, `projects/12345/global/networks/myVPC`. [Format](/compute/docs/reference/rest/v1/networks/insert) is of the form `projects/{project}/global/networks/{network}`. Where {project} is a project number, as in `12345`, and {network} is a network name. To specify this field, you must have already [configured VPC Network Peering for Vertex AI](https://cloud.google.com/vertex-ai/docs/general/vpc-peering). If this field is left unspecified, the job is not peered with any network.""" + persistent_resource_id: Optional[str] + """Optional. The ID of the PersistentResource in the same Project and Location which to run If this is specified, the job will be run on existing machines held by the PersistentResource instead of on-demand short-live machines. The network and CMEK configs on the job should be consistent with those on the PersistentResource, otherwise, the job will be rejected.""" -class EnvVar(_common.BaseModel): - """Represents an environment variable present in a Container or Python Module.""" + protected_artifact_location_id: Optional[str] + """The ID of the location to store protected artifacts. e.g. us-central1. Populate only when the location is different than CustomJob location. List of supported locations: https://cloud.google.com/vertex-ai/docs/general/locations""" - name: Optional[str] = Field( - default=None, - description="""Required. Name of the environment variable. Must be a valid C identifier.""", - ) - value: Optional[str] = Field( - default=None, - description="""Required. Variables that reference a $(VAR_NAME) are expanded using the previous defined environment variables in the container and any service environment variables. If a variable cannot be resolved, the reference in the input string will be unchanged. The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped references will never be expanded, regardless of whether the variable exists or not.""", - ) + psc_interface_config: Optional[PscInterfaceConfigDict] + """Optional. Configuration for PSC-I for CustomJob.""" + reserved_ip_ranges: Optional[list[str]] + """Optional. A list of names for the reserved ip ranges under the VPC network that can be used for this job. If set, we will deploy the job within the provided ip ranges. Otherwise, the job will be deployed to any ip ranges under the provided VPC network. Example: ['vertex-ai-ip-range'].""" -class EnvVarDict(TypedDict, total=False): - """Represents an environment variable present in a Container or Python Module.""" + scheduling: Optional[SchedulingDict] + """Scheduling options for a CustomJob.""" - name: Optional[str] - """Required. Name of the environment variable. Must be a valid C identifier.""" + service_account: Optional[str] + """Specifies the service account for workload run-as account. Users submitting jobs must have act-as permission on this run-as account. If unspecified, the [Vertex AI Custom Code Service Agent](https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents) for the CustomJob's project is used.""" - value: Optional[str] - """Required. Variables that reference a $(VAR_NAME) are expanded using the previous defined environment variables in the container and any service environment variables. If a variable cannot be resolved, the reference in the input string will be unchanged. The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped references will never be expanded, regardless of whether the variable exists or not.""" + tensorboard: Optional[str] + """Optional. The name of a Vertex AI Tensorboard resource to which this CustomJob will upload Tensorboard logs. Format: `projects/{project}/locations/{location}/tensorboards/{tensorboard}`""" + worker_pool_specs: Optional[list[WorkerPoolSpecDict]] + """Required. The spec of the worker pools including machine type and Docker image. All worker pools except the first one are optional and can be skipped by providing an empty value.""" -EnvVarOrDict = Union[EnvVar, EnvVarDict] +CustomJobSpecOrDict = Union[CustomJobSpec, CustomJobSpecDict] -class ContainerSpec(_common.BaseModel): - """The spec of a Container.""" - args: Optional[list[str]] = Field( +class CustomJob(_common.BaseModel): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" + + display_name: Optional[str] = Field( default=None, - description="""The arguments to be passed when starting the container.""", + description="""Required. The display name of the CustomJob. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", ) - command: Optional[list[str]] = Field( + job_spec: Optional[CustomJobSpec] = Field( + default=None, description="""Required. Job spec.""" + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( default=None, - description="""The command to be invoked when the container is started. It overrides the entrypoint instruction in Dockerfile when provided.""", + description="""Customer-managed encryption key options for a CustomJob. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key.""", ) - env: Optional[list[EnvVar]] = Field( + state: Optional[genai_types.JobState] = Field( + default=None, description="""Output only. The detailed state of the job.""" + ) + error: Optional[genai_types.GoogleRpcStatus] = Field( default=None, - description="""Environment variables to be passed to the container. Maximum limit is 100.""", + description="""Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`.""", ) - image_uri: Optional[str] = Field( + create_time: Optional[datetime.datetime] = Field( default=None, - description="""Required. The URI of a container image in the Container Registry that is to be run on each worker replica.""", + description="""Output only. Time when the CustomJob was created.""", + ) + end_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Time when the CustomJob entered any of the following states: `JOB_STATE_SUCCEEDED`, `JOB_STATE_FAILED`, `JOB_STATE_CANCELLED`.""", + ) + labels: Optional[dict[str, str]] = Field( + default=None, + description="""The labels with user-defined metadata to organize CustomJobs. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", + ) + name: Optional[str] = Field( + default=None, description="""Output only. Resource name of a CustomJob.""" + ) + satisfies_pzi: Optional[bool] = Field( + default=None, description="""Output only. Reserved for future use.""" + ) + satisfies_pzs: Optional[bool] = Field( + default=None, description="""Output only. Reserved for future use.""" + ) + start_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state.""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Time when the CustomJob was most recently updated.""", + ) + web_access_uris: Optional[dict[str, str]] = Field( + default=None, + description="""Output only. URIs for accessing [interactive shells](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) (one URI for each training node). Only available if job_spec.enable_web_access is `true`. The keys are names of each node in the training job; for example, `workerpool0-0` for the primary node, `workerpool1-0` for the first node in the second worker pool, and `workerpool1-1` for the second node in the second worker pool. The values are the URIs for each node's interactive shell.""", + ) + placeholder_to_make_trial_info_oneof_work_in_generated_proto: Optional[int] = Field( + default=None, + description="""Place holder to make oneof work in generated v1alpha1/v1beta1,v1 proto.""", + ) + description: Optional[str] = Field( + default=None, description="""The description of the CustomJob.""" + ) + pubsub_topic: Optional[str] = Field( + default=None, + description="""Immutable. The `Topic.name` of the Pub/Sub topic to which to publish the update when this job finishes. Must be of the format: `projects/{project}/topics/{topic}`. If not provided, such an update won't be sent, but the job state can still be polled via JobService.GetCustomJob. If a non-existing topic, or a topic to which Vertex AI is not allowed to publish to is provided, the job creation will succeed, but the update still won't be sent. The update message contains as its data a JSON in the following format: ``` { "title": "Job state update", "type": "object", "properties": { "job": { "type": "string", "description": "The resource name of the job, in \"projects/{project}/locations/{location}/custom_jobs/{custom_job}\" format, this update pertains to." }, "status": { "type": "string", "enum": ["succeeded", "failed", "cancelled", "expired"], "description": "The status in which the job finished its execution." } } } ```""", ) -class ContainerSpecDict(TypedDict, total=False): - """The spec of a Container.""" - - args: Optional[list[str]] - """The arguments to be passed when starting the container.""" +class CustomJobDict(TypedDict, total=False): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - command: Optional[list[str]] - """The command to be invoked when the container is started. It overrides the entrypoint instruction in Dockerfile when provided.""" + 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.""" - env: Optional[list[EnvVarDict]] - """Environment variables to be passed to the container. Maximum limit is 100.""" + job_spec: Optional[CustomJobSpecDict] + """Required. Job spec.""" - image_uri: Optional[str] - """Required. The URI of a container image in the Container Registry that is to be run on each worker replica.""" + 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.""" -ContainerSpecOrDict = Union[ContainerSpec, ContainerSpecDict] + 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 DiskSpec(_common.BaseModel): - """Represents the spec of disk options.""" + 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`.""" - boot_disk_size_gb: Optional[int] = Field( - default=None, description="""Size in GB of the boot disk (default is 100GB).""" - ) - boot_disk_type: Optional[str] = Field( - default=None, - description="""Type of the boot disk. For non-A3U machines, the default value is "pd-ssd", for A3U machines, the default value is "hyperdisk-balanced". Valid values: "pd-ssd" (Persistent Disk Solid State Drive), "pd-standard" (Persistent Disk Hard Disk Drive) or "hyperdisk-balanced".""", - ) + 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 DiskSpecDict(TypedDict, total=False): - """Represents the spec of disk options.""" + satisfies_pzi: Optional[bool] + """Output only. Reserved for future use.""" - boot_disk_size_gb: Optional[int] - """Size in GB of the boot disk (default is 100GB).""" + satisfies_pzs: Optional[bool] + """Output only. Reserved for future use.""" - 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".""" + start_time: Optional[datetime.datetime] + """Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state.""" + update_time: Optional[datetime.datetime] + """Output only. Time when the CustomJob was most recently updated.""" -DiskSpecOrDict = Union[DiskSpec, DiskSpecDict] + 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.""" + placeholder_to_make_trial_info_oneof_work_in_generated_proto: Optional[int] + """Place holder to make oneof work in generated v1alpha1/v1beta1,v1 proto.""" -class LustreMount(_common.BaseModel): - """Represents a mount configuration for Lustre file system.""" + description: Optional[str] + """The description of the CustomJob.""" - filesystem: Optional[str] = Field( - default=None, description="""Required. The name of the Lustre filesystem.""" - ) - instance_ip: Optional[str] = Field( - default=None, description="""Required. IP address of the Lustre instance.""" - ) - mount_point: Optional[str] = Field( - default=None, - description="""Required. Destination mount path. The Lustre file system will be mounted for the user under /mnt/lustre/""", - ) - volume_handle: Optional[str] = Field( - default=None, - description="""Required. The unique identifier of the Lustre volume.""", - ) + pubsub_topic: Optional[str] + """Immutable. The `Topic.name` of the Pub/Sub topic to which to publish the update when this job finishes. Must be of the format: `projects/{project}/topics/{topic}`. If not provided, such an update won't be sent, but the job state can still be polled via JobService.GetCustomJob. If a non-existing topic, or a topic to which Vertex AI is not allowed to publish to is provided, the job creation will succeed, but the update still won't be sent. The update message contains as its data a JSON in the following format: ``` { "title": "Job state update", "type": "object", "properties": { "job": { "type": "string", "description": "The resource name of the job, in \"projects/{project}/locations/{location}/custom_jobs/{custom_job}\" format, this update pertains to." }, "status": { "type": "string", "enum": ["succeeded", "failed", "cancelled", "expired"], "description": "The status in which the job finished its execution." } } } ```""" -class LustreMountDict(TypedDict, total=False): - """Represents a mount configuration for Lustre file system.""" +CustomJobOrDict = Union[CustomJob, CustomJobDict] - filesystem: Optional[str] - """Required. The name of the Lustre filesystem.""" - instance_ip: Optional[str] - """Required. IP address of the Lustre instance.""" +class VertexBaseConfig(_common.BaseModel): + """Base config for Vertex AI.""" - mount_point: Optional[str] - """Required. Destination mount path. The Lustre file system will be mounted for the user under /mnt/lustre/""" - - volume_handle: Optional[str] - """Required. The unique identifier of the Lustre volume.""" - - -LustreMountOrDict = Union[LustreMount, LustreMountDict] - - -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.""", - ) - reservation_affinity_type: Optional[Type] = Field( - default=None, - description="""Required. Specifies the reservation affinity type.""", - ) - values: Optional[list[str]] = Field( - default=None, - description="""Optional. Corresponds to the label values of a reservation resource. This must be the full resource name of the reservation or reservation block.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -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 VertexBaseConfigDict(TypedDict, total=False): + """Base config for Vertex AI.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - reservation_affinity_type: Optional[Type] - """Required. Specifies the reservation affinity type.""" - values: Optional[list[str]] - """Optional. Corresponds to the label values of a reservation resource. This must be the full resource name of the reservation or reservation block.""" +VertexBaseConfigOrDict = Union[VertexBaseConfig, VertexBaseConfigDict] -ReservationAffinityOrDict = Union[ReservationAffinity, ReservationAffinityDict] +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 MachineSpec(_common.BaseModel): - """Specification of a single machine.""" - accelerator_count: Optional[int] = Field( - default=None, - description="""The number of accelerators to attach to the machine. For [accelerator optimized machine types](https://cloud.google.com/compute/docs/accelerator-optimized-machines), One may set the accelerator_count from 1 to N for machine with N GPUs. If accelerator_count is less than or equal to N / 2, Agent Platform co-schedules the replicas of the model into the same VM to save cost. For example, if the machine type is a3-highgpu-8g, which has 8 H100 GPUs, one can set accelerator_count to 1 to 8. If accelerator_count is 1, 2, 3, or 4, Agent Platform co-schedules 8, 4, 2, or 2 replicas of the model into the same VM to save cost. When co-scheduling, CPU, memory and storage on the VM will be distributed to replicas on the VM. For example, one can expect a co-scheduled replica requesting 2 GPUs out of a 8-GPU VM will receive 25% of the CPU, memory and storage of the VM. Note that the feature is not compatible with multihost_gpu_node_count. When multihost_gpu_node_count is set, the co-scheduling will not be enabled.""", - ) - accelerator_type: Optional[AcceleratorType] = Field( - default=None, - description="""Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""", - ) - gpu_partition_size: Optional[str] = Field( - default=None, - description="""Optional. Immutable. The Nvidia GPU partition size. When specified, the requested accelerators will be partitioned into smaller GPU partitions. For example, if the request is for 8 units of NVIDIA A100 GPUs, and gpu_partition_size="1g.10gb", the service will create 8 * 7 = 56 partitioned MIG instances. The partition size must be a value supported by the requested accelerator. Refer to [Nvidia GPU Partitioning](https://cloud.google.com/kubernetes-engine/docs/how-to/gpus-multi#multi-instance_gpu_partitions) for the available partition sizes. If set, the accelerator_count should be set to 1.""", - ) - machine_type: Optional[str] = Field( - default=None, - description="""Immutable. The type of the machine. See the [list of machine types supported for prediction](https://cloud.google.com/gemini-enterprise-agent-platform/machine-learning/predictions/configure-compute#machine-types) See the [list of machine types supported for custom training](https://cloud.google.com/gemini-enterprise-agent-platform/machine-learning/training/configure-compute#machine-types). For DeployedModel this field is optional, and the default value is `n1-standard-2`. For BatchPredictionJob or as part of WorkerPoolSpec this field is required.""", - ) - min_gpu_driver_version: Optional[str] = Field( - default=None, - description="""Optional. Immutable. The minimum GPU driver version that this machine requires. For example, "535.104.06". If not specified, the default GPU driver version will be used by the underlying infrastructure.""", - ) - multihost_gpu_node_count: Optional[int] = Field( - default=None, - description="""Optional. Immutable. The number of nodes per replica for multihost GPU deployments.""", - ) - reservation_affinity: Optional[ReservationAffinity] = Field( - default=None, - description="""Optional. Immutable. Configuration controlling how this resource pool consumes reservation.""", - ) - 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 _CustomJobParametersDict(TypedDict, total=False): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" + custom_job: Optional[CustomJobDict] + """""" -class MachineSpecDict(TypedDict, total=False): - """Specification of a single machine.""" + config: Optional[VertexBaseConfigDict] + """""" - 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.""" - accelerator_type: Optional[AcceleratorType] - """Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""" +_CustomJobParametersOrDict = Union[_CustomJobParameters, _CustomJobParametersDict] - 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.""" +class _GetCustomJobParameters(_common.BaseModel): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - 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.""" + name: Optional[str] = Field(default=None, description="""""") + config: Optional[VertexBaseConfig] = Field(default=None, description="""""") - 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.""" +class _GetCustomJobParametersDict(TypedDict, total=False): + """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - tpu_topology: Optional[str] - """Immutable. The topology of the TPUs. Corresponds to the TPU topologies available from GKE. (Example: tpu_topology: "2x2x1").""" + name: Optional[str] + """""" + + config: Optional[VertexBaseConfigDict] + """""" -MachineSpecOrDict = Union[MachineSpec, MachineSpecDict] +_GetCustomJobParametersOrDict = Union[ + _GetCustomJobParameters, _GetCustomJobParametersDict +] -class NfsMount(_common.BaseModel): - """Represents a mount configuration for Network File System (NFS) to mount.""" +class CancelQueryJobRuntimeConfig(_common.BaseModel): + """Config for canceling async querying agent runtimes.""" - mount_point: Optional[str] = Field( - default=None, - description="""Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - path: Optional[str] = Field( + operation_name: Optional[str] = Field( default=None, - description="""Required. Source path exported from NFS server. Has to start with '/', and combined with the ip address, it indicates the source mount path in the form of `server:path`""", - ) - server: Optional[str] = Field( - default=None, description="""Required. IP address of the NFS server.""" + description="""Name of the longrunning operation returned from run_query_job.""", ) -class NfsMountDict(TypedDict, total=False): - """Represents a mount configuration for Network File System (NFS) to mount.""" - - mount_point: Optional[str] - """Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""" +class CancelQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for canceling async querying agent runtimes.""" - 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`""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - server: Optional[str] - """Required. IP address of the NFS server.""" + operation_name: Optional[str] + """Name of the longrunning operation returned from run_query_job.""" -NfsMountOrDict = Union[NfsMount, NfsMountDict] +CancelQueryJobRuntimeConfigOrDict = Union[ + CancelQueryJobRuntimeConfig, CancelQueryJobRuntimeConfigDict +] -class PythonPackageSpec(_common.BaseModel): - """The spec of a Python packaged code.""" +class _CancelQueryJobRuntimeRequestParameters(_common.BaseModel): + """Parameters for canceling async querying agent runtimes.""" - args: Optional[list[str]] = Field( - default=None, - description="""Command line arguments to be passed to the Python task.""", - ) - env: Optional[list[EnvVar]] = Field( - default=None, - description="""Environment variables to be passed to the python module. Maximum limit is 100.""", - ) - executor_image_uri: Optional[str] = Field( - default=None, - description="""Required. The URI of a container image in Artifact Registry that will run the provided Python package. Vertex AI provides a wide range of executor images with pre-installed packages to meet users' various use cases. See the list of [pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers). You must use an image from this list.""", - ) - package_uris: Optional[list[str]] = Field( - default=None, - description="""Required. The Google Cloud Storage location of the Python package files which are the training program and its dependent packages. The maximum number of package URIs is 100.""", + name: Optional[str] = Field( + default=None, description="""Name of the reasoning engine resource.""" ) - python_module: Optional[str] = Field( - default=None, - description="""Required. The Python module name to run after installing the packages.""", + config: Optional[CancelQueryJobRuntimeConfig] = Field( + default=None, description="""""" ) -class PythonPackageSpecDict(TypedDict, total=False): - """The spec of a Python packaged code.""" +class _CancelQueryJobRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for canceling async querying agent runtimes.""" - args: Optional[list[str]] - """Command line arguments to be passed to the Python task.""" + name: Optional[str] + """Name of the reasoning engine resource.""" - env: Optional[list[EnvVarDict]] - """Environment variables to be passed to the python module. Maximum limit is 100.""" + config: Optional[CancelQueryJobRuntimeConfigDict] + """""" - executor_image_uri: Optional[str] - """Required. The URI of a container image in Artifact Registry that will run the provided Python package. Vertex AI provides a wide range of executor images with pre-installed packages to meet users' various use cases. See the list of [pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers). You must use an image from this list.""" - package_uris: Optional[list[str]] - """Required. The Google Cloud Storage location of the Python package files which are the training program and its dependent packages. The maximum number of package URIs is 100.""" +_CancelQueryJobRuntimeRequestParametersOrDict = Union[ + _CancelQueryJobRuntimeRequestParameters, _CancelQueryJobRuntimeRequestParametersDict +] - python_module: Optional[str] - """Required. The Python module name to run after installing the packages.""" +class CancelQueryJobResult(_common.BaseModel): + """Result of canceling a query job.""" -PythonPackageSpecOrDict = Union[PythonPackageSpec, PythonPackageSpecDict] + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) -class WorkerPoolSpec(_common.BaseModel): - """Represents the spec of a worker pool in a job.""" +class CancelQueryJobResultDict(TypedDict, total=False): + """Result of canceling a query 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( - 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.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +CancelQueryJobResultOrDict = Union[CancelQueryJobResult, CancelQueryJobResultDict] + + +class CheckQueryJobRuntimeConfig(_common.BaseModel): + """Config for async querying agent runtimes.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - replica_count: Optional[int] = Field( + retrieve_result: Optional[bool] = Field( default=None, - description="""Optional. The number of worker replicas to use for this worker pool.""", + description="""Whether to retrieve the results of the query job.""", ) -class WorkerPoolSpecDict(TypedDict, total=False): - """Represents the spec of a worker pool in a job.""" +class CheckQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for async querying agent runtimes.""" - container_spec: Optional[ContainerSpecDict] - """The custom container task.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - disk_spec: Optional[DiskSpecDict] - """Disk spec.""" + retrieve_result: Optional[bool] + """Whether to retrieve the results of the query job.""" - lustre_mounts: Optional[list[LustreMountDict]] - """Optional. List of Lustre mounts.""" - machine_spec: Optional[MachineSpecDict] - """Optional. Immutable. The specification of a single machine.""" +CheckQueryJobRuntimeConfigOrDict = Union[ + CheckQueryJobRuntimeConfig, CheckQueryJobRuntimeConfigDict +] - nfs_mounts: Optional[list[NfsMountDict]] - """Optional. List of NFS mount spec.""" - python_package_spec: Optional[PythonPackageSpecDict] - """The Python packaged task.""" +class _CheckQueryJobRuntimeRequestParameters(_common.BaseModel): + """Parameters for async querying agent runtimes.""" - replica_count: Optional[int] - """Optional. The number of worker replicas to use for this worker pool.""" + name: Optional[str] = Field(default=None, description="""Name of the query job.""") + config: Optional[CheckQueryJobRuntimeConfig] = Field( + default=None, description="""""" + ) -WorkerPoolSpecOrDict = Union[WorkerPoolSpec, WorkerPoolSpecDict] +class _CheckQueryJobRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for async querying agent runtimes.""" + name: Optional[str] + """Name of the query job.""" -class CustomJobSpec(_common.BaseModel): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" + config: Optional[CheckQueryJobRuntimeConfigDict] + """""" - 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'].""", + +_CheckQueryJobRuntimeRequestParametersOrDict = Union[ + _CheckQueryJobRuntimeRequestParameters, _CheckQueryJobRuntimeRequestParametersDict +] + + +class CheckQueryJobResult(_common.BaseModel): + """Result of checking a query job.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - scheduling: Optional[Scheduling] = Field( - default=None, description="""Scheduling options for a CustomJob.""" + operation_name: Optional[str] = Field( + default=None, description="""Name of the agent runtime operation.""" ) - 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.""", + output_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the output file.""" ) - 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}`""", + status: Optional[str] = Field( + default=None, description="""Status of the operation.""" ) - 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.""", + result: Optional[str] = Field( + default=None, description="""JSON result of the operation.""" ) -class CustomJobSpecDict(TypedDict, total=False): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +class CheckQueryJobResultDict(TypedDict, total=False): + """Result of checking a query job.""" - base_output_directory: Optional[genai_types.GcsDestination] - """The Cloud Storage location to store the output of this CustomJob or HyperparameterTuningJob. For HyperparameterTuningJob, the baseOutputDirectory of each child CustomJob backing a Trial is set to a subdirectory of name id under its parent HyperparameterTuningJob's baseOutputDirectory. The following Vertex AI environment variables will be passed to containers or python modules when this field is set: For CustomJob: * AIP_MODEL_DIR = `/model/` * AIP_CHECKPOINT_DIR = `/checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `/logs/` For CustomJob backing a Trial of HyperparameterTuningJob: * AIP_MODEL_DIR = `//model/` * AIP_CHECKPOINT_DIR = `//checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `//logs/`""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - enable_dashboard_access: Optional[bool] - """Optional. Whether you want Vertex AI to enable access to the customized dashboard in training chief container. If set to `true`, you can access the dashboard at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""" + operation_name: Optional[str] + """Name of the agent runtime operation.""" - enable_web_access: Optional[bool] - """Optional. Whether you want Vertex AI to enable [interactive shell access](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) to training containers. If set to `true`, you can access interactive shells at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""" + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" - experiment: Optional[str] - """Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`""" + status: Optional[str] + """Status of the operation.""" - 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}`""" + result: Optional[str] + """JSON result of the operation.""" - models: Optional[list[str]] - """Optional. The name of the Model resources for which to generate a mapping to artifact URIs. Applicable only to some of the Google-provided custom jobs. Format: `projects/{project}/locations/{location}/models/{model}` In order to retrieve a specific version of the model, also provide the version ID or version alias. Example: `projects/{project}/locations/{location}/models/{model}@2` or `projects/{project}/locations/{location}/models/{model}@golden` If no version ID or alias is specified, the "default" version will be returned. The "default" version alias is created for the first version of the model, and can be moved to other versions later on. There will be exactly one default version.""" - network: Optional[str] - """Optional. The full name of the Compute Engine [network](/compute/docs/networks-and-firewalls#networks) to which the Job should be peered. For example, `projects/12345/global/networks/myVPC`. [Format](/compute/docs/reference/rest/v1/networks/insert) is of the form `projects/{project}/global/networks/{network}`. Where {project} is a project number, as in `12345`, and {network} is a network name. To specify this field, you must have already [configured VPC Network Peering for Vertex AI](https://cloud.google.com/vertex-ai/docs/general/vpc-peering). If this field is left unspecified, the job is not peered with any network.""" +CheckQueryJobResultOrDict = Union[CheckQueryJobResult, CheckQueryJobResultDict] - 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 _RunQueryJobRuntimeConfig(_common.BaseModel): + """Config for running a query job on an agent runtime.""" - psc_interface_config: Optional[PscInterfaceConfigDict] - """Optional. Configuration for PSC-I for CustomJob.""" + 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.""" + ) + output_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the output file.""" + ) - 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.""" +class _RunQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for running a query job on an agent runtime.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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}`""" + input_gcs_uri: Optional[str] + """The GCS URI of the input file.""" - 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.""" + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" -CustomJobSpecOrDict = Union[CustomJobSpec, CustomJobSpecDict] +_RunQueryJobRuntimeConfigOrDict = Union[ + _RunQueryJobRuntimeConfig, _RunQueryJobRuntimeConfigDict +] -class CustomJob(_common.BaseModel): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +class _RunQueryJobRuntimeRequestParameters(_common.BaseModel): + """Parameters for running a query job on an agent runtime.""" - display_name: Optional[str] = Field( - default=None, - description="""Required. The display name of the CustomJob. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", - ) - job_spec: Optional[CustomJobSpec] = Field( - default=None, description="""Required. Job spec.""" - ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, - description="""Customer-managed encryption key options for a CustomJob. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key.""", - ) - state: Optional[genai_types.JobState] = Field( - default=None, description="""Output only. The detailed state of the job.""" - ) - error: Optional[genai_types.GoogleRpcStatus] = Field( - default=None, - description="""Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`.""", - ) - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Time when the CustomJob was created.""", - ) - end_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Time when the CustomJob entered any of the following states: `JOB_STATE_SUCCEEDED`, `JOB_STATE_FAILED`, `JOB_STATE_CANCELLED`.""", - ) - labels: Optional[dict[str, str]] = Field( - default=None, - description="""The labels with user-defined metadata to organize CustomJobs. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", - ) name: Optional[str] = Field( - default=None, description="""Output only. Resource name of a CustomJob.""" - ) - satisfies_pzi: Optional[bool] = Field( - default=None, description="""Output only. Reserved for future use.""" - ) - satisfies_pzs: Optional[bool] = Field( - default=None, description="""Output only. Reserved for future use.""" - ) - start_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state.""", - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Time when the CustomJob was most recently updated.""", + default=None, description="""Name of the agent runtime.""" ) - web_access_uris: Optional[dict[str, str]] = Field( - default=None, - description="""Output only. URIs for accessing [interactive shells](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) (one URI for each training node). Only available if job_spec.enable_web_access is `true`. The keys are names of each node in the training job; for example, `workerpool0-0` for the primary node, `workerpool1-0` for the first node in the second worker pool, and `workerpool1-1` for the second node in the second worker pool. The values are the URIs for each node's interactive shell.""", + config: Optional[_RunQueryJobRuntimeConfig] = Field( + default=None, description="""""" ) -class CustomJobDict(TypedDict, total=False): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - - display_name: Optional[str] - """Required. The display name of the CustomJob. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" - - job_spec: Optional[CustomJobSpecDict] - """Required. Job spec.""" - - encryption_spec: Optional[genai_types.EncryptionSpec] - """Customer-managed encryption key options for a CustomJob. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key.""" - - state: Optional[genai_types.JobState] - """Output only. The detailed state of the job.""" +class _RunQueryJobRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for running a query job on an agent runtime.""" - error: Optional[genai_types.GoogleRpcStatus] - """Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`.""" + name: Optional[str] + """Name of the agent runtime.""" - create_time: Optional[datetime.datetime] - """Output only. Time when the CustomJob was created.""" + config: Optional[_RunQueryJobRuntimeConfigDict] + """""" - end_time: Optional[datetime.datetime] - """Output only. Time when the CustomJob entered any of the following states: `JOB_STATE_SUCCEEDED`, `JOB_STATE_FAILED`, `JOB_STATE_CANCELLED`.""" - labels: Optional[dict[str, str]] - """The labels with user-defined metadata to organize CustomJobs. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""" +_RunQueryJobRuntimeRequestParametersOrDict = Union[ + _RunQueryJobRuntimeRequestParameters, _RunQueryJobRuntimeRequestParametersDict +] - name: Optional[str] - """Output only. Resource name of a CustomJob.""" - satisfies_pzi: Optional[bool] - """Output only. Reserved for future use.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent( + _common.BaseModel +): + """The conversation source event for generating memories.""" - satisfies_pzs: Optional[bool] - """Output only. Reserved for future use.""" + content: Optional[genai_types.Content] = Field( + default=None, description="""Required. Represents the content of the event.""" + ) - start_time: Optional[datetime.datetime] - """Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state.""" - update_time: Optional[datetime.datetime] - """Output only. Time when the CustomJob was most recently updated.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict( + TypedDict, total=False +): + """The conversation source event for generating memories.""" - web_access_uris: Optional[dict[str, str]] - """Output only. URIs for accessing [interactive shells](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) (one URI for each training node). Only available if job_spec.enable_web_access is `true`. The keys are names of each node in the training job; for example, `workerpool0-0` for the primary node, `workerpool1-0` for the first node in the second worker pool, and `workerpool1-1` for the second node in the second worker pool. The values are the URIs for each node's interactive shell.""" + content: Optional[genai_types.Content] + """Required. Represents the content of the event.""" -CustomJobOrDict = Union[CustomJob, CustomJobDict] +MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventOrDict = ( + Union[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent, + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict, + ] +) -class VertexBaseConfig(_common.BaseModel): - """Base config for Vertex AI.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource( + _common.BaseModel +): + """A conversation source for the example. This is similar to `DirectContentsSource`.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + events: Optional[ + list[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent + ] + ] = Field( + default=None, + description="""Optional. Represents the input conversation events for the example.""", ) -class VertexBaseConfigDict(TypedDict, total=False): - """Base config for Vertex AI.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict( + TypedDict, total=False +): + """A conversation source for the example. This is similar to `DirectContentsSource`.""" + events: Optional[ + list[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict + ] + ] + """Optional. Represents the input conversation events for the example.""" -VertexBaseConfigOrDict = Union[VertexBaseConfig, VertexBaseConfigDict] +MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceOrDict = Union[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource, + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict, +] -class _CustomJobParameters(_common.BaseModel): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - custom_job: Optional[CustomJob] = Field(default=None, description="""""") - config: Optional[VertexBaseConfig] = Field(default=None, description="""""") +class MemoryTopicId(_common.BaseModel): + """The topic ID for a memory.""" + custom_memory_topic_label: Optional[str] = Field( + default=None, + description="""Optional. Represents the custom memory topic label.""", + ) + managed_memory_topic: Optional[ManagedTopicEnum] = Field( + default=None, description="""Optional. Represents the managed memory topic.""" + ) -class _CustomJobParametersDict(TypedDict, total=False): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" - custom_job: Optional[CustomJobDict] - """""" +class MemoryTopicIdDict(TypedDict, total=False): + """The topic ID for a memory.""" - config: Optional[VertexBaseConfigDict] - """""" + custom_memory_topic_label: Optional[str] + """Optional. Represents the custom memory topic label.""" + managed_memory_topic: Optional[ManagedTopicEnum] + """Optional. Represents the managed memory topic.""" -_CustomJobParametersOrDict = Union[_CustomJobParameters, _CustomJobParametersDict] +MemoryTopicIdOrDict = Union[MemoryTopicId, MemoryTopicIdDict] -class _GetCustomJobParameters(_common.BaseModel): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory( + _common.BaseModel +): + """A memory generated by the operation.""" - name: Optional[str] = Field(default=None, description="""""") - config: Optional[VertexBaseConfig] = Field(default=None, description="""""") + fact: Optional[str] = Field( + default=None, + description="""Required. Represents the fact to generate a memory from.""", + ) + topics: Optional[list[MemoryTopicId]] = Field( + default=None, + description="""Optional. Represents the list of topics that the memory should be associated with. For example, use `custom_memory_topic_label = "jargon"` if the extracted memory is an example of memory extraction for the custom topic `jargon`.""", + ) -class _GetCustomJobParametersDict(TypedDict, total=False): - """Represents a job that runs custom workloads such as a Docker container or a Python package.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict( + TypedDict, total=False +): + """A memory generated by the operation.""" - name: Optional[str] - """""" + fact: Optional[str] + """Required. Represents the fact to generate a memory from.""" - config: Optional[VertexBaseConfigDict] - """""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. Represents the list of topics that the memory should be associated with. For example, use `custom_memory_topic_label = "jargon"` if the extracted memory is an example of memory extraction for the custom topic `jargon`.""" -_GetCustomJobParametersOrDict = Union[ - _GetCustomJobParameters, _GetCustomJobParametersDict +MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryOrDict = Union[ + MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory, + MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict, ] -class CancelQueryJobAgentEngineConfig(_common.BaseModel): - """Config for canceling async querying agent engines.""" +class MemoryBankCustomizationConfigGenerateMemoriesExample(_common.BaseModel): + """An example of how to generate memories for a particular scope.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - operation_name: Optional[str] = Field( + conversation_source: Optional[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource + ] = Field(default=None, description="""A conversation source for the example.""") + generated_memories: Optional[ + list[MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory] + ] = Field( default=None, - description="""Name of the longrunning operation returned from run_query_job.""", + description="""Optional. Represents the memories that are expected to be generated from the input conversation. An empty list indicates that no memories are expected to be generated for the input conversation.""", ) -class CancelQueryJobAgentEngineConfigDict(TypedDict, total=False): - """Config for canceling async querying agent engines.""" +class MemoryBankCustomizationConfigGenerateMemoriesExampleDict(TypedDict, total=False): + """An example of how to generate memories for a particular scope.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + conversation_source: Optional[ + MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict + ] + """A conversation source for the example.""" - operation_name: Optional[str] - """Name of the longrunning operation returned from run_query_job.""" + generated_memories: Optional[ + list[MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict] + ] + """Optional. Represents the memories that are expected to be generated from the input conversation. An empty list indicates that no memories are expected to be generated for the input conversation.""" -CancelQueryJobAgentEngineConfigOrDict = Union[ - CancelQueryJobAgentEngineConfig, CancelQueryJobAgentEngineConfigDict +MemoryBankCustomizationConfigGenerateMemoriesExampleOrDict = Union[ + MemoryBankCustomizationConfigGenerateMemoriesExample, + MemoryBankCustomizationConfigGenerateMemoriesExampleDict, ] -class _CancelQueryJobAgentEngineRequestParameters(_common.BaseModel): - """Parameters for canceling async querying agent engines.""" +class MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic(_common.BaseModel): + """A custom memory topic defined by the developer.""" - name: Optional[str] = Field( - default=None, description="""Name of the reasoning engine resource.""" + label: Optional[str] = Field( + default=None, description="""Required. Represents the label of the topic.""" ) - config: Optional[CancelQueryJobAgentEngineConfig] = Field( - default=None, description="""""" + description: Optional[str] = Field( + default=None, + description="""Required. Represents the description of the memory topic. This should explain what information should be extracted for this topic.""", ) -class _CancelQueryJobAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for canceling async querying agent engines.""" +class MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict( + TypedDict, total=False +): + """A custom memory topic defined by the developer.""" - name: Optional[str] - """Name of the reasoning engine resource.""" + label: Optional[str] + """Required. Represents the label of the topic.""" - config: Optional[CancelQueryJobAgentEngineConfigDict] - """""" + description: Optional[str] + """Required. Represents the description of the memory topic. This should explain what information should be extracted for this topic.""" -_CancelQueryJobAgentEngineRequestParametersOrDict = Union[ - _CancelQueryJobAgentEngineRequestParameters, - _CancelQueryJobAgentEngineRequestParametersDict, +MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicOrDict = Union[ + MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic, + MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict, ] -class CancelQueryJobResult(_common.BaseModel): - """Result of canceling a query job.""" +class MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic(_common.BaseModel): + """A managed memory topic defined by the system.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + managed_topic_enum: Optional[ManagedTopicEnum] = Field( + default=None, description="""Required. Represents the managed topic.""" ) -class CancelQueryJobResultDict(TypedDict, total=False): - """Result of canceling a query job.""" +class MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict( + TypedDict, total=False +): + """A managed memory topic defined by the system.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + managed_topic_enum: Optional[ManagedTopicEnum] + """Required. Represents the managed topic.""" -CancelQueryJobResultOrDict = Union[CancelQueryJobResult, CancelQueryJobResultDict] +MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicOrDict = Union[ + MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic, + MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict, +] -class CheckQueryJobAgentEngineConfig(_common.BaseModel): - """Config for async querying agent engines.""" +class MemoryBankCustomizationConfigMemoryTopic(_common.BaseModel): + """A topic of information that should be extracted from conversations and stored as memories.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + custom_memory_topic: Optional[ + MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic + ] = Field( + default=None, description="""A custom memory topic defined by the developer.""" ) - retrieve_result: Optional[bool] = Field( - default=None, - description="""Whether to retrieve the results of the query job.""", + managed_memory_topic: Optional[ + MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic + ] = Field( + default=None, description="""A managed memory topic defined by Memory Bank.""" ) -class CheckQueryJobAgentEngineConfigDict(TypedDict, total=False): - """Config for async querying agent engines.""" +class MemoryBankCustomizationConfigMemoryTopicDict(TypedDict, total=False): + """A topic of information that should be extracted from conversations and stored as memories.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + custom_memory_topic: Optional[ + MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict + ] + """A custom memory topic defined by the developer.""" - retrieve_result: Optional[bool] - """Whether to retrieve the results of the query job.""" + managed_memory_topic: Optional[ + MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict + ] + """A managed memory topic defined by Memory Bank.""" -CheckQueryJobAgentEngineConfigOrDict = Union[ - CheckQueryJobAgentEngineConfig, CheckQueryJobAgentEngineConfigDict +MemoryBankCustomizationConfigMemoryTopicOrDict = Union[ + MemoryBankCustomizationConfigMemoryTopic, + MemoryBankCustomizationConfigMemoryTopicDict, ] -class _CheckQueryJobAgentEngineRequestParameters(_common.BaseModel): - """Parameters for async querying agent engines.""" +class MemoryBankCustomizationConfigConsolidationConfig(_common.BaseModel): + """Represents configuration for customizing how memories are consolidated.""" - name: Optional[str] = Field(default=None, description="""Name of the query job.""") - config: Optional[CheckQueryJobAgentEngineConfig] = Field( - default=None, description="""""" + revisions_per_candidate_count: Optional[int] = Field( + default=None, + description="""Optional. Represents the maximum number of revisions to consider for each candidate memory. If not set, then the default value (1) will be used, which means that only the latest revision will be considered.""", ) -class _CheckQueryJobAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for async querying agent engines.""" - - name: Optional[str] - """Name of the query job.""" +class MemoryBankCustomizationConfigConsolidationConfigDict(TypedDict, total=False): + """Represents configuration for customizing how memories are consolidated.""" - config: Optional[CheckQueryJobAgentEngineConfigDict] - """""" + revisions_per_candidate_count: Optional[int] + """Optional. Represents the maximum number of revisions to consider for each candidate memory. If not set, then the default value (1) will be used, which means that only the latest revision will be considered.""" -_CheckQueryJobAgentEngineRequestParametersOrDict = Union[ - _CheckQueryJobAgentEngineRequestParameters, - _CheckQueryJobAgentEngineRequestParametersDict, +MemoryBankCustomizationConfigConsolidationConfigOrDict = Union[ + MemoryBankCustomizationConfigConsolidationConfig, + MemoryBankCustomizationConfigConsolidationConfigDict, ] -class CheckQueryJobResult(_common.BaseModel): - """Result of checking a query job.""" +class MemoryBankCustomizationConfig(_common.BaseModel): + """Represents configuration for organizing natural language memories for a particular scope.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + enable_third_person_memories: Optional[bool] = Field( + default=None, + description="""Optional. Indicates whether the memories will be generated in the third person (i.e. "The user generates memories with Memory Bank."). By default, the memories will be generated in the first person (i.e. "I generate memories with Memory Bank.")""", ) - operation_name: Optional[str] = Field( - default=None, description="""Name of the agent engine operation.""" + generate_memories_examples: Optional[ + list[MemoryBankCustomizationConfigGenerateMemoriesExample] + ] = Field( + default=None, + description="""Optional. Provides examples of how to generate memories for a particular scope.""", ) - output_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the output file.""" + memory_topics: Optional[list[MemoryBankCustomizationConfigMemoryTopic]] = Field( + default=None, + description="""Optional. Represents topics of information that should be extracted from conversations and stored as memories. If not set, then Memory Bank's default topics will be used.""", ) - status: Optional[str] = Field( - default=None, description="""Status of the operation.""" + scope_keys: Optional[list[str]] = Field( + default=None, + description="""Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""", ) - result: Optional[str] = Field( - default=None, description="""JSON result of the operation.""" + consolidation_config: Optional[MemoryBankCustomizationConfigConsolidationConfig] = ( + Field( + default=None, + description="""Optional. Represents configuration for customizing how memories are consolidated together.""", + ) + ) + disable_natural_language_memories: Optional[bool] = Field( + default=None, + description="""Optional. Indicates whether natural language memory generation should be disabled for all requests. By default, natural language memory generation is enabled. Set this to `true` when you only want to generate structured memories.""", ) -class CheckQueryJobResultDict(TypedDict, total=False): - """Result of checking a query job.""" +class MemoryBankCustomizationConfigDict(TypedDict, total=False): + """Represents configuration for organizing natural language memories for a particular scope.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + enable_third_person_memories: Optional[bool] + """Optional. Indicates whether the memories will be generated in the third person (i.e. "The user generates memories with Memory Bank."). By default, the memories will be generated in the first person (i.e. "I generate memories with Memory Bank.")""" - operation_name: Optional[str] - """Name of the agent engine operation.""" + generate_memories_examples: Optional[ + list[MemoryBankCustomizationConfigGenerateMemoriesExampleDict] + ] + """Optional. Provides examples of how to generate memories for a particular scope.""" - output_gcs_uri: Optional[str] - """The GCS URI of the output file.""" + memory_topics: Optional[list[MemoryBankCustomizationConfigMemoryTopicDict]] + """Optional. Represents topics of information that should be extracted from conversations and stored as memories. If not set, then Memory Bank's default topics will be used.""" - status: Optional[str] - """Status of the operation.""" + scope_keys: Optional[list[str]] + """Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""" - result: Optional[str] - """JSON result of the operation.""" + consolidation_config: Optional[MemoryBankCustomizationConfigConsolidationConfigDict] + """Optional. Represents configuration for customizing how memories are consolidated together.""" + disable_natural_language_memories: Optional[bool] + """Optional. Indicates whether natural language memory generation should be disabled for all requests. By default, natural language memory generation is enabled. Set this to `true` when you only want to generate structured memories.""" -CheckQueryJobResultOrDict = Union[CheckQueryJobResult, CheckQueryJobResultDict] + +MemoryBankCustomizationConfigOrDict = Union[ + MemoryBankCustomizationConfig, MemoryBankCustomizationConfigDict +] -class _RunQueryJobAgentEngineConfig(_common.BaseModel): - """Config for running a query job on an agent engine.""" +class MemoryGenerationTriggerConfigGenerationTriggerRule(_common.BaseModel): + """Represents the active rule that determines when to flush the buffer.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + event_count: Optional[int] = Field( + default=None, + description="""Optional. Specifies to trigger generation when the event count reaches this limit.""", ) - input_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the input file.""" + fixed_interval: Optional[str] = Field( + default=None, + description="""Optional. Specifies to trigger generation at a fixed interval. The duration must have a minute-level granularity.""", ) - output_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the output file.""" + idle_duration: Optional[str] = Field( + default=None, + description="""Optional. Specifies to trigger generation if the stream is inactive for the specified duration after the most recent event. The duration must have a minute-level granularity.""", + ) + overlap_event_count: Optional[int] = Field( + default=None, + description="""Optional. Re-include the last N already-processed events in the next window.""", + ) + token_limit: Optional[int] = Field( + default=None, + description="""Optional. Specifies to trigger generation when the token count reaches this limit.""", ) -class _RunQueryJobAgentEngineConfigDict(TypedDict, total=False): - """Config for running a query job on an agent engine.""" +class MemoryGenerationTriggerConfigGenerationTriggerRuleDict(TypedDict, total=False): + """Represents the active rule that determines when to flush the buffer.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + event_count: Optional[int] + """Optional. Specifies to trigger generation when the event count reaches this limit.""" - input_gcs_uri: Optional[str] - """The GCS URI of the input file.""" + fixed_interval: Optional[str] + """Optional. Specifies to trigger generation at a fixed interval. The duration must have a minute-level granularity.""" - output_gcs_uri: Optional[str] - """The GCS URI of the output file.""" + idle_duration: Optional[str] + """Optional. Specifies to trigger generation if the stream is inactive for the specified duration after the most recent event. The duration must have a minute-level granularity.""" + + overlap_event_count: Optional[int] + """Optional. Re-include the last N already-processed events in the next window.""" + + token_limit: Optional[int] + """Optional. Specifies to trigger generation when the token count reaches this limit.""" -_RunQueryJobAgentEngineConfigOrDict = Union[ - _RunQueryJobAgentEngineConfig, _RunQueryJobAgentEngineConfigDict +MemoryGenerationTriggerConfigGenerationTriggerRuleOrDict = Union[ + MemoryGenerationTriggerConfigGenerationTriggerRule, + MemoryGenerationTriggerConfigGenerationTriggerRuleDict, ] -class _RunQueryJobAgentEngineRequestParameters(_common.BaseModel): - """Parameters for running a query job on an agent engine.""" +class MemoryGenerationTriggerConfig(_common.BaseModel): + """The configuration for triggering memory generation for ingested events.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[_RunQueryJobAgentEngineConfig] = Field( - default=None, description="""""" + generation_rule: Optional[MemoryGenerationTriggerConfigGenerationTriggerRule] = ( + Field( + default=None, + description="""Optional. Represents the active rule that determines when to flush the buffer. If not set, then the stream will be force flushed immediately.""", + ) ) -class _RunQueryJobAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for running a query job on an agent engine.""" - - name: Optional[str] - """Name of the agent engine.""" +class MemoryGenerationTriggerConfigDict(TypedDict, total=False): + """The configuration for triggering memory generation for ingested events.""" - config: Optional[_RunQueryJobAgentEngineConfigDict] - """""" + generation_rule: Optional[MemoryGenerationTriggerConfigGenerationTriggerRuleDict] + """Optional. Represents the active rule that determines when to flush the buffer. If not set, then the stream will be force flushed immediately.""" -_RunQueryJobAgentEngineRequestParametersOrDict = Union[ - _RunQueryJobAgentEngineRequestParameters, - _RunQueryJobAgentEngineRequestParametersDict, +MemoryGenerationTriggerConfigOrDict = Union[ + MemoryGenerationTriggerConfig, MemoryGenerationTriggerConfigDict ] -class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent( - _common.BaseModel -): - """The conversation source event for generating memories.""" +class ReasoningEngineContextSpecMemoryBankConfigGenerationConfig(_common.BaseModel): + """Configuration for how to generate memories.""" - content: Optional[genai_types.Content] = Field( - default=None, description="""Required. Represents the content of the event.""" + model: Optional[str] = Field( + default=None, + description="""Optional. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", + ) + generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + default=None, + description="""Optional. Specifies the default trigger configuration for generating memories using `IngestEvents`.""", + ) + memory_extraction_instructions: Optional[str] = Field( + default=None, + description="""Optional. A custom prompt to use for extracting memories from conversations. If not set, a default prompt will be used.""", ) -class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict( +class ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict( TypedDict, total=False ): - """The conversation source event for generating memories.""" + """Configuration for how to generate memories.""" - content: Optional[genai_types.Content] - """Required. Represents the content of the event.""" + model: Optional[str] + """Optional. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] + """Optional. Specifies the default trigger configuration for generating memories using `IngestEvents`.""" -MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventOrDict = ( - Union[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent, - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict, - ] -) + memory_extraction_instructions: Optional[str] + """Optional. A custom prompt to use for extracting memories from conversations. If not set, a default prompt will be used.""" -class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource( +ReasoningEngineContextSpecMemoryBankConfigGenerationConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigGenerationConfig, + ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict, +] + + +class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig( _common.BaseModel ): - """A conversation source for the example. This is similar to `DirectContentsSource`.""" + """Configuration for how to perform similarity search on memories.""" - events: Optional[ - list[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEvent - ] - ] = Field( + embedding_model: Optional[str] = Field( default=None, - description="""Optional. Represents the input conversation events for the example.""", + description="""Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", ) -class MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict( +class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict( TypedDict, total=False ): - """A conversation source for the example. This is similar to `DirectContentsSource`.""" + """Configuration for how to perform similarity search on memories.""" - events: Optional[ - list[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceEventDict - ] - ] - """Optional. Represents the input conversation events for the example.""" + embedding_model: Optional[str] + """Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" -MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceOrDict = Union[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource, - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict, +ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig, + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict, ] -class MemoryTopicId(_common.BaseModel): - """The topic ID for a memory.""" - - custom_memory_topic_label: Optional[str] = Field( - default=None, - description="""Optional. Represents the custom memory topic label.""", - ) - managed_memory_topic: Optional[ManagedTopicEnum] = Field( - default=None, description="""Optional. Represents the managed memory topic.""" - ) - - -class MemoryTopicIdDict(TypedDict, total=False): - """The topic ID for a memory.""" - - custom_memory_topic_label: Optional[str] - """Optional. Represents the custom memory topic label.""" - - managed_memory_topic: Optional[ManagedTopicEnum] - """Optional. Represents the managed memory topic.""" - - -MemoryTopicIdOrDict = Union[MemoryTopicId, MemoryTopicIdDict] - - -class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory( +class ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig( _common.BaseModel ): - """A memory generated by the operation.""" + """Configuration for TTL of the memories in the Memory Bank based on the action that created or updated the memory.""" - fact: Optional[str] = Field( + create_ttl: Optional[str] = Field( default=None, - description="""Required. Represents the fact to generate a memory from.""", + description="""Optional. The TTL duration for memories uploaded via CreateMemory.""", ) - topics: Optional[list[MemoryTopicId]] = Field( + generate_created_ttl: Optional[str] = Field( default=None, - description="""Optional. Represents the list of topics that the memory should be associated with. For example, use `custom_memory_topic_label = "jargon"` if the extracted memory is an example of memory extraction for the custom topic `jargon`.""", + description="""Optional. The TTL duration for memories newly generated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.CREATED).""", + ) + generate_updated_ttl: Optional[str] = Field( + default=None, + description="""Optional. The TTL duration for memories updated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.UPDATED). In the case of an UPDATE action, the `expire_time` of the existing memory will be updated to the new value (now + TTL).""", ) -class MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict( +class ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict( TypedDict, total=False ): - """A memory generated by the operation.""" + """Configuration for TTL of the memories in the Memory Bank based on the action that created or updated the memory.""" - fact: Optional[str] - """Required. Represents the fact to generate a memory from.""" + create_ttl: Optional[str] + """Optional. The TTL duration for memories uploaded via CreateMemory.""" - topics: Optional[list[MemoryTopicIdDict]] - """Optional. Represents the list of topics that the memory should be associated with. For example, use `custom_memory_topic_label = "jargon"` if the extracted memory is an example of memory extraction for the custom topic `jargon`.""" + generate_created_ttl: Optional[str] + """Optional. The TTL duration for memories newly generated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.CREATED).""" + generate_updated_ttl: Optional[str] + """Optional. The TTL duration for memories updated via GenerateMemories (GenerateMemoriesResponse.GeneratedMemory.Action.UPDATED). In the case of an UPDATE action, the `expire_time` of the existing memory will be updated to the new value (now + TTL).""" -MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryOrDict = Union[ - MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory, - MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict, + +ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig, + ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict, ] -class MemoryBankCustomizationConfigGenerateMemoriesExample(_common.BaseModel): - """An example of how to generate memories for a particular scope.""" +class ReasoningEngineContextSpecMemoryBankConfigTtlConfig(_common.BaseModel): + """Configuration for automatically setting the TTL ("time-to-live") of the memories in the Memory Bank.""" - conversation_source: Optional[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSource - ] = Field(default=None, description="""A conversation source for the example.""") - generated_memories: Optional[ - list[MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemory] + default_ttl: Optional[str] = Field( + default=None, + description="""Optional. The default TTL duration of the memories in the Memory Bank. This applies to all operations that create or update a memory.""", + ) + granular_ttl_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig ] = Field( default=None, - description="""Optional. Represents the memories that are expected to be generated from the input conversation. An empty list indicates that no memories are expected to be generated for the input conversation.""", + description="""Optional. The granular TTL configuration of the memories in the Memory Bank.""", + ) + memory_revision_default_ttl: Optional[str] = Field( + default=None, + description="""Optional. The default TTL duration of the memory revisions in the Memory Bank. This applies to all operations that create a memory revision. If not set, a default TTL of 365 days will be used.""", ) -class MemoryBankCustomizationConfigGenerateMemoriesExampleDict(TypedDict, total=False): - """An example of how to generate memories for a particular scope.""" +class ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict(TypedDict, total=False): + """Configuration for automatically setting the TTL ("time-to-live") of the memories in the Memory Bank.""" - conversation_source: Optional[ - MemoryBankCustomizationConfigGenerateMemoriesExampleConversationSourceDict - ] - """A conversation source for the example.""" + default_ttl: Optional[str] + """Optional. The default TTL duration of the memories in the Memory Bank. This applies to all operations that create or update a memory.""" - generated_memories: Optional[ - list[MemoryBankCustomizationConfigGenerateMemoriesExampleGeneratedMemoryDict] + granular_ttl_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict ] - """Optional. Represents the memories that are expected to be generated from the input conversation. An empty list indicates that no memories are expected to be generated for the input conversation.""" + """Optional. The granular TTL configuration of the memories in the Memory Bank.""" + memory_revision_default_ttl: Optional[str] + """Optional. The default TTL duration of the memory revisions in the Memory Bank. This applies to all operations that create a memory revision. If not set, a default TTL of 365 days will be used.""" -MemoryBankCustomizationConfigGenerateMemoriesExampleOrDict = Union[ - MemoryBankCustomizationConfigGenerateMemoriesExample, - MemoryBankCustomizationConfigGenerateMemoriesExampleDict, -] +ReasoningEngineContextSpecMemoryBankConfigTtlConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfigTtlConfig, + ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict, +] -class MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic(_common.BaseModel): - """A custom memory topic defined by the developer.""" - label: Optional[str] = Field( - default=None, description="""Required. Represents the label of the topic.""" +class StructuredMemorySchemaConfig(_common.BaseModel): + """Represents the OpenAPI schema of the structured memories.""" + + memory_schema: Optional[genai_types.Schema] = Field( + default=None, + description="""Required. Represents the OpenAPI schema of the structured memories.""", ) - description: Optional[str] = Field( + id: Optional[str] = Field( default=None, - description="""Required. Represents the description of the memory topic. This should explain what information should be extracted for this topic.""", + description="""Required. Represents the ID of the schema. Must be 1-63 characters, start with a lowercase letter, and consist of lowercase letters, numbers, and hyphens.""", + ) + memory_type: Optional[MemoryType] = Field( + default=None, + description="""Optional. Represents the type of the structured memories associated with the schema. If not set, then `STRUCTURED_PROFILE` will be used.""", ) -class MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict( - TypedDict, total=False -): - """A custom memory topic defined by the developer.""" +class StructuredMemorySchemaConfigDict(TypedDict, total=False): + """Represents the OpenAPI schema of the structured memories.""" - label: Optional[str] - """Required. Represents the label of the topic.""" + memory_schema: Optional[genai_types.Schema] + """Required. Represents the OpenAPI schema of the structured memories.""" - description: Optional[str] - """Required. Represents the description of the memory topic. This should explain what information should be extracted for this topic.""" + id: Optional[str] + """Required. Represents the ID of the schema. Must be 1-63 characters, start with a lowercase letter, and consist of lowercase letters, numbers, and hyphens.""" + memory_type: Optional[MemoryType] + """Optional. Represents the type of the structured memories associated with the schema. If not set, then `STRUCTURED_PROFILE` will be used.""" -MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicOrDict = Union[ - MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic, - MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict, + +StructuredMemorySchemaConfigOrDict = Union[ + StructuredMemorySchemaConfig, StructuredMemorySchemaConfigDict ] -class MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic(_common.BaseModel): - """A managed memory topic defined by the system.""" +class StructuredMemoryConfig(_common.BaseModel): + """Configuration for organizing structured memories within a scope.""" - managed_topic_enum: Optional[ManagedTopicEnum] = Field( - default=None, description="""Required. Represents the managed topic.""" + schema_configs: Optional[list[StructuredMemorySchemaConfig]] = Field( + default=None, + description="""Optional. Represents configuration of the structured memories' schemas.""", + ) + scope_keys: Optional[list[str]] = Field( + default=None, + description="""Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""", ) -class MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict( - TypedDict, total=False -): - """A managed memory topic defined by the system.""" +class StructuredMemoryConfigDict(TypedDict, total=False): + """Configuration for organizing structured memories within a scope.""" - managed_topic_enum: Optional[ManagedTopicEnum] - """Required. Represents the managed topic.""" + schema_configs: Optional[list[StructuredMemorySchemaConfigDict]] + """Optional. Represents configuration of the structured memories' schemas.""" + scope_keys: Optional[list[str]] + """Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""" -MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicOrDict = Union[ - MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic, - MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict, -] +StructuredMemoryConfigOrDict = Union[StructuredMemoryConfig, StructuredMemoryConfigDict] -class MemoryBankCustomizationConfigMemoryTopic(_common.BaseModel): - """A topic of information that should be extracted from conversations and stored as memories.""" - custom_memory_topic: Optional[ - MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopic +class ReasoningEngineContextSpecMemoryBankConfig(_common.BaseModel): + """Specification for a Memory Bank.""" + + customization_configs: Optional[list[MemoryBankCustomizationConfig]] = Field( + default=None, + description="""Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""", + ) + disable_memory_revisions: Optional[bool] = Field( + default=None, + description="""If true, no memory revisions will be created for any requests to the Memory Bank.""", + ) + generation_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigGenerationConfig ] = Field( - default=None, description="""A custom memory topic defined by the developer.""" + default=None, + description="""Optional. Configuration for how to generate memories for the Memory Bank.""", ) - managed_memory_topic: Optional[ - MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopic + similarity_search_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig ] = Field( - default=None, description="""A managed memory topic defined by Memory Bank.""" + default=None, + description="""Optional. Configuration for how to perform similarity search on memories. If not set, the Memory Bank will use the default embedding model `text-embedding-005`.""", + ) + ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfig] = Field( + default=None, + description="""Optional. Configuration for automatic TTL ("time-to-live") of the memories in the Memory Bank. If not set, TTL will not be applied automatically. The TTL can be explicitly set by modifying the `expire_time` of each Memory resource.""", + ) + structured_memory_configs: Optional[list[StructuredMemoryConfig]] = Field( + default=None, + description="""Optional. Configuration for organizing structured memories for a particular scope.""", ) -class MemoryBankCustomizationConfigMemoryTopicDict(TypedDict, total=False): - """A topic of information that should be extracted from conversations and stored as memories.""" +class ReasoningEngineContextSpecMemoryBankConfigDict(TypedDict, total=False): + """Specification for a Memory Bank.""" - custom_memory_topic: Optional[ - MemoryBankCustomizationConfigMemoryTopicCustomMemoryTopicDict + customization_configs: Optional[list[MemoryBankCustomizationConfigDict]] + """Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""" + + disable_memory_revisions: Optional[bool] + """If true, no memory revisions will be created for any requests to the Memory Bank.""" + + generation_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict ] - """A custom memory topic defined by the developer.""" + """Optional. Configuration for how to generate memories for the Memory Bank.""" - managed_memory_topic: Optional[ - MemoryBankCustomizationConfigMemoryTopicManagedMemoryTopicDict + similarity_search_config: Optional[ + ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict ] - """A managed memory topic defined by Memory Bank.""" + """Optional. Configuration for how to perform similarity search on memories. If not set, the Memory Bank will use the default embedding model `text-embedding-005`.""" + ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict] + """Optional. Configuration for automatic TTL ("time-to-live") of the memories in the Memory Bank. If not set, TTL will not be applied automatically. The TTL can be explicitly set by modifying the `expire_time` of each Memory resource.""" -MemoryBankCustomizationConfigMemoryTopicOrDict = Union[ - MemoryBankCustomizationConfigMemoryTopic, - MemoryBankCustomizationConfigMemoryTopicDict, + structured_memory_configs: Optional[list[StructuredMemoryConfigDict]] + """Optional. Configuration for organizing structured memories for a particular scope.""" + + +ReasoningEngineContextSpecMemoryBankConfigOrDict = Union[ + ReasoningEngineContextSpecMemoryBankConfig, + ReasoningEngineContextSpecMemoryBankConfigDict, ] -class MemoryBankCustomizationConfigConsolidationConfig(_common.BaseModel): - """Represents configuration for customizing how memories are consolidated.""" +class ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfig( + _common.BaseModel +): + """Configuration for how to perform similarity search on examples.""" - revisions_per_candidate_count: Optional[int] = Field( + embedding_model: Optional[str] = Field( default=None, - description="""Optional. Represents the maximum number of revisions to consider for each candidate memory. If not set, then the default value (1) will be used, which means that only the latest revision will be considered.""", + description="""Required. The Gemini model used to generate embeddings to lookup similar examples. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", ) -class MemoryBankCustomizationConfigConsolidationConfigDict(TypedDict, total=False): - """Represents configuration for customizing how memories are consolidated.""" +class ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfigDict( + TypedDict, total=False +): + """Configuration for how to perform similarity search on examples.""" - 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.""" + embedding_model: Optional[str] + """Required. The Gemini model used to generate embeddings to lookup similar examples. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" -MemoryBankCustomizationConfigConsolidationConfigOrDict = Union[ - MemoryBankCustomizationConfigConsolidationConfig, - MemoryBankCustomizationConfigConsolidationConfigDict, +ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfigOrDict = Union[ + ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfig, + ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfigDict, ] -class MemoryBankCustomizationConfig(_common.BaseModel): - """Represents configuration for organizing natural language memories for a particular scope.""" +class ReasoningEngineContextSpecExampleStoreConfig(_common.BaseModel): + """Specification for an Example Store.""" - 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] + similarity_search_config: Optional[ + ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfig ] = Field( default=None, - description="""Optional. Provides examples of how to generate memories for a particular scope.""", - ) - memory_topics: Optional[list[MemoryBankCustomizationConfigMemoryTopic]] = Field( - default=None, - description="""Optional. Represents topics of information that should be extracted from conversations and stored as memories. If not set, then Memory Bank's default topics will be used.""", + description="""Optional. Configuration for how to perform similarity search on examples. If not set, the Example Store will use the default embedding model `text-embedding-005`.""", ) - scope_keys: Optional[list[str]] = Field( + + +class ReasoningEngineContextSpecExampleStoreConfigDict(TypedDict, total=False): + """Specification for an Example Store.""" + + similarity_search_config: Optional[ + ReasoningEngineContextSpecExampleStoreConfigSimilaritySearchConfigDict + ] + """Optional. Configuration for how to perform similarity search on examples. If not set, the Example Store will use the default embedding model `text-embedding-005`.""" + + +ReasoningEngineContextSpecExampleStoreConfigOrDict = Union[ + ReasoningEngineContextSpecExampleStoreConfig, + ReasoningEngineContextSpecExampleStoreConfigDict, +] + + +class ReasoningEngineContextSpec(_common.BaseModel): + """Configuration for how Agent Engine sub-resources should manage context.""" + + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfig] = Field( default=None, - description="""Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""", + description="""Optional. Specification for a Memory Bank, which manages memories for the Agent Engine.""", ) - consolidation_config: Optional[MemoryBankCustomizationConfigConsolidationConfig] = ( + example_store_config: Optional[ReasoningEngineContextSpecExampleStoreConfig] = ( Field( default=None, - description="""Optional. Represents configuration for customizing how memories are consolidated together.""", + description="""Optional. Specification for an Example Store, which manages few-shot examples for the Agent Engine.""", ) ) - disable_natural_language_memories: Optional[bool] = Field( - default=None, - description="""Optional. Indicates whether natural language memory generation should be disabled for all requests. By default, natural language memory generation is enabled. Set this to `true` when you only want to generate structured memories.""", - ) - - -class MemoryBankCustomizationConfigDict(TypedDict, total=False): - """Represents configuration for organizing natural language memories for a particular scope.""" - - enable_third_person_memories: Optional[bool] - """Optional. Indicates whether the memories will be generated in the third person (i.e. "The user generates memories with Memory Bank."). By default, the memories will be generated in the first person (i.e. "I generate memories with Memory Bank.")""" - - generate_memories_examples: Optional[ - list[MemoryBankCustomizationConfigGenerateMemoriesExampleDict] - ] - """Optional. Provides examples of how to generate memories for a particular scope.""" - memory_topics: Optional[list[MemoryBankCustomizationConfigMemoryTopicDict]] - """Optional. Represents topics of information that should be extracted from conversations and stored as memories. If not set, then Memory Bank's default topics will be used.""" - scope_keys: Optional[list[str]] - """Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""" +class ReasoningEngineContextSpecDict(TypedDict, total=False): + """Configuration for how Agent Engine sub-resources should manage context.""" - consolidation_config: Optional[MemoryBankCustomizationConfigConsolidationConfigDict] - """Optional. Represents configuration for customizing how memories are consolidated together.""" + memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfigDict] + """Optional. Specification for a Memory Bank, which manages memories for the Agent Engine.""" - 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.""" + example_store_config: Optional[ReasoningEngineContextSpecExampleStoreConfigDict] + """Optional. Specification for an Example Store, which manages few-shot examples for the Agent Engine.""" -MemoryBankCustomizationConfigOrDict = Union[ - MemoryBankCustomizationConfig, MemoryBankCustomizationConfigDict +ReasoningEngineContextSpecOrDict = Union[ + ReasoningEngineContextSpec, ReasoningEngineContextSpecDict ] -class MemoryGenerationTriggerConfigGenerationTriggerRule(_common.BaseModel): - """Represents the active rule that determines when to flush the buffer.""" +class SecretRef(_common.BaseModel): + """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" - event_count: Optional[int] = Field( - default=None, - description="""Optional. Specifies to trigger generation when the event count reaches this limit.""", - ) - fixed_interval: Optional[str] = Field( - default=None, - description="""Optional. Specifies to trigger generation at a fixed interval. The duration must have a minute-level granularity.""", - ) - idle_duration: Optional[str] = Field( + secret: Optional[str] = Field( default=None, - description="""Optional. Specifies to trigger generation if the stream is inactive for the specified duration after the most recent event. The duration must have a minute-level granularity.""", + description="""Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""", ) - overlap_event_count: Optional[int] = Field( + version: Optional[str] = Field( default=None, - description="""Optional. Re-include the last N already-processed events in the next window.""", + 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 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.""" +class SecretRefDict(TypedDict, total=False): + """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" - fixed_interval: Optional[str] - """Optional. Specifies to trigger generation at a fixed interval. The duration must have a minute-level granularity.""" + secret: Optional[str] + """Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""" - idle_duration: Optional[str] - """Optional. Specifies to trigger generation if the stream is inactive for the specified duration after the most recent event. The duration must have a minute-level granularity.""" - - overlap_event_count: Optional[int] - """Optional. Re-include the last N already-processed events in the next window.""" - - -MemoryGenerationTriggerConfigGenerationTriggerRuleOrDict = Union[ - MemoryGenerationTriggerConfigGenerationTriggerRule, - MemoryGenerationTriggerConfigGenerationTriggerRuleDict, -] - - -class MemoryGenerationTriggerConfig(_common.BaseModel): - """The configuration for triggering memory generation for ingested events.""" - - generation_rule: Optional[MemoryGenerationTriggerConfigGenerationTriggerRule] = ( - Field( - default=None, - description="""Optional. Represents the active rule that determines when to flush the buffer. If not set, then the stream will be force flushed immediately.""", - ) - ) - - -class MemoryGenerationTriggerConfigDict(TypedDict, total=False): - """The configuration for triggering memory generation for ingested events.""" - - 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.""" + 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.""" -MemoryGenerationTriggerConfigOrDict = Union[ - MemoryGenerationTriggerConfig, MemoryGenerationTriggerConfigDict -] +SecretRefOrDict = Union[SecretRef, SecretRefDict] -class ReasoningEngineContextSpecMemoryBankConfigGenerationConfig(_common.BaseModel): - """Configuration for how to generate memories.""" +class SecretEnvVar(_common.BaseModel): + """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" - model: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Optional. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", + description="""Required. Name of the secret environment variable.""", ) - generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + secret_ref: Optional[SecretRef] = Field( default=None, - description="""Optional. Specifies the default trigger configuration for generating memories using `IngestEvents`.""", + description="""Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""", ) -class ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict( - TypedDict, total=False -): - """Configuration for how to generate memories.""" +class SecretEnvVarDict(TypedDict, total=False): + """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" - model: Optional[str] - """Optional. The model used to generate memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + name: Optional[str] + """Required. Name of the secret environment variable.""" - generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] - """Optional. Specifies the default trigger configuration for generating memories using `IngestEvents`.""" + secret_ref: Optional[SecretRefDict] + """Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" -ReasoningEngineContextSpecMemoryBankConfigGenerationConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigGenerationConfig, - ReasoningEngineContextSpecMemoryBankConfigGenerationConfigDict, -] +SecretEnvVarOrDict = Union[SecretEnvVar, SecretEnvVarDict] -class ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig( +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig( _common.BaseModel ): - """Configuration for how to perform similarity search on memories.""" + """Configuration for traffic originating from a Reasoning Engine.""" - embedding_model: Optional[str] = Field( + agent_gateway: Optional[str] = Field( default=None, - description="""Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""", + description="""Required. The resource name of the Agent Gateway 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 ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict( +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict( TypedDict, total=False ): - """Configuration for how to perform similarity search on memories.""" + """Configuration for traffic originating from a Reasoning Engine.""" - embedding_model: Optional[str] - """Required. The model used to generate embeddings to lookup similar memories. Format: `projects/{project}/locations/{location}/publishers/google/models/{model}`.""" + agent_gateway: Optional[str] + """Required. The resource name of the Agent Gateway 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}`""" -ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig, - ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfigDict, +ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigOrDict = Union[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig, + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict, ] -class ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig( +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig( _common.BaseModel ): - """Configuration for TTL of the memories in the Memory Bank based on the action that created or updated the memory.""" + """Configuration for traffic targeting a Reasoning Engine.""" - 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( + agent_gateway: 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).""", + 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 ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict( +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict( 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).""" + """Configuration for traffic targeting a Reasoning Engine.""" - 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).""" + 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}`""" -ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfig, - ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict, +ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigOrDict = Union[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig, + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict, ] -class ReasoningEngineContextSpecMemoryBankConfigTtlConfig(_common.BaseModel): - """Configuration for automatically setting the TTL ("time-to-live") of the memories in the Memory Bank.""" +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfig(_common.BaseModel): + """Agent Gateway configuration for a Reasoning Engine deployment.""" - 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 + agent_to_anywhere_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig ] = Field( default=None, - description="""Optional. The granular TTL configuration of the memories in the Memory Bank.""", + description="""Optional. Configuration for traffic originating from the Reasoning Engine. When unset, outgoing traffic is not routed through an Agent Gateway.""", ) - memory_revision_default_ttl: Optional[str] = Field( + client_to_agent_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig + ] = Field( default=None, - description="""Optional. The default TTL duration of the memory revisions in the Memory Bank. This applies to all operations that create a memory revision. If not set, a default TTL of 365 days will be used.""", + description="""Optional. Configuration for traffic targeting the Reasoning Engine. When unset, incoming traffic is not routed through an Agent Gateway.""", ) -class ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict(TypedDict, total=False): - """Configuration for automatically setting the TTL ("time-to-live") of the memories in the Memory Bank.""" - - default_ttl: Optional[str] - """Optional. The default TTL duration of the memories in the Memory Bank. This applies to all operations that create or update a memory.""" +class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict(TypedDict, total=False): + """Agent Gateway configuration for a Reasoning Engine deployment.""" - granular_ttl_config: Optional[ - ReasoningEngineContextSpecMemoryBankConfigTtlConfigGranularTtlConfigDict + agent_to_anywhere_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict ] - """Optional. The granular TTL configuration of the memories in the Memory Bank.""" + """Optional. Configuration for traffic originating from the Reasoning Engine. When unset, outgoing traffic is not routed through an Agent Gateway.""" - 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.""" + 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.""" -ReasoningEngineContextSpecMemoryBankConfigTtlConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfigTtlConfig, - ReasoningEngineContextSpecMemoryBankConfigTtlConfigDict, +ReasoningEngineSpecDeploymentSpecAgentGatewayConfigOrDict = Union[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig, + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict, ] -class StructuredMemorySchemaConfig(_common.BaseModel): - """Represents the OpenAPI schema of the structured memories.""" +class KeepAliveProbeHttpGet(_common.BaseModel): + """Specifies the HTTP GET configuration for the probe.""" - memory_schema: Optional[genai_types.Schema] = Field( - default=None, - description="""Required. Represents the OpenAPI schema of the structured memories.""", - ) - id: Optional[str] = Field( + path: 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.""", + description="""Required. Specifies the path of the HTTP GET request (e.g., `"/is_busy"`).""", ) - memory_type: Optional[MemoryType] = Field( + port: Optional[int] = Field( default=None, - description="""Optional. Represents the type of the structured memories associated with the schema. If not set, then `STRUCTURED_PROFILE` will be used.""", + description="""Optional. Specifies the port number on the container to which the request is sent.""", ) -class StructuredMemorySchemaConfigDict(TypedDict, total=False): - """Represents the OpenAPI schema of the structured memories.""" - - memory_schema: Optional[genai_types.Schema] - """Required. Represents the OpenAPI schema of the structured memories.""" +class KeepAliveProbeHttpGetDict(TypedDict, total=False): + """Specifies the HTTP GET configuration for the probe.""" - 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.""" + path: Optional[str] + """Required. Specifies the path of the HTTP GET request (e.g., `"/is_busy"`).""" - 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.""" + port: Optional[int] + """Optional. Specifies the port number on the container to which the request is sent.""" -StructuredMemorySchemaConfigOrDict = Union[ - StructuredMemorySchemaConfig, StructuredMemorySchemaConfigDict -] +KeepAliveProbeHttpGetOrDict = Union[KeepAliveProbeHttpGet, KeepAliveProbeHttpGetDict] -class StructuredMemoryConfig(_common.BaseModel): - """Configuration for organizing structured memories within a scope.""" +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.""" - schema_configs: Optional[list[StructuredMemorySchemaConfig]] = Field( + http_get: Optional[KeepAliveProbeHttpGet] = Field( default=None, - description="""Optional. Represents configuration of the structured memories' schemas.""", + description="""Optional. Specifies the HTTP GET configuration for the probe.""", ) - scope_keys: Optional[list[str]] = Field( + max_seconds: Optional[int] = Field( default=None, - description="""Optional. Represents the scope keys (i.e. 'user_id') for which to use this config. A request's scope must include all of the provided keys for the config to be used (order does not matter). If empty, then the config will be used for all requests that do not have a more specific config. Only one default config is allowed per Memory Bank.""", + description="""Optional. Specifies the maximum duration (in seconds) to keep the instance alive via this probe. Can be a maximum of 3600 seconds (1 hour).""", ) -class StructuredMemoryConfigDict(TypedDict, total=False): - """Configuration for organizing structured memories within a scope.""" +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.""" - schema_configs: Optional[list[StructuredMemorySchemaConfigDict]] - """Optional. Represents configuration of the structured memories' schemas.""" + http_get: Optional[KeepAliveProbeHttpGetDict] + """Optional. Specifies the HTTP GET configuration for the probe.""" - 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.""" + 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).""" -StructuredMemoryConfigOrDict = Union[StructuredMemoryConfig, StructuredMemoryConfigDict] +KeepAliveProbeOrDict = Union[KeepAliveProbe, KeepAliveProbeDict] -class ReasoningEngineContextSpecMemoryBankConfig(_common.BaseModel): - """Specification for a Memory Bank.""" +class ReasoningEngineSpecDeploymentSpec(_common.BaseModel): + """The specification of a Reasoning Engine deployment.""" - customization_configs: Optional[list[MemoryBankCustomizationConfig]] = Field( - default=None, - description="""Optional. Configuration for how to customize Memory Bank behavior for a particular scope.""", + agent_server_mode: Optional[AgentServerMode] = Field( + default=None, description="""The agent server mode.""" ) - disable_memory_revisions: Optional[bool] = Field( + container_concurrency: Optional[int] = Field( default=None, - description="""If true, no memory revisions will be created for any requests to the Memory Bank.""", + description="""Optional. Concurrency for each container and agent server. Recommended value: 2 * cpu + 1. Defaults to 9.""", ) - generation_config: Optional[ - ReasoningEngineContextSpecMemoryBankConfigGenerationConfig - ] = Field( + env: Optional[list[EnvVar]] = Field( default=None, - description="""Optional. Configuration for how to generate memories for the Memory Bank.""", + description="""Optional. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API.""", ) - similarity_search_config: Optional[ - ReasoningEngineContextSpecMemoryBankConfigSimilaritySearchConfig - ] = Field( + max_instances: Optional[int] = 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 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].""", ) - ttl_config: Optional[ReasoningEngineContextSpecMemoryBankConfigTtlConfig] = Field( + min_instances: Optional[int] = Field( default=None, - description="""Optional. Configuration for automatic TTL ("time-to-live") of the memories in the Memory Bank. If not set, TTL will not be applied automatically. The TTL can be explicitly set by modifying the `expire_time` of each Memory resource.""", + description="""Optional. The minimum number of application instances that will be kept running at all times. Defaults to 1. Range: [0, 75].""", ) - structured_memory_configs: Optional[list[StructuredMemoryConfig]] = Field( + psc_interface_config: Optional[PscInterfaceConfig] = Field( + default=None, description="""Optional. Configuration for PSC-I.""" + ) + resource_limits: Optional[dict[str, str]] = Field( default=None, - description="""Optional. Configuration for organizing structured memories for a particular scope.""", + 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.""", + ) + dedicated_ingress_endpoint_enabled: Optional[bool] = Field( + default=None, + description="""Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""", ) -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.""" +class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False): + """The specification of a Reasoning Engine deployment.""" - 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`.""" + agent_server_mode: Optional[AgentServerMode] + """The agent server mode.""" - 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.""" + container_concurrency: Optional[int] + """Optional. Concurrency for each container and agent server. Recommended value: 2 * cpu + 1. Defaults to 9.""" - structured_memory_configs: Optional[list[StructuredMemoryConfigDict]] - """Optional. Configuration for organizing structured memories for a particular scope.""" + 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].""" -ReasoningEngineContextSpecMemoryBankConfigOrDict = Union[ - ReasoningEngineContextSpecMemoryBankConfig, - ReasoningEngineContextSpecMemoryBankConfigDict, -] + 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].""" + psc_interface_config: Optional[PscInterfaceConfigDict] + """Optional. Configuration for PSC-I.""" -class ReasoningEngineContextSpec(_common.BaseModel): - """Configuration for how Agent Engine sub-resources should manage context.""" + 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""" - memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfig] = Field( - default=None, - description="""Optional. Specification for a Memory Bank, which manages memories for the Agent Engine.""", - ) + secret_env: Optional[list[SecretEnvVarDict]] + """Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""" + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] + """Optional. Agent Gateway configuration for the Reasoning Engine deployment.""" -class ReasoningEngineContextSpecDict(TypedDict, total=False): - """Configuration for how Agent Engine sub-resources should manage context.""" + 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.""" - memory_bank_config: Optional[ReasoningEngineContextSpecMemoryBankConfigDict] - """Optional. Specification for a Memory Bank, which manages memories for the Agent Engine.""" + dedicated_ingress_endpoint_enabled: Optional[bool] + """Optional. If true, the Reasoning Engine will be deployed with a dedicated ingress endpoint.""" -ReasoningEngineContextSpecOrDict = Union[ - ReasoningEngineContextSpec, ReasoningEngineContextSpecDict +ReasoningEngineSpecDeploymentSpecOrDict = Union[ + ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict ] -class SecretRef(_common.BaseModel): - """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" +class ReasoningEngineSpecPackageSpec(_common.BaseModel): + """User-provided package specification, containing pickled object and package requirements.""" - secret: Optional[str] = Field( + dependency_files_gcs_uri: Optional[str] = Field( default=None, - description="""Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""", + description="""Optional. The Cloud Storage URI of the dependency files in tar.gz format.""", ) - version: Optional[str] = Field( + pickle_object_gcs_uri: Optional[str] = Field( default=None, - description="""The Cloud Secret Manager secret version. Can be 'latest' for the latest version, an integer for a specific version, or a version alias.""", + description="""Optional. The Cloud Storage URI of the pickled python object.""", ) - - -class SecretRefDict(TypedDict, total=False): - """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" - - secret: Optional[str] - """Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""" - - 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.""" - - -SecretRefOrDict = Union[SecretRef, SecretRefDict] - - -class SecretEnvVar(_common.BaseModel): - """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" - - name: Optional[str] = Field( + python_version: Optional[str] = Field( default=None, - description="""Required. Name of the secret environment variable.""", + 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.""", ) - secret_ref: Optional[SecretRef] = Field( + requirements_gcs_uri: Optional[str] = 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="""Optional. The Cloud Storage URI of the `requirements.txt` file""", ) -class SecretEnvVarDict(TypedDict, total=False): - """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" +class ReasoningEngineSpecPackageSpecDict(TypedDict, total=False): + """User-provided package specification, containing pickled object and package requirements.""" - name: Optional[str] - """Required. Name of the secret environment variable.""" + dependency_files_gcs_uri: Optional[str] + """Optional. The Cloud Storage URI of the dependency files in tar.gz format.""" - secret_ref: Optional[SecretRefDict] - """Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" + 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.""" -SecretEnvVarOrDict = Union[SecretEnvVar, SecretEnvVarDict] + requirements_gcs_uri: Optional[str] + """Optional. The Cloud Storage URI of the `requirements.txt` file""" -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig( - _common.BaseModel -): - """Configuration for traffic originating from a Reasoning Engine.""" +ReasoningEngineSpecPackageSpecOrDict = Union[ + ReasoningEngineSpecPackageSpec, ReasoningEngineSpecPackageSpecDict +] - agent_gateway: Optional[str] = Field( + +class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig(_common.BaseModel): + """Configuration for the Agent Development Kit (ADK).""" + + json_config: Optional[dict[str, Any]] = Field( default=None, - description="""Required. The resource name of the Agent Gateway for outbound traffic. It must be set to a Google-managed gateway whose `governed_access_path` is `AGENT_TO_ANYWHERE`. Format: `projects/{project}/locations/{location}/agentGateways/{agent_gateway}`""", + description="""Required. The value of the ADK config in JSON format.""", ) -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict( +class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict( TypedDict, total=False ): - """Configuration for traffic originating from a Reasoning Engine.""" + """Configuration for the Agent Development Kit (ADK).""" - 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}`""" + json_config: Optional[dict[str, Any]] + """Required. The value of the ADK config in JSON format.""" -ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigOrDict = Union[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfig, - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict, +ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig, + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict, ] -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig( - _common.BaseModel -): - """Configuration for traffic targeting a Reasoning Engine.""" +class ReasoningEngineSpecSourceCodeSpecInlineSource(_common.BaseModel): + """Specifies source code provided as a byte stream.""" - agent_gateway: Optional[str] = Field( + source_archive: Optional[bytes] = Field( default=None, - description="""Required. The resource name of the Agent Gateway 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="""Required. Input only. The application source code archive. It must be a compressed tarball (.tar.gz) file.""", ) -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict( - TypedDict, total=False -): - """Configuration for traffic targeting a Reasoning Engine.""" +class ReasoningEngineSpecSourceCodeSpecInlineSourceDict(TypedDict, total=False): + """Specifies source code provided as a byte stream.""" - agent_gateway: Optional[str] - """Required. The resource name of the Agent Gateway 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}`""" + source_archive: Optional[bytes] + """Required. Input only. The application source code archive. It must be a compressed tarball (.tar.gz) file.""" -ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigOrDict = Union[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig, - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfigDict, +ReasoningEngineSpecSourceCodeSpecInlineSourceOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecInlineSource, + ReasoningEngineSpecSourceCodeSpecInlineSourceDict, ] -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfig(_common.BaseModel): - """Agent Gateway configuration for a Reasoning Engine deployment.""" +class ReasoningEngineSpecSourceCodeSpecAgentConfigSource(_common.BaseModel): + """Specification for the deploying from agent config.""" - 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.""", - ) - client_to_agent_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigClientToAgentConfig - ] = Field( + adk_config: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig + ] = Field(default=None, description="""Required. The ADK configuration.""") + inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSource] = Field( default=None, - description="""Optional. Configuration for traffic targeting the Reasoning Engine. When unset, incoming traffic is not routed through an Agent Gateway.""", + description="""Optional. Any additional files needed to interpret the config. If a `requirements.txt` file is present in the `inline_source`, the corresponding packages will be installed. If no `requirements.txt` file is present in `inline_source`, then the latest version of `google-adk` will be installed for interpreting the ADK config.""", ) -class ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict(TypedDict, total=False): - """Agent Gateway configuration for a Reasoning Engine deployment.""" +class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict(TypedDict, total=False): + """Specification for the deploying from agent config.""" - agent_to_anywhere_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigAgentToAnywhereConfigDict + adk_config: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict ] - """Optional. Configuration for traffic originating from the Reasoning Engine. When unset, outgoing traffic is not routed through an Agent Gateway.""" + """Required. The ADK configuration.""" - 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.""" + 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.""" -ReasoningEngineSpecDeploymentSpecAgentGatewayConfigOrDict = Union[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig, - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict, +ReasoningEngineSpecSourceCodeSpecAgentConfigSourceOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSource, + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict, ] -class KeepAliveProbeHttpGet(_common.BaseModel): - """Specifies the HTTP GET configuration for the probe.""" +class ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig(_common.BaseModel): + """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. - path: Optional[str] = Field( + This includes the repository, revision, and directory to use. + """ + + git_repository_link: Optional[str] = Field( default=None, - description="""Required. Specifies the path of the HTTP GET request (e.g., `"/is_busy"`).""", + description="""Required. The Developer Connect Git repository link, formatted as `projects/{project_id}/locations/{location_id}/connections/{connection_id}/gitRepositoryLink/{repository_link_id}`.""", ) - port: Optional[int] = Field( + dir: Optional[str] = Field( default=None, - description="""Optional. Specifies the port number on the container to which the request is sent.""", + 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 KeepAliveProbeHttpGetDict(TypedDict, total=False): - """Specifies the HTTP GET configuration for the probe.""" - - path: Optional[str] - """Required. Specifies the path of the HTTP GET request (e.g., `"/is_busy"`).""" +class ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict( + TypedDict, total=False +): + """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. - port: Optional[int] - """Optional. Specifies the port number on the container to which the request is sent.""" + This includes the repository, revision, and directory to use. + """ + git_repository_link: Optional[str] + """Required. The Developer Connect Git repository link, formatted as `projects/{project_id}/locations/{location_id}/connections/{connection_id}/gitRepositoryLink/{repository_link_id}`.""" -KeepAliveProbeHttpGetOrDict = Union[KeepAliveProbeHttpGet, KeepAliveProbeHttpGetDict] + dir: Optional[str] + """Required. Directory, relative to the source root, in which to run the build.""" + 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.""" -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( - default=None, - description="""Optional. Specifies the HTTP GET configuration for the probe.""", - ) - max_seconds: Optional[int] = Field( +ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig, + ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict, +] + + +class ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource(_common.BaseModel): + """Specifies source code to be fetched from a Git repository managed through the Developer Connect service.""" + + config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig] = 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="""Required. The Developer Connect configuration that defines the specific repository, revision, and directory to use as the source code root.""", ) -class KeepAliveProbeDict(TypedDict, total=False): - """Represents the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""" - - http_get: Optional[KeepAliveProbeHttpGetDict] - """Optional. Specifies the HTTP GET configuration for the probe.""" +class ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict( + TypedDict, total=False +): + """Specifies source code to be fetched from a Git repository managed through the Developer Connect service.""" - 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).""" + config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict] + """Required. The Developer Connect configuration that defines the specific repository, revision, and directory to use as the source code root.""" -KeepAliveProbeOrDict = Union[KeepAliveProbe, KeepAliveProbeDict] +ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource, + ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict, +] -class ReasoningEngineSpecDeploymentSpec(_common.BaseModel): - """The specification of a Reasoning Engine deployment.""" +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.""" - 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( + build_args: Optional[dict[str, str]] = Field( default=None, - description="""Optional. The minimum number of application instances that will be kept running at all times. Defaults to 1. Range: [0, 75].""", - ) - psc_interface_config: Optional[PscInterfaceConfig] = Field( - default=None, description="""Optional. Configuration for PSC-I.""" + description="""Optional. Build arguments to be used. They will be passed through --build-arg flags.""", ) - resource_limits: Optional[dict[str, str]] = Field( + + +class ReasoningEngineSpecSourceCodeSpecImageSpecDict(TypedDict, total=False): + """The image spec for building an image (within a single build step), based on the config file (i.e. Dockerfile) in the source directory.""" + + build_args: Optional[dict[str, str]] + """Optional. Build arguments to be used. They will be passed through --build-arg flags.""" + + +ReasoningEngineSpecSourceCodeSpecImageSpecOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecImageSpec, + ReasoningEngineSpecSourceCodeSpecImageSpecDict, +] + + +class ReasoningEngineSpecSourceCodeSpecPythonSpec(_common.BaseModel): + """Specification for running a Python application from source.""" + + entrypoint_module: Optional[str] = Field( default=None, - description="""Optional. Resource limits for each container. Only 'cpu' and 'memory' keys are supported. Defaults to {"cpu": "4", "memory": "4Gi"}. * The only supported values for CPU are '1', '2', '4', '6' and '8'. For more information, go to https://cloud.google.com/run/docs/configuring/cpu. * The only supported values for memory are '1Gi', '2Gi', ... '32 Gi'. * For required cpu on different memory values, go to https://cloud.google.com/run/docs/configuring/memory-limits""", + description="""Optional. The 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`.""", ) - secret_env: Optional[list[SecretEnvVar]] = Field( + entrypoint_object: Optional[str] = Field( default=None, - description="""Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""", + description="""Optional. The 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`.""", ) - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig - ] = Field( + requirements_file: Optional[str] = Field( default=None, - description="""Optional. Agent Gateway configuration for the Reasoning Engine deployment.""", + description="""Optional. The path to the requirements file, relative to the source root. If not specified, defaults to "requirements.txt".""", ) - keep_alive_probe: Optional[KeepAliveProbe] = Field( + version: Optional[str] = Field( default=None, - description="""Optional. Specifies the configuration for keep-alive probe. Contains configuration on a specified endpoint that a deployment host should use to keep the container alive based on the probe settings.""", + description="""Optional. 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 ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False): - """The specification of a Reasoning Engine deployment.""" - - 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.""" - - 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].""" - - 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].""" - - psc_interface_config: Optional[PscInterfaceConfigDict] - """Optional. Configuration for PSC-I.""" +class ReasoningEngineSpecSourceCodeSpecPythonSpecDict(TypedDict, total=False): + """Specification for running a Python application from source.""" - 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""" + 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`.""" - 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.""" + 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`.""" - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict - ] - """Optional. Agent Gateway configuration for the Reasoning Engine deployment.""" + requirements_file: Optional[str] + """Optional. The path to the requirements file, relative to the source root. If not specified, defaults to "requirements.txt".""" - 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.""" + 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.""" -ReasoningEngineSpecDeploymentSpecOrDict = Union[ - ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict +ReasoningEngineSpecSourceCodeSpecPythonSpecOrDict = Union[ + ReasoningEngineSpecSourceCodeSpecPythonSpec, + ReasoningEngineSpecSourceCodeSpecPythonSpecDict, ] -class ReasoningEngineSpecPackageSpec(_common.BaseModel): - """User-provided package specification, containing pickled object and package requirements.""" +class ReasoningEngineSpecSourceCodeSpec(_common.BaseModel): + """Specification for deploying from source code.""" - dependency_files_gcs_uri: Optional[str] = Field( - default=None, - description="""Optional. The Cloud Storage URI of the dependency files in tar.gz format.""", + agent_config_source: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSource + ] = Field( + default=None, description="""Source code is generated from the agent config.""" ) - pickle_object_gcs_uri: Optional[str] = Field( + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource + ] = Field( default=None, - description="""Optional. The Cloud Storage URI of the pickled python object.""", + description="""Source code is in a Git repository managed by Developer Connect.""", ) - python_version: Optional[str] = Field( + image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpec] = Field( default=None, - description="""Optional. The Python version. Supported values are 3.10, 3.11, 3.12, 3.13, 3.14. If not specified, the default value is 3.10.""", + description="""Optional. Configuration for building an image with custom config file.""", ) - requirements_gcs_uri: Optional[str] = Field( - default=None, - description="""Optional. The Cloud Storage URI of the `requirements.txt` file""", + inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSource] = Field( + default=None, description="""Source code is provided directly in the request.""" + ) + python_spec: Optional[ReasoningEngineSpecSourceCodeSpecPythonSpec] = Field( + default=None, description="""Configuration for a Python application.""" ) -class ReasoningEngineSpecPackageSpecDict(TypedDict, total=False): - """User-provided package specification, containing pickled object and package requirements.""" +class ReasoningEngineSpecSourceCodeSpecDict(TypedDict, total=False): + """Specification for deploying from source code.""" - dependency_files_gcs_uri: Optional[str] - """Optional. The Cloud Storage URI of the dependency files in tar.gz format.""" + agent_config_source: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict + ] + """Source code is generated from the agent config.""" - pickle_object_gcs_uri: Optional[str] - """Optional. The Cloud Storage URI of the pickled python object.""" + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict + ] + """Source code is in a Git repository managed by Developer Connect.""" - 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.""" + image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict] + """Optional. Configuration for building an image with custom config file.""" - requirements_gcs_uri: Optional[str] - """Optional. The Cloud Storage URI of the `requirements.txt` file""" + inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSourceDict] + """Source code is provided directly in the request.""" + python_spec: Optional[ReasoningEngineSpecSourceCodeSpecPythonSpecDict] + """Configuration for a Python application.""" -ReasoningEngineSpecPackageSpecOrDict = Union[ - ReasoningEngineSpecPackageSpec, ReasoningEngineSpecPackageSpecDict + +ReasoningEngineSpecSourceCodeSpecOrDict = Union[ + ReasoningEngineSpecSourceCodeSpec, ReasoningEngineSpecSourceCodeSpecDict ] -class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig(_common.BaseModel): - """Configuration for the Agent Development Kit (ADK).""" +class ReasoningEngineSpecContainerSpec(_common.BaseModel): + """Specification for deploying from a container image.""" - json_config: Optional[dict[str, Any]] = Field( + image_uri: Optional[str] = Field( default=None, - description="""Required. The value of the ADK config in JSON format.""", + description="""Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""", ) -class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict( - TypedDict, total=False -): - """Configuration for the Agent Development Kit (ADK).""" +class ReasoningEngineSpecContainerSpecDict(TypedDict, total=False): + """Specification for deploying from a container image.""" - json_config: Optional[dict[str, Any]] - """Required. The value of the ADK config in JSON format.""" + 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.""" -ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfig, - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict, +ReasoningEngineSpecContainerSpecOrDict = Union[ + ReasoningEngineSpecContainerSpec, ReasoningEngineSpecContainerSpecDict ] -class ReasoningEngineSpecSourceCodeSpecInlineSource(_common.BaseModel): - """Specifies source code provided as a byte stream.""" +class ReasoningEngineSpecBuildSpec(_common.BaseModel): + """Specification for building container image.""" - source_archive: Optional[bytes] = Field( + worker_pool: Optional[str] = Field( default=None, - description="""Required. Input only. The application source code archive. It must be a compressed tarball (.tar.gz) file.""", + 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 ReasoningEngineSpecSourceCodeSpecInlineSourceDict(TypedDict, total=False): - """Specifies source code provided as a byte stream.""" +class ReasoningEngineSpecBuildSpecDict(TypedDict, total=False): + """Specification for building container image.""" - 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 ReasoningEngineSpecSourceCodeSpecAgentConfigSource(_common.BaseModel): - """Specification for the deploying from agent config.""" - - 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.""", - ) - - -class ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict(TypedDict, total=False): - """Specification for the deploying from agent config.""" - - adk_config: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceAdkConfigDict - ] - """Required. The ADK configuration.""" - - inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSourceDict] - """Optional. Any additional files needed to interpret the config. If a `requirements.txt` file is present in the `inline_source`, the corresponding packages will be installed. If no `requirements.txt` file is present in `inline_source`, then the latest version of `google-adk` will be installed for interpreting the ADK config.""" - - -ReasoningEngineSpecSourceCodeSpecAgentConfigSourceOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSource, - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict, -] - - -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. - """ - - git_repository_link: Optional[str] = Field( - default=None, - description="""Required. The Developer Connect Git repository link, formatted as `projects/{project_id}/locations/{location_id}/connections/{connection_id}/gitRepositoryLink/{repository_link_id}`.""", - ) - dir: Optional[str] = Field( - default=None, - description="""Required. Directory, relative to the source root, in which to run the build.""", - ) - revision: Optional[str] = Field( - default=None, - description="""Required. The revision to fetch from the Git repository such as a branch, a tag, a commit SHA, or any Git ref.""", - ) - - -class ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict( - TypedDict, total=False -): - """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. - - This includes the repository, revision, and directory to use. - """ - - git_repository_link: Optional[str] - """Required. The Developer Connect Git repository link, formatted as `projects/{project_id}/locations/{location_id}/connections/{connection_id}/gitRepositoryLink/{repository_link_id}`.""" - - dir: Optional[str] - """Required. Directory, relative to the source root, in which to run the build.""" - - 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 ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource(_common.BaseModel): - """Specifies source code to be fetched from a Git repository managed through the Developer Connect service.""" - - 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 ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict( - TypedDict, total=False -): - """Specifies source code to be fetched from a Git repository managed through the Developer Connect service.""" - - config: Optional[ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict] - """Required. The Developer Connect configuration that defines the specific repository, revision, and directory to use as the source code root.""" - - -ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource, - ReasoningEngineSpecSourceCodeSpecDeveloperConnectSourceDict, -] - - -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.""" - - 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 ReasoningEngineSpecSourceCodeSpecImageSpecDict(TypedDict, total=False): - """The image spec for building an image (within a single build step), based on the config file (i.e. Dockerfile) in the source directory.""" - - build_args: Optional[dict[str, str]] - """Optional. Build arguments to be used. They will be passed through --build-arg flags.""" - - -ReasoningEngineSpecSourceCodeSpecImageSpecOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecImageSpec, - ReasoningEngineSpecSourceCodeSpecImageSpecDict, -] + 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}`""" -class ReasoningEngineSpecSourceCodeSpecPythonSpec(_common.BaseModel): - """Specification for running a Python application from source.""" - - 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".""", - ) - version: 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.""", - ) - - -class ReasoningEngineSpecSourceCodeSpecPythonSpecDict(TypedDict, total=False): - """Specification for running a Python application from source.""" - - entrypoint_module: Optional[str] - """Optional. The Python module to load as the entrypoint, specified as a fully qualified module name. For example: path.to.agent. If not specified, defaults to "agent". The project root will be added to Python sys.path, allowing imports to be specified relative to the root. This field should not be set if the source is `agent_config_source`.""" - - entrypoint_object: Optional[str] - """Optional. The name of the callable object within the `entrypoint_module` to use as the application If not specified, defaults to "root_agent". This field should not be set if the source is `agent_config_source`.""" - - requirements_file: Optional[str] - """Optional. The path to the requirements file, relative to the source root. If not specified, defaults to "requirements.txt".""" - - version: Optional[str] - """Optional. The version of Python to use. Supported versions include 3.10, 3.11, 3.12, 3.13, 3.14. If not specified, default value is 3.10.""" - - -ReasoningEngineSpecSourceCodeSpecPythonSpecOrDict = Union[ - ReasoningEngineSpecSourceCodeSpecPythonSpec, - ReasoningEngineSpecSourceCodeSpecPythonSpecDict, -] - - -class ReasoningEngineSpecSourceCodeSpec(_common.BaseModel): - """Specification for deploying from source code.""" - - agent_config_source: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSource - ] = Field( - default=None, description="""Source code is generated from the agent config.""" - ) - developer_connect_source: Optional[ - ReasoningEngineSpecSourceCodeSpecDeveloperConnectSource - ] = Field( - default=None, - description="""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.""" - ) - - -class ReasoningEngineSpecSourceCodeSpecDict(TypedDict, total=False): - """Specification for deploying from source code.""" - - agent_config_source: Optional[ - ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict - ] - """Source code is generated from the agent config.""" - - 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.""" - - inline_source: Optional[ReasoningEngineSpecSourceCodeSpecInlineSourceDict] - """Source code is provided directly in the request.""" - - python_spec: Optional[ReasoningEngineSpecSourceCodeSpecPythonSpecDict] - """Configuration for a Python application.""" - - -ReasoningEngineSpecSourceCodeSpecOrDict = Union[ - ReasoningEngineSpecSourceCodeSpec, ReasoningEngineSpecSourceCodeSpecDict -] - - -class ReasoningEngineSpecContainerSpec(_common.BaseModel): - """Specification for deploying from a container image.""" - - image_uri: Optional[str] = Field( - default=None, - description="""Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""", - ) - - -class ReasoningEngineSpecContainerSpecDict(TypedDict, total=False): - """Specification for deploying from a container image.""" - - image_uri: Optional[str] - """Required. The Artifact Registry Docker image URI (e.g., us-central1-docker.pkg.dev/my-project/my-repo/my-image:tag) of the container image that is to be run on each worker replica.""" - - -ReasoningEngineSpecContainerSpecOrDict = Union[ - ReasoningEngineSpecContainerSpec, ReasoningEngineSpecContainerSpecDict +ReasoningEngineSpecBuildSpecOrDict = Union[ + ReasoningEngineSpecBuildSpec, ReasoningEngineSpecBuildSpecDict ] class ReasoningEngineSpec(_common.BaseModel): - """The specification of an agent engine.""" + """The specification of an agent runtime.""" agent_card: Optional[dict[str, Any]] = Field( default=None, @@ -8448,7 +8228,7 @@ class ReasoningEngineSpec(_common.BaseModel): class ReasoningEngineSpecDict(TypedDict, total=False): - """The specification of an agent engine.""" + """The specification of an agent runtime.""" 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.""" @@ -8480,6 +8260,15 @@ class ReasoningEngineSpecDict(TypedDict, total=False): container_spec: Optional[ReasoningEngineSpecContainerSpecDict] """Deploy from a container image with a defined entrypoint and commands.""" + example_store: Optional[str] + """Optional. The resource name of the linked ExampleStore. At query time, examples can be used to guide the performance of the agent by providing the expected response or demonstrating when and how tools should be called. Format: `projects/{project}/locations/{location}/exampleStores/{example_store}`""" + + runtime_active: Optional[bool] + """Output only. Boolean signifying if this agent engine has a runtime currently or not.""" + + build_spec: Optional[ReasoningEngineSpecBuildSpecDict] + """Optional. Configuration for building container image.""" + ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict] @@ -8589,18 +8378,116 @@ class ReasoningEngineTrafficConfigDict(TypedDict, total=False): ] -class ReasoningEngine(_common.BaseModel): - """An agent engine.""" +class ExperimentConfig(_common.BaseModel): + """The experiment config to control feature experiments.""" - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( - default=None, - description="""Customer-managed encryption key spec for a ReasoningEngine. If set, this ReasoningEngine and all sub-resources of this ReasoningEngine will be secured by this key.""", - ) - context_spec: Optional[ReasoningEngineContextSpec] = Field( - default=None, - description="""Optional. Configuration for how Agent Engine sub-resources should manage context.""", - ) - create_time: Optional[datetime.datetime] = Field( + pass + + +class ExperimentConfigDict(TypedDict, total=False): + """The experiment config to control feature experiments.""" + + pass + + +ExperimentConfigOrDict = Union[ExperimentConfig, ExperimentConfigDict] + + +class ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollection( + _common.BaseModel +): + """Performs no automatic garbage collection on Runtime Revisions.""" + + pass + + +class ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollectionDict( + TypedDict, total=False +): + """Performs no automatic garbage collection on Runtime Revisions.""" + + pass + + +ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollectionOrDict = Union[ + ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollection, + ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollectionDict, +] + + +class ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatest(_common.BaseModel): + """Keeps only the latest N Runtime Revisions active.""" + + max_active_revisions: Optional[int] = Field( + default=None, + description="""Required. Specifies the maximum number of Runtime Revisions to keep active. If an update to Reasoning Engine would result in exceeding this number of active Runtime Revisions, a new Runtime Revision will be created, while the oldest Runtime Revision will be automatically deprecated, providing it's not configured to serve traffic via `traffic_config`. If the oldest Runtime Revision is configured to serve traffic, the update will fail validation. No changes will be made to the Reasoning Engine, existing Runtime Revisions, and no new Runtime Revision will be created.""", + ) + + +class ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatestDict( + TypedDict, total=False +): + """Keeps only the latest N Runtime Revisions active.""" + + max_active_revisions: Optional[int] + """Required. Specifies the maximum number of Runtime Revisions to keep active. If an update to Reasoning Engine would result in exceeding this number of active Runtime Revisions, a new Runtime Revision will be created, while the oldest Runtime Revision will be automatically deprecated, providing it's not configured to serve traffic via `traffic_config`. If the oldest Runtime Revision is configured to serve traffic, the update will fail validation. No changes will be made to the Reasoning Engine, existing Runtime Revisions, and no new Runtime Revision will be created.""" + + +ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatestOrDict = Union[ + ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatest, + ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatestDict, +] + + +class ReasoningEngineRevisionGarbageCollectionStrategy(_common.BaseModel): + """Configures garbage collection of Runtime Revisions.""" + + no_garbage_collection: Optional[ + ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollection + ] = Field( + default=None, + description="""Optional. Performs no automatic garbage collection on Runtime Revisions.""", + ) + keep_n_latest: Optional[ + ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatest + ] = Field( + default=None, + description="""Optional. Keeps only the latest N Runtime Revisions active.""", + ) + + +class ReasoningEngineRevisionGarbageCollectionStrategyDict(TypedDict, total=False): + """Configures garbage collection of Runtime Revisions.""" + + no_garbage_collection: Optional[ + ReasoningEngineRevisionGarbageCollectionStrategyNoGarbageCollectionDict + ] + """Optional. Performs no automatic garbage collection on Runtime Revisions.""" + + keep_n_latest: Optional[ + ReasoningEngineRevisionGarbageCollectionStrategyKeepNLatestDict + ] + """Optional. Keeps only the latest N Runtime Revisions active.""" + + +ReasoningEngineRevisionGarbageCollectionStrategyOrDict = Union[ + ReasoningEngineRevisionGarbageCollectionStrategy, + ReasoningEngineRevisionGarbageCollectionStrategyDict, +] + + +class ReasoningEngine(_common.BaseModel): + """An agent runtime.""" + + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""Customer-managed encryption key spec for a ReasoningEngine. If set, this ReasoningEngine and all sub-resources of this ReasoningEngine will be secured by this key.""", + ) + context_spec: Optional[ReasoningEngineContextSpec] = Field( + default=None, + description="""Optional. Configuration for how Agent Engine sub-resources should manage context.""", + ) + create_time: Optional[datetime.datetime] = Field( default=None, description="""Output only. Timestamp when this ReasoningEngine was created.""", ) @@ -8634,10 +8521,23 @@ class ReasoningEngine(_common.BaseModel): default=None, description="""Optional. Traffic distribution configuration for the Reasoning Engine.""", ) + experiment_config: Optional[ExperimentConfig] = Field( + default=None, + description="""Optional. The experiment config used to control which features to enable in this API call.""", + ) + revision_garbage_collection_strategy: Optional[ + ReasoningEngineRevisionGarbageCollectionStrategy + ] = Field( + default=None, + description="""Optional. Configures garbage collection of Runtime Revisions.""", + ) + url: Optional[str] = Field( + default=None, description="""Output only. The URL of the reasoning engine.""" + ) class ReasoningEngineDict(TypedDict, total=False): - """An agent engine.""" + """An agent runtime.""" 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.""" @@ -8672,12 +8572,23 @@ class ReasoningEngineDict(TypedDict, total=False): traffic_config: Optional[ReasoningEngineTrafficConfigDict] """Optional. Traffic distribution configuration for the Reasoning Engine.""" + experiment_config: Optional[ExperimentConfigDict] + """Optional. The experiment config used to control which features to enable in this API call.""" + + revision_garbage_collection_strategy: Optional[ + ReasoningEngineRevisionGarbageCollectionStrategyDict + ] + """Optional. Configures garbage collection of Runtime Revisions.""" + + url: Optional[str] + """Output only. The URL of the reasoning engine.""" + ReasoningEngineOrDict = Union[ReasoningEngine, ReasoningEngineDict] -class AgentEngineOperation(_common.BaseModel): - """Operation that has an agent engine as a response.""" +class RuntimeOperation(_common.BaseModel): + """Operation that has an agent runtime as a response.""" name: Optional[str] = Field( default=None, @@ -8696,12 +8607,12 @@ class AgentEngineOperation(_common.BaseModel): description="""The error result of the operation in case of failure or cancellation.""", ) response: Optional[ReasoningEngine] = Field( - default=None, description="""The created Agent Engine.""" + default=None, description="""The created Agent Runtime.""" ) -class AgentEngineOperationDict(TypedDict, total=False): - """Operation that has an agent engine as a response.""" +class RuntimeOperationDict(TypedDict, total=False): + """Operation that has an agent runtime 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}`.""" @@ -8716,57 +8627,63 @@ class AgentEngineOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" response: Optional[ReasoningEngineDict] - """The created Agent Engine.""" + """The created Agent Runtime.""" -AgentEngineOperationOrDict = Union[AgentEngineOperation, AgentEngineOperationDict] +RuntimeOperationOrDict = Union[RuntimeOperation, RuntimeOperationDict] -class CreateAgentEngineConfig(_common.BaseModel): - """Config for create agent engine.""" +class CreateRuntimeConfig(_common.BaseModel): + """Config for create agent runtime.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) display_name: Optional[str] = Field( default=None, - description="""The user-defined name of the Agent Engine. + description="""The user-defined name of the Agent Runtime. The display name can be up to 128 characters long and can comprise any UTF-8 characters. """, ) description: Optional[str] = Field( - default=None, description="""The description of the Agent Engine.""" + default=None, description="""The description of the Agent Runtime.""" ) spec: Optional[ReasoningEngineSpec] = Field( - default=None, description="""Optional. Configurations of the Agent Engine.""" + default=None, description="""Optional. Configurations of the Agent Runtime.""" ) context_spec: Optional[ReasoningEngineContextSpec] = Field( default=None, - description="""Optional. The context spec to be used for the Agent Engine.""", + description="""Optional. The context spec to be used for the Agent Runtime.""", ) psc_interface_config: Optional[PscInterfaceConfig] = Field( default=None, description="""Optional. The PSC interface config for PSC-I to be used for the - Agent Engine.""", + Agent Runtime.""", + ) + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + ] = Field( + default=None, + description="""Agent Gateway configuration for a Reasoning Engine deployment.""", ) min_instances: Optional[int] = Field( default=None, - description="""The minimum number of instances to run for the Agent Engine. + description="""The minimum number of instances to run for the Agent Runtime. Defaults to 1. Range: [0, 10]. """, ) max_instances: Optional[int] = Field( default=None, - description="""The maximum number of instances to run for the Agent Engine. + description="""The maximum number of instances to run for the Agent Runtime. Defaults to 100. Range: [1, 1000]. If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. """, ) resource_limits: Optional[dict[str, str]] = Field( default=None, - description="""The resource limits to be applied to the Agent Engine. + description="""The resource limits to be applied to the Agent Runtime. Required keys: 'cpu' and 'memory'. Supported values for 'cpu': '1', '2', '4', '6', '8'. Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. @@ -8774,20 +8691,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 @@ -8798,8 +8721,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 @@ -8821,19 +8744,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. """, ) @@ -8841,7 +8764,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". @@ -8852,14 +8775,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 @@ -8868,77 +8791,75 @@ class CreateAgentEngineConfig(_common.BaseModel): subdirectory and the path must be added to `extra_packages`. """, ) - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig - ] = Field( - default=None, - description="""Agent Gateway configuration for a Reasoning Engine deployment.""", - ) - keep_alive_probe: Optional[KeepAliveProbe] = Field( - default=None, - description="""Optional. Specifies the configuration for keep-alive probe. - Contains configuration on a specified endpoint that a deployment host - should use to keep the container alive based on the probe settings.""", - ) -class CreateAgentEngineConfigDict(TypedDict, total=False): - """Config for create agent engine.""" +class CreateRuntimeConfigDict(TypedDict, total=False): + """Config for create 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 @@ -8948,8 +8869,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 @@ -8969,24 +8890,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". @@ -8996,13 +8917,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 @@ -9011,93 +8932,79 @@ 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.""" - - 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.""" - -CreateAgentEngineConfigOrDict = Union[ - CreateAgentEngineConfig, CreateAgentEngineConfigDict -] +CreateRuntimeConfigOrDict = Union[CreateRuntimeConfig, CreateRuntimeConfigDict] -class _CreateAgentEngineRequestParameters(_common.BaseModel): - """Parameters for creating agent engines.""" +class _CreateRuntimeRequestParameters(_common.BaseModel): + """Parameters for creating agent runtimes.""" - config: Optional[CreateAgentEngineConfig] = Field(default=None, description="""""") + config: Optional[CreateRuntimeConfig] = Field(default=None, description="""""") -class _CreateAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for creating agent engines.""" +class _CreateRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for creating agent runtimes.""" - config: Optional[CreateAgentEngineConfigDict] + config: Optional[CreateRuntimeConfigDict] """""" -_CreateAgentEngineRequestParametersOrDict = Union[ - _CreateAgentEngineRequestParameters, _CreateAgentEngineRequestParametersDict +_CreateRuntimeRequestParametersOrDict = Union[ + _CreateRuntimeRequestParameters, _CreateRuntimeRequestParametersDict ] -class DeleteAgentEngineConfig(_common.BaseModel): - """Config for deleting agent engine.""" +class DeleteRuntimeConfig(_common.BaseModel): + """Config for deleting agent runtime.""" 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 DeleteRuntimeConfigDict(TypedDict, total=False): + """Config for deleting agent runtime.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteAgentEngineConfigOrDict = Union[ - DeleteAgentEngineConfig, DeleteAgentEngineConfigDict -] +DeleteRuntimeConfigOrDict = Union[DeleteRuntimeConfig, DeleteRuntimeConfigDict] -class _DeleteAgentEngineRequestParameters(_common.BaseModel): - """Parameters for deleting agent engines.""" +class _DeleteRuntimeRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtimes.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + 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[DeleteAgentEngineConfig] = Field(default=None, description="""""") + config: Optional[DeleteRuntimeConfig] = Field(default=None, description="""""") -class _DeleteAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engines.""" +class _DeleteRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtimes.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime.""" force: Optional[bool] """If set to true, any child resources will also be deleted.""" - config: Optional[DeleteAgentEngineConfigDict] + config: Optional[DeleteRuntimeConfigDict] """""" -_DeleteAgentEngineRequestParametersOrDict = Union[ - _DeleteAgentEngineRequestParameters, _DeleteAgentEngineRequestParametersDict +_DeleteRuntimeRequestParametersOrDict = Union[ + _DeleteRuntimeRequestParameters, _DeleteRuntimeRequestParametersDict ] -class DeleteAgentEngineOperation(_common.BaseModel): - """Operation for deleting agent engines.""" +class DeleteRuntimeOperation(_common.BaseModel): + """Operation for deleting agent runtimes.""" name: Optional[str] = Field( default=None, @@ -9117,8 +9024,8 @@ class DeleteAgentEngineOperation(_common.BaseModel): ) -class DeleteAgentEngineOperationDict(TypedDict, total=False): - """Operation for deleting agent engines.""" +class DeleteRuntimeOperationDict(TypedDict, total=False): + """Operation for deleting agent runtimes.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -9133,55 +9040,53 @@ class DeleteAgentEngineOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -DeleteAgentEngineOperationOrDict = Union[ - DeleteAgentEngineOperation, DeleteAgentEngineOperationDict -] +DeleteRuntimeOperationOrDict = Union[DeleteRuntimeOperation, DeleteRuntimeOperationDict] -class GetAgentEngineConfig(_common.BaseModel): - """Config for create agent engine.""" +class GetRuntimeConfig(_common.BaseModel): + """Config for create agent runtime.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetAgentEngineConfigDict(TypedDict, total=False): - """Config for create agent engine.""" +class GetRuntimeConfigDict(TypedDict, total=False): + """Config for create agent runtime.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetAgentEngineConfigOrDict = Union[GetAgentEngineConfig, GetAgentEngineConfigDict] +GetRuntimeConfigOrDict = Union[GetRuntimeConfig, GetRuntimeConfigDict] -class _GetAgentEngineRequestParameters(_common.BaseModel): - """Parameters for getting agent engines.""" +class _GetRuntimeRequestParameters(_common.BaseModel): + """Parameters for getting agent runtimes.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime.""" ) - config: Optional[GetAgentEngineConfig] = Field(default=None, description="""""") + config: Optional[GetRuntimeConfig] = Field(default=None, description="""""") -class _GetAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for getting agent engines.""" +class _GetRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for getting agent runtimes.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime.""" - config: Optional[GetAgentEngineConfigDict] + config: Optional[GetRuntimeConfigDict] """""" -_GetAgentEngineRequestParametersOrDict = Union[ - _GetAgentEngineRequestParameters, _GetAgentEngineRequestParametersDict +_GetRuntimeRequestParametersOrDict = Union[ + _GetRuntimeRequestParameters, _GetRuntimeRequestParametersDict ] -class ListAgentEngineConfig(_common.BaseModel): - """Config for listing agent engines.""" +class ListRuntimeConfig(_common.BaseModel): + """Config for listing agent runtimes.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -9195,8 +9100,8 @@ class ListAgentEngineConfig(_common.BaseModel): ) -class ListAgentEngineConfigDict(TypedDict, total=False): - """Config for listing agent engines.""" +class ListRuntimeConfigDict(TypedDict, total=False): + """Config for listing agent runtimes.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -9212,29 +9117,29 @@ class ListAgentEngineConfigDict(TypedDict, total=False): For field names both snake_case and camelCase are supported.""" -ListAgentEngineConfigOrDict = Union[ListAgentEngineConfig, ListAgentEngineConfigDict] +ListRuntimeConfigOrDict = Union[ListRuntimeConfig, ListRuntimeConfigDict] -class _ListAgentEngineRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" +class _ListRuntimeRequestParameters(_common.BaseModel): + """Parameters for listing agent runtimes.""" - config: Optional[ListAgentEngineConfig] = Field(default=None, description="""""") + config: Optional[ListRuntimeConfig] = Field(default=None, description="""""") -class _ListAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" +class _ListRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtimes.""" - config: Optional[ListAgentEngineConfigDict] + config: Optional[ListRuntimeConfigDict] """""" -_ListAgentEngineRequestParametersOrDict = Union[ - _ListAgentEngineRequestParameters, _ListAgentEngineRequestParametersDict +_ListRuntimeRequestParametersOrDict = Union[ + _ListRuntimeRequestParameters, _ListRuntimeRequestParametersDict ] class ListReasoningEnginesResponse(_common.BaseModel): - """Response for listing agent engines.""" + """Response for listing agent runtimes.""" sdk_http_response: Optional[genai_types.HttpResponse] = Field( default=None, description="""Used to retain the full HTTP response.""" @@ -9242,13 +9147,13 @@ class ListReasoningEnginesResponse(_common.BaseModel): next_page_token: Optional[str] = Field(default=None, description="""""") reasoning_engines: Optional[list[ReasoningEngine]] = Field( default=None, - description="""List of agent engines. + description="""List of agent runtimes. """, ) class ListReasoningEnginesResponseDict(TypedDict, total=False): - """Response for listing agent engines.""" + """Response for listing agent runtimes.""" sdk_http_response: Optional[genai_types.HttpResponse] """Used to retain the full HTTP response.""" @@ -9257,7 +9162,7 @@ class ListReasoningEnginesResponseDict(TypedDict, total=False): """""" reasoning_engines: Optional[list[ReasoningEngineDict]] - """List of agent engines. + """List of agent runtimes. """ @@ -9266,52 +9171,52 @@ class ListReasoningEnginesResponseDict(TypedDict, total=False): ] -class GetAgentEngineOperationConfig(_common.BaseModel): +class GetRuntimeOperationConfig(_common.BaseModel): http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetAgentEngineOperationConfigDict(TypedDict, total=False): +class GetRuntimeOperationConfigDict(TypedDict, total=False): http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetAgentEngineOperationConfigOrDict = Union[ - GetAgentEngineOperationConfig, GetAgentEngineOperationConfigDict +GetRuntimeOperationConfigOrDict = Union[ + GetRuntimeOperationConfig, GetRuntimeOperationConfigDict ] -class _GetAgentEngineOperationParameters(_common.BaseModel): - """Parameters for getting an operation with an agent engine as a response.""" +class _GetRuntimeOperationParameters(_common.BaseModel): + """Parameters for getting an operation with an agent runtime as a response.""" operation_name: Optional[str] = Field( default=None, description="""The server-assigned name for the operation.""" ) - config: Optional[GetAgentEngineOperationConfig] = Field( + config: Optional[GetRuntimeOperationConfig] = 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 _GetRuntimeOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with an agent runtime as a response.""" operation_name: Optional[str] """The server-assigned name for the operation.""" - config: Optional[GetAgentEngineOperationConfigDict] + config: Optional[GetRuntimeOperationConfigDict] """Used to override the default configuration.""" -_GetAgentEngineOperationParametersOrDict = Union[ - _GetAgentEngineOperationParameters, _GetAgentEngineOperationParametersDict +_GetRuntimeOperationParametersOrDict = Union[ + _GetRuntimeOperationParameters, _GetRuntimeOperationParametersDict ] -class QueryAgentEngineConfig(_common.BaseModel): - """Config for querying agent engines.""" +class QueryRuntimeConfig(_common.BaseModel): + """Config for querying agent runtimes.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -9325,8 +9230,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 QueryRuntimeConfigDict(TypedDict, total=False): + """Config for querying agent runtimes.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -9341,35 +9246,35 @@ class QueryAgentEngineConfigDict(TypedDict, total=False): """""" -QueryAgentEngineConfigOrDict = Union[QueryAgentEngineConfig, QueryAgentEngineConfigDict] +QueryRuntimeConfigOrDict = Union[QueryRuntimeConfig, QueryRuntimeConfigDict] -class _QueryAgentEngineRequestParameters(_common.BaseModel): - """Parameters for querying agent engines.""" +class _QueryRuntimeRequestParameters(_common.BaseModel): + """Parameters for querying agent runtimes.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime.""" ) - config: Optional[QueryAgentEngineConfig] = Field(default=None, description="""""") + config: Optional[QueryRuntimeConfig] = Field(default=None, description="""""") -class _QueryAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for querying agent engines.""" +class _QueryRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for querying agent runtimes.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime.""" - config: Optional[QueryAgentEngineConfigDict] + config: Optional[QueryRuntimeConfigDict] """""" -_QueryAgentEngineRequestParametersOrDict = Union[ - _QueryAgentEngineRequestParameters, _QueryAgentEngineRequestParametersDict +_QueryRuntimeRequestParametersOrDict = Union[ + _QueryRuntimeRequestParameters, _QueryRuntimeRequestParametersDict ] class QueryReasoningEngineResponse(_common.BaseModel): - """The response for querying an agent engine.""" + """The response for querying an agent runtime.""" output: Optional[Any] = Field( default=None, @@ -9378,7 +9283,7 @@ class QueryReasoningEngineResponse(_common.BaseModel): class QueryReasoningEngineResponseDict(TypedDict, total=False): - """The response for querying an agent engine.""" + """The response for querying an agent runtime.""" output: Optional[Any] """Response provided by users in JSON object format.""" @@ -9389,51 +9294,57 @@ class QueryReasoningEngineResponseDict(TypedDict, total=False): ] -class UpdateAgentEngineConfig(_common.BaseModel): - """Config for updating agent engine.""" +class UpdateRuntimeConfig(_common.BaseModel): + """Config for updating agent runtime.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) display_name: Optional[str] = Field( default=None, - description="""The user-defined name of the Agent Engine. + description="""The user-defined name of the Agent Runtime. The display name can be up to 128 characters long and can comprise any UTF-8 characters. """, ) description: Optional[str] = Field( - default=None, description="""The description of the Agent Engine.""" + default=None, description="""The description of the Agent Runtime.""" ) spec: Optional[ReasoningEngineSpec] = Field( - default=None, description="""Optional. Configurations of the Agent Engine.""" + default=None, description="""Optional. Configurations of the Agent Runtime.""" ) context_spec: Optional[ReasoningEngineContextSpec] = Field( default=None, - description="""Optional. The context spec to be used for the Agent Engine.""", + description="""Optional. The context spec to be used for the Agent Runtime.""", ) psc_interface_config: Optional[PscInterfaceConfig] = Field( default=None, description="""Optional. The PSC interface config for PSC-I to be used for the - Agent Engine.""", + Agent Runtime.""", + ) + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfig + ] = Field( + default=None, + description="""Agent Gateway configuration for a Reasoning Engine deployment.""", ) min_instances: Optional[int] = Field( default=None, - description="""The minimum number of instances to run for the Agent Engine. + description="""The minimum number of instances to run for the Agent Runtime. Defaults to 1. Range: [0, 10]. """, ) max_instances: Optional[int] = Field( default=None, - description="""The maximum number of instances to run for the Agent Engine. + description="""The maximum number of instances to run for the Agent Runtime. Defaults to 100. Range: [1, 1000]. If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. """, ) resource_limits: Optional[dict[str, str]] = Field( default=None, - description="""The resource limits to be applied to the Agent Engine. + description="""The resource limits to be applied to the Agent Runtime. Required keys: 'cpu' and 'memory'. Supported values for 'cpu': '1', '2', '4', '6', '8'. Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. @@ -9441,20 +9352,26 @@ class UpdateAgentEngineConfig(_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 @@ -9465,8 +9382,8 @@ class UpdateAgentEngineConfig(_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 @@ -9488,19 +9405,19 @@ class UpdateAgentEngineConfig(_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. """, ) @@ -9508,7 +9425,7 @@ class UpdateAgentEngineConfig(_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". @@ -9519,14 +9436,14 @@ class UpdateAgentEngineConfig(_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 @@ -9535,18 +9452,6 @@ class UpdateAgentEngineConfig(_common.BaseModel): subdirectory and the path must be added to `extra_packages`. """, ) - agent_gateway_config: Optional[ - ReasoningEngineSpecDeploymentSpecAgentGatewayConfig - ] = Field( - default=None, - description="""Agent Gateway configuration for a Reasoning Engine deployment.""", - ) - keep_alive_probe: Optional[KeepAliveProbe] = Field( - default=None, - description="""Optional. Specifies the configuration for keep-alive probe. - Contains configuration on a specified endpoint that a deployment host - should use to keep the container alive based on the probe settings.""", - ) update_mask: Optional[str] = Field( default=None, description="""The update mask to apply. For the `FieldMask` definition, see @@ -9558,63 +9463,73 @@ class UpdateAgentEngineConfig(_common.BaseModel): ) -class UpdateAgentEngineConfigDict(TypedDict, total=False): - """Config for updating 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 @@ -9624,8 +9539,8 @@ class UpdateAgentEngineConfigDict(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 @@ -9645,24 +9560,24 @@ class UpdateAgentEngineConfigDict(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". @@ -9672,13 +9587,13 @@ class UpdateAgentEngineConfigDict(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 @@ -9687,16 +9602,6 @@ class UpdateAgentEngineConfigDict(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.""" - - 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.""" - update_mask: Optional[str] """The update mask to apply. For the `FieldMask` definition, see https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" @@ -9705,382 +9610,270 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False): """Traffic distribution configuration for the Reasoning Engine.""" -UpdateAgentEngineConfigOrDict = Union[ - UpdateAgentEngineConfig, UpdateAgentEngineConfigDict -] +UpdateRuntimeConfigOrDict = Union[UpdateRuntimeConfig, UpdateRuntimeConfigDict] -class _UpdateAgentEngineRequestParameters(_common.BaseModel): - """Parameters for updating agent engines.""" +class _UpdateRuntimeRequestParameters(_common.BaseModel): + """Parameters for updating agent runtimes.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime.""" ) - config: Optional[UpdateAgentEngineConfig] = Field(default=None, description="""""") + config: Optional[UpdateRuntimeConfig] = Field(default=None, description="""""") -class _UpdateAgentEngineRequestParametersDict(TypedDict, total=False): - """Parameters for updating agent engines.""" +class _UpdateRuntimeRequestParametersDict(TypedDict, total=False): + """Parameters for updating agent runtimes.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime.""" - config: Optional[UpdateAgentEngineConfigDict] + config: Optional[UpdateRuntimeConfigDict] """""" -_UpdateAgentEngineRequestParametersOrDict = Union[ - _UpdateAgentEngineRequestParameters, _UpdateAgentEngineRequestParametersDict +_UpdateRuntimeRequestParametersOrDict = Union[ + _UpdateRuntimeRequestParameters, _UpdateRuntimeRequestParametersDict ] -class MemoryMetadataValue(_common.BaseModel): - """The metadata values for memories.""" +class GetRuntimeRevisionConfig(_common.BaseModel): + """Config for getting an Agent Runtime Runtime Revision.""" - 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) -class MemoryMetadataValueDict(TypedDict, total=False): - """The metadata values for memories.""" +class GetRuntimeRevisionConfigDict(TypedDict, total=False): + """Config for getting an Agent Runtime Runtime Revision.""" - bool_value: Optional[bool] - """Represents a boolean value.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - double_value: Optional[float] - """Represents a double value.""" - string_value: Optional[str] - """Represents a string value.""" +GetRuntimeRevisionConfigOrDict = Union[ + GetRuntimeRevisionConfig, GetRuntimeRevisionConfigDict +] - timestamp_value: Optional[datetime.datetime] - """Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""" +class _GetRuntimeRevisionRequestParameters(_common.BaseModel): + """Parameters for getting an agent runtime runtime revision.""" -MemoryMetadataValueOrDict = Union[MemoryMetadataValue, MemoryMetadataValueDict] + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime runtime revision.""" + ) + config: Optional[GetRuntimeRevisionConfig] = Field(default=None, description="""""") -class AgentEngineMemoryConfig(_common.BaseModel): - """Config for creating a Memory.""" +class _GetRuntimeRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for getting an agent runtime runtime 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( - default=None, - description="""Optional. Input only. The TTL for this resource. + name: Optional[str] + """Name of the agent runtime runtime revision.""" - The expiration time is computed: now + TTL.""", - ) - expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", - ) - revision_expire_time: Optional[datetime.datetime] = Field( + config: Optional[GetRuntimeRevisionConfigDict] + """""" + + +_GetRuntimeRevisionRequestParametersOrDict = Union[ + _GetRuntimeRevisionRequestParameters, _GetRuntimeRevisionRequestParametersDict +] + + +class ReasoningEngineRuntimeRevision(_common.BaseModel): + """A runtime revision.""" + + create_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. Timestamp when this ReasoningEngineRuntimeRevision was created.""", ) - revision_ttl: Optional[str] = Field( + name: Optional[str] = Field( default=None, - description="""Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""", + description="""Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""", ) - disable_memory_revisions: Optional[bool] = Field( + spec: Optional[ReasoningEngineSpec] = 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="""Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""", ) - 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.""", + state: Optional[State] = Field( + default=None, description="""Output only. The state of the revision.""" ) -class AgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for creating a Memory.""" +class ReasoningEngineRuntimeRevisionDict(TypedDict, total=False): + """A runtime revision.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this ReasoningEngineRuntimeRevision was created.""" - display_name: Optional[str] - """The display name of the memory.""" + name: Optional[str] + """Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""" - description: Optional[str] - """The description of the memory.""" + spec: Optional[ReasoningEngineSpecDict] + """Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + state: Optional[State] + """Output only. The state of the revision.""" - ttl: Optional[str] - """Optional. Input only. The TTL for this resource. - The expiration time is computed: now + TTL.""" +ReasoningEngineRuntimeRevisionOrDict = Union[ + ReasoningEngineRuntimeRevision, ReasoningEngineRuntimeRevisionDict +] - expire_time: Optional[datetime.datetime] - """Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""" - revision_expire_time: Optional[datetime.datetime] - """Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" +class ListRuntimeRevisionsConfig(_common.BaseModel): + """Config for listing agent runtime revisions.""" - revision_ttl: Optional[str] - """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" + 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.""", + ) - 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 ListRuntimeRevisionsConfigDict(TypedDict, total=False): + """Config for listing agent runtime 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] + """Used to override HTTP request options.""" - 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.""" + 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.""" -AgentEngineMemoryConfigOrDict = Union[ - AgentEngineMemoryConfig, AgentEngineMemoryConfigDict +ListRuntimeRevisionsConfigOrDict = Union[ + ListRuntimeRevisionsConfig, ListRuntimeRevisionsConfigDict ] -class _CreateAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for creating Agent Engine Memories.""" +class _ListRuntimeRevisionsRequestParameters(_common.BaseModel): + """Parameters for listing agent runtime revisions.""" name: Optional[str] = Field( - default=None, - description="""Name of the agent engine to create the memory under.""", - ) - fact: Optional[str] = Field( - default=None, - description="""The fact of the memory. - - This is the semantic knowledge extracted from the source content).""", + default=None, description="""Name of the reasoning engine.""" ) - 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[ListRuntimeRevisionsConfig] = Field( + default=None, description="""""" ) - config: Optional[AgentEngineMemoryConfig] = Field(default=None, description="""""") -class _CreateAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for creating Agent Engine Memories.""" +class _ListRuntimeRevisionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtime revisions.""" name: Optional[str] - """Name of the agent engine to create the memory under.""" - - fact: Optional[str] - """The fact of the memory. - - This is the semantic knowledge extracted from the source content).""" - - scope: Optional[dict[str, str]] - """The scope of the memory. - - Memories are isolated within their scope. The scope is defined when - creating or generating memories. Up to 5 key-value pairs are accepted, - and scope values cannot contain the wildcard character '*'.""" + """Name of the reasoning engine.""" - config: Optional[AgentEngineMemoryConfigDict] + config: Optional[ListRuntimeRevisionsConfigDict] """""" -_CreateAgentEngineMemoryRequestParametersOrDict = Union[ - _CreateAgentEngineMemoryRequestParameters, - _CreateAgentEngineMemoryRequestParametersDict, +_ListRuntimeRevisionsRequestParametersOrDict = Union[ + _ListRuntimeRevisionsRequestParameters, _ListRuntimeRevisionsRequestParametersDict ] -class MemoryStructuredContent(_common.BaseModel): - """Represents the structured value of the memory.""" +class ListReasoningEnginesRuntimeRevisionsResponse(_common.BaseModel): + """Response for listing agent runtime runtime revisions.""" - data: Optional[dict[str, Any]] = Field( - default=None, - description="""Required. Represents the structured value of the memory.""", - ) - schema_id: Optional[str] = Field( - default=None, - description="""Required. Represents the schema ID for which this structured memory belongs to.""", + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" ) + next_page_token: Optional[str] = Field(default=None, description="""""") + reasoning_engine_runtime_revisions: Optional[ + list[ReasoningEngineRuntimeRevision] + ] = Field(default=None, description="""List of agent runtime revisions.""") -class MemoryStructuredContentDict(TypedDict, total=False): - """Represents the structured value of the memory.""" +class ListReasoningEnginesRuntimeRevisionsResponseDict(TypedDict, total=False): + """Response for listing agent runtime runtime revisions.""" - data: Optional[dict[str, Any]] - """Required. Represents the structured value of the memory.""" + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" - schema_id: Optional[str] - """Required. Represents the schema ID for which this structured memory belongs to.""" + next_page_token: Optional[str] + """""" + reasoning_engine_runtime_revisions: Optional[ + list[ReasoningEngineRuntimeRevisionDict] + ] + """List of agent runtime revisions.""" -MemoryStructuredContentOrDict = Union[ - MemoryStructuredContent, MemoryStructuredContentDict + +ListReasoningEnginesRuntimeRevisionsResponseOrDict = Union[ + ListReasoningEnginesRuntimeRevisionsResponse, + ListReasoningEnginesRuntimeRevisionsResponseDict, ] -class Memory(_common.BaseModel): - """A memory.""" +class DeleteRuntimeRevisionConfig(_common.BaseModel): + """Config for deleting an Agent Runtime Runtime Revision.""" - create_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Represents the timestamp when this Memory was created.""", - ) - description: Optional[str] = Field( - default=None, - description="""Optional. Represents the description of the Memory.""", - ) - disable_memory_revisions: Optional[bool] = Field( - default=None, - description="""Optional. Input only. Indicates whether no revision will be created for this request.""", - ) - display_name: Optional[str] = Field( - default=None, - description="""Optional. Represents the display name of the Memory.""", - ) - expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. Represents the timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""", - ) - fact: Optional[str] = Field( - default=None, - description="""Optional. Represents semantic knowledge extracted from the source content.""", - ) - metadata: Optional[dict[str, MemoryMetadataValue]] = Field( - default=None, - description="""Optional. Represents user-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""", - ) - name: Optional[str] = Field( - default=None, - description="""Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""", - ) - revision_expire_time: Optional[datetime.datetime] = Field( - default=None, - description="""Optional. Input only. Represents the timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", - ) - 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. 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.""" - ) - ttl: Optional[str] = Field( - default=None, - description="""Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""", - ) - update_time: Optional[datetime.datetime] = Field( - default=None, - description="""Output only. Represents the timestamp when this Memory was most recently updated.""", - ) - memory_type: Optional[MemoryType] = Field( - default=None, - description="""Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - structured_content: Optional[MemoryStructuredContent] = Field( - default=None, - description="""Optional. Represents the structured content of the memory.""", + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Waits for the operation to complete before returning.""", ) -class MemoryDict(TypedDict, total=False): - """A memory.""" - - create_time: Optional[datetime.datetime] - """Output only. Represents the timestamp when this Memory was created.""" - - description: Optional[str] - """Optional. Represents the description of the Memory.""" - - disable_memory_revisions: Optional[bool] - """Optional. Input only. Indicates whether no revision will be created for this request.""" - - display_name: Optional[str] - """Optional. Represents the display name of the Memory.""" - - expire_time: Optional[datetime.datetime] - """Optional. Represents the timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""" - - fact: Optional[str] - """Optional. Represents semantic knowledge extracted from the source content.""" +class DeleteRuntimeRevisionConfigDict(TypedDict, total=False): + """Config for deleting an Agent Runtime Runtime Revision.""" - 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.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - name: Optional[str] - """Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""" + wait_for_completion: Optional[bool] + """Waits for the operation to complete before returning.""" - 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.""" +DeleteRuntimeRevisionConfigOrDict = Union[ + DeleteRuntimeRevisionConfig, DeleteRuntimeRevisionConfigDict +] - revision_ttl: Optional[str] - """Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""" - scope: Optional[dict[str, str]] - """Required. Immutable. Represents the scope of the Memory. Memories are isolated within their scope. The scope is defined when creating or generating memories. Scope values cannot contain the wildcard character '*'.""" +class _DeleteRuntimeRevisionRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtime runtime revisions.""" - topics: Optional[list[MemoryTopicIdDict]] - """Optional. Represents the Topics of the Memory.""" + name: Optional[str] = Field( + default=None, + description="""Name of the agent runtime runtime revision to delete.""", + ) + config: Optional[DeleteRuntimeRevisionConfig] = Field( + default=None, description="""""" + ) - ttl: Optional[str] - """Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""" - update_time: Optional[datetime.datetime] - """Output only. Represents the timestamp when this Memory was most recently updated.""" +class _DeleteRuntimeRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtime runtime revisions.""" - 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.""" + name: Optional[str] + """Name of the agent runtime runtime revision to delete.""" - structured_content: Optional[MemoryStructuredContentDict] - """Optional. Represents the structured content of the memory.""" + config: Optional[DeleteRuntimeRevisionConfigDict] + """""" -MemoryOrDict = Union[Memory, MemoryDict] +_DeleteRuntimeRevisionRequestParametersOrDict = Union[ + _DeleteRuntimeRevisionRequestParameters, _DeleteRuntimeRevisionRequestParametersDict +] -class AgentEngineMemoryOperation(_common.BaseModel): - """Operation that has an agent engine memory as a response.""" +class DeleteRuntimeRevisionOperation(_common.BaseModel): + """Operation for deleting agent runtime revisions.""" name: Optional[str] = Field( default=None, @@ -10098,13 +9891,10 @@ class AgentEngineMemoryOperation(_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 Agent Engine Memory.""" - ) -class AgentEngineMemoryOperationDict(TypedDict, total=False): - """Operation that has an agent engine memory as a response.""" +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}`.""" @@ -10118,261 +9908,181 @@ class AgentEngineMemoryOperationDict(TypedDict, total=False): error: Optional[dict[str, Any]] """The error result of the operation in case of failure or cancellation.""" - response: Optional[MemoryDict] - """The Agent Engine Memory.""" - -AgentEngineMemoryOperationOrDict = Union[ - AgentEngineMemoryOperation, AgentEngineMemoryOperationDict +DeleteRuntimeRevisionOperationOrDict = Union[ + DeleteRuntimeRevisionOperation, DeleteRuntimeRevisionOperationDict ] -class DeleteAgentEngineMemoryConfig(_common.BaseModel): - """Config for deleting an Agent Engine Memory.""" +class GetDeleteRuntimeRevisionOperationConfig(_common.BaseModel): http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class DeleteAgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Memory.""" +class GetDeleteRuntimeRevisionOperationConfigDict(TypedDict, total=False): http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteAgentEngineMemoryConfigOrDict = Union[ - DeleteAgentEngineMemoryConfig, DeleteAgentEngineMemoryConfigDict +GetDeleteRuntimeRevisionOperationConfigOrDict = Union[ + GetDeleteRuntimeRevisionOperationConfig, GetDeleteRuntimeRevisionOperationConfigDict ] -class _DeleteAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for deleting agent engines.""" +class _GetDeleteRuntimeRevisionOperationParameters(_common.BaseModel): + """Parameters for getting an operation that deletes an agent runtime revision.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine memory to delete.""" + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" ) - config: Optional[DeleteAgentEngineMemoryConfig] = Field( - default=None, description="""""" + config: Optional[GetDeleteRuntimeRevisionOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" ) -class _DeleteAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engines.""" +class _GetDeleteRuntimeRevisionOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation that deletes an agent runtime revision.""" - name: Optional[str] - """Name of the agent engine memory to delete.""" + operation_name: Optional[str] + """The server-assigned name for the operation.""" - config: Optional[DeleteAgentEngineMemoryConfigDict] - """""" + config: Optional[GetDeleteRuntimeRevisionOperationConfigDict] + """Used to override the default configuration.""" -_DeleteAgentEngineMemoryRequestParametersOrDict = Union[ - _DeleteAgentEngineMemoryRequestParameters, - _DeleteAgentEngineMemoryRequestParametersDict, +_GetDeleteRuntimeRevisionOperationParametersOrDict = Union[ + _GetDeleteRuntimeRevisionOperationParameters, + _GetDeleteRuntimeRevisionOperationParametersDict, ] -class DeleteAgentEngineMemoryOperation(_common.BaseModel): - """Operation for deleting agent engines.""" +class QueryRuntimeRevisionConfig(_common.BaseModel): + """Config for querying 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.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - 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.""", + class_method: Optional[str] = Field( + default=None, description="""The class method to call.""" ) - error: Optional[dict[str, Any]] = Field( - default=None, - description="""The error result of the operation in case of failure or cancellation.""", + 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 DeleteAgentEngineMemoryOperationDict(TypedDict, total=False): - """Operation for deleting agent engines.""" +class QueryRuntimeRevisionConfigDict(TypedDict, total=False): + """Config for querying 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}`.""" + 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.""" + class_method: Optional[str] + """The class method to call.""" - 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.""" + input: Optional[dict[str, Any]] + """The input to the class method.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + include_all_fields: Optional[bool] + """""" -DeleteAgentEngineMemoryOperationOrDict = Union[ - DeleteAgentEngineMemoryOperation, DeleteAgentEngineMemoryOperationDict +QueryRuntimeRevisionConfigOrDict = Union[ + QueryRuntimeRevisionConfig, QueryRuntimeRevisionConfigDict ] -class GenerateMemoriesRequestVertexSessionSource(_common.BaseModel): - """The vertex session source for generating memories.""" +class _QueryRuntimeRevisionRequestParameters(_common.BaseModel): + """Parameters for querying agent runtime revisions.""" - 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}`""", + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime revision.""" ) - 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.""", + config: Optional[QueryRuntimeRevisionConfig] = Field( + default=None, description="""""" ) -class GenerateMemoriesRequestVertexSessionSourceDict(TypedDict, total=False): - """The vertex session source for generating memories.""" - - end_time: Optional[datetime.datetime] - """Optional. End time (exclusive) of the time range. If not set, the end time is unbounded.""" +class _QueryRuntimeRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for querying agent runtime revisions.""" - session: Optional[str] - """Required. The resource name of the Session to generate memories for. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/sessions/{session}`""" + name: Optional[str] + """Name of the agent runtime revision.""" - start_time: Optional[datetime.datetime] - """Optional. Time range to define which session events should be used to generate memories. Start time (inclusive) of the time range. If not set, the start time is unbounded.""" + config: Optional[QueryRuntimeRevisionConfigDict] + """""" -GenerateMemoriesRequestVertexSessionSourceOrDict = Union[ - GenerateMemoriesRequestVertexSessionSource, - GenerateMemoriesRequestVertexSessionSourceDict, +_QueryRuntimeRevisionRequestParametersOrDict = Union[ + _QueryRuntimeRevisionRequestParameters, _QueryRuntimeRevisionRequestParametersDict ] -class GenerateMemoriesRequestDirectContentsSourceEvent(_common.BaseModel): +class MemoryMetadataValue(_common.BaseModel): + """The metadata values for memories.""" - content: Optional[genai_types.Content] = Field( - default=None, - description="""Required. A single piece of content from which to generate memories.""", + bool_value: Optional[bool] = Field( + default=None, description="""Represents a boolean value.""" ) - - -class GenerateMemoriesRequestDirectContentsSourceEventDict(TypedDict, total=False): - - content: Optional[genai_types.Content] - """Required. A single piece of content from which to generate memories.""" - - -GenerateMemoriesRequestDirectContentsSourceEventOrDict = Union[ - GenerateMemoriesRequestDirectContentsSourceEvent, - GenerateMemoriesRequestDirectContentsSourceEventDict, -] - - -class GenerateMemoriesRequestDirectContentsSource(_common.BaseModel): - """The direct contents source for generating memories.""" - - events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEvent]] = Field( - default=None, - description="""Required. The source content (i.e. chat history) to generate memories from.""", + double_value: Optional[float] = Field( + default=None, description="""Represents a double value.""" ) - - -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.""" - - -GenerateMemoriesRequestDirectContentsSourceOrDict = Union[ - GenerateMemoriesRequestDirectContentsSource, - GenerateMemoriesRequestDirectContentsSourceDict, -] - - -class GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(_common.BaseModel): - """A direct memory to upload to Memory Bank.""" - - fact: Optional[str] = Field( - default=None, - description="""Required. The fact to consolidate with existing memories.""", + string_value: Optional[str] = Field( + default=None, description="""Represents a string value.""" ) - topics: Optional[list[MemoryTopicId]] = Field( + timestamp_value: Optional[datetime.datetime] = Field( default=None, - description="""Optional. The topics that the consolidated memories should be associated with.""", + description="""Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""", ) -class GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict( - TypedDict, total=False -): - """A direct memory to upload to Memory Bank.""" - - fact: Optional[str] - """Required. The fact to consolidate with existing memories.""" - - topics: Optional[list[MemoryTopicIdDict]] - """Optional. The topics that the consolidated memories should be associated with.""" - - -GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryOrDict = Union[ - GenerateMemoriesRequestDirectMemoriesSourceDirectMemory, - GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict, -] - - -class GenerateMemoriesRequestDirectMemoriesSource(_common.BaseModel): - """The direct memories source for generating memories.""" +class MemoryMetadataValueDict(TypedDict, total=False): + """The metadata values for 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.""", - ) + bool_value: Optional[bool] + """Represents a boolean value.""" + double_value: Optional[float] + """Represents a double value.""" -class GenerateMemoriesRequestDirectMemoriesSourceDict(TypedDict, total=False): - """The direct memories source for generating memories.""" + string_value: Optional[str] + """Represents a string value.""" - direct_memories: Optional[ - list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict] - ] - """Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""" + timestamp_value: Optional[datetime.datetime] + """Represents a timestamp value. When filtering on timestamp values, only the seconds field will be compared.""" -GenerateMemoriesRequestDirectMemoriesSourceOrDict = Union[ - GenerateMemoriesRequestDirectMemoriesSource, - GenerateMemoriesRequestDirectMemoriesSourceDict, -] +MemoryMetadataValueOrDict = Union[MemoryMetadataValue, MemoryMetadataValueDict] -class GenerateAgentEngineMemoriesConfig(_common.BaseModel): - """Config for generating memories.""" +class RuntimeMemoryConfig(_common.BaseModel): + """Config for creating a Memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - 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.""", + 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.""", ) - revision_labels: Optional[dict[str, str]] = Field( + ttl: 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. Input only. The TTL for this resource. + + The expiration time is computed: now + TTL.""", + ) + expire_time: Optional[datetime.datetime] = Field( + default=None, + description="""Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""", ) revision_expire_time: Optional[datetime.datetime] = Field( default=None, @@ -10384,41 +10094,43 @@ class GenerateAgentEngineMemoriesConfig(_common.BaseModel): ) disable_memory_revisions: Optional[bool] = Field( default=None, - description="""Optional. Input only. If true, no revisions will be created for this request.""", + description="""Optional. Input only. If true, no revision 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.""", + topics: Optional[list[MemoryTopicId]] = Field( + default=None, description="""Optional. The topics of the memory.""" ) - metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] = Field( + metadata: Optional[dict[str, MemoryMetadataValue]] = Field( default=None, - description="""Optional. The strategy to use when applying metadata to existing memories.""", + 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.""", ) - allowed_topics: Optional[list[MemoryTopicId]] = Field( + memory_id: Optional[str] = Field( default=None, - description="""Optional. Restricts memory generation to a subset of memory topics.""", + 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 GenerateAgentEngineMemoriesConfigDict(TypedDict, total=False): - """Config for generating memories.""" +class RuntimeMemoryConfigDict(TypedDict, total=False): + """Config for creating a Memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - disable_consolidation: Optional[bool] - """Whether to disable consolidation of memories. + display_name: Optional[str] + """The display name of the memory.""" - If true, generated memories will not be consolidated with existing - memories; all generated memories will be added as new memories regardless - of whether they are duplicates of or contradictory to existing memories. - By default, memory consolidation is enabled.""" + description: Optional[str] + """The description of the memory.""" wait_for_completion: Optional[bool] """Waits for the operation to complete before returning.""" - 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.""" + ttl: Optional[str] + """Optional. Input only. The TTL for this resource. + + The expiration time is computed: now + TTL.""" + + expire_time: Optional[datetime.datetime] + """Optional. Timestamp of when this resource is considered expired. This is *always* provided on output, regardless of what `expiration` was sent on input.""" revision_expire_time: Optional[datetime.datetime] """Optional. Input only. Timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" @@ -10427,156 +10139,239 @@ class GenerateAgentEngineMemoriesConfigDict(TypedDict, total=False): """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" disable_memory_revisions: Optional[bool] - """Optional. Input only. If true, no revisions will be created for this request.""" + """Optional. Input only. If true, no revision 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.""" + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The topics of the memory.""" - metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] - """Optional. The strategy to use when applying metadata to existing memories.""" + 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.""" - allowed_topics: Optional[list[MemoryTopicIdDict]] - """Optional. Restricts memory generation to a subset of memory topics.""" + 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.""" -GenerateAgentEngineMemoriesConfigOrDict = Union[ - GenerateAgentEngineMemoriesConfig, GenerateAgentEngineMemoriesConfigDict -] +RuntimeMemoryConfigOrDict = Union[RuntimeMemoryConfig, RuntimeMemoryConfigDict] -class _GenerateAgentEngineMemoriesRequestParameters(_common.BaseModel): - """Parameters for generating agent engine memories.""" +class _CreateRuntimeMemoryRequestParameters(_common.BaseModel): + """Parameters for creating Agent Runtime Memories.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to generate memories for.""", + description="""Name of the agent runtime to create the memory under.""", ) - vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSource] = Field( + fact: Optional[str] = Field( default=None, - description="""The vertex session source of the memories that should be generated.""", - ) - direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSource] = ( - Field( - default=None, - description="""The direct contents source of the memories that should be generated.""", - ) - ) - direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSource] = ( - Field( - default=None, - description="""The direct memories source of the memories that should be generated.""", - ) + 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 memories that should be generated. + description="""The scope of the memory. - 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[GenerateAgentEngineMemoriesConfig] = Field( - default=None, description="""""" + Memories are isolated within their scope. The scope is defined when + creating or generating memories. Up to 5 key-value pairs are accepted, + and scope values cannot contain the wildcard character '*'.""", ) + config: Optional[RuntimeMemoryConfig] = Field(default=None, description="""""") -class _GenerateAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for generating agent engine memories.""" +class _CreateRuntimeMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for creating Agent Runtime Memories.""" name: Optional[str] - """Name of the agent engine to generate memories for.""" + """Name of the agent runtime to create the memory under.""" - 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.""" + fact: Optional[str] + """The fact of the memory. - direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSourceDict] - """The direct memories source of the memories that should be generated.""" + This is the semantic knowledge extracted from the source content).""" scope: Optional[dict[str, str]] - """The scope of the memories that should be generated. + """The scope of the memory. - 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 '*'.""" + 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[GenerateAgentEngineMemoriesConfigDict] + config: Optional[RuntimeMemoryConfigDict] """""" -_GenerateAgentEngineMemoriesRequestParametersOrDict = Union[ - _GenerateAgentEngineMemoriesRequestParameters, - _GenerateAgentEngineMemoriesRequestParametersDict, +_CreateRuntimeMemoryRequestParametersOrDict = Union[ + _CreateRuntimeMemoryRequestParameters, _CreateRuntimeMemoryRequestParametersDict ] -class GenerateMemoriesResponseGeneratedMemory(_common.BaseModel): - """A memory that was generated.""" +class MemoryStructuredContent(_common.BaseModel): + """Represents the structured value of the memory.""" - memory: Optional[Memory] = Field( - default=None, description="""The generated memory.""" - ) - action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] = Field( - default=None, description="""The action to take.""" + data: Optional[dict[str, Any]] = Field( + default=None, + description="""Required. Represents the structured value of the memory.""", ) - previous_revision: Optional[str] = Field( + schema_id: 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}`""", + description="""Required. Represents the schema ID for which this structured memory belongs to.""", ) -class GenerateMemoriesResponseGeneratedMemoryDict(TypedDict, total=False): - """A memory that was generated.""" - - memory: Optional[MemoryDict] - """The generated memory.""" +class MemoryStructuredContentDict(TypedDict, total=False): + """Represents the structured value of the memory.""" - action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] - """The action to take.""" + data: Optional[dict[str, Any]] + """Required. Represents the structured value of the memory.""" - 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}`""" + schema_id: Optional[str] + """Required. Represents the schema ID for which this structured memory belongs to.""" -GenerateMemoriesResponseGeneratedMemoryOrDict = Union[ - GenerateMemoriesResponseGeneratedMemory, GenerateMemoriesResponseGeneratedMemoryDict +MemoryStructuredContentOrDict = Union[ + MemoryStructuredContent, MemoryStructuredContentDict ] -class GenerateMemoriesResponse(_common.BaseModel): - """The response for generating memories.""" +class Memory(_common.BaseModel): + """A memory.""" - generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemory]] = Field( - default=None, description="""The generated memories.""" + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Represents the timestamp when this Memory was created.""", + ) + description: Optional[str] = Field( + default=None, + description="""Optional. Represents the description of the Memory.""", + ) + disable_memory_revisions: Optional[bool] = Field( + default=None, + description="""Optional. Input only. Indicates whether no revision will be created for this request.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""Optional. Represents the display name of the Memory.""", + ) + expire_time: Optional[datetime.datetime] = Field( + default=None, + description="""Optional. Represents the timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""", + ) + fact: Optional[str] = Field( + default=None, + description="""Optional. Represents semantic knowledge extracted from the source content.""", + ) + metadata: Optional[dict[str, MemoryMetadataValue]] = Field( + default=None, + description="""Optional. Represents user-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""", + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""", + ) + revision_expire_time: Optional[datetime.datetime] = Field( + default=None, + description="""Optional. Input only. Represents the timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""", + ) + 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. 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.""" + ) + ttl: Optional[str] = Field( + default=None, + description="""Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Represents the timestamp when this Memory was most recently updated.""", + ) + memory_type: Optional[MemoryType] = Field( + default=None, + description="""Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""", + ) + structured_content: Optional[MemoryStructuredContent] = Field( + default=None, + description="""Optional. Represents the structured content of the memory.""", + ) + structured_data: Optional[dict[str, Any]] = Field( + default=None, + description="""Optional. Deprecated: Use `structured_content` instead.""", ) -class GenerateMemoriesResponseDict(TypedDict, total=False): - """The response for generating memories.""" +class MemoryDict(TypedDict, total=False): + """A memory.""" - generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemoryDict]] - """The generated memories.""" + 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.""" -GenerateMemoriesResponseOrDict = Union[ - GenerateMemoriesResponse, GenerateMemoriesResponseDict -] + disable_memory_revisions: Optional[bool] + """Optional. Input only. Indicates whether no revision will be created for this request.""" + + display_name: Optional[str] + """Optional. Represents the display name of the Memory.""" + + expire_time: Optional[datetime.datetime] + """Optional. Represents the timestamp of when this resource is considered expired. This is *always* provided on output when `expiration` is set on input, regardless of whether `expire_time` or `ttl` was provided.""" + + fact: Optional[str] + """Optional. Represents semantic knowledge extracted from the source content.""" + + metadata: Optional[dict[str, MemoryMetadataValueDict]] + """Optional. Represents user-provided metadata for the Memory. This information was provided when creating, updating, or generating the Memory. It was not generated by Memory Bank.""" + + name: Optional[str] + """Identifier. Represents the resource name of the Memory. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}`""" + + revision_expire_time: Optional[datetime.datetime] + """Optional. Input only. Represents the timestamp of when the revision is considered expired. If not set, the memory revision will be kept until manually deleted.""" + + revision_labels: Optional[dict[str, str]] + """Optional. Input only. Represents the labels to apply to the Memory Revision created as a result of this request.""" + + revision_ttl: Optional[str] + """Optional. Input only. Represents the TTL for the revision. The expiration time is computed: now + TTL.""" + + scope: Optional[dict[str, str]] + """Required. Immutable. Represents the scope of the Memory. Memories are isolated within their scope. The scope is defined when creating or generating memories. Scope values cannot contain the wildcard character '*'.""" + + topics: Optional[list[MemoryTopicIdDict]] + """Optional. Represents the Topics of the Memory.""" + + ttl: Optional[str] + """Optional. Input only. Represents the TTL for this resource. The expiration time is computed: now + TTL.""" + + update_time: Optional[datetime.datetime] + """Output only. Represents the timestamp when this Memory was most recently updated.""" + + memory_type: Optional[MemoryType] + """Optional. Represents the type of the memory. If not set, the `NATURAL_LANGUAGE_COLLECTION` type is used. If `STRUCTURED_COLLECTION` or `STRUCTURED_PROFILE` is used, then `structured_data` must be provided.""" + + structured_content: Optional[MemoryStructuredContentDict] + """Optional. Represents the structured content of the memory.""" + + structured_data: Optional[dict[str, Any]] + """Optional. Deprecated: Use `structured_content` instead.""" -class AgentEngineGenerateMemoriesOperation(_common.BaseModel): - """Operation that generates memories for an agent engine.""" +MemoryOrDict = Union[Memory, MemoryDict] + + +class RuntimeMemoryOperation(_common.BaseModel): + """Operation that has an agent runtime memory as a response.""" name: Optional[str] = Field( default=None, @@ -10594,13 +10389,13 @@ class AgentEngineGenerateMemoriesOperation(_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.""" + response: Optional[Memory] = Field( + default=None, description="""The Agent Runtime Memory.""" ) -class AgentEngineGenerateMemoriesOperationDict(TypedDict, total=False): - """Operation that generates memories for an agent engine.""" +class RuntimeMemoryOperationDict(TypedDict, total=False): + """Operation that has an agent runtime memory as a response.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -10614,129 +10409,260 @@ class AgentEngineGenerateMemoriesOperationDict(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.""" + response: Optional[MemoryDict] + """The Agent Runtime Memory.""" -AgentEngineGenerateMemoriesOperationOrDict = Union[ - AgentEngineGenerateMemoriesOperation, AgentEngineGenerateMemoriesOperationDict -] +RuntimeMemoryOperationOrDict = Union[RuntimeMemoryOperation, RuntimeMemoryOperationDict] -class GetAgentEngineMemoryConfig(_common.BaseModel): - """Config for getting an Agent Engine Memory.""" +class DeleteRuntimeMemoryConfig(_common.BaseModel): + """Config for deleting an Agent Runtime Memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetAgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Memory.""" +class DeleteRuntimeMemoryConfigDict(TypedDict, total=False): + """Config for deleting an Agent Runtime Memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetAgentEngineMemoryConfigOrDict = Union[ - GetAgentEngineMemoryConfig, GetAgentEngineMemoryConfigDict +DeleteRuntimeMemoryConfigOrDict = Union[ + DeleteRuntimeMemoryConfig, DeleteRuntimeMemoryConfigDict ] -class _GetAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine.""" +class _DeleteRuntimeMemoryRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtimes.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime memory to delete.""" ) - config: Optional[GetAgentEngineMemoryConfig] = Field( + config: Optional[DeleteRuntimeMemoryConfig] = Field( default=None, description="""""" ) -class _GetAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine.""" +class _DeleteRuntimeMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtimes.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime memory to delete.""" - config: Optional[GetAgentEngineMemoryConfigDict] + config: Optional[DeleteRuntimeMemoryConfigDict] """""" -_GetAgentEngineMemoryRequestParametersOrDict = Union[ - _GetAgentEngineMemoryRequestParameters, _GetAgentEngineMemoryRequestParametersDict +_DeleteRuntimeMemoryRequestParametersOrDict = Union[ + _DeleteRuntimeMemoryRequestParameters, _DeleteRuntimeMemoryRequestParametersDict ] -class IngestionDirectContentsSourceEvent(_common.BaseModel): - """The direct contents source event for ingesting events.""" +class DeleteRuntimeMemoryOperation(_common.BaseModel): + """Operation for deleting agent runtimes.""" - content: Optional[genai_types.Content] = Field( - default=None, description="""Required. The content of the event.""" + 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}`.""", ) - event_id: Optional[str] = Field( + metadata: Optional[dict[str, Any]] = Field( default=None, - description="""Optional. A unique identifier for the event. If an event with the same event_id is ingested multiple times, it will be de-duplicated.""", + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", ) - event_time: Optional[datetime.datetime] = Field( + done: Optional[bool] = Field( default=None, - description="""Optional. The time at which the event occurred. If provided, this timestamp will be used for ordering events within a stream. If not provided, the server-side ingestion time will be used.""", + description="""If the 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 IngestionDirectContentsSourceEventDict(TypedDict, total=False): - """The direct contents source event for ingesting events.""" +class DeleteRuntimeMemoryOperationDict(TypedDict, total=False): + """Operation for deleting agent runtimes.""" - content: Optional[genai_types.Content] - """Required. The content of the 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}`.""" - 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.""" + 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.""" + + +DeleteRuntimeMemoryOperationOrDict = Union[ + DeleteRuntimeMemoryOperation, DeleteRuntimeMemoryOperationDict +] + + +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 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.""" + + session: Optional[str] + """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] + """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.""" + + +GenerateMemoriesRequestVertexSessionSourceOrDict = Union[ + GenerateMemoriesRequestVertexSessionSource, + GenerateMemoriesRequestVertexSessionSourceDict, +] + + +class GenerateMemoriesRequestDirectContentsSourceEvent(_common.BaseModel): + + content: Optional[genai_types.Content] = Field( + default=None, + description="""Required. A single piece of content from which to generate memories.""", + ) + event_time: Optional[datetime.datetime] = Field( + default=None, description="""Optional. The time at which the event occurred.""" + ) + + +class GenerateMemoriesRequestDirectContentsSourceEventDict(TypedDict, total=False): + + content: Optional[genai_types.Content] + """Required. A single piece of content from which to generate memories.""" 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.""" + """Optional. The time at which the event occurred.""" -IngestionDirectContentsSourceEventOrDict = Union[ - IngestionDirectContentsSourceEvent, IngestionDirectContentsSourceEventDict +GenerateMemoriesRequestDirectContentsSourceEventOrDict = Union[ + GenerateMemoriesRequestDirectContentsSourceEvent, + GenerateMemoriesRequestDirectContentsSourceEventDict, ] -class IngestionDirectContentsSource(_common.BaseModel): - """The direct contents source for ingesting events.""" +class GenerateMemoriesRequestDirectContentsSource(_common.BaseModel): + """The direct contents source for generating memories.""" - events: Optional[list[IngestionDirectContentsSourceEvent]] = Field( - default=None, description="""Required. The events to ingest.""" + events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEvent]] = Field( + default=None, + description="""Required. The source content (i.e. chat history) to generate memories from.""", ) -class IngestionDirectContentsSourceDict(TypedDict, total=False): - """The direct contents source for ingesting events.""" +class GenerateMemoriesRequestDirectContentsSourceDict(TypedDict, total=False): + """The direct contents source for generating memories.""" - events: Optional[list[IngestionDirectContentsSourceEventDict]] - """Required. The events to ingest.""" + events: Optional[list[GenerateMemoriesRequestDirectContentsSourceEventDict]] + """Required. The source content (i.e. chat history) to generate memories from.""" -IngestionDirectContentsSourceOrDict = Union[ - IngestionDirectContentsSource, IngestionDirectContentsSourceDict +GenerateMemoriesRequestDirectContentsSourceOrDict = Union[ + GenerateMemoriesRequestDirectContentsSource, + GenerateMemoriesRequestDirectContentsSourceDict, ] -class IngestEventsConfig(_common.BaseModel): - """Config for ingesting events.""" +class GenerateMemoriesRequestDirectMemoriesSourceDirectMemory(_common.BaseModel): + """A direct memory to upload to Memory Bank.""" + + 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.""", + ) + + +class GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict( + TypedDict, total=False +): + """A direct memory to upload to Memory Bank.""" + + fact: Optional[str] + """Required. The fact to consolidate with existing memories.""" + + topics: Optional[list[MemoryTopicIdDict]] + """Optional. The topics that the consolidated memories should be associated with.""" + + +GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryOrDict = Union[ + GenerateMemoriesRequestDirectMemoriesSourceDirectMemory, + GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict, +] + + +class GenerateMemoriesRequestDirectMemoriesSource(_common.BaseModel): + """The direct memories source for generating memories.""" + + direct_memories: Optional[ + list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemory] + ] = Field( + default=None, + description="""Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""", + ) + + +class GenerateMemoriesRequestDirectMemoriesSourceDict(TypedDict, total=False): + """The direct memories source for generating memories.""" + + direct_memories: Optional[ + list[GenerateMemoriesRequestDirectMemoriesSourceDirectMemoryDict] + ] + """Required. The direct memories to upload to Memory Bank. At most 5 direct memories are allowed per request.""" + + +GenerateMemoriesRequestDirectMemoriesSourceOrDict = Union[ + GenerateMemoriesRequestDirectMemoriesSource, + GenerateMemoriesRequestDirectMemoriesSourceDict, +] + + +class GenerateRuntimeMemoriesConfig(_common.BaseModel): + """Config for generating memories.""" 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( + disable_consolidation: 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="""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, @@ -10762,20 +10688,28 @@ class IngestEventsConfig(_common.BaseModel): default=None, description="""Optional. The strategy to use when applying metadata to existing memories.""", ) + allowed_topics: Optional[list[MemoryTopicId]] = Field( + default=None, + description="""Optional. Restricts memory generation to a subset of memory topics.""", + ) -class IngestEventsConfigDict(TypedDict, total=False): - """Config for ingesting events.""" +class GenerateRuntimeMemoriesConfigDict(TypedDict, total=False): + """Config for generating memories.""" 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.""" + disable_consolidation: Optional[bool] + """Whether to disable consolidation of memories. - 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`.""" + 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.""" 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.""" @@ -10795,69 +10729,148 @@ class IngestEventsConfigDict(TypedDict, total=False): metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] """Optional. The strategy to use when applying metadata to existing memories.""" + allowed_topics: Optional[list[MemoryTopicIdDict]] + """Optional. Restricts memory generation to a subset of memory topics.""" -IngestEventsConfigOrDict = Union[IngestEventsConfig, IngestEventsConfigDict] +GenerateRuntimeMemoriesConfigOrDict = Union[ + GenerateRuntimeMemoriesConfig, GenerateRuntimeMemoriesConfigDict +] -class _IngestEventsRequestParameters(_common.BaseModel): - """Parameters for purging agent engine memories.""" + +class _GenerateRuntimeMemoriesRequestParameters(_common.BaseModel): + """Parameters for generating agent runtime memories.""" name: Optional[str] = Field( - default=None, description="""Name of the Agent Engine to ingest events into.""" - ) - stream_id: Optional[str] = Field( - default=None, description="""The ID of the stream to ingest events into.""" + default=None, + description="""Name of the agent runtime to generate memories for.""", ) - direct_contents_source: Optional[IngestionDirectContentsSource] = Field( + vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSource] = Field( default=None, - description="""The direct memories source of the events that should be ingested.""", + description="""The vertex session source of the memories that should be generated.""", + ) + direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSource] = ( + Field( + default=None, + description="""The direct contents source of the memories that should be generated.""", + ) + ) + direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSource] = ( + Field( + default=None, + description="""The direct memories source of the memories that should be generated.""", + ) ) scope: Optional[dict[str, str]] = Field( default=None, - description="""The scope of the memories that should be generated from the stream. + description="""The scope of the memories that should be generated. - Memories will be consolidated across memories with the same scope. Scope + 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 '*'.""", ) - generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( - default=None, - description="""The configuration for the memory generation trigger.""", + config: Optional[GenerateRuntimeMemoriesConfig] = Field( + default=None, description="""""" ) - config: Optional[IngestEventsConfig] = Field(default=None, description="""""") -class _IngestEventsRequestParametersDict(TypedDict, total=False): - """Parameters for purging agent engine memories.""" +class _GenerateRuntimeMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for generating agent runtime memories.""" name: Optional[str] - """Name of the Agent Engine to ingest events into.""" + """Name of the agent runtime to generate memories for.""" - stream_id: Optional[str] - """The ID of the stream to ingest events into.""" + vertex_session_source: Optional[GenerateMemoriesRequestVertexSessionSourceDict] + """The vertex session source of the memories that should be generated.""" - direct_contents_source: Optional[IngestionDirectContentsSourceDict] - """The direct memories source of the events that should be ingested.""" + direct_contents_source: Optional[GenerateMemoriesRequestDirectContentsSourceDict] + """The direct contents source of the memories that should be generated.""" + + direct_memories_source: Optional[GenerateMemoriesRequestDirectMemoriesSourceDict] + """The direct memories source of the memories that should be generated.""" scope: Optional[dict[str, str]] - """The scope of the memories that should be generated from the stream. + """The scope of the memories that should be generated. - Memories will be consolidated across memories with the same scope. Scope + 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 '*'.""" - generation_trigger_config: Optional[MemoryGenerationTriggerConfigDict] - """The configuration for the memory generation trigger.""" - - config: Optional[IngestEventsConfigDict] + config: Optional[GenerateRuntimeMemoriesConfigDict] """""" -_IngestEventsRequestParametersOrDict = Union[ - _IngestEventsRequestParameters, _IngestEventsRequestParametersDict +_GenerateRuntimeMemoriesRequestParametersOrDict = Union[ + _GenerateRuntimeMemoriesRequestParameters, + _GenerateRuntimeMemoriesRequestParametersDict, ] -class MemoryBankIngestEventsOperation(_common.BaseModel): - """Operation that ingests events into a memory bank.""" +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 GenerateMemoriesResponseGeneratedMemoryDict(TypedDict, total=False): + """A memory that was generated.""" + + memory: Optional[MemoryDict] + """The generated memory.""" + + action: Optional[GenerateMemoriesResponseGeneratedMemoryAction] + """The action to take.""" + + previous_revision: Optional[str] + """The previous revision of the Memory before the action was performed. This + field is only set if the action is `UPDATED` or `DELETED`. You can use + this to rollback the Memory to the previous revision, undoing the action. + Format: + `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/memories/{memory}/revisions/{revision}`""" + + +GenerateMemoriesResponseGeneratedMemoryOrDict = Union[ + GenerateMemoriesResponseGeneratedMemory, GenerateMemoriesResponseGeneratedMemoryDict +] + + +class GenerateMemoriesResponse(_common.BaseModel): + """The response for generating memories.""" + + generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemory]] = Field( + default=None, description="""The generated memories.""" + ) + + +class GenerateMemoriesResponseDict(TypedDict, total=False): + """The response for generating memories.""" + + generated_memories: Optional[list[GenerateMemoriesResponseGeneratedMemoryDict]] + """The generated memories.""" + + +GenerateMemoriesResponseOrDict = Union[ + GenerateMemoriesResponse, GenerateMemoriesResponseDict +] + + +class RuntimeGenerateMemoriesOperation(_common.BaseModel): + """Operation that generates memories for an agent runtime.""" name: Optional[str] = Field( default=None, @@ -10875,10 +10888,13 @@ class MemoryBankIngestEventsOperation(_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 MemoryBankIngestEventsOperationDict(TypedDict, total=False): - """Operation that ingests events into a memory bank.""" +class RuntimeGenerateMemoriesOperationDict(TypedDict, total=False): + """Operation that generates memories for an agent runtime.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -10892,207 +10908,478 @@ class MemoryBankIngestEventsOperationDict(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.""" -MemoryBankIngestEventsOperationOrDict = Union[ - MemoryBankIngestEventsOperation, MemoryBankIngestEventsOperationDict + +RuntimeGenerateMemoriesOperationOrDict = Union[ + RuntimeGenerateMemoriesOperation, RuntimeGenerateMemoriesOperationDict ] -class ListAgentEngineMemoryConfig(_common.BaseModel): - """Config for listing agent engine memories.""" +class GetRuntimeMemoryConfig(_common.BaseModel): + """Config for getting an Agent Runtime Memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) - 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`""", - ) -class ListAgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for listing agent engine memories.""" +class GetRuntimeMemoryConfigDict(TypedDict, total=False): + """Config for getting an Agent Runtime Memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - page_size: Optional[int] - """""" - - page_token: Optional[str] - """""" - - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" - - order_by: Optional[str] - """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`""" - -ListAgentEngineMemoryConfigOrDict = Union[ - ListAgentEngineMemoryConfig, ListAgentEngineMemoryConfigDict -] +GetRuntimeMemoryConfigOrDict = Union[GetRuntimeMemoryConfig, GetRuntimeMemoryConfigDict] -class _ListAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" +class _GetRuntimeMemoryRequestParameters(_common.BaseModel): + """Parameters for getting an agent runtime.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" - ) - config: Optional[ListAgentEngineMemoryConfig] = Field( - default=None, description="""""" + default=None, description="""Name of the agent runtime.""" ) + config: Optional[GetRuntimeMemoryConfig] = Field(default=None, description="""""") -class _ListAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" +class _GetRuntimeMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for getting an agent runtime.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime.""" - config: Optional[ListAgentEngineMemoryConfigDict] + config: Optional[GetRuntimeMemoryConfigDict] """""" -_ListAgentEngineMemoryRequestParametersOrDict = Union[ - _ListAgentEngineMemoryRequestParameters, _ListAgentEngineMemoryRequestParametersDict +_GetRuntimeMemoryRequestParametersOrDict = Union[ + _GetRuntimeMemoryRequestParameters, _GetRuntimeMemoryRequestParametersDict ] -class ListReasoningEnginesMemoriesResponse(_common.BaseModel): - """Response for listing agent engine memories.""" +class IngestionDirectContentsSourceEvent(_common.BaseModel): + """The direct contents source event for ingesting events.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + content: Optional[genai_types.Content] = Field( + default=None, description="""Required. The content of the event.""" ) - next_page_token: Optional[str] = Field(default=None, description="""""") - memories: Optional[list[Memory]] = Field( - default=None, description="""List of agent engine memories.""" + 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.""", ) -class ListReasoningEnginesMemoriesResponseDict(TypedDict, total=False): - """Response for listing agent engine memories.""" +class IngestionDirectContentsSourceEventDict(TypedDict, total=False): + """The direct contents source event for ingesting events.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + content: Optional[genai_types.Content] + """Required. The content of the event.""" - next_page_token: Optional[str] - """""" + 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.""" - memories: Optional[list[MemoryDict]] - """List of agent engine memories.""" + 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.""" -ListReasoningEnginesMemoriesResponseOrDict = Union[ - ListReasoningEnginesMemoriesResponse, ListReasoningEnginesMemoriesResponseDict +IngestionDirectContentsSourceEventOrDict = Union[ + IngestionDirectContentsSourceEvent, IngestionDirectContentsSourceEventDict ] -class _GetAgentEngineMemoryOperationParameters(_common.BaseModel): - """Parameters for getting an operation with a memory as a response.""" +class IngestionDirectContentsSource(_common.BaseModel): + """The direct contents source for ingesting events.""" - 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.""" + events: Optional[list[IngestionDirectContentsSourceEvent]] = Field( + default=None, description="""Required. The events to ingest.""" ) -class _GetAgentEngineMemoryOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with a memory as a response.""" - - operation_name: Optional[str] - """The server-assigned name for the operation.""" +class IngestionDirectContentsSourceDict(TypedDict, total=False): + """The direct contents source for ingesting events.""" - config: Optional[GetAgentEngineOperationConfigDict] - """Used to override the default configuration.""" + events: Optional[list[IngestionDirectContentsSourceEventDict]] + """Required. The events to ingest.""" -_GetAgentEngineMemoryOperationParametersOrDict = Union[ - _GetAgentEngineMemoryOperationParameters, - _GetAgentEngineMemoryOperationParametersDict, +IngestionDirectContentsSourceOrDict = Union[ + IngestionDirectContentsSource, IngestionDirectContentsSourceDict ] -class _GetAgentEngineGenerateMemoriesOperationParameters(_common.BaseModel): - """Parameters for getting an operation with generated memories as a response.""" +class IngestEventsConfig(_common.BaseModel): + """Config for ingesting events.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - config: Optional[GetAgentEngineOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" + 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 _GetAgentEngineGenerateMemoriesOperationParametersDict(TypedDict, total=False): - """Parameters for getting an operation with generated memories as a response.""" +class IngestEventsConfigDict(TypedDict, total=False): + """Config for ingesting events.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - config: Optional[GetAgentEngineOperationConfigDict] - """Used to override the default configuration.""" + 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`.""" -_GetAgentEngineGenerateMemoriesOperationParametersOrDict = Union[ - _GetAgentEngineGenerateMemoriesOperationParameters, - _GetAgentEngineGenerateMemoriesOperationParametersDict, -] + 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 RetrieveMemoriesRequestSimilaritySearchParams(_common.BaseModel): - """The parameters for semantic similarity search based retrieval.""" + revision_ttl: Optional[str] + """Optional. Input only. The TTL for the revision. The expiration time is computed: now + TTL.""" - 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( - 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.""", - ) + 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 RetrieveMemoriesRequestSimilaritySearchParamsDict(TypedDict, total=False): - """The parameters for semantic similarity search based retrieval.""" + metadata_merge_strategy: Optional[MemoryMetadataMergeStrategy] + """Optional. The strategy to use when applying metadata to existing 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.""" - 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.""" +IngestEventsConfigOrDict = Union[IngestEventsConfig, IngestEventsConfigDict] -RetrieveMemoriesRequestSimilaritySearchParamsOrDict = Union[ - RetrieveMemoriesRequestSimilaritySearchParams, - RetrieveMemoriesRequestSimilaritySearchParamsDict, -] +class _IngestEventsRequestParameters(_common.BaseModel): + """Parameters for purging agent runtime memories.""" + + name: Optional[str] = Field( + default=None, description="""Name of the Agent Runtime to ingest events into.""" + ) + stream_id: Optional[str] = Field( + default=None, description="""The ID of the stream to ingest events into.""" + ) + 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. + + Memories will be consolidated across memories with the same scope. Scope + values cannot contain the wildcard character '*'.""", + ) + generation_trigger_config: Optional[MemoryGenerationTriggerConfig] = Field( + default=None, + description="""The configuration for the memory generation trigger.""", + ) + config: Optional[IngestEventsConfig] = Field(default=None, description="""""") + + +class _IngestEventsRequestParametersDict(TypedDict, total=False): + """Parameters for purging agent runtime memories.""" + + name: Optional[str] + """Name of the Agent Runtime to ingest events into.""" + + stream_id: Optional[str] + """The ID of the stream to ingest events into.""" + + 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. + + 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, + 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 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}`.""" + + 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.""" + + +MemoryBankIngestEventsOperationOrDict = Union[ + MemoryBankIngestEventsOperation, MemoryBankIngestEventsOperationDict +] + + +class ListRuntimeMemoryConfig(_common.BaseModel): + """Config for listing agent runtime memories.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + 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`""", + ) + + +class ListRuntimeMemoryConfigDict(TypedDict, total=False): + """Config for listing agent runtime 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.""" + + 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`""" + + +ListRuntimeMemoryConfigOrDict = Union[ + ListRuntimeMemoryConfig, ListRuntimeMemoryConfigDict +] + + +class _ListRuntimeMemoryRequestParameters(_common.BaseModel): + """Parameters for listing agent runtimes.""" + + name: Optional[str] = Field( + default=None, description="""Name of the agent runtime.""" + ) + config: Optional[ListRuntimeMemoryConfig] = Field(default=None, description="""""") + + +class _ListRuntimeMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtimes.""" + + name: Optional[str] + """Name of the agent runtime.""" + + config: Optional[ListRuntimeMemoryConfigDict] + """""" + + +_ListRuntimeMemoryRequestParametersOrDict = Union[ + _ListRuntimeMemoryRequestParameters, _ListRuntimeMemoryRequestParametersDict +] + + +class ListReasoningEnginesMemoriesResponse(_common.BaseModel): + """Response for listing agent runtime memories.""" + + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" + ) + next_page_token: Optional[str] = Field(default=None, description="""""") + memories: Optional[list[Memory]] = Field( + default=None, description="""List of agent runtime memories.""" + ) + + +class ListReasoningEnginesMemoriesResponseDict(TypedDict, total=False): + """Response for listing agent runtime memories.""" + + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + + next_page_token: Optional[str] + """""" + + memories: Optional[list[MemoryDict]] + """List of agent runtime memories.""" + + +ListReasoningEnginesMemoriesResponseOrDict = Union[ + ListReasoningEnginesMemoriesResponse, ListReasoningEnginesMemoriesResponseDict +] + + +class _GetRuntimeMemoryOperationParameters(_common.BaseModel): + """Parameters for getting an operation with a memory as a response.""" + + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetRuntimeOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" + ) + + +class _GetRuntimeMemoryOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with a memory as a response.""" + + operation_name: Optional[str] + """The server-assigned name for the operation.""" + + config: Optional[GetRuntimeOperationConfigDict] + """Used to override the default configuration.""" + + +_GetRuntimeMemoryOperationParametersOrDict = Union[ + _GetRuntimeMemoryOperationParameters, _GetRuntimeMemoryOperationParametersDict +] + + +class _GetRuntimeGenerateMemoriesOperationParameters(_common.BaseModel): + """Parameters for getting an operation with generated memories as a response.""" + + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetRuntimeOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" + ) + + +class _GetRuntimeGenerateMemoriesOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation with generated memories as a response.""" + + operation_name: Optional[str] + """The server-assigned name for the operation.""" + + config: Optional[GetRuntimeOperationConfigDict] + """Used to override the default configuration.""" + + +_GetRuntimeGenerateMemoriesOperationParametersOrDict = Union[ + _GetRuntimeGenerateMemoriesOperationParameters, + _GetRuntimeGenerateMemoriesOperationParametersDict, +] + + +class RetrieveMemoriesRequestSimilaritySearchParams(_common.BaseModel): + """The parameters for semantic similarity search based retrieval.""" + + 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( + default=None, + description="""Optional. The maximum number of memories to return. The service may return fewer than this value. If unspecified, at most 3 memories will be returned. The maximum value is 100; values above 100 will be coerced to 100.""", + ) + + +class 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.""" + + 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.""" + + +RetrieveMemoriesRequestSimilaritySearchParamsOrDict = Union[ + RetrieveMemoriesRequestSimilaritySearchParams, + RetrieveMemoriesRequestSimilaritySearchParamsDict, +] class RetrieveMemoriesRequestSimpleRetrievalParams(_common.BaseModel): @@ -11183,7 +11470,7 @@ class MemoryConjunctionFilterDict(TypedDict, total=False): ] -class RetrieveAgentEngineMemoriesConfig(_common.BaseModel): +class RetrieveRuntimeMemoriesConfig(_common.BaseModel): """Config for retrieving memories.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -11226,7 +11513,7 @@ class RetrieveAgentEngineMemoriesConfig(_common.BaseModel): ) -class RetrieveAgentEngineMemoriesConfigDict(TypedDict, total=False): +class RetrieveRuntimeMemoriesConfigDict(TypedDict, total=False): """Config for retrieving memories.""" http_options: Optional[genai_types.HttpOptions] @@ -11265,17 +11552,17 @@ class RetrieveAgentEngineMemoriesConfigDict(TypedDict, total=False): retrieve memories matching any of the specified `MemoryType` values.""" -RetrieveAgentEngineMemoriesConfigOrDict = Union[ - RetrieveAgentEngineMemoriesConfig, RetrieveAgentEngineMemoriesConfigDict +RetrieveRuntimeMemoriesConfigOrDict = Union[ + RetrieveRuntimeMemoriesConfig, RetrieveRuntimeMemoriesConfigDict ] -class _RetrieveAgentEngineMemoriesRequestParameters(_common.BaseModel): - """Parameters for retrieving agent engine memories.""" +class _RetrieveRuntimeMemoriesRequestParameters(_common.BaseModel): + """Parameters for retrieving agent runtime memories.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to retrieve memories from.""", + description="""Name of the agent runtime to retrieve memories from.""", ) scope: Optional[dict[str, str]] = Field( default=None, @@ -11297,16 +11584,16 @@ class _RetrieveAgentEngineMemoriesRequestParameters(_common.BaseModel): description="""Parameters for simple (non-similarity search) retrieval.""", ) ) - config: Optional[RetrieveAgentEngineMemoriesConfig] = Field( + config: Optional[RetrieveRuntimeMemoriesConfig] = Field( default=None, description="""""" ) -class _RetrieveAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving agent engine memories.""" +class _RetrieveRuntimeMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for retrieving agent runtime memories.""" name: Optional[str] - """Name of the agent engine to retrieve memories from.""" + """Name of the agent runtime to retrieve memories from.""" scope: Optional[dict[str, str]] """The scope of the memories to retrieve. @@ -11323,13 +11610,13 @@ class _RetrieveAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): simple_retrieval_params: Optional[RetrieveMemoriesRequestSimpleRetrievalParamsDict] """Parameters for simple (non-similarity search) retrieval.""" - config: Optional[RetrieveAgentEngineMemoriesConfigDict] + config: Optional[RetrieveRuntimeMemoriesConfigDict] """""" -_RetrieveAgentEngineMemoriesRequestParametersOrDict = Union[ - _RetrieveAgentEngineMemoriesRequestParameters, - _RetrieveAgentEngineMemoriesRequestParametersDict, +_RetrieveRuntimeMemoriesRequestParametersOrDict = Union[ + _RetrieveRuntimeMemoriesRequestParameters, + _RetrieveRuntimeMemoriesRequestParametersDict, ] @@ -11408,11 +11695,11 @@ class RetrieveMemoryProfilesConfigDict(TypedDict, total=False): class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): - """Parameters for retrieving agent engine memory profiles.""" + """Parameters for retrieving agent runtime memory profiles.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to retrieve memory profiles from.""", + description="""Name of the agent runtime to retrieve memory profiles from.""", ) scope: Optional[dict[str, str]] = Field( default=None, @@ -11428,10 +11715,10 @@ class _RetrieveMemoryProfilesRequestParameters(_common.BaseModel): class _RetrieveMemoryProfilesRequestParametersDict(TypedDict, total=False): - """Parameters for retrieving agent engine memory profiles.""" + """Parameters for retrieving agent runtime memory profiles.""" name: Optional[str] - """Name of the agent engine to retrieve memory profiles from.""" + """Name of the agent runtime to retrieve memory profiles from.""" scope: Optional[dict[str, str]] """The scope of the memories to retrieve. @@ -11502,7 +11789,7 @@ class RetrieveProfilesResponseDict(TypedDict, total=False): ] -class RollbackAgentEngineMemoryConfig(_common.BaseModel): +class RollbackRuntimeMemoryConfig(_common.BaseModel): """Config for rolling back a memory.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -11514,7 +11801,7 @@ class RollbackAgentEngineMemoryConfig(_common.BaseModel): ) -class RollbackAgentEngineMemoryConfigDict(TypedDict, total=False): +class RollbackRuntimeMemoryConfigDict(TypedDict, total=False): """Config for rolling back a memory.""" http_options: Optional[genai_types.HttpOptions] @@ -11524,45 +11811,44 @@ class RollbackAgentEngineMemoryConfigDict(TypedDict, total=False): """Waits for the operation to complete before returning.""" -RollbackAgentEngineMemoryConfigOrDict = Union[ - RollbackAgentEngineMemoryConfig, RollbackAgentEngineMemoryConfigDict +RollbackRuntimeMemoryConfigOrDict = Union[ + RollbackRuntimeMemoryConfig, RollbackRuntimeMemoryConfigDict ] -class _RollbackAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for generating agent engine memories.""" +class _RollbackRuntimeMemoryRequestParameters(_common.BaseModel): + """Parameters for generating agent runtime memories.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine memory to rollback.""" + default=None, description="""Name of the agent runtime memory to rollback.""" ) target_revision_id: Optional[str] = Field( default=None, description="""The ID of the revision to rollback to.""" ) - config: Optional[RollbackAgentEngineMemoryConfig] = Field( + config: Optional[RollbackRuntimeMemoryConfig] = Field( default=None, description="""""" ) -class _RollbackAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for generating agent engine memories.""" +class _RollbackRuntimeMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for generating agent runtime memories.""" name: Optional[str] - """Name of the agent engine memory to rollback.""" + """Name of the agent runtime memory to rollback.""" target_revision_id: Optional[str] """The ID of the revision to rollback to.""" - config: Optional[RollbackAgentEngineMemoryConfigDict] + config: Optional[RollbackRuntimeMemoryConfigDict] """""" -_RollbackAgentEngineMemoryRequestParametersOrDict = Union[ - _RollbackAgentEngineMemoryRequestParameters, - _RollbackAgentEngineMemoryRequestParametersDict, +_RollbackRuntimeMemoryRequestParametersOrDict = Union[ + _RollbackRuntimeMemoryRequestParameters, _RollbackRuntimeMemoryRequestParametersDict ] -class AgentEngineRollbackMemoryOperation(_common.BaseModel): +class RuntimeRollbackMemoryOperation(_common.BaseModel): """Operation that rolls back a memory.""" name: Optional[str] = Field( @@ -11583,7 +11869,7 @@ class AgentEngineRollbackMemoryOperation(_common.BaseModel): ) -class AgentEngineRollbackMemoryOperationDict(TypedDict, total=False): +class RuntimeRollbackMemoryOperationDict(TypedDict, total=False): """Operation that rolls back a memory.""" name: Optional[str] @@ -11599,13 +11885,13 @@ class AgentEngineRollbackMemoryOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -AgentEngineRollbackMemoryOperationOrDict = Union[ - AgentEngineRollbackMemoryOperation, AgentEngineRollbackMemoryOperationDict +RuntimeRollbackMemoryOperationOrDict = Union[ + RuntimeRollbackMemoryOperation, RuntimeRollbackMemoryOperationDict ] -class UpdateAgentEngineMemoryConfig(_common.BaseModel): - """Config for updating agent engine memory.""" +class UpdateRuntimeMemoryConfig(_common.BaseModel): + """Config for updating agent runtime memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -11660,8 +11946,8 @@ class UpdateAgentEngineMemoryConfig(_common.BaseModel): ) -class UpdateAgentEngineMemoryConfigDict(TypedDict, total=False): - """Config for updating agent engine memory.""" +class UpdateRuntimeMemoryConfigDict(TypedDict, total=False): + """Config for updating agent runtime memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -11706,16 +11992,16 @@ class UpdateAgentEngineMemoryConfigDict(TypedDict, total=False): https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" -UpdateAgentEngineMemoryConfigOrDict = Union[ - UpdateAgentEngineMemoryConfig, UpdateAgentEngineMemoryConfigDict +UpdateRuntimeMemoryConfigOrDict = Union[ + UpdateRuntimeMemoryConfig, UpdateRuntimeMemoryConfigDict ] -class _UpdateAgentEngineMemoryRequestParameters(_common.BaseModel): - """Parameters for updating agent engine memories.""" +class _UpdateRuntimeMemoryRequestParameters(_common.BaseModel): + """Parameters for updating agent runtime memories.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine memory to update.""" + default=None, description="""Name of the agent runtime memory to update.""" ) fact: Optional[str] = Field( default=None, @@ -11731,16 +12017,16 @@ class _UpdateAgentEngineMemoryRequestParameters(_common.BaseModel): creating or generating memories. Up to 5 key-value pairs are accepted, and scope values cannot contain the wildcard character '*'.""", ) - config: Optional[UpdateAgentEngineMemoryConfig] = Field( + config: Optional[UpdateRuntimeMemoryConfig] = Field( default=None, description="""""" ) -class _UpdateAgentEngineMemoryRequestParametersDict(TypedDict, total=False): - """Parameters for updating agent engine memories.""" +class _UpdateRuntimeMemoryRequestParametersDict(TypedDict, total=False): + """Parameters for updating agent runtime memories.""" name: Optional[str] - """Name of the agent engine memory to update.""" + """Name of the agent runtime memory to update.""" fact: Optional[str] """The updated fact of the memory. @@ -11754,17 +12040,16 @@ class _UpdateAgentEngineMemoryRequestParametersDict(TypedDict, total=False): creating or generating memories. Up to 5 key-value pairs are accepted, and scope values cannot contain the wildcard character '*'.""" - config: Optional[UpdateAgentEngineMemoryConfigDict] + config: Optional[UpdateRuntimeMemoryConfigDict] """""" -_UpdateAgentEngineMemoryRequestParametersOrDict = Union[ - _UpdateAgentEngineMemoryRequestParameters, - _UpdateAgentEngineMemoryRequestParametersDict, +_UpdateRuntimeMemoryRequestParametersOrDict = Union[ + _UpdateRuntimeMemoryRequestParameters, _UpdateRuntimeMemoryRequestParametersDict ] -class PurgeAgentEngineMemoriesConfig(_common.BaseModel): +class PurgeRuntimeMemoriesConfig(_common.BaseModel): """Config for purging memories.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -11776,7 +12061,7 @@ class PurgeAgentEngineMemoriesConfig(_common.BaseModel): ) -class PurgeAgentEngineMemoriesConfigDict(TypedDict, total=False): +class PurgeRuntimeMemoriesConfigDict(TypedDict, total=False): """Config for purging memories.""" http_options: Optional[genai_types.HttpOptions] @@ -11786,16 +12071,17 @@ class PurgeAgentEngineMemoriesConfigDict(TypedDict, total=False): """Waits for the operation to complete before returning.""" -PurgeAgentEngineMemoriesConfigOrDict = Union[ - PurgeAgentEngineMemoriesConfig, PurgeAgentEngineMemoriesConfigDict +PurgeRuntimeMemoriesConfigOrDict = Union[ + PurgeRuntimeMemoriesConfig, PurgeRuntimeMemoriesConfigDict ] -class _PurgeAgentEngineMemoriesRequestParameters(_common.BaseModel): - """Parameters for purging agent engine memories.""" +class _PurgeRuntimeMemoriesRequestParameters(_common.BaseModel): + """Parameters for purging agent runtime memories.""" name: Optional[str] = Field( - default=None, description="""Name of the Agent Engine to purge memories from.""" + default=None, + description="""Name of the Agent Runtime to purge memories from.""", ) filter: Optional[str] = Field( default=None, @@ -11823,16 +12109,16 @@ class _PurgeAgentEngineMemoriesRequestParameters(_common.BaseModel): default=None, description="""If true, the memories will actually be purged. If false, the purge request will be validated but not executed.""", ) - config: Optional[PurgeAgentEngineMemoriesConfig] = Field( + config: Optional[PurgeRuntimeMemoriesConfig] = Field( default=None, description="""""" ) -class _PurgeAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): - """Parameters for purging agent engine memories.""" +class _PurgeRuntimeMemoriesRequestParametersDict(TypedDict, total=False): + """Parameters for purging agent runtime memories.""" name: Optional[str] - """Name of the Agent Engine to purge memories from.""" + """Name of the Agent Runtime to purge memories from.""" filter: Optional[str] """The standard list filter to determine which memories to purge. @@ -11857,13 +12143,12 @@ class _PurgeAgentEngineMemoriesRequestParametersDict(TypedDict, total=False): force: Optional[bool] """If true, the memories will actually be purged. If false, the purge request will be validated but not executed.""" - config: Optional[PurgeAgentEngineMemoriesConfigDict] + config: Optional[PurgeRuntimeMemoriesConfigDict] """""" -_PurgeAgentEngineMemoriesRequestParametersOrDict = Union[ - _PurgeAgentEngineMemoriesRequestParameters, - _PurgeAgentEngineMemoriesRequestParametersDict, +_PurgeRuntimeMemoriesRequestParametersOrDict = Union[ + _PurgeRuntimeMemoriesRequestParameters, _PurgeRuntimeMemoriesRequestParametersDict ] @@ -11885,8 +12170,8 @@ class PurgeMemoriesResponseDict(TypedDict, total=False): PurgeMemoriesResponseOrDict = Union[PurgeMemoriesResponse, PurgeMemoriesResponseDict] -class AgentEnginePurgeMemoriesOperation(_common.BaseModel): - """Operation that purges memories from an agent engine.""" +class RuntimePurgeMemoriesOperation(_common.BaseModel): + """Operation that purges memories from an agent runtime.""" name: Optional[str] = Field( default=None, @@ -11909,8 +12194,8 @@ class AgentEnginePurgeMemoriesOperation(_common.BaseModel): ) -class AgentEnginePurgeMemoriesOperationDict(TypedDict, total=False): - """Operation that purges memories from an agent engine.""" +class RuntimePurgeMemoriesOperationDict(TypedDict, total=False): + """Operation that purges memories from an agent runtime.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -11928,55 +12213,55 @@ class AgentEnginePurgeMemoriesOperationDict(TypedDict, total=False): """The response for purging memories.""" -AgentEnginePurgeMemoriesOperationOrDict = Union[ - AgentEnginePurgeMemoriesOperation, AgentEnginePurgeMemoriesOperationDict +RuntimePurgeMemoriesOperationOrDict = Union[ + RuntimePurgeMemoriesOperation, RuntimePurgeMemoriesOperationDict ] -class GetAgentEngineMemoryRevisionConfig(_common.BaseModel): - """Config for getting an Agent Engine Memory Revision.""" +class GetRuntimeMemoryRevisionConfig(_common.BaseModel): + """Config for getting an Agent Runtime Memory Revision.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetAgentEngineMemoryRevisionConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Memory Revision.""" +class GetRuntimeMemoryRevisionConfigDict(TypedDict, total=False): + """Config for getting an Agent Runtime Memory Revision.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetAgentEngineMemoryRevisionConfigOrDict = Union[ - GetAgentEngineMemoryRevisionConfig, GetAgentEngineMemoryRevisionConfigDict +GetRuntimeMemoryRevisionConfigOrDict = Union[ + GetRuntimeMemoryRevisionConfig, GetRuntimeMemoryRevisionConfigDict ] -class _GetAgentEngineMemoryRevisionRequestParameters(_common.BaseModel): - """Parameters for getting an Agent Engine memory revision.""" +class _GetRuntimeMemoryRevisionRequestParameters(_common.BaseModel): + """Parameters for getting an Agent Runtime memory revision.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime.""" ) - config: Optional[GetAgentEngineMemoryRevisionConfig] = Field( + config: Optional[GetRuntimeMemoryRevisionConfig] = Field( default=None, description="""""" ) -class _GetAgentEngineMemoryRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for getting an Agent Engine memory revision.""" +class _GetRuntimeMemoryRevisionRequestParametersDict(TypedDict, total=False): + """Parameters for getting an Agent Runtime memory revision.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime.""" - config: Optional[GetAgentEngineMemoryRevisionConfigDict] + config: Optional[GetRuntimeMemoryRevisionConfigDict] """""" -_GetAgentEngineMemoryRevisionRequestParametersOrDict = Union[ - _GetAgentEngineMemoryRevisionRequestParameters, - _GetAgentEngineMemoryRevisionRequestParametersDict, +_GetRuntimeMemoryRevisionRequestParametersOrDict = Union[ + _GetRuntimeMemoryRevisionRequestParameters, + _GetRuntimeMemoryRevisionRequestParametersDict, ] @@ -11995,6 +12280,10 @@ class IntermediateExtractedMemory(_common.BaseModel): default=None, description="""Output only. Represents the structured value of the extracted memory.""", ) + is_explicit: Optional[bool] = Field( + default=None, + description="""Output only. Indicates that the extracted memory originated from an explicit instruction to remember or forget information.""", + ) class IntermediateExtractedMemoryDict(TypedDict, total=False): @@ -12009,6 +12298,9 @@ class IntermediateExtractedMemoryDict(TypedDict, total=False): structured_data: Optional[dict[str, Any]] """Output only. Represents the structured value of the extracted memory.""" + is_explicit: Optional[bool] + """Output only. Indicates that the extracted memory originated from an explicit instruction to remember or forget information.""" + IntermediateExtractedMemoryOrDict = Union[ IntermediateExtractedMemory, IntermediateExtractedMemoryDict @@ -12076,8 +12368,8 @@ class MemoryRevisionDict(TypedDict, total=False): MemoryRevisionOrDict = Union[MemoryRevision, MemoryRevisionDict] -class ListAgentEngineMemoryRevisionsConfig(_common.BaseModel): - """Config for listing Agent Engine memory revisions.""" +class ListRuntimeMemoryRevisionsConfig(_common.BaseModel): + """Config for listing Agent Runtime memory revisions.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -12091,8 +12383,8 @@ class ListAgentEngineMemoryRevisionsConfig(_common.BaseModel): ) -class ListAgentEngineMemoryRevisionsConfigDict(TypedDict, total=False): - """Config for listing Agent Engine memory revisions.""" +class ListRuntimeMemoryRevisionsConfigDict(TypedDict, total=False): + """Config for listing Agent Runtime memory revisions.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -12108,40 +12400,40 @@ class ListAgentEngineMemoryRevisionsConfigDict(TypedDict, total=False): For field names both snake_case and camelCase are supported.""" -ListAgentEngineMemoryRevisionsConfigOrDict = Union[ - ListAgentEngineMemoryRevisionsConfig, ListAgentEngineMemoryRevisionsConfigDict +ListRuntimeMemoryRevisionsConfigOrDict = Union[ + ListRuntimeMemoryRevisionsConfig, ListRuntimeMemoryRevisionsConfigDict ] -class _ListAgentEngineMemoryRevisionsRequestParameters(_common.BaseModel): - """Parameters for listing Agent Engine memory revisions.""" +class _ListRuntimeMemoryRevisionsRequestParameters(_common.BaseModel): + """Parameters for listing Agent Runtime memory revisions.""" name: Optional[str] = Field( - default=None, description="""Name of the Agent Engine memory""" + default=None, description="""Name of the Agent Runtime memory""" ) - config: Optional[ListAgentEngineMemoryRevisionsConfig] = Field( + config: Optional[ListRuntimeMemoryRevisionsConfig] = Field( default=None, description="""""" ) -class _ListAgentEngineMemoryRevisionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing Agent Engine memory revisions.""" +class _ListRuntimeMemoryRevisionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing Agent Runtime memory revisions.""" name: Optional[str] - """Name of the Agent Engine memory""" + """Name of the Agent Runtime memory""" - config: Optional[ListAgentEngineMemoryRevisionsConfigDict] + config: Optional[ListRuntimeMemoryRevisionsConfigDict] """""" -_ListAgentEngineMemoryRevisionsRequestParametersOrDict = Union[ - _ListAgentEngineMemoryRevisionsRequestParameters, - _ListAgentEngineMemoryRevisionsRequestParametersDict, +_ListRuntimeMemoryRevisionsRequestParametersOrDict = Union[ + _ListRuntimeMemoryRevisionsRequestParameters, + _ListRuntimeMemoryRevisionsRequestParametersDict, ] -class ListAgentEngineMemoryRevisionsResponse(_common.BaseModel): - """Response for listing agent engine memory revisions.""" +class ListRuntimeMemoryRevisionsResponse(_common.BaseModel): + """Response for listing agent runtime memory revisions.""" sdk_http_response: Optional[genai_types.HttpResponse] = Field( default=None, description="""Used to retain the full HTTP response.""" @@ -12152,8 +12444,8 @@ class ListAgentEngineMemoryRevisionsResponse(_common.BaseModel): ) -class ListAgentEngineMemoryRevisionsResponseDict(TypedDict, total=False): - """Response for listing agent engine memory revisions.""" +class ListRuntimeMemoryRevisionsResponseDict(TypedDict, total=False): + """Response for listing agent runtime memory revisions.""" sdk_http_response: Optional[genai_types.HttpResponse] """Used to retain the full HTTP response.""" @@ -12165,8 +12457,8 @@ class ListAgentEngineMemoryRevisionsResponseDict(TypedDict, total=False): """List of memory revisions.""" -ListAgentEngineMemoryRevisionsResponseOrDict = Union[ - ListAgentEngineMemoryRevisionsResponse, ListAgentEngineMemoryRevisionsResponseDict +ListRuntimeMemoryRevisionsResponseOrDict = Union[ + ListRuntimeMemoryRevisionsResponse, ListRuntimeMemoryRevisionsResponseDict ] @@ -12228,6 +12520,13 @@ class RagQuery(_common.BaseModel): default=None, description="""Optional. The query in text format to get relevant contexts.""", ) + vector_distance_threshold: Optional[float] = Field( + default=None, + description="""Optional. Only return contexts with vector distance smaller than the threshold. This field is only used internally to set the RetrieveRagStoreContextsRequest. External users should use the public vector_distance_threshold field inside the VertexRagStore.""", + ) + metadata_filter: Optional[str] = Field( + default=None, description="""Optional. String for metadata filtering.""" + ) class RagQueryDict(TypedDict, total=False): @@ -12245,6 +12544,12 @@ class RagQueryDict(TypedDict, total=False): text: Optional[str] """Optional. The query in text format to get relevant contexts.""" + vector_distance_threshold: Optional[float] + """Optional. Only return contexts with vector distance smaller than the threshold. This field is only used internally to set the RetrieveRagStoreContextsRequest. External users should use the public vector_distance_threshold field inside the VertexRagStore.""" + + metadata_filter: Optional[str] + """Optional. String for metadata filtering.""" + RagQueryOrDict = Union[RagQuery, RagQueryDict] @@ -12428,6 +12733,10 @@ class RagFileParsingConfigLlmParser(_common.BaseModel): default=None, description="""The name of a LLM model used for parsing. Format: * `projects/{project_id}/locations/{location}/publishers/{publisher}/models/{model}`""", ) + memory_mode: Optional[bool] = Field( + default=None, + description="""Whether the llm parser is for memory parasing or not. If set to true, the llm parser will use the memory mode prompt. If set to false or not set, the llm parser will use the default prompt.""", + ) class RagFileParsingConfigLlmParserDict(TypedDict, total=False): @@ -12445,6 +12754,9 @@ class RagFileParsingConfigLlmParserDict(TypedDict, total=False): model_name: Optional[str] """The name of a LLM model used for parsing. Format: * `projects/{project_id}/locations/{location}/publishers/{publisher}/models/{model}`""" + memory_mode: Optional[bool] + """Whether the llm parser is for memory parasing or not. If set to true, the llm parser will use the memory mode prompt. If set to false or not set, the llm parser will use the default prompt.""" + RagFileParsingConfigLlmParserOrDict = Union[ RagFileParsingConfigLlmParser, RagFileParsingConfigLlmParserDict @@ -15197,6 +15509,17 @@ class UploadRagFileConfig(_common.BaseModel): default=None, description="""Specifies the transformation config for RagFiles.""", ) + inline_metadata: Optional[str] = Field( + default=None, description="""Inline metadata for the RagFile.""" + ) + 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 upload 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.""", + ) + 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.""", + ) class UploadRagFileConfigDict(TypedDict, total=False): @@ -15214,6 +15537,15 @@ class UploadRagFileConfigDict(TypedDict, total=False): rag_file_transformation_config: Optional[RagFileTransformationConfigDict] """Specifies the transformation config for RagFiles.""" + inline_metadata: Optional[str] + """Inline metadata for the RagFile.""" + + 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 upload 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.""" + + 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.""" + UploadRagFileConfigOrDict = Union[UploadRagFileConfig, UploadRagFileConfigDict] @@ -15258,428 +15590,832 @@ class _UploadRagFileParametersDict(TypedDict, total=False): ] -class UploadRagFileResponse(_common.BaseModel): - """Response for uploading a Rag File.""" +class GdataObjectId(_common.BaseModel): + """This is a copy of the tech.blob.ObjectId proto, which could not be used directly here due to transitive closure issues with JavaScript support; see http://b/8801763.""" - error: Optional[genai_types.GoogleRpcStatus] = Field( + bucket_name: Optional[str] = Field( default=None, - description="""The error that occurred while processing the RagFile.""", + description="""The name of the bucket to which this object belongs.""", ) - rag_file: Optional[RagFile] = Field( + object_name: Optional[str] = Field( + default=None, description="""The name of the object.""" + ) + generation: Optional[int] = Field( default=None, - description="""The RagFile that had been uploaded into the RagCorpus.""", + description="""Generation of the object. Generations are monotonically increasing across writes, allowing them to be be compared to determine which generation is newer. If this is omitted in a request, then you are requesting the live object. See http://go/bigstore-versions""", ) -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.""" +class GdataObjectIdDict(TypedDict, total=False): + """This is a copy of the tech.blob.ObjectId proto, which could not be used directly here due to transitive closure issues with JavaScript support; see http://b/8801763.""" + bucket_name: Optional[str] + """The name of the bucket to which this object belongs.""" -UploadRagFileResponseOrDict = Union[UploadRagFileResponse, UploadRagFileResponseDict] + object_name: Optional[str] + """The name of the object.""" + generation: Optional[int] + """Generation of the object. Generations are monotonically increasing across writes, allowing them to be be compared to determine which generation is newer. If this is omitted in a request, then you are requesting the live object. See http://go/bigstore-versions""" -class GetAgentEngineRuntimeRevisionConfig(_common.BaseModel): - """Config for getting an Agent Engine Runtime Revision.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) +GdataObjectIdOrDict = Union[GdataObjectId, GdataObjectIdDict] -class GetAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Runtime Revision.""" +class GdataBlobstore2Info(_common.BaseModel): + """Information to read/write to blobstore2.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + blob_id: Optional[str] = Field( + default=None, + description="""The blob id, e.g., /blobstore/prod/playground/scotty""", + ) + blob_generation: Optional[int] = Field( + default=None, description="""The blob generation id.""" + ) + read_token: Optional[str] = Field( + default=None, + description="""The blob read token. Needed to read blobs that have not been replicated. Might not be available until the final call.""", + ) + upload_metadata_container: Optional[bytes] = Field( + default=None, + description="""Metadata passed from Blobstore -> Scotty for a new GCS upload. This is a signed, serialized blobstore2.BlobMetadataContainer proto which must never be consumed outside of Bigstore, and is not applicable to non-GCS media uploads.""", + ) + download_read_handle: Optional[bytes] = Field( + default=None, + description="""Read handle passed from Bigstore -> Scotty for a GCS download. This is a signed, serialized blobstore2.ReadHandle proto which must never be set outside of Bigstore, and is not applicable to non-GCS media downloads.""", + ) + download_external_read_token: Optional[bytes] = Field( + default=None, + description="""A serialized External Read Token passed from Bigstore -> Scotty for a GCS download. This field must never be consumed outside of Bigstore, and is not applicable to non-GCS media uploads.""", + ) + upload_fragment_list_creation_info: Optional[bytes] = Field( + default=None, + description="""A serialized Object Fragment List Creation Info passed from Bigstore -> Scotty for a GCS upload. This field must never be consumed outside of Bigstore, and is not applicable to non-GCS media uploads.""", + ) -GetAgentEngineRuntimeRevisionConfigOrDict = Union[ - GetAgentEngineRuntimeRevisionConfig, GetAgentEngineRuntimeRevisionConfigDict -] +class GdataBlobstore2InfoDict(TypedDict, total=False): + """Information to read/write to blobstore2.""" + blob_id: Optional[str] + """The blob id, e.g., /blobstore/prod/playground/scotty""" -class _GetAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine runtime revision.""" + blob_generation: Optional[int] + """The blob generation id.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine runtime revision.""" - ) - config: Optional[GetAgentEngineRuntimeRevisionConfig] = Field( - default=None, description="""""" - ) + read_token: Optional[str] + """The blob read token. Needed to read blobs that have not been replicated. Might not be available until the final call.""" + upload_metadata_container: Optional[bytes] + """Metadata passed from Blobstore -> Scotty for a new GCS upload. This is a signed, serialized blobstore2.BlobMetadataContainer proto which must never be consumed outside of Bigstore, and is not applicable to non-GCS media uploads.""" -class _GetAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine runtime revision.""" + download_read_handle: Optional[bytes] + """Read handle passed from Bigstore -> Scotty for a GCS download. This is a signed, serialized blobstore2.ReadHandle proto which must never be set outside of Bigstore, and is not applicable to non-GCS media downloads.""" - name: Optional[str] - """Name of the agent engine runtime revision.""" + download_external_read_token: Optional[bytes] + """A serialized External Read Token passed from Bigstore -> Scotty for a GCS download. This field must never be consumed outside of Bigstore, and is not applicable to non-GCS media uploads.""" - config: Optional[GetAgentEngineRuntimeRevisionConfigDict] - """""" + upload_fragment_list_creation_info: Optional[bytes] + """A serialized Object Fragment List Creation Info passed from Bigstore -> Scotty for a GCS upload. This field must never be consumed outside of Bigstore, and is not applicable to non-GCS media uploads.""" -_GetAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ - _GetAgentEngineRuntimeRevisionRequestParameters, - _GetAgentEngineRuntimeRevisionRequestParametersDict, -] +GdataBlobstore2InfoOrDict = Union[GdataBlobstore2Info, GdataBlobstore2InfoDict] -class ReasoningEngineRuntimeRevision(_common.BaseModel): - """A runtime revision.""" +class GdataCompositeMedia(_common.BaseModel): + """A sequence of media data references representing composite data. Introduced to support Bigstore composite objects. For details, visit http://go/bigstore-composites.""" - create_time: Optional[datetime.datetime] = Field( + length: Optional[int] = Field( + default=None, description="""Size of the data, in bytes""" + ) + reference_type: Optional[ + Literal["PATH", "BLOB_REF", "INLINE", "BIGSTORE_REF", "COSMO_BINARY_REFERENCE"] + ] = Field( + default=None, description="""Describes what the field reference contains.""" + ) + path: Optional[str] = Field( + default=None, description="""Path to the data, set if reference_type is PATH""" + ) + blob_ref: Optional[bytes] = Field( default=None, - description="""Output only. Timestamp when this ReasoningEngineRuntimeRevision was created.""", + description="""Blobstore v1 reference, set if reference_type is BLOBSTORE_REF This should be the byte representation of a blobstore.BlobRef. Since Blobstore is deprecating v1, use blobstore2_info instead. For now, any v2 blob will also be represented in this field as v1 BlobRef.""", ) - name: Optional[str] = Field( + inline: Optional[bytes] = Field( + default=None, description="""Media data, set if reference_type is INLINE""" + ) + object_id: Optional[GdataObjectId] = Field( default=None, - description="""Identifier. The resource name of the ReasoningEngineRuntimeRevision. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}/runtimeRevisions/{runtime_revision}`""", + description="""Reference to a TI Blob, set if reference_type is BIGSTORE_REF.""", ) - spec: Optional[ReasoningEngineSpec] = Field( + blobstore2_info: Optional[GdataBlobstore2Info] = Field( default=None, - description="""Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""", + description="""Blobstore v2 info, set if reference_type is BLOBSTORE_REF and it refers to a v2 blob.""", ) - state: Optional[State] = Field( - default=None, description="""Output only. The state of the revision.""" + cosmo_binary_reference: Optional[bytes] = Field( + default=None, + description="""A binary data reference for a media download. Serves as a technology-agnostic binary reference in some Google infrastructure. This value is a serialized storage_cosmo.BinaryReference proto. Storing it as bytes is a hack to get around the fact that the cosmo proto (as well as others it includes) doesn't support JavaScript. This prevents us from including the actual type of this field.""", + ) + crc32c_hash: Optional[int] = Field( + default=None, description="""crc32.c hash for the payload.""" + ) + md5_hash: Optional[bytes] = Field( + default=None, description="""MD5 hash for the payload.""" + ) + sha1_hash: Optional[bytes] = Field( + default=None, description="""SHA-1 hash for the payload.""" ) -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}`""" - - spec: Optional[ReasoningEngineSpecDict] - """Immutable. Configurations of the ReasoningEngineRuntimeRevision. Contains only revision specific fields.""" - - state: Optional[State] - """Output only. The state of the revision.""" +class GdataCompositeMediaDict(TypedDict, total=False): + """A sequence of media data references representing composite data. Introduced to support Bigstore composite objects. For details, visit http://go/bigstore-composites.""" + length: Optional[int] + """Size of the data, in bytes""" -ReasoningEngineRuntimeRevisionOrDict = Union[ - ReasoningEngineRuntimeRevision, ReasoningEngineRuntimeRevisionDict -] + reference_type: Optional[ + Literal["PATH", "BLOB_REF", "INLINE", "BIGSTORE_REF", "COSMO_BINARY_REFERENCE"] + ] + """Describes what the field reference contains.""" + path: Optional[str] + """Path to the data, set if reference_type is PATH""" -class ListAgentEngineRuntimeRevisionsConfig(_common.BaseModel): - """Config for listing reasoning engine runtime revisions.""" + blob_ref: Optional[bytes] + """Blobstore v1 reference, set if reference_type is BLOBSTORE_REF This should be the byte representation of a blobstore.BlobRef. Since Blobstore is deprecating v1, use blobstore2_info instead. For now, any v2 blob will also be represented in this field as v1 BlobRef.""" - 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.""", - ) + inline: Optional[bytes] + """Media data, set if reference_type is INLINE""" + object_id: Optional[GdataObjectIdDict] + """Reference to a TI Blob, set if reference_type is BIGSTORE_REF.""" -class ListAgentEngineRuntimeRevisionsConfigDict(TypedDict, total=False): - """Config for listing reasoning engine runtime revisions.""" + blobstore2_info: Optional[GdataBlobstore2InfoDict] + """Blobstore v2 info, set if reference_type is BLOBSTORE_REF and it refers to a v2 blob.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + cosmo_binary_reference: Optional[bytes] + """A binary data reference for a media download. Serves as a technology-agnostic binary reference in some Google infrastructure. This value is a serialized storage_cosmo.BinaryReference proto. Storing it as bytes is a hack to get around the fact that the cosmo proto (as well as others it includes) doesn't support JavaScript. This prevents us from including the actual type of this field.""" - page_size: Optional[int] - """""" + crc32c_hash: Optional[int] + """crc32.c hash for the payload.""" - page_token: Optional[str] - """""" + md5_hash: Optional[bytes] + """MD5 hash for the payload.""" - filter: Optional[str] - """An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported.""" + sha1_hash: Optional[bytes] + """SHA-1 hash for the payload.""" -ListAgentEngineRuntimeRevisionsConfigOrDict = Union[ - ListAgentEngineRuntimeRevisionsConfig, ListAgentEngineRuntimeRevisionsConfigDict -] +GdataCompositeMediaOrDict = Union[GdataCompositeMedia, GdataCompositeMediaDict] -class _ListAgentEngineRuntimeRevisionsRequestParameters(_common.BaseModel): - """Parameters for listing reasoning engine runtime revisions.""" +class GdataDiffVersionResponse(_common.BaseModel): + """Backend response for a Diff get version response. For details on the Scotty Diff protocol, visit http://go/scotty-diff-protocol.""" - name: Optional[str] = Field( - default=None, description="""Name of the reasoning engine.""" + object_version: Optional[str] = Field( + default=None, description="""The version of the object stored at the server.""" ) - config: Optional[ListAgentEngineRuntimeRevisionsConfig] = Field( - default=None, description="""""" + object_size_bytes: Optional[int] = Field( + default=None, description="""The total size of the server object.""" ) -class _ListAgentEngineRuntimeRevisionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing reasoning engine runtime revisions.""" +class GdataDiffVersionResponseDict(TypedDict, total=False): + """Backend response for a Diff get version response. For details on the Scotty Diff protocol, visit http://go/scotty-diff-protocol.""" - name: Optional[str] - """Name of the reasoning engine.""" + object_version: Optional[str] + """The version of the object stored at the server.""" - config: Optional[ListAgentEngineRuntimeRevisionsConfigDict] - """""" + object_size_bytes: Optional[int] + """The total size of the server object.""" -_ListAgentEngineRuntimeRevisionsRequestParametersOrDict = Union[ - _ListAgentEngineRuntimeRevisionsRequestParameters, - _ListAgentEngineRuntimeRevisionsRequestParametersDict, +GdataDiffVersionResponseOrDict = Union[ + GdataDiffVersionResponse, GdataDiffVersionResponseDict ] -class ListReasoningEnginesRuntimeRevisionsResponse(_common.BaseModel): - """Response for listing agent engine runtime revisions.""" +class GdataDiffChecksumsResponse(_common.BaseModel): + """Backend response for a Diff get checksums response. For details on the Scotty Diff protocol, visit http://go/scotty-diff-protocol.""" - sdk_http_response: Optional[genai_types.HttpResponse] = Field( - default=None, description="""Used to retain the full HTTP response.""" + object_version: Optional[str] = Field( + default=None, + description="""The object version of the object the checksums are being returned for.""", ) - 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.""" + object_size_bytes: Optional[int] = Field( + default=None, description="""The total size of the server object.""" + ) + chunk_size_bytes: Optional[int] = Field( + default=None, + description="""The chunk size of checksums. Must be a multiple of 256KB.""", + ) + checksums_location: Optional[GdataCompositeMedia] = Field( + default=None, + description="""Exactly one of these fields must be populated. If checksums_location is filled, the server will return the corresponding contents to the user. If object_location is filled, the server will calculate the checksums based on the content there and return that to the user. For details on the format of the checksums, see http://go/scotty-diff-protocol.""", + ) + object_location: Optional[GdataCompositeMedia] = Field( + default=None, + description="""If set, calculate the checksums based on the contents and return them to the caller.""", ) -class ListReasoningEnginesRuntimeRevisionsResponseDict(TypedDict, total=False): - """Response for listing agent engine runtime revisions.""" +class GdataDiffChecksumsResponseDict(TypedDict, total=False): + """Backend response for a Diff get checksums response. For details on the Scotty Diff protocol, visit http://go/scotty-diff-protocol.""" - sdk_http_response: Optional[genai_types.HttpResponse] - """Used to retain the full HTTP response.""" + object_version: Optional[str] + """The object version of the object the checksums are being returned for.""" - next_page_token: Optional[str] - """""" + object_size_bytes: Optional[int] + """The total size of the server object.""" - reasoning_engine_runtime_revisions: Optional[ - list[ReasoningEngineRuntimeRevisionDict] - ] - """List of reasoning engine runtime revisions.""" + chunk_size_bytes: Optional[int] + """The chunk size of checksums. Must be a multiple of 256KB.""" + checksums_location: Optional[GdataCompositeMediaDict] + """Exactly one of these fields must be populated. If checksums_location is filled, the server will return the corresponding contents to the user. If object_location is filled, the server will calculate the checksums based on the content there and return that to the user. For details on the format of the checksums, see http://go/scotty-diff-protocol.""" -ListReasoningEnginesRuntimeRevisionsResponseOrDict = Union[ - ListReasoningEnginesRuntimeRevisionsResponse, - ListReasoningEnginesRuntimeRevisionsResponseDict, + object_location: Optional[GdataCompositeMediaDict] + """If set, calculate the checksums based on the contents and return them to the caller.""" + + +GdataDiffChecksumsResponseOrDict = Union[ + GdataDiffChecksumsResponse, GdataDiffChecksumsResponseDict ] -class DeleteAgentEngineRuntimeRevisionConfig(_common.BaseModel): - """Config for deleting an Agent Engine Runtime Revision.""" +class GdataDiffDownloadResponse(_common.BaseModel): + """Backend response for a Diff download response. For details on the Scotty Diff protocol, visit http://go/scotty-diff-protocol.""" - 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.""", + object_location: Optional[GdataCompositeMedia] = Field( + default=None, description="""The original object location.""" ) -class DeleteAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Runtime Revision.""" - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" +class GdataDiffDownloadResponseDict(TypedDict, total=False): + """Backend response for a Diff download response. For details on the Scotty Diff protocol, visit http://go/scotty-diff-protocol.""" - wait_for_completion: Optional[bool] - """Waits for the operation to complete before returning.""" + object_location: Optional[GdataCompositeMediaDict] + """The original object location.""" -DeleteAgentEngineRuntimeRevisionConfigOrDict = Union[ - DeleteAgentEngineRuntimeRevisionConfig, DeleteAgentEngineRuntimeRevisionConfigDict +GdataDiffDownloadResponseOrDict = Union[ + GdataDiffDownloadResponse, GdataDiffDownloadResponseDict ] -class _DeleteAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): - """Parameters for deleting agent engine runtime revisions.""" +class GdataDiffUploadRequest(_common.BaseModel): + """A Diff upload request. For details on the Scotty Diff protocol, visit http://go/scotty-diff-protocol.""" - name: Optional[str] = Field( + object_version: Optional[str] = Field( default=None, - description="""Name of the agent engine runtime revision to delete.""", + description="""The object version of the object that is the base version the incoming diff script will be applied to. This field will always be filled in.""", ) - config: Optional[DeleteAgentEngineRuntimeRevisionConfig] = Field( - default=None, description="""""" + object_info: Optional[GdataCompositeMedia] = Field( + default=None, + description="""The location of the new object. Agents must clone the object located here, as the upload server will delete the contents once a response is received.""", + ) + checksums_info: Optional[GdataCompositeMedia] = Field( + default=None, + description="""The location of the checksums for the new object. Agents must clone the object located here, as the upload server will delete the contents once a response is received. For details on the format of the checksums, see http://go/scotty-diff-protocol.""", ) -class _DeleteAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engine runtime revisions.""" +class GdataDiffUploadRequestDict(TypedDict, total=False): + """A Diff upload request. For details on the Scotty Diff protocol, visit http://go/scotty-diff-protocol.""" - name: Optional[str] - """Name of the agent engine runtime revision to delete.""" + object_version: Optional[str] + """The object version of the object that is the base version the incoming diff script will be applied to. This field will always be filled in.""" - config: Optional[DeleteAgentEngineRuntimeRevisionConfigDict] - """""" + object_info: Optional[GdataCompositeMediaDict] + """The location of the new object. Agents must clone the object located here, as the upload server will delete the contents once a response is received.""" + checksums_info: Optional[GdataCompositeMediaDict] + """The location of the checksums for the new object. Agents must clone the object located here, as the upload server will delete the contents once a response is received. For details on the format of the checksums, see http://go/scotty-diff-protocol.""" -_DeleteAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ - _DeleteAgentEngineRuntimeRevisionRequestParameters, - _DeleteAgentEngineRuntimeRevisionRequestParametersDict, -] +GdataDiffUploadRequestOrDict = Union[GdataDiffUploadRequest, GdataDiffUploadRequestDict] -class DeleteAgentEngineRuntimeRevisionOperation(_common.BaseModel): - """Operation for deleting agent engine 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( +class GdataDiffUploadResponse(_common.BaseModel): + """Backend response for a Diff upload request. For details on the Scotty Diff protocol, visit http://go/scotty-diff-protocol.""" + + object_version: 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 object version of the object at the server. Must be included in the end notification response. The version in the end notification response must correspond to the new version of the object that is now stored at the server, after the upload.""", ) - error: Optional[dict[str, Any]] = Field( + original_object: Optional[GdataCompositeMedia] = Field( default=None, - description="""The error result of the operation in case of failure or cancellation.""", + description="""The location of the original file for a diff upload request. Must be filled in if responding to an upload start notification.""", ) -class DeleteAgentEngineRuntimeRevisionOperationDict(TypedDict, total=False): - """Operation for deleting agent engine 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.""" +class GdataDiffUploadResponseDict(TypedDict, total=False): + """Backend response for a Diff upload request. For details on the Scotty Diff protocol, visit http://go/scotty-diff-protocol.""" - 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.""" + object_version: Optional[str] + """The object version of the object at the server. Must be included in the end notification response. The version in the end notification response must correspond to the new version of the object that is now stored at the server, after the upload.""" - error: Optional[dict[str, Any]] - """The error result of the operation in case of failure or cancellation.""" + original_object: Optional[GdataCompositeMediaDict] + """The location of the original file for a diff upload request. Must be filled in if responding to an upload start notification.""" -DeleteAgentEngineRuntimeRevisionOperationOrDict = Union[ - DeleteAgentEngineRuntimeRevisionOperation, - DeleteAgentEngineRuntimeRevisionOperationDict, +GdataDiffUploadResponseOrDict = Union[ + GdataDiffUploadResponse, GdataDiffUploadResponseDict ] -class GetDeleteAgentEngineRuntimeRevisionOperationConfig(_common.BaseModel): +class GdataContentTypeInfo(_common.BaseModel): + """Detailed Content-Type information from Scotty. The Content-Type of the media will typically be filled in by the header or Scotty's best_guess, but this extended information provides the backend with more information so that it can make a better decision if needed. This is only used on media upload requests from Scotty.""" - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" + best_guess: Optional[str] = Field( + default=None, + description="""Scotty's best guess of what the content type of the file is.""", + ) + from_header: Optional[str] = Field( + default=None, + description="""The content type of the file as specified in the request headers, multipart headers, or RUPIO start request.""", + ) + from_file_name: Optional[str] = Field( + default=None, + description="""The content type of the file derived from the file extension of the original file name used by the client.""", + ) + from_url_path: Optional[str] = Field( + default=None, + description="""The content type of the file derived from the file extension of the URL path. The URL path is assumed to represent a file name (which is typically only true for agents that are providing a REST API).""", + ) + from_bytes: Optional[str] = Field( + default=None, + description="""The content type of the file derived by looking at specific bytes (i.e. "magic bytes") of the actual file.""", + ) + from_fusion_id: Optional[str] = Field( + default=None, + description="""The content type of the file detected by Fusion ID. go/fusionid""", + ) + fusion_id_detection_metadata: Optional[bytes] = Field( + default=None, + description="""Metadata information from Fusion ID detection. Serialized FusionIdDetectionMetadata proto. Only set if from_fusion_id is set.""", ) -class GetDeleteAgentEngineRuntimeRevisionOperationConfigDict(TypedDict, total=False): - - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - -GetDeleteAgentEngineRuntimeRevisionOperationConfigOrDict = Union[ - GetDeleteAgentEngineRuntimeRevisionOperationConfig, - GetDeleteAgentEngineRuntimeRevisionOperationConfigDict, -] +class GdataContentTypeInfoDict(TypedDict, total=False): + """Detailed Content-Type information from Scotty. The Content-Type of the media will typically be filled in by the header or Scotty's best_guess, but this extended information provides the backend with more information so that it can make a better decision if needed. This is only used on media upload requests from Scotty.""" + best_guess: Optional[str] + """Scotty's best guess of what the content type of the file is.""" -class _GetDeleteAgentEngineRuntimeRevisionOperationParameters(_common.BaseModel): - """Parameters for getting an operation that deletes a agent engine runtime revision.""" + from_header: Optional[str] + """The content type of the file as specified in the request headers, multipart headers, or RUPIO start request.""" - operation_name: Optional[str] = Field( - default=None, description="""The server-assigned name for the operation.""" - ) - config: Optional[GetDeleteAgentEngineRuntimeRevisionOperationConfig] = Field( - default=None, description="""Used to override the default configuration.""" - ) + from_file_name: Optional[str] + """The content type of the file derived from the file extension of the original file name used by the client.""" + from_url_path: Optional[str] + """The content type of the file derived from the file extension of the URL path. The URL path is assumed to represent a file name (which is typically only true for agents that are providing a REST API).""" -class _GetDeleteAgentEngineRuntimeRevisionOperationParametersDict( - TypedDict, total=False -): - """Parameters for getting an operation that deletes a agent engine runtime revision.""" + from_bytes: Optional[str] + """The content type of the file derived by looking at specific bytes (i.e. "magic bytes") of the actual file.""" - operation_name: Optional[str] - """The server-assigned name for the operation.""" + from_fusion_id: Optional[str] + """The content type of the file detected by Fusion ID. go/fusionid""" - config: Optional[GetDeleteAgentEngineRuntimeRevisionOperationConfigDict] - """Used to override the default configuration.""" + fusion_id_detection_metadata: Optional[bytes] + """Metadata information from Fusion ID detection. Serialized FusionIdDetectionMetadata proto. Only set if from_fusion_id is set.""" -_GetDeleteAgentEngineRuntimeRevisionOperationParametersOrDict = Union[ - _GetDeleteAgentEngineRuntimeRevisionOperationParameters, - _GetDeleteAgentEngineRuntimeRevisionOperationParametersDict, -] +GdataContentTypeInfoOrDict = Union[GdataContentTypeInfo, GdataContentTypeInfoDict] -class QueryAgentEngineRuntimeRevisionConfig(_common.BaseModel): - """Config for querying agent engine runtime revisions.""" +class GdataDownloadParameters(_common.BaseModel): + """Parameters specific to media downloads.""" - 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.""" + allow_gzip_compression: Optional[bool] = Field( + default=None, + description="""A boolean to be returned in the response to Scotty. Allows/disallows gzip encoding of the payload content when the server thinks it's advantageous (hence, does not guarantee compression) which allows Scotty to GZip the response to the client.""", ) - input: Optional[dict[str, Any]] = Field( - default=None, description="""The input to the class method.""" + ignore_range: Optional[bool] = Field( + default=None, + description="""Determining whether or not Apiary should skip the inclusion of any Content-Range header on its response to Scotty.""", ) - include_all_fields: Optional[bool] = Field(default=False, description="""""") - -class QueryAgentEngineRuntimeRevisionConfigDict(TypedDict, total=False): - """Config for querying agent engine runtime revisions.""" - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - - class_method: Optional[str] - """The class method to call.""" +class GdataDownloadParametersDict(TypedDict, total=False): + """Parameters specific to media downloads.""" - input: Optional[dict[str, Any]] - """The input to the class method.""" + allow_gzip_compression: Optional[bool] + """A boolean to be returned in the response to Scotty. Allows/disallows gzip encoding of the payload content when the server thinks it's advantageous (hence, does not guarantee compression) which allows Scotty to GZip the response to the client.""" - include_all_fields: Optional[bool] - """""" + ignore_range: Optional[bool] + """Determining whether or not Apiary should skip the inclusion of any Content-Range header on its response to Scotty.""" -QueryAgentEngineRuntimeRevisionConfigOrDict = Union[ - QueryAgentEngineRuntimeRevisionConfig, QueryAgentEngineRuntimeRevisionConfigDict +GdataDownloadParametersOrDict = Union[ + GdataDownloadParameters, GdataDownloadParametersDict ] -class _QueryAgentEngineRuntimeRevisionRequestParameters(_common.BaseModel): - """Parameters for querying agent engine runtime revisions.""" +class GdataMedia(_common.BaseModel): + """A reference to data stored on the filesystem, on GFS or in blobstore.""" - name: Optional[str] = Field( - default=None, description="""Name of the agent engine runtime revision.""" + content_type: Optional[str] = Field( + default=None, description="""MIME type of the data""" ) - config: Optional[QueryAgentEngineRuntimeRevisionConfig] = Field( - default=None, description="""""" + timestamp: Optional[int] = Field( + default=None, + description="""Time at which the media data was last updated, in milliseconds since UNIX epoch""", + ) + token: Optional[str] = Field( + default=None, + description="""A unique fingerprint/version id for the media data""", + ) + length: Optional[int] = Field( + default=None, description="""Size of the data, in bytes""" + ) + filename: Optional[str] = Field(default=None, description="""Original file name""") + reference_type: Optional[ + Literal[ + "PATH", + "BLOB_REF", + "INLINE", + "GET_MEDIA", + "COMPOSITE_MEDIA", + "BIGSTORE_REF", + "DIFF_VERSION_RESPONSE", + "DIFF_CHECKSUMS_RESPONSE", + "DIFF_DOWNLOAD_RESPONSE", + "DIFF_UPLOAD_REQUEST", + "DIFF_UPLOAD_RESPONSE", + "COSMO_BINARY_REFERENCE", + "ARBITRARY_BYTES", + ] + ] = Field( + default=None, description="""Describes what the field reference contains.""" + ) + path: Optional[str] = Field( + default=None, description="""Path to the data, set if reference_type is PATH""" + ) + blob_ref: Optional[bytes] = Field( + default=None, + description="""Blobstore v1 reference, set if reference_type is BLOBSTORE_REF This should be the byte representation of a blobstore.BlobRef. Since Blobstore is deprecating v1, use blobstore2_info instead. For now, any v2 blob will also be represented in this field as v1 BlobRef.""", + ) + inline: Optional[bytes] = Field( + default=None, description="""Media data, set if reference_type is INLINE""" + ) + media_id: Optional[bytes] = Field( + default=None, + description="""Media id to forward to the operation GetMedia. Can be set if reference_type is GET_MEDIA.""", + ) + hash: Optional[str] = Field( + default=None, + description="""Deprecated, use one of explicit hash type fields instead. These two hash related fields will only be populated on Scotty based media uploads and will contain the content of the hash group in the NotificationRequest: http://cs/#google3/blobstore2/api/scotty/service/proto/upload_listener.proto&q=class:Hash Hex encoded hash value of the uploaded media.""", + ) + algorithm: Optional[str] = Field( + default=None, + description="""Deprecated, use one of explicit hash type fields instead. Algorithm used for calculating the hash. As of 2011/01/21, "MD5" is the only possible value for this field. New values may be added at any time.""", + ) + composite_media: Optional[list[GdataCompositeMedia]] = Field( + default=None, + description="""A composite media composed of one or more media objects, set if reference_type is COMPOSITE_MEDIA. The media length field must be set to the sum of the lengths of all composite media objects. Note: All composite media must have length specified.""", + ) + bigstore_object_ref: Optional[bytes] = Field( + default=None, description="""Use object_id instead.""" + ) + object_id: Optional[GdataObjectId] = Field( + default=None, + description="""Reference to a TI Blob, set if reference_type is BIGSTORE_REF.""", + ) + blobstore2_info: Optional[GdataBlobstore2Info] = Field( + default=None, + description="""Blobstore v2 info, set if reference_type is BLOBSTORE_REF and it refers to a v2 blob.""", + ) + diff_version_response: Optional[GdataDiffVersionResponse] = Field( + default=None, description="""Set if reference_type is DIFF_VERSION_RESPONSE.""" + ) + diff_checksums_response: Optional[GdataDiffChecksumsResponse] = Field( + default=None, + description="""Set if reference_type is DIFF_CHECKSUMS_RESPONSE.""", + ) + diff_download_response: Optional[GdataDiffDownloadResponse] = Field( + default=None, description="""Set if reference_type is DIFF_DOWNLOAD_RESPONSE.""" + ) + diff_upload_request: Optional[GdataDiffUploadRequest] = Field( + default=None, description="""Set if reference_type is DIFF_UPLOAD_REQUEST.""" + ) + diff_upload_response: Optional[GdataDiffUploadResponse] = Field( + default=None, description="""Set if reference_type is DIFF_UPLOAD_RESPONSE.""" + ) + content_type_info: Optional[GdataContentTypeInfo] = Field( + default=None, + description="""Extended content type information provided for Scotty uploads.""", + ) + download_parameters: Optional[GdataDownloadParameters] = Field( + default=None, description="""Parameters for a media download.""" + ) + crc32c_hash: Optional[int] = Field( + default=None, + description="""For Scotty Uploads: Scotty-provided hashes for uploads For Scotty Downloads: (WARNING: DO NOT USE WITHOUT PERMISSION FROM THE SCOTTY TEAM.) A Hash provided by the agent to be used to verify the data being downloaded. Currently only supported for inline payloads. Further, only crc32c_hash is currently supported.""", + ) + md5_hash: Optional[bytes] = Field( + default=None, description="""Scotty-provided MD5 hash for an upload.""" + ) + sha1_hash: Optional[bytes] = Field( + default=None, description="""Scotty-provided SHA1 hash for an upload.""" + ) + sha256_hash: Optional[bytes] = Field( + default=None, description="""Scotty-provided SHA256 hash for an upload.""" + ) + sha512_hash: Optional[bytes] = Field( + default=None, description="""Scotty-provided SHA512 hash for an upload.""" + ) + is_potential_retry: Optional[bool] = Field( + default=None, + description="""|is_potential_retry| is set false only when Scotty is certain that it has not sent the request before. When a client resumes an upload, this field must be set true in agent calls, because Scotty cannot be certain that it has never sent the request before due to potential failure in the session state persistence.""", + ) + cosmo_binary_reference: Optional[bytes] = Field( + default=None, + description="""A binary data reference for a media download. Serves as a technology-agnostic binary reference in some Google infrastructure. This value is a serialized storage_cosmo.BinaryReference proto. Storing it as bytes is a hack to get around the fact that the cosmo proto (as well as others it includes) doesn't support JavaScript. This prevents us from including the actual type of this field.""", + ) + hash_verified: Optional[bool] = Field( + default=None, + description="""For Scotty uploads only. If a user sends a hash code and the backend has requested that Scotty verify the upload against the client hash, Scotty will perform the check on behalf of the backend and will reject it if the hashes don't match. This is set to true if Scotty performed this verification.""", ) -class _QueryAgentEngineRuntimeRevisionRequestParametersDict(TypedDict, total=False): - """Parameters for querying agent engine runtime revisions.""" - - name: Optional[str] - """Name of the agent engine runtime revision.""" - - config: Optional[QueryAgentEngineRuntimeRevisionConfigDict] - """""" - +class GdataMediaDict(TypedDict, total=False): + """A reference to data stored on the filesystem, on GFS or in blobstore.""" + + content_type: Optional[str] + """MIME type of the data""" + + timestamp: Optional[int] + """Time at which the media data was last updated, in milliseconds since UNIX epoch""" + + token: Optional[str] + """A unique fingerprint/version id for the media data""" + + length: Optional[int] + """Size of the data, in bytes""" + + filename: Optional[str] + """Original file name""" + + reference_type: Optional[ + Literal[ + "PATH", + "BLOB_REF", + "INLINE", + "GET_MEDIA", + "COMPOSITE_MEDIA", + "BIGSTORE_REF", + "DIFF_VERSION_RESPONSE", + "DIFF_CHECKSUMS_RESPONSE", + "DIFF_DOWNLOAD_RESPONSE", + "DIFF_UPLOAD_REQUEST", + "DIFF_UPLOAD_RESPONSE", + "COSMO_BINARY_REFERENCE", + "ARBITRARY_BYTES", + ] + ] + """Describes what the field reference contains.""" + + path: Optional[str] + """Path to the data, set if reference_type is PATH""" + + blob_ref: Optional[bytes] + """Blobstore v1 reference, set if reference_type is BLOBSTORE_REF This should be the byte representation of a blobstore.BlobRef. Since Blobstore is deprecating v1, use blobstore2_info instead. For now, any v2 blob will also be represented in this field as v1 BlobRef.""" + + inline: Optional[bytes] + """Media data, set if reference_type is INLINE""" + + media_id: Optional[bytes] + """Media id to forward to the operation GetMedia. Can be set if reference_type is GET_MEDIA.""" + + hash: Optional[str] + """Deprecated, use one of explicit hash type fields instead. These two hash related fields will only be populated on Scotty based media uploads and will contain the content of the hash group in the NotificationRequest: http://cs/#google3/blobstore2/api/scotty/service/proto/upload_listener.proto&q=class:Hash Hex encoded hash value of the uploaded media.""" + + algorithm: Optional[str] + """Deprecated, use one of explicit hash type fields instead. Algorithm used for calculating the hash. As of 2011/01/21, "MD5" is the only possible value for this field. New values may be added at any time.""" + + composite_media: Optional[list[GdataCompositeMediaDict]] + """A composite media composed of one or more media objects, set if reference_type is COMPOSITE_MEDIA. The media length field must be set to the sum of the lengths of all composite media objects. Note: All composite media must have length specified.""" + + bigstore_object_ref: Optional[bytes] + """Use object_id instead.""" + + object_id: Optional[GdataObjectIdDict] + """Reference to a TI Blob, set if reference_type is BIGSTORE_REF.""" + + blobstore2_info: Optional[GdataBlobstore2InfoDict] + """Blobstore v2 info, set if reference_type is BLOBSTORE_REF and it refers to a v2 blob.""" + + diff_version_response: Optional[GdataDiffVersionResponseDict] + """Set if reference_type is DIFF_VERSION_RESPONSE.""" + + diff_checksums_response: Optional[GdataDiffChecksumsResponseDict] + """Set if reference_type is DIFF_CHECKSUMS_RESPONSE.""" + + diff_download_response: Optional[GdataDiffDownloadResponseDict] + """Set if reference_type is DIFF_DOWNLOAD_RESPONSE.""" + + diff_upload_request: Optional[GdataDiffUploadRequestDict] + """Set if reference_type is DIFF_UPLOAD_REQUEST.""" + + diff_upload_response: Optional[GdataDiffUploadResponseDict] + """Set if reference_type is DIFF_UPLOAD_RESPONSE.""" + + content_type_info: Optional[GdataContentTypeInfoDict] + """Extended content type information provided for Scotty uploads.""" + + download_parameters: Optional[GdataDownloadParametersDict] + """Parameters for a media download.""" + + crc32c_hash: Optional[int] + """For Scotty Uploads: Scotty-provided hashes for uploads For Scotty Downloads: (WARNING: DO NOT USE WITHOUT PERMISSION FROM THE SCOTTY TEAM.) A Hash provided by the agent to be used to verify the data being downloaded. Currently only supported for inline payloads. Further, only crc32c_hash is currently supported.""" + + md5_hash: Optional[bytes] + """Scotty-provided MD5 hash for an upload.""" + + sha1_hash: Optional[bytes] + """Scotty-provided SHA1 hash for an upload.""" + + sha256_hash: Optional[bytes] + """Scotty-provided SHA256 hash for an upload.""" + + sha512_hash: Optional[bytes] + """Scotty-provided SHA512 hash for an upload.""" + + is_potential_retry: Optional[bool] + """|is_potential_retry| is set false only when Scotty is certain that it has not sent the request before. When a client resumes an upload, this field must be set true in agent calls, because Scotty cannot be certain that it has never sent the request before due to potential failure in the session state persistence.""" + + cosmo_binary_reference: Optional[bytes] + """A binary data reference for a media download. Serves as a technology-agnostic binary reference in some Google infrastructure. This value is a serialized storage_cosmo.BinaryReference proto. Storing it as bytes is a hack to get around the fact that the cosmo proto (as well as others it includes) doesn't support JavaScript. This prevents us from including the actual type of this field.""" + + hash_verified: Optional[bool] + """For Scotty uploads only. If a user sends a hash code and the backend has requested that Scotty verify the upload against the client hash, Scotty will perform the check on behalf of the backend and will reject it if the hashes don't match. This is set to true if Scotty performed this verification.""" + + +GdataMediaOrDict = Union[GdataMedia, GdataMediaDict] + + +class ApiservingMediaResponseInfo(_common.BaseModel): + """This message is for backends to pass their scotty media specific fields to ESF. Backend will include this in their response message to ESF. Example: ExportFile is an rpc defined for upload using scotty from ESF. rpc ExportFile(ExportFileRequest) returns (ExportFileResponse) Message ExportFileResponse will include apiserving.MediaResponseInfo to tell ESF about data like dynamic_dropzone it needs to pass to Scotty. message ExportFileResponse { optional gdata.Media blob = 1; optional apiserving.MediaResponseInfo media_response_info = 2 }""" + + dynamic_dropzone: Optional[str] = Field( + default=None, + description="""Specifies the Scotty dropzone to use for uploads.""", + ) + custom_data: Optional[str] = Field( + default=None, + description="""Data to copy from backend response to the next backend requests. Custom data is returned to Scotty in the agent_state field, which Scotty will then provide in subsequent upload notifications.""", + ) + dynamic_drop_target: Optional[bytes] = Field( + default=None, + description="""Specifies the Scotty Drop Target to use for uploads. If present in a media response, Scotty does not upload to a standard drop zone. Instead, Scotty saves the upload directly to the location specified in this drop target. Unlike drop zones, the drop target is the final storage location for an upload. So, the agent does not need to clone the blob at the end of the upload. The agent is responsible for garbage collecting any orphaned blobs that may occur due to aborted uploads. For more information, see the drop target design doc here: http://goto/ScottyDropTarget This field will be preferred to dynamicDropzone. If provided, the identified field in the response must be of the type uploader.agent.DropTarget.""", + ) + data_storage_transform: Optional[bytes] = Field( + default=None, + description="""Specifies any transformation to be applied to data before persisting it or retrieving from storage. E.g., encryption options for blobstore2. This should be of the form uploader_service.DataStorageTransform.""", + ) + traffic_class_field: Optional[ + Literal[ + "BE1", + "AF1", + "AF2", + "AF3", + "AF4", + "NC1", + "NC0", + "BE0", + "LLQ", + "LLQ1", + "LLQ2", + "LLQ3", + ] + ] = Field( + default=None, + description="""Specifies the TrafficClass that Scotty should use for any RPCs to fetch the response bytes. Will override the traffic class GTOS of the incoming http request. This is a temporary field to facilitate whitelisting and experimentation by the bigstore agent only. For instance, this does not apply to RTMP reads. WARNING: DO NOT USE WITHOUT PERMISSION FROM THE SCOTTY TEAM.""", + ) + media_for_diff: Optional[GdataMedia] = Field( + default=None, + description="""Diff Updates must respond to a START notification with this Media proto to tell Scotty to decode the diff encoded payload and apply the diff against this field. If the request was diff encoded, but this field is not set, Scotty will treat the encoding as identity. This is corresponding to Apiary's DiffUploadResponse.original_object (//depot/google3/gdata/rosy/proto/data.proto?l=413). See go/esf-scotty-diff-upload for more information.""", + ) + verify_hash_from_header: Optional[bool] = Field( + default=None, + description="""Tells Scotty to verify hashes on the agent's behalf by parsing out the X-Goog-Hash header.""", + ) + request_class: Optional[UnknownR] = Field( + default=None, + description="""Request class to use for all Blobstore operations for this request.""", + ) + scotty_customer_log: Optional[bytes] = Field( + default=None, + description="""Customer-specific data to be recorded in the Scotty logs type is logs_proto_scotty.CustomerLog""", + ) + scotty_agent_user_id: Optional[int] = Field( + default=None, + description="""Requester ID passed along to be recorded in the Scotty logs""", + ) + original_object_blob_mint_index: Optional[int] = Field( + default=None, + description="""For the first notification of a |diff_encoded| HttpRequestInfo, this is the index of the blob mint that Scotty should use when reading the original blob. This field is optional. It's not required ever, even if `destination_blob_mint_index` is set. In situations like that, we will use the destination blob's mint for the destination blob and regular blob ACL checks for the original object. Note: This field is only for use by Drive API for diff uploads.""", + ) + destination_blob_mint_index: Optional[int] = Field( + default=None, + description="""For the first notification of a |diff_encoded| HttpRequestInfo, this is the index of the blob mint that Scotty should use when writing the resulting blob. This field is optional. It's not required ever, even if `original_object_blob_mint_index` is set. In situations like that, we will use the destination blob's mint for the destination blob and regular blob ACL checks for the original object. Note: This field is only for use by Drive API for diff uploads.""", + ) + + +class ApiservingMediaResponseInfoDict(TypedDict, total=False): + """This message is for backends to pass their scotty media specific fields to ESF. Backend will include this in their response message to ESF. Example: ExportFile is an rpc defined for upload using scotty from ESF. rpc ExportFile(ExportFileRequest) returns (ExportFileResponse) Message ExportFileResponse will include apiserving.MediaResponseInfo to tell ESF about data like dynamic_dropzone it needs to pass to Scotty. message ExportFileResponse { optional gdata.Media blob = 1; optional apiserving.MediaResponseInfo media_response_info = 2 }""" + + dynamic_dropzone: Optional[str] + """Specifies the Scotty dropzone to use for uploads.""" + + custom_data: Optional[str] + """Data to copy from backend response to the next backend requests. Custom data is returned to Scotty in the agent_state field, which Scotty will then provide in subsequent upload notifications.""" + + dynamic_drop_target: Optional[bytes] + """Specifies the Scotty Drop Target to use for uploads. If present in a media response, Scotty does not upload to a standard drop zone. Instead, Scotty saves the upload directly to the location specified in this drop target. Unlike drop zones, the drop target is the final storage location for an upload. So, the agent does not need to clone the blob at the end of the upload. The agent is responsible for garbage collecting any orphaned blobs that may occur due to aborted uploads. For more information, see the drop target design doc here: http://goto/ScottyDropTarget This field will be preferred to dynamicDropzone. If provided, the identified field in the response must be of the type uploader.agent.DropTarget.""" + + data_storage_transform: Optional[bytes] + """Specifies any transformation to be applied to data before persisting it or retrieving from storage. E.g., encryption options for blobstore2. This should be of the form uploader_service.DataStorageTransform.""" + + traffic_class_field: Optional[ + Literal[ + "BE1", + "AF1", + "AF2", + "AF3", + "AF4", + "NC1", + "NC0", + "BE0", + "LLQ", + "LLQ1", + "LLQ2", + "LLQ3", + ] + ] + """Specifies the TrafficClass that Scotty should use for any RPCs to fetch the response bytes. Will override the traffic class GTOS of the incoming http request. This is a temporary field to facilitate whitelisting and experimentation by the bigstore agent only. For instance, this does not apply to RTMP reads. WARNING: DO NOT USE WITHOUT PERMISSION FROM THE SCOTTY TEAM.""" + + media_for_diff: Optional[GdataMediaDict] + """Diff Updates must respond to a START notification with this Media proto to tell Scotty to decode the diff encoded payload and apply the diff against this field. If the request was diff encoded, but this field is not set, Scotty will treat the encoding as identity. This is corresponding to Apiary's DiffUploadResponse.original_object (//depot/google3/gdata/rosy/proto/data.proto?l=413). See go/esf-scotty-diff-upload for more information.""" + + verify_hash_from_header: Optional[bool] + """Tells Scotty to verify hashes on the agent's behalf by parsing out the X-Goog-Hash header.""" + + request_class: Optional[UnknownR] + """Request class to use for all Blobstore operations for this request.""" + + scotty_customer_log: Optional[bytes] + """Customer-specific data to be recorded in the Scotty logs type is logs_proto_scotty.CustomerLog""" -_QueryAgentEngineRuntimeRevisionRequestParametersOrDict = Union[ - _QueryAgentEngineRuntimeRevisionRequestParameters, - _QueryAgentEngineRuntimeRevisionRequestParametersDict, + scotty_agent_user_id: Optional[int] + """Requester ID passed along to be recorded in the Scotty logs""" + + original_object_blob_mint_index: Optional[int] + """For the first notification of a |diff_encoded| HttpRequestInfo, this is the index of the blob mint that Scotty should use when reading the original blob. This field is optional. It's not required ever, even if `destination_blob_mint_index` is set. In situations like that, we will use the destination blob's mint for the destination blob and regular blob ACL checks for the original object. Note: This field is only for use by Drive API for diff uploads.""" + + destination_blob_mint_index: Optional[int] + """For the first notification of a |diff_encoded| HttpRequestInfo, this is the index of the blob mint that Scotty should use when writing the resulting blob. This field is optional. It's not required ever, even if `original_object_blob_mint_index` is set. In situations like that, we will use the destination blob's mint for the destination blob and regular blob ACL checks for the original object. Note: This field is only for use by Drive API for diff uploads.""" + + +ApiservingMediaResponseInfoOrDict = Union[ + ApiservingMediaResponseInfo, ApiservingMediaResponseInfoDict ] +class UploadRagFileResponse(_common.BaseModel): + """Response for uploading a Rag File.""" + + error: Optional[genai_types.GoogleRpcStatus] = Field( + default=None, + description="""The error that occurred while processing the RagFile.""", + ) + rag_file: Optional[RagFile] = Field( + default=None, + description="""The RagFile that had been uploaded into the RagCorpus.""", + ) + media_response_info: Optional[ApiservingMediaResponseInfo] = Field( + default=None, + description="""Additional metadata information provided by Scotty Server. It will be propagated by Scotty Agent automatically.""", + ) + + +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.""" + + media_response_info: Optional[ApiservingMediaResponseInfoDict] + """Additional metadata information provided by Scotty Server. It will be propagated by Scotty Agent automatically.""" + + +UploadRagFileResponseOrDict = Union[UploadRagFileResponse, UploadRagFileResponseDict] + + class SandboxEnvironmentSpecCodeExecutionEnvironment(_common.BaseModel): """The code execution environment with customized settings.""" @@ -15712,13 +16448,17 @@ class SandboxEnvironmentSpecCodeExecutionEnvironmentDict(TypedDict, total=False) class SandboxEnvironmentSpecComputerUseEnvironment(_common.BaseModel): """The computer use environment with customized settings.""" - pass + image_uri: Optional[str] = Field( + default=None, + description="""Optional. The Artifact Registry Docker image URI (e.g., us-docker.dev.pkg/my-proj/my-repo/my-image:my-tag). When unspecified, the default Computer Use container image will be used.""", + ) class SandboxEnvironmentSpecComputerUseEnvironmentDict(TypedDict, total=False): """The computer use environment with customized settings.""" - pass + image_uri: Optional[str] + """Optional. The Artifact Registry Docker image URI (e.g., us-docker.dev.pkg/my-proj/my-repo/my-image:my-tag). When unspecified, the default Computer Use container image will be used.""" SandboxEnvironmentSpecComputerUseEnvironmentOrDict = Union[ @@ -15753,7 +16493,7 @@ class SandboxEnvironmentSpecDict(TypedDict, total=False): SandboxEnvironmentSpecOrDict = Union[SandboxEnvironmentSpec, SandboxEnvironmentSpecDict] -class CreateAgentEngineSandboxConfig(_common.BaseModel): +class CreateRuntimeSandboxConfig(_common.BaseModel): """Config for creating a Sandbox.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -15789,7 +16529,7 @@ class CreateAgentEngineSandboxConfig(_common.BaseModel): ) -class CreateAgentEngineSandboxConfigDict(TypedDict, total=False): +class CreateRuntimeSandboxConfigDict(TypedDict, total=False): """Config for creating a Sandbox.""" http_options: Optional[genai_types.HttpOptions] @@ -15819,42 +16559,41 @@ class CreateAgentEngineSandboxConfigDict(TypedDict, total=False): """Owner information for this sandbox environment. A sandbox can only be restored from a snapshot belonging to the same owner.""" -CreateAgentEngineSandboxConfigOrDict = Union[ - CreateAgentEngineSandboxConfig, CreateAgentEngineSandboxConfigDict +CreateRuntimeSandboxConfigOrDict = Union[ + CreateRuntimeSandboxConfig, CreateRuntimeSandboxConfigDict ] -class _CreateAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for creating Agent Engine Sandboxes.""" +class _CreateRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for creating Agent Runtime Sandboxes.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to create the sandbox under.""", + description="""Name of the Agent Runtime to create the sandbox under.""", ) spec: Optional[SandboxEnvironmentSpec] = Field( default=None, description="""The specification of the sandbox.""" ) - config: Optional[CreateAgentEngineSandboxConfig] = Field( + config: Optional[CreateRuntimeSandboxConfig] = Field( default=None, description="""""" ) -class _CreateAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for creating Agent Engine Sandboxes.""" +class _CreateRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for creating Agent Runtime Sandboxes.""" name: Optional[str] - """Name of the agent engine to create the sandbox under.""" + """Name of the Agent Runtime to create the sandbox under.""" spec: Optional[SandboxEnvironmentSpecDict] """The specification of the sandbox.""" - config: Optional[CreateAgentEngineSandboxConfigDict] + config: Optional[CreateRuntimeSandboxConfigDict] """""" -_CreateAgentEngineSandboxRequestParametersOrDict = Union[ - _CreateAgentEngineSandboxRequestParameters, - _CreateAgentEngineSandboxRequestParametersDict, +_CreateRuntimeSandboxRequestParametersOrDict = Union[ + _CreateRuntimeSandboxRequestParameters, _CreateRuntimeSandboxRequestParametersDict ] @@ -15876,6 +16615,10 @@ class SandboxEnvironmentConnectionInfo(_common.BaseModel): default=None, description="""Output only. The routing token for the SandboxEnvironment.""", ) + sandbox_hostname: Optional[str] = Field( + default=None, + description="""Output only. The hostname of the SandboxEnvironment.""", + ) class SandboxEnvironmentConnectionInfoDict(TypedDict, total=False): @@ -15893,6 +16636,9 @@ class SandboxEnvironmentConnectionInfoDict(TypedDict, total=False): routing_token: Optional[str] """Output only. The routing token for the SandboxEnvironment.""" + sandbox_hostname: Optional[str] + """Output only. The hostname of the SandboxEnvironment.""" + SandboxEnvironmentConnectionInfoOrDict = Union[ SandboxEnvironmentConnectionInfo, SandboxEnvironmentConnectionInfoDict @@ -16003,8 +16749,8 @@ class SandboxEnvironmentDict(TypedDict, total=False): SandboxEnvironmentOrDict = Union[SandboxEnvironment, SandboxEnvironmentDict] -class AgentEngineSandboxOperation(_common.BaseModel): - """Operation that has an agent engine sandbox as a response.""" +class RuntimeSandboxOperation(_common.BaseModel): + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] = Field( default=None, @@ -16023,12 +16769,12 @@ class AgentEngineSandboxOperation(_common.BaseModel): description="""The error result of the operation in case of failure or cancellation.""", ) response: Optional[SandboxEnvironment] = Field( - default=None, description="""The Agent Engine Sandbox.""" + default=None, description="""The Agent Runtime Sandbox.""" ) -class AgentEngineSandboxOperationDict(TypedDict, total=False): - """Operation that has an agent engine sandbox as a response.""" +class RuntimeSandboxOperationDict(TypedDict, total=False): + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -16043,63 +16789,62 @@ class AgentEngineSandboxOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" response: Optional[SandboxEnvironmentDict] - """The Agent Engine Sandbox.""" + """The Agent Runtime Sandbox.""" -AgentEngineSandboxOperationOrDict = Union[ - AgentEngineSandboxOperation, AgentEngineSandboxOperationDict +RuntimeSandboxOperationOrDict = Union[ + RuntimeSandboxOperation, RuntimeSandboxOperationDict ] -class DeleteAgentEngineSandboxConfig(_common.BaseModel): - """Config for deleting an Agent Engine Sandbox.""" +class DeleteRuntimeSandboxConfig(_common.BaseModel): + """Config for deleting an Agent Runtime Sandbox.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class DeleteAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Sandbox.""" +class DeleteRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for deleting an Agent Runtime Sandbox.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteAgentEngineSandboxConfigOrDict = Union[ - DeleteAgentEngineSandboxConfig, DeleteAgentEngineSandboxConfigDict +DeleteRuntimeSandboxConfigOrDict = Union[ + DeleteRuntimeSandboxConfig, DeleteRuntimeSandboxConfigDict ] -class _DeleteAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for deleting agent engines.""" +class _DeleteRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtimes.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine sandbox to delete.""" + default=None, description="""Name of the agent runtime sandbox to delete.""" ) - config: Optional[DeleteAgentEngineSandboxConfig] = Field( + config: Optional[DeleteRuntimeSandboxConfig] = Field( default=None, description="""""" ) -class _DeleteAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engines.""" +class _DeleteRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtimes.""" name: Optional[str] - """Name of the agent engine sandbox to delete.""" + """Name of the agent runtime sandbox to delete.""" - config: Optional[DeleteAgentEngineSandboxConfigDict] + config: Optional[DeleteRuntimeSandboxConfigDict] """""" -_DeleteAgentEngineSandboxRequestParametersOrDict = Union[ - _DeleteAgentEngineSandboxRequestParameters, - _DeleteAgentEngineSandboxRequestParametersDict, +_DeleteRuntimeSandboxRequestParametersOrDict = Union[ + _DeleteRuntimeSandboxRequestParameters, _DeleteRuntimeSandboxRequestParametersDict ] -class DeleteAgentEngineSandboxOperation(_common.BaseModel): - """Operation for deleting agent engines.""" +class DeleteRuntimeSandboxOperation(_common.BaseModel): + """Operation for deleting agent runtimes.""" name: Optional[str] = Field( default=None, @@ -16119,8 +16864,8 @@ class DeleteAgentEngineSandboxOperation(_common.BaseModel): ) -class DeleteAgentEngineSandboxOperationDict(TypedDict, total=False): - """Operation for deleting agent engines.""" +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}`.""" @@ -16135,8 +16880,8 @@ class DeleteAgentEngineSandboxOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -DeleteAgentEngineSandboxOperationOrDict = Union[ - DeleteAgentEngineSandboxOperation, DeleteAgentEngineSandboxOperationDict +DeleteRuntimeSandboxOperationOrDict = Union[ + DeleteRuntimeSandboxOperation, DeleteRuntimeSandboxOperationDict ] @@ -16191,57 +16936,57 @@ class ChunkDict(TypedDict, total=False): ChunkOrDict = Union[Chunk, ChunkDict] -class ExecuteCodeAgentEngineSandboxConfig(_common.BaseModel): - """Config for executing code in an Agent Engine sandbox.""" +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 ExecuteCodeAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for executing code in an Agent Engine sandbox.""" +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.""" -ExecuteCodeAgentEngineSandboxConfigOrDict = Union[ - ExecuteCodeAgentEngineSandboxConfig, ExecuteCodeAgentEngineSandboxConfigDict +ExecuteCodeRuntimeSandboxConfigOrDict = Union[ + ExecuteCodeRuntimeSandboxConfig, ExecuteCodeRuntimeSandboxConfigDict ] -class _ExecuteCodeAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for executing code in an agent engine sandbox.""" +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 sandbox to execute code in.""", + 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[ExecuteCodeAgentEngineSandboxConfig] = Field( + config: Optional[ExecuteCodeRuntimeSandboxConfig] = Field( default=None, description="""""" ) -class _ExecuteCodeAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for executing code in an agent engine sandbox.""" +class _ExecuteCodeRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for executing code in an agent runtime sandbox.""" name: Optional[str] - """Name of the agent engine sandbox to execute code in.""" + """Name of the agent runtime sandbox to execute code in.""" inputs: Optional[list[ChunkDict]] """Inputs to the code execution.""" - config: Optional[ExecuteCodeAgentEngineSandboxConfigDict] + config: Optional[ExecuteCodeRuntimeSandboxConfigDict] """""" -_ExecuteCodeAgentEngineSandboxRequestParametersOrDict = Union[ - _ExecuteCodeAgentEngineSandboxRequestParameters, - _ExecuteCodeAgentEngineSandboxRequestParametersDict, +_ExecuteCodeRuntimeSandboxRequestParametersOrDict = Union[ + _ExecuteCodeRuntimeSandboxRequestParameters, + _ExecuteCodeRuntimeSandboxRequestParametersDict, ] @@ -16265,54 +17010,52 @@ class ExecuteSandboxEnvironmentResponseDict(TypedDict, total=False): ] -class GetAgentEngineSandboxConfig(_common.BaseModel): - """Config for getting an Agent Engine Memory.""" +class GetRuntimeSandboxConfig(_common.BaseModel): + """Config for getting an Agent Runtime Memory.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) -class GetAgentEngineSandboxConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Memory.""" +class GetRuntimeSandboxConfigDict(TypedDict, total=False): + """Config for getting an Agent Runtime Memory.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetAgentEngineSandboxConfigOrDict = Union[ - GetAgentEngineSandboxConfig, GetAgentEngineSandboxConfigDict +GetRuntimeSandboxConfigOrDict = Union[ + GetRuntimeSandboxConfig, GetRuntimeSandboxConfigDict ] -class _GetAgentEngineSandboxRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine sandbox.""" +class _GetRuntimeSandboxRequestParameters(_common.BaseModel): + """Parameters for getting an agent runtime sandbox.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine sandbox.""" - ) - config: Optional[GetAgentEngineSandboxConfig] = Field( - default=None, description="""""" + default=None, description="""Name of the agent runtime sandbox.""" ) + config: Optional[GetRuntimeSandboxConfig] = Field(default=None, description="""""") -class _GetAgentEngineSandboxRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine sandbox.""" +class _GetRuntimeSandboxRequestParametersDict(TypedDict, total=False): + """Parameters for getting an agent runtime sandbox.""" name: Optional[str] - """Name of the agent engine sandbox.""" + """Name of the agent runtime sandbox.""" - config: Optional[GetAgentEngineSandboxConfigDict] + config: Optional[GetRuntimeSandboxConfigDict] """""" -_GetAgentEngineSandboxRequestParametersOrDict = Union[ - _GetAgentEngineSandboxRequestParameters, _GetAgentEngineSandboxRequestParametersDict +_GetRuntimeSandboxRequestParametersOrDict = Union[ + _GetRuntimeSandboxRequestParameters, _GetRuntimeSandboxRequestParametersDict ] -class ListAgentEngineSandboxesConfig(_common.BaseModel): - """Config for listing agent engine sandboxes.""" +class ListRuntimeSandboxesConfig(_common.BaseModel): + """Config for listing agent runtime sandboxes.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" @@ -16326,8 +17069,8 @@ class ListAgentEngineSandboxesConfig(_common.BaseModel): ) -class ListAgentEngineSandboxesConfigDict(TypedDict, total=False): - """Config for listing agent engine sandboxes.""" +class ListRuntimeSandboxesConfigDict(TypedDict, total=False): + """Config for listing agent runtime sandboxes.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -16343,52 +17086,51 @@ class ListAgentEngineSandboxesConfigDict(TypedDict, total=False): For field names both snake_case and camelCase are supported.""" -ListAgentEngineSandboxesConfigOrDict = Union[ - ListAgentEngineSandboxesConfig, ListAgentEngineSandboxesConfigDict +ListRuntimeSandboxesConfigOrDict = Union[ + ListRuntimeSandboxesConfig, ListRuntimeSandboxesConfigDict ] -class _ListAgentEngineSandboxesRequestParameters(_common.BaseModel): - """Parameters for listing agent engine sandboxes.""" +class _ListRuntimeSandboxesRequestParameters(_common.BaseModel): + """Parameters for listing agent runtime sandboxes.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime.""" ) - config: Optional[ListAgentEngineSandboxesConfig] = Field( + config: Optional[ListRuntimeSandboxesConfig] = Field( default=None, description="""""" ) -class _ListAgentEngineSandboxesRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engine sandboxes.""" +class _ListRuntimeSandboxesRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtime sandboxes.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime.""" - config: Optional[ListAgentEngineSandboxesConfigDict] + config: Optional[ListRuntimeSandboxesConfigDict] """""" -_ListAgentEngineSandboxesRequestParametersOrDict = Union[ - _ListAgentEngineSandboxesRequestParameters, - _ListAgentEngineSandboxesRequestParametersDict, +_ListRuntimeSandboxesRequestParametersOrDict = Union[ + _ListRuntimeSandboxesRequestParameters, _ListRuntimeSandboxesRequestParametersDict ] -class ListAgentEngineSandboxesResponse(_common.BaseModel): - """Response for listing agent engine sandboxes.""" +class ListRuntimeSandboxesResponse(_common.BaseModel): + """Response for listing agent runtime sandboxes.""" sdk_http_response: Optional[genai_types.HttpResponse] = Field( default=None, description="""Used to retain the full HTTP response.""" ) next_page_token: Optional[str] = Field(default=None, description="""""") sandbox_environments: Optional[list[SandboxEnvironment]] = Field( - default=None, description="""List of agent engine sandboxes.""" + default=None, description="""List of agent runtime sandboxes.""" ) -class ListAgentEngineSandboxesResponseDict(TypedDict, total=False): - """Response for listing agent engine sandboxes.""" +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.""" @@ -16397,38 +17139,37 @@ class ListAgentEngineSandboxesResponseDict(TypedDict, total=False): """""" sandbox_environments: Optional[list[SandboxEnvironmentDict]] - """List of agent engine sandboxes.""" + """List of agent runtime sandboxes.""" -ListAgentEngineSandboxesResponseOrDict = Union[ - ListAgentEngineSandboxesResponse, ListAgentEngineSandboxesResponseDict +ListRuntimeSandboxesResponseOrDict = Union[ + ListRuntimeSandboxesResponse, ListRuntimeSandboxesResponseDict ] -class _GetAgentEngineSandboxOperationParameters(_common.BaseModel): +class _GetRuntimeSandboxOperationParameters(_common.BaseModel): """Parameters for getting an operation with a sandbox 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 _GetAgentEngineSandboxOperationParametersDict(TypedDict, total=False): +class _GetRuntimeSandboxOperationParametersDict(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.""" - config: Optional[GetAgentEngineOperationConfigDict] + config: Optional[GetRuntimeOperationConfigDict] """Used to override the default configuration.""" -_GetAgentEngineSandboxOperationParametersOrDict = Union[ - _GetAgentEngineSandboxOperationParameters, - _GetAgentEngineSandboxOperationParametersDict, +_GetRuntimeSandboxOperationParametersOrDict = Union[ + _GetRuntimeSandboxOperationParameters, _GetRuntimeSandboxOperationParametersDict ] @@ -16660,7 +17401,7 @@ class _CreateSandboxEnvironmentTemplateRequestParameters(_common.BaseModel): name: Optional[str] = Field( default=None, - description="""Name of the agent engine to create the template under.""", + description="""Name of the agent runtime to create the template under.""", ) config: Optional[CreateSandboxEnvironmentTemplateConfig] = Field( default=None, description="""""" @@ -16674,7 +17415,7 @@ class _CreateSandboxEnvironmentTemplateRequestParametersDict(TypedDict, total=Fa """Parameters for creating Sandbox Environment Templates.""" name: Optional[str] - """Name of the agent engine to create the template under.""" + """Name of the agent runtime to create the template under.""" config: Optional[CreateSandboxEnvironmentTemplateConfigDict] """""" @@ -16788,7 +17529,7 @@ class SandboxEnvironmentTemplateDict(TypedDict, total=False): class SandboxEnvironmentTemplateOperation(_common.BaseModel): - """Operation that has an agent engine sandbox as a response.""" + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] = Field( default=None, @@ -16807,12 +17548,12 @@ class SandboxEnvironmentTemplateOperation(_common.BaseModel): 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.""" + default=None, description="""The Agent Runtime Sandbox Template.""" ) class SandboxEnvironmentTemplateOperationDict(TypedDict, total=False): - """Operation that has an agent engine sandbox as a response.""" + """Operation that has an agent runtime sandbox as a response.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -16827,7 +17568,7 @@ class SandboxEnvironmentTemplateOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" response: Optional[SandboxEnvironmentTemplateDict] - """The Agent Engine Sandbox Template.""" + """The Agent Runtime Sandbox Template.""" SandboxEnvironmentTemplateOperationOrDict = Union[ @@ -17011,7 +17752,7 @@ class _ListSandboxEnvironmentTemplatesRequestParameters(_common.BaseModel): """Parameters for listing sandbox templates.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime.""" ) config: Optional[ListSandboxEnvironmentTemplatesConfig] = Field( default=None, description="""""" @@ -17022,7 +17763,7 @@ class _ListSandboxEnvironmentTemplatesRequestParametersDict(TypedDict, total=Fal """Parameters for listing sandbox templates.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime.""" config: Optional[ListSandboxEnvironmentTemplatesConfigDict] """""" @@ -17070,7 +17811,7 @@ class _GetSandboxEnvironmentTemplateOperationParameters(_common.BaseModel): 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.""" ) @@ -17081,7 +17822,7 @@ class _GetSandboxEnvironmentTemplateOperationParametersDict(TypedDict, total=Fal operation_name: Optional[str] """The server-assigned name for the operation.""" - config: Optional[GetAgentEngineOperationConfigDict] + config: Optional[GetRuntimeOperationConfigDict] """Used to override the default configuration.""" @@ -17091,7 +17832,7 @@ class _GetSandboxEnvironmentTemplateOperationParametersDict(TypedDict, total=Fal ] -class CreateAgentEngineSandboxSnapshotConfig(_common.BaseModel): +class CreateRuntimeSandboxSnapshotConfig(_common.BaseModel): """Config for creating a Sandbox Environment Snapshot.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -17113,7 +17854,7 @@ class CreateAgentEngineSandboxSnapshotConfig(_common.BaseModel): ) -class CreateAgentEngineSandboxSnapshotConfigDict(TypedDict, total=False): +class CreateRuntimeSandboxSnapshotConfigDict(TypedDict, total=False): """Config for creating a Sandbox Environment Snapshot.""" http_options: Optional[genai_types.HttpOptions] @@ -17132,8 +17873,8 @@ class CreateAgentEngineSandboxSnapshotConfigDict(TypedDict, total=False): """Waits for the operation to complete before returning.""" -CreateAgentEngineSandboxSnapshotConfigOrDict = Union[ - CreateAgentEngineSandboxSnapshotConfig, CreateAgentEngineSandboxSnapshotConfigDict +CreateRuntimeSandboxSnapshotConfigOrDict = Union[ + CreateRuntimeSandboxSnapshotConfig, CreateRuntimeSandboxSnapshotConfigDict ] @@ -17143,7 +17884,7 @@ class _CreateSandboxEnvironmentSnapshotRequestParameters(_common.BaseModel): source_sandbox_environment_name: Optional[str] = Field( default=None, description="""Name of the sandbox environment to snapshot.""" ) - config: Optional[CreateAgentEngineSandboxSnapshotConfig] = Field( + config: Optional[CreateRuntimeSandboxSnapshotConfig] = Field( default=None, description="""""" ) @@ -17154,7 +17895,7 @@ class _CreateSandboxEnvironmentSnapshotRequestParametersDict(TypedDict, total=Fa source_sandbox_environment_name: Optional[str] """Name of the sandbox environment to snapshot.""" - config: Optional[CreateAgentEngineSandboxSnapshotConfigDict] + config: Optional[CreateRuntimeSandboxSnapshotConfigDict] """""" @@ -17257,8 +17998,8 @@ class SandboxEnvironmentSnapshotDict(TypedDict, total=False): ] -class AgentEngineSandboxSnapshotOperation(_common.BaseModel): - """Operation that has an agent engine sandbox snapshot as a response.""" +class RuntimeSandboxSnapshotOperation(_common.BaseModel): + """Operation that has an agent runtime sandbox snapshot as a response.""" name: Optional[str] = Field( default=None, @@ -17277,12 +18018,12 @@ class AgentEngineSandboxSnapshotOperation(_common.BaseModel): 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.""" + default=None, description="""The Agent Runtime Sandbox Snapshot.""" ) -class AgentEngineSandboxSnapshotOperationDict(TypedDict, total=False): - """Operation that has an agent engine sandbox snapshot as a response.""" +class RuntimeSandboxSnapshotOperationDict(TypedDict, total=False): + """Operation that has an agent runtime sandbox snapshot as a response.""" name: Optional[str] """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" @@ -17297,11 +18038,11 @@ class AgentEngineSandboxSnapshotOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" response: Optional[SandboxEnvironmentSnapshotDict] - """The Agent Engine Sandbox Snapshot.""" + """The Agent Runtime Sandbox Snapshot.""" -AgentEngineSandboxSnapshotOperationOrDict = Union[ - AgentEngineSandboxSnapshotOperation, AgentEngineSandboxSnapshotOperationDict +RuntimeSandboxSnapshotOperationOrDict = Union[ + RuntimeSandboxSnapshotOperation, RuntimeSandboxSnapshotOperationDict ] @@ -17536,34 +18277,34 @@ class ListSandboxEnvironmentSnapshotsResponseDict(TypedDict, total=False): ] -class _GetAgentEngineSandboxSnapshotOperationParameters(_common.BaseModel): +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 _GetAgentEngineSandboxSnapshotOperationParametersDict(TypedDict, total=False): +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.""" -_GetAgentEngineSandboxSnapshotOperationParametersOrDict = Union[ - _GetAgentEngineSandboxSnapshotOperationParameters, - _GetAgentEngineSandboxSnapshotOperationParametersDict, +_GetRuntimeSandboxSnapshotOperationParametersOrDict = Union[ + _GetRuntimeSandboxSnapshotOperationParameters, + _GetRuntimeSandboxSnapshotOperationParametersDict, ] -class CreateAgentEngineSessionConfig(_common.BaseModel): +class CreateRuntimeSessionConfig(_common.BaseModel): """Config for creating a Session.""" http_options: Optional[genai_types.HttpOptions] = Field( @@ -17600,7 +18341,7 @@ class CreateAgentEngineSessionConfig(_common.BaseModel): ) -class CreateAgentEngineSessionConfigDict(TypedDict, total=False): +class CreateRuntimeSessionConfigDict(TypedDict, total=False): """Config for creating a Session.""" http_options: Optional[genai_types.HttpOptions] @@ -17630,42 +18371,41 @@ class CreateAgentEngineSessionConfigDict(TypedDict, total=False): """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.""" -CreateAgentEngineSessionConfigOrDict = Union[ - CreateAgentEngineSessionConfig, CreateAgentEngineSessionConfigDict +CreateRuntimeSessionConfigOrDict = Union[ + CreateRuntimeSessionConfig, CreateRuntimeSessionConfigDict ] -class _CreateAgentEngineSessionRequestParameters(_common.BaseModel): - """Parameters for creating Agent Engine Sessions.""" +class _CreateRuntimeSessionRequestParameters(_common.BaseModel): + """Parameters for creating Agent Runtime Sessions.""" name: Optional[str] = Field( default=None, - description="""Name of the agent engine to create the session under.""", + 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[CreateAgentEngineSessionConfig] = Field( + config: Optional[CreateRuntimeSessionConfig] = Field( default=None, description="""""" ) -class _CreateAgentEngineSessionRequestParametersDict(TypedDict, total=False): - """Parameters for creating Agent Engine Sessions.""" +class _CreateRuntimeSessionRequestParametersDict(TypedDict, total=False): + """Parameters for creating Agent Runtime Sessions.""" name: Optional[str] - """Name of the agent engine to create the session under.""" + """Name of the agent runtime to create the session under.""" user_id: Optional[str] """The user ID of the session.""" - config: Optional[CreateAgentEngineSessionConfigDict] + config: Optional[CreateRuntimeSessionConfigDict] """""" -_CreateAgentEngineSessionRequestParametersOrDict = Union[ - _CreateAgentEngineSessionRequestParameters, - _CreateAgentEngineSessionRequestParametersDict, +_CreateRuntimeSessionRequestParametersOrDict = Union[ + _CreateRuntimeSessionRequestParameters, _CreateRuntimeSessionRequestParametersDict ] @@ -17707,6 +18447,20 @@ class Session(_common.BaseModel): default=None, description="""Required. Immutable. String id provided by the user""", ) + agents: Optional[list[str]] = Field( + default=None, + description="""Output only. The resource names of the Agents that interacted within the session. Each resource name has the format: `projects/{project}/locations/{location}/agents/{agent}`.""", + ) + default_run_step_limit: Optional[int] = Field( + default=None, + description="""The default limit of the number of steps each child Run can contain. Default to 20 if not specified.""", + ) + state: Optional[State] = Field( + default=None, description="""Output only. The state of the session.""" + ) + custom_metadata: Optional[dict[str, Any]] = Field( + default=None, description="""Optional. Custom metadata for the session.""" + ) class SessionDict(TypedDict, total=False): @@ -17739,12 +18493,24 @@ class SessionDict(TypedDict, total=False): user_id: Optional[str] """Required. Immutable. String id provided by the user""" + agents: Optional[list[str]] + """Output only. The resource names of the Agents that interacted within the session. Each resource name has the format: `projects/{project}/locations/{location}/agents/{agent}`.""" + + default_run_step_limit: Optional[int] + """The default limit of the number of steps each child Run can contain. Default to 20 if not specified.""" + + state: Optional[State] + """Output only. The state of the session.""" + + custom_metadata: Optional[dict[str, Any]] + """Optional. Custom metadata for the session.""" + SessionOrDict = Union[Session, SessionDict] -class AgentEngineSessionOperation(_common.BaseModel): - """Operation that has an agent engine session as a response.""" +class RuntimeSessionOperation(_common.BaseModel): + """Operation that has an agent runtime session as a response.""" name: Optional[str] = Field( default=None, @@ -17763,12 +18529,12 @@ class AgentEngineSessionOperation(_common.BaseModel): description="""The error result of the operation in case of failure or cancellation.""", ) response: Optional[Session] = Field( - default=None, description="""The Agent Engine Session.""" + default=None, description="""The Agent Runtime Session.""" ) -class AgentEngineSessionOperationDict(TypedDict, total=False): - """Operation that has an agent engine session 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}`.""" @@ -17783,63 +18549,62 @@ class AgentEngineSessionOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" response: Optional[SessionDict] - """The Agent Engine Session.""" + """The Agent Runtime Session.""" -AgentEngineSessionOperationOrDict = Union[ - AgentEngineSessionOperation, AgentEngineSessionOperationDict +RuntimeSessionOperationOrDict = Union[ + RuntimeSessionOperation, RuntimeSessionOperationDict ] -class DeleteAgentEngineSessionConfig(_common.BaseModel): - """Config for deleting an Agent Engine Session.""" +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 DeleteAgentEngineSessionConfigDict(TypedDict, total=False): - """Config for deleting an Agent Engine Session.""" +class DeleteRuntimeSessionConfigDict(TypedDict, total=False): + """Config for deleting an Agent Runtime Session.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -DeleteAgentEngineSessionConfigOrDict = Union[ - DeleteAgentEngineSessionConfig, DeleteAgentEngineSessionConfigDict +DeleteRuntimeSessionConfigOrDict = Union[ + DeleteRuntimeSessionConfig, DeleteRuntimeSessionConfigDict ] -class _DeleteAgentEngineSessionRequestParameters(_common.BaseModel): - """Parameters for deleting agent engine sessions.""" +class _DeleteRuntimeSessionRequestParameters(_common.BaseModel): + """Parameters for deleting agent runtime sessions.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine session to delete.""" + default=None, description="""Name of the agent runtime session to delete.""" ) - config: Optional[DeleteAgentEngineSessionConfig] = Field( + config: Optional[DeleteRuntimeSessionConfig] = Field( default=None, description="""""" ) -class _DeleteAgentEngineSessionRequestParametersDict(TypedDict, total=False): - """Parameters for deleting agent engine sessions.""" +class _DeleteRuntimeSessionRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent runtime sessions.""" name: Optional[str] - """Name of the agent engine session to delete.""" + """Name of the agent runtime session to delete.""" - config: Optional[DeleteAgentEngineSessionConfigDict] + config: Optional[DeleteRuntimeSessionConfigDict] """""" -_DeleteAgentEngineSessionRequestParametersOrDict = Union[ - _DeleteAgentEngineSessionRequestParameters, - _DeleteAgentEngineSessionRequestParametersDict, +_DeleteRuntimeSessionRequestParametersOrDict = Union[ + _DeleteRuntimeSessionRequestParameters, _DeleteRuntimeSessionRequestParametersDict ] -class DeleteAgentEngineSessionOperation(_common.BaseModel): - """Operation for deleting agent engine sessions.""" +class DeleteRuntimeSessionOperation(_common.BaseModel): + """Operation for deleting agent runtime sessions.""" name: Optional[str] = Field( default=None, @@ -17859,8 +18624,8 @@ class DeleteAgentEngineSessionOperation(_common.BaseModel): ) -class DeleteAgentEngineSessionOperationDict(TypedDict, total=False): - """Operation for deleting agent engine sessions.""" +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}`.""" @@ -17875,59 +18640,57 @@ class DeleteAgentEngineSessionOperationDict(TypedDict, total=False): """The error result of the operation in case of failure or cancellation.""" -DeleteAgentEngineSessionOperationOrDict = Union[ - DeleteAgentEngineSessionOperation, DeleteAgentEngineSessionOperationDict +DeleteRuntimeSessionOperationOrDict = Union[ + DeleteRuntimeSessionOperation, DeleteRuntimeSessionOperationDict ] -class GetAgentEngineSessionConfig(_common.BaseModel): - """Config for getting an Agent Engine Session.""" +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 GetAgentEngineSessionConfigDict(TypedDict, total=False): - """Config for getting an Agent Engine Session.""" +class GetRuntimeSessionConfigDict(TypedDict, total=False): + """Config for getting an Agent Runtime Session.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" -GetAgentEngineSessionConfigOrDict = Union[ - GetAgentEngineSessionConfig, GetAgentEngineSessionConfigDict +GetRuntimeSessionConfigOrDict = Union[ + GetRuntimeSessionConfig, GetRuntimeSessionConfigDict ] -class _GetAgentEngineSessionRequestParameters(_common.BaseModel): - """Parameters for getting an agent engine session.""" +class _GetRuntimeSessionRequestParameters(_common.BaseModel): + """Parameters for getting an agent runtime session.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine session.""" - ) - config: Optional[GetAgentEngineSessionConfig] = Field( - default=None, description="""""" + default=None, description="""Name of the agent runtime session.""" ) + config: Optional[GetRuntimeSessionConfig] = Field(default=None, description="""""") -class _GetAgentEngineSessionRequestParametersDict(TypedDict, total=False): - """Parameters for getting an agent engine session.""" +class _GetRuntimeSessionRequestParametersDict(TypedDict, total=False): + """Parameters for getting an agent runtime session.""" name: Optional[str] - """Name of the agent engine session.""" + """Name of the agent runtime session.""" - config: Optional[GetAgentEngineSessionConfigDict] + config: Optional[GetRuntimeSessionConfigDict] """""" -_GetAgentEngineSessionRequestParametersOrDict = Union[ - _GetAgentEngineSessionRequestParameters, _GetAgentEngineSessionRequestParametersDict +_GetRuntimeSessionRequestParametersOrDict = Union[ + _GetRuntimeSessionRequestParameters, _GetRuntimeSessionRequestParametersDict ] -class ListAgentEngineSessionsConfig(_common.BaseModel): - """Config for listing agent engine sessions.""" +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.""" @@ -17941,8 +18704,8 @@ class ListAgentEngineSessionsConfig(_common.BaseModel): ) -class ListAgentEngineSessionsConfigDict(TypedDict, total=False): - """Config for listing agent engine sessions.""" +class ListRuntimeSessionsConfigDict(TypedDict, total=False): + """Config for listing agent runtime sessions.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -17958,52 +18721,51 @@ class ListAgentEngineSessionsConfigDict(TypedDict, total=False): For field names both snake_case and camelCase are supported.""" -ListAgentEngineSessionsConfigOrDict = Union[ - ListAgentEngineSessionsConfig, ListAgentEngineSessionsConfigDict +ListRuntimeSessionsConfigOrDict = Union[ + ListRuntimeSessionsConfig, ListRuntimeSessionsConfigDict ] -class _ListAgentEngineSessionsRequestParameters(_common.BaseModel): - """Parameters for listing agent engines.""" +class _ListRuntimeSessionsRequestParameters(_common.BaseModel): + """Parameters for listing agent runtimes.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine.""" + default=None, description="""Name of the agent runtime.""" ) - config: Optional[ListAgentEngineSessionsConfig] = Field( + config: Optional[ListRuntimeSessionsConfig] = Field( default=None, description="""""" ) -class _ListAgentEngineSessionsRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engines.""" +class _ListRuntimeSessionsRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtimes.""" name: Optional[str] - """Name of the agent engine.""" + """Name of the agent runtime.""" - config: Optional[ListAgentEngineSessionsConfigDict] + config: Optional[ListRuntimeSessionsConfigDict] """""" -_ListAgentEngineSessionsRequestParametersOrDict = Union[ - _ListAgentEngineSessionsRequestParameters, - _ListAgentEngineSessionsRequestParametersDict, +_ListRuntimeSessionsRequestParametersOrDict = Union[ + _ListRuntimeSessionsRequestParameters, _ListRuntimeSessionsRequestParametersDict ] class ListReasoningEnginesSessionsResponse(_common.BaseModel): - """Response for listing agent engine sessions.""" + """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="""""") sessions: Optional[list[Session]] = Field( - default=None, description="""List of agent engine sessions.""" + default=None, description="""List of agent runtime sessions.""" ) class ListReasoningEnginesSessionsResponseDict(TypedDict, total=False): - """Response for listing agent engine sessions.""" + """Response for listing agent runtime sessions.""" sdk_http_response: Optional[genai_types.HttpResponse] """Used to retain the full HTTP response.""" @@ -18012,7 +18774,7 @@ class ListReasoningEnginesSessionsResponseDict(TypedDict, total=False): """""" sessions: Optional[list[SessionDict]] - """List of agent engine sessions.""" + """List of agent runtime sessions.""" ListReasoningEnginesSessionsResponseOrDict = Union[ @@ -18020,35 +18782,34 @@ class ListReasoningEnginesSessionsResponseDict(TypedDict, total=False): ] -class _GetAgentEngineSessionOperationParameters(_common.BaseModel): +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 _GetAgentEngineSessionOperationParametersDict(TypedDict, total=False): +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.""" -_GetAgentEngineSessionOperationParametersOrDict = Union[ - _GetAgentEngineSessionOperationParameters, - _GetAgentEngineSessionOperationParametersDict, +_GetRuntimeSessionOperationParametersOrDict = Union[ + _GetRuntimeSessionOperationParameters, _GetRuntimeSessionOperationParametersDict ] -class UpdateAgentEngineSessionConfig(_common.BaseModel): - """Config for updating agent engine 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.""" @@ -18088,12 +18849,12 @@ class UpdateAgentEngineSessionConfig(_common.BaseModel): https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""", ) user_id: Optional[str] = Field( - default=None, description="""User ID of the agent engine session to update.""" + default=None, description="""User ID of the agent runtime session to update.""" ) -class UpdateAgentEngineSessionConfigDict(TypedDict, total=False): - """Config for updating agent engine session.""" +class UpdateRuntimeSessionConfigDict(TypedDict, total=False): + """Config for updating agent runtime session.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -18126,38 +18887,37 @@ class UpdateAgentEngineSessionConfigDict(TypedDict, total=False): https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""" user_id: Optional[str] - """User ID of the agent engine session to update.""" + """User ID of the agent runtime session to update.""" -UpdateAgentEngineSessionConfigOrDict = Union[ - UpdateAgentEngineSessionConfig, UpdateAgentEngineSessionConfigDict +UpdateRuntimeSessionConfigOrDict = Union[ + UpdateRuntimeSessionConfig, UpdateRuntimeSessionConfigDict ] -class _UpdateAgentEngineSessionRequestParameters(_common.BaseModel): - """Parameters for updating 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 session to update.""" + default=None, description="""Name of the agent runtime session to update.""" ) - config: Optional[UpdateAgentEngineSessionConfig] = Field( + config: Optional[UpdateRuntimeSessionConfig] = Field( default=None, description="""""" ) -class _UpdateAgentEngineSessionRequestParametersDict(TypedDict, total=False): - """Parameters for updating agent engine sessions.""" +class _UpdateRuntimeSessionRequestParametersDict(TypedDict, total=False): + """Parameters for updating agent runtime sessions.""" name: Optional[str] - """Name of the agent engine session to update.""" + """Name of the agent runtime session to update.""" - config: Optional[UpdateAgentEngineSessionConfigDict] + config: Optional[UpdateRuntimeSessionConfigDict] """""" -_UpdateAgentEngineSessionRequestParametersOrDict = Union[ - _UpdateAgentEngineSessionRequestParameters, - _UpdateAgentEngineSessionRequestParametersDict, +_UpdateRuntimeSessionRequestParametersOrDict = Union[ + _UpdateRuntimeSessionRequestParameters, _UpdateRuntimeSessionRequestParametersDict ] @@ -18188,6 +18948,21 @@ class EventActions(_common.BaseModel): default=None, description="""Optional. If set, the event transfers to the specified agent.""", ) + transfer_to_agent: Optional[str] = Field( + default=None, + description="""Optional. If set, the event transfers to the specified agent. This field is intended to replace 'transfer_agent'. Not in use pending data migration.""", + ) + requested_tool_confirmations: Optional[dict[str, Any]] = Field( + default=None, + description="""Optional. A dict of tool confirmation requested by this event, keyed by function call id.""", + ) + end_of_agent: Optional[bool] = Field( + default=None, + description="""Optional. If true, the current agent has finished its current run.""", + ) + agent_state: Optional[dict[str, Any]] = Field( + default=None, description="""Optional. The agent state at the current event.""" + ) class EventActionsDict(TypedDict, total=False): @@ -18211,6 +18986,18 @@ class EventActionsDict(TypedDict, total=False): transfer_agent: Optional[str] """Optional. If set, the event transfers to the specified agent.""" + transfer_to_agent: Optional[str] + """Optional. If set, the event transfers to the specified agent. This field is intended to replace 'transfer_agent'. Not in use pending data migration.""" + + requested_tool_confirmations: Optional[dict[str, Any]] + """Optional. A dict of tool confirmation requested by this event, keyed by function call id.""" + + end_of_agent: Optional[bool] + """Optional. If true, the current agent has finished its current run.""" + + agent_state: Optional[dict[str, Any]] + """Optional. The agent state at the current event.""" + EventActionsOrDict = Union[EventActions, EventActionsDict] @@ -18251,6 +19038,12 @@ class EventMetadata(_common.BaseModel): output_transcription: Optional[genai_types.Transcription] = Field( default=None, description="""Optional. Audio transcription of model output.""" ) + citation_metadata: Optional[genai_types.CitationMetadata] = Field( + default=None, description="""Optional. Citation metadata for the response.""" + ) + usage_metadata: Optional[genai_types.UsageMetadata] = Field( + default=None, description="""Optional. Usage metadata for the response.""" + ) class EventMetadataDict(TypedDict, total=False): @@ -18283,12 +19076,18 @@ class EventMetadataDict(TypedDict, total=False): output_transcription: Optional[genai_types.Transcription] """Optional. Audio transcription of model output.""" + citation_metadata: Optional[genai_types.CitationMetadataDict] + """Optional. Citation metadata for the response.""" + + usage_metadata: Optional[genai_types.UsageMetadataDict] + """Optional. Usage metadata for the response.""" + EventMetadataOrDict = Union[EventMetadata, EventMetadataDict] -class AppendAgentEngineSessionEventConfig(_common.BaseModel): - """Config for appending agent engine session event.""" +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.""" @@ -18315,8 +19114,8 @@ class AppendAgentEngineSessionEventConfig(_common.BaseModel): ) -class AppendAgentEngineSessionEventConfigDict(TypedDict, total=False): - """Config for appending agent engine session event.""" +class AppendRuntimeSessionEventConfigDict(TypedDict, total=False): + """Config for appending agent runtime session event.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -18340,75 +19139,75 @@ class AppendAgentEngineSessionEventConfigDict(TypedDict, total=False): """Weakly typed raw event data in proto struct format.""" -AppendAgentEngineSessionEventConfigOrDict = Union[ - AppendAgentEngineSessionEventConfig, AppendAgentEngineSessionEventConfigDict +AppendRuntimeSessionEventConfigOrDict = Union[ + AppendRuntimeSessionEventConfig, AppendRuntimeSessionEventConfigDict ] -class _AppendAgentEngineSessionEventRequestParameters(_common.BaseModel): - """Parameters for appending agent engines.""" +class _AppendRuntimeSessionEventRequestParameters(_common.BaseModel): + """Parameters for appending agent runtimes.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine session.""" + default=None, description="""Name of the agent runtime session.""" ) author: Optional[str] = Field( - default=None, description="""Author of the agent engine session event.""" + default=None, description="""Author of the agent runtime session event.""" ) invocation_id: Optional[str] = Field( - default=None, description="""Invocation ID of the agent engine.""" + default=None, description="""Invocation ID of the agent runtime.""" ) timestamp: Optional[datetime.datetime] = Field( default=None, description="""Timestamp indicating when the event was created.""" ) - config: Optional[AppendAgentEngineSessionEventConfig] = Field( + config: Optional[AppendRuntimeSessionEventConfig] = Field( default=None, description="""""" ) -class _AppendAgentEngineSessionEventRequestParametersDict(TypedDict, total=False): - """Parameters for appending agent engines.""" +class _AppendRuntimeSessionEventRequestParametersDict(TypedDict, total=False): + """Parameters for appending agent runtimes.""" name: Optional[str] - """Name of the agent engine session.""" + """Name of the agent runtime session.""" author: Optional[str] - """Author of the agent engine session event.""" + """Author of the agent runtime session event.""" invocation_id: Optional[str] - """Invocation ID of the agent engine.""" + """Invocation ID of the agent runtime.""" timestamp: Optional[datetime.datetime] """Timestamp indicating when the event was created.""" - config: Optional[AppendAgentEngineSessionEventConfigDict] + config: Optional[AppendRuntimeSessionEventConfigDict] """""" -_AppendAgentEngineSessionEventRequestParametersOrDict = Union[ - _AppendAgentEngineSessionEventRequestParameters, - _AppendAgentEngineSessionEventRequestParametersDict, +_AppendRuntimeSessionEventRequestParametersOrDict = Union[ + _AppendRuntimeSessionEventRequestParameters, + _AppendRuntimeSessionEventRequestParametersDict, ] -class AppendAgentEngineSessionEventResponse(_common.BaseModel): - """Response for appending agent engine session event.""" +class AppendRuntimeSessionEventResponse(_common.BaseModel): + """Response for appending agent runtime session event.""" pass -class AppendAgentEngineSessionEventResponseDict(TypedDict, total=False): - """Response for appending agent engine session event.""" +class AppendRuntimeSessionEventResponseDict(TypedDict, total=False): + """Response for appending agent runtime session event.""" pass -AppendAgentEngineSessionEventResponseOrDict = Union[ - AppendAgentEngineSessionEventResponse, AppendAgentEngineSessionEventResponseDict +AppendRuntimeSessionEventResponseOrDict = Union[ + AppendRuntimeSessionEventResponse, AppendRuntimeSessionEventResponseDict ] -class ListAgentEngineSessionEventsConfig(_common.BaseModel): - """Config for listing agent engine session events.""" +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.""" @@ -18422,8 +19221,8 @@ class ListAgentEngineSessionEventsConfig(_common.BaseModel): ) -class ListAgentEngineSessionEventsConfigDict(TypedDict, total=False): - """Config for listing agent engine session events.""" +class ListRuntimeSessionEventsConfigDict(TypedDict, total=False): + """Config for listing agent runtime session events.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" @@ -18439,35 +19238,35 @@ class ListAgentEngineSessionEventsConfigDict(TypedDict, total=False): For field names both snake_case and camelCase are supported.""" -ListAgentEngineSessionEventsConfigOrDict = Union[ - ListAgentEngineSessionEventsConfig, ListAgentEngineSessionEventsConfigDict +ListRuntimeSessionEventsConfigOrDict = Union[ + ListRuntimeSessionEventsConfig, ListRuntimeSessionEventsConfigDict ] -class _ListAgentEngineSessionEventsRequestParameters(_common.BaseModel): - """Parameters for listing agent engine session events.""" +class _ListRuntimeSessionEventsRequestParameters(_common.BaseModel): + """Parameters for listing agent runtime session events.""" name: Optional[str] = Field( - default=None, description="""Name of the agent engine session.""" + default=None, description="""Name of the agent runtime session.""" ) - config: Optional[ListAgentEngineSessionEventsConfig] = Field( + config: Optional[ListRuntimeSessionEventsConfig] = Field( default=None, description="""""" ) -class _ListAgentEngineSessionEventsRequestParametersDict(TypedDict, total=False): - """Parameters for listing agent engine session events.""" +class _ListRuntimeSessionEventsRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent runtime session events.""" name: Optional[str] - """Name of the agent engine session.""" + """Name of the agent runtime session.""" - config: Optional[ListAgentEngineSessionEventsConfigDict] + config: Optional[ListRuntimeSessionEventsConfigDict] """""" -_ListAgentEngineSessionEventsRequestParametersOrDict = Union[ - _ListAgentEngineSessionEventsRequestParameters, - _ListAgentEngineSessionEventsRequestParametersDict, +_ListRuntimeSessionEventsRequestParametersOrDict = Union[ + _ListRuntimeSessionEventsRequestParameters, + _ListRuntimeSessionEventsRequestParametersDict, ] @@ -18551,8 +19350,8 @@ class SessionEventDict(TypedDict, total=False): SessionEventOrDict = Union[SessionEvent, SessionEventDict] -class ListAgentEngineSessionEventsResponse(_common.BaseModel): - """Response for listing agent engine session events.""" +class ListRuntimeSessionEventsResponse(_common.BaseModel): + """Response for listing agent runtime session events.""" sdk_http_response: Optional[genai_types.HttpResponse] = Field( default=None, description="""Used to retain the full HTTP response.""" @@ -18563,8 +19362,8 @@ class ListAgentEngineSessionEventsResponse(_common.BaseModel): ) -class ListAgentEngineSessionEventsResponseDict(TypedDict, total=False): - """Response for listing agent engine session events.""" +class ListRuntimeSessionEventsResponseDict(TypedDict, total=False): + """Response for listing agent runtime session events.""" sdk_http_response: Optional[genai_types.HttpResponse] """Used to retain the full HTTP response.""" @@ -18576,8 +19375,8 @@ class ListAgentEngineSessionEventsResponseDict(TypedDict, total=False): """List of session events.""" -ListAgentEngineSessionEventsResponseOrDict = Union[ - ListAgentEngineSessionEventsResponse, ListAgentEngineSessionEventsResponseDict +ListRuntimeSessionEventsResponseOrDict = Union[ + ListRuntimeSessionEventsResponse, ListRuntimeSessionEventsResponseDict ] @@ -21536,6 +22335,10 @@ class Skill(_common.BaseModel): skill_source: Optional[SkillSource] = Field( default=None, description="""Output only. The source of the Skill.""" ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""Optional. Customer-managed encryption key spec for a Skill. If set, this Skill and all sub-resources of this Skill will be secured by this key.""", + ) class SkillDict(TypedDict, total=False): @@ -21580,6 +22383,9 @@ class SkillDict(TypedDict, total=False): skill_source: Optional[SkillSource] """Output only. The source of the Skill.""" + encryption_spec: Optional[genai_types.EncryptionSpecDict] + """Optional. Customer-managed encryption key spec for a Skill. If set, this Skill and all sub-resources of this Skill will be secured by this key.""" + SkillOrDict = Union[Skill, SkillDict] @@ -22748,6 +23554,32 @@ class ProbeDict(TypedDict, total=False): ProbeOrDict = Union[Probe, ProbeDict] +class AcceleratorRequirement(_common.BaseModel): + """Describes an accelerator requirement for a container.""" + + type: Optional[AcceleratorType] = Field( + default=None, + description="""Type of the accelerator needed to run the container.""", + ) + count: Optional[int] = Field( + default=None, + description="""Number of accelerators of the specified type needed to run the container.""", + ) + + +class AcceleratorRequirementDict(TypedDict, total=False): + """Describes an accelerator requirement for a container.""" + + type: Optional[AcceleratorType] + """Type of the accelerator needed to run the container.""" + + count: Optional[int] + """Number of accelerators of the specified type needed to run the container.""" + + +AcceleratorRequirementOrDict = Union[AcceleratorRequirement, AcceleratorRequirementDict] + + class ModelContainerSpec(_common.BaseModel): """Specification of a container for serving predictions. Some fields in this message correspond to fields in the [Kubernetes Container v1 core specification](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core).""" @@ -22807,6 +23639,10 @@ class ModelContainerSpec(_common.BaseModel): default=None, description="""Immutable. Specification for Kubernetes startup probe.""", ) + accelerator_requirements: Optional[list[AcceleratorRequirement]] = Field( + default=None, + description="""Immutable. Accelerators required to run the container. This changes how containers are started. For example, GPU requirements are used to set the `resources` field in the Kubernetes [Container v1 core](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core) spec, as described on the [Schedule GPUs](https://kubernetes.io/docs/tasks/manage-gpus/scheduling-gpus/) page of the Kubernetes documentation. Currently, this field is only used when deploying to EdgeDevices.""", + ) class ModelContainerSpecDict(TypedDict, total=False): @@ -22854,6 +23690,9 @@ class ModelContainerSpecDict(TypedDict, total=False): startup_probe: Optional[ProbeDict] """Immutable. Specification for Kubernetes startup probe.""" + accelerator_requirements: Optional[list[AcceleratorRequirementDict]] + """Immutable. Accelerators required to run the container. This changes how containers are started. For example, GPU requirements are used to set the `resources` field in the Kubernetes [Container v1 core](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.23/#container-v1-core) spec, as described on the [Schedule GPUs](https://kubernetes.io/docs/tasks/manage-gpus/scheduling-gpus/) page of the Kubernetes documentation. Currently, this field is only used when deploying to EdgeDevices.""" + ModelContainerSpecOrDict = Union[ModelContainerSpec, ModelContainerSpecDict] @@ -22891,6 +23730,25 @@ class AutoscalingMetricSpecDict(TypedDict, total=False): AutoscalingMetricSpecOrDict = Union[AutoscalingMetricSpec, AutoscalingMetricSpecDict] +class NodeRecycling(_common.BaseModel): + """Node recycling config for FlexStart. To help ensure a smooth transition of nodes and prevent downtime for your running jobs, flex-start supports node recycling. When a node reaches the end of its duration, GKE automatically replaces the node with a new one to preserve your running workloads.""" + + lead_time_seconds: Optional[int] = Field( + default=None, + description="""This parameter specifies how early (in seconds) before a node reaches the end of its seven-day duration for a new node provisioning process should start to substitute it. A longer lead time increases the probability that the new node is ready before the old one is removed, but might incur additional costs.""", + ) + + +class NodeRecyclingDict(TypedDict, total=False): + """Node recycling config for FlexStart. To help ensure a smooth transition of nodes and prevent downtime for your running jobs, flex-start supports node recycling. When a node reaches the end of its duration, GKE automatically replaces the node with a new one to preserve your running workloads.""" + + lead_time_seconds: Optional[int] + """This parameter specifies how early (in seconds) before a node reaches the end of its seven-day duration for a new node provisioning process should start to substitute it. A longer lead time increases the probability that the new node is ready before the old one is removed, but might incur additional costs.""" + + +NodeRecyclingOrDict = Union[NodeRecycling, NodeRecyclingDict] + + class FlexStart(_common.BaseModel): """FlexStart is used to schedule the deployment workload on DWS resource. It contains the max duration of the deployment.""" @@ -22898,6 +23756,9 @@ class FlexStart(_common.BaseModel): default=None, description="""The max duration of the deployment is max_runtime_duration. The deployment will be terminated after the duration. The max_runtime_duration can be set up to 7 days.""", ) + node_recycling: Optional[NodeRecycling] = Field( + default=None, description="""If this field is set, node recycling is enabled.""" + ) class FlexStartDict(TypedDict, total=False): @@ -22906,6 +23767,9 @@ class FlexStartDict(TypedDict, total=False): max_runtime_duration: Optional[str] """The max duration of the deployment is max_runtime_duration. The deployment will be terminated after the duration. The max_runtime_duration can be set up to 7 days.""" + node_recycling: Optional[NodeRecyclingDict] + """If this field is set, node recycling is enabled.""" + FlexStartOrDict = Union[FlexStart, FlexStartDict] @@ -22938,23 +23802,95 @@ class DedicatedResourcesScaleToZeroSpecDict(TypedDict, total=False): ] -class DedicatedResources(_common.BaseModel): - """A description of resources that are dedicated to a DeployedModel or DeployedIndex, and that need a higher degree of manual configuration.""" +class ProvisioningPriority(_common.BaseModel): + """Defines a single compute resource preference.""" - autoscaling_metric_specs: Optional[list[AutoscalingMetricSpec]] = Field( - default=None, - description="""Immutable. The metric specifications that overrides a resource utilization metric (CPU utilization, accelerator's duty cycle, and so on) target value (default to 60 if not set). At most one entry is allowed per metric. If machine_spec.accelerator_count is above 0, the autoscaling will be based on both CPU utilization and accelerator's duty cycle metrics and scale up when either metrics exceeds its target value while scale down if both metrics are under their target value. The default target value is 60 for both metrics. If machine_spec.accelerator_count is 0, the autoscaling will be based on CPU utilization metric only with default target value 60 if not explicitly set. For example, in the case of Online Prediction, if you want to override target CPU utilization to 80, you should set autoscaling_metric_specs.metric_name to `aiplatform.googleapis.com/prediction/online/cpu/utilization` and autoscaling_metric_specs.target to `80`.""", + provisioning_mode: Optional[ProvisionMode] = Field( + default=None, description="""The predefined provisioning mode.""" ) - flex_start: Optional[FlexStart] = Field( + reservation: Optional[str] = Field( default=None, - description="""Optional. Immutable. If set, use DWS resource to schedule the deployment workload. reference: (https://cloud.google.com/blog/products/compute/introducing-dynamic-workload-scheduler)""", + description="""The full resource name of a SPECIFIC reservation. This field is only valid when the provisioning_mode is SPECIFIC_RESERVATION.""", ) - initial_replica_count: Optional[int] = Field( + + +class ProvisioningPriorityDict(TypedDict, total=False): + """Defines a single compute resource preference.""" + + provisioning_mode: Optional[ProvisionMode] + """The predefined provisioning mode.""" + + reservation: Optional[str] + """The full resource name of a SPECIFIC reservation. This field is only valid when the provisioning_mode is SPECIFIC_RESERVATION.""" + + +ProvisioningPriorityOrDict = Union[ProvisioningPriority, ProvisioningPriorityDict] + + +class ProvisioningModelPriority(_common.BaseModel): + """A priority level, which contains one or more provisioning priorities that are considered to be of the same priority.""" + + priorities: Optional[list[ProvisioningPriority]] = Field( default=None, - description="""Immutable. Number of initial replicas being deployed on when scaling the workload up from zero or when creating the workload in case min_replica_count = 0. When min_replica_count > 0 (meaning that the scale-to-zero feature is not enabled), initial_replica_count should not be set. When min_replica_count = 0 (meaning that the scale-to-zero feature is enabled), initial_replica_count should be larger than zero, but no greater than max_replica_count.""", + description="""A list of provisioning priorities within a single priority level.""", ) - machine_spec: Optional[MachineSpec] = Field( - default=None, + + +class ProvisioningModelPriorityDict(TypedDict, total=False): + """A priority level, which contains one or more provisioning priorities that are considered to be of the same priority.""" + + priorities: Optional[list[ProvisioningPriorityDict]] + """A list of provisioning priorities within a single priority level.""" + + +ProvisioningModelPriorityOrDict = Union[ + ProvisioningModelPriority, ProvisioningModelPriorityDict +] + + +class ProvisioningConfig(_common.BaseModel): + """Defines a provisioning config for scheduling and sourcing compute resources.""" + + strategy_type: Optional[ProvisioningStrategyType] = Field( + default=None, + description="""A predefined, simple strategy. If this is set, the provisioning_model_priorities will be ignored.""", + ) + priority_levels: Optional[list[ProvisioningModelPriority]] = Field( + default=None, + description="""A list of compute resource preferences for custom strategies. When strategy_type is "MINIMIZE_COST" or "MAXIMIZE_AVAILABILITY", we allow only one priority_levels and inside it, allows only SPECIFIC_RESERVATION (one or many).""", + ) + + +class ProvisioningConfigDict(TypedDict, total=False): + """Defines a provisioning config for scheduling and sourcing compute resources.""" + + strategy_type: Optional[ProvisioningStrategyType] + """A predefined, simple strategy. If this is set, the provisioning_model_priorities will be ignored.""" + + priority_levels: Optional[list[ProvisioningModelPriorityDict]] + """A list of compute resource preferences for custom strategies. When strategy_type is "MINIMIZE_COST" or "MAXIMIZE_AVAILABILITY", we allow only one priority_levels and inside it, allows only SPECIFIC_RESERVATION (one or many).""" + + +ProvisioningConfigOrDict = Union[ProvisioningConfig, ProvisioningConfigDict] + + +class DedicatedResources(_common.BaseModel): + """A description of resources that are dedicated to a DeployedModel or DeployedIndex, and that need a higher degree of manual configuration.""" + + autoscaling_metric_specs: Optional[list[AutoscalingMetricSpec]] = Field( + default=None, + description="""Immutable. The metric specifications that overrides a resource utilization metric (CPU utilization, accelerator's duty cycle, and so on) target value (default to 60 if not set). At most one entry is allowed per metric. If machine_spec.accelerator_count is above 0, the autoscaling will be based on both CPU utilization and accelerator's duty cycle metrics and scale up when either metrics exceeds its target value while scale down if both metrics are under their target value. The default target value is 60 for both metrics. If machine_spec.accelerator_count is 0, the autoscaling will be based on CPU utilization metric only with default target value 60 if not explicitly set. For example, in the case of Online Prediction, if you want to override target CPU utilization to 80, you should set autoscaling_metric_specs.metric_name to `aiplatform.googleapis.com/prediction/online/cpu/utilization` and autoscaling_metric_specs.target to `80`.""", + ) + flex_start: Optional[FlexStart] = Field( + default=None, + description="""Optional. Immutable. If set, use DWS resource to schedule the deployment workload. reference: (https://cloud.google.com/blog/products/compute/introducing-dynamic-workload-scheduler)""", + ) + initial_replica_count: Optional[int] = Field( + default=None, + description="""Immutable. Number of initial replicas being deployed on when scaling the workload up from zero or when creating the workload in case min_replica_count = 0. When min_replica_count > 0 (meaning that the scale-to-zero feature is not enabled), initial_replica_count should not be set. When min_replica_count = 0 (meaning that the scale-to-zero feature is enabled), initial_replica_count should be larger than zero, but no greater than max_replica_count.""", + ) + machine_spec: Optional[MachineSpec] = Field( + default=None, description="""Required. Immutable. The specification of a single machine being used.""", ) max_replica_count: Optional[int] = Field( @@ -22977,6 +23913,14 @@ class DedicatedResources(_common.BaseModel): default=None, description="""Optional. If true, schedule the deployment workload on [spot VMs](https://cloud.google.com/kubernetes-engine/docs/concepts/spot-vms).""", ) + deployment_resource_pool_id: Optional[str] = Field( + default=None, + description="""Optional. Immutable. Specifies the deployment resource pool to use for this deployment. This must refer to a deployment resource pool with a pool type of `RESOURCE_POOL_TYPE_RESERVED_CAPACITY`. If not specified, an internal deployment resource pool will be used. Format: `projects/{project}/locations/{location}/deploymentResourcePools/{deployment_resource_pool}`""", + ) + provisioning_config: Optional[ProvisioningConfig] = Field( + default=None, + description="""Optional. Configuration for mixed provisioning modes. Once set, the DedicatedResources.MachineSpec.ReservationAffinity will be ignored.""", + ) class DedicatedResourcesDict(TypedDict, total=False): @@ -23009,6 +23953,12 @@ class DedicatedResourcesDict(TypedDict, total=False): spot: Optional[bool] """Optional. If true, schedule the deployment workload on [spot VMs](https://cloud.google.com/kubernetes-engine/docs/concepts/spot-vms).""" + deployment_resource_pool_id: Optional[str] + """Optional. Immutable. Specifies the deployment resource pool to use for this deployment. This must refer to a deployment resource pool with a pool type of `RESOURCE_POOL_TYPE_RESERVED_CAPACITY`. If not specified, an internal deployment resource pool will be used. Format: `projects/{project}/locations/{location}/deploymentResourcePools/{deployment_resource_pool}`""" + + provisioning_config: Optional[ProvisioningConfigDict] + """Optional. Configuration for mixed provisioning modes. Once set, the DedicatedResources.MachineSpec.ReservationAffinity will be ignored.""" + DedicatedResourcesOrDict = Union[DedicatedResources, DedicatedResourcesDict] @@ -23294,6 +24244,325 @@ class PublisherModelCallToActionViewRestApiDict(TypedDict, total=False): ] +class PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResources( + _common.BaseModel +): + """Resources message.""" + + machine_spec: Optional[MachineSpec] = Field( + default=None, description="""Required. Immutable. Machine Specifications.""" + ) + replica_count: Optional[int] = Field( + default=None, description="""Required. Count of replicas.""" + ) + + +class PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResourcesDict( + TypedDict, total=False +): + """Resources message.""" + + machine_spec: Optional[MachineSpecDict] + """Required. Immutable. Machine Specifications.""" + + replica_count: Optional[int] + """Required. Count of replicas.""" + + +PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResourcesOrDict = Union[ + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResources, + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResourcesDict, +] + + +class PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModel( + _common.BaseModel +): + """Stores metadata for Coprocessor Resource. Ex. Trust And Safety Resource.""" + + coprocessor_model_id: Optional[str] = Field( + default=None, description="""Required. Coprocessor Model Id.""" + ) + resources: Optional[ + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResources + ] = Field( + default=None, + description="""Optional. Resources required for CoprocessorModel.""", + ) + parent_model_reference: Optional[PublisherModelResourceReference] = Field( + default=None, + description="""Required. Reference to a Publisher uploaded T&S model.""", + ) + dedicated_resources: Optional[DedicatedResources] = Field( + default=None, + description="""Required. Resources required for CoprocessorModel. Use this field instead of resources.""", + ) + + +class PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModelDict( + TypedDict, total=False +): + """Stores metadata for Coprocessor Resource. Ex. Trust And Safety Resource.""" + + coprocessor_model_id: Optional[str] + """Required. Coprocessor Model Id.""" + + resources: Optional[ + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResourcesDict + ] + """Optional. Resources required for CoprocessorModel.""" + + parent_model_reference: Optional[PublisherModelResourceReferenceDict] + """Required. Reference to a Publisher uploaded T&S model.""" + + dedicated_resources: Optional[DedicatedResourcesDict] + """Required. Resources required for CoprocessorModel. Use this field instead of resources.""" + + +PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModelOrDict = Union[ + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModel, + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModelDict, +] + + +class PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnit( + _common.BaseModel +): + """A Monetizable Deployment Unit. Contains Marketplace resource information and deployment configuration.""" + + deployment_unit_id: Optional[str] = Field( + default=None, + description="""Output only. The resource name of Deployment Unit.""", + ) + resources_per_unit: Optional[ + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResources + ] = Field(default=None, description="""Optional. Resources per unit.""") + coprocessor_models: Optional[ + list[ + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModel + ] + ] = Field( + default=None, + description="""Optional. Coprocessor Resources per unit map. Ex. Trust and Safety.""", + ) + dedicated_resources: Optional[DedicatedResources] = Field( + default=None, + description="""Required. Dedicated resources. This field should be set instead of resources_per_unit.""", + ) + + +class PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitDict( + TypedDict, total=False +): + """A Monetizable Deployment Unit. Contains Marketplace resource information and deployment configuration.""" + + deployment_unit_id: Optional[str] + """Output only. The resource name of Deployment Unit.""" + + resources_per_unit: Optional[ + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitResourcesDict + ] + """Optional. Resources per unit.""" + + coprocessor_models: Optional[ + list[ + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitCoprocessorModelDict + ] + ] + """Optional. Coprocessor Resources per unit map. Ex. Trust and Safety.""" + + dedicated_resources: Optional[DedicatedResourcesDict] + """Required. Dedicated resources. This field should be set instead of resources_per_unit.""" + + +PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitOrDict = Union[ + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnit, + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitDict, +] + + +class PublisherModelCallToActionDeployMonetizedModelMarketplaceListing( + _common.BaseModel +): + """Marketplace listing ID and Plan ID. This is seeded by Marketplace when a new Offer is accepted.""" + + service_id: Optional[str] = Field( + default=None, description="""Required. Marketplace Listing Id.""" + ) + service_level: Optional[str] = Field( + default=None, description="""Required. Marketplace Listing Plan Id.""" + ) + disable_deployment_billing: Optional[bool] = Field( + default=None, + description="""Optional. Field to disable billing for model deployments.""", + ) + + +class PublisherModelCallToActionDeployMonetizedModelMarketplaceListingDict( + TypedDict, total=False +): + """Marketplace listing ID and Plan ID. This is seeded by Marketplace when a new Offer is accepted.""" + + service_id: Optional[str] + """Required. Marketplace Listing Id.""" + + service_level: Optional[str] + """Required. Marketplace Listing Plan Id.""" + + disable_deployment_billing: Optional[bool] + """Optional. Field to disable billing for model deployments.""" + + +PublisherModelCallToActionDeployMonetizedModelMarketplaceListingOrDict = Union[ + PublisherModelCallToActionDeployMonetizedModelMarketplaceListing, + PublisherModelCallToActionDeployMonetizedModelMarketplaceListingDict, +] + + +class PublisherModelCallToActionDeployMonetizedModelSingleTenant(_common.BaseModel): + """Single Tenant Model configuration. This is used for MG Lite models.""" + + parent_model: Optional[PublisherModelResourceReference] = Field( + default=None, description="""Required. Parent model reference.""" + ) + + +class PublisherModelCallToActionDeployMonetizedModelSingleTenantDict( + TypedDict, total=False +): + """Single Tenant Model configuration. This is used for MG Lite models.""" + + parent_model: Optional[PublisherModelResourceReferenceDict] + """Required. Parent model reference.""" + + +PublisherModelCallToActionDeployMonetizedModelSingleTenantOrDict = Union[ + PublisherModelCallToActionDeployMonetizedModelSingleTenant, + PublisherModelCallToActionDeployMonetizedModelSingleTenantDict, +] + + +class PublisherModelCallToActionDeployMonetizedModelMonetizationConfig( + _common.BaseModel +): + """Configurations related to Marketplace Monetizable models.""" + + monetizable_deployment_units: Optional[ + list[PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnit] + ] = Field( + default=None, + description="""Optional. Deployment Units for a Marketplace Published model.""", + ) + marketplace_listings: Optional[ + list[PublisherModelCallToActionDeployMonetizedModelMarketplaceListing] + ] = Field( + default=None, + description="""Required. Listing Id and Listing Plan Id, seeded by Marketplace.""", + ) + single_tenant: Optional[ + PublisherModelCallToActionDeployMonetizedModelSingleTenant + ] = Field( + default=None, + description="""Optional. Single Tenant configuration, used for MG Lite models.""", + ) + + +class PublisherModelCallToActionDeployMonetizedModelMonetizationConfigDict( + TypedDict, total=False +): + """Configurations related to Marketplace Monetizable models.""" + + monetizable_deployment_units: Optional[ + list[ + PublisherModelCallToActionDeployMonetizedModelMonetizableDeploymentUnitDict + ] + ] + """Optional. Deployment Units for a Marketplace Published model.""" + + marketplace_listings: Optional[ + list[PublisherModelCallToActionDeployMonetizedModelMarketplaceListingDict] + ] + """Required. Listing Id and Listing Plan Id, seeded by Marketplace.""" + + single_tenant: Optional[ + PublisherModelCallToActionDeployMonetizedModelSingleTenantDict + ] + """Optional. Single Tenant configuration, used for MG Lite models.""" + + +PublisherModelCallToActionDeployMonetizedModelMonetizationConfigOrDict = Union[ + PublisherModelCallToActionDeployMonetizedModelMonetizationConfig, + PublisherModelCallToActionDeployMonetizedModelMonetizationConfigDict, +] + + +class PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfig( + _common.BaseModel +): + """Configurations for Consumer level Access.""" + + allow_consumer_console_logging: Optional[bool] = Field( + default=None, + description="""Optional. A boolean to allower consume to access console loggginig.""", + ) + + +class PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfigDict( + TypedDict, total=False +): + """Configurations for Consumer level Access.""" + + allow_consumer_console_logging: Optional[bool] + """Optional. A boolean to allower consume to access console loggginig.""" + + +PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfigOrDict = Union[ + PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfig, + PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfigDict, +] + + +class PublisherModelCallToActionDeployMonetizedModel(_common.BaseModel): + """Deploy a Monetized Model to Marketplace.""" + + monetization_config: Optional[ + PublisherModelCallToActionDeployMonetizedModelMonetizationConfig + ] = Field(default=None, description="""Optional. Specific to Project Pinnacle.""") + consumer_access_config: Optional[ + PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfig + ] = Field( + default=None, + description="""Optional. Configurations for Consumer level Access.""", + ) + service_account: Optional[str] = Field( + default=None, + description="""Optional. Publisher provided service account to launch model container with.""", + ) + + +class PublisherModelCallToActionDeployMonetizedModelDict(TypedDict, total=False): + """Deploy a Monetized Model to Marketplace.""" + + monetization_config: Optional[ + PublisherModelCallToActionDeployMonetizedModelMonetizationConfigDict + ] + """Optional. Specific to Project Pinnacle.""" + + consumer_access_config: Optional[ + PublisherModelCallToActionDeployMonetizedModelConsumerAccessConfigDict + ] + """Optional. Configurations for Consumer level Access.""" + + service_account: Optional[str] + """Optional. Publisher provided service account to launch model container with.""" + + +PublisherModelCallToActionDeployMonetizedModelOrDict = Union[ + PublisherModelCallToActionDeployMonetizedModel, + PublisherModelCallToActionDeployMonetizedModelDict, +] + + class PublisherModelCallToAction(_common.BaseModel): """Actions could take on this Publisher Model.""" @@ -23360,6 +24629,12 @@ class PublisherModelCallToAction(_common.BaseModel): view_rest_api: Optional[PublisherModelCallToActionViewRestApi] = Field( default=None, description="""Optional. To view Rest API docs.""" ) + deploy_monetized_model: Optional[PublisherModelCallToActionDeployMonetizedModel] = ( + Field( + default=None, + description="""Optional. Deploy action for Marketplace model.""", + ) + ) class PublisherModelCallToActionDict(TypedDict, total=False): @@ -23419,6 +24694,9 @@ class PublisherModelCallToActionDict(TypedDict, total=False): view_rest_api: Optional[PublisherModelCallToActionViewRestApiDict] """Optional. To view Rest API docs.""" + deploy_monetized_model: Optional[PublisherModelCallToActionDeployMonetizedModelDict] + """Optional. Deploy action for Marketplace model.""" + PublisherModelCallToActionOrDict = Union[ PublisherModelCallToAction, PublisherModelCallToActionDict @@ -23880,7 +25158,7 @@ class FeedbackEntry(_common.BaseModel): ) feedback_labels: Optional[list[str]] = Field( default=None, - description="""Optional. Specific labels for feedback (non-factual, offensive, etc.).""", + description="""Optional. Specific labels for feedback. At this point, this field is accepted only the restricted set of values: "inaccurate", "incomplete", "unhelpful", "off_topic", "too_slow", "unsafe", "hallucination", "tool_error".""", ) feedback_text: Optional[str] = Field( default=None, @@ -23923,7 +25201,7 @@ class FeedbackEntryDict(TypedDict, total=False): """Required. The ID of the event to which the feedback relates to.""" feedback_labels: Optional[list[str]] - """Optional. Specific labels for feedback (non-factual, offensive, etc.).""" + """Optional. Specific labels for feedback. At this point, this field is accepted only the restricted set of values: "inaccurate", "incomplete", "unhelpful", "off_topic", "too_slow", "unsafe", "hallucination", "tool_error".""" feedback_text: Optional[str] """Optional. Qualitative free-form comments provided by the user.""" @@ -25464,1113 +26742,1317 @@ 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 - -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 BatchPredictionResourceUsageAssessmentResult(_common.BaseModel): + """Result of batch prediction resource usage assessment.""" -RubricVerdict = evals_types.RubricVerdict -RubricVerdictDict = evals_types.RubricVerdictDict -RubricVerdictDictOrDict = evals_types.RubricVerdictOrDict + 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.""" + ) -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 BatchPredictionResourceUsageAssessmentResultDict(TypedDict, total=False): + """Result of batch prediction resource usage assessment.""" -Message = evals_types.Message -MessageDict = evals_types.MessageDict -MessageDictOrDict = evals_types.MessageOrDict + token_count: Optional[int] + """Number of tokens in the dataset.""" -Importance = evals_types.Importance + audio_token_count: Optional[int] + """Number of audio tokens in the dataset.""" -class AgentEngineDict(TypedDict, total=False): - """An agent engine instance.""" +BatchPredictionResourceUsageAssessmentResultOrDict = Union[ + BatchPredictionResourceUsageAssessmentResult, + BatchPredictionResourceUsageAssessmentResultDict, +] - api_client: Optional[Any] - """The underlying API client.""" - api_async_client: Optional[Any] - """The underlying API client for asynchronous operations.""" +class BatchPredictionValidationAssessmentResult(_common.BaseModel): + """Result of batch prediction validation assessment.""" - api_resource: Optional[ReasoningEngineDict] - """The underlying API resource (i.e. ReasoningEngine).""" + errors: Optional[list[str]] = Field( + default=None, description="""The list of errors found in the dataset.""" + ) -AgentEngineOrDict = Union[AgentEngine, AgentEngineDict] +class BatchPredictionValidationAssessmentResultDict(TypedDict, total=False): + """Result of batch prediction validation assessment.""" + errors: Optional[list[str]] + """The list of errors found in the dataset.""" -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. +BatchPredictionValidationAssessmentResultOrDict = Union[ + BatchPredictionValidationAssessmentResult, + BatchPredictionValidationAssessmentResultDict, +] - It must be a valid GCS bucket name, e.g. "gs://bucket-name". It is - required if `agent_engine` is specified.""", - ) - requirements: Optional[Any] = Field( - default=None, - description="""The set of PyPI dependencies needed. - It can either be the path to a single file (requirements.txt), or an - ordered list of strings corresponding to each line of the requirements - file.""", - ) - display_name: Optional[str] = Field( - default=None, - description="""The user-defined name of the Agent Engine. +class TuningResourceUsageAssessmentResult(_common.BaseModel): + """Result of tuning resource usage assessment.""" - 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.""", + token_count: Optional[int] = Field( + default=None, description="""The number of tokens in the dataset.""" ) - extra_packages: Optional[list[str]] = Field( + billable_character_count: Optional[int] = Field( default=None, - description="""The set of extra user-provided packages (if any).""", + description="""The number of billable characters in the dataset.""", ) - 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.""" - ) +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.""" -AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict] + prompt_data: Optional["PromptDataDict"] + """""" -class RunQueryJobAgentEngineConfig(_common.BaseModel): - """Config for checking a query job on an agent engine.""" +PromptOrDict = Union[Prompt, PromptDict] - http_options: Optional[genai_types.HttpOptions] = Field( - default=None, description="""Used to override HTTP request options.""" - ) - query: Optional[str] = Field( - default=None, description="""The query to send to the agent engine.""" - ) - output_gcs_uri: Optional[str] = Field( - default=None, - description="""The GCS URI to use for the output. - If it is a file, the system use this file to store the response. - If it represents a directory, the system automatically generate a file - for the response. - In both cases, the input query will be stored in the same directory under - the same file name prefix as the output file.""", - ) +class SchemaPromptInstanceVariableValue(_common.BaseModel): + """Represents a prompt instance variable.""" -class RunQueryJobAgentEngineConfigDict(TypedDict, total=False): - """Config for checking a query job on an agent engine.""" + part_list: Optional[SchemaPromptSpecPartList] = Field( + default=None, description="""The parts of the variable value.""" + ) - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" - query: Optional[str] - """The query to send to the agent engine.""" +class SchemaPromptInstanceVariableValueDict(TypedDict, total=False): + """Represents a prompt instance variable.""" - 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.""" + part_list: Optional[SchemaPromptSpecPartListDict] + """The parts of the variable value.""" -RunQueryJobAgentEngineConfigOrDict = Union[ - RunQueryJobAgentEngineConfig, RunQueryJobAgentEngineConfigDict +SchemaPromptInstanceVariableValueOrDict = Union[ + SchemaPromptInstanceVariableValue, SchemaPromptInstanceVariableValueDict ] -class RunQueryJobResult(_common.BaseModel): - """Result of running a query job.""" +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.""" ) - job_name: Optional[str] = Field( + prompt_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. 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.""" + 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.""", + ) + 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.""", ) -class RunQueryJobResultDict(TypedDict, total=False): - """Result of running a query job.""" +class CreatePromptConfigDict(TypedDict, total=False): + """Config for creating a prompt.""" 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.""" + prompt_display_name: Optional[str] + """The display name for the prompt. If not set, a default name with a timestamp will be used.""" - input_gcs_uri: Optional[str] - """The GCS URI of the input file.""" + timeout: Optional[int] + """The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""" - output_gcs_uri: Optional[str] - """The GCS URI of the output file.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""" + version_display_name: Optional[str] + """The display name for the prompt version. If not set, a default name with a timestamp will be used.""" -RunQueryJobResultOrDict = Union[RunQueryJobResult, RunQueryJobResultDict] +CreatePromptConfigOrDict = Union[CreatePromptConfig, CreatePromptConfigDict] -class CheckQueryJobResponse(_common.BaseModel): - """Response from LRO.""" + +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.""" ) - output_gcs_uri: Optional[str] = Field( - default=None, description="""The GCS URI of the output file.""" + version_display_name: Optional[str] = Field( + default=None, + description="""The display name for the prompt version. If not set, a default name with a timestamp will be used.""", + ) + timeout: Optional[int] = Field( + default=90, + description="""The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""", + ) + prompt_display_name: Optional[str] = Field( + default=None, + description="""The display name for the prompt. If not set, a default name with a timestamp will be used.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""", ) -class CheckQueryJobResponseDict(TypedDict, total=False): - """Response from LRO.""" +class CreatePromptVersionConfigDict(TypedDict, total=False): + """Config for creating a prompt version.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" - output_gcs_uri: Optional[str] - """The GCS URI of the output file.""" + 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.""" -CheckQueryJobResponseOrDict = Union[CheckQueryJobResponse, CheckQueryJobResponseDict] + 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 AssembleDataset(_common.BaseModel): - """Represents the assembled dataset.""" - bigquery_destination: Optional[str] = Field( - default=None, - description="""The BigQuery destination of the assembled dataset.""", +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 AssembleDatasetDict(TypedDict, total=False): - """Represents the assembled dataset.""" +class GetPromptConfigDict(TypedDict, total=False): + """Config for getting a prompt.""" - bigquery_destination: Optional[str] - """The BigQuery destination of the assembled dataset.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" -AssembleDatasetOrDict = Union[AssembleDataset, AssembleDatasetDict] +GetPromptConfigOrDict = Union[GetPromptConfig, GetPromptConfigDict] -class BatchPredictionResourceUsageAssessmentResult(_common.BaseModel): - """Result of batch prediction resource usage assessment.""" +class PromptRef(_common.BaseModel): + """Reference to a prompt.""" - 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.""" - ) + prompt_id: Optional[str] = Field(default=None, description="""""") + model: Optional[str] = Field(default=None, description="""""") -class BatchPredictionResourceUsageAssessmentResultDict(TypedDict, total=False): - """Result of batch prediction resource usage assessment.""" +class PromptRefDict(TypedDict, total=False): + """Reference to a prompt.""" - token_count: Optional[int] - """Number of tokens in the dataset.""" + prompt_id: Optional[str] + """""" - audio_token_count: Optional[int] - """Number of audio tokens in the dataset.""" + model: Optional[str] + """""" -BatchPredictionResourceUsageAssessmentResultOrDict = Union[ - BatchPredictionResourceUsageAssessmentResult, - BatchPredictionResourceUsageAssessmentResultDict, -] +PromptRefOrDict = Union[PromptRef, PromptRefDict] -class BatchPredictionValidationAssessmentResult(_common.BaseModel): - """Result of batch prediction validation assessment.""" +class PromptVersionRef(_common.BaseModel): + """Reference to a prompt version.""" - errors: Optional[list[str]] = Field( - default=None, description="""The list of errors found in the dataset.""" - ) + prompt_id: Optional[str] = Field(default=None, description="""""") + version_id: Optional[str] = Field(default=None, description="""""") + model: Optional[str] = Field(default=None, description="""""") -class BatchPredictionValidationAssessmentResultDict(TypedDict, total=False): - """Result of batch prediction validation assessment.""" +class PromptVersionRefDict(TypedDict, total=False): + """Reference to a prompt version.""" - errors: Optional[list[str]] - """The list of errors found in the dataset.""" + prompt_id: Optional[str] + """""" + version_id: Optional[str] + """""" -BatchPredictionValidationAssessmentResultOrDict = Union[ - BatchPredictionValidationAssessmentResult, - BatchPredictionValidationAssessmentResultDict, -] + model: Optional[str] + """""" -class TuningResourceUsageAssessmentResult(_common.BaseModel): - """Result of tuning resource usage assessment.""" +PromptVersionRefOrDict = Union[PromptVersionRef, PromptVersionRefDict] - token_count: Optional[int] = Field( - default=None, description="""The number of tokens in the dataset.""" + +class OptimizeJobConfig(_common.BaseModel): + """VAPO Prompt Optimizer Config.""" + + config_path: Optional[str] = Field( + default=None, + description="""The gcs path to the config file, e.g. gs://bucket/config.json.""", ) - 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 TuningValidationAssessmentResultDict(TypedDict, total=False): - """The result of a tuning validation assessment.""" +class ListDeployableModelsConfig(_common.BaseModel): + """Config for listing deployable models.""" - errors: Optional[list[str]] - """The list of errors found in the dataset.""" + 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.""" + ) -TuningValidationAssessmentResultOrDict = Union[ - TuningValidationAssessmentResult, TuningValidationAssessmentResultDict +class ListDeployableModelsConfigDict(TypedDict, total=False): + """Config for listing deployable models.""" + + include_hugging_face_models: Optional[bool] + """Whether to list Hugging Face models.""" + + model_filter: Optional[str] + """Optional. A string to filter the models by.""" + + +ListDeployableModelsConfigOrDict = Union[ + ListDeployableModelsConfig, ListDeployableModelsConfigDict ] -class Prompt(_common.BaseModel): - """Represents a prompt.""" +class ListModelGardenModelsConfig(_common.BaseModel): + """Config for listing Model Garden models.""" - prompt_data: Optional["PromptData"] = Field(default=None, description="""""") - _dataset: Optional["Dataset"] = PrivateAttr(default=None) - _dataset_version: Optional["DatasetVersion"] = PrivateAttr(default=None) + 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(self) -> "Dataset": - return self._dataset # type: ignore[return-value] - @property - def dataset_version(self) -> "DatasetVersion": - return self._dataset_version # type: ignore[return-value] +class ListModelGardenModelsConfigDict(TypedDict, total=False): + """Config for listing Model Garden models.""" - @property - def prompt_id(self) -> Optional[str]: - """Returns the ID associated with the prompt resource.""" - if self._dataset and self._dataset.name: - return self._dataset.name.split("/")[-1] - elif not self._dataset and ( - self._dataset_version and self._dataset_version.name - ): - return self._dataset_version.name.split("datasets/")[1].split("/")[0] - return None + include_hugging_face_models: Optional[bool] + """Whether to list Hugging Face 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 + model_filter: Optional[str] + """Optional. A string to filter the models by.""" - def assemble_contents(self) -> list[genai_types.Content]: - """Transforms a Prompt object into a list with a single genai_types.Content object. - 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. +ListModelGardenModelsConfigOrDict = Union[ + ListModelGardenModelsConfig, ListModelGardenModelsConfigDict +] - This can be used to call generate_content() in the Gen AI SDK. - Example usage: +class ListPublisherModelDeployOptionsConfig(_common.BaseModel): + """Config for listing the deploy options of a publisher model.""" - 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"), - }, - ], - ), - ) + 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.""", + ) - from google import genai - 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 ListPublisherModelDeployOptionsConfigDict(TypedDict, total=False): + """Config for listing the deploy options of a publisher model.""" - 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 [] + 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 not self.prompt_data.variables: - return self.prompt_data.contents + 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.""" - 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." - ) + 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.""" - 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] + 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.""" - 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 - if not has_placeholders: - return [genai_types.Content(parts=parts_to_process)] +ListPublisherModelDeployOptionsConfigOrDict = Union[ + ListPublisherModelDeployOptionsConfig, ListPublisherModelDeployOptionsConfigDict +] - all_rendered_parts: list[genai_types.Part] = [] - for var_dict in self.prompt_data.variables: - for template_item in parts_to_process: - template_part = ( - template_item - if isinstance(template_item, genai_types.Part) - else genai_types.Part(text=str(template_item)) - ) - if template_part.text: - rendered_text = template_part.text +class ListCustomModelDeployOptionsConfig(_common.BaseModel): + """Config for listing custom model deploy options.""" - for key, value in var_dict.items(): - placeholder = f"{{{key}}}" - replacement_text = None + filter_by_user_quota: Optional[bool] = Field( + default=True, + description="""Whether to filter recommendations to regions with user quota. - 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")] + 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. + + 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). + """, + ) -PromptData = SchemaPromptSpecPromptMessage -PromptDataDict = SchemaPromptSpecPromptMessageDict -PromptDataOrDict = Union[PromptData, PromptDataDict] +class ListCustomModelDeployOptionsConfigDict(TypedDict, total=False): + """Config for listing custom model deploy options.""" + + filter_by_user_quota: Optional[bool] + """Whether to filter recommendations to regions with user quota. + + 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. + + 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). + """ + + +ListCustomModelDeployOptionsConfigOrDict = Union[ + ListCustomModelDeployOptionsConfig, ListCustomModelDeployOptionsConfigDict +] + + +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.""" + ) + accelerator_type: Optional[str] = Field( + default=None, description="""The accelerator type.""" + ) + accelerator_count: Optional[int] = Field( + default=None, description="""The number of accelerators.""" + ) + + +class DeployOptionDict(TypedDict, total=False): + """A verified deploy option for a model.""" + + option_name: Optional[str] + """The name of the deploy task.""" + + serving_container_image_uri: Optional[str] + """The URI of the serving container.""" + + machine_type: Optional[str] + """The machine type.""" + + accelerator_type: Optional[str] + """The accelerator type.""" + + accelerator_count: Optional[int] + """The number of accelerators.""" + + +DeployOptionOrDict = Union[DeployOption, DeployOptionDict] + + +class Runtime(_common.BaseModel): + """An agent runtime instance.""" + + api_client: Optional[Any] = Field( + default=None, description="""The underlying API client.""" + ) + api_async_client: Optional[Any] = Field( + default=None, + description="""The underlying API client for asynchronous operations.""", + ) + api_resource: Optional[ReasoningEngine] = Field( + default=None, + description="""The underlying API resource (i.e. ReasoningEngine).""", + ) + + # Allows dynamic binding of methods based on the registered operations. + model_config = ConfigDict(extra="allow") + + def __repr__(self) -> str: + return ( + f"Runtime(api_resource.name='{self.api_resource.name}')" + if self.api_resource is not None + else "Runtime(api_resource.name=None)" + ) + + def operation_schemas(self) -> Optional[list[Dict[str, Any]]]: + """Returns the schemas of all registered operations for the agent.""" + if not isinstance(self.api_resource, ReasoningEngine): + raise ValueError("api_resource is not initialized.") + if not self.api_resource.spec: + raise ValueError("api_resource.spec is not initialized.") + return self.api_resource.spec.class_methods + + def delete( + self, + force: bool = False, + config: Optional[DeleteRuntimeConfigOrDict] = None, + ) -> None: + """Deletes the agent engine. + + Args: + force (bool): + Optional. If set to True, child resources will also be deleted. + Otherwise, the request will fail with FAILED_PRECONDITION error when + the Agent Engine has undeleted child resources. Defaults to False. + config (DeleteRuntimeConfig): + Optional. Additional configurations for deleting the Agent Engine. + """ + if not isinstance(self.api_resource, ReasoningEngine): + raise ValueError("api_resource is not initialized.") + self.api_client.delete(name=self.api_resource.name, force=force, config=config) # type: ignore[union-attr] + + +RubricContentProperty = evals_types.RubricContentProperty +RubricContentPropertyDict = evals_types.RubricContentPropertyDict +RubricContentPropertyDictOrDict = evals_types.RubricContentPropertyOrDict + +RubricContent = evals_types.RubricContent +RubricContentDict = evals_types.RubricContentDict +RubricContentDictOrDict = evals_types.RubricContentOrDict + +Rubric = evals_types.Rubric +RubricDict = evals_types.RubricDict +RubricDictOrDict = evals_types.RubricOrDict + +RubricVerdict = evals_types.RubricVerdict +RubricVerdictDict = evals_types.RubricVerdictDict +RubricVerdictDictOrDict = evals_types.RubricVerdictOrDict + +CandidateResult = evals_types.CandidateResult +CandidateResultDict = evals_types.CandidateResultDict +CandidateResultDictOrDict = evals_types.CandidateResultOrDict + +Event = evals_types.Event +EventDict = evals_types.EventDict +EventDictOrDict = evals_types.EventOrDict + +Message = evals_types.Message +MessageDict = evals_types.MessageDict +MessageDictOrDict = evals_types.MessageOrDict + +Importance = evals_types.Importance + + +class RuntimeDict(TypedDict, total=False): + """An agent runtime instance.""" + + api_client: Optional[Any] + """The underlying API client.""" + + api_async_client: Optional[Any] + """The underlying API client for asynchronous operations.""" + + api_resource: Optional[ReasoningEngineDict] + """The underlying API resource (i.e. ReasoningEngine).""" + + +RuntimeOrDict = Union[Runtime, RuntimeDict] + + +class RuntimeConfig(_common.BaseModel): + """Config for agent runtime methods.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + staging_bucket: Optional[str] = Field( + default=None, + description="""The GCS bucket to use for staging the artifacts needed. + + It must be a valid GCS bucket name, e.g. "gs://bucket-name". It is + required if `agent_engine` is specified.""", + ) + requirements: Optional[Any] = Field( + default=None, + description="""The set of PyPI dependencies needed. + + It can either be the path to a single file (requirements.txt), or an + ordered list of strings corresponding to each line of the requirements + file.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""The user-defined name of the Agent 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. + + If it is a dictionary, the keys are the environment variable names, and + the values are the corresponding values.""", + ) + service_account: Optional[str] = Field( + default=None, + description="""The service account to be used for the Agent Runtime. + + If not specified, the default Reasoning Engine P6SA service agent will be used.""", + ) + identity_type: Optional[IdentityType] = Field( + default=None, description="""The identity type to use for the Agent Runtime.""" + ) + context_spec: Optional[ReasoningEngineContextSpec] = Field( + default=None, + description="""The context spec to be used for the Agent Runtime.""", + ) + psc_interface_config: Optional[PscInterfaceConfig] = Field( + default=None, + description="""The PSC interface config for PSC-I to be used for the Agent Runtime.""", + ) + min_instances: Optional[int] = Field( + default=None, + description="""The minimum number of instances to run for the Agent Runtime. + Defaults to 1. Range: [0, 10]. + """, + ) + max_instances: Optional[int] = Field( + default=None, + description="""The maximum number of instances to run for the Agent Runtime. + Defaults to 100. Range: [1, 1000]. + If VPC-SC or PSC-I is enabled, the acceptable range is [1, 100]. + """, + ) + resource_limits: Optional[dict[str, str]] = Field( + default=None, + description="""The resource limits to be applied to the Agent Runtime. + Required keys: 'cpu' and 'memory'. + Supported values for 'cpu': '1', '2', '4', '6', '8'. + Supported values for 'memory': '1Gi', '2Gi', ..., '32Gi'. + """, + ) + container_concurrency: Optional[int] = Field( + default=None, + description="""The container concurrency to be used for the Agent Runtime. + Recommended value: 2 * cpu + 1. Defaults to 9. + """, + ) + keep_alive_probe: Optional[KeepAliveProbe] = Field( + default=None, + description="""Optional. Specifies the configuration for keep-alive probe. + Contains configuration on a specified endpoint that a deployment host + should use to keep the container alive based on the probe settings.""", + ) + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""The encryption spec to be used for the Agent Runtime.""", + ) + labels: Optional[dict[str, str]] = Field( + default=None, description="""The labels to be used for the Agent Runtime.""" + ) + agent_server_mode: Optional[AgentServerMode] = Field( + default=None, description="""The agent server mode to use for deployment.""" + ) + class_methods: Optional[list[dict[str, Any]]] = Field( + default=None, + description="""The class methods to be used for the Agent Runtime. + If specified, they'll override the class methods that are autogenerated by + default. By default, methods are generated by inspecting the agent object + and generating a corresponding method for each method defined on the + agent class. + """, + ) + source_packages: Optional[list[str]] = Field( + default=None, + description="""The user-provided paths to the source packages (if any). + If specified, the files in the source packages will be packed into a + a tarball file, uploaded to Agent Runtime's API, and deployed to the + Agent Runtime. + The following fields will be ignored: + - agent + - extra_packages + - staging_bucket + - requirements + The following fields will be used to install and use the agent from the + source packages: + - entrypoint_module (required) + - entrypoint_object (required) + - requirements_file (optional) + - class_methods (required) + """, + ) + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfig + ] = Field( + default=None, + description="""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.""", + ) + + +class RuntimeConfigDict(TypedDict, total=False): + """Config for agent runtime methods.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + staging_bucket: Optional[str] + """The GCS bucket to use for staging the artifacts needed. + + It must be a valid GCS bucket name, e.g. "gs://bucket-name". It is + required if `agent_engine` is specified.""" + + requirements: Optional[Any] + """The set of PyPI dependencies needed. + + It can either be the path to a single file (requirements.txt), or an + ordered list of strings corresponding to each line of the requirements + file.""" + + display_name: Optional[str] + """The user-defined name of the Agent Runtime. + + The name can be up to 128 characters long and can comprise any UTF-8 + character.""" + + description: Optional[str] + """The description of the Agent Runtime.""" + + gcs_dir_name: Optional[str] + """The GCS bucket directory under `staging_bucket` to use for staging + the artifacts needed.""" + + extra_packages: Optional[list[str]] + """The set of extra user-provided packages (if any).""" + + env_vars: Optional[Any] + """The environment variables to be set when running the Agent Runtime. + + If it is a dictionary, the keys are the environment variable names, and + the values are the corresponding values.""" + + service_account: Optional[str] + """The service account to be used for the Agent Runtime. + + If not specified, the default Reasoning Engine P6SA service agent will be used.""" + + identity_type: Optional[IdentityType] + """The identity type to use for the Agent Runtime.""" -ParsedResponseUnion = Union[ - prompts_types.ParsedResponse, prompts_types.ParsedResponseFewShot -] -ParsedResponseUnionDict = Union[ - prompts_types.ParsedResponseDict, prompts_types.ParsedResponseFewShotDict -] + context_spec: Optional[ReasoningEngineContextSpecDict] + """The context spec to be used for the Agent Runtime.""" + psc_interface_config: Optional[PscInterfaceConfigDict] + """The PSC interface config for PSC-I to be used for the Agent Runtime.""" -class PromptDict(TypedDict, total=False): - """Represents a prompt.""" + min_instances: Optional[int] + """The minimum number of instances to run for the Agent Runtime. + Defaults to 1. Range: [0, 10]. + """ - prompt_data: Optional["PromptDataDict"] - """""" + 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'. + """ -PromptOrDict = Union[Prompt, PromptDict] + container_concurrency: Optional[int] + """The container concurrency to be used for the Agent Runtime. + Recommended value: 2 * cpu + 1. Defaults to 9. + """ + keep_alive_probe: Optional[KeepAliveProbeDict] + """Optional. Specifies the configuration for keep-alive probe. + Contains configuration on a specified endpoint that a deployment host + should use to keep the container alive based on the probe settings.""" -class SchemaPromptInstanceVariableValue(_common.BaseModel): - """Represents a prompt instance variable.""" + encryption_spec: Optional[genai_types.EncryptionSpec] + """The encryption spec to be used for the Agent Runtime.""" - part_list: Optional[SchemaPromptSpecPartList] = Field( - default=None, description="""The parts of the variable value.""" - ) + 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 SchemaPromptInstanceVariableValueDict(TypedDict, total=False): - """Represents a prompt instance variable.""" + 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. + """ - part_list: Optional[SchemaPromptSpecPartListDict] - """The parts of the variable value.""" + source_packages: Optional[list[str]] + """The user-provided paths to the source packages (if any). + If specified, the files in the source packages will be packed into a + a tarball file, uploaded to Agent Runtime's API, and deployed to the + Agent Runtime. + The following fields will be ignored: + - agent + - extra_packages + - staging_bucket + - requirements + The following fields will be used to install and use the agent from the + source packages: + - entrypoint_module (required) + - entrypoint_object (required) + - requirements_file (optional) + - class_methods (required) + """ + developer_connect_source: Optional[ + ReasoningEngineSpecSourceCodeSpecDeveloperConnectConfigDict + ] + """Specifies the configuration for fetching source code from a Git repository that is managed by Developer Connect. This includes the repository, revision, and directory to use.""" -SchemaPromptInstanceVariableValueOrDict = Union[ - SchemaPromptInstanceVariableValue, SchemaPromptInstanceVariableValueDict -] + entrypoint_module: Optional[str] + """The entrypoint module to be used for the Agent Runtime + This field only used when source_packages is specified.""" + entrypoint_object: Optional[str] + """The entrypoint object to be used for the Agent Runtime. + This field only used when source_packages is specified.""" -class CreatePromptConfig(_common.BaseModel): - """Config for creating a prompt.""" + 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. + """ - 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.""", - ) - 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.""", - ) - 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.""", - ) + agent_framework: Optional[ + Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + ] + """The agent framework to be used for the Agent Runtime. + The OSS agent framework used to develop the agent. + Currently supported values: "google-adk", "langchain", "langgraph", + "ag2", "llama-index", "custom". + If not specified: + - If `agent` is specified, the agent framework will be auto-detected. + - If `source_packages` is specified, the agent framework will + default to "custom".""" + python_version: Optional[Literal["3.10", "3.11", "3.12", "3.13", "3.14"]] + """The Python version to be used for the Agent Runtime. + If not specified, it will use the current Python version of the environment. + Supported versions: "3.10", "3.11", "3.12", "3.13", "3.14". + """ -class CreatePromptConfigDict(TypedDict, total=False): - """Config for creating a prompt.""" + 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`. + """ - http_options: Optional[genai_types.HttpOptions] - """Used to override HTTP request options.""" + image_spec: Optional[ReasoningEngineSpecSourceCodeSpecImageSpecDict] + """The image spec for the Agent Runtime.""" - prompt_display_name: Optional[str] - """The display name for the prompt. If not set, a default name with a timestamp will be used.""" + agent_config_source: Optional[ + ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict + ] + """The agent config source for the Agent Runtime.""" - timeout: Optional[int] - """The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""" + traffic_config: Optional[ReasoningEngineTrafficConfigDict] + """The traffic config for the Agent Runtime.""" - 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.""" + container_spec: Optional[ReasoningEngineSpecContainerSpecDict] + """The container spec for the Agent Runtime.""" - version_display_name: Optional[str] - """The display name for the prompt version. If not set, a default name with a timestamp will be used.""" + agent_gateway_config: Optional[ + ReasoningEngineSpecDeploymentSpecAgentGatewayConfigDict + ] + """Agent Gateway configuration for a Agent Runtime deployment.""" -CreatePromptConfigOrDict = Union[CreatePromptConfig, CreatePromptConfigDict] +RuntimeConfigOrDict = Union[RuntimeConfig, RuntimeConfigDict] -class CreatePromptVersionConfig(_common.BaseModel): - """Config for creating a prompt version.""" +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.""" ) - version_display_name: Optional[str] = Field( - default=None, - description="""The display name for the prompt version. If not set, a default name with a timestamp will be used.""", - ) - timeout: Optional[int] = Field( - default=90, - description="""The timeout for the create_version request in seconds. If not set, the default timeout is 90 seconds.""", - ) - prompt_display_name: Optional[str] = Field( - default=None, - description="""The display name for the prompt. If not set, a default name with a timestamp will be used.""", + query: Optional[str] = Field( + default=None, description="""The query to send to the agent runtime.""" ) - encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + output_gcs_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="""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 CreatePromptVersionConfigDict(TypedDict, total=False): - """Config for creating a prompt version.""" +class RunQueryJobRuntimeConfigDict(TypedDict, total=False): + """Config for checking a query job on an agent runtime.""" 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.""" + query: Optional[str] + """The query to send to the agent runtime.""" - encryption_spec: Optional[genai_types.EncryptionSpec] - """Customer-managed encryption key spec for a prompt dataset. If set, this prompt dataset and all sub-resources of this prompt dataset will be secured by this key.""" + output_gcs_uri: Optional[str] + """The GCS URI 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.""" -CreatePromptVersionConfigOrDict = Union[ - CreatePromptVersionConfig, CreatePromptVersionConfigDict +RunQueryJobRuntimeConfigOrDict = Union[ + RunQueryJobRuntimeConfig, RunQueryJobRuntimeConfigDict ] -class GetPromptConfig(_common.BaseModel): - """Config for getting a prompt.""" +class RunQueryJobResult(_common.BaseModel): + """Result of running a query job.""" http_options: Optional[genai_types.HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) + job_name: Optional[str] = Field( + default=None, + description="""Name of the agent runtime operation to later check for status.""", + ) + input_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the input file.""" + ) + output_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the output file.""" + ) -class GetPromptConfigDict(TypedDict, total=False): - """Config for getting a prompt.""" +class RunQueryJobResultDict(TypedDict, total=False): + """Result of running a query job.""" http_options: Optional[genai_types.HttpOptions] """Used to override HTTP request options.""" + job_name: Optional[str] + """Name of the agent runtime operation to later check for status.""" -GetPromptConfigOrDict = Union[GetPromptConfig, GetPromptConfigDict] - - -class PromptRef(_common.BaseModel): - """Reference to a prompt.""" - - prompt_id: Optional[str] = Field(default=None, description="""""") - model: Optional[str] = Field(default=None, description="""""") - - -class PromptRefDict(TypedDict, total=False): - """Reference to a prompt.""" - - prompt_id: Optional[str] - """""" - - model: Optional[str] - """""" - - -PromptRefOrDict = Union[PromptRef, PromptRefDict] - - -class PromptVersionRef(_common.BaseModel): - """Reference to a prompt version.""" - - prompt_id: Optional[str] = Field(default=None, description="""""") - version_id: Optional[str] = Field(default=None, description="""""") - model: Optional[str] = Field(default=None, description="""""") - - -class PromptVersionRefDict(TypedDict, total=False): - """Reference to a prompt version.""" - - prompt_id: Optional[str] - """""" - - version_id: Optional[str] - """""" + input_gcs_uri: Optional[str] + """The GCS URI of the input file.""" - model: Optional[str] - """""" + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" -PromptVersionRefOrDict = Union[PromptVersionRef, PromptVersionRefDict] +RunQueryJobResultOrDict = Union[RunQueryJobResult, RunQueryJobResultDict] -class OptimizeJobConfig(_common.BaseModel): - """VAPO Prompt Optimizer Config.""" +class CheckQueryJobResponse(_common.BaseModel): + """Response from LRO.""" - config_path: Optional[str] = Field( - default=None, - description="""The gcs path to the config file, e.g. gs://bucket/config.json.""", - ) - service_account: Optional[str] = Field( - default=None, - description="""The service account to use for the custom job. Cannot be provided at the same time as service_account_project_number.""", - ) - service_account_project_number: Optional[Union[int, str]] = Field( - default=None, - description="""The project number used to construct the default service account:{service_account_project_number}-compute@developer.gserviceaccount.comCannot be provided at the same time as "service_account".""", - ) - wait_for_completion: Optional[bool] = Field( - default=True, - description="""Whether to wait for the job tocomplete. Ignored for async jobs.""", + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" ) - 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.""", + output_gcs_uri: Optional[str] = Field( + default=None, description="""The GCS URI of the output file.""" ) -class OptimizeJobConfigDict(TypedDict, total=False): - """VAPO Prompt Optimizer Config.""" - - config_path: Optional[str] - """The gcs path to the config file, e.g. gs://bucket/config.json.""" - - service_account: Optional[str] - """The service account to use for the custom job. Cannot be provided at the same time as service_account_project_number.""" - - service_account_project_number: Optional[Union[int, str]] - """The project number used to construct the default service account:{service_account_project_number}-compute@developer.gserviceaccount.comCannot be provided at the same time as "service_account".""" +class CheckQueryJobResponseDict(TypedDict, total=False): + """Response from LRO.""" - wait_for_completion: Optional[bool] - """Whether to wait for the job tocomplete. Ignored for async jobs.""" + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" - 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.""" + output_gcs_uri: Optional[str] + """The GCS URI of the output file.""" -OptimizeJobConfigOrDict = Union[OptimizeJobConfig, OptimizeJobConfigDict] +CheckQueryJobResponseOrDict = Union[CheckQueryJobResponse, CheckQueryJobResponseDict] -class AgentEngineRuntimeRevision(_common.BaseModel): - """An agent engine runtime revision instance.""" +class RuntimeRevision(_common.BaseModel): + """An agent runtime revision instance.""" api_client: Optional[Any] = Field( default=None, description="""The underlying API client.""" @@ -26589,9 +28071,9 @@ class AgentEngineRuntimeRevision(_common.BaseModel): def __repr__(self) -> str: return ( - f"AgentEngineRuntimeRevision(api_resource.name='{self.api_resource.name}')" + f"RuntimeRevision(api_resource.name='{self.api_resource.name}')" if self.api_resource is not None - else "AgentEngineRuntimeRevision(api_resource.name=None)" + else "RuntimeRevision(api_resource.name=None)" ) def operation_schemas(self) -> Optional[list[Dict[str, Any]]]: @@ -26604,12 +28086,12 @@ def operation_schemas(self) -> Optional[list[Dict[str, Any]]]: def delete( self, - config: Optional[DeleteAgentEngineRuntimeRevisionConfigOrDict] = None, + config: Optional[DeleteRuntimeRevisionConfigOrDict] = None, ) -> None: """Deletes the agent engine runtime revision. Args: - config (DeleteAgentEngineRuntimeRevisionConfig): + config (DeleteRuntimeRevisionConfig): Optional. Additional configurations for deleting the Agent Engine Runtime Revision. """ if not isinstance(self.api_resource, ReasoningEngineRuntimeRevision): @@ -26617,8 +28099,8 @@ def delete( 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.""" +class RuntimeRevisionDict(TypedDict, total=False): + """An agent runtime revision instance.""" api_client: Optional[Any] """The underlying API client.""" @@ -26630,210 +28112,4 @@ class AgentEngineRuntimeRevisionDict(TypedDict, total=False): """The underlying API resource (i.e. ReasoningEngineRuntimeRevision).""" -AgentEngineRuntimeRevisionOrDict = Union[ - AgentEngineRuntimeRevision, AgentEngineRuntimeRevisionDict -] - - -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 ListDeployableModelsConfigDict(TypedDict, total=False): - """Config for listing deployable models.""" - - include_hugging_face_models: Optional[bool] - """Whether to list Hugging Face models.""" - - model_filter: Optional[str] - """Optional. A string to filter the models by.""" - - -ListDeployableModelsConfigOrDict = Union[ - ListDeployableModelsConfig, ListDeployableModelsConfigDict -] - - -class ListModelGardenModelsConfig(_common.BaseModel): - """Config for listing Model Garden models.""" - - include_hugging_face_models: Optional[bool] = Field( - default=None, description="""Whether to list Hugging Face models.""" - ) - model_filter: Optional[str] = Field( - default=None, description="""Optional. A string to filter the models by.""" - ) - - -class ListModelGardenModelsConfigDict(TypedDict, total=False): - """Config for listing Model Garden models.""" - - include_hugging_face_models: Optional[bool] - """Whether to list Hugging Face models.""" - - model_filter: Optional[str] - """Optional. A string to filter the models by.""" - - -ListModelGardenModelsConfigOrDict = Union[ - ListModelGardenModelsConfig, ListModelGardenModelsConfigDict -] - - -class ListPublisherModelDeployOptionsConfig(_common.BaseModel): - """Config for listing the deploy options of a publisher model.""" - - machine_type_filter: Optional[Union[str, list[str]]] = Field( - default=None, - description="""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.""", - ) - - -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.""" - - -ListPublisherModelDeployOptionsConfigOrDict = Union[ - ListPublisherModelDeployOptionsConfig, ListPublisherModelDeployOptionsConfigDict -] - - -class ListCustomModelDeployOptionsConfig(_common.BaseModel): - """Config for listing custom model deploy options.""" - - filter_by_user_quota: Optional[bool] = Field( - default=True, - description="""Whether to filter recommendations to regions with user quota. - - 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. - - 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 ListCustomModelDeployOptionsConfigDict(TypedDict, total=False): - """Config for listing custom model deploy options.""" - - filter_by_user_quota: Optional[bool] - """Whether to filter recommendations to regions with user quota. - - 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. - - 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). - """ - - -ListCustomModelDeployOptionsConfigOrDict = Union[ - ListCustomModelDeployOptionsConfig, ListCustomModelDeployOptionsConfigDict -] - - -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.""" - ) - accelerator_type: Optional[str] = Field( - default=None, description="""The accelerator type.""" - ) - accelerator_count: Optional[int] = Field( - default=None, description="""The number of accelerators.""" - ) - - -class DeployOptionDict(TypedDict, total=False): - """A verified deploy option for a model.""" - - option_name: Optional[str] - """The name of the deploy task.""" - - serving_container_image_uri: Optional[str] - """The URI of the serving container.""" - - machine_type: Optional[str] - """The machine type.""" - - accelerator_type: Optional[str] - """The accelerator type.""" - - accelerator_count: Optional[int] - """The number of accelerators.""" - - -DeployOptionOrDict = Union[DeployOption, DeployOptionDict] +RuntimeRevisionOrDict = Union[RuntimeRevision, RuntimeRevisionDict] diff --git a/agentplatform/agent_engines/__init__.py b/agentplatform/agent_engines/__init__.py deleted file mode 100644 index ac4802a9b2..0000000000 --- a/agentplatform/agent_engines/__init__.py +++ /dev/null @@ -1,428 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -"""Classes and functions for working with agent engines.""" - -from typing import Dict, Iterable, Optional, Sequence, Union - -from google.cloud.aiplatform import base -from google.cloud.aiplatform import initializer -from google.cloud.aiplatform import utils as aip_utils -from google.cloud.aiplatform_v1 import types as aip_types - -# We just want to re-export certain classes -# pylint: disable=g-multiple-import,g-importing-member -from agentplatform.agent_engines._agent_engines import ( - _AgentEngineInterface, - AgentEngine, - Cloneable, - ModuleAgent, - OperationRegistrable, - Queryable, - AsyncQueryable, - StreamQueryable, - AsyncStreamQueryable, -) -from agentplatform.agent_engines.templates.adk import ( - AdkApp, -) -from agentplatform.agent_engines.templates.ag2 import ( - AG2Agent, -) -from agentplatform.agent_engines.templates.langchain import ( - LangchainAgent, -) -from agentplatform.agent_engines.templates.langgraph import ( - LanggraphAgent, -) -from agentplatform.agent_engines.templates.llama_index import ( - LlamaIndexQueryPipelineAgent, -) - - -_LOGGER = base.Logger(__name__) - - -def get(resource_name: str) -> AgentEngine: - """Retrieves an Agent Engine resource. - - Args: - resource_name (str): - Required. A fully-qualified resource name or ID such as - "projects/123/locations/us-central1/reasoningEngines/456" or - "456" when project and location are initialized or passed. - """ - return AgentEngine(resource_name) - - -def create( - agent_engine: Optional[_AgentEngineInterface] = None, - *, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - build_options: Optional[Dict[str, Sequence[str]]] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[aip_types.EncryptionSpec] = None, -) -> AgentEngine: - """Creates a new Agent Engine. - - The Agent Engine will be an instance of the `agent_engine` that - was passed in, running remotely on Vertex AI. - - Sample ``src_dir`` contents (e.g. ``./user_src_dir``): - - .. code-block:: python - - user_src_dir/ - |-- main.py - |-- requirements.txt - |-- user_code/ - | |-- utils.py - | |-- ... - |-- installation_scripts/ - | |-- install_package.sh - | |-- ... - |-- ... - - To build an Agent Engine with the above files, run: - - .. code-block:: python - - remote_agent = agent_engines.create( - agent_engine=local_agent, - requirements=[ - # I.e. the PyPI dependencies listed in requirements.txt - "google-cloud-aiplatform==1.25.0", - "langchain==0.0.242", - ... - ], - extra_packages=[ - "./user_src_dir/main.py", # a single file - "./user_src_dir/user_code", # a directory - ... - ], - build_options={ - "installation": [ - "./user_src_dir/installation_scripts/install_package.sh", - ... - ], - }, - ) - - Args: - agent_engine (AgentEngineInterface): - Required. The Agent Engine to be created. - requirements (Union[str, Sequence[str]]): - Optional. The set of PyPI dependencies needed. It can either be - the path to a single file (requirements.txt), or an ordered list - of strings corresponding to each line of the requirements file. - display_name (str): - Optional. The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any - UTF-8 character. - description (str): - Optional. The description of the Agent Engine. - gcs_dir_name (str): - Optional. The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - extra_packages (Sequence[str]): - Optional. The set of extra user-provided packages (if any). - env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): - Optional. The environment variables to be set when running the - Agent Engine. If it is a list of strings, each string should be - a valid key to `os.environ`. If it is a dictionary, the keys are - the environment variable names, and the values are the - corresponding values. - build_options (Dict[str, Sequence[str]]): - Optional. The build options for the Agent Engine. This includes - options such as installation scripts. - service_account (str): - Optional. The service account to be used for the Agent Engine. If - not specified, the default reasoning engine service agent service - account will be used. - psc_interface_config (PscInterfaceConfig): - Optional. The PSC interface config for the Agent Engine. If not - specified, the default PSC interface config will be used. - min_instances (int): - Optional. The minimum number of instances to run the Agent Engine. - If not specified, the default value will be used. - max_instances (int): - Optional. The maximum number of instances to run the Agent Engine. - If not specified, the default value will be used. - resource_limits (Dict[str, str]): - Optional. The resource limits for the Agent Engine. If not - specified, the default value will be used. - container_concurrency (int): - Optional. The container concurrency for the Agent Engine. If not - specified, the default value will be used. - encryption_spec (EncryptionSpec): - Optional. The encryption spec for the Agent Engine. If not - specified, the default encryption spec will be used. - - Returns: - AgentEngine: The Agent Engine that was created. - - Raises: - ValueError: If the `project` was not set using `agentplatform.init`. - ValueError: If the `location` was not set using `agentplatform.init`. - ValueError: If the `staging_bucket` was not set using agentplatform.init. - ValueError: If the `staging_bucket` does not start with "gs://". - FileNotFoundError: If `extra_packages` includes a file or directory - that does not exist. - IOError: If requirements is a string that corresponds to a - nonexistent file. - """ - return AgentEngine.create( - agent_engine=agent_engine, - requirements=requirements, - display_name=display_name, - description=description, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - env_vars=env_vars, - build_options=build_options, - service_account=service_account, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - encryption_spec=encryption_spec, - ) - - -def list(*, filter: str = "") -> Iterable[AgentEngine]: - """List all instances of Agent Engine matching the filter. - - Example Usage: - - .. code-block:: python - import agentplatform - from agentplatform import agent_engines - - agentplatform.init(project="my_project", location="us-central1") - agent_engines.list(filter='display_name="My Custom Agent"') - - Args: - filter (str): - Optional. An expression for filtering the results of the request. - For field names both snake_case and camelCase are supported. - - Returns: - Iterable[AgentEngine]: An iterable of Agent Engines matching the filter. - """ - api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineClientWithOverride, - ) - for agent in api_client.list_reasoning_engines( - request=aip_types.ListReasoningEnginesRequest( - parent=initializer.global_config.common_location_path(), - filter=filter, - ) - ): - yield AgentEngine(agent.name) - - -def delete( - resource_name: str, - *, - force: bool = False, - **kwargs, -) -> None: - """Delete an Agent Engine resource. - - Args: - resource_name (str): - Required. The name of the Agent Engine to be deleted. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}` - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error - when the Agent Engine has undeleted child resources. Defaults to - False. - **kwargs (dict[str, Any]): - Optional. Additional keyword arguments to pass to the - delete_reasoning_engine method. - """ - api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineClientWithOverride, - ) - _LOGGER.info(f"Deleting AgentEngine resource: {resource_name}") - operation_future = api_client.delete_reasoning_engine( - request=aip_types.DeleteReasoningEngineRequest( - name=resource_name, - force=force, - **(kwargs or {}), - ) - ) - _LOGGER.info(f"Delete AgentEngine backing LRO: {operation_future.operation.name}") - operation_future.result() - _LOGGER.info(f"AgentEngine resource deleted: {resource_name}") - - -def update( - resource_name: str, - *, - agent_engine: Optional[Union[Queryable, OperationRegistrable]] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - build_options: Optional[Dict[str, Sequence[str]]] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[aip_types.EncryptionSpec] = None, -) -> "AgentEngine": - """Updates an existing Agent Engine. - - This method updates the configuration of a deployed Agent Engine, identified - by its resource name. Unlike the `create` function which requires an - `agent_engine` object, all arguments in this method are optional. This - method allows you to modify individual aspects of the configuration by - providing any of the optional arguments. - - Args: - resource_name (str): - Required. The name of the Agent Engine to be updated. Format: - `projects/{project}/locations/{location}/reasoningEngines/{resource_id}`. - agent_engine (AgentEngineInterface): - Optional. The instance to be used as the updated Agent Engine. If it - is not specified, the existing instance will be used. - requirements (Union[str, Sequence[str]]): - Optional. The set of PyPI dependencies needed. It can either be - the path to a single file (requirements.txt), or an ordered list - of strings corresponding to each line of the requirements file. - If it is not specified, the existing requirements will be used. - If it is set to an empty string or list, the existing - requirements will be removed. - display_name (str): - Optional. The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any - UTF-8 character. - description (str): - Optional. The description of the Agent Engine. - gcs_dir_name (str): - Optional. The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - extra_packages (Sequence[str]): - Optional. The set of extra user-provided packages (if any). If - it is not specified, the existing extra packages will be used. - If it is set to an empty list, the existing extra packages will - be removed. - env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): - Optional. The environment variables to be set when running the - Agent Engine. If it is a list of strings, each string should be - a valid key to `os.environ`. If it is a dictionary, the keys are - the environment variable names, and the values are the - corresponding values. - build_options (Dict[str, Sequence[str]]): - Optional. The build options for the Agent Engine. This includes - options such as installation scripts. - service_account (str): - Optional. The service account to be used for the Agent Engine. If - not specified, the default reasoning engine service agent service - account will be used. - min_instances (int): - Optional. The minimum number of instances to run the Agent Engine. - If not specified, the default value will be used. - max_instances (int): - Optional. The maximum number of instances to run the Agent Engine. - If not specified, the default value will be used. - resource_limits (Dict[str, str]): - Optional. The resource limits for the Agent Engine. If not - specified, the default value will be used. - container_concurrency (int): - Optional. The container concurrency for the Agent Engine. If not - specified, the default value will be used. - encryption_spec (EncryptionSpec): - Optional. The encryption spec for the Agent Engine. If not - specified, the default encryption spec will be used. - - Returns: - AgentEngine: The Agent Engine that was updated. - - Raises: - ValueError: If the `staging_bucket` was not set using agentplatform.init. - ValueError: If the `staging_bucket` does not start with "gs://". - FileNotFoundError: If `extra_packages` includes a file or directory - that does not exist. - ValueError: if none of `display_name`, `description`, - `requirements`, `extra_packages`, `agent_engine`, or `build_options` - were specified. - IOError: If requirements is a string that corresponds to a - nonexistent file. - """ - agent = get(resource_name) - return agent.update( - agent_engine=agent_engine, - requirements=requirements, - display_name=display_name, - description=description, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - env_vars=env_vars, - build_options=build_options, - service_account=service_account, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - encryption_spec=encryption_spec, - ) - - -__all__ = ( - # Resources - "AgentEngine", - # Protocols - "Cloneable", - "OperationRegistrable", - "Queryable", - "AsyncQueryable", - "StreamQueryable", - "AsyncStreamQueryable", - # Methods - "create", - "delete", - "get", - "list", - "update", - # Templates - "AdkApp", - "ModuleAgent", - "LangchainAgent", - "LanggraphAgent", - "AG2Agent", - "LlamaIndexQueryPipelineAgent", -) diff --git a/agentplatform/agent_engines/_agent_engines.py b/agentplatform/agent_engines/_agent_engines.py deleted file mode 100644 index 35333b82a4..0000000000 --- a/agentplatform/agent_engines/_agent_engines.py +++ /dev/null @@ -1,2028 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -import abc -import inspect -import io -import json -import logging -import os -import sys -import tarfile -import types -import typing -from typing import ( - Any, - AsyncIterable, - Callable, - Coroutine, - Dict, - Iterable, - List, - Optional, - Protocol, - Sequence, - Tuple, - Union, -) - -from google.api_core import exceptions -from google.cloud import storage -from google.cloud.aiplatform import base -from google.cloud.aiplatform import initializer -from google.cloud.aiplatform import utils as aip_utils -from google.cloud.aiplatform_v1 import types as aip_types -from google.cloud.aiplatform_v1.types import reasoning_engine_service -from agentplatform._genai import _agent_engines_utils -import httpx -import proto - -from google.protobuf import field_mask_pb2 - - -_LOGGER = base.Logger("agentplatform.agent_engines") - -_SUPPORTED_PYTHON_VERSIONS = ("3.10", "3.11", "3.12", "3.13", "3.14") -_DEFAULT_GCS_DIR_NAME = "agent_engine" -_BLOB_FILENAME = "agent_engine.pkl" -_REQUIREMENTS_FILE = "requirements.txt" -_EXTRA_PACKAGES_FILE = "dependencies.tar.gz" -_STANDARD_API_MODE = "" -_ASYNC_API_MODE = "async" -_STREAM_API_MODE = "stream" -_ASYNC_STREAM_API_MODE = "async_stream" -_BIDI_STREAM_API_MODE = "bidi_stream" -_A2A_EXTENSION_MODE = "a2a_extension" -_A2A_AGENT_CARD = "a2a_agent_card" -_MODE_KEY_IN_SCHEMA = "api_mode" -_METHOD_NAME_KEY_IN_SCHEMA = "name" -_DEFAULT_METHOD_NAME = "query" -_DEFAULT_ASYNC_METHOD_NAME = "async_query" -_DEFAULT_STREAM_METHOD_NAME = "stream_query" -_DEFAULT_ASYNC_STREAM_METHOD_NAME = "async_stream_query" -_DEFAULT_METHOD_RETURN_TYPE = "dict[str, Any]" -_DEFAULT_ASYNC_METHOD_RETURN_TYPE = "Coroutine[Any, Any, Any]" -_DEFAULT_STREAM_METHOD_RETURN_TYPE = "Iterable[Any]" -_DEFAULT_ASYNC_STREAM_METHOD_RETURN_TYPE = "AsyncIterable[Any]" -_DEFAULT_METHOD_DOCSTRING_TEMPLATE = """ - Runs the Agent Engine to serve the user request. - This will be based on the `.{method_name}(...)` of the python object that - was passed in when creating the Agent Engine. The method will invoke the - `{default_method_name}` API client of the python object. - Args: - **kwargs: - Optional. The arguments of the `.{method_name}(...)` method. - Returns: - {return_type}: The response from serving the user request. -""" -_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE = ( - "Failed to register API methods. Please follow the guide to " - "register the API methods: " - "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#custom-methods. " - "Error: {%s}" -) -_AGENT_FRAMEWORK_ATTR = "agent_framework" -_DEFAULT_AGENT_FRAMEWORK = "custom" -_BUILD_OPTIONS_INSTALLATION = "installation_scripts" -_DEFAULT_METHOD_NAME_MAP = { - _STANDARD_API_MODE: _DEFAULT_METHOD_NAME, - _ASYNC_API_MODE: _DEFAULT_ASYNC_METHOD_NAME, - _STREAM_API_MODE: _DEFAULT_STREAM_METHOD_NAME, - _ASYNC_STREAM_API_MODE: _DEFAULT_ASYNC_STREAM_METHOD_NAME, -} -_DEFAULT_METHOD_RETURN_TYPE_MAP = { - _STANDARD_API_MODE: _DEFAULT_METHOD_RETURN_TYPE, - _ASYNC_API_MODE: _DEFAULT_ASYNC_METHOD_RETURN_TYPE, - _STREAM_API_MODE: _DEFAULT_STREAM_METHOD_RETURN_TYPE, - _ASYNC_STREAM_API_MODE: _DEFAULT_ASYNC_STREAM_METHOD_RETURN_TYPE, -} - - -try: - from google.adk.agents import BaseAgent - - ADKAgent = BaseAgent -except (ImportError, AttributeError): - ADKAgent = None - -try: - from a2a.types import ( - AgentCard, - AgentInterface, - Message, - TaskIdParams, - TaskQueryParams, - ) - from a2a.utils.constants import TransportProtocol, PROTOCOL_VERSION_CURRENT - from a2a.client import ClientConfig, ClientFactory - - AgentCard = AgentCard - AgentInterface = AgentInterface - TransportProtocol = TransportProtocol - PROTOCOL_VERSION_CURRENT = PROTOCOL_VERSION_CURRENT - Message = Message - ClientConfig = ClientConfig - ClientFactory = ClientFactory - TaskIdParams = TaskIdParams - TaskQueryParams = TaskQueryParams -except (ImportError, AttributeError): - AgentCard = None - AgentInterface = None - TransportProtocol = None - PROTOCOL_VERSION_CURRENT = None - Message = None - ClientConfig = None - ClientFactory = None - TaskIdParams = None - TaskQueryParams = None - - -@typing.runtime_checkable -class Queryable(Protocol): - """Protocol for Agent Engines that can be queried.""" - - @abc.abstractmethod - def query(self, **kwargs) -> Any: - """Runs the Agent Engine to serve the user query.""" - - -@typing.runtime_checkable -class AsyncQueryable(Protocol): - """Protocol for Agent Engines that can be queried asynchronously.""" - - @abc.abstractmethod - def async_query(self, **kwargs) -> Coroutine[Any, Any, Any]: - """Runs the Agent Engine to serve the user query asynchronously.""" - - -@typing.runtime_checkable -class AsyncStreamQueryable(Protocol): - """Protocol for Agent Engines that can stream responses asynchronously.""" - - @abc.abstractmethod - async def async_stream_query(self, **kwargs) -> AsyncIterable[Any]: - """Asynchronously stream responses to serve the user query.""" - - -@typing.runtime_checkable -class StreamQueryable(Protocol): - """Protocol for Agent Engines that can stream responses.""" - - @abc.abstractmethod - def stream_query(self, **kwargs) -> Iterable[Any]: - """Stream responses to serve the user query.""" - - -@typing.runtime_checkable -class BidiStreamQueryable(Protocol): - """Protocol for Agent Engines that can stream requests and responses.""" - - @abc.abstractmethod - async def bidi_stream_query(self, **kwargs) -> AsyncIterable[Any]: - """Asynchronously stream requests and responses to serve the user query.""" - - -@typing.runtime_checkable -class Cloneable(Protocol): - """Protocol for Agent Engines that can be cloned.""" - - @abc.abstractmethod - def clone(self) -> Any: - """Return a clone of the object.""" - - -@typing.runtime_checkable -class OperationRegistrable(Protocol): - """Protocol for agents that have registered operations.""" - - @abc.abstractmethod - def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]: - """Register the user provided operations (modes and methods).""" - - -_AgentEngineInterface = Union[ - ADKAgent, - AsyncQueryable, - AsyncStreamQueryable, - BidiStreamQueryable, - OperationRegistrable, - Queryable, - StreamQueryable, -] - - -def _wrap_agent_operation(agent: Any, operation: str): - """Wraps an agent operation into a method (works for all API modes).""" - - def _method(self, **kwargs): - if not self._tmpl_attrs.get("agent"): - self.set_up() - return getattr(self._tmpl_attrs["agent"], operation)(**kwargs) - - _method.__name__ = operation - _method.__doc__ = getattr(agent, operation).__doc__ - return _method - - -class ModuleAgent(Cloneable, OperationRegistrable): - """Agent that is defined by a module and an agent name. - - This agent is instantiated by importing a module and instantiating an agent - from that module. It also allows to register operations that are defined in - the agent. - """ - - def __init__( - self, - *, - module_name: str, - agent_name: str, - register_operations: Dict[str, Sequence[str]], - sys_paths: Optional[Sequence[str]] = None, - agent_framework: Optional[str] = None, - ): - """Initializes a module-based agent. - - Args: - module_name (str): - Required. The name of the module to import. - agent_name (str): - Required. The name of the agent in the module to instantiate. - register_operations (Dict[str, Sequence[str]]): - Required. A dictionary of API modes to a list of method names. - sys_paths (Sequence[str]): - Optional. The system paths to search for the module. It should - be relative to the directory where the code will be running. - I.e. it should correspond to the directory being passed to - `extra_packages=...` in the create method. It will be appended - to the system path in the sequence being specified here, and - only be appended if it is not already in the system path. - """ - self.agent_framework = agent_framework - self._tmpl_attrs = { - "module_name": module_name, - "agent_name": agent_name, - "register_operations": register_operations, - "sys_paths": sys_paths, - } - - def clone(self): - """Return a clone of the agent.""" - return ModuleAgent( - module_name=self._tmpl_attrs.get("module_name"), - agent_name=self._tmpl_attrs.get("agent_name"), - register_operations=self._tmpl_attrs.get("register_operations"), - sys_paths=self._tmpl_attrs.get("sys_paths"), - agent_framework=self.agent_framework, - ) - - def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]: - return self._tmpl_attrs.get("register_operations") - - def set_up(self) -> None: - """Sets up the agent for execution of queries at runtime. - - It runs the code to import the agent from the module, and registers the - operations of the agent. - """ - if self._tmpl_attrs.get("sys_paths"): - import sys - - for sys_path in self._tmpl_attrs.get("sys_paths"): - abs_path = os.path.abspath(sys_path) - if abs_path not in sys.path: - sys.path.append(abs_path) - - import importlib - - module = importlib.import_module(self._tmpl_attrs.get("module_name")) - try: - importlib.reload(module) - except Exception as e: - _LOGGER.warning( - f"Failed to reload module {self._tmpl_attrs.get('module_name')}: {e}" - ) - agent_name = self._tmpl_attrs.get("agent_name") - try: - agent = getattr(module, agent_name) - except AttributeError as e: - raise AttributeError( - f"Agent {agent_name} not found in module " - f"{self._tmpl_attrs.get('module_name')}" - ) from e - if not self.agent_framework: - self.agent_framework = _get_agent_framework(agent) - self._tmpl_attrs["agent"] = agent - if hasattr(agent, "set_up"): - agent.set_up() - for operations in self.register_operations().values(): - for operation in operations: - op = _wrap_agent_operation(agent, operation) - setattr(self, operation, types.MethodType(op, self)) - - -class AgentEngine(base.VertexAiResourceNounWithFutureManager): - """Represents a Vertex AI Agent Engine resource.""" - - client_class = aip_utils.AgentEngineClientWithOverride - _resource_noun = "reasoning_engine" - _getter_method = "get_reasoning_engine" - _list_method = "list_reasoning_engines" - _delete_method = "delete_reasoning_engine" - _parse_resource_name_method = "parse_reasoning_engine_path" - _format_resource_name_method = "reasoning_engine_path" - - def __init__(self, resource_name: str): - """Retrieves an Agent Engine resource. - - Args: - resource_name (str): - Required. A fully-qualified resource name or ID such as - "projects/123/locations/us-central1/reasoningEngines/456" or - "456" when project and location are initialized or passed. - """ - super().__init__(resource_name=resource_name) - self.execution_api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionClientWithOverride, - ) - self.execution_async_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionAsyncClientWithOverride, - ) - self._gca_resource = self._get_gca_resource(resource_name=resource_name) - try: - _register_api_methods_or_raise(self) - except Exception as e: - _LOGGER.warning(_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e) - self._operation_schemas = None - - @property - def resource_name(self) -> str: - """Fully-qualified resource name.""" - return self._gca_resource.name - - @classmethod - def create( - cls, - agent_engine: Optional[_AgentEngineInterface] = None, - *, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - build_options: Optional[Dict[str, Sequence[str]]] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[aip_types.EncryptionSpec] = None, - ) -> "AgentEngine": - """Creates a new Agent Engine. - - The Agent Engine will be an instance of the `agent_engine` that - was passed in, running remotely on Vertex AI. - - Sample `src_dir` contents (e.g. `./user_src_dir`): - - .. code-block:: python - - user_src_dir/ - |-- main.py - |-- requirements.txt - |-- user_code/ - | |-- utils.py - | |-- ... - |-- installation_scripts/ - | |-- install_package.sh - | |-- ... - |-- ... - - To build an Agent Engine with the above files, run: - - .. code-block:: python - - remote_agent = agent_engines.create( - agent_engine=local_agent, - requirements=[ - # I.e. the PyPI dependencies listed in requirements.txt - "google-cloud-aiplatform==1.25.0", - "langchain==0.0.242", - ... - ], - extra_packages=[ - "./user_src_dir/main.py", # a single file - "./user_src_dir/user_code", # a directory - ... - ], - build_options={ - "installation_scripts": [ - "./user_src_dir/installation_scripts/install_package.sh", - ... - ], - }, - ) - - Args: - agent_engine (AgentEngineInterface): - Optional. The Agent Engine to be created. - requirements (Union[str, Sequence[str]]): - Optional. The set of PyPI dependencies needed. It can either be - the path to a single file (requirements.txt), or an ordered list - of strings corresponding to each line of the requirements file. - display_name (str): - Optional. The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any - UTF-8 character. - description (str): - Optional. The description of the Agent Engine. - gcs_dir_name (str): - Optional. The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - extra_packages (Sequence[str]): - Optional. The set of extra user-provided packages (if any). - env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): - Optional. The environment variables to be set when running the - Agent Engine. If it is a list of strings, each string should be - a valid key to `os.environ`. If it is a dictionary, the keys are - the environment variable names, and the values are the - corresponding values. - build_options (Dict[str, Sequence[str]]): - Optional. The build options for the Agent Engine. - The following keys are supported: - - installation_scripts: - Optional. The paths to the installation scripts to be - executed in the Docker image. - The scripts must be located in the `installation_scripts` - subdirectory and the path must be added to `extra_packages`. - service_account (str): - Optional. The service account to be used for the Agent Engine. - If not specified, the default reasoning engine service agent - service account will be used. - psc_interface_config (aip_types.PscInterfaceConfig): - Optional. The Private Service Connect interface config for the - Agent Engine. - min_instances (int): - Optional. The minimum number of instances to be running for the - Agent Engine. - max_instances (int): - Optional. The maximum number of instances to be running for the - Agent Engine. - resource_limits (Dict[str, str]): - Optional. The resource limits for the Agent Engine. - container_concurrency (int): - Optional. The container concurrency for the Agent Engine. - encryption_spec (aip_types.EncryptionSpec): - Optional. The Cloud KMS resource identifier of the customer - managed encryption key used to protect the model. Has the - form: - `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`. - The key needs to be in the same region as the model. - - Returns: - AgentEngine: The Agent Engine that was created. - - Raises: - ValueError: If the `project` was not set using `agentplatform.init`. - ValueError: If the `location` was not set using `agentplatform.init`. - ValueError: If the `staging_bucket` was not set using agentplatform.init. - ValueError: If the `staging_bucket` does not start with "gs://". - ValueError: If `extra_packages` is specified but `agent_engine` is None. - ValueError: If `requirements` is specified but `agent_engine` is None. - ValueError: If `env_vars` has a dictionary entry that does not - correspond to a SecretRef. - ValueError: If `env_vars` is a list which contains a string that - does not exist in `os.environ`. - TypeError: If `env_vars` is not a list of strings or a dictionary. - TypeError: If `env_vars` has a value that is not a string or SecretRef. - FileNotFoundError: If `extra_packages` includes a file or directory - that does not exist. - IOError: If requirements is a string that corresponds to a - nonexistent file. - """ - sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" - _validate_sys_version_or_raise(sys_version) - gcs_dir_name = gcs_dir_name or _DEFAULT_GCS_DIR_NAME - staging_bucket = initializer.global_config.staging_bucket - - if agent_engine is not None: - agent_engine = _validate_agent_engine_or_raise(agent_engine) - staging_bucket = _validate_staging_bucket_or_raise(staging_bucket) - if _is_adk_agent(None, agent_engine): - env_vars = _add_telemetry_enablement_env(env_vars=env_vars) - - if agent_engine is None: - if requirements is not None: - raise ValueError("requirements must be None if agent_engine is None.") - if extra_packages is not None: - raise ValueError("extra_packages must be None if agent_engine is None.") - requirements = _validate_requirements_or_raise( - agent_engine=agent_engine, - requirements=requirements, - ) - extra_packages = _validate_extra_packages_or_raise( - extra_packages=extra_packages, - build_options=build_options, - ) - - sdk_resource = cls.__new__(cls) - base.VertexAiResourceNounWithFutureManager.__init__(sdk_resource) - - # Prepares the Agent Engine for creation in Vertex AI. - # This involves packaging and uploading the artifacts for - # agent_engine, requirements and extra_packages to - # `staging_bucket/gcs_dir_name`. - _prepare( - agent_engine=agent_engine, - requirements=requirements, - project=sdk_resource.project, - location=sdk_resource.location, - staging_bucket=staging_bucket, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - ) - reasoning_engine = aip_types.ReasoningEngine( - display_name=display_name, - description=description, - encryption_spec=encryption_spec, - ) - if agent_engine is not None: - # Update the package spec. - package_spec = aip_types.ReasoningEngineSpec.PackageSpec( - python_version=sys_version, - pickle_object_gcs_uri="{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _BLOB_FILENAME, - ), - ) - if extra_packages: - package_spec.dependency_files_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _EXTRA_PACKAGES_FILE, - ) - if requirements: - package_spec.requirements_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _REQUIREMENTS_FILE, - ) - agent_engine_spec = aip_types.ReasoningEngineSpec( - package_spec=package_spec, - ) - if ( - env_vars - or psc_interface_config - or min_instances is not None - or max_instances is not None - or resource_limits - or container_concurrency is not None - ): - deployment_spec, _ = _generate_deployment_spec_or_raise( - env_vars=env_vars, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - ) - agent_engine_spec.deployment_spec = deployment_spec - class_methods_spec = _generate_class_methods_spec_or_raise( - agent_engine=agent_engine, - operations=_get_registered_operations(agent_engine), - ) - agent_engine_spec.class_methods.extend(class_methods_spec) - if service_account: - agent_engine_spec.service_account = service_account - reasoning_engine.spec = agent_engine_spec - reasoning_engine.spec.agent_framework = _get_agent_framework(agent_engine) - operation_future = sdk_resource.api_client.create_reasoning_engine( - parent=initializer.global_config.common_location_path( - project=sdk_resource.project, location=sdk_resource.location - ), - reasoning_engine=reasoning_engine, - ) - _LOGGER.log_create_with_lro(cls, operation_future) - _LOGGER.info( - f"View progress and logs at https://console.cloud.google.com/logs/query?project={sdk_resource.project}" - ) - created_resource = operation_future.result() - _LOGGER.info(f"{cls.__name__} created. Resource name: {created_resource.name}") - _LOGGER.info(f"To use this {cls.__name__} in another session:") - _LOGGER.info( - f"agent_engine = agentplatform.agent_engines.get('{created_resource.name}')" - ) - # We use `._get_gca_resource(...)` instead of `created_resource` to - # fully instantiate the attributes of the agent engine. - sdk_resource._gca_resource = sdk_resource._get_gca_resource( - resource_name=created_resource.name - ) - sdk_resource.execution_api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionClientWithOverride, - credentials=sdk_resource.credentials, - location_override=sdk_resource.location, - ) - sdk_resource.execution_async_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionAsyncClientWithOverride, - credentials=sdk_resource.credentials, - location_override=sdk_resource.location, - ) - if agent_engine is not None: - try: - _register_api_methods_or_raise(sdk_resource) - except Exception as e: - _LOGGER.warning(_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e) - sdk_resource._operation_schemas = None - return sdk_resource - - def update( - self, - *, - agent_engine: Optional[_AgentEngineInterface] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - gcs_dir_name: Optional[str] = None, - extra_packages: Optional[Sequence[str]] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - build_options: Optional[Dict[str, Sequence[str]]] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[aip_types.EncryptionSpec] = None, - ) -> "AgentEngine": - """Updates an existing Agent Engine. - - This method updates the configuration of an existing Agent Engine - running remotely, which is identified by its resource name. - Unlike the `create` function which requires a `agent_engine` object, - all arguments in this method are optional. - This method allows you to modify individual aspects of the configuration - by providing any of the optional arguments. - - Args: - agent_engine (AgentEngineInterface): - Optional. The instance to be used as the updated Agent Engine. - If it is not specified, the existing instance will be used. - requirements (Union[str, Sequence[str]]): - Optional. The set of PyPI dependencies needed. It can either be - the path to a single file (requirements.txt), or an ordered list - of strings corresponding to each line of the requirements file. - If it is not specified, the existing requirements will be used. - If it is set to an empty string or list, the existing - requirements will be removed. - display_name (str): - Optional. The user-defined name of the Agent Engine. - The name can be up to 128 characters long and can comprise any - UTF-8 character. - description (str): - Optional. The description of the Agent Engine. - gcs_dir_name (str): - Optional. The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - extra_packages (Sequence[str]): - Optional. The set of extra user-provided packages (if any). If - it is not specified, the existing extra packages will be used. - If it is set to an empty list, the existing extra packages will - be removed. - env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): - Optional. The environment variables to be set when running the - Agent Engine. If it is a list of strings, each string should be - a valid key to `os.environ`. If it is a dictionary, the keys are - the environment variable names, and the values are the - corresponding values. - build_options (Dict[str, Sequence[str]]): - Optional. The build options for the Agent Engine. - The following keys are supported: - - installation_scripts: - Optional. The paths to the installation scripts to be - executed in the Docker image. - The scripts must be located in the `installation_scripts` - subdirectory and the path must be added to `extra_packages`. - service_account (str): - Optional. The service account to be used for the Agent Engine. - If not specified, the default reasoning engine service agent - service account will be used. - psc_interface_config (aip_types.PscInterfaceConfig): - Optional. The Private Service Connect interface config for the - Agent Engine. - min_instances (int): - Optional. The minimum number of instances to be running for the - Agent Engine. - max_instances (int): - Optional. The maximum number of instances to be running for the - Agent Engine. - resource_limits (Dict[str, str]): - Optional. The resource limits for the Agent Engine. - container_concurrency (int): - Optional. The container concurrency for the Agent Engine. - encryption_spec (aip_types.EncryptionSpec): - Optional. The Cloud KMS resource identifier of the customer - managed encryption key used to protect the model. Has the - form: - `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`. - The key needs to be in the same region as the model. - - Returns: - AgentEngine: The Agent Engine that was updated. - - Raises: - ValueError: If the `staging_bucket` was not set using agentplatform.init. - ValueError: If the `staging_bucket` does not start with "gs://". - ValueError: If `env_vars` has a dictionary entry that does not - correspond to a SecretRef. - ValueError: If `env_vars` is a list which contains a string that - does not exist in `os.environ`. - TypeError: If `env_vars` is not a list of strings or a dictionary. - TypeError: If `env_vars` has a value that is not a string or SecretRef. - FileNotFoundError: If `extra_packages` includes a file or directory - that does not exist. - ValueError: if none of `display_name`, `description`, `requirements`, - `extra_packages`, `env_vars`, or `agent_engine` were specified. - IOError: If requirements is a string that corresponds to a - nonexistent file. - """ - staging_bucket = initializer.global_config.staging_bucket - staging_bucket = _validate_staging_bucket_or_raise(staging_bucket) - historical_operation_schemas = self.operation_schemas() - gcs_dir_name = gcs_dir_name or _DEFAULT_GCS_DIR_NAME - - # Validate the arguments. - if not any( - [ - agent_engine, - requirements, - extra_packages, - display_name, - description, - env_vars, - build_options, - service_account, - psc_interface_config, - min_instances is not None, - max_instances is not None, - resource_limits, - container_concurrency is not None, - encryption_spec, - ] - ): - raise ValueError( - "At least one of `agent_engine`, `requirements`, " - "`extra_packages`, `display_name`, `description`, " - "`env_vars`, `build_options`, `service_account`, " - "`psc_interface_config`, `min_instances`, `max_instances`, " - "`resource_limits`, `container_concurrency`, or " - "`encryption_spec` must be specified." - ) - if requirements is not None: - requirements = _validate_requirements_or_raise( - agent_engine=agent_engine, - requirements=requirements, - ) - if extra_packages is not None: - extra_packages = _validate_extra_packages_or_raise( - extra_packages=extra_packages, - build_options=build_options, - ) - if agent_engine is not None: - agent_engine = _validate_agent_engine_or_raise(agent_engine) - - if _is_adk_agent(self, agent_engine): - env_vars = _add_telemetry_enablement_env(env_vars=env_vars) - - # Prepares the Agent Engine for update in Vertex AI. This involves - # packaging and uploading the artifacts for agent_engine, requirements - # and extra_packages to `staging_bucket/gcs_dir_name`. - _prepare( - agent_engine=agent_engine, - requirements=requirements, - project=self.project, - location=self.location, - staging_bucket=staging_bucket, - gcs_dir_name=gcs_dir_name, - extra_packages=extra_packages, - ) - update_request = _generate_update_request_or_raise( - resource_name=self.resource_name, - staging_bucket=staging_bucket, - gcs_dir_name=gcs_dir_name, - agent_engine=agent_engine, - requirements=requirements, - extra_packages=extra_packages, - display_name=display_name, - description=description, - env_vars=env_vars, - service_account=service_account, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - encryption_spec=encryption_spec, - ) - operation_future = self.api_client.update_reasoning_engine( - request=update_request - ) - _LOGGER.info( - f"Update Agent Engine backing LRO: {operation_future.operation.name}" - ) - created_resource = operation_future.result() - _LOGGER.info(f"Agent Engine updated. Resource name: {created_resource.name}") - self._operation_schemas = None - self.execution_api_client = initializer.global_config.create_client( - client_class=aip_utils.AgentEngineExecutionClientWithOverride, - ) - # We use `._get_gca_resource(...)` instead of `created_resource` to - # fully instantiate the attributes of the agent engine. - self._gca_resource = self._get_gca_resource(resource_name=self.resource_name) - - if ( - agent_engine is None - or historical_operation_schemas == self.operation_schemas() - ): - # The operations of the agent engine are unchanged, so we return it. - return self - - # If the agent engine has changed and the historical operation - # schemas are different from the current operation schemas, we need to - # unregister the historical operation schemas and register the current - # operation schemas. - _unregister_api_methods(self, historical_operation_schemas) - try: - _register_api_methods_or_raise(self) - except Exception as e: - _LOGGER.warning(_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e) - return self - - def delete( - self, - *, - force: bool = False, - **kwargs, - ) -> None: - """Deletes the ReasoningEngine. - - Args: - force (bool): - Optional. If set to True, child resources will also be deleted. - Otherwise, the request will fail with FAILED_PRECONDITION error - when the Agent Engine has undeleted child resources. Defaults to - False. - **kwargs (dict[str, Any]): - Optional. Additional keyword arguments to pass to the - delete_reasoning_engine method. - """ - kwargs = kwargs or {} - operation_future = self.api_client.delete_reasoning_engine( - request=aip_types.DeleteReasoningEngineRequest( - name=self.resource_name, - force=force, - **kwargs, - ), - ) - _LOGGER.info( - f"Delete Agent Engine backing LRO: {operation_future.operation.name}" - ) - operation_future.result() - _LOGGER.info(f"Agent Engine deleted. Resource name: {self.resource_name}") - - def operation_schemas(self) -> Sequence[_agent_engines_utils.JsonDict]: - """Returns the (Open)API schemas for the Agent Engine.""" - spec = _agent_engines_utils._to_dict(self._gca_resource.spec) - if not hasattr(self, "_operation_schemas") or self._operation_schemas is None: - self._operation_schemas = spec.get("class_methods", []) - return self._operation_schemas - - -def _validate_sys_version_or_raise(sys_version: str) -> None: - """Tries to validate the python system version.""" - if sys_version not in _SUPPORTED_PYTHON_VERSIONS: - raise ValueError( - f"Unsupported python version: {sys_version}. AgentEngine " - f"only supports {_SUPPORTED_PYTHON_VERSIONS} at the moment." - ) - if sys_version != f"{sys.version_info.major}.{sys.version_info.minor}": - _LOGGER.warning( - f"{sys_version=} is inconsistent with {sys.version_info=}. " - "This might result in issues with deployment, and should only " - "be used as a workaround for advanced cases." - ) - - -def _validate_staging_bucket_or_raise(staging_bucket: Optional[str]) -> str: - """Tries to validate the staging bucket.""" - if not staging_bucket: - raise ValueError( - "Please provide a `staging_bucket` in `agentplatform.init(...)`" - ) - if not staging_bucket.startswith("gs://"): - raise ValueError(f"{staging_bucket=} must start with `gs://`") - return staging_bucket - - -def _validate_agent_engine_or_raise( - agent_engine: _AgentEngineInterface, - logger: base.Logger = _LOGGER, -) -> _AgentEngineInterface: - """Tries to validate the agent engine. - - The agent engine must have one of the following: - * a callable method named `query` - * a callable method named `stream_query` - * a callable method named `async_stream_query` - * a callable method named `bidi_stream_query` - * a callable method named `register_operations` - - Args: - agent_engine: The agent engine to be validated. - logger: The logger to use for logging. - - Returns: - The validated agent engine. - - Raises: - TypeError: If `agent_engine` has no callable method named `query`, - `stream_query` or `register_operations`. - ValueError: If `agent_engine` has an invalid `query`, `stream_query` or - `register_operations` signature. - """ - try: - from google.adk.agents import BaseAgent - - if isinstance(agent_engine, BaseAgent): - logger.info("Deploying google.adk.agents.Agent as an application.") - from agentplatform import agent_engines - - agent_engine = agent_engines.AdkApp(agent=agent_engine) - except Exception: - pass - is_queryable = isinstance(agent_engine, Queryable) and callable(agent_engine.query) - is_async_queryable = isinstance(agent_engine, AsyncQueryable) and callable( - agent_engine.async_query - ) - is_stream_queryable = isinstance(agent_engine, StreamQueryable) and callable( - agent_engine.stream_query - ) - is_async_stream_queryable = isinstance( - agent_engine, AsyncStreamQueryable - ) and callable(agent_engine.async_stream_query) - is_bidi_stream_queryable = isinstance( - agent_engine, BidiStreamQueryable - ) and callable(agent_engine.bidi_stream_query) - is_operation_registrable = isinstance( - agent_engine, OperationRegistrable - ) and callable(agent_engine.register_operations) - - if not ( - is_queryable - or is_async_queryable - or is_stream_queryable - or is_operation_registrable - or is_async_stream_queryable - or is_bidi_stream_queryable - ): - raise TypeError( - "agent_engine has none of the following callable methods: " - "`query`, `async_query`, `stream_query`, `async_stream_query`, " - "`bidi_stream_query` or `register_operations`." - ) - - if is_queryable: - try: - inspect.signature(getattr(agent_engine, "query")) - except ValueError as err: - raise ValueError( - "Invalid query signature. This might be due to a missing " - "`self` argument in the agent_engine.query method." - ) from err - - if is_async_queryable: - try: - inspect.signature(getattr(agent_engine, "async_query")) - except ValueError as err: - raise ValueError( - "Invalid async_query signature. This might be due to a missing " - "`self` argument in the agent_engine.async_query method." - ) from err - - if is_stream_queryable: - try: - inspect.signature(getattr(agent_engine, "stream_query")) - except ValueError as err: - raise ValueError( - "Invalid stream_query signature. This might be due to a missing" - " `self` argument in the agent_engine.stream_query method." - ) from err - - if is_async_stream_queryable: - try: - inspect.signature(getattr(agent_engine, "async_stream_query")) - except ValueError as err: - raise ValueError( - "Invalid async_stream_query signature. This might be due to a " - " missing `self` argument in the " - "agent_engine.async_stream_query method." - ) from err - - if is_bidi_stream_queryable: - try: - inspect.signature(getattr(agent_engine, "bidi_stream_query")) - except ValueError as err: - raise ValueError( - "Invalid bidi_stream_query signature. This might be due to a " - " missing `self` argument in the " - "agent_engine.bidi_stream_query method." - ) from err - - if is_operation_registrable: - try: - inspect.signature(getattr(agent_engine, "register_operations")) - except ValueError as err: - raise ValueError( - "Invalid register_operations signature. This might be due to a " - "missing `self` argument in the " - "agent_engine.register_operations method." - ) from err - - if isinstance(agent_engine, Cloneable): - # Avoid undeployable states. - agent_engine = agent_engine.clone() - return agent_engine - - -def _is_adk_agent( - agent_engine_to_update: Optional[AgentEngine], - new_agent_engine: Optional[_AgentEngineInterface], -) -> bool: - """Checks if the agent engine is an ADK agent. - - Args: - agent_engine_to_update: Existing agent engine, None if creating new one. - new_agent_engine: The new agent engine to deploy. Can be None during an update, if the Python agent implementation is not provided, and should remain unchanged. - - Returns: - True if the agent after the create/update operation, will be an ADK agent. - """ - - from agentplatform.agent_engines.templates import adk - - if new_agent_engine is not None: - return ( - getattr(new_agent_engine, "agent_framework", None) - == adk.AdkApp.agent_framework - ) - if agent_engine_to_update is not None: - return ( - agent_engine_to_update.gca_resource.spec.agent_framework - == adk.AdkApp.agent_framework - ) - return False - - -EnvVars = Optional[Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]] - - -def _add_telemetry_enablement_env(*, env_vars: EnvVars) -> EnvVars: - """Adds telemetry enablement env var to the env vars. - - This is in order to achieve default-on telemetry. - If the telemetry enablement env var is already set, we do not override it. - - Args: - env_vars: The env vars to add the telemetry enablement env var to. - - Returns: - The env vars with the telemetry enablement env var added. - """ - - GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = ( - "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY" - ) - - if env_vars is None: - return {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"} - if isinstance(env_vars, dict): - return ( - env_vars - if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in env_vars - else env_vars | {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"} - ) - if isinstance(env_vars, list) or isinstance(env_vars, tuple): - if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY not in os.environ: - os.environ[GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY] = "unspecified" - - if isinstance(env_vars, list): - return env_vars + [GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY] - else: - return env_vars + (GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY,) - - raise TypeError( - f"env_vars must be a list, tuple or a dict, but got {type(env_vars)}." - ) - - -def _validate_requirements_or_raise( - *, - agent_engine: _AgentEngineInterface, - requirements: Optional[Sequence[str]] = None, - logger: logging.getLoggerClass() = _LOGGER, -) -> Sequence[str]: - """Tries to validate the requirements.""" - if requirements is None: - requirements = [] - elif isinstance(requirements, str): - try: - logger.info(f"Reading requirements from {requirements=}") - with open(requirements) as f: - requirements = f.read().splitlines() - logger.info(f"Read the following lines: {requirements}") - except IOError as err: - raise IOError(f"Failed to read requirements from {requirements=}") from err - requirements = _agent_engines_utils._validate_requirements_or_warn( - obj=agent_engine, - requirements=requirements, - ) - logger.info(f"The final list of requirements: {requirements}") - return requirements - - -def _validate_extra_packages_or_raise( - extra_packages: Optional[Sequence[str]], - build_options: Optional[Dict[str, Sequence[str]]] = None, -) -> Sequence[str]: - """Tries to validates the extra packages.""" - extra_packages = extra_packages or [] - if build_options and _BUILD_OPTIONS_INSTALLATION in build_options: - _agent_engines_utils._validate_installation_scripts_or_raise( - script_paths=build_options[_BUILD_OPTIONS_INSTALLATION], - packages=extra_packages, - ) - for extra_package in extra_packages: - if not os.path.exists(extra_package): - raise FileNotFoundError( - f"Extra package specified but not found: {extra_package=}" - ) - return extra_packages - - -def _get_gcs_bucket( - *, - project: str, - location: str, - staging_bucket: str, - logger: base.Logger = _LOGGER, -) -> storage.Bucket: - """Gets or creates the GCS bucket.""" - storage = _agent_engines_utils._import_cloud_storage_or_raise() - storage_client = storage.Client(project=project) - staging_bucket = staging_bucket.replace("gs://", "") - try: - gcs_bucket = storage_client.get_bucket(staging_bucket) - logger.info(f"Using bucket {staging_bucket}") - except exceptions.NotFound: - new_bucket = storage_client.bucket(staging_bucket) - gcs_bucket = storage_client.create_bucket(new_bucket, location=location) - logger.info(f"Creating bucket {staging_bucket} in {location=}") - return gcs_bucket - - -def _upload_agent_engine( - *, - agent_engine: _AgentEngineInterface, - gcs_bucket: storage.Bucket, - gcs_dir_name: str, - logger: base.Logger = _LOGGER, -) -> None: - """Uploads the agent engine to GCS.""" - cloudpickle = _agent_engines_utils._import_cloudpickle_or_raise() - blob = gcs_bucket.blob(f"{gcs_dir_name}/{_BLOB_FILENAME}") - with blob.open("wb") as f: - try: - cloudpickle.dump(agent_engine, f) - except Exception as e: - url = "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#deployment-considerations" - error_msg = f"Failed to serialize agent engine. Visit {url} for details." - if "google._upb._message" in str(e) or "Descriptor" in str(e): - error_msg += ( - " This is often caused by protobuf objects (like Part, AgentCard) " - "being imported at the global module level. Please move these " - "imports inside the functions or methods where they are used. " - "Alternatively, you can import the entire module: " - "`from a2a import types as a2a_types`." - ) - raise TypeError(error_msg) from e - with blob.open("rb") as f: - try: - _ = cloudpickle.load(f) - except Exception as e: - raise TypeError("Agent engine serialized to an invalid format") from e - dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" - logger.info(f"Wrote to {dir_name}/{_BLOB_FILENAME}") - - -def _upload_requirements( - *, - requirements: Sequence[str], - gcs_bucket: storage.Bucket, - gcs_dir_name: str, - logger: base.Logger = _LOGGER, -) -> None: - """Uploads the requirements file to GCS.""" - blob = gcs_bucket.blob(f"{gcs_dir_name}/{_REQUIREMENTS_FILE}") - blob.upload_from_string("\n".join(requirements)) - dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" - logger.info(f"Writing to {dir_name}/{_REQUIREMENTS_FILE}") - - -def _upload_extra_packages( - *, - extra_packages: Sequence[str], - gcs_bucket: storage.Bucket, - gcs_dir_name: str, - logger: base.Logger = _LOGGER, -) -> None: - """Uploads extra packages to GCS.""" - logger.info("Creating in-memory tarfile of extra_packages") - tar_fileobj = io.BytesIO() - with tarfile.open(fileobj=tar_fileobj, mode="w|gz") as tar: - for file in extra_packages: - tar.add(file) - tar_fileobj.seek(0) - blob = gcs_bucket.blob(f"{gcs_dir_name}/{_EXTRA_PACKAGES_FILE}") - blob.upload_from_string(tar_fileobj.read()) - dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}" - logger.info(f"Writing to {dir_name}/{_EXTRA_PACKAGES_FILE}") - - -def _prepare( - agent_engine: Optional[_AgentEngineInterface], - requirements: Optional[Sequence[str]], - extra_packages: Optional[Sequence[str]], - project: str, - location: str, - staging_bucket: str, - gcs_dir_name: str, - logger: base.Logger = _LOGGER, -) -> None: - """Prepares the agent engine for creation or updates in Vertex AI. - - This involves packaging and uploading artifacts to Cloud Storage. Note that - 1. This does not actually update the Agent Engine in Vertex AI. - 2. This will only generate and upload a pickled object if specified. - 3. This will only generate and upload the dependencies.tar.gz file if - extra_packages is non-empty. - - Args: - agent_engine: The agent engine to be prepared. - requirements (Sequence[str]): The set of PyPI dependencies needed. - extra_packages (Sequence[str]): The set of extra user-provided packages. - project (str): The project for the staging bucket. - location (str): The location for the staging bucket. - staging_bucket (str): The staging bucket name in the form "gs://...". - gcs_dir_name (str): The GCS bucket directory under `staging_bucket` to - use for staging the artifacts needed. - """ - if agent_engine is None: - return - gcs_bucket = _get_gcs_bucket( - project=project, - location=location, - staging_bucket=staging_bucket, - logger=logger, - ) - _upload_agent_engine( - agent_engine=agent_engine, - gcs_bucket=gcs_bucket, - gcs_dir_name=gcs_dir_name, - logger=logger, - ) - if requirements is not None: - _upload_requirements( - requirements=requirements, - gcs_bucket=gcs_bucket, - gcs_dir_name=gcs_dir_name, - logger=logger, - ) - if extra_packages is not None: - _upload_extra_packages( - extra_packages=extra_packages, - gcs_bucket=gcs_bucket, - gcs_dir_name=gcs_dir_name, - logger=logger, - ) - - -def _update_deployment_spec_with_env_vars_dict_or_raise( - *, - deployment_spec: aip_types.ReasoningEngineSpec.DeploymentSpec, - env_vars: Dict[str, Union[str, aip_types.SecretRef]], -) -> None: - for key, value in env_vars.items(): - if isinstance(value, Dict): - try: - secret_ref = _agent_engines_utils._to_proto( - value, aip_types.SecretRef() - ) - except Exception as e: - raise ValueError(f"Failed to convert to secret ref: {value}") from e - deployment_spec.secret_env.append( - aip_types.SecretEnvVar(name=key, secret_ref=secret_ref) - ) - elif isinstance(value, aip_types.SecretRef): - deployment_spec.secret_env.append( - aip_types.SecretEnvVar(name=key, secret_ref=value) - ) - elif isinstance(value, str): - deployment_spec.env.append(aip_types.EnvVar(name=key, value=value)) - else: - raise TypeError( - f"Unknown value type in env_vars for {key}. " - f"Must be a str or SecretRef: {value}" - ) - - -def _update_deployment_spec_with_env_vars_list_or_raise( - *, - deployment_spec: aip_types.ReasoningEngineSpec.DeploymentSpec, - env_vars: Sequence[str], -) -> None: - for env_var in env_vars: - if env_var not in os.environ: - raise ValueError(f"Env var not found in os.environ: {env_var}.") - deployment_spec.env.append( - aip_types.EnvVar(name=env_var, value=os.environ[env_var]) - ) - - -def _generate_deployment_spec_or_raise( - *, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, -) -> Tuple[aip_types.ReasoningEngineSpec.DeploymentSpec, List[str]]: - deployment_spec = aip_types.ReasoningEngineSpec.DeploymentSpec() - update_masks = [] - if env_vars: - deployment_spec.env = [] - deployment_spec.secret_env = [] - if isinstance(env_vars, Dict): - _update_deployment_spec_with_env_vars_dict_or_raise( - deployment_spec=deployment_spec, - env_vars=env_vars, - ) - elif isinstance(env_vars, Sequence): - _update_deployment_spec_with_env_vars_list_or_raise( - deployment_spec=deployment_spec, - env_vars=env_vars, - ) - else: - raise TypeError( - f"env_vars must be a list, tuple or a dict, but got {type(env_vars)}." - ) - if deployment_spec.env: - update_masks.append("spec.deployment_spec.env") - if deployment_spec.secret_env: - update_masks.append("spec.deployment_spec.secret_env") - if psc_interface_config: - deployment_spec.psc_interface_config = psc_interface_config - update_masks.append("spec.deployment_spec.psc_interface_config") - if min_instances is not None: - deployment_spec.min_instances = min_instances - update_masks.append("spec.deployment_spec.min_instances") - if max_instances is not None: - deployment_spec.max_instances = max_instances - update_masks.append("spec.deployment_spec.max_instances") - if resource_limits: - deployment_spec.resource_limits = resource_limits - update_masks.append("spec.deployment_spec.resource_limits") - if container_concurrency is not None: - deployment_spec.container_concurrency = container_concurrency - update_masks.append("spec.deployment_spec.container_concurrency") - return deployment_spec, update_masks - - -def _get_agent_framework( - agent_engine: _AgentEngineInterface, -) -> str: - if ( - hasattr(agent_engine, _AGENT_FRAMEWORK_ATTR) - and getattr(agent_engine, _AGENT_FRAMEWORK_ATTR) is not None - and isinstance(getattr(agent_engine, _AGENT_FRAMEWORK_ATTR), str) - ): - return getattr(agent_engine, _AGENT_FRAMEWORK_ATTR) - return _DEFAULT_AGENT_FRAMEWORK - - -def _generate_update_request_or_raise( - *, - resource_name: str, - staging_bucket: str, - gcs_dir_name: str = _DEFAULT_GCS_DIR_NAME, - agent_engine: Optional[_AgentEngineInterface] = None, - requirements: Optional[Union[str, Sequence[str]]] = None, - extra_packages: Optional[Sequence[str]] = None, - display_name: Optional[str] = None, - description: Optional[str] = None, - env_vars: Optional[ - Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]] - ] = None, - service_account: Optional[str] = None, - psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None, - min_instances: Optional[int] = None, - max_instances: Optional[int] = None, - resource_limits: Optional[Dict[str, str]] = None, - container_concurrency: Optional[int] = None, - encryption_spec: Optional[str] = None, -) -> reasoning_engine_service.UpdateReasoningEngineRequest: - """Tries to generate the update request for the agent engine.""" - is_spec_update = False - update_masks: List[str] = [] - agent_engine_spec = aip_types.ReasoningEngineSpec() - package_spec = aip_types.ReasoningEngineSpec.PackageSpec() - if requirements is not None: - is_spec_update = True - update_masks.append("spec.package_spec.requirements_gcs_uri") - package_spec.requirements_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _REQUIREMENTS_FILE, - ) - if extra_packages is not None: - is_spec_update = True - update_masks.append("spec.package_spec.dependency_files_gcs_uri") - package_spec.dependency_files_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _EXTRA_PACKAGES_FILE, - ) - if agent_engine is not None: - is_spec_update = True - update_masks.append("spec.package_spec.pickle_object_gcs_uri") - package_spec.pickle_object_gcs_uri = "{}/{}/{}".format( - staging_bucket, - gcs_dir_name, - _BLOB_FILENAME, - ) - class_methods_spec = _generate_class_methods_spec_or_raise( - agent_engine=agent_engine, - operations=_get_registered_operations(agent_engine), - ) - agent_engine_spec.class_methods.extend(class_methods_spec) - update_masks.append("spec.class_methods") - agent_engine_spec.agent_framework = _get_agent_framework(agent_engine) - update_masks.append("spec.agent_framework") - if ( - env_vars is not None - or psc_interface_config - or min_instances is not None - or max_instances is not None - or resource_limits - or container_concurrency is not None - ): - is_spec_update = True - deployment_spec, deployment_update_masks = _generate_deployment_spec_or_raise( - env_vars=env_vars, - psc_interface_config=psc_interface_config, - min_instances=min_instances, - max_instances=max_instances, - resource_limits=resource_limits, - container_concurrency=container_concurrency, - ) - update_masks.extend(deployment_update_masks) - agent_engine_spec.deployment_spec = deployment_spec - if service_account is not None: - is_spec_update = True - update_masks.append("spec.service_account") - agent_engine_spec.service_account = service_account - - agent_engine_message = aip_types.ReasoningEngine(name=resource_name) - if is_spec_update: - if package_spec: - agent_engine_spec.package_spec = package_spec - agent_engine_message.spec = agent_engine_spec - if display_name: - agent_engine_message.display_name = display_name - update_masks.append("display_name") - if description: - agent_engine_message.description = description - update_masks.append("description") - if encryption_spec: - agent_engine_message.encryption_spec = encryption_spec - update_masks.append("encryption_spec") - if not update_masks: - raise ValueError( - "At least one of `agent_engine`, `requirements`, `extra_packages`, " - "`display_name`, `description`, `env_vars`, or " - "`encryption_spec` must be specified." - ) - return reasoning_engine_service.UpdateReasoningEngineRequest( - reasoning_engine=agent_engine_message, - update_mask=field_mask_pb2.FieldMask(paths=update_masks), - ) - - -def _wrap_query_operation( - method_name: str, -) -> Callable[..., _agent_engines_utils.JsonDict]: - """Wraps an Agent Engine method, creating a callable for `query` API. - - This function creates a callable object that executes the specified - Agent Engine method using the `query` API. It handles the creation of - the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - doc: Documentation string for the method. - - Returns: - A callable object that executes the method on the Agent Engine via - the `query` API. - """ - - def _method(self, **kwargs) -> _agent_engines_utils.JsonDict: - response = self.execution_api_client.query_reasoning_engine( - request=aip_types.QueryReasoningEngineRequest( - name=self.resource_name, - input=kwargs, - class_method=method_name, - ), - ) - output = _agent_engines_utils._to_dict(response) - return output.get("output", output) - - return _method - - -def _wrap_async_query_operation(method_name: str) -> Callable[..., Coroutine]: - """Wraps an Agent Engine method, creating an async callable for `query` API. - - This function creates a callable object that executes the specified - Agent Engine method asynchronously using the `query` API. It handles the - creation of the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - doc: Documentation string for the method. - - Returns: - A callable object that executes the method on the Agent Engine via - the `query` API. - """ - - async def _method(self, **kwargs) -> _agent_engines_utils.JsonDict: - response = await self.execution_async_client.query_reasoning_engine( - request=aip_types.QueryReasoningEngineRequest( - name=self.resource_name, - input=kwargs, - class_method=method_name, - ), - ) - output = _agent_engines_utils._to_dict(response) - return output.get("output", output) - - return _method - - -def _wrap_stream_query_operation(*, method_name: str) -> Callable[..., Iterable[Any]]: - """Wraps an Agent Engine method, creating a callable for `stream_query` API. - - This function creates a callable object that executes the specified - Agent Engine method using the `stream_query` API. It handles the - creation of the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - doc: Documentation string for the method. - - Returns: - A callable object that executes the method on the Agent Engine via - the `stream_query` API. - """ - - def _method(self, **kwargs) -> Iterable[Any]: - response = self.execution_api_client.stream_query_reasoning_engine( - request=aip_types.StreamQueryReasoningEngineRequest( - name=self.resource_name, - input=kwargs, - class_method=method_name, - ), - ) - for chunk in response: - for parsed_json in _agent_engines_utils._yield_parsed_json_from_httpbody( - chunk - ): - if parsed_json is not None: - yield parsed_json - - return _method - - -def _wrap_async_stream_query_operation( - *, method_name: str -) -> Callable[..., AsyncIterable[Any]]: - """Wraps an Agent Engine method, creating an async callable for `stream_query` API. - - This function creates a callable object that executes the specified - Agent Engine method using the `stream_query` API. It handles the - creation of the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - doc: Documentation string for the method. - - Returns: - A callable object that executes the method on the Agent Engine via - the `stream_query` API. - """ - - async def _method(self, **kwargs) -> AsyncIterable[Any]: - response = self.execution_api_client.stream_query_reasoning_engine( - request=aip_types.StreamQueryReasoningEngineRequest( - name=self.resource_name, - input=kwargs, - class_method=method_name, - ), - ) - for chunk in response: - for parsed_json in _agent_engines_utils._yield_parsed_json_from_httpbody( - chunk - ): - if parsed_json is not None: - yield parsed_json - - return _method - - -def _wrap_bidi_stream_query_operation( - *, method_name: str -) -> Callable[..., AsyncIterable[Any]]: - """Wraps an Agent Engine method, creating an async callable for `bidi_stream_query` API. - - This function creates a callable object that executes the specified - Agent Engine method using the `bidi_stream_query` API. It handles the - creation of the API request and the processing of the API response. - - Args: - method_name: The name of the Agent Engine method to call. - - Returns: - A callable object that executes the method on the Agent Engine via - the `bidi_stream_query` API. - """ - - async def _method(self, **kwargs) -> AsyncIterable[Any]: - # Agent Engine bidi streaming query execution should use GenAI SDK Agent - # Engine live API client directly. - raise NotImplementedError( - f"{method_name} is not implemented, please use GenAI SDK Agent " - "Enginve live API client instead." - ) - - -def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list]: - """Wraps an Agent Engine method, creating a callable for A2A API. - - Args: - method_name: The name of the Agent Engine method to call. - agent_card: The agent card to use for the A2A API call. - Example: - {'additionalInterfaces': None, - 'capabilities': {'extensions': None, - 'pushNotifications': None, - 'stateTransitionHistory': None, - 'streaming': False}, - 'defaultInputModes': ['text'], - 'defaultOutputModes': ['text'], - 'description': ( - 'A helpful assistant agent that can answer questions.' - ), - 'documentationUrl': None, - 'iconUrl': None, - 'name': 'Q&A Agent', - 'preferredTransport': 'JSONRPC', - 'protocolVersion': '0.3.0', - 'provider': None, - 'security': None, - 'securitySchemes': None, - 'signatures': None, - 'skills': [{ - 'description': ( - 'A helpful assistant agent that can answer questions.' - ), - 'examples': ['Who is leading 2025 F1 Standings?', - 'Where can i find an active volcano?'], - 'id': 'question_answer', - 'inputModes': None, - 'name': 'Q&A Agent', - 'outputModes': None, - 'security': None, - 'tags': ['Question-Answer']}], - 'supportsAuthenticatedExtendedCard': True, - 'url': 'http://localhost:8080/', - 'version': '1.0.0'} - Returns: - A callable object that executes the method on the Agent Engine via - the A2A API. - """ - - async def _method(self, **kwargs) -> Any: - """Wraps an Agent Engine method, creating a callable for A2A API.""" - a2a_agent_card = AgentCard(**json.loads(agent_card)) - - # A2A + AE integration currently only supports Rest API. - if ( - a2a_agent_card.supported_interfaces - and a2a_agent_card.supported_interfaces[0].protocol_binding - != TransportProtocol.HTTP_JSON - ): - raise ValueError( - "Only HTTP+JSON is supported for primary interface on agent card " - ) - - # Set primary interface to HTTP+JSON if not set. - if not a2a_agent_card.supported_interfaces: - a2a_agent_card.supported_interfaces = [] - a2a_agent_card.supported_interfaces.append( - AgentInterface( - protocol_binding=TransportProtocol.HTTP_JSON, - protocol_version=PROTOCOL_VERSION_CURRENT, - ) - ) - - # AE cannot support streaming yet. Turn off streaming for now. - if a2a_agent_card.capabilities and a2a_agent_card.capabilities.streaming: - raise ValueError( - "Streaming is not supported in Agent Engine, please change " - "a2a_agent_card.capabilities.streaming to False." - ) - - if not hasattr(a2a_agent_card.capabilities, "streaming"): - a2a_agent_card.capabilities.streaming = False - - # agent_card is set on the class_methods before set_up is invoked. - # Ensure that the agent_card url is set correctly before the client is created. - url = f"https://{initializer.global_config.api_endpoint}/v1beta1/{self.resource_name}/a2a" - a2a_agent_card.supported_interfaces[0].url = url - - # Using a2a client, inject the auth token from the global config. - config = ClientConfig( - supported_transports=[ - TransportProtocol.HTTP_JSON, - ], - use_client_preference=True, - httpx_client=httpx.AsyncClient( - headers={ - "Authorization": ( - f"Bearer {initializer.global_config.credentials.token}" - ) - } - ), - ) - factory = ClientFactory(config) - client = factory.create(a2a_agent_card) - - # kokoro job uses python 3.9, replaced match with if else. - if method_name == "on_message_send": - response = client.send_message(Message(**kwargs)) - chunks = [] - async for chunk in response: - chunks.append(chunk) - return chunks - elif method_name == "on_get_task": - response = await client.get_task(TaskQueryParams(**kwargs)) - elif method_name == "on_cancel_task": - response = await client.cancel_task(TaskIdParams(**kwargs)) - elif method_name == "handle_authenticated_agent_card": - response = await client.get_card() - else: - raise ValueError(f"Unknown method name: {method_name}") - - return response - - return _method - - -def _unregister_api_methods( - obj: "AgentEngine", operation_schemas: Sequence[_agent_engines_utils.JsonDict] -): - """Unregisters Agent Engine API methods based on operation schemas. - - This function iterates through operation schemas provided by the - AgentEngine object. Each schema defines an API mode and method name. - It dynamically unregisters methods on the AgentEngine object. This - should only be used when updating the object. - - Args: - obj: The AgentEngine object to augment with API methods. - operation_schemas: The operation schemas to use for method unregistration. - """ - for operation_schema in operation_schemas: - if "name" in operation_schema: - method_name = operation_schema.get("name") - if hasattr(obj, method_name): - delattr(obj, method_name) - - -def _register_api_methods_or_raise( - obj: "AgentEngine", - wrap_operation_fn: Optional[ - dict[str, Callable[[str, str], Callable[..., Any]]] - ] = None, -): - """Registers Agent Engine API methods based on operation schemas. - - This function iterates through operation schemas provided by the - AgentEngine object. Each schema defines an API mode and method name. - It dynamically creates and registers methods on the AgentEngine object - to handle API calls based on the specified API mode. - Currently, only standard API mode `` is supported. - - Args: - obj: The AgentEngine object to augment with API methods. - wrap_operation_fn: A dictionary of API modes and method wrapping - functions. - - Raises: - ValueError: If the API mode is not supported or if the operation schema - is missing any required fields (e.g. `api_mode` or `name`). - """ - for operation_schema in obj.operation_schemas(): - if _MODE_KEY_IN_SCHEMA not in operation_schema: - raise ValueError( - f"Operation schema {operation_schema} does not" - f" contain an `{_MODE_KEY_IN_SCHEMA}` field." - ) - api_mode = operation_schema.get(_MODE_KEY_IN_SCHEMA) - if _METHOD_NAME_KEY_IN_SCHEMA not in operation_schema: - raise ValueError( - f"Operation schema {operation_schema} does not" - f" contain a `{_METHOD_NAME_KEY_IN_SCHEMA}` field." - ) - method_name = operation_schema.get(_METHOD_NAME_KEY_IN_SCHEMA) - method_description = operation_schema.get( - "description", - _DEFAULT_METHOD_DOCSTRING_TEMPLATE.format( - method_name=method_name, - default_method_name=_DEFAULT_METHOD_NAME_MAP.get( - api_mode, _DEFAULT_METHOD_NAME - ), - return_type=_DEFAULT_METHOD_RETURN_TYPE_MAP.get( - api_mode, - _DEFAULT_METHOD_RETURN_TYPE, - ), - ), - ) - _wrap_operation_map = { - _STANDARD_API_MODE: _wrap_query_operation, - _ASYNC_API_MODE: _wrap_async_query_operation, - _STREAM_API_MODE: _wrap_stream_query_operation, - _ASYNC_STREAM_API_MODE: _wrap_async_stream_query_operation, - _BIDI_STREAM_API_MODE: _wrap_bidi_stream_query_operation, - _A2A_EXTENSION_MODE: _wrap_a2a_operation, - } - if isinstance(wrap_operation_fn, dict) and api_mode in wrap_operation_fn: - # Override the default function with user-specified function if it exists. - _wrap_operation = wrap_operation_fn[api_mode] - elif api_mode in _wrap_operation_map: - _wrap_operation = _wrap_operation_map[api_mode] - else: - supported_api_modes = ", ".join( - f"`{mode}`" for mode in sorted(_wrap_operation_map.keys()) - ) - raise ValueError( - f"Unsupported api mode: `{api_mode}`," - f" Supported modes are: {supported_api_modes}." - ) - - # Bind the method to the object. - if api_mode == _A2A_EXTENSION_MODE: - agent_card = operation_schema.get(_A2A_AGENT_CARD) - method = _wrap_operation(method_name=method_name, agent_card=agent_card) - else: - method = _wrap_operation(method_name=method_name) - method.__name__ = method_name - method.__doc__ = method_description - setattr(obj, method_name, types.MethodType(method, obj)) - - -def _get_registered_operations( - agent_engine: _AgentEngineInterface, -) -> Dict[str, List[str]]: - """Retrieves registered operations for a AgentEngine.""" - if isinstance(agent_engine, OperationRegistrable): - return agent_engine.register_operations() - - operations = {} - if isinstance(agent_engine, Queryable): - operations[_STANDARD_API_MODE] = [_DEFAULT_METHOD_NAME] - if isinstance(agent_engine, AsyncQueryable): - operations[_ASYNC_API_MODE] = [_DEFAULT_ASYNC_METHOD_NAME] - if isinstance(agent_engine, StreamQueryable): - operations[_STREAM_API_MODE] = [_DEFAULT_STREAM_METHOD_NAME] - if isinstance(agent_engine, AsyncStreamQueryable): - operations[_ASYNC_STREAM_API_MODE] = [_DEFAULT_ASYNC_STREAM_METHOD_NAME] - return operations - - -def _generate_class_methods_spec_or_raise( - *, - agent_engine: _AgentEngineInterface, - operations: Dict[str, List[str]], - logger: base.Logger = _LOGGER, -) -> List[proto.Message]: - """Generates a ReasoningEngineSpec based on the registered operations. - - Args: - agent_engine: The AgentEngine instance. - operations: A dictionary of API modes and method names. - - Returns: - A list of ReasoningEngineSpec.ClassMethod messages. - - Raises: - ValueError: If a method defined in `register_operations` is not found on - the AgentEngine. - """ - if isinstance(agent_engine, ModuleAgent): - # We do a dry-run of setting up the agent engine to have the operations - # needed for registration. - agent_engine = agent_engine.clone() - try: - agent_engine.set_up() - except Exception as e: - raise ValueError( - f"Failed to set up agent engine {agent_engine}: {e}" - ) from e - class_methods_spec = [] - for mode, method_names in operations.items(): - for method_name in method_names: - if mode == _BIDI_STREAM_API_MODE: - _LOGGER.warning( - "Bidi stream API mode is not supported yet in Vertex SDK, " - "please use the GenAI SDK instead. Skipping " - f"method {method_name}." - ) - continue - if not hasattr(agent_engine, method_name): - raise ValueError( - f"Method `{method_name}` defined in `register_operations`" - " not found on AgentEngine." - ) - - method = getattr(agent_engine, method_name) - try: - schema_dict = _agent_engines_utils._generate_schema( - method, schema_name=method_name - ) - except Exception as e: - logger.warning(f"failed to generate schema for {method_name}: {e}") - continue - - class_method = _agent_engines_utils._to_proto(schema_dict) - class_method[_MODE_KEY_IN_SCHEMA] = mode - # A2A agent card is a special case, when running in A2A mode, - if hasattr(agent_engine, "agent_card"): - from google.protobuf import json_format - - class_method[_A2A_AGENT_CARD] = json_format.MessageToJson( - getattr(agent_engine, "agent_card") - ) - class_methods_spec.append(class_method) - - return class_methods_spec - - -def _class_methods_to_class_methods_spec( - class_methods: List[dict[str, Any]], -) -> List[proto.Message]: - """Converts a list of class methods to a list of ReasoningEngineSpec.ClassMethod messages.""" - return [ - _agent_engines_utils._to_proto(class_method) for class_method in class_methods - ] diff --git a/agentplatform/frameworks/__init__.py b/agentplatform/frameworks/__init__.py new file mode 100644 index 0000000000..346cdc96c9 --- /dev/null +++ b/agentplatform/frameworks/__init__.py @@ -0,0 +1,39 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Classes for working with agent platform.""" + +from agentplatform.frameworks import a2a +from agentplatform.frameworks import adk +from agentplatform.frameworks import ag2 +from agentplatform.frameworks import langchain +from agentplatform.frameworks import langgraph +from agentplatform.frameworks import llama_index + + +A2aAgent = a2a.A2aAgent +AdkApp = adk.AdkApp +AG2Agent = ag2.AG2Agent +LangchainAgent = langchain.LangchainAgent +LanggraphAgent = langgraph.LanggraphAgent +LlamaIndexQueryPipelineAgent = llama_index.LlamaIndexQueryPipelineAgent + +__all__ = ( + "A2aAgent", + "AdkApp", + "AG2Agent", + "LangchainAgent", + "LanggraphAgent", + "LlamaIndexQueryPipelineAgent", +) diff --git a/agentplatform/agent_engines/templates/a2a.py b/agentplatform/frameworks/a2a.py similarity index 97% rename from agentplatform/agent_engines/templates/a2a.py rename to agentplatform/frameworks/a2a.py index 705d57c428..79d313886d 100644 --- a/agentplatform/agent_engines/templates/a2a.py +++ b/agentplatform/frameworks/a2a.py @@ -259,7 +259,6 @@ def __init__( ): """Initializes the A2A agent.""" # pylint: disable=g-import-not-at-top - from google.cloud.aiplatform import initializer from a2a.utils.constants import TransportProtocol, PROTOCOL_VERSION_CURRENT if ( @@ -276,8 +275,6 @@ def __init__( ) self._tmpl_attrs: dict[str, Any] = { - "project": initializer.global_config.project, - "location": initializer.global_config.location, "agent_card": agent_card, "agent_executor": None, "agent_executor_kwargs": agent_executor_kwargs or {}, @@ -318,10 +315,17 @@ 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 + + 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_id = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID", "test-agent-engine") version = "v1beta1" diff --git a/agentplatform/agent_engines/templates/adk.py b/agentplatform/frameworks/adk.py similarity index 98% rename from agentplatform/agent_engines/templates/adk.py rename to agentplatform/frameworks/adk.py index e7b1b60001..08fb1bb465 100644 --- a/agentplatform/agent_engines/templates/adk.py +++ b/agentplatform/frameworks/adk.py @@ -252,13 +252,13 @@ def __init__(self, **kwargs): # The session ID. def dump(self) -> Dict[str, Any]: - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils result = {} if self.events: result["events"] = [] for event in self.events: - event_dict = _agent_engines_utils.dump_event_for_json(event) + event_dict = _runtimes_utils.dump_event_for_json(event) event_dict["invocation_id"] = event_dict.get("invocation_id", "") result["events"].append(event_dict) if self.artifacts: @@ -477,9 +477,9 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]: # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils - if _agent_engines_utils.is_noop_or_proxy_tracer_provider(tracer_provider): + if _runtimes_utils.is_noop_or_proxy_tracer_provider(tracer_provider): tracer_provider = opentelemetry.sdk.trace.TracerProvider(resource=resource) opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -736,7 +736,6 @@ def __init__( This parameter is ignored if `enable_tracing` is False. """ import os - from google.cloud.aiplatform import initializer adk_version = get_adk_version() if not is_version_sufficient("1.5.0"): @@ -763,8 +762,6 @@ def __init__( ) self._tmpl_attrs: Dict[str, Any] = { - "project": initializer.global_config.project, - "location": initializer.global_config.location, "agent": agent, "app": app, "app_name": app_name, @@ -775,9 +772,7 @@ def __init__( "memory_service_builder": memory_service_builder, "credential_service_builder": credential_service_builder, "instrumentor_builder": instrumentor_builder, - "express_mode_api_key": ( - initializer.global_config.api_key or os.environ.get("GOOGLE_API_KEY") - ), + "express_mode_api_key": os.environ.get("GOOGLE_API_KEY"), } def _serialize(self, obj: Any) -> Any: @@ -947,19 +942,16 @@ def set_up(self): ) os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "1" - project = self._tmpl_attrs.get("project") - if project: - os.environ["GOOGLE_CLOUD_PROJECT"] = project - location = self._tmpl_attrs.get("location") + project = os.environ.get("GOOGLE_CLOUD_PROJECT") + location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") or os.getenv( + "GOOGLE_CLOUD_LOCATION" + ) if location: if "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION" not in os.environ: os.environ["GOOGLE_CLOUD_AGENT_ENGINE_LOCATION"] = location if "GOOGLE_CLOUD_LOCATION" not in os.environ: os.environ["GOOGLE_CLOUD_LOCATION"] = location - agent_engine_location = os.environ.get( - "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", # the runtime env var (if set) - location, # the location set in the AdkApp template - ) + agent_engine_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 @@ -1010,12 +1002,13 @@ def set_up(self): ), ) + project_id = self._get_project_id(project) if custom_instrumentor and self._tracing_enabled(): - self._tmpl_attrs["instrumentor"] = custom_instrumentor(self.project_id()) + self._tmpl_attrs["instrumentor"] = custom_instrumentor(project_id) if not custom_instrumentor: self._tmpl_attrs["instrumentor"] = _default_instrumentor_builder( - self.project_id(), + project_id, enable_tracing=self._tracing_enabled(), enable_logging=enable_logging, ) @@ -1184,7 +1177,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): @@ -1244,7 +1237,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( @@ -1294,7 +1287,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): @@ -1321,7 +1314,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, @@ -1329,7 +1322,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. @@ -2138,8 +2131,8 @@ def _tracing_enabled(self) -> bool: and is_version_sufficient("1.17.0") ) - def project_id(self) -> Optional[str]: - if project := self._tmpl_attrs.get("project"): + def _get_project_id(self, project: str) -> Optional[str]: + if project: try: from google.cloud.aiplatform.utils import ( resource_manager_utils, diff --git a/agentplatform/agent_engines/templates/ag2.py b/agentplatform/frameworks/ag2.py similarity index 89% rename from agentplatform/agent_engines/templates/ag2.py rename to agentplatform/frameworks/ag2.py index ee12e49d8f..2d456b7e75 100644 --- a/agentplatform/agent_engines/templates/ag2.py +++ b/agentplatform/frameworks/ag2.py @@ -23,6 +23,8 @@ Sequence, Union, ) +import os +import copy if TYPE_CHECKING: try: @@ -89,15 +91,13 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): - from agentplatform._genai import _agent_engines_utils - - cloud_trace_exporter = _agent_engines_utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _agent_engines_utils._import_cloud_trace_v2_or_warn() - openinference_autogen = _agent_engines_utils._import_openinference_autogen_or_warn() - opentelemetry = _agent_engines_utils._import_opentelemetry_or_warn() - opentelemetry_sdk_trace = ( - _agent_engines_utils._import_opentelemetry_sdk_trace_or_warn() - ) + from agentplatform._genai import _runtimes_utils + + cloud_trace_exporter = _runtimes_utils._import_cloud_trace_exporter_or_warn() + cloud_trace_v2 = _runtimes_utils._import_cloud_trace_v2_or_warn() + openinference_autogen = _runtimes_utils._import_openinference_autogen_or_warn() + opentelemetry = _runtimes_utils._import_opentelemetry_or_warn() + opentelemetry_sdk_trace = _runtimes_utils._import_opentelemetry_sdk_trace_or_warn() if all( ( cloud_trace_exporter, @@ -144,7 +144,7 @@ def _default_instrumentor_builder(project_id: str): # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. - if _agent_engines_utils.is_noop_or_proxy_tracer_provider(tracer_provider): + if _runtimes_utils.is_noop_or_proxy_tracer_provider(tracer_provider): tracer_provider = opentelemetry_sdk_trace.TracerProvider() opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -337,11 +337,7 @@ def __init__( If not provided, a default instrumentor builder will be used. This parameter is ignored if `enable_tracing` is False. """ - from google.cloud.aiplatform import initializer - self._tmpl_attrs: dict[str, Any] = { - "project": initializer.global_config.project, - "location": initializer.global_config.location, "model_name": model, "api_type": api_type or "google", "system_instruction": system_instruction, @@ -353,45 +349,51 @@ def __init__( "instrumentor": None, "instrumentor_builder": instrumentor_builder, "enable_tracing": enable_tracing, + "provided_llm_config": copy.deepcopy(llm_config), + "provided_runnable_kwargs": copy.deepcopy(runnable_kwargs), } - self._tmpl_attrs["llm_config"] = llm_config or { + if tools: + _validate_tools(tools) + self._tmpl_attrs["tools"] = tools + + def set_up(self): + """Sets up the agent for execution of queries at runtime. + + It initializes the runnable, binds the runnable with tools. + Project and Location are sourced from environment variables. + """ + project = os.environ.get("GOOGLE_CLOUD_PROJECT") + location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") or os.getenv( + "GOOGLE_CLOUD_LOCATION" + ) + + llm_config = { "config_list": [ { - "project_id": self._tmpl_attrs.get("project"), - "location": self._tmpl_attrs.get("location"), + "project_id": project, + "location": location, "model": self._tmpl_attrs.get("model_name"), "api_type": self._tmpl_attrs.get("api_type"), } ] } - self._tmpl_attrs["runnable_kwargs"] = _prepare_runnable_kwargs( - runnable_kwargs=runnable_kwargs, - llm_config=self._tmpl_attrs.get("llm_config"), + if self._tmpl_attrs.get("provided_llm_config"): + llm_config = self._tmpl_attrs.get("provided_llm_config") + + runnable_kwargs = _prepare_runnable_kwargs( + runnable_kwargs=self._tmpl_attrs.get("provided_runnable_kwargs"), + llm_config=llm_config, system_instruction=self._tmpl_attrs.get("system_instruction"), runnable_name=self._tmpl_attrs.get("runnable_name"), ) - if tools: - # We validate tools at initialization for actionable feedback before - # they are deployed. - _validate_tools(tools) - self._tmpl_attrs["tools"] = tools - def set_up(self): - """Sets up the agent for execution of queries at runtime. - - It initializes the runnable, binds the runnable with tools. - - This method should not be called for an object that being passed to - the ReasoningEngine service for deployment, as it initializes clients - that can not be serialized. - """ if self._tmpl_attrs.get("enable_tracing"): instrumentor_builder = ( self._tmpl_attrs.get("instrumentor_builder") or _default_instrumentor_builder ) self._tmpl_attrs["instrumentor"] = instrumentor_builder( - project_id=self._tmpl_attrs.get("project") + project_id=project, ) # Set up tools. @@ -399,10 +401,10 @@ def set_up(self): ag2_tool_objects = self._tmpl_attrs.get("ag2_tool_objects") if tools and not ag2_tool_objects: from agentplatform._genai import ( - _agent_engines_utils, + _runtimes_utils, ) - autogen_tools = _agent_engines_utils._import_autogen_tools_or_warn() + autogen_tools = _runtimes_utils._import_autogen_tools_or_warn() if autogen_tools: for tool in tools: ag2_tool_objects.append(autogen_tools.Tool(func_or_tool=tool)) @@ -411,22 +413,21 @@ def set_up(self): runnable_builder = ( self._tmpl_attrs.get("runnable_builder") or _default_runnable_builder ) - self._tmpl_attrs["runnable"] = runnable_builder( - **self._tmpl_attrs.get("runnable_kwargs") - ) + self._tmpl_attrs["runnable"] = runnable_builder(**runnable_kwargs) def clone(self) -> "AG2Agent": """Returns a clone of the AG2Agent.""" - import copy return AG2Agent( model=self._tmpl_attrs.get("model_name"), api_type=self._tmpl_attrs.get("api_type"), - llm_config=copy.deepcopy(self._tmpl_attrs.get("llm_config")), + llm_config=copy.deepcopy(self._tmpl_attrs.get("provided_llm_config")), system_instruction=self._tmpl_attrs.get("system_instruction"), runnable_name=self._tmpl_attrs.get("runnable_name"), tools=copy.deepcopy(self._tmpl_attrs.get("tools")), - runnable_kwargs=copy.deepcopy(self._tmpl_attrs.get("runnable_kwargs")), + runnable_kwargs=copy.deepcopy( + self._tmpl_attrs.get("provided_runnable_kwargs") + ), runnable_builder=self._tmpl_attrs.get("runnable_builder"), enable_tracing=self._tmpl_attrs.get("enable_tracing"), instrumentor_builder=self._tmpl_attrs.get("instrumentor_builder"), @@ -488,6 +489,6 @@ def query( **kwargs, ) - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils - return _agent_engines_utils.to_json_serializable_autogen_object(response) + return _runtimes_utils.to_json_serializable_autogen_object(response) diff --git a/agentplatform/agent_engines/templates/langchain.py b/agentplatform/frameworks/langchain.py similarity index 95% rename from agentplatform/agent_engines/templates/langchain.py rename to agentplatform/frameworks/langchain.py index 17751524ba..188b6f0964 100644 --- a/agentplatform/agent_engines/templates/langchain.py +++ b/agentplatform/frameworks/langchain.py @@ -114,15 +114,11 @@ def _default_model_builder( ) return model except ImportError: - import agentplatform - from google.cloud.aiplatform import initializer from langchain_google_vertexai import ChatVertexAI - current_project = initializer.global_config.project - current_location = initializer.global_config.location - agentplatform.init(project=project, location=location) - model = ChatVertexAI(model_name=model_name, **model_kwargs) - agentplatform.init(project=current_project, location=current_location) + model = ChatVertexAI( + model_name=model_name, project=project, location=location, **model_kwargs + ) return model @@ -192,17 +188,13 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils - cloud_trace_exporter = _agent_engines_utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _agent_engines_utils._import_cloud_trace_v2_or_warn() - openinference_langchain = ( - _agent_engines_utils._import_openinference_langchain_or_warn() - ) - opentelemetry = _agent_engines_utils._import_opentelemetry_or_warn() - opentelemetry_sdk_trace = ( - _agent_engines_utils._import_opentelemetry_sdk_trace_or_warn() - ) + cloud_trace_exporter = _runtimes_utils._import_cloud_trace_exporter_or_warn() + cloud_trace_v2 = _runtimes_utils._import_cloud_trace_v2_or_warn() + openinference_langchain = _runtimes_utils._import_openinference_langchain_or_warn() + opentelemetry = _runtimes_utils._import_opentelemetry_or_warn() + opentelemetry_sdk_trace = _runtimes_utils._import_opentelemetry_sdk_trace_or_warn() if all( ( cloud_trace_exporter, @@ -249,7 +241,7 @@ def _default_instrumentor_builder(project_id: str): # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. - if _agent_engines_utils.is_noop_or_proxy_tracer_provider(tracer_provider): + if _runtimes_utils.is_noop_or_proxy_tracer_provider(tracer_provider): tracer_provider = opentelemetry_sdk_trace.TracerProvider() opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -541,11 +533,7 @@ def __init__( TypeError: If there is an invalid tool (e.g. function with an input that did not specify its type). """ - from google.cloud.aiplatform import initializer - self._tmpl_attrs: dict[str, Any] = { - "project": initializer.global_config.project, - "location": initializer.global_config.location, "tools": [], "model_name": model, "system_instruction": system_instruction, @@ -586,20 +574,26 @@ def set_up(self): service for deployment, as it might initialize clients that can not be serialized. """ + import os + + project = os.environ.get("GOOGLE_CLOUD_PROJECT") + location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") or os.getenv( + "GOOGLE_CLOUD_LOCATION" + ) if self._tmpl_attrs.get("enable_tracing"): instrumentor_builder = ( self._tmpl_attrs.get("instrumentor_builder") or _default_instrumentor_builder ) self._tmpl_attrs["instrumentor"] = instrumentor_builder( - project_id=self._tmpl_attrs.get("project") + project_id=project, ) model_builder = self._tmpl_attrs.get("model_builder") or _default_model_builder self._tmpl_attrs["model"] = model_builder( model_name=self._tmpl_attrs.get("model_name"), model_kwargs=self._tmpl_attrs.get("model_kwargs"), - project=self._tmpl_attrs.get("project"), - location=self._tmpl_attrs.get("location"), + project=project, + location=location, ) runnable_builder = ( self._tmpl_attrs.get("runnable_builder") or _default_runnable_builder diff --git a/agentplatform/agent_engines/templates/langgraph.py b/agentplatform/frameworks/langgraph.py similarity index 95% rename from agentplatform/agent_engines/templates/langgraph.py rename to agentplatform/frameworks/langgraph.py index 0ebdb690cc..c460e0a518 100644 --- a/agentplatform/agent_engines/templates/langgraph.py +++ b/agentplatform/frameworks/langgraph.py @@ -105,15 +105,11 @@ def _default_model_builder( ) return model except ImportError: - import agentplatform - from google.cloud.aiplatform import initializer from langchain_google_vertexai import ChatVertexAI - current_project = initializer.global_config.project - current_location = initializer.global_config.location - agentplatform.init(project=project, location=location) - model = ChatVertexAI(model_name=model_name, **model_kwargs) - agentplatform.init(project=current_project, location=current_location) + model = ChatVertexAI( + model_name=model_name, project=project, location=location, **model_kwargs + ) return model @@ -168,17 +164,13 @@ def _default_runnable_builder( def _default_instrumentor_builder(project_id: str): - from agentplatform._genai import _agent_engines_utils + from agentplatform._genai import _runtimes_utils - cloud_trace_exporter = _agent_engines_utils._import_cloud_trace_exporter_or_warn() - cloud_trace_v2 = _agent_engines_utils._import_cloud_trace_v2_or_warn() - openinference_langchain = ( - _agent_engines_utils._import_openinference_langchain_or_warn() - ) - opentelemetry = _agent_engines_utils._import_opentelemetry_or_warn() - opentelemetry_sdk_trace = ( - _agent_engines_utils._import_opentelemetry_sdk_trace_or_warn() - ) + cloud_trace_exporter = _runtimes_utils._import_cloud_trace_exporter_or_warn() + cloud_trace_v2 = _runtimes_utils._import_cloud_trace_v2_or_warn() + openinference_langchain = _runtimes_utils._import_openinference_langchain_or_warn() + opentelemetry = _runtimes_utils._import_opentelemetry_or_warn() + opentelemetry_sdk_trace = _runtimes_utils._import_opentelemetry_sdk_trace_or_warn() if all( ( cloud_trace_exporter, @@ -224,7 +216,7 @@ def _default_instrumentor_builder(project_id: str): # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. - if _agent_engines_utils.is_noop_or_proxy_tracer_provider(tracer_provider): + if _runtimes_utils.is_noop_or_proxy_tracer_provider(tracer_provider): tracer_provider = opentelemetry_sdk_trace.TracerProvider() opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -465,11 +457,7 @@ def checkpointer_builder(**kwargs): TypeError: If there is an invalid tool (e.g. function with an input that did not specify its type). """ - from google.cloud.aiplatform import initializer - self._tmpl_attrs: dict[str, Any] = { - "project": initializer.global_config.project, - "location": initializer.global_config.location, "tools": [], "model_name": model, "model_kwargs": model_kwargs, @@ -502,20 +490,26 @@ def set_up(self): the ReasoningEngine service for deployment, as it initializes clients that can not be serialized. """ + import os + + project = os.environ.get("GOOGLE_CLOUD_PROJECT") + location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") or os.getenv( + "GOOGLE_CLOUD_LOCATION" + ) if self._tmpl_attrs.get("enable_tracing"): instrumentor_builder = ( self._tmpl_attrs.get("instrumentor_builder") or _default_instrumentor_builder ) self._tmpl_attrs["instrumentor"] = instrumentor_builder( - project_id=self._tmpl_attrs.get("project") + project_id=project, ) model_builder = self._tmpl_attrs.get("model_builder") or _default_model_builder self._tmpl_attrs["model"] = model_builder( model_name=self._tmpl_attrs.get("model_name"), model_kwargs=self._tmpl_attrs.get("model_kwargs"), - project=self._tmpl_attrs.get("project"), - location=self._tmpl_attrs.get("location"), + project=project, + location=location, ) checkpointer_builder = self._tmpl_attrs.get("checkpointer_builder") if checkpointer_builder: diff --git a/agentplatform/agent_engines/templates/llama_index.py b/agentplatform/frameworks/llama_index.py similarity index 93% rename from agentplatform/agent_engines/templates/llama_index.py rename to agentplatform/frameworks/llama_index.py index d9147fccef..3f68e93e39 100644 --- a/agentplatform/agent_engines/templates/llama_index.py +++ b/agentplatform/frameworks/llama_index.py @@ -58,8 +58,6 @@ def _default_model_builder( model_kwargs: Optional[Mapping[str, Any]] = None, ) -> "FunctionCallingLLM": """Creates a default model builder for LlamaIndex.""" - import agentplatform - from google.cloud.aiplatform import initializer from llama_index.llms import google_genai model_kwargs = model_kwargs or {} @@ -68,9 +66,6 @@ def _default_model_builder( vertexai_config={"project": project, "location": location}, **model_kwargs, ) - current_project = initializer.global_config.project - current_location = initializer.global_config.location - agentplatform.init(project=current_project, location=current_location) return model @@ -346,10 +341,6 @@ def __init__( enable_tracing (bool): Optional. Whether to enable tracing. Defaults to False. """ - from google.cloud.aiplatform import initializer - - self._project = initializer.global_config.project - self._location = initializer.global_config.location self._model_name = model self._system_instruction = system_instruction self._prompt = prompt @@ -383,21 +374,28 @@ def set_up(self): the ReasoningEngine service for deployment, as it initializes clients that can not be serialized. """ + import os + + project = os.environ.get("GOOGLE_CLOUD_PROJECT") + location = os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_LOCATION") or os.getenv( + "GOOGLE_CLOUD_LOCATION" + ) + if self._enable_tracing: - from agentplatform._genai.agent_engines import ( - _agent_engines_utils, + from agentplatform._genai import ( + _runtimes_utils, ) cloud_trace_exporter = ( - _agent_engines_utils._import_cloud_trace_exporter_or_warn() + _runtimes_utils._import_cloud_trace_exporter_or_warn() ) - cloud_trace_v2 = _agent_engines_utils._import_cloud_trace_v2_or_warn() + cloud_trace_v2 = _runtimes_utils._import_cloud_trace_v2_or_warn() openinference_llama_index = ( - _agent_engines_utils._import_openinference_llama_index_or_warn() + _runtimes_utils._import_openinference_llama_index_or_warn() ) - opentelemetry = _agent_engines_utils._import_opentelemetry_or_warn() + opentelemetry = _runtimes_utils._import_opentelemetry_or_warn() opentelemetry_sdk_trace = ( - _agent_engines_utils._import_opentelemetry_sdk_trace_or_warn() + _runtimes_utils._import_opentelemetry_sdk_trace_or_warn() ) if all( ( @@ -412,9 +410,9 @@ def set_up(self): credentials, _ = google.auth.default() span_exporter = cloud_trace_exporter.CloudTraceSpanExporter( - project_id=self._project, + project_id=project, client=cloud_trace_v2.TraceServiceClient( - credentials=credentials.with_quota_project(self._project), + credentials=credentials.with_quota_project(project), ), ) span_processor: SpanProcessor = ( @@ -447,9 +445,7 @@ def set_up(self): # Avoids AttributeError: # 'ProxyTracerProvider' and 'NoOpTracerProvider' objects has no # attribute 'add_span_processor'. - if _agent_engines_utils.is_noop_or_proxy_tracer_provider( - tracer_provider - ): + if _runtimes_utils.is_noop_or_proxy_tracer_provider(tracer_provider): tracer_provider = opentelemetry_sdk_trace.TracerProvider() opentelemetry.trace.set_tracer_provider(tracer_provider) # Avoids OpenTelemetry client already exists error. @@ -482,8 +478,8 @@ def set_up(self): self._model = model_builder( model_name=self._model_name, model_kwargs=self._model_kwargs, - project=self._project, - location=self._location, + project=project, + location=location, ) if self._retriever_builder: @@ -546,8 +542,8 @@ def query( Returns: The output of querying the Agent with the given input and config. """ - from agentplatform._genai.agent_engines import ( - _agent_engines_utils, + from agentplatform._genai import ( + _runtimes_utils, ) if isinstance(input, str): @@ -557,9 +553,9 @@ def query( self.set_up() if kwargs.get("batch"): - nest_asyncio = _agent_engines_utils._import_nest_asyncio_or_warn() + nest_asyncio = _runtimes_utils._import_nest_asyncio_or_warn() nest_asyncio.apply() - return _agent_engines_utils.to_json_serializable_llama_index_object( + return _runtimes_utils.to_json_serializable_llama_index_object( self._runnable.run(**input, **kwargs) ) diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py b/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py index 4fa0d87aff..49dc705a54 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_a2a.py @@ -13,63 +13,25 @@ # limitations under the License. # +import importlib import os -import sys -import tempfile from unittest import mock -import pytest -import cloudpickle -import pydantic from google import auth -from google.api_core import operation as ga_operation from google.auth import credentials as auth_credentials -from google.cloud import storage -from google.cloud import aiplatform -from google.cloud.aiplatform import base - -from google.cloud.aiplatform_v1 import types -from google.cloud.aiplatform_v1.services import reasoning_engine_service -from vertexai import agent_engines -from vertexai.agent_engines import _agent_engines -from vertexai.agent_engines import _utils -from google.protobuf import struct_pb2 - - -class CapitalizeEngine: - """A sample Agent Engine.""" - - def query(self, unused_arbitrary_string_name: str) -> str: - """Runs the engine.""" - return unused_arbitrary_string_name.upper() - - -class CapitalizeEngineWithCard(CapitalizeEngine): - - def __init__(self, card): - self.agent_card = card - - def __getstate__(self): - state = self.__dict__.copy() - if hasattr(self.agent_card, "DESCRIPTOR"): - state["agent_card"] = None - return state +import agentplatform +from agentplatform._genai import _runtimes_utils +from agentplatform._genai import types as _genai_types +from google.genai import types as genai_types +import pytest - def __setstate__(self, state): - self.__dict__.update(state) +from a2a.types import AgentSkill -def _create_empty_fake_package(package_name: str) -> str: - temp_dir = tempfile.mkdtemp() - package_dir = os.path.join(temp_dir, package_name) - os.makedirs(package_dir) - init_path = os.path.join(package_dir, "__init__.py") - open(init_path, "w").close() - return temp_dir +from agentplatform.frameworks.a2a import A2aAgent +from agentplatform.frameworks.a2a import create_agent_card -_TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials()) -_TEST_STAGING_BUCKET = "gs://test-bucket" _TEST_LOCATION = "us-central1" _TEST_PROJECT = "test-project" _TEST_RESOURCE_ID = "1028944691210842416" @@ -77,65 +39,25 @@ def _create_empty_fake_package(package_name: str) -> str: _TEST_AGENT_ENGINE_RESOURCE_NAME = ( f"{_TEST_PARENT}/reasoningEngines/{_TEST_RESOURCE_ID}" ) -_TEST_AGENT_ENGINE_DISPLAY_NAME = "Agent Engine Display Name" -_TEST_GCS_DIR_NAME = _agent_engines._DEFAULT_GCS_DIR_NAME -_TEST_BLOB_FILENAME = _agent_engines._BLOB_FILENAME -_TEST_REQUIREMENTS_FILE = _agent_engines._REQUIREMENTS_FILE -_TEST_EXTRA_PACKAGES_FILE = _agent_engines._EXTRA_PACKAGES_FILE -_TEST_STANDARD_API_MODE = _agent_engines._STANDARD_API_MODE -_TEST_DEFAULT_METHOD_NAME = _agent_engines._DEFAULT_METHOD_NAME -_TEST_MODE_KEY_IN_SCHEMA = _agent_engines._MODE_KEY_IN_SCHEMA - -_TEST_AGENT_ENGINE_EXTRA_PACKAGE = "fake.py" - -_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH = _create_empty_fake_package( - _TEST_AGENT_ENGINE_EXTRA_PACKAGE +_TEST_AGENT_ENGINE_DISPLAY_NAME = "test-a2a-agent" +_TEST_STAGING_BUCKET = "gs://test-bucket" +_TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials()) +_TEST_SKILL = AgentSkill( + id="hello_world", + name="Returns hello world", + description="just returns hello world", + tags=["hello world"], + examples=["hi", "hello world"], ) -_TEST_AGENT_ENGINE_REQUIREMENTS = [ - "google-cloud-aiplatform==1.29.0", - "langchain", -] - -_TEST_AGENT_ENGINE_GCS_URI = "{}/{}/{}".format( - _TEST_STAGING_BUCKET, - _TEST_GCS_DIR_NAME, - _TEST_BLOB_FILENAME, -) -_TEST_AGENT_ENGINE_DEPENDENCY_FILES_GCS_URI = "{}/{}/{}".format( - _TEST_STAGING_BUCKET, - _TEST_GCS_DIR_NAME, - _TEST_EXTRA_PACKAGES_FILE, -) -_TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI = "{}/{}/{}".format( - _TEST_STAGING_BUCKET, - _TEST_GCS_DIR_NAME, - _TEST_REQUIREMENTS_FILE, -) -_TEST_AGENT_ENGINE_QUERY_SCHEMA = _utils.to_proto( - _utils.generate_schema( - CapitalizeEngine().query, - schema_name=_TEST_DEFAULT_METHOD_NAME, +def _make_agent() -> A2aAgent: + card = create_agent_card( + agent_name="Test", + description="Test", + skills=[_TEST_SKILL], ) -) -_TEST_AGENT_ENGINE_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STANDARD_API_MODE - -_TEST_AGENT_ENGINE_PACKAGE_SPEC = types.ReasoningEngineSpec.PackageSpec( - python_version=f"{sys.version_info.major}.{sys.version_info.minor}", - pickle_object_gcs_uri=_TEST_AGENT_ENGINE_GCS_URI, - dependency_files_gcs_uri=_TEST_AGENT_ENGINE_DEPENDENCY_FILES_GCS_URI, - requirements_gcs_uri=_TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI, -) - -_TEST_AGENT_ENGINE_OBJ = types.ReasoningEngine( - name=_TEST_AGENT_ENGINE_RESOURCE_NAME, - spec=types.ReasoningEngineSpec( - package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC, - agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK, - ), -) -_TEST_AGENT_ENGINE_OBJ.spec.class_methods.append(_TEST_AGENT_ENGINE_QUERY_SCHEMA) + return A2aAgent(agent_card=card) @pytest.fixture(scope="module") @@ -148,128 +70,109 @@ def google_auth_mock(): yield google_auth_mock -@pytest.fixture(scope="module") -def cloud_storage_create_bucket_mock(): - with mock.patch.object(storage, "Client") as cloud_storage_mock: - bucket_mock = mock.Mock(spec=storage.Bucket) - bucket_mock.blob.return_value.open.return_value = "blob_file" - bucket_mock.blob.return_value.upload_from_filename.return_value = None - bucket_mock.blob.return_value.upload_from_string.return_value = None - - cloud_storage_mock.get_bucket = mock.Mock( - side_effect=ValueError("bucket not found") - ) - cloud_storage_mock.bucket.return_value = bucket_mock - cloud_storage_mock.create_bucket.return_value = bucket_mock - - yield cloud_storage_mock - - -@pytest.fixture(scope="module") -def cloudpickle_load_mock(): - with mock.patch.object(cloudpickle, "load") as cloudpickle_load_mock: - yield cloudpickle_load_mock - - -@pytest.fixture(scope="module") -def create_agent_engine_mock(): - with mock.patch.object( - reasoning_engine_service.ReasoningEngineServiceClient, - "create_reasoning_engine", - ) as create_agent_engine_mock: - create_agent_engine_lro_mock = mock.Mock(spec=ga_operation.Operation) - create_agent_engine_lro_mock.result.return_value = _TEST_AGENT_ENGINE_OBJ - create_agent_engine_mock.return_value = create_agent_engine_lro_mock - yield create_agent_engine_mock - - -@pytest.fixture(scope="function") -def get_gca_resource_mock(): - with mock.patch.object( - base.VertexAiResourceNoun, - "_get_gca_resource", - ) as get_gca_resource_mock: - get_gca_resource_mock.return_value = _TEST_AGENT_ENGINE_OBJ - yield get_gca_resource_mock - - @pytest.mark.usefixtures("google_auth_mock") -class TestAgentEngineA2A: +class TestA2aAgentCreate: def setup_method(self): - aiplatform.init( + importlib.reload(agentplatform) + os.environ["GOOGLE_CLOUD_PROJECT"] = _TEST_PROJECT + os.environ["GOOGLE_CLOUD_LOCATION"] = _TEST_LOCATION + self.client = agentplatform.Client( project=_TEST_PROJECT, location=_TEST_LOCATION, credentials=_TEST_CREDENTIALS, - staging_bucket=_TEST_STAGING_BUCKET, ) - def test_create_agent_engine_with_protobuf_agent_card( + def teardown_method(self): + for key in [ + "GOOGLE_CLOUD_PROJECT", + "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", + "GOOGLE_CLOUD_LOCATION", + "GOOGLE_GENAI_USE_VERTEXAI", + ]: + os.environ.pop(key, None) + + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") + @mock.patch.object( + _runtimes_utils, + "_get_reasoning_engine_id", + return_value=_TEST_RESOURCE_ID, + ) + def test_create_agent_with_agent_card( self, - create_agent_engine_mock, - cloud_storage_create_bucket_mock, - cloudpickle_load_mock, - get_gca_resource_mock, + mock_get_reasoning_engine_id, + mock_await_operation, + mock_prepare, ): - a2a_pb2 = None - # fmt: off - try: - try: - from a2a.compat.v0_3 import a2a_v0_3_pb2 as a2a_pb2 - except ImportError: - from a2a.grpc import a2a_pb2 - has_a2a_pb2 = True - except (ImportError, TypeError): - has_a2a_pb2 = False - # fmt: on - - if not has_a2a_pb2: - pytest.skip("a2a_pb2 could not be imported.") - - card = a2a_pb2.AgentCard(name="test_agent_card") - agent = CapitalizeEngineWithCard(card) - - agent_engines.create( - agent, - display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, - requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, - extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], + mock_await_operation.return_value = _genai_types.RuntimeOperation( + response=_genai_types.ReasoningEngine( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + ) ) + agent = _make_agent() + with mock.patch.object( + self.client.runtimes._api_client, "request" + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="") + self.client.runtimes.create( + agent=agent, + config=_genai_types.RuntimeConfig( + display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, + staging_bucket=_TEST_STAGING_BUCKET, + ), + ) - expected_reasoning_engine = types.ReasoningEngine( - display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, - spec=types.ReasoningEngineSpec( - package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC, - agent_framework=_agent_engines._DEFAULT_AGENT_FRAMEWORK, - ), - ) - from google.protobuf import json_format + request_mock.assert_called_once() + _, _, request_dict, _ = request_mock.call_args[0] + class_methods = request_dict["spec"]["class_methods"] + assert any("a2a_agent_card" in method for method in class_methods) - expected_class_method = struct_pb2.Struct() - expected_class_method.CopyFrom(_TEST_AGENT_ENGINE_QUERY_SCHEMA) - expected_class_method["a2a_agent_card"] = json_format.MessageToJson(card) - expected_reasoning_engine.spec.class_methods.append(expected_class_method) - create_agent_engine_mock.assert_called_with( - parent=_TEST_PARENT, - reasoning_engine=expected_reasoning_engine, - ) +class TestA2aLocationResolution: + @pytest.fixture(autouse=True) + def patch_create_rest_routes(self, monkeypatch): + """Patch to avoid optional dependencies not actually used in these tests.""" + import sys + from unittest.mock import MagicMock - def test_create_agent_engine_with_invalid_agent_card( - self, - create_agent_engine_mock, - cloud_storage_create_bucket_mock, - cloudpickle_load_mock, - get_gca_resource_mock, - ): - agent = CapitalizeEngineWithCard(card="invalid_card_type_string") + mock_rest_routes = MagicMock() + mock_rest_routes.create_rest_routes = MagicMock() + monkeypatch.setitem( + sys.modules, "a2a.server.routes.rest_routes", mock_rest_routes + ) - with pytest.raises( - TypeError, - match="Unsupported AgentCard type", + def test_location_from_agent_engine_env_var(self): + agent = _make_agent() + with mock.patch.dict( + os.environ, + {"GOOGLE_CLOUD_AGENT_ENGINE_LOCATION": "us-east1"}, + clear=True, ): - agent_engines.create( - agent, - display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, - requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, - extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], - ) + agent.set_up() + assert os.environ.get("GOOGLE_CLOUD_LOCATION") == "us-east1" + + if hasattr(agent.agent_card, "url"): + url = agent.agent_card.url + else: + url = agent.agent_card.supported_interfaces[0].url + assert "us-east1" in url + + def test_location_env_var_precedence(self): + agent = _make_agent() + with mock.patch.dict( + os.environ, + { + "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION": "us-east1", + "GOOGLE_CLOUD_LOCATION": "us-west1", + }, + clear=True, + ): + agent.set_up() + # Should not overwrite existing GOOGLE_CLOUD_LOCATION. + assert os.environ.get("GOOGLE_CLOUD_LOCATION") == "us-west1" + + if hasattr(agent.agent_card, "url"): + url = agent.agent_card.url + else: + url = agent.agent_card.supported_interfaces[0].url + assert "us-east1" in url diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py index 7e3d5dba6f..2e8592aafc 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py @@ -34,9 +34,9 @@ from google.cloud.aiplatform import initializer from google.cloud.aiplatform_v1 import types as aip_types from google.cloud.aiplatform_v1.services import reasoning_engine_service -from agentplatform._genai import agent_engines -from agentplatform._genai import _agent_engines_utils -from agentplatform.agent_engines.templates import ( +from agentplatform._genai import runtimes +from agentplatform._genai import _runtimes_utils +from agentplatform.frameworks import ( adk as adk_template, ) from google.genai import types @@ -414,12 +414,13 @@ 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) + # The agent frameworks source project and location from environment + # variables (not the global initializer). + os.environ["GOOGLE_CLOUD_PROJECT"] = _TEST_PROJECT + os.environ["GOOGLE_CLOUD_LOCATION"] = _TEST_LOCATION def teardown_method(self): - initializer.global_pool.shutdown(wait=True) for key in [ "GOOGLE_CLOUD_PROJECT", "GOOGLE_CLOUD_AGENT_ENGINE_LOCATION", @@ -430,8 +431,10 @@ 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 + # Project and location are no longer stored as template attributes; they + # are sourced from environment variables during set_up(). + assert app._tmpl_attrs.get("project") is None + assert app._tmpl_attrs.get("location") is None assert app._tmpl_attrs.get("runner") is None def test_set_up( @@ -559,9 +562,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() @@ -1214,8 +1215,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( @@ -1322,7 +1323,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 @@ -1477,7 +1478,7 @@ def update_agent_engine_mock(): "get_gca_resource_mock", "get_agent_engine_mock", ) -class TestAgentEngines: +class TestRuntimes: def setup_method(self): importlib.reload(initializer) @@ -1514,8 +1515,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, @@ -1530,13 +1531,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}, ) @@ -1569,8 +1570,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, @@ -1585,13 +1586,13 @@ def test_update_default_telemetry_enablement( "projects/test-project/locations/us-central1/reasoningEngines/123456/operations/789" ) mock_update.return_value = mock_operation - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, ) ) client = agentplatform.Client(project=_TEST_PROJECT, location=_TEST_LOCATION) - client.agent_engines.update( + client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, agent=adk_template.AdkApp(agent=_TEST_AGENT), config={ diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_ag2.py b/tests/unit/agentplatform/frameworks/test_frameworks_ag2.py index b8e2337f51..3861f8d4bb 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_ag2.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_ag2.py @@ -15,14 +15,14 @@ import dataclasses import importlib import json +import os from typing import Optional from unittest import mock from google import auth import agentplatform -from google.cloud.aiplatform import initializer -from agentplatform import agent_engines -from agentplatform._genai import _agent_engines_utils +from agentplatform import frameworks +from agentplatform._genai import _runtimes_utils import pytest @@ -94,7 +94,7 @@ def dataclasses_is_dataclass_mock(): @pytest.fixture def to_json_serializable_autogen_object_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "to_json_serializable_autogen_object", ) as to_json_serializable_autogen_object_mock: to_json_serializable_autogen_object_mock.return_value = {} @@ -104,7 +104,7 @@ def to_json_serializable_autogen_object_mock(): @pytest.fixture def cloud_trace_exporter_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_cloud_trace_exporter_or_warn", ) as cloud_trace_exporter_mock: yield cloud_trace_exporter_mock @@ -127,7 +127,7 @@ def simple_span_processor_mock(): @pytest.fixture def autogen_instrumentor_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_autogen_or_warn", ) as autogen_instrumentor_mock: yield autogen_instrumentor_mock @@ -136,7 +136,7 @@ def autogen_instrumentor_mock(): @pytest.fixture def autogen_instrumentor_none_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_autogen_or_warn", ) as autogen_instrumentor_mock: autogen_instrumentor_mock.return_value = None @@ -146,7 +146,7 @@ def autogen_instrumentor_none_mock(): @pytest.fixture def autogen_tools_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_autogen_tools_or_warn", ) as autogen_tools_mock: autogen_tools_mock.return_value = mock.MagicMock() @@ -170,24 +170,30 @@ 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, - ) + # The agent frameworks source project and location from environment + # variables (not the global initializer). + 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 + # Project and location are no longer stored as template attributes; they + # are sourced from environment variables during set_up(). + assert agent._tmpl_attrs.get("project") is None + assert agent._tmpl_attrs.get("location") is None assert agent._tmpl_attrs.get("runnable") is None def test_initialization_with_tools(self, autogen_tools_mock): @@ -195,7 +201,7 @@ def test_initialization_with_tools(self, autogen_tools_mock): place_tool_query, place_photo_query, ] - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, system_instruction=_TEST_SYSTEM_INSTRUCTION, @@ -210,7 +216,7 @@ def test_initialization_with_tools(self, autogen_tools_mock): assert agent._tmpl_attrs.get("ag2_tool_objects") def test_set_up(self): - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, runnable_builder=lambda **kwargs: kwargs, @@ -220,7 +226,7 @@ def test_set_up(self): assert agent._tmpl_attrs.get("runnable") is not None def test_clone(self): - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, runnable_builder=lambda **kwargs: kwargs, @@ -234,7 +240,7 @@ def test_clone(self): assert agent_clone._tmpl_attrs.get("runnable") is not None def test_query(self, to_json_serializable_autogen_object_mock): - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, ) @@ -262,7 +268,7 @@ def test_enable_tracing( simple_span_processor_mock, autogen_instrumentor_mock, ): - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, enable_tracing=True, @@ -275,7 +281,7 @@ def test_enable_tracing( @pytest.mark.usefixtures("caplog") def test_enable_tracing_warning(self, caplog, autogen_instrumentor_none_mock): - agent = agent_engines.AG2Agent( + agent = frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, enable_tracing=True, @@ -294,7 +300,7 @@ def _return_input_no_typing(input_): class TestConvertToolsOrRaiseErrors: def test_raise_untyped_input_args(self, agentplatform_init_mock): with pytest.raises(TypeError, match=r"has untyped input_arg"): - agent_engines.AG2Agent( + frameworks.AG2Agent( model=_TEST_MODEL, runnable_name=_TEST_RUNNABLE_NAME, tools=[_return_input_no_typing], @@ -302,23 +308,23 @@ def test_raise_untyped_input_args(self, agentplatform_init_mock): class TestToJsonSerializableAutoGenObject: - """Tests for `_agent_engines_utils.to_json_serializable_autogen_object`.""" + """Tests for `_runtimes_utils.to_json_serializable_autogen_object`.""" def test_autogen_chat_result( self, dataclasses_asdict_mock, dataclasses_is_dataclass_mock, ): - mock_chat_result: _agent_engines_utils.AutogenChatResult = mock.Mock( - spec=_agent_engines_utils.AutogenChatResult + mock_chat_result: _runtimes_utils.AutogenChatResult = mock.Mock( + spec=_runtimes_utils.AutogenChatResult ) - _agent_engines_utils.to_json_serializable_autogen_object(mock_chat_result) + _runtimes_utils.to_json_serializable_autogen_object(mock_chat_result) dataclasses_is_dataclass_mock.assert_called_once_with(mock_chat_result) dataclasses_asdict_mock.assert_called_once_with(mock_chat_result) def test_autogen_run_response(self): - mock_response: _agent_engines_utils.AutogenRunResponse = mock.Mock( - spec=_agent_engines_utils.AutogenRunResponse + mock_response: _runtimes_utils.AutogenRunResponse = mock.Mock( + spec=_runtimes_utils.AutogenRunResponse ) mock_agent = MockAgent( name="TestAgent", @@ -342,13 +348,13 @@ def test_autogen_run_response(self): }, "cost": {"total_cost": 5.5}, } - got = _agent_engines_utils.to_json_serializable_autogen_object(mock_response) + got = _runtimes_utils.to_json_serializable_autogen_object(mock_response) mock_response.process.assert_called_once() assert got == want def test_autogen_empty_run_response(self): - mock_response: _agent_engines_utils.AutogenRunResponse = mock.Mock( - spec=_agent_engines_utils.AutogenRunResponse + mock_response: _runtimes_utils.AutogenRunResponse = mock.Mock( + spec=_runtimes_utils.AutogenRunResponse ) mock_response.summary = None mock_response.messages = [] @@ -362,7 +368,7 @@ def test_autogen_empty_run_response(self): "last_speaker": None, "cost": None, } - got = _agent_engines_utils.to_json_serializable_autogen_object(mock_response) + got = _runtimes_utils.to_json_serializable_autogen_object(mock_response) assert got == want @@ -377,7 +383,7 @@ class SimpleDataClass: instance = SimpleDataClass(field1="value1", field2=123) want = {"field1": "value1", "field2": 123} - got = _agent_engines_utils._dataclass_to_dict_or_raise(instance) + got = _runtimes_utils._dataclass_to_dict_or_raise(instance) assert got == want def test_not_a_dataclass_raises_type_error(self): @@ -386,4 +392,4 @@ class NotADataclass: instance = NotADataclass() with pytest.raises(TypeError, match="Object is not a dataclass"): - _agent_engines_utils._dataclass_to_dict_or_raise(instance) + _runtimes_utils._dataclass_to_dict_or_raise(instance) diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_langchain.py b/tests/unit/agentplatform/frameworks/test_frameworks_langchain.py index 8d50518aff..e174ee5443 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_langchain.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_langchain.py @@ -13,15 +13,15 @@ # limitations under the License. # import importlib +import os from typing import Optional from unittest import mock from google import auth import agentplatform -from google.cloud.aiplatform import initializer -from agentplatform import agent_engines +from agentplatform import frameworks -from agentplatform._genai import _agent_engines_utils +from agentplatform._genai import _runtimes_utils import pytest @@ -90,7 +90,7 @@ def langchain_dump_mock(): @pytest.fixture def cloud_trace_exporter_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_cloud_trace_exporter_or_warn", ) as cloud_trace_exporter_mock: yield cloud_trace_exporter_mock @@ -113,7 +113,7 @@ def simple_span_processor_mock(): @pytest.fixture def langchain_instrumentor_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_langchain_or_warn", ) as langchain_instrumentor_mock: yield langchain_instrumentor_mock @@ -122,7 +122,7 @@ def langchain_instrumentor_mock(): @pytest.fixture def langchain_instrumentor_none_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_langchain_or_warn", ) as langchain_instrumentor_mock: langchain_instrumentor_mock.return_value = None @@ -132,12 +132,11 @@ 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, - ) + # The agent frameworks source project and location from environment + # variables (not the global initializer). + os.environ["GOOGLE_CLOUD_PROJECT"] = _TEST_PROJECT + os.environ["GOOGLE_CLOUD_LOCATION"] = _TEST_LOCATION self.prompt = { "input": lambda x: x["input"], "agent_scratchpad": ( @@ -152,13 +151,20 @@ 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 + # Project and location are no longer stored as template attributes; they + # are sourced from environment variables during set_up(). + assert agent._tmpl_attrs.get("project") is None + assert agent._tmpl_attrs.get("location") is None assert agent._tmpl_attrs.get("runnable") is None def test_initialization_with_tools(self): @@ -166,7 +172,7 @@ def test_initialization_with_tools(self): place_tool_query, StructuredTool.from_function(place_photo_query), ] - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, system_instruction=_TEST_SYSTEM_INSTRUCTION, tools=tools, @@ -180,7 +186,7 @@ def test_initialization_with_tools(self): assert agent._tmpl_attrs.get("runnable") is not None def test_set_up(self): - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, prompt=self.prompt, output_parser=self.output_parser, @@ -192,7 +198,7 @@ def test_set_up(self): assert agent._tmpl_attrs.get("runnable") is not None def test_clone(self): - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, prompt=self.prompt, output_parser=self.output_parser, @@ -208,7 +214,7 @@ def test_clone(self): assert agent_clone._tmpl_attrs.get("runnable") is not None def test_query(self, langchain_dump_mock): - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, prompt=self.prompt, output_parser=self.output_parser, @@ -222,7 +228,7 @@ def test_query(self, langchain_dump_mock): ) def test_stream_query(self, langchain_dump_mock): - agent = agent_engines.LangchainAgent(model=_TEST_MODEL) + agent = frameworks.LangchainAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].stream.return_value = [] list(agent.stream_query(input="test stream query")) @@ -240,7 +246,7 @@ def test_enable_tracing( simple_span_processor_mock, langchain_instrumentor_mock, ): - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, prompt=self.prompt, output_parser=self.output_parser, @@ -257,7 +263,7 @@ def test_enable_tracing( @pytest.mark.usefixtures("caplog") def test_enable_tracing_warning(self, caplog, langchain_instrumentor_none_mock): - agent = agent_engines.LangchainAgent( + agent = frameworks.LangchainAgent( model=_TEST_MODEL, prompt=self.prompt, output_parser=self.output_parser, @@ -277,7 +283,7 @@ def _return_input_no_typing(input_): class TestConvertToolsOrRaiseErrors: def test_raise_untyped_input_args(self, agentplatform_init_mock): with pytest.raises(TypeError, match=r"has untyped input_arg"): - agent_engines.LangchainAgent( + frameworks.LangchainAgent( model=_TEST_MODEL, tools=[_return_input_no_typing], ) @@ -289,7 +295,7 @@ def test_raise_both_system_instruction_and_prompt_error(self, agentplatform_init ValueError, match=r"Only one of `prompt` or `system_instruction` should be specified.", ): - agent_engines.LangchainAgent( + frameworks.LangchainAgent( model=_TEST_MODEL, system_instruction=_TEST_SYSTEM_INSTRUCTION, prompt=prompts.ChatPromptTemplate.from_messages( diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_langgraph.py b/tests/unit/agentplatform/frameworks/test_frameworks_langgraph.py index 82fa32d53b..dafe9cc2ec 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_langgraph.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_langgraph.py @@ -13,14 +13,14 @@ # limitations under the License. # import importlib +import os from typing import Any, Dict, List, Optional from unittest import mock from google import auth import agentplatform -from google.cloud.aiplatform import initializer -from agentplatform import agent_engines -from agentplatform._genai import _agent_engines_utils +from agentplatform import frameworks +from agentplatform._genai import _runtimes_utils import pytest from langchain_core import runnables @@ -101,7 +101,7 @@ def langchain_dump_mock(): @pytest.fixture def cloud_trace_exporter_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_cloud_trace_exporter_or_warn", ) as cloud_trace_exporter_mock: yield cloud_trace_exporter_mock @@ -124,7 +124,7 @@ def simple_span_processor_mock(): @pytest.fixture def langchain_instrumentor_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_langchain_or_warn", ) as langchain_instrumentor_mock: yield langchain_instrumentor_mock @@ -133,7 +133,7 @@ def langchain_instrumentor_mock(): @pytest.fixture def langchain_instrumentor_none_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_langchain_or_warn", ) as langchain_instrumentor_mock: langchain_instrumentor_mock.return_value = None @@ -143,21 +143,27 @@ 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, - ) + # The agent frameworks source project and location from environment + # variables (not the global initializer). + 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 + # Project and location are no longer stored as template attributes; they + # are sourced from environment variables during set_up(). + assert agent._tmpl_attrs.get("project") is None + assert agent._tmpl_attrs.get("location") is None assert agent._tmpl_attrs.get("runnable") is None def test_initialization_with_tools(self): @@ -165,7 +171,7 @@ def test_initialization_with_tools(self): place_tool_query, StructuredTool.from_function(place_photo_query), ] - agent = agent_engines.LanggraphAgent( + agent = frameworks.LanggraphAgent( model=_TEST_MODEL, tools=tools, model_builder=lambda **kwargs: kwargs, @@ -178,7 +184,7 @@ def test_initialization_with_tools(self): assert agent._tmpl_attrs.get("runnable") is not None def test_set_up(self): - agent = agent_engines.LanggraphAgent( + agent = frameworks.LanggraphAgent( model=_TEST_MODEL, model_builder=lambda **kwargs: kwargs, runnable_builder=lambda **kwargs: kwargs, @@ -188,7 +194,7 @@ def test_set_up(self): assert agent._tmpl_attrs.get("runnable") is not None def test_clone(self): - agent = agent_engines.LanggraphAgent( + agent = frameworks.LanggraphAgent( model=_TEST_MODEL, model_builder=lambda **kwargs: kwargs, runnable_builder=lambda **kwargs: kwargs, @@ -202,7 +208,7 @@ def test_clone(self): assert agent_clone._tmpl_attrs.get("runnable") is not None def test_query(self, langchain_dump_mock): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() mocks = mock.Mock() mocks.attach_mock(mock=agent._tmpl_attrs.get("runnable"), attribute="invoke") @@ -217,7 +223,7 @@ def test_query(self, langchain_dump_mock): ) def test_stream_query(self, langchain_dump_mock): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].stream.return_value = [] list(agent.stream_query(input="test stream query")) @@ -238,7 +244,7 @@ def test_enable_tracing( simple_span_processor_mock, langchain_instrumentor_mock, ): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL, enable_tracing=True) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL, enable_tracing=True) assert agent._tmpl_attrs.get("instrumentor") is None # TODO(b/384730642): Re-enable this test once the parent issue is fixed. # agent.set_up() @@ -250,21 +256,21 @@ def test_enable_tracing( @pytest.mark.usefixtures("caplog") def test_enable_tracing_warning(self, caplog, langchain_instrumentor_none_mock): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL, enable_tracing=True) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL, enable_tracing=True) assert agent._tmpl_attrs.get("instrumentor") is None # TODO(b/383923584): Re-enable this test once the parent issue is fixed. # agent.set_up() # assert "enable_tracing=True but proceeding with tracing disabled" in caplog.text def test_get_state_history_empty(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].get_state_history.return_value = [] history = list(agent.get_state_history()) assert history == [] def test_get_state_history(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].get_state_history.return_value = [ mock.Mock(), @@ -283,7 +289,7 @@ def test_get_state_history(self): ] def test_get_state_history_with_config(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].get_state_history.return_value = [ mock.Mock(), @@ -302,7 +308,7 @@ def test_get_state_history_with_config(self): ] def test_get_state(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].get_state.return_value = mock.Mock() agent._tmpl_attrs["runnable"].get_state.return_value._asdict.return_value = { @@ -312,7 +318,7 @@ def test_get_state(self): assert state == {"test_key": "test_value"} def test_get_state_with_config(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent._tmpl_attrs["runnable"].get_state.return_value = mock.Mock() agent._tmpl_attrs["runnable"].get_state.return_value._asdict.return_value = { @@ -322,13 +328,13 @@ def test_get_state_with_config(self): assert state == {"test_key": "test_value"} def test_update_state(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent.update_state() agent._tmpl_attrs["runnable"].update_state.assert_called_once() def test_update_state_with_config(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent.update_state(config=_TEST_CONFIG) agent._tmpl_attrs["runnable"].update_state.assert_called_once_with( @@ -336,7 +342,7 @@ def test_update_state_with_config(self): ) def test_update_state_with_config_and_kwargs(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) agent._tmpl_attrs["runnable"] = mock.Mock() agent.update_state(config=_TEST_CONFIG, test_key="test_value") agent._tmpl_attrs["runnable"].update_state.assert_called_once_with( @@ -344,7 +350,7 @@ def test_update_state_with_config_and_kwargs(self): ) def test_register_operations(self): - agent = agent_engines.LanggraphAgent(model=_TEST_MODEL) + agent = frameworks.LanggraphAgent(model=_TEST_MODEL) expected_operations = { "": ["query", "get_state", "update_state"], "stream": ["stream_query", "get_state_history"], @@ -360,6 +366,6 @@ def _return_input_no_typing(input_): class TestConvertToolsOrRaiseErrors: def test_raise_untyped_input_args(self, agentplatform_init_mock): with pytest.raises(TypeError, match=r"has untyped input_arg"): - agent_engines.LanggraphAgent( + frameworks.LanggraphAgent( model=_TEST_MODEL, tools=[_return_input_no_typing] ) diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py b/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py index aaef28749b..77f09a5cef 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_llama_index.py @@ -14,15 +14,15 @@ # import importlib import json +import os from unittest import mock from google import auth import agentplatform -from google.cloud.aiplatform import initializer -from agentplatform.agent_engines.templates import ( +from agentplatform.frameworks import ( llama_index, ) -from agentplatform._genai.agent_engines import _agent_engines_utils +from agentplatform._genai import _runtimes_utils from llama_index.core import prompts from llama_index.core.base.llms import types @@ -71,7 +71,7 @@ def model_builder_mock(): @pytest.fixture def cloud_trace_exporter_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_cloud_trace_exporter_or_warn", ) as cloud_trace_exporter_mock: yield cloud_trace_exporter_mock @@ -94,7 +94,7 @@ def simple_span_processor_mock(): @pytest.fixture def llama_index_instrumentor_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_llama_index_or_warn", ) as llama_index_instrumentor_mock: yield llama_index_instrumentor_mock @@ -103,7 +103,7 @@ def llama_index_instrumentor_mock(): @pytest.fixture def llama_index_instrumentor_none_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_openinference_llama_index_or_warn", ) as llama_index_instrumentor_mock: llama_index_instrumentor_mock.return_value = None @@ -113,7 +113,7 @@ def llama_index_instrumentor_none_mock(): @pytest.fixture def nest_asyncio_apply_mock(): with mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_import_nest_asyncio_or_warn", ) as nest_asyncio_apply_mock: yield nest_asyncio_apply_mock @@ -122,12 +122,11 @@ 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, - ) + # The agent frameworks source project and location from environment + # variables (not the global initializer). + os.environ["GOOGLE_CLOUD_PROJECT"] = _TEST_PROJECT + os.environ["GOOGLE_CLOUD_LOCATION"] = _TEST_LOCATION self.prompt = prompts.ChatPromptTemplate( message_templates=[ types.ChatMessage( @@ -142,13 +141,18 @@ 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 + # Project and location are no longer stored on the agent; they are + # sourced from environment variables during set_up(). assert agent._runnable is None def test_set_up(self): @@ -264,17 +268,17 @@ class TestToJsonSerializableLlamaIndexObject: """Tests for `_utils.to_json_serializable_llama_index_object`.""" def test_llama_index_response(self): - mock_response: _agent_engines_utils.LlamaIndexResponse = mock.Mock( - spec=_agent_engines_utils.LlamaIndexResponse + mock_response: _runtimes_utils.LlamaIndexResponse = mock.Mock( + spec=_runtimes_utils.LlamaIndexResponse ) mock_response.response = "test response" mock_response.source_nodes = [ mock.Mock( - spec=_agent_engines_utils.LlamaIndexBaseModel, + spec=_runtimes_utils.LlamaIndexBaseModel, model_dump_json=lambda: '{"name": "model1"}', ), mock.Mock( - spec=_agent_engines_utils.LlamaIndexBaseModel, + spec=_runtimes_utils.LlamaIndexBaseModel, model_dump_json=lambda: '{"name": "model2"}', ), ] @@ -285,61 +289,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..32f059ff2b 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 @@ -164,7 +164,7 @@ def mock_agent_engine_create_path_exists(): def mock_agent_engine_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 @@ -174,7 +174,7 @@ def mock_agent_engine_create_base64_encoded_tarball(): def mock_agent_engine_create_docker_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_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 diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py index e4b33ad2d5..6f7227ec62 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_delete.py @@ -19,14 +19,14 @@ def test_delete_memory(client): - ae_memory_operation = client.agent_engines.memories.delete( + ae_memory_operation = client.runtimes.memories.delete( name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/5605466683931099136", ) - assert isinstance(ae_memory_operation, types.DeleteAgentEngineMemoryOperation) + assert isinstance(ae_memory_operation, types.DeleteRuntimeMemoryOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.delete", + test_method="runtimes.memories.delete", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_get.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_get.py index bca6fca8c9..707479f147 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_get.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_get.py @@ -20,7 +20,7 @@ def test_get_memory(client): memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688" - ae_memory = client.agent_engines.memories.get(name=memory_name) + ae_memory = client.runtimes.memories.get(name=memory_name) assert isinstance(ae_memory, types.Memory) assert ae_memory.name == memory_name @@ -28,5 +28,5 @@ def test_get_memory(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.get", + test_method="runtimes.memories.get", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py index 5ca959abe2..9a8ecfb8c3 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_create.py @@ -19,16 +19,16 @@ def test_private_create_memory(client): - ae_memory_operation = client.agent_engines.memories._create( + ae_memory_operation = client.runtimes.memories._create( name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584", fact="memory_fact", scope={"user_id": "123"}, ) - assert isinstance(ae_memory_operation, types.AgentEngineMemoryOperation) + assert isinstance(ae_memory_operation, types.RuntimeMemoryOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._create", + test_method="runtimes.memories._create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_generate.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_generate.py index 7b294a01cf..f9c417ec1d 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_generate.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_generate.py @@ -19,17 +19,17 @@ def test_private_generate_memory(client): - ae_memory_operation = client.agent_engines.memories._generate( + ae_memory_operation = client.runtimes.memories._generate( name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584", vertex_session_source=types.GenerateMemoriesRequestVertexSessionSource( session="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/sessions/6922431337672474624" ), ) - assert isinstance(ae_memory_operation, types.AgentEngineGenerateMemoriesOperation) + assert isinstance(ae_memory_operation, types.RuntimeGenerateMemoriesOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._generate", + test_method="runtimes.memories._generate", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_generate_memories_operation.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_generate_memories_operation.py index bdcd2f14a4..d9feda4c64 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_generate_memories_operation.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_generate_memories_operation.py @@ -19,14 +19,14 @@ def test_private_get_generate_memories_operation(client): - memory_operation = client.agent_engines.memories._get_generate_memories_operation( + memory_operation = client.runtimes.memories._get_generate_memories_operation( operation_name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/operations/5669315676343369728" ) - assert isinstance(memory_operation, types.AgentEngineGenerateMemoriesOperation) + assert isinstance(memory_operation, types.RuntimeGenerateMemoriesOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._get_generate_memories_operation", + test_method="runtimes.memories._get_generate_memories_operation", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py index 23c92f7821..0c9c74a684 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_get_memory_operation.py @@ -19,14 +19,14 @@ def test_private_get_memory_operation(client): - memory_operation = client.agent_engines.memories._get_memory_operation( + memory_operation = client.runtimes.memories._get_memory_operation( operation_name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688/operations/1044963283964002304" ) - assert isinstance(memory_operation, types.AgentEngineMemoryOperation) + assert isinstance(memory_operation, types.RuntimeMemoryOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._get_memory_operation", + test_method="runtimes.memories._get_memory_operation", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py index 662ad020fb..24f608dc6c 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_list.py @@ -20,7 +20,7 @@ def test_private_list_memory(client): ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584" - memory_list = client.agent_engines.memories._list(name=ae_name) + memory_list = client.runtimes.memories._list(name=ae_name) assert isinstance(memory_list, types.ListReasoningEnginesMemoriesResponse) assert isinstance(memory_list.memories[0], types.Memory) @@ -28,5 +28,5 @@ def test_private_list_memory(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._list", + test_method="runtimes.memories._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py index d77cb3a795..025bcfb1a9 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_purge.py @@ -19,17 +19,17 @@ def test_private_purge(client): - ae_memory_purge_operation = client.agent_engines.memories._purge( + ae_memory_purge_operation = client.runtimes.memories._purge( name="projects/964831358985/locations/us-central1/reasoningEngines/6086402690647064576", filter="scope.user_id=123", ) assert isinstance( - ae_memory_purge_operation, types.AgentEnginePurgeMemoriesOperation + ae_memory_purge_operation, types.RuntimePurgeMemoriesOperation ) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._purge", + test_method="runtimes.memories._purge", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py index 4098b11d9b..754a3430ee 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_retrieve.py @@ -20,7 +20,7 @@ def test_private_retrieve(client): ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584" - retrieved_memories = client.agent_engines.memories._retrieve( + retrieved_memories = client.runtimes.memories._retrieve( name=ae_name, scope={"user_id": "123"}, ) @@ -34,5 +34,5 @@ def test_private_retrieve(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._retrieve", + test_method="runtimes.memories._retrieve", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_rollback.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_rollback.py index 85329da47e..898ec79aee 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_rollback.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_rollback.py @@ -19,15 +19,15 @@ def test_private_rollback(client): - rollback_operation = client.agent_engines.memories._rollback( + rollback_operation = client.runtimes.memories._rollback( name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688", target_revision_id="3001207491565453312", ) - assert isinstance(rollback_operation, types.AgentEngineRollbackMemoryOperation) + assert isinstance(rollback_operation, types.RuntimeRollbackMemoryOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._rollback", + test_method="runtimes.memories._rollback", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py index c6dbc1f095..40ef680b5b 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memories_private_update.py @@ -20,16 +20,16 @@ def test_private_update_memory(client): memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688" - memory_update_operation = client.agent_engines.memories._update( + memory_update_operation = client.runtimes.memories._update( name=memory_name, fact="memory_fact_updated", scope={"user_id": "123"}, ) - assert isinstance(memory_update_operation, types.AgentEngineMemoryOperation) + assert isinstance(memory_update_operation, types.RuntimeMemoryOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories._update", + test_method="runtimes.memories._update", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py b/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py index f58e365fe7..54ce8d2bb8 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_get.py @@ -20,7 +20,7 @@ def test_get_memory_revisions(client): memory_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688/revisions/516064922187071488" - memory_revision = client.agent_engines.memories.revisions.get(name=memory_name) + memory_revision = client.runtimes.memories.revisions.get(name=memory_name) assert isinstance(memory_revision, types.MemoryRevision) assert memory_revision.name == memory_name @@ -28,5 +28,5 @@ def test_get_memory_revisions(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.revisions.get", + test_method="runtimes.memories.revisions.get", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py b/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py index c827a57a06..76eb966f38 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_memory_revisions_private_list.py @@ -20,13 +20,13 @@ def test_private_list_memory_revisions(client): ae_name = "projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584/memories/3858070028511346688" - memory_list = client.agent_engines.memories.revisions._list(name=ae_name) - assert isinstance(memory_list, types.ListAgentEngineMemoryRevisionsResponse) + memory_list = client.runtimes.memories.revisions._list(name=ae_name) + assert isinstance(memory_list, types.ListRuntimeMemoryRevisionsResponse) assert isinstance(memory_list.memory_revisions[0], types.MemoryRevision) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.revisions._list", + test_method="runtimes.memories.revisions._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_create.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_create.py index fc776d9874..aa46a18a0c 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_create.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_create.py @@ -21,7 +21,7 @@ def test_create_sandbox_snapshot(client): - snapshot = client.agent_engines.sandboxes.snapshots._create( + snapshot = client.sandboxes.snapshots._create( source_sandbox_environment_name="projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironments/525190525100228608", config={ "display_name": "test_snapshot", @@ -30,11 +30,11 @@ def test_create_sandbox_snapshot(client): }, ) - assert isinstance(snapshot, types.AgentEngineSandboxSnapshotOperation) + assert isinstance(snapshot, types.RuntimeSandboxSnapshotOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.snapshots._create", + test_method="sandboxes.snapshots._create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_delete.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_delete.py index 730f9b262a..e6ac2d5d39 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_delete.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_delete.py @@ -21,7 +21,7 @@ def test_delete_sandbox_snapshot(client): - result = client.agent_engines.sandboxes.snapshots._delete( + result = client.sandboxes.snapshots._delete( name="projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironmentSnapshots/421086565159141376", ) @@ -31,5 +31,5 @@ def test_delete_sandbox_snapshot(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.snapshots._delete", + test_method="sandboxes.snapshots._delete", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_get.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_get.py index b9d5075c8b..7f6c35c68c 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_get.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_get.py @@ -22,7 +22,7 @@ def test_get_sandbox_snapshot(client): snapshot_name = "projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472/sandboxEnvironmentSnapshots/2433069698686910464" - fetched_snapshot = client.agent_engines.sandboxes.snapshots._get( + fetched_snapshot = client.sandboxes.snapshots._get( name=snapshot_name, ) @@ -33,5 +33,5 @@ def test_get_sandbox_snapshot(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.snapshots._get", + test_method="sandboxes.snapshots._get", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_list.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_list.py index a22bb1c244..1f2dfc84c3 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_snapshots_list.py @@ -21,7 +21,7 @@ def test_list_sandbox_snapshots(client): - snapshots_list_operation = client.agent_engines.sandboxes.snapshots._list( + snapshots_list_operation = client.sandboxes.snapshots._list( name="projects/802583348448/locations/us-central1/reasoningEngines/6130241318758121472", ) @@ -37,5 +37,5 @@ def test_list_sandbox_snapshots(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.snapshots._list", + test_method="sandboxes.snapshots._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_byoc_create.py b/tests/unit/agentplatform/genai/replays/test_ae_sandbox_templates_byoc_create.py index 47d686b972..bac8650ffd 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" ), @@ -63,5 +63,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 f1356297ad..80e3259d86 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" ), @@ -42,5 +42,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..b937456122 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( + agent_engine_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(agent_engine_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..7760aef6ba 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( + agent_engine_sandbox_delete_operation = client.sandboxes._delete( name=( "reasoningEngines/2886612747586371584/sandboxEnvironments/6068475153556176896" ), ) assert isinstance( - agent_engine_sandbox_delete_operation, types.DeleteAgentEngineSandboxOperation + agent_engine_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..7fd4ae6f9d 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,7 +22,7 @@ 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) + agent_engine_sandbox = client.sandboxes._get(name=sandbox_name) assert isinstance(agent_engine_sandbox, types.SandboxEnvironment) assert agent_engine_sandbox.name == sandbox_name @@ -30,5 +30,5 @@ def test_private_get(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_ae_sandboxes_private_get_sandbox_operation.py b/tests/unit/agentplatform/genai/replays/test_ae_sandboxes_private_get_sandbox_operation.py index 1de905140d..fce97745c6 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( + agent_engine_sandbox = client.sandboxes._get_sandbox_operation( operation_name=operation_name ) - assert isinstance(agent_engine_sandbox, types.AgentEngineSandboxOperation) + assert isinstance(agent_engine_sandbox, types.RuntimeSandboxOperation) assert agent_engine_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..fd8a5988f8 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,7 +19,7 @@ def test_private_list(client): - agent_engine_sandbox_list_operation = client.agent_engines.sandboxes._list( + agent_engine_sandbox_list_operation = client.sandboxes._list( name=("reasoningEngines/2886612747586371584"), ) assert isinstance( @@ -27,12 +27,12 @@ def test_private_list(client): types.SandboxEnvironment, ) assert isinstance( - agent_engine_sandbox_list_operation, types.ListAgentEngineSandboxesResponse + agent_engine_sandbox_list_operation, types.ListRuntimeSandboxesResponse ) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes._list", + test_method="sandboxes._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_delete.py b/tests/unit/agentplatform/genai/replays/test_ae_session_delete.py index 73358920ef..7b44a6e775 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_delete.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_delete.py @@ -19,14 +19,14 @@ def test_delete_session_non_blocking(client): - ae_session_operation = client.agent_engines.sessions.delete( + ae_session_operation = client.sessions.delete( name=("reasoningEngines/2886612747586371584/sessions/8521561049109889024"), ) - assert isinstance(ae_session_operation, types.DeleteAgentEngineSessionOperation) + assert isinstance(ae_session_operation, types.DeleteRuntimeSessionOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions.delete", + test_method="sessions.delete", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_events_append.py b/tests/unit/agentplatform/genai/replays/test_ae_session_events_append.py index aef176ea86..62d15a7094 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_events_append.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_events_append.py @@ -21,17 +21,17 @@ def test_append_session_event(client): - session_event = client.agent_engines.sessions.events.append( + session_event = client.sessions.events.append( name="reasoningEngines/2886612747586371584/sessions/6922431337672474624", author="test-user-123", invocation_id="test-invocation-id", timestamp=datetime.datetime.fromtimestamp(1234567860, tz=datetime.timezone.utc), ) - assert isinstance(session_event, types.AppendAgentEngineSessionEventResponse) + assert isinstance(session_event, types.AppendRuntimeSessionEventResponse) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions.events.append", + test_method="sessions.events.append", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_events_private_list.py b/tests/unit/agentplatform/genai/replays/test_ae_session_events_private_list.py index eb3153ab17..f681f3c335 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_events_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_events_private_list.py @@ -19,11 +19,11 @@ def test_private_list_session_events(client): - session_event_list_response = client.agent_engines.sessions.events._list( + session_event_list_response = client.sessions.events._list( name="reasoningEngines/2886612747586371584/sessions/6922431337672474624", ) assert isinstance( - session_event_list_response, types.ListAgentEngineSessionEventsResponse + session_event_list_response, types.ListRuntimeSessionEventsResponse ) assert len(session_event_list_response.session_events) == 1 assert isinstance(session_event_list_response.session_events[0], types.SessionEvent) @@ -32,5 +32,5 @@ def test_private_list_session_events(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions.events._list", + test_method="sessions.events._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_private_create.py b/tests/unit/agentplatform/genai/replays/test_ae_session_private_create.py index 28dd62cac2..f693c6e164 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_private_create.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_private_create.py @@ -19,15 +19,15 @@ def test_private_create_session(client): - ae_session_operation = client.agent_engines.sessions._create( + ae_session_operation = client.sessions._create( name="projects/964831358985/locations/us-central1/reasoningEngines/2886612747586371584", user_id="test-user-id", ) - assert isinstance(ae_session_operation, types.AgentEngineSessionOperation) + assert isinstance(ae_session_operation, types.RuntimeSessionOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions._create", + test_method="sessions._create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_private_get.py b/tests/unit/agentplatform/genai/replays/test_ae_session_private_get.py index 35bfdd7b63..c1100873b6 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_private_get.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_private_get.py @@ -19,14 +19,14 @@ def test_private_get_session_operation(client): - ae_session_operation = client.agent_engines.sessions._get_session_operation( + ae_session_operation = client.sessions._get_session_operation( operation_name="reasoningEngines/2886612747586371584/sessions/3080649749292908544/operations/758783840595476480", ) - assert isinstance(ae_session_operation, types.AgentEngineSessionOperation) + assert isinstance(ae_session_operation, types.RuntimeSessionOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions._get_session_operation", + test_method="sessions._get_session_operation", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_private_list.py b/tests/unit/agentplatform/genai/replays/test_ae_session_private_list.py index f8a2fe338e..b01cf859ef 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_private_list.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_private_list.py @@ -19,7 +19,7 @@ def test_private_list_session(client): - session_list_response = client.agent_engines.sessions._list( + session_list_response = client.sessions._list( name="reasoningEngines/2886612747586371584", ) assert isinstance(session_list_response, types.ListReasoningEnginesSessionsResponse) @@ -29,5 +29,5 @@ def test_private_list_session(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions._list", + test_method="sessions._list", ) diff --git a/tests/unit/agentplatform/genai/replays/test_ae_session_private_update.py b/tests/unit/agentplatform/genai/replays/test_ae_session_private_update.py index 3f62dee2f0..90bf70aec9 100644 --- a/tests/unit/agentplatform/genai/replays/test_ae_session_private_update.py +++ b/tests/unit/agentplatform/genai/replays/test_ae_session_private_update.py @@ -19,18 +19,18 @@ def test_private_update_session(client): - agent_engine_session_operation = client.agent_engines.sessions._update( + agent_engine_session_operation = client.sessions._update( name="reasoningEngines/2886612747586371584/sessions/3080649749292908544", - config=types.UpdateAgentEngineSessionConfig( + config=types.UpdateRuntimeSessionConfig( display_name="test-agent-engine-session-updated", user_id="test-user-id", ), ) - assert isinstance(agent_engine_session_operation, types.AgentEngineSessionOperation) + assert isinstance(agent_engine_session_operation, types.RuntimeSessionOperation) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sessions._update", + test_method="sessions._update", ) diff --git a/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_methods.py b/tests/unit/agentplatform/genai/replays/test_agent_engine_a2a_methods.py index f9665236c0..1934763308 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( + agent_engine = client.runtimes.get( name="projects/932854658080/locations/us-central1/reasoningEngines/857830725653626880", ) - assert isinstance(agent_engine, types.AgentEngine) + assert isinstance(agent_engine, types.Runtime) message_data = { "messageId": "msg-123", @@ -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..6589853d12 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( + agent_engine = client.runtimes.get( name="projects/964831358985/locations/us-central1/reasoningEngines/6859679872613089280", ) - assert isinstance(agent_engine, types.AgentEngine) + assert isinstance(agent_engine, types.Runtime) with mock.patch( "httpx.AsyncClient", spec=httpx.AsyncClient @@ -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..775a731beb 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( + agent_engine_operation = client.runtimes._create( config={"labels": labels}, ) - assert isinstance(agent_engine_operation, types.AgentEngineOperation) + assert isinstance(agent_engine_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..d3678ac1d4 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( + agent_engine_operation = client.runtimes._delete( name="reasoningEngines/7571341522470174720", ) - assert isinstance(agent_engine_operation, types.DeleteAgentEngineOperation) + assert isinstance(agent_engine_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..45aa58b201 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,7 +19,7 @@ def test_private_get(client): - agent_engine = client.agent_engines._get( + agent_engine = client.runtimes._get( name="reasoningEngines/2886612747586371584", ) assert isinstance(agent_engine, types.ReasoningEngine) @@ -29,5 +29,5 @@ def test_private_get(client): 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..38113a1d78 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( + agent_engine_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(agent_engine_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..69eee5daed 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( + agent_engine = client.runtimes.create( config={"labels": labels}, ) assert agent_engine.api_resource.labels == labels # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.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( + agent_engine = client.runtimes.create( config={ "context_spec": { "memory_bank_config": { @@ -117,7 +117,7 @@ def test_create_with_context_spec(client): "http_options": {"api_version": "v1beta1"}, }, ) - agent_engine = client.agent_engines.get(name=agent_engine.api_resource.name) + agent_engine = client.runtimes.get(name=agent_engine.api_resource.name) memory_bank_config = agent_engine.api_resource.context_spec.memory_bank_config assert memory_bank_config.generation_config.model == generation_model assert ( @@ -132,7 +132,7 @@ 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=agent_engine.api_resource.name, force=True) def test_create_with_source_packages( @@ -171,7 +171,7 @@ def _update_ver(obj): mock_agent_engine_create_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - agent_engine = client.agent_engines.create( + agent_engine = client.runtimes.create( config={ "display_name": "test-agent-engine-source-packages", "source_packages": [ @@ -189,12 +189,12 @@ def _update_ver(obj): ) assert agent_engine.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=agent_engine.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( + agent_engine = client.runtimes.create( config={ "identity_type": types.IdentityType.AGENT_IDENTITY, "http_options": {"api_version": "v1beta1"}, @@ -208,11 +208,11 @@ def test_create_with_identity_type(client): agent_engine.api_resource.spec.effective_identity ) # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.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..5524be7d03 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 @@ -48,13 +48,13 @@ def test_create_a2a_agent(client, is_replay_mode): # that fails when cloudpickle.load expects bytes. We mock _upload_agent_engine # 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_agent_engine") if is_replay_mode else contextlib.nullcontext() ) with upload_patch: - agent_engine = client.agent_engines.create( + agent_engine = 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 isinstance(agent_engine, types.Runtime) assert agent_engine.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=agent_engine.api_resource.name, force=True) pytestmark = pytest_helper.setup( diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a_task.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a_task.py deleted file mode 100644 index 6d6355cdf5..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_a2a_task.py +++ /dev/null @@ -1,173 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -from google.genai import types as genai_types -import pytest - - -def test_create_simple_a2a_task(client): - # Use the autopush environment. - client._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client._api_client._http_options.api_version = "internal" - - task = client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig( - context_id="context123", - metadata={ - "key": "value", - "key2": [{"key3": "value3", "key4": "value4"}], - }, - status_details=types.TaskStatusDetails( - task_message=types.TaskMessage( - role="user", - message_id="message123", - parts=[ - genai_types.Part( - text="hello123", - ) - ], - metadata={ - "key42": "value42", - }, - ), - ), - output=types.TaskOutput( - artifacts=[ - types.TaskArtifact( - artifact_id="artifact123", - display_name="display_name123", - description="description123", - parts=[ - genai_types.Part( - text="hello456", - ) - ], - ) - ], - ), - ), - ) - - assert isinstance(task, types.A2aTask) - assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" - assert task.context_id == "context123" - assert task.state == types.State.SUBMITTED - assert task.status_details.task_message.role == "user" - assert task.status_details.task_message.message_id == "message123" - assert task.status_details.task_message.parts[0].text == "hello123" - assert task.status_details.task_message.metadata["key42"] == "value42" - assert task.output.artifacts[0].artifact_id == "artifact123" - assert task.output.artifacts[0].display_name == "display_name123" - assert task.output.artifacts[0].description == "description123" - assert task.output.artifacts[0].parts[0].text == "hello456" - assert task.metadata == { - "key": "value", - "key2": [{"key3": "value3", "key4": "value4"}], - } - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), -) - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_create_simple_a2a_task_async(client): - # Use the autopush environment. - client.aio._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client.aio._api_client._http_options.api_version = "internal" - - task = await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig( - context_id="context123", - metadata={ - "key": "value", - "key2": [{"key3": "value3", "key4": "value4"}], - }, - status_details=types.TaskStatusDetails( - task_message=types.TaskMessage( - role="user", - message_id="message123", - parts=[ - genai_types.Part( - text="hello123", - ) - ], - metadata={ - "key42": "value42", - }, - ), - ), - output=types.TaskOutput( - artifacts=[ - types.TaskArtifact( - artifact_id="artifact123", - display_name="display_name123", - description="description123", - parts=[ - genai_types.Part( - text="hello456", - ) - ], - ) - ], - ), - ), - ) - - assert isinstance(task, types.A2aTask) - assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" - assert task.context_id == "context123" - assert task.state == types.State.SUBMITTED - assert task.status_details.task_message.role == "user" - assert task.status_details.task_message.message_id == "message123" - assert task.status_details.task_message.parts[0].text == "hello123" - assert task.status_details.task_message.metadata["key42"] == "value42" - assert task.output.artifacts[0].artifact_id == "artifact123" - assert task.output.artifacts[0].display_name == "display_name123" - assert task.output.artifacts[0].description == "description123" - assert task.output.artifacts[0].parts[0].text == "hello456" - assert task.metadata == { - "key": "value", - "key2": [{"key3": "value3", "key4": "value4"}], - } - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_developer_connect.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_developer_connect.py index f26d0eb715..3177c334dd 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( + agent_engine = client.runtimes.create( config={ "display_name": "test-agent-engine-dev-connect", "developer_connect_source": developer_connect_source_config, @@ -84,11 +84,11 @@ def _update_ver(obj): == developer_connect_source_config.dir ) # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.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..944f7a6792 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 @@ -60,7 +60,7 @@ def _update_ver(obj): mock_agent_engine_create_docker_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - agent_engine = client.agent_engines.create( + agent_engine = client.runtimes.create( config={ "display_name": "test-agent-engine-docker", "description": "test agent engine with docker spec", @@ -76,11 +76,11 @@ def _update_ver(obj): ) assert agent_engine.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=agent_engine.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.create", + test_method="runtimes.create", ) diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_memory.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_memory.py index f321a921c0..fd0ef83be1 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_memory.py +++ b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_memory.py @@ -21,8 +21,8 @@ def test_create_memory_with_ttl(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) metadata = { @@ -36,17 +36,17 @@ def test_create_memory_with_ttl(client): ), } - operation = client.agent_engines.memories.create( + operation = client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, - config=types.AgentEngineMemoryConfig( + config=types.RuntimeMemoryConfig( display_name="my_memory_fact", ttl="120s", metadata=metadata, ), ) - assert isinstance(operation, types.AgentEngineMemoryOperation) + assert isinstance(operation, types.RuntimeMemoryOperation) assert operation.response.fact == "memory_fact" assert operation.response.scope == {"user_id": "123"} assert operation.response.name.startswith(agent_engine.api_resource.name) @@ -59,48 +59,48 @@ def test_create_memory_with_ttl(client): ) assert operation.response.metadata == metadata # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) def test_create_memory_with_expire_time(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) expire_time = datetime.datetime( 2027, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc ) - operation = client.agent_engines.memories.create( + operation = client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, - config=types.AgentEngineMemoryConfig( + config=types.RuntimeMemoryConfig( display_name="my_memory_fact", expire_time=expire_time ), ) - assert isinstance(operation, types.AgentEngineMemoryOperation) + assert isinstance(operation, types.RuntimeMemoryOperation) assert operation.response.fact == "memory_fact" assert operation.response.scope == {"user_id": "123"} assert operation.response.name.startswith(agent_engine.api_resource.name) assert operation.response.expire_time == expire_time # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) def test_create_memory_with_custom_memory_id(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - operation = client.agent_engines.memories.create( + operation = client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, - config=types.AgentEngineMemoryConfig( + config=types.RuntimeMemoryConfig( display_name="my_memory_fact", memory_id="my-memory-id" ), ) - assert isinstance(operation, types.AgentEngineMemoryOperation) + assert isinstance(operation, types.RuntimeMemoryOperation) assert operation.response.fact == "memory_fact" assert operation.response.scope == {"user_id": "123"} assert ( @@ -108,11 +108,11 @@ def test_create_memory_with_custom_memory_id(client): == f"{agent_engine.api_resource.name}/memories/my-memory-id" ) # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.create_memory", + test_method="runtimes.create_memory", ) diff --git a/tests/unit/agentplatform/genai/replays/test_create_agent_engine_sandbox.py b/tests/unit/agentplatform/genai/replays/test_create_agent_engine_sandbox.py index c8257565d7..f272ea71dd 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,11 +19,11 @@ def test_create_sandbox(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - operation = client.agent_engines.sandboxes.create( + operation = client.sandboxes.create( name=agent_engine.api_resource.name, poll_interval_seconds=1, spec={ @@ -31,11 +31,11 @@ def test_create_sandbox(client): "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 @@ -47,5 +47,5 @@ def test_create_sandbox(client): 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..7c0f5b6d3b 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,22 +21,22 @@ def test_create_session_with_ttl(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() try: - assert isinstance(agent_engine, types.AgentEngine) + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - operation = client.agent_engines.create_session( + operation = client.sessions.create( name=agent_engine.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" @@ -52,28 +52,28 @@ 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=agent_engine.api_resource.name, force=True) def test_create_session_with_expire_time(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() try: - assert isinstance(agent_engine, types.AgentEngine) + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) expire_time = datetime.datetime( 2028, 1, 1, 12, 30, 00, tzinfo=datetime.timezone.utc ) - operation = client.agent_engines.sessions.create( + operation = client.sessions.create( name=agent_engine.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" @@ -82,25 +82,25 @@ def test_create_session_with_expire_time(client): assert operation.done finally: # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) def test_create_session_with_custom_session_id(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() try: - assert isinstance(agent_engine, types.AgentEngine) + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - operation = client.agent_engines.sessions.create( + operation = client.sessions.create( name=agent_engine.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" @@ -111,11 +111,11 @@ def test_create_session_with_custom_session_id(client): assert operation.done finally: # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.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..c66fa01ae6 100644 --- a/tests/unit/agentplatform/genai/replays/test_create_feedback_entry.py +++ b/tests/unit/agentplatform/genai/replays/test_create_feedback_entry.py @@ -32,8 +32,8 @@ def test_create(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) try: @@ -68,7 +68,7 @@ def test_create(client): } finally: # Clean up resources. - client.agent_engines.delete( + client.runtimes.delete( name=agent_engine.api_resource.name, force=True, ) @@ -76,8 +76,8 @@ def test_create(client): @pytest.mark.asyncio async def test_create_async(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) try: @@ -112,7 +112,7 @@ async def test_create_async(client): } finally: # Clean up resources. - await client.aio.agent_engines.delete( + await client.aio.runtimes.delete( name=agent_engine.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..37f566ecf2 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 @@ -39,7 +39,7 @@ def test_delete_runtime_revision( mock_agent_engine_create_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - agent_engine = client.agent_engines.create( + agent_engine = client.runtimes.create( config={ "display_name": "test-agent-engine-delete-runtime-revision", "source_packages": [ @@ -61,7 +61,7 @@ def test_delete_runtime_revision( mock_agent_engine_create_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - updated_agent_engine = client.agent_engines.update( + updated_agent_engine = client.runtimes.update( name=agent_engine.api_resource.name, config={ "display_name": "test-agent-engine-update-traffic-with-agent-after-update", @@ -79,18 +79,18 @@ def test_delete_runtime_revision( }, ) - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.runtimes.revisions.list( name=updated_agent_engine.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( + runtime_revisions_iter = client.runtimes.revisions.list( name=updated_agent_engine.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) @@ -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_agent_engine.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", ) @@ -127,7 +127,7 @@ async def test_delete_runtime_revision_async( mock_agent_engine_create_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - agent_engine = client.agent_engines.create( + agent_engine = client.runtimes.create( config={ "display_name": "test-agent-engine-delete-runtime-revision", "source_packages": [ @@ -149,7 +149,7 @@ async def test_delete_runtime_revision_async( mock_agent_engine_create_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - updated_agent_engine = client.agent_engines.update( + updated_agent_engine = client.runtimes.update( name=agent_engine.api_resource.name, config={ "display_name": "test-agent-engine-update-traffic-with-agent-after-update", @@ -167,7 +167,7 @@ async def test_delete_runtime_revision_async( }, ) - runtime_revisions_iter = client.aio.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.aio.runtimes.revisions.list( name=updated_agent_engine.api_resource.name, ) runtime_revisions_list = [] @@ -175,10 +175,10 @@ async def test_delete_runtime_revision_async( 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( + assert isinstance(operation, types.DeleteRuntimeRevisionOperation) + await client.aio.runtimes.delete( name=updated_agent_engine.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..cef6e2fa81 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 + agent_engine = client.runtimes.create() + operation = client.runtimes.delete(name=agent_engine.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( + agent_engine = client.runtimes.create() + operation = await client.aio.runtimes.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 + assert isinstance(operation, types.DeleteRuntimeOperation) + assert "Deleting Runtime resource" in caplog.text + assert f"Started Runtime delete operation: {operation.name}" in caplog.text diff --git a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_a2a_task.py b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_a2a_task.py deleted file mode 100644 index b2859dea24..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_a2a_task.py +++ /dev/null @@ -1,96 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -from google.genai import errors -import pytest - - -def test_delete_a2a_task(client): - # Use the autopush environment. - client._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client._api_client._http_options.api_version = "internal" - - created_task = client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - assert isinstance(created_task, types.A2aTask) - - # Test validity check. - task = client.agent_engines.a2a_tasks.get( - name=created_task.name, - ) - assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" - - client.agent_engines.a2a_tasks.delete(name=created_task.name) - - with pytest.raises(errors.ClientError, match="404 NOT_FOUND"): - client.agent_engines.a2a_tasks.get(name=created_task.name) - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), -) - - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_delete_a2a_task_async(client): - # Use the autopush environment. - client.aio._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client.aio._api_client._http_options.api_version = "internal" - - created_task = await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - assert isinstance(created_task, types.A2aTask) - - # Test validity check. - task = await client.aio.agent_engines.a2a_tasks.get( - name=created_task.name, - ) - assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" - - await client.aio.agent_engines.a2a_tasks.delete(name=created_task.name) - - with pytest.raises(errors.ClientError, match="404 NOT_FOUND"): - await client.aio.agent_engines.a2a_tasks.get(name=created_task.name) - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_memory.py b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_memory.py index 9d3dea4c72..61a35e3378 100644 --- a/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_memory.py +++ b/tests/unit/agentplatform/genai/replays/test_delete_agent_engine_memory.py @@ -22,23 +22,23 @@ def test_delete_memory(client): - agent_engine = client.agent_engines.create() - operation = client.agent_engines.memories.create( + agent_engine = client.runtimes.create() + operation = client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, ) memory = operation.response - operation = client.agent_engines.memories.delete(name=memory.name) - assert isinstance(operation, types.DeleteAgentEngineMemoryOperation) + operation = client.runtimes.memories.delete(name=memory.name) + assert isinstance(operation, types.DeleteRuntimeMemoryOperation) assert operation.name.startswith(memory.name + "/operations/") - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.delete_memory", + test_method="runtimes.delete_memory", ) @@ -47,16 +47,16 @@ def test_delete_memory(client): @pytest.mark.asyncio async def test_delete_memory_async(client): - agent_engine = client.agent_engines.create() - operation = await client.aio.agent_engines.memories.create( + agent_engine = client.runtimes.create() + operation = await client.aio.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, ) memory = operation.response - operation = await client.aio.agent_engines.memories.delete(name=memory.name) - assert isinstance(operation, types.DeleteAgentEngineMemoryOperation) + operation = await client.aio.runtimes.memories.delete(name=memory.name) + assert isinstance(operation, types.DeleteRuntimeMemoryOperation) assert operation.name.startswith(memory.name + "/operations/") - await client.aio.agent_engines.delete( + await client.aio.runtimes.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..0edfbe8338 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) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - operation = client.agent_engines.sandboxes.create( + operation = client.sandboxes.create( name=agent_engine.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..0b20d3552f 100644 --- a/tests/unit/agentplatform/genai/replays/test_delete_feedback_entry.py +++ b/tests/unit/agentplatform/genai/replays/test_delete_feedback_entry.py @@ -33,8 +33,8 @@ def test_delete(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) try: @@ -59,7 +59,7 @@ def test_delete(client): client.feedback_entries.get(name=create_operation.response.name) finally: # Clean up resources. - client.agent_engines.delete( + client.runtimes.delete( name=agent_engine.api_resource.name, force=True, ) @@ -67,8 +67,8 @@ def test_delete(client): @pytest.mark.asyncio async def test_delete_async(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) try: @@ -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( + await client.aio.runtimes.delete( name=agent_engine.api_resource.name, force=True, ) 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..32b6a4d736 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) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - operation = client.agent_engines.sandboxes.create( + operation = client.sandboxes.create( name=agent_engine.api_resource.name, spec={ "code_execution_environment": { "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB" } }, - config=types.CreateAgentEngineSandboxConfig(display_name="test_sandbox"), + config=types.CreateRuntimeSandboxConfig(display_name="test_sandbox"), ) - assert isinstance(operation, types.AgentEngineSandboxOperation) + assert isinstance(operation, types.RuntimeSandboxOperation) code = """ with open("test.txt", "r") as input: @@ -50,7 +50,7 @@ def test_execute_code_sandbox(client): } ], } - response = client.agent_engines.sandboxes.execute_code( + response = client.sandboxes.execute_code( name=operation.response.name, input_data=input_data, ) @@ -62,5 +62,5 @@ def test_execute_code_sandbox(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.sandboxes.execute_code", + test_method="sandboxes.execute_code", ) diff --git a/tests/unit/agentplatform/genai/replays/test_generate_agent_engine_memories.py b/tests/unit/agentplatform/genai/replays/test_generate_agent_engine_memories.py index 4942e3ac9e..3efae031e4 100644 --- a/tests/unit/agentplatform/genai/replays/test_generate_agent_engine_memories.py +++ b/tests/unit/agentplatform/genai/replays/test_generate_agent_engine_memories.py @@ -24,15 +24,15 @@ def test_generate_and_rollback_memories(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() assert not list( - client.agent_engines.memories.list( + client.runtimes.memories.list( name=agent_engine.api_resource.name, ) ) # Generate memories using source content. This result is non-deterministic, # because an LLM is used to generate the memories. - client.agent_engines.memories.generate( + client.runtimes.memories.generate( name=agent_engine.api_resource.name, scope={"user_id": "test-user-id"}, direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource( @@ -49,12 +49,12 @@ def test_generate_and_rollback_memories(client): ) ] ), - config=types.GenerateAgentEngineMemoriesConfig( + config=types.GenerateRuntimeMemoriesConfig( revision_labels={"key": "value"} ), ) memories = list( - client.agent_engines.memories.list( + client.runtimes.memories.list( name=agent_engine.api_resource.name, ) ) @@ -62,7 +62,7 @@ def test_generate_and_rollback_memories(client): # Every action that modifies a memory creates a new revision. memory_revisions = list( - client.agent_engines.memories.revisions.list( + client.runtimes.memories.revisions.list( name=memories[0].name, ) ) @@ -72,28 +72,28 @@ def test_generate_and_rollback_memories(client): revision_name = memory_revisions[0].name # Update the memory. - client.agent_engines.memories._update( + client.runtimes.memories._update( name=memories[0].name, fact="This is temporary", scope={"user_id": "test-user-id"}, ) - memory = client.agent_engines.memories.get(name=memories[0].name) + memory = client.runtimes.memories.get(name=memories[0].name) assert memory.fact == "This is temporary" # Rollback to the revision with the original fact that was created by the # generation request. - client.agent_engines.memories.rollback( + client.runtimes.memories.rollback( name=memories[0].name, target_revision_id=revision_name.split("/")[-1], ) - memory = client.agent_engines.memories.get(name=memories[0].name) + memory = client.runtimes.memories.get(name=memories[0].name) assert memory.fact == memory_revisions[0].fact # Update the memory again using generation. We use the original source # content to ensure that the original memory is updated. The response should # refer to the previous revision. pre_extracted_fact = "I am a software engineer focusing in security" - response = client.agent_engines.memories.generate( + response = client.runtimes.memories.generate( name=agent_engine.api_resource.name, scope={"user_id": "test-user-id"}, direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource( @@ -107,18 +107,18 @@ def test_generate_and_rollback_memories(client): # The memory was updated, so the previous revision is set. assert response.response.generated_memories[0].previous_revision is not None memory_revisions = list( - client.agent_engines.memories.revisions.list(name=memories[0].name) + client.runtimes.memories.revisions.list(name=memories[0].name) ) # Memory Revisions are returned in descending order by revision create time. # We can't make an assertion on the actual value, since it's # generated and thus non-deterministic. assert memory_revisions[0].extracted_memories[0].fact is not None - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) def test_generate_memories_direct_memories_source(client): - agent_engine = client.agent_engines.create() - client.agent_engines.memories.generate( + agent_engine = client.runtimes.create() + client.runtimes.memories.generate( name=agent_engine.api_resource.name, scope={"user_id": "test-user-id"}, direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource( @@ -131,23 +131,23 @@ def test_generate_memories_direct_memories_source(client): ), ] ), - config=types.GenerateAgentEngineMemoriesConfig(wait_for_completion=True), + config=types.GenerateRuntimeMemoriesConfig(wait_for_completion=True), ) assert ( len( list( - client.agent_engines.memories.list( + client.runtimes.memories.list( name=agent_engine.api_resource.name, ) ) ) >= 1 ) - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) def test_generate_memories_with_metadata(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() metadata = { "my_string_key": types.MemoryMetadataValue(string_value="my_string_value"), "my_double_key": types.MemoryMetadataValue(double_value=123.456), @@ -169,14 +169,14 @@ def test_generate_memories_with_metadata(client): ) scope = {"user_id": "test-user-id"} - operation = client.agent_engines.memories.generate( + operation = client.runtimes.memories.generate( name=agent_engine.api_resource.name, scope=scope, direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig(metadata=metadata), + config=types.GenerateRuntimeMemoriesConfig(metadata=metadata), ) assert len(operation.response.generated_memories) >= 1 - memory = client.agent_engines.memories.get( + memory = client.runtimes.memories.get( name=operation.response.generated_memories[0].memory.name ) assert memory.metadata == metadata @@ -185,11 +185,11 @@ def test_generate_memories_with_metadata(client): overwrite_metadata = { "my_string_key": types.MemoryMetadataValue(string_value="new_value"), } - operation = client.agent_engines.memories.generate( + operation = client.runtimes.memories.generate( name=agent_engine.api_resource.name, scope=scope, direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig( + config=types.GenerateRuntimeMemoriesConfig( metadata=overwrite_metadata, metadata_merge_strategy=types.MemoryMetadataMergeStrategy.OVERWRITE, ), @@ -199,7 +199,7 @@ def test_generate_memories_with_metadata(client): operation.response.generated_memories[0].action == types.GenerateMemoriesResponseGeneratedMemoryAction.UPDATED ) - memory = client.agent_engines.memories.get( + memory = client.runtimes.memories.get( name=operation.response.generated_memories[0].memory.name ) assert memory.metadata == overwrite_metadata @@ -208,11 +208,11 @@ def test_generate_memories_with_metadata(client): new_metadata = { "my_double_key": types.MemoryMetadataValue(double_value=123.456), } - operation = client.agent_engines.memories.generate( + operation = client.runtimes.memories.generate( name=agent_engine.api_resource.name, scope=scope, direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig( + config=types.GenerateRuntimeMemoriesConfig( metadata=new_metadata, metadata_merge_strategy=types.MemoryMetadataMergeStrategy.MERGE, ), @@ -222,7 +222,7 @@ def test_generate_memories_with_metadata(client): operation.response.generated_memories[0].action == types.GenerateMemoriesResponseGeneratedMemoryAction.UPDATED ) - memory = client.agent_engines.memories.get( + memory = client.runtimes.memories.get( name=operation.response.generated_memories[0].memory.name ) assert memory.metadata == {**overwrite_metadata, **new_metadata} @@ -233,11 +233,11 @@ def test_generate_memories_with_metadata(client): restricted_metadata = { "my_string_key": types.MemoryMetadataValue(string_value="new_value2"), } - operation = client.agent_engines.memories.generate( + operation = client.runtimes.memories.generate( name=agent_engine.api_resource.name, scope=scope, direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig( + config=types.GenerateRuntimeMemoriesConfig( metadata=restricted_metadata, metadata_merge_strategy="REQUIRE_EXACT_MATCH", ), @@ -248,18 +248,18 @@ def test_generate_memories_with_metadata(client): operation.response.generated_memories[0].action == types.GenerateMemoriesResponseGeneratedMemoryAction.CREATED ) - memory = client.agent_engines.memories.get( + memory = client.runtimes.memories.get( name=operation.response.generated_memories[0].memory.name ) assert memory.metadata == restricted_metadata # Send a second request where the metadata matches only one of the existing # memories. - operation = client.agent_engines.memories.generate( + operation = client.runtimes.memories.generate( name=agent_engine.api_resource.name, scope=scope, direct_memories_source=direct_memories_source, - config=types.GenerateAgentEngineMemoriesConfig( + config=types.GenerateRuntimeMemoriesConfig( metadata=restricted_metadata, metadata_merge_strategy="REQUIRE_EXACT_MATCH", ), @@ -271,13 +271,13 @@ def test_generate_memories_with_metadata(client): ) assert operation.response.generated_memories[0].memory.name == memory.name - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.generate_memories", + test_method="runtimes.generate_memories", ) @@ -286,8 +286,8 @@ def test_generate_memories_with_metadata(client): @pytest.mark.asyncio async def test_generate_and_rollback_memories_async(client): - agent_engine = client.agent_engines.create() - await client.aio.agent_engines.memories.generate( + agent_engine = client.runtimes.create() + await client.aio.runtimes.memories.generate( name=agent_engine.api_resource.name, scope={"user_id": "test-user-id"}, direct_memories_source=types.GenerateMemoriesRequestDirectMemoriesSource( @@ -300,15 +300,15 @@ async def test_generate_and_rollback_memories_async(client): ), ] ), - config=types.GenerateAgentEngineMemoriesConfig(wait_for_completion=True), + config=types.GenerateRuntimeMemoriesConfig(wait_for_completion=True), ) - memories_pager = await client.aio.agent_engines.memories.list( + memories_pager = await client.aio.runtimes.memories.list( name=agent_engine.api_resource.name ) memory_list = [item async for item in memories_pager] assert len(memory_list) >= 1 - revisions_pager = await client.aio.agent_engines.memories.revisions.list( + revisions_pager = await client.aio.runtimes.memories.revisions.list( name=memory_list[0].name ) memory_revisions = [item async for item in revisions_pager] @@ -316,31 +316,31 @@ async def test_generate_and_rollback_memories_async(client): revision_name = memory_revisions[0].name # Update the memory. - client.agent_engines.memories._update( + client.runtimes.memories._update( name=memory_list[0].name, fact="This is temporary", scope={"user_id": "test-user-id"}, ) - memory = await client.aio.agent_engines.memories.get(name=memory_list[0].name) + memory = await client.aio.runtimes.memories.get(name=memory_list[0].name) assert memory.fact == "This is temporary" # Rollback to the revision with the original fact that was created by the # generation request. - await client.aio.agent_engines.memories.rollback( + await client.aio.runtimes.memories.rollback( name=memory_list[0].name, target_revision_id=revision_name.split("/")[-1], ) - memory = await client.aio.agent_engines.memories.get(name=memory_list[0].name) + memory = await client.aio.runtimes.memories.get(name=memory_list[0].name) assert memory.fact == memory_revisions[0].fact - await client.aio.agent_engines.delete( + await client.aio.runtimes.delete( name=agent_engine.api_resource.name, force=True ) def test_generate_memories_with_allowed_topics(client): - agent_engine = client.agent_engines.create() - client.agent_engines.memories.generate( + agent_engine = client.runtimes.create() + client.runtimes.memories.generate( name=agent_engine.api_resource.name, scope={"user_id": "test-user-id"}, direct_contents_source=types.GenerateMemoriesRequestDirectContentsSource( @@ -359,7 +359,7 @@ def test_generate_memories_with_allowed_topics(client): ), ] ), - config=types.GenerateAgentEngineMemoriesConfig( + config=types.GenerateRuntimeMemoriesConfig( allowed_topics=[ types.MemoryTopicId( managed_memory_topic=types.ManagedTopicEnum.USER_PREFERENCES @@ -371,11 +371,11 @@ def test_generate_memories_with_allowed_topics(client): assert ( len( list( - client.agent_engines.memories.list( + client.runtimes.memories.list( name=agent_engine.api_resource.name, ) ) ) == 1 ) - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_get_ae_runtime_revision.py b/tests/unit/agentplatform/genai/replays/test_get_ae_runtime_revision.py index ef3413b5e7..4b3e597bca 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 @@ -38,7 +38,7 @@ def test_get_runtime_revisions( mock_agent_engine_create_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - agent_engine = client.agent_engines.create( + agent_engine = client.runtimes.create( config={ "display_name": "test-agent-engine-get-runtime-revisions", "source_packages": [ @@ -58,18 +58,18 @@ def test_get_runtime_revisions( agent_engine.api_resource.display_name == "test-agent-engine-get-runtime-revisions" ) - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.runtimes.revisions.list( name=agent_engine.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) @@ -78,7 +78,7 @@ def test_get_runtime_revisions( 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",) @@ -99,7 +99,7 @@ async def test_async_get_runtime_revisions( mock_agent_engine_create_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - agent_engine = client.agent_engines.create( + agent_engine = client.runtimes.create( config={ "display_name": "test-agent-engine-get-runtime-revisions", "source_packages": [ @@ -119,21 +119,21 @@ async def test_async_get_runtime_revisions( agent_engine.api_resource.display_name == "test-agent-engine-get-runtime-revisions" ) - runtime_revisions_iter = client.aio.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.aio.runtimes.revisions.list( name=agent_engine.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( + await client.aio.runtimes.delete( name=agent_engine.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_a2a_task.py b/tests/unit/agentplatform/genai/replays/test_get_agent_engine_a2a_task.py deleted file mode 100644 index 6ffb4b4b50..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_a2a_task.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -import pytest - - -def test_get_a2a_task(client): - # Use the autopush environment. - client._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client._api_client._http_options.api_version = "internal" - - created_task = client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - assert isinstance(created_task, types.A2aTask) - - task = client.agent_engines.a2a_tasks.get( - name=created_task.name, - ) - assert isinstance(task, types.A2aTask) - assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" - assert task.context_id == "context123" - assert task.state == types.State.SUBMITTED - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), -) - - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_get_a2a_task_async(client): - # Use the autopush environment. - client.aio._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client.aio._api_client._http_options.api_version = "internal" - - created_task = await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - assert isinstance(created_task, types.A2aTask) - - task = await client.aio.agent_engines.a2a_tasks.get( - name=created_task.name, - ) - assert isinstance(task, types.A2aTask) - assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" - assert task.context_id == "context123" - assert task.state == types.State.SUBMITTED - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py b/tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py index 3ec3093562..9da221e571 100644 --- a/tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py +++ b/tests/unit/agentplatform/genai/replays/test_get_agent_engine_memory.py @@ -21,25 +21,25 @@ def test_get_memory(client): - agent_engine = client.agent_engines.create() - operation = client.agent_engines.memories.create( + agent_engine = client.runtimes.create() + operation = client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, ) - assert isinstance(operation, types.AgentEngineMemoryOperation) - memory = client.agent_engines.memories.get( + assert isinstance(operation, types.RuntimeMemoryOperation) + memory = client.runtimes.memories.get( name=operation.response.name, ) assert isinstance(memory, types.Memory) assert memory.name == operation.response.name - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.get", + test_method="runtimes.memories.get", ) @@ -48,18 +48,18 @@ def test_get_memory(client): @pytest.mark.asyncio async def test_get_memory_async(client): - agent_engine = client.agent_engines.create() - operation = await client.aio.agent_engines.memories.create( + agent_engine = client.runtimes.create() + operation = await client.aio.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, ) - assert isinstance(operation, types.AgentEngineMemoryOperation) - memory = await client.aio.agent_engines.memories.get( + assert isinstance(operation, types.RuntimeMemoryOperation) + memory = await client.aio.runtimes.memories.get( name=operation.response.name, ) assert isinstance(memory, types.Memory) assert memory.name == operation.response.name - await client.aio.agent_engines.delete( + await client.aio.runtimes.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..690139318c 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) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - operation = client.agent_engines.sandboxes.create( + operation = client.sandboxes.create( name=agent_engine.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..744b793239 100644 --- a/tests/unit/agentplatform/genai/replays/test_get_feedback_entry.py +++ b/tests/unit/agentplatform/genai/replays/test_get_feedback_entry.py @@ -32,8 +32,8 @@ def test_get(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) try: @@ -66,13 +66,13 @@ 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=agent_engine.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) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) try: @@ -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( + await client.aio.runtimes.delete( name=agent_engine.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_ingest_events_memory_bank.py b/tests/unit/agentplatform/genai/replays/test_ingest_events_memory_bank.py index ff7ad46583..68135ee286 100644 --- a/tests/unit/agentplatform/genai/replays/test_ingest_events_memory_bank.py +++ b/tests/unit/agentplatform/genai/replays/test_ingest_events_memory_bank.py @@ -18,16 +18,16 @@ def test_ingest_events(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() assert not list( - client.agent_engines.memories.list( + client.runtimes.memories.list( name=agent_engine.api_resource.name, ) ) scope = {"user_id": "test-user-id"} # Generate memories using source content. This result is non-deterministic, # because an LLM is used to generate the memories. - client.agent_engines.memories.ingest_events( + client.runtimes.memories.ingest_events( name=agent_engine.api_resource.name, scope=scope, direct_contents_source={ @@ -48,7 +48,7 @@ def test_ingest_events(client): }, ) memories = list( - client.agent_engines.memories.retrieve( + client.runtimes.memories.retrieve( name=agent_engine.api_resource.name, scope=scope, ) @@ -58,7 +58,7 @@ def test_ingest_events(client): # of inactivity. assert len(memories) == 0 - client.agent_engines.memories.ingest_events( + client.runtimes.memories.ingest_events( name=agent_engine.api_resource.name, scope=scope, direct_contents_source={ @@ -84,7 +84,7 @@ def test_ingest_events(client): }, ) memories = list( - client.agent_engines.memories.retrieve( + client.runtimes.memories.retrieve( name=agent_engine.api_resource.name, scope=scope, simple_retrieval_params={ @@ -101,18 +101,18 @@ def test_ingest_events(client): # The user-provided `revision_labels` are applied to the generated memory's # revision (not the Memory itself), so list the memory's revisions to verify. revisions = list( - client.agent_engines.memories.revisions.list( + client.runtimes.memories.revisions.list( name=memories[0].memory.name, ) ) assert revisions assert revisions[0].labels == {"source": "ingest-events-test"} - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.ingest_events", + test_method="runtimes.memories.ingest_events", ) diff --git a/tests/unit/agentplatform/genai/replays/test_list_ae_runtime_revisions.py b/tests/unit/agentplatform/genai/replays/test_list_ae_runtime_revisions.py index 49f9251479..74cf9e1338 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 @@ -38,7 +38,7 @@ def test_list_runtime_revisions( mock_agent_engine_create_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - agent_engine = client.agent_engines.create( + agent_engine = client.runtimes.create( config={ "display_name": "test-agent-engine-list-runtime-revisions", "source_packages": [ @@ -58,12 +58,12 @@ def test_list_runtime_revisions( agent_engine.api_resource.display_name == "test-agent-engine-list-runtime-revisions" ) - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.runtimes.revisions.list( name=agent_engine.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 ) @@ -74,7 +74,7 @@ def test_list_runtime_revisions( 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",) @@ -95,7 +95,7 @@ async def test_async_list_runtime_revisions( mock_agent_engine_create_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - agent_engine = client.agent_engines.create( + agent_engine = client.runtimes.create( config={ "display_name": "test-agent-engine-list-runtime-revisions", "source_packages": [ @@ -115,18 +115,18 @@ async def test_async_list_runtime_revisions( agent_engine.api_resource.display_name == "test-agent-engine-list-runtime-revisions" ) - runtime_revisions_iter = client.aio.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.aio.runtimes.revisions.list( name=agent_engine.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( + await client.aio.runtimes.delete( name=agent_engine.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_task_events.py b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_task_events.py deleted file mode 100644 index 58b26bc3c8..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_task_events.py +++ /dev/null @@ -1,160 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -import pytest - - -def test_list_simple_a2a_task_events(client): - # Use the autopush environment. - client._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client._api_client._http_options.api_version = "internal" - task = client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task999", - config=types.CreateAgentEngineTaskConfig(context_id="context999"), - ) - assert isinstance(task, types.A2aTask) - - events = list( - client.agent_engines.a2a_tasks.events.list( - name=task.name, - ) - ) - assert len(events) == 1 - assert events[0].event_data.state_change.new_state == "SUBMITTED" - - client.agent_engines.a2a_tasks.events.append( - name=task.name, - task_events=[ - types.TaskEvent( - event_data=types.TaskEventData( - metadata_change=types.TaskMetadataChange( - new_metadata={"key1": "value1"} - ) - ), - event_sequence_number=1, - ), - types.TaskEvent( - event_data=types.TaskEventData( - metadata_change=types.TaskMetadataChange( - new_metadata={"key2": "value2"} - ) - ), - event_sequence_number=2, - ), - ], - ) - - result = list( - client.agent_engines.a2a_tasks.events.list( - name=task.name, - config=types.ListAgentEngineTaskEventsConfig( - order_by="create_time desc", - ), - ) - ) - - assert len(result) == 3 - assert result[0].event_sequence_number == 2 - assert result[0].event_data.metadata_change.new_metadata == {"key2": "value2"} - assert result[1].event_sequence_number == 1 - assert result[1].event_data.metadata_change.new_metadata == {"key1": "value1"} - assert result[2].event_data.state_change.new_state == "SUBMITTED" - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), -) - - -@pytest.mark.asyncio -async def test_list_simple_a2a_task_events_async(client): - # Use the autopush environment. - client.aio._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client.aio._api_client._http_options.api_version = "internal" - task = await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task999", - config=types.CreateAgentEngineTaskConfig(context_id="context999"), - ) - assert isinstance(task, types.A2aTask) - - events = list( - await client.aio.agent_engines.a2a_tasks.events.list( - name=task.name, - ) - ) - assert len(events) == 1 - assert events[0].event_data.state_change.new_state == "SUBMITTED" - - await client.aio.agent_engines.a2a_tasks.events.append( - name=task.name, - task_events=[ - types.TaskEvent( - event_data=types.TaskEventData( - metadata_change=types.TaskMetadataChange( - new_metadata={"key1": "value1"} - ) - ), - event_sequence_number=1, - ), - types.TaskEvent( - event_data=types.TaskEventData( - metadata_change=types.TaskMetadataChange( - new_metadata={"key2": "value2"} - ) - ), - event_sequence_number=2, - ), - ], - ) - - result = list( - await client.aio.agent_engines.a2a_tasks.events.list( - name=task.name, - config=types.ListAgentEngineTaskEventsConfig( - order_by="create_time desc", - ), - ) - ) - - assert len(result) == 3 - assert result[0].event_sequence_number == 2 - assert result[0].event_data.metadata_change.new_metadata == {"key2": "value2"} - assert result[1].event_sequence_number == 1 - assert result[1].event_data.metadata_change.new_metadata == {"key1": "value1"} - assert result[2].event_data.state_change.new_state == "SUBMITTED" - - # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_tasks.py b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_tasks.py deleted file mode 100644 index 1f580531b2..0000000000 --- a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_a2a_tasks.py +++ /dev/null @@ -1,117 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# pylint: disable=protected-access,bad-continuation,missing-function-docstring - -from tests.unit.agentplatform.genai.replays import pytest_helper -from agentplatform._genai import types -import pytest - - -def test_list_a2a_tasks(client): - # Use the autopush environment. - client._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client._api_client._http_options.api_version = "internal" - - assert not list( - client.agent_engines.a2a_tasks.list( - name=agent_engine.api_resource.name, - ) - ) - client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - client.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task456", - config=types.CreateAgentEngineTaskConfig(context_id="context456"), - ) - a2a_tasks_list = client.agent_engines.a2a_tasks.list( - name=agent_engine.api_resource.name, - config=types.ListAgentEngineTasksConfig( - page_size=1, - order_by="create_time asc", - ), - ) - assert len(a2a_tasks_list) == 1 - assert isinstance(a2a_tasks_list[0], types.A2aTask) - assert a2a_tasks_list[0].name == ( - f"{agent_engine.api_resource.name}/a2aTasks/task123" - ) - assert a2a_tasks_list[0].context_id == "context123" - - # Clean up resources. - agent_engine.delete(force=True) - - -pytestmark = pytest_helper.setup( - file=__file__, - globals_for_file=globals(), -) - - -pytest_plugins = ("pytest_asyncio",) - - -@pytest.mark.asyncio -async def test_list_a2a_tasks_async(client): - # Use the autopush environment. - client.aio._api_client._http_options.base_url = ( - "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/" - ) - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) - assert isinstance(agent_engine.api_resource, types.ReasoningEngine) - # Use the internal API version for internal API access. - client.aio._api_client._http_options.api_version = "internal" - - assert not list( - await client.aio.agent_engines.a2a_tasks.list( - name=agent_engine.api_resource.name, - ) - ) - await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task123", - config=types.CreateAgentEngineTaskConfig(context_id="context123"), - ) - await client.aio.agent_engines.a2a_tasks.create( - name=agent_engine.api_resource.name, - a2a_task_id="task456", - config=types.CreateAgentEngineTaskConfig(context_id="context456"), - ) - a2a_tasks_list = await client.aio.agent_engines.a2a_tasks.list( - name=agent_engine.api_resource.name, - config=types.ListAgentEngineTasksConfig( - page_size=1, - order_by="create_time asc", - ), - ) - assert len(a2a_tasks_list) == 1 - assert isinstance(a2a_tasks_list[0], types.A2aTask) - assert a2a_tasks_list[0].name == ( - f"{agent_engine.api_resource.name}/a2aTasks/task123" - ) - assert a2a_tasks_list[0].context_id == "context123" - - # Clean up resources. - agent_engine.delete(force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_memories.py b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_memories.py index 4373b4bf25..c1191a9c15 100644 --- a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_memories.py +++ b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_memories.py @@ -21,13 +21,13 @@ def test_list_memories(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() assert not list( - client.agent_engines.memories.list( + client.runtimes.memories.list( name=agent_engine.api_resource.name, ) ) - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, @@ -35,7 +35,7 @@ def test_list_memories(client): "wait_for_completion": True, }, ) - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_2", scope={"user_id": "456"}, @@ -43,9 +43,9 @@ def test_list_memories(client): "wait_for_completion": True, }, ) - memory_list = client.agent_engines.memories.list( + memory_list = client.runtimes.memories.list( name=agent_engine.api_resource.name, - config=types.ListAgentEngineMemoryConfig( + config=types.ListRuntimeMemoryConfig( page_size=1, order_by="create_time asc", ), @@ -61,7 +61,7 @@ def test_list_memories(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.list", + test_method="runtimes.memories.list", ) @@ -70,13 +70,13 @@ def test_list_memories(client): @pytest.mark.asyncio async def test_async_list_memories(client): - agent_engine = client.agent_engines.create() - pager = await client.aio.agent_engines.memories.list( + agent_engine = client.runtimes.create() + pager = await client.aio.runtimes.memories.list( name=agent_engine.api_resource.name ) assert not [item async for item in pager] - await client.aio.agent_engines.memories.create( + await client.aio.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_2", scope={"user_id": "456"}, @@ -84,13 +84,13 @@ async def test_async_list_memories(client): "wait_for_completion": True, }, ) - pager = await client.aio.agent_engines.memories.list( + pager = await client.aio.runtimes.memories.list( name=agent_engine.api_resource.name ) memory_list = [item async for item in pager] assert len(memory_list) == 1 assert isinstance(memory_list[0], types.Memory) - await client.aio.agent_engines.delete( + await client.aio.runtimes.delete( name=agent_engine.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_list_agent_engine_sandboxes.py b/tests/unit/agentplatform/genai/replays/test_list_agent_engine_sandboxes.py index 0e141495d1..bd7f8288c2 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,25 +19,25 @@ def test_list_sandboxes(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() assert not list( - client.agent_engines.sandboxes.list( + client.sandboxes.list( name=agent_engine.api_resource.name, ) ) - operation = client.agent_engines.sandboxes.create( + operation = client.sandboxes.create( name=agent_engine.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( + sandbox_list = client.sandboxes.list( name=agent_engine.api_resource.name, ) assert len(sandbox_list) == 1 @@ -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..376ed00564 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( + agent_engine = client.runtimes.create() + operation = client.sessions.create( name=agent_engine.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( + agent_engine = client.runtimes.create() + operation = await client.aio.sessions.create( name=agent_engine.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=agent_engine.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..a0c656b1fd 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() + agent_engine = client.runtimes.create() assert not list( - client.agent_engines.sessions.list( + client.sessions.list( name=agent_engine.api_resource.name, ) ) - client.agent_engines.sessions.create( + client.sessions.create( name=agent_engine.api_resource.name, user_id="test-user-123", ) - session_list = client.agent_engines.sessions.list( + session_list = client.sessions.list( name=agent_engine.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=agent_engine.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( + agent_engine = client.runtimes.create() + pager = await client.aio.sessions.list( name=agent_engine.api_resource.name ) assert not [item async for item in pager] - await client.aio.agent_engines.sessions.create( + await client.aio.sessions.create( name=agent_engine.api_resource.name, user_id="test-user-123", ) - pager = await client.aio.agent_engines.sessions.list( + pager = await client.aio.sessions.list( name=agent_engine.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=agent_engine.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..ba7fc57beb 100644 --- a/tests/unit/agentplatform/genai/replays/test_list_feedback_entries.py +++ b/tests/unit/agentplatform/genai/replays/test_list_feedback_entries.py @@ -32,8 +32,8 @@ def test_list(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) try: @@ -104,7 +104,7 @@ def test_list(client): assert thumbs_down_entry.custom_metadata == {"key_a": "val_a"} finally: # Clean up resources. - client.agent_engines.delete( + client.runtimes.delete( name=agent_engine.api_resource.name, force=True, ) @@ -112,8 +112,8 @@ def test_list(client): @pytest.mark.asyncio async def test_list_async(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) try: @@ -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( + await client.aio.runtimes.delete( name=agent_engine.api_resource.name, force=True, ) diff --git a/tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py b/tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py index e123597953..e9c12e04aa 100644 --- a/tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py +++ b/tests/unit/agentplatform/genai/replays/test_purge_agent_engine_memories.py @@ -22,27 +22,27 @@ def test_purge_memories(client): """Tests purging memories.""" - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() try: - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_1", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_2", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_3", scope={"user_id": "456"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_4", scope={"user_id": "456"}, @@ -51,7 +51,7 @@ def test_purge_memories(client): "metadata": {"my_key": {"string_value": "my_value"}}, }, ) - operation = client.agent_engines.memories.purge( + operation = client.runtimes.memories.purge( name=agent_engine.api_resource.name, filter="scope.user_id=123", config={"wait_for_completion": True}, @@ -62,7 +62,7 @@ def test_purge_memories(client): assert ( len( list( - client.agent_engines.memories.list( + client.runtimes.memories.list( name=agent_engine.api_resource.name ) ) @@ -70,7 +70,7 @@ def test_purge_memories(client): == 4 ) # Now, actually purge the memories. - operation = client.agent_engines.memories.purge( + operation = client.runtimes.memories.purge( name=agent_engine.api_resource.name, filter="scope.user_id=123", force=True, @@ -81,7 +81,7 @@ def test_purge_memories(client): assert ( len( list( - client.agent_engines.memories.list( + client.runtimes.memories.list( name=agent_engine.api_resource.name ) ) @@ -89,7 +89,7 @@ def test_purge_memories(client): == 2 ) # Purge memories using filter groups. - operation = client.agent_engines.memories.purge( + operation = client.runtimes.memories.purge( name=agent_engine.api_resource.name, force=True, filter_groups=[ @@ -104,7 +104,7 @@ def test_purge_memories(client): assert ( len( list( - client.agent_engines.memories.list( + client.runtimes.memories.list( name=agent_engine.api_resource.name ) ) @@ -112,13 +112,13 @@ def test_purge_memories(client): == 1 ) finally: - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.purge", + test_method="runtimes.memories.purge", ) @@ -127,28 +127,28 @@ def test_purge_memories(client): @pytest.mark.asyncio async def test_purge_memories_async(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() try: - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_1", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_2", scope={"user_id": "123"}, config={"wait_for_completion": True}, ) - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_3", scope={"user_id": "456"}, config={"wait_for_completion": True}, ) - operation = await client.aio.agent_engines.memories.purge( + operation = await client.aio.runtimes.memories.purge( name=agent_engine.api_resource.name, filter="scope.user_id=123", config={"wait_for_completion": True}, @@ -159,7 +159,7 @@ async def test_purge_memories_async(client): assert ( len( list( - client.agent_engines.memories.list( + client.runtimes.memories.list( name=agent_engine.api_resource.name ) ) @@ -167,7 +167,7 @@ async def test_purge_memories_async(client): == 3 ) # Now, actually purge the memories. - operation = await client.aio.agent_engines.memories.purge( + operation = await client.aio.runtimes.memories.purge( name=agent_engine.api_resource.name, filter="scope.user_id=123", force=True, @@ -178,7 +178,7 @@ async def test_purge_memories_async(client): assert ( len( list( - client.agent_engines.memories.list( + client.runtimes.memories.list( name=agent_engine.api_resource.name ) ) @@ -186,4 +186,4 @@ async def test_purge_memories_async(client): == 1 ) finally: - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) diff --git a/tests/unit/agentplatform/genai/replays/test_retrieve_agent_engine_memories.py b/tests/unit/agentplatform/genai/replays/test_retrieve_agent_engine_memories.py index 02a3e034d0..6817741b09 100644 --- a/tests/unit/agentplatform/genai/replays/test_retrieve_agent_engine_memories.py +++ b/tests/unit/agentplatform/genai/replays/test_retrieve_agent_engine_memories.py @@ -24,9 +24,9 @@ def test_retrieve_memories_with_similarity_search_params(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() assert not list( - client.agent_engines.memories.retrieve( + client.runtimes.memories.retrieve( name=agent_engine.api_resource.name, scope={"user_id": "123"}, similarity_search_params=types.RetrieveMemoriesRequestSimilaritySearchParams( @@ -34,7 +34,7 @@ def test_retrieve_memories_with_similarity_search_params(client): ), ) ) - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_1", scope={"user_id": "123"}, @@ -42,7 +42,7 @@ def test_retrieve_memories_with_similarity_search_params(client): assert ( len( list( - client.agent_engines.memories.retrieve( + client.runtimes.memories.retrieve( name=agent_engine.api_resource.name, scope={"user_id": "123"}, ) @@ -51,12 +51,12 @@ def test_retrieve_memories_with_similarity_search_params(client): == 1 ) assert not list( - client.agent_engines.memories.retrieve( + client.runtimes.memories.retrieve( name=agent_engine.api_resource.name, scope={"user_id": "456"}, ) ) - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_2", scope={"user_id": "123"}, @@ -64,7 +64,7 @@ def test_retrieve_memories_with_similarity_search_params(client): assert ( len( list( - client.agent_engines.memories.retrieve( + client.runtimes.memories.retrieve( name=agent_engine.api_resource.name, scope={"user_id": "123"}, ) @@ -77,13 +77,13 @@ def test_retrieve_memories_with_similarity_search_params(client): def test_retrieve_memories_with_simple_retrieval_params(client): - agent_engine = client.agent_engines.create() - client.agent_engines.memories.create( + agent_engine = client.runtimes.create() + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_1", scope={"user_id": "123"}, ) - memories = client.agent_engines.memories.retrieve( + memories = client.runtimes.memories.retrieve( name=agent_engine.api_resource.name, scope={"user_id": "123"}, simple_retrieval_params=types.RetrieveMemoriesRequestSimpleRetrievalParams( @@ -94,17 +94,17 @@ def test_retrieve_memories_with_simple_retrieval_params(client): assert isinstance(memories.page[0], types.RetrieveMemoriesResponseRetrievedMemory) assert memories.page_size == 1 - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_2", scope={"user_id": "123"}, ) - memories = client.agent_engines.memories.retrieve( + memories = client.runtimes.memories.retrieve( name=agent_engine.api_resource.name, scope={"user_id": "123"} ) assert memories.page_size == 2 - memories = client.agent_engines.memories.retrieve( + memories = client.runtimes.memories.retrieve( name=agent_engine.api_resource.name, scope={"user_id": "123"}, config={"filter": 'fact="memory_fact_2"'}, @@ -117,7 +117,7 @@ def test_retrieve_memories_with_simple_retrieval_params(client): def test_retrieve_memories_with_metadata(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() metadata = { "my_string_key": types.MemoryMetadataValue(string_value="my_string_value"), "my_double_key": types.MemoryMetadataValue(double_value=123.456), @@ -129,12 +129,12 @@ def test_retrieve_memories_with_metadata(client): ), } scope = {"user_id": "123"} - client.agent_engines.memories.create( + client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_1", scope=scope, ) - operation = client.agent_engines.memories.create( + operation = client.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact_2", scope=scope, @@ -142,7 +142,7 @@ def test_retrieve_memories_with_metadata(client): ) memory_name2 = operation.response.name - results = client.agent_engines.memories.retrieve( + results = client.runtimes.memories.retrieve( name=agent_engine.api_resource.name, scope=scope, config={ @@ -168,7 +168,7 @@ def test_retrieve_memories_with_metadata(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.memories.retrieve", + test_method="runtimes.memories.retrieve", ) @@ -177,20 +177,20 @@ def test_retrieve_memories_with_metadata(client): @pytest.mark.asyncio async def test_retrieve_memories_async(client): - agent_engine = client.agent_engines.create() - operation = await client.aio.agent_engines.memories.create( + agent_engine = client.runtimes.create() + operation = await client.aio.runtimes.memories.create( name=agent_engine.api_resource.name, fact="memory_fact", scope={"user_id": "123"}, ) - assert isinstance(operation, types.AgentEngineMemoryOperation) - pager = await client.aio.agent_engines.memories.retrieve( + assert isinstance(operation, types.RuntimeMemoryOperation) + pager = await client.aio.runtimes.memories.retrieve( name=agent_engine.api_resource.name, scope={"user_id": "123"}, ) memories = [item async for item in pager] assert len(memories) == 1 assert isinstance(memories[0], types.RetrieveMemoriesResponseRetrievedMemory) - await client.aio.agent_engines.delete( + await client.aio.runtimes.delete( name=agent_engine.api_resource.name, force=True ) diff --git a/tests/unit/agentplatform/genai/replays/test_run_inference.py b/tests/unit/agentplatform/genai/replays/test_run_inference.py index 0c4afb9eaf..670b192e81 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( + agent_engine = client.runtimes.get( name="projects/977012026409/locations/us-central1" "/reasoningEngines/7188347537655332864" ) diff --git a/tests/unit/agentplatform/genai/replays/test_structured_memories.py b/tests/unit/agentplatform/genai/replays/test_structured_memories.py index 5d1c2a39a5..b3b2f98c95 100644 --- a/tests/unit/agentplatform/genai/replays/test_structured_memories.py +++ b/tests/unit/agentplatform/genai/replays/test_structured_memories.py @@ -44,7 +44,7 @@ def test_generate_and_retrieve_profile(client): structured_memory_config_obj = types.StructuredMemoryConfig( **structured_memory_config ) - agent_engine = client.agent_engines.create( + agent_engine = client.runtimes.create( config={ "context_spec": { "memory_bank_config": { @@ -56,7 +56,7 @@ def test_generate_and_retrieve_profile(client): }, ) try: - agent_engine = client.agent_engines.get(name=agent_engine.api_resource.name) + agent_engine = client.runtimes.get(name=agent_engine.api_resource.name) memory_bank_config = agent_engine.api_resource.context_spec.memory_bank_config assert memory_bank_config.customization_configs == [ memory_bank_customization_config @@ -66,7 +66,7 @@ def test_generate_and_retrieve_profile(client): ] scope = {"user_id": "123"} - client.agent_engines.memories.generate( + client.runtimes.memories.generate( name=agent_engine.api_resource.name, scope=scope, direct_contents_source={ @@ -74,7 +74,7 @@ def test_generate_and_retrieve_profile(client): }, ) memories = list( - client.agent_engines.memories.retrieve( + client.runtimes.memories.retrieve( name=agent_engine.api_resource.name, scope=scope, config={"memory_types": ["STRUCTURED_PROFILE"]}, @@ -83,18 +83,18 @@ def test_generate_and_retrieve_profile(client): assert len(memories) >= 1 assert memories[0].memory.structured_content is not None - response = client.agent_engines.memories.retrieve_profiles( + response = client.runtimes.memories.retrieve_profiles( name=agent_engine.api_resource.name, scope=scope ) assert len(response.profiles) == 1 finally: # Clean up resources. - client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) + client.runtimes.delete(name=agent_engine.api_resource.name, force=True) pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.retrieve_profiles", + test_method="runtimes.retrieve_profiles", ) diff --git a/tests/unit/agentplatform/genai/replays/test_update_agent_engine.py b/tests/unit/agentplatform/genai/replays/test_update_agent_engine.py index 1829e79a02..00c6d8cf56 100644 --- a/tests/unit/agentplatform/genai/replays/test_update_agent_engine.py +++ b/tests/unit/agentplatform/genai/replays/test_update_agent_engine.py @@ -20,17 +20,17 @@ def test_agent_engines_update(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() assert agent_engine.api_resource.display_name is None - updated_agent_engine = client.agent_engines.update( + updated_agent_engine = client.runtimes.update( name=agent_engine.api_resource.name, - config=types.AgentEngineConfig( + config=types.RuntimeConfig( display_name="updated_display_name", description="updated description", ), ) - assert isinstance(updated_agent_engine, types.AgentEngine) + assert isinstance(updated_agent_engine, types.Runtime) assert updated_agent_engine.api_resource.name == agent_engine.api_resource.name assert updated_agent_engine.api_resource.display_name == "updated_display_name" @@ -40,5 +40,5 @@ def test_agent_engines_update(client): pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), - test_method="agent_engines.update", + test_method="runtimes.update", ) diff --git a/tests/unit/agentplatform/genai/replays/test_update_feedback_context.py b/tests/unit/agentplatform/genai/replays/test_update_feedback_context.py index 2b19237aaa..6715cc7cd9 100644 --- a/tests/unit/agentplatform/genai/replays/test_update_feedback_context.py +++ b/tests/unit/agentplatform/genai/replays/test_update_feedback_context.py @@ -33,7 +33,7 @@ def test_update_and_get(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() try: operation = client.feedback_entries.create( @@ -110,7 +110,7 @@ def test_update_and_get(client): finally: # Clean up resources. - client.agent_engines.delete( + client.runtimes.delete( name=agent_engine.api_resource.name, force=True, ) @@ -118,7 +118,7 @@ def test_update_and_get(client): @pytest.mark.asyncio async def test_update_and_get_async(client): - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() try: operation = await client.aio.feedback_entries.create( @@ -195,7 +195,7 @@ async def test_update_and_get_async(client): finally: # Clean up resources. - await client.aio.agent_engines.delete( + await client.aio.runtimes.delete( name=agent_engine.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..e49f608d03 100644 --- a/tests/unit/agentplatform/genai/replays/test_update_feedback_entry.py +++ b/tests/unit/agentplatform/genai/replays/test_update_feedback_entry.py @@ -32,8 +32,8 @@ def test_update(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) try: @@ -113,7 +113,7 @@ def test_update(client): finally: # Clean up resources. - client.agent_engines.delete( + client.runtimes.delete( name=agent_engine.api_resource.name, force=True, ) @@ -121,8 +121,8 @@ def test_update(client): @pytest.mark.asyncio async def test_update_async(client): - agent_engine = client.agent_engines.create() - assert isinstance(agent_engine, types.AgentEngine) + agent_engine = client.runtimes.create() + assert isinstance(agent_engine, types.Runtime) assert isinstance(agent_engine.api_resource, types.ReasoningEngine) try: @@ -202,7 +202,7 @@ async def test_update_async(client): finally: # Clean up resources. - await client.aio.agent_engines.delete( + await client.aio.runtimes.delete( name=agent_engine.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..4dbe42be68 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,15 +30,15 @@ def test_agent_engines_update_traffic_to_always_latest(client): ) client._api_client._http_options.api_version = "v1beta1" - agent_engine = client.agent_engines.create() + agent_engine = client.runtimes.create() traffic_config = types.ReasoningEngineTrafficConfig( traffic_split_always_latest=types.ReasoningEngineTrafficConfigTrafficSplitAlwaysLatest(), ) - updated_agent_engine = client.agent_engines.update( + updated_agent_engine = client.runtimes.update( name=agent_engine.api_resource.name, - config=types.AgentEngineConfig( + config=types.RuntimeConfig( traffic_config=traffic_config, ), ) @@ -63,7 +63,7 @@ def test_agent_engines_update_traffic_to_manual_split( mock_agent_engine_create_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - agent_engine = client.agent_engines.create( + agent_engine = client.runtimes.create( config={ "display_name": "test-agent-engine-update-traffic-to-manual-split", "source_packages": [ @@ -80,7 +80,7 @@ def test_agent_engines_update_traffic_to_manual_split( }, ) - runtime_revisions_iter = client.agent_engines.runtimes.revisions.list( + runtime_revisions_iter = client.runtimes.revisions.list( name=agent_engine.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) @@ -101,9 +101,9 @@ def test_agent_engines_update_traffic_to_manual_split( ), ) - updated_agent_engine = client.agent_engines.update( + updated_agent_engine = client.runtimes.update( name=agent_engine.api_resource.name, - config=types.AgentEngineConfig( + config=types.RuntimeConfig( traffic_config=traffic_config, ), ) @@ -127,7 +127,7 @@ def test_agent_engines_update_traffic_with_agent_update( mock_agent_engine_create_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - agent_engine = client.agent_engines.create( + agent_engine = client.runtimes.create( config={ "display_name": "test-agent-engine-update-traffic-with-agent-before-update", "source_packages": [ @@ -148,7 +148,7 @@ def test_agent_engines_update_traffic_with_agent_update( == "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( + runtime_revisions_iter = client.runtimes.revisions.list( name=agent_engine.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) @@ -173,7 +173,7 @@ def test_agent_engines_update_traffic_with_agent_update( mock_agent_engine_create_base64_encoded_tarball, mock_agent_engine_create_path_exists, ): - updated_agent_engine = client.agent_engines.update( + updated_agent_engine = client.runtimes.update( name=agent_engine.api_resource.name, config={ "display_name": "test-agent-engine-update-traffic-with-agent-after-update", @@ -197,7 +197,7 @@ def test_agent_engines_update_traffic_with_agent_update( == "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( + runtime_revisions_iter = client.runtimes.revisions.list( name=agent_engine.api_resource.name, ) runtime_revisions_list = list(runtime_revisions_iter) @@ -219,5 +219,5 @@ def test_agent_engines_update_traffic_with_agent_update( 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..93840f0225 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( + agent_engine_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( + agent_engine_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( @@ -1032,10 +1032,10 @@ def test_query_runtime_revision(self): ) def test_query_agent_engine_async(self): - agent = self.client.agent_engines.runtimes.revisions._register_api_methods( - agent_engine_runtime_revision=_genai_types.AgentEngineRuntimeRevision( + agent = self.client.runtimes.revisions._register_api_methods( + agent_engine_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)) @@ -1065,11 +1065,11 @@ def test_query_agent_engine_async(self): def test_query_agent_engine_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( + agent_engine_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( @@ -1097,13 +1097,13 @@ 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( + agent_engine_runtime_revision=_genai_types.RuntimeRevision( + api_client=self.client.runtimes.revisions, api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -1154,7 +1154,7 @@ def teardown_method(self): def test_delete_runtime_revision(self): with mock.patch.object( - self.client.aio.agent_engines.runtimes.revisions._api_client, + self.client.aio.runtimes.revisions._api_client, "async_request", ) as request_mock: request_mock.side_effect = [ @@ -1171,7 +1171,7 @@ def test_delete_runtime_revision(self): ), ] asyncio.run( - self.client.aio.agent_engines.runtimes.revisions.delete( + self.client.aio.runtimes.revisions.delete( name=_TEST_AGENT_ENGINE_RUNTIME_REVISION_RESOURCE_NAME ) ) diff --git a/tests/unit/agentplatform/genai/test_agent_engines.py b/tests/unit/agentplatform/genai/test_agent_engines.py index e92f121e8d..8b6f90331a 100644 --- a/tests/unit/agentplatform/genai/test_agent_engines.py +++ b/tests/unit/agentplatform/genai/test_agent_engines.py @@ -31,9 +31,9 @@ from google.cloud import aiplatform import agentplatform from google.cloud.aiplatform import initializer -from agentplatform.agent_engines.templates import adk -from agentplatform._genai import _agent_engines_utils -from agentplatform._genai import agent_engines +from agentplatform.frameworks import adk +from agentplatform._genai import _runtimes_utils +from agentplatform._genai import runtimes from agentplatform._genai import types as _genai_types from google.genai import client as genai_client from google.genai import types as genai_types @@ -305,7 +305,7 @@ def clone(self): return self def register_operations(self) -> Dict[str, List[str]]: - # Registered method `missing_method` is not a method of the AgentEngine. + # Registered method `missing_method` is not a method of the Runtime. return { _TEST_STANDARD_API_MODE: [ _TEST_DEFAULT_METHOD_NAME, @@ -323,7 +323,7 @@ def method_to_be_unregistered(self, unused_arbitrary_string_name: str) -> str: return unused_arbitrary_string_name.upper() def register_operations(self) -> Dict[str, List[str]]: - # Registered method `missing_method` is not a method of the AgentEngine. + # Registered method `missing_method` is not a method of the Runtime. return {_TEST_STANDARD_API_MODE: [_TEST_METHOD_TO_BE_UNREGISTERED_NAME]} @@ -341,30 +341,30 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_AGENT_ENGINE_DISPLAY_NAME = "Agent Engine Display Name" _TEST_AGENT_ENGINE_DESCRIPTION = "Agent Engine Description" _TEST_AGENT_ENGINE_LIST_FILTER = f'display_name="{_TEST_AGENT_ENGINE_DISPLAY_NAME}"' -_TEST_GCS_DIR_NAME = _agent_engines_utils._DEFAULT_GCS_DIR_NAME -_TEST_BLOB_FILENAME = _agent_engines_utils._BLOB_FILENAME -_TEST_REQUIREMENTS_FILE = _agent_engines_utils._REQUIREMENTS_FILE -_TEST_EXTRA_PACKAGES_FILE = _agent_engines_utils._EXTRA_PACKAGES_FILE -_TEST_STANDARD_API_MODE = _agent_engines_utils._STANDARD_API_MODE -_TEST_ASYNC_API_MODE = _agent_engines_utils._ASYNC_API_MODE -_TEST_STREAM_API_MODE = _agent_engines_utils._STREAM_API_MODE -_TEST_ASYNC_STREAM_API_MODE = _agent_engines_utils._ASYNC_STREAM_API_MODE -_TEST_BIDI_STREAM_API_MODE = _agent_engines_utils._BIDI_STREAM_API_MODE -_TEST_DEFAULT_METHOD_NAME = _agent_engines_utils._DEFAULT_METHOD_NAME -_TEST_DEFAULT_ASYNC_METHOD_NAME = _agent_engines_utils._DEFAULT_ASYNC_METHOD_NAME -_TEST_DEFAULT_STREAM_METHOD_NAME = _agent_engines_utils._DEFAULT_STREAM_METHOD_NAME +_TEST_GCS_DIR_NAME = _runtimes_utils._DEFAULT_GCS_DIR_NAME +_TEST_BLOB_FILENAME = _runtimes_utils._BLOB_FILENAME +_TEST_REQUIREMENTS_FILE = _runtimes_utils._REQUIREMENTS_FILE +_TEST_EXTRA_PACKAGES_FILE = _runtimes_utils._EXTRA_PACKAGES_FILE +_TEST_STANDARD_API_MODE = _runtimes_utils._STANDARD_API_MODE +_TEST_ASYNC_API_MODE = _runtimes_utils._ASYNC_API_MODE +_TEST_STREAM_API_MODE = _runtimes_utils._STREAM_API_MODE +_TEST_ASYNC_STREAM_API_MODE = _runtimes_utils._ASYNC_STREAM_API_MODE +_TEST_BIDI_STREAM_API_MODE = _runtimes_utils._BIDI_STREAM_API_MODE +_TEST_DEFAULT_METHOD_NAME = _runtimes_utils._DEFAULT_METHOD_NAME +_TEST_DEFAULT_ASYNC_METHOD_NAME = _runtimes_utils._DEFAULT_ASYNC_METHOD_NAME +_TEST_DEFAULT_STREAM_METHOD_NAME = _runtimes_utils._DEFAULT_STREAM_METHOD_NAME _TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME = ( - _agent_engines_utils._DEFAULT_ASYNC_STREAM_METHOD_NAME + _runtimes_utils._DEFAULT_ASYNC_STREAM_METHOD_NAME ) _TEST_DEFAULT_BIDI_STREAM_METHOD_NAME = ( - _agent_engines_utils._DEFAULT_BIDI_STREAM_METHOD_NAME + _runtimes_utils._DEFAULT_BIDI_STREAM_METHOD_NAME ) _TEST_CAPITALIZE_ENGINE_METHOD_DOCSTRING = "Runs the engine." _TEST_STREAM_METHOD_DOCSTRING = "Runs the stream engine." _TEST_ASYNC_STREAM_METHOD_DOCSTRING = "Runs the async stream engine." _TEST_BIDI_STREAM_METHOD_DOCSTRING = "Runs the bidi stream engine." -_TEST_MODE_KEY_IN_SCHEMA = _agent_engines_utils._MODE_KEY_IN_SCHEMA -_TEST_METHOD_NAME_KEY_IN_SCHEMA = _agent_engines_utils._METHOD_NAME_KEY_IN_SCHEMA +_TEST_MODE_KEY_IN_SCHEMA = _runtimes_utils._MODE_KEY_IN_SCHEMA +_TEST_METHOD_NAME_KEY_IN_SCHEMA = _runtimes_utils._METHOD_NAME_KEY_IN_SCHEMA _TEST_CUSTOM_METHOD_NAME = "custom_method" _TEST_CUSTOM_ASYNC_METHOD_NAME = "custom_async_method" _TEST_CUSTOM_STREAM_METHOD_NAME = "custom_stream_method" @@ -454,14 +454,14 @@ def register_operations(self) -> Dict[str, List[str]]: "lib", "main.py", ] -_TEST_AGENT_ENGINE_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_QUERY_SCHEMA = _runtimes_utils._generate_schema( CapitalizeEngine().query, schema_name=_TEST_DEFAULT_METHOD_NAME, ) _TEST_AGENT_ENGINE_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STANDARD_API_MODE _TEST_PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}" _TEST_PYTHON_VERSION_OVERRIDE = "3.11" -_TEST_AGENT_ENGINE_FRAMEWORK = _agent_engines_utils._DEFAULT_AGENT_FRAMEWORK +_TEST_AGENT_ENGINE_FRAMEWORK = _runtimes_utils._DEFAULT_AGENT_FRAMEWORK _TEST_AGENT_ENGINE_CLASS_METHOD_1 = { "description": "Runs the engine.", "name": "query", @@ -590,38 +590,38 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE = [{"output": "hello"}, {"output": "world"}] _TEST_AGENT_ENGINE_OPERATION_SCHEMAS = [] _TEST_AGENT_ENGINE_EXTRA_PACKAGE = "fake.py" -_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA = _runtimes_utils._generate_schema( AsyncQueryEngine().async_query, schema_name=_TEST_DEFAULT_ASYNC_METHOD_NAME, ) _TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_ASYNC_API_MODE -_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA = _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_method, schema_name=_TEST_CUSTOM_METHOD_NAME, ) _TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( _TEST_STANDARD_API_MODE ) -_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA = _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_async_method, schema_name=_TEST_CUSTOM_ASYNC_METHOD_NAME, ) _TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( _TEST_ASYNC_API_MODE ) -_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA = _runtimes_utils._generate_schema( StreamQueryEngine().stream_query, schema_name=_TEST_DEFAULT_STREAM_METHOD_NAME, ) _TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STREAM_API_MODE -_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA = _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_stream_method, schema_name=_TEST_CUSTOM_STREAM_METHOD_NAME, ) _TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( _TEST_STREAM_API_MODE ) -_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA = _runtimes_utils._generate_schema( AsyncStreamQueryEngine().async_stream_query, schema_name=_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME, ) @@ -629,7 +629,7 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_ASYNC_STREAM_API_MODE ) _TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA = ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_async_stream_method, schema_name=_TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME, ) @@ -637,7 +637,7 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = ( _TEST_ASYNC_STREAM_API_MODE ) -_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA = _runtimes_utils._generate_schema( OperationRegistrableEngine().bidi_stream_query, schema_name=_TEST_DEFAULT_BIDI_STREAM_METHOD_NAME, ) @@ -645,7 +645,7 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_BIDI_STREAM_API_MODE ) _TEST_AGENT_ENGINE_CUSTOM_BIDI_STREAM_QUERY_SCHEMA = ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_bidi_stream_method, schema_name=_TEST_CUSTOM_BIDI_STREAM_METHOD_NAME, ) @@ -675,7 +675,7 @@ def register_operations(self) -> Dict[str, List[str]]: _TEST_NO_OPERATION_REGISTRABLE_SCHEMAS = [ _TEST_AGENT_ENGINE_QUERY_SCHEMA, ] -_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA = _agent_engines_utils._generate_schema( +_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA = _runtimes_utils._generate_schema( MethodToBeUnregisteredEngine().method_to_be_unregistered, schema_name=_TEST_METHOD_TO_BE_UNREGISTERED_NAME, ) @@ -849,7 +849,7 @@ class FakeObject: @pytest.mark.usefixtures("google_auth_mock") -class TestAgentEngineHelpers: +class TestRuntimeHelpers: def setup_method(self): importlib.reload(initializer) importlib.reload(aiplatform) @@ -866,9 +866,9 @@ def setup_method(self): def teardown_method(self): initializer.global_pool.shutdown(wait=True) - @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_prepare") def test_create_agent_engine_config_lightweight(self, mock_prepare): - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", staging_bucket=_TEST_STAGING_BUCKET, display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, @@ -879,7 +879,7 @@ def test_create_agent_engine_config_lightweight(self, mock_prepare): "description": _TEST_AGENT_ENGINE_DESCRIPTION, } - @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_prepare") @pytest.mark.parametrize( "env_vars,expected_env_vars", [ @@ -912,7 +912,7 @@ def test_agent_engine_adk_telemetry_enablement( agent.clone = lambda: agent agent.register_operations = lambda: {} - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", agent=agent, staging_bucket=_TEST_STAGING_BUCKET, @@ -927,11 +927,11 @@ def test_agent_engine_adk_telemetry_enablement( ] @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_create_base64_encoded_tarball", return_value="test_tarball", ) - @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_prepare") @pytest.mark.parametrize( "env_vars,expected_env_vars", [ @@ -965,7 +965,7 @@ def test_agent_engine_adk_telemetry_enablement_through_source_packages( test_file_path = os.path.join(tmpdir, "test_file.txt") with open(test_file_path, "w") as f: f.write("test content") - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -982,9 +982,9 @@ def test_agent_engine_adk_telemetry_enablement_through_source_packages( {"name": key, "value": value} for key, value in expected_env_vars.items() ] - @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_prepare") def test_create_agent_engine_config_full(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, @@ -1047,7 +1047,7 @@ def test_create_agent_engine_config_full(self, mock_prepare): ) @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_create_base64_encoded_tarball", return_value="test_tarball", ) @@ -1062,7 +1062,7 @@ def test_create_agent_engine_config_with_source_packages( with open(requirements_file_path, "w") as f: f.write("requests==2.0.0") - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1107,7 +1107,7 @@ def test_create_agent_engine_config_with_developer_connect_source(self): "revision": "main", "dir": "agent", } - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1140,7 +1140,7 @@ 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", ) @@ -1151,7 +1151,7 @@ def test_create_agent_engine_config_with_empty_keep_alive_probe( 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, @@ -1169,7 +1169,7 @@ def test_create_agent_engine_config_with_agent_config_source_and_requirements_fi with open(requirements_file_path, "w") as f: f.write("requests==2.0.0") - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1194,7 +1194,7 @@ def test_create_agent_engine_config_with_agent_config_source_and_entrypoint_modu ): caplog.set_level(logging.WARNING, logger="vertexai_genai.agentengines") - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1219,7 +1219,7 @@ 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", ) @@ -1235,7 +1235,7 @@ def test_create_agent_engine_config_with_source_packages_and_image_spec_raises( f.write("requests==2.0.0") with pytest.raises(ValueError) as excinfo: - self.client.agent_engines._create_config( + self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1252,7 +1252,7 @@ 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", ) @@ -1268,7 +1268,7 @@ def test_create_agent_engine_config_with_agent_config_source_and_image_spec_rais f.write("requests==2.0.0") with pytest.raises(ValueError) as excinfo: - self.client.agent_engines._create_config( + self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1282,7 +1282,7 @@ 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( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1306,7 +1306,7 @@ 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", ) @@ -1321,7 +1321,7 @@ def test_create_agent_engine_config_with_source_packages_and_agent_config_source with open(requirements_file_path, "w") as f: f.write("requests==2.0.0") - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1353,7 +1353,7 @@ def test_create_agent_engine_config_with_source_packages_and_agent_config_source def test_create_agent_engine_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, @@ -1375,7 +1375,7 @@ def test_create_agent_engine_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, @@ -1400,7 +1400,7 @@ def test_create_agent_engine_config_with_container_spec_and_keep_alive_probe( def test_create_agent_engine_config_with_container_spec_and_others_raises(self): container_spec = {"image_uri": "gcr.io/test-project/test-image"} with pytest.raises(ValueError) as excinfo: - self.client.agent_engines._create_config( + self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1410,7 +1410,7 @@ def test_create_agent_engine_config_with_container_spec_and_others_raises(self): assert "please do not specify `agent`" in str(excinfo.value) with pytest.raises(ValueError) as excinfo: - self.client.agent_engines._create_config( + self.client.runtimes._create_config( mode="create", display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, @@ -1420,11 +1420,11 @@ 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") + @mock.patch.object(_runtimes_utils, "_validate_packages_or_raise") def test_create_agent_engine_config_with_source_packages_and_build_options( self, mock_validate_packages, mock_create_base64_encoded_tarball ): @@ -1438,7 +1438,7 @@ def test_create_agent_engine_config_with_source_packages_and_build_options( source_packages = [test_file_path, "installation_scripts/install.sh"] mock_validate_packages.return_value = source_packages - self.client.agent_engines._create_config( + self.client.runtimes._create_config( mode="create", source_packages=source_packages, entrypoint_module="main", @@ -1451,15 +1451,15 @@ def test_create_agent_engine_config_with_source_packages_and_build_options( build_options=build_options, ) - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_validate_packages_or_raise") + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_validate_packages_or_raise") def test_create_agent_engine_config_with_build_options( self, mock_validate_packages, mock_prepare ): build_options = {"installation_scripts": ["install.sh"]} extra_packages = ["install.sh"] - self.client.agent_engines._create_config( + self.client.runtimes._create_config( mode="create", agent=self.test_agent, staging_bucket=_TEST_STAGING_BUCKET, @@ -1473,9 +1473,9 @@ def test_create_agent_engine_config_with_build_options( build_options=build_options, ) - @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_prepare") def test_update_agent_engine_config_full(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, @@ -1538,9 +1538,9 @@ def test_update_agent_engine_config_full(self, mock_prepare): ] ) - @mock.patch.object(_agent_engines_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_prepare") def test_update_agent_engine_clear_service_account(self, mock_prepare): - config = self.client.agent_engines._create_config( + config = self.client.runtimes._create_config( mode="update", service_account="", identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT, @@ -1559,7 +1559,7 @@ def test_update_agent_engine_clear_service_account(self, mock_prepare): def test_get_agent_operation(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body=json.dumps( @@ -1570,7 +1570,7 @@ def test_get_agent_operation(self): } ), ) - operation = self.client.agent_engines._get_agent_operation( + operation = self.client.runtimes._get_agent_operation( operation_name=_TEST_AGENT_ENGINE_OPERATION_NAME, ) request_mock.assert_called_with( @@ -1579,13 +1579,13 @@ def test_get_agent_operation(self): {"_url": {"operationName": _TEST_AGENT_ENGINE_OPERATION_NAME}}, None, ) - assert isinstance(operation, _genai_types.AgentEngineOperation) + assert isinstance(operation, _genai_types.RuntimeOperation) assert operation.done assert isinstance(operation.response, _genai_types.ReasoningEngine) def test_await_operation(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body=json.dumps( @@ -1596,9 +1596,9 @@ def test_await_operation(self): } ), ) - operation = _agent_engines_utils._await_operation( + operation = _runtimes_utils._await_operation( operation_name=_TEST_AGENT_ENGINE_OPERATION_NAME, - get_operation_fn=self.client.agent_engines._get_agent_operation, + get_operation_fn=self.client.runtimes._get_agent_operation, ) request_mock.assert_called_with( "get", @@ -1606,12 +1606,12 @@ def test_await_operation(self): {"_url": {"operationName": _TEST_AGENT_ENGINE_OPERATION_NAME}}, None, ) - assert isinstance(operation, _genai_types.AgentEngineOperation) + assert isinstance(operation, _genai_types.RuntimeOperation) def test_register_api_methods(self): - agent = self.client.agent_engines._register_api_methods( - agent_engine=_genai_types.AgentEngine( - api_client=self.client.agent_engines._api_client, + agent = self.client.runtimes._register_api_methods( + agent_engine=_genai_types.Runtime( + api_client=self.client.runtimes._api_client, api_resource=_genai_types.ReasoningEngine( spec=_genai_types.ReasoningEngineSpec( class_methods=[ @@ -1627,7 +1627,7 @@ def test_register_api_methods(self): @pytest.mark.usefixtures("caplog") def test_invalid_requirement_warning(self, caplog): - _agent_engines_utils._parse_constraints( + _runtimes_utils._parse_constraints( constraints=["invalid requirement line"], ) assert "Failed to parse constraint" in caplog.text @@ -1638,7 +1638,7 @@ def test_requirements_with_whl_files(self): "/content/wxPython-4.2.3-cp39-cp39-macosx_12_0_x86_64.whl", "https://wxpython.org/Phoenix/snapshot-builds/wxPython-4.2.2-cp38-cp38-macosx_12_0_x86_64.whl", ] - result = _agent_engines_utils._parse_constraints( + result = _runtimes_utils._parse_constraints( constraints=whl_files, ) assert result == { @@ -1650,7 +1650,7 @@ def test_requirements_with_whl_files(self): def test_compare_requirements_with_required_packages(self): requirements = {"requests": "2.0.0"} constraints = ["requests==1.0.0"] - result = _agent_engines_utils._compare_requirements( + result = _runtimes_utils._compare_requirements( requirements=requirements, constraints=constraints, ) @@ -1666,7 +1666,7 @@ def test_compare_requirements_with_required_packages(self): def test_scan_simple_object(self): """Test scanning an object importing a known third-party package.""" fake_obj = _create_fake_object_with_module("requests") - requirements = _agent_engines_utils._scan_requirements( + requirements = _runtimes_utils._scan_requirements( obj=fake_obj, package_distributions=_TEST_PACKAGE_DISTRIBUTIONS, ) @@ -1680,7 +1680,7 @@ def test_scan_simple_object(self): def test_scan_object_with_stdlib_module(self): """Test that stdlib modules are ignored by default.""" fake_obj_stdlib = _create_fake_object_with_module("json") - requirements = _agent_engines_utils._scan_requirements( + requirements = _runtimes_utils._scan_requirements( obj=fake_obj_stdlib, package_distributions=_TEST_PACKAGE_DISTRIBUTIONS, ) @@ -1695,13 +1695,13 @@ def test_scan_object_with_stdlib_module(self): def test_scan_with_default_ignore_modules(self, monkeypatch): """Test implicitly ignoring a module.""" fake_obj = _create_fake_object_with_module("requests") - original_base = _agent_engines_utils._BASE_MODULES + original_base = _runtimes_utils._BASE_MODULES monkeypatch.setattr( - _agent_engines_utils, + _runtimes_utils, "_BASE_MODULES", set(original_base) | {"requests"}, ) - requirements = _agent_engines_utils._scan_requirements( + requirements = _runtimes_utils._scan_requirements( obj=fake_obj, package_distributions=_TEST_PACKAGE_DISTRIBUTIONS, ) @@ -1716,7 +1716,7 @@ def test_scan_with_default_ignore_modules(self, monkeypatch): def test_scan_with_explicit_ignore_modules(self): """Test explicitly ignoring a module.""" fake_obj = _create_fake_object_with_module("requests") - requirements = _agent_engines_utils._scan_requirements( + requirements = _runtimes_utils._scan_requirements( obj=fake_obj, ignore_modules=["requests"], package_distributions=_TEST_PACKAGE_DISTRIBUTIONS, @@ -1765,7 +1765,7 @@ def test_scan_with_explicit_ignore_modules(self): ], ) def test_to_parsed_json(self, obj, expected): - for got, want in zip(_agent_engines_utils._yield_parsed_json(obj), expected): + for got, want in zip(_runtimes_utils._yield_parsed_json(obj), expected): assert got == want # pytest does not allow absl.testing.parameterized.named_parameters. @@ -1802,12 +1802,12 @@ def test_to_parsed_json(self, obj, expected): ], ) def test_yield_parsed_json_from_httpbody(self, obj, expected): - got = list(_agent_engines_utils._yield_parsed_json_from_httpbody(obj)) + got = list(_runtimes_utils._yield_parsed_json_from_httpbody(obj)) assert got == expected def test_yield_parsed_json_from_httpbody_non_json_content_type(self): body = httpbody_pb2.HttpBody(content_type="text/plain", data=b"hello") - assert list(_agent_engines_utils._yield_parsed_json_from_httpbody(body)) == [ + assert list(_runtimes_utils._yield_parsed_json_from_httpbody(body)) == [ body ] @@ -1820,7 +1820,7 @@ def test_create_base64_encoded_tarball(self): origin_dir = os.getcwd() try: os.chdir(tmpdir) - encoded_tarball = _agent_engines_utils._create_base64_encoded_tarball( + encoded_tarball = _runtimes_utils._create_base64_encoded_tarball( source_packages=["test_file.txt"] ) finally: @@ -1843,18 +1843,18 @@ 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_agent_engine") + @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, @@ -1866,7 +1866,7 @@ def test_prepare_with_creds( 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, @@ -1882,12 +1882,12 @@ def test_prepare_with_creds( 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_agent_engine") + @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, @@ -1900,7 +1900,7 @@ def test_prepare_without_creds( 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, @@ -1992,13 +1992,13 @@ def test_get_reasoning_engine_id( ): if expected_exception: with pytest.raises(expected_exception) as excinfo: - _agent_engines_utils._get_reasoning_engine_id( + _runtimes_utils._get_reasoning_engine_id( operation_name=operation_name, resource_name=resource_name ) assert expected_message in str(excinfo.value) else: assert ( - _agent_engines_utils._get_reasoning_engine_id( + _runtimes_utils._get_reasoning_engine_id( operation_name=operation_name, resource_name=resource_name ) == expected_id @@ -2006,7 +2006,7 @@ def test_get_reasoning_engine_id( @pytest.mark.usefixtures("google_auth_mock") -class TestAgentEngine: +class TestRuntime: def setup_method(self): importlib.reload(initializer) importlib.reload(aiplatform) @@ -2025,10 +2025,10 @@ def teardown_method(self): def test_get_agent_engine(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, @@ -2038,11 +2038,11 @@ def test_get_agent_engine(self): def test_list_agent_engine(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") expected_query_params = {"filter": _TEST_AGENT_ENGINE_LIST_FILTER} - list(self.client.agent_engines.list(config=expected_query_params)) + list(self.client.runtimes.list(config=expected_query_params)) request_mock.assert_called_with( "get", f"reasoningEngines?{urlencode(expected_query_params)}", @@ -2051,10 +2051,10 @@ 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, ) @@ -2065,7 +2065,7 @@ def test_create_agent_engine( mock_prepare, caplog, ): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, @@ -2073,12 +2073,12 @@ def test_create_agent_engine( ) caplog.set_level(logging.INFO, logger="vertexai_genai.agentengines") with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, @@ -2106,16 +2106,16 @@ 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"agent_engine=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, ) @@ -2125,22 +2125,22 @@ def test_create_agent_engine_lightweight( 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.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, ) @@ -2155,10 +2155,10 @@ 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, ) @@ -2181,19 +2181,19 @@ def test_create_agent_engine_with_env_vars_dict( "agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK, }, } - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], @@ -2258,10 +2258,10 @@ 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, ) @@ -2286,19 +2286,19 @@ def test_create_agent_engine_with_custom_service_account( "agent_framework": _TEST_AGENT_ENGINE_FRAMEWORK, }, } - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], @@ -2364,10 +2364,10 @@ 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, ) @@ -2392,19 +2392,19 @@ def test_create_agent_engine_with_experimental_mode( "class_methods": [_TEST_AGENT_ENGINE_CLASS_METHOD_1], }, } - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], @@ -2470,13 +2470,13 @@ 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, ) @@ -2486,7 +2486,7 @@ def test_create_agent_engine_with_source_packages( mock_await_operation, mock_create_base64_encoded_tarball, ): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, @@ -2501,11 +2501,11 @@ def test_create_agent_engine_with_source_packages( f.write("requests==2.0.0") with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( - config=_genai_types.AgentEngineConfig( + self.client.runtimes.create( + config=_genai_types.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, source_packages=[test_file_path], @@ -2541,10 +2541,10 @@ 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, ) @@ -2566,19 +2566,19 @@ def test_create_agent_engine_with_class_methods( "class_methods": _TEST_AGENT_ENGINE_CLASS_METHODS, }, } - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], @@ -2640,10 +2640,10 @@ 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, ) @@ -2666,19 +2666,19 @@ def test_create_agent_engine_with_agent_framework( "agent_framework": _TEST_AGENT_FRAMEWORK, }, } - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], @@ -2742,12 +2742,12 @@ def test_create_agent_engine_with_agent_framework( ) @pytest.mark.usefixtures("caplog") - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") def test_update_agent_engine_requirements( self, mock_await_operation, mock_prepare, caplog ): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, @@ -2755,13 +2755,13 @@ def test_update_agent_engine_requirements( ) caplog.set_level(logging.INFO, logger="vertexai_genai.agentengines") with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.update( + self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.RuntimeConfig( staging_bucket=_TEST_STAGING_BUCKET, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, ), @@ -2793,31 +2793,31 @@ def test_update_agent_engine_requirements( }, None, ) - assert "Agent Engine updated. To use it in another session:" in caplog.text + assert "Agent Runtime updated. To use it in another session:" in caplog.text assert ( - f"agent_engine=client.agent_engines.get(" + f"agent_engine=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") + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") def test_update_agent_engine_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.RuntimeConfig( staging_bucket=_TEST_STAGING_BUCKET, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH], @@ -2853,38 +2853,38 @@ def test_update_agent_engine_extra_packages( None, ) - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") def test_update_agent_engine_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.RuntimeConfig( env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT ), ) - @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") def test_update_agent_engine_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.RuntimeConfig( staging_bucket=_TEST_STAGING_BUCKET, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT, @@ -2923,25 +2923,25 @@ def test_update_agent_engine_env_vars( None, ) - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") def test_update_agent_engine_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.RuntimeConfig( staging_bucket=_TEST_STAGING_BUCKET, keep_alive_probe={}, ), @@ -2976,11 +2976,11 @@ def test_update_agent_engine_with_empty_keep_alive_probe( None, ) - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_await_operation") def test_update_agent_engine_with_container_spec_and_keep_alive_probe( self, mock_await_operation ): - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, @@ -2988,12 +2988,12 @@ def test_update_agent_engine_with_container_spec_and_keep_alive_probe( ) container_spec = {"image_uri": "gcr.io/test-project/test-image"} with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") - self.client.agent_engines.update( + self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, - config=_genai_types.AgentEngineConfig( + config=_genai_types.RuntimeConfig( container_spec=container_spec, keep_alive_probe=_TEST_AGENT_ENGINE_KEEP_ALIVE_PROBE, class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS, @@ -3026,21 +3026,21 @@ def test_update_agent_engine_with_container_spec_and_keep_alive_probe( None, ) - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_await_operation") def test_update_agent_engine_display_name(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, ) ) 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.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, ), ) @@ -3055,21 +3055,21 @@ def test_update_agent_engine_display_name(self, mock_await_operation): None, ) - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_await_operation") def test_update_agent_engine_description(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, ) ) 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.RuntimeConfig( description=_TEST_AGENT_ENGINE_DESCRIPTION, ), ) @@ -3086,10 +3086,10 @@ def test_update_agent_engine_description(self, mock_await_operation): def test_delete_agent_engine(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, @@ -3099,10 +3099,10 @@ def test_delete_agent_engine(self): def test_delete_agent_engine_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, ) @@ -3134,12 +3134,12 @@ def test_query_agent_engine(self, http_options_arg, expected_http_options): 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( + agent_engine=_genai_types.Runtime( + api_client=self.client.runtimes, api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -3163,11 +3163,11 @@ def test_query_agent_engine(self, http_options_arg, expected_http_options): ) @mock.patch("google.cloud.storage.Client") - @mock.patch.object(agent_engines.AgentEngines, "_get") + @mock.patch.object(runtimes.Runtimes, "_get") @mock.patch("uuid.uuid4") def test_run_query_job_agent_engine(self, mock_uuid, get_mock, mock_storage_client): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body='{"name": "projects/123/locations/us-central1/reasoningEngines/456/operations/789"}' @@ -3194,7 +3194,7 @@ def test_run_query_job_agent_engine(self, mock_uuid, get_mock, mock_storage_clie ), ) - result = self.client.agent_engines.run_query_job( + result = self.client.runtimes.run_query_job( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config={ "query": _TEST_QUERY_PROMPT, @@ -3228,7 +3228,7 @@ def test_run_query_job_agent_engine_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/"}, ) @@ -3237,19 +3237,19 @@ def test_run_query_job_agent_engine_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( self, mock_uuid, get_mock, mock_storage_client ): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body='{"name": "projects/123/locations/us-central1/reasoningEngines/456/operations/789"}' @@ -3279,7 +3279,7 @@ def test_run_query_job_agent_engine_bucket_creation_forbidden( with pytest.raises( ValueError, match="Permission denied to check existence of bucket" ): - self.client.agent_engines.run_query_job( + self.client.runtimes.run_query_job( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config={ "query": _TEST_QUERY_PROMPT, @@ -3288,13 +3288,13 @@ def test_run_query_job_agent_engine_bucket_creation_forbidden( ) @mock.patch("google.cloud.storage.Client") - @mock.patch.object(agent_engines.AgentEngines, "_get") + @mock.patch.object(runtimes.Runtimes, "_get") @mock.patch("uuid.uuid4") def test_run_query_job_agent_engine_file_uri( self, mock_uuid, get_mock, mock_storage_client ): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body='{"name": "projects/123/locations/us-central1/reasoningEngines/456/operations/789"}' @@ -3315,7 +3315,7 @@ def test_run_query_job_agent_engine_file_uri( ), ) - result = self.client.agent_engines.run_query_job( + result = self.client.runtimes.run_query_job( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config={ "query": _TEST_QUERY_PROMPT, @@ -3333,13 +3333,13 @@ def test_run_query_job_agent_engine_file_uri( ) @mock.patch("google.cloud.storage.Client") - @mock.patch.object(agent_engines.AgentEngines, "_get") + @mock.patch.object(runtimes.Runtimes, "_get") @mock.patch("uuid.uuid4") def test_run_query_job_agent_engine_directory_no_slash( self, mock_uuid, get_mock, mock_storage_client ): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( body='{"name": "projects/123/locations/us-central1/reasoningEngines/456/operations/789"}' @@ -3362,7 +3362,7 @@ def test_run_query_job_agent_engine_directory_no_slash( ), ) - result = self.client.agent_engines.run_query_job( + result = self.client.runtimes.run_query_job( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, config={ "query": _TEST_QUERY_PROMPT, @@ -3400,10 +3400,10 @@ def test_query_agent_engine_async(self, http_options_arg, expected_http_options) 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( + agent_engine=_genai_types.Runtime( + api_async_client=runtimes.AsyncRuntimes( + api_client_=self.client.runtimes._api_client ), api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, @@ -3416,7 +3416,7 @@ def test_query_agent_engine_async(self, http_options_arg, expected_http_options) ) ) with mock.patch.object( - self.client.agent_engines._api_client, "async_request" + self.client.runtimes._api_client, "async_request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse(body="") asyncio.run(agent.async_query(**kwargs)) @@ -3433,11 +3433,11 @@ def test_query_agent_engine_async(self, http_options_arg, expected_http_options) def test_cancel_query_job_agent_engine(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}, ) @@ -3455,7 +3455,7 @@ def test_cancel_query_job_agent_engine(self): def test_check_query_job_agent_engine(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( headers={}, @@ -3473,7 +3473,7 @@ def test_check_query_job_agent_engine(self): mock_bucket.blob.return_value = mock_blob mock_storage_client.return_value.bucket.return_value = mock_bucket - result = self.client.agent_engines.check_query_job( + result = self.client.runtimes.check_query_job( name="projects/123/locations/us-central1/reasoningEngines/456/operations/789", config={"retrieve_result": True}, ) @@ -3492,7 +3492,7 @@ def test_check_query_job_agent_engine(self): def test_check_query_job_agent_engine_running(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( headers={}, @@ -3503,7 +3503,7 @@ def test_check_query_job_agent_engine_running(self): ), ) - result = self.client.agent_engines.check_query_job( + result = self.client.runtimes.check_query_job( name="projects/123/locations/us-central1/reasoningEngines/456/operations/789", config={"retrieve_result": True}, ) @@ -3517,14 +3517,14 @@ def test_check_query_job_agent_engine_running(self): def test_check_query_job_agent_engine_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}, ) @@ -3538,7 +3538,7 @@ def test_check_query_job_agent_engine_failed(self): def test_check_query_job_agent_engine_no_retrieve(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( headers={}, @@ -3549,7 +3549,7 @@ def test_check_query_job_agent_engine_no_retrieve(self): ), ) - result = self.client.agent_engines.check_query_job( + result = self.client.runtimes.check_query_job( name="projects/123/locations/us-central1/reasoningEngines/456/operations/789", config={"retrieve_result": False}, ) @@ -3563,7 +3563,7 @@ def test_check_query_job_agent_engine_no_retrieve(self): def test_check_query_job_agent_engine_blob_not_exists(self): with mock.patch.object( - self.client.agent_engines._api_client, "request" + self.client.runtimes._api_client, "request" ) as request_mock: request_mock.return_value = genai_types.HttpResponse( headers={}, @@ -3585,7 +3585,7 @@ def test_check_query_job_agent_engine_blob_not_exists(self): ValueError, match="Failed to retrieve blob results for gs://my-output-bucket/output.json", ): - self.client.agent_engines.check_query_job( + self.client.runtimes.check_query_job( name="projects/123/locations/us-central1/reasoningEngines/456/operations/789", config={"retrieve_result": True}, ) @@ -3611,11 +3611,11 @@ def test_query_agent_engine_stream(self, http_options_arg, expected_http_options 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( + agent_engine=_genai_types.Runtime( + api_client=self.client.runtimes, api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -3666,12 +3666,12 @@ async def mock_async_generator(): yield genai_types.HttpResponse(body=b"") with mock.patch.object( - self.client.agent_engines._api_client, "async_request_streamed" + self.client.runtimes._api_client, "async_request_streamed" ) as request_mock: request_mock.return_value = mock_async_generator() - agent = self.client.agent_engines._register_api_methods( - agent_engine=_genai_types.AgentEngine( - api_client=self.client.agent_engines, + agent = self.client.runtimes._register_api_methods( + agent_engine=_genai_types.Runtime( + api_client=self.client.runtimes, api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -3708,7 +3708,7 @@ async def consume(): _TEST_NO_OPERATION_REGISTRABLE_SCHEMAS, [ ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( CapitalizeEngine().query, schema_name=_TEST_DEFAULT_METHOD_NAME, ), @@ -3721,70 +3721,70 @@ async def consume(): _TEST_OPERATION_REGISTRABLE_SCHEMAS, [ ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().query, schema_name=_TEST_DEFAULT_METHOD_NAME, ), _TEST_STANDARD_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_method, schema_name=_TEST_CUSTOM_METHOD_NAME, ), _TEST_STANDARD_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().async_query, schema_name=_TEST_DEFAULT_ASYNC_METHOD_NAME, ), _TEST_ASYNC_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_async_method, schema_name=_TEST_CUSTOM_ASYNC_METHOD_NAME, ), _TEST_ASYNC_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().stream_query, schema_name=_TEST_DEFAULT_STREAM_METHOD_NAME, ), _TEST_STREAM_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_stream_method, schema_name=_TEST_CUSTOM_STREAM_METHOD_NAME, ), _TEST_STREAM_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().async_stream_query, schema_name=_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME, ), _TEST_ASYNC_STREAM_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_async_stream_method, schema_name=_TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME, ), _TEST_ASYNC_STREAM_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().bidi_stream_query, schema_name=_TEST_DEFAULT_BIDI_STREAM_METHOD_NAME, ), _TEST_BIDI_STREAM_API_MODE, ), ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationRegistrableEngine().custom_bidi_stream_method, schema_name=_TEST_CUSTOM_BIDI_STREAM_METHOD_NAME, ), @@ -3797,7 +3797,7 @@ async def consume(): _TEST_OPERATION_NOT_REGISTERED_SCHEMAS, [ ( - _agent_engines_utils._generate_schema( + _runtimes_utils._generate_schema( OperationNotRegisteredEngine().custom_method, schema_name=_TEST_CUSTOM_METHOD_NAME, ), @@ -3808,7 +3808,7 @@ async def consume(): ], ) @mock.patch.object(genai_client.Client, "_get_api_client") - @mock.patch.object(agent_engines.AgentEngines, "_get") + @mock.patch.object(runtimes.Runtimes, "_get") def test_operation_schemas( self, mock_get, @@ -3817,7 +3817,7 @@ def test_operation_schemas( test_class_methods_spec, want_operation_schema_api_modes, ): - test_agent_engine = _genai_types.AgentEngine( + test_agent_engine = _genai_types.Runtime( api_resource=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec( @@ -3831,11 +3831,11 @@ def test_operation_schemas( want_operation_schemas.append(want_operation_schema) assert test_agent_engine.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, ) @@ -3849,15 +3849,15 @@ def test_create_agent_engine_with_creds( mock_operation = mock.Mock() mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME mock_create.return_value = mock_operation - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, ) ) - self.client.agent_engines.create( + self.client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, staging_bucket=_TEST_STAGING_BUCKET, ), @@ -3871,12 +3871,12 @@ 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, ) @@ -3891,7 +3891,7 @@ def test_create_agent_engine_without_creds( mock_operation = mock.Mock() mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME mock_create.return_value = mock_operation - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, @@ -3903,9 +3903,9 @@ def test_create_agent_engine_without_creds( client = agentplatform.Client( project=_TEST_PROJECT, location=_TEST_LOCATION, credentials=mock_creds ) - client.agent_engines.create( + client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, staging_bucket=_TEST_STAGING_BUCKET, ), @@ -3919,11 +3919,11 @@ 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, ) @@ -3937,7 +3937,7 @@ def test_create_agent_engine_with_no_creds_in_client( mock_operation = mock.Mock() mock_operation.name = _TEST_AGENT_ENGINE_OPERATION_NAME mock_create.return_value = mock_operation - mock_await_operation.return_value = _genai_types.AgentEngineOperation( + mock_await_operation.return_value = _genai_types.RuntimeOperation( response=_genai_types.ReasoningEngine( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_TEST_AGENT_ENGINE_SPEC, @@ -3946,9 +3946,9 @@ def test_create_agent_engine_with_no_creds_in_client( client = agentplatform.Client( project=_TEST_PROJECT, location=_TEST_LOCATION, credentials=None ) - client.agent_engines.create( + client.runtimes.create( agent=self.test_agent, - config=_genai_types.AgentEngineConfig( + config=_genai_types.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, staging_bucket=_TEST_STAGING_BUCKET, ), @@ -3964,7 +3964,7 @@ def test_create_agent_engine_with_no_creds_in_client( @pytest.mark.usefixtures("google_auth_mock") -class TestAgentEngineErrors: +class TestRuntimeErrors: def setup_method(self): importlib.reload(initializer) importlib.reload(aiplatform) @@ -3976,27 +3976,27 @@ def setup_method(self): ) self.test_agent = CapitalizeEngine() - @mock.patch.object(_agent_engines_utils, "_prepare") - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_prepare") + @mock.patch.object(_runtimes_utils, "_await_operation") @mock.patch.object( - _agent_engines_utils, + _runtimes_utils, "_get_reasoning_engine_id", return_value=_TEST_RESOURCE_ID, ) def test_create_agent_engine_error( 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.RuntimeConfig( display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME, description=_TEST_AGENT_ENGINE_DESCRIPTION, requirements=_TEST_AGENT_ENGINE_REQUIREMENTS, @@ -4008,19 +4008,19 @@ def test_create_agent_engine_error( ) assert "Failed to create agent engine" in str(excinfo.value) - @mock.patch.object(_agent_engines_utils, "_await_operation") + @mock.patch.object(_runtimes_utils, "_await_operation") def test_update_agent_engine_description(self, mock_await_operation): - 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.update( + self.client.runtimes.update( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, - config=_genai_types.AgentEngineConfig( + config=_genai_types.RuntimeConfig( description=_TEST_AGENT_ENGINE_DESCRIPTION, ), ) @@ -4078,7 +4078,7 @@ def test_update_agent_engine_description(self, mock_await_operation): ], ) @pytest.mark.usefixtures("caplog") - @mock.patch.object(agent_engines.AgentEngines, "_get") + @mock.patch.object(runtimes.Runtimes, "_get") def test_invalid_operation_schema( self, mock_get, @@ -4091,7 +4091,7 @@ def test_invalid_operation_schema( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec=_genai_types.ReasoningEngineSpec(class_methods=test_operation_schemas), ) - self.client.agent_engines.get(name=_TEST_AGENT_ENGINE_RESOURCE_NAME) + self.client.runtimes.get(name=_TEST_AGENT_ENGINE_RESOURCE_NAME) assert want_log_output in caplog.text @pytest.mark.parametrize( @@ -4158,12 +4158,12 @@ def test_validate_resource_limits_or_raise( self, resource_limits, expected_exception, expected_message ): with pytest.raises(expected_exception) as excinfo: - _agent_engines_utils._validate_resource_limits_or_raise(resource_limits) + _runtimes_utils._validate_resource_limits_or_raise(resource_limits) assert expected_message in str(excinfo.value) @pytest.mark.usefixtures("google_auth_mock") -class TestAsyncAgentEngine: +class TestAsyncRuntime: def setup_method(self): importlib.reload(initializer) importlib.reload(aiplatform) @@ -4182,11 +4182,11 @@ def teardown_method(self): def test_delete_agent_engine(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 ) ) @@ -4199,11 +4199,11 @@ def test_delete_agent_engine(self): def test_delete_agent_engine_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 b93b6db39b..92821ec1fe 100644 --- a/tests/unit/agentplatform/genai/test_evals.py +++ b/tests/unit/agentplatform/genai/test_evals.py @@ -3349,7 +3349,7 @@ 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_agentplatform_client.return_value.runtimes.get.return_value = ( mock_agent_engine ) @@ -3359,7 +3359,7 @@ 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( @@ -3458,7 +3458,7 @@ 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_agentplatform_client.return_value.runtimes.get.return_value = ( mock_agent_engine ) @@ -3468,7 +3468,7 @@ 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( @@ -3550,7 +3550,7 @@ def test_run_inference_with_agent_engine_with_response_column_raises_error( ) mock_agent_engine = mock.Mock() - mock_agentplatform_client.return_value.agent_engines.get.return_value = ( + mock_agentplatform_client.return_value.runtimes.get.return_value = ( mock_agent_engine ) @@ -3622,7 +3622,7 @@ def test_run_inference_with_agent_engine_falls_back_to_managed_sessions_api( }, ] mock_agent_engine.stream_query.return_value = iter(stream_query_return_value) - mock_agentplatform_client.return_value.agent_engines.get.return_value = ( + mock_agentplatform_client.return_value.runtimes.get.return_value = ( mock_agent_engine ) @@ -3635,7 +3635,7 @@ def test_run_inference_with_agent_engine_falls_back_to_managed_sessions_api( mock_agent_engine.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"}, ), ) @@ -4301,7 +4301,7 @@ def test_run_inference_non_gemini_agent_routes_to_agent_engine( } ] ) - mock_agentplatform_client.return_value.agent_engines.get.return_value = ( + mock_agentplatform_client.return_value.runtimes.get.return_value = ( mock_agent_engine ) @@ -4310,7 +4310,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() diff --git a/tests/unit/agentplatform/genai/test_live_agent_engines.py b/tests/unit/agentplatform/genai/test_live_agent_engines.py index 8af10749a3..2dba7590b7 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,7 +67,7 @@ async def test_async_live_agent_engines_connect( json.dumps({"output": "WORLD"}).encode("utf-8"), ] - async with test_client.aio.live.agent_engines.connect( + async with test_client.aio.live.runtimes.connect( agent_engine="test-agent-engine", config={"class_method": "bidi_stream_query", "input": {"query": "hello"}}, ) as session: diff --git a/tests/unit/agentplatform/genai/test_sandbox.py b/tests/unit/agentplatform/genai/test_sandbox.py index 8f3f331d2c..32a4315c50 100644 --- a/tests/unit/agentplatform/genai/test_sandbox.py +++ b/tests/unit/agentplatform/genai/test_sandbox.py @@ -84,7 +84,7 @@ def test_send_command(self, mock_get_api_client): body=b"{}", headers={} ) - self.client.agent_engines.sandboxes.send_command( + self.client.sandboxes.send_command( http_method="GET", access_token="test_token", sandbox_environment=mock_sandbox, @@ -120,7 +120,7 @@ def test_generate_browser_ws_headers( body=b'{"endpoint": "test/endpoint"}', headers={} ) ws_url, headers = ( - self.client.agent_engines.sandboxes.generate_browser_ws_headers( + self.client.sandboxes.generate_browser_ws_headers( sandbox_environment=mock_sandbox, service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL, timeout=3600,