core/agent/task.py

from .base_agent import BaseAgent
from prompt.constants import modeling_methods
from prompt.template import (TASK_ANALYSIS_PROMPT, TASK_RESULT_PROMPT, TASK_ANSWER_PROMPT, 
                             TASK_FORMULAS_PROMPT, TASK_FORMULAS_CRITIQUE_PROMPT, TASK_FORMULAS_IMPROVEMENT_PROMPT, 
                             TASK_MODELING_PROMPT, TASK_MODELING_CRITIQUE_PROMPT, TASK_MODELING_IMPROVEMENT_PROMPT,
                             TASK_CODING_PROMPT, TASK_CODING_DEBUG_PROMPT, CODE_STRUCTURE_PROMPT, 
                             TASK_RESULT_WITH_CODE_PROMPT, COO_PROMPT, TASK_CODING_WO_COO_PROMPT)
import sys
import os
import subprocess
import selectors
import tiktoken
import json


class EnvException(Exception):
    def __init__(self, message):
        self.message = message 
    def __str__(self):
        return self.message
    

def execute_script(script_path, work_dir):
    try:
        device = 0
        python = "python"
        cmd = f"CUDA_VISIBLE_DEVICES={device} {python} -u {script_path}"
        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, shell=True, cwd=work_dir)

        stdout_lines = []
        stderr_lines = []

        selector = selectors.DefaultSelector()
        selector.register(process.stdout, selectors.EVENT_READ)
        selector.register(process.stderr, selectors.EVENT_READ)

        while process.poll() is None and selector.get_map():
            events = selector.select(timeout=1)

            for key, _ in events:
                line = key.fileobj.readline()
                if key.fileobj == process.stdout:
                    print("STDOUT:", line, end =" ")
                    stdout_lines.append(line)
                else:
                    print("STDERR:", line, end =" ")
                    stderr_lines.append(line)

        for line in process.stdout:
            line = line
            print("STDOUT:", line, end =" ")
            stdout_lines.append(line)
        for line in process.stderr:
            line = line
            print("STDERR:", line, end =" ")
            stderr_lines.append(line)

        return_code = process.returncode

        if return_code != 0:
            observation = "".join(stderr_lines)
        else:
            observation = "".join(stdout_lines)
        if observation == "" and return_code == 0:
            # printed to stderr only
            observation = "".join(stderr_lines)
        return "The script has been executed. Here is the output:\n" + observation
    except Exception as e:
        print("++++", "Wrong!")
        raise EnvException(f"Something went wrong in executing {script_path}: {e}. Please check if it is ready to be executed.")


