edit branch name

Author: nanjiangwill

agent/cli.py

@@ -163,7 +163,7 @@ def config(
 @agent_app.command()
 def run(
     branch: str = typer.Argument(
-        ...,
+        None,
         help="Branch name of current run",
     ),
     override_previous_changes: bool = typer.Option(

agent/display.py

@@ -135,7 +135,8 @@ def __init__(self, total_repos: int):
         )
         self.layout["info"]["other_info"].split_row(
             Layout(name="backend", ratio=1),
-            Layout(name="log_dir", ratio=1),
+            Layout(name="branch", ratio=2),
+            Layout(name="log_dir", ratio=2),
         )
         self.layout["info"]["agent_info"].split_row(
             Layout(name="agent_name", ratio=1),
@@ -151,6 +152,10 @@ def __init__(self, total_repos: int):
         self.layout["info"]["other_info"]["backend"].update(
             Panel(self.backend_display, title="Backend", border_style="blue")
         )
+        self.branch_display = Text("", justify="center")
+        self.layout["info"]["other_info"]["branch"].update(
+            Panel(self.branch_display, title="Branch", border_style="blue")
+        )
         self.log_dir_display = Text("", justify="center")
         self.layout["info"]["other_info"]["log_dir"].update(
             Panel(self.log_dir_display, title="Log Directory", border_style="blue")
@@ -227,9 +232,16 @@ def update_time_display(self, time_in_seconds: int) -> None:
             Panel(self.time_display, title="Time Taken", border_style="blue")
         )
 
+    def update_branch_display(self, branch: str) -> None:
+        """Update the branch display with the given branch."""
+        self.branch_display = Text(f"{branch}", justify="center")
+        self.layout["info"]["other_info"]["branch"].update(
+            Panel(self.branch_display, title="Branch", border_style="blue")
+        )
+
     def update_backend_display(self, backend: str) -> None:
         """Update the backend display with the given backend."""
-        self.backend_display = Text(f"Backend Using: {backend}", justify="center")
+        self.backend_display = Text(f"{backend}", justify="center")
         self.layout["info"]["other_info"]["backend"].update(
             Panel(self.backend_display, title="Backend", border_style="blue")
         )

agent/run_agent.py

@@ -48,7 +48,7 @@ def run_agent_for_repo(
     agent_config: AgentConfig,
     example: RepoInstance,
     update_queue: multiprocessing.Queue,
-    branch: Optional[str] = None,
+    branch: str,
     override_previous_changes: bool = False,
     backend: str = "modal",
     log_dir: str = str(RUN_AGENT_LOG_DIR.resolve()),
@@ -88,10 +88,6 @@ def run_agent_for_repo(
             f"{agent_config.agent_name} is not implemented; please add your implementations in baselines/agents.py."
         )
 
-    # 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
@@ -173,7 +169,7 @@ def run_agent_for_repo(
 
 
 def run_agent(
-    branch: str,
+    branch: Optional[str],
     override_previous_changes: bool,
     backend: str,
     agent_config_file: str,
@@ -208,6 +204,10 @@ def run_agent(
     # if len(filtered_dataset) > 1:
     #     sys.stdout = open(os.devnull, "w")
 
+    # if branch_name is not provided, create a new branch name based on agent_config
+    if branch is None:
+        branch = args2string(agent_config)
+
     with TerminalDisplay(len(filtered_dataset)) as display:
         not_started_repos = [
             cast(RepoInstance, example)["repo"].split("/")[-1]
@@ -231,6 +231,7 @@ def run_agent(
             agent_config.use_spec_info,
             agent_config.use_lint_info,
         )
+        display.update_branch_display(branch)
         with multiprocessing.Manager() as manager:
             update_queue = manager.Queue()
             with multiprocessing.Pool(processes=max_parallel_repos) as pool: