update agent ablation code

Author: nanjiangwill

.gitignore

@@ -162,9 +162,10 @@ cython_debug/
 #.idea/
 
 logs/
-repos/
+repos*/
 config.yml
 hydra_outputs/
 .commit0*
 .agent*
-docs/analysis*.md
+docs/analysis*.md
+agent/run_agent_no_rich.py

agent/agent_utils.py

@@ -9,16 +9,22 @@
 from import_deps import ModuleSet
 from graphlib import TopologicalSorter, CycleError
 import yaml
-
+from rank_bm25 import BM25Okapi
 from agent.class_types import AgentConfig
+import subprocess
 
 PROMPT_HEADER = ">>> Here is the Task:\n"
+FUNCTION_HEADER = "\n\n>>> Here are all functions in the file, complete the implementations for all functions (i.e., those with pass statements):\n"
 REFERENCE_HEADER = "\n\n>>> Here is the Reference for you to finish the task:\n"
 REPO_INFO_HEADER = "\n\n>>> Here is the Repository Information:\n"
 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"
+FUNCTION_BY_FUNCTION_HEADER = """"\nYour task is to implement function {unimplemented_functions} by replacing the pass statement with actual functional code.
+Please note that there could be multiple occurrences of {unimplemented_functions}, and you need to implement them all.
+Do not change the names of existing functions or classes, as they may be referenced from other code like unit tests, etc.
+When 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."""
 # prefix components:
 space = "    "
 branch = "│   "
@@ -123,6 +129,32 @@ def get_file_info(file_path: Path, prefix: str = "") -> str:
     return "\n".join(filter(None, tree_string))
 
 
+def get_unimplemented_functions(file_path: Path) -> List[str]:
+    """Get all the functions in a file."""
+    with open(file_path, "r") as f:
+        content = f.read()
+
+    # Find all function definitions with their bodies
+    pattern = r"def\s+(\w+)\s*\([^)]*\)[^:]*:(?:\s*(?:'''[\s\S]*?'''|\"\"\"[\s\S]*?\"\"\"))?\s*((?:(?!\ndef\s+).)*?)(?=\s*def\s+|\s*$)"
+    matches = re.finditer(pattern, content, re.DOTALL)
+
+    # Keep only functions that have just 'pass'
+    # List to store unimplemented function definitions
+    unimplemented_functions = []
+    for match in matches:
+        func_name = match.group(1)
+        func_body = match.group(2).strip()
+        # Check if function only contains 'pass' statement
+        if "pass" in func_body:
+            unimplemented_functions.append(f"def {func_name}()")
+            # # Find the full function definition using regex pattern
+            # func_pattern = rf"def\s+{func_name}\s*\([^)]*\)[^:]*:"
+            # func_match = re.search(func_pattern, content)
+            # if func_match:
+            #     unimplemented.append(func_match.group(0))
+    return unimplemented_functions
+
+
 def collect_test_files(directory: str) -> list[str]:
     """Collect all the test files in the directory."""
     test_files = []
