Update app.py

app.py

@@ -112,6 +112,22 @@ ALLOWED_FILE_TYPES = [
     "audio/wav",
     "audio/ogg",
 ]
+ALLOWED_EXTENSIONS = [
+    ".pdf",
+    ".docx",
+    ".txt",
+    ".md",
+    ".json",
+    ".png",
+    ".webp",
+    ".jpeg",
+    ".jpg",
+    ".gif",
+    ".mp4",
+    ".mpeg",
+    ".wav",
+    ".ogg",
+]
 
 
 def setup_environment():
@@ -341,12 +357,17 @@ class GradioUI:
     def __init__(self, file_upload_folder: str | None = None):
         """Initialize the Gradio UI with optional file upload functionality."""
         self.file_upload_folder = file_upload_folder
+        self.allowed_extensions = ALLOWED_EXTENSIONS  # Use the constant
 
         if self.file_upload_folder is not None:
             os.makedirs(self.file_upload_folder, exist_ok=True)
 
     def interact_with_agent(
-        self, prompt: str, messages: List[Dict], session_state: Dict
+        self,
+        prompt: str,
+        messages: List[Dict],
+        session_state: Dict,
+        uploaded_files: List[str],
     ) -> List[Dict]:  # Type hints
         """Main interaction handler with the agent."""
 
@@ -372,12 +393,39 @@ class GradioUI:
         messages.append(gr.ChatMessage(role="user", content=prompt))
         yield messages
 
+        # Process files, linking them to the message
+        file_message = ""
+        if uploaded_files:
+            file_info = {}
+            for file_path in uploaded_files:
+                ext = os.path.splitext(file_path)[1].lower()
+                if ext in [".jpg", ".jpeg", ".png", ".gif", ".webp"]:
+                    category = "images"
+                elif ext in [".mp3", ".wav", ".ogg"]:
+                    category = "audio"
+                else:
+                    category = "documents"
+
+                if category not in file_info:
+                    file_info[category] = []
+                file_info[category].append(os.path.basename(file_path))
+
+            # Format file information for the agent
+            file_message = "\nYou have been provided with these files:\n"
+            for category, files in file_info.items():
+                file_message += f"- {category.capitalize()}: {', '.join(files)}\n"
+
+        prompt_with_files = prompt + file_message
+        prompt_with_files += "\nUse inspect_file_as_text for documents, visualizer for images, and the appropriate tools for audio files."
+
         try:
             # Check if agent should be reset (e.g., if too many requests)
             reset_needed = session_state["request_count"] > 15
 
             for msg in stream_to_gradio(
-                session_state["agent"], task=prompt, reset_agent_memory=reset_needed
+                session_state["agent"],
+                task=prompt_with_files,
+                reset_agent_memory=reset_needed,
             ):
                 messages.append(msg)
                 yield messages
@@ -395,88 +443,63 @@ class GradioUI:
             )
             yield messages
 
-    def upload_file(self, file, file_uploads_log):
-        """Handle file uploads with validation, security, and clear feedback."""
-        if file is None:
-            return gr.Textbox("No file uploaded", visible=True), file_uploads_log
-
-        try:
-            # Get file size and check limit before processing
-            file_size_mb = os.path.getsize(file.name) / (1024 * 1024)  # Size in MB
-            max_file_size_mb = 50  # Define the limit
-
-            if file_size_mb > max_file_size_mb:
-                return (
-                    gr.Textbox(
-                        f"❌ File size ({file_size_mb:.1f} MB) exceeds {max_file_size_mb} MB limit.",
-                        visible=True,
-                    ),
-                    file_uploads_log,
-                )
-
-            # Check MIME type
-            mime_type, _ = mimetypes.guess_type(file.name)  # Correct unpacking
-            if mime_type not in ALLOWED_FILE_TYPES:
-                return (
-                    gr.Textbox(
-                        f"❌ File type '{mime_type or 'unknown'}' is not allowed. Supported types: {', '.join(t.split('/')[-1] for t in ALLOWED_FILE_TYPES)}",
-                        visible=True,
-                    ),
-                    file_uploads_log,
-                )
+    def log_user_message(self, text_input: str) -> Tuple[str, gr.Textbox, gr.Button]:
+        """Process user message log files."""
 
-            # Sanitize file name with better pattern
-            original_name = os.path.basename(file.name)
-            sanitized_name = re.sub(r"[^\w\-.]", "", original_name)
+        return (
+            text_input,
+            gr.Textbox(value="", interactive=False, placeholder="Processing..."),
+            gr.Button(interactive=False),
+        )
 
