diff --git a/contributing/samples/integrations/eventarc/domain_specific_agent/README.md b/contributing/samples/integrations/eventarc/domain_specific_agent/README.md index 8738bbbd1b..3284152c7e 100644 --- a/contributing/samples/integrations/eventarc/domain_specific_agent/README.md +++ b/contributing/samples/integrations/eventarc/domain_specific_agent/README.md @@ -46,6 +46,12 @@ gcloud eventarc message-buses create my-bus \ *(Make sure to update the `BUS_NAME` variable in `agent.py` to match your actual bus URI).* +3. Install the GCP extra dependency (required for Eventarc publishing): + +```bash +pip install "google-adk[gcp]" +``` + `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) @@ -89,12 +95,15 @@ complete_outreach_dynamic_tool = toolset.create_publish_tool( ### Example C: Lambda Execution & Mixed Custom Attributes -The developer uses Python callables to generate IDs dynamically at runtime. +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. ```python def get_custom_trace_id(payload: OutreachContext) -> str: return f"trace-{payload.customer_id}-{uuid.uuid4().hex[:8]}" +def get_source_from_session(ctx: Context) -> str: + return f"//my-agent/outreach/{ctx.session_id}" + complete_outreach_lambda_tool = toolset.create_publish_tool( name="complete_outreach_lambda", description="Logs a completed outreach attempt.", @@ -102,8 +111,8 @@ complete_outreach_lambda_tool = toolset.create_publish_tool( 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 + source=get_source_from_session, # Evaluated against runtime Context + id=get_custom_trace_id, # Evaluated against event payload custom_attributes={ "environment": "production", # Statically bound "priority": AgentProvided("The priority of the outreach: 'high' or 'low'") @@ -114,9 +123,9 @@ complete_outreach_lambda_tool = toolset.create_publish_tool( **What the Agent Sees:** `complete_outreach_lambda(event_data: OutreachContext, priority: str)` -### Example D: Empty Payloads & Dynamic Defaults +### Example D: Empty Payloads, Omit Headers & Dynamic Defaults -The developer wants to emit a simple signal (no business payload). If the agent omits the priority, it is dynamically calculated. +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. ```python def default_priority(_: None) -> str: @@ -130,6 +139,7 @@ ping_system_tool = toolset.create_publish_tool( ce_attributes_binding=CloudEventAttributesBinding( type="system.ping", source="//my-agent/ping", + time=OMIT, # Omits time attribute from event custom_attributes={ "retry": AgentProvided("Whether to retry on failure", default="false"), "priority": AgentProvided("The priority of the ping", default=default_priority) diff --git a/contributing/samples/integrations/eventarc/domain_specific_agent/agent.py b/contributing/samples/integrations/eventarc/domain_specific_agent/agent.py index 2582a140c5..7d7b085604 100644 --- a/contributing/samples/integrations/eventarc/domain_specific_agent/agent.py +++ b/contributing/samples/integrations/eventarc/domain_specific_agent/agent.py @@ -17,6 +17,7 @@ import uuid from google.adk.agents import llm_agent +from google.adk.agents.context import Context from google.adk.auth import auth_credential from google.adk.integrations.eventarc import AgentProvided from google.adk.integrations.eventarc import CloudEventAttributesBinding @@ -104,11 +105,16 @@ class OutreachContext(pydantic.BaseModel): # Example C: Lambda Execution & Mixed Custom Attributes -# The developer uses Python callables to generate IDs dynamically at runtime. +# The developer uses Python callables to generate attributes dynamically at runtime. +# Callables can inspect the event payload, the runtime Context, or both. def get_custom_trace_id(payload: OutreachContext) -> str: return f"trace-{payload.customer_id}-{uuid.uuid4().hex[:8]}" +def get_source_from_session(ctx: Context) -> str: + return f"//my-agent/outreach/{ctx.session_id}" + + complete_outreach_lambda_tool = toolset.create_publish_tool( name="complete_outreach_lambda", description=( @@ -119,7 +125,7 @@ def get_custom_trace_id(payload: OutreachContext) -> str: bus=f"projects/{PROJECT_ID}/locations/us-central1/messageBuses/{BUS_NAME}", ce_attributes_binding=CloudEventAttributesBinding( type="vendor_outreach.completed", - source="//my-agent/outreach", + source=get_source_from_session, id=get_custom_trace_id, custom_attributes={ "environment": "production", @@ -145,6 +151,7 @@ def default_priority(_: None) -> str: ce_attributes_binding=CloudEventAttributesBinding( type="system.ping", source="//my-agent/ping", + time=OMIT, # Omits time attribute from event custom_attributes={ "retry": AgentProvided( "Whether to retry on failure", default="false" diff --git a/contributing/samples/integrations/eventarc/generic_agent/README.md b/contributing/samples/integrations/eventarc/generic_agent/README.md index 4d66262427..8f16d62700 100644 --- a/contributing/samples/integrations/eventarc/generic_agent/README.md +++ b/contributing/samples/integrations/eventarc/generic_agent/README.md @@ -39,6 +39,12 @@ gcloud eventarc message-buses create my-bus \ *(Make sure to update the `BUS_NAME` variable in `agent.py` to match your actual bus URI).* +3. Install the GCP extra dependency (required for Eventarc publishing): + +```bash +pip install "google-adk[gcp]" +``` + 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` diff --git a/src/google/adk/integrations/eventarc/_domain_specific_publish.py b/src/google/adk/integrations/eventarc/_domain_specific_publish.py index 71ded396d0..73d6ecdd0e 100644 --- a/src/google/adk/integrations/eventarc/_domain_specific_publish.py +++ b/src/google/adk/integrations/eventarc/_domain_specific_publish.py @@ -24,6 +24,7 @@ from google.adk.agents.context import Context from google.adk.tools.google_tool import GoogleTool +from google.adk.utils.context_utils import find_context_parameter import google.auth.credentials import pydantic @@ -44,6 +45,17 @@ class OmitSentinel: OMIT = OmitSentinel() +_CONTEXT_PARAM_NAMES: frozenset[str] = frozenset( + {"ctx", "context", "tool_context"} +) + + +def _is_context_param(func: Any, param_name: str) -> bool: + return ( + param_name in _CONTEXT_PARAM_NAMES + or find_context_parameter(func) == param_name + ) + @dataclass class AgentProvided: @@ -72,7 +84,18 @@ class AgentProvided: @dataclass class CloudEventAttributesBinding: - """Configuration for binding CloudEvent attributes to static values, lambdas, or AgentProvided fields.""" + """Configuration for binding CloudEvent attributes to static values, lambdas, or AgentProvided fields. + + Lambda/callable bindings can accept: + - 1 parameter for the event payload (`lambda p: ...`) + - 1 parameter for the runtime context (`lambda ctx: ...` or type-annotated with `Context`) + - 2 parameters for both (`lambda p, ctx: ...`) + - 0 parameters (`lambda: ...`) + + Setting optional attributes (`time`, `datacontenttype`, `subject`, + `custom_attributes`) to `OMIT` omits them from the published CloudEvent. + Required attributes (`type`, `source`, `id`, `specversion`) cannot be `OMIT`. + """ type: AttributeBinding source: AttributeBinding @@ -92,7 +115,11 @@ def build_domain_specific_tool( ce_attributes_binding: CloudEventAttributesBinding, payload_schema: type[pydantic.BaseModel] | None = None, ) -> GoogleTool: - """Dynamically builds a GoogleTool wrapping publish_message with specific bindings.""" + """Dynamically builds a GoogleTool wrapping publish_message with specific bindings. + + Callable bindings in `ce_attributes_binding` can inspect the event payload, the + runtime `Context` (`tool_context`), or both. + """ # 1. Validation mandatory_fields = ["type", "source"] @@ -118,6 +145,13 @@ def build_domain_specific_tool( if bus is None: raise TypeError("The 'bus' parameter is mandatory and cannot be None.") + for field in ("id", "specversion"): + val = getattr(ce_attributes_binding, field) + if val is OMIT: + raise TypeError( + f"CloudEvent field '{field}' is mandatory and cannot be OMIT." + ) + reserved_attributes = { "type", "source", @@ -300,7 +334,25 @@ def resolve_attr(key: str, binding: Any, is_mandatory: bool) -> Any: # Evaluate lambdas if callable(val): - val = val(payload) + tool_context = kwargs.get("tool_context") + try: + sig = inspect.signature(val) + params = list(sig.parameters.values()) + if len(params) == 2: + first_param = params[0] + if _is_context_param(val, first_param.name): + val = val(tool_context, payload) + else: + val = val(payload, tool_context) + elif len(params) == 1: + if _is_context_param(val, params[0].name): + val = val(tool_context) + else: + val = val(payload) + else: + val = val() + except (ValueError, TypeError): + val = val(payload) if val is OMIT: if is_mandatory: @@ -325,7 +377,14 @@ def resolve_attr(key: str, binding: Any, is_mandatory: bool) -> Any: val = resolve_attr( field, getattr(ce_attributes_binding, field), is_mandatory ) - if val is not OMIT and val is not None: + if val is OMIT: + if field in ("time", "datacontenttype"): + publish_kwargs[field] = "" + elif field in ("id", "specversion"): + raise ValueError( + f"CloudEvent attribute '{field}' is mandatory and cannot be OMIT." + ) + elif val is not None: publish_kwargs[field] = val # Resolve custom attributes diff --git a/tests/unittests/integrations/eventarc/test_domain_specific_publish.py b/tests/unittests/integrations/eventarc/test_domain_specific_publish.py index 51dbc50b13..73d1478fb0 100644 --- a/tests/unittests/integrations/eventarc/test_domain_specific_publish.py +++ b/tests/unittests/integrations/eventarc/test_domain_specific_publish.py @@ -18,6 +18,7 @@ import inspect from unittest import mock +from google.adk.agents.context import Context from google.adk.integrations.eventarc import _config as config from google.adk.integrations.eventarc import _domain_specific_publish as domain_specific_publish from google.adk.integrations.eventarc import _eventarc_toolset as eventarc_toolset @@ -178,7 +179,61 @@ async def test_runtime_execution_with_payload(mock_publish, toolset): assert kwargs["type"] == "action.login" assert kwargs["source"] == "my-source" assert kwargs["subject"] == "user123" - assert "time" not in kwargs + assert kwargs["time"] == "" + assert kwargs["data"] == {"user_id": "user123", "action": "login"} + + +@pytest.mark.asyncio +@mock.patch.object(domain_specific_publish, "publish_message", autospec=True) +async def test_runtime_execution_with_context_and_payload_lambdas( + mock_publish, toolset +): + def get_custom_id(c: Context) -> str: + return f"id-{c.session_id}" + + tool = domain_specific_publish.build_domain_specific_tool( + toolset=toolset, + name="test_tool", + description="desc", + bus="my-bus", + ce_attributes_binding=domain_specific_publish.CloudEventAttributesBinding( + type=lambda p: f"action.{p.action}", + source=lambda ctx: f"//agent/{ctx.session_id}", + subject=lambda payload, ctx: f"{payload.user_id}-{ctx.session_id}", + id=get_custom_id, + specversion=lambda: "1.0", + custom_attributes={ + "ordertest": lambda ctx, payload: ( + f"{ctx.session_id}:{payload.action}" + ), + }, + time=domain_specific_publish.OMIT, + ), + payload_schema=DummyPayload, + ) + + payload = DummyPayload(user_id="user123", action="login") + mock_ctx = mock.Mock(spec=Context) + mock_ctx.session_id = "session456" + + await tool.func( + event_data=payload, + credentials=None, + settings=config.EventarcToolConfig(), + tool_context=mock_ctx, + ) + + mock_publish.assert_called_once() + kwargs = mock_publish.call_args.kwargs + + assert kwargs["bus"] == "my-bus" + assert kwargs["type"] == "action.login" + assert kwargs["source"] == "//agent/session456" + assert kwargs["subject"] == "user123-session456" + assert kwargs["id"] == "id-session456" + assert kwargs["specversion"] == "1.0" + assert kwargs["custom_attributes"] == {"ordertest": "session456:login"} + assert kwargs["time"] == "" assert kwargs["data"] == {"user_id": "user123", "action": "login"} @@ -416,3 +471,57 @@ def test_custom_attribute_missing_raises_typeerror(toolset): custom_attributes={"mykey": domain_specific_publish.MISSING}, ), ) + + +@pytest.mark.asyncio +@mock.patch.object(domain_specific_publish, "publish_message", autospec=True) +async def test_time_and_datacontenttype_omit_pass_empty_string( + mock_publish, toolset +): + tool = domain_specific_publish.build_domain_specific_tool( + toolset=toolset, + name="test_tool", + description="desc", + bus="my-bus", + ce_attributes_binding=domain_specific_publish.CloudEventAttributesBinding( + type="my-type", + source="my-source", + time=domain_specific_publish.OMIT, + datacontenttype=domain_specific_publish.OMIT, + ), + payload_schema=DummyPayload, + ) + + await tool.func( + event_data=DummyPayload(user_id="u1", action="a1"), + credentials=None, + settings=config.EventarcToolConfig(), + tool_context=mock.Mock(), + ) + + mock_publish.assert_called_once() + kwargs = mock_publish.call_args.kwargs + assert kwargs["time"] == "" + assert kwargs["datacontenttype"] == "" + + +@pytest.mark.parametrize("field", ["id", "specversion"]) +def test_id_and_specversion_omit_raise_typeerror(toolset, field): + binding_kwargs = { + "type": "my-type", + "source": "my-source", + field: domain_specific_publish.OMIT, + } + with pytest.raises( + TypeError, + match=f"CloudEvent field '{field}' is mandatory and cannot be OMIT.", + ): + domain_specific_publish.build_domain_specific_tool( + toolset=toolset, + name="test_tool", + description="desc", + bus="my-bus", + ce_attributes_binding=domain_specific_publish.CloudEventAttributesBinding( + **binding_kwargs + ), + )