Update app.py

app.py

@@ -639,38 +639,43 @@ class XylariaChat:
             
         def streaming_response(message, chat_history, image_filepath, math_ocr_image_path):
             if message.strip().lower() == "/image":
-                chat_history.append([message, loading_svg]) # Display loading SVG
-                yield "", chat_history, None, None 
-                
+                # Extract image generation prompt from the last user message before /image
                 image_prompt = ""
-                # Find the last user message before the /image command for context
-                for i in range(len(chat_history) - 2, -1, -1):
+                for i in range(len(chat_history) - 1, -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"
 
+                # Generate image based on the extracted prompt
                 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 in the AI's chat bubble
-                    chat_history[-1][1] = image_html
-                    
+                    # Append the /image command and generated image to chat history
+                    chat_history.append([message, ""]) 
+                    chat_history.append(["", image_html])
+
+                    # Update conversation history
                     self.conversation_history.append(ChatMessage(role="user", content=message).to_dict())
-                    self.conversation_history.append(ChatMessage(role="assistant", content=image_html).to_dict()) # Correctly append image to conversation history
+                    self.conversation_history.append(ChatMessage(role="assistant", content=image_html).to_dict())
 
-                    yield "", chat_history, None, None
+                    # Update the chat history file
+                    self.save_chat()
+                    all_chats = self.load_all_chats()
+                    chat_titles = [f"{chat['timestamp']}: {chat['conversation'][0]['content'][:30]}..." if len(chat['conversation']) > 0 and chat['conversation'][0]['content'] else f"{chat['timestamp']}: Empty Chat" for chat in all_chats]
+                    
+                    yield "", chat_history, None, None, gr.update(choices=chat_titles, visible=True)
                 else:
-                    chat_history[-1][1] = image_bytes # Error message
-                    yield "", chat_history, None, None
+                    chat_history.append([message, image_bytes])
+                    yield "", chat_history, None, None, None
                 return
             
-            ocr_text = ""
+                        ocr_text = ""
             if math_ocr_image_path:
                 ocr_text = self.perform_math_ocr(math_ocr_image_path)
                 if ocr_text.startswith("Error"):