-            # Save the uploaded file
-            file_path = os.path.join(self.file_upload_folder, sanitized_name)
-            shutil.copy(file.name, file_path)
+    def upload_file(self, files: List[str]) -> Tuple[str, List[str]]:
+        """Handle file uploads with validation, security, and clear feedback."""
+        if not files:
+            return "No file uploaded", []
 
-            return gr.Textbox(
-                f"✓ File uploaded successfully: {os.path.basename(file_path)} ({file_size_mb:.1f} MB)",
-                visible=True,
-            ), file_uploads_log + [file_path]
+        uploaded_files = []
+        error_message = None
 
-        except Exception as e:
-            return (
-                gr.Textbox(f"❌ Upload error: {str(e)}", visible=True),
-                file_uploads_log,
-            )
-
-    def log_user_message(self, text_input, file_uploads_log):
-        """Process user message and handle file references with proper agent types."""
-        message = text_input
+        for file_path in files:  # Iterate through the list of uploaded files
+            try:
+                # Check file extension
+                file_extension = os.path.splitext(file_path)[1].lower()
+                if file_extension not in self.allowed_extensions:
+                    error_message = (
+                        f"❌ File type '{file_extension}' is not allowed. "
+                        f"Supported types: {', '.join(ALLOWED_EXTENSIONS)}"  # Use Constant
+                    )
 
-        if len(file_uploads_log) > 0:
-            # Group files by type for better agent processing
-            file_info = {}
-            for file_path in file_uploads_log:
-                ext = os.path.splitext(file_path)[1].lower()
-                if ext in [".jpg", ".jpeg", ".png", ".gif", ".webp"]:
-                    category = "images"
-                elif ext in [".mp3", ".wav", ".ogg"]:
-                    category = "audio"
+                    return error_message, []  # Return immediately on first error
                 else:
-                    category = "documents"
-
-                if category not in file_info:
-                    file_info[category] = []
-                file_info[category].append(os.path.basename(file_path))
-
-            # Format file information for the agent
-            file_message = "\nYou have been provided with these files:\n"
-            for category, files in file_info.items():
-                file_message += f"- {category.capitalize()}: {', '.join(files)}\n"
+                    file_size_mb = os.path.getsize(file_path) / (
+                        1024 * 1024
+                    )  # Size in MB
+                    max_file_size_mb = 50  # Define the limit
+
+                    if file_size_mb > max_file_size_mb:
+                        error_message = f"❌ File size ({file_size_mb:.1f} MB) exceeds {max_file_size_mb} MB limit."
+                        return error_message, []
+
+                    sanitized_name = re.sub(
+                        r"[^\w\-.]", "", os.path.basename(file_path)
+                    )  # Sanitize
+                    dest_path = os.path.join(self.file_upload_folder, sanitized_name)
+                    shutil.copy(file_path, dest_path)
+                    uploaded_files.append(dest_path)
+                    print(f"Uploaded {file_path} to {dest_path}")
 
-            message += file_message
-            message += "\nUse inspect_file_as_text for documents, visualizer for images, and the appropriate tools for audio files."
+            except Exception as e:
+                error_message = f"❌ Upload error: {str(e)}"
+                return error_message, []  # Immediately return on error
 
-        return (
-            message,
-            gr.Textbox(value="", interactive=False, placeholder="Processing..."),
-            gr.Button(interactive=False),
-        )
+        if error_message:
+            return error_message, []
+        else:
+            return (
+                f"✓ Files uploaded successfully: {', '.join([os.path.basename(f) for f in uploaded_files])}",
+                uploaded_files,
+            )
 
     def detect_device(self, request: gr.Request):
         """Detect whether the user is on mobile or desktop device."""
@@ -550,42 +573,17 @@ class GradioUI:
                 if self.file_upload_folder is not None:
                     with gr.Group():
                         gr.Markdown("📎 Upload Documents")
