Set up generator class and auth for Data Engineering Agent by minzznguyen · Pull Request #404 · GoogleCloudPlatform/evalbench

minzznguyen · 2026-05-26T21:59:44Z

No description provided.

… with GCP ADC Task 1.1: Define DataEngineeringAgentGenerator inheriting from QueryGenerator in data_engineering_agent.py. Task 1.2: Integrate A2A SDK and configure GcpAdcCredentialService to use GCP Application Default Credentials (ADC) for authentication. TAG=agy CONV=aa927cc7-418a-41e3-b658-9b82915e18eb

IsmailMehdi

Code review

Scaffolds a new DataEngineeringAgentGenerator plus a GcpAdcCredentialService for the A2A SDK's AuthInterceptor. Per the WIP label, auth is the focus; messaging in generate_internal is stubbed.

Correctness & functional issues

Generator is not registered. evalbench/generators/models/__init__.py's get_generator() dispatch has no data_engineering_agent branch and the class isn't imported. Even with a config pointing at it, instantiation falls through to ValueError("Unknown Generator ..."). Either add the branch now or call this out as a follow-up in the PR description.
Missing dependency declaration. a2a (the A2A SDK) is not in pyproject.toml. The repo recently migrated to uv (d4501e4) — please uv add a2a so the lockfile is updated. Without this, the import fails anywhere outside the author's local env.
Blocking sync call inside async def. credentials.refresh(Request()) does a synchronous HTTPS round-trip and is called from async def get_credentials. Inside an asyncio event loop this blocks every other task. Use await asyncio.to_thread(credentials.refresh, Request()), or an async-compatible auth flow.
Token fetched on every call, no caching. ADC tokens are valid ~1 hour. google.auth.default() + refresh() runs on every get_credentials invocation — wasteful and slow on the hot path. Cache credentials on self and only refresh when not credentials.valid (the library exposes credentials.expired / valid).
Unused import. from a2a.types import a2a_pb2 as pb — pb is never referenced. Drop until generate_internal actually needs it.
self.logger is set but never used. All output goes through print(...). Pick one — and given the rest of the codebase, that should be logging (see gemini.py, claude.py, generator.py).

Security

Logging token bytes, even truncated. print(f\"... GCP ADC token: {token[:10]}...\") writes the first 10 chars of an access token to stdout. Even prefixes are sensitive and end up in CI logs / shared terminals. Replace with self.logger.debug(\"Retrieved GCP ADC token\") — no bytes of the token at all.
security_scheme_name is ignored. The service hands the same bearer token back regardless of scheme. If multiple schemes are ever registered, this is silently wrong. At minimum add a comment that this provider intentionally only services one scheme; ideally branch (or raise) on unknown names.
No error handling around google.auth.default() / refresh(). Both can raise (DefaultCredentialsError, RefreshError). Wrap with try/except and log clearly — saves real debugging time later.

Style / conventions

Emoji-laden print statements (🔑, ✅) don't match the rest of the codebase, which uses standard logging. Remove emoji and route through the logger.
Trailing whitespace on the line after self.logger = logging.getLogger(__name__).
The closest existing pattern to mirror is claude.py (concise constructor, single client attribute, plain logging).

Tests

None added. Understandable for a WIP that doesn't do anything yet, but at least one unit test that constructs the generator with a fake CredentialService (no network call) would catch the registration gap and the unused-import problem in CI. Plan to add tests alongside Task 1.3.

Recommended order

Replace the token-prefix print with logger.debug — no token bytes.
Move credentials.refresh(Request()) off the event loop (asyncio.to_thread) and cache the credentials object.
Add a2a via uv add so lockfile/pyproject.toml are updated.
Register DataEngineeringAgentGenerator in generators/models/__init__.py once the class is real — or leave a TODO in the PR body.
Remove the unused a2a_pb2 as pb import; pick logging over print.
Wrap google.auth.default() / refresh() with try/except.

Risk summary

Low blast radius today — the generator can't be selected via config and generate_internal raises. Main risks are landing dead/insecure scaffolding into main: (a) token-prefix logging, (b) missing dep declaration breaking imports if anything reaches it, (c) blocking sync auth call that will silently hurt throughput once Task 1.3 lands. Address (a) and (c) before this merges even as WIP scaffolding.

- Add `a2a-sdk>=1.0.3` to dependencies in `pyproject.toml` and update `uv.lock`. - Update `DataEngineeringAgentGenerator` to configure `endpoint`, `target_workspace` and clean up GCP ADC Credential Service logic with event loop safety (async refresh). - Register `DataEngineeringAgentGenerator` in the generators factory. - Add `evalbench/test/data_engineering_agent_test.py` to verify correct generator setup and configurations. TAG=agy CONV=aa927cc7-418a-41e3-b658-9b82915e18eb

IsmailMehdi · 2026-05-27T20:24:13Z

+        self.name = "data_engineering_agent"
+        self.endpoint = querygenerator_config.get(
+            "endpoint",
+            "https://geminidataanalytics.googleapis.com/v1/a2a/projects/bq-dataworkeragent-test/locations/us-west4/agents/dataengineeringagent"


no hardcoded endpoints, please have a config with defaults.

IsmailMehdi · 2026-05-27T20:24:23Z

+        )
+        self.target_workspace = querygenerator_config.get(
+            "target_workspace",
+            "projects/bq-dataworkeragent-test/locations/us-west4/repositories/agent_demo_dataform/workspaces/james_test_123"


https://github.com/GoogleCloudPlatform/evalbench/pull/404/changes#r3313631524:~:text=no%20hardcoded%20endpoints%2C%20please%20have%20a%20config%20with%20defaults.

- Remove hardcoded test URLs for `endpoint` and `target_workspace` from `DataEngineeringAgentGenerator`. - Raise `ValueError` if either configuration key is missing or empty. TAG=agy CONV=aa927cc7-418a-41e3-b658-9b82915e18eb

TAG=agy CONV=aa927cc7-418a-41e3-b658-9b82915e18eb

- Add a new unit test `test_get_credentials_invalid_scheme` verifying that `GcpAdcCredentialService` raises `ValueError` for unsupported auth schemes. - Replace absolute generator import in `evalbench/generators/models/__init__.py` with a relative one. - Clean up unused imports in `data_engineering_agent.py`. TAG=agy CONV=aa927cc7-418a-41e3-b658-9b82915e18eb

- Revert relative import of `QueryGenerator` back to absolute import `generators.models.generator.QueryGenerator` in `evalbench/generators/models/__init__.py` to prevent runtime package resolution issues. TAG=agy CONV=aa927cc7-418a-41e3-b658-9b82915e18eb

graceqi-g · 2026-05-28T18:32:30Z

+                "Credentials: %s",
+                e,
+            )
+            return None


are we able to talk to the endpoint if we fail to get the credentials? if not, let's allow the error propagate to the caller instead of suppressing it here (same for the general exception below). i.e. remove the try-except

graceqi-g · 2026-05-28T18:33:06Z

+)
+
+
+def test_data_engineering_agent_generator_setup():


