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,71 @@ 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
+^^^^^^^^^^^^^^^^^^^^^^
+
+The operators below manage `Vertex AI Agent Engine
+<https://docs.cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview>`__ resources.
+
+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.RunQueryJobOperator`.
+The operator uses the public ``run_query_job`` SDK method. The ``config`` parameter
+can include ``query`` and ``output_gcs_uri``. The SDK writes query input and output
+through Google Cloud Storage. By default, the operator waits for the query job to
+complete and returns the serialized query job result. Set ``retrieve_result`` to
+``True`` in ``check_config`` to return the query job result from Google Cloud Storage.
+
+.. 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_run_query_job_operator]
+    :end-before: [END how_to_cloud_vertex_ai_run_query_job_operator]
+
+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_run_query_job_operator_deferrable]
+    :end-before: [END how_to_cloud_vertex_ai_run_query_job_operator_deferrable]
+
+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.
+
+.. 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

@@ -615,6 +615,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
@@ -901,6 +902,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