new file

Author: MaxWhitton25

.gitignore

@@ -167,4 +167,5 @@ config.yml
 hydra_outputs/
 .commit0*
 .agent*
-docs/analysis*.md
+docs/analysis*.md
+.aider*

BRANCH

@@ -57,6 +57,101 @@ def get_money_cost(self) -> float:
                         last_cost = float(match.group(1))
         return last_cost
 
+class AgentTeams(Agents):
+    def __init__(self, max_iteration: int, model_name: str):
+        super().__init__(max_iteration)
+        self.model = Model(model_name)
+    def run(
+        self,
+        message: str,
+        test_cmd: str,
+        lint_cmd: str,
+        fnames: list[str],
+        log_dir: Path,
+        test_first: bool = False,
+    ) -> AgentReturn:
+        """Start agent team"""
+        if test_cmd:
+            auto_test = True
+        else:
+            auto_test = False
+        if lint_cmd:
+            auto_lint = True
+        else:
+            auto_lint = False
+        log_dir = log_dir.resolve()
+        log_dir.mkdir(parents=True, exist_ok=True)
+        input_history_file = log_dir / ".team.input.history"
+        chat_history_file = log_dir / ".team.chat.history.md"
+
+        # Set up logging
+        log_file = log_dir / "team.log"
+        logging.basicConfig(
+            filename=log_file,
+            level=logging.INFO,
+            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+        )
+
+        # Redirect print statements to the log file
+        sys.stdout = open(log_file, "a")
+        sys.stderr = open(log_file, "a")
+
+        # Log the message
+        agent_message_log_file = log_dir / "agent_message.log"
+        with open(agent_message_log_file, "a") as f:
+            f.write(f"Message Sent: {message}\n\n")
+
+        # Configure httpx and backoff logging
+        handle_logging("httpx", log_file)
+        handle_logging("backoff", log_file)
+
+        io = InputOutput(
+            yes=True,
+            input_history_file=input_history_file,
+            chat_history_file=chat_history_file,
+        )
+        manager = Coder.create(
+            main_model=self.model,
+            fnames=fnames,
+            auto_lint=auto_lint,
+            auto_test=auto_test,
+            lint_cmds={"python": lint_cmd},
+            test_cmd=test_cmd,
+            io=io,
+        )
+        manager.max_reflection = 1
+        manager.stream = True
+
+        # Run the agent
+        manager_message = "First, add every file in the repo to your conversation. Second, write a plan of attack to implement the entire repo and return this plan."
+        manager.run(manager_message)
+
+
+        coder = Coder.create(
+            main_model=self.model,
+            fnames=fnames,
+            auto_lint=auto_lint,
+            auto_test=auto_test,
+            lint_cmds={"python": lint_cmd},
+            test_cmd=test_cmd,
+            io=io,
+        )
+        coder.max_reflection = self.max_iteration
+        coder.stream = True
+
+        # Run the agent
+        with open(chat_history_file, 'r', encoding='utf-8') as file:
+            plan = file.read()
+        coder_message = "follow this implementation plan: "+plan
+        coder.run(coder_message)
+
+        sys.stdout.close()
+        sys.stderr.close()
+        # Restore original stdout and stderr
+        sys.stdout = sys.__stdout__
+        sys.stderr = sys.__stderr__
+
+        return AgentReturn(log_file)
 
 class AiderAgents(Agents):
     def __init__(self, max_iteration: int, model_name: str):

agent/agents.py

@@ -1,4 +1,5 @@
 import os
+from typing import Any
 import yaml
 import multiprocessing
 from datasets import load_dataset
@@ -14,6 +15,7 @@
 )
 import subprocess
 from agent.agents import AiderAgents
+from agent.agents import AgentTeams
 from typing import Optional, Type, cast
 from types import TracebackType
 from agent.class_types import AgentConfig
@@ -44,7 +46,6 @@ def __exit__(
     ) -> None:
         os.chdir(self.cwd)
 
