🚀 multiple improvements

- booth sync/async methods
- status output
- code/output logging in response

Author: shroominic

codeinterpreterapi/session.py

@@ -22,20 +22,22 @@
 from langchain.tools import BaseTool, StructuredTool
 
 from codeinterpreterapi.agents import OpenAIFunctionsAgent
-from codeinterpreterapi.chains import get_file_modifications, remove_download_link
+from codeinterpreterapi.chains import (
+    aget_file_modifications,
+    aremove_download_link,
+    get_file_modifications,
+    remove_download_link,
+)
 from codeinterpreterapi.config import settings
+from codeinterpreterapi.parser import CodeAgentOutputParser, CodeChatAgentOutputParser
 from codeinterpreterapi.prompts import code_interpreter_system_message
 from codeinterpreterapi.schema import (
     CodeInput,
     CodeInterpreterResponse,
     File,
+    SessionStatus,
     UserRequest,
 )
-from codeinterpreterapi.utils import (
-    CodeAgentOutputParser,
-    CodeCallbackHandler,
-    CodeChatAgentOutputParser,
-)
 
 
 class CodeInterpreterSession:
@@ -45,19 +47,20 @@ def __init__(
         additional_tools: list[BaseTool] = [],
         **kwargs,
     ) -> None:
-        self.codebox = CodeBox()
+        self.codebox = CodeBox(**kwargs)
         self.verbose = kwargs.get("verbose", settings.VERBOSE)
         self.tools: list[BaseTool] = self._tools(additional_tools)
         self.llm: BaseLanguageModel = llm or self._choose_llm(**kwargs)
         self.agent_executor: AgentExecutor = self._agent_executor()
         self.input_files: list[File] = []
         self.output_files: list[File] = []
+        self.code_log: list[tuple[str, str]] = []
 
-    def start(self) -> None:
-        self.codebox.start()
+    def start(self) -> SessionStatus:
+        return SessionStatus.from_codebox_status(self.codebox.start())
 
-    async def astart(self) -> None:
-        await self.codebox.astart()
+    async def astart(self) -> SessionStatus:
+        return SessionStatus.from_codebox_status(await self.codebox.astart())
 
     def _tools(self, additional_tools: list[BaseTool]) -> list[BaseTool]:
         return additional_tools + [
@@ -128,7 +131,6 @@ def _choose_agent(self) -> BaseSingleActionAgent:
     def _agent_executor(self) -> AgentExecutor:
         return AgentExecutor.from_agent_and_tools(
             agent=self._choose_agent(),
-            callbacks=[CodeCallbackHandler(self)],
             max_iterations=9,
             tools=self.tools,
             verbose=self.verbose,
@@ -137,18 +139,67 @@ def _agent_executor(self) -> AgentExecutor:
             ),
         )
 
