diff --git a/rdagent/scenarios/data_science/dev/runner/__init__.py b/rdagent/scenarios/data_science/dev/runner/__init__.py index 5c1820a6..d5dd500e 100644 --- a/rdagent/scenarios/data_science/dev/runner/__init__.py +++ b/rdagent/scenarios/data_science/dev/runner/__init__.py @@ -83,12 +83,21 @@ class DSRunnerMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy): hyperparameter_tuning_suggestion=prev_task_feedback.hyperparameter_tuning_suggestion, ) - batch_edit = extract_output_fn( - APIBackend().build_messages_and_create_chat_completion( - user_prompt=user_prompt, - system_prompt=system_prompt, + if self.settings.diff_mode: + batch_edit = extract_output_fn( + APIBackend().build_messages_and_create_chat_completion( + user_prompt=user_prompt, + system_prompt=system_prompt, + ), + prefix=workspace.workspace_path, + ) + else: + batch_edit = extract_output_fn( + APIBackend().build_messages_and_create_chat_completion( + user_prompt=user_prompt, + system_prompt=system_prompt, + ) ) - ) batch_edit = {k: v for k, v in batch_edit.items() if k in workspace.file_dict.keys()} diff --git a/rdagent/utils/agent/apply_patch.py b/rdagent/utils/agent/apply_patch.py index 3e7ab52d..edf2c27a 100644 --- a/rdagent/utils/agent/apply_patch.py +++ b/rdagent/utils/agent/apply_patch.py @@ -12,6 +12,7 @@ import pathlib from collections.abc import Callable from dataclasses import dataclass, field from enum import Enum +from pathlib import Path # --------------------------------------------------------------------------- # @@ -76,6 +77,7 @@ class Parser: index: int = 0 patch: Patch = field(default_factory=Patch) fuzz: int = 0 + prefix: Path | None = None # ------------- low-level helpers -------------------------------------- # def _cur_line(self) -> str: @@ -123,6 +125,8 @@ class Parser: while not self.is_done(("*** End Patch",)): # ---------- UPDATE ---------- # path = self.read_str("*** Update File: ") + if self.prefix: + path = str(self.prefix / path) if path: if path in self.patch.actions: raise DiffError(f"Duplicate update for file: {path}") @@ -137,6 +141,8 @@ class Parser: # ---------- DELETE ---------- # path = self.read_str("*** Delete File: ") + if self.prefix: + path = str(self.prefix / path) if path: if path in self.patch.actions: raise DiffError(f"Duplicate delete for file: {path}") @@ -147,6 +153,8 @@ class Parser: # ---------- ADD ---------- # path = self.read_str("*** Add File: ") + if self.prefix: + path = str(self.prefix / path) if path: if path in self.patch.actions: raise DiffError(f"Duplicate add for file: {path}") @@ -403,7 +411,7 @@ def patch_to_commit(patch: Patch, orig: dict[str, str]) -> Commit: # --------------------------------------------------------------------------- # # User-facing helpers # --------------------------------------------------------------------------- # -def text_to_patch(text: str, orig: dict[str, str]) -> tuple[Patch, int]: +def text_to_patch(text: str, orig: dict[str, str], prefix: Path | None = None) -> tuple[Patch, int]: lines = text.splitlines() # preserves blank lines, no strip() if ( len(lines) < 2 @@ -412,21 +420,31 @@ def text_to_patch(text: str, orig: dict[str, str]) -> tuple[Patch, int]: ): raise DiffError("Invalid patch text - missing sentinels") - parser = Parser(current_files=orig, lines=lines, index=1) + parser = Parser(current_files=orig, lines=lines, index=1, prefix=prefix) parser.parse() return parser.patch, parser.fuzz -def identify_files_needed(text: str) -> list[str]: +def identify_files_needed(text: str, prefix: Path | None = None) -> 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: ") - ] + update_files = [line[len("*** Update File: ") :] for line in lines if line.startswith("*** Update File: ")] + delete_files = [line[len("*** Delete File: ") :] for line in lines if line.startswith("*** Delete File: ")] + all_files = update_files + delete_files + + if prefix is None: + return all_files + else: + return [str(prefix / file) for file in all_files] -def identify_files_added(text: str) -> list[str]: +def identify_files_added(text: str, prefix: Path | None = None) -> list[str]: lines = text.splitlines() - return [line[len("*** Add File: ") :] for line in lines if line.startswith("*** Add File: ")] + added_files = [line[len("*** Add File: ") :] for line in lines if line.startswith("*** Add File: ")] + + if prefix is None: + return added_files + else: + return [str(prefix / file) for file in added_files] # --------------------------------------------------------------------------- # @@ -468,12 +486,13 @@ def process_patch( write_fn: Callable[[str, str], None], remove_fn: Callable[[str], None], inplace: bool = False, + prefix: Path | None = None, ) -> str: if not text.startswith("*** Begin Patch"): raise DiffError("Patch text must start with *** Begin Patch") - paths = identify_files_needed(text) + paths = identify_files_needed(text, prefix) orig = load_files(paths, open_fn) - patch, _fuzz = text_to_patch(text, orig) + patch, _fuzz = text_to_patch(text, orig, prefix) commit = patch_to_commit(patch, orig) batch_edit = apply_commit(commit, write_fn, remove_fn, inplace) return batch_edit @@ -501,13 +520,13 @@ def remove_file(path: str) -> None: # --------------------------------------------------------------------------- # # CLI entry-point # --------------------------------------------------------------------------- # -def apply_patch_from_text(patch_text: str, inplace: bool = False) -> str: +def apply_patch_from_text(patch_text: str, inplace: bool = False, prefix: Path | None = None) -> 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) + result = process_patch(patch_text, open_file, write_file, remove_file, inplace, prefix) return result except DiffError as exc: raise exc diff --git a/rdagent/utils/agent/ret.py b/rdagent/utils/agent/ret.py index ba698fd3..85b143ac 100644 --- a/rdagent/utils/agent/ret.py +++ b/rdagent/utils/agent/ret.py @@ -7,6 +7,7 @@ We think this part can be shared. import json import re from abc import abstractclassmethod +from pathlib import Path from typing import Any from rdagent.utils.agent.apply_patch import apply_patch_from_text @@ -91,12 +92,13 @@ class PythonBatchPatchOut(AgentOut): return T(".tpl:PythonBatchPatchOut").r() @classmethod - def extract_output(cls, resp: str) -> str: + def extract_output(cls, resp: str, prefix: Path | None = None) -> str: + code_blocks = {} # 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() + matches = patch_pattern.findall(resp) + for match in matches: + code_blocks.update(apply_patch_from_text(match[0], inplace=False, prefix=prefix)) # Step 2: apply the patch, this will modify the file in place - return apply_patch_from_text(resp, inplace=False) + return code_blocks