test: fix test_rag_agent test (#1215)

Summary:

Test Plan:
LLAMA_STACK_CONFIG=fireworks pytest -s -v
tests/client-sdk/agents/test_agents.py::test_rag_agent --safety-shield
meta-llama/Llama-Guard-3-8B
This commit is contained in:
ehhuang 2025-02-21 15:24:28 -08:00 committed by GitHub
parent e7d261ef4a
commit bf38d0aba0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -496,10 +496,11 @@ def test_rag_agent(llama_stack_client, agent_config):
stream=False,
)
# rag is called
assert response.steps[0].tool_calls[0].tool_name == "query_from_memory"
tool_execution_step = next(step for step in response.steps if step.step_type == "tool_execution")
assert tool_execution_step.tool_calls[0].tool_name == "query_from_memory"
# document ids are present in metadata
assert "num-0" in response.steps[0].tool_responses[0].metadata["document_ids"]
assert expected_kw in response.output_message.content
assert "num-0" in tool_execution_step.tool_responses[0].metadata["document_ids"]
assert expected_kw in response.output_message.content.lower()
def test_rag_and_code_agent(llama_stack_client, agent_config):