react_code_history.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 
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(
        "--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}

        """
        error_history = []

        # (Assuming 'CI_agent_prompt' is already defined for the first iteration)

        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)
            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:
                # --- START OF MODIFIED 'else' BLOCK ---
                
                # 1. Get the current error and append it to the history list
                current_error = current_results["error"]
                error_history.append(current_error)
                
                # Format the history for clean insertion into the prompt
                # This turns the list into a numbered string
                formatted_error_history = "\n".join(
                    [f"Attempt {idx+1} Error: {err}" for idx, err in enumerate(error_history)]
                )

                # 2. Add the complete error history into the prompt
                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_error)