phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-02 08:44:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add more attributes for inference, shields and memory

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-25 13:31:30 -08:00
parent b3021ea2da
commit 54bc5f2d55
1 changed files with 6 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

8
llama_stack/providers/inline/agents/meta_reference/agent_instance.py
View file

@ -423,7 +423,7 @@ class ChatAgent(ShieldRunnerMixin):
content = "" content = ""
stop_reason = None stop_reason = None
with tracing.span("inference"): with tracing.span("inference") as span:
async for chunk in await self.inference_api.chat_completion( async for chunk in await self.inference_api.chat_completion(
self.agent_config.model, self.agent_config.model,
input_messages, input_messages,
@ -443,7 +443,6 @@ class ChatAgent(ShieldRunnerMixin):
if isinstance(delta, ToolCallDelta): if isinstance(delta, ToolCallDelta):
if delta.parse_status == ToolCallParseStatus.success: if delta.parse_status == ToolCallParseStatus.success:
tool_calls.append(delta.content) tool_calls.append(delta.content)
if stream: if stream:
yield AgentTurnResponseStreamChunk( yield AgentTurnResponseStreamChunk(
event=AgentTurnResponseEvent( event=AgentTurnResponseEvent(
@ -473,6 +472,11 @@ class ChatAgent(ShieldRunnerMixin):
if event.stop_reason is not None: if event.stop_reason is not None:
stop_reason = event.stop_reason 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 stop_reason = stop_reason or StopReason.out_of_tokens