@@ -216,7 +216,9 @@ def __init__(
216216 self ._response_index = 0
217217 self ._candidate_index = 0
218218
219- def start_span_as_current_span (self , model_name , function_name , end_on_exit = True ):
219+ def start_span_as_current_span (
220+ self , model_name , function_name , end_on_exit = True
221+ ):
220222 return self ._otel_wrapper .start_as_current_span (
221223 f"{ _GENERATE_CONTENT_OP_NAME } { model_name } " ,
222224 start_time = self ._start_time ,
@@ -226,7 +228,7 @@ def start_span_as_current_span(self, model_name, function_name, end_on_exit=True
226228 gen_ai_attributes .GEN_AI_REQUEST_MODEL : self ._genai_request_model ,
227229 gen_ai_attributes .GEN_AI_OPERATION_NAME : _GENERATE_CONTENT_OP_NAME ,
228230 },
229- end_on_exit = end_on_exit
231+ end_on_exit = end_on_exit ,
230232 )
231233
232234 def process_request (
@@ -637,21 +639,22 @@ async def instrumented_generate_content_stream(
637639 with helper .start_span_as_current_span (
638640 model ,
639641 "google.genai.AsyncModels.generate_content_stream" ,
640- end_on_exit = False ) as span :
641- helper .process_request (contents , config )
642+ end_on_exit = False ,
643+ ) as span :
644+ helper .process_request (contents , config )
642645 try :
643646 response_async_generator = await wrapped_func (
644- self ,
645- model = model ,
646- contents = contents ,
647- config = config ,
648- ** kwargs ,
649- )
647+ self ,
648+ model = model ,
649+ contents = contents ,
650+ config = config ,
651+ ** kwargs ,
652+ )
650653 except Exception as error :
651- helper .process_error (error )
652- helper .finalize_processing ()
653- with trace .use_span (span , end_on_exit = True ):
654- raise
654+ helper .process_error (error )
655+ helper .finalize_processing ()
656+ with trace .use_span (span , end_on_exit = True ):
657+ raise
655658
656659 async def _response_async_generator_wrapper ():
657660 with trace .use_span (span , end_on_exit = True ):
@@ -666,7 +669,6 @@ async def _response_async_generator_wrapper():
666669 helper .finalize_processing ()
667670
668671 class _GeneratorProvider :
669-
670672 def __aiter__ (self ):
671673 return _response_async_generator_wrapper ()
672674
0 commit comments