@@ -81,9 +81,57 @@ async def _create_invocation_context(
81
81
82
82
83
83
@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 (
85
132
monkeypatch , mock_span_fixture
86
133
):
134
+ """Test trace_call_llm handles binary content serialization correctly."""
87
135
monkeypatch .setattr (
88
136
'opentelemetry.trace.get_current_span' , lambda : mock_span_fixture
89
137
)
@@ -123,11 +171,14 @@ async def test_trace_call_llm_function_response_includes_part_from_bytes(
123
171
llm_response = LlmResponse (turn_complete = True )
124
172
trace_call_llm (invocation_context , 'test_event_id' , llm_request , llm_response )
125
173
174
+ # Verify basic telemetry attributes are set
126
175
expected_calls = [
127
176
mock .call ('gen_ai.system' , 'gcp.vertex.agent' ),
128
177
]
129
178
assert mock_span_fixture .set_attribute .call_count == 7
130
179
mock_span_fixture .set_attribute .assert_has_calls (expected_calls )
180
+
181
+ # Verify binary content is replaced with '<not serializable>' in JSON
131
182
llm_request_json_str = None
132
183
for call_obj in mock_span_fixture .set_attribute .call_args_list :
133
184
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(
141
192
assert llm_request_json_str .count ('<not serializable>' ) == 2
142
193
143
194
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
-
176
195
def test_trace_tool_call_with_scalar_response (
177
196
monkeypatch , mock_span_fixture , mock_tool_fixture , mock_event_fixture
178
197
):
0 commit comments