From 54bc5f2d556edf5bbb0b7ce9c911253324b121b2 Mon Sep 17 00:00:00 2001 From: Dinesh Yeduguru Date: Mon, 25 Nov 2024 13:31:30 -0800 Subject: [PATCH] add more attributes for inference, shields and memory --- .../inline/agents/meta_reference/agent_instance.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py index 2c26b2d9e..7e2979c24 100644 --- a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py +++ b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py @@ -423,7 +423,7 @@ class ChatAgent(ShieldRunnerMixin): content = "" stop_reason = None - with tracing.span("inference"): + with tracing.span("inference") as span: async for chunk in await self.inference_api.chat_completion( self.agent_config.model, input_messages, @@ -443,7 +443,6 @@ class ChatAgent(ShieldRunnerMixin): if isinstance(delta, ToolCallDelta): if delta.parse_status == ToolCallParseStatus.success: tool_calls.append(delta.content) - if stream: yield AgentTurnResponseStreamChunk( event=AgentTurnResponseEvent( @@ -473,6 +472,11 @@ class ChatAgent(ShieldRunnerMixin): if event.stop_reason is not None: stop_reason = event.stop_reason + span.set_attribute("stop_reason", stop_reason) + span.set_attribute("content", content) + span.set_attribute( + "tool_calls", [tc.model_dump_json() for tc in tool_calls] + ) stop_reason = stop_reason or StopReason.out_of_tokens