coding_agent.py

# This file is adapted from https://github.com/jennyzzt/dgm.

import argparse
import logging
import os
import subprocess
import threading
from logging.handlers import RotatingFileHandler
from time import time

from llm_withtools import (CLAUDE_MODEL, OPENAI_MODEL, chat_with_agent,
                           convert_msg_history)
from utils.eval_utils import (get_report_score, msg_history_to_report,
                              score_tie_breaker)
from utils.git_utils import apply_patch, diff_versus_commit, reset_to_commit

# Thread-local storage for logger instances
thread_local = threading.local()


def get_thread_logger():
    """
    Get the logger instance specific to the current thread.
    Returns None if no logger has been set for this thread.
    """
    return getattr(thread_local, "logger", None)


def set_thread_logger(logger):
    """
    Set the logger instance for the current thread.
    """
    thread_local.logger = logger


def setup_logger(log_file="./chat_history.md", level=logging.INFO):
    """
    Set up a logger with both file and console handlers.
    """
    # Create logger with a unique name based on thread ID
    logger = logging.getLogger(f"AgenticSystem-{threading.get_ident()}")
    logger.setLevel(level)

    # Remove existing handlers to avoid duplicates
    logger.handlers = []

    # Create formatters
    file_formatter = logging.Formatter("%(message)s")

    # Create and set up file handler
    os.makedirs(os.path.dirname(log_file), exist_ok=True)
    file_handler = RotatingFileHandler(
        log_file, maxBytes=10 * 1024 * 1024, backupCount=5
    )
    file_handler.setLevel(level)
    file_handler.setFormatter(file_formatter)

    # Add handlers to logger
    logger.addHandler(file_handler)

    # Store logger in thread-local storage
    set_thread_logger(logger)

    return logger


def safe_log(message, level=logging.INFO):
    """
    Thread-safe logging function that ensures messages go to the correct logger.
    """
    logger = get_thread_logger()
    if logger:
        logger.log(level, message)
    else:
        print(f"Warning: No logger found for thread {threading.get_ident()}")


class AgenticSystem:
    def __init__(
        self,
        problem_statement,
        git_tempdir,
        base_commit,
        chat_history_file="./chat_history.md",
        test_description=None,
        self_improve=False,
        instance_id=None,
        model=CLAUDE_MODEL,
    ):
        self.problem_statement = problem_statement
        self.git_tempdir = git_tempdir
        self.base_commit = base_commit
        self.chat_history_file = chat_history_file
        self.test_description = test_description
        self.self_improve = self_improve
        self.instance_id = instance_id if not self_improve else "hgm"
        self.code_model = model

        # Initialize logger and store it in thread-local storage
        self.logger = setup_logger(chat_history_file)

        # Clear the log file
        with open(chat_history_file, "w") as f:
            f.write("")

    def get_current_edits(self):
        diff = str(diff_versus_commit(self.git_tempdir, self.base_commit))
        return diff

    def get_regression_tests(self):
        """
        Get the regression tests from the repository.
        """
        instruction = f"""I have uploaded a Python code repository in the directory {self.git_tempdir}.

<problem_description>
{self.problem_statement}
</problem_description>

<test_description>
{self.test_description}
</test_description>

Your task is to identify regression tests in the {self.git_tempdir} directory that should pass both before and after addressing the <problem_description>. I have already taken care of the required dependencies.
At the end, please provide a summary that includes where the regression tests are located, what they are testing, and how they can be executed.
"""

        new_msg_history, _ = chat_with_agent(
            instruction, model=self.code_model, msg_history=[], logging=safe_log
        )
        new_msg_history = convert_msg_history(new_msg_history, self.code_model)
        regression_tests_summary = new_msg_history[-1]
        try:
            regression_tests_summary = regression_tests_summary["content"]
        except:
            try:
                regression_tests_summary = str(regression_tests_summary)
            except:
                pass
        return regression_tests_summary

    def run_regression_tests(self, regression_tests_summary):
        """
        Run the regression tests and get the test report.
        """
        code_diff = self.get_current_edits()
        instruction = f"""I have uploaded a Python code repository in the directory {self.git_tempdir}. There is an attempt to address the problem statement. Please review the changes and run the regression tests.

<problem_description>
{self.problem_statement}
</problem_description>

<attempted_solution>
{code_diff}
</attempted_solution>

<test_description>
{self.test_description}
</test_description>

<regression_tests_summary>
{regression_tests_summary}
</regression_tests_summary>

Your task is to run the regression tests in the {self.git_tempdir} directory to ensure that the changes made to the code address the <problem_description>.
"""
        new_msg_history, _ = chat_with_agent(
            instruction, model=self.code_model, msg_history=[], logging=safe_log
        )
        test_report = msg_history_to_report(
            self.instance_id, new_msg_history, model=self.code_model
        )
        return test_report

    def forward(self, timeout=3600):
        timeout -= 60
        start_time = time()
        """
        The forward function for the AgenticSystem.
        """
        instruction = f"""I have uploaded a Python code repository in the directory {self.git_tempdir}. Help solve the following problem.

<problem_description>
{self.problem_statement}
</problem_description>

<test_description>
{self.test_description}
</test_description>

Your task is to make changes to the files in the {self.git_tempdir} directory to address the <problem_description>. I have already taken care of the required dependencies.
"""
        chat_history, n_llm_calls_used = chat_with_agent(
            instruction,
            model=self.code_model,
            msg_history=[],
            logging=safe_log,
            timeout=timeout - (time() - start_time),
        )
        chat_history_str = str(chat_history)


def main():
    parser = argparse.ArgumentParser(
        description="Process repository with an agentic system."
    )
    parser.add_argument(
        "--problem_statement", required=True, help="The problem statement to process"
    )
    parser.add_argument(
        "--git_dir", required=True, help="Path to git repository directory"
    )
    parser.add_argument(
        "--base_commit", required=True, help="Base commit hash to compare against"
    )
    parser.add_argument(
        "--chat_history_file", required=True, help="Path to chat history file"
    )
    parser.add_argument(
        "--outdir", required=False, default="/hgm/", help="Output directory"
    )
    parser.add_argument(
        "--test_description",
        default=None,
        required=False,
        help="Description of how to test the repository",
    )
    parser.add_argument(
        "--self_improve",
        default=False,
        action="store_true",
        help="Whether to self-improve the repository or solving swe",
    )
    parser.add_argument("--instance_id", default=None, help="Instance ID for SWE issue")
    parser.add_argument(
        "--model",
        required=False,
        default=OPENAI_MODEL,
        help="LLM model to use for processing",
    )
    parser.add_argument(
        "--timeout", type=int, default=3600, help="Timeout for LLM calls in seconds"
    )
    args = parser.parse_args()

    # Process the repository
    agentic_system = AgenticSystem(
        problem_statement=args.problem_statement,
        git_tempdir=args.git_dir,
        base_commit=args.base_commit,
        chat_history_file=args.chat_history_file,
        test_description=args.test_description,
        self_improve=args.self_improve,
        instance_id=args.instance_id,
        model=args.model,
    )

    # Run the agentic system to try to solve the problem
    agentic_system.forward(args.timeout)

    # Get code diff and save to model_patch.diff
    model_patch = diff_versus_commit(args.git_dir, args.base_commit)
    model_patch_outfile = (
        os.path.join(args.outdir, "model_patch.diff")
        if args.outdir
        else "model_patch.diff"
    )
    with open(model_patch_outfile, "w") as f:
        f.write(model_patch)


if __name__ == "__main__":
    main()