🧹 black cleanup

codeinterpreterapi/callbacks.py

@@ -11,7 +11,7 @@ class CodeCallbackHandler(AsyncIteratorCallbackHandler):
     def __init__(self, session: "CodeInterpreterSession"):
         self.session = session
         super().__init__()
-    
+
     async def on_agent_action(
         self,
         action: AgentAction,
@@ -26,4 +26,4 @@ class CodeCallbackHandler(AsyncIteratorCallbackHandler):
                 f"⚙️ Running code: ```python\n{action.tool_input['code']}\n```"  # type: ignore
             )
         else:
-            raise ValueError(f"Unknown action: {action.tool}")
+            raise ValueError(f"Unknown action: {action.tool}")

codeinterpreterapi/chains/functions_agent.py

@@ -105,9 +105,11 @@ def _format_intermediate_steps(
         messages.extend(_convert_agent_action_to_messages(agent_action, observation))
 
     return messages
-    
 
-async def _parse_ai_message(message: BaseMessage, llm: BaseLanguageModel) -> Union[AgentAction, AgentFinish]:
+
+async def _parse_ai_message(
+    message: BaseMessage, llm: BaseLanguageModel
+) -> Union[AgentAction, AgentFinish]:
     """Parse an AI message."""
     if not isinstance(message, AIMessage):
         raise TypeError(f"Expected an AI message got {type(message)}")
@@ -117,7 +119,7 @@ async def _parse_ai_message(message: BaseMessage, llm: BaseLanguageModel) -> Uni
     if function_call:
         function_call = message.additional_kwargs["function_call"]
         function_name = function_call["name"]
-        try:    
+        try:
             _tool_input = json.loads(function_call["arguments"])
         except JSONDecodeError:
             if function_name == "python":
@@ -199,8 +201,9 @@ class OpenAIFunctionsAgent(BaseSingleActionAgent):
     def functions(self) -> List[dict]:
         return [dict(format_tool_to_openai_function(t)) for t in self.tools]
 
-    def plan(self): raise NotImplementedError
-    
+    def plan(self):
+        raise NotImplementedError
+
     async def aplan(
         self,
         intermediate_steps: List[Tuple[AgentAction, str]],
@@ -229,7 +232,7 @@ class OpenAIFunctionsAgent(BaseSingleActionAgent):
         )
         agent_decision = await _parse_ai_message(predicted_message, self.llm)
         return agent_decision
-    
+
     @classmethod
     def create_prompt(
         cls,

codeinterpreterapi/chains/modifications_check.py

@@ -18,28 +18,27 @@ from langchain.schema import (
 prompt = ChatPromptTemplate(
     input_variables=["code"],
     messages=[
-        SystemMessage(content=
-            "The user will input some code and you will need to determine if the code makes any changes to the file system. \n"
+        SystemMessage(
+            content="The user will input some code and you will need to determine if the code makes any changes to the file system. \n"
             "With changes it means creating new files or modifying exsisting ones.\n"
             "Answer with a function call `determine_modifications` and list them inside.\n"
             "If the code does not make any changes to the file system, still answer with the function call but return an empty list.\n",
         ),
-        HumanMessagePromptTemplate.from_template("{code}")
-    ]
+        HumanMessagePromptTemplate.from_template("{code}"),
+    ],
 )
 
 functions = [
     {
         "name": "determine_modifications",
-        "description": 
-            "Based on code of the user determine if the code makes any changes to the file system. \n"
-            "With changes it means creating new files or modifying exsisting ones.\n",
+        "description": "Based on code of the user determine if the code makes any changes to the file system. \n"
+        "With changes it means creating new files or modifying exsisting ones.\n",
         "parameters": {
             "type": "object",
             "properties": {
                 "modifications": {
                     "type": "array",
-                    "items": { "type": "string" },
+                    "items": {"type": "string"},
                     "description": "The filenames that are modified by the code.",
                 },
             },
@@ -50,30 +49,30 @@ functions = [
 
 
 async def get_file_modifications(
-    code: str, 
+    code: str,
     llm: BaseLanguageModel,
     retry: int = 2,
 ) -> List[str] | None:
-    if retry < 1: 
+    if retry < 1:
         return None
     messages = prompt.format_prompt(code=code).to_messages()
     message = await llm.apredict_messages(messages, functions=functions)
-    
+
     if not isinstance(message, AIMessage):
         raise OutputParserException("Expected an AIMessage")
-    
+
     function_call = message.additional_kwargs.get("function_call", None)
-    
+
     if function_call is None:
-        return await get_file_modifications(code, llm, retry=retry-1)
-    else: 
+        return await get_file_modifications(code, llm, retry=retry - 1)
+    else:
         function_call = json.loads(function_call["arguments"])
         return function_call["modifications"]
-    
+
 
 async def test():
     llm = ChatOpenAI(model="gpt-3.5-turbo-0613")  # type: ignore
-    
+
     code = """
     import matplotlib.pyplot as plt
 
@@ -89,15 +88,16 @@ async def test():
     """
 
     code2 = "import pandas as pd\n\n# Read the Excel file\ndata = pd.read_excel('Iris.xlsx')\n\n# Convert the data to CSV\ndata.to_csv('Iris.csv', index=False)"
-    
+
     modifications = await get_file_modifications(code2, llm)
-    
+
     print(modifications)
 
 
 if __name__ == "__main__":
     import asyncio
     from dotenv import load_dotenv
+
     load_dotenv()
-    
+
     asyncio.run(test())

codeinterpreterapi/chains/remove_download_link.py

@@ -8,51 +8,54 @@ from langchain.schema import (
     AIMessage,
     OutputParserException,
     SystemMessage,
-    HumanMessage
+    HumanMessage,
 )
 
 
 prompt = ChatPromptTemplate(
     input_variables=["input_response"],
     messages=[
-        SystemMessage(content=
-            "The user will send you a response and you need to remove the download link from it.\n"
+        SystemMessage(
+            content="The user will send you a response and you need to remove the download link from it.\n"
             "Reformat the remaining message so no whitespace or half sentences are still there.\n"
             "If the response does not contain a download link, return the response as is.\n"
         ),
-        HumanMessage(content="The dataset has been successfully converted to CSV format. You can download the converted file [here](sandbox:/Iris.csv)."),
+        HumanMessage(
+            content="The dataset has been successfully converted to CSV format. You can download the converted file [here](sandbox:/Iris.csv)."
+        ),
         AIMessage(content="The dataset has been successfully converted to CSV format."),
-        HumanMessagePromptTemplate.from_template("{input_response}")
-    ]
+        HumanMessagePromptTemplate.from_template("{input_response}"),
+    ],
 )
 
 
 async def remove_download_link(
-    input_response: str, 
+    input_response: str,
     llm: BaseLanguageModel,
 ) -> str:
     messages = prompt.format_prompt(input_response=input_response).to_messages()
     message = await llm.apredict_messages(messages)
-    
+
     if not isinstance(message, AIMessage):
         raise OutputParserException("Expected an AIMessage")
-    
+
     return message.content
-    
+
 
 async def test():
     llm = ChatOpenAI(model="gpt-3.5-turbo-0613")  # type: ignore
-    
+
     example = "I have created the plot to your dataset.\n\nLink to the file [here](sandbox:/plot.png)."
-    
+
     modifications = await remove_download_link(example, llm)
-    
+
     print(modifications)
 
 
 if __name__ == "__main__":
     import asyncio
     import dotenv
+
     dotenv.load_dotenv()
-    
+
     asyncio.run(test())

codeinterpreterapi/config.py

@@ -9,10 +9,11 @@ class CodeInterpreterAPISettings(BaseSettings):
     """
     CodeInterpreter API Config
     """
+
     VERBOSE: bool = False
-    
+
     CODEBOX_API_KEY: str | None = None
     OPENAI_API_KEY: str | None = None
-    
+
 
 settings = CodeInterpreterAPISettings()

codeinterpreterapi/prompts/__init__.py

@@ -1 +1 @@
-from .system_message import system_message as code_interpreter_system_message
+from .system_message import system_message as code_interpreter_system_message

codeinterpreterapi/prompts/system_message.py

@@ -1,7 +1,8 @@
 from langchain.schema import SystemMessage
 
 
-system_message = SystemMessage(content="""
+system_message = SystemMessage(
+    content="""
 Assistant is designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. 
 As a language model, Assistant is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.
 Assistant is constantly learning and improving, and its capabilities are constantly evolving. 
@@ -13,4 +14,5 @@ The human also maybe thinks this code interpreter is for writing code but it is
 Tell the human if they use the code interpreter incorrectly.
 Already installed packages are: (numpy pandas matplotlib seaborn scikit-learn yfinance scipy statsmodels sympy bokeh plotly dash networkx).
 If you encounter an error, try again and fix the code.
-""")
+"""
+)

codeinterpreterapi/schema/file.py

@@ -5,53 +5,58 @@ from pydantic import BaseModel
 class File(BaseModel):
     name: str
     content: bytes
-    
+
     @classmethod
     def from_path(cls, path: str):
         with open(path, "rb") as f:
             path = path.split("/")[-1]
             return cls(name=path, content=f.read())
-      
-    @classmethod  
+
+    @classmethod
     async def afrom_path(cls, path: str):
-        await asyncio.to_thread(cls.from_path, path)    
-    
+        await asyncio.to_thread(cls.from_path, path)
+
     @classmethod
     def from_url(cls, url: str):
         import requests  # type: ignore
+
         r = requests.get(url)
         return cls(name=url.split("/")[-1], content=r.content)
-    
+
     @classmethod
     async def afrom_url(cls, url: str):
         import aiohttp
+
         async with aiohttp.ClientSession() as session:
             async with session.get(url) as r:
                 return cls(name=url.split("/")[-1], content=await r.read())
-    
+
     def save(self, path: str):
         with open(path, "wb") as f:
             f.write(self.content)
-    
+
     async def asave(self, path: str):
         await asyncio.to_thread(self.save, path)
-    
+
     def show_image(self):
         try:
             from PIL import Image  # type: ignore
         except ImportError:
-            print("Please install it with `pip install codeinterpreterapi[image_support]` to display images.")
+            print(
+                "Please install it with `pip install codeinterpreterapi[image_support]` to display images."
+            )
             exit(1)
-        
+
         from io import BytesIO
+
         img_io = BytesIO(self.content)
         img = Image.open(img_io)
-        
+
         # Display the image
         img.show()
-    
+
     def __str__(self):
         return self.name
-    
+
     def __repr__(self):
-        return f"File(name={self.name})"
+        return f"File(name={self.name})"

codeinterpreterapi/schema/input.py

@@ -1,9 +1,9 @@
 from pydantic import BaseModel
 
 
-class CodeInput(BaseModel): 
+class CodeInput(BaseModel):
     code: str
-    
 
-class FileInput(BaseModel): 
+
+class FileInput(BaseModel):
     filename: str

codeinterpreterapi/schema/response.py

@@ -4,21 +4,21 @@ from .file import File
 
 class UserRequest(HumanMessage):
     files: list[File] = []
-    
+
     def __str__(self):
         return self.content
-    
+
     def __repr__(self):
         return f"UserRequest(content={self.content}, files={self.files})"
 
 
 class CodeInterpreterResponse(AIMessage):
     files: list[File] = []
-    # final_code: str = ""  TODO: implement 
+    # final_code: str = ""  TODO: implement
     # final_output: str = ""  TODO: implement
-    
+
     def __str__(self):
         return self.content
-    
+
     def __repr__(self):
-        return f"CodeInterpreterResponse(content={self.content}, files={self.files})"
+        return f"CodeInterpreterResponse(content={self.content}, files={self.files})"

codeinterpreterapi/session.py

@@ -9,7 +9,7 @@ from langchain.prompts.chat import MessagesPlaceholder
 from langchain.agents import AgentExecutor, BaseSingleActionAgent
 from langchain.memory import ConversationBufferMemory
 
-from codeinterpreterapi.schema import CodeInterpreterResponse, CodeInput, File, UserRequest # type: ignore
+from codeinterpreterapi.schema import CodeInterpreterResponse, CodeInput, File, UserRequest  # type: ignore
 from codeinterpreterapi.config import settings
 from codeinterpreterapi.chains.functions_agent import OpenAIFunctionsAgent
 from codeinterpreterapi.prompts import code_interpreter_system_message
@@ -18,8 +18,7 @@ from codeinterpreterapi.chains.modifications_check import get_file_modifications
 from codeinterpreterapi.chains.remove_download_link import remove_download_link
 
 
-class CodeInterpreterSession():
-    
+class CodeInterpreterSession:
     def __init__(self, model=None, openai_api_key=None) -> None:
         self.codebox = CodeBox()
         self.tools: list[StructuredTool] = self._tools()
@@ -27,46 +26,46 @@ class CodeInterpreterSession():
         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()
-    
+
     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,
+                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
-        
+
         return ChatOpenAI(
             temperature=0.03,
-            model=model, 
+            model=model,
             openai_api_key=openai_api_key,
             max_retries=3,
-            request_timeout=60*3,
+            request_timeout=60 * 3,
         )  # type: ignore
-    
+
     def _agent(self) -> BaseSingleActionAgent:
         return OpenAIFunctionsAgent.from_llm_and_tools(
             llm=self.llm,
@@ -74,7 +73,7 @@ class CodeInterpreterSession():
             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(),
@@ -84,23 +83,23 @@ class CodeInterpreterSession():
             verbose=settings.VERBOSE,
             memory=ConversationBufferMemory(memory_key="memory", return_messages=True),
         )
-    
+
     async def show_code(self, code: str) -> None:
-        """ Callback function to show code to the user. """
+        """Callback function to show code to the user."""
         if settings.VERBOSE:
             print(code)
-    
-    def run_handler(self, code: str): 
+
+    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 """
+        """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))
@@ -108,58 +107,64 @@ class CodeInterpreterSession():
             self.output_files.append(File(name=filename, content=file_buffer.read()))
             return f"Image {filename} got send to the user."
 
-        elif output.type == "error": 
+        elif output.type == "error":
             if "ModuleNotFoundError" in output.content:
-                if package := re.search(r"ModuleNotFoundError: No module named '(.*)'", 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."
             # TODO: preanalyze error to optimize next code generation
             print("Error:", output.content)
-        
-        elif (modifications := await get_file_modifications(code, self.llm)):
+
+        elif modifications := await get_file_modifications(code, self.llm):
             for filename in modifications:
-                if filename in [file.name for file in self.input_files]: 
+                if filename in [file.name for file in self.input_files]:
                     continue
                 fileb = await self.codebox.adownload(filename)
-                if not fileb.content: 
+                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()))
-            
+                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: 
+        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 = (
+                "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 """
+        """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_msg: str, 
+        self,
+        user_msg: str,
         files: list[File] = [],
         detailed_error: bool = False,
     ) -> CodeInterpreterResponse:
-        """ Generate a Code Interpreter response based on the user's input."""
+        """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)
@@ -168,20 +173,21 @@ class CodeInterpreterSession():
         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}"
+                return CodeInterpreterResponse(
+                    content=f"Error in CodeInterpreterSession: {e.__class__.__name__}  - {e}"
                 )
             else:
-                return CodeInterpreterResponse(content=
-                    "Sorry, something went while generating your response."
+                return CodeInterpreterResponse(
+                    content="Sorry, something went while generating your response."
                     "Please try again or restart the session."
                 )
 
     async def __aenter__(self) -> "CodeInterpreterSession":
         await self._init()
         return self
-    
+
     async def __aexit__(self, exc_type, exc_value, traceback) -> None:
         await self._close()

examples/convert_file.py

@@ -8,13 +8,14 @@ async def main():
         files = [
             File.from_path("examples/iris.csv"),
         ]
-        
+
         output = await session.generate_response(user_request, files=files)
         file = output.files[0]
-        
-        file.save("examples/iris.xlsx")     
+
+        file.save("examples/iris.xlsx")
 
 
 if __name__ == "__main__":
     import asyncio
+
     asyncio.run(main())

examples/plot_sin_wave.py

@@ -4,12 +4,14 @@ from codeinterpreterapi import CodeInterpreterSession
 async def main():
     async with CodeInterpreterSession() as session:
         user_request = "Plot a sin wave and show it to me."
-        
+
         output = await session.generate_response(user_request)
-        
+
         file = output.files[0]
         file.show_image()
 
+
 if __name__ == "__main__":
     import asyncio
+
     asyncio.run(main())

examples/show_bitcoin_chart.py

@@ -6,9 +6,9 @@ async def main():
     async with CodeInterpreterSession() as session:
         currentdate = datetime.now().strftime("%Y-%m-%d")
         user_request = f"Plot the bitcoin chart of 2023 YTD (today is {currentdate})"
-        
+
         output = await session.generate_response(user_request)
-        
+
         file = output.files[0]
         file.show_image()
 
@@ -21,4 +21,5 @@ async def main():
 
 if __name__ == "__main__":
     import asyncio
+
     asyncio.run(main())