Stop samples that loop unproductively

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

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 = [

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