11
11
12
12
from strands .telemetry .tracer import JSONEncoder , Tracer , get_tracer , serialize
13
13
from strands .types .content import ContentBlock
14
- from strands .types .streaming import StopReason , Usage
14
+ from strands .types .streaming import Metrics , StopReason , Usage
15
15
16
16
17
17
@pytest .fixture (autouse = True )
@@ -205,9 +205,10 @@ def test_end_model_invoke_span(mock_span):
205
205
tracer = Tracer ()
206
206
message = {"role" : "assistant" , "content" : [{"text" : "Response" }]}
207
207
usage = Usage (inputTokens = 10 , outputTokens = 20 , totalTokens = 30 )
208
+ metrics = Metrics (latencyMs = 20 , timeToFirstByteMs = 10 )
208
209
stop_reason : StopReason = "end_turn"
209
210
210
- tracer .end_model_invoke_span (mock_span , message , usage , stop_reason )
211
+ tracer .end_model_invoke_span (mock_span , message , usage , metrics , stop_reason )
211
212
212
213
mock_span .set_attribute .assert_any_call ("gen_ai.usage.prompt_tokens" , 10 )
213
214
mock_span .set_attribute .assert_any_call ("gen_ai.usage.input_tokens" , 10 )
@@ -216,6 +217,8 @@ def test_end_model_invoke_span(mock_span):
216
217
mock_span .set_attribute .assert_any_call ("gen_ai.usage.total_tokens" , 30 )
217
218
mock_span .set_attribute .assert_any_call ("gen_ai.usage.cache_read_input_tokens" , 0 )
218
219
mock_span .set_attribute .assert_any_call ("gen_ai.usage.cache_write_input_tokens" , 0 )
220
+ mock_span .set_attribute .assert_any_call ("gen_ai.server.request.duration" , 20 )
221
+ mock_span .set_attribute .assert_any_call ("gen_ai.server.time_to_first_token" , 10 )
219
222
mock_span .add_event .assert_called_with (
220
223
"gen_ai.choice" ,
221
224
attributes = {"message" : json .dumps (message ["content" ]), "finish_reason" : "end_turn" },
@@ -231,9 +234,10 @@ def test_end_model_invoke_span_latest_conventions(mock_span):
231
234
tracer .use_latest_genai_conventions = True
232
235
message = {"role" : "assistant" , "content" : [{"text" : "Response" }]}
233
236
usage = Usage (inputTokens = 10 , outputTokens = 20 , totalTokens = 30 )
237
+ metrics = Metrics (latencyMs = 20 , timeToFirstByteMs = 10 )
234
238
stop_reason : StopReason = "end_turn"
235
239
236
- tracer .end_model_invoke_span (mock_span , message , usage , stop_reason )
240
+ tracer .end_model_invoke_span (mock_span , message , usage , metrics , stop_reason )
237
241
238
242
mock_span .set_attribute .assert_any_call ("gen_ai.usage.prompt_tokens" , 10 )
239
243
mock_span .set_attribute .assert_any_call ("gen_ai.usage.input_tokens" , 10 )
@@ -242,6 +246,8 @@ def test_end_model_invoke_span_latest_conventions(mock_span):
242
246
mock_span .set_attribute .assert_any_call ("gen_ai.usage.total_tokens" , 30 )
243
247
mock_span .set_attribute .assert_any_call ("gen_ai.usage.cache_read_input_tokens" , 0 )
244
248
mock_span .set_attribute .assert_any_call ("gen_ai.usage.cache_write_input_tokens" , 0 )
249
+ mock_span .set_attribute .assert_any_call ("gen_ai.server.time_to_first_token" , 10 )
250
+ mock_span .set_attribute .assert_any_call ("gen_ai.server.request.duration" , 20 )
245
251
mock_span .add_event .assert_called_with (
246
252
"gen_ai.client.inference.operation.details" ,
247
253
attributes = {
@@ -766,8 +772,9 @@ def test_end_model_invoke_span_with_cache_metrics(mock_span):
766
772
cacheWriteInputTokens = 3 ,
767
773
)
768
774
stop_reason : StopReason = "end_turn"
775
+ metrics = Metrics (latencyMs = 10 , timeToFirstByteMs = 5 )
769
776
770
- tracer .end_model_invoke_span (mock_span , message , usage , stop_reason )
777
+ tracer .end_model_invoke_span (mock_span , message , usage , metrics , stop_reason )
771
778
772
779
mock_span .set_attribute .assert_any_call ("gen_ai.usage.prompt_tokens" , 10 )
773
780
mock_span .set_attribute .assert_any_call ("gen_ai.usage.input_tokens" , 10 )
@@ -776,6 +783,8 @@ def test_end_model_invoke_span_with_cache_metrics(mock_span):
776
783
mock_span .set_attribute .assert_any_call ("gen_ai.usage.total_tokens" , 30 )
777
784
mock_span .set_attribute .assert_any_call ("gen_ai.usage.cache_read_input_tokens" , 5 )
778
785
mock_span .set_attribute .assert_any_call ("gen_ai.usage.cache_write_input_tokens" , 3 )
786
+ mock_span .set_attribute .assert_any_call ("gen_ai.server.request.duration" , 10 )
787
+ mock_span .set_attribute .assert_any_call ("gen_ai.server.time_to_first_token" , 5 )
779
788
mock_span .set_status .assert_called_once_with (StatusCode .OK )
780
789
mock_span .end .assert_called_once ()
781
790
0 commit comments