@@ -122,25 +122,34 @@ def is_metrics_enabled() -> bool:
122122 return (os .getenv ("TRACELOOP_METRICS_ENABLED" ) or "true" ).lower () == "true"
123123
124124
125- def _set_input_attributes (span , instance , kwargs ):
125+ def _set_input_attributes (span , args = None , kwargs = None ):
126126 if not span .is_recording ():
127127 return
128128
129- if should_send_prompts () and kwargs is not None and len (kwargs ) > 0 :
130- prompt = kwargs .get ("prompt" )
131- if isinstance (prompt , list ):
132- for index , input in enumerate (prompt ):
133- _set_span_attribute (
134- span ,
135- f"{ SpanAttributes .LLM_PROMPTS } .{ index } .user" ,
136- input ,
137- )
138- else :
129+ args = args or ()
130+ kwargs = kwargs or {}
131+ prompt = kwargs .get ("prompt" )
132+ if not prompt and args :
133+ first_arg = args [0 ]
134+ if isinstance (first_arg , (str , list )):
135+ prompt = first_arg
136+
137+ if not prompt :
138+ return
139+
140+ if isinstance (prompt , list ):
141+ for index , input_text in enumerate (prompt ):
139142 _set_span_attribute (
140143 span ,
141- f"{ SpanAttributes .LLM_PROMPTS } .0 .user" ,
142- prompt ,
144+ f"{ SpanAttributes .LLM_PROMPTS } .{ index } .user" ,
145+ str ( input_text ). strip () ,
143146 )
147+ else :
148+ _set_span_attribute (
149+ span ,
150+ f"{ SpanAttributes .LLM_PROMPTS } .0.user" ,
151+ str (prompt ).strip (),
152+ )
144153
145154
146155def set_model_input_attributes (span , instance ):
@@ -480,11 +489,11 @@ def _handle_input(span, event_logger, name, instance, response_counter, args, kw
480489
481490 if "generate" in name :
482491 set_model_input_attributes (span , instance )
492+ if should_send_prompts ():
493+ _set_input_attributes (span , args = args , kwargs = kwargs )
483494
484495 if should_emit_events () and event_logger :
485496 _emit_input_events (args , kwargs , event_logger )
486- elif "generate" in name :
487- _set_input_attributes (span , instance , kwargs )
488497
489498
490499@dont_throw
0 commit comments