From f8ea1be6f5874d9266ab595ce6f070af74988f3a Mon Sep 17 00:00:00 2001
From: Yue Pan <79363355+dcloud347@users.noreply.github.com>
Date: Fri, 4 Jul 2025 23:24:19 +0800
Subject: [PATCH 01/10] Update README.md to change PROMETHEUS_POSTGRES_URI to
 PROMETHEUS_POSTGRES_URL

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 5d573ff1..5ac265b4 100644
--- a/README.md
+++ b/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`
 
 ---
 

From ec1d298de668c9a8e065e97c0952577550deaad5 Mon Sep 17 00:00:00 2001
From: Yue Pan <79363355+dcloud347@users.noreply.github.com>
Date: Sat, 5 Jul 2025 15:10:55 +0800
Subject: [PATCH 02/10] Add final_patch_name parameter to
 FinalPatchSelectionNode and update related logic

---
 .../nodes/final_patch_selection_node.py       |  9 ++--
 .../issue_not_verified_bug_subgraph.py        | 41 ++++++++++++++++++-
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/prometheus/lang_graph/nodes/final_patch_selection_node.py b/prometheus/lang_graph/nodes/final_patch_selection_node.py
index 260ddb1d..37386dd2 100644
--- a/prometheus/lang_graph/nodes/final_patch_selection_node.py
+++ b/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]}
diff --git a/prometheus/lang_graph/subgraphs/issue_not_verified_bug_subgraph.py b/prometheus/lang_graph/subgraphs/issue_not_verified_bug_subgraph.py
index 6346ea29..76ad66b6 100644
--- a/prometheus/lang_graph/subgraphs/issue_not_verified_bug_subgraph.py
+++ b/prometheus/lang_graph/subgraphs/issue_not_verified_bug_subgraph.py
@@ -22,6 +22,17 @@
 
 
 class IssueNotVerifiedBugSubgraph:
+    """
+    This class defines a LangGraph-based subgraph to address GitHub issues that are suspected bugs
+    but have not yet been verified. It performs the following high-level steps:
+
+    1. Retrieves context from the knowledge graph based on the issue description.
+    2. Uses an LLM to analyze the issue and generate possible fix strategies.
+    3. Applies edits iteratively to the codebase using tool-enabled LLM actions.
+    4. Creates Git diffs and evaluates patches.
+    5. Selects a final patch based on quality and completeness.
+    """
+
     def __init__(
         self,
         advanced_model: BaseChatModel,
@@ -31,6 +42,7 @@ def __init__(
         neo4j_driver: neo4j.Driver,
         max_token_per_neo4j_result: int,
     ):
+        # Step 1: Prepare the issue-related context
         issue_bug_context_message_node = IssueBugContextMessageNode()
         context_retrieval_subgraph_node = ContextRetrievalSubgraphNode(
             model=base_model,
@@ -41,9 +53,11 @@ def __init__(
             context_key_name="bug_fix_context",
         )
 
+        # Step 2: Analyze the issue to identify potential bug-fix strategies
         issue_bug_analyzer_message_node = IssueBugAnalyzerMessageNode()
         issue_bug_analyzer_node = IssueBugAnalyzerNode(advanced_model)
 
+        # Step 3: Generate and apply candidate patches
         edit_message_node = EditMessageNode()
         edit_node = EditNode(advanced_model, kg)
         edit_tools = ToolNode(
@@ -51,16 +65,22 @@ def __init__(
             name="edit_tools",
             messages_key="edit_messages",
         )
+
+        # Step 4: Generate Git diffs and evaluate patches
         git_diff_node = GitDiffNode(git_repo, "edit_patches", return_list=True)
 
+        # Step 5: Reset git state and messages if the patch is insufficient
         git_reset_node = GitResetNode(git_repo)
         reset_issue_bug_analyzer_messages_node = ResetMessagesNode("issue_bug_analyzer_messages")
         reset_edit_messages_node = ResetMessagesNode("edit_messages")
 
-        final_patch_selection_node = FinalPatchSelectionNode(advanced_model)
+        # Step 6: Select the final patch from candidates
+        final_patch_selection_node = FinalPatchSelectionNode(advanced_model, "final_patch")
 
+        # Construct the LangGraph workflow
         workflow = StateGraph(IssueNotVerifiedBugState)
 
+        # Add nodes to the graph
         workflow.add_node("issue_bug_context_message_node", issue_bug_context_message_node)
         workflow.add_node("context_retrieval_subgraph_node", context_retrieval_subgraph_node)
 
@@ -80,6 +100,7 @@ def __init__(
 
         workflow.add_node("final_patch_selection_node", final_patch_selection_node)
 
+        # Define control flow between nodes
         workflow.set_entry_point("issue_bug_context_message_node")
         workflow.add_edge("issue_bug_context_message_node", "context_retrieval_subgraph_node")
         workflow.add_edge("context_retrieval_subgraph_node", "issue_bug_analyzer_message_node")
@@ -87,6 +108,8 @@ def __init__(
         workflow.add_edge("issue_bug_analyzer_node", "edit_message_node")
 
         workflow.add_edge("edit_message_node", "edit_node")
+
+        # Conditional path: if tool usage is needed, go to ToolNode; else proceed to diff
         workflow.add_conditional_edges(
             "edit_node",
             functools.partial(tools_condition, messages_key="edit_messages"),
@@ -94,16 +117,19 @@ def __init__(
         )
         workflow.add_edge("edit_tools", "edit_node")
 
+        # If not enough patches generated yet, reset and continue the loop
         workflow.add_conditional_edges(
             "git_diff_node",
             lambda state: len(state["edit_patches"]) < state["number_of_candidate_patch"],
             {True: "git_reset_node", False: "final_patch_selection_node"},
         )
 
+        # Reset loop for next candidate patch generation
         workflow.add_edge("git_reset_node", "reset_issue_bug_analyzer_messages_node")
         workflow.add_edge("reset_issue_bug_analyzer_messages_node", "reset_edit_messages_node")
         workflow.add_edge("reset_edit_messages_node", "issue_bug_analyzer_message_node")
 
+        # Final termination
         workflow.add_edge("final_patch_selection_node", END)
 
         self.subgraph = workflow.compile()
@@ -116,6 +142,19 @@ def invoke(
         number_of_candidate_patch: int,
         recursion_limit: int = 999,
     ):
+        """
+        Run the bug-fix subgraph on a given GitHub issue.
+
+        Args:
+            issue_title: The title of the GitHub issue.
+            issue_body: The body/description of the issue.
+            issue_comments: A list of comments on the issue for additional context.
+            number_of_candidate_patch: How many patch candidates to attempt before finalizing.
+            recursion_limit: Max iterations to allow in the graph loop (safety mechanism).
+
+        Returns:
+            Dict with the selected 'final_patch'.
+        """
         config = {"recursion_limit": recursion_limit}
 
         input_state = {

From af74bced8725c35f084a124f78802954809ba34b Mon Sep 17 00:00:00 2001
From: Yue Pan <79363355+dcloud347@users.noreply.github.com>
Date: Sat, 5 Jul 2025 15:18:06 +0800
Subject: [PATCH 03/10] Add candidate_patch_number parameter and enhance patch
 selection logic in IssueVerifiedBugSubgraph

---
 .../subgraphs/issue_verified_bug_state.py     |   4 +
 .../subgraphs/issue_verified_bug_subgraph.py  | 101 +++++++++---------
 2 files changed, 54 insertions(+), 51 deletions(-)

diff --git a/prometheus/lang_graph/subgraphs/issue_verified_bug_state.py b/prometheus/lang_graph/subgraphs/issue_verified_bug_state.py
index 227f8291..7278eba7 100644
--- a/prometheus/lang_graph/subgraphs/issue_verified_bug_state.py
+++ b/prometheus/lang_graph/subgraphs/issue_verified_bug_state.py
@@ -1,3 +1,4 @@
+from operator import add
 from typing import Annotated, Mapping, Sequence, TypedDict
 
 from langchain_core.messages import BaseMessage
@@ -22,6 +23,7 @@ class IssueVerifiedBugState(TypedDict):
     edit_messages: Annotated[Sequence[BaseMessage], add_messages]
 
     edit_patch: str
+    edit_patches: Annotated[Sequence[str], add]
 
     reproducing_test_fail_log: str
 
@@ -35,3 +37,5 @@ class IssueVerifiedBugState(TypedDict):
 
     max_refined_query_loop: int
     refined_query: str
+
+    number_of_candidate_patch: int
diff --git a/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py b/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py
index e151116a..9614e889 100644
--- a/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py
+++ b/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py
@@ -16,30 +16,22 @@
 from prometheus.lang_graph.nodes.context_retrieval_subgraph_node import ContextRetrievalSubgraphNode
 from prometheus.lang_graph.nodes.edit_message_node import EditMessageNode
 from prometheus.lang_graph.nodes.edit_node import EditNode
+from prometheus.lang_graph.nodes.final_patch_selection_node import FinalPatchSelectionNode
 from prometheus.lang_graph.nodes.git_diff_node import GitDiffNode
+from prometheus.lang_graph.nodes.git_reset_node import GitResetNode
 from prometheus.lang_graph.nodes.issue_bug_analyzer_message_node import IssueBugAnalyzerMessageNode
 from prometheus.lang_graph.nodes.issue_bug_analyzer_node import IssueBugAnalyzerNode
 from prometheus.lang_graph.nodes.issue_bug_context_message_node import IssueBugContextMessageNode
 from prometheus.lang_graph.nodes.noop_node import NoopNode
+from prometheus.lang_graph.nodes.reset_messages_node import ResetMessagesNode
 from prometheus.lang_graph.nodes.update_container_node import UpdateContainerNode
 from prometheus.lang_graph.subgraphs.issue_verified_bug_state import IssueVerifiedBugState
 
 
 class IssueVerifiedBugSubgraph:
     """
-    A LangGraph-based subgraph that handles verified bug issues by generating,
-    applying, and validating patch candidates.
-
-    This subgraph executes the following phases:
-    1. Context construction and retrieval from knowledge graph and codebase
-    2. Semantic analysis of the bug using advanced LLM
-    3. Patch generation via LLM and optional tool invocations
-    4. Patch application with Git diff visualization
-    5. Build and test the modified code in a containerized environment
-    6. Iterative refinement if verification fails
-
-    Attributes:
-        subgraph (StateGraph): The compiled LangGraph workflow to handle verified bugs.
+    LangGraph subgraph for resolving verified bugs with iterative patch generation
+    and final selection from multiple candidates.
     """
 
     def __init__(
@@ -53,23 +45,11 @@ def __init__(
         max_token_per_neo4j_result: int,
         build_commands: Optional[Sequence[str]] = None,
         test_commands: Optional[Sequence[str]] = None,
+        candidate_patch_number: int = 5,
     ):
-        """
-        Initialize the verified bug fix subgraph.
-
-        Args:
-            advanced_model (BaseChatModel): A strong LLM used for bug understanding and patch generation.
-            base_model (BaseChatModel): A smaller, less expensive LLM used for context retrieval and test verification.
-            container (BaseContainer): A build/test container to run code validations.
-            kg (KnowledgeGraph): A knowledge graph used for context-aware retrieval of relevant code entities.
-            git_repo (GitRepository): Git interface to apply patches and get diffs.
-            neo4j_driver (neo4j.Driver): Neo4j driver for executing graph-based semantic queries.
-            max_token_per_neo4j_result (int): Maximum tokens to limit output from Neo4j query results.
-            build_commands (Optional[Sequence[str]]): Commands to build the project inside the container.
-            test_commands (Optional[Sequence[str]]): Commands to test the project inside the container.
-        """
-
-        # Phase 1: Retrieve context related to the bug
+        self.candidate_patch_number = candidate_patch_number
+
+        # Step 1: Context setup
         issue_bug_context_message_node = IssueBugContextMessageNode()
         context_retrieval_subgraph_node = ContextRetrievalSubgraphNode(
             model=base_model,
@@ -80,11 +60,11 @@ def __init__(
             context_key_name="bug_fix_context",
         )
 
-        # Phase 2: Analyze the bug and generate hypotheses
+        # Step 2: Bug analysis
         issue_bug_analyzer_message_node = IssueBugAnalyzerMessageNode()
         issue_bug_analyzer_node = IssueBugAnalyzerNode(advanced_model)
 
-        # Phase 3: Generate code edits and optionally apply toolchains
+        # Step 3: Patch generation
         edit_message_node = EditMessageNode()
         edit_node = EditNode(advanced_model, kg)
         edit_tools = ToolNode(
@@ -93,47 +73,55 @@ def __init__(
             messages_key="edit_messages",
         )
 
-        # Phase 4: Apply patch, diff changes, and update the container
-        git_diff_node = GitDiffNode(git_repo, "edit_patch", "reproduced_bug_file")
+        # Step 4: Git diff accumulation
+        git_diff_node = GitDiffNode(
+            git_repo, "edit_patches", "reproduced_bug_file", return_list=True
+        )
+
+        # Reset & loop control
+        git_reset_node = GitResetNode(git_repo)
+        reset_issue_bug_analyzer_messages_node = ResetMessagesNode("issue_bug_analyzer_messages")
+        reset_edit_messages_node = ResetMessagesNode("edit_messages")
+
+        # Step 5: Patch selection and update container
+        patches_selection_node = FinalPatchSelectionNode(
+            final_patch_name="edit_patch", model=advanced_model
+        )
         update_container_node = UpdateContainerNode(container, git_repo)
 
-        # Phase 5: Re-run test case that reproduces the bug
+        # Step 6: Bug test
         bug_fix_verification_subgraph_node = BugFixVerificationSubgraphNode(
-            base_model,
-            container,
+            base_model, container
         )
 
-        # Phase 6: Optionally run full build and test after fix
+        # Step 7: Optional full build/test
         build_or_test_branch_node = NoopNode()
         build_and_test_subgraph_node = BuildAndTestSubgraphNode(
-            container,
-            advanced_model,
-            kg,
-            build_commands,
-            test_commands,
+            container, advanced_model, kg, build_commands, test_commands
         )
 
-        # Build the LangGraph workflow
+        # Build graph
         workflow = StateGraph(IssueVerifiedBugState)
 
-        # Add nodes to graph
+        # Add nodes
         workflow.add_node("issue_bug_context_message_node", issue_bug_context_message_node)
         workflow.add_node("context_retrieval_subgraph_node", context_retrieval_subgraph_node)
-
         workflow.add_node("issue_bug_analyzer_message_node", issue_bug_analyzer_message_node)
         workflow.add_node("issue_bug_analyzer_node", issue_bug_analyzer_node)
-
         workflow.add_node("edit_message_node", edit_message_node)
         workflow.add_node("edit_node", edit_node)
         workflow.add_node("edit_tools", edit_tools)
         workflow.add_node("git_diff_node", git_diff_node)
+        workflow.add_node("git_reset_node", git_reset_node)
+        workflow.add_node("reset_issue_bug_analyzer_messages_node", reset_issue_bug_analyzer_messages_node)
+        workflow.add_node("reset_edit_messages_node", reset_edit_messages_node)
+        workflow.add_node("patches_selection_node", patches_selection_node)
         workflow.add_node("update_container_node", update_container_node)
-
         workflow.add_node("bug_fix_verification_subgraph_node", bug_fix_verification_subgraph_node)
         workflow.add_node("build_or_test_branch_node", build_or_test_branch_node)
         workflow.add_node("build_and_test_subgraph_node", build_and_test_subgraph_node)
 
-        # Define edges for full flow
+        # Graph transitions
         workflow.set_entry_point("issue_bug_context_message_node")
         workflow.add_edge("issue_bug_context_message_node", "context_retrieval_subgraph_node")
         workflow.add_edge("context_retrieval_subgraph_node", "issue_bug_analyzer_message_node")
@@ -141,15 +129,25 @@ def __init__(
         workflow.add_edge("issue_bug_analyzer_node", "edit_message_node")
         workflow.add_edge("edit_message_node", "edit_node")
 
-        # Conditionally invoke tools or continue to diffing
         workflow.add_conditional_edges(
             "edit_node",
             functools.partial(tools_condition, messages_key="edit_messages"),
             {"tools": "edit_tools", END: "git_diff_node"},
         )
-
         workflow.add_edge("edit_tools", "edit_node")
-        workflow.add_edge("git_diff_node", "update_container_node")
+
+        # Loop if not enough patches
+        workflow.add_conditional_edges(
+            "git_diff_node",
+            lambda state: len(state["edit_patches"]) < state["number_of_candidate_patch"],
+            {True: "git_reset_node", False: "patches_selection_node"},
+        )
+
+        workflow.add_edge("git_reset_node", "reset_issue_bug_analyzer_messages_node")
+        workflow.add_edge("reset_issue_bug_analyzer_messages_node", "reset_edit_messages_node")
+        workflow.add_edge("reset_edit_messages_node", "issue_bug_analyzer_message_node")
+
+        workflow.add_edge("patches_selection_node", "update_container_node")
         workflow.add_edge("update_container_node", "bug_fix_verification_subgraph_node")
 
         # If test still fails, loop back to reanalyze the bug
@@ -197,6 +195,7 @@ def invoke(
             "run_existing_test": run_existing_test,
             "reproduced_bug_file": reproduced_bug_file,
             "reproduced_bug_commands": reproduced_bug_commands,
+            "number_of_candidate_patch": self.candidate_patch_number,
             "max_refined_query_loop": 3,
         }
 

From 078b2876ed95fa0e122a9d02fdf945b3106cc861 Mon Sep 17 00:00:00 2001
From: Yue Pan <79363355+dcloud347@users.noreply.github.com>
Date: Sat, 5 Jul 2025 15:18:21 +0800
Subject: [PATCH 04/10] Refactor BugFixVerificationSubgraphNode instantiation
 and format workflow node addition

---
 .../lang_graph/subgraphs/issue_verified_bug_subgraph.py   | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py b/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py
index 9614e889..ba076fab 100644
--- a/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py
+++ b/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py
@@ -90,9 +90,7 @@ def __init__(
         update_container_node = UpdateContainerNode(container, git_repo)
 
         # Step 6: Bug test
-        bug_fix_verification_subgraph_node = BugFixVerificationSubgraphNode(
-            base_model, container
-        )
+        bug_fix_verification_subgraph_node = BugFixVerificationSubgraphNode(base_model, container)
 
         # Step 7: Optional full build/test
         build_or_test_branch_node = NoopNode()
@@ -113,7 +111,9 @@ def __init__(
         workflow.add_node("edit_tools", edit_tools)
         workflow.add_node("git_diff_node", git_diff_node)
         workflow.add_node("git_reset_node", git_reset_node)
-        workflow.add_node("reset_issue_bug_analyzer_messages_node", reset_issue_bug_analyzer_messages_node)
+        workflow.add_node(
+            "reset_issue_bug_analyzer_messages_node", reset_issue_bug_analyzer_messages_node
+        )
         workflow.add_node("reset_edit_messages_node", reset_edit_messages_node)
         workflow.add_node("patches_selection_node", patches_selection_node)
         workflow.add_node("update_container_node", update_container_node)

From 1f6b1dd47e66bd46d1c162a8d8c1cafac94e04d8 Mon Sep 17 00:00:00 2001
From: Yue Pan <79363355+dcloud347@users.noreply.github.com>
Date: Sat, 5 Jul 2025 16:28:20 +0800
Subject: [PATCH 05/10] Add patch parameter to create_branch method for
 applying patches

---
 prometheus/git/git_repository.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/prometheus/git/git_repository.py b/prometheus/git/git_repository.py
index aa5412da..7cbfbd29 100644
--- a/prometheus/git/git_repository.py
+++ b/prometheus/git/git_repository.py
@@ -139,6 +139,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)

From dd5a77cac642c2a9a133afdf59ea3877263267f5 Mon Sep 17 00:00:00 2001
From: Yue Pan <79363355+dcloud347@users.noreply.github.com>
Date: Sun, 6 Jul 2025 01:21:28 +0800
Subject: [PATCH 06/10] Add max_output_tokens parameter to get_model function
 calls in tests

---
 tests/app/services/test_llm_service.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/tests/app/services/test_llm_service.py b/tests/app/services/test_llm_service.py
index 06a5f43f..f5556e65 100644
--- a/tests/app/services/test_llm_service.py
+++ b/tests/app/services/test_llm_service.py
@@ -38,6 +38,7 @@ def test_llm_service_init(mock_custom_chat_openai, mock_chat_anthropic):
         openai_format_api_key="openai-key",
         openai_format_base_url="https://api.openai.com/v1",
         anthropic_api_key="anthropic-key",
+        max_output_tokens=settings.MAX_OUTPUT_TOKENS,
     )
 
     # Verify
@@ -66,6 +67,7 @@ def test_get_openai_format_model(mock_custom_chat_openai):
         model_name="openrouter/model",
         openai_format_api_key="openrouter-key",
         openai_format_base_url="https://openrouter.ai/api/v1",
+        max_output_tokens=settings.MAX_OUTPUT_TOKENS,
     )
 
     # Verify
@@ -81,7 +83,11 @@ def test_get_openai_format_model(mock_custom_chat_openai):
 
 def test_get_model_claude(mock_chat_anthropic):
     # Exercise
-    get_model(model_name="claude-2.1", anthropic_api_key="anthropic-key")
+    get_model(
+        model_name="claude-2.1",
+        anthropic_api_key="anthropic-key",
+        max_output_tokens=settings.MAX_OUTPUT_TOKENS,
+    )
 
     # Verify
     mock_chat_anthropic.assert_called_once_with(
@@ -95,7 +101,11 @@ def test_get_model_claude(mock_chat_anthropic):
 
 def test_get_model_gemini(mock_chat_google):
     # Exercise
-    get_model(model_name="gemini-pro", gemini_api_key="gemini-key")
+    get_model(
+        model_name="gemini-pro",
+        gemini_api_key="gemini-key",
+        max_output_tokens=settings.MAX_OUTPUT_TOKENS,
+    )
 
     # Verify
     mock_chat_google.assert_called_once_with(

From 6dc52628df4c7eb8f257e92c8a00230173c4caad Mon Sep 17 00:00:00 2001
From: Yue Pan <79363355+dcloud347@users.noreply.github.com>
Date: Sun, 6 Jul 2025 01:21:44 +0800
Subject: [PATCH 07/10] Update test command in README to remove unnecessary
 filter

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 5ac265b4..9dd1ea13 100644
--- a/README.md
+++ b/README.md
@@ -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:

From d39b222db72799944bac72fb0738e91a92edbf5e Mon Sep 17 00:00:00 2001
From: Yue Pan <79363355+dcloud347@users.noreply.github.com>
Date: Sun, 6 Jul 2025 17:37:24 +0800
Subject: [PATCH 08/10] Refactor git repository fixture to ignore .git
 directory and enhance cleanup process

---
 tests/git/test_git_repository.py | 43 ++++++++++++++++++++++++++++++++
 tests/test_utils/fixtures.py     | 10 +++++---
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/tests/git/test_git_repository.py b/tests/git/test_git_repository.py
index a48e04c1..7e724e79 100644
--- a/tests/git/test_git_repository.py
+++ b/tests/git/test_git_repository.py
@@ -133,3 +133,46 @@ def test_remove_repository(git_repo_fixture):  # noqa: F811
 
         mock_rmtree.assert_called_once_with(local_path)
         assert git_repo.repo is None
+
+
+@pytest.mark.skipif(
+    sys.platform.startswith("win"),
+    reason="Test fails on Windows because of cptree in git_repo_fixture",
+)
+@pytest.mark.git
+def test_reset_repository_with_exclusions(git_repo_fixture):  # noqa: F811
+    repo_path = Path(git_repo_fixture.working_dir).absolute()
+    tracked_file = repo_path / "test.c"
+    untracked_file = repo_path / "temp_untracked.txt"
+    excluded_file = repo_path / "keep_me.txt"
+
+    # Create GitRepository instance
+    git_repo = GitRepository(
+        address=str(repo_path),
+        working_directory=Path("/foo/bar"),
+        copy_to_working_dir=False,
+    )
+
+    # Initialize the repository and create a tracked file
+    git_repo.repo.git.add(tracked_file)
+    git_repo.repo.git.commit("-m", "initial commit", "--allow-empty")
+
+    # 1. Modify a tracked file
+    original_content = tracked_file.read_text()
+    tracked_file.write_text("int main() { return 0; }\n")
+
+    # 2. Create untracked and excluded files
+    untracked_file.write_text("This is a temp file")
+    excluded_file.write_text("Don't delete me")
+
+    # 3. Run reset with exclusions
+    rel_excluded_file = str(excluded_file.relative_to(repo_path))
+    git_repo.reset_repository(excluded_files=[rel_excluded_file])
+
+    # 4. Assertions
+    assert tracked_file.read_text() == original_content
+    assert not untracked_file.exists()
+    assert excluded_file.exists()
+
+    # Cleanup
+    excluded_file.unlink(missing_ok=True)
diff --git a/tests/test_utils/fixtures.py b/tests/test_utils/fixtures.py
index 4073a844..b4aefdab 100644
--- a/tests/test_utils/fixtures.py
+++ b/tests/test_utils/fixtures.py
@@ -73,10 +73,12 @@ def git_repo_fixture():
     original_project_path = test_project_paths.TEST_PROJECT_PATH
 
     try:
-        shutil.copytree(original_project_path, temp_project_dir)
-        shutil.move(temp_project_dir / test_project_paths.GIT_DIR.name, temp_project_dir / ".git")
+        shutil.copytree(
+            original_project_path, temp_project_dir, ignore=shutil.ignore_patterns(".git")
+        )
+
+        repo = Repo.init(temp_project_dir)
 
-        repo = Repo(temp_project_dir)
         yield repo
     finally:
-        shutil.rmtree(temp_project_dir)
+        shutil.rmtree(temp_dir, ignore_errors=True)

From 08157fc58e40e4701199756b0e1f66a44d6b62dd Mon Sep 17 00:00:00 2001
From: Yue Pan <79363355+dcloud347@users.noreply.github.com>
Date: Sun, 6 Jul 2025 18:03:22 +0800
Subject: [PATCH 09/10] Enhance BugReproductionSubgraph with detailed
 docstrings and structured workflow nodes

---
 .../subgraphs/bug_reproduction_subgraph.py    | 91 ++++++++++++++-----
 1 file changed, 67 insertions(+), 24 deletions(-)

diff --git a/prometheus/lang_graph/subgraphs/bug_reproduction_subgraph.py b/prometheus/lang_graph/subgraphs/bug_reproduction_subgraph.py
index 8a4893ff..e8d22c4e 100644
--- a/prometheus/lang_graph/subgraphs/bug_reproduction_subgraph.py
+++ b/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,

From a670cc10556c2cb55894a763a12ea6d8dd443502 Mon Sep 17 00:00:00 2001
From: Yue Pan <79363355+dcloud347@users.noreply.github.com>
Date: Sun, 6 Jul 2025 19:38:30 +0800
Subject: [PATCH 10/10] Add support for excluding files in repository reset
 functionality

---
 prometheus/git/git_repository.py              | 25 ++++++++++++++++---
 prometheus/lang_graph/nodes/git_reset_node.py | 24 +++++++++++++-----
 .../subgraphs/issue_verified_bug_subgraph.py  |  2 +-
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/prometheus/git/git_repository.py b/prometheus/git/git_repository.py
index 7cbfbd29..8f3fcf4c 100644
--- a/prometheus/git/git_repository.py
+++ b/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:
diff --git a/prometheus/lang_graph/nodes/git_reset_node.py b/prometheus/lang_graph/nodes/git_reset_node.py
index b9d63b30..86648257 100644
--- a/prometheus/lang_graph/nodes/git_reset_node.py
+++ b/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)
diff --git a/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py b/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py
index ba076fab..52abaa23 100644
--- a/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py
+++ b/prometheus/lang_graph/subgraphs/issue_verified_bug_subgraph.py
@@ -79,7 +79,7 @@ def __init__(
         )
 
         # Reset & loop control
-        git_reset_node = GitResetNode(git_repo)
+        git_reset_node = GitResetNode(git_repo, exclude_files_key="reproduced_bug_file")
         reset_issue_bug_analyzer_messages_node = ResetMessagesNode("issue_bug_analyzer_messages")
         reset_edit_messages_node = ResetMessagesNode("edit_messages")