refactor: Inject upload_hook in Instrumentor.

Author: argaj

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

@@ -54,7 +54,7 @@
     InputMessage,
     OutputMessage,
 )
-from opentelemetry.util.genai.upload_hook import load_upload_hook
+from opentelemetry.util.genai.upload_hook import UploadHook
 
 from .allowlist_util import AllowList
 from .custom_semconv import GCP_GENAI_OPERATION_CONFIG
@@ -290,12 +290,14 @@ def __init__(
         models_object: Union[Models, AsyncModels],
         otel_wrapper: OTelWrapper,
         model: str,
+        upload_hook: UploadHook,
         generate_content_config_key_allowlist: Optional[AllowList] = None,
     ):
         self._start_time = time.time_ns()
         self._otel_wrapper = otel_wrapper
         self._genai_system = _determine_genai_system(models_object)
         self._genai_request_model = model
+        self.upload_hook = upload_hook
         self._finish_reasons_set = set()
         self._error_type = None
         self._input_tokens = 0
@@ -490,15 +492,13 @@ def _maybe_log_completion_details(
             event = Event(
                 name="gen_ai.client.inference.operation.details", attributes=attributes
             )
-            hook = load_upload_hook()
-            hook.upload(
+            self.upload_hook.upload(
                 inputs=input_messages,
                 outputs=output_messages,
                 system_instruction=system_instructions,
                 span=span,
                 log_record=event,
             )
-            hook.shutdown()
             self._otel_wrapper.log_completion_details(event=event)
 
     def _maybe_log_system_instruction(
@@ -689,6 +689,7 @@ def _record_duration_metric(self):
 def _create_instrumented_generate_content(
     snapshot: _MethodsSnapshot,
     otel_wrapper: OTelWrapper,
+    upload_hook: UploadHook,
     generate_content_config_key_allowlist: Optional[AllowList] = None,
 ):
     wrapped_func = snapshot.generate_content
@@ -706,6 +707,7 @@ def instrumented_generate_content(
             self,
             otel_wrapper,
             model,
+            upload_hook,
             generate_content_config_key_allowlist=generate_content_config_key_allowlist,
         )
         with helper.start_span_as_current_span(
@@ -746,6 +748,7 @@ def instrumented_generate_content(
 def _create_instrumented_generate_content_stream(
     snapshot: _MethodsSnapshot,
     otel_wrapper: OTelWrapper,
+    upload_hook: UploadHook,
     generate_content_config_key_allowlist: Optional[AllowList] = None,
 ):
     wrapped_func = snapshot.generate_content_stream
@@ -763,6 +766,7 @@ def instrumented_generate_content_stream(
             self,
             otel_wrapper,
             model,
+            upload_hook,
             generate_content_config_key_allowlist=generate_content_config_key_allowlist,
         )
         with helper.start_span_as_current_span(
@@ -803,6 +807,7 @@ def instrumented_generate_content_stream(
 def _create_instrumented_async_generate_content(
     snapshot: _MethodsSnapshot,
     otel_wrapper: OTelWrapper,
+    upload_hook: UploadHook,
     generate_content_config_key_allowlist: Optional[AllowList] = None,
 ):
     wrapped_func = snapshot.async_generate_content
@@ -820,6 +825,7 @@ async def instrumented_generate_content(
             self,
             otel_wrapper,
             model,
+            upload_hook,
             generate_content_config_key_allowlist=generate_content_config_key_allowlist,
         )
         with helper.start_span_as_current_span(
@@ -861,6 +867,7 @@ async def instrumented_generate_content(
 def _create_instrumented_async_generate_content_stream(  # type: ignore
     snapshot: _MethodsSnapshot,
     otel_wrapper: OTelWrapper,
+    upload_hook: UploadHook,
     generate_content_config_key_allowlist: Optional[AllowList] = None,
 ):
     wrapped_func = snapshot.async_generate_content_stream
@@ -878,6 +885,7 @@ async def instrumented_generate_content_stream(
             self,
             otel_wrapper,
             model,
+            upload_hook,
             generate_content_config_key_allowlist=generate_content_config_key_allowlist,
         )
         with helper.start_span_as_current_span(
@@ -936,28 +944,33 @@ def uninstrument_generate_content(snapshot: object):
 
 def instrument_generate_content(
     otel_wrapper: OTelWrapper,
+    upload_hook: UploadHook,
     generate_content_config_key_allowlist: Optional[AllowList] = None,
 ) -> object:
     snapshot = _MethodsSnapshot()
     Models.generate_content = _create_instrumented_generate_content(
         snapshot,
         otel_wrapper,
+        upload_hook,
         generate_content_config_key_allowlist=generate_content_config_key_allowlist,
     )
     Models.generate_content_stream = _create_instrumented_generate_content_stream(
         snapshot,
         otel_wrapper,
+        upload_hook,
         generate_content_config_key_allowlist=generate_content_config_key_allowlist,
     )
     AsyncModels.generate_content = _create_instrumented_async_generate_content(
         snapshot,
         otel_wrapper,
+        upload_hook,
         generate_content_config_key_allowlist=generate_content_config_key_allowlist,
     )
     AsyncModels.generate_content_stream = (
         _create_instrumented_async_generate_content_stream(
             snapshot,
             otel_wrapper,
+            upload_hook,
             generate_content_config_key_allowlist=generate_content_config_key_allowlist,
         )
     )

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

@@ -18,6 +18,7 @@
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.metrics import get_meter_provider
 from opentelemetry.trace import get_tracer_provider
+from opentelemetry.util.genai.upload_hook import load_upload_hook
 
 from .allowlist_util import AllowList
 from .generate_content import (
@@ -58,8 +59,10 @@ def _instrument(self, **kwargs: Any):
             event_logger_provider=event_logger_provider,
             meter_provider=meter_provider,
         )
+        upload_hook = kwargs.get("upload_hook") or load_upload_hook()
         self._generate_content_snapshot = instrument_generate_content(
             otel_wrapper,
+            upload_hook,
             generate_content_config_key_allowlist=self._generate_content_config_key_allowlist,
         )