From a48896aa1e25c7d5985228d859ad9d8a11f374c5 Mon Sep 17 00:00:00 2001 From: jinno Date: Wed, 29 May 2024 01:39:21 +0900 Subject: [PATCH] fix: split python.py from tools.py --- src/codeinterpreterapi/config.py | 3 + src/codeinterpreterapi/session.py | 139 +------------------------- src/codeinterpreterapi/tools/tools.py | 28 +----- 3 files changed, 7 insertions(+), 163 deletions(-) diff --git a/src/codeinterpreterapi/config.py b/src/codeinterpreterapi/config.py index 01325e37..96c3e84d 100644 --- a/src/codeinterpreterapi/config.py +++ b/src/codeinterpreterapi/config.py @@ -44,6 +44,9 @@ class CodeInterpreterAPISettings(BaseSettings): # deprecated VERBOSE: bool = DEBUG + # Environment + WORK_DIR: str = "/app/work" + class Config: env_file = "./.env" extra = "ignore" diff --git a/src/codeinterpreterapi/session.py b/src/codeinterpreterapi/session.py index 080b80cf..c17add40 100644 --- a/src/codeinterpreterapi/session.py +++ b/src/codeinterpreterapi/session.py @@ -1,15 +1,10 @@ -import base64 import re -import subprocess -import tempfile import traceback -from io import BytesIO from types import TracebackType from typing import Any, Optional, Type -from uuid import UUID, uuid4 +from uuid import UUID from codeboxapi import CodeBox # type: ignore -from codeboxapi.schema import CodeBoxOutput # type: ignore from gui_agent_loop_core.schema.schema import GuiAgentInterpreterChatResponseStr from langchain.agents import AgentExecutor from langchain.callbacks.base import Callbacks @@ -22,12 +17,7 @@ from langchain_core.tools import BaseTool from codeinterpreterapi.agents.agents import CodeInterpreterAgent -from codeinterpreterapi.chains import ( - aget_file_modifications, - aremove_download_link, - get_file_modifications, - remove_download_link, -) +from codeinterpreterapi.chains import aremove_download_link, remove_download_link from codeinterpreterapi.chat_history import CodeBoxChatMessageHistory from codeinterpreterapi.config import settings from codeinterpreterapi.llm.llm import CodeInterpreterLlm @@ -184,131 +174,6 @@ def _history_backend(self) -> BaseChatMessageHistory: ) ) - def show_code(self, code: str) -> None: - if self.verbose: - print(code) - - async def ashow_code(self, code: str) -> None: - """Callback function to show code to the user.""" - if self.verbose: - print(code) - - def _get_handler_local_command(self, code: str): - with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".py") as temp_file: - temp_file.write(code) - temp_file_path = temp_file.name - - command = f"cd src/codeinterpreterapi/invoke_tasks && invoke -c python run-code-file '{temp_file_path}'" - return command - - def _run_handler_local(self, code: str): - print("_run_handler_local code=", code) - command = self._get_handler_local_command(code) - try: - output_content = subprocess.check_output(command, shell=True, universal_newlines=True) - self.code_log.append((code, output_content)) - return output_content - except subprocess.CalledProcessError as e: - print(f"An error occurred: {e}") - return None - - async def _arun_handler_local(self, code: str): - print("_arun_handler_local code=", code) - command = self._get_handler_local_command(code) - try: - output_content = await subprocess.check_output(command, shell=True, universal_newlines=True) - self.code_log.append((code, output_content)) - return output_content - except subprocess.CalledProcessError as e: - print(f"An error occurred: {e}") - return None - - def _run_handler(self, code: str) -> str: - """Run code in container and send the output to the user""" - self.show_code(code) - output: CodeBoxOutput = self.codebox.run(code) - self.code_log.append((code, output.content)) - - if not isinstance(output.content, str): - raise TypeError("Expected output.content to be a string.") - - if output.type == "image/png": - filename = f"image-{uuid4()}.png" - file_buffer = BytesIO(base64.b64decode(output.content)) - file_buffer.name = filename - self.output_files.append(File(name=filename, content=file_buffer.read())) - return f"Image {filename} got send to the user." - - elif output.type == "error": - if "ModuleNotFoundError" in output.content: - if package := re.search( - r"ModuleNotFoundError: No module named '(.*)'", - output.content, - ): - self.codebox.install(package.group(1)) - return f"{package.group(1)} was missing but " "got installed now. Please try again." - else: - # TODO: pre-analyze error to optimize next code generation - pass - if self.verbose: - print("Error:", output.content) - - elif modifications := get_file_modifications(code, self.llm): - for filename in modifications: - if filename in [file.name for file in self.input_files]: - continue - fileb = self.codebox.download(filename) - if not fileb.content: - continue - file_buffer = BytesIO(fileb.content) - file_buffer.name = filename - self.output_files.append(File(name=filename, content=file_buffer.read())) - - return output.content - - async def _arun_handler(self, code: str) -> str: - """Run code in container and send the output to the user""" - await self.ashow_code(code) - output: CodeBoxOutput = await self.codebox.arun(code) - self.code_log.append((code, output.content)) - - if not isinstance(output.content, str): - raise TypeError("Expected output.content to be a string.") - - if output.type == "image/png": - filename = f"image-{uuid4()}.png" - file_buffer = BytesIO(base64.b64decode(output.content)) - file_buffer.name = filename - self.output_files.append(File(name=filename, content=file_buffer.read())) - return f"Image {filename} got send to the user." - - elif output.type == "error": - if "ModuleNotFoundError" in output.content: - if package := re.search( - r"ModuleNotFoundError: No module named '(.*)'", - output.content, - ): - await self.codebox.ainstall(package.group(1)) - return f"{package.group(1)} was missing but " "got installed now. Please try again." - else: - # TODO: pre-analyze error to optimize next code generation - pass - if self.verbose: - print("Error:", output.content) - - elif modifications := await aget_file_modifications(code, self.llm): - for filename in modifications: - if filename in [file.name for file in self.input_files]: - continue - fileb = await self.codebox.adownload(filename) - if not fileb.content: - continue - file_buffer = BytesIO(fileb.content) - file_buffer.name = filename - self.output_files.append(File(name=filename, content=file_buffer.read())) - - return output.content - def _input_handler(self, request: UserRequest) -> None: """Callback function to handle user input.""" if not request.files: diff --git a/src/codeinterpreterapi/tools/tools.py b/src/codeinterpreterapi/tools/tools.py index 94e77bf6..e4ce2fb3 100644 --- a/src/codeinterpreterapi/tools/tools.py +++ b/src/codeinterpreterapi/tools/tools.py @@ -1,10 +1,8 @@ from langchain_community.tools.shell.tool import BaseTool, ShellTool from langchain_community.tools.tavily_search import TavilySearchResults from langchain_core.language_models import BaseLanguageModel -from langchain_core.tools import StructuredTool -from codeinterpreterapi.config import settings -from codeinterpreterapi.schema import CodeInput +from codeinterpreterapi.tools.python import PythonTools class CodeInterpreterTools: @@ -21,33 +19,11 @@ def __init__( self._llm = llm def get_all_tools(self) -> list[BaseTool]: - self.add_tools_python() + self._additional_tools.extend(PythonTools.get_tools_python()) self.add_tools_shell() self.add_tools_web_search() return self._additional_tools - def add_tools_python(self) -> None: - tools = [ - StructuredTool( - name="python", - description="Input a string of code to a ipython interpreter.\n" - "Write the entire code in a single string.\n" - "This string can be really long.\n" - "Do not start your code with a line break.\n" - "For example, do 'import numpy', not '\\nimport numpy'." - "Variables are preserved between runs. " - + ( - ("You can use all default python packages " f"specifically also these: {settings.CUSTOM_PACKAGES}") - if settings.CUSTOM_PACKAGES - else "" - ), # TODO: or include this in the system message - func=self._run_handler_func, - coroutine=self._arun_handler_func, - args_schema=CodeInput, # type: ignore - ), - ] - self._additional_tools += tools - def add_tools_shell(self) -> None: shell_tool = ShellTool() shell_tool.description = shell_tool.description + f"args {shell_tool.args}".replace("{", "{{").replace(