From 67b35613bb5e0f7e5222571e05d1c8c133494254 Mon Sep 17 00:00:00 2001 From: Dinesh Yeduguru Date: Wed, 8 Jan 2025 14:05:28 -0800 Subject: [PATCH] test turn overrides in unit tests --- .../agents/meta_reference/tests/test_chat_agent.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/llama_stack/providers/inline/agents/meta_reference/tests/test_chat_agent.py b/llama_stack/providers/inline/agents/meta_reference/tests/test_chat_agent.py index 6e789bf19..6b8a846ee 100644 --- a/llama_stack/providers/inline/agents/meta_reference/tests/test_chat_agent.py +++ b/llama_stack/providers/inline/agents/meta_reference/tests/test_chat_agent.py @@ -12,6 +12,7 @@ from llama_models.llama3.api.datatypes import BuiltinTool from llama_stack.apis.agents import ( AgentConfig, + AgentToolGroupWithArgs, AgentTurnCreateRequest, AgentTurnResponseTurnCompletePayload, StepType, @@ -478,3 +479,15 @@ async def test_chat_agent_tools( assert MEMORY_QUERY_TOOL in tool_defs if expected_code_interpreter: assert BuiltinTool.code_interpreter in tool_defs + if expected_memory and expected_code_interpreter: + # override the tools for turn + new_tool_defs, _ = await chat_agent._get_tool_defs( + toolgroups_for_turn=[ + AgentToolGroupWithArgs( + name=MEMORY_TOOLGROUP, + args={"memory_banks": ["test_memory_bank"]}, + ) + ] + ) + assert MEMORY_QUERY_TOOL in new_tool_defs + assert BuiltinTool.code_interpreter not in new_tool_defs