Skip to content

Commit a3bd111

Browse files
milenvkcopybara-github
authored andcommitted
feat(eventarc): Support Context callables and correct OMIT behavior
Merge #6600 **1\. Link to an existing issue (if applicable):** - Related: google/adk-docs\#2045 (Addresses technical verification feedback in [https://github.com/google/adk-docs/pull/2045\#issuecomment-5184438669](https://github.com/google/adk-docs/pull/2045#issuecomment-5184438669)) **2\. Or, if no issue exists, describe the change:** **Problem:** 1. Callable attribute bindings in `CloudEventAttributesBinding` were always evaluated against the event `payload`. This prevented developers from correlating CloudEvents with ADK runtime telemetry (such as session IDs or invocation IDs from `Context`). 2. Setting `time=OMIT` or `datacontenttype=OMIT` in `CloudEventAttributesBinding` skipped adding keyword arguments when calling `publish_message`. Because `publish_message` auto-generates default UTC timestamps and content types when arguments are `None` or omitted, `time=OMIT` generated a timestamp instead of omitting the header. 3. Sample READMEs omitted the required `pip install "google-adk[gcp]"` prerequisite step needed for Eventarc publishing. **Solution:** 1. Added automatic signature inspection (`0-arg`, `1-arg`, and `2-arg` callables) to `CloudEventAttributesBinding` so callables can receive the event `payload`, the runtime `Context` (`tool_context`), or both, while preserving full backward compatibility with existing payload callbacks. 2. Setting `time=OMIT` or `datacontenttype=OMIT` now explicitly passes empty string (`""`) to `publish_message` so attributes are omitted from published CloudEvents. Explicitly setting required CloudEvent specification headers (`id=OMIT`, `specversion=OMIT`) now raises a `TypeError` at tool build time. 3. Updated sample agents and sample READMEs (`domain_specific_agent` and `generic_agent`) to demonstrate `Context` callables, `time=OMIT`, and GCP extra prerequisites. ### Testing Plan **Unit Tests:** - I have added or updated unit tests for my change. - All unit tests pass locally. Summary of passed `pytest` results: ``` uv run --all-extras pytest tests/unittests/integrations/eventarc -v ======================== 63 passed, 4 warnings, 13 subtests passed in 3.14s ======================== ``` - Added `test_runtime_execution_with_context_and_payload_lambdas` to verify 1-parameter (`payload` or `Context`) and 2-parameter callables. - Added `test_time_and_datacontenttype_omit_pass_empty_string` to verify omission of `time` and `datacontenttype`. - Added `test_id_and_specversion_omit_raise_typeerror` to verify static validation against omitting mandatory CloudEvent specification headers. **Manual End-to-End (E2E) Tests:** - Verified that sample agent tools in `contributing/samples/integrations/eventarc/domain_specific_agent/agent.py` build and run correctly. - Confirmed that `complete_outreach_lambda_tool` correctly injects `Context.session_id` into the event source and that `ping_system_tool` emits events without a timestamp header when configured with `time=OMIT`. ### Checklist - [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document. - [x] I have performed a self-review of my own code. - [x] I have commented my code, particularly in hard-to-understand areas. - [x] I have added tests that prove my fix is effective or that my feature works. - [x] New and existing unit tests pass locally with my changes. - [x] I have manually tested my changes end-to-end. - [x] Any dependent changes have been merged and published in downstream modules. Additional context Addresses the technical verification report on google/adk-docs#2045 comment. COPYBARA_INTEGRATE_REVIEW=#6600 from milenvk:main 647744f PiperOrigin-RevId: 964084439
1 parent bddbb3d commit a3bd111

5 files changed

Lines changed: 219 additions & 18 deletions

File tree

contributing/samples/integrations/eventarc/domain_specific_agent/README.md

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,12 @@ gcloud eventarc message-buses create my-bus \
4646

4747
*(Make sure to update the `BUS_NAME` variable in `agent.py` to match your actual bus URI).*
4848

