Multi-Patch Selection Feature Complete

README.md

@@ -141,7 +141,7 @@ Set the following variables in your `.env` file:
 * `PROMETHEUS_GITHUB_ACCESS_TOKEN`
 * `PROMETHEUS_KNOWLEDGE_GRAPH_MAX_AST_DEPTH`
 * `PROMETHEUS_NEO4J_BATCH_SIZE`
-* `PROMETHEUS_POSTGRES_URI`
+* `PROMETHEUS_POSTGRES_URL`
 
 ---
 
@@ -164,7 +164,7 @@ Set the following variables in your `.env` file:
 2. Run tests:
 
    ```bash
-   coverage run --source=prometheus -m pytest -v -s -m "not git"
+   coverage run --source=prometheus -m pytest -v -s
    ```
 
 3. Generate coverage report:

prometheus/git/git_repository.py

@@ -120,9 +120,28 @@ def get_diff(self, excluded_files: Optional[Sequence[str]] = None) -> str:
     def get_working_directory(self) -> Path:
         return Path(self.repo.working_dir).absolute()
 
-    def reset_repository(self):
-        self.repo.git.reset("--hard")
-        self.repo.git.clean("-fd")
+    def reset_repository(self, excluded_files: Optional[Sequence[str]] = None):
+        """
+        Reset all modified and untracked files in the repository,
+        excluding specified files.
+        """
+        # Normalize excluded_files to a set of relative strings
+        if excluded_files is None:
+            excluded_files = set()
+        elif isinstance(excluded_files, (str, Path)):
+            excluded_files = set(excluded_files)
+
+        # 1. Reset modified tracked files
+        modified_files = self.repo.git.diff("--name-only").splitlines()
+        reset_files = [f for f in modified_files if f not in excluded_files]
+        if reset_files:
+            self.repo.git.restore("--staged", "--worktree", "--", *reset_files)
+
+        # 2. Clean untracked files
+        untracked_files = self.repo.git.ls_files("--others", "--exclude-standard").splitlines()
+        clean_files = [f for f in untracked_files if f not in excluded_files]
+        if clean_files:
+            self.repo.git.clean("-fd", "--", *clean_files)
 
     def remove_repository(self):
         if self.repo is not None:
