@@ -110,6 +110,7 @@ def __init__(
110
110
persistence_store : KVStore ,
111
111
created_at : str ,
112
112
policy : list [AccessRule ],
113
+ telemetry_enabled : bool = False ,
113
114
):
114
115
self .agent_id = agent_id
115
116
self .agent_config = agent_config
@@ -120,6 +121,7 @@ def __init__(
120
121
self .tool_runtime_api = tool_runtime_api
121
122
self .tool_groups_api = tool_groups_api
122
123
self .created_at = created_at
124
+ self .telemetry_enabled = telemetry_enabled
123
125
124
126
ShieldRunnerMixin .__init__ (
125
127
self ,
@@ -188,28 +190,30 @@ async def get_messages_from_turns(self, turns: list[Turn]) -> list[Message]:
188
190
189
191
async def create_and_execute_turn (self , request : AgentTurnCreateRequest ) -> AsyncGenerator :
190
192
turn_id = str (uuid .uuid4 ())
191
- span = tracing .get_current_span ()
192
- if span :
193
- span .set_attribute ("session_id" , request .session_id )
194
- span .set_attribute ("agent_id" , self .agent_id )
195
- span .set_attribute ("request" , request .model_dump_json ())
196
- span .set_attribute ("turn_id" , turn_id )
197
- if self .agent_config .name :
198
- span .set_attribute ("agent_name" , self .agent_config .name )
193
+ if self .telemetry_enabled :
194
+ span = tracing .get_current_span ()
195
+ if span is not None :
196
+ span .set_attribute ("session_id" , request .session_id )
197
+ span .set_attribute ("agent_id" , self .agent_id )
198
+ span .set_attribute ("request" , request .model_dump_json ())
199
+ span .set_attribute ("turn_id" , turn_id )
200
+ if self .agent_config .name :
201
+ span .set_attribute ("agent_name" , self .agent_config .name )
199
202
200
203
await self ._initialize_tools (request .toolgroups )
201
204
async for chunk in self ._run_turn (request , turn_id ):
202
205
yield chunk
203
206
204
207
async def resume_turn (self , request : AgentTurnResumeRequest ) -> AsyncGenerator :
205
- span = tracing .get_current_span ()
206
- if span :
207
- span .set_attribute ("agent_id" , self .agent_id )
208
- span .set_attribute ("session_id" , request .session_id )
209
- span .set_attribute ("request" , request .model_dump_json ())
210
- span .set_attribute ("turn_id" , request .turn_id )
211
- if self .agent_config .name :
212
- span .set_attribute ("agent_name" , self .agent_config .name )
208
+ if self .telemetry_enabled :
209
+ span = tracing .get_current_span ()
210
+ if span is not None :
211
+ span .set_attribute ("agent_id" , self .agent_id )
212
+ span .set_attribute ("session_id" , request .session_id )
213
+ span .set_attribute ("request" , request .model_dump_json ())
214
+ span .set_attribute ("turn_id" , request .turn_id )
215
+ if self .agent_config .name :
216
+ span .set_attribute ("agent_name" , self .agent_config .name )
213
217
214
218
await self ._initialize_tools ()
215
219
async for chunk in self ._run_turn (request ):
@@ -395,9 +399,12 @@ async def run_multiple_shields_wrapper(
395
399
touchpoint : str ,
396
400
) -> AsyncGenerator :
397
401
async with tracing .span ("run_shields" ) as span :
398
- span .set_attribute ("input" , [m .model_dump_json () for m in messages ])
402
+ if self .telemetry_enabled and span is not None :
403
+ span .set_attribute ("input" , [m .model_dump_json () for m in messages ])
404
+ if len (shields ) == 0 :
405
+ span .set_attribute ("output" , "no shields" )
406
+
399
407
if len (shields ) == 0 :
400
- span .set_attribute ("output" , "no shields" )
401
408
return
402
409
403
410
step_id = str (uuid .uuid4 ())
@@ -430,7 +437,8 @@ async def run_multiple_shields_wrapper(
430
437
)
431
438
)
432
439
)
433
- span .set_attribute ("output" , e .violation .model_dump_json ())
440
+ if self .telemetry_enabled and span is not None :
441
+ span .set_attribute ("output" , e .violation .model_dump_json ())
434
442
435
443
yield CompletionMessage (
436
444
content = str (e ),
@@ -453,7 +461,8 @@ async def run_multiple_shields_wrapper(
453
461
)
454
462
)
455
463
)
456
- span .set_attribute ("output" , "no violations" )
464
+ if self .telemetry_enabled and span is not None :
465
+ span .set_attribute ("output" , "no violations" )
457
466
458
467
async def _run (
459
468
self ,
@@ -518,8 +527,9 @@ async def _run(
518
527
stop_reason : StopReason | None = None
519
528
520
529
async with tracing .span ("inference" ) as span :
521
- if self .agent_config .name :
522
- span .set_attribute ("agent_name" , self .agent_config .name )
530
+ if self .telemetry_enabled and span is not None :
531
+ if self .agent_config .name :
532
+ span .set_attribute ("agent_name" , self .agent_config .name )
523
533
524
534
def _serialize_nested (value ):
525
535
"""Recursively serialize nested Pydantic models to dicts."""
@@ -637,18 +647,19 @@ def _add_type(openai_msg: dict) -> OpenAIMessageParam:
637
647
else :
638
648
raise ValueError (f"Unexpected delta type { type (delta )} " )
639
649
640
- span .set_attribute ("stop_reason" , stop_reason or StopReason .end_of_turn )
641
- span .set_attribute (
642
- "input" ,
643
- json .dumps ([json .loads (m .model_dump_json ()) for m in input_messages ]),
644
- )
645
- output_attr = json .dumps (
646
- {
647
- "content" : content ,
648
- "tool_calls" : [json .loads (t .model_dump_json ()) for t in tool_calls ],
649
- }
650
- )
651
- span .set_attribute ("output" , output_attr )
650
+ if self .telemetry_enabled and span is not None :
651
+ span .set_attribute ("stop_reason" , stop_reason or StopReason .end_of_turn )
652
+ span .set_attribute (
653
+ "input" ,
654
+ json .dumps ([json .loads (m .model_dump_json ()) for m in input_messages ]),
655
+ )
656
+ output_attr = json .dumps (
657
+ {
658
+ "content" : content ,
659
+ "tool_calls" : [json .loads (t .model_dump_json ()) for t in tool_calls ],
660
+ }
661
+ )
662
+ span .set_attribute ("output" , output_attr )
652
663
653
664
n_iter += 1
654
665
await self .storage .set_num_infer_iters_in_turn (session_id , turn_id , n_iter )
@@ -756,7 +767,9 @@ def _add_type(openai_msg: dict) -> OpenAIMessageParam:
756
767
{
757
768
"tool_name" : tool_call .tool_name ,
758
769
"input" : message .model_dump_json (),
759
- },
770
+ }
771
+ if self .telemetry_enabled
772
+ else {},
760
773
) as span :
761
774
tool_execution_start_time = datetime .now (UTC ).isoformat ()
762
775
tool_result = await self .execute_tool_call_maybe (
@@ -771,7 +784,8 @@ def _add_type(openai_msg: dict) -> OpenAIMessageParam:
771
784
call_id = tool_call .call_id ,
772
785
content = tool_result .content ,
773
786
)
774
- span .set_attribute ("output" , result_message .model_dump_json ())
787
+ if self .telemetry_enabled and span is not None :
788
+ span .set_attribute ("output" , result_message .model_dump_json ())
775
789
776
790
# Store tool execution step
777
791
tool_execution_step = ToolExecutionStep (
0 commit comments