🧹 black cleanup

Author: shroominic

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 @@ async def on_agent_action(
                 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 @@ def input_keys(self) -> List[str]:
     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 @@ async def aplan(
         )
         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 @@
 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 @@
 
 
 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 @@
     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 @@
 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 @@
 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 @@
 
 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})"