chenxc1029/Local-Code-Interpreter

1import json2import openai3import os4import copy5import shutil6from jupyter_backend import *7from typing import *8 9functions = [10    {11        "name": "execute_code",12        "description": "This function allows you to execute Python code and retrieve the terminal output. If the code "13                       "generates image output, the function will return the text '[image]'. The code is sent to a "14                       "Jupyter kernel for execution. The kernel will remain active after execution, retaining all "15                       "variables in memory.",16        "parameters": {17            "type": "object",18            "properties": {19                "code": {20                    "type": "string",21                    "description": "The code text"22                }23            },24            "required": ["code"],25        }26    }27]28 29system_msg = '''You are an AI code interpreter.30Your goal is to help users do a variety of jobs by executing Python code.31 32You should:331. Comprehend the user's requirements carefully & 34 35 36 37to the letter. 382. Give a brief description for what you plan to do & call the execute_code function to run code393. Provide results analysis based on the execution output. 404. If error occurred, try to fix it.41 42Note: If the user uploads a file, you will receive a system message "User uploaded a file: filename". Use the filename as the path in the code. '''43 44with open('config.json') as f:45    config = json.load(f)46 47if not config['API_KEY']:48    config['API_KEY'] = os.getenv('OPENAI_API_KEY')49    os.unsetenv('OPENAI_API_KEY')50 51 52def get_config():53    return config54 55 56def config_openai_api(api_type, api_base, api_version, api_key):57    openai.api_type = api_type58    openai.api_base = api_base59    openai.api_version = api_version60    openai.api_key = api_key61 62 63class GPTResponseLog:64    def __init__(self):65        self.assistant_role_name = ''66        self.content = ''67        self.function_name = None68        self.function_args_str = ''69        self.display_code_block = ''70        self.finish_reason = 'stop'71        self.bot_history = None72 73    def reset_gpt_response_log_values(self, exclude=None):74        if exclude is None:75            exclude = []76 77        attributes = {'assistant_role_name': '',78                      'content': '',79                      'function_name': None,80                      'function_args_str': '',81                      'display_code_block': '',82                      'finish_reason': 'stop',83                      'bot_history': None}84 85        for attr_name in exclude:86            del attributes[attr_name]87        for attr_name, value in attributes.items():88            setattr(self, attr_name, value)89 90    def set_assistant_role_name(self, assistant_role_name: str):91        self.assistant_role_name = assistant_role_name92 93    def add_content(self, content: str):94        self.content += content95 96    def set_function_name(self, function_name: str):97        self.function_name = function_name98 99    def copy_current_bot_history(self, bot_history: List):100        self.bot_history = copy.deepcopy(bot_history)101 102    def add_function_args_str(self, function_args_str: str):103        self.function_args_str += function_args_str104 105    def update_display_code_block(self, display_code_block):106        self.display_code_block = display_code_block107 108    def update_finish_reason(self, finish_reason: str):109        self.finish_reason = finish_reason110 111 112class BotBackend(GPTResponseLog):113    def __init__(self):114        super().__init__()115        self.unique_id = hash(id(self))116        self.jupyter_work_dir = f'cache/work_dir_{self.unique_id}'117        self.jupyter_kernel = JupyterKernel(work_dir=self.jupyter_work_dir)118        self.gpt_model_choice = "GPT-3.5"119        self.revocable_files = []120        self._init_conversation()121        self._init_api_config()122        self._init_kwargs_for_chat_completion()123 124    def _init_conversation(self):125        first_system_msg = {'role': 'system', 'content': system_msg}126        if hasattr(self, 'conversation'):127            self.conversation.clear()128            self.conversation.append(first_system_msg)129        else:130            self.conversation: List[Dict] = [first_system_msg]131 132    def _init_api_config(self):133        self.config = get_config()134        api_type = self.config['API_TYPE']135        api_base = self.config['API_base']136        api_version = self.config['API_VERSION']137        api_key = config['API_KEY']138        config_openai_api(api_type, api_base, api_version, api_key)139 140    def _init_kwargs_for_chat_completion(self):141        self.kwargs_for_chat_completion = {142            'stream': True,143            'messages': self.conversation,144            'functions': functions,145            'function_call': 'auto'146        }147 148        model_name = self.config['model'][self.gpt_model_choice]['model_name']149 150        if self.config['API_TYPE'] == 'azure':151            self.kwargs_for_chat_completion['engine'] = model_name152        else:153            self.kwargs_for_chat_completion['model'] = model_name154 155    def _clear_all_files_in_work_dir(self):156        for filename in os.listdir(self.jupyter_work_dir):157            os.remove(158                os.path.join(self.jupyter_work_dir, filename)159            )160 161    def add_gpt_response_content_message(self):162        self.conversation.append(163            {'role': self.assistant_role_name, 'content': self.content}164        )165 166    def add_text_message(self, user_text):167        self.conversation.append(168            {'role': 'user', 'content': user_text}169        )170        self.revocable_files.clear()171        self.update_finish_reason(finish_reason='new_input')172 173    def add_file_message(self, path, bot_msg):174        filename = os.path.basename(path)175        work_dir = self.jupyter_work_dir176 177        shutil.copy(path, work_dir)178 179        gpt_msg = {'role': 'system', 'content': f'User uploaded a file: {filename}'}180        self.conversation.append(gpt_msg)181        self.revocable_files.append(182            {183                'bot_msg': bot_msg,184                'gpt_msg': gpt_msg,185                'path': os.path.join(work_dir, filename)186            }187        )188 189    def add_function_call_response_message(self, function_response: str, save_tokens=True):190        self.conversation.append(191            {192                "role": self.assistant_role_name,193                "name": self.function_name,194                "content": self.function_args_str195            }196        )197 198        if save_tokens and len(function_response) > 500:199            function_response = f'{function_response[:200]}\n[Output too much, the middle part output is omitted]\n ' \200                                f'End part of output:\n{function_response[-200:]}'201        self.conversation.append(202            {203                "role": "function",204                "name": self.function_name,205                "content": function_response,206            }207        )208 209    def revoke_file(self):210        if self.revocable_files:211            file = self.revocable_files[-1]212            bot_msg = file['bot_msg']213            gpt_msg = file['gpt_msg']214            path = file['path']215 216            assert self.conversation[-1] is gpt_msg217            del self.conversation[-1]218 219            os.remove(path)220 221            del self.revocable_files[-1]222 223            return bot_msg224        else:225            return None226 227    def update_gpt_model_choice(self, model_choice):228        self.gpt_model_choice = model_choice229        self._init_kwargs_for_chat_completion()230 231    def restart(self):232        self._clear_all_files_in_work_dir()233        self.revocable_files.clear()234        self._init_conversation()235        self.reset_gpt_response_log_values()236        self.jupyter_kernel.restart_jupyter_kernel()237