patcher: handle recursion errors and better reflection (#752)

* patcher: handle recursion errors and better reflection

* patcher: fix patch-strategy reflection state + add draw graph py

* fix lint

* add img

patcher/README.md

@@ -1 +1,3 @@
 # Buttercup Patcher
+
+![Patcher State Machine](patcher_state_machine.png)

patcher/src/buttercup/patcher/agents/common.py

@@ -155,6 +155,7 @@ class ExecutionInfo(BaseModel):
     """Execution info"""
 
     root_cause_analysis_tries: int = Field(default=0)
+    patch_strategy_tries: int = Field(default=0)
     tests_tries: int = Field(default=0)
     reflection_decision: PatcherAgentName | None = None
     reflection_guidance: str | None = None

patcher/src/buttercup/patcher/agents/config.py

@@ -15,6 +15,7 @@ class PatcherConfig(BaseModel):
     max_last_failure_retries: int = Field(default=3)
     max_minutes_run_povs: int = Field(default=30)
     max_root_cause_analysis_retries: int = Field(default=3)
+    max_patch_strategy_retries: int = Field(default=3)
     max_tests_retries: int = Field(default=5)
     ctx_retriever_recursion_limit: int = Field(default=80)
     n_initial_stackframes: int = Field(default=4)

patcher/src/buttercup/patcher/agents/draw_png.py

@@ -0,0 +1,88 @@
+"""Module to visualize the patcher agent state machine."""
+
+from pathlib import Path
+from typing import Optional
+
+import graphviz
+
+
+def draw_state_machine(output_path: Optional[Path] = None) -> None:
+    """Draw the patcher agent state machine and save it as a PNG file.
+
+    Args:
+        output_path: Optional path where to save the PNG file. If None, saves in current directory.
+    """
+    # Create a new directed graph
+    dot = graphviz.Digraph(comment="Patcher Agent State Machine")
+    dot.attr(rankdir="LR")  # Left to right layout
+
+    # Add nodes for each state
+    states = [
+        "START",
+        "INPUT_PROCESSING",
+        "FIND_TESTS",
+        "INITIAL_CODE_SNIPPET_REQUESTS",
+        "ROOT_CAUSE_ANALYSIS",
+        "PATCH_STRATEGY",
+        "CREATE_PATCH",
+        "BUILD_PATCH",
+        "RUN_POV",
+        "RUN_TESTS",
+        "REFLECTION",
+        "CONTEXT_RETRIEVER",
+        "END",
+    ]
+
+    # Add nodes with styling
+    for state in states:
+        if state == "START":
+            dot.node(state, state, shape="oval", style="filled", fillcolor="green")
+        elif state == "END":
+            dot.node(state, state, shape="oval", style="filled", fillcolor="red")
+        else:
+            dot.node(state, state, shape="box", style="rounded,filled", fillcolor="lightblue")
+
+    # Add edges to show the flow
+    # Start and End transitions
+    dot.edge("START", "INPUT_PROCESSING", color="blue", penwidth="2.0")
+    dot.edge("RUN_TESTS", "END", color="blue", penwidth="2.0")
+    dot.edge("REFLECTION", "END", color="gray")
+
+    # Main flow (colored in blue)
+    dot.edge("INPUT_PROCESSING", "FIND_TESTS", color="blue", penwidth="2.0")
+    dot.edge("INPUT_PROCESSING", "INITIAL_CODE_SNIPPET_REQUESTS", color="blue", penwidth="2.0")
+    dot.edge("FIND_TESTS", "ROOT_CAUSE_ANALYSIS", color="blue", penwidth="2.0")
+    dot.edge("INITIAL_CODE_SNIPPET_REQUESTS", "ROOT_CAUSE_ANALYSIS", color="blue", penwidth="2.0")
+    dot.edge("ROOT_CAUSE_ANALYSIS", "PATCH_STRATEGY", color="blue", penwidth="2.0")
+    dot.edge("PATCH_STRATEGY", "CREATE_PATCH", color="blue", penwidth="2.0")
+    dot.edge("CREATE_PATCH", "BUILD_PATCH", color="blue", penwidth="2.0")
+    dot.edge("BUILD_PATCH", "RUN_POV", color="blue", penwidth="2.0")
+    dot.edge("RUN_POV", "RUN_TESTS", color="blue", penwidth="2.0")
+
+    # Add transitions to REFLECTION (in gray)
+    dot.edge("ROOT_CAUSE_ANALYSIS", "REFLECTION", color="gray")
+    dot.edge("PATCH_STRATEGY", "REFLECTION", color="gray")
+    dot.edge("CREATE_PATCH", "REFLECTION", color="gray")
+    dot.edge("BUILD_PATCH", "REFLECTION", color="gray")
+    dot.edge("RUN_POV", "REFLECTION", color="gray")
+    dot.edge("RUN_TESTS", "REFLECTION", color="gray")
+
+    # Add transitions from REFLECTION (in gray)
+    dot.edge("REFLECTION", "ROOT_CAUSE_ANALYSIS", color="gray")
+    dot.edge("REFLECTION", "PATCH_STRATEGY", color="gray")
+    dot.edge("REFLECTION", "CREATE_PATCH", color="gray")
+    dot.edge("REFLECTION", "CONTEXT_RETRIEVER", color="gray")
+
+    # Add transitions from CONTEXT_RETRIEVER (in gray)
+    dot.edge("CONTEXT_RETRIEVER", "REFLECTION", color="gray")
+    dot.edge("CONTEXT_RETRIEVER", "PATCH_STRATEGY", color="gray")
+    dot.edge("CONTEXT_RETRIEVER", "ROOT_CAUSE_ANALYSIS", color="gray")
+
+    # Save the graph
+    if output_path is None:
+        output_path = Path("patcher_state_machine")
+    dot.render(str(output_path), format="png", cleanup=True)
+
+
+if __name__ == "__main__":
+    draw_state_machine()

patcher/src/buttercup/patcher/agents/reflection.py

@@ -24,6 +24,7 @@ from buttercup.patcher.agents.common import (
     PatchAttempt,
     CodeSnippetRequest,
     PatchAnalysis,
+    PatchStrategy,
 )
 from buttercup.patcher.agents.config import PatcherConfig
 from buttercup.common.llm import ButtercupLLM, create_default_llm
@@ -680,7 +681,16 @@ class ReflectionAgent(PatcherAgentBase):
                 state.context.task_id,
                 state.context.submission_index,
             )
