mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-28 13:08:46 +00:00
6a5b73ca7c
Allow users to name an agent and use the name in telemetry instead of relying on randomly generated agent_ids. This improves the developer experience by making it easier to find specific agents in telemetry logs. Closes #1832
117 lines
3.8 KiB
Python
117 lines
3.8 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
import time
|
|
from uuid import uuid4
|
|
|
|
from llama_stack_client import Agent
|
|
|
|
|
|
def test_agent_query_spans(llama_stack_client, text_model_id):
|
|
agent = Agent(llama_stack_client, model=text_model_id, instructions="You are a helpful assistant")
|
|
session_id = agent.create_session(f"test-session-{uuid4()}")
|
|
agent.create_turn(
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": "Give me a sentence that contains the word: hello",
|
|
}
|
|
],
|
|
session_id=session_id,
|
|
stream=False,
|
|
)
|
|
|
|
# Wait for the span to be logged
|
|
time.sleep(2)
|
|
|
|
agent_logs = []
|
|
|
|
for span in llama_stack_client.telemetry.query_spans(
|
|
attribute_filters=[
|
|
{"key": "session_id", "op": "eq", "value": session_id},
|
|
],
|
|
attributes_to_return=["input", "output"],
|
|
):
|
|
if span.attributes["output"] != "no shields":
|
|
agent_logs.append(span.attributes)
|
|
|
|
assert len(agent_logs) == 1
|
|
assert "Give me a sentence that contains the word: hello" in agent_logs[0]["input"]
|
|
assert "hello" in agent_logs[0]["output"].lower()
|
|
|
|
|
|
def test_agent_name_filtering(llama_stack_client, text_model_id):
|
|
# Create an agent with a specific name
|
|
agent_name = f"test-agent-{uuid4()}"
|
|
print(f"Using agent_name: {agent_name}")
|
|
|
|
agent = Agent(
|
|
llama_stack_client,
|
|
model=text_model_id,
|
|
instructions="You are a helpful assistant",
|
|
name=agent_name,
|
|
)
|
|
session_id = agent.create_session(f"test-session-{uuid4()}")
|
|
print(f"Created session_id: {session_id}")
|
|
|
|
agent.create_turn(
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": "Give me a sentence that contains the word: hello",
|
|
}
|
|
],
|
|
session_id=session_id,
|
|
stream=False,
|
|
)
|
|
|
|
# Wait for the span to be logged - increase the time to ensure it's processed
|
|
time.sleep(5)
|
|
|
|
# Query spans filtered by session_id to see what's available
|
|
all_spans = []
|
|
for span in llama_stack_client.telemetry.query_spans(
|
|
attribute_filters=[
|
|
{"key": "session_id", "op": "eq", "value": session_id},
|
|
],
|
|
attributes_to_return=["input", "output", "agent_name", "agent_id", "session_id"],
|
|
):
|
|
all_spans.append(span.attributes)
|
|
|
|
print(f"All spans for session {session_id}:")
|
|
for span in all_spans:
|
|
print(f"Span attributes: {span}")
|
|
|
|
# Query all spans to see if any have the agent_name attribute
|
|
agent_name_spans = []
|
|
for span in llama_stack_client.telemetry.query_spans(
|
|
attribute_filters=[],
|
|
attributes_to_return=["agent_name"],
|
|
):
|
|
if "agent_name" in span.attributes:
|
|
agent_name_spans.append(span.attributes)
|
|
|
|
print("All spans with agent_name attribute:")
|
|
for span in agent_name_spans:
|
|
print(f"Span with agent_name: {span}")
|
|
|
|
# Query spans filtered by agent name
|
|
agent_logs = []
|
|
for span in llama_stack_client.telemetry.query_spans(
|
|
attribute_filters=[
|
|
{"key": "agent_name", "op": "eq", "value": agent_name},
|
|
],
|
|
attributes_to_return=["input", "output", "agent_name"],
|
|
):
|
|
if "output" in span.attributes and span.attributes["output"] != "no shields":
|
|
agent_logs.append(span.attributes)
|
|
|
|
print(f"Found {len(agent_logs)} spans filtered by agent_name")
|
|
|
|
assert len(agent_logs) == 1
|
|
assert agent_logs[0]["agent_name"] == agent_name
|
|
assert "Give me a sentence that contains the word: hello" in agent_logs[0]["input"]
|
|
assert "hello" in agent_logs[0]["output"].lower()
|