@@ -56,11 +56,8 @@ def traced_method(wrapped, instance, args, kwargs):
56
56
attributes = span_attributes ,
57
57
end_on_exit = False ,
58
58
) as span :
59
- if span .is_recording ():
60
- for message in kwargs .get ("messages" , []):
61
- event_logger .emit (
62
- message_to_event (message , capture_content )
63
- )
59
+ for message in kwargs .get ("messages" , []):
60
+ event_logger .emit (message_to_event (message , capture_content ))
64
61
65
62
start = default_timer ()
66
63
result = None
@@ -76,6 +73,9 @@ def traced_method(wrapped, instance, args, kwargs):
76
73
_set_response_attributes (
77
74
span , result , event_logger , capture_content
78
75
)
76
+ for choice in getattr (result , "choices" , []):
77
+ event_logger .emit (choice_to_event (choice , capture_content ))
78
+
79
79
span .end ()
80
80
return result
81
81
@@ -114,11 +114,8 @@ async def traced_method(wrapped, instance, args, kwargs):
114
114
attributes = span_attributes ,
115
115
end_on_exit = False ,
116
116
) as span :
117
- if span .is_recording ():
118
- for message in kwargs .get ("messages" , []):
119
- event_logger .emit (
120
- message_to_event (message , capture_content )
121
- )
117
+ for message in kwargs .get ("messages" , []):
118
+ event_logger .emit (message_to_event (message , capture_content ))
122
119
123
120
start = default_timer ()
124
121
result = None
@@ -134,6 +131,9 @@ async def traced_method(wrapped, instance, args, kwargs):
134
131
_set_response_attributes (
135
132
span , result , event_logger , capture_content
136
133
)
134
+ for choice in getattr (result , "choices" , []):
135
+ event_logger .emit (choice_to_event (choice , capture_content ))
136
+
137
137
span .end ()
138
138
return result
139
139
@@ -228,12 +228,8 @@ def _set_response_attributes(
228
228
)
229
229
230
230
if getattr (result , "choices" , None ):
231
- choices = result .choices
232
- for choice in choices :
233
- event_logger .emit (choice_to_event (choice , capture_content ))
234
-
235
231
finish_reasons = []
236
- for choice in choices :
232
+ for choice in result . choices :
237
233
finish_reasons .append (choice .finish_reason or "error" )
238
234
239
235
set_span_attribute (
@@ -333,42 +329,43 @@ def setup(self):
333
329
334
330
def cleanup (self ):
335
331
if self ._span_started :
336
- if self .response_model :
332
+ if self .span .is_recording ():
333
+ if self .response_model :
334
+ set_span_attribute (
335
+ self .span ,
336
+ GenAIAttributes .GEN_AI_RESPONSE_MODEL ,
337
+ self .response_model ,
338
+ )
339
+
340
+ if self .response_id :
341
+ set_span_attribute (
342
+ self .span ,
343
+ GenAIAttributes .GEN_AI_RESPONSE_ID ,
344
+ self .response_id ,
345
+ )
346
+
337
347
set_span_attribute (
338
348
self .span ,
339
- GenAIAttributes .GEN_AI_RESPONSE_MODEL ,
340
- self .response_model ,
349
+ GenAIAttributes .GEN_AI_USAGE_INPUT_TOKENS ,
350
+ self .prompt_tokens ,
341
351
)
342
-
343
- if self .response_id :
344
352
set_span_attribute (
345
353
self .span ,
346
- GenAIAttributes .GEN_AI_RESPONSE_ID ,
347
- self .response_id ,
354
+ GenAIAttributes .GEN_AI_USAGE_OUTPUT_TOKENS ,
355
+ self .completion_tokens ,
348
356
)
349
357
350
- set_span_attribute (
351
- self .span ,
352
- GenAIAttributes .GEN_AI_USAGE_INPUT_TOKENS ,
353
- self .prompt_tokens ,
354
- )
355
- set_span_attribute (
356
- self .span ,
357
- GenAIAttributes .GEN_AI_USAGE_OUTPUT_TOKENS ,
358
- self .completion_tokens ,
359
- )
360
-
361
- set_span_attribute (
362
- self .span ,
363
- GenAIAttributes .GEN_AI_OPENAI_RESPONSE_SERVICE_TIER ,
364
- self .service_tier ,
365
- )
358
+ set_span_attribute (
359
+ self .span ,
360
+ GenAIAttributes .GEN_AI_OPENAI_RESPONSE_SERVICE_TIER ,
361
+ self .service_tier ,
362
+ )
366
363
367
- set_span_attribute (
368
- self .span ,
369
- GenAIAttributes .GEN_AI_RESPONSE_FINISH_REASONS ,
370
- self .finish_reasons ,
371
- )
364
+ set_span_attribute (
365
+ self .span ,
366
+ GenAIAttributes .GEN_AI_RESPONSE_FINISH_REASONS ,
367
+ self .finish_reasons ,
368
+ )
372
369
373
370
for idx , choice in enumerate (self .choice_buffers ):
374
371
message = {"role" : "assistant" }
0 commit comments