airflow.providers.google.cloud.hooks.vertex_ai.agent_engine¶

This module contains a Google Cloud Vertex AI Agent Engine hook.

Attributes¶

`VERTEX_AI_AGENT_ENGINE_API_VERSION`
`VERTEX_AI_AGENT_ENGINE_OPERATION_URL`
`DEFAULT_AGENT_ENGINE_OPERATION_REQUEST_TIMEOUT`

Classes¶

`AgentEngineHook`	Hook for Google Cloud Vertex AI Agent Engine APIs.
`AgentEngineAsyncHook`	Async hook for Google Cloud Vertex AI Agent Engine APIs.

Functions¶

`extract_operation_id`(operation_name)	Extract the operation ID from a fully qualified operation name.
`serialize_value`(value)	Recursively convert SDK model objects to JSON-serializable types.

Module Contents¶

airflow.providers.google.cloud.hooks.vertex_ai.agent_engine.VERTEX_AI_AGENT_ENGINE_API_VERSION = 'v1beta1'[source]¶

airflow.providers.google.cloud.hooks.vertex_ai.agent_engine.VERTEX_AI_AGENT_ENGINE_OPERATION_URL = 'https://{location}-aiplatform.googleapis.com/{api_version}/{operation_name}'[source]¶

airflow.providers.google.cloud.hooks.vertex_ai.agent_engine.DEFAULT_AGENT_ENGINE_OPERATION_REQUEST_TIMEOUT = 60.0[source]¶

airflow.providers.google.cloud.hooks.vertex_ai.agent_engine.extract_operation_id(operation_name)[source]¶

Extract the operation ID from a fully qualified operation name.

airflow.providers.google.cloud.hooks.vertex_ai.agent_engine.serialize_value(value)[source]¶

Recursively convert SDK model objects to JSON-serializable types.

class airflow.providers.google.cloud.hooks.vertex_ai.agent_engine.AgentEngineHook(gcp_conn_id='google_cloud_default', impersonation_chain=None, **kwargs)[source]¶

Bases: airflow.providers.google.common.hooks.base_google.GoogleBaseHook

Hook for Google Cloud Vertex AI Agent Engine APIs.

Wraps the agent_engines module of the Vertex AI SDK client and the Reasoning Engine Execution Service GAPIC client: https://docs.cloud.google.com/python/docs/reference/agentplatform/latest/vertexai._genai.agent_engines.AgentEngines

get_agent_engine_client(project_id, location)[source]¶

Return the Vertex AI Agent Engine client.

get_reasoning_engine_execution_service_client(location=None)[source]¶

Return the Reasoning Engine Execution Service client.

static build_agent_engine_name(project_id, location, agent_engine_id)[source]¶

Build a fully qualified Agent Engine resource name.

static build_operation_name(project_id, location, operation_id)[source]¶

Build a fully qualified Agent Engine operation name.

