mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: add usage span attributes to telemetry (#356)
Merge https://github.com/google/adk-python/pull/1079 Fixes part of #356 Add usage attributes to span. Note: Since the handling of GenAI event bodies in OpenTelemetry has not yet been determined, I have temporarily added only attributes related to usage. COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/1079 from soundTricker:feature/356-support-more-opentelemetry-semantics 99a9d0352b4bca165baa645440e39ce7199f072b PiperOrigin-RevId: 774834279
This commit is contained in:
committed by
Copybara-Service
parent
f033e405c1
commit
ea69c9093a
@@ -195,6 +195,16 @@ def trace_call_llm(
|
||||
llm_response_json,
|
||||
)
|
||||
|
||||
if llm_response.usage_metadata is not None:
|
||||
span.set_attribute(
|
||||
'gen_ai.usage.input_tokens',
|
||||
llm_response.usage_metadata.prompt_token_count,
|
||||
)
|
||||
span.set_attribute(
|
||||
'gen_ai.usage.output_tokens',
|
||||
llm_response.usage_metadata.total_token_count,
|
||||
)
|
||||
|
||||
|
||||
def trace_send_data(
|
||||
invocation_context: InvocationContext,
|
||||
|
||||
@@ -141,6 +141,36 @@ async def test_trace_call_llm_function_response_includes_part_from_bytes(
|
||||
assert llm_request_json_str.count('<not serializable>') == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_call_llm_usage_metadata(monkeypatch, mock_span_fixture):
|
||||
monkeypatch.setattr(
|
||||
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
|
||||
)
|
||||
|
||||
agent = LlmAgent(name='test_agent')
|
||||
invocation_context = await _create_invocation_context(agent)
|
||||
llm_request = LlmRequest(
|
||||
config=types.GenerateContentConfig(system_instruction=''),
|
||||
)
|
||||
llm_response = LlmResponse(
|
||||
turn_complete=True,
|
||||
usage_metadata=types.GenerateContentResponseUsageMetadata(
|
||||
total_token_count=100, prompt_token_count=50
|
||||
),
|
||||
)
|
||||
trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response)
|
||||
|
||||
expected_calls = [
|
||||
mock.call('gen_ai.system', 'gcp.vertex.agent'),
|
||||
mock.call('gen_ai.usage.input_tokens', 50),
|
||||
mock.call('gen_ai.usage.output_tokens', 100),
|
||||
]
|
||||
assert mock_span_fixture.set_attribute.call_count == 9
|
||||
mock_span_fixture.set_attribute.assert_has_calls(
|
||||
expected_calls, any_order=True
|
||||
)
|
||||
|
||||
|
||||
def test_trace_tool_call_with_scalar_response(
|
||||
monkeypatch, mock_span_fixture, mock_tool_fixture, mock_event_fixture
|
||||
):
|
||||
|
||||
Reference in New Issue
Block a user