react_code_log_dec7_tool_use.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[7]):
        print(f"--- Question {i+1} ---")
        q = all_questions[7]
        # Access individual fields using their keys
        f_name = q['function_name']
        prompt = q['question_prompt']
        code   = q['provided_code']
        ## 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 EHT observation operator (a function named as build_eht_observation_operator)

        You are also provided a tool to resolve this problem. The tool has been implemented in advance, so you don't have to reimplement the tool and you could just import the tool and directly use it!!
        the provided tool is for {provided_tool_explain}.
        the detail i/o of this tool is {provided_tool_docstring}. 
        You could use this tool by following this example:
        {provided_tool_usage}
        You must first use this tool!
        If you don't use this tool, give out your reason first!
        """
        logging.info(CI_agent_prompt)
        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 EHT observation operator (a function named as build_eht_observation_operator)
            You are also provided a tool to resolve this problem, the provided tool is for {provided_tool_explain} and the detail i/o of this tool is {provided_tool_docstring}. You could use this tool by following this example{provided_tool_usage}
            You must first use this tool!
            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.
            """
                logging.info(CI_agent_prompt)
            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 EHT observation operator (a function named as build_eht_observation_operator)
            You are also provided a tool to resolve this problem, the provided tool is for {provided_tool_explain} and the detail i/o of this tool is {provided_tool_docstring}. You could use this tool by following this example{provided_tool_usage}
            You must first use this tool!
            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.
            """
                logging.info(CI_agent_prompt)

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