react_code_log.py

## change the step_by_step or direct code in code_instructions in /fs-computility-new/UPDZ02_sunhe/chensiyi.p/ReAct/scripts/formatter.py
## in the /fs-computility-new/UPDZ02_sunhe/chensiyi.p/ReAct/scripts/formatter.py, if only dealing with the function generation code, remeber to enable the sanitize, but not necessary!
import argparse
import aiofiles
import asyncio
import json
import logging
import os
from datetime import datetime
from code_dev_env import code_development_env
from scripts.async_llm import LLMsConfig
from scripts.async_llm import create_llm_instance
import workspace.InverseProb.workflows.template.operator as operator

def load_data(problem_description_path_json, specific_indices = None):
        data = []
        with open(problem_description_path_json, mode="r", encoding="utf-8") as file:
            
                data = json.load(file)
        if specific_indices is not None:
            filtered_data = [data[i] for i in specific_indices if i < len(data)]
            return filtered_data
        return data

def parse_args():
    parser = argparse.ArgumentParser(description="react_code")
    parser.add_argument(
        "--working_file_path",
        type=str,
        required=True,
        help="write the test code given by the llm",
    )
    parser.add_argument("--command", 
                    nargs='*',  # <--- This is the correct way
                    default=[
                        "python", 
                        "/fs-computility-new/UPDZ02_sunhe/chensiyi.p/AFlow/code_development/obs_arg/eval_v3.py", 
                        "--npix", "32", 
                        "--obspath", "/fs-computility-new/UPDZ02_sunhe/shared/eht_imaging/DPI/dataset/interferometry1/obs.uvfits"
                    ], 
                    help="the command used to run the test env")
    parser.add_argument(
        "--working_folder_location",
        type=str,
        default="./",
        help="the folder to perform test command and save the results",
    )
    parser.add_argument(
        "--problem_description_path_json",
        type=str,
        default="workspace",
        help="Optimized result save path",
    )
    parser.add_argument("--max_rounds", type=int, default=20, help="Max iteration rounds")
    parser.add_argument(
        "--model_name",
        type=str,
        default="deepseek-r1-250528",
        help="Specifies the name of the model used for optimization tasks.",
    )
    # --- New Arguments for Logging ---
    parser.add_argument(
        "--log_dir",
        type=str,
        default="./logs",
        help="Directory to save log files."
    )
    parser.add_argument(
        "--log_name",
        type=str,
        default="react_run",
        help="Base name for the log file. A timestamp will be appended."
    )
    return parser.parse_args()

def setup_logging(log_dir, log_name):
    """Configures the logging system."""
    # Ensure the log directory exists
    os.makedirs(log_dir, exist_ok=True)
    
    # Create a timestamped log filename
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file_name = f"{log_name}_{timestamp}.log"
    log_file_path = os.path.join(log_dir, log_file_name)
    
    # Get the root logger
    logger = logging.getLogger()
    logger.setLevel(logging.INFO) # Set the minimum level to log
    
    # Create a formatter
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    
    # Create File Handler (to write to file)
    file_handler = logging.FileHandler(log_file_path)
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    
    # Create Console Handler (to print to console)
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)
    
    logger.info(f"Logging initialized. Log file: {log_file_path}")
    return logger

if __name__ == "__main__":
    args = parse_args()
    
    # --- Setup the logger as the first step ---
    setup_logging(args.log_dir, args.log_name)
    
    models_config = LLMsConfig.default()
    models_config = models_config.get(args.model_name)
    custom_code_generate = operator.CustomCodeGenerate(create_llm_instance(models_config))
    problem_description = load_data(args.problem_description_path_json)
    
    ## here define a Markov decision process
    for question_id, question in enumerate(problem_description):
        logging.info(f"--- Starting new problem: Question ID {question_id} ---")
        logging.info(f"Problem Description: {problem_description[question]}")
        previous_code = ""
        prevous_results = ""
        CI_agent_prompt=f"""
        You are a expert in computational imaging expert and you're provided with code development task.
        You should write code to finish the job and your code would be written into a certain file and tested to check whether you produced the same results like the provided expert code
        The problem is {problem_description[question]}
        The previous trail code is {previous_code}
        The previous test results of the trail code is {prevous_results}

        """
        error_history = []
        for i in range(args.max_rounds):
            logging.info(f"--- Round {i} ---")
            logging.debug(f"Sending prompt: {CI_agent_prompt}") # Debug level for long prompts
            
            solution = asyncio.run(custom_code_generate(problem=CI_agent_prompt, instruction=""))
            logging.debug(f"Received solution: {solution['response']}") # Debug level for code
            
            current_results, current_code = code_development_env(solution['response'], args.working_file_path, args.command, args.working_folder_location)
            

            
            if current_results["status"]== "SUCCESS":
                logging.info("Question solved!")
                logging.info(f"answer code {current_code}")
                break
            elif (current_results["status"]!= "SUCCESS") and (i==0):
                logging.error(f"Test failed. Error: {current_results['error']}")
                logging.error(f"error code: {current_code}")
                CI_agent_prompt = f"""
            You are a expert in computational imaging expert and you're provided with code development task.
            You should write code to finish the job and your code would be written into a certain file and tested to check whether you produced the same results like the provided expert code
            The problem is {question}
            The previous trail code is {solution['response']}
            The previous test error of the trail code is {current_results["error"]}
            You should revised the previous trail code with the provided test error info!
            """
               
            else:
                logging.error(f"Test failed. Error: {current_results['error']}")
                logging.error(f"error code: {current_code}")
                formatted_error_history = "\n".join(
                    [f"Attempt {idx+1} Error: {err}" for idx, err in enumerate(error_history)]
                )
                CI_agent_prompt = f"""
            You are a expert in computational imaging expert and you're provided with code development task.
            You should write code to finish the job and your code would be written into a certain file and tested to check whether you produced the same results like the provided expert code
            The problem is {question}
            The previous trail code is {solution['response']}
            The previous test error of the trail code is {current_results["error"]}
            You should revised the previous trail code with the provided test error info!
            Here is the full history of errors from all past attempts:
            {formatted_error_history}
            You should revised the previous trail code, taking into account the **latest test error** and the **full error history** to avoid repeating mistakes.
            """

            error_history.append(current_results["error"])
        with open(args.working_file_path, 'w', encoding='utf-8') as f:
                    f.write(current_code)