let's also add tests for the bad paths, i.e. when endpoint and/or target_workspace are missing and ensure the value error is being raised

similarly, let's also try to cover the bad paths in GcpAdcCredentialService.get_credentials

- Update `GcpAdcCredentialService` to propagate `DefaultCredentialsError` and `RefreshError` up instead of catching them silently and returning `None`. - Add `test_generator_setup_missing_endpoint` and `test_generator_setup_missing_workspace` in `data_engineering_agent_test.py`. - Add `test_get_credentials_error_resiliency_default` and `test_get_credentials_error_resiliency_refresh` to verify error propagation. TAG=agy CONV=aa927cc7-418a-41e3-b658-9b82915e18eb

James Nguyen added 2 commits May 26, 2026 21:53

Merge remote-tracking branch 'origin/main' into data_engineering_agent

f9eb773

minzznguyen requested a review from IsmailMehdi as a code owner May 26, 2026 21:59

IsmailMehdi reviewed May 27, 2026

View reviewed changes

IsmailMehdi requested changes May 27, 2026

View reviewed changes

James Nguyen added 4 commits May 27, 2026 20:49

style: remove comments from data_engineering_agent.py

e1bc087

TAG=agy CONV=aa927cc7-418a-41e3-b658-9b82915e18eb

minzznguyen changed the title ~~[WIP] Set up generator class and auth for Data Engineering Agent~~ Set up generator class and auth for Data Engineering Agent May 28, 2026

graceqi-g requested changes May 28, 2026

View reviewed changes

graceqi-g approved these changes May 28, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Set up generator class and auth for Data Engineering Agent #404

Set up generator class and auth for Data Engineering Agent #404
minzznguyen wants to merge 8 commits into
GoogleCloudPlatform:mainfrom
minzznguyen:data_engineering_agent

minzznguyen commented May 26, 2026

Uh oh!

IsmailMehdi left a comment

Uh oh!

IsmailMehdi May 27, 2026

Uh oh!

IsmailMehdi May 27, 2026

Uh oh!

graceqi-g May 28, 2026

Uh oh!

graceqi-g May 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		)


		def test_data_engineering_agent_generator_setup():

Conversation

minzznguyen commented May 26, 2026

Uh oh!

IsmailMehdi left a comment

Choose a reason for hiding this comment

Code review

Correctness & functional issues

Security

Style / conventions

Tests

Recommended order

Risk summary

Uh oh!

IsmailMehdi May 27, 2026

Choose a reason for hiding this comment

Uh oh!

IsmailMehdi May 27, 2026

Choose a reason for hiding this comment

Uh oh!

graceqi-g May 28, 2026

Choose a reason for hiding this comment

Uh oh!

graceqi-g May 28, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants