Skip to content

Add synchronous Vertex AI Agent Engine query operator - #70933

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

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

Conversation

@AlejandroMorgante

@AlejandroMorgante AlejandroMorgante commented Aug 2, 2026

Copy link
Copy Markdown
Contributor

Add RunAgentQueryOperator 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.

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.
@AlejandroMorgante
AlejandroMorgante force-pushed the add-agent-engine-query-operator branch from 5ed968a to dfdd1a9 Compare August 2, 2026 22:43
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)

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 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(

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 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):

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

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