Skip to content

Commit 936fc4f

Browse files
hillhackhillhack
committed
test: update tests for dynamic tool filtering
1 parent 2423184 commit 936fc4f

File tree

5 files changed

+46
-31
lines changed

5 files changed

+46
-31
lines changed

‎tests/test_integration/test_memory_reasoning.py‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,8 @@ def make_mock_agent(memory_instance, *, step_prompt="You are an agent in a simul
3838
agent.step_prompt = step_prompt
3939
agent.llm = Mock()
4040
agent.tool_manager = Mock()
41-
agent.tool_manager.get_all_tools_schema.return_value = {}
41+
agent.tool_manager.get_annotated_tools_schema.return_value = {}
42+
agent.tool_manager.get_feasible_tools_schema.return_value = {}
4243
agent._step_display_data = {}
4344

4445
# Wire memory
@@ -169,7 +170,8 @@ def _setup(self, monkeypatch):
169170
agent.step_prompt = "You are an agent"
170171
agent.llm = Mock()
171172
agent.tool_manager = Mock()
172-
agent.tool_manager.get_all_tools_schema.return_value = {}
173+
agent.tool_manager.get_annotated_tools_schema.return_value = {}
174+
agent.tool_manager.get_feasible_tools_schema.return_value = {}
173175
agent._step_display_data = {}
174176

175177
memory = STLTMemory(
@@ -244,7 +246,8 @@ def _setup(self, monkeypatch):
244246
agent.step_prompt = "You are an agent"
245247
agent.llm = Mock()
246248
agent.tool_manager = Mock()
247-
agent.tool_manager.get_all_tools_schema.return_value = {}
249+
agent.tool_manager.get_annotated_tools_schema.return_value = {}
250+
agent.tool_manager.get_feasible_tools_schema.return_value = {}
248251
agent._step_display_data = {}
249252

250253
memory = EpisodicMemory(
@@ -328,7 +331,8 @@ def _setup(self, monkeypatch):
328331
agent.step_prompt = "You are an agent"
329332
agent.llm = Mock()
330333
agent.tool_manager = Mock()
331-
agent.tool_manager.get_all_tools_schema.return_value = {}
334+
agent.tool_manager.get_annotated_tools_schema.return_value = {}
335+
agent.tool_manager.get_feasible_tools_schema.return_value = {}
332336

333337
memory = STLTMemory(
334338
agent=agent,
@@ -444,7 +448,8 @@ def _setup(self, monkeypatch):
444448
agent.step_prompt = "You are an agent"
445449
agent.llm = Mock()
446450
agent.tool_manager = Mock()
447-
agent.tool_manager.get_all_tools_schema.return_value = {}
451+
agent.tool_manager.get_annotated_tools_schema.return_value = {}
452+
agent.tool_manager.get_feasible_tools_schema.return_value = {}
448453

449454
memory = EpisodicMemory(
450455
agent=agent,
@@ -556,7 +561,7 @@ def _setup(self, monkeypatch):
556561
agent.step_prompt = "You are an agent"
557562
agent.llm = Mock()
558563
agent.tool_manager = Mock()
559-
agent.tool_manager.get_all_tools_schema.return_value = {}
564+
agent.tool_manager.get_annotated_tools_schema.return_value = {}
560565
agent._step_display_data = {}
561566
default_obs = Observation(step=1, self_state={}, local_state={})
562567
agent.generate_obs = Mock(return_value=default_obs)
@@ -637,7 +642,7 @@ def _setup(self, monkeypatch):
637642
agent.step_prompt = "You are an agent"
638643
agent.llm = Mock()
639644
agent.tool_manager = Mock()
640-
agent.tool_manager.get_all_tools_schema.return_value = {}
645+
agent.tool_manager.get_annotated_tools_schema.return_value = {}
641646
agent._step_display_data = {}
642647
default_obs = Observation(step=1, self_state={}, local_state={})
643648
agent.generate_obs = Mock(return_value=default_obs)

‎tests/test_reasoning/test_cot.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,8 @@ def test_plan_with_selected_tools(self):
107107
mock_agent.memory.add_to_memory = Mock()
108108
mock_agent.llm = Mock()
109109
mock_agent.tool_manager = Mock()
110-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
110+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
111+
mock_agent.tool_manager.get_feasible_tools_schema.return_value = {}
111112
mock_agent._step_display_data = {} # Use real dict instead of Mock
112113
# Mock the LLM response for planning
113114
mock_plan_response = Mock()
@@ -132,7 +133,8 @@ def test_plan_with_selected_tools(self):
132133
assert isinstance(result, Plan)
133134
assert result.ttl == 3
134135
# Check that tool schema was called with selected tools
135-
assert mock_agent.tool_manager.get_all_tools_schema.call_count == 2
136+
assert mock_agent.tool_manager.get_annotated_tools_schema.call_count == 1
137+
assert mock_agent.tool_manager.get_feasible_tools_schema.call_count == 1
136138

137139
def test_plan_no_prompt_error(self):
138140
"""Test plan method raises error when no prompt is provided."""
@@ -158,7 +160,8 @@ def test_aplan_uses_step_prompt_when_no_prompt_given(self):
158160
mock_agent.memory.aadd_to_memory = AsyncMock()
159161
mock_agent.llm = Mock()
160162
mock_agent.tool_manager = Mock()
161-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
163+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
164+
mock_agent.tool_manager.get_feasible_tools_schema.return_value = {}
162165
mock_agent._step_display_data = {}
163166

164167
mock_plan_response = Mock()
@@ -192,7 +195,8 @@ def test_aplan_async_version(self):
192195
mock_agent.memory.aadd_to_memory = AsyncMock()
193196
mock_agent.llm = Mock()
194197
mock_agent.tool_manager = Mock()
195-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
198+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
199+
mock_agent.tool_manager.get_feasible_tools_schema.return_value = {}
196200
mock_agent._step_display_data = {} # Use real dict instead of Mock
197201

198202
# Mock the async LLM response for planning

‎tests/test_reasoning/test_react.py‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ def test_plan_with_prompt(self):
8282
mock_agent.memory.add_to_memory = Mock()
8383
mock_agent.llm = Mock()
8484
mock_agent.tool_manager = Mock()
85-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
85+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
8686

8787
# Mock the LLM response
8888
mock_response = Mock()
@@ -117,7 +117,7 @@ def test_plan_with_selected_tools(self):
117117
mock_agent.memory.add_to_memory = Mock()
118118
mock_agent.llm = Mock()
119119
mock_agent.tool_manager = Mock()
120-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
120+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
121121

122122
# Mock the LLM response
123123
mock_response = Mock()
@@ -137,7 +137,9 @@ def test_plan_with_selected_tools(self):
137137
result = reasoning.plan(obs=obs, ttl=3, selected_tools=selected_tools)
138138

139139
assert result == mock_plan
140-
mock_agent.tool_manager.get_all_tools_schema.assert_called_with(selected_tools)
140+
mock_agent.tool_manager.get_annotated_tools_schema.assert_called_with(
141+
agent=mock_agent, selected_tools=selected_tools
142+
)
141143
reasoning.execute_tool_call.assert_called_once_with(
142144
"test_action",
143145
selected_tools=selected_tools,
@@ -171,7 +173,7 @@ def test_aplan_async_version(self):
171173
mock_agent.memory.aadd_to_memory = AsyncMock()
172174
mock_agent.llm = Mock()
173175
mock_agent.tool_manager = Mock()
174-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
176+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
175177

176178
# Mock the async LLM response
177179
mock_response = Mock()

‎tests/test_reasoning/test_reasoning.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ def test_execute_tool_call_generates_plan(self):
5656
mock_agent.llm.generate.return_value = mock_llm_response
5757

5858
# Mock the Tool Manager
59-
mock_agent.tool_manager.get_all_tools_schema.return_value = [
59+
mock_agent.tool_manager.get_feasible_tools_schema.return_value = [
6060
{"schema": "example"}
6161
]
6262

@@ -81,8 +81,8 @@ def plan(self, prompt=None, obs=None, ttl=1, selected_tools=None):
8181
tool_choice="required",
8282
)
8383
# Assert that the tool manager was asked for the correct schema
84-
mock_agent.tool_manager.get_all_tools_schema.assert_called_once_with(
85-
selected_tools=["tool1"]
84+
mock_agent.tool_manager.get_feasible_tools_schema.assert_called_once_with(
85+
agent=mock_agent, selected_tools=["tool1"]
8686
)
8787
# Assert that the output is a correctly formed Plan object
8888
assert isinstance(result_plan, Plan)
@@ -98,7 +98,7 @@ def test_execute_tool_call_propagates_ttl(self):
9898
mock_llm_response.choices = [Mock()]
9999
mock_llm_response.choices[0].message = "Final LLM message"
100100
mock_agent.llm.generate.return_value = mock_llm_response
101-
mock_agent.tool_manager.get_all_tools_schema.return_value = []
101+
mock_agent.tool_manager.get_feasible_tools_schema.return_value = []
102102

103103
class ConcreteReasoning(Reasoning):
104104
def plan(self, prompt=None, obs=None, ttl=1, selected_tools=None):

‎tests/test_reasoning/test_rewoo.py‎

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ def test_plan_new_plan_generation(self):
8585
mock_agent.memory.add_to_memory = Mock()
8686
mock_agent.llm = Mock()
8787
mock_agent.tool_manager = Mock()
88-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
88+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
8989

9090
# Mock the LLM response for planning
9191
mock_plan_response = Mock()
@@ -139,7 +139,7 @@ def test_plan_with_custom_prompt(self):
139139
mock_agent.memory.add_to_memory = Mock()
140140
mock_agent.llm = Mock()
141141
mock_agent.tool_manager = Mock()
142-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
142+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
143143

144144
# Mock the LLM response for planning
145145
mock_plan_response = Mock()
@@ -177,7 +177,7 @@ def test_plan_with_selected_tools(self):
177177
mock_agent.memory.add_to_memory = Mock()
178178
mock_agent.llm = Mock()
179179
mock_agent.tool_manager = Mock()
180-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
180+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
181181

182182
# Mock the LLM response for planning
183183
mock_plan_response = Mock()
@@ -201,7 +201,9 @@ def test_plan_with_selected_tools(self):
201201
result = reasoning.plan(selected_tools=selected_tools)
202202

203203
assert isinstance(result, Plan)
204-
mock_agent.tool_manager.get_all_tools_schema.assert_called_with(selected_tools)
204+
mock_agent.tool_manager.get_annotated_tools_schema.assert_called_with(
205+
agent=mock_agent, selected_tools=selected_tools
206+
)
205207

206208
def test_plan_no_prompt_error(self):
207209
"""Test plan method raises error when no prompt is provided."""
@@ -232,7 +234,7 @@ def test_plan_with_no_tool_calls(self):
232234
mock_agent.memory.add_to_memory = Mock()
233235
mock_agent.llm = Mock()
234236
mock_agent.tool_manager = Mock()
235-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
237+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
236238

237239
# Mock the LLM response for planning
238240
mock_plan_response = Mock()
@@ -295,7 +297,7 @@ def test_aplan_new_plan_generation(self):
295297
mock_agent.memory.add_to_memory = Mock()
296298
mock_agent.llm = Mock()
297299
mock_agent.tool_manager = Mock()
298-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
300+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
299301

300302
mock_plan_response = Mock()
301303
mock_plan_response.choices = [Mock()]
@@ -345,7 +347,7 @@ def test_plan_uses_provided_obs_without_regeneration(self):
345347
mock_agent.memory.add_to_memory = Mock()
346348
mock_agent.llm = Mock()
347349
mock_agent.tool_manager = Mock()
348-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
350+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
349351

350352
mock_plan_response = Mock()
351353
mock_plan_response.choices = [Mock()]
@@ -385,7 +387,7 @@ def test_aplan_uses_provided_obs_without_regeneration(self):
385387
mock_agent.memory.add_to_memory = Mock()
386388
mock_agent.llm = Mock()
387389
mock_agent.tool_manager = Mock()
388-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
390+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
389391

390392
mock_plan_response = Mock()
391393
mock_plan_response.choices = [Mock()]
@@ -429,7 +431,7 @@ def test_aplan_with_selected_tools(self):
429431
mock_agent.memory.add_to_memory = Mock()
430432
mock_agent.llm = Mock()
431433
mock_agent.tool_manager = Mock()
432-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
434+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
433435

434436
mock_plan_response = Mock()
435437
mock_plan_response.choices = [Mock()]
@@ -455,7 +457,9 @@ def test_aplan_with_selected_tools(self):
455457
)
456458

457459
assert isinstance(result, Plan)
458-
mock_agent.tool_manager.get_all_tools_schema.assert_called_with(selected_tools)
460+
mock_agent.tool_manager.get_annotated_tools_schema.assert_called_with(
461+
agent=mock_agent, selected_tools=selected_tools
462+
)
459463

460464
def test_aplan_with_no_tool_calls(self):
461465
"""Test aplan method when execution returns no tool calls."""
@@ -469,7 +473,7 @@ def test_aplan_with_no_tool_calls(self):
469473
mock_agent.memory.add_to_memory = Mock()
470474
mock_agent.llm = Mock()
471475
mock_agent.tool_manager = Mock()
472-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
476+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
473477

474478
mock_plan_response = Mock()
475479
mock_plan_response.choices = [Mock()]
@@ -507,7 +511,7 @@ def test_aplan_uses_step_prompt_when_no_prompt_given(self):
507511
mock_agent.memory.add_to_memory = Mock()
508512
mock_agent.llm = Mock()
509513
mock_agent.tool_manager = Mock()
510-
mock_agent.tool_manager.get_all_tools_schema.return_value = {}
514+
mock_agent.tool_manager.get_annotated_tools_schema.return_value = {}
511515

512516
mock_plan_response = Mock()
513517
mock_plan_response.choices = [Mock()]

0 commit comments

Comments
 (0)