3333from opentelemetry .trace import Span , SpanKind , Tracer
3434from opentelemetry .trace .propagation import set_span_in_context
3535from opentelemetry .util .genai .handler import TelemetryHandler
36+ from opentelemetry .util .genai .invocation import InferenceInvocation
3637from opentelemetry .util .genai .types import (
3738 ContentCapturingMode ,
3839 Error ,
39- LLMInvocation , # pylint: disable=no-name-in-module # TODO: migrate to InferenceInvocation
4040 OutputMessage ,
4141 Text ,
4242 ToolCallRequest ,
@@ -68,7 +68,9 @@ def traced_method(wrapped, instance, args, kwargs):
6868 ** get_llm_request_attributes (kwargs , instance , False )
6969 }
7070
71- span_name = f"{ span_attributes [GenAIAttributes .GEN_AI_OPERATION_NAME ]} { span_attributes [GenAIAttributes .GEN_AI_REQUEST_MODEL ]} "
71+ operation_name = span_attributes [GenAIAttributes .GEN_AI_OPERATION_NAME ]
72+ model = span_attributes .get (GenAIAttributes .GEN_AI_REQUEST_MODEL )
73+ span_name = f"{ operation_name } { model } " if model else operation_name
7274 with tracer .start_as_current_span (
7375 name = span_name ,
7476 kind = SpanKind .CLIENT ,
@@ -128,10 +130,8 @@ def chat_completions_create_v_new(
128130 capture_content = content_capturing_mode != ContentCapturingMode .NO_CONTENT
129131
130132 def traced_method (wrapped , instance , args , kwargs ):
131- chat_invocation = handler .start_llm (
132- create_chat_invocation (
133- kwargs , instance , capture_content = capture_content
134- )
133+ chat_invocation = create_chat_invocation (
134+ handler , kwargs , instance , capture_content = capture_content
135135 )
136136
137137 try :
@@ -143,18 +143,16 @@ def traced_method(wrapped, instance, args, kwargs):
143143 parsed_result = result
144144 if is_streaming (kwargs ):
145145 return ChatStreamWrapper (
146- parsed_result , handler , chat_invocation , capture_content
146+ parsed_result , chat_invocation , capture_content
147147 )
148148
149149 _set_response_properties (
150150 chat_invocation , parsed_result , capture_content
151151 )
152- handler . stop_llm ( chat_invocation )
152+ chat_invocation . stop ( )
153153 return result
154154 except Exception as error :
155- handler .fail_llm (
156- chat_invocation , Error (type = type (error ), message = str (error ))
157- )
155+ chat_invocation .fail (Error (type = type (error ), message = str (error )))
158156 raise
159157
160158 return traced_method
@@ -173,7 +171,9 @@ async def traced_method(wrapped, instance, args, kwargs):
173171 ** get_llm_request_attributes (kwargs , instance , False )
174172 }
175173
176- span_name = f"{ span_attributes [GenAIAttributes .GEN_AI_OPERATION_NAME ]} { span_attributes [GenAIAttributes .GEN_AI_REQUEST_MODEL ]} "
174+ operation_name = span_attributes [GenAIAttributes .GEN_AI_OPERATION_NAME ]
175+ model = span_attributes .get (GenAIAttributes .GEN_AI_REQUEST_MODEL )
176+ span_name = f"{ operation_name } { model } " if model else operation_name
177177 with tracer .start_as_current_span (
178178 name = span_name ,
179179 kind = SpanKind .CLIENT ,
@@ -232,10 +232,8 @@ def async_chat_completions_create_v_new(
232232 capture_content = content_capturing_mode != ContentCapturingMode .NO_CONTENT
233233
234234 async def traced_method (wrapped , instance , args , kwargs ):
235- chat_invocation = handler .start_llm (
236- create_chat_invocation (
237- kwargs , instance , capture_content = capture_content
238- )
235+ chat_invocation = create_chat_invocation (
236+ handler , kwargs , instance , capture_content = capture_content
239237 )
240238
241239 try :
@@ -247,19 +245,17 @@ async def traced_method(wrapped, instance, args, kwargs):
247245 parsed_result = result
248246 if is_streaming (kwargs ):
249247 return ChatStreamWrapper (
250- parsed_result , handler , chat_invocation , capture_content
248+ parsed_result , chat_invocation , capture_content
251249 )
252250
253251 _set_response_properties (
254252 chat_invocation , parsed_result , capture_content
255253 )
256- handler . stop_llm ( chat_invocation )
254+ chat_invocation . stop ( )
257255 return result
258256
259257 except Exception as error :
260- handler .fail_llm (
261- chat_invocation , Error (type = type (error ), message = str (error ))
262- )
258+ chat_invocation .fail (Error (type = type (error ), message = str (error )))
263259 raise
264260
265261 return traced_method
@@ -373,7 +369,9 @@ async def traced_method(wrapped, instance, args, kwargs):
373369
374370def _get_embeddings_span_name (span_attributes ):
375371 """Get span name for embeddings operations."""
376- return f"{ span_attributes [GenAIAttributes .GEN_AI_OPERATION_NAME ]} { span_attributes [GenAIAttributes .GEN_AI_REQUEST_MODEL ]} "
372+ operation_name = span_attributes [GenAIAttributes .GEN_AI_OPERATION_NAME ]
373+ model = span_attributes .get (GenAIAttributes .GEN_AI_REQUEST_MODEL )
374+ return f"{ operation_name } { model } " if model else operation_name
377375
378376
379377def _record_metrics (
@@ -495,8 +493,8 @@ def _set_response_attributes(span, result):
495493
496494
497495def _set_response_properties (
498- chat_invocation : LLMInvocation , result , capture_content : bool
499- ) -> LLMInvocation :
496+ chat_invocation : InferenceInvocation , result , capture_content : bool
497+ ) -> InferenceInvocation :
500498 if getattr (result , "model" , None ):
501499 chat_invocation .response_model_name = result .model
502500
@@ -868,8 +866,7 @@ def cleanup(self, error: Optional[BaseException] = None):
868866
869867
870868class ChatStreamWrapper (BaseStreamWrapper ):
871- handler : TelemetryHandler
872- invocation : LLMInvocation
869+ invocation : InferenceInvocation
873870 response_id : Optional [str ] = None
874871 response_model : Optional [str ] = None
875872 service_tier : Optional [str ] = None
@@ -880,13 +877,11 @@ class ChatStreamWrapper(BaseStreamWrapper):
880877 def __init__ (
881878 self ,
882879 stream : Stream ,
883- handler : TelemetryHandler ,
884- invocation : LLMInvocation ,
880+ invocation : InferenceInvocation ,
885881 capture_content : bool ,
886882 ):
887883 super ().__init__ (stream , capture_content = capture_content )
888884 self .stream = stream
889- self .handler = handler
890885 self .invocation = invocation
891886 self .choice_buffers = []
892887
@@ -944,9 +939,7 @@ def cleanup(self, error: Optional[BaseException] = None):
944939 self ._set_output_messages ()
945940
946941 if error :
947- self .handler .fail_llm (
948- self .invocation , Error (type = type (error ), message = str (error ))
949- )
942+ self .invocation .fail (Error (type = type (error ), message = str (error )))
950943 else :
951- self .handler . stop_llm ( self . invocation )
944+ self .invocation . stop ( )
952945 self ._started = False
0 commit comments