Add synchronous Vertex AI Agent Engine query operator - #70933
Add synchronous Vertex AI Agent Engine query operator#70933AlejandroMorgante wants to merge 5 commits into
Conversation
SameerMesiah97
left a comment
There was a problem hiding this comment.
Just a few minor comments. Very clean overall.
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.
The public hook exposes configurable request controls, so its unit test should detect regressions that replace caller-supplied values with defaults.
5ed968a to
dfdd1a9
Compare
| 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) |
There was a problem hiding this comment.
@AlejandroMorgante as I see in the code 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?
| return client.get(name=name, config=config) | ||
|
|
||
| @GoogleBaseHook.fallback_to_default_project_id | ||
| def query_agent_engine( |
There was a problem hiding this 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
| return result | ||
|
|
||
|
|
||
| class RunAgentQueryOperator(GoogleCloudBaseOperator): |
There was a problem hiding this 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
Add
RunAgentQueryOperatorfor direct synchronous Agent Engine invocations through the publicReasoningEngineExecutionServiceGAPIC. This complementsRunQueryJobOperatorwhen 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 complete JSON-serializable GAPIC response, 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
Validation
44 passed.4973 passed, 59 skipped.1 passed(details and reproduction steps below).pr-management-code-review, dry-run) atdfdd1a9: no blocking, major, minor, or nit findings. The public GAPIC surface was cross-checked against the minimum supportedgoogle-cloud-aiplatform==1.155.0. No GitHub review was posted.The complete system test was also run end-to-end against a real GCP environment. It can be reproduced from the Airflow checkout with GCP credentials forwarded to Breeze,
SYSTEM_TESTS_GCP_PROJECTconfigured infiles/airflow-breeze-config/environment_variables.env, and a unique environment ID:Latest result:
1 passed in 747.16s. The run exercised create, get, synchronous query, synchronous and deferrable query jobs, update, and delete flows, and completed resource teardown successfully.The one-time GCP project setup required by the test remains documented in the system test module docstring.
Was generative AI tooling used to co-author this PR?
Generated-by: Codex (GPT-5) following the guidelines