Merge pull request #33 from commit-0/aider_reorg

Aider reorg

Author: wenting-zhao

baselines/agents.py

@@ -0,0 +1,54 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+from aider.coders import Coder
+from aider.models import Model
+from aider.io import InputOutput
+
+
+class Agents(ABC):
+    @abstractmethod
+    def run(self) -> None:
+        """Start agent"""
+        raise NotImplementedError
+
+
+class AiderAgents(Agents):
+    def __init__(self, model_name: str):
+        self.model = Model(model_name)
+
+    def run(
+        self,
+        message: str,
+        test_cmd: str,
+        lint_cmd: str,
+        fnames: list[str],
+        log_dir: Path,
+    ) -> None:
+        """Start aider agent"""
+        if test_cmd:
+            auto_test = True
+        else:
+            auto_test = False
+        if lint_cmd:
+            auto_lint = True
+        else:
+            auto_lint = False
+        log_dir.mkdir(parents=True, exist_ok=True)
+        input_history_file = log_dir / ".aider.input.history"
+        chat_history_file = log_dir / ".aider.chat.history.md"
+        io = InputOutput(
+            yes=True,
+            input_history_file=input_history_file,
+            chat_history_file=chat_history_file,
+        )
+        coder = Coder.create(
+            main_model=self.model,
+            fnames=fnames,
+            auto_lint=auto_lint,
+            auto_test=auto_test,
+            lint_cmds=lint_cmd,
+            test_cmd=test_cmd,
+            io=io,
+        )
+        coder.run(message)

baselines/class_types.py

@@ -11,8 +11,9 @@ class Commit0Config:
 
 
 @dataclass
-class AiderConfig:
-    llm_name: str
+class AgentConfig:
+    agent_name: str
+    model_name: str
     use_user_prompt: bool
     user_prompt: str
     use_repo_info: bool

baselines/baseline_utils.py renamed to baselines/commit0_utils.py

@@ -2,9 +2,9 @@
 import re
 import subprocess
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import List
 
-from baselines.class_types import AiderConfig
+from baselines.class_types import AgentConfig
 
 PROMPT_HEADER = ">>> Here is the Task:\n"
 REFERENCE_HEADER = "\n\n>>> Here is the Reference for you to finish the task:\n"
@@ -116,7 +116,7 @@ def get_file_info(file_path: Path, prefix: str = "") -> str:
     return "\n".join(filter(None, tree_string))
 
 
-def get_target_edit_files_cmd_args(target_dir: str) -> str:
+def get_target_edit_files(target_dir: str) -> list[str]:
     """Find the files with the error 'NotImplementedError('IMPLEMENT ME
     HERE')'.
     """
@@ -135,54 +135,43 @@ def get_target_edit_files_cmd_args(target_dir: str) -> str:
     # Only keep python files
     files = [file for file in files if file.endswith(".py")]
 
-    return " ".join(files)
+    return files
 
 
-def get_message_to_aider(
-    aider_config: AiderConfig,
-    target_edit_files_cmd_args: str,
+def get_message(
+    agent_config: AgentConfig,
     repo_path: str,
-    ds: Dict[str, Any],
+    test_dir: str,
 ) -> str:
     """Get the message to Aider."""
-    prompt = f"{PROMPT_HEADER} " + aider_config.user_prompt
+    prompt = f"{PROMPT_HEADER}" + agent_config.user_prompt
 
-    if aider_config.use_unit_tests_info and ds["test"]["test_dir"]:
+    if agent_config.use_unit_tests_info and test_dir:
         unit_tests_info = (
             f"\n{UNIT_TESTS_INFO_HEADER} "
             + get_dir_info(
-                dir_path=Path(os.path.join(repo_path, ds["test"]["test_dir"])),
+                dir_path=Path(os.path.join(repo_path, test_dir)),
                 prefix="",
                 include_stubs=True,
-            )[: aider_config.max_unit_tests_info_length]
+            )[: agent_config.max_unit_tests_info_length]
         )
     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 = (
-            f"\n{REFERENCE_HEADER} "
-            + get_reference(ds["specification"])[
-                : aider_config.max_reference_info_length
-            ]
-        )
-    else:
-        reference = ""
-
-    if aider_config.use_repo_info:
+    if agent_config.use_repo_info:
         repo_info = (
             f"\n{REPO_INFO_HEADER} "
             + get_dir_info(
                 dir_path=Path(repo_path), prefix="", max_depth=2, include_stubs=False
-            )[: aider_config.max_repo_info_length]
+            )[: agent_config.max_repo_info_length]
         )
     else:
         repo_info = ""
 
