react_code.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
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=50, 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.",
    )
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_args()
    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):
        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}

        """
        for i in range(args.max_rounds):
            # print("prompt", CI_agent_prompt)
            # solution = asyncio.run(custom_code_generate(problem=CI_agent_prompt, instruction=""))
            # current_results, current_code = code_development_env(solution['response'], args.working_file_path, args.command, args.working_folder_location)
            current_results, current_code = code_development_env("assert 1==0", args.working_file_path, args.command, args.working_folder_location)
            print("*"*12)
            print(f"rounds {i}")
            # print("solution", solution)
            print("current_results", current_results)
            if current_results["status"]== "SUCCESS":
                print("question solved!")
                with open(args.working_file_path, 'w', encoding='utf-8') as f:
                    f.write(current_code)
                break
            else:
                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!
            """