2424 BlockedReason ,
2525 Candidate ,
2626 Content ,
27- ContentUnion ,
28- ContentUnionDict ,
2927 ContentListUnion ,
3028 ContentListUnionDict ,
29+ ContentUnion ,
30+ ContentUnionDict ,
3131 GenerateContentConfigOrDict ,
3232 GenerateContentResponse ,
3333)
@@ -196,12 +196,11 @@ def _get_top_p(config: Optional[GenerateContentConfigOrDict]):
196196def _to_dict (value : object ):
197197 if isinstance (value , dict ):
198198 return value
199- if hasattr (value , ' model_dump' ):
199+ if hasattr (value , " model_dump" ):
200200 return value .model_dump ()
201201 return json .loads (json .dumps (value ))
202202
203203
204-
205204class _GenerateContentInstrumentationHelper :
206205 def __init__ (
207206 self ,
@@ -345,19 +344,19 @@ def _maybe_log_user_prompt(
345344 self , contents : Union [ContentListUnion , ContentListUnionDict ]
346345 ):
347346 if isinstance (contents , list ):
348- total = len (contents )
349- index = 0
347+ total = len (contents )
348+ index = 0
350349 for entry in contents :
351- self ._maybe_log_single_user_prompt (entry , index = index , total = total )
350+ self ._maybe_log_single_user_prompt (
351+ entry , index = index , total = total
352+ )
352353 index += 1
353354 else :
354355 self ._maybe_log_single_user_prompt (contents )
355356
356357 def _maybe_log_single_user_prompt (
357- self ,
358- contents : Union [ContentUnion , ContentUnionDict ],
359- index = 0 ,
360- total = 1 ):
358+ self , contents : Union [ContentUnion , ContentUnionDict ], index = 0 , total = 1
359+ ):
361360 # TODO: figure out how to report the index in a manner that is
362361 # aligned with the OTel semantic conventions.
363362 attributes = {
@@ -401,7 +400,9 @@ def _maybe_log_response_stats(self, response: GenerateContentResponse):
401400 #
402401 pass
403402
404- def _maybe_log_response_safety_ratings (self , response : GenerateContentResponse ):
403+ def _maybe_log_response_safety_ratings (
404+ self , response : GenerateContentResponse
405+ ):
405406 # TODO: Determine if there is a way that we can log
406407 # the "prompt_feedback". This would be especially useful
407408 # in the case where the response is blocked.
@@ -418,7 +419,8 @@ def _maybe_log_response(self, response: GenerateContentResponse):
418419 candidate ,
419420 flat_candidate_index = self ._candidate_index ,
420421 candidate_in_response_index = candidate_in_response_index ,
421- response_index = self ._response_index )
422+ response_index = self ._response_index ,
423+ )
422424 self ._candidate_index += 1
423425 candidate_in_response_index += 1
424426
@@ -427,10 +429,11 @@ def _maybe_log_response_candidate(
427429 candidate : Candidate ,
428430 flat_candidate_index : int ,
429431 candidate_in_response_index : int ,
430- response_index : int ):
432+ response_index : int ,
433+ ):
431434 # TODO: Determine if there might be a way to report the
432435 # response index and candidate response index.
433- attributes = {
436+ attributes = {
434437 gen_ai_attributes .GEN_AI_SYSTEM : self ._genai_system ,
435438 }
436439 # TODO: determine if "role" should be reported here or not and, if so,
@@ -445,7 +448,7 @@ def _maybe_log_response_candidate(
445448 # "citation_metadata", "grounding_metadata", "logprobs_result", etc.
446449 #
447450 # See also: "TODOS.md"
448- body = {
451+ body = {
449452 "index" : flat_candidate_index ,
450453 }
451454 if self ._content_recording_enabled :
0 commit comments