Add Vertex AI Agent Engine operators (Create, Get, Query, Update, Delete)

providers/google/README.rst

@@ -76,6 +76,7 @@ PIP package                                 Version required
 ``google-api-python-client``                ``>=2.0.2``
 ``google-auth``                             ``>=2.29.0``
 ``google-auth-httplib2``                    ``>=0.0.1``
+``google-genai``                            ``>=2.8.0``
 ``google-cloud-aiplatform[evaluation]``     ``>=1.155.0``
 ``ray[default]``                            ``>=2.42.0; python_version < "3.13"``
 ``ray[default]``                            ``>=2.49.0; python_version >= "3.13" and python_version < "3.14"``

providers/google/docs/index.rst

@@ -129,6 +129,7 @@ PIP package                                 Version required
 ``google-api-python-client``                ``>=2.0.2``
 ``google-auth``                             ``>=2.29.0``
 ``google-auth-httplib2``                    ``>=0.0.1``
+``google-genai``                            ``>=2.8.0``
 ``google-cloud-aiplatform[evaluation]``     ``>=1.155.0``
 ``ray[default]``                            ``>=2.42.0; python_version < "3.13"``
 ``ray[default]``                            ``>=2.49.0; python_version >= "3.13" and python_version < "3.14"``

providers/google/docs/operators/cloud/vertex_ai.rst

@@ -26,6 +26,74 @@ With Vertex AI, both AutoML training and custom training are available options.
 Whichever option you choose for training, you can save models, deploy models, and
 request predictions with Vertex AI.
 
+Managing Agent Engines
+^^^^^^^^^^^^^^^^^^^^^^
+
+To create a Vertex AI Agent Engine you can use
+:class:`~airflow.providers.google.cloud.operators.vertex_ai.agent_engine.CreateAgentEngineOperator`.
+
+.. exampleinclude:: /../../google/tests/system/google/cloud/vertex_ai/example_vertex_ai_agent_engine.py
+    :language: python
+    :dedent: 4
+    :start-after: [START how_to_cloud_vertex_ai_create_agent_engine_operator]
+    :end-before: [END how_to_cloud_vertex_ai_create_agent_engine_operator]
+
+To get an Agent Engine you can use
+:class:`~airflow.providers.google.cloud.operators.vertex_ai.agent_engine.GetAgentEngineOperator`.
+
+.. exampleinclude:: /../../google/tests/system/google/cloud/vertex_ai/example_vertex_ai_agent_engine.py
+    :language: python
+    :dedent: 4
+    :start-after: [START how_to_cloud_vertex_ai_get_agent_engine_operator]
+    :end-before: [END how_to_cloud_vertex_ai_get_agent_engine_operator]
+
+To run a query job on an Agent Engine you can use
+:class:`~airflow.providers.google.cloud.operators.vertex_ai.agent_engine.QueryAgentEngineOperator`.
+The operator uses the public ``run_query_job`` SDK method. The ``config`` parameter
+is required and must include ``query`` and ``output_gcs_uri``. The SDK writes query
+input and output through Google Cloud Storage and the operator returns the serialized
+query job result metadata, including the job name and GCS URIs. This operator starts
+the query job and does not wait for the query job output.
+
+.. exampleinclude:: /../../google/tests/system/google/cloud/vertex_ai/example_vertex_ai_agent_engine.py
+    :language: python
+    :dedent: 4
+    :start-after: [START how_to_cloud_vertex_ai_query_agent_engine_operator]
+    :end-before: [END how_to_cloud_vertex_ai_query_agent_engine_operator]
+
+To wait for a query job and retrieve its output you can use
+:class:`~airflow.providers.google.cloud.operators.vertex_ai.agent_engine.CheckQueryAgentEngineOperator`.
+The operator uses the public ``check_query_job`` SDK method. The ``config`` parameter
+is optional. Set ``retrieve_result`` to ``True`` in ``config`` to return the query job
+result from Google Cloud Storage. The same operation can be performed in the
+deferrable mode.
+
+.. exampleinclude:: /../../google/tests/system/google/cloud/vertex_ai/example_vertex_ai_agent_engine.py
+    :language: python
+    :dedent: 4
+    :start-after: [START how_to_cloud_vertex_ai_check_query_agent_engine_operator]
+    :end-before: [END how_to_cloud_vertex_ai_check_query_agent_engine_operator]
+
+To update an Agent Engine you can use
+:class:`~airflow.providers.google.cloud.operators.vertex_ai.agent_engine.UpdateAgentEngineOperator`.
+
+.. exampleinclude:: /../../google/tests/system/google/cloud/vertex_ai/example_vertex_ai_agent_engine.py
+    :language: python
+    :dedent: 4
+    :start-after: [START how_to_cloud_vertex_ai_update_agent_engine_operator]
+    :end-before: [END how_to_cloud_vertex_ai_update_agent_engine_operator]
+
+To delete an Agent Engine you can use
+:class:`~airflow.providers.google.cloud.operators.vertex_ai.agent_engine.DeleteAgentEngineOperator`.
+By default, the operator waits until the delete operation completes. The same operation
+can be performed in the deferrable mode.
+
+.. exampleinclude:: /../../google/tests/system/google/cloud/vertex_ai/example_vertex_ai_agent_engine.py
+    :language: python
+    :dedent: 4
+    :start-after: [START how_to_cloud_vertex_ai_delete_agent_engine_operator]
+    :end-before: [END how_to_cloud_vertex_ai_delete_agent_engine_operator]
+
 Creating Datasets
 ^^^^^^^^^^^^^^^^^
 

providers/google/provider.yaml

@@ -614,6 +614,7 @@ operators:
       - airflow.providers.google.cloud.operators.vertex_ai.generative_model
       - airflow.providers.google.cloud.operators.vertex_ai.feature_store
       - airflow.providers.google.cloud.operators.vertex_ai.ray
+      - airflow.providers.google.cloud.operators.vertex_ai.agent_engine
   - integration-name: Google Data Studio
     python-modules:
       - airflow.providers.google.cloud.operators.looker
@@ -900,6 +901,7 @@ hooks:
       - airflow.providers.google.cloud.hooks.vertex_ai.generative_model
       - airflow.providers.google.cloud.hooks.vertex_ai.prediction_service
       - airflow.providers.google.cloud.hooks.vertex_ai.feature_store
+      - airflow.providers.google.cloud.hooks.vertex_ai.agent_engine
       - airflow.providers.google.cloud.hooks.vertex_ai.ray
   - integration-name: Google Data Studio
     python-modules:

providers/google/pyproject.toml

@@ -80,6 +80,7 @@ dependencies = [
     "google-api-python-client>=2.0.2",
     "google-auth>=2.29.0",
     "google-auth-httplib2>=0.0.1",
+    "google-genai>=2.8.0",
     # google-cloud-aiplatform doesn't install ray for python 3.12 (issue: https://github.com/googleapis/python-aiplatform/issues/5252).
     # Temporarily lock in ray 2.42.0 which is compatible with python 3.12 until linked issue is solved.
     # Remove the ray dependency as well as google-cloud-bigquery-storage once linked issue is fixed