49+
3. Install the GCP extra dependency (required for Eventarc publishing):
50+
51+
```bash
52+
pip install "google-adk[gcp]"
53+
```
54+
4955
`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`.
5056

5157
### Example A: Fully Statically Bound (Safest)
@@ -89,21 +95,24 @@ complete_outreach_dynamic_tool = toolset.create_publish_tool(
8995

9096
### Example C: Lambda Execution & Mixed Custom Attributes
9197

92-
The developer uses Python callables to generate IDs dynamically at runtime.
98+
The developer uses Python callables to generate attributes dynamically at runtime. Callables can inspect the event payload, the agent's runtime `Context` (`tool_context`), or both.
9399

94100
```python
95101
def get_custom_trace_id(payload: OutreachContext) -> str:
96102
return f"trace-{payload.customer_id}-{uuid.uuid4().hex[:8]}"
97103

104+
def get_source_from_session(ctx: Context) -> str:
105+
return f"//my-agent/outreach/{ctx.session_id}"
106+
98107
complete_outreach_lambda_tool = toolset.create_publish_tool(
99108
name="complete_outreach_lambda",
100109
description="Logs a completed outreach attempt.",
101110
payload_schema=OutreachContext,
102111
bus=f"projects/{PROJECT_ID}/locations/us-central1/messageBuses/{BUS_NAME}",
103112
ce_attributes_binding=CloudEventAttributesBinding(
104113
type="vendor_outreach.completed",
105-
source="//my-agent/outreach",
106-
id=get_custom_trace_id, # Executed at runtime
114+
source=get_source_from_session, # Evaluated against runtime Context
115+
id=get_custom_trace_id, # Evaluated against event payload
107116
custom_attributes={
108117
"environment": "production", # Statically bound
109118
"priority": AgentProvided("The priority of the outreach: 'high' or 'low'")
@@ -114,9 +123,9 @@ complete_outreach_lambda_tool = toolset.create_publish_tool(
114123

115124
**What the Agent Sees:** `complete_outreach_lambda(event_data: OutreachContext, priority: str)`
116125

117-
### Example D: Empty Payloads & Dynamic Defaults
126+
### Example D: Empty Payloads, Omit Headers & Dynamic Defaults
118127

119-
The developer wants to emit a simple signal (no business payload). If the agent omits the priority, it is dynamically calculated.
128+
The developer wants to emit a simple signal (no business payload) without a timestamp header (`time=OMIT`). If the agent omits the priority, it is dynamically calculated.
120129

121130
```python
122131
def default_priority(_: None) -> str:
@@ -130,6 +139,7 @@ ping_system_tool = toolset.create_publish_tool(
130139
ce_attributes_binding=CloudEventAttributesBinding(
131140
type="system.ping",
132141
source="//my-agent/ping",
142+
time=OMIT, # Omits time attribute from event
133143
custom_attributes={
134144
"retry": AgentProvided("Whether to retry on failure", default="false"),
135145
"priority": AgentProvided("The priority of the ping", default=default_priority)

contributing/samples/integrations/eventarc/domain_specific_agent/agent.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import uuid
1818

1919
from google.adk.agents import llm_agent
20+
from google.adk.agents.context import Context
2021
from google.adk.auth import auth_credential
2122
from google.adk.integrations.eventarc import AgentProvided
2223
from google.adk.integrations.eventarc import CloudEventAttributesBinding
@@ -104,11 +105,16 @@ class OutreachContext(pydantic.BaseModel):
104105

105106

106107
# Example C: Lambda Execution & Mixed Custom Attributes
107-
# The developer uses Python callables to generate IDs dynamically at runtime.
108+
# The developer uses Python callables to generate attributes dynamically at runtime.
109+
# Callables can inspect the event payload, the runtime Context, or both.
108110
def get_custom_trace_id(payload: OutreachContext) -> str:
109111
return f"trace-{payload.customer_id}-{uuid.uuid4().hex[:8]}"
110112

111113

114+
def get_source_from_session(ctx: Context) -> str:
115+
return f"//my-agent/outreach/{ctx.session_id}"
116+
117+
112118
complete_outreach_lambda_tool = toolset.create_publish_tool(
113119
name="complete_outreach_lambda",
114120
description=(
@@ -119,7 +125,7 @@ def get_custom_trace_id(payload: OutreachContext) -> str:
119125
bus=f"projects/{PROJECT_ID}/locations/us-central1/messageBuses/{BUS_NAME}",
120126
ce_attributes_binding=CloudEventAttributesBinding(
121127
type="vendor_outreach.completed",
122-
source="//my-agent/outreach",
128+
source=get_source_from_session,
123129
id=get_custom_trace_id,
124130
custom_attributes={
125131
"environment": "production",
@@ -145,6 +151,7 @@ def default_priority(_: None) -> str:
145151
ce_attributes_binding=CloudEventAttributesBinding(
146152
type="system.ping",
147153
source="//my-agent/ping",
154+
time=OMIT, # Omits time attribute from event
148155
custom_attributes={
149156
"retry": AgentProvided(
150157
"Whether to retry on failure", default="false"

contributing/samples/integrations/eventarc/generic_agent/README.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,12 @@ gcloud eventarc message-buses create my-bus \
3939

4040
*(Make sure to update the `BUS_NAME` variable in `agent.py` to match your actual bus URI).*
4141

42+
3. Install the GCP extra dependency (required for Eventarc publishing):
43+
44+
```bash
45+
pip install "google-adk[gcp]"
46+
```
47+
4248
Set up environment variables in your `.env` file for using Google AI Studio or Google Cloud Vertex AI for the LLM service. For example:
4349

4450
- `GOOGLE_GENAI_USE_ENTERPRISE=FALSE`

src/google/adk/integrations/eventarc/_domain_specific_publish.py

Lines changed: 66 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424

2525
from google.adk.agents.context import Context
2626
from google.adk.tools.google_tool import GoogleTool
27+
from google.adk.utils.context_utils import find_context_parameter
2728
import google.auth.credentials
2829
import pydantic
2930

@@ -45,6 +46,10 @@ class OmitSentinel:
4546
OMIT = OmitSentinel()
4647

4748

49+
def _is_context_param(func: Any, param_name: str) -> bool:
50+
return find_context_parameter(func) == param_name
51+
52+
4853
@dataclass
4954
class AgentProvided:
5055
"""Indicates that a CloudEvent attribute should be provided by the LLM."""
@@ -53,26 +58,37 @@ class AgentProvided:
5358
default: Any | OmitSentinel | MissingSentinel = MISSING
5459

5560

56-
AttributeBinding = str | Callable[[Any], str] | AgentProvided
61+
AttributeBinding = str | Callable[..., str] | AgentProvided
5762
OptionalAttributeBinding = (
5863
str
59-
| Callable[[Any], str | OmitSentinel]
64+
| Callable[..., str | OmitSentinel]
6065
| AgentProvided
6166
| OmitSentinel
6267
| MissingSentinel
6368
| None
6469
)
6570
CustomAttributeBinding = (
66-
str | Callable[[Any], str | OmitSentinel] | AgentProvided | OmitSentinel
71+
str | Callable[..., str | OmitSentinel] | AgentProvided | OmitSentinel
6772
)
6873
SpecVersionBinding = (
69-
str | Callable[[Any], str] | AgentProvided | MissingSentinel | None
74+
str | Callable[..., str] | AgentProvided | MissingSentinel | None
7075
)
7176

7277

7378
@dataclass
7479
class CloudEventAttributesBinding:
75-
"""Configuration for binding CloudEvent attributes to static values, lambdas, or AgentProvided fields."""
80+
"""Configuration for binding CloudEvent attributes to static values, lambdas, or AgentProvided fields.
81+
82+
Lambda/callable bindings can accept:
83+
- 1 parameter for the event payload (`lambda p: ...`)
84+
- 1 parameter for the runtime context (`lambda ctx: ...` or type-annotated with `Context`)
85+
- 2 parameters for both (`lambda p, ctx: ...`)
86+
- 0 parameters (`lambda: ...`)
87+
88+
Setting optional attributes (`time`, `datacontenttype`, `subject`,
89+
`custom_attributes`) to `OMIT` omits them from the published CloudEvent.
90+
Required attributes (`type`, `source`, `id`, `specversion`) cannot be `OMIT`.
91+
"""
7692

7793
type: AttributeBinding
7894
source: AttributeBinding
@@ -92,17 +108,21 @@ def build_domain_specific_tool(
92108
ce_attributes_binding: CloudEventAttributesBinding,
93109
payload_schema: type[pydantic.BaseModel] | None = None,
94110
) -> GoogleTool:
95-
"""Dynamically builds a GoogleTool wrapping publish_message with specific bindings."""
111+
"""Dynamically builds a GoogleTool wrapping publish_message with specific bindings.
112+
113+
Callable bindings in `ce_attributes_binding` can inspect the event payload, the
114+
runtime `Context` (`tool_context`), or both.
115+
"""
96116

97117
# 1. Validation
98118
mandatory_fields = ["type", "source"]
99119
for field in mandatory_fields:
100120
val = getattr(ce_attributes_binding, field)
101-
if val is MISSING:
121+
if val is MISSING: # type: ignore[comparison-overlap]
102122
raise TypeError(
103123
f"CloudEventAttributesBinding requires '{field}' to be provided."
104124
)
105-
if val is OMIT:
125+
if val is OMIT: # type: ignore[comparison-overlap]
106126
raise TypeError(
107127
f"CloudEvent field '{field}' is mandatory and cannot be OMIT."
108128
)
@@ -118,6 +138,13 @@ def build_domain_specific_tool(
118138
if bus is None:
119139
raise TypeError("The 'bus' parameter is mandatory and cannot be None.")
120140

141+
for field in ("id", "specversion"):
142+
val = getattr(ce_attributes_binding, field)
143+
if val is OMIT: # type: ignore[comparison-overlap]
144+
raise TypeError(
145+
f"CloudEvent field '{field}' is mandatory and cannot be OMIT."
146+
)
147+
121148
reserved_attributes = {
122149
"type",
123150
"source",
@@ -300,7 +327,29 @@ def resolve_attr(key: str, binding: Any, is_mandatory: bool) -> Any:
300327

301328
# Evaluate lambdas
302329
if callable(val):
303-
val = val(payload)
330+
tool_context = kwargs.get("tool_context")
331+
try:
332+
sig = inspect.signature(val)
333+
except (ValueError, TypeError):
334+
sig = None
335+
336+
if sig is not None:
337+
params = list(sig.parameters.values())
338+
if len(params) == 2:
339+
first_param = params[0]
340+
if _is_context_param(val, first_param.name):
341+
val = val(tool_context, payload)
342+
else:
343+
val = val(payload, tool_context)
344+
elif len(params) == 1:
345+
if _is_context_param(val, params[0].name):
346+
val = val(tool_context)
347+
else:
348+
val = val(payload)
349+
else:
350+
val = val()
351+
else:
352+
val = val(payload)
304353

305354
if val is OMIT:
306355
if is_mandatory:
@@ -325,7 +374,14 @@ def resolve_attr(key: str, binding: Any, is_mandatory: bool) -> Any:
325374
val = resolve_attr(
326375
field, getattr(ce_attributes_binding, field), is_mandatory
327376
)
328-
if val is not OMIT and val is not None:
377+
if val is OMIT:
378+
if field in ("time", "datacontenttype"):
379+
publish_kwargs[field] = ""
380+
elif field in ("id", "specversion"):
381+
raise ValueError(
382+
f"CloudEvent attribute '{field}' is mandatory and cannot be OMIT."
383+
)
384+
elif val is not None:
329385
publish_kwargs[field] = val
330386

331387
# Resolve custom attributes

0 commit comments

Comments
 (0)