-    message_to_aider = prompt + reference + repo_info + unit_tests_info
+    message_to_agent = prompt + repo_info + unit_tests_info
 
-    return message_to_aider
+    return message_to_agent
 
 
 def get_reference(specification_pdf_path: str) -> str:

baselines/configs/aider.yaml renamed to baselines/configs/agent.yaml

@@ -3,11 +3,14 @@ defaults:
   - base
   - _self_
 
-aider_config:
+commit0_config:
+  repo_split: minitorch
+
+agent_config:
   use_user_prompt: false
   use_repo_info: false
   use_unit_tests_info: false
   use_reference_info: false
-  use_lint_info: true
+  use_lint_info: false
   pre_commit_config_path: .pre-commit-config.yaml
-  run_tests: true
+  run_tests: false

baselines/configs/base.yaml

@@ -10,10 +10,11 @@ commit0_config:
   repo_split: "simpy"
   num_workers: 10
 
-aider_config:
-  llm_name: "claude-3-5-sonnet-20240620"
+agent_config:
+  agent_name: "aider"
+  model_name: "claude-3-5-sonnet-20240620"
   use_user_prompt: false
-  user_prompt: "Here is the Task:\n Your task is to iteratively implement the each function that is 'NotImplementedError('IMPLEMENT ME HERE')' in these files until there are no more 'NotImplementedError('IMPLEMENT ME HERE')' and pass the unit tests.\nMake sure you read the files carefully.\nYour output should be the edited code files.\nUse the above instructions to modify the supplied files.\nDo not change the names of existing functions or classes, as they may be referenced from other code like unit tests, etc.\nOnly use standard python libraries, do not suggest installing any packages."
+  user_prompt: "Here is your task:\nYou need to implement all functions with 'NotImplementedError('IMPLEMENT ME HERE')' and pass the unit tests.\nDo not change the names of existing functions or classes, as they may be referenced from other code like unit tests, etc.\nWhen you generate code, you must maintain the original formatting of the function stubs (such as whitespaces), otherwise we will not able to search/replace blocks for code modifications, and therefore you will receive a score of 0 for your generated code."
   use_repo_info: false
   use_unit_tests_info: false
   use_reference_info: false
@@ -27,4 +28,4 @@ aider_config:
 
 hydra:
   run:
-    dir: ./hydra_outputs/${now:%Y-%m-%d}/${now:%H-%M-%S}
+    dir: ./hydra_outputs/${now:%Y-%m-%d}/${now:%H-%M-%S}

baselines/run_agent.py

