Skip to content

Add synchronous Vertex AI Agent Engine query operator - #70933

Open
AlejandroMorgante wants to merge 7 commits into
apache:mainfrom
AlejandroMorgante:add-agent-engine-query-operator
Open

Add synchronous Vertex AI Agent Engine query operator#70933
AlejandroMorgante wants to merge 7 commits into
apache:mainfrom
AlejandroMorgante:add-agent-engine-query-operator

Conversation

@AlejandroMorgante

@AlejandroMorgante AlejandroMorgante commented Aug 2, 2026

Copy link
Copy Markdown
Contributor

Add RunReasoningEngineQueryOperator 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 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

  • Targeted unit tests: 44 passed.
  • Full Google provider suite: 4973 passed, 59 skipped.
  • Provider mypy checks passed.
  • Pre-commit and manual prek checks passed.
  • Google provider documentation build and spellcheck passed.
  • End-to-end GCP system test: 1 passed (details and reproduction steps below).
  • AI-assisted Magpie self-review (pr-management-code-review, dry-run) at dfdd1a9: no blocking, major, minor, or nit findings. The public GAPIC surface was cross-checked against the minimum supported google-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_PROJECT configured in files/airflow-breeze-config/environment_variables.env, and a unique environment ID:

SYSTEM_TESTS_ENV_ID=<unique-id> \
BREEZE_INIT_COMMAND='export GOOGLE_CLOUD_PROJECT="${SYSTEM_TESTS_GCP_PROJECT}"' \
breeze testing system-tests \
  --forward-credentials \
  --test-timeout 2400 \
  providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_agent_engine.py \
  -q

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?
  • Yes — Codex (GPT-5)

Generated-by: Codex (GPT-5) following the guidelines

@boring-cyborg boring-cyborg Bot added area:providers kind:documentation provider:google Google (including GCP) related issues labels Aug 2, 2026

@SameerMesiah97 SameerMesiah97 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just a few minor comments. Very clean overall.

@AlejandroMorgante
AlejandroMorgante force-pushed the add-agent-engine-query-operator branch from 5ed968a to dfdd1a9 Compare August 2, 2026 22:43
@AlejandroMorgante
AlejandroMorgante force-pushed the add-agent-engine-query-operator branch from dfdd1a9 to 98f5b06 Compare August 4, 2026 23:18
@AlejandroMorgante

Copy link
Copy Markdown
Contributor Author

@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.

@AlejandroMorgante
AlejandroMorgante force-pushed the add-agent-engine-query-operator branch 2 times, most recently from f92b8fa to 197d3a8 Compare August 5, 2026 15:18
@AlejandroMorgante
AlejandroMorgante marked this pull request as draft August 6, 2026 02:05
@AlejandroMorgante
AlejandroMorgante marked this pull request as ready for review August 6, 2026 02:05

@MaksYermak MaksYermak left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@AlejandroMorgante could you provide the screenshots of passing system tests for this new operator?

Comment thread providers/google/docs/operators/cloud/vertex_ai.rst Outdated
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.
Users familiar with ReasoningEngineExecutionService should be able to find the equivalent Airflow API without translating between resource vocabularies.
Matching the GAPIC resource vocabulary makes the operator easier for users to discover and understand.
@AlejandroMorgante
AlejandroMorgante force-pushed the add-agent-engine-query-operator branch from 197d3a8 to ed7f8c2 Compare August 7, 2026 19:56
@AlejandroMorgante

AlejandroMorgante commented Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

@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.

image image image

@MaksYermak MaksYermak left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:providers kind:documentation provider:google Google (including GCP) related issues

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants