import json import copy import shutil from jupyter_backend import * from tools import * from typing import * from notebook_serializer import add_markdown_to_notebook, add_code_cell_to_notebook functions = [ { "name": "execute_code", "description": "This function allows you to execute Python code and retrieve the terminal output. If the code " "generates image output, the function will return the text '[image]'. The code is sent to a " "Jupyter kernel for execution. The kernel will remain active after execution, retaining all " "variables in memory.", "parameters": { "type": "object", "properties": { "code": { "type": "string", "description": "The code text" } }, "required": ["code"], } }, ] system_msg = '''You are an AI code interpreter. Your goal is to help users do a variety of jobs by executing Python code. You should: 1. Comprehend the user's requirements carefully & to the letter. 2. Give a brief description for what you plan to do & call the provided function to run code. 3. Provide results analysis based on the execution output. 4. If error occurred, try to fix it. 5. Response in the same language as the user. Note: 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. ''' with open('config.json') as f: config = json.load(f) if not config['API_KEY']: config['API_KEY'] = os.getenv('OPENAI_API_KEY') os.unsetenv('OPENAI_API_KEY') def get_config(): return config def config_openai_api(api_type, api_base, api_version, api_key): openai.api_type = api_type openai.api_base = api_base openai.api_version = api_version openai.api_key = api_key class GPTResponseLog: def __init__(self): self.assistant_role_name = '' self.content = '' self.function_name = None self.function_args_str = '' self.code_str = '' self.display_code_block = '' self.finish_reason = 'stop' self.bot_history = None self.stop_generating = False self.code_executing = False self.interrupt_signal_sent = False def reset_gpt_response_log_values(self, exclude=None): if exclude is None: exclude = [] attributes = {'assistant_role_name': '', 'content': '', 'function_name': None, 'function_args_str': '', 'code_str': '', 'display_code_block': '', 'finish_reason': 'stop', 'bot_history': None, 'stop_generating': False, 'code_executing': False, 'interrupt_signal_sent': False} for attr_name in exclude: del attributes[attr_name] for attr_name, value in attributes.items(): setattr(self, attr_name, value) def set_assistant_role_name(self, assistant_role_name: str): self.assistant_role_name = assistant_role_name def add_content(self, content: str): self.content += content def set_function_name(self, function_name: str): self.function_name = function_name def copy_current_bot_history(self, bot_history: List): self.bot_history = copy.deepcopy(bot_history) def add_function_args_str(self, function_args_str: str): self.function_args_str += function_args_str def update_code_str(self, code_str: str): self.code_str = code_str def update_display_code_block(self, display_code_block): self.display_code_block = display_code_block def update_finish_reason(self, finish_reason: str): self.finish_reason = finish_reason def update_stop_generating_state(self, stop_generating: bool): self.stop_generating = stop_generating def update_code_executing_state(self, code_executing: bool): self.code_executing = code_executing def update_interrupt_signal_sent(self, interrupt_signal_sent: bool): self.interrupt_signal_sent = interrupt_signal_sent class BotBackend(GPTResponseLog): def __init__(self): super().__init__() self.unique_id = hash(id(self)) self.jupyter_work_dir = f'cache/work_dir_{self.unique_id}' self.tool_log = f'cache/tool_{self.unique_id}.log' self.jupyter_kernel = JupyterKernel(work_dir=self.jupyter_work_dir) self.gpt_model_choice = "GPT-3.5" self.revocable_files = [] self.system_msg = system_msg self.functions = copy.deepcopy(functions) self._init_api_config() self._init_tools() self._init_conversation() self._init_kwargs_for_chat_completion() def _init_conversation(self): first_system_msg = {'role': 'system', 'content': self.system_msg} self.context_window_tokens = 0 # num of tokens actually sent to GPT self.sliced = False # whether the conversion is sliced if hasattr(self, 'conversation'): self.conversation.clear() self.conversation.append(first_system_msg) else: self.conversation: List[Dict] = [first_system_msg] def _init_api_config(self): self.config = get_config() api_type = self.config['API_TYPE'] api_base = self.config['API_base'] api_version = self.config['API_VERSION'] api_key = config['API_KEY'] config_openai_api(api_type, api_base, api_version, api_key) def _init_tools(self): self.additional_tools = {} tool_datas = get_available_tools(self.config) if tool_datas: self.system_msg += '\n\nAdditional tools:' for tool_data in tool_datas: system_prompt = tool_data['system_prompt'] tool_name = tool_data['tool_name'] tool_description = tool_data['tool_description'] self.system_msg += f'\n{tool_name}: {system_prompt}' self.functions.append(tool_description) self.additional_tools[tool_name] = { 'tool': tool_data['tool'], 'additional_parameters': copy.deepcopy(tool_data['additional_parameters']) } for parameter, value in self.additional_tools[tool_name]['additional_parameters'].items(): if callable(value): self.additional_tools[tool_name]['additional_parameters'][parameter] = value(self) def _init_kwargs_for_chat_completion(self): self.kwargs_for_chat_completion = { 'stream': True, 'messages': self.conversation, 'functions': self.functions, 'function_call': 'auto' } model_name = self.config['model'][self.gpt_model_choice]['model_name'] if self.config['API_TYPE'] == 'azure': self.kwargs_for_chat_completion['engine'] = model_name else: self.kwargs_for_chat_completion['model'] = model_name def _backup_all_files_in_work_dir(self): count = 1 backup_dir = f'cache/backup_{self.unique_id}' while os.path.exists(backup_dir): count += 1 backup_dir = f'cache/backup_{self.unique_id}_{count}' shutil.copytree(src=self.jupyter_work_dir, dst=backup_dir) def _clear_all_files_in_work_dir(self, backup=True): if backup: self._backup_all_files_in_work_dir() for filename in os.listdir(self.jupyter_work_dir): path = os.path.join(self.jupyter_work_dir, filename) if os.path.isdir(path): shutil.rmtree(path) else: os.remove(path) def _save_tool_log(self, tool_response): with open(self.tool_log, 'a', encoding='utf-8') as log_file: log_file.write(f'Previous conversion: {self.conversation}\n') log_file.write(f'Model choice: {self.gpt_model_choice}\n') log_file.write(f'Tool name: {self.function_name}\n') log_file.write(f'Parameters: {self.function_args_str}\n') log_file.write(f'Response: {tool_response}\n') log_file.write('----------\n\n') def add_gpt_response_content_message(self): self.conversation.append( {'role': self.assistant_role_name, 'content': self.content} ) add_markdown_to_notebook(self.content, title="Assistant") def add_text_message(self, user_text): self.conversation.append( {'role': 'user', 'content': user_text} ) self.revocable_files.clear() self.update_finish_reason(finish_reason='new_input') add_markdown_to_notebook(user_text, title="User") def add_file_message(self, path, bot_msg): filename = os.path.basename(path) work_dir = self.jupyter_work_dir shutil.copy(path, work_dir) gpt_msg = {'role': 'system', 'content': f'User uploaded a file: {filename}'} self.conversation.append(gpt_msg) self.revocable_files.append( { 'bot_msg': bot_msg, 'gpt_msg': gpt_msg, 'path': os.path.join(work_dir, filename) } ) def add_function_call_response_message(self, function_response: Union[str, None], save_tokens=True): if self.code_str is not None: add_code_cell_to_notebook(self.code_str) self.conversation.append( { "role": self.assistant_role_name, "name": self.function_name, "content": self.function_args_str } ) if function_response is not None: if save_tokens and len(function_response) > 500: function_response = f'{function_response[:200]}\n[Output too much, the middle part output is omitted]\n ' \ f'End part of output:\n{function_response[-200:]}' self.conversation.append( { "role": "function", "name": self.function_name, "content": function_response, } ) self._save_tool_log(tool_response=function_response) def append_system_msg(self, prompt): self.conversation.append( {'role': 'system', 'content': prompt} ) def revoke_file(self): if self.revocable_files: file = self.revocable_files[-1] bot_msg = file['bot_msg'] gpt_msg = file['gpt_msg'] path = file['path'] assert self.conversation[-1] is gpt_msg del self.conversation[-1] os.remove(path) del self.revocable_files[-1] return bot_msg else: return None def update_gpt_model_choice(self, model_choice): self.gpt_model_choice = model_choice self._init_kwargs_for_chat_completion() def update_token_count(self, num_tokens): self.__setattr__('context_window_tokens', num_tokens) def update_sliced_state(self, sliced): self.__setattr__('sliced', sliced) def send_interrupt_signal(self): self.jupyter_kernel.send_interrupt_signal() self.update_interrupt_signal_sent(interrupt_signal_sent=True) def restart(self): self.revocable_files.clear() self._init_conversation() self.reset_gpt_response_log_values() self.jupyter_kernel.restart_jupyter_kernel() self._clear_all_files_in_work_dir()