5454 InputMessage ,
5555 OutputMessage ,
5656)
57- from opentelemetry .util .genai .upload_hook import UploadHook
57+ from opentelemetry .util .genai .completion_hook import CompletionHook
5858
5959from .allowlist_util import AllowList
6060from .custom_semconv import GCP_GENAI_OPERATION_CONFIG
@@ -290,14 +290,14 @@ def __init__(
290290 models_object : Union [Models , AsyncModels ],
291291 otel_wrapper : OTelWrapper ,
292292 model : str ,
293- upload_hook : UploadHook ,
293+ completion_hook : CompletionHook ,
294294 generate_content_config_key_allowlist : Optional [AllowList ] = None ,
295295 ):
296296 self ._start_time = time .time_ns ()
297297 self ._otel_wrapper = otel_wrapper
298298 self ._genai_system = _determine_genai_system (models_object )
299299 self ._genai_request_model = model
300- self .upload_hook = upload_hook
300+ self .completion_hook = completion_hook
301301 self ._finish_reasons_set = set ()
302302 self ._error_type = None
303303 self ._input_tokens = 0
@@ -492,7 +492,7 @@ def _maybe_log_completion_details(
492492 event = Event (
493493 name = "gen_ai.client.inference.operation.details" , attributes = attributes
494494 )
495- self .upload_hook . upload (
495+ self .completion_hook . on_completion (
496496 inputs = input_messages ,
497497 outputs = output_messages ,
498498 system_instruction = system_instructions ,
@@ -689,7 +689,7 @@ def _record_duration_metric(self):
689689def _create_instrumented_generate_content (
690690 snapshot : _MethodsSnapshot ,
691691 otel_wrapper : OTelWrapper ,
692- upload_hook : UploadHook ,
692+ completion_hook : CompletionHook ,
693693 generate_content_config_key_allowlist : Optional [AllowList ] = None ,
694694):
695695 wrapped_func = snapshot .generate_content
@@ -707,7 +707,7 @@ def instrumented_generate_content(
707707 self ,
708708 otel_wrapper ,
709709 model ,
710- upload_hook ,
710+ completion_hook ,
711711 generate_content_config_key_allowlist = generate_content_config_key_allowlist ,
712712 )
713713 with helper .start_span_as_current_span (
@@ -748,7 +748,7 @@ def instrumented_generate_content(
748748def _create_instrumented_generate_content_stream (
749749 snapshot : _MethodsSnapshot ,
750750 otel_wrapper : OTelWrapper ,
751- upload_hook : UploadHook ,
751+ completion_hook : CompletionHook ,
752752 generate_content_config_key_allowlist : Optional [AllowList ] = None ,
753753):
754754 wrapped_func = snapshot .generate_content_stream
@@ -766,7 +766,7 @@ def instrumented_generate_content_stream(
766766 self ,
767767 otel_wrapper ,
768768 model ,
769- upload_hook ,
769+ completion_hook ,
770770 generate_content_config_key_allowlist = generate_content_config_key_allowlist ,
771771 )
772772 with helper .start_span_as_current_span (
@@ -807,7 +807,7 @@ def instrumented_generate_content_stream(
807807def _create_instrumented_async_generate_content (
808808 snapshot : _MethodsSnapshot ,
809809 otel_wrapper : OTelWrapper ,
810- upload_hook : UploadHook ,
810+ completion_hook : CompletionHook ,
811811 generate_content_config_key_allowlist : Optional [AllowList ] = None ,
812812):
813813 wrapped_func = snapshot .async_generate_content
@@ -825,7 +825,7 @@ async def instrumented_generate_content(
825825 self ,
826826 otel_wrapper ,
827827 model ,
828- upload_hook ,
828+ completion_hook ,
829829 generate_content_config_key_allowlist = generate_content_config_key_allowlist ,
830830 )
831831 with helper .start_span_as_current_span (
@@ -867,7 +867,7 @@ async def instrumented_generate_content(
867867def _create_instrumented_async_generate_content_stream ( # type: ignore
868868 snapshot : _MethodsSnapshot ,
869869 otel_wrapper : OTelWrapper ,
870- upload_hook : UploadHook ,
870+ completion_hook : CompletionHook ,
871871 generate_content_config_key_allowlist : Optional [AllowList ] = None ,
872872):
873873 wrapped_func = snapshot .async_generate_content_stream
@@ -885,7 +885,7 @@ async def instrumented_generate_content_stream(
885885 self ,
886886 otel_wrapper ,
887887 model ,
888- upload_hook ,
888+ completion_hook ,
889889 generate_content_config_key_allowlist = generate_content_config_key_allowlist ,
890890 )
891891 with helper .start_span_as_current_span (
@@ -944,33 +944,33 @@ def uninstrument_generate_content(snapshot: object):
944944
945945def instrument_generate_content (
946946 otel_wrapper : OTelWrapper ,
947- upload_hook : UploadHook ,
947+ completion_hook : CompletionHook ,
948948 generate_content_config_key_allowlist : Optional [AllowList ] = None ,
949949) -> object :
950950 snapshot = _MethodsSnapshot ()
951951 Models .generate_content = _create_instrumented_generate_content (
952952 snapshot ,
953953 otel_wrapper ,
954- upload_hook ,
954+ completion_hook ,
955955 generate_content_config_key_allowlist = generate_content_config_key_allowlist ,
956956 )
957957 Models .generate_content_stream = _create_instrumented_generate_content_stream (
958958 snapshot ,
959959 otel_wrapper ,
960- upload_hook ,
960+ completion_hook ,
961961 generate_content_config_key_allowlist = generate_content_config_key_allowlist ,
962962 )
963963 AsyncModels .generate_content = _create_instrumented_async_generate_content (
964964 snapshot ,
965965 otel_wrapper ,
966- upload_hook ,
966+ completion_hook ,
967967 generate_content_config_key_allowlist = generate_content_config_key_allowlist ,
968968 )
969969 AsyncModels .generate_content_stream = (
970970 _create_instrumented_async_generate_content_stream (
971971 snapshot ,
972972 otel_wrapper ,
973- upload_hook ,
973+ completion_hook ,
974974 generate_content_config_key_allowlist = generate_content_config_key_allowlist ,
975975 )
976976 )
0 commit comments