feat: support user labels in RunConfig

Allow callers to supply custom user labels via RunConfig to propagate billing and attribution labels down to LlmRequest and Vertex API requests.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 952960713
This commit is contained in:
Shangjie Chen
2026-07-23 14:31:51 -07:00
committed by Copybara-Service
parent 19df9b9b9b
commit 13bef9cf93
3 changed files with 36 additions and 0 deletions
+3
View File
@@ -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."""
+7
View File
@@ -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
@@ -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',
}