gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-08-01 09:27:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: refine the prompt to force complete code & refine the logic of running (#1069)

Browse Source 
* change refine prompt for full code

* fix: fix the logic of running

* refine prompt

* fix some bugs

* fix

* add two guidelines

* refactor the code

* make costeer evaluator more logical

* refine eval prompt

* make costeer eval prompt markdown

* update code diff prompt

* correct pipeline

* feat: add apply_patch utility and update ret.py with patch functionality (#1071)

* restore to the right version

* fix the docstring

* fix extract_output fcn

* add inplace parameter to apply patch

* remove enable_runner_iteration and make the eval prompt same as main

* refine runner eval prompt based on main

* Update rdagent/scenarios/data_science/dev/runner/prompts.yaml

* add wait_retry

* refactor: move enable_runner_code_diff to DSRunnerCoSTEERSettings as diff_mode

* reformat and remove enable_runner_code_diff

---------

Co-authored-by: yuanteli <1957922024@qq.com>
Co-authored-by: Xu <v-xuminrui@microsoft.com>
Co-authored-by: Jensen Lee <91518020+Jensen246@users.noreply.github.com>
Co-authored-by: you-n-g <you-n-g@users.noreply.github.com>
Co-authored-by: Qizheng Li <jenssenlee@163.com>
This commit is contained in:
amstrongzyf
2025-07-17 11:59:27 +08:00
committed by GitHub
parent 6469a0679a
commit e3679d6c0d
7 changed files with 730 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rdagent/core/experiment.py
+1 -1
View File
@@ -184,7 +184,7 @@ class FBWorkspace(Workspace):
workspace_data_file_path = workspace_path / data_file_path.name
if workspace_data_file_path.exists():
workspace_data_file_path.unlink()
if platform.system() == "Linux":
if platform.system() in ("Linux", "Darwin"):
os.symlink(data_file_path, workspace_data_file_path)
if platform.system() == "Windows":
os.link(data_file_path, workspace_data_file_path)
rdagent/scenarios/data_science/dev/runner/__init__.py
+28 -10
View File
@@ -21,8 +21,9 @@ from rdagent.core.scenario import Scenario
from rdagent.log import rdagent_logger as logger
from rdagent.oai.llm_utils import APIBackend, md5_hash
from rdagent.scenarios.data_science.dev.runner.eval import DSCoSTEERCoSTEEREvaluator
from rdagent.utils.agent.ret import PythonBatchEditOut
from rdagent.utils.agent.ret import PythonBatchEditOut, PythonBatchPatchOut
from rdagent.utils.agent.tpl import T
from rdagent.utils.workflow import wait_retry
class DSRunnerCoSTEERSettings(DSCoderCoSTEERSettings):
@@ -33,10 +34,12 @@ class DSRunnerCoSTEERSettings(DSCoderCoSTEERSettings):
max_seconds: int = DS_RD_SETTING.full_timeout
env_type: str = "docker"
diff_mode: bool = False
# TODO: extract a function for env and conf.
class DSRunnerMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
@wait_retry(retry_n=5)
def implement_one_task(
self,
target_task: CoSTEERTask,
@@ -44,29 +47,44 @@ class DSRunnerMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
workspace: FBWorkspace | None = None,
prev_task_feedback: CoSTEERSingleFeedback | None = None,
) -> dict[str, str]:
if prev_task_feedback is None:
# if no prev_tak_feedback, it is the first loop; we do not make any changes and goto evaluators directly.
# if no prev_task_feedback, it is the first loop; we do not make any changes and goto evaluators directly.
return {}
# Output Agent Map
output_map = {
True: (PythonBatchPatchOut.get_spec(), PythonBatchPatchOut.extract_output),
False: (
PythonBatchEditOut.get_spec(with_del=False),
PythonBatchEditOut.extract_output,
),
}
output_spec, extract_output_fn = output_map[self.settings.diff_mode]
if prev_task_feedback.hyperparameter_tuning_decision:
task_information_str = target_task.get_task_information()
# 1. code
# Use system_refine for hyperparameter tuning
system_prompt = T(".prompts:DSCoSTEER.system_refine").r(
out_spec=PythonBatchEditOut.get_spec(with_del=False),
out_spec=output_spec,
diff_mode=self.settings.diff_mode,
)
else:
task_information_str = target_task.get_task_information()
# 1. code
system_prompt = T(".prompts:DSCoSTEER.system_debugger").r(
# Use system_debugger for error fixing and debugging
system_prompt = T(".prompts:DSCoSTEER.system_refine").r(
task_desc=task_information_str,
out_spec=PythonBatchEditOut.get_spec(with_del=False),
out_spec=output_spec,
diff_mode=self.settings.diff_mode,
)
# Generate user prompt for both cases
user_prompt = T(".prompts:DSCoSTEER.user").r(
code=workspace.all_codes,
feedback=prev_task_feedback,
hyperparameter_tuning_suggestion=prev_task_feedback.hyperparameter_tuning_suggestion,
)
batch_edit = PythonBatchEditOut.extract_output(
batch_edit = extract_output_fn(
APIBackend().build_messages_and_create_chat_completion(
user_prompt=user_prompt,
system_prompt=system_prompt,
@@ -112,7 +130,7 @@ class DSCoSTEERRunner(CoSTEER):
settings = DSRunnerCoSTEERSettings()
es = DSRunnerMultiProcessEvolvingStrategy(scen=scen, settings=settings)
# In runner, we don't need very big loops, so we set max_loop to 3
# In runner, we don't need very big loops, so we set max_loop to runner_max_loop
super().__init__(
*args,
settings=settings,
rdagent/scenarios/data_science/dev/runner/prompts.yaml
+90 -59
View File
@@ -1,57 +1,60 @@
DSCoSTEER_eval:
system: |-
You are a data scientist responsible for evaluating all the code.
{% include "scenarios.data_science.share:scen.role" %}
{% if is_sub_enabled %}
You will be provided with:
1. `Code base`: The code base of the solution
2. `The stdout of code execution and testing`: The generated stdout when executing the code base and corresponding testing
3, `The time spent on code execution`: The time spent on the code execution
4. `The timeout of code execution`: the time limitation of the code execution
5. `The percent of timeout used`: the percentage of the time limitation used
Your task is to perform the following evaluation(s):
## Target Task Description
The user is trying to build a data science solution in the following scenario:
# Evalution 1: Code Correctness
## Scenario
The code is focusing on the following scenario:
{{ scenario }}
The task is as follows:
## Target Task Description
The code is focusing on the following task
{{ task_desc }}
## Runtime Environment
You have following environment to run the code:
{{ runtime_environment }}
The whole workflow includes multiple stages, such as:
- Data loading
- Feature engineering
- Model training
- Ensembling
## You'll be provided with the following information about a solution to the Target Task
`code base`: The code base of the solution
`the stdout of code execution and testing`: The generated stdout when executing the code base and corresponding testing
`the time spent on code execution`: The time spent on the code execution
`the timeout of code execution`: the time limitation of the code execution
`the percent of timeout used`: the percentage of the time limitation used
## Your task is to provide feedback on the solution to the Target Task
In the feedback response,
Evaluate the code base based on several aspects, including execution, return checking, and code quality. After your evaluation, make a clear decision to either accept or reject the solution in the `final_decision` section.
## Evaluation Guidelines
1. Evaluate the code base based on several aspects, including execution correctness, return checking, and code quality.
2. Ensure the code does not contain any incorrect, fabricated, or deceptive operations, such as mocking data, scores, or results.
3. Confirm that the prediction file (`submission.csv`) is generated using only the test dataset, and its format matches the sample submission.
If the code does not satisfy the requirements:
- Set "final_decision" to false.
- set "hyperparameter_tuning_decision" to false.
- Set "hyperparameter_tuning_suggestion" to an empty string.
If the code satisfy the requirements:
- Proceed to the next evaluation.
# Evaluation 2: Hyperparameter
## Evaluation Description
The user will provide you the time spent on the whole code execution and the timeout of the code execution. You should decide whether the hyperparameter is reasonable based on the time.
For example, if the code uses only a very small portion of the allowed time, and hyperparameters like `n_estimators` or `epochs` have low values, with early stopping not being triggered and possible signs of underfitting, you should suggest increasing these hyperparameters.
You should also notice other resources utilization hyper-parameters,
For example, if you are using a GPU with large memory, and the batch size is set very low, you should suggest increasing the batch size if it is not reasonable.
Please provide your feedback in two key-value pairs:
"hyperparameter_tuning_decision": <true/false>
"hyperparameter_tuning_suggestion": <suggestion in plain text for hyperparameter tuning, e.g., increase n_estimators to 1000, increase epochs to 100, increase batch size to 64, give an empty string if decide not to tune the hyperparameter>
[Notice]
- You should only suggest the hyperparameter tuning if the code applies early stopping strategy because increasing the training time blindly may lead to overfitting. Once you found the code didn't apply early stopping strategy, you should not suggest to tune the hyperparameter.
- Your suggestion should be reasonable and include not only the target hyperparameter but also the hyperparameter sets.
- Your suggestion should have a strong chance of improving the model's performance. Focus on the most obvious and impactful opportunities for quick improvement by leveraging more training time. Don't explore hyperparameters with low confidence. If there are no obvious and impactful opportunities and the code runs well, please accept it.
- Once you decide to tune the hyperparameter you should set "final_decision" to false.
[Format]
- "hyperparameter_tuning_suggestion" should begin with a clear observation, followed by your suggestion. For example: "[Observation] The maximum number of epochs was reached, but the validation loss is still going down and early stopping was not activated. Only 15% of the allowed time was used. [Suggestion] We recommend increasing epochs to 100 to avoid underfitting and further improve model performance."
## Evaluation Guidelines
1. The code execution time or resource utilization suggest that there is room for improvement in the hyperparameters.
2. The code must apply early stopping strategy already (in order to prevent overfitting).
3. Your suggestion should have a strong chance of improving the model's performance. Focus on the most obvious and impactful opportunities for quick improvement by leveraging more training time. Don't explore hyperparameters with low confidence. If there are no obvious and impactful opportunities and the code runs well, please accept it.
If the code satisfy the requirements:
- Set "hyperparameter_tuning_decision" to true.
- Set "final_decision" to false.
- Provide a reasonable suggestion in "hyperparameter_tuning_suggestion". The "hyperparameter_tuning_suggestion" should begin with a clear observation, followed by your suggestion. For example: "[Observation] The maximum number of epochs was reached, but the validation loss is still going down and early stopping was not activated. Only 15% of the allowed time was used. [Suggestion] We recommend increasing epochs to 100 to avoid underfitting and further improve model performance."
If the code does not satisfy the requirements:
- Set "hyperparameter_tuning_decision" to false.
- Set "hyperparameter_tuning_suggestion" to an empty string.
{% if is_sub_enabled %}
The user will provide you the whole code base, some logs generated during the execution of the whole workflow. Your evaluation scope includes whether the workflow code:
1. Executes successfully, correctly organizing components and generating a final submission.
2. Generates predictions in the correct format, ensuring they align with the **sample submission** structure!
Please respond with your feedback in the following JSON format and order
## Output format
Please respond with your feedback in the following JSON format and order without anything else:
```json
{
"execution": "Describe whether the whole code base executed successfully and generating the final submission. Include any errors or issues encountered, and retain all error messages and traceback details.",
@@ -81,62 +84,90 @@ DSCoSTEER_eval:
# NOTE: when is_sub_enabled == False, we don't have any checking about the return. So it is just placeholder currently
user: |-
--------- code base ---------
# Code base
{{ code }}
--------- the stdout of code execution and testing ---------
## Stdout of code execution and testing
{{ stdout }}
--------- the time spent on code execution ---------
# The time spend on code execution and timeout
{{ time_spent }}
--------- the timeout of code execution ---------
## The timeout of code execution
{{ timeout }}
--------- the percent of timeout used ---------
## The percent of timeout used
{{ percent_of_timeout_used }}
DSCoSTEER:
system_debugger: |-
You are a world-class data scientist and machine learning engineer with deep expertise in statistics, mathematics, and computer science.
You have finished the implementation of the whole workflow which has executed well on a sampled dataset. However, the user has reported that the workflow failed to execute on the full dataset.
{% include "scenarios.data_science.share:scen.role" %}
You have finished the implementation of the whole workflow which has executed well on a sampled dataset. Now we are working on the full dataset.
The user has reported that the workflow failed to execute on the full dataset.
Your will be provided with:
1. Code base.
2. Task description, which is the task the code is trying to solve.
3. Feedback generated during the execution of the whole workflow.
4. Suggestions for hyperparameter tuning.
Your job is to debug the whole code base, try to correct the errors, and ensure that the workflow can execute successfully on the full dataset.
Your current job is to debug the whole code base, try to correct the errors, and ensure that the workflow can execute successfully on the full dataset.
The user will provide your the whole code base and some feedback generated during the execution of the whole workflow. Please identify the issues and provide the corrected code.
Task description:
## Task description
{{ task_desc }}
Your modified code should follow the minimal changes principle. You should only modify the code that is necessary to fix the issues but not affect any other parts of the code. Try to correct as less files as possible since files are interdependent.
## Instructions
1. Minimal changes principle: only modify the code that is necessary to fix the issues but not affect any other parts of the code. Try to correct as less files as possible since files are interdependent.
{% if diff_mode %}
2. You must output in Code Diff format. The detailed format specification is as follows.
{% else %}
2. You must output the COMPLETE and FULL code. Do not truncate, summarize, or omit any parts of the code. Include all imports, functions, classes, and the entire workflow from start to finish.
{% endif %}
## Output Format
{% if out_spec %}
{{ out_spec }}
{% else %}
Please response the code in the following json format. Here is an example structure for the JSON output:
Please response the code in the following JSON format without anything else.
{
"code": "The Python code as a string."
}
{% endif %}
system_refine: |-
You are a world-class data scientist and machine learning engineer with deep expertise in statistics, mathematics, and computer science.
You have finished the implementation of the whole workflow which has executed well on a sampled dataset. However, the user has reported that the hyperparameters are not reasonable and the code didn't make the best use of the time limit.
Your current job is to refine the whole code base, try to refine the hyperparameters.
{% include "scenarios.data_science.share:scen.role" %}
You have finished the implementation of the whole workflow which has executed well on a sampled dataset. Now we are working on the full dataset.
The user has reported that the hyperparameters are not reasonable and the code didn't make the best use of the time limit.
Your will be provided with:
1. Code base.
2. Feedback generated during the execution of the whole workflow.
3. Suggestions for hyperparameter tuning.
Your task is to refine the code base and modify the hyperparameters based on the feedback and suggestions.
The user will provide your the whole code base, some feedback generated during the execution of the whole workflow and some suggestions for hyperparameter tuning.
Your modified code should follow the minimal changes principle. Only modify the hyperparameters that is necessary.
## Instructions
1. Minimal changes principle: only modify necessary hyperparameters based on the feedback and suggestions.
{% if diff_mode %}
2. You must output in Code Diff format. The detailed format specification is as follows.
{% else %}
2. You must output the COMPLETE and FULL code. Do not truncate, summarize, or omit any parts of the code. Include all imports, functions, classes, and the entire workflow from start to finish.
{% endif %}
## Output Format
{% if out_spec %}
{{ out_spec }}
{% else %}
Please response the code in the following json format. Here is an example structure for the JSON output:
Please response the code in the following JSON format without anything else.
{
"code": "The Python code as a string."
}
{% endif %}
user: |-
--------- code base ---------
# Code Base
{{ code }}
--------- feedback ---------
## Feedback
{{ feedback }}
{% if hyperparameter_tuning_suggestion is not none %}
--------- hyperparameter tuning suggestions ---------
## Hyperparameter Tuning Suggestion
{{ hyperparameter_tuning_suggestion }}
{% endif %}
rdagent/utils/agent/apply_patch.py
+532
View File
@@ -0,0 +1,532 @@
#!/usr/bin/env python3
# The following code is modified from https://cookbook.openai.com/examples/gpt4-1_prompting_guide
"""
A self-contained **pure-Python 3.9+** utility for applying human-readable
“pseudo-diff” patch files to a collection of text files.
"""
from __future__ import annotations
import pathlib
from collections.abc import Callable
from dataclasses import dataclass, field
from enum import Enum
# --------------------------------------------------------------------------- #
# Domain objects
# --------------------------------------------------------------------------- #
class ActionType(str, Enum):
ADD = "add"
DELETE = "delete"
UPDATE = "update"
@dataclass
class FileChange:
type: ActionType
old_content: str | None = None
new_content: str | None = None
move_path: str | None = None
@dataclass
class Commit:
changes: dict[str, FileChange] = field(default_factory=dict)
# --------------------------------------------------------------------------- #
# Exceptions
# --------------------------------------------------------------------------- #
class DiffError(ValueError):
"""Any problem detected while parsing or applying a patch."""
# --------------------------------------------------------------------------- #
# Helper dataclasses used while parsing patches
# --------------------------------------------------------------------------- #
@dataclass
class Chunk:
orig_index: int = -1
del_lines: list[str] = field(default_factory=list)
ins_lines: list[str] = field(default_factory=list)
@dataclass
class PatchAction:
type: ActionType
new_file: str | None = None
chunks: list[Chunk] = field(default_factory=list)
move_path: str | None = None
@dataclass
class Patch:
actions: dict[str, PatchAction] = field(default_factory=dict)
# --------------------------------------------------------------------------- #
# Patch text parser
# --------------------------------------------------------------------------- #
@dataclass
class Parser:
current_files: dict[str, str]
lines: list[str]
index: int = 0
patch: Patch = field(default_factory=Patch)
fuzz: int = 0
# ------------- low-level helpers -------------------------------------- #
def _cur_line(self) -> str:
if self.index >= len(self.lines):
raise DiffError("Unexpected end of input while parsing patch")
return self.lines[self.index]
@staticmethod
def _norm(line: str) -> str:
"""Strip CR so comparisons work for both LF and CRLF input."""
return line.rstrip("\r")
# ------------- scanning convenience ----------------------------------- #
def is_done(self, prefixes: tuple[str, ...] | None = None) -> bool:
if self.index >= len(self.lines):
return True
if prefixes and len(prefixes) > 0 and self._norm(self._cur_line()).startswith(prefixes):
return True
return False
def startswith(self, prefix: str | tuple[str, ...]) -> bool:
return self._norm(self._cur_line()).startswith(prefix)
def read_str(self, prefix: str) -> str:
"""
Consume the current line if it starts with *prefix* and return the text
**after** the prefix. Raises if prefix is empty.
"""
if prefix == "":
raise ValueError("read_str() requires a non-empty prefix")
if self._norm(self._cur_line()).startswith(prefix):
text = self._cur_line()[len(prefix) :]
self.index += 1
return text
return ""
def read_line(self) -> str:
"""Return the current raw line and advance."""
line = self._cur_line()
self.index += 1
return line
# ------------- public entry point -------------------------------------- #
def parse(self) -> None:
while not self.is_done(("*** End Patch",)):
# ---------- UPDATE ---------- #
path = self.read_str("*** Update File: ")
if path:
if path in self.patch.actions:
raise DiffError(f"Duplicate update for file: {path}")
move_to = self.read_str("*** Move to: ")
if path not in self.current_files:
raise DiffError(f"Update File Error - missing file: {path}")
text = self.current_files[path]
action = self._parse_update_file(text)
action.move_path = move_to or None
self.patch.actions[path] = action
continue
# ---------- DELETE ---------- #
path = self.read_str("*** Delete File: ")
if path:
if path in self.patch.actions:
raise DiffError(f"Duplicate delete for file: {path}")
if path not in self.current_files:
raise DiffError(f"Delete File Error - missing file: {path}")
self.patch.actions[path] = PatchAction(type=ActionType.DELETE)
continue
# ---------- ADD ---------- #
path = self.read_str("*** Add File: ")
if path:
if path in self.patch.actions:
raise DiffError(f"Duplicate add for file: {path}")
if path in self.current_files:
raise DiffError(f"Add File Error - file already exists: {path}")
self.patch.actions[path] = self._parse_add_file()
continue
raise DiffError(f"Unknown line while parsing: {self._cur_line()}")
if not self.startswith("*** End Patch"):
raise DiffError("Missing *** End Patch sentinel")
self.index += 1 # consume sentinel
# ------------- section parsers ---------------------------------------- #
def _parse_update_file(self, text: str) -> PatchAction:
action = PatchAction(type=ActionType.UPDATE)
lines = text.split("\n")
index = 0
while not self.is_done(
(
"*** End Patch",
"*** Update File:",
"*** Delete File:",
"*** Add File:",
"*** End of File",
),
):
def_str = self.read_str("@@ ")
section_str = ""
if not def_str and self._norm(self._cur_line()) == "@@":
section_str = self.read_line()
if not (def_str or section_str or index == 0):
raise DiffError(f"Invalid line in update section:\n{self._cur_line()}")
if def_str.strip():
found = False
if def_str not in lines[:index]:
for i, s in enumerate(lines[index:], index):
if s == def_str:
index = i + 1
found = True
break
if not found and def_str.strip() not in [s.strip() for s in lines[:index]]:
for i, s in enumerate(lines[index:], index):
if s.strip() == def_str.strip():
index = i + 1
self.fuzz += 1
found = True
break
next_ctx, chunks, end_idx, eof = peek_next_section(self.lines, self.index)
new_index, fuzz = find_context(lines, next_ctx, index, eof)
if new_index == -1:
ctx_txt = "\n".join(next_ctx)
raise DiffError(
f"Invalid {'EOF ' if eof else ''}context at {index}:\n{ctx_txt}",
)
self.fuzz += fuzz
for ch in chunks:
ch.orig_index += new_index
action.chunks.append(ch)
index = new_index + len(next_ctx)
self.index = end_idx
return action
def _parse_add_file(self) -> PatchAction:
lines: list[str] = []
while not self.is_done(
("*** End Patch", "*** Update File:", "*** Delete File:", "*** Add File:"),
):
s = self.read_line()
if not s.startswith("+"):
raise DiffError(f"Invalid Add File line (missing '+'): {s}")
lines.append(s[1:]) # strip leading '+'
return PatchAction(type=ActionType.ADD, new_file="\n".join(lines))
# --------------------------------------------------------------------------- #
# Helper functions
# --------------------------------------------------------------------------- #
def find_context_core(
lines: list[str],
context: list[str],
start: int,
) -> tuple[int, int]:
if not context:
return start, 0
for i in range(start, len(lines)):
if lines[i : i + len(context)] == context:
return i, 0
for i in range(start, len(lines)):
if [s.rstrip() for s in lines[i : i + len(context)]] == [s.rstrip() for s in context]:
return i, 1
for i in range(start, len(lines)):
if [s.strip() for s in lines[i : i + len(context)]] == [s.strip() for s in context]:
return i, 100
return -1, 0
def find_context(
lines: list[str],
context: list[str],
start: int,
eof: bool,
) -> tuple[int, int]:
if eof:
new_index, fuzz = find_context_core(lines, context, len(lines) - len(context))
if new_index != -1:
return new_index, fuzz
new_index, fuzz = find_context_core(lines, context, start)
return new_index, fuzz + 10_000
return find_context_core(lines, context, start)
def peek_next_section(
lines: list[str],
index: int,
) -> tuple[list[str], list[Chunk], int, bool]:
old: list[str] = []
del_lines: list[str] = []
ins_lines: list[str] = []
chunks: list[Chunk] = []
mode = "keep"
orig_index = index
while index < len(lines):
s = lines[index]
if s.startswith(
(
"@@",
"*** End Patch",
"*** Update File:",
"*** Delete File:",
"*** Add File:",
"*** End of File",
),
):
break
if s == "***":
break
if s.startswith("***"):
raise DiffError(f"Invalid Line: {s}")
index += 1
last_mode = mode
if s == "":
s = " "
if s[0] == "+":
mode = "add"
elif s[0] == "-":
mode = "delete"
elif s[0] == " ":
mode = "keep"
else:
raise DiffError(f"Invalid Line: {s}")
s = s[1:]
if mode == "keep" and last_mode != mode:
if ins_lines or del_lines:
chunks.append(
Chunk(
orig_index=len(old) - len(del_lines),
del_lines=del_lines,
ins_lines=ins_lines,
),
)
del_lines, ins_lines = [], []
if mode == "delete":
del_lines.append(s)
old.append(s)
elif mode == "add":
ins_lines.append(s)
elif mode == "keep":
old.append(s)
if ins_lines or del_lines:
chunks.append(
Chunk(
orig_index=len(old) - len(del_lines),
del_lines=del_lines,
ins_lines=ins_lines,
),
)
if index < len(lines) and lines[index] == "*** End of File":
index += 1
return old, chunks, index, True
if index == orig_index:
raise DiffError("Nothing in this section")
return old, chunks, index, False
# --------------------------------------------------------------------------- #
# Patch → Commit and Commit application
# --------------------------------------------------------------------------- #
def _get_updated_file(text: str, action: PatchAction, path: str) -> str:
if action.type is not ActionType.UPDATE:
raise DiffError("_get_updated_file called with non-update action")
orig_lines = text.split("\n")
dest_lines: list[str] = []
orig_index = 0
for chunk in action.chunks:
if chunk.orig_index > len(orig_lines):
raise DiffError(
f"{path}: chunk.orig_index {chunk.orig_index} exceeds file length",
)
if orig_index > chunk.orig_index:
raise DiffError(
f"{path}: overlapping chunks at {orig_index} > {chunk.orig_index}",
)
dest_lines.extend(orig_lines[orig_index : chunk.orig_index])
orig_index = chunk.orig_index
dest_lines.extend(chunk.ins_lines)
orig_index += len(chunk.del_lines)
dest_lines.extend(orig_lines[orig_index:])
return "\n".join(dest_lines)
def patch_to_commit(patch: Patch, orig: dict[str, str]) -> Commit:
commit = Commit()
for path, action in patch.actions.items():
if action.type is ActionType.DELETE:
commit.changes[path] = FileChange(
type=ActionType.DELETE,
old_content=orig[path],
)
elif action.type is ActionType.ADD:
if action.new_file is None:
raise DiffError("ADD action without file content")
commit.changes[path] = FileChange(
type=ActionType.ADD,
new_content=action.new_file,
)
elif action.type is ActionType.UPDATE:
new_content = _get_updated_file(orig[path], action, path)
commit.changes[path] = FileChange(
type=ActionType.UPDATE,
old_content=orig[path],
new_content=new_content,
move_path=action.move_path,
)
return commit
# --------------------------------------------------------------------------- #
# User-facing helpers
# --------------------------------------------------------------------------- #
def text_to_patch(text: str, orig: dict[str, str]) -> tuple[Patch, int]:
lines = text.splitlines() # preserves blank lines, no strip()
if (
len(lines) < 2
or not Parser._norm(lines[0]).startswith("*** Begin Patch")
or Parser._norm(lines[-1]) != "*** End Patch"
):
raise DiffError("Invalid patch text - missing sentinels")
parser = Parser(current_files=orig, lines=lines, index=1)
parser.parse()
return parser.patch, parser.fuzz
def identify_files_needed(text: str) -> list[str]:
lines = text.splitlines()
return [line[len("*** Update File: ") :] for line in lines if line.startswith("*** Update File: ")] + [
line[len("*** Delete File: ") :] for line in lines if line.startswith("*** Delete File: ")
]
def identify_files_added(text: str) -> list[str]:
lines = text.splitlines()
return [line[len("*** Add File: ") :] for line in lines if line.startswith("*** Add File: ")]
# --------------------------------------------------------------------------- #
# File-system helpers
# --------------------------------------------------------------------------- #
def load_files(paths: list[str], open_fn: Callable[[str], str]) -> dict[str, str]:
return {path: open_fn(path) for path in paths}
def apply_commit(
commit: Commit,
write_fn: Callable[[str, str], None],
remove_fn: Callable[[str], None],
inplace: bool = False,
) -> None | dict:
batch_edit = {}
for path, change in commit.changes.items():
if change.type is ActionType.DELETE:
remove_fn(path)
elif change.type is ActionType.ADD:
if change.new_content is None:
raise DiffError(f"ADD change for {path} has no content")
write_fn(path, change.new_content)
elif change.type is ActionType.UPDATE:
if change.new_content is None:
raise DiffError(f"UPDATE change for {path} has no new content")
if inplace:
target = change.move_path or path
write_fn(target, change.new_content)
if change.move_path:
remove_fn(path)
batch_edit[path] = change.new_content
return batch_edit
def process_patch(
text: str,
open_fn: Callable[[str], str],
write_fn: Callable[[str, str], None],
remove_fn: Callable[[str], None],
inplace: bool = False,
) -> str:
if not text.startswith("*** Begin Patch"):
raise DiffError("Patch text must start with *** Begin Patch")
paths = identify_files_needed(text)
orig = load_files(paths, open_fn)
patch, _fuzz = text_to_patch(text, orig)
commit = patch_to_commit(patch, orig)
batch_edit = apply_commit(commit, write_fn, remove_fn, inplace)
return batch_edit
# --------------------------------------------------------------------------- #
# Default FS helpers
# --------------------------------------------------------------------------- #
def open_file(path: str) -> str:
with open(path, encoding="utf-8") as fh:
return fh.read()
def write_file(path: str, content: str) -> None:
target = pathlib.Path(path)
target.parent.mkdir(parents=True, exist_ok=True)
with target.open("wt", encoding="utf-8") as fh:
fh.write(content)
def remove_file(path: str) -> None:
pathlib.Path(path).unlink(missing_ok=True)
# --------------------------------------------------------------------------- #
# CLI entry-point
# --------------------------------------------------------------------------- #
def apply_patch_from_text(patch_text: str, inplace: bool = False) -> str:
"""Apply patch text to filesystem, same as main() but with parameter input"""
if not patch_text:
raise DiffError("Patch text cannot be empty")
try:
result = process_patch(patch_text, open_file, write_file, remove_file, inplace)
return result
except DiffError as exc:
raise exc
def main() -> None:
import sys
patch_text = sys.stdin.read()
if not patch_text:
print("Please pass patch text through stdin", file=sys.stderr)
return
try:
result = process_patch(patch_text, open_file, write_file, remove_file)
except DiffError as exc:
print(exc, file=sys.stderr)
return
print(result)
if __name__ == "__main__":
main()
rdagent/utils/agent/ret.py
+19 -1
View File
@@ -9,6 +9,7 @@ import re
from abc import abstractclassmethod
from typing import Any
from rdagent.utils.agent.apply_patch import apply_patch_from_text
from rdagent.utils.agent.tpl import T
@@ -17,7 +18,7 @@ class AgentOut:
@abstractclassmethod
def get_spec(cls, **context: Any) -> str:
raise NotImplementedError(f"Please implement the `get_spec` method")
raise NotImplementedError("Please implement the `get_spec` method")
@classmethod
def extract_output(cls, resp: str) -> Any:
@@ -82,3 +83,20 @@ class PythonBatchEditOut(AgentOut):
code_blocks[file_name.strip()] = code.strip()
return code_blocks
class PythonBatchPatchOut(AgentOut):
@classmethod
def get_spec(cls):
return T(".tpl:PythonBatchPatchOut").r()
@classmethod
def extract_output(cls, resp: str) -> str:
# Step 1: extract patch by pattern
patch_pattern = re.compile(r"(\*\*\* Begin Patch\s*(.*?)\s*\*\*\* End Patch)", re.DOTALL)
match = patch_pattern.search(resp)
if match:
resp = match.group(1).rstrip()
# Step 2: apply the patch, this will modify the file in place
return apply_patch_from_text(resp, inplace=False)
rdagent/utils/agent/tpl.py
+2 -2
View File
@@ -52,8 +52,8 @@ def load_content(uri: str, caller_dir: Path | None = None, ftype: str = "yaml")
for file_path in file_path_l:
try:
if ftype == "yaml":
with file_path.open() as file:
# Load the YAML file
# Parse the UTF-8 encoded YAML configuration for cross-platform compatibility
with file_path.open(encoding="utf-8") as file:
yaml_content = yaml.safe_load(file)
# Traverse the YAML content to get the desired template
for key in yaml_trace:
rdagent/utils/agent/tpl.yaml
+58
View File
@@ -53,3 +53,61 @@ PythonBatchEditOut: |-
- To explicitly remove a file, provide only `__DEL__` within the code block for that file.
- To replace a file with a new one, first provide ` __DEL__` for the original file, then include a separate entry with new file name and the new code.
{% endif %}
# The following prompt is modified from https://cookbook.openai.com/examples/gpt4-1_prompting_guide
PythonBatchPatchOut: |-
This is a custom utility that makes it more convenient to add, remove, move, or edit code files. `apply_patch` effectively allows you to execute a diff/patch against a file, but the format of the diff specification is unique to this task, so pay careful attention to these instructions. To use the `apply_patch` command, you should pass a message of the following structure as "input":
%%bash
apply_patch <<"EOF"
*** Begin Patch
[YOUR_PATCH]
*** End Patch
EOF
Where [YOUR_PATCH] is the actual content of your patch, specified in the following V4A diff format.
*** [ACTION] File: [path/to/file] -> ACTION can be one of Add, Update, or Delete.
For each snippet of code that needs to be changed, repeat the following:
[context_before] -> See below for further instructions on context.
- [old_code] -> Precede the old code with a minus sign.
+ [new_code] -> Precede the new, replacement code with a plus sign.
[context_after] -> See below for further instructions on context.
For instructions on [context_before] and [context_after]:
- By default, show 3 lines of code immediately above and 3 lines immediately below each change. If a change is within 3 lines of a previous change, do NOT duplicate the first change’s [context_after] lines in the second change’s [context_before] lines.
- If 3 lines of context is insufficient to uniquely identify the snippet of code within the file, use the @@ operator to indicate the class or function to which the snippet belongs. For instance, we might have:
@@ class BaseClass
[3 lines of pre-context]
- [old_code]
+ [new_code]
[3 lines of post-context]
- If a code block is repeated so many times in a class or function such that even a single @@ statement and 3 lines of context cannot uniquely identify the snippet of code, you can use multiple `@@` statements to jump to the right context. For instance:
@@ class BaseClass
@@ def method():
[3 lines of pre-context]
- [old_code]
+ [new_code]
[3 lines of post-context]
Note, then, that we do not use line numbers in this diff format, as the context is enough to uniquely identify code. An example of a message that you might pass as "input" to this function, in order to apply a patch, is shown below.
%%bash
apply_patch <<"EOF"
*** Begin Patch
*** Update File: pygorithm/searching/binary_search.py
@@ class BaseClass
@@ def search():
- pass
+ raise NotImplementedError()
@@ class Subclass
@@ def search():
- pass
+ raise NotImplementedError()
*** End Patch
EOF