class Task(BaseAgent):
    def __init__(self, llm, coo=True, rag=True):
        super().__init__(llm)
        self.coo = coo
        self.rag = rag
        if coo:
            self.coo_prompt = COO_PROMPT
        else:
            self.coo_prompt = ""

    def analysis(self, prompt: str, task_description: str, user_prompt: str = ''):
        prompt = TASK_ANALYSIS_PROMPT.format(prompt=prompt, coo_prompt=self.coo_prompt, task_description=task_description, user_prompt=user_prompt).strip()
        return self.llm.generate(prompt)
    
    def formulas_actor(self, prompt: str, data_summary: str, task_description: str, task_analysis: str, modeling_methods: str, user_prompt: str = ''):
        prompt = TASK_FORMULAS_PROMPT.format(prompt=prompt, coo_prompt=self.coo_prompt, data_summary=data_summary, task_description=task_description, task_analysis=task_analysis, modeling_methods=modeling_methods, user_prompt=user_prompt).strip()
        return self.llm.generate(prompt)

    def formulas_critic(self, data_summary: str, task_description: str, task_analysis: str, modeling_formulas: str):
        prompt = TASK_FORMULAS_CRITIQUE_PROMPT.format(data_summary=data_summary, task_description=task_description, task_analysis=task_analysis, modeling_formulas=modeling_formulas).strip()
        return self.llm.generate(prompt)
    
    def formulas_improvement(self, data_summary: str, task_description: str, task_analysis: str, modeling_formulas: str, modeling_formulas_critique: str, user_prompt: str = ''):
        prompt = TASK_FORMULAS_IMPROVEMENT_PROMPT.format(data_summary=data_summary, task_description=task_description, task_analysis=task_analysis, modeling_formulas=modeling_formulas, modeling_formulas_critique=modeling_formulas_critique, user_prompt=user_prompt).strip()
        return self.llm.generate(prompt)

    def formulas(self, prompt: str, data_summary: str, task_description: str, task_analysis: str, modeling_methods: str, round: int = 1, user_prompt: str = ''):
        formulas = self.formulas_actor(prompt, data_summary, task_description, task_analysis, modeling_methods, user_prompt)
        if self.rag:
            for i in range(round):
                print(f'FORMULAS Round {i+1}')
                formulas_critique = self.formulas_critic(data_summary, task_description, task_analysis, formulas)
                formulas = self.formulas_improvement(data_summary, task_description, task_analysis, formulas, formulas_critique, user_prompt)
        
        return formulas

    def modeling_actor(self, prompt: str, data_summary: str, task_description: str, task_analysis: str, formulas: str, user_prompt: str = ''):
        prompt = TASK_MODELING_PROMPT.format(prompt=prompt, coo_prompt=self.coo_prompt, data_summary=data_summary, task_description=task_description, task_analysis=task_analysis, modeling_formulas=formulas, user_prompt=user_prompt).strip()
        return self.llm.generate(prompt)

    # def modeling_critic(self, task_description: str, task_analysis: str, data_summary: str, formulas: str, modeling_process: str):
    #     prompt = TASK_MODELING_CRITIQUE_PROMPT.format(task_description=task_description, task_analysis=task_analysis, data_summary=data_summary, modeling_formulas=formulas, modeling_process=modeling_process).strip()
    #     return self.llm.generate(prompt)
    
    # def modeling_improvement(self, task_description: str, task_analysis: str, data_summary: str, formulas: str, modeling_process: str, modeling_process_critique: str):
    #     prompt = TASK_MODELING_IMPROVEMENT_PROMPT.format(task_description=task_description, task_analysis=task_analysis, data_summary=data_summary, modeling_formulas=formulas, modeling_process=modeling_process, modeling_process_critique=modeling_process_critique).strip()
    #     return self.llm.generate(prompt)

    # def modeling(self, task_description: str, task_analysis: str, data_summary: str, formulas: str, round: int = 1):
    #     process = self.modeling_actor(task_description, task_analysis, data_summary, formulas)
    #     for i in range(round):
    #         print(f'MODELING Round {i+1}')
    #         process_critique = self.modeling_critic(task_description, task_analysis, data_summary, formulas, process)
    #         process = self.modeling_improvement(task_description, task_analysis, data_summary, formulas, process, process_critique)
    #     return process

    def modeling(self, prompt: str, data_summary: str, task_description: str, task_analysis: str, formulas: str, round: int = 1, user_prompt: str = ''):
        return self.modeling_actor(prompt, data_summary, task_description, task_analysis, formulas, user_prompt)
    
    def modeling_actor(self, prompt: str, data_summary: str, task_description: str, task_analysis: str, formulas: str, modeling: str, user_prompt: str = ''):
        prompt = TASK_MODELING_PROMPT.format(prompt=prompt, coo_prompt=self.coo_prompt, data_summary=data_summary, task_description=task_description, task_analysis=task_analysis, modeling_formulas=formulas, modeling_methods=modeling, user_prompt=user_prompt).strip()
        return self.llm.generate(prompt)
    
    def coding_actor(self, data_file, data_summary, variable_description, task_description: str, task_analysis: str, formulas: str, modeling: str, dependent_file_prompt: str, code_template: str, script_name: str, work_dir: str, user_prompt: str = ''):
        if self.coo:
            prompt = TASK_CODING_PROMPT.format(data_file=data_file, data_summary=data_summary, variable_description=variable_description, task_description=task_description, task_analysis=task_analysis, modeling_formulas=formulas, modeling_process=modeling, dependent_file_prompt=dependent_file_prompt, code_template=code_template, user_prompt=user_prompt).strip()
        else:
            prompt = TASK_CODING_WO_COO_PROMPT.format(data_file=data_file, data_summary=data_summary, variable_description=variable_description, task_description=task_description, task_analysis=task_analysis, modeling_formulas=formulas, modeling_process=modeling, code_template=code_template, user_prompt=user_prompt).strip()  
        max_retry = 0
        while max_retry < 5:
            max_retry += 1
            try:
                completion = self.llm.generate(prompt)
                new_content = completion.split("```python")[1].split("```")[0].strip()
                break  
            except Exception as e:
                # Format control.
                print(f"Retry! The code does not start with ```python")
                continue

        with open(os.path.join(work_dir, script_name), "w") as f:
            f.write(new_content)
        
        # Execute the script.
        try:
            observation = execute_script(script_name, work_dir)
            ## If observation is too long, we only keep the last ~2k tokens.
            enc = tiktoken.get_encoding("cl100k_base")
            tokens = len(enc.encode(observation))
            if tokens >= 2000:
                observation = observation[:2000]
                tokens = len(enc.encode(observation))
        except Exception as e:
            print(e)
            input("Ah oh, Got stuck! Press any key to continue.")

        return new_content, observation
    
    def coding_debugger(self, code_template: str, modeling: str, code: str, observation: str, script_name: str, work_dir: str, user_prompt: str = ''):
        
        prompt = TASK_CODING_DEBUG_PROMPT.format(code_template=code_template, modeling_process=modeling, code=code, observation=observation, user_prompt=user_prompt).strip()
        
        max_retry = 0
        while max_retry < 5:
            max_retry += 1
            try:
                completion = self.llm.generate(prompt)
                new_content = completion.split("```python")[1].split("```")[0].strip()
                break  
            except Exception as e:
                # Format control.
                print(f"Retry! The code does not start with ```python")
                continue

        with open(os.path.join(work_dir, script_name), "w") as f:
            f.write(new_content)
        
        # Execute the script.
        try:
            observation = execute_script(script_name, work_dir)
            ## If observation is too long, we only keep the last ~2k tokens.
            enc = tiktoken.get_encoding("cl100k_base")
            tokens = len(enc.encode(observation))
            if tokens >= 2000:
                observation = observation[:2000]
                tokens = len(enc.encode(observation))
        except Exception as e:
            print(e)
            input("Ah oh, Got stuck! Press any key to continue.")

        return new_content, observation
    
    def coding(self, data_file, data_summary, variable_description, task_description: str, task_analysis: str, formulas: str, modeling: str, dependent_file_prompt: str, code_template: str, script_name: str, work_dir: str, try_num: int = 5, round: int = 1, user_prompt: str = ''):
        for i in range(try_num):
            print("="*10 + f" Try: {i + 1} " + "="*10)
            iteration = 0
            max_iteration = 3
            while iteration < max_iteration:
                print("="*10 + f" Iteration: {iteration + 1} " + "="*10)
                if iteration == 0:
                    code, observation = self.coding_actor(data_file, data_summary, variable_description, task_description, task_analysis, formulas, modeling, dependent_file_prompt, code_template, script_name, work_dir, user_prompt)
                    # If the script has been successfully executed: Exit.
                    if "Traceback (most recent call last):" not in observation and "SyntaxError: invalid syntax" not in observation and "IndentationError" not in observation:
                        return code, True, observation.split("The script has been executed. Here is the output:\n")[1]
                else:
                    code, observation = self.coding_debugger(code_template, modeling, code, observation, script_name, work_dir, user_prompt)
                    # If the script has been successfully executed: Exit.
                    if "Traceback (most recent call last):" not in observation and "SyntaxError: invalid syntax" not in observation and "IndentationError" not in observation:
                        return code, True, observation.split("The script has been executed. Here is the output:\n")[1]
                iteration += 1

        return code, False, None

    def result(self, task_description: str, task_analysis: str, task_formulas: str, task_modeling: str, user_prompt: str = '', execution_result: str = ''):
        if execution_result == '':
            prompt = TASK_RESULT_PROMPT.format(task_description=task_description, task_analysis=task_analysis, task_formulas=task_formulas, task_modeling=task_modeling, user_prompt=user_prompt).strip()
        else:
            prompt = TASK_RESULT_WITH_CODE_PROMPT.format(task_description=task_description, task_analysis=task_analysis, task_formulas=task_formulas, task_modeling=task_modeling, user_prompt=user_prompt, execution_result=execution_result).strip()
        return self.llm.generate(prompt)

    def answer(self, task_description: str, task_analysis: str, task_formulas: str, task_modeling: str, task_result: str, user_prompt: str = ''):
        prompt = TASK_ANSWER_PROMPT.format(task_description=task_description, task_analysis=task_analysis, task_formulas=task_formulas, task_modeling=task_modeling, task_result=task_result, user_prompt=user_prompt).strip()
        return self.llm.generate(prompt)

    def extract_code_structure(self, task_id, code: str, save_path: str):
        prompt = CODE_STRUCTURE_PROMPT.format(code=code, save_path=save_path)
        count = 0
        for i in range(5):
            try:
                strucutre = self.llm.generate(prompt)
                structure_string = strucutre.strip('```json\n').strip('```')
                structure_json = json.loads(structure_string)
                for i in range(len(structure_json['file_outputs'])):
                    structure_json['file_outputs'][i]['file_description'] = 'This file is generated by code for Task {}. '.format(task_id) + structure_json['file_outputs'][i]['file_description']
                return structure_json
            except:
                continue
        if count == 5:
            sys.exit("Fail at extract_code_structure")