diff --git a/apn/limits.py b/apn/limits.py new file mode 100644 index 0000000..17511cf --- /dev/null +++ b/apn/limits.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from inspect_ai.agent import AgentContinue, AgentState +from inspect_boltons.limits import NoToolCallLimit, RepeatedTextLimit + +UNPRODUCTIVE_LOOP_LIMITS = ( + NoToolCallLimit(turns=100, unproductive_tools=["resources"]), + RepeatedTextLimit(turns=30), +) + + +def continue_unless_looping(message: str) -> AgentContinue: + async def on_continue(state: AgentState) -> bool | str: + for limit in UNPRODUCTIVE_LOOP_LIMITS: + limit.check(state) + if state.output.message.tool_calls: + return True + return message + + return on_continue diff --git a/apn/solver.py b/apn/solver.py index 0952257..d7783a2 100644 --- a/apn/solver.py +++ b/apn/solver.py @@ -6,6 +6,7 @@ from inspect_ai.agent import ( Agent, AgentAttempts, + AgentContinue, AgentState, AgentSubmit, as_solver, @@ -26,6 +27,7 @@ from apn.checker import Claim from apn.layout import ENTRY_PATH +from apn.limits import continue_unless_looping from apn.prompts import user_prompt from apn.scorer import CLAIM_STORE_KEY from apn.tools import bash @@ -119,7 +121,7 @@ def build_agent( tools: Sequence[Tool | ToolDef | ToolSource], attempts: AgentAttempts, submit: AgentSubmit, - on_continue: str, + on_continue: AgentContinue, compaction: CompactionStrategy, ) -> Agent: """Construct the configured agent loop. @@ -210,7 +212,7 @@ async def solve(state: TaskState, generate: Generate) -> TaskState: ), # The default continue message is very generic ("proceed to the next step"), this one # might be better at avoiding doom loops. - on_continue="Continue working on the problem.", + on_continue=continue_unless_looping("Continue working on the problem."), compaction=CompactionSummary(threshold=300_000), ) state.messages = [ diff --git a/tests/test_agent.py b/tests/test_agent.py index 2acd4fd..48ed3b0 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -19,6 +19,7 @@ from inspect_ai.scorer import INCORRECT, Score import apn.solver as solver_module +from apn.limits import continue_unless_looping from apn.solver import ( INCORRECT_MESSAGE, RESOURCE_INCORRECT_MESSAGE, @@ -92,7 +93,7 @@ def fake(**kwargs: Any) -> object: tools=[], attempts=AgentAttempts(attempts=1), submit=AgentSubmit(), - on_continue="go", + on_continue=continue_unless_looping("go"), compaction=CompactionSummary(), ) return seen