Update app.py

app.py

@@ -3,6 +3,9 @@ import os
 import re
 import shutil
 from typing import Optional
+import io
+import zipfile
+from datetime import datetime
 
 from dotenv import load_dotenv
 from huggingface_hub import login
@@ -146,7 +149,7 @@ class ToolRegistry:
     def load_web_tools(model, browser, text_limit=20000):
         """Initialize and return web-related tools."""
         return [
-            GoogleSearchTool(provider="serper"),  
+            GoogleSearchTool(provider="serper"),
             VisitTool(browser),
             PageUpTool(browser),
             PageDownTool(browser),
@@ -176,8 +179,8 @@ def create_agent():
     # Initialize model
     model = LiteLLMModel(
         custom_role_conversions=custom_role_conversions,
-        model_id="openrouter/perplexity/r1-1776", # currently serving: 
-    )   # DEEPSEEK = openrouter/perplexity/r1-1776 <--- boss model
+        model_id="openrouter/perplexity/r1-1776",  # currently serving:
+    )  # DEEPSEEK = openrouter/perplexity/r1-1776 <--- boss model
 
     # Initialize tools
     text_limit = 30000
@@ -278,6 +281,7 @@ class GradioUI:
             for msg in stream_to_gradio(
                 session_state["agent"], task=prompt, reset_agent_memory=False
             ):
+                msg = self.track_generated_files(msg, session_state)
                 messages.append(msg)
                 yield messages  # Yield messages after each step
             yield messages  # Yield messages one last time
@@ -286,6 +290,48 @@ class GradioUI:
             print(f"Error in interaction: {str(e)}")
             raise
 
+    def track_generated_files(self, message, session_state):
+        """Track file paths in session state for later export"""
+        if isinstance(message, gr.ChatMessage) and isinstance(message.content, dict):
+            if "image" in message.content:
+                file_path = message.content["image"]
+                session_state.setdefault("generated_files", []).append(file_path)
+            elif "audio" in message.content:
+                file_path = message.content["audio"]
+                session_state.setdefault("generated_files", []).append(file_path)
+        return message
+
+    def export_chat(self, chat_history):
+        """Export conversation history as markdown file"""
+        md_content = "# Conversation History\n\n"
+        for msg in chat_history:
+            role = "User" if msg.role == "user" else "Assistant"
+            content = msg.content
+            if isinstance(content, dict):
+                if "image" in content:
+                    content = f"![Generated Image]({content['image']})"
+                elif "audio" in content:
+                    content = f"🔊 [Audio File]({content['audio']})"
+            md_content += f"**{role}**: {content}\n\n"
+
+        filename = f"conversation-{datetime.now().strftime('%Y%m%d-%H%M%S')}.md"
+        return md_content.encode("utf-8"), filename
+
+    def export_files(self, session_state):
+        """Package generated media files into a zip"""
+        if not session_state.get("generated_files"):
+            return None
+
+        zip_buffer = io.BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
+            for file_path in session_state["generated_files"]:
+                if os.path.exists(file_path):
+                    zf.write(file_path, os.path.basename(file_path))
+
+        zip_buffer.seek(0)
+        filename = f"generated-files-{datetime.now().strftime('%Y%m%d-%H%M%S')}.zip"
+        return zip_buffer.getvalue(), filename
+
     def upload_file(
         self,
         file,
@@ -385,30 +431,14 @@ class GradioUI:
         # Default case if no clear indicators
         return "Desktop"
 
-    def launch(self, **kwargs):
-        """Launch the Gradio UI with responsive layout."""
-        with gr.Blocks(theme="ocean", fill_height=True) as demo:
-            # Different layouts for mobile and computer devices
-            @gr.render()
-            def layout(request: gr.Request):
-                device = self.detect_device(request)
-                print(f"device - {device}")
-                # Render layout with sidebar
-                if device == "Desktop":
-                    return self._create_desktop_layout()
-                else:
-                    return self._create_mobile_layout()
-
-        demo.queue(max_size=20).launch(
-            debug=True, **kwargs
-        )  # Add queue with reasonable size
-
     def _create_desktop_layout(self):
         """Create the desktop layout with sidebar."""
         with gr.Blocks(fill_height=True) as sidebar_demo:
             with gr.Sidebar():
-                gr.Markdown("""#OpenDeepResearch - 3theSmolagents!
-                Model_id: R1-1776""")
+                gr.Markdown(
+                    """#OpenDeepResearch - 3theSmolagents!
+                Model_id: R1-1776"""
+                )
                 with gr.Group():
                     gr.Markdown("**What's on your mind mate?**", container=True)
                     text_input = gr.Textbox(
@@ -432,6 +462,13 @@ class GradioUI:
                         [upload_status, file_uploads_log],
                     )
 
+                # Add download buttons row
+                with gr.Row():
+                    self.chat_download = gr.DownloadButton("💾 Save Chat", visible=True)
+                    self.file_download = gr.DownloadButton(
+                        "📦 Export Media", visible=True
+                    )
+
                 gr.HTML("<br><br><h4><center>Powered by:</center></h4>")
                 with gr.Row():
                     gr.HTML(
@@ -514,6 +551,11 @@ class GradioUI:
             )
             launch_research_btn = gr.Button("Run", variant="primary")
 
+            # Add compact mobile download buttons
+            with gr.Row():
+                self.chat_download = gr.DownloadButton("💾 Chat", visible=True)
+                self.file_download = gr.DownloadButton("📦 Media", visible=True)
+
             self._connect_event_handlers(
                 text_input,
                 launch_research_btn,
@@ -577,6 +619,36 @@ class GradioUI:
             [text_input, launch_research_btn],
         )
 
+    def launch(self, **kwargs):
+        """Launch the Gradio UI with responsive layout."""
+        with gr.Blocks(theme="ocean", fill_height=True) as demo:
+            # Different layouts for mobile and computer devices
+            @gr.render()
+            def layout(request: gr.Request):
+                device = self.detect_device(request)
+                print(f"device - {device}")
+                # Render layout with sidebar
+                if device == "Desktop":
+                    return self._create_desktop_layout()
+                else:
+                    return self._create_mobile_layout()
+
+            # Configure download button click handlers
+            self.chat_download.click(
+                self.export_chat,
+                inputs=[chatbot],
+                outputs=gr.File(label="Chat History"),
+            )
+            self.file_download.click(
+                self.export_files,
+                inputs=[session_state],
+                outputs=gr.File(label="Media Files"),
+            )
+
+        demo.queue(max_size=20).launch(
+            debug=True, **kwargs
+        )  # Add queue with reasonable size
+
 
 # ------------------------ Execution ------------------------
 def main():