agent with rich display

Author: nanjiangwill

agent/commit0_utils.py renamed to agent/agent_utils.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 from typing import List
 import fitz
+import yaml
 
 from agent.class_types import AgentConfig
 
@@ -118,7 +119,7 @@ def get_file_info(file_path: Path, prefix: str = "") -> str:
     return "\n".join(filter(None, tree_string))
 
 
-def get_target_edit_files(target_dir: str) -> list[str]:
+def get_target_edit_files(target_dir: str, src_prefix: str) -> list[str]:
     """Find the files with functions with the pass statement."""
     files = []
     for root, _, filenames in os.walk(target_dir):
@@ -131,7 +132,7 @@ def get_target_edit_files(target_dir: str) -> list[str]:
 
     # Remove the base_dir prefix
     files = [file.replace(target_dir, "").lstrip("/") for file in files]
-
+    files = [src_prefix + file for file in files]
     # Only keep python files
     files = [file for file in files if file.endswith(".py")]
 
@@ -308,3 +309,17 @@ def get_lint_cmd(repo_name: str, use_lint_info: bool) -> str:
     else:
         lint_cmd = ""
     return lint_cmd
+
+
+def write_agent_config(agent_config_file: str, agent_config: dict) -> None:
+    """Write the agent config to the file."""
+    with open(agent_config_file, "w") as f:
+        yaml.dump(agent_config, f)
+
+
+def read_yaml_config(config_file: str) -> dict:
+    """Read the yaml config from the file."""
+    if not os.path.exists(config_file):
+        raise FileNotFoundError(f"The config file '{config_file}' does not exist.")
+    with open(config_file, "r") as f:
+        return yaml.load(f, Loader=yaml.FullLoader)

agent/agents.py

@@ -1,22 +1,16 @@
 import sys
-import os
 from abc import ABC, abstractmethod
 from pathlib import Path
 import logging
 
 from aider.coders import Coder
 from aider.models import Model
 from aider.io import InputOutput
-from tenacity import retry, wait_exponential, RetryCallState, retry_if_exception_type
+import re
 
 
-class APIError(Exception):
-    def __init__(self, status_code: int, message: str):
-        self.status_code = status_code
-        self.message = message
-        super().__init__(f"API Error: {status_code} - {message}")
-
-def handle_logging(logging_name: str, log_file: Path):
+def handle_logging(logging_name: str, log_file: Path) -> None:
+    """Handle logging for agent"""
     logger = logging.getLogger(logging_name)
     logger.setLevel(logging.INFO)
     logger.propagate = False
@@ -26,17 +20,32 @@ def handle_logging(logging_name: str, log_file: Path):
     )
     logger.addHandler(logger_handler)
 
+
+class AgentReturn(ABC):
+    def __init__(self, log_file: Path):
+        self.log_file = log_file
+        self.last_cost = self.get_money_cost()
+
+    def get_money_cost(self) -> float:
+        """Get accumulated money cost from log file"""
+        last_cost = 0.0
+        with open(self.log_file, "r") as file:
+            for line in file:
+                if "Tokens:" in line and "Cost:" in line:
+                    match = re.search(
+                        r"Cost: \$\d+\.\d+ message, \$(\d+\.\d+) session", line
+                    )
+                    if match:
+                        last_cost = float(match.group(1))
+        return last_cost
+
+
 class Agents(ABC):
-    def __init__(self, max_iteration: int, retry_if_api_error_codes: tuple[int, ...] = (429, 503, 529)):
+    def __init__(self, max_iteration: int):
         self.max_iteration = max_iteration
 
-        # error code 429 is rate limit exceeded for openai and anthropic
-        # error code 503 is service overloaded for openai
-        # error code 529 is service overloaded for anthropic
-        self.retry_if_api_error_codes = retry_if_api_error_codes
-
     @abstractmethod
-    def run(self) -> None:
+    def run(self) -> AgentReturn:
         """Start agent"""
         raise NotImplementedError
 
@@ -46,84 +55,81 @@ def __init__(self, max_iteration: int, model_name: str):
         super().__init__(max_iteration)
         self.model = Model(model_name)
 
-    @retry(
-        wait=wait_exponential(multiplier=1, min=4, max=10),
-        retry=retry_if_exception_type(APIError)
-    )
     def run(
         self,
         message: str,
         test_cmd: str,
         lint_cmd: str,
         fnames: list[str],
         log_dir: Path,
-    ) -> None:
+    ) -> AgentReturn:
         """Start aider agent"""
-        try:
-            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 / ".aider.input.history"
-            chat_history_file = log_dir / ".aider.chat.history.md"
-
-            print(
-                f"check {os.path.abspath(chat_history_file)} for prompts and lm generations",
-                file=sys.stderr,
-            )
-            # Set up logging
-            log_file = log_dir / "aider.log"
-            logging.basicConfig(
-                filename=log_file,
-                level=logging.INFO,
-                format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-            )
+        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 / ".aider.input.history"
+        chat_history_file = log_dir / ".aider.chat.history.md"
 
-            # Redirect print statements to the log file
-            sys.stdout = open(log_file, "a")
-            sys.stderr = open(log_file, "a")
+        # Set up logging
+        log_file = log_dir / "aider.log"
+        logging.basicConfig(
+            filename=log_file,
+            level=logging.INFO,
+            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+        )
 
-            # Configure httpx and backoff logging
-            handle_logging("httpx", log_file)
-            handle_logging("backoff", log_file)
+        # Redirect print statements to the log file
+        sys.stdout = open(log_file, "a")
+        sys.stderr = open(log_file, "a")
 
-            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={"python": lint_cmd},
-                test_cmd=test_cmd,
-                io=io,
+        # 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,
+        )
+        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
+        # coder.run(message)
+
+        #### TMP
+        import time
+        import random
+
+        time.sleep(random.random() * 5)
+        n = random.random() / 10
+        with open(log_file, "a") as f:
+            f.write(
+                f"> Tokens: 33k sent, 1.3k received. Cost: $0.12 message, ${n} session.  \n"
             )
-            coder.max_reflection = self.max_iteration
-            coder.stream = False
-
-            # Run the agent
-            raise Exception("test")
-            coder.run(message)
-        
-        except Exception as e:
-            # If the exception is related to API errors, raise an APIError
-            if hasattr(e, 'status_code') and e.status_code in self.retry_if_api_error_codes:
-                raise APIError(e.status_code, str(e))
-            # For other exceptions, re-raise them
-            raise
-        finally:
-            # Close redirected stdout and stderr
-            sys.stdout.close()
-            sys.stderr.close()
-            # Restore original stdout and stderr
-            sys.stdout = sys.__stdout__
-            sys.stderr = sys.__stderr__
+        #### TMP
+
+        # Close redirected stdout and stderr
+        sys.stdout.close()
+        sys.stderr.close()
+        # Restore original stdout and stderr
+        sys.stdout = sys.__stdout__
+        sys.stderr = sys.__stderr__
+
+        return AgentReturn(log_file)

agent/class_types.py

@@ -5,7 +5,6 @@
 class AgentConfig:
     agent_name: str
     model_name: str
-    backend: str
     use_user_prompt: bool
     user_prompt: str
     use_repo_info: bool