@@ -139,6 +158,7 @@ def create_and_push_branch(self, branch_name: str, commit_message: str, patch: s
         Args:
             branch_name: Name of the new branch to create.
             commit_message: Message for the commit.
+            patch: The patch content to apply to the new branch.
         """
         with tempfile.NamedTemporaryFile(mode="w", suffix=".patch") as tmp_file:
             tmp_file.write(patch)

prometheus/lang_graph/nodes/final_patch_selection_node.py

@@ -119,13 +119,14 @@ class FinalPatchSelectionNode:
 {patches}
 """
 
-    def __init__(self, model: BaseChatModel, max_retries: int = 2):
+    def __init__(self, model: BaseChatModel, final_patch_name: str, max_retries: int = 2):
         self.max_retries = max_retries
         prompt = ChatPromptTemplate.from_messages(
             [("system", self.SYS_PROMPT), ("human", "{human_prompt}")]
         )
         structured_llm = model.with_structured_output(FinalPatchSelectionStructuredOutput)
         self.model = prompt | structured_llm
+        self.final_patch_name = final_patch_name
         self._logger = logging.getLogger("prometheus.lang_graph.nodes.final_patch_selection_node")
 
     def format_human_message(self, state: Dict):
@@ -149,10 +150,10 @@ def __call__(self, state: Dict):
             response = self.model.invoke({"human_prompt": human_prompt})
             self._logger.info(f"FinalPatchSelectionNode response at {try_index} try:\n{response}")
 
-            if response.patch_index >= 0 and response.patch_index < len(state["edit_patches"]):
-                return {"final_patch": state["edit_patches"][response.patch_index]}
+            if 0 <= response.patch_index < len(state["edit_patches"]):
+                return {self.final_patch_name: state["edit_patches"][response.patch_index]}
 
         self._logger.info(
             "FinalPatchSelectionNode failed to select a patch with correct index, defaulting to 0"
         )
-        return {"final_patch": state["edit_patches"][0]}
+        return {self.final_patch_name: state["edit_patches"][0]}

prometheus/lang_graph/nodes/git_reset_node.py

@@ -1,16 +1,28 @@
 import logging
+from typing import Dict, Sequence
 
 from prometheus.git.git_repository import GitRepository
 
 
 class GitResetNode:
-    def __init__(
-        self,
-        git_repo: GitRepository,
-    ):
+    def __init__(self, git_repo: GitRepository, exclude_files_key: str = None):
         self.git_repo = git_repo
         self._logger = logging.getLogger("prometheus.lang_graph.nodes.git_reset_node")
+        self.exclude_files_key = exclude_files_key
 
-    def __call__(self, _):
+    def __call__(self, state: Dict):
         self._logger.debug("Resetting the git repository")
-        self.git_repo.reset_repository()
+        excluded_files = []
+        if (
+            self.exclude_files_key
+            and self.exclude_files_key in state
+            and state[self.exclude_files_key]
+        ):
+            excluded_files = state[self.exclude_files_key]
+            if not isinstance(excluded_files, Sequence):
+                excluded_files = [excluded_files]
+            excluded_files = [str(f) for f in excluded_files]
+            self._logger.debug(
+                f"Excluding the following files when resetting the repository: {excluded_files}"
+            )
+        self.git_repo.reset_repository(excluded_files)

prometheus/lang_graph/subgraphs/bug_reproduction_subgraph.py

@@ -29,6 +29,16 @@
 
 
 class BugReproductionSubgraph:
+    """
+    A LangGraph-based subgraph that attempts to reproduce a bug reported in an issue
+    by synthesizing context, modifying files, running tests, and observing behavior.
+
+    This subgraph integrates multiple nodes for message generation, code modification,
+    file handling, Git diffing/resetting, container updating, test execution, and
+    structured evaluation. It forms a cyclic workflow that retries reproduction until
+    successful or recursion limits are reached.
+    """
+
     def __init__(
         self,
         advanced_model: BaseChatModel,
@@ -42,7 +52,10 @@ def __init__(
     ):
         self.git_repo = git_repo
 
+        # Node that generates initial bug context message from issue info
         issue_bug_reproduction_context_message_node = IssueBugReproductionContextMessageNode()
+
+        # Node that retrieves contextual code and AST info from Neo4j KG
         context_retrieval_subgraph_node = ContextRetrievalSubgraphNode(
             base_model,
             kg,
@@ -52,46 +65,70 @@ def __init__(
             "bug_reproducing_context",
         )
 
+        # Node that generates write instructions in natural language
         bug_reproducing_write_message_node = BugReproducingWriteMessageNode()
+
+        # Node that synthesizes code changes using advanced LLM
         bug_reproducing_write_node = BugReproducingWriteNode(advanced_model, kg)
+
+        # ToolNode wrapping write tools (e.g., insert, replace, etc.)
         bug_reproducing_write_tools = ToolNode(
             tools=bug_reproducing_write_node.tools,
             name="bug_reproducing_write_tools",
             messages_key="bug_reproducing_write_messages",
         )
+
+        # Node that inspects/modifies files (e.g., test configs, Dockerfiles)
         bug_reproducing_file_node = BugReproducingFileNode(base_model, kg)
+
+        # ToolNode wrapping file operation tools
         bug_reproducing_file_tools = ToolNode(
             tools=bug_reproducing_file_node.tools,
             name="bug_reproducing_file_tools",
             messages_key="bug_reproducing_file_messages",
         )
+
+        # Node that generates a Git diff from the modifications
         git_diff_node = GitDiffNode(git_repo, "bug_reproducing_patch")
+
+        # Node that rebuilds the container with the updated code
         update_container_node = UpdateContainerNode(container, git_repo)
+
+        # Node that runs test commands inside the container
         bug_reproducing_execute_node = BugReproducingExecuteNode(
             base_model, container, test_commands
         )
+
+        # ToolNode wrapping test execution tools (e.g., rerun, patch env, etc.)
         bug_reproducing_execute_tools = ToolNode(
             tools=bug_reproducing_execute_node.tools,
             name="bug_reproducing_execute_tools",
             messages_key="bug_reproducing_execute_messages",
         )
+
+        # Node that parses test outputs and determines if bug is reproduced
         bug_reproducing_structured_node = BugReproducingStructuredNode(advanced_model)
+
+        # Reset message buffers for file and execution steps before retry
         reset_bug_reproducing_file_messages_node = ResetMessagesNode(
             "bug_reproducing_file_messages"
         )
         reset_bug_reproducing_execute_messages_node = ResetMessagesNode(
             "bug_reproducing_execute_messages"
         )
+
+        # Reset Git state before retry
         git_reset_node = GitResetNode(git_repo)
 
+        # Define the LangGraph workflow using StateGraph
         workflow = StateGraph(BugReproductionState)
 
+        # Add all nodes to the graph
         workflow.add_node(
             "issue_bug_reproduction_context_message_node",
             issue_bug_reproduction_context_message_node,
         )
         workflow.add_node("context_retrieval_subgraph_node", context_retrieval_subgraph_node)
-
         workflow.add_node("bug_reproducing_write_message_node", bug_reproducing_write_message_node)
         workflow.add_node("bug_reproducing_write_node", bug_reproducing_write_node)
         workflow.add_node("bug_reproducing_write_tools", bug_reproducing_write_tools)
@@ -102,7 +139,6 @@ def __init__(
         workflow.add_node("bug_reproducing_execute_node", bug_reproducing_execute_node)
         workflow.add_node("bug_reproducing_execute_tools", bug_reproducing_execute_tools)
         workflow.add_node("bug_reproducing_structured_node", bug_reproducing_structured_node)
-
         workflow.add_node(
             "reset_bug_reproducing_file_messages_node", reset_bug_reproducing_file_messages_node
         )
@@ -112,61 +148,55 @@ def __init__(
         )
         workflow.add_node("git_reset_node", git_reset_node)
 
+        # Define transitions between nodes
         workflow.set_entry_point("issue_bug_reproduction_context_message_node")
         workflow.add_edge(
             "issue_bug_reproduction_context_message_node", "context_retrieval_subgraph_node"
         )
         workflow.add_edge("context_retrieval_subgraph_node", "bug_reproducing_write_message_node")
-
         workflow.add_edge("bug_reproducing_write_message_node", "bug_reproducing_write_node")
+
+        # Conditional loop through write tools
         workflow.add_conditional_edges(
             "bug_reproducing_write_node",
             functools.partial(tools_condition, messages_key="bug_reproducing_write_messages"),
-            {
-                "tools": "bug_reproducing_write_tools",
-                END: "bug_reproducing_file_node",
-            },
+            {"tools": "bug_reproducing_write_tools", END: "bug_reproducing_file_node"},
         )
         workflow.add_edge("bug_reproducing_write_tools", "bug_reproducing_write_node")
+
+        # Conditional loop through file tools
         workflow.add_conditional_edges(
             "bug_reproducing_file_node",
             functools.partial(tools_condition, messages_key="bug_reproducing_file_messages"),
-            {
-                "tools": "bug_reproducing_file_tools",
-                END: "git_diff_node",
-            },
+            {"tools": "bug_reproducing_file_tools", END: "git_diff_node"},
         )
         workflow.add_edge("bug_reproducing_file_tools", "bug_reproducing_file_node")
+
         workflow.add_edge("git_diff_node", "update_container_node")
         workflow.add_edge("update_container_node", "bug_reproducing_execute_node")
+
+        # Conditional loop through execution tools
         workflow.add_conditional_edges(
             "bug_reproducing_execute_node",
             functools.partial(tools_condition, messages_key="bug_reproducing_execute_messages"),
-            {
-                "tools": "bug_reproducing_execute_tools",
-                END: "bug_reproducing_structured_node",
-            },
+            {"tools": "bug_reproducing_execute_tools", END: "bug_reproducing_structured_node"},
         )
         workflow.add_edge("bug_reproducing_execute_tools", "bug_reproducing_execute_node")
+
+        # Final conditional edge: bug reproduced or retry
         workflow.add_conditional_edges(
             "bug_reproducing_structured_node",
             lambda state: state["reproduced_bug"],
             {True: END, False: "reset_bug_reproducing_file_messages_node"},
         )
-
         workflow.add_edge(
             "reset_bug_reproducing_file_messages_node",
             "reset_bug_reproducing_execute_messages_node",
         )
-        workflow.add_edge(
-            "reset_bug_reproducing_execute_messages_node",
-            "git_reset_node",
-        )
-        workflow.add_edge(
-            "git_reset_node",
-            "bug_reproducing_write_message_node",
-        )
+        workflow.add_edge("reset_bug_reproducing_execute_messages_node", "git_reset_node")
+        workflow.add_edge("git_reset_node", "bug_reproducing_write_message_node")
 
+        # Compile the subgraph for use
         self.subgraph = workflow.compile()
 
     def invoke(
@@ -176,6 +206,18 @@ def invoke(
         issue_comments: Sequence[Mapping[str, str]],
         recursion_limit: int = 50,
     ):
+        """
+        Invoke the bug reproduction subgraph on a given issue.
+
+        Args:
+            issue_title (str): The title of the issue.
+            issue_body (str): The issue description.
+            issue_comments (Sequence[Mapping[str, str]]): List of GitHub issue comments.
+            recursion_limit (int): Max iterations before aborting.
+
+        Returns:
+            Dict[str, Any]: Reproduction results including file and command info.
+        """
         config = {"recursion_limit": recursion_limit}
 
         input_state = {
@@ -193,6 +235,7 @@ def invoke(
                 "reproduced_bug_commands": output_state["reproduced_bug_commands"],
             }
         except GraphRecursionError:
+            # Fall back to safe state on failure
             self.git_repo.reset_repository()
             return {
                 "reproduced_bug": False,