This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 08974228440 Add synchronous Vertex AI Agent Engine query operator
(#70933)
08974228440 is described below
commit 089742284403a93106e4ad0ca08bed1c8707be2c
Author: Alejandro Morgante <[email protected]>
AuthorDate: Thu Aug 13 15:40:54 2026 -0300
Add synchronous Vertex AI Agent Engine query operator (#70933)
* Add synchronous Vertex AI Agent Engine queries
Direct Agent Engine invocations currently require a Google Cloud
Storage-backed query job in the provider. A synchronous public GAPIC path
supports request-response workflows without that intermediate storage layer.
* Handle omitted Agent Engine query input
* Cover custom Agent Engine query options in hook test
The public hook exposes configurable request controls, so its unit test
should detect regressions that replace caller-supplied values with defaults.
* Address Agent Engine query review feedback
* Return complete Agent Engine query responses
* Align synchronous Agent Engine query names with GAPIC
Users familiar with ReasoningEngineExecutionService should be able to find
the equivalent Airflow API without translating between resource vocabularies.
* Clarify synchronous Vertex AI Reasoning Engine terminology
Matching the GAPIC resource vocabulary makes the operator easier for users
to discover and understand.
* Preserve Agent Engine query retry behavior
The GAPIC retry sentinel type is private and cannot be relied upon by
provider code. This RPC has no default retry policy, so an explicit None keeps
existing behavior while leaving custom Retry policies available.
---
.../google/docs/operators/cloud/vertex_ai.rst | 11 +++
.../google/cloud/hooks/vertex_ai/agent_engine.py | 63 +++++++++++++++-
.../cloud/operators/vertex_ai/agent_engine.py | 87 ++++++++++++++++++++++
.../vertex_ai/example_vertex_ai_agent_engine.py | 13 ++++
.../cloud/hooks/vertex_ai/test_agent_engine.py | 87 ++++++++++++++++++++++
.../cloud/operators/vertex_ai/test_agent_engine.py | 63 ++++++++++++++++
6 files changed, 323 insertions(+), 1 deletion(-)
diff --git a/providers/google/docs/operators/cloud/vertex_ai.rst
b/providers/google/docs/operators/cloud/vertex_ai.rst
index f11b8787ef6..dd69e361de2 100644
--- a/providers/google/docs/operators/cloud/vertex_ai.rst
+++ b/providers/google/docs/operators/cloud/vertex_ai.rst
@@ -50,6 +50,17 @@ To get an Agent Engine you can use
: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 query a Reasoning Engine synchronously you can use
+:class:`~airflow.providers.google.cloud.operators.vertex_ai.agent_engine.RunReasoningEngineQueryOperator`.
+The operator calls the public ``query_reasoning_engine`` GAPIC method and
returns
+the serialized response without using 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_reasoning_engine_query_operator]
+ :end-before: [END
how_to_cloud_vertex_ai_run_reasoning_engine_query_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
diff --git
a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py
b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py
index 5471572d3ce..ab6c130df7d 100644
---
a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py
+++
b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py
@@ -25,8 +25,10 @@ from typing import TYPE_CHECKING, Any
import google.auth.transport.requests
from asgiref.sync import sync_to_async
+from google.cloud.aiplatform_v1 import ReasoningEngineExecutionServiceClient
from vertexai import Client
+from airflow.providers.google.common.consts import CLIENT_INFO
from airflow.providers.google.common.hooks.base_google import (
PROVIDE_PROJECT_ID,
GoogleBaseAsyncHook,
@@ -34,6 +36,8 @@ from airflow.providers.google.common.hooks.base_google import
(
)
if TYPE_CHECKING:
+ from google.api_core.retry import Retry
+ from google.cloud.aiplatform_v1.types import QueryReasoningEngineResponse
from vertexai._genai import types
@@ -66,7 +70,8 @@ class AgentEngineHook(GoogleBaseHook):
"""
Hook for Google Cloud Vertex AI Agent Engine APIs.
- Wraps the ``agent_engines`` module of the Vertex AI SDK client:
+ 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
"""
@@ -90,6 +95,23 @@ class AgentEngineHook(GoogleBaseHook):
credentials=self.get_credentials(),
).agent_engines
+ def _get_api_endpoint(self, location: str | None = None) -> str | None:
+ if location and location != "global" and self.is_default_universe():
+ return f"{location}-aiplatform.googleapis.com:443"
+ return None
+
+ def get_reasoning_engine_execution_service_client(
+ self, location: str | None = None
+ ) -> ReasoningEngineExecutionServiceClient:
+ """Return the Reasoning Engine Execution Service client."""
+ return ReasoningEngineExecutionServiceClient(
+ credentials=self.get_credentials(),
+ client_info=CLIENT_INFO,
+ client_options=self.get_client_options(
+ api_endpoint_override=self._get_api_endpoint(location=location)
+ ),
+ )
+
@staticmethod
def build_agent_engine_name(project_id: str, location: str,
agent_engine_id: str) -> str:
"""Build a fully qualified Agent Engine resource name."""
@@ -141,6 +163,45 @@ class AgentEngineHook(GoogleBaseHook):
name = self.build_agent_engine_name(project_id, location,
agent_engine_id)
return client.get(name=name, config=config)
+ @GoogleBaseHook.fallback_to_default_project_id
+ def query_reasoning_engine(
+ self,
+ location: str,
+ reasoning_engine_id: str,
+ input_data: dict[str, Any] | None = None,
+ class_method: str = "query",
+ retry: Retry | None = None,
+ timeout: float | None = None,
+ metadata: Sequence[tuple[str, str]] = (),
+ project_id: str = PROVIDE_PROJECT_ID,
+ ) -> QueryReasoningEngineResponse:
+ """
+ Query a Reasoning Engine synchronously.
+
+ :param location: Required. The ID of the Google Cloud location that
the service belongs to.
+ :param reasoning_engine_id: Required. The Reasoning Engine resource ID.
+ :param input_data: Optional. Input for the Reasoning Engine class
method in JSON object format.
+ Defaults to ``None``.
+ :param class_method: Optional. The Reasoning Engine class method to
invoke. Defaults to ``query``.
+ :param retry: Designation of what errors, if any, should be retried.
Defaults to ``None``.
+ :param timeout: The timeout for this request. Defaults to ``None``.
+ :param metadata: Strings which should be sent along with the request
as metadata. Defaults
+ to an empty tuple.
+ :param project_id: Optional. The ID of the Google Cloud project.
Defaults to the project
+ configured in the connection.
+ """
+ client =
self.get_reasoning_engine_execution_service_client(location=location)
+ name = client.reasoning_engine_path(project_id, location,
reasoning_engine_id)
+ request: dict[str, Any] = {"name": name, "class_method": class_method}
+ if input_data is not None:
+ request["input"] = input_data
+ return client.query_reasoning_engine(
+ request=request,
+ retry=retry,
+ timeout=timeout,
+ metadata=metadata,
+ )
+
@GoogleBaseHook.fallback_to_default_project_id
def run_query_job(
self,
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py
b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py
index 3b6ea3a833d..96637e17d16 100644
---
a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py
+++
b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py
@@ -23,6 +23,8 @@ from collections.abc import Sequence
from functools import cached_property
from typing import TYPE_CHECKING, Any
+from google.cloud.aiplatform_v1.types import QueryReasoningEngineResponse
+
from airflow.providers.common.compat.sdk import conf
from airflow.providers.google.cloud.hooks.vertex_ai.agent_engine import (
AgentEngineHook,
@@ -33,6 +35,8 @@ from airflow.providers.google.cloud.operators.cloud_base
import GoogleCloudBaseO
from airflow.providers.google.cloud.triggers.vertex_ai import
AgentEngineQueryJobTrigger
if TYPE_CHECKING:
+ from google.api_core.retry import Retry
+ from pydantic import JsonValue
from vertexai._genai import types
from airflow.providers.common.compat.sdk import Context
@@ -165,6 +169,89 @@ class GetAgentEngineOperator(GoogleCloudBaseOperator):
return result
+class RunReasoningEngineQueryOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Reasoning Engine synchronously.
+
+ :param project_id: Required (templated). The ID of the Google Cloud
project that the service
+ belongs to.
+ :param location: Required (templated). The ID of the Google Cloud location
that the service
+ belongs to.
+ :param reasoning_engine_id: Required (templated). The Reasoning Engine
resource ID.
+ :param input_data: Optional (templated). Input for the Reasoning Engine
class method in JSON
+ object format. Defaults to ``None``.
+ :param class_method: Optional (templated). The Reasoning Engine class
method to invoke. Defaults
+ to ``query``.
+ :param retry: Designation of what errors, if any, should be retried.
Defaults to ``None``.
+ :param timeout: The timeout for this request. Defaults to ``None``.
+ :param metadata: Strings which should be sent along with the request as
metadata. Defaults to
+ an empty tuple.
+ :param gcp_conn_id: The connection ID to use connecting to Google Cloud
(templated). Defaults
+ to ``google_cloud_default``.
+ :param impersonation_chain: Optional service account to impersonate using
short-term credentials
+ (templated). Defaults to ``None``.
+ """
+
+ template_fields = (
+ "project_id",
+ "location",
+ "reasoning_engine_id",
+ "input_data",
+ "class_method",
+ "gcp_conn_id",
+ "impersonation_chain",
+ )
+
+ def __init__(
+ self,
+ *,
+ project_id: str,
+ location: str,
+ reasoning_engine_id: str,
+ input_data: dict[str, Any] | None = None,
+ class_method: str = "query",
+ retry: Retry | None = None,
+ timeout: float | None = None,
+ metadata: Sequence[tuple[str, str]] = (),
+ gcp_conn_id: str = "google_cloud_default",
+ impersonation_chain: str | Sequence[str] | None = None,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.project_id = project_id
+ self.location = location
+ self.reasoning_engine_id = reasoning_engine_id
+ self.input_data = input_data
+ self.class_method = class_method
+ self.retry = retry
+ self.timeout = timeout
+ self.metadata = metadata
+ self.gcp_conn_id = gcp_conn_id
+ self.impersonation_chain = impersonation_chain
+
+ @cached_property
+ def hook(self) -> AgentEngineHook:
+ return AgentEngineHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ )
+
+ def execute(self, context: Context) -> dict[str, JsonValue]:
+ self.log.info("Querying Reasoning Engine %s.",
self.reasoning_engine_id)
+ response = self.hook.query_reasoning_engine(
+ project_id=self.project_id,
+ location=self.location,
+ reasoning_engine_id=self.reasoning_engine_id,
+ input_data=self.input_data,
+ class_method=self.class_method,
+ retry=self.retry,
+ timeout=self.timeout,
+ metadata=self.metadata,
+ )
+ self.log.info("Reasoning Engine %s returned a response.",
self.reasoning_engine_id)
+ return QueryReasoningEngineResponse.to_dict(response)
+
+
class RunQueryJobOperator(GoogleCloudBaseOperator):
"""
Run a query job on a Vertex AI Agent Engine.
diff --git
a/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_agent_engine.py
b/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_agent_engine.py
index 1f7ebb3bdfb..f93a12d4a70 100644
---
a/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_agent_engine.py
+++
b/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_agent_engine.py
@@ -61,6 +61,7 @@ from
airflow.providers.google.cloud.operators.vertex_ai.agent_engine import (
DeleteAgentEngineOperator,
GetAgentEngineOperator,
RunQueryJobOperator,
+ RunReasoningEngineQueryOperator,
UpdateAgentEngineOperator,
)
@@ -191,6 +192,17 @@ with DAG(
)
# [END how_to_cloud_vertex_ai_get_agent_engine_operator]
+ # [START how_to_cloud_vertex_ai_run_reasoning_engine_query_operator]
+ run_reasoning_engine_query = RunReasoningEngineQueryOperator(
+ task_id="run_reasoning_engine_query",
+ project_id=PROJECT_ID,
+ location=LOCATION,
+ reasoning_engine_id=AGENT_ENGINE_ID,
+ input_data={"input": "hello from Airflow"},
+ timeout=300,
+ )
+ # [END how_to_cloud_vertex_ai_run_reasoning_engine_query_operator]
+
# [START how_to_cloud_vertex_ai_run_query_job_operator]
run_query_job = RunQueryJobOperator(
task_id="run_query_job",
@@ -259,6 +271,7 @@ with DAG(
[create_bucket, build_agent_image]
>> create_agent_engine
>> get_agent_engine
+ >> run_reasoning_engine_query
>> run_query_job
>> run_query_job_deferrable
>> update_agent_engine
diff --git
a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_agent_engine.py
b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_agent_engine.py
index 54574865612..25f7ff46b8f 100644
---
a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_agent_engine.py
+++
b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_agent_engine.py
@@ -22,6 +22,7 @@ from unittest import mock
import pytest
from airflow.providers.google.cloud.hooks.vertex_ai.agent_engine import
AgentEngineAsyncHook, AgentEngineHook
+from airflow.providers.google.common.consts import CLIENT_INFO
from unit.google.cloud.utils.base_gcp_mock import
mock_base_gcp_hook_default_project_id
@@ -39,6 +40,7 @@ OPERATION_ID = "delete-123"
QUERY_OPERATION_ID = "query-123"
CONFIG = {"display_name": "test-agent-engine"}
QUERY_CONFIG = {"query": "hello", "output_gcs_uri":
"gs://test-bucket/query-output/"}
+QUERY_INPUT = {"input": "hello"}
CHECK_QUERY_CONFIG = {"retrieve_result": True}
@@ -62,6 +64,38 @@ class TestAgentEngineHookWithDefaultProjectId:
)
assert result == mock_client.return_value.agent_engines
+ @pytest.mark.parametrize(
+ ("location", "is_default_universe", "expected"),
+ [
+ (GCP_LOCATION, True,
f"{GCP_LOCATION}-aiplatform.googleapis.com:443"),
+ ("global", True, None),
+ (GCP_LOCATION, False, None),
+ ],
+ )
+ def test_get_api_endpoint(self, location, is_default_universe, expected):
+ self.hook.is_default_universe = mock.create_autospec(
+ AgentEngineHook.is_default_universe,
return_value=is_default_universe
+ )
+
+ assert self.hook._get_api_endpoint(location=location) == expected
+
+
@mock.patch(AGENT_ENGINE_STRING.format("ReasoningEngineExecutionServiceClient"),
autospec=True)
+ def test_get_reasoning_engine_execution_service_client(self, mock_client):
+ self.hook.get_credentials =
mock.Mock(return_value=mock.sentinel.credentials, spec=())
+ self.hook.get_client_options =
mock.Mock(return_value=mock.sentinel.client_options, spec=())
+
+ result =
self.hook.get_reasoning_engine_execution_service_client(location=GCP_LOCATION)
+
+ self.hook.get_client_options.assert_called_once_with(
+
api_endpoint_override=f"{GCP_LOCATION}-aiplatform.googleapis.com:443"
+ )
+ mock_client.assert_called_once_with(
+ credentials=mock.sentinel.credentials,
+ client_info=CLIENT_INFO,
+ client_options=mock.sentinel.client_options,
+ )
+ assert result == mock_client.return_value
+
@mock.patch(AGENT_ENGINE_STRING.format("AgentEngineHook.get_agent_engine_client"),
autospec=True)
def test_create_agent_engine(self, mock_get_client):
result = self.hook.create_agent_engine(
@@ -100,6 +134,59 @@ class TestAgentEngineHookWithDefaultProjectId:
mock_get_client.return_value.get.assert_called_once_with(name=AGENT_ENGINE_NAME,
config=CONFIG)
assert result == mock_get_client.return_value.get.return_value
+ @mock.patch(
+
AGENT_ENGINE_STRING.format("AgentEngineHook.get_reasoning_engine_execution_service_client"),
+ autospec=True,
+ )
+ def test_query_reasoning_engine(self, mock_get_client):
+ result = self.hook.query_reasoning_engine(
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ reasoning_engine_id=AGENT_ENGINE_ID,
+ input_data=QUERY_INPUT,
+ class_method="custom_query",
+ retry=mock.sentinel.retry,
+ timeout=60,
+ metadata=(("key", "value"),),
+ )
+
+ mock_get_client.assert_called_once_with(self.hook,
location=GCP_LOCATION)
+
mock_get_client.return_value.reasoning_engine_path.assert_called_once_with(
+ GCP_PROJECT, GCP_LOCATION, AGENT_ENGINE_ID
+ )
+
mock_get_client.return_value.query_reasoning_engine.assert_called_once_with(
+ request={
+ "name":
mock_get_client.return_value.reasoning_engine_path.return_value,
+ "input": QUERY_INPUT,
+ "class_method": "custom_query",
+ },
+ retry=mock.sentinel.retry,
+ timeout=60,
+ metadata=(("key", "value"),),
+ )
+ assert result ==
mock_get_client.return_value.query_reasoning_engine.return_value
+
+ @mock.patch(
+
AGENT_ENGINE_STRING.format("AgentEngineHook.get_reasoning_engine_execution_service_client"),
+ autospec=True,
+ )
+ def test_query_reasoning_engine_without_input_data(self, mock_get_client):
+ self.hook.query_reasoning_engine(
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ reasoning_engine_id=AGENT_ENGINE_ID,
+ )
+
+
mock_get_client.return_value.query_reasoning_engine.assert_called_once_with(
+ request={
+ "name":
mock_get_client.return_value.reasoning_engine_path.return_value,
+ "class_method": "query",
+ },
+ retry=None,
+ timeout=None,
+ metadata=(),
+ )
+
@mock.patch(AGENT_ENGINE_STRING.format("AgentEngineHook.get_agent_engine_client"),
autospec=True)
def test_run_query_job(self, mock_get_client):
result = self.hook.run_query_job(
diff --git
a/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_agent_engine.py
b/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_agent_engine.py
index 586eead4c77..8d65ff1c32b 100644
---
a/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_agent_engine.py
+++
b/providers/google/tests/unit/google/cloud/operators/vertex_ai/test_agent_engine.py
@@ -20,6 +20,7 @@ from __future__ import annotations
from unittest import mock
import pytest
+from google.cloud.aiplatform_v1.types import QueryReasoningEngineResponse
from airflow.providers.common.compat.sdk import TaskDeferred
from airflow.providers.google.cloud.operators.vertex_ai.agent_engine import (
@@ -27,6 +28,7 @@ from
airflow.providers.google.cloud.operators.vertex_ai.agent_engine import (
DeleteAgentEngineOperator,
GetAgentEngineOperator,
RunQueryJobOperator,
+ RunReasoningEngineQueryOperator,
UpdateAgentEngineOperator,
)
@@ -41,6 +43,7 @@ AGENT_ENGINE_ID = "123"
AGENT_ENGINE_NAME =
"projects/test-project/locations/us-central1/reasoningEngines/123"
CONFIG = {"display_name": "test-agent-engine"}
QUERY_CONFIG = {"query": "hello", "output_gcs_uri":
"gs://test-bucket/query-output/"}
+QUERY_INPUT = {"input": "hello"}
CHECK_QUERY_CONFIG = {"retrieve_result": True}
OPERATION = {"name": "operations/delete-123", "done": False}
QUERY_OPERATION_NAME = "operations/query-123"
@@ -148,6 +151,66 @@ class TestGetAgentEngineOperator:
)
+class TestRunReasoningEngineQueryOperator:
+ @mock.patch(AGENT_ENGINE_PATH.format("AgentEngineHook"), autospec=True)
+ def test_execute(self, mock_hook, context):
+ query_output = {"message": "Hello from Agent Engine"}
+ mock_hook.return_value.query_reasoning_engine.return_value =
QueryReasoningEngineResponse(
+ output=query_output
+ )
+ op = RunReasoningEngineQueryOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ reasoning_engine_id=AGENT_ENGINE_ID,
+ input_data=QUERY_INPUT,
+ class_method="custom_query",
+ retry=mock.sentinel.retry,
+ timeout=60,
+ metadata=(("key", "value"),),
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ result = op.execute(context=context)
+
+ assert_hook_created(mock_hook)
+ mock_hook.return_value.query_reasoning_engine.assert_called_once_with(
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ reasoning_engine_id=AGENT_ENGINE_ID,
+ input_data=QUERY_INPUT,
+ class_method="custom_query",
+ retry=mock.sentinel.retry,
+ timeout=60,
+ metadata=(("key", "value"),),
+ )
+ assert result == {"output": query_output}
+
+ @mock.patch(AGENT_ENGINE_PATH.format("AgentEngineHook"), autospec=True)
+ def test_execute_uses_no_retry_by_default(self, mock_hook, context):
+ mock_hook.return_value.query_reasoning_engine.return_value =
QueryReasoningEngineResponse()
+ op = RunReasoningEngineQueryOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ reasoning_engine_id=AGENT_ENGINE_ID,
+ )
+
+ op.execute(context=context)
+
+ mock_hook.return_value.query_reasoning_engine.assert_called_once_with(
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ reasoning_engine_id=AGENT_ENGINE_ID,
+ input_data=None,
+ class_method="query",
+ retry=None,
+ timeout=None,
+ metadata=(),
+ )
+
+
class TestRunQueryJobOperator:
@mock.patch(AGENT_ENGINE_PATH.format("AgentEngineHook"), autospec=True)
def test_execute(self, mock_hook, context):
