@@ -282,9 +282,9 @@ def _set_output_data(span, response, kwargs, integration, finish_span=True):
282
282
283
283
def new_iterator ():
284
284
# type: () -> Iterator[ChatCompletionChunk]
285
- with capture_internal_exceptions ():
286
- count_tokens_manually = True
287
- for x in old_iterator :
285
+ count_tokens_manually = True
286
+ for x in old_iterator :
287
+ with capture_internal_exceptions () :
288
288
# OpenAI chat completion API
289
289
if hasattr (x , "choices" ):
290
290
choice_index = 0
@@ -315,8 +315,9 @@ def new_iterator():
315
315
)
316
316
count_tokens_manually = False
317
317
318
- yield x
318
+ yield x
319
319
320
+ with capture_internal_exceptions ():
320
321
if len (data_buf ) > 0 :
321
322
all_responses = ["" .join (chunk ) for chunk in data_buf ]
322
323
if should_send_default_pii () and integration .include_prompts :
@@ -337,9 +338,9 @@ def new_iterator():
337
338
338
339
async def new_iterator_async ():
339
340
# type: () -> AsyncIterator[ChatCompletionChunk]
340
- with capture_internal_exceptions ():
341
- count_tokens_manually = True
342
- async for x in old_iterator :
341
+ count_tokens_manually = True
342
+ async for x in old_iterator :
343
+ with capture_internal_exceptions () :
343
344
# OpenAI chat completion API
344
345
if hasattr (x , "choices" ):
345
346
choice_index = 0
@@ -370,8 +371,9 @@ async def new_iterator_async():
370
371
)
371
372
count_tokens_manually = False
372
373
373
- yield x
374
+ yield x
374
375
376
+ with capture_internal_exceptions ():
375
377
if len (data_buf ) > 0 :
376
378
all_responses = ["" .join (chunk ) for chunk in data_buf ]
377
379
if should_send_default_pii () and integration .include_prompts :
0 commit comments