3030from opentelemetry import baggage
3131from opentelemetry .trace import Span
3232from opentelemetry .trace .status import StatusCode
33+ import os
3334
3435
3536def estimate_tokens (prompt ):
@@ -42,6 +43,9 @@ def estimate_tokens(prompt):
4243
4344
4445def set_event_completion_chunk (span : Span , chunk ):
46+ enabled = os .environ .get ("TRACE_PROMPT_COMPLETION_DATA" , "true" )
47+ if enabled .lower () == "false" :
48+ return
4549 span .add_event (
4650 name = SpanAttributes .LLM_CONTENT_COMPLETION_CHUNK ,
4751 attributes = {
@@ -203,6 +207,9 @@ def get_tool_calls(item):
203207
204208
205209def set_event_completion (span : Span , result_content ):
210+ enabled = os .environ .get ("TRACE_PROMPT_COMPLETION_DATA" , "true" )
211+ if enabled .lower () == "false" :
212+ return
206213
207214 span .add_event (
208215 name = SpanAttributes .LLM_CONTENT_COMPLETION ,
@@ -352,15 +359,9 @@ def process_chunk(self, chunk):
352359 )
353360 self .completion_tokens += token_counts
354361 content .append (tool_call .function .arguments )
355- self .span .add_event (
356- Event .STREAM_OUTPUT .value ,
357- {
358- SpanAttributes .LLM_CONTENT_COMPLETION_CHUNK : (
359- "" .join (content )
360- if len (content ) > 0 and content [0 ] is not None
361- else ""
362- )
363- },
362+ set_event_completion_chunk (
363+ self .span ,
364+ "" .join (content ) if len (content ) > 0 and content [0 ] is not None else "" ,
364365 )
365366 if content :
366367 self .result_content .append (content [0 ])
@@ -369,16 +370,11 @@ def process_chunk(self, chunk):
369370 token_counts = estimate_tokens (chunk .text )
370371 self .completion_tokens += token_counts
371372 content = [chunk .text ]
372- self .span .add_event (
373- Event .STREAM_OUTPUT .value ,
374- {
375- SpanAttributes .LLM_CONTENT_COMPLETION_CHUNK : (
376- "" .join (content )
377- if len (content ) > 0 and content [0 ] is not None
378- else ""
379- )
380- },
373+ set_event_completion_chunk (
374+ self .span ,
375+ "" .join (content ) if len (content ) > 0 and content [0 ] is not None else "" ,
381376 )
377+
382378 if content :
383379 self .result_content .append (content [0 ])
384380
0 commit comments