From 217a90a2e6c9725aeaac3dffedca8d63c25037fd Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Wed, 22 Jul 2026 16:04:27 -0700 Subject: [PATCH] feat(eventarc): add Eventarc Advanced toolset for ADK This adds a new integration for Google Cloud Eventarc Advanced. Provides `eventarc_toolset` which allows LLM agents to publish structured CloudEvents. It strictly validates and sanitizes attributes according to the CloudEvents 1.0 specification and supports dynamically resolving fields like `time` and `id` at runtime using lambdas. Includes comprehensive testing and documentation. PiperOrigin-RevId: 952384562 --- .../samples/integrations/eventarc/README.md | 115 ----------- .../samples/integrations/eventarc/__init__.py | 13 -- .../eventarc/domain_specific_agent/README.md | 151 +++++++++++++++ .../eventarc/domain_specific_agent/agent.py | 181 ++++++++++++++++++ .../eventarc/generic_agent/README.md | 92 +++++++++ .../eventarc/{ => generic_agent}/agent.py | 14 +- pyproject.toml | 3 + .../browser/assets/config/runtime-config.json | 2 +- .../adk/integrations/eventarc/__init__.py | 9 + .../adk/integrations/eventarc/_client.py | 50 +++-- .../eventarc/_domain_specific_publish.py | 76 +++----- .../eventarc/_eventarc_toolset.py | 51 +++++ .../integrations/eventarc/_message_tool.py | 14 +- .../integrations/eventarc/test_client.py | 56 +++--- .../eventarc/test_eventarc_toolset.py | 13 ++ .../eventarc/test_message_tool.py | 142 +++++++------- 16 files changed, 682 insertions(+), 300 deletions(-) delete mode 100644 contributing/samples/integrations/eventarc/README.md delete mode 100644 contributing/samples/integrations/eventarc/__init__.py create mode 100644 contributing/samples/integrations/eventarc/domain_specific_agent/README.md create mode 100644 contributing/samples/integrations/eventarc/domain_specific_agent/agent.py create mode 100644 contributing/samples/integrations/eventarc/generic_agent/README.md rename contributing/samples/integrations/eventarc/{ => generic_agent}/agent.py (90%) diff --git a/contributing/samples/integrations/eventarc/README.md b/contributing/samples/integrations/eventarc/README.md deleted file mode 100644 index c9ac0c0b..00000000 --- a/contributing/samples/integrations/eventarc/README.md +++ /dev/null @@ -1,115 +0,0 @@ -# Eventarc Tools Sample - -## Introduction - -This sample agent demonstrates the Eventarc first-party tool in ADK, -distributed via the `google.adk.integrations.eventarc` module. This tool suite currently includes: - -1. `publish_message` - -Publishes a structured event in CloudEvents format to a Google Cloud Eventarc message bus. - -## How to use - -Set up environment variables in your `.env` file for using -[Google AI Studio](https://google.github.io/adk-docs/get-started/quickstart/#gemini---google-ai-studio) -or -[Google Cloud Vertex AI](https://google.github.io/adk-docs/get-started/quickstart/#gemini---google-cloud-vertex-ai) -for the LLM service for your agent. For example, for using Google AI Studio you -would set: - -- GOOGLE_GENAI_USE_VERTEXAI=FALSE -- GOOGLE_API_KEY={your api key} - -### With Application Default Credentials - -This mode is useful for quick development when the agent builder is the only -user interacting with the agent. The tools are run with these credentials. - -1. Create application default credentials on the machine where the agent would - be running by following https://cloud.google.com/docs/authentication/provide-credentials-adc. - -1. Set `CREDENTIALS_TYPE=None` in `agent.py` - -1. Run the agent - -### With Service Account Keys - -This mode is useful for quick development when the agent builder wants to run -the agent with service account credentials. The tools are run with these -credentials. - -1. Create service account key by following https://cloud.google.com/iam/docs/service-account-creds#user-managed-keys. - -1. Set `CREDENTIALS_TYPE=AuthCredentialTypes.SERVICE_ACCOUNT` in `agent.py` - -1. Download the key file and replace `"service_account_key.json"` with the path - -1. Run the agent - -### With Interactive OAuth - -1. Follow - https://developers.google.com/identity/protocols/oauth2#1.-obtain-oauth-2.0-credentials-from-the-dynamic_data.setvar.console_name. - to get your client id and client secret. Be sure to choose "web" as your client - type. - -1. Follow https://developers.google.com/workspace/guides/configure-oauth-consent to add scope "https://www.googleapis.com/auth/cloud-platform". - -1. Follow https://developers.google.com/identity/protocols/oauth2/web-server#creatingcred to add http://localhost/dev-ui/ to "Authorized redirect URIs". - -Note: localhost here is just a hostname that you use to access the dev ui, -replace it with the actual hostname you use to access the dev ui. - -1. For 1st run, allow popup for localhost in Chrome. - -1. Configure your `.env` file to add two more variables before running the agent: - -- OAUTH_CLIENT_ID={your client id} -- OAUTH_CLIENT_SECRET={your client secret} - -Note: don't create a separate .env, instead put it to the same .env file that -stores your Vertex AI or Dev ML credentials - -1. Set `CREDENTIALS_TYPE=AuthCredentialTypes.OAUTH2` in `agent.py` and run the agent - -### With Agent Identity (in Agent Runtime / Vertex AI Reasoning Engine) - -When deploying this agent to Agent Runtime, it can use its unique SPIFFE-based Agent Identity to authenticate. This is the recommended security best practice. - -1. **Configure Deployment**: Create a `.agent_engine_config.json` file in this directory to specify the identity type: - - ```json - { - "identity_type": "AGENT_IDENTITY" - } - ``` - -1. **Use Default Credentials**: Leave `CREDENTIALS_TYPE = None` in `agent.py` (which is the default). This configures the agent to use Application Default Credentials (ADC), which automatically resolves to the Agent Identity in the container runtime environment. - -1. **Deploy the Agent**: Deploy your agent using the ADK CLI: - - ```bash - uv run adk deploy agent_engine \ - --project=YOUR_PROJECT_ID \ - --region=YOUR_REGION \ - --display_name=eventarc-agent-test \ - contributing/samples/integrations/eventarc - ``` - - Take note of the generated **Reasoning Engine ID** (e.g., `1234567890`) and the **Project Number** of your project. - -1. **Grant IAM Permissions**: Grant the Eventarc Message Bus User role (`roles/eventarc.messageBusUser`) to the Agent Identity principal at the project level: - - ```bash - gcloud projects add-iam-policy-binding YOUR_PROJECT_ID \ - --member="principal://agents.global.org-171145599760.system.id.goog/resources/aiplatform/projects/YOUR_PROJECT_NUMBER/locations/YOUR_REGION/reasoningEngines/YOUR_REASONING_ENGINE_ID" \ - --role="roles/eventarc.messageBusUser" - ``` - - *Note: Eventarc Advanced message buses require `roles/eventarc.messageBusUser` for publishing, rather than `roles/eventarc.publisher`.* - -## Sample prompts - -- "Publish an event of type 'com.example.hello' to bus 'projects/my-project/locations/global/messageBuses/my-bus' with data 'Hello World' and source '//my/agent'" -- "Send a JSON payload to Eventarc bus 'projects/my-project/locations/global/messageBuses/my-bus' representing a user sign-up event" diff --git a/contributing/samples/integrations/eventarc/__init__.py b/contributing/samples/integrations/eventarc/__init__.py deleted file mode 100644 index 58d482ea..00000000 --- a/contributing/samples/integrations/eventarc/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# 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. diff --git a/contributing/samples/integrations/eventarc/domain_specific_agent/README.md b/contributing/samples/integrations/eventarc/domain_specific_agent/README.md new file mode 100644 index 00000000..8738bbbd --- /dev/null +++ b/contributing/samples/integrations/eventarc/domain_specific_agent/README.md @@ -0,0 +1,151 @@ +# Eventarc Domain-Specific Agent Sample + +## Overview + +This sample agent demonstrates the `create_publish_tool` factory from the Eventarc first-party tool suite in ADK (`google.adk.integrations.eventarc`). It shows how to create domain-specific, strict-schema publishing tools. This allows you to lock down event routing parameters (e.g. `bus`, `type`, `source`) using `CloudEventAttributesBinding` to static values, runtime lambdas, or selectively-exposed agent fields (`AgentProvided`), while binding the event payload to a strictly validated Pydantic model. This prevents hallucinated routing destinations and guarantees structured JSON event data matching your business domain. + +## Sample Inputs + +- `We just successfully completed a vendor outreach call with customer CUST-883. Resolution notes: All issues resolved.` + +- `Log a dynamic outreach event for customer CUST-992. It should go to the message bus 'projects/your_project/locations/us-central1/messageBuses/outreach-bus' and the subject is 'urgent-outreach'.` + +- `We just successfully completed a vendor outreach call with customer CUST-123. Resolution notes: All issues resolved. This is a high priority outreach.` + +- `Please ping the system with high priority, do not retry.` + +## Graph + +```mermaid +graph TD + DomainAgent[adk_sample_domain_eventarc_agent] -->|calls| StaticTool(complete_outreach_static) + DomainAgent -->|calls| DynamicTool(complete_outreach_dynamic) + DomainAgent -->|calls| LambdaTool(complete_outreach_lambda) + DomainAgent -->|calls| PingTool(ping_system) +``` + +## How To + +### Prerequisites: Set up Eventarc + +Before running the agent, you must enable the Eventarc APIs and create a target Message Bus in your Google Cloud Project. + +1. Enable the Eventarc APIs: + +```bash +gcloud services enable eventarc.googleapis.com eventarcpublishing.googleapis.com +``` + +2. Create a Message Bus: + +```bash +gcloud eventarc message-buses create my-bus \ + --location=us-central1 \ + --logging-config=DEBUG +``` + +*(Make sure to update the `BUS_NAME` variable in `agent.py` to match your actual bus URI).* + +`create_publish_tool` is highly flexible. It uses `pydantic.create_model` to construct the LLM's function signature, encapsulating the `payload_schema` inside an `event_data` parameter and appending any parameter marked with `AgentProvided`. + +### Example A: Fully Statically Bound (Safest) + +The developer locks down all routing. The agent only provides the business data. + +```python +complete_outreach_static_tool = toolset.create_publish_tool( + name="complete_outreach_static", + description="Logs a completed outreach attempt (statically bound routing).", + payload_schema=OutreachContext, + bus=f"projects/{PROJECT_ID}/locations/us-central1/messageBuses/{BUS_NAME}", + ce_attributes_binding=CloudEventAttributesBinding( + type="vendor_outreach.completed", + source="//my-agent/outreach", + ) +) +``` + +**What the Agent Sees:** `complete_outreach_static(event_data: OutreachContext)` + +### Example B: Agent-Provided Attributes (Dynamic) + +The developer forces the agent to decide the routing bus and the event subject based on the conversation context. + +```python +complete_outreach_dynamic_tool = toolset.create_publish_tool( + name="complete_outreach_dynamic", + description="Logs a completed outreach attempt (dynamic routing).", + payload_schema=OutreachContext, + bus=AgentProvided("The full regional bus name"), + ce_attributes_binding=CloudEventAttributesBinding( + type="vendor_outreach.completed", + source="//my-agent/outreach", + subject=AgentProvided("The unique Customer ID being reached out to.") + ) +) +``` + +**What the Agent Sees:** `complete_outreach_dynamic(event_data: OutreachContext, bus: str, subject: str)` + +### Example C: Lambda Execution & Mixed Custom Attributes + +The developer uses Python callables to generate IDs dynamically at runtime. + +```python +def get_custom_trace_id(payload: OutreachContext) -> str: + return f"trace-{payload.customer_id}-{uuid.uuid4().hex[:8]}" + +complete_outreach_lambda_tool = toolset.create_publish_tool( + name="complete_outreach_lambda", + description="Logs a completed outreach attempt.", + payload_schema=OutreachContext, + bus=f"projects/{PROJECT_ID}/locations/us-central1/messageBuses/{BUS_NAME}", + ce_attributes_binding=CloudEventAttributesBinding( + type="vendor_outreach.completed", + source="//my-agent/outreach", + id=get_custom_trace_id, # Executed at runtime + custom_attributes={ + "environment": "production", # Statically bound + "priority": AgentProvided("The priority of the outreach: 'high' or 'low'") + } + ) +) +``` + +**What the Agent Sees:** `complete_outreach_lambda(event_data: OutreachContext, priority: str)` + +### Example D: Empty Payloads & Dynamic Defaults + +The developer wants to emit a simple signal (no business payload). If the agent omits the priority, it is dynamically calculated. + +```python +def default_priority(_: None) -> str: + return "low" + +ping_system_tool = toolset.create_publish_tool( + name="ping_system", + description="Pings the system. No data required.", + payload_schema=None, # No payload! + bus=f"projects/{PROJECT_ID}/locations/us-central1/messageBuses/{BUS_NAME}", + ce_attributes_binding=CloudEventAttributesBinding( + type="system.ping", + source="//my-agent/ping", + custom_attributes={ + "retry": AgentProvided("Whether to retry on failure", default="false"), + "priority": AgentProvided("The priority of the ping", default=default_priority) + } + ) +) +``` + +**What the Agent Sees:** `ping_system(retry: str = "false", priority: str | None = None)` + +### Example E: Handling Reserved Keywords and Strict Validation + +`create_publish_tool` strictly validates custom attributes to ensure they comply with the CloudEvent specification. If your custom attribute name collides with a Python reserved keyword, an implicit pointer, or starts with a digit, the tool factory automatically handles this. It securely modifies the exposed parameter name for the LLM (e.g., `self_` or `_123foo`) to avoid collisions, and translates it back during runtime execution. + +## Next Steps: Building Event-Driven AI Workflows + +Publishing an event to a Message Bus is only the first half of the journey. To route these events to other agents or microservices, you will need to set up Eventarc Pipelines and Enrollments. + +To learn how to connect multiple AI agents together using Eventarc, check out the official codelab: **[Build Event-Driven AI Agents with Eventarc, Cloud Run and ADK](https://codelabs.devsite.corp.google.com/eventarc-ai-agents#0)**. diff --git a/contributing/samples/integrations/eventarc/domain_specific_agent/agent.py b/contributing/samples/integrations/eventarc/domain_specific_agent/agent.py new file mode 100644 index 00000000..2582a140 --- /dev/null +++ b/contributing/samples/integrations/eventarc/domain_specific_agent/agent.py @@ -0,0 +1,181 @@ +# 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. + +import os +import textwrap +import uuid + +from google.adk.agents import llm_agent +from google.adk.auth import auth_credential +from google.adk.integrations.eventarc import AgentProvided +from google.adk.integrations.eventarc import CloudEventAttributesBinding +from google.adk.integrations.eventarc import EventarcCredentialsConfig +from google.adk.integrations.eventarc import EventarcToolConfig +from google.adk.integrations.eventarc import EventarcToolset +from google.adk.integrations.eventarc import OMIT +import google.auth +import pydantic + +PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT", "your_gcp_project_id") +BUS_NAME = os.getenv("EVENTARC_BUS_NAME", "outreach-bus") + +# Define the desired credential type. +# By default use Application Default Credentials (ADC) from the local +# environment, which can be set up by following +# https://cloud.google.com/docs/authentication/provide-credentials-adc. +CREDENTIALS_TYPE = None + +# Define an appropriate application name +EVENTARC_DOMAIN_AGENT_NAME = "adk_sample_domain_eventarc_agent" + + +# Define Eventarc tool config. +tool_config = EventarcToolConfig(project_id=os.getenv("GOOGLE_CLOUD_PROJECT")) + +if CREDENTIALS_TYPE == auth_credential.AuthCredentialTypes.OAUTH2: + credentials_config = EventarcCredentialsConfig( + client_id=os.getenv("OAUTH_CLIENT_ID"), + client_secret=os.getenv("OAUTH_CLIENT_SECRET"), + ) +elif CREDENTIALS_TYPE == auth_credential.AuthCredentialTypes.SERVICE_ACCOUNT: + creds, _ = google.auth.load_credentials_from_file("service_account_key.json") + credentials_config = EventarcCredentialsConfig(credentials=creds) +else: + application_default_credentials, _ = google.auth.default() + credentials_config = EventarcCredentialsConfig( + credentials=application_default_credentials + ) + +toolset = EventarcToolset( + credentials_config=credentials_config, tool_config=tool_config +) + +# --------------------------------------------------------------------------- +# Create Domain-Specific Publish Tools +# --------------------------------------------------------------------------- + + +class OutreachContext(pydantic.BaseModel): + customer_id: str + resolution_notes: str + successful: bool + + +# Example A: Fully Statically Bound (Safest) +# The developer locks down all routing. The agent only provides the business data. +complete_outreach_static_tool = toolset.create_publish_tool( + name="complete_outreach_static", + description="Logs a completed outreach attempt (statically bound routing).", + payload_schema=OutreachContext, + bus=f"projects/{PROJECT_ID}/locations/us-central1/messageBuses/{BUS_NAME}", + ce_attributes_binding=CloudEventAttributesBinding( + type="vendor_outreach.completed", + source="//my-agent/outreach", + ), +) + +# Example B: Agent-Provided Attributes (Dynamic) +# The developer forces the agent to decide the routing bus and the event subject. +complete_outreach_dynamic_tool = toolset.create_publish_tool( + name="complete_outreach_dynamic", + description="Logs a completed outreach attempt (dynamic routing).", + payload_schema=OutreachContext, + bus=AgentProvided( + "The full regional bus name: e.g.," + f" 'projects/{PROJECT_ID}/locations/us-central1/messageBuses/{BUS_NAME}'" + ), + ce_attributes_binding=CloudEventAttributesBinding( + type="vendor_outreach.completed", + source="//my-agent/outreach", + subject=AgentProvided("The unique Customer ID being reached out to."), + ), +) + + +# Example C: Lambda Execution & Mixed Custom Attributes +# The developer uses Python callables to generate IDs dynamically at runtime. +def get_custom_trace_id(payload: OutreachContext) -> str: + return f"trace-{payload.customer_id}-{uuid.uuid4().hex[:8]}" + + +complete_outreach_lambda_tool = toolset.create_publish_tool( + name="complete_outreach_lambda", + description=( + "Logs a completed outreach attempt (with lambda executions and custom" + " attributes)." + ), + payload_schema=OutreachContext, + bus=f"projects/{PROJECT_ID}/locations/us-central1/messageBuses/{BUS_NAME}", + ce_attributes_binding=CloudEventAttributesBinding( + type="vendor_outreach.completed", + source="//my-agent/outreach", + id=get_custom_trace_id, + custom_attributes={ + "environment": "production", + "priority": AgentProvided( + "The priority of the outreach: 'high' or 'low'" + ), + }, + ), +) + + +# Example D: Empty Payloads & Dynamic Defaults +# Emit a simple signal (no business payload). The agent optionally decides the priority. +def default_priority(_: None) -> str: + return "low" + + +ping_system_tool = toolset.create_publish_tool( + name="ping_system", + description="Pings the system. No data required.", + payload_schema=None, + bus=f"projects/{PROJECT_ID}/locations/us-central1/messageBuses/{BUS_NAME}", + ce_attributes_binding=CloudEventAttributesBinding( + type="system.ping", + source="//my-agent/ping", + custom_attributes={ + "retry": AgentProvided( + "Whether to retry on failure", default="false" + ), + "priority": AgentProvided( + "The priority of the ping", default=default_priority + ), + }, + ), +) + +root_agent = llm_agent.LlmAgent( + name=EVENTARC_DOMAIN_AGENT_NAME, + description=( + "Agent configured with domain-specific Eventarc publishing tools." + ), + instruction=textwrap.dedent( + """ You are an e-commerce outreach agent. You can publish specific business events. + You have four tools showing different configurations: + 1. `complete_outreach_static`: Fully static routing. Just pass `event_data`. + 2. `complete_outreach_dynamic`: Provide the `bus`, `subject` and `event_data`. + 3. `complete_outreach_lambda`: Provide the `priority` and `event_data`. + 4. `ping_system`: No payload! Optionally provide `retry` and `priority`. + + When a user gives you an instruction, determine which tool to use and execute it. + """ + ), + tools=[ + complete_outreach_static_tool, + complete_outreach_dynamic_tool, + complete_outreach_lambda_tool, + ping_system_tool, + ], +) diff --git a/contributing/samples/integrations/eventarc/generic_agent/README.md b/contributing/samples/integrations/eventarc/generic_agent/README.md new file mode 100644 index 00000000..4d662624 --- /dev/null +++ b/contributing/samples/integrations/eventarc/generic_agent/README.md @@ -0,0 +1,92 @@ +# Eventarc Generic Agent Sample + +## Overview + +This sample agent demonstrates the Eventarc first-party tool in ADK, distributed via the `google.adk.integrations.eventarc` module. It uses the `publish_message` tool to publish a structured event in CloudEvents format asynchronously to a Google Cloud Eventarc message bus. This exposes the full CloudEvent spec to the agent with connection pooling and caching across calls. + +## Sample Inputs + +- `Publish an event of type 'com.example.hello' to bus 'projects/my-project/locations/global/messageBuses/my-bus' with data 'Hello World' and source '//my/agent'` + +- `Send a JSON payload to Eventarc bus 'projects/my-project/locations/global/messageBuses/my-bus' representing a user sign-up event` + +## Graph + +```mermaid +graph TD + GenericAgent[adk_sample_eventarc_agent] -->|calls| PublishMessageTool(EventarcToolset) +``` + +## How To + +### Prerequisites: Set up Eventarc + +Before running the agent, you must enable the Eventarc APIs and create a target Message Bus in your Google Cloud Project. + +1. Enable the Eventarc APIs: + +```bash +gcloud services enable eventarc.googleapis.com eventarcpublishing.googleapis.com +``` + +2. Create a Message Bus: + +```bash +gcloud eventarc message-buses create my-bus \ + --location=us-central1 \ + --logging-config=DEBUG +``` + +*(Make sure to update the `BUS_NAME` variable in `agent.py` to match your actual bus URI).* + +Set up environment variables in your `.env` file for using Google AI Studio or Google Cloud Vertex AI for the LLM service. For example: + +- `GOOGLE_GENAI_USE_VERTEXAI=FALSE` +- `GOOGLE_API_KEY={your api key}` + +### With Application Default Credentials + +This mode is useful for quick development when the agent builder is the only user interacting with the agent. + +1. Create application default credentials on the machine where the agent would be running (https://cloud.google.com/docs/authentication/provide-credentials-adc). +1. Set `CREDENTIALS_TYPE=None` in `agent.py`. +1. Run the agent. + +### With Service Account Keys + +This mode is useful for running the agent with service account credentials. + +1. Create a service account key (https://cloud.google.com/iam/docs/service-account-creds#user-managed-keys). +1. Set `CREDENTIALS_TYPE=AuthCredentialTypes.SERVICE_ACCOUNT` in `agent.py`. +1. Download the key file and replace `"service_account_key.json"` with the path. +1. Run the agent. + +### With Interactive OAuth + +1. Obtain OAuth 2.0 credentials from the Google Cloud Console. Choose "web" as your client type. +1. Configure OAuth consent to add scope "https://www.googleapis.com/auth/cloud-platform". +1. Add `http://localhost/dev-ui/` to "Authorized redirect URIs". +1. Configure your `.env` file with `OAUTH_CLIENT_ID` and `OAUTH_CLIENT_SECRET`. +1. Set `CREDENTIALS_TYPE=AuthCredentialTypes.OAUTH2` in `agent.py` and run the agent. + +### With Agent Identity (in Agent Runtime / Vertex AI Reasoning Engine) + +When deploying this agent to Agent Runtime, it can use its unique SPIFFE-based Agent Identity to authenticate. + +1. **Configure Deployment**: Create a `.agent_engine_config.json` file in the specific agent's directory to specify `"identity_type": "AGENT_IDENTITY"`. +1. **Use Default Credentials**: Leave `CREDENTIALS_TYPE = None` in `agent.py`. +1. **Deploy the Agent**: Deploy your agent using the ADK CLI: + ```bash + uv run adk deploy agent_engine \ + --project=YOUR_PROJECT_ID \ + --region=YOUR_REGION \ + --display_name=eventarc-agent-test \ + contributing/samples/integrations/eventarc/generic_agent + ``` +1. **Grant IAM Permissions**: Grant the Eventarc Message Bus User role (`roles/eventarc.messageBusUser`) to the Agent Identity principal at the project level. + +## Next Steps: Building Event-Driven AI Workflows + +Publishing an event to a Message Bus is only the first half of the journey. To route these events to other agents or microservices, you will need to set up Eventarc Pipelines and Enrollments. + +To learn how to connect multiple AI agents together using Eventarc, check out the official codelab: **[Build Event-Driven AI Agents with Eventarc, Cloud Run and ADK](https://codelabs.devsite.corp.google.com/eventarc-ai-agents#0)**. diff --git a/contributing/samples/integrations/eventarc/agent.py b/contributing/samples/integrations/eventarc/generic_agent/agent.py similarity index 90% rename from contributing/samples/integrations/eventarc/agent.py rename to contributing/samples/integrations/eventarc/generic_agent/agent.py index 64cd4736..f2bccdcd 100644 --- a/contributing/samples/integrations/eventarc/agent.py +++ b/contributing/samples/integrations/eventarc/generic_agent/agent.py @@ -15,8 +15,8 @@ import os import textwrap -from google.adk.agents.llm_agent import LlmAgent -from google.adk.auth.auth_credential import AuthCredentialTypes +from google.adk.agents import llm_agent +from google.adk.auth import auth_credential from google.adk.integrations.eventarc import EventarcCredentialsConfig from google.adk.integrations.eventarc import EventarcToolConfig from google.adk.integrations.eventarc import EventarcToolset @@ -36,7 +36,7 @@ EVENTARC_AGENT_NAME = "adk_sample_eventarc_agent" # You can optionally set the project_id here, or let the agent infer it from context/user input. tool_config = EventarcToolConfig(project_id=os.getenv("GOOGLE_CLOUD_PROJECT")) -if CREDENTIALS_TYPE == AuthCredentialTypes.OAUTH2: +if CREDENTIALS_TYPE == auth_credential.AuthCredentialTypes.OAUTH2: # Initialize the tools to do interactive OAuth # The environment variables OAUTH_CLIENT_ID and OAUTH_CLIENT_SECRET # must be set @@ -44,7 +44,7 @@ if CREDENTIALS_TYPE == AuthCredentialTypes.OAUTH2: client_id=os.getenv("OAUTH_CLIENT_ID"), client_secret=os.getenv("OAUTH_CLIENT_SECRET"), ) -elif CREDENTIALS_TYPE == AuthCredentialTypes.SERVICE_ACCOUNT: +elif CREDENTIALS_TYPE == auth_credential.AuthCredentialTypes.SERVICE_ACCOUNT: # Initialize the tools to use the credentials in the service account key. # If this flow is enabled, make sure to replace the file path with your own # service account key file @@ -59,13 +59,13 @@ else: credentials=application_default_credentials ) -eventarc_toolset = EventarcToolset( +toolset = EventarcToolset( credentials_config=credentials_config, tool_config=tool_config ) # The variable name `root_agent` determines what your root agent is for the # debug CLI -root_agent = LlmAgent( +root_agent = llm_agent.LlmAgent( name=EVENTARC_AGENT_NAME, description=( "Agent to publish structured CloudEvents to Google Cloud Eventarc." @@ -74,5 +74,5 @@ root_agent = LlmAgent( You are a cloud engineer agent with access to Google Cloud Eventarc tools. You can publish CloudEvents structured messages to Eventarc message buses. """), - tools=[eventarc_toolset], + tools=[toolset], ) diff --git a/pyproject.toml b/pyproject.toml index 5e18fd92..8c9436ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,6 +75,7 @@ optional-dependencies.all = [ "google-cloud-bigtable>=2.39.1", "google-cloud-dataplex>=1.7,<3", "google-cloud-discoveryengine>=0.13.12,<0.14", + "google-cloud-eventarc-publishing>=0.10,<1", "google-cloud-parametermanager>=0.4,<1", "google-cloud-pubsub>=2,<3", "google-cloud-resource-manager>=1.12,<2", @@ -173,6 +174,7 @@ optional-dependencies.gcp = [ "google-cloud-bigtable>=2.39.1", "google-cloud-dataplex>=1.7,<3", "google-cloud-discoveryengine>=0.13.12,<0.14", + "google-cloud-eventarc-publishing>=0.10,<1", "google-cloud-parametermanager>=0.4,<1", "google-cloud-pubsub>=2,<3", "google-cloud-resource-manager>=1.12,<2", @@ -217,6 +219,7 @@ optional-dependencies.test = [ "google-cloud-bigtable>=2.39.1", "google-cloud-dataplex>=1.7,<3", "google-cloud-discoveryengine>=0.13.12,<0.14", + "google-cloud-eventarc-publishing>=0.10,<1", "google-cloud-firestore>=2.11,<3", "google-cloud-iamconnectorcredentials>=0.1,<0.2", "google-cloud-parametermanager>=0.4,<1", diff --git a/src/google/adk/cli/browser/assets/config/runtime-config.json b/src/google/adk/cli/browser/assets/config/runtime-config.json index c8f49d88..e2628ca7 100644 --- a/src/google/adk/cli/browser/assets/config/runtime-config.json +++ b/src/google/adk/cli/browser/assets/config/runtime-config.json @@ -1,3 +1,3 @@ { "backendUrl": "" -} \ No newline at end of file +} diff --git a/src/google/adk/integrations/eventarc/__init__.py b/src/google/adk/integrations/eventarc/__init__.py index 61d0fd79..931e8676 100644 --- a/src/google/adk/integrations/eventarc/__init__.py +++ b/src/google/adk/integrations/eventarc/__init__.py @@ -16,10 +16,19 @@ from ._config import EventarcCredentialsConfig from ._config import EventarcToolConfig +# pylint: disable=g-importing-member +from ._domain_specific_publish import AgentProvided +from ._domain_specific_publish import CloudEventAttributesBinding +from ._domain_specific_publish import MISSING +from ._domain_specific_publish import OMIT from ._eventarc_toolset import EventarcToolset __all__ = [ + "AgentProvided", + "CloudEventAttributesBinding", "EventarcCredentialsConfig", "EventarcToolConfig", "EventarcToolset", + "MISSING", + "OMIT", ] diff --git a/src/google/adk/integrations/eventarc/_client.py b/src/google/adk/integrations/eventarc/_client.py index e83d05a1..f7564975 100644 --- a/src/google/adk/integrations/eventarc/_client.py +++ b/src/google/adk/integrations/eventarc/_client.py @@ -18,6 +18,7 @@ from __future__ import annotations import collections import hashlib +import inspect import os import threading import time @@ -27,15 +28,17 @@ from google.api_core.gapic_v1 import client_info if typing.TYPE_CHECKING: from google.cloud import eventarc_publishing_v1 # type: ignore - from google.cloud.eventarc_publishing_v1 import PublisherClient # type: ignore + from google.cloud.eventarc_publishing_v1 import PublisherAsyncClient # type: ignore else: try: from google.cloud import eventarc_publishing_v1 # type: ignore - PublisherClient = eventarc_publishing_v1.PublisherClient + PublisherAsyncClient = getattr( + eventarc_publishing_v1, "PublisherAsyncClient", typing.Any + ) except ImportError: eventarc_publishing_v1 = None - PublisherClient = typing.Any + PublisherAsyncClient = typing.Any try: from google.adk import version # type: ignore @@ -48,17 +51,19 @@ _CACHE_TTL = 1800 # 30 minutes _CACHE_MAX_SIZE = 10 _publisher_client_cache: collections.OrderedDict[ - tuple[str | None, str, int, str], tuple[PublisherClient, float] + tuple[str | None, str, int, str], tuple[PublisherAsyncClient, float] ] = collections.OrderedDict() _publisher_client_lock = threading.Lock() -def _close_client(client: typing.Any) -> None: +async def _close_client(client: typing.Any) -> None: """Explicitly closes the gRPC transport channel of the client.""" transport = getattr(client, "transport", None) if transport is not None and hasattr(transport, "close"): try: - transport.close() + res = transport.close() + if inspect.isawaitable(res): + await res except Exception: # pylint: disable=broad-except pass @@ -131,12 +136,12 @@ def _get_cache_key( return (project_id, final_user_agent, os.getpid(), cred_id) -def get_publisher_client( +async def get_publisher_client( *, credentials: typing.Any, user_agent: str | None = None, project_id: str | None = None, -) -> PublisherClient: +) -> PublisherAsyncClient: """Gets or creates a publisher client for Eventarc.""" if eventarc_publishing_v1 is None: raise RuntimeError("google-cloud-eventarc-publishing is not installed") @@ -148,6 +153,7 @@ def get_publisher_client( project_id=project_id, ) current_time = time.time() + old_client_to_close = None with _publisher_client_lock: client_entry = _publisher_client_cache.get(cache_key) if client_entry is not None: @@ -161,22 +167,25 @@ def get_publisher_client( info = client_info.ClientInfo(user_agent=final_user_agent) # type: ignore[no-untyped-call] client = typing.cast( - PublisherClient, - eventarc_publishing_v1.PublisherClient( + PublisherAsyncClient, + eventarc_publishing_v1.PublisherAsyncClient( credentials=credentials, client_info=info, ), ) if len(_publisher_client_cache) >= _CACHE_MAX_SIZE: - _, (old_client, _) = _publisher_client_cache.popitem(last=False) - _close_client(old_client) + _, (old_client_to_close, _) = _publisher_client_cache.popitem(last=False) _publisher_client_cache[cache_key] = (client, current_time + _CACHE_TTL) - return client + + if old_client_to_close is not None: + await _close_client(old_client_to_close) + + return client -def remove_publisher_client( +async def remove_publisher_client( *, credentials: typing.Any, user_agent: str | None = None, @@ -192,5 +201,14 @@ def remove_publisher_client( with _publisher_client_lock: entry = _publisher_client_cache.pop(cache_key, None) - if entry is not None: - _close_client(entry[0]) + if entry is not None: + await _close_client(entry[0]) + + +async def cleanup_clients() -> None: + """Cleans up all cached publisher clients.""" + with _publisher_client_lock: + clients = list(_publisher_client_cache.values()) + _publisher_client_cache.clear() + for client, _ in clients: + await _close_client(client) diff --git a/src/google/adk/integrations/eventarc/_domain_specific_publish.py b/src/google/adk/integrations/eventarc/_domain_specific_publish.py index 3d10f57b..71ded396 100644 --- a/src/google/adk/integrations/eventarc/_domain_specific_publish.py +++ b/src/google/adk/integrations/eventarc/_domain_specific_publish.py @@ -53,64 +53,42 @@ class AgentProvided: default: Any | OmitSentinel | MissingSentinel = MISSING +AttributeBinding = str | Callable[[Any], str] | AgentProvided +OptionalAttributeBinding = ( + str + | Callable[[Any], str | OmitSentinel] + | AgentProvided + | OmitSentinel + | MissingSentinel + | None +) +CustomAttributeBinding = ( + str | Callable[[Any], str | OmitSentinel] | AgentProvided | OmitSentinel +) +SpecVersionBinding = ( + str | Callable[[Any], str] | AgentProvided | MissingSentinel | None +) + + @dataclass class CloudEventAttributesBinding: """Configuration for binding CloudEvent attributes to static values, lambdas, or AgentProvided fields.""" - type: str | Callable[[Any], str] | AgentProvided - source: str | Callable[[Any], str] | AgentProvided - datacontenttype: ( - str - | Callable[[Any], str | OmitSentinel] - | AgentProvided - | OmitSentinel - | MissingSentinel - | None - ) = MISSING - subject: ( - str - | Callable[[Any], str | OmitSentinel] - | AgentProvided - | OmitSentinel - | MissingSentinel - | None - ) = MISSING - time: ( - str - | Callable[[Any], str | OmitSentinel] - | AgentProvided - | OmitSentinel - | MissingSentinel - | None - ) = MISSING - specversion: ( - str | Callable[[Any], str] | AgentProvided | MissingSentinel | None - ) = MISSING - id: ( - str - | Callable[[Any], str | OmitSentinel] - | AgentProvided - | OmitSentinel - | MissingSentinel - | None - ) = MISSING - custom_attributes: ( - dict[ - str, - str - | Callable[[Any], str | OmitSentinel] - | AgentProvided - | OmitSentinel, - ] - | None - ) = None + type: AttributeBinding + source: AttributeBinding + datacontenttype: OptionalAttributeBinding = MISSING + subject: OptionalAttributeBinding = MISSING + time: OptionalAttributeBinding = MISSING + specversion: SpecVersionBinding = MISSING + id: OptionalAttributeBinding = MISSING + custom_attributes: dict[str, CustomAttributeBinding] | None = None def build_domain_specific_tool( toolset: Any, # Typed as Any to avoid circular import with EventarcToolset name: str, description: str, - bus: str | Callable[[Any], str] | AgentProvided, + bus: AttributeBinding, ce_attributes_binding: CloudEventAttributesBinding, payload_schema: type[pydantic.BaseModel] | None = None, ) -> GoogleTool: @@ -361,7 +339,7 @@ def build_domain_specific_tool( if custom_attr_dict: publish_kwargs["custom_attributes"] = custom_attr_dict - return publish_message(**publish_kwargs) # type: ignore[arg-type] + return await publish_message(**publish_kwargs) # type: ignore[arg-type] # Attach signature and annotations _execute.__signature__ = inspect.Signature(parameters=parameters) # type: ignore[attr-defined] diff --git a/src/google/adk/integrations/eventarc/_eventarc_toolset.py b/src/google/adk/integrations/eventarc/_eventarc_toolset.py index 15da29c0..2ba60462 100644 --- a/src/google/adk/integrations/eventarc/_eventarc_toolset.py +++ b/src/google/adk/integrations/eventarc/_eventarc_toolset.py @@ -17,7 +17,13 @@ from __future__ import annotations from typing import Any +from typing import Callable +import pydantic +from typing_extensions import override + +from . import _client as eventarc_client +from . import _domain_specific_publish as domain_specific_publish from ...features import experimental from ...features import FeatureName from ...tools.base_tool import BaseTool @@ -73,3 +79,48 @@ class EventarcToolset(BaseToolset): for tool in self._tools if self._is_tool_selected(tool, readonly_context) ] + + def create_publish_tool( + self, + *, + name: str, + description: str, + bus: str | Callable[[Any], str] | domain_specific_publish.AgentProvided, + ce_attributes_binding: domain_specific_publish.CloudEventAttributesBinding, + payload_schema: type[pydantic.BaseModel] | None = None, + ) -> GoogleTool: + """Creates a domain-specific publish tool with static or dynamic bindings. + + This acts as a wrapper around the generic `publish_message` tool, allowing + developers to lock down specific CloudEvent attributes (like `bus`, `type`, + or `source`) or make them dynamically generated based on the payload. + + Args: + name: The name of the tool as exposed to the LLM agent. + description: A prompt-friendly description of what this tool does. + bus: The GCP Eventarc Advanced bus resource name. Can be static, + callable, or agent-provided. + ce_attributes_binding: The configuration mapping CloudEvent attributes + to static values, runtime lambdas, or AgentProvided fields. + payload_schema: An optional Pydantic BaseModel representing the expected + structured data payload. If provided, the LLM will be forced to + provide this structured data. + + Returns: + A GoogleTool instance that can be attached to an agent. + """ + tool = domain_specific_publish.build_domain_specific_tool( + toolset=self, + name=name, + description=description, + bus=bus, + ce_attributes_binding=ce_attributes_binding, + payload_schema=payload_schema, + ) + self._tools.append(tool) + return tool + + @override + async def close(self) -> None: + """Clean up resources used by the toolset.""" + await eventarc_client.cleanup_clients() diff --git a/src/google/adk/integrations/eventarc/_message_tool.py b/src/google/adk/integrations/eventarc/_message_tool.py index e7727622..d0a3c050 100644 --- a/src/google/adk/integrations/eventarc/_message_tool.py +++ b/src/google/adk/integrations/eventarc/_message_tool.py @@ -44,7 +44,7 @@ from . import _client as eventarc_client from . import _config as config -def publish_message( +async def publish_message( *, bus: str, type: str, @@ -251,6 +251,9 @@ def publish_message( if time_attr: custom_attr["time"] = time_attr + if subject: + custom_attr["subject"] = subject + # Prepare CloudEvent attributes attributes = {} for k, v in custom_attr.items(): @@ -269,9 +272,6 @@ def publish_message( "attributes": attributes, } - if subject: - event_kwargs["subject"] = subject - if text_data is not None: event_kwargs["text_data"] = text_data if binary_data is not None: @@ -282,7 +282,7 @@ def publish_message( project_id = settings.project_id if settings else None try: - client = eventarc_client.get_publisher_client( + client = await eventarc_client.get_publisher_client( credentials=credentials, project_id=project_id ) @@ -290,11 +290,11 @@ def publish_message( message_bus=bus, proto_message=event ) timeout = settings.publish_timeout if settings else 15.0 - client.publish(request=request, timeout=timeout) + await client.publish(request=request, timeout=timeout) return {"status": "SUCCESS", "message_id": id} except Exception as e: - eventarc_client.remove_publisher_client( + await eventarc_client.remove_publisher_client( credentials=credentials, project_id=project_id ) return {"status": "ERROR", "error_details": repr(e)} diff --git a/tests/unittests/integrations/eventarc/test_client.py b/tests/unittests/integrations/eventarc/test_client.py index df466940..5e404521 100644 --- a/tests/unittests/integrations/eventarc/test_client.py +++ b/tests/unittests/integrations/eventarc/test_client.py @@ -28,7 +28,7 @@ import google.oauth2.credentials import google.oauth2.service_account -class TestEventarcClient(unittest.TestCase): +class TestEventarcClient(unittest.IsolatedAsyncioTestCase): def test_get_credential_id(self): # Service Account @@ -225,7 +225,7 @@ class TestEventarcClient(unittest.TestCase): ) @mock.patch.object(client, "eventarc_publishing_v1", autospec=True) - def test_get_publisher_client_cache(self, mock_eventarc_publishing): + async def test_get_publisher_client_cache(self, mock_eventarc_publishing): # Reset cache client._publisher_client_cache.clear() @@ -234,31 +234,31 @@ class TestEventarcClient(unittest.TestCase): ) mock_client_cls = mock.Mock() - mock_eventarc_publishing.PublisherClient = mock_client_cls + mock_eventarc_publishing.PublisherAsyncClient = mock_client_cls # Return a new mock instance each time mock_client_cls.side_effect = lambda **kwargs: mock.Mock() # First call creates the client - c1 = client.get_publisher_client(credentials=creds, project_id="p1") + c1 = await client.get_publisher_client(credentials=creds, project_id="p1") mock_client_cls.assert_called_once() # Second call returns cached client - c2 = client.get_publisher_client(credentials=creds, project_id="p1") + c2 = await client.get_publisher_client(credentials=creds, project_id="p1") mock_client_cls.assert_called_once() self.assertIs(c1, c2) # Different project creates new client - c3 = client.get_publisher_client(credentials=creds, project_id="p2") + c3 = await client.get_publisher_client(credentials=creds, project_id="p2") self.assertEqual(mock_client_cls.call_count, 2) self.assertIsNot(c1, c3) @mock.patch.object(client, "eventarc_publishing_v1", autospec=True) - def test_remove_publisher_client(self, mock_eventarc_publishing): + async def test_remove_publisher_client(self, mock_eventarc_publishing): client._publisher_client_cache.clear() mock_client_cls = mock.Mock() - mock_eventarc_publishing.PublisherClient = mock_client_cls + mock_eventarc_publishing.PublisherAsyncClient = mock_client_cls mock_client = mock.Mock() mock_client.transport = mock.Mock() mock_client_cls.return_value = mock_client @@ -266,24 +266,26 @@ class TestEventarcClient(unittest.TestCase): creds = mock.create_autospec( google.auth.credentials.Credentials, instance=True ) - c1 = client.get_publisher_client(credentials=creds, project_id="p1") + c1 = await client.get_publisher_client(credentials=creds, project_id="p1") self.assertEqual(len(client._publisher_client_cache), 1) # Remove client - client.remove_publisher_client(credentials=creds, project_id="p1") + await client.remove_publisher_client(credentials=creds, project_id="p1") self.assertEqual(len(client._publisher_client_cache), 0) mock_client.transport.close.assert_called_once() # Remove again is safe - client.remove_publisher_client(credentials=creds, project_id="p1") + await client.remove_publisher_client(credentials=creds, project_id="p1") @mock.patch.object(client, "eventarc_publishing_v1", autospec=True) - def test_publisher_client_cache_lru_eviction(self, mock_eventarc_publishing): + async def test_publisher_client_cache_lru_eviction( + self, mock_eventarc_publishing + ): """Verifies LRU eviction and transport closing when cache is full.""" client._publisher_client_cache.clear() mock_client_cls = mock.Mock() - mock_eventarc_publishing.PublisherClient = mock_client_cls + mock_eventarc_publishing.PublisherAsyncClient = mock_client_cls # Track created mock clients and mock their transports clients_list = [] @@ -302,15 +304,17 @@ class TestEventarcClient(unittest.TestCase): # Fill cache to MAX_SIZE for i in range(client._CACHE_MAX_SIZE): - client.get_publisher_client(credentials=creds, project_id=f"project-{i}") + await client.get_publisher_client( + credentials=creds, project_id=f"project-{i}" + ) # Hit project-0 to make it recently used - client.get_publisher_client(credentials=creds, project_id="project-0") + await client.get_publisher_client(credentials=creds, project_id="project-0") # Now project-1 should be the oldest. # Add another client to trigger eviction next_proj = f"project-{client._CACHE_MAX_SIZE}" - client.get_publisher_client(credentials=creds, project_id=next_proj) + await client.get_publisher_client(credentials=creds, project_id=next_proj) self.assertEqual( len(client._publisher_client_cache), client._CACHE_MAX_SIZE @@ -321,21 +325,21 @@ class TestEventarcClient(unittest.TestCase): mock_client_cls.reset_mock() # project-1 should be evicted - client.get_publisher_client(credentials=creds, project_id="project-1") + await client.get_publisher_client(credentials=creds, project_id="project-1") mock_client_cls.assert_called_once() mock_client_cls.reset_mock() # project-0 should still be in cache - client.get_publisher_client(credentials=creds, project_id="project-0") + await client.get_publisher_client(credentials=creds, project_id="project-0") mock_client_cls.assert_not_called() @mock.patch.object(client, "eventarc_publishing_v1", autospec=True) - def test_get_publisher_client_cache_external_account( + async def test_get_publisher_client_cache_external_account( self, mock_eventarc_publishing ): client._publisher_client_cache.clear() mock_client_cls = mock.Mock() - mock_eventarc_publishing.PublisherClient = mock_client_cls + mock_eventarc_publishing.PublisherAsyncClient = mock_client_cls mock_client_cls.side_effect = lambda **kwargs: mock.Mock() creds1 = google.auth.identity_pool.Credentials( @@ -349,21 +353,21 @@ class TestEventarcClient(unittest.TestCase): credential_source={"file": "path1"}, ) - c1 = client.get_publisher_client(credentials=creds1, project_id="p1") + c1 = await client.get_publisher_client(credentials=creds1, project_id="p1") mock_client_cls.assert_called_once() - c2 = client.get_publisher_client(credentials=creds2, project_id="p1") + c2 = await client.get_publisher_client(credentials=creds2, project_id="p1") # Should be a cache hit, so call count remains 1 mock_client_cls.assert_called_once() self.assertIs(c1, c2) @mock.patch.object(client, "eventarc_publishing_v1", autospec=True) - def test_get_publisher_client_cache_user_credentials( + async def test_get_publisher_client_cache_user_credentials( self, mock_eventarc_publishing ): client._publisher_client_cache.clear() mock_client_cls = mock.Mock() - mock_eventarc_publishing.PublisherClient = mock_client_cls + mock_eventarc_publishing.PublisherAsyncClient = mock_client_cls mock_client_cls.side_effect = lambda **kwargs: mock.Mock() creds1 = google.oauth2.credentials.Credentials( @@ -381,10 +385,10 @@ class TestEventarcClient(unittest.TestCase): client_secret="secret1", ) - c1 = client.get_publisher_client(credentials=creds1, project_id="p1") + c1 = await client.get_publisher_client(credentials=creds1, project_id="p1") mock_client_cls.assert_called_once() - c2 = client.get_publisher_client(credentials=creds2, project_id="p1") + c2 = await client.get_publisher_client(credentials=creds2, project_id="p1") # Should be a cache hit, so call count remains 1 mock_client_cls.assert_called_once() self.assertIs(c1, c2) diff --git a/tests/unittests/integrations/eventarc/test_eventarc_toolset.py b/tests/unittests/integrations/eventarc/test_eventarc_toolset.py index 95c6f963..152722c3 100644 --- a/tests/unittests/integrations/eventarc/test_eventarc_toolset.py +++ b/tests/unittests/integrations/eventarc/test_eventarc_toolset.py @@ -15,9 +15,11 @@ import unittest +from unittest import mock import warnings from google.adk.features._feature_registry import _WARNED_FEATURES +from google.adk.integrations.eventarc import _eventarc_toolset as eventarc_toolset from google.adk.integrations.eventarc import EventarcCredentialsConfig from google.adk.integrations.eventarc import EventarcToolConfig from google.adk.integrations.eventarc import EventarcToolset @@ -58,6 +60,17 @@ class TestEventarcToolset(unittest.IsolatedAsyncioTestCase): self.assertEqual(len(tools), 1) self.assertEqual(tools[0].name, "publish_message") + @mock.patch.object(eventarc_toolset, "eventarc_client", autospec=True) + async def test_close_cleans_up_clients(self, mock_client): + toolset = EventarcToolset( + credentials_config=EventarcCredentialsConfig( + credentials=google.oauth2.credentials.Credentials(token="fake") + ) + ) + mock_client.cleanup_clients = mock.AsyncMock() + await toolset.close() + mock_client.cleanup_clients.assert_called_once() + def test_eventarc_toolset_experimental_warning(self): _WARNED_FEATURES.clear() with warnings.catch_warnings(record=True) as w: diff --git a/tests/unittests/integrations/eventarc/test_message_tool.py b/tests/unittests/integrations/eventarc/test_message_tool.py index 7bfeea11..7ed689c6 100644 --- a/tests/unittests/integrations/eventarc/test_message_tool.py +++ b/tests/unittests/integrations/eventarc/test_message_tool.py @@ -26,16 +26,18 @@ from google.adk.integrations.eventarc import _message_tool as message_tool import google.oauth2.credentials -class TestMessageTool(unittest.TestCase): +class TestMessageTool(unittest.IsolatedAsyncioTestCase): def setUp(self): self.mock_client_module = mock.patch.object( message_tool, "eventarc_client", autospec=True ).start() self.mock_publisher_client = mock.MagicMock(spec=["publish"]) - self.mock_client_module.get_publisher_client.return_value = ( - self.mock_publisher_client + self.mock_publisher_client.publish = mock.AsyncMock() + self.mock_client_module.get_publisher_client = mock.AsyncMock( + return_value=self.mock_publisher_client ) + self.mock_client_module.remove_publisher_client = mock.AsyncMock() self.mock_eventarc_v1 = mock.patch.object( message_tool, "eventarc_publishing_v1", autospec=True ).start() @@ -46,8 +48,8 @@ class TestMessageTool(unittest.TestCase): def tearDown(self): mock.patch.stopall() - def test_publish_message_success_text(self): - res = message_tool.publish_message( + async def test_publish_message_success_text(self): + res = await message_tool.publish_message( bus="projects/test/locations/global/messageBuses/my-bus", type="com.example.test", source="//test/source", @@ -63,11 +65,11 @@ class TestMessageTool(unittest.TestCase): credentials=self.credentials, project_id="test-project" ) - def test_publish_message_custom_timeout(self): + async def test_publish_message_custom_timeout(self): custom_settings = config.EventarcToolConfig( project_id="test-project", publish_timeout=30.0 ) - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="projects/test/locations/global/messageBuses/my-bus", type="com.example.test", source="//test/source", @@ -80,8 +82,8 @@ class TestMessageTool(unittest.TestCase): call_kwargs = self.mock_publisher_client.publish.call_args.kwargs self.assertEqual(call_kwargs.get("timeout"), 30.0) - def test_publish_message_success_json(self): - res = message_tool.publish_message( + async def test_publish_message_success_json(self): + res = await message_tool.publish_message( bus="projects/test/locations/global/messageBuses/my-bus", type="com.example.test", source="//test/source", @@ -91,9 +93,9 @@ class TestMessageTool(unittest.TestCase): ) self.assertEqual(res["status"], "SUCCESS") - def test_publish_message_base64_encoded(self): + async def test_publish_message_base64_encoded(self): encoded_data = base64.b64encode(b"binary data").decode("utf-8") - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="projects/test/locations/global/messageBuses/my-bus", type="com.example.test", source="//test/source", @@ -104,8 +106,8 @@ class TestMessageTool(unittest.TestCase): ) self.assertEqual(res["status"], "SUCCESS") - def test_publish_message_invalid_base64(self): - res = message_tool.publish_message( + async def test_publish_message_invalid_base64(self): + res = await message_tool.publish_message( bus="projects/test/locations/global/messageBuses/my-bus", type="com.example.test", source="//test/source", @@ -117,11 +119,11 @@ class TestMessageTool(unittest.TestCase): self.assertEqual(res["status"], "ERROR") self.assertIn("Invalid base64", res["error_details"]) - def test_publish_message_unserializable_json(self): + async def test_publish_message_unserializable_json(self): class CustomClass: pass - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="projects/test/locations/global/messageBuses/my-bus", type="com.example.test", source="//test/source", @@ -132,7 +134,7 @@ class TestMessageTool(unittest.TestCase): self.assertEqual(res["status"], "ERROR") self.assertIn("Failed to serialize data", res["error_details"]) - def test_publish_message_invalid_inputs(self): + async def test_publish_message_invalid_inputs(self): cases = [ { "name": "invalid_type", @@ -188,11 +190,11 @@ class TestMessageTool(unittest.TestCase): "settings": self.settings, } kwargs.update(case["update_kwargs"]) - res = message_tool.publish_message(**kwargs) + res = await message_tool.publish_message(**kwargs) self.assertEqual(res["status"], "ERROR") self.assertIn(case["expected_error"], res["error_details"]) - def test_publish_message_time_valid_rfc3339(self): + async def test_publish_message_time_valid_rfc3339(self): valid_times = [ "2026-06-03T12:00:00Z", "2026-06-03T12:00:00.123456Z", @@ -203,7 +205,7 @@ class TestMessageTool(unittest.TestCase): for valid_time in valid_times: with self.subTest(time=valid_time): self.mock_eventarc_v1.reset_mock() - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -219,9 +221,9 @@ class TestMessageTool(unittest.TestCase): ce_string=valid_time ) - def test_publish_message_exception_eviction(self): + async def test_publish_message_exception_eviction(self): self.mock_publisher_client.publish.side_effect = RuntimeError("API failed") - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="projects/test/locations/global/messageBuses/my-bus", type="com.example.test", source="//test/source", @@ -237,7 +239,7 @@ class TestMessageTool(unittest.TestCase): ) @mock.patch.object(message_tool, "opentelemetry", autospec=True) - def test_publish_message_tracing(self, mock_opentelemetry): + async def test_publish_message_tracing(self, mock_opentelemetry): def inject_mock(carrier): carrier["traceparent"] = "00-testtrace-testid-01" carrier["tracestate"] = "teststate=1" @@ -246,7 +248,7 @@ class TestMessageTool(unittest.TestCase): inject_mock ) - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="projects/test/locations/global/messageBuses/my-bus", type="com.example.test", source="//test/source", @@ -265,9 +267,9 @@ class TestMessageTool(unittest.TestCase): ce_string="00-testtrace-testid-01" ) - def test_publish_message_empty_string_data(self): + async def test_publish_message_empty_string_data(self): # Act - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="projects/test/locations/global/messageBuses/my-bus", type="com.example.test", source="//test/source", @@ -278,9 +280,9 @@ class TestMessageTool(unittest.TestCase): # Assert self.assertEqual(res["status"], "SUCCESS") - def test_publish_message_empty_dict_data(self): + async def test_publish_message_empty_dict_data(self): # Act - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="projects/test/locations/global/messageBuses/my-bus", type="com.example.test", source="//test/source", @@ -291,9 +293,9 @@ class TestMessageTool(unittest.TestCase): # Assert self.assertEqual(res["status"], "SUCCESS") - def test_publish_message_missing_library(self): + async def test_publish_message_missing_library(self): with mock.patch.object(message_tool, "eventarc_publishing_v1", None): - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -303,8 +305,8 @@ class TestMessageTool(unittest.TestCase): self.assertEqual(res["status"], "ERROR") self.assertIn("not installed", res["error_details"]) - def test_publish_message_time_empty_string(self): - res = message_tool.publish_message( + async def test_publish_message_time_empty_string(self): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -316,8 +318,8 @@ class TestMessageTool(unittest.TestCase): event_kwargs = self.mock_eventarc_v1.types.CloudEvent.call_args.kwargs self.assertNotIn("time", event_kwargs.get("attributes", {})) - def test_publish_message_explicit_datacontenttype(self): - res = message_tool.publish_message( + async def test_publish_message_explicit_datacontenttype(self): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -335,10 +337,10 @@ class TestMessageTool(unittest.TestCase): ce_string="application/xml" ) - def test_publish_message_image_payload(self): + async def test_publish_message_image_payload(self): # Simulate an agent sending an image # "iVBORw0KGgo=" is a valid base64 snippet (e.g. PNG header) - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -357,8 +359,10 @@ class TestMessageTool(unittest.TestCase): ce_string="image/png" ) - def test_publish_message_explicit_datacontenttype_json_with_binary_data(self): - res = message_tool.publish_message( + async def test_publish_message_explicit_datacontenttype_json_with_binary_data( + self, + ): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -377,8 +381,10 @@ class TestMessageTool(unittest.TestCase): ce_string="application/json" ) - def test_publish_message_explicit_datacontenttype_xml_with_dict_data(self): - res = message_tool.publish_message( + async def test_publish_message_explicit_datacontenttype_xml_with_dict_data( + self, + ): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -396,8 +402,8 @@ class TestMessageTool(unittest.TestCase): ce_string="application/xml" ) - def test_publish_message_empty_datacontenttype(self): - res = message_tool.publish_message( + async def test_publish_message_empty_datacontenttype(self): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -410,8 +416,8 @@ class TestMessageTool(unittest.TestCase): event_kwargs = self.mock_eventarc_v1.types.CloudEvent.call_args.kwargs self.assertNotIn("datacontenttype", event_kwargs.get("attributes", {})) - def test_publish_message_with_subject(self): - res = message_tool.publish_message( + async def test_publish_message_with_subject(self): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -422,10 +428,14 @@ class TestMessageTool(unittest.TestCase): ) self.assertEqual(res["status"], "SUCCESS") event_kwargs = self.mock_eventarc_v1.types.CloudEvent.call_args.kwargs - self.assertEqual(event_kwargs.get("subject"), "test-subject") + self.assertNotIn("subject", event_kwargs) + self.assertIn("subject", event_kwargs.get("attributes", {})) + self.mock_eventarc_v1.types.CloudEvent.CloudEventAttributeValue.assert_any_call( + ce_string="test-subject" + ) - def test_publish_message_data_integer(self): - res = message_tool.publish_message( + async def test_publish_message_data_integer(self): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -440,8 +450,8 @@ class TestMessageTool(unittest.TestCase): ce_string="application/json" ) - def test_publish_message_data_boolean(self): - res = message_tool.publish_message( + async def test_publish_message_data_boolean(self): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -456,8 +466,8 @@ class TestMessageTool(unittest.TestCase): ce_string="application/json" ) - def test_publish_message_data_list_of_dicts(self): - res = message_tool.publish_message( + async def test_publish_message_data_list_of_dicts(self): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -472,8 +482,8 @@ class TestMessageTool(unittest.TestCase): ce_string="application/json" ) - def test_publish_message_data_unicode(self): - res = message_tool.publish_message( + async def test_publish_message_data_unicode(self): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -488,8 +498,8 @@ class TestMessageTool(unittest.TestCase): ce_string="text/plain" ) - def test_publish_message_custom_attributes_type_casting(self): - res = message_tool.publish_message( + async def test_publish_message_custom_attributes_type_casting(self): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -505,8 +515,8 @@ class TestMessageTool(unittest.TestCase): ce_string="42" ) - def test_publish_message_explicit_specversion(self): - res = message_tool.publish_message( + async def test_publish_message_explicit_specversion(self): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -518,8 +528,8 @@ class TestMessageTool(unittest.TestCase): event_kwargs = self.mock_eventarc_v1.types.CloudEvent.call_args.kwargs self.assertEqual(event_kwargs.get("spec_version"), "1.1") - def test_publish_message_explicit_id(self): - res = message_tool.publish_message( + async def test_publish_message_explicit_id(self): + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -532,9 +542,9 @@ class TestMessageTool(unittest.TestCase): self.assertEqual(event_kwargs.get("id"), "custom-event-id-99") self.assertEqual(res["message_id"], "custom-event-id-99") - def test_publish_message_base64_without_datacontenttype(self): + async def test_publish_message_base64_without_datacontenttype(self): # Simulate an agent sending base64 but forgetting the datacontenttype - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -550,7 +560,7 @@ class TestMessageTool(unittest.TestCase): ce_string="application/octet-stream" ) - def test_publish_message_data_deeply_nested_dict(self): + async def test_publish_message_data_deeply_nested_dict(self): nested_data = { "user": { "id": 101, @@ -574,7 +584,7 @@ class TestMessageTool(unittest.TestCase): "version": [1, 2, {"build": "rc1"}], }, } - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -589,7 +599,7 @@ class TestMessageTool(unittest.TestCase): ce_string="application/json" ) - def test_publish_message_data_deeply_nested_list(self): + async def test_publish_message_data_deeply_nested_list(self): nested_list = [ [1, 2, [3, 4, [5, {"six": 6}]]], {"seven": [8, 9]}, @@ -598,7 +608,7 @@ class TestMessageTool(unittest.TestCase): None, [{"eleven": {"twelve": [13, 14]}}], ] - res = message_tool.publish_message( + res = await message_tool.publish_message( bus="bus", type="type", source="source", @@ -613,8 +623,8 @@ class TestMessageTool(unittest.TestCase): ce_string="application/json" ) - def test_publish_message_auto_generated_attributes(self): - res = message_tool.publish_message( + async def test_publish_message_auto_generated_attributes(self): + res = await message_tool.publish_message( bus="projects/test/locations/global/messageBuses/my-bus", type="com.example.test", source="//test/source",