create_agent_engine(location, agent=None, config=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Create an Agent Engine.

Parameters:

location (str) – Required. The ID of the Google Cloud location that the service belongs to.
agent (Any | None) – Optional. The agent object to deploy.
config (vertexai._genai.types.AgentEngineConfigOrDict | None) – Optional. Configuration for the Agent Engine.
project_id (str) – Optional. The ID of the Google Cloud project. Defaults to the project configured in the connection.

get_agent_engine(location, agent_engine_id, config=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Get an Agent Engine.

Parameters:

location (str) – Required. The ID of the Google Cloud location that the service belongs to.
agent_engine_id (str) – Required. The Agent Engine ID.
config (vertexai._genai.types.GetAgentEngineConfigOrDict | None) – Optional. Configuration for getting the Agent Engine.
project_id (str) – Optional. The ID of the Google Cloud project. Defaults to the project configured in the connection.

query_reasoning_engine(location, reasoning_engine_id, input_data=None, class_method='query', retry=None, timeout=None, metadata=(), project_id=PROVIDE_PROJECT_ID)[source]¶

Query a Reasoning Engine synchronously.

Parameters:

location (str) – Required. The ID of the Google Cloud location that the service belongs to.
reasoning_engine_id (str) – Required. The Reasoning Engine resource ID.
input_data (dict[str, Any] | None) – Optional. Input for the Reasoning Engine class method in JSON object format. Defaults to None.
class_method (str) – Optional. The Reasoning Engine class method to invoke. Defaults to query.
retry (google.api_core.retry.Retry | None) – Designation of what errors, if any, should be retried. Defaults to None.
timeout (float | None) – The timeout for this request. Defaults to None.
metadata (collections.abc.Sequence[tuple[str, str]]) – Strings which should be sent along with the request as metadata. Defaults to an empty tuple.
project_id (str) – Optional. The ID of the Google Cloud project. Defaults to the project configured in the connection.

run_query_job(location, agent_engine_id, config=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Run a query job on an Agent Engine.

Parameters:

location (str) – Required. The ID of the Google Cloud location that the service belongs to.
agent_engine_id (str) – Required. The Agent Engine ID.
config (vertexai._genai.types.RunQueryJobAgentEngineConfigOrDict | None) – Optional. Configuration for the query job (query, output_gcs_uri).
project_id (str) – Optional. The ID of the Google Cloud project. Defaults to the project configured in the connection.

check_query_agent_engine_job(location, operation_id, config=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Check a query job on an Agent Engine.

Parameters:

location (str) – Required. The ID of the Google Cloud location that the service belongs to.
operation_id (str) – Required. The query job operation ID.
config (vertexai._genai.types.CheckQueryJobAgentEngineConfigOrDict | None) – Optional. Configuration for checking the query job.
project_id (str) – Optional. The ID of the Google Cloud project. Defaults to the project configured in the connection.

wait_for_query_agent_engine_job(location, operation_id, config=None, poll_interval=30, timeout=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Wait until an Agent Engine query job completes.

Parameters:

location (str) – Required. The ID of the Google Cloud location that the service belongs to.
operation_id (str) – Required. The query job operation ID.
config (vertexai._genai.types.CheckQueryJobAgentEngineConfigOrDict | None) – Optional. Configuration for checking the query job.
poll_interval (float) – Time, in seconds, to wait between checks.
timeout (float | None) – Optional timeout, in seconds.
project_id (str) – Optional. The ID of the Google Cloud project. Defaults to the project configured in the connection.

update_agent_engine(location, agent_engine_id, config, agent=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Update an Agent Engine.

Parameters:

location (str) – Required. The ID of the Google Cloud location that the service belongs to.
agent_engine_id (str) – Required. The Agent Engine ID.
config (vertexai._genai.types.AgentEngineConfigOrDict) – Required. Configuration for the Agent Engine update.
agent (Any | None) – Optional. The updated agent object to deploy.
project_id (str) – Optional. The ID of the Google Cloud project. Defaults to the project configured in the connection.

delete_agent_engine(location, agent_engine_id, force=None, config=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Delete an Agent Engine.

Parameters:

location (str) – Required. The ID of the Google Cloud location that the service belongs to.
agent_engine_id (str) – Required. The Agent Engine ID.
force (bool | None) – Optional. Whether to forcefully delete child resources. Defaults to False when not specified.
config (vertexai._genai.types.DeleteAgentEngineConfigOrDict | None) – Optional. Additional deletion configuration.
project_id (str) – Optional. The ID of the Google Cloud project. Defaults to the project configured in the connection.

get_agent_engine_operation(location, operation_id, request_timeout=DEFAULT_AGENT_ENGINE_OPERATION_REQUEST_TIMEOUT, project_id=PROVIDE_PROJECT_ID)[source]¶

Return a Vertex AI Agent Engine long-running operation.

Parameters:

location (str) – The ID of the Google Cloud location that the service belongs to.
operation_id (str) – The Agent Engine operation ID.
request_timeout (float | None) – Optional timeout, in seconds, for the operation request.
project_id (str) – Optional. The ID of the Google Cloud project. Defaults to the project configured in the connection.

wait_for_agent_engine_operation(location, operation_id, poll_interval=30, timeout=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Wait until an Agent Engine operation completes.

Parameters:

location (str) – The ID of the Google Cloud location that the service belongs to.
operation_id (str) – The Agent Engine operation ID.
poll_interval (float) – Time, in seconds, to wait between checks.
timeout (float | None) – Optional timeout, in seconds.
project_id (str) – Optional. The ID of the Google Cloud project. Defaults to the project configured in the connection.

class airflow.providers.google.cloud.hooks.vertex_ai.agent_engine.AgentEngineAsyncHook(gcp_conn_id='google_cloud_default', impersonation_chain=None, **kwargs)[source]¶

Bases: airflow.providers.google.common.hooks.base_google.GoogleBaseAsyncHook

Async hook for Google Cloud Vertex AI Agent Engine APIs.

sync_hook_class[source]¶

async check_query_agent_engine_job(location, operation_id, config=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Check a query job on an Agent Engine.