aider update message to aider logic

Author: nanjiangwill

baselines/baseline_utils.py

@@ -125,7 +125,9 @@ def get_prompt(file_list: str) -> str:
 
 
 def find_files_with_error(target_dir: str) -> List[str]:
-    """Find the files with the error 'NotImplementedError('IMPLEMENT ME HERE')'."""
+    """Find the files with the error 'NotImplementedError('IMPLEMENT ME
+    HERE')'.
+    """
     # The grep command
     command = f"grep -R -l \"NotImplementedError('IMPLEMENT ME HERE')\" {target_dir}"
 

baselines/run_aider.py

@@ -5,6 +5,10 @@
 from pathlib import Path
 
 import hydra
+from datasets import load_dataset
+from omegaconf import OmegaConf
+from tqdm.contrib.concurrent import thread_map
+
 from baselines.baseline_utils import (
     PROMPT_HEADER,
     REFERENCE_HEADER,
@@ -16,9 +20,6 @@
     get_reference,
 )
 from baselines.class_types import AiderConfig, BaselineConfig, Commit0Config
-from datasets import load_dataset
-from omegaconf import OmegaConf
-from tqdm.contrib.concurrent import thread_map
 
 # from aider.run_aider import get_aider_cmd
 
@@ -30,23 +31,12 @@
 
 def get_aider_cmd(
     model: str,
-    file: str,
-    prompt: str,
-    reference: str,
-    repo_info: str,
-    unit_tests_info: str,
+    files: str,
+    message_to_aider: str,
     test_cmd: str,
 ) -> str:
     """Get the Aider command based on the given context."""
-    message = f"{PROMPT_HEADER} {prompt}"
-    if reference:
-        message += f"\n{REFERENCE_HEADER} {reference}"
-    if repo_info:
-        message += f"\n{REPO_INFO_HEADER} {repo_info}"
-    if unit_tests_info:
-        message += f"\n{UNIT_TESTS_INFO_HEADER} {unit_tests_info}"
-
-    aider_cmd = f"aider --model {model} --file {file} --message \"{message}\" --auto-test --test --test-cmd '{test_cmd}' --yes"
+    aider_cmd = f"aider --model {model} --file {files} --message \"{message_to_aider}\" --auto-test --test --test-cmd '{test_cmd}' --yes"
 
     return aider_cmd
 
@@ -57,54 +47,54 @@ def run_aider_for_repo(
     ds: dict,
 ) -> None:
     """Run Aider for a given repository."""
-    if commit0_config is None or aider_config is None or ds is None:
+    if commit0_config is None or aider_config is None:
         raise ValueError("Invalid input")
 
     # get repo info
     _, repo_name = ds["repo"].split("/")
 
-    # TODO: assuming we have all test_files
+    # TODO: assuming we have all test_files, which we currently do not have
     test_files = ds["test_files"]
 
     repo_path = os.path.join(commit0_config.base_dir, repo_name)
     target_edit_files = find_files_with_error(repo_path)
 
     target_edit_files_cmd_args = " ".join(target_edit_files)
 
-    prompt = get_prompt(target_edit_files_cmd_args)
+    prompt = f"{PROMPT_HEADER} " + get_prompt(target_edit_files_cmd_args)
 
     # support context for aider
     if aider_config.use_unit_tests_info and ds["test"]["test_dir"]:
-        unit_tests_info = get_dir_info(
+        unit_tests_info = f"\n{UNIT_TESTS_INFO_HEADER} " + get_dir_info(
             dir_path=Path(os.path.join(repo_path, ds["test"]["test_dir"])),
             prefix="",
             include_stubs=True,
         )
     else:
         unit_tests_info = ""
 
+    # TODO: assuming we have specification, which we currently do not have
     if aider_config.use_reference_info and ds["specification"]:
-        reference = get_reference(ds["specification"])
+        reference = f"\n{REFERENCE_HEADER} " + get_reference(ds["specification"])
     else:
         reference = ""
 
     if aider_config.use_repo_info:
-        repo_info = get_dir_info(
+        repo_info = f"\n{REPO_INFO_HEADER} " + get_dir_info(
             dir_path=Path(repo_path), prefix="", max_depth=2, include_stubs=False
         )
     else:
         repo_info = ""
 
+    message_to_aider = prompt + reference + repo_info + unit_tests_info
+
     for test_file in test_files:
         test_cmd = f"python -m commit0.harness.run_pytest_ids --repo {repo_name} --test_ids {test_file} --branch_name aider"
 
         aider_cmd = get_aider_cmd(
             aider_config.llm_name,
             target_edit_files_cmd_args,
-            prompt,
-            reference,
-            repo_info,
-            unit_tests_info,
+            message_to_aider,
             test_cmd,
         )
 
@@ -128,7 +118,10 @@ def run_aider_for_repo(
 
 @hydra.main(version_base=None, config_path="config", config_name="aider")
 def main(config: BaselineConfig) -> None:
-    """Main function to run Aider for a given repository. Will run in parallel for each repo."""
+    """Main function to run Aider for a given repository.
+
+    Will run in parallel for each repo.
+    """
     config = BaselineConfig(config=OmegaConf.to_object(config))
     commit0_config = config.commit0_config
     aider_config = config.aider_config