test: fix test_rag_agent test

Summary:

Test Plan:
This commit is contained in:
Eric Huang 2025-02-21 15:19:26 -08:00
parent e7d261ef4a
commit 01aef1382c

View file

@ -496,10 +496,11 @@ def test_rag_agent(llama_stack_client, agent_config):
stream=False, stream=False,
) )
# rag is called # 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 # document ids are present in metadata
assert "num-0" in response.steps[0].tool_responses[0].metadata["document_ids"] assert "num-0" in tool_execution_step.tool_responses[0].metadata["document_ids"]
assert expected_kw in response.output_message.content assert expected_kw in response.output_message.content.lower()
def test_rag_and_code_agent(llama_stack_client, agent_config): def test_rag_and_code_agent(llama_stack_client, agent_config):