update file finding logic

Author: nanjiangwill

agent/agent_utils.py

@@ -18,6 +18,7 @@
 UNIT_TESTS_INFO_HEADER = "\n\n>>> Here are the Unit Tests Information:\n"
 LINT_INFO_HEADER = "\n\n>>> Here is the Lint Information:\n"
 SPEC_INFO_HEADER = "\n\n>>> Here is the Specification Information:\n"
+IMPORT_DEPENDENCIES_HEADER = "\n\n>>> Here are the Import Dependencies:\n"
 # prefix components:
 space = "    "
 branch = "│   "
@@ -208,7 +209,9 @@ def ignore_cycles(graph: dict) -> list[str]:
         return ignore_cycles(graph)
 
 
-def topological_sort_based_on_dependencies(pkg_paths: list[str]) -> list[str]:
+def topological_sort_based_on_dependencies(
+    pkg_paths: list[str],
+) -> tuple[list[str], dict]:
     """Topological sort based on dependencies."""
     module_set = ModuleSet([str(p) for p in pkg_paths])
 
@@ -224,7 +227,7 @@ def topological_sort_based_on_dependencies(pkg_paths: list[str]) -> list[str]:
 
     import_dependencies_files = ignore_cycles(import_dependencies)
 
-    return import_dependencies_files
+    return import_dependencies_files, import_dependencies
 
 
 def get_target_edit_files(
@@ -233,7 +236,7 @@ def get_target_edit_files(
     test_dir: str,
     latest_commit: str,
     reference_commit: str,
-) -> list[str]:
+) -> tuple[list[str], dict]:
     """Find the files with functions with the pass statement."""
     target_dir = str(local_repo.working_dir)
     files = _find_files_to_edit(target_dir, src_dir, test_dir)
@@ -248,7 +251,9 @@ def get_target_edit_files(
     # Change to reference commit to get the correct dependencies
     local_repo.git.checkout(reference_commit)
 
-    topological_sort_files = topological_sort_based_on_dependencies(filtered_files)
+    topological_sort_files, import_dependencies = (
+        topological_sort_based_on_dependencies(filtered_files)
+    )
     if len(topological_sort_files) != len(filtered_files):
         if len(topological_sort_files) < len(filtered_files):
             # Find the missing elements
@@ -271,7 +276,14 @@ def get_target_edit_files(
         file.replace(target_dir, "").lstrip("/") for file in topological_sort_files
     ]
 
-    return topological_sort_files
+    # Remove the base_dir prefix from import dependencies
+    import_dependencies_without_prefix = {}
+    for key, value in import_dependencies.items():
+        key_without_prefix = key.replace(target_dir, "").lstrip("/")
+        value_without_prefix = [v.replace(target_dir, "").lstrip("/") for v in value]
+        import_dependencies_without_prefix[key_without_prefix] = value_without_prefix
+
+    return topological_sort_files, import_dependencies_without_prefix
 
 
 def get_message(
@@ -331,6 +343,20 @@ def get_message(
     return message_to_agent
 
 
+def update_message_with_dependencies(message: str, dependencies: list[str]) -> str:
+    """Update the message with the dependencies."""
+    if len(dependencies) == 0:
+        return message
+    import_dependencies_info = f"\n{IMPORT_DEPENDENCIES_HEADER}"
+    for dependency in dependencies:
+        with open(dependency, "r") as file:
+            import_dependencies_info += (
+                f"\nHere is the content of the file {dependency}:\n{file.read()}"
+            )
+    message += import_dependencies_info
+    return message
+
+
 def get_specification(specification_pdf_path: Path) -> str:
     """Get the reference for a given specification PDF path."""
     # TODO: after pdf_to_text is available, use it to extract the text from the PDF

agent/cli.py

@@ -178,6 +178,10 @@ def run(
         ".agent.yaml",
         help="Path to the agent config file",
     ),
+    commit0_config_file: str = typer.Option(
+        ".commit0.yaml",
+        help="Path to the commit0 config file",
+    ),
     log_dir: str = typer.Option(
         str(RUN_AGENT_LOG_DIR.resolve()),
         help="Log directory to store the logs",
@@ -202,6 +206,7 @@ def run(
             override_previous_changes,
             backend,
             agent_config_file,
+            commit0_config_file,
             log_dir,
             max_parallel_repos,
             display_repo_progress_num,
@@ -212,6 +217,7 @@ def run(
             override_previous_changes,
             backend,
             agent_config_file,
+            commit0_config_file,
             log_dir,
             max_parallel_repos,
         )

agent/run_agent.py

@@ -7,6 +7,7 @@
     create_branch,
     get_message,
     get_target_edit_files,
+    update_message_with_dependencies,
     get_lint_cmd,
     read_yaml_config,
 )
@@ -91,7 +92,7 @@ def run_agent_for_repo(
     if latest_commit.hexsha != example["base_commit"] and override_previous_changes:
         local_repo.git.reset("--hard", example["base_commit"])
 
-    target_edit_files = get_target_edit_files(
+    target_edit_files, import_dependencies = get_target_edit_files(
         local_repo,
         example["src_dir"],
         example["test"]["test_dir"],
@@ -161,6 +162,8 @@ def run_agent_for_repo(
             )
             for f in target_edit_files:
                 update_queue.put(("set_current_file", (repo_name, f)))
+                dependencies = import_dependencies[f]
+                message = update_message_with_dependencies(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)
@@ -179,6 +182,7 @@ def run_agent(
     override_previous_changes: bool,
     backend: str,
     agent_config_file: str,
+    commit0_config_file: str,
     log_dir: str,
     max_parallel_repos: int,
     display_repo_progress_num: int,
@@ -188,7 +192,7 @@ def run_agent(
 
     agent_config = AgentConfig(**config)
 
-    commit0_config = read_commit0_dot_file(".commit0.yaml")
+    commit0_config = read_commit0_dot_file(commit0_config_file)
 
     dataset = load_dataset(
         commit0_config["dataset_name"], split=commit0_config["dataset_split"]

agent/run_agent_no_rich.py

@@ -9,6 +9,7 @@
     create_branch,
     get_message,
     get_target_edit_files,
+    update_message_with_dependencies,
     get_lint_cmd,
     read_yaml_config,
 )
@@ -88,7 +89,7 @@ def run_agent_for_repo(
         local_repo.git.reset("--hard", example["base_commit"])
 
     # get target files to edit and test files to run
-    target_edit_files = get_target_edit_files(
+    target_edit_files, import_dependencies = get_target_edit_files(
         local_repo,
         example["src_dir"],
         example["test"]["test_dir"],
@@ -144,6 +145,8 @@ def run_agent_for_repo(
                 agent_config, repo_path, test_dir=example["test"]["test_dir"]
             )
             for f in target_edit_files:
+                dependencies = import_dependencies[f]
+                message = update_message_with_dependencies(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)
@@ -156,6 +159,7 @@ def run_agent(
     override_previous_changes: bool,
     backend: str,
     agent_config_file: str,
+    commit0_config_file: str,
     log_dir: str,
     max_parallel_repos: int,
 ) -> None:
@@ -167,7 +171,7 @@ def run_agent(
 
     agent_config = AgentConfig(**config)
 
-    commit0_config = read_commit0_dot_file(".commit0.yaml")
+    commit0_config = read_commit0_dot_file(commit0_config_file)
 
     dataset = load_dataset(
         commit0_config["dataset_name"], split=commit0_config["dataset_split"]