-                        upload_file = gr.File(
+                        # File input with multiple enabled with the allowed extensions
+                        file_upload = gr.File(
                             label="Upload files for analysis",
-                            file_types=[
-                                "pdf",
-                                "docx",
-                                "txt",
-                                "md",
-                                "csv",
-                                "xlsx",
-                                "jpg",
-                                "png",
-                            ],
+                            file_types=self.allowed_extensions,
                             file_count="multiple",
                         )
+
                         upload_status = gr.Textbox(
                             label="Upload Status", interactive=False, visible=False
                         )
-                        file_uploads_log = gr.State([])
-
-                        # Show uploaded files list
-                        uploaded_files_display = gr.Markdown("No files uploaded yet")
-
-                        upload_file.change(
-                            self.upload_file,
-                            [upload_file, file_uploads_log],
-                            [upload_status, file_uploads_log],
-                        ).then(
-                            lambda files: (
-                                "Uploaded Files:\n"
-                                + "\n".join([f"- {os.path.basename(f)}" for f in files])
-                                if files
-                                else "No files uploaded yet"
-                            ),
-                            [file_uploads_log],
-                            [uploaded_files_display],
-                        )
+                        uploaded_files_state = gr.State([])
 
                 gr.HTML("<br><hr><h4><center>Powered by:</center></h4>")
                 with gr.Row():
@@ -604,8 +602,6 @@ class GradioUI:
             # Main chat area with improved styling
             session_state = gr.State({})
             stored_messages = gr.State([])
-            if "file_uploads_log" not in locals():
-                file_uploads_log = gr.State([])
 
             chatbot = gr.Chatbot(
                 label="OpenDeepResearch Assistant",
@@ -623,19 +619,26 @@ class GradioUI:
 
             # Connect clear button
             clear_btn.click(
-                lambda: ([], [], {"agent": session_state.get("agent")}),
+                lambda: ([], [], {"agent": session_state.get("agent")}, []),
                 None,
-                [chatbot, stored_messages, session_state],
+                [chatbot, stored_messages, session_state, uploaded_files_state],
             )
 
+            if self.file_upload_folder is not None:
+                file_upload.change(
+                    self.upload_file,
+                    [file_upload],
+                    [upload_status, uploaded_files_state],
+                )
+
             # Connect event handlers
             self._connect_event_handlers(
                 text_input,
                 launch_research_btn,
-                file_uploads_log,
                 stored_messages,
                 chatbot,
                 session_state,
+                uploaded_files_state,
             )
 
             return sidebar_demo
@@ -647,7 +650,13 @@ class GradioUI:
             # Add session state to store session-specific data
             session_state = gr.State({})
             stored_messages = gr.State([])
-            file_uploads_log = gr.State([])
+            # File input with multiple enabled with the allowed extensions
+            file_upload = gr.File(
+                label="Upload files for analysis",
+                file_types=self.allowed_extensions,
+                file_count="multiple",
+            )
+            uploaded_files_state = gr.State([])
 
             chatbot = gr.Chatbot(
                 label="open-Deep-Research",
@@ -662,14 +671,13 @@ class GradioUI:
 
             # If an upload folder is provided, enable the upload feature
             if self.file_upload_folder is not None:
-                upload_file = gr.File(label="Upload a file")
                 upload_status = gr.Textbox(
                     label="Upload Status", interactive=False, visible=False
                 )
-                upload_file.change(
+                file_upload.change(
                     self.upload_file,
-                    [upload_file, file_uploads_log],
-                    [upload_status, file_uploads_log],
+                    [file_upload],
+                    [upload_status, uploaded_files_state],
                 )
 
             text_input = gr.Textbox(
@@ -682,10 +690,10 @@ class GradioUI:
             self._connect_event_handlers(
                 text_input,
                 launch_research_btn,
-                file_uploads_log,
                 stored_messages,
                 chatbot,
                 session_state,
+                uploaded_files_state,
             )
 
             return simple_demo
@@ -694,20 +702,20 @@ class GradioUI:
         self,
         text_input,
         launch_research_btn,
-        file_uploads_log,
         stored_messages,
         chatbot,
         session_state,
+        uploaded_files_state,
     ):
         """Connect the event handlers for input elements."""
         # Connect text input submit event
         text_input.submit(
             self.log_user_message,
-            [text_input, file_uploads_log],
+            [text_input],
             [stored_messages, text_input, launch_research_btn],
         ).then(
             self.interact_with_agent,
-            [stored_messages, chatbot, session_state],
+            [stored_messages, chatbot, session_state, uploaded_files_state],
             [chatbot],
         ).then(
             lambda: (
@@ -724,11 +732,11 @@ class GradioUI:
         # Connect button click event
         launch_research_btn.click(
             self.log_user_message,
-            [text_input, file_uploads_log],
+            [text_input],
             [stored_messages, text_input, launch_research_btn],
         ).then(
             self.interact_with_agent,
-            [stored_messages, chatbot, session_state],
+            [stored_messages, chatbot, session_state, uploaded_files_state],
             [chatbot],
         ).then(
             lambda: (