-            return Command(goto=PatcherAgentName.PATCH_STRATEGY.value)
+            root_cause = state.root_cause
+            if root_cause is None:
+                root_cause = "No root cause found, figure out a root cause"
+
+            return Command(
+                update={
+                    "root_cause": root_cause,
+                },
+                goto=PatcherAgentName.PATCH_STRATEGY.value,
+            )
 
         execution_info = state.execution_info
         execution_info.root_cause_analysis_tries += 1
@@ -691,7 +701,35 @@ class ReflectionAgent(PatcherAgentBase):
             update={
                 "execution_info": execution_info,
             },
-            goto=PatcherAgentName.INPUT_PROCESSING.value,
+            goto=PatcherAgentName.ROOT_CAUSE_ANALYSIS.value,
+        )
+
+    def _patch_strategy_failed(self, state: PatcherAgentState, configuration: PatcherConfig) -> Command:
+        """Patch strategy failed, reflect on the failure."""
+        if state.execution_info.patch_strategy_tries >= configuration.max_patch_strategy_retries:
+            logger.warning(
+                "[%s / %s] Reached max patch strategy failures, just move forward with what we have",
+                state.context.task_id,
+                state.context.submission_index,
+            )
+            strategy = state.patch_strategy
+            if strategy is None:
+                strategy = PatchStrategy(full="Figure out a patch strategy", summary="Figure out a patch strategy")
+
+            return Command(
+                update={
+                    "patch_strategy": strategy,
+                },
+                goto=PatcherAgentName.CREATE_PATCH.value,
+            )
+
+        execution_info = state.execution_info
+        execution_info.patch_strategy_tries += 1
+        return Command(
+            update={
+                "execution_info": execution_info,
+            },
+            goto=PatcherAgentName.PATCH_STRATEGY.value,
         )
 
     def reflect_on_patch(self, state: PatcherAgentState, config: RunnableConfig) -> Command:
@@ -738,8 +776,7 @@ class ReflectionAgent(PatcherAgentBase):
                 goto=PatcherAgentName.CONTEXT_RETRIEVER.value,
             )
 
