@@ -282,9 +282,9 @@ def _set_output_data(span, response, kwargs, integration, finish_span=True):
282282
283283 def new_iterator ():
284284 # 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 () :
288288 # OpenAI chat completion API
289289 if hasattr (x , "choices" ):
290290 choice_index = 0
@@ -315,8 +315,9 @@ def new_iterator():
315315 )
316316 count_tokens_manually = False
317317
318- yield x
318+ yield x
319319
320+ with capture_internal_exceptions ():
320321 if len (data_buf ) > 0 :
321322 all_responses = ["" .join (chunk ) for chunk in data_buf ]
322323 if should_send_default_pii () and integration .include_prompts :
@@ -337,9 +338,9 @@ def new_iterator():
337338
338339 async def new_iterator_async ():
339340 # 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 () :
343344 # OpenAI chat completion API
344345 if hasattr (x , "choices" ):
345346 choice_index = 0
@@ -370,8 +371,9 @@ async def new_iterator_async():
370371 )
371372 count_tokens_manually = False
372373
373- yield x
374+ yield x
374375
376+ with capture_internal_exceptions ():
375377 if len (data_buf ) > 0 :
376378 all_responses = ["" .join (chunk ) for chunk in data_buf ]
377379 if should_send_default_pii () and integration .include_prompts :
0 commit comments