Always call completion hook

Author: aabmass

instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/generate_content.py

@@ -164,7 +164,9 @@ def _to_dict(value: object):
 
 
 def _add_request_options_to_span(
-    span: Span, config: Optional[GenerateContentConfigOrDict], allow_list: AllowList
+    span: Span,
+    config: Optional[GenerateContentConfigOrDict],
+    allow_list: AllowList,
 ):
     if config is None:
         return
@@ -208,7 +210,9 @@ def _add_request_options_to_span(
         },
     )
     for key, value in attributes.items():
-        if key.startswith(GCP_GENAI_OPERATION_CONFIG) and not allow_list.allowed(key):
+        if key.startswith(
+            GCP_GENAI_OPERATION_CONFIG
+        ) and not allow_list.allowed(key):
             # The allowlist is used to control inclusion of the dynamic keys.
             continue
         span.set_attribute(key, value)
@@ -244,7 +248,9 @@ def _wrapped_config_with_tools(
     if not config.tools:
         return config
     result = copy.copy(config)
-    result.tools = [wrapped_tool(tool, otel_wrapper, **kwargs) for tool in config.tools]
+    result.tools = [
+        wrapped_tool(tool, otel_wrapper, **kwargs) for tool in config.tools
+    ]
     return result
 
 
@@ -267,10 +273,12 @@ def _create_completion_details_attributes(
 ) -> dict[str, Any]:
     attributes: dict[str, Any] = {
         gen_ai_attributes.GEN_AI_INPUT_MESSAGES: [
-            dataclasses.asdict(input_message) for input_message in input_messages
+            dataclasses.asdict(input_message)
+            for input_message in input_messages
         ],
         gen_ai_attributes.GEN_AI_OUTPUT_MESSAGES: [
-            dataclasses.asdict(output_message) for output_message in output_messages
+            dataclasses.asdict(output_message)
+            for output_message in output_messages
         ],
     }
     if system_instructions:
@@ -463,43 +471,54 @@ def _maybe_log_completion_details(
             system_instructions = to_system_instructions(
                 content=transformers.t_contents(system_content)[0]
             )
-        input_messages = to_input_messages(contents=transformers.t_contents(request))
-        output_messages = to_output_messages(candidates=response.candidates or [])
+        input_messages = to_input_messages(
+            contents=transformers.t_contents(request)
+        )
+        output_messages = to_output_messages(
+            candidates=response.candidates or []
+        )
 
         span = None
         if self._content_recording_enabled in [
             ContentCapturingMode.SPAN_ONLY,
             ContentCapturingMode.SPAN_AND_EVENT,
         ]:
-            completion_details_attributes = _create_completion_details_attributes(
-                input_messages,
-                output_messages,
-                system_instructions,
-                as_str=True,
+            completion_details_attributes = (
+                _create_completion_details_attributes(
+                    input_messages,
+                    output_messages,
+                    system_instructions,
+                    as_str=True,
+                )
             )
             span = trace.get_current_span()
             span.set_attributes(completion_details_attributes)
+
+        event = Event(
+            name="gen_ai.client.inference.operation.details",
+            attributes=attributes,
+        )
         if self._content_recording_enabled in [
             ContentCapturingMode.EVENT_ONLY,
             ContentCapturingMode.SPAN_AND_EVENT,
         ]:
-            completion_details_attributes = _create_completion_details_attributes(
-                input_messages,
-                output_messages,
-                system_instructions,
-            )
-            attributes.update(completion_details_attributes)
-            event = Event(
-                name="gen_ai.client.inference.operation.details", attributes=attributes
-            )
-            self.completion_hook.on_completion(
-                inputs=input_messages,
-                outputs=output_messages,
-                system_instruction=system_instructions,
-                span=span,
-                log_record=event,
+            completion_details_attributes = (
+                _create_completion_details_attributes(
+                    input_messages,
+                    output_messages,
+                    system_instructions,
+                )
             )
-            self._otel_wrapper.log_completion_details(event=event)
+            event.attributes.update(completion_details_attributes)
+
+        self.completion_hook.on_completion(
+            inputs=input_messages,
+            outputs=output_messages,
+            system_instruction=system_instructions,
+            span=span,
+            log_record=event,
+        )
+        self._otel_wrapper.log_completion_details(event=event)
 
     def _maybe_log_system_instruction(
         self, config: Optional[GenerateContentConfigOrDict] = None
@@ -538,7 +557,9 @@ def _maybe_log_user_prompt(
             total = len(contents)
             index = 0
             for entry in contents:
-                self._maybe_log_single_user_prompt(entry, index=index, total=total)
+                self._maybe_log_single_user_prompt(
+                    entry, index=index, total=total
+                )
                 index += 1
         else:
             self._maybe_log_single_user_prompt(contents)
@@ -644,7 +665,9 @@ def _maybe_log_response_stats(self, response: GenerateContentResponse):
         #
         pass
 
-    def _maybe_log_response_safety_ratings(self, response: GenerateContentResponse):
+    def _maybe_log_response_safety_ratings(
+        self, response: GenerateContentResponse
+    ):
         # TODO: Determine if there is a way that we can log
         # the "prompt_feedback". This would be especially useful
         # in the case where the response is blocked.
@@ -914,13 +937,18 @@ async def _response_async_generator_wrapper():
             with trace.use_span(span, end_on_exit=True):
                 try:
                     async for response in response_async_generator:
-                        if helper.sem_conv_opt_in_mode == _StabilityMode.DEFAULT:
+                        if (
+                            helper.sem_conv_opt_in_mode
+                            == _StabilityMode.DEFAULT
+                        ):
                             helper.process_response(response)
                         elif (
                             helper.sem_conv_opt_in_mode
                             == _StabilityMode.GEN_AI_LATEST_EXPERIMENTAL
                         ):
-                            helper.process_completion(contents, response, config)
+                            helper.process_completion(
+                                contents, response, config
+                            )
                         else:
                             raise ValueError(
                                 f"Sem Conv opt in mode {helper.sem_conv_opt_in_mode} not supported."
@@ -966,12 +994,10 @@ def instrument_generate_content(
         completion_hook,
         generate_content_config_key_allowlist=generate_content_config_key_allowlist,
     )
-    AsyncModels.generate_content_stream = (
-        _create_instrumented_async_generate_content_stream(
-            snapshot,
-            otel_wrapper,
-            completion_hook,
-            generate_content_config_key_allowlist=generate_content_config_key_allowlist,
-        )
+    AsyncModels.generate_content_stream = _create_instrumented_async_generate_content_stream(
+        snapshot,
+        otel_wrapper,
+        completion_hook,
+        generate_content_config_key_allowlist=generate_content_config_key_allowlist,
     )
     return snapshot