react_code_log_dec7.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
import shutil

Hardware_info= """
Hardware Environment Specification:

Device: 1x NVIDIA A800-SXM4
VRAM: 80 GB (High Availability)
CUDA Driver: 12.2 (Driver 535.129.03)
Utilization: Currently idle (0% usage)
"""

def load_data(problem_description_path_json, specific_indices = None):
        with open(problem_description_path_json, 'r') as f:
            
                data = json.load(f)
        all_questions = data['questions']
        return all_questions

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)
    
    # copting the working_folder

    # 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()
    # copy the working folder into an agent test environment! Nor ruin the older version of old test environment folder
    source_folder = args.working_folder_location
    # Removes trailing slash if present (e.g., 'folder/' -> 'folder') to get the clean name
    clean_source = source_folder.rstrip('/')
    parent_dir = os.path.dirname(clean_source)
    folder_name = os.path.basename(clean_source)

    # Construct new name automatically (e.g., DPItorch -> DPItorch_agent)
    destination_folder = os.path.join(parent_dir, f"{folder_name}_agent")
    if os.path.exists(destination_folder):
        pass
    else:
        shutil.copytree(source_folder, destination_folder, dirs_exist_ok=True)
    working_folder_location = destination_folder
    # --- 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))
    all_questions = load_data(args.problem_description_path_json)
    
    ## here define a Markov decision process
    # for question_id, question in enumerate(problem_description):

###### TODO: we are testing the tool used and new input output info provided!
#     with open("/fs-computility-new/UPDZ02_sunhe/shared/auto_tutorial_flow_2/auto_tutorial_flow/auto_flow/generated_docs_forward_operator.json", 'r', encoding='utf-8') as f:
#             reloaded_docs = json.load(f)
#     function_docstring = reloaded_docs["function_docstring"]
# ########### TODO: using tool!
#     provided_tool_explain = reloaded_docs["function_explanation"]
#     provided_tool_docstring = reloaded_docs["function_docstring"]
#     provided_tool_usage = reloaded_docs['usage_example']
####################

    for i, q in enumerate(all_questions):
        print(f"--- Question {i+1} ---")
        
        # Access individual fields using their keys
        f_name = q['function_name']
        if f_name == "Loss_angle_diff":
            continue
        prompt = q['question_prompt']
        code   = q['provided_code']
        with open(f"/fs-computility-new/UPDZ02_sunhe/shared/auto_tutorial_flow_2/auto_tutorial_flow/auto_flow/generated_docs_{f_name}.json", 'r', encoding='utf-8') as f:
            reloaded_docs = json.load(f)
            function_docstring = reloaded_docs["function_docstring"]
        ########### TODO: using tool!
            provided_tool_explain = reloaded_docs["function_explanation"]
            provided_tool_docstring = reloaded_docs["function_docstring"]
            provided_tool_usage = reloaded_docs['usage_example']
        ## set up working file path:
        gt_save_path =  os.path.join(working_folder_location, f"gt_{f_name}.py")
        agent_code_path = os.path.join(working_folder_location, f"agent_{f_name}.py")
        if os.path.exists(gt_save_path):
            working_file_path = agent_code_path
        else:
            with open(agent_code_path, 'r', encoding='utf-8', errors='ignore') as f:
                gt_code = f.read()
            with open(gt_save_path, 'w', encoding='utf-8', errors='ignore') as f:
                f.write(gt_code)
            working_file_path = agent_code_path
        logging.info(f"--- Starting new problem: Question function {f_name} ---")
        logging.info(f"Problem Description: {prompt} and the provided code {code}, The detailed input/output info is {function_docstring}")
        previous_code = ""
        prevous_results = ""
        CI_agent_prompt=f"""
        You are a expert in computational imaging expert and you're provided with code development task tailed for developing scientific functions.
        You should write code to complete the function 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
        You have access to these hardware:{Hardware_info}
        The target function is {f_name}
        The problem is {prompt}
        And the provided code is {code}. The detailed input/output info is {function_docstring}. using the packages imported and other helper function defined in the provided code.
        Your answer should contain the completed function and the other imported packages and provided helper function.
        """
        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'], working_file_path, args.command, working_folder_location, f_name)
            

            
            if current_results["status"]== "SUCCESS" or current_results["status"]== "Wronganswer":
                current_status = current_results["status"]
                logging.info(f"Question status {current_status}!")
                logging.info(f"answer code {current_code}")
                break
            elif current_results["status"]==  "FAIL" 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 tailed for developing scientific functions.
            You should write code to complete the function 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
            You have access to these hardware:{Hardware_info}
            The target function is {f_name}
            The problem is {prompt}
            And the provided code is {code}, The detailed input/output info is {function_docstring}. using the packages imported and other helper function defined in the provided code.
            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!
            Your answer should contain the completed function and the other imported packages and provided helper function.
            """
               
            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 tailed for developing scientific functions.
            You should write code to complete the function 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
            You have access to these hardware:{Hardware_info}
            The target function is {f_name}
            The problem is {prompt}
            And the provided code is {code}, The detailed input/output info is {function_docstring}. using the packages imported and other helper function defined in the provided code.
            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(working_file_path, 'w', encoding='utf-8') as f:
                    f.write(current_code)