243 lines
10 KiB
Python
Executable File
243 lines
10 KiB
Python
Executable File
import re
|
|
import shlex
|
|
|
|
from trae_agent.tools import tools_registry
|
|
from trae_agent.tools.base import Tool, ToolResult
|
|
from trae_agent.utils.config import ModelConfig
|
|
from trae_agent.utils.llm_clients.llm_basics import LLMMessage, LLMResponse
|
|
from trae_agent.utils.llm_clients.llm_client import LLMClient
|
|
from trae_agent.utils.trajectory_recorder import TrajectoryRecorder
|
|
|
|
from .sandbox import Sandbox
|
|
|
|
|
|
class CandidatePatch:
|
|
def __init__(self, id, patch, cleaned_patch, is_success_regression, is_success_patch):
|
|
self.id = id
|
|
self.patch = patch
|
|
self.cleaned_patch = cleaned_patch
|
|
self.is_success_regression = is_success_regression
|
|
self.is_success_patch = is_success_patch
|
|
|
|
|
|
def build_system_prompt(candidate_length: int) -> str:
|
|
init_prompt = f"""\
|
|
# ROLE: Act as an expert code evaluator. Given a codebase, an github issue and **{candidate_length} candidate patches** proposed by your colleagues, your responsibility is to **select the correct one** to solve the issue.
|
|
|
|
# WORK PROCESS:
|
|
You are given a software issue and multiple candidate patches. Your goal is to identify the patch that correctly resolves the issue.
|
|
|
|
Follow these steps methodically:
|
|
|
|
**1. Understand the Issue and Codebase**
|
|
Carefully read the issue description to comprehend the problem. You may need to examine the codebase for context, including:
|
|
(1) Code referenced in the issue description;
|
|
(2) The original code modified by each patch;
|
|
(3) Unchanged parts of the same file;
|
|
(4) Related files, functions, or modules that interact with the affected code.
|
|
|
|
**2. Analyze the Candidate Patches**
|
|
For each patch, analyze its logic and intended fix. Consider whether the changes align with the issue description and coding conventions.
|
|
|
|
**3. Validate Functionality (Optional but Recommended)**
|
|
If needed, write and run unit tests to evaluate the correctness and potential side effects of each patch.
|
|
|
|
**4. Select the Best Patch**
|
|
Choose the patch that best resolves the issue with minimal risk of introducing new problems.
|
|
|
|
# FINAL REPORT: If you have successfully selected the correct patch, submit your answer in the following format:
|
|
### Status: succeed
|
|
### Result: Patch-x
|
|
### Analysis: [Explain why Patch-x is correct.]
|
|
|
|
# IMPORTANT TIPS:
|
|
1. Never avoid making a selection.
|
|
2. Do not propose new patches.
|
|
3. There must be at least one correct patch.
|
|
"""
|
|
return init_prompt
|
|
|
|
|
|
def parse_tool_response(answer: LLMResponse, finish_reason: str, sandbox_session):
|
|
result: list[LLMMessage] = []
|
|
print("finish_reason:", finish_reason)
|
|
if answer.tool_calls and len(answer.tool_calls) > 0:
|
|
for tool_call in answer.tool_calls:
|
|
tool_call_id = tool_call.call_id
|
|
tool_name = tool_call.name
|
|
|
|
if tool_name == "str_replace_based_edit_tool":
|
|
cmd = "cd /home/swe-bench/tools/ && /home/swe-bench/py312/bin/python3 execute_str_replace_editor.py"
|
|
elif tool_name == "bash":
|
|
cmd = (
|
|
"cd /home/swe-bench/tools/ && /home/swe-bench/py312/bin/python3 execute_bash.py"
|
|
)
|
|
else:
|
|
tool_message = LLMMessage(
|
|
role="user",
|
|
content="The tool name you provided is not in the list. Please choose one from `str_replace_editor` or `bash`!",
|
|
tool_result=ToolResult(
|
|
call_id=tool_call_id,
|
|
name=tool_name,
|
|
success=False,
|
|
error="The tool name you provided is not in the list. Please choose one from `str_replace_editor` or `bash`!",
|
|
),
|
|
)
|
|
result.append(tool_message)
|
|
continue
|
|
|
|
all_arguments_valid = True
|
|
tool_arguments = tool_call.arguments
|
|
for key in tool_arguments:
|
|
if isinstance(tool_arguments[key], list):
|
|
try:
|
|
tool_arguments[key] = str([int(factor) for factor in tool_arguments[key]])
|
|
cmd += f" --{key} {shlex.quote(tool_arguments[key])}"
|
|
except Exception:
|
|
pass
|
|
elif isinstance(tool_arguments[key], (int, bool)):
|
|
cmd += f" --{key} {tool_arguments[key]}"
|
|
elif isinstance(tool_arguments[key], dict):
|
|
all_arguments_valid = False
|
|
break
|
|
else:
|
|
cmd += f" --{key} {shlex.quote(tool_arguments[key])}"
|
|
|
|
if not all_arguments_valid:
|
|
print("Tool Call Status: -1")
|
|
tool_message = LLMMessage(
|
|
role="user",
|
|
content="Failed call tool. One of the arguments is dict type, you need to check the definition the tool.",
|
|
tool_result=ToolResult(
|
|
call_id=tool_call_id,
|
|
name=tool_name,
|
|
success=False,
|
|
error="Failed call tool. One of the arguments is dict type, you need to check the definition the tool.",
|
|
),
|
|
)
|
|
result.append(tool_message)
|
|
continue
|
|
|
|
cmd += " > /home/swe-bench/tools/log.out 2>&1"
|
|
print(repr(cmd))
|
|
_ = sandbox_session.execute(cmd)
|
|
sandbox_res = sandbox_session.execute("cat /home/swe-bench/tools/log.out")
|
|
status = ""
|
|
status_line_index = -1
|
|
sandbox_res_str_list = sandbox_res.split("\n")
|
|
for index, line in enumerate(sandbox_res_str_list):
|
|
if line.strip().startswith("Tool Call Status:"):
|
|
status = line
|
|
status_line_index = index
|
|
break
|
|
if status_line_index != -1:
|
|
sandbox_res_str_list.pop(status_line_index)
|
|
res_content = "\n".join(sandbox_res_str_list)
|
|
print(status)
|
|
tool_message = LLMMessage(
|
|
role="user",
|
|
content=res_content,
|
|
tool_result=ToolResult(
|
|
call_id=tool_call_id,
|
|
name=tool_name,
|
|
success=status != "Tool Call Status: -1",
|
|
result=res_content,
|
|
error=None if status != "Tool Call Status: -1" else res_content,
|
|
),
|
|
)
|
|
result.append(tool_message)
|
|
|
|
return result
|
|
|
|
|
|
class SelectorAgent:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
llm_config: ModelConfig,
|
|
sandbox: Sandbox,
|
|
project_path: str,
|
|
issue_description: str,
|
|
trajectory_file_name: str,
|
|
candidate_list: list[CandidatePatch],
|
|
max_turn: int = 50,
|
|
):
|
|
self.llm_config = llm_config
|
|
self.max_turn = max_turn
|
|
self.sandbox = sandbox
|
|
self.sandbox_session = self.sandbox.get_session()
|
|
self.sandbox_session.execute("git reset --hard HEAD")
|
|
self.initial_messages: list[LLMMessage] = []
|
|
self.candidate_list: list[CandidatePatch] = candidate_list
|
|
self.project_path: str = project_path
|
|
self.issue_description: str = issue_description
|
|
self.tools: list[Tool] = [
|
|
tools_registry[tool_name](model_provider=llm_config.model_provider.provider)
|
|
for tool_name in ["bash", "str_replace_based_edit_tool"]
|
|
]
|
|
self.llm_client = LLMClient(llm_config)
|
|
self.trajectory_recorder: TrajectoryRecorder = TrajectoryRecorder(trajectory_file_name)
|
|
|
|
self.initial_messages.append(
|
|
LLMMessage(role="system", content=build_system_prompt(len(candidate_list)))
|
|
)
|
|
user_prompt = f"\n[Codebase path]:\n{project_path}\n\n[Github issue description]:\n```\n{issue_description}\n```\n\n[Candidate Patches]:"
|
|
for idx in range(0, len(candidate_list)):
|
|
user_prompt += f"\nPatch-{idx + 1}:\n```\n{candidate_list[idx].patch}\n```"
|
|
user_message = LLMMessage(role="user", content=user_prompt)
|
|
self.initial_messages.append(user_message)
|
|
|
|
def run(self):
|
|
print(f"max_turn: {self.max_turn}")
|
|
print(f"### User Prompt:\n{self.initial_messages[1].content}\n")
|
|
|
|
turn = 0
|
|
final_id, final_patch = self.candidate_list[0].id, self.candidate_list[0].patch
|
|
messages = self.initial_messages
|
|
while turn < self.max_turn:
|
|
turn += 1
|
|
llm_response = self.llm_client.chat(messages, self.llm_config, self.tools)
|
|
self.trajectory_recorder.record_llm_interaction(
|
|
messages,
|
|
llm_response,
|
|
self.llm_config.model_provider.provider,
|
|
self.llm_config.model,
|
|
self.tools,
|
|
)
|
|
answer_content = llm_response.content
|
|
print(f"\n### Selector's Answer({turn})\n", answer_content)
|
|
messages: list[LLMMessage] = []
|
|
match = re.search(
|
|
r"(?:###\s*)?Status:\s*(success|succeed|successfully|successful)\s*\n\s*(?:###\s*)?Result:",
|
|
answer_content,
|
|
)
|
|
|
|
if match:
|
|
print("Match-1:", match.group(1).strip())
|
|
match = re.search(
|
|
r"(?:###\s*)?Result:\s*(.+?)\s*(?:###\s*)?Analysis:", answer_content
|
|
)
|
|
if match:
|
|
result = match.group(1).strip().split("Patch-")[-1]
|
|
print("Match-2:", result)
|
|
if result in [str(_ + 1) for _ in range(len(self.candidate_list))]:
|
|
final_id = self.candidate_list[int(result) - 1].id
|
|
final_patch = self.candidate_list[int(result) - 1].patch
|
|
else:
|
|
final_id = self.candidate_list[0].id
|
|
final_patch = self.candidate_list[0].patch
|
|
break
|
|
else:
|
|
messages += parse_tool_response(
|
|
llm_response, llm_response.finish_reason or "", self.sandbox_session
|
|
)
|
|
if messages[-1].content and " seconds. Partial output:" in messages[-1].content:
|
|
self.sandbox_session = self.sandbox.get_session()
|
|
|
|
print(f"\n### System Response({turn})\n", messages)
|
|
self.trajectory_recorder.finalize_recording(True, final_patch)
|
|
self.sandbox_session.execute("git reset --hard HEAD")
|
|
self.sandbox_session.close()
|
|
|
|
return final_id, final_patch
|