Update app.py

app.py

@@ -5,8 +5,7 @@ import os
 import re
 import shutil
 import datetime
-import mimetypes
-from typing import Optional, List, Dict, Tuple
+from typing import Optional, List, Dict, Tuple, Generator, Union
 
 from dotenv import load_dotenv
 from huggingface_hub import login
@@ -150,7 +149,7 @@ class ModelManager:
     @staticmethod
     def load_model(
         chosen_inference: str, model_id: str, key_manager: Optional[object] = None
-    ):
+    ) -> Union[HfApiModel, LiteLLMModel, OpenAIServerModel, TransformersModel]:
         """Load the specified model with appropriate configuration.
 
         Args:
@@ -275,15 +274,16 @@ def create_agent() -> CodeAgent:
 
         try:
             doc_tools = ToolRegistry.load_document_tools()
-        except AssertionError as e:
+        except Exception as e:
             print(f"Warning: Error loading document tools: {str(e)}")
             print("Attempting to continue with available tools...")
+            doc_tools = []
 
         image_generator = ToolRegistry.load_image_generation_tools()
 
         # Combine available tools (filter out None values)
         all_tools = [visualizer] + web_tools + doc_tools
-        if image_generator:  # Add only if it's not None
+        if image_generator:
             all_tools.append(image_generator)
 
         # Log available tools
@@ -305,11 +305,11 @@ def create_agent() -> CodeAgent:
 
 
 def stream_to_gradio(
-    agent,
+    agent: CodeAgent,
     task: str,
     reset_agent_memory: bool = False,
     additional_args: Optional[dict] = None,
-):
+) -> Generator[gr.ChatMessage, None, None]:
     """Runs an agent with the given task and streams messages as Gradio ChatMessages."""
     try:
         for step_log in agent.run(
@@ -354,10 +354,10 @@ def stream_to_gradio(
 class GradioUI:
     """A one-line interface to launch your agent in Gradio."""
 
-    def __init__(self, file_upload_folder: str | None = None):
+    def __init__(self, file_upload_folder: Optional[str] = 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
+        self.allowed_extensions = ALLOWED_EXTENSIONS
 
         if self.file_upload_folder:
             os.makedirs(self.file_upload_folder, exist_ok=True)
@@ -385,7 +385,7 @@ class GradioUI:
                     )
                 )
                 yield messages
-                return  # Exit if can't create agent
+                return
 
         session_state["request_count"] += 1
 
@@ -419,7 +419,7 @@ class GradioUI:
         except Exception as e:
             prompt_with_files = prompt
             print(
-                f"WARNING: Error processing files: {e}.  Continuing without file info."
+                f"WARNING: Error processing files: {e}. Continuing without file info."
             )
 
         try:
@@ -455,7 +455,14 @@ class GradioUI:
         )
 
     def upload_file(self, files: List[str]) -> Tuple[str, List[str]]:
-        """Handle file uploads with validation, security, and clear feedback."""
+        """Handle file uploads with validation, security, and clear feedback.
+
+        Args:
+            files: List of file paths to upload
+
+        Returns:
+            Tuple of (status message, list of uploaded file paths)
+        """
         if not files:
             return "No file uploaded", []
 
@@ -523,52 +530,14 @@ class GradioUI:
     def launch(self, **kwargs):
         """Launch the Gradio UI with responsive layout."""
         with gr.Blocks(theme="ocean", fill_height=True) as demo:
-            # JavaScript-based layout
-            demo.load(
-                None,  # No server-side function on load
-                [],
-                [],
-                _js="""
-                () => {
-                    const isMobile = /Android|webOS|iPhone|iPad|iPod|BlackBerry|IEMobile|Opera Mini/i.test(navigator.userAgent);
-                    if (isMobile) {
-                        // Mobile-specific logic (e.g., hide/show certain elements)
-                        document.body.classList.add('mobile-layout'); // Example: Add a CSS class
-                        console.log('Mobile layout applied.');
-                    } else {
-                        // Desktop-specific logic
-                         document.body.classList.add('desktop-layout');
-                        console.log('Desktop layout applied.');
-                    }
-                }
-                """,
-            )
-
-            desktop_layout = self._create_desktop_layout()
-            mobile_layout = self._create_mobile_layout()
-
-            # Conditional visibility of the layout
-            js_function = """
-            () => {
-                const isMobile = /Android|webOS|iPhone|iPad|iPod|BlackBerry|IEMobile|Opera Mini/i.test(navigator.userAgent);
-                if (isMobile) {
-                   return ['block','none'];
-                } else {
-                    return ['none','block'];
-                }
-            }
-            """
-
-            # Set the visibility of the layouts based on the clients device
-            mobile_layout.style = {"display": "none"}  # Hide by default
-            desktop_layout.style = {"display": "none"}  # Hide by default
-
-            gr.on(
-                [],
-                [],
-                [mobile_layout, desktop_layout],
-                js_function,
-            )
+            # Different layouts for mobile and computer devices
+            @gr.render()
+            def layout(request: gr.Request):
+                device = self.detect_device(request)
+                print(f"Device detected: {device}")
+                if device == "Desktop":
+                    return self._create_desktop_layout()
+                return self._create_mobile_layout()
 
         demo.queue(max_size=20).launch(debug=True, **kwargs)
 
@@ -642,16 +611,14 @@ class GradioUI:
 
             clear_btn.click(
                 lambda: ([], [], {"agent": session_state.get("agent")}, []),
-                None,
-                [chatbot, stored_messages, session_state, uploaded_files_state],
+                outputs=[chatbot, stored_messages, session_state, uploaded_files_state],
             )
 
-            if self.file_upload_folder:
-                file_upload.change(
-                    self.upload_file,
-                    [file_upload],
-                    [upload_status, uploaded_files_state],
-                )
+            file_upload.change(
+                self.upload_file,
+                inputs=[file_upload],
+                outputs=[upload_status, uploaded_files_state],
+            )
 
             self._connect_event_handlers(
                 text_input,
@@ -667,7 +634,7 @@ class GradioUI:
     def _create_mobile_layout(self):
         """Create the mobile layout (simpler without sidebar)."""
         with gr.Column(visible=False) as simple_demo:
-            gr.Markdown("""#OpenDeepResearch - free the AI agents!""")
+            gr.Markdown("# OpenDeepResearch - free the AI agents!")
             session_state = gr.State({})
             stored_messages = gr.State([])
             file_upload = gr.File(
@@ -678,7 +645,7 @@ class GradioUI:
             uploaded_files_state = gr.State([])
 
             chatbot = gr.Chatbot(
-                label="open-Deep-Research",
+                label="OpenDeepResearch",
                 type="messages",
                 avatar_images=(
                     None,
@@ -718,52 +685,51 @@ class GradioUI:
 
     def _connect_event_handlers(
         self,
-        text_input,
-        launch_research_btn,
-        stored_messages,
-        chatbot,
-        session_state,
-        uploaded_files_state,
-    ):
-        """Connect the event handlers for input elements."""
+        text_input: gr.Textbox,
+        launch_btn: gr.Button,
+        stored_messages: gr.State,
+        chatbot: gr.Chatbot,
+        session_state: gr.State,
+        uploaded_files_state: gr.State,
+    ) -> None:
+        """Connect the event handlers for input elements with proper flow control."""
+
+        # Text input submission handler
         text_input.submit(
             self.log_user_message,
-            [text_input],
-            [stored_messages, text_input, launch_research_btn],
+            inputs=[text_input],
+            outputs=[text_input, text_input, launch_btn],
         ).then(
             self.interact_with_agent,
-            [stored_messages, chatbot, session_state, uploaded_files_state],
-            [chatbot],
+            inputs=[text_input, stored_messages, session_state, uploaded_files_state],
+            outputs=[chatbot],
         ).then(
             lambda: (
                 gr.Textbox(
-                    interactive=True,
-                    placeholder="Enter your prompt here and press the button",
+                    value="", interactive=True, placeholder="Enter your prompt here..."
                 ),
                 gr.Button(interactive=True),
             ),
-            None,
-            [text_input, launch_research_btn],
+            outputs=[text_input, launch_btn],
         )
 
-        launch_research_btn.click(
+        # Button click handler (same flow)
+        launch_btn.click(
             self.log_user_message,
-            [text_input],
-            [stored_messages, text_input, launch_research_btn],
+            inputs=[text_input],
+            outputs=[text_input, text_input, launch_btn],
         ).then(
             self.interact_with_agent,
-            [stored_messages, chatbot, session_state, uploaded_files_state],
-            [chatbot],
+            inputs=[text_input, stored_messages, session_state, uploaded_files_state],
+            outputs=[chatbot],
         ).then(
             lambda: (
                 gr.Textbox(
-                    interactive=True,
-                    placeholder="Enter your prompt here and press the button",
+                    value="", interactive=True, placeholder="Enter your prompt here..."
                 ),
                 gr.Button(interactive=True),
             ),
-            None,
-            [text_input, launch_research_btn],
+            outputs=[text_input, launch_btn],
         )
 
 
@@ -772,7 +738,14 @@ def main():
     """Main entry point for the application."""
     setup_environment()
     os.makedirs(f"./{BROWSER_CONFIG['downloads_folder']}", exist_ok=True)
-    GradioUI(file_upload_folder="uploaded_files").launch()
+
+    # Set up file upload directory
+    upload_folder = os.path.join(os.getcwd(), "uploads")
+    os.makedirs(upload_folder, exist_ok=True)
+
+    # Initialize and launch the UI
+    ui = GradioUI(file_upload_folder=upload_folder)
+    ui.launch(share=True)
 
 
 if __name__ == "__main__":