diff --git a/src/google/adk/agents/run_config.py b/src/google/adk/agents/run_config.py index 6bce7a73..d35a9ce5 100644 --- a/src/google/adk/agents/run_config.py +++ b/src/google/adk/agents/run_config.py @@ -199,6 +199,9 @@ class RunConfig(BaseModel): http_options: Optional[types.HttpOptions] = None """HTTP options for the agent execution (e.g. custom headers).""" + labels: Optional[dict[str, str]] = None + """User labels for the current invocation (e.g. for billing/attribution).""" + response_modalities: Optional[list[types.Modality]] = None """The output modalities. If not set, it's default to AUDIO.""" diff --git a/src/google/adk/flows/llm_flows/basic.py b/src/google/adk/flows/llm_flows/basic.py index b0ca2400..61d22d77 100644 --- a/src/google/adk/flows/llm_flows/basic.py +++ b/src/google/adk/flows/llm_flows/basic.py @@ -81,6 +81,13 @@ def _build_basic_request( if run_config_http_options: _merge_run_config_http_options(llm_request.config, run_config_http_options) + + # Merge per-invocation user labels from RunConfig into the request config + # (e.g. for billing, telemetry, and revenue attribution across services). + if invocation_context.run_config and invocation_context.run_config.labels: + if llm_request.config.labels is None: + llm_request.config.labels = {} + llm_request.config.labels.update(invocation_context.run_config.labels) # Only set output_schema if no tools are specified. as of now, model don't # support output_schema and tools together. we have a workaround to support # both output_schema and tools at the same time. see diff --git a/tests/unittests/flows/llm_flows/test_basic_processor.py b/tests/unittests/flows/llm_flows/test_basic_processor.py index ff3eae58..35923d72 100644 --- a/tests/unittests/flows/llm_flows/test_basic_processor.py +++ b/tests/unittests/flows/llm_flows/test_basic_processor.py @@ -373,3 +373,29 @@ class TestBasicLlmRequestProcessor: assert ( llm_request.config.http_options.headers['Agent-Header'] == 'agent-val' ) + + @pytest.mark.asyncio + async def test_merges_run_config_labels(self): + """RunConfig labels are merged into llm_request.config.labels.""" + agent = LlmAgent( + name='test_agent', + model='gemini-1.5-flash', + generate_content_config=types.GenerateContentConfig( + labels={'agent_label': 'val1'} + ), + ) + + invocation_context = await _create_invocation_context(agent) + invocation_context.run_config = RunConfig( + labels={'goog-originating-logical-product-id': 'prod1'} + ) + llm_request = LlmRequest() + + processor = _BasicLlmRequestProcessor() + async for _ in processor.run_async(invocation_context, llm_request): + pass + + assert llm_request.config.labels == { + 'agent_label': 'val1', + 'goog-originating-logical-product-id': 'prod1', + }