@@ -0,0 +1,156 @@
+import os
+import hydra
+from datasets import load_dataset
+import traceback
+from baselines.commit0_utils import (
+    get_message,
+    get_target_edit_files,
+)
+from baselines.agents import AiderAgents
+from typing import Optional, Type
+from types import TracebackType
+from hydra.core.config_store import ConfigStore
+from baselines.class_types import AgentConfig, Commit0Config
+from commit0.harness.constants import SPLIT
+from commit0.harness.get_pytest_ids import main as get_tests
+from commit0.harness.constants import RUN_AIDER_LOG_DIR, RepoInstance
+from tqdm import tqdm
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+
+class DirContext:
+    def __init__(self, d: str):
+        self.dir = d
+        self.cwd = os.getcwd()
+
+    def __enter__(self):
+        os.chdir(self.dir)
+
+    def __exit__(
+        self,
+        exctype: Optional[Type[BaseException]],
+        excinst: Optional[BaseException],
+        exctb: Optional[TracebackType],
+    ) -> None:
+        os.chdir(self.cwd)
+
+
+def run_agent_for_repo(
+    commit0_config: Commit0Config,
+    agent_config: AgentConfig,
+    example: RepoInstance,
+) -> None:
+    """Run Aider for a given repository."""
+    # get repo info
+    _, repo_name = example["repo"].split("/")
+
+    repo_name = repo_name.lower()
+    repo_name = repo_name.replace(".", "-")
+
+    # Call the commit0 get-tests command to retrieve test files
+    test_files_str = get_tests(repo_name, stdout=False)
+    test_files = sorted(list(set([i.split(":")[0] for i in test_files_str])))
+
+    repo_path = os.path.join(commit0_config.base_dir, repo_name)
+    repo_path = os.path.abspath(repo_path)
+
+    target_edit_files = get_target_edit_files(repo_path)
+
+    if agent_config.agent_name == "aider":
+        agent = AiderAgents(agent_config.model_name)
+    else:
+        raise NotImplementedError(
+            f"{agent_config.agent_name} is not implemented; please add your implementations in baselines/agents.py."
+        )
+
+    with DirContext(repo_path):
+        if commit0_config is None or agent_config is None:
+            raise ValueError("Invalid input")
+
+        message = get_message(agent_config, repo_path, example["test"]["test_dir"])
+
+        if agent_config.use_lint_info:
+            lint_cmd = "pre-commit run --config ../../.pre-commit-config.yaml --files"
+        else:
+            lint_cmd = ""
+
+        if agent_config.run_tests:
+            # when unit test feedback is available, iterate over test files
+            for test_file in test_files:
+                test_cmd = f"python -m commit0 test {repo_path} {test_file}"
+                test_file_name = test_file.replace(".py", "").replace("/", "__")
+                log_dir = RUN_AIDER_LOG_DIR / "with_tests" / test_file_name
+
+                agent.run(
+                    message,
+                    test_cmd,
+                    lint_cmd,
+                    target_edit_files,
+                    log_dir,
+                )
+        else:
+            # when unit test feedback is not available, iterate over target files to edit
+            for f in target_edit_files:
+                file_name = f.replace(".py", "").replace("/", "__")
+                log_dir = RUN_AIDER_LOG_DIR / "no_tests" / file_name
+
+                agent.run(message, "", lint_cmd, [f], log_dir)
+
+
+def main() -> None:
+    """Main function to run Aider for a given repository.
+
+    Will run in parallel for each repo.
+    """
+    cs = ConfigStore.instance()
+    cs.store(name="user", node=Commit0Config)
+    cs.store(name="user", node=AgentConfig)
+    hydra.initialize(version_base=None, config_path="configs")
+    config = hydra.compose(config_name="agent")
+    commit0_config = Commit0Config(**config.commit0_config)
+    agent_config = AgentConfig(**config.agent_config)
+
+    dataset = load_dataset(
+        commit0_config.dataset_name, split=commit0_config.dataset_split
+    )
+    filtered_dataset = [
+        example
+        for example in dataset
+        if commit0_config.repo_split == "all"
+        or (
+            isinstance(example, dict)
+            and "repo" in example
+            and isinstance(example["repo"], str)
+            and example["repo"].split("/")[-1]
+            in SPLIT.get(commit0_config.repo_split, [])
+        )
+    ]
+    assert len(filtered_dataset) > 0, "No examples available"
+
+    with tqdm(
+        total=len(filtered_dataset), smoothing=0, desc="Running Aider for repos"
+    ) as pbar:
+        with ThreadPoolExecutor(max_workers=commit0_config.num_workers) as executor:
+            # Create a future for running Aider for each repo
+            futures = {
+                executor.submit(
+                    run_agent_for_repo,
+                    commit0_config,
+                    agent_config,
+                    example,  # type: ignore
+                ): example
+                for example in filtered_dataset
+            }
+            # Wait for each future to complete
+            for future in as_completed(futures):
+                pbar.update(1)
+                try:
+                    # Update progress bar, check if Aider ran successfully
+                    future.result()
+                except Exception:
+                    traceback.print_exc()
+                    continue
+
+
+if __name__ == "__main__":
+    main()