-
 def run_agent_for_repo(
     repo_base_dir: str,
     agent_config: AgentConfig,
@@ -56,9 +57,53 @@ def run_agent_for_repo(
     log_dir: str = str(RUN_AGENT_LOG_DIR.resolve()),
     commit0_config_file: str = "",
 ) -> None:
-    """Run Aider for a given repository."""
-    # get repo info
+    
     commit0_config = read_commit0_dot_file(commit0_config_file)
+    if agent_config.agent_name == "aider":
+        run_aider_for_repo(
+                            commit0_config,
+                            commit0_config["base_dir"],
+                            agent_config,
+                            cast(RepoInstance, example),
+                            update_queue,
+                            branch,
+                            override_previous_changes,
+                            backend,
+                            log_dir,
+                            commit0_config_file,
+                        )
+    elif agent_config.agent_name == "aider_team":
+        run_team_for_repo(
+                            commit0_config,
+                            commit0_config["base_dir"],
+                            agent_config,
+                            cast(RepoInstance, example),
+                            update_queue,
+                            branch,
+                            override_previous_changes,
+                            backend,
+                            log_dir,
+                            commit0_config_file,
+                        )
+    else:
+        raise NotImplementedError(
+            f"{agent_config.agent_name} is not implemented; please add your implementations in baselines/agents.py."
+        )
+        
+def run_aider_for_repo(
+    commit0_config: Any,
+    repo_base_dir: str,
+    agent_config: AgentConfig,
+    example: RepoInstance,
+    update_queue: multiprocessing.Queue,
+    branch: str,
+    override_previous_changes: bool = False,
+    backend: str = "modal",
+    log_dir: str = str(RUN_AGENT_LOG_DIR.resolve()),
+    commit0_config_file: str = "",
+) -> None:
+    
+    """Run Aider for a given repository."""
 
     assert "commit0" in commit0_config["dataset_name"]
     _, repo_name = example["repo"].split("/")
@@ -79,12 +124,7 @@ def run_agent_for_repo(
             f"{repo_path} is not a git repo. Check if base_dir is correctly specified."
         )
 
-    if agent_config.agent_name == "aider":
-        agent = AiderAgents(agent_config.max_iteration, agent_config.model_name)
-    else:
-        raise NotImplementedError(
-            f"{agent_config.agent_name} is not implemented; please add your implementations in baselines/agents.py."
-        )
+    agent = AiderAgents(agent_config.max_iteration, agent_config.model_name)
 
     # Check if there are changes in the current branch
     if local_repo.is_dirty():
@@ -219,6 +259,146 @@ def run_agent_for_repo(
                 )
     update_queue.put(("finish_repo", repo_name))
 
+def run_team_for_repo(
+    commit0_config: Any,
+    repo_base_dir: str,
+    agent_config: AgentConfig,
+    example: RepoInstance,
+    update_queue: multiprocessing.Queue,
+    branch: str,
+    override_previous_changes: bool = False,
+    backend: str = "modal",
+    log_dir: str = str(RUN_AGENT_LOG_DIR.resolve()),
+    commit0_config_file: str = "",
+) -> None:
+    """Run Aider Team for a given repository."""
+    # get repo info
+    assert "commit0" in commit0_config["dataset_name"]
+    _, repo_name = example["repo"].split("/")
+
+    # before starting, display all information to terminal
+    update_queue.put(("start_repo", (repo_name, 0)))
+
+    # repo_name = repo_name.lower()
+    # repo_name = repo_name.replace(".", "-")
+
+    repo_path = os.path.join(repo_base_dir, repo_name)
+    repo_path = os.path.abspath(repo_path)
+
+    try:
+        local_repo = Repo(repo_path)
+    except Exception:
+        raise Exception(
+            f"{repo_path} is not a git repo. Check if base_dir is correctly specified."
+        )
+
+    manager = AiderAgents(1, agent_config.model_name)
+    coder = AiderAgents(agent_config.max_iteration, agent_config.model_name)
+
+     # Check if there are changes in the current branch
+    if local_repo.is_dirty():
+        # Stage all changes
+        local_repo.git.add(A=True)
+        # Commit changes with the message "left from last change"
+        local_repo.index.commit("left from last change")
+
+    # # if branch_name is not provided, create a new branch name based on agent_config
+    # if branch is None:
+    #     branch = args2string(agent_config)
+    create_branch(local_repo, branch, example["base_commit"])
+
+    # in cases where the latest commit of branch is not commit 0
+    # set it back to commit 0
+    latest_commit = local_repo.commit(branch)
+    if latest_commit.hexsha != example["base_commit"] and override_previous_changes:
+        local_repo.git.reset("--hard", example["base_commit"])
+
+    # get target files to edit and test files to run
+    target_edit_files, import_dependencies = get_target_edit_files(
+        local_repo,
+        example["src_dir"],
+        example["test"]["test_dir"],
+        branch,
+        example["reference_commit"],
+        agent_config.use_topo_sort_dependencies,
+    )
+
+    lint_files = get_changed_files_from_commits(
+        local_repo, "HEAD", example["base_commit"]
+    )
+    # Call the commit0 get-tests command to retrieve test files
+    test_files_str = get_tests(repo_name, verbose=0)
+    test_files = sorted(list(set([i.split(":")[0] for i in test_files_str])))
+
+    # prepare the log dir
+    experiment_log_dir = (
+        Path(log_dir)
+        / repo_name
+        / branch
+        / datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    )
+    experiment_log_dir.mkdir(parents=True, exist_ok=True)
+
+    # write agent_config to .agent.yaml in the log_dir for record
+    agent_config_log_file = experiment_log_dir / ".agent.yaml"
+    with open(agent_config_log_file, "w") as agent_config_file:
+        yaml.dump(agent_config, agent_config_file)
+
+    manager_message = "Write a concise plan of attack to implement the entire repo, but don't actually do any coding. The plan should not include any reccommendations to add files and should be a maximum of 500 words."
+    
+    with DirContext(repo_path):
+        if agent_config is None:
+            raise ValueError("Invalid input")
+        else:
+            # when unit test feedback is not available, iterate over target files to edit
+
+            update_queue.put(
+                ("start_repo", (repo_name, len(target_edit_files)))
+            )
+            
+            #TODO: add support for unit test / lint feedback
+            
+            for f in target_edit_files:
+                update_queue.put(("set_current_file", (repo_name, f)))
+                dependencies = import_dependencies[f]
+            file_name = "all"
+            file_log_dir = experiment_log_dir / file_name
+            lint_cmd = get_lint_cmd(repo_name, agent_config.use_lint_info)
+                
+            
+            agent_return = manager.run(manager_message, "", lint_cmd, target_edit_files, file_log_dir)
+            with open(agent_return.log_file, 'r', encoding='utf-8') as file:
+                plan = file.read()
+            coder_message = "follow this implementation plan: "+plan
+
+            agent_return = coder.run(coder_message, "", lint_cmd, target_edit_files, file_log_dir)
+            
+            # for f in target_edit_files:
+            #     update_queue.put(("set_current_file", (repo_name, f)))
+            #     dependencies = import_dependencies[f]
+            #     message = update_message_with_dependencies(coder_message, dependencies)
+            #     file_name = f.replace(".py", "").replace("/", "__")
+            #     file_log_dir = experiment_log_dir / file_name
+            #     lint_cmd = get_lint_cmd(repo_name, agent_config.use_lint_info)
+            #     agent_return = coder.run(message, "", lint_cmd, [f], file_log_dir)
+            #     update_queue.put(
+            #         (
+            #             "update_money_display",
+            #             (repo_name, file_name, agent_return.last_cost),
+            #         )
+            #     )
+            update_queue.put(
+                    (
+                        "update_money_display",
+                        (repo_name, file_name, agent_return.last_cost),
+                    )
+                )
+    
+    
+    
+    update_queue.put(("finish_repo", repo_name))
+
+
 
 def run_agent(
     branch: str,
@@ -303,6 +483,7 @@ def run_agent(
                     result = pool.apply_async(
                         run_agent_for_repo,
                         args=(
+                            commit0_config,
                             commit0_config["base_dir"],
                             agent_config,
                             cast(RepoInstance, example),