Update app.py

app.py

@@ -585,26 +585,35 @@ class XylariaChat:
             
         def streaming_response(message, chat_history, image_filepath, math_ocr_image_path):
             if message.strip().lower() == "/image":
-                chat_history.append([message, ""])
-                yield "", chat_history, None, None, None
-
-                image_prompt = self.conversation_history[-2]["content"] if len(self.conversation_history) >= 2 else ""
+                chat_history.append([message, loading_svg]) # Display loading SVG
+                yield "", chat_history, None, None 
+                
+                image_prompt = ""
+                # Find the last user message before the /image command for context
+                for i in range(len(chat_history) - 2, -1, -1):
+                    if chat_history[i][0] and chat_history[i][0].strip().lower() != "/image":
+                        image_prompt = chat_history[i][0]
+                        break
+                
                 if not image_prompt:
                     image_prompt = "A realistic image"
 
                 image_bytes = self.generate_image(image_prompt)
+
                 if isinstance(image_bytes, bytes):
                     base64_image = base64.b64encode(image_bytes).decode("utf-8")
                     image_html = f'<img src="data:image/png;base64,{base64_image}" alt="Generated Image" style="max-width: 100%; max-height: 400px;">'
+                    
+                    # Replace loading SVG with the generated image
                     chat_history[-1][1] = image_html
                     
                     self.conversation_history.append(ChatMessage(role="user", content=message).to_dict())
                     self.conversation_history.append(ChatMessage(role="assistant", content=image_html).to_dict())
-                    
-                    yield "", chat_history, None, None, None
+
+                    yield "", chat_history, None, None
                 else:
-                    chat_history[-1][1] = image_bytes
-                    yield "", chat_history, None, None, None
+                    chat_history[-1][1] = image_bytes # Error message
+                    yield "", chat_history, None, None
                 return
             
             ocr_text = ""