clean ReAct framework

Author: ZFancy

prototype/agingbench/cli/runners.py

@@ -392,8 +392,14 @@ def _run_s2(sut_cfg: dict, scenario_cfg: dict, output_dir: Path,
         from agingbench.generators.s2_generator import S2Generator
         from agingbench.generators.pressure_config import PressureConfig
         gen_n = gen_sessions if gen_sessions > 0 else n_sessions
-        generated_data = S2Generator(seed=sut_cfg.get("seed", 42),
-                                     pressure=_resolve_pressure(sut_cfg, scenario_cfg)).generate(gen_n)
+        generated_data = S2Generator(
+            seed=sut_cfg.get("seed", 42),
+            pressure=_resolve_pressure(sut_cfg, scenario_cfg),
+            dense_accumulator=sut_cfg.get(
+                "dense_accumulator",
+                scenario_cfg.get("dense_accumulator", False),
+            ),
+        ).generate(gen_n)
         n_sessions = gen_n
 
     with TraceLogger(str(trace_path)) as tracer:

prototype/agingbench/core/adapters/react_file_adapter.py

@@ -17,7 +17,7 @@
 from typing import Optional
 
 from ..agent_adapter import AgentAdapter, AgentResponse
-from ..agent import ReferenceAgent, strip_thinking
+from ..agent import ReferenceAgent, _ParseFailure, strip_thinking
 from ..llm import BaseLLM
 from ..tools import ToolSpec, ToolRegistry
 from ..memory.workspace import WorkspaceMemoryPolicy
@@ -116,15 +116,11 @@ def _build_tools(self) -> ToolRegistry:
         return registry
 
     def _write_file(self, args: dict) -> str:
-        # Defensive: handle various argument formats
         path_str = args.get("path") or args.get("filename") or args.get("file") or "notes.txt"
         content = args.get("content") or args.get("text") or args.get("data") or str(args)
-        # Llama / small models sometimes emit numeric or boolean `content`
-        # values (e.g. {"content": 42}); pathlib.write_text requires str.
         if not isinstance(content, str):
             content = str(content)
         path = self.workspace_dir / path_str
-        # Sanitize path (prevent directory traversal)
         try:
             path.resolve().relative_to(self.workspace_dir.resolve())
         except ValueError:
@@ -183,19 +179,33 @@ def send_message(self, message: str) -> AgentResponse:
         total_input = 0
         total_output = 0
 
+        # Same-call cache: identical (tool, args) reuse the result so
+        # repeated retries don't redundantly hit the filesystem and
+        # consume the turn budget.
+        tool_cache: dict[tuple, object] = {}
+
+        def _tool_cache_key(name: str, args: dict) -> tuple:
+            try:
+                return (name, json.dumps(args, sort_keys=True, default=str))
+            except Exception:
+                return (name, repr(args))
+
         # ReAct loop
         for turn in range(self.max_turns):
             response = self.llm.chat_with_usage(messages)
             total_input += response.input_tokens
             total_output += response.output_tokens
-            text = response.text
+            # Parse the visible (stripped) content, not the raw text — otherwise
+            # Actions or Final Answers inside <think> blocks of a thinking
+            # model (Gemma 4, DeepSeek R1) get treated as committed protocol
+            # output. Visible text is also what we add to history.
+            text = strip_thinking(response.text, self.llm)
 
             # Check for Final Answer
             final = ReferenceAgent._parse_final_answer(text)
             if final is not None:
-                # Save to conversation history
                 self._conversation_history.append({"role": "user", "content": message})
-                self._conversation_history.append({"role": "assistant", "content": strip_thinking(final, self.llm)})
+                self._conversation_history.append({"role": "assistant", "content": final})
                 return AgentResponse(
                     text=final,
                     tool_calls=tool_calls,
@@ -209,26 +219,53 @@ def send_message(self, message: str) -> AgentResponse:
             action = ReferenceAgent._parse_action(text)
             if action is not None:
                 tool_name, tool_args = action
+                # Malformed Action Input — surface as an Observation
+                # instead of dispatching the tool with a sentinel that
+                # would crash any args.get(...) inside the tool fn.
+                if isinstance(tool_args, _ParseFailure):
+                    messages.append({"role": "assistant", "content": text})
+                    messages.append({
+                        "role": "user",
+                        "content": f"Observation: ERROR: {tool_args.reason}",
+                    })
+                    continue
                 tool_spec = self._tool_registry.get(tool_name)
                 if tool_spec:
-                    result = tool_spec.call(tool_args)
+                    cache_key = _tool_cache_key(tool_name, tool_args)
+                    if cache_key in tool_cache:
+                        result = tool_cache[cache_key]
+                        obs = (
+                            f"Observation: {result} "
+                            f"[repeated call — same result as before; "
+                            f"try different arguments or commit a Final Answer]"
+                        )
+                    else:
+                        # Catch tool exceptions so a raising tool fn doesn't
+                        # crash the entire run (the file-tool layer can raise
+                        # TypeError on numeric content, OSError on path
+                        # collisions, etc).
+                        try:
+                            result = tool_spec.call(tool_args)
+                        except Exception as e:
+                            result = f"ERROR: {type(e).__name__}: {e}"
+                        tool_cache[cache_key] = result
+                        obs = f"Observation: {result}"
                     tool_calls.append({
                         "tool": tool_name,
                         "input": tool_args,
                         "result": str(result)[:500],
                     })
                     if tool_name == "write_file":
                         files_changed.append(tool_args.get("path", ""))
-                    # Add observation to messages
-                    messages.append({"role": "assistant", "content": strip_thinking(text, self.llm)})
-                    messages.append({"role": "user", "content": f"Observation: {result}"})
+                    messages.append({"role": "assistant", "content": text})
+                    messages.append({"role": "user", "content": obs})
                 else:
-                    messages.append({"role": "assistant", "content": strip_thinking(text, self.llm)})
+                    messages.append({"role": "assistant", "content": text})
                     messages.append({"role": "user", "content": f"Observation: Unknown tool '{tool_name}'. Available: {', '.join(t.name for t in self._tool_registry)}"})
             else:
                 # No action or final answer — treat as final answer
                 self._conversation_history.append({"role": "user", "content": message})
-                self._conversation_history.append({"role": "assistant", "content": strip_thinking(text, self.llm)})
+                self._conversation_history.append({"role": "assistant", "content": text})
                 return AgentResponse(
                     text=text,
                     tool_calls=tool_calls,

prototype/agingbench/core/agent.py

@@ -17,6 +17,17 @@
 _THINK_BLOCK = re.compile(r"<think>.*?(?:</think>|\Z)\s*", re.DOTALL | re.IGNORECASE)
 
 
+class _ParseFailure:
+    """Sentinel returned by ``_parse_action`` when the action name parsed but
+    the JSON args did not. The run loop surfaces ``reason`` as an error
+    Observation so the model can re-emit valid JSON next turn.
+    """
+    __slots__ = ("reason",)
+
+    def __init__(self, reason: str):
+        self.reason = reason
+
+
 def _llm_model_id(llm: BaseLLM) -> str:
     return getattr(llm, "model_id", None) or getattr(llm, "model", "") or ""
 
@@ -127,50 +138,140 @@ def run_session(self, task: str, session_id: int = 0) -> dict:
         ]
         tool_calls = []
 
+        # Local cache: identical (tool, args) calls reuse the result so the
+        # model is not punished with extra LLM turns for re-issuing the
+        # same call. Discarded when run_session returns.
+        tool_cache: dict[tuple, object] = {}
+
+        def _tool_cache_key(name: str, args: dict) -> tuple:
+            try:
+                return (name, json.dumps(args, sort_keys=True, default=str))
+            except Exception:
+                return (name, repr(args))
+
         for turn in range(self.max_turns):
             response = self.llm.chat(messages)
-            messages.append({"role": "assistant", "content": strip_thinking(response, self.llm)})
-
-            final = self._parse_final_answer(response)
+            # Parse the visible (stripped) content, not the raw response —
+            # otherwise Actions or Final Answers inside <think> blocks of a
+            # thinking model (Gemma 4, DeepSeek R1) get treated as committed
+            # protocol output. Visible is also what we add to history.
+            visible = strip_thinking(response, self.llm)
+            messages.append({"role": "assistant", "content": visible})
+
+            final = self._parse_final_answer(visible)
             if final is not None:
                 return {"output": final, "tool_calls": tool_calls, "turns": turn + 1}
 
-            parsed = self._parse_action(response)
+            parsed = self._parse_action(visible)
             if parsed:
                 tool_name, tool_input = parsed
+                if isinstance(tool_input, _ParseFailure):
+                    messages.append({
+                        "role": "user",
+                        "content": f"Observation: ERROR: {tool_input.reason}",
+                    })
+                    continue
                 spec = self.tools.get(tool_name)
                 if spec:
-                    result = spec.call(tool_input)
+                    cache_key = _tool_cache_key(tool_name, tool_input)
+                    if cache_key in tool_cache:
+                        result = tool_cache[cache_key]
+                        cached_hint = (
+                            " [repeated call — same result as before; "
+                            "try a different argument or commit a Final Answer]"
+                        )
+                        obs_text = f"Observation: {result}{cached_hint}"
+                    else:
+                        try:
+                            result = spec.call(tool_input)
+                        except Exception as e:
+                            result = f"ERROR: {type(e).__name__}: {e}"
+                        tool_cache[cache_key] = result
+                        obs_text = f"Observation: {result}"
                     tool_calls.append({"tool": tool_name, "version": spec.version,
                                        "input": tool_input, "result": result})
-                    messages.append({"role": "user", "content": f"Observation: {result}"})
+                    messages.append({"role": "user", "content": obs_text})
                 else:
                     messages.append(
                         {"role": "user", "content": f"Error: unknown tool '{tool_name}'"}
                     )
             else:
-                # No action or final answer — nudge the model to continue.
-                # Ensures messages end with a user turn (required by some providers).
+                # No Action and no Final Answer; ensure messages end with a
+                # user turn (required by some providers) and nudge.
                 messages.append(
                     {"role": "user", "content": "Continue. Use the Action/Final Answer format."}
                 )
 
-        return {"output": response, "tool_calls": tool_calls, "turns": self.max_turns}
+        # max_turns exhausted without a Final Answer; return the stripped
+        # last response so downstream scoring doesn't see <think> blocks.
+        return {"output": visible, "tool_calls": tool_calls, "turns": self.max_turns}
 
     # ---------------------------------------------------------------- helpers
 
     @staticmethod
-    def _parse_action(text: str) -> Optional[tuple[str, dict]]:
-        action = re.search(r"Action:\s*(\w+)", text)
-        inp = re.search(r"Action Input:\s*(\{.*?\})", text, re.DOTALL)
-        if action and inp:
-            try:
-                return action.group(1), json.loads(inp.group(1))
-            except json.JSONDecodeError:
-                return action.group(1), {"raw": inp.group(1)}
+    def _extract_balanced_json(text: str, start: int) -> Optional[str]:
+        """Find a balanced ``{...}`` JSON object at or after ``start``.
+
+        Walks the string respecting string literals and escapes so a ``}``
+        inside a quoted value does not close the object early. Returns
+        None if no balanced object is found.
+        """
+        depth = 0
+        in_string = False
+        escape = False
+        open_idx = None
+        for i in range(start, len(text)):
+            ch = text[i]
+            if in_string:
+                if escape:
+                    escape = False
+                elif ch == "\\":
+                    escape = True
+                elif ch == '"':
+                    in_string = False
+                continue
+            if ch == '"':
+                in_string = True
+                continue
+            if ch == "{":
+                if open_idx is None:
+                    open_idx = i
+                depth += 1
+            elif ch == "}":
+                depth -= 1
+                if depth == 0 and open_idx is not None:
+                    return text[open_idx : i + 1]
+                if depth < 0:
+                    return None
         return None
 
+    @staticmethod
+    def _parse_action(text: str) -> Optional[tuple[str, dict]]:
+        action = re.search(r"Action:\s*([\w-]+)", text)
+        if not action:
+            return None
+        inp_marker = re.search(r"Action Input:\s*", text)
+        if inp_marker is None:
+            return None
+        balanced = ReferenceAgent._extract_balanced_json(text, inp_marker.end())
+        if balanced is None:
+            return action.group(1), _ParseFailure(
+                "Action Input did not contain a parseable JSON object."
+            )
+        try:
+            return action.group(1), json.loads(balanced)
+        except json.JSONDecodeError as e:
+            return action.group(1), _ParseFailure(
+                f"Action Input JSON was malformed: {e}. Retry with valid JSON."
+            )
+
     @staticmethod
     def _parse_final_answer(text: str) -> Optional[str]:
-        m = re.search(r"Final Answer:\s*(.*)", text, re.DOTALL)
+        # Stop at the next ReAct marker so commit-then-retry scaffolding
+        # does not leak into the scored output.
+        m = re.search(
+            r"Final Answer:\s*(.*?)(?=\n\s*(?:Thought|Action|Observation)\s*:|\Z)",
+            text,
+            re.DOTALL,
+        )
         return m.group(1).strip() if m else None