-    async def show_code(self, code: str) -> None:
+    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 _run_handler(self, code: str):
-        raise NotImplementedError("Use arun_handler for now.")
+        """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-{uuid.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: preanalyze 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):
         """Run code in container and send the output to the user"""
-        print("Running code in container...", code)
+        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.")
@@ -176,7 +227,7 @@ async def _arun_handler(self, code: str):
             if self.verbose:
                 print("Error:", output.content)
 
-        elif modifications := await get_file_modifications(code, self.llm):
+        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
@@ -191,7 +242,22 @@ async def _arun_handler(self, code: str):
 
         return output.content
 
-    async def _input_handler(self, request: UserRequest):
+    def _input_handler(self, request: UserRequest) -> None:
+        """Callback function to handle user input."""
+        if not request.files:
+            return
+        if not request.content:
+            request.content = (
+                "I uploaded, just text me back and confirm that you got the file(s)."
+            )
+        request.content += "\n**The user uploaded the following files: **\n"
+        for file in request.files:
+            self.input_files.append(file)
+            request.content += f"[Attachment: {file.name}]\n"
+            self.codebox.upload(file.name, file.content)
+        request.content += "**File(s) are now available in the cwd. **\n"
+
+    async def _ainput_handler(self, request: UserRequest):
         # TODO: variables as context to the agent
         # TODO: current files as context to the agent
         if not request.files:
@@ -207,7 +273,7 @@ async def _input_handler(self, request: UserRequest):
             await self.codebox.aupload(file.name, file.content)
         request.content += "**File(s) are now available in the cwd. **\n"
 
-    async def _output_handler(self, final_response: str) -> CodeInterpreterResponse:
+    def _output_handler(self, final_response: str) -> CodeInterpreterResponse:
         """Embed images in the response"""
         for file in self.output_files:
             if str(file.name) in final_response:
@@ -216,25 +282,98 @@ async def _output_handler(self, final_response: str) -> CodeInterpreterResponse:
 
         if self.output_files and re.search(r"\n\[.*\]\(.*\)", final_response):
             try:
-                final_response = await remove_download_link(final_response, self.llm)
+                final_response = remove_download_link(final_response, self.llm)
             except Exception as e:
                 if self.verbose:
                     print("Error while removing download links:", e)
 
-        return CodeInterpreterResponse(content=final_response, files=self.output_files)
+        output_files = self.output_files
+        code_log = self.code_log
+        self.output_files = []
+        self.code_log = []
+
+        return CodeInterpreterResponse(
+            content=final_response, files=output_files, code_log=code_log
+        )
+
+    async def _aoutput_handler(self, final_response: str) -> CodeInterpreterResponse:
+        """Embed images in the response"""
+        for file in self.output_files:
+            if str(file.name) in final_response:
+                # rm ![Any](file.name) from the response
+                final_response = re.sub(r"\n\n!\[.*\]\(.*\)", "", final_response)
+
+        if self.output_files and re.search(r"\n\[.*\]\(.*\)", final_response):
+            try:
+                final_response = await aremove_download_link(final_response, self.llm)
+            except Exception as e:
+                if self.verbose:
+                    print("Error while removing download links:", e)
+
+        output_files = self.output_files
+        code_log = self.code_log
+        self.output_files = []
+        self.code_log = []
+
+        return CodeInterpreterResponse(
+            content=final_response, files=output_files, code_log=code_log
+        )
+
+    def generate_response_sync(
+        self,
+        user_msg: str,
+        files: list[File] = [],
+        detailed_error: bool = False,
+    ) -> CodeInterpreterResponse:
+        """Generate a Code Interpreter response based on the user's input."""
+        user_request = UserRequest(content=user_msg, files=files)
+        try:
+            self._input_handler(user_request)
+            response = self.agent_executor.run(input=user_request.content)
+            return self._output_handler(response)
+        except Exception as e:
+            if self.verbose:
+                traceback.print_exc()
+            if detailed_error:
+                return CodeInterpreterResponse(
+                    content="Error in CodeInterpreterSession: "
+                    f"{e.__class__.__name__}  - {e}"
+                )
+            else:
+                return CodeInterpreterResponse(
+                    content="Sorry, something went while generating your response."
+                    "Please try again or restart the session."
+                )
 
     async def generate_response(
         self,
         user_msg: str,
         files: list[File] = [],
         detailed_error: bool = False,
+    ) -> CodeInterpreterResponse:
+        print(
+            "DEPRECATION WARNING: Use agenerate_response for async generation.\n"
+            "This function will be converted to sync in the future.\n"
+            "You can use generate_response_sync for now.",
+        )
+        return await self.agenerate_response(
+            user_msg=user_msg,
+            files=files,
+            detailed_error=detailed_error,
+        )
+
+    async def agenerate_response(
+        self,
+        user_msg: str,
+        files: list[File] = [],
+        detailed_error: bool = False,
     ) -> CodeInterpreterResponse:
         """Generate a Code Interpreter response based on the user's input."""
         user_request = UserRequest(content=user_msg, files=files)
         try:
-            await self._input_handler(user_request)
+            await self._ainput_handler(user_request)
             response = await self.agent_executor.arun(input=user_request.content)
-            return await self._output_handler(response)
+            return await self._aoutput_handler(response)
         except Exception as e:
             if self.verbose:
                 traceback.print_exc()
@@ -249,11 +388,24 @@ async def generate_response(
                     "Please try again or restart the session."
                 )
 
-    async def is_running(self) -> bool:
+    def is_running(self) -> bool:
+        return self.codebox.status() == "running"
+
+    async def ais_running(self) -> bool:
         return await self.codebox.astatus() == "running"
 
-    async def astop(self) -> None:
-        await self.codebox.astop()
+    def stop(self) -> SessionStatus:
+        return SessionStatus.from_codebox_status(self.codebox.stop())
+
+    async def astop(self) -> SessionStatus:
+        return SessionStatus.from_codebox_status(await self.codebox.astop())
+
+    def __enter__(self) -> "CodeInterpreterSession":
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        self.stop()
 
     async def __aenter__(self) -> "CodeInterpreterSession":
         await self.astart()