-        current_patch_attempt = state.get_last_patch_attempt()
-        if not current_patch_attempt or state.execution_info.prev_node == PatcherAgentName.ROOT_CAUSE_ANALYSIS:
+        if state.execution_info.prev_node == PatcherAgentName.ROOT_CAUSE_ANALYSIS:
             logger.warning(
                 "[%s / %s] No patch attempt found, the root cause analysis is probably wrong",
                 state.context.task_id,
@@ -747,4 +784,21 @@ class ReflectionAgent(PatcherAgentBase):
             )
             return self._root_cause_analysis_failed(state, configuration)
 
+        if state.execution_info.prev_node == PatcherAgentName.PATCH_STRATEGY:
+            logger.warning(
+                "[%s / %s] Patch strategy failed, reflecting on it",
+                state.context.task_id,
+                state.context.submission_index,
+            )
+            return self._patch_strategy_failed(state, configuration)
+
+        current_patch_attempt = state.get_last_patch_attempt()
+        if not current_patch_attempt:
+            logger.error(
+                "[%s / %s] No patch attempt found, this should never happen, going back to input processing",
+                state.context.task_id,
+                state.context.submission_index,
+            )
+            return Command(goto=PatcherAgentName.INPUT_PROCESSING.value)
+
         return self._analyze_failure(state, configuration, current_patch_attempt)

patcher/src/buttercup/patcher/agents/rootcause.py

@@ -2,6 +2,7 @@
 
 import logging
 import re
+import langgraph.errors
 from typing import Annotated, Literal
 from langgraph.prebuilt import InjectedState
 from langchain_core.prompts import MessagesPlaceholder
@@ -225,6 +226,18 @@ class RootCauseAgent(PatcherAgentBase):
 
         try:
             root_cause_dict = self.root_cause_chain.invoke(state)
+        except langgraph.errors.GraphRecursionError:
+            logger.error(
+                "Reached recursion limit for root cause analysis in Challenge Task %s/%s",
+                state.context.task_id,
+                self.challenge.name,
+            )
+            return Command(
+                update={
+                    "execution_info": execution_info,
+                },
+                goto=PatcherAgentName.REFLECTION.value,
+            )
         except Exception as e:
             logger.exception("Error parsing root cause: %s", e)
             return Command(

patcher/src/buttercup/patcher/agents/swe.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import difflib
 import uuid
+import langgraph.errors
 from langchain_core.messages import BaseMessage
 from langgraph.prebuilt import InjectedState
 from langchain_core.prompts import MessagesPlaceholder
@@ -700,12 +701,26 @@ class SWEAgent(PatcherAgentBase):
         configurable = {
             "thread_id": str(uuid.uuid4()),
         }
-        strategy_state_dict = self.patch_strategy_chain.invoke(
-            state,
-            config=RunnableConfig(
-                configurable=configurable,
-            ),
-        )
+        try:
+            strategy_state_dict = self.patch_strategy_chain.invoke(
+                state,
+                config=RunnableConfig(
+                    configurable=configurable,
+                ),
+            )
+        except langgraph.errors.GraphRecursionError:
+            logger.error(
+                "Reached recursion limit for patch strategy in Challenge Task %s/%s",
+                state.context.task_id,
+                self.challenge.name,
+            )
+            return Command(
+                update={
+                    "execution_info": execution_info,
+                },
+                goto=PatcherAgentName.REFLECTION.value,
+            )
+
         try:
             strategy_state = PatcherAgentState.model_validate(strategy_state_dict)
         except ValidationError as e:
@@ -729,13 +744,13 @@ class SWEAgent(PatcherAgentBase):
         patch_strategy_str = str(strategy_state.messages[-1].content)
         if "<request_information>" in patch_strategy_str:
             new_code_snippet_requests = self._parse_code_snippet_requests(patch_strategy_str)
+            execution_info.code_snippet_requests = new_code_snippet_requests
             logger.info(
                 "[%s / %s] Requesting additional information", state.context.task_id, state.context.submission_index
             )
             return Command(
                 update={
                     "execution_info": execution_info,
-                    "code_snippet_requests": new_code_snippet_requests,
                 },
                 goto=PatcherAgentName.REFLECTION.value,
             )