Update app.py

app.py

@@ -10,9 +10,22 @@ from PIL import Image
 import io
 import matplotlib as plt
 
+from PIL import Image
+
 api_key = os.getenv("OPENAI_API_KEY")
 assistant_id = os.getenv("assistant_id")
 
+def handle_and_display_image(file_id):
+    api_response = client.files.with_raw_response.retrieve_content(file_id)
+    if api_response.status_code == 200:
+        content = api_response.content
+        with open('image.png', 'wb') as f:
+            f.write(content)
+        return Image.open('image.png')
+    else:
+        print('Failed to download file.')
+        return None
+
 client = openai.Client()
 
 MODEL = "gpt-4-1106-preview"
@@ -59,7 +72,6 @@ elif hasattr(st.session_state.run, 'status') and st.session_state.run.status ==
         thread_id=st.session_state.thread.id
     )
 
-    
     # Display messages
     for message in reversed(st.session_state.messages.data):
         if message.role in ["user", "assistant"]:
@@ -68,16 +80,11 @@ elif hasattr(st.session_state.run, 'status') and st.session_state.run.status ==
                     if content_part.type == 'text':  # For text responses
                         message_text = content_part.text.value
                         st.markdown(message_text)
-                    elif content_part.type == 'image_url':  # For image URLs
-                        image_url = content_part.image_url  # Adjust based on actual response structure
-                        st.image(image_url)
-                    elif content_part.type == 'base64_image':  # For base64-encoded images
-                        base64_data = content_part.base64_image  # Adjust based on actual response structure
-                        image_data = base64.b64decode(base64_data)
-                        image = Image.open(io.BytesIO(image_data))
-                        st.image(image)
-                   
-    
+                    # New code to handle image files
+                    elif hasattr(content_part, 'image_file') and content_part.image_file:  
+                        image = handle_and_display_image(content_part.image_file.file_id)
+                        if image:
+                            st.image(image, caption="Generated Image")
 
 
 if prompt := st.chat_input("Wie kann ich dir helfen?"):