Removing useless functions from file.py (from_bytes) + UI update + is now possible to download the processed files and to provide any kind of file

codeinterpreterapi/config.py

@@ -14,7 +14,7 @@ class CodeInterpreterAPISettings(BaseSettings):
     VERBOSE: bool = False
 
     CODEBOX_API_KEY: Optional[str] = None
-    OPENAI_API_KEY: Optional[str] = None
+    OPENAI_API_KEY: Optional[str] = "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"
 
 
 settings = CodeInterpreterAPISettings()

codeinterpreterapi/schema/file.py

@@ -6,10 +6,6 @@ class File(BaseModel):
     name: str
     content: bytes
 
-    @classmethod
-    def from_bytes(cls, name: str, content: bytes):
-        return cls(name=name, content=content)
-
     @classmethod
     def from_path(cls, path: str):
         with open(path, "rb") as f:

frontend/{demo.py → app.py}

@@ -5,6 +5,9 @@ import streamlit as st
 from codeinterpreterapi import File
 from frontend.utils import get_images
 
+# Page configuration
+st.set_page_config(layout="wide")
+
 st.title('Code Interpreter API 🚀')
 
 # This will create a sidebar
@@ -14,18 +17,17 @@ st.sidebar.markdown("[Github Repo](https://github.com/shroominic/codeinterpreter
 
 # This will create a textbox where you can input text
 input_text = st.text_area("Write your prompt")
-uploaded_files = st.file_uploader("Choose a CSV file", accept_multiple_files=True,
-                                  type=".csv")
+uploaded_files = st.file_uploader("Upload your files", accept_multiple_files=True)
 
 uploaded_files_list = []
 for uploaded_file in uploaded_files:
     bytes_data = uploaded_file.read()
-    uploaded_files_list.append(File.from_bytes(name=uploaded_file.name,
-                                               content=bytes_data))
+    uploaded_files_list.append(File(name=uploaded_file.name,
+                                    content=bytes_data))
 
 # This will create a button
 button_pressed = st.button('Run code interpreter', use_container_width=True)
 
-# This will display the image only when the button is pressed
+# This will display the images only when the button is pressed
 if button_pressed and input_text != "":
-    asyncio.run(get_images(input_text, files=uploaded_files_list))
+    asyncio.run(get_images(input_text, files=uploaded_files_list))

frontend/utils.py

@@ -1,5 +1,16 @@
 from codeinterpreterapi import CodeInterpreterSession
 import streamlit as st
+import tempfile
+import os
+import shutil
+
+
+def create_temp_folder() -> str:
+    """
+    Creates a temp folder
+    """
+    temp_folder = tempfile.mkdtemp()
+    return temp_folder
 
 
 async def get_images(prompt: str, files: list = None):
@@ -17,5 +28,24 @@ async def get_images(prompt: str, files: list = None):
             with st.chat_message("assistant"):
                 st.write(response.content)
 
-                for file in response.files:
-                    st.image(file.get_image(), caption=prompt, use_column_width=True)
+                # Showing Results
+                for _file in response.files:
+                    st.image(_file.get_image(), caption=prompt, use_column_width=True)
+
+                # Allowing the download of the results
+                if len(response.files) == 1:
+                    st.download_button('Download Results', response.files[0].content,
+                                       file_name=response.files[0].name,
+                                       use_container_width=True)
+                else:
+                    target_path = tempfile.mkdtemp()
+                    for _file in response.files:
+                        _file.save(os.path.join(target_path, _file.name))
+
+                    zip_path = os.path.join(os.path.dirname(target_path), "archive")
+                    shutil.make_archive(zip_path, 'zip', target_path)
+
+                    with open(zip_path + ".zip", 'rb') as f:
+                        st.download_button('Download Results', f,
+                                           file_name="archive.zip",
+                                           use_container_width=True)