diff --git a/geminidataanalytics/dataengineeringagent/README.md b/geminidataanalytics/dataengineeringagent/README.md new file mode 100644 index 00000000000..d135b7b4ed1 --- /dev/null +++ b/geminidataanalytics/dataengineeringagent/README.md @@ -0,0 +1,140 @@ +# Data Engineering Agent A2A Client Example + +This directory contains a sample Python implementation of an +[Agent-to-Agent (A2A)](https://a2a-protocol.org/) client designed to interact +with the **Google Cloud Data Engineering Agent (DEA)**. + +## Background + +The Data Engineering Agent is a BigQuery and Dataform ELT expert capable of +building, managing, and troubleshooting data pipelines. To enable +interoperability across different platforms and agents, it exposes an interface +following the A2A protocol. + +Official Documentation: +[Data Engineering Agent API Overview](https://docs.cloud.google.com/gemini/data-agents/data-engineering-agent/api-overview) + +This example demonstrates how to use the open-source +[A2A Python SDK](https://github.com/a2aproject/a2a-python) to: 1. **Discover** +the agent's capabilities via its Agent Card. 2. **Authenticate** using Google +Application Default Credentials (ADC). 3. **Maintain State** across multi-turn +conversations using Conversation Tokens persisted to a local file. 4. **Handle +Complex Tasks** by automatically resuming execution when agent finished with +`DEADLINE_EXCEEDED`. 5. **Configure Extensions** like `Instruction` to customize +agent behavior. + +## Features + +- **Native A2A SDK Usage:** Uses `A2ACardResolver` and `create_client` for + idiomatic protocol interaction. +- **Streaming-only Execution:** Hardcoded to use response streaming for lowest + real-time latency and optimal interaction patterns. +- **Session Persistence:** Automatically saves and loads the + `conversationToken` from a local file, allowing multi-turn conversations via + repeated script executions. +- **Configurable Parameters:** Exposes `gcp_resource_id` for flexibility, + automatically extracting project and location details. +- **Instruction Loading:** Automatically reads custom instructions from local + files or directories, using filenames as the instruction name and file + content as the definition. +- **Extension Header Support:** Uses `ServiceParametersFactory` to correctly + set the `A2A-Extensions` HTTP header required by the agent. +- **Automated Resumption:** Detects `DEADLINE_EXCEEDED` via the + `finish_reason` extension and transparently continues the task. + +## Prerequisites + +- Python 3.10 or higher. +- [Google Cloud SDK (gcloud)](https://cloud.google.com/sdk/docs/install) + installed and configured. +- Enable required APIs + [Use the Data Engineering Agent to build and modify data pipelines](https://docs.cloud.google.com/bigquery/docs/data-engineering-agent-pipelines#required-apis) + +## Setup + +1. **Create and activate a virtual environment:** + + ```bash + python3 -m venv .dea + source .dea/bin/activate + ``` + +2. **Install dependencies:** + + ```bash + pip install -r requirements.txt + ``` + +3. **Authenticate with Google Cloud:** + + ```bash + gcloud auth application-default login + ``` + +## Usage + +### Single Message Mode + +Sends a single message and exits. `gcp_resource_id` and `message` are required. + +``` +python3 dea_a2a_client.py \ + --gcp_resource_id projects/my-project/locations/us-central1/repositories/my-repo/workspaces/default \ + --message "List my Dataform tables" +``` + +### Multi-turn Conversation (State Persistence) + +To maintain a conversation across multiple calls, use the +`--conversation_token_path` argument. The script will save the conversation +token to this file and reload it in subsequent calls. + +``` +# First turn (starts session) +python3 dea_a2a_client.py \ + --gcp_resource_id projects/my-project/locations/us-central1/repositories/my-repo/workspaces/default \ + --message "hi" \ + --conversation_token_path ./token.txt + +# Second turn (continues previous context) +python3 dea_a2a_client.py \ + --gcp_resource_id projects/my-project/locations/us-central1/repositories/my-repo/workspaces/default \ + --message "Explain the first table" \ + --conversation_token_path ./token.txt +``` + +### Advanced: Providing Local Instructions + +You can point the client to local files (e.g., SQL style guides) to influence +the agent's behavior. + +``` +python3 dea_a2a_client.py \ + --gcp_resource_id projects/my-project/locations/us-central1/repositories/my-repo/workspaces/default \ + --message "List my tables" \ + --instruction_path ./style_guide.md +``` + +#### Command-line Arguments + +Argument | Required | Description +:-------------------------- | :------- | :---------- +`--gcp_resource_id` | **Yes** | The target resource ID. Supported formats `projects/{p}/locations/{l}/repositories/{r}/workspaces/{w}` (Dataform) +`--message` | **Yes** | The message to send to the agent. +`--conversation_token_path` | No | Path to a local file to persist conversation token. Allows for multi-turn conversations. +`--instruction_path` | No | Path to a file or directory containing instructions. Can be repeated. + +### Running Unit Tests + +Make sure the virtual environment is activated, then run: + +```bash +python3 dea_a2a_client_test.py +``` + +Alternatively, if the virtual environment is not activated, you can run it +directly using the venv python: + +```bash +./.dea/bin/python3 dea_a2a_client_test.py +``` diff --git a/geminidataanalytics/dataengineeringagent/dea_a2a_client.py b/geminidataanalytics/dataengineeringagent/dea_a2a_client.py new file mode 100644 index 00000000000..2e48ab590b4 --- /dev/null +++ b/geminidataanalytics/dataengineeringagent/dea_a2a_client.py @@ -0,0 +1,358 @@ +#!/usr/bin/env python + +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Example A2A client for the Data Engineering Agent. + +This script demonstrates how to interact with the Data Engineering Agent (DEA) +using the A2A Python SDK. It handles Agent Card resolution, multi-turn +conversations using conversation tokens persisted to a local file, and +automatic retries on deadline exceeded. + +The script is compatible with DEA v0.3 protocol flavor. +""" + +from __future__ import annotations + +import argparse +import asyncio +import os +import pathlib +import re +from typing import Any, Dict, List, Optional +import uuid + +from a2a import client as a2a_client +from a2a import types as a2a_types +from a2a.client import service_parameters as a2a_service_params +from a2a.utils import constants as a2a_constants +from google.auth import default +from google.auth.transport import requests as auth_requests +from google.protobuf import json_format as pb_json_format +import httpx + +import logging + + +# --- Constants --- + +DEA_EXT_PREFIX = "https://geminidataanalytics.googleapis.com/a2a/extensions" +GCP_RESOURCE_EXT_URI = f"{DEA_EXT_PREFIX}/gcpresource/v1" +CONVERSATION_TOKEN_EXT_URI = f"{DEA_EXT_PREFIX}/conversationtoken/v1" +FINISH_REASON_EXT_URI = f"{DEA_EXT_PREFIX}/finishreason/v1" +INSTRUCTION_EXT_URI = f"{DEA_EXT_PREFIX}/instruction/v1" +MESSAGE_LEVEL_EXT_URI = f"{DEA_EXT_PREFIX}/messagelevel/v1" + +AGENT_NAME = "dataengineeringagent" +AGENT_BASE_URL_TEMPLATE = ( + "https://geminidataanalytics.googleapis.com/v1/a2a/projects/" + "{project_id}/locations/{location}/agents/{agent_name}" +) + + +# --- Helper Functions --- + + +async def force_alt_sse(request: httpx.Request) -> None: + """HTTP request hook that dynamically forces GFE SSE transcoding for streams.""" + if request.url.path.endswith((":stream", ":subscribe")): + request.url = request.url.copy_merge_params({"alt": "sse"}) + + +def get_bearer_token() -> Optional[str]: + """Fetches a Google Cloud bearer token using Application Default Credentials (ADC).""" + try: + credentials, _ = default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + auth_request = auth_requests.Request() + credentials.refresh(auth_request) + return credentials.token + except Exception as e: # pylint: disable=broad-exception-caught + logging.error("Error getting credentials: %s", e) + logging.info( + "Please ensure you have authenticated with 'gcloud auth" + " application-default login'." + ) + return None + + +def load_instructions(paths: List[str]) -> List[Dict[str, str]]: + """Reads instructions from files or directories.""" + instructions = [] + for path_str in paths: + path = pathlib.Path(path_str) + if not path.exists(): + raise FileNotFoundError(f"Instruction path does not exist: {path_str}") + if path.is_file(): + try: + instructions.append( + {"name": path.name, "definition": path.read_text(encoding="utf-8")} + ) + except UnicodeDecodeError: + logging.warning("Skipping non-UTF-8 file: %s", path_str) + elif path.is_dir(): + for file_path in path.iterdir(): + if file_path.is_file() and not file_path.name.startswith("."): + try: + instructions.append({ + "name": file_path.name, + "definition": file_path.read_text(encoding="utf-8"), + }) + except UnicodeDecodeError: + logging.warning("Skipping non-UTF-8 file: %s", file_path) + return instructions + + +# --- Core Interaction Snippet --- + + +# [START geminidataanalytics_a2a_client] +async def interact_with_data_engineering_agent( + auth_token: str, + gcp_resource_id: str, + message: str, + conversation_token: Optional[str] = None, + instructions: Optional[List[Dict[str, str]]] = None, + max_retries: int = 5, +) -> Optional[str]: + """Sends a query to the Data Engineering Agent via the A2A protocol. + + This function implements the complete client communication lifecycle: + 1. Resolves the Agent Card to discover endpoint protocols. + 2. Establishes the A2A client session over HTTP/SSE. + 3. Attaches required metadata extensions (GCP Resource ID, Conversation Token, + instructions). + 4. Invokes the agent and streams responses. + 5. Handles DEADLINE_EXCEEDED finish reason by auto-resuming the turn. + + Args: + auth_token: A Google Cloud access token for authentication. + gcp_resource_id: The target GCP resource ID, formatted as + 'projects/{project}/locations/{location}/...' + message: The natural language query to send to the agent. + conversation_token: Optional conversation token to resume an existing + session. + instructions: Optional custom agent instructions. + + Returns: + The updated conversation token for subsequent turns. + """ + # Extract project_id and location from gcp_resource_id + match = re.search(r"projects/([^/]+)/locations/([^/]+)", gcp_resource_id) + if not match: + raise ValueError( + f"Invalid gcp_resource_id format: {gcp_resource_id}. " + "Expected 'projects/{project}/locations/{location}/...'" + ) + project_id, location = match.groups() + agent_base_url = AGENT_BASE_URL_TEMPLATE.format( + project_id=project_id, + location=location, + agent_name=AGENT_NAME, + ) + + headers = {"Authorization": f"Bearer {auth_token}"} + async with httpx.AsyncClient( + headers=headers, + timeout=600.0, + event_hooks={"request": [force_alt_sse]}, + ) as httpx_client: + + # Resolve the Agent Card from the endpoint + resolver = a2a_client.A2ACardResolver( + httpx_client=httpx_client, base_url=agent_base_url + ) + agent_card = await resolver.get_agent_card(relative_card_path="v1/card") + + # Create the A2A protocol client + config = a2a_client.ClientConfig( + supported_protocol_bindings=[a2a_constants.TransportProtocol.HTTP_JSON], + streaming=True, + httpx_client=httpx_client, + ) + client = await a2a_client.create_client( + agent=agent_card, client_config=config + ) + + try: + # Helper to build request metadata + def build_metadata(token_val: Optional[str]) -> Dict[str, Any]: + meta = {GCP_RESOURCE_EXT_URI: {"gcpResourceId": gcp_resource_id}} + if token_val: + meta[CONVERSATION_TOKEN_EXT_URI] = token_val + if instructions: + meta[INSTRUCTION_EXT_URI] = {"agentInstructions": instructions} + return meta + + context_id = str(uuid.uuid4()) + current_token = conversation_token + query_to_send = message + is_retry = False + retry_count = 0 + + while True: + if is_retry: + logging.info("Resuming session due to deadline exceeded...") + + user_message = a2a_types.Message( + message_id=str(uuid.uuid4()), + role=a2a_types.Role.ROLE_USER, + parts=[a2a_types.Part(text=query_to_send)], + context_id=context_id, + ) + + request = a2a_types.SendMessageRequest( + message=user_message, + metadata=build_metadata(current_token), + ) + + extensions = [ + MESSAGE_LEVEL_EXT_URI, + INSTRUCTION_EXT_URI, + GCP_RESOURCE_EXT_URI, + CONVERSATION_TOKEN_EXT_URI, + FINISH_REASON_EXT_URI, + ] + + service_params = a2a_service_params.ServiceParametersFactory.create( + [a2a_service_params.with_a2a_extensions(extensions)] + ) + context = a2a_client.ClientCallContext( + service_parameters=service_params, timeout=600.0 + ) + + finish_reason: Optional[str] = None + responses = client.send_message(request, context=context) + + async for response_obj in responses: + # Process and print outcomes to console + # (teaches user how to handle output) + print(response_obj) + event = pb_json_format.MessageToDict( + response_obj, preserving_proto_field_name=True + ) + events = event if isinstance(event, list) else [event] + for e in events: + # Extract metadata from different possible A2A event layers + meta = None + if "task" in e: + meta = e["task"].get("metadata", {}) + elif "status_update" in e: + meta = e["status_update"].get("metadata", {}) + elif "message" in e: + meta = e["message"].get("metadata", {}) + + if meta: + # Update conversation token + token_val = meta.get(CONVERSATION_TOKEN_EXT_URI) + if token_val: + extracted_token = ( + token_val.get("conversationToken") + if isinstance(token_val, dict) + else token_val + ) + if extracted_token: + current_token = extracted_token + # Extract finish reason + reason_val = meta.get(FINISH_REASON_EXT_URI) + if reason_val: + finish_reason = ( + reason_val.get("finishReason") + if isinstance(reason_val, dict) + else reason_val + ) + + # Handle deadline exceeded by repeating the turn + if finish_reason == "DEADLINE_EXCEEDED": + retry_count += 1 + if retry_count > max_retries: + logging.error( + "Max retries (%d) exceeded due to repeated DEADLINE_EXCEEDED.", + max_retries, + ) + break + logging.warning("Agent execution hit deadline. Resuming...") + query_to_send = "Please continue." + is_retry = True + else: + break + + return current_token + + finally: + await client.close() +# [END geminidataanalytics_a2a_client] + + +# --- Main CLI Entrypoint --- + + +async def main() -> None: + parser = argparse.ArgumentParser( + description="Data Engineering Agent A2A Client" + ) + parser.add_argument( + "--gcp_resource_id", required=True, help="Target GCP resource ID." + ) + parser.add_argument("--message", required=True, help="The message to send.") + parser.add_argument( + "--conversation_token_path", + help="Local file to persist conversation token.", + ) + parser.add_argument( + "--instruction_path", + action="append", + help="Path to instruction files/dirs.", + ) + + args = parser.parse_args() + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + + token = os.getenv("GOOGLE_ACCESS_TOKEN") or get_bearer_token() + if not token: + logging.error("Authentication failed. No Google access token found.") + return + + conversation_token = None + if args.conversation_token_path: + token_file = pathlib.Path(args.conversation_token_path) + if token_file.exists(): + conversation_token = token_file.read_text(encoding="utf-8").strip() + + try: + instructions = load_instructions(args.instruction_path or []) + new_token = await interact_with_data_engineering_agent( + auth_token=token, + gcp_resource_id=args.gcp_resource_id, + message=args.message, + conversation_token=conversation_token, + instructions=instructions, + ) + + if args.conversation_token_path and new_token: + pathlib.Path(args.conversation_token_path).write_text( + new_token, encoding="utf-8" + ) + + except Exception as e: # pylint: disable=broad-exception-caught + logging.exception("Operation failed: %s", e) + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except KeyboardInterrupt: + pass diff --git a/geminidataanalytics/dataengineeringagent/dea_a2a_client_test.py b/geminidataanalytics/dataengineeringagent/dea_a2a_client_test.py new file mode 100644 index 00000000000..fb4146548c9 --- /dev/null +++ b/geminidataanalytics/dataengineeringagent/dea_a2a_client_test.py @@ -0,0 +1,278 @@ +#!/usr/bin/env python + +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the Data Engineering Agent A2A client.""" + +import unittest +from unittest import mock + +import dea_a2a_client + + +class DeaA2AClientTest(unittest.IsolatedAsyncioTestCase): + + @mock.patch("dea_a2a_client.default") + @mock.patch("dea_a2a_client.auth_requests.Request") + def test_get_bearer_token_success(self, unused_auth_request, mock_default): + mock_creds = mock.MagicMock() + mock_creds.token = "test-token" + mock_default.return_value = (mock_creds, "test-project") + token = dea_a2a_client.get_bearer_token() + self.assertEqual(token, "test-token") + mock_creds.refresh.assert_called_once() + + @mock.patch("dea_a2a_client.default") + def test_get_bearer_token_failure(self, mock_default): + mock_default.side_effect = Exception("Auth failed") + token = dea_a2a_client.get_bearer_token() + self.assertIsNone(token) + + @mock.patch("dea_a2a_client.a2a_client.A2ACardResolver") + @mock.patch("dea_a2a_client.a2a_client.create_client") + @mock.patch("httpx.AsyncClient") + @mock.patch("dea_a2a_client.pb_json_format.MessageToDict") + async def test_multi_turn_and_retry( + self, + mock_message_to_dict, + mock_httpx_cls, + mock_create_client, + mock_resolver_cls, + ): + mock_message_to_dict.side_effect = lambda x, **kwargs: x + + # 1. Mock Agent Card + mock_card = mock.MagicMock() + mock_card.name = "DEA" + mock_interface = mock.MagicMock() + mock_card.supported_interfaces = [mock_interface] + + mock_resolver = mock_resolver_cls.return_value + mock_resolver.get_agent_card = mock.AsyncMock(return_value=mock_card) + + # 2. Mock Client + mock_client = mock.MagicMock() + mock_client.close = mock.AsyncMock() + mock_create_client.return_value = mock_client + + # 3. Define Response Behavior + # Turn 1: Returns DEADLINE_EXCEEDED and a token + token_v1 = "token-v1" + response1 = { + "task": { + "id": "t1", + "metadata": { + dea_a2a_client.CONVERSATION_TOKEN_EXT_URI: { + "conversationToken": token_v1 + }, + dea_a2a_client.FINISH_REASON_EXT_URI: { + "finishReason": "DEADLINE_EXCEEDED" + }, + }, + "history": [ + {"role": "agent", "content": {"text": "Part 1 of response"}} + ], + } + } + + # Turn 2 (Retry of Turn 1): Returns FINISHED + response2 = { + "task": { + "id": "t2", + "metadata": { + dea_a2a_client.CONVERSATION_TOKEN_EXT_URI: { + "conversationToken": "token-v2" + }, + dea_a2a_client.FINISH_REASON_EXT_URI: { + "finishReason": "FINISHED" + }, + }, + "history": [ + {"role": "agent", "content": {"text": "Part 2 of response"}} + ], + } + } + + class AsyncIter: + + def __init__(self, items): + self.items = items + + def __aiter__(self): + return self + + async def __anext__(self): + if not self.items: + raise StopAsyncIteration + return self.items.pop(0) + + # Set side effect to return different responses sequentially + mock_client.send_message.side_effect = [ + AsyncIter([response1]), + AsyncIter([response2]), + ] + + # 4. Run Interaction + token = "fake-token" + gcp_res_id = "projects/tp/locations/tl/repositories/tr/workspaces/default" + result_token = await dea_a2a_client.interact_with_data_engineering_agent( + auth_token=token, + gcp_resource_id=gcp_res_id, + message="Initial Query", + ) + + # 5. Verify + self.assertEqual(result_token, "token-v2") + + # Verify Authorization header was passed in httpx client initialization + mock_httpx_cls.assert_called_once() + called_headers = mock_httpx_cls.call_args[1].get("headers") + self.assertEqual(called_headers, {"Authorization": f"Bearer {token}"}) + + @mock.patch("dea_a2a_client.a2a_client.A2ACardResolver") + @mock.patch("dea_a2a_client.a2a_client.create_client") + @mock.patch("httpx.AsyncClient") + @mock.patch("dea_a2a_client.pb_json_format.MessageToDict") + async def test_session_persistence( + self, + mock_message_to_dict, + unused_httpx, + mock_create_client, + mock_resolver_cls, + ): + mock_message_to_dict.side_effect = lambda x, **kwargs: x + + # Mock Client + mock_client = mock.MagicMock() + mock_client.close = mock.AsyncMock() + mock_create_client.return_value = mock_client + + mock_card = mock.MagicMock() + mock_card.supported_interfaces = [mock.MagicMock()] + mock_resolver_cls.return_value.get_agent_card = mock.AsyncMock( + return_value=mock_card + ) + + response = {"message": {"role": "agent", "content": {"text": "response"}}} + + class AsyncIter: + + def __init__(self, item): + self.item = item + self.used = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self.used: + raise StopAsyncIteration + self.used = True + return self.item + + mock_client.send_message.return_value = AsyncIter(response) + + # Run session with existing token state + token = "fake-token" + gcp_resource_id = "projects/p/locations/l/repositories/r/workspaces/w" + existing_token = "old-token" + await dea_a2a_client.interact_with_data_engineering_agent( + auth_token=token, + gcp_resource_id=gcp_resource_id, + message="query", + conversation_token=existing_token, + ) + + # Verify existing state was sent + call_request = mock_client.send_message.call_args[0][0] + self.assertEqual( + call_request.metadata[dea_a2a_client.CONVERSATION_TOKEN_EXT_URI], + existing_token, + ) + + @mock.patch("dea_a2a_client.a2a_client.A2ACardResolver") + @mock.patch("dea_a2a_client.a2a_client.create_client") + @mock.patch("httpx.AsyncClient") + @mock.patch("dea_a2a_client.pb_json_format.MessageToDict") + async def test_retry_limit_exceeded( + self, + mock_message_to_dict, + unused_httpx, + mock_create_client, + mock_resolver_cls, + ): + mock_message_to_dict.side_effect = lambda x, **kwargs: x + + # Mock Agent Card + mock_card = mock.MagicMock() + mock_card.supported_interfaces = [mock.MagicMock()] + mock_resolver_cls.return_value.get_agent_card = mock.AsyncMock( + return_value=mock_card + ) + + # Mock Client + mock_client = mock.MagicMock() + mock_client.close = mock.AsyncMock() + mock_create_client.return_value = mock_client + + # Define Response (always returns DEADLINE_EXCEEDED) + response = { + "task": { + "id": "t1", + "metadata": { + dea_a2a_client.FINISH_REASON_EXT_URI: { + "finishReason": "DEADLINE_EXCEEDED" + }, + }, + "history": [{"role": "agent", "content": {"text": "Waiting..."}}], + } + } + + class AsyncIter: + + def __init__(self, item): + self.item = item + self.used = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self.used: + raise StopAsyncIteration + self.used = True + return self.item + + # The client will always return a deadline exceeded response + mock_client.send_message.side_effect = lambda *args, **kwargs: AsyncIter(response) + + # Run Interaction with max_retries = 2 + token = "fake-token" + gcp_res_id = "projects/tp/locations/tl/repositories/tr/workspaces/default" + result_token = await dea_a2a_client.interact_with_data_engineering_agent( + auth_token=token, + gcp_resource_id=gcp_res_id, + message="Initial Query", + max_retries=2, + ) + + # Verify it stopped and returned None/original token + self.assertIsNone(result_token) + # The client should have been called 3 times total (1 initial + 2 retries) + self.assertEqual(mock_client.send_message.call_count, 3) + + +if __name__ == "__main__": + unittest.main()