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:
committed by
Copybara-Service
parent
19df9b9b9b
commit
13bef9cf93
@@ -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."""
|
||||
|
||||
|
||||
@@ -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',
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user