2525 AsyncIterator ,
2626 Awaitable ,
2727 Iterator ,
28- Mapping ,
2928 Optional ,
3029 Union ,
3130)
@@ -189,7 +188,7 @@ def _to_dict(value: object):
189188def _create_request_attributes (
190189 config : Optional [GenerateContentConfigOrDict ],
191190 allow_list : AllowList ,
192- ) -> dict [str , Any ]:
191+ ) -> dict [str , AttributeValue ]:
193192 if not config :
194193 return {}
195194 config = _to_dict (config )
@@ -291,8 +290,8 @@ def _create_completion_details_attributes(
291290 output_messages : list [OutputMessage ],
292291 system_instructions : list [MessagePart ],
293292 as_str : bool = False ,
294- ) -> dict [str , Any ]:
295- attributes : dict [str , Any ] = {
293+ ) -> dict [str , AttributeValue ]:
294+ attributes : dict [str , AttributeValue ] = {
296295 gen_ai_attributes .GEN_AI_INPUT_MESSAGES : [
297296 dataclasses .asdict (input_message )
298297 for input_message in input_messages
@@ -310,10 +309,11 @@ def _create_completion_details_attributes(
310309 return attributes
311310
312311
313- def _get_extra_generate_content_attributes () -> Optional [
314- Mapping [str , AttributeValue ]
315- ]:
316- return context_api .get_value (GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY )
312+ def _get_extra_generate_content_attributes () -> dict [str , AttributeValue ]:
313+ attrs = context_api .get_value (
314+ GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY
315+ )
316+ return dict (attrs or {})
317317
318318
319319class _GenerateContentInstrumentationHelper :
@@ -371,7 +371,7 @@ def start_span_as_current_span(
371371 end_on_exit = end_on_exit ,
372372 )
373373
374- def create_final_attributes (self ) -> dict [str , Any ]:
374+ def create_final_attributes (self ) -> dict [str , AttributeValue ]:
375375 final_attributes = {
376376 gen_ai_attributes .GEN_AI_USAGE_INPUT_TOKENS : self ._input_tokens ,
377377 gen_ai_attributes .GEN_AI_USAGE_OUTPUT_TOKENS : self ._output_tokens ,
@@ -463,8 +463,9 @@ def _maybe_update_error_type(self, response: GenerateContentResponse):
463463
464464 def _maybe_log_completion_details (
465465 self ,
466- request_attributes : dict [str , Any ],
467- final_attributes : dict [str , Any ],
466+ extra_attributes : dict [str , AttributeValue ],
467+ request_attributes : dict [str , AttributeValue ],
468+ final_attributes : dict [str , AttributeValue ],
468469 request : Union [ContentListUnion , ContentListUnionDict ],
469470 candidates : list [Candidate ],
470471 config : Optional [GenerateContentConfigOrDict ] = None ,
@@ -487,7 +488,9 @@ def _maybe_log_completion_details(
487488 span = trace .get_current_span ()
488489 event = LogRecord (
489490 event_name = "gen_ai.client.inference.operation.details" ,
490- attributes = request_attributes | final_attributes ,
491+ attributes = extra_attributes
492+ | request_attributes
493+ | final_attributes ,
491494 )
492495 self .completion_hook .on_completion (
493496 inputs = input_messages ,
@@ -742,9 +745,8 @@ def instrumented_generate_content(
742745 with helper .start_span_as_current_span (
743746 model , "google.genai.Models.generate_content"
744747 ) as span :
745- if extra_attributes := _get_extra_generate_content_attributes ():
746- span .set_attributes (extra_attributes )
747- span .set_attributes (request_attributes )
748+ extra_attributes = _get_extra_generate_content_attributes ()
749+ span .set_attributes (extra_attributes | request_attributes )
748750 if helper .sem_conv_opt_in_mode == _StabilityMode .DEFAULT :
749751 helper .process_request (contents , config , span )
750752 try :
@@ -773,6 +775,7 @@ def instrumented_generate_content(
773775 final_attributes = helper .create_final_attributes ()
774776 span .set_attributes (final_attributes )
775777 helper ._maybe_log_completion_details (
778+ extra_attributes ,
776779 request_attributes ,
777780 final_attributes ,
778781 contents ,
@@ -817,9 +820,8 @@ def instrumented_generate_content_stream(
817820 with helper .start_span_as_current_span (
818821 model , "google.genai.Models.generate_content_stream"
819822 ) as span :
820- if extra_attributes := _get_extra_generate_content_attributes ():
821- span .set_attributes (extra_attributes )
822- span .set_attributes (request_attributes )
823+ extra_attributes = _get_extra_generate_content_attributes ()
824+ span .set_attributes (extra_attributes | request_attributes )
823825 if helper .sem_conv_opt_in_mode == _StabilityMode .DEFAULT :
824826 helper .process_request (contents , config , span )
825827 try :
@@ -848,6 +850,7 @@ def instrumented_generate_content_stream(
848850 final_attributes = helper .create_final_attributes ()
849851 span .set_attributes (final_attributes )
850852 helper ._maybe_log_completion_details (
853+ extra_attributes ,
851854 request_attributes ,
852855 final_attributes ,
853856 contents ,
@@ -892,9 +895,8 @@ async def instrumented_generate_content(
892895 with helper .start_span_as_current_span (
893896 model , "google.genai.AsyncModels.generate_content"
894897 ) as span :
895- if extra_attributes := _get_extra_generate_content_attributes ():
896- span .set_attributes (extra_attributes )
897- span .set_attributes (request_attributes )
898+ extra_attributes = _get_extra_generate_content_attributes ()
899+ span .set_attributes (extra_attributes | request_attributes )
898900 if helper .sem_conv_opt_in_mode == _StabilityMode .DEFAULT :
899901 helper .process_request (contents , config , span )
900902 try :
@@ -922,6 +924,7 @@ async def instrumented_generate_content(
922924 final_attributes = helper .create_final_attributes ()
923925 span .set_attributes (final_attributes )
924926 helper ._maybe_log_completion_details (
927+ extra_attributes ,
925928 request_attributes ,
926929 final_attributes ,
927930 contents ,
@@ -968,9 +971,8 @@ async def instrumented_generate_content_stream(
968971 "google.genai.AsyncModels.generate_content_stream" ,
969972 end_on_exit = False ,
970973 ) as span :
971- if extra_attributes := _get_extra_generate_content_attributes ():
972- span .set_attributes (extra_attributes )
973- span .set_attributes (request_attributes )
974+ extra_attributes = _get_extra_generate_content_attributes ()
975+ span .set_attributes (extra_attributes | request_attributes )
974976 if (
975977 not helper .sem_conv_opt_in_mode
976978 == _StabilityMode .GEN_AI_LATEST_EXPERIMENTAL
@@ -990,6 +992,7 @@ async def instrumented_generate_content_stream(
990992 final_attributes = helper .create_final_attributes ()
991993 span .set_attributes (final_attributes )
992994 helper ._maybe_log_completion_details (
995+ extra_attributes ,
993996 request_attributes ,
994997 final_attributes ,
995998 contents ,
@@ -1023,6 +1026,7 @@ async def _response_async_generator_wrapper():
10231026 final_attributes = helper .create_final_attributes ()
10241027 span .set_attributes (final_attributes )
10251028 helper ._maybe_log_completion_details (
1029+ extra_attributes ,
10261030 request_attributes ,
10271031 final_attributes ,
10281032 contents ,
0 commit comments