You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "AlejandroMorgante (via GitHub)" <gi...@apache.org> on 2026/08/02 02:07:59 UTC
[PR] Add synchronous Vertex AI Agent Engine query operator [airflow]
AlejandroMorgante opened a new pull request, #70933:
URL: https://github.com/apache/airflow/pull/70933
Add `QueryAgentEngineOperator` for direct synchronous Agent Engine invocations through the public `ReasoningEngineExecutionService` GAPIC. This complements `RunQueryJobOperator` when a Google Cloud Storage-backed asynchronous job is not needed.
The hook configures the regional endpoint and forwards retry, timeout, metadata, input, and class method options. The operator returns the JSON-serializable Agent Engine output directly, without first fetching the runtime or creating a query job.
The new path is covered by hook and operator unit tests, provider documentation, and the existing Agent Engine system-test example.
related: #68479
Testing:
- `uv run --project providers/google pytest providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_agent_engine.py providers/google/tests/unit/google/cloud/operators/vertex_ai/test_agent_engine.py -q` (`43 passed`)
- `breeze testing providers-tests --test-type "Providers[google]"` (`4972 passed, 59 skipped`)
- `breeze run mypy providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py`
- `prek run --from-ref upstream/main --stage pre-commit`
- `prek run --from-ref upstream/main --stage manual`
- `breeze build-docs google --one-pass-only` (documentation and spellcheck)
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes — Codex (GPT-5)
Generated-by: Codex (GPT-5) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add synchronous Vertex AI Agent Engine query operator [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal merged PR #70933:
URL: https://github.com/apache/airflow/pull/70933
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add synchronous Vertex AI Agent Engine query operator [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #70933:
URL: https://github.com/apache/airflow/pull/70933#discussion_r3700285584
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +169,82 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class QueryAgentEngineOperator(GoogleCloudBaseOperator):
Review Comment:
I would rename this to `RunAgentQueryOperator` to make the pairing with `RunQueryJobOperator` more explicit.
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +169,82 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class QueryAgentEngineOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent Engine synchronously.
+
+ :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+ :param location: Required. The ID of the Google Cloud location that the service belongs to.
+ :param agent_engine_id: Required. The Agent Engine ID.
+ :param input_data: Optional. Input for the Agent Engine class method in JSON object format.
+ :param class_method: Optional. The Agent Engine class method to invoke. Defaults to ``query``.
+ :param retry: Designation of what errors, if any, should be retried.
+ :param timeout: The timeout for this request.
+ :param metadata: Strings which should be sent along with the request as metadata.
+ :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+ :param impersonation_chain: Optional service account to impersonate using short-term credentials.
+ """
+
+ template_fields = (
+ "project_id",
+ "location",
+ "agent_engine_id",
+ "input_data",
+ "class_method",
+ "gcp_conn_id",
+ "impersonation_chain",
+ )
+
+ def __init__(
+ self,
+ *,
+ project_id: str,
+ location: str,
+ agent_engine_id: str,
+ input_data: dict[str, Any] | None = None,
+ class_method: str = "query",
+ retry: Retry | _MethodDefault = DEFAULT,
+ 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.agent_engine_id = agent_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) -> Any:
Review Comment:
Could you be more specific with the return type? `Any` is not ideal.
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +169,82 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class QueryAgentEngineOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent Engine synchronously.
+
+ :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+ :param location: Required. The ID of the Google Cloud location that the service belongs to.
+ :param agent_engine_id: Required. The Agent Engine ID.
+ :param input_data: Optional. Input for the Agent Engine class method in JSON object format.
+ :param class_method: Optional. The Agent Engine class method to invoke. Defaults to ``query``.
+ :param retry: Designation of what errors, if any, should be retried.
+ :param timeout: The timeout for this request.
+ :param metadata: Strings which should be sent along with the request as metadata.
+ :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+ :param impersonation_chain: Optional service account to impersonate using short-term credentials.
Review Comment:
1) Add the default values for each of the parameters to the docstring.
2) Indicate whether each parameter is templated by adding `(templated)` if it is.
##########
providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py:
##########
@@ -141,6 +164,43 @@ def get_agent_engine(
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_agent_engine(
+ self,
+ location: str,
+ agent_engine_id: str,
+ input_data: dict[str, Any] | None = None,
+ class_method: str = "query",
+ retry: Retry | _MethodDefault = DEFAULT,
+ timeout: float | None = None,
+ metadata: Sequence[tuple[str, str]] = (),
+ project_id: str = PROVIDE_PROJECT_ID,
+ ) -> QueryReasoningEngineResponse:
+ """
+ Query an Agent Engine synchronously.
+
+ :param location: Required. The ID of the Google Cloud location that the service belongs to.
+ :param agent_engine_id: Required. The Agent Engine ID.
+ :param input_data: Optional. Input for the Agent Engine class method in JSON object format.
+ :param class_method: Optional. The Agent Engine class method to invoke. Defaults to ``query``.
+ :param retry: Designation of what errors, if any, should be retried.
+ :param timeout: The timeout for this request.
+ :param metadata: Strings which should be sent along with the request as metadata.
+ :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project
+ configured in the connection.
Review Comment:
I think some of the docstring entries do not have the default values mentioned. Please include them.
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +169,82 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class QueryAgentEngineOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent Engine synchronously.
+
+ :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+ :param location: Required. The ID of the Google Cloud location that the service belongs to.
+ :param agent_engine_id: Required. The Agent Engine ID.
+ :param input_data: Optional. Input for the Agent Engine class method in JSON object format.
+ :param class_method: Optional. The Agent Engine class method to invoke. Defaults to ``query``.
+ :param retry: Designation of what errors, if any, should be retried.
+ :param timeout: The timeout for this request.
+ :param metadata: Strings which should be sent along with the request as metadata.
+ :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+ :param impersonation_chain: Optional service account to impersonate using short-term credentials.
+ """
+
+ template_fields = (
+ "project_id",
+ "location",
+ "agent_engine_id",
+ "input_data",
+ "class_method",
+ "gcp_conn_id",
+ "impersonation_chain",
+ )
+
+ def __init__(
+ self,
+ *,
+ project_id: str,
+ location: str,
+ agent_engine_id: str,
+ input_data: dict[str, Any] | None = None,
+ class_method: str = "query",
+ retry: Retry | _MethodDefault = DEFAULT,
+ 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.agent_engine_id = agent_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) -> Any:
+ self.log.info("Querying Agent Engine %s.", self.agent_engine_id)
+ response = self.hook.query_agent_engine(
+ project_id=self.project_id,
+ location=self.location,
+ agent_engine_id=self.agent_engine_id,
+ input_data=self.input_data,
+ class_method=self.class_method,
+ retry=self.retry,
+ timeout=self.timeout,
+ metadata=self.metadata,
+ )
+ self.log.info("Agent Engine %s returned a response.", self.agent_engine_id)
+ return QueryReasoningEngineResponse.to_dict(response).get("output")
Review Comment:
Is there a reason we're returning only the output field rather than the full serialized response?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add synchronous Vertex AI Agent Engine query operator [airflow]
Posted by "AlejandroMorgante (via GitHub)" <gi...@apache.org>.
AlejandroMorgante commented on code in PR #70933:
URL: https://github.com/apache/airflow/pull/70933#discussion_r3700454222
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +169,82 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class QueryAgentEngineOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent Engine synchronously.
+
+ :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+ :param location: Required. The ID of the Google Cloud location that the service belongs to.
+ :param agent_engine_id: Required. The Agent Engine ID.
+ :param input_data: Optional. Input for the Agent Engine class method in JSON object format.
+ :param class_method: Optional. The Agent Engine class method to invoke. Defaults to ``query``.
+ :param retry: Designation of what errors, if any, should be retried.
+ :param timeout: The timeout for this request.
+ :param metadata: Strings which should be sent along with the request as metadata.
+ :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+ :param impersonation_chain: Optional service account to impersonate using short-term credentials.
+ """
+
+ template_fields = (
+ "project_id",
+ "location",
+ "agent_engine_id",
+ "input_data",
+ "class_method",
+ "gcp_conn_id",
+ "impersonation_chain",
+ )
+
+ def __init__(
+ self,
+ *,
+ project_id: str,
+ location: str,
+ agent_engine_id: str,
+ input_data: dict[str, Any] | None = None,
+ class_method: str = "query",
+ retry: Retry | _MethodDefault = DEFAULT,
+ 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.agent_engine_id = agent_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) -> Any:
+ self.log.info("Querying Agent Engine %s.", self.agent_engine_id)
+ response = self.hook.query_agent_engine(
+ project_id=self.project_id,
+ location=self.location,
+ agent_engine_id=self.agent_engine_id,
+ input_data=self.input_data,
+ class_method=self.class_method,
+ retry=self.retry,
+ timeout=self.timeout,
+ metadata=self.metadata,
+ )
+ self.log.info("Agent Engine %s returned a response.", self.agent_engine_id)
+ return QueryReasoningEngineResponse.to_dict(response).get("output")
Review Comment:
Done, good catch!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add synchronous Vertex AI Agent Engine query operator [airflow]
Posted by "MaksYermak (via GitHub)" <gi...@apache.org>.
MaksYermak commented on code in PR #70933:
URL: https://github.com/apache/airflow/pull/70933#discussion_r3734465685
##########
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 an Agent Engine synchronously you can use
Review Comment:
@AlejandroMorgante could you please update documentation and mentioned that is ReasoningEngine operator?
##########
providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py:
##########
@@ -141,6 +164,45 @@ def get_agent_engine(
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 | _MethodDefault = DEFAULT,
+ timeout: float | None = None,
+ metadata: Sequence[tuple[str, str]] = (),
+ project_id: str = PROVIDE_PROJECT_ID,
+ ) -> QueryReasoningEngineResponse:
+ """
+ Query an Agent Engine synchronously.
Review Comment:
@AlejandroMorgante could you also change the docstring here to mention that it is Reasoning Engine?
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +170,90 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class RunReasoningEngineQueryOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent 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 for the
+ Agent Engine.
+ :param input_data: Optional (templated). Input for the Agent Engine class method in JSON object
+ format. Defaults to ``None``.
+ :param class_method: Optional (templated). The Agent Engine class method to invoke. Defaults to
+ ``query``.
+ :param retry: Designation of what errors, if any, should be retried. Defaults to ``DEFAULT``.
+ :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 | _MethodDefault = DEFAULT,
+ 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 Agent Engine %s.", self.reasoning_engine_id)
Review Comment:
@AlejandroMorgante could you update the log message to "Querying Reasoning Engine"?
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +170,90 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class RunReasoningEngineQueryOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent Engine synchronously.
Review Comment:
@AlejandroMorgante could you change that it is query to reasoning engine?
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +170,90 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class RunReasoningEngineQueryOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent 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 for the
+ Agent Engine.
+ :param input_data: Optional (templated). Input for the Agent Engine class method in JSON object
+ format. Defaults to ``None``.
+ :param class_method: Optional (templated). The Agent Engine class method to invoke. Defaults to
+ ``query``.
+ :param retry: Designation of what errors, if any, should be retried. Defaults to ``DEFAULT``.
+ :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 | _MethodDefault = DEFAULT,
+ 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 Agent 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("Agent Engine %s returned a response.", self.reasoning_engine_id)
Review Comment:
@AlejandroMorgante the same is here could you change "Agent" to "Reasoning"?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add synchronous Vertex AI Agent Engine query operator [airflow]
Posted by "AlejandroMorgante (via GitHub)" <gi...@apache.org>.
AlejandroMorgante commented on PR #70933:
URL: https://github.com/apache/airflow/pull/70933#issuecomment-5222482371
@MaksYermak
> @AlejandroMorgante could you provide the screenshots of passing system tests for this new operator?
Hello! Here are the E2E test screenshots showing the synchronous Reasoning Engine query operator, the successful task run, and the response stored in XCom.
<img width="3018" height="1636" alt="image" src="https://github.com/user-attachments/assets/8f032faa-9702-48db-92f6-7a0380b07948" />
<img width="3018" height="1636" alt="image" src="https://github.com/user-attachments/assets/d91447e4-528b-4b4d-b56d-87cd901b2e22" />
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add synchronous Vertex AI Agent Engine query operator [airflow]
Posted by "MaksYermak (via GitHub)" <gi...@apache.org>.
MaksYermak commented on code in PR #70933:
URL: https://github.com/apache/airflow/pull/70933#discussion_r3702978401
##########
providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py:
##########
@@ -141,6 +164,45 @@ def get_agent_engine(
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_agent_engine(
+ self,
+ location: str,
+ agent_engine_id: str,
+ input_data: dict[str, Any] | None = None,
+ class_method: str = "query",
+ retry: Retry | _MethodDefault = DEFAULT,
+ timeout: float | None = None,
+ metadata: Sequence[tuple[str, str]] = (),
+ project_id: str = PROVIDE_PROJECT_ID,
+ ) -> QueryReasoningEngineResponse:
+ """
+ Query an Agent Engine synchronously.
+
+ :param location: Required. The ID of the Google Cloud location that the service belongs to.
+ :param agent_engine_id: Required. The Agent Engine ID.
+ :param input_data: Optional. Input for the Agent Engine class method in JSON object format.
+ Defaults to ``None``.
+ :param class_method: Optional. The Agent Engine class method to invoke. Defaults to ``query``.
+ :param retry: Designation of what errors, if any, should be retried. Defaults to ``DEFAULT``.
+ :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, agent_engine_id)
Review Comment:
@AlejandroMorgante as I see in the [code](https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/client.py#L269) it is `reasoning_engine` parameter not the `agent_engine`. Could you please change `agent_engine_id` to `reasoning_engine_id` to be consist with the Client code?
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +170,89 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class RunAgentQueryOperator(GoogleCloudBaseOperator):
Review Comment:
@AlejandroMorgante the same question as for Hook method, why do you use the `RunAgentQueryOperator` name when under the hood the Client will call `query_reasoning_engine` method? I think it is better to use `RunReasoningEngineQueryOperator` name
##########
providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py:
##########
@@ -141,6 +164,45 @@ def get_agent_engine(
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_agent_engine(
Review Comment:
@AlejandroMorgante could you please clarify for me why do you use `query_agent_engine` name here but later in the code the `Client` call `query_reasoning_engine` method? In my opinion it is make sense to use the `query_reasoning_engine` name for the Hook method
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add synchronous Vertex AI Agent Engine query operator [airflow]
Posted by "AlejandroMorgante (via GitHub)" <gi...@apache.org>.
AlejandroMorgante commented on PR #70933:
URL: https://github.com/apache/airflow/pull/70933#issuecomment-5185693327
@MaksYermak I applied the suggested renames-thanks for the feedback! My original intention was to follow the pattern we had been developing, with a synchronous Agent Engine query operator alongside the existing query-job operator, and keep their naming consistent. However, I agree that, from an Airflow user's perspective, matching the underlying client resource and method names makes the API easier to discover and understand.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add synchronous Vertex AI Agent Engine query operator [airflow]
Posted by "AlejandroMorgante (via GitHub)" <gi...@apache.org>.
AlejandroMorgante commented on code in PR #70933:
URL: https://github.com/apache/airflow/pull/70933#discussion_r3761866089
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -23,6 +23,9 @@
from functools import cached_property
from typing import TYPE_CHECKING, Any
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
Review Comment:
I removed the private _MethodDefault import and now use None, preserving the existing behavior while still allowing users to provide an explicit Retry policy.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add synchronous Vertex AI Agent Engine query operator [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #70933:
URL: https://github.com/apache/airflow/pull/70933#discussion_r3761312365
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -23,6 +23,9 @@
from functools import cached_property
from typing import TYPE_CHECKING, Any
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
Review Comment:
importing private classes/functions from 3rd party packages is not allowed.
Either google needs to make this part of their public api or avoid using it in Airflow code
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org