From 1e250e9ec0b02378a937bf4ccc744cf0726754e7 Mon Sep 17 00:00:00 2001 From: Michael Kouremetis Date: Wed, 17 Dec 2025 12:11:43 -0500 Subject: [PATCH 1/6] tool call tracking --- dreadnode/agent/agent.py | 16 ++++++++++++++-- dreadnode/agent/error.py | 7 +++++++ dreadnode/agent/result.py | 2 +- 3 files changed, 22 insertions(+), 3 deletions(-) diff --git a/dreadnode/agent/agent.py b/dreadnode/agent/agent.py index a680b877..5420c808 100644 --- a/dreadnode/agent/agent.py +++ b/dreadnode/agent/agent.py @@ -14,7 +14,7 @@ from rigging.message import inject_system_content from ulid import ULID # can't access via rg -from dreadnode.agent.error import MaxStepsError +from dreadnode.agent.error import MaxStepsError, MaxToolCallsError from dreadnode.agent.events import ( AgentEnd, AgentError, @@ -89,7 +89,9 @@ class Agent(Model): ) """The agent's core instructions.""" max_steps: int = Config(default=10) - """The maximum number of steps (generation + tool calls).""" + """The maximum number of steps (generations).""" + max_tool_calls: int = Config(default=-1) + """The maximum number of tool calls. Defaults to infinite.""" caching: rg.caching.CacheMode | None = Config(default=None, repr=False) """How to handle cache_control entries on inference messages.""" @@ -566,6 +568,7 @@ async def _process_tool_call( # Core step loop step = 0 + tool_calls = 0 error: Exception | str | None = None while step < self.max_steps: @@ -662,6 +665,7 @@ async def _process_tool_call( messages.append(event.message) if stopped_by_tool_call is None and event.stop: stopped_by_tool_call = event.tool_call + tool_calls += 1 yield event if stopped_by_tool_call: @@ -670,6 +674,11 @@ async def _process_tool_call( f"{stopped_by_tool_call.id} requested to stop the agent." ) + if self.max_tool_calls != -1 and tool_calls >= self.max_tool_calls: + raise Finish( + reason="Reached maximum allowed tool calls." + ) + # Check for stop conditions (again) if any(cond(events) for cond in stop_conditions): @@ -690,6 +699,9 @@ async def _process_tool_call( if step >= self.max_steps: error = MaxStepsError(max_steps=self.max_steps) stop_reason = "max_steps_reached" + elif tool_calls >= self.max_tool_calls: + error = MaxToolCallsError(max_tool_calls=self.max_tool_calls) + stop_reason = "max_tool_calls_reached" elif error is not None: stop_reason = "error" elif events and isinstance(events[-1], AgentStalled): diff --git a/dreadnode/agent/error.py b/dreadnode/agent/error.py index feda14c7..f0ca4483 100644 --- a/dreadnode/agent/error.py +++ b/dreadnode/agent/error.py @@ -4,3 +4,10 @@ class MaxStepsError(Exception): def __init__(self, max_steps: int): super().__init__(f"Maximum steps reached ({max_steps}).") self.max_steps = max_steps + +class MaxToolCallsError(Exception): + """Raise from a hook to stop the agent's run due to reaching the maximum number of tool calls.""" + + def __init__(self, max_tool_calls: int): + super().__init__(f"Maximum tool calls reached ({max_tool_calls}).") + self.max_tool_calls = max_tool_calls diff --git a/dreadnode/agent/result.py b/dreadnode/agent/result.py index f8d952c1..9edcc7a7 100644 --- a/dreadnode/agent/result.py +++ b/dreadnode/agent/result.py @@ -8,7 +8,7 @@ if t.TYPE_CHECKING: from dreadnode.agent.agent import Agent -AgentStopReason = t.Literal["finished", "max_steps_reached", "error", "stalled"] +AgentStopReason = t.Literal["finished", "max_steps_reached", "max_tool_calls_reached", "error", "stalled"] @dataclass(config=ConfigDict(arbitrary_types_allowed=True)) From 8ca20cb9394cbc4300102e688f893abc444eeaef Mon Sep 17 00:00:00 2001 From: Michael Kouremetis Date: Wed, 17 Dec 2025 12:16:52 -0500 Subject: [PATCH 2/6] tool call tracking --- dreadnode/agent/agent.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/dreadnode/agent/agent.py b/dreadnode/agent/agent.py index 5420c808..9d52e3d3 100644 --- a/dreadnode/agent/agent.py +++ b/dreadnode/agent/agent.py @@ -675,9 +675,7 @@ async def _process_tool_call( ) if self.max_tool_calls != -1 and tool_calls >= self.max_tool_calls: - raise Finish( - reason="Reached maximum allowed tool calls." - ) + raise Finish("Reached maximum allowed tool calls.") # Check for stop conditions (again) @@ -699,7 +697,7 @@ async def _process_tool_call( if step >= self.max_steps: error = MaxStepsError(max_steps=self.max_steps) stop_reason = "max_steps_reached" - elif tool_calls >= self.max_tool_calls: + elif self.max_tool_calls != -1 and tool_calls >= self.max_tool_calls: error = MaxToolCallsError(max_tool_calls=self.max_tool_calls) stop_reason = "max_tool_calls_reached" elif error is not None: From 8a7a312b33179316590f0d97ef9e083341f572f7 Mon Sep 17 00:00:00 2001 From: Michael Kouremetis Date: Wed, 17 Dec 2025 13:16:46 -0500 Subject: [PATCH 3/6] better placement for tool counting, will now fail on the exact tool call that is past max --- dreadnode/agent/agent.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/dreadnode/agent/agent.py b/dreadnode/agent/agent.py index 9d52e3d3..3387376d 100644 --- a/dreadnode/agent/agent.py +++ b/dreadnode/agent/agent.py @@ -490,10 +490,17 @@ async def _dispatch(event: AgentEvent) -> t.AsyncIterator[AgentEvent]: # noqa: raise winning_reaction # Tool calling - + tool_calls = 0 + async def _process_tool_call( tool_call: "rg.tools.ToolCall", ) -> t.AsyncGenerator[AgentEvent, None]: + + nonlocal tool_calls + + if self.max_tool_calls != -1 and tool_calls >= self.max_tool_calls: + raise Finish("Reached maximum allowed tool calls.") + async for event in _dispatch( ToolStart( session_id=session_id, @@ -515,6 +522,7 @@ async def _process_tool_call( tool = next((t for t in self.all_tools if t.name == tool_call.name), None) if tool is not None: + tool_calls += 1 try: message, stop = await tool.handle_tool_call(tool_call) except Reaction: @@ -568,7 +576,6 @@ async def _process_tool_call( # Core step loop step = 0 - tool_calls = 0 error: Exception | str | None = None while step < self.max_steps: @@ -665,7 +672,6 @@ async def _process_tool_call( messages.append(event.message) if stopped_by_tool_call is None and event.stop: stopped_by_tool_call = event.tool_call - tool_calls += 1 yield event if stopped_by_tool_call: @@ -674,9 +680,6 @@ async def _process_tool_call( f"{stopped_by_tool_call.id} requested to stop the agent." ) - if self.max_tool_calls != -1 and tool_calls >= self.max_tool_calls: - raise Finish("Reached maximum allowed tool calls.") - # Check for stop conditions (again) if any(cond(events) for cond in stop_conditions): From cdc4482171c7814891505aae293061e0f4b07c83 Mon Sep 17 00:00:00 2001 From: Michael Kouremetis Date: Wed, 17 Dec 2025 13:27:27 -0500 Subject: [PATCH 4/6] added test case --- tests/test_agent.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/test_agent.py b/tests/test_agent.py index 5ffe3ded..dcb41ff7 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -8,7 +8,7 @@ from rigging.generator.base import GeneratedMessage from dreadnode.agent.agent import Agent, TaskAgent -from dreadnode.agent.error import MaxStepsError +from dreadnode.agent.error import MaxStepsError, MaxToolCallsError from dreadnode.agent.events import AgentEnd, AgentEvent, AgentStalled, Reacted, ToolStart from dreadnode.agent.hooks.base import retry_with_feedback from dreadnode.agent.reactions import RetryWithFeedback @@ -298,6 +298,31 @@ async def test_run_stops_on_max_steps(mock_generator: MockGenerator, simple_tool assert result.steps == 1 +@pytest.mark.asyncio +async def test_run_stops_on_max_tool_calls( + mock_generator: MockGenerator, simple_tool: AnyTool +) -> None: + """Ensure the agent run terminates with a MaxToolCallsError when exceeding max_tool_calls.""" + # The agent will just keep calling the tool. + mock_generator._responses = [ + MockGenerator.tool_response("get_weather", {"city": "A"}), + MockGenerator.tool_response("get_weather", {"city": "B"}), + MockGenerator.tool_response("get_weather", {"city": "C"}), + ] + + agent = Agent( + name="MaxToolCallsAgent", + model=mock_generator, + tools=[simple_tool], + max_tool_calls=2, + ) + result = await agent.run("...") + + assert result.failed + assert result.stop_reason == "max_tool_calls_reached" + assert isinstance(result.error, MaxToolCallsError) + + @pytest.mark.asyncio async def test_run_stops_on_stop_condition( mock_generator: MockGenerator, simple_tool: AnyTool From 3cb5b4a7af719e4a287e20d942431fe3d3e8bbd2 Mon Sep 17 00:00:00 2001 From: Michael Kouremetis Date: Wed, 17 Dec 2025 14:26:23 -0500 Subject: [PATCH 5/6] linting --- dreadnode/agent/agent.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dreadnode/agent/agent.py b/dreadnode/agent/agent.py index 3387376d..512e2841 100644 --- a/dreadnode/agent/agent.py +++ b/dreadnode/agent/agent.py @@ -491,13 +491,13 @@ async def _dispatch(event: AgentEvent) -> t.AsyncIterator[AgentEvent]: # noqa: # Tool calling tool_calls = 0 - + async def _process_tool_call( tool_call: "rg.tools.ToolCall", ) -> t.AsyncGenerator[AgentEvent, None]: - + nonlocal tool_calls - + if self.max_tool_calls != -1 and tool_calls >= self.max_tool_calls: raise Finish("Reached maximum allowed tool calls.") From 3ee2a86fc6879055cc8a94722b787fe414b5b493 Mon Sep 17 00:00:00 2001 From: Michael Kouremetis Date: Wed, 17 Dec 2025 15:55:46 -0500 Subject: [PATCH 6/6] linting --- dreadnode/agent/agent.py | 1 - dreadnode/agent/error.py | 1 + dreadnode/agent/result.py | 4 +++- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/dreadnode/agent/agent.py b/dreadnode/agent/agent.py index 512e2841..79fbb965 100644 --- a/dreadnode/agent/agent.py +++ b/dreadnode/agent/agent.py @@ -495,7 +495,6 @@ async def _dispatch(event: AgentEvent) -> t.AsyncIterator[AgentEvent]: # noqa: async def _process_tool_call( tool_call: "rg.tools.ToolCall", ) -> t.AsyncGenerator[AgentEvent, None]: - nonlocal tool_calls if self.max_tool_calls != -1 and tool_calls >= self.max_tool_calls: diff --git a/dreadnode/agent/error.py b/dreadnode/agent/error.py index f0ca4483..f132cce0 100644 --- a/dreadnode/agent/error.py +++ b/dreadnode/agent/error.py @@ -5,6 +5,7 @@ def __init__(self, max_steps: int): super().__init__(f"Maximum steps reached ({max_steps}).") self.max_steps = max_steps + class MaxToolCallsError(Exception): """Raise from a hook to stop the agent's run due to reaching the maximum number of tool calls.""" diff --git a/dreadnode/agent/result.py b/dreadnode/agent/result.py index 9edcc7a7..9a5fe01d 100644 --- a/dreadnode/agent/result.py +++ b/dreadnode/agent/result.py @@ -8,7 +8,9 @@ if t.TYPE_CHECKING: from dreadnode.agent.agent import Agent -AgentStopReason = t.Literal["finished", "max_steps_reached", "max_tool_calls_reached", "error", "stalled"] +AgentStopReason = t.Literal[ + "finished", "max_steps_reached", "max_tool_calls_reached", "error", "stalled" +] @dataclass(config=ConfigDict(arbitrary_types_allowed=True))