feat: add usage span attributes to telemetry (#356)

Merge #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=#1079 from soundTricker:feature/356-support-more-opentelemetry-semantics 99a9d03
PiperOrigin-RevId: 774834279

Author: soundTricker

src/google/adk/telemetry.py

@@ -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,

tests/unittests/test_telemetry.py

@@ -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
 ):