|
| 1 | +# Data Engineering Agent A2A Client Example |
| 2 | + |
| 3 | +This directory contains a sample Python implementation of an |
| 4 | +[Agent-to-Agent (A2A)](https://a2a-protocol.org/) client designed to interact |
| 5 | +with the **Google Cloud Data Engineering Agent (DEA)**. |
| 6 | + |
| 7 | +## Background |
| 8 | + |
| 9 | +The Data Engineering Agent is a BigQuery and Dataform ELT expert capable of |
| 10 | +building, managing, and troubleshooting data pipelines. To enable |
| 11 | +interoperability across different platforms and agents, it exposes an interface |
| 12 | +following the A2A protocol. |
| 13 | + |
| 14 | +Official Documentation: |
| 15 | +[Data Engineering Agent API Overview](https://docs.cloud.google.com/gemini/data-agents/data-engineering-agent/api-overview) |
| 16 | + |
| 17 | +This example demonstrates how to use the open-source |
| 18 | +[A2A Python SDK](https://github.com/a2aproject/a2a-python) to: 1. **Discover** |
| 19 | +the agent's capabilities via its Agent Card. 2. **Authenticate** using Google |
| 20 | +Application Default Credentials (ADC). 3. **Maintain State** across multi-turn |
| 21 | +conversations using Conversation Tokens persisted to a local file. 4. **Handle |
| 22 | +Complex Tasks** by automatically resuming execution when agent finished with |
| 23 | +`DEADLINE_EXCEEDED`. 5. **Configure Extensions** like `Instruction` to customize |
| 24 | +agent behavior. |
| 25 | + |
| 26 | +## Features |
| 27 | + |
| 28 | +- **Native A2A SDK Usage:** Uses `A2ACardResolver` and `create_client` for |
| 29 | + idiomatic protocol interaction. |
| 30 | +- **Streaming-only Execution:** Hardcoded to use response streaming for lowest |
| 31 | + real-time latency and optimal interaction patterns. |
| 32 | +- **Session Persistence:** Automatically saves and loads the |
| 33 | + `conversationToken` from a local file, allowing multi-turn conversations via |
| 34 | + repeated script executions. |
| 35 | +- **Configurable Parameters:** Exposes `gcp_resource_id` for flexibility, |
| 36 | + automatically extracting project and location details. |
| 37 | +- **Instruction Loading:** Automatically reads custom instructions from local |
| 38 | + files or directories, using filenames as the instruction name and file |
| 39 | + content as the definition. |
| 40 | +- **Extension Header Support:** Uses `ServiceParametersFactory` to correctly |
| 41 | + set the `A2A-Extensions` HTTP header required by the agent. |
| 42 | +- **Automated Resumption:** Detects `DEADLINE_EXCEEDED` via the |
| 43 | + `finish_reason` extension and transparently continues the task. |
| 44 | + |
| 45 | +## Prerequisites |
| 46 | + |
| 47 | +- Python 3.10 or higher. |
| 48 | +- [Google Cloud SDK (gcloud)](https://cloud.google.com/sdk/docs/install) |
| 49 | + installed and configured. |
| 50 | +- Enable required APIs |
| 51 | + [Use the Data Engineering Agent to build and modify data pipelines](https://docs.cloud.google.com/bigquery/docs/data-engineering-agent-pipelines#required-apis) |
| 52 | + |
| 53 | +## Setup |
| 54 | + |
| 55 | +1. **Create and activate a virtual environment:** |
| 56 | + |
| 57 | + ```bash |
| 58 | + python3 -m venv .dea |
| 59 | + source .dea/bin/activate |
| 60 | + ``` |
| 61 | + |
| 62 | +2. **Install dependencies:** |
| 63 | + |
| 64 | + ```bash |
| 65 | + pip install -r requirements.txt |
| 66 | + ``` |
| 67 | + |
| 68 | +3. **Authenticate with Google Cloud:** |
| 69 | + |
| 70 | + ```bash |
| 71 | + gcloud auth application-default login |
| 72 | + ``` |
| 73 | + |
| 74 | +## Usage |
| 75 | + |
| 76 | +### Single Message Mode |
| 77 | + |
| 78 | +Sends a single message and exits. `gcp_resource_id` and `message` are required. |
| 79 | + |
| 80 | +``` |
| 81 | +python3 dea_a2a_client.py \ |
| 82 | + --gcp_resource_id projects/my-project/locations/us-central1/repositories/my-repo/workspaces/default \ |
| 83 | + --message "List my Dataform tables" |
| 84 | +``` |
| 85 | +
|
| 86 | +### Multi-turn Conversation (State Persistence) |
| 87 | +
|
| 88 | +To maintain a conversation across multiple calls, use the |
| 89 | +`--conversation_token_path` argument. The script will save the conversation |
| 90 | +token to this file and reload it in subsequent calls. |
| 91 | +
|
| 92 | +``` |
| 93 | +# First turn (starts session) |
| 94 | +python3 dea_a2a_client.py \ |
| 95 | + --gcp_resource_id projects/my-project/locations/us-central1/repositories/my-repo/workspaces/default \ |
| 96 | + --message "hi" \ |
| 97 | + --conversation_token_path ./token.txt |
| 98 | + |
| 99 | +# Second turn (continues previous context) |
| 100 | +python3 dea_a2a_client.py \ |
| 101 | + --gcp_resource_id projects/my-project/locations/us-central1/repositories/my-repo/workspaces/default \ |
| 102 | + --message "Explain the first table" \ |
| 103 | + --conversation_token_path ./token.txt |
| 104 | +``` |
| 105 | +
|
| 106 | +### Advanced: Providing Local Instructions |
| 107 | +
|
| 108 | +You can point the client to local files (e.g., SQL style guides) to influence |
| 109 | +the agent's behavior. |
| 110 | +
|
| 111 | +``` |
| 112 | +python3 dea_a2a_client.py \ |
| 113 | + --gcp_resource_id projects/my-project/locations/us-central1/repositories/my-repo/workspaces/default \ |
| 114 | + --message "List my tables" \ |
| 115 | + --instruction_path ./style_guide.md |
| 116 | +``` |
| 117 | +
|
| 118 | +#### Command-line Arguments |
| 119 | +
|
| 120 | +Argument | Required | Description |
| 121 | +:-------------------------- | :------- | :---------- |
| 122 | +`--gcp_resource_id` | **Yes** | The target resource ID. Supported formats `projects/{p}/locations/{l}/repositories/{r}/workspaces/{w}` (Dataform) |
| 123 | +`--message` | **Yes** | The message to send to the agent. |
| 124 | +`--conversation_token_path` | No | Path to a local file to persist conversation token. Allows for multi-turn conversations. |
| 125 | +`--instruction_path` | No | Path to a file or directory containing instructions. Can be repeated. |
| 126 | +
|
| 127 | +### Running Unit Tests |
| 128 | +
|
| 129 | +Make sure the virtual environment is activated, then run: |
| 130 | +
|
| 131 | +```bash |
| 132 | +python3 dea_a2a_client_test.py |
| 133 | +``` |
| 134 | + |
| 135 | +Alternatively, if the virtual environment is not activated, you can run it |
| 136 | +directly using the venv python: |
| 137 | + |
| 138 | +```bash |
| 139 | +./.dea/bin/python3 dea_a2a_client_test.py |
| 140 | +``` |
0 commit comments