implement Session

codeinterpreterapi/session.py

@@ -1,46 +1,176 @@
-from codeboxapi import CodeBox
-from promptkit import ChatSession
+import uuid, base64, re
+from io import BytesIO
+from codeboxapi import CodeBox  # type: ignore
+from codeboxapi.schema import CodeBoxOutput  # type: ignore
+from langchain.tools import StructuredTool
+from langchain.chat_models import ChatOpenAI
+from langchain.chat_models.base import BaseChatModel
+from langchain.prompts.chat import MessagesPlaceholder
+from langchain.agents import AgentExecutor, BaseSingleActionAgent
+from langchain.memory import ConversationBufferMemory
+
+from codeinterpreterapi.schemas import CodeInterpreterResponse, CodeInput, File, UserRequest # type: ignore
+from codeinterpreterapi.config import settings
+from codeinterpreterapi.functions_agent import OpenAIFunctionsAgent
+from codeinterpreterapi.prompts import code_interpreter_system_message
+from codeinterpreterapi.callbacks import CodeCallbackHandler
+from codeinterpreterapi.chains.modifications_check import get_file_modifications
+from codeinterpreterapi.chains.remove_download_link import remove_download_link
 
 
 class CodeInterpreterSession():
     
-    def __init__(self):
-        self.chatgpt = ChatSession()
+    def __init__(self, model=None, openai_api_key=None) -> None:
         self.codebox = CodeBox()
+        self.tools: list[StructuredTool] = self._tools()
+        self.llm: BaseChatModel = self._llm(model, openai_api_key)
+        self.agent_executor: AgentExecutor = self._agent_executor()
+        self.input_files: list[File] = []
+        self.output_files: list[File] = []
+    
+    async def _init(self) -> None:
+        await self.codebox.astart()
+    
+    async def _close(self) -> None:
+        await self.codebox.astop()
     
-    async def _init(self):
-        await self.codebox.start()
+    def _tools(self) -> list[StructuredTool]:
+        return [
+            StructuredTool(
+                name = "python",
+                description = 
+                    # TODO: variables as context to the agent
+                    # TODO: current files as context to the agent
+                    "Input a string of code to a python interpreter (jupyter kernel). "
+                    "Variables are preserved between runs. ",
+                func = self.codebox.run,
+                coroutine = self.arun_handler,
+                args_schema = CodeInput,
+            ),
+        ]
+        
+    def _llm(self, model: str | None, openai_api_key: str | None) -> BaseChatModel:
+        if model is None:
+            model = "gpt-4"
+        
+        if openai_api_key is None:
+            if settings.OPENAI_API_KEY is None:
+                raise ValueError("OpenAI API key missing.")
+            else:
+                openai_api_key = settings.OPENAI_API_KEY
         
-    async def _close(self):
-        await self.codebox.stop()
+        return ChatOpenAI(
+            temperature=0.03,
+            model=model, 
+            openai_api_key=openai_api_key,
+            max_retries=3,
+            request_timeout=60*3,
+        )  # type: ignore
     
-    async def code_decision(self, user_request: str):
-        # check if the user wants something that requires python code execution
-        # if yes, return "code"
-        # if no, return "default"
-        pass
+    def _agent(self) -> BaseSingleActionAgent:
+        return OpenAIFunctionsAgent.from_llm_and_tools(
+            llm=self.llm,
+            tools=self.tools,
+            system_message=code_interpreter_system_message,
+            extra_prompt_messages=[MessagesPlaceholder(variable_name="memory")],
+        )
+        
+    def _agent_executor(self) -> AgentExecutor:
+        return AgentExecutor.from_agent_and_tools(
+            agent=self._agent(),
+            callbacks=[CodeCallbackHandler(self)],
+            max_iterations=9,
+            tools=self.tools,
+            verbose=settings.VERBOSE,
+            memory=ConversationBufferMemory(memory_key="memory", return_messages=True),
+        )
     
-    async def generate_response(self, text: str, files: list[dict[str, bytes]]):  # list of "file_name" x "file_content"
-        """ Generate a Code Interpreter response based on the user's input."""
-        if self.code_decision() == "code":
-            pass
-            # plan what code to write (potentially multiple steps)
-            # code = chatgpt.run(code generation template)
-            # codebox.run(code)
-                # on error
-                    # check if package is required
-                        # if yes, install package
-                    # ask for analysis if the error can be fixed
-                        # if yes, continue code generation
-                        # if no, return AssistantResponse
-                # on success
-                    # check if to output files to the user
-                        # if yes, return AssistantResponse with files
-            # write a response based on the code execution
-            # return AssistantResponse
-        else:
-            pass
-            # return AssistantResponse
-        pass
+    async def show_code(self, code: str) -> None:
+        """ Callback function to show code to the user. """
+        if settings.VERBOSE:
+            print(code)
+    
+    def run_handler(self, code: str): 
+        raise NotImplementedError("Use arun_handler for now.")
+    
+    async def arun_handler(self, code: str):
+        """ Run code in container and send the output to the user """
+        # TODO: upload files
+        output: CodeBoxOutput = await self.codebox.arun(code)
+        
+        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(discord.File(path_like_file, filename)) TODO: add to output_files
+            return f"Image {filename} got send to the user."
 
-    
+        elif output.type == "error": 
+            # TODO: check if package install is required
+            # TODO: preanalyze error to optimize next code generation
+            print("Error:", output.content)
+        
+        elif (modifications := await get_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
+    
+    async def input_handler(self, request: UserRequest):
+        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"
+            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:
+        """ 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(rf"\n\n!\[.*\]\(.*\)", "", final_response)
+        
+        if self.output_files and re.search(rf"\n\[.*\]\(.*\)", final_response):
+            final_response = await remove_download_link(final_response, self.llm)
+        
+        return CodeInterpreterResponse(content=final_response, files=self.output_files)
+    
+    async def generate_response(
+        self, 
+        user_request: UserRequest, 
+        files: list[File], # list of "file_name" x "file_content"
+        detailed_error: bool = False,
+    ) -> CodeInterpreterResponse:
+        """ Generate a Code Interpreter response based on the user's input."""
+        try:
+            await self.input_handler(user_request)
+            response = await self.agent_executor.arun(input=user_request.content)
+            return await self.output_handler(response)
+        except Exception as e:
+            if settings.VERBOSE:
+                import traceback
+                traceback.print_exc()
+            if detailed_error:
+                return CodeInterpreterResponse(content=
+                    f"Error in CodeInterpreterSession: {e.__class__.__name__}  - {e}"
+                )
+            else:
+                return CodeInterpreterResponse(content=
+                    "Sorry, something went while generating your response."
+                    "Please try again or restart the session."
+                )