From 178a5c3b936ac423dc770b15bf3cae3e93a910a9 Mon Sep 17 00:00:00 2001 From: reluctantfuturist Date: Thu, 10 Apr 2025 10:30:10 -0700 Subject: [PATCH] moved the test from test_telemetry to test_agents --- tests/integration/agents/test_agents.py | 55 ++++++++++++++ tests/integration/telemetry/test_telemetry.py | 74 ------------------- 2 files changed, 55 insertions(+), 74 deletions(-) diff --git a/tests/integration/agents/test_agents.py b/tests/integration/agents/test_agents.py index 7def55291..117d1079f 100644 --- a/tests/integration/agents/test_agents.py +++ b/tests/integration/agents/test_agents.py @@ -115,6 +115,61 @@ def test_agent_simple(llama_stack_client_with_mocked_inference, agent_config): assert "I can't" in logs_str +def test_agent_name(llama_stack_client, text_model_id): + agent_name = f"test-agent-{uuid4()}" + + 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()}") + + agent.create_turn( + messages=[ + { + "role": "user", + "content": "Give me a sentence that contains the word: hello", + } + ], + session_id=session_id, + stream=False, + ) + + 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) + + 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) + + 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) + + 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() + + def test_tool_config(llama_stack_client_with_mocked_inference, agent_config): common_params = dict( model="meta-llama/Llama-3.2-3B-Instruct", diff --git a/tests/integration/telemetry/test_telemetry.py b/tests/integration/telemetry/test_telemetry.py index 9f6d06b99..c46de3742 100644 --- a/tests/integration/telemetry/test_telemetry.py +++ b/tests/integration/telemetry/test_telemetry.py @@ -41,77 +41,3 @@ def test_agent_query_spans(llama_stack_client, text_model_id): 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()