Skip to content

Commit 9ba8eec

Browse files
Jacksunweicopybara-github
authored andcommitted
feat(tracing): Adds more OpenTelemetry convention attributes: gen_ai.request.max_tokens, gen_ai.request.top_p and gen_ai.response.finish_reasons
Fixes #1234 PiperOrigin-RevId: 795082815
1 parent 9cfe433 commit 9ba8eec

File tree

3 files changed

+74
-34
lines changed

3 files changed

+74
-34
lines changed

src/google/adk/models/llm_response.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,9 @@ class LlmResponse(BaseModel):
6767
Only used for streaming mode.
6868
"""
6969

70+
finish_reason: Optional[types.FinishReason] = None
71+
"""The finish reason of the response."""
72+
7073
error_code: Optional[str] = None
7174
"""Error code if the response is an error. Code varies by model."""
7275

@@ -97,7 +100,7 @@ class LlmResponse(BaseModel):
97100
@staticmethod
98101
def create(
99102
generate_content_response: types.GenerateContentResponse,
100-
) -> 'LlmResponse':
103+
) -> LlmResponse:
101104
"""Creates an LlmResponse from a GenerateContentResponse.
102105
103106
Args:
@@ -115,12 +118,14 @@ def create(
115118
content=candidate.content,
116119
grounding_metadata=candidate.grounding_metadata,
117120
usage_metadata=usage_metadata,
121+
finish_reason=candidate.finish_reason,
118122
)
119123
else:
120124
return LlmResponse(
121125
error_code=candidate.finish_reason,
122126
error_message=candidate.finish_message,
123127
usage_metadata=usage_metadata,
128+
finish_reason=candidate.finish_reason,
124129
)
125130
else:
126131
if generate_content_response.prompt_feedback:

src/google/adk/telemetry.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,17 @@ def trace_call_llm(
184184
_safe_json_serialize(_build_llm_request_for_trace(llm_request)),
185185
)
186186
# Consider removing once GenAI SDK provides a way to record this info.
187+
if llm_request.config:
188+
if llm_request.config.top_p:
189+
span.set_attribute(
190+
'gen_ai.request.top_p',
191+
llm_request.config.top_p,
192+
)
193+
if llm_request.config.max_output_tokens:
194+
span.set_attribute(
195+
'gen_ai.request.max_tokens',
196+
llm_request.config.max_output_tokens,
197+
)
187198

188199
try:
189200
llm_response_json = llm_response.model_dump_json(exclude_none=True)
@@ -204,6 +215,11 @@ def trace_call_llm(
204215
'gen_ai.usage.output_tokens',
205216
llm_response.usage_metadata.candidates_token_count,
206217
)
218+
if llm_response.finish_reason:
219+
span.set_attribute(
220+
'gen_ai.response.finish_reasons',
221+
[llm_response.finish_reason.value.lower()],
222+
)
207223

208224

209225
def trace_send_data(

tests/unittests/test_telemetry.py

Lines changed: 52 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -81,9 +81,57 @@ async def _create_invocation_context(
8181

8282

8383
@pytest.mark.asyncio
84-
async def test_trace_call_llm_function_response_includes_part_from_bytes(
84+
async def test_trace_call_llm(monkeypatch, mock_span_fixture):
85+
"""Test trace_call_llm sets all telemetry attributes correctly with normal content."""
86+
monkeypatch.setattr(
87+
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
88+
)
89+
90+
agent = LlmAgent(name='test_agent')
91+
invocation_context = await _create_invocation_context(agent)
92+
llm_request = LlmRequest(
93+
contents=[
94+
types.Content(
95+
role='user',
96+
parts=[types.Part(text='Hello, how are you?')],
97+
),
98+
],
99+
config=types.GenerateContentConfig(
100+
system_instruction='You are a helpful assistant.',
101+
top_p=0.95,
102+
max_output_tokens=1024,
103+
),
104+
)
105+
llm_response = LlmResponse(
106+
turn_complete=True,
107+
finish_reason=types.FinishReason.STOP,
108+
usage_metadata=types.GenerateContentResponseUsageMetadata(
109+
total_token_count=100,
110+
prompt_token_count=50,
111+
candidates_token_count=50,
112+
),
113+
)
114+
trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response)
115+
116+
expected_calls = [
117+
mock.call('gen_ai.system', 'gcp.vertex.agent'),
118+
mock.call('gen_ai.request.top_p', 0.95),
119+
mock.call('gen_ai.request.max_tokens', 1024),
120+
mock.call('gen_ai.usage.input_tokens', 50),
121+
mock.call('gen_ai.usage.output_tokens', 50),
122+
mock.call('gen_ai.response.finish_reasons', ['stop']),
123+
]
124+
assert mock_span_fixture.set_attribute.call_count == 12
125+
mock_span_fixture.set_attribute.assert_has_calls(
126+
expected_calls, any_order=True
127+
)
128+
129+
130+
@pytest.mark.asyncio
131+
async def test_trace_call_llm_with_binary_content(
85132
monkeypatch, mock_span_fixture
86133
):
134+
"""Test trace_call_llm handles binary content serialization correctly."""
87135
monkeypatch.setattr(
88136
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
89137
)
@@ -123,11 +171,14 @@ async def test_trace_call_llm_function_response_includes_part_from_bytes(
123171
llm_response = LlmResponse(turn_complete=True)
124172
trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response)
125173

174+
# Verify basic telemetry attributes are set
126175
expected_calls = [
127176
mock.call('gen_ai.system', 'gcp.vertex.agent'),
128177
]
129178
assert mock_span_fixture.set_attribute.call_count == 7
130179
mock_span_fixture.set_attribute.assert_has_calls(expected_calls)
180+
181+
# Verify binary content is replaced with '<not serializable>' in JSON
131182
llm_request_json_str = None
132183
for call_obj in mock_span_fixture.set_attribute.call_args_list:
133184
if call_obj.args[0] == 'gcp.vertex.agent.llm_request':
@@ -141,38 +192,6 @@ async def test_trace_call_llm_function_response_includes_part_from_bytes(
141192
assert llm_request_json_str.count('<not serializable>') == 2
142193

143194

144-
@pytest.mark.asyncio
145-
async def test_trace_call_llm_usage_metadata(monkeypatch, mock_span_fixture):
146-
monkeypatch.setattr(
147-
'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
148-
)
149-
150-
agent = LlmAgent(name='test_agent')
151-
invocation_context = await _create_invocation_context(agent)
152-
llm_request = LlmRequest(
153-
config=types.GenerateContentConfig(system_instruction=''),
154-
)
155-
llm_response = LlmResponse(
156-
turn_complete=True,
157-
usage_metadata=types.GenerateContentResponseUsageMetadata(
158-
total_token_count=100,
159-
prompt_token_count=50,
160-
candidates_token_count=50,
161-
),
162-
)
163-
trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response)
164-
165-
expected_calls = [
166-
mock.call('gen_ai.system', 'gcp.vertex.agent'),
167-
mock.call('gen_ai.usage.input_tokens', 50),
168-
mock.call('gen_ai.usage.output_tokens', 50),
169-
]
170-
assert mock_span_fixture.set_attribute.call_count == 9
171-
mock_span_fixture.set_attribute.assert_has_calls(
172-
expected_calls, any_order=True
173-
)
174-
175-
176195
def test_trace_tool_call_with_scalar_response(
177196
monkeypatch, mock_span_fixture, mock_tool_fixture, mock_event_fixture
178197
):

0 commit comments

Comments
 (0)