@@ -265,9 +297,9 @@ def get_target_edit_files(
             raise ValueError(
                 "topological_sort_files should not be longer than filtered_files"
             )
-    assert len(topological_sort_files) == len(
-        filtered_files
-    ), "all files should be included"
+    assert len(topological_sort_files) == len(filtered_files), (
+        "all files should be included"
+    )
 
     # change to latest commit
     local_repo.git.checkout(branch)
@@ -324,9 +356,9 @@ def get_target_edit_files_from_patch(
                 raise ValueError(
                     "topological_sort_files should not be longer than target_files_list"
                 )
-        assert len(topological_sort_files) == len(
-            target_files_list
-        ), "all files should be included"
+        assert len(topological_sort_files) == len(target_files_list), (
+            "all files should be included"
+        )
 
         topological_sort_files = [
             file.replace(working_dir, "").lstrip("/") for file in topological_sort_files
@@ -347,6 +379,7 @@ def get_message(
     agent_config: AgentConfig,
     repo_path: str,
     test_files: list[str] | None = None,
+    input_file: str | None = None,
 ) -> str:
     """Get the message to Aider."""
     prompt = f"{PROMPT_HEADER}" + agent_config.user_prompt
@@ -383,11 +416,11 @@ def get_message(
         with bz2.open("spec.pdf.bz2", "rb") as in_file:
             with open("spec.pdf", "wb") as out_file:
                 out_file.write(in_file.read())
-        spec_info = (
-            f"\n{SPEC_INFO_HEADER} "
-            + get_specification(specification_pdf_path=Path(repo_path, "spec.pdf"))[
-                : agent_config.max_spec_info_length
-            ]
+        spec_info = f"\n{SPEC_INFO_HEADER} " + get_specification(
+            specification_pdf_path=Path(repo_path, "spec.pdf"),
+            use_retrieval=True,
+            query=input_file if input_file else "",
+            top_k=10,
         )
     else:
         spec_info = ""
@@ -397,6 +430,39 @@ def get_message(
     return message_to_agent
 
 
+def get_message_function_by_function(
+    agent_config: AgentConfig,
+    repo_path: str,
+    input_file: str,
+    test_files: list[str] | None = None,
+) -> list[str]:
+    """Get the message to Aider."""
+    context = get_message(agent_config, repo_path, test_files)
+
+    if agent_config.implementation_strategy == "module_by_module":
+        function_info = []
+    elif agent_config.implementation_strategy == "function_by_function":
+        function_info = []
+        unimplemented_functions = get_unimplemented_functions(
+            file_path=Path(os.path.join(repo_path, input_file))
+        )
+        # Get the original function stubs and filter out implemented functions
+        for i in range(len(unimplemented_functions)):
+            function_info.append(
+                FUNCTION_BY_FUNCTION_HEADER.format(
+                    unimplemented_functions=unimplemented_functions[i]
+                )
+            )
+    else:
+        raise ValueError(
+            f"Invalid implementation strategy: {agent_config.implementation_strategy}"
+        )
+
+    messages_to_agent = [context + uf for uf in unimplemented_functions]
+
+    return messages_to_agent
+
+
 def update_message_with_dependencies(message: str, dependencies: list[str]) -> str:
     """Update the message with the dependencies."""
     if len(dependencies) == 0:
@@ -411,19 +477,43 @@ def update_message_with_dependencies(message: str, dependencies: list[str]) -> s
     return message
 
 
-def get_specification(specification_pdf_path: Path) -> str:
+def get_specification(
+    specification_pdf_path: Path,
+    use_retrieval: bool = True,
+    query: str = "",
+    top_k: int = 20,
+) -> 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
     # Open the specified PDF file
+
     document = fitz.open(specification_pdf_path)
-    text = ""
+    corpus = []
 
+    # current_trunk = ""
     # Iterate through the pages
     for page_num in range(len(document)):
         page = document.load_page(page_num)  # loads the specified page
-        text += page.get_text()  # type: ignore
 
-    return text
+        current_page_text = page.get_text()  # type: ignore
+        # Cut page text into chunks of 1000 characters
+        text_chunks = [
+            current_page_text[i : i + 1000]
+            for i in range(0, len(current_page_text), 1000)
+        ]
+        corpus.extend(text_chunks)
+        # corpus.append(page.get_text())  # type: ignore
+    if not use_retrieval:
+        return "\n".join(corpus)
+
+    assert query != "", "query should not be empty"
+    query = open(query).read()
+    tokenized_corpus = [doc.split(" ") for doc in corpus]
+    bm25 = BM25Okapi(tokenized_corpus)
+    doc_scores = bm25.get_scores(query)
+    sorted_doc_scores = sorted(enumerate(doc_scores), key=lambda x: x[1], reverse=True)
+    sorted_doc_indices = [i for i, _ in sorted_doc_scores]
+    return "\n".join(corpus[i] for i in sorted_doc_indices[:top_k])
 
 
 def create_branch(repo: git.Repo, branch: str, from_commit: str) -> None:
@@ -486,6 +576,21 @@ def get_changed_files_from_commits(
         return []
 
 
+def run_eval_after_each_commit(
+    branch: str, backend: str, commit0_config_file: str, repo_name: str
+) -> str:
+    """Run the eval command after each commit."""
+    eval_cmd = f"python -m commit0 evaluate --branch {branch} --backend {backend} --commit0-config-file {commit0_config_file} --timeout 100"
+    try:
+        result = subprocess.run(
+            eval_cmd, shell=True, capture_output=True, text=True, check=True
+        )
+        return result.stdout
+    except subprocess.CalledProcessError as e:
+        print(f"Error running eval command: {e}")
+        return e.stdout if e.stdout else str(e)
+
+
 def args2string(agent_config: AgentConfig) -> str:
     """Converts specific fields from an `AgentConfig` object into a formatted string.
 

agent/agents.py

@@ -27,6 +27,8 @@ def __init__(self, log_file: Path):
         self.log_file = log_file
 
         self.last_cost = 0.0
+        self.total_token_in = 0
+        self.total_token_out = 0
 
 
 class Agents(ABC):
@@ -43,6 +45,8 @@ class AiderReturn(AgentReturn):
     def __init__(self, log_file: Path):
         super().__init__(log_file)
         self.last_cost = self.get_money_cost()
+        self.total_token_in = self.get_total_token_in()
+        self.total_token_out = self.get_total_token_out()
 
     def get_money_cost(self) -> float:
         """Get accumulated money cost from log file"""
@@ -57,18 +61,54 @@ def get_money_cost(self) -> float:
                         last_cost = float(match.group(1))
         return last_cost
 
+    def get_total_token_in(self) -> int:
+        """Get total token in from log file"""
+        total_tokens = 0
+        with open(self.log_file, "r") as file:
+            for line in file:
+                if "Tokens:" in line:
+                    match = re.search(r"Tokens: ([\d.]+k?) sent", line)
+                    if match:
+                        token_str = match.group(1)
+                        if token_str.endswith("k"):
+                            total_tokens = int(float(token_str[:-1]) * 1000)
+                        else:
+                            total_tokens = int(float(token_str))
+        return total_tokens
+
+    def get_total_token_out(self) -> int:
+        """Get total token out from log file"""
+        total_tokens = 0
+        with open(self.log_file, "r") as file:
+            for line in file:
+                if "Tokens:" in line:
+                    match = re.search(r"(\d+) received", line)
+                    if match:
+                        total_str = match.group(1)
+                        if total_str.endswith("k"):
+                            total_tokens = int(float(total_str[:-1]) * 1000)
+                        else:
+                            total_tokens = int(float(total_str))
+        return total_tokens
+
 
 class AiderAgents(Agents):
     def __init__(self, max_iteration: int, model_name: str):
         super().__init__(max_iteration)
         self.model = Model(model_name)
         # Check if API key is set for the model
-        if "gpt" in model_name:
+        if "openrouter" in model_name:
+            api_key = os.environ.get("OPENROUTER_API_KEY", None)
+        elif "gpt" in model_name:
             api_key = os.environ.get("OPENAI_API_KEY", None)
         elif "claude" in model_name:
             api_key = os.environ.get("ANTHROPIC_API_KEY", None)
         elif "gemini" in model_name:
-            api_key = os.environ.get("API_KEY", None)
+            api_key = os.environ.get("GEMINI_API_KEY", None)
+        elif "deepseek" in model_name:
+            api_key = os.environ.get("DEEPSEEK_API_KEY", None)
+        elif "mistral" in model_name:
+            api_key = os.environ.get("MISTRAL_API_KEY", None)
         else:
             raise ValueError(f"Unsupported model: {model_name}")
 
@@ -87,6 +127,7 @@ def run(
         log_dir: Path,
         test_first: bool = False,
         lint_first: bool = False,
+        current_attempt: int = 0,
     ) -> AgentReturn:
         """Start aider agent"""
         if test_cmd:
@@ -99,11 +140,22 @@ def run(
             auto_lint = False
         log_dir = log_dir.resolve()
         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"
-
+        input_history_file = (
+            log_dir / ".aider.input.history"
+            if current_attempt == 0
+            else log_dir / f".aider_{current_attempt}.input.history"
+        )
+        chat_history_file = (
+            log_dir / ".aider.chat.history.md"
+            if current_attempt == 0
+            else log_dir / f".aider_{current_attempt}.chat.history.md"
+        )
         # Set up logging
-        log_file = log_dir / "aider.log"
+        log_file = (
+            log_dir / "aider.log"
+            if current_attempt == 0
+            else log_dir / f"aider_{current_attempt}.log"
+        )
         logging.basicConfig(
             filename=log_file,
             level=logging.INFO,
@@ -133,7 +185,7 @@ def run(
             io=io,
         )
         coder.max_reflections = self.max_iteration
-        coder.stream = True
+        coder.stream = False
 
         # Run the agent
         if test_first:

agent/class_types.py

@@ -22,3 +22,5 @@ class AgentConfig:
     run_tests: bool
     max_iteration: int
     record_test_for_each_commit: bool
+    implementation_strategy: str
+    repeat_times_for_each_inquiry: int

agent/cli.py

@@ -135,6 +135,14 @@ def config(
         False,
         help="Run the lint on the entire directory",
     ),
+    implementation_strategy: str = typer.Option(
+        "module_by_module",
+        help="Implementation strategy to use",
+    ),
+    repeat_times_for_each_inquiry: int = typer.Option(
+        1,
+        help="Repeat times for each inquiry",
+    ),
     record_test_for_each_commit: bool = typer.Option(
         False,
         help="Record the test for each commit",
@@ -173,6 +181,8 @@ def config(
         "use_lint_info": use_lint_info,
         "max_lint_info_length": max_lint_info_length,
         "run_entire_dir_lint": run_entire_dir_lint,
+        "implementation_strategy": implementation_strategy,
+        "repeat_times_for_each_inquiry": repeat_times_for_each_inquiry,
         "pre_commit_config_path": pre_commit_config_path,
         "record_test_for_each_commit": record_test_for_each_commit,
     }