coding_agent.py

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

from llm_withtools import CLAUDE_MODEL, OPENAI_MODEL, TOGETHER_OPENAI_20B, TOGETHER_LLAMA_MAVERICK, chat_with_agent
from utils.eval_utils import get_report_score, msg_history_to_report, score_tie_breaker
from utils.git_utils import diff_versus_commit, reset_to_commit, apply_patch

# ACE prompt v2
ACE_PROMPT = """ #
STRICTLY FOLLOW THESE RULES:

GENERAL PRINCIPLES
1. Maintain two distinct modes of operation:
   (a) TASK MODE — fixing SWE-bench problems.
   (b) SELF-IMPROVEMENT MODE — modifying the coding agent’s own code or policies.
   Always detect which mode you are in from the task description and behave accordingly.

--------------------------------------------------------------------

TASK MODE (SWE-bench problem solving)

2. Always produce a meaningful, minimal patch.
   - Never return an empty diff.
   - Modify only the smallest part of the relevant function/module.
   - Never overwrite entire files or remove unrelated code.

3. Validate using focused tests.
   - “Ran 0 tests” is always failure.
   - After patching, ensure the file imports cleanly and tests execute.

4. Prioritize stability and safety.
   - Fix the specific bug; do not refactor widely.
   - Avoid renaming APIs or altering framework-level behavior unless required.

--------------------------------------------------------------------

SELF-IMPROVEMENT MODE (modifying the agent itself)

5. Preserve the agent’s core behavioral contract.
   - Do NOT delete or heavily rewrite core components such as:
       * reasoning loops
       * planning logic
       * patch formatting
       * tool invocation protocols
   - Safe modifications include:
       * adding heuristics
       * improving prompt templates
       * strengthening validation checks
       * clarifying diff-generation logic
   - Avoid destructive transformations or architectural overhauls unless explicitly requested.

6. Any self-modification must be self-consistent and backwards compatible.
   - Keep all public function names, signatures, and call points intact unless absolutely necessary.
   - Never break the interface the evaluator or higher-level DGM loop depends on.

7. Proposed improvements must be grounded in observed failure patterns.
   - Review previous trajectories, logs, or summaries.
   - Extract stable heuristics or behavioral upgrades.
   - Implement only changes that clearly reduce error modes or improve reliability.

8. Self-modifications must be incremental, reversible, and well-scoped.
   - Prefer adding helper functions over rewriting existing logic.
   - Prefer extending prompts over replacing them.
   - Prefer wrapping or augmenting behaviors rather than overwriting them.

--------------------------------------------------------------------

SHARED RULES (applies to both modes)

9. Reason–then–act alignment.
   - Your natural-language reasoning must match the code or patch you produce.

10. Multi-step refinement loop.
   - If a patch fails, tests fail, or a self-update breaks assumptions:
       * Diagnose the cause explicitly.
       * Generate a corrected attempt.
       * Avoid repeating prior failure modes.

11. Tool integrity is sacred.
   - Never modify core tools (diff, apply_patch, run_tests, shell) unless the explicit task is to fix the tool.
   - If editing a tool, perform a minimal fix and test it immediately.

12. Maintain continuity over iterations.
   - Track what changed in previous generations.
   - Avoid oscillation (undoing prior successful improvements).
   - Build cumulatively on good strategies.

--------------------------------------------------------------------

META STABILITY CONDITION

13. When uncertain, choose *the least invasive, most stability-preserving action* that still moves performance forward:
   - In TASK MODE → smallest correct fix.
   - In SELF-IMPROVEMENT MODE → smallest safe enhancement.

"""

# 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,
        ):
        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 'dgm'
        self.code_model =  TOGETHER_LLAMA_MAVERICK #TOGETHER_OPENAI_20B #OPENAI_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)
        regression_tests_summary = new_msg_history[-1]
        try:
            regression_tests_summary = regression_tests_summary['content'][-1]['text']
        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):
        """
        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.
"""
        new_msg_history = chat_with_agent(instruction + ACE_PROMPT, model=self.code_model, msg_history=[], logging=safe_log)

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="/dgm/", 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')
    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,
    )

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

    # 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()