Update app.py

app.py

@@ -3,7 +3,7 @@ import re
 import shutil
 import datetime
 import mimetypes
-from typing import Optional, List, Dict, Tuple  # More specific typing
+from typing import Optional, List, Dict, Tuple
 
 from dotenv import load_dotenv
 from huggingface_hub import login
@@ -35,7 +35,6 @@ from smolagents import (
 from smolagents.agent_types import AgentText, AgentImage, AgentAudio
 from smolagents.gradio_ui import pull_messages_from_step, handle_agent_output_types
 
-
 # ------------------------ Configuration and Setup ------------------------
 AUTHORIZED_IMPORTS = [
     "requests",
@@ -131,33 +130,30 @@ ALLOWED_EXTENSIONS = [
 
 
 def setup_environment():
-    """
-    Initialize environment variables and authenticate with Hugging Face Hub.
-    """
+    """Initialize environment variables and authenticate with Hugging Face Hub."""
     load_dotenv(override=True)
-    hf_token = os.getenv("HF_TOKEN")  # Get token once
+    hf_token = os.getenv("HF_TOKEN")
     if hf_token:
         login(hf_token)
-        print("HF_TOKEN (last 10 characters):", hf_token[-10:])
+        print(f"HF_TOKEN (last 10 characters): {hf_token[-10:]}")
     else:
         print("HF_TOKEN not found in environment variables.")
 
 
 # ------------------------ Model and Tool Management ------------------------
 class ModelManager:
-    """
-    Manages model loading and initialization.
-    """
+    """Manages model loading and initialization."""
 
     @staticmethod
-    def load_model(chosen_inference: str, model_id: str, key_manager=None):
-        """
-        Load the specified model with appropriate configuration.
+    def load_model(
+        chosen_inference: str, model_id: str, key_manager: Optional[object] = None
+    ):
+        """Load the specified model with appropriate configuration.
 
         Args:
-            chosen_inference (str): The type of inference to use (e.g., "hf_api", "openai").
-            model_id (str): The ID of the model to load.
-            key_manager:  (Optional) Key manager for API keys. Required for OpenAI.
+            chosen_inference: The type of inference to use (e.g., "hf_api", "openai").
+            model_id: The ID of the model to load.
+            key_manager: Key manager for API keys (required for OpenAI).
 
         Returns:
             An instance of the specified model class.
@@ -170,24 +166,23 @@ class ModelManager:
         try:
             if chosen_inference == "hf_api":
                 return HfApiModel(model_id=model_id)
-            elif chosen_inference == "hf_api_provider":
+            if chosen_inference == "hf_api_provider":
                 return HfApiModel(provider="together")
-            elif chosen_inference == "litellm":
+            if chosen_inference == "litellm":
                 return LiteLLMModel(model_id=model_id)
-            elif chosen_inference == "openai":
+            if chosen_inference == "openai":
                 if not key_manager:
                     raise ValueError("Key manager required for OpenAI model")
                 return OpenAIServerModel(
                     model_id=model_id, api_key=key_manager.get_key("openai_api_key")
                 )
-            elif chosen_inference == "transformers":
+            if chosen_inference == "transformers":
                 return TransformersModel(
                     model_id="HuggingFaceTB/SmolLM2-1.7B-Instruct",
                     device_map="auto",
                     max_new_tokens=1000,
                 )
-            else:
-                raise ValueError(f"Invalid inference type: {chosen_inference}")
+            raise ValueError(f"Invalid inference type: {chosen_inference}")
         except Exception as e:
             print(f"✗ Couldn't load model: {e}")
             raise
@@ -197,14 +192,13 @@ class ToolRegistry:
     """Manages tool initialization and organization."""
 
     @staticmethod
-    def load_web_tools(model, browser, text_limit=20000) -> List[Tool]:
-        """
-        Initialize and return web-related tools.
+    def load_web_tools(model, browser, text_limit: int = 20000) -> List[Tool]:
+        """Initialize and return web-related tools.
 
         Args:
             model: The language model to use.
             browser: The web browser instance.
-            text_limit (int): The maximum text length for the text inspector tool.
+            text_limit: The maximum text length for the text inspector tool.
 
         Returns:
             A list of web-related tools.
@@ -222,8 +216,7 @@ class ToolRegistry:
 
     @staticmethod
     def load_document_tools() -> List[Tool]:
-        """
-        Initialize and return document processing tools.
+        """Initialize and return document processing tools.
 
         Returns:
             List of document tools.
@@ -232,8 +225,7 @@ class ToolRegistry:
 
     @staticmethod
     def load_image_generation_tools() -> Optional[Tool]:
-        """
-        Initialize and return image generation tools.
+        """Initialize and return image generation tools.
 
         Returns:
             The image generation tool or None if initialization fails.
@@ -251,8 +243,7 @@ class ToolRegistry:
 
 # ------------------------ Agent Creation and Execution ------------------------
 def create_agent() -> CodeAgent:
-    """
-    Creates a fresh agent instance with configured tools.
+    """Creates a fresh agent instance with configured tools.
 
     Returns:
         CodeAgent: Configured agent ready for use.
@@ -309,7 +300,10 @@ def create_agent() -> CodeAgent:
 
 
 def stream_to_gradio(
-    agent, task: str, reset_agent_memory: bool = False, additional_args=None
+    agent,
+    task: str,
+    reset_agent_memory: bool = False,
+    additional_args: Optional[dict] = None,
 ):
     """Runs an agent with the given task and streams messages as Gradio ChatMessages."""
     try:
@@ -344,9 +338,10 @@ def stream_to_gradio(
                 role="assistant", content=f"Final answer: {str(final_answer)}"
             )
     except Exception as e:
+        error_message = f"Error occurred during processing: {str(e)}\n\nPlease try again with a different query or check your inputs."
         yield gr.ChatMessage(
             role="assistant",
-            content=f"Error occurred during processing: {str(e)}\n\nPlease try again with a different query or check your inputs.",
+            content=error_message,
         )
 
 
@@ -359,7 +354,7 @@ class GradioUI:
         self.file_upload_folder = file_upload_folder
         self.allowed_extensions = ALLOWED_EXTENSIONS  # Use the constant
 
-        if self.file_upload_folder is not None:
+        if self.file_upload_folder:
             os.makedirs(self.file_upload_folder, exist_ok=True)
 
     def interact_with_agent(
@@ -368,57 +363,61 @@ class GradioUI:
         messages: List[Dict],
         session_state: Dict,
         uploaded_files: List[str],
-    ) -> List[Dict]:  # Type hints
+    ) -> List[Dict]:
         """Main interaction handler with the agent."""
 
-        # Get or create session-specific agent with cache persistence
         if "agent" not in session_state:
             try:
                 session_state["agent"] = create_agent()
                 session_state["creation_time"] = datetime.datetime.now()
                 session_state["request_count"] = 0
             except Exception as e:
+                error_message = f"Error initializing agent: {str(e)}\n\nPlease refresh the page and try again."
                 messages.append(
                     gr.ChatMessage(
                         role="assistant",
-                        content=f"Error initializing agent: {str(e)}\n\nPlease refresh the page and try again.",
+                        content=error_message,
                     )
                 )
                 yield messages
-                return
+                return  # Exit if can't create agent
 
         session_state["request_count"] += 1
 
-        # Add user message
         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
+        try:
+            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))
+
+                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
+            else:
+                prompt_with_files = prompt
+        except Exception as e:
+            prompt_with_files = prompt
+            print(
+                f"WARNING: Error processing files: {e}.  Continuing without file info."
+            )
 
         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(
@@ -429,22 +428,21 @@ class GradioUI:
                 messages.append(msg)
                 yield messages
 
-            # If we reset the agent memory, update the request count
             if reset_needed:
                 session_state["request_count"] = 1
 
         except Exception as e:
+            error_message = f"Error processing your request: {str(e)}\n\nPlease try again with a different query."
             messages.append(
                 gr.ChatMessage(
                     role="assistant",
-                    content=f"Error processing your request: {str(e)}\n\nPlease try again with a different query.",
+                    content=error_message,
                 )
             )
             yield messages
 
     def log_user_message(self, text_input: str) -> Tuple[str, gr.Textbox, gr.Button]:
         """Process user message log files."""
-
         return (
             text_input,
             gr.Textbox(value="", interactive=False, placeholder="Processing..."),
@@ -459,65 +457,55 @@ class GradioUI:
         uploaded_files = []
         error_message = None
 
-        for file_path in files:  # Iterate through the list of uploaded files
+        for file_path in 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
+                        f"Supported types: {', '.join(ALLOWED_EXTENSIONS)}"
                     )
+                    return error_message, []
+
+                file_size_mb = os.path.getsize(file_path) / (1024 * 1024)
+                max_file_size_mb = 50
 
-                    return error_message, []  # Return immediately on first error
-                else:
-                    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}")
+                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))
+                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}")
 
             except Exception as e:
                 error_message = f"❌ Upload error: {str(e)}"
-                return error_message, []  # Immediately return on error
+                return error_message, []
 
         if error_message:
             return error_message, []
-        else:
-            return (
-                f"✓ Files uploaded successfully: {', '.join([os.path.basename(f) for f in uploaded_files])}",
-                uploaded_files,
-            )
+
+        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."""
         if not request:
-            return "Unknown device"  # Handle case where request is none.
+            return "Unknown device"
 
-        # Method 1: Check sec-ch-ua-mobile header
         is_mobile_header = request.headers.get("sec-ch-ua-mobile")
         if is_mobile_header:
             return "Mobile" if "?1" in is_mobile_header else "Desktop"
 
-        # Method 2: Check user-agent string
         user_agent = request.headers.get("user-agent", "").lower()
         mobile_keywords = ["android", "iphone", "ipad", "mobile", "phone"]
-
         if any(keyword in user_agent for keyword in mobile_keywords):
             return "Mobile"
 
-        # Method 3: Check platform
         platform = request.headers.get("sec-ch-ua-platform", "").lower()
         if platform:
             if platform in ['"android"', '"ios"']:
@@ -525,25 +513,21 @@ class GradioUI:
             if platform in ['"windows"', '"macos"', '"linux"']:
                 return "Desktop"
 
-        # 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()
                 return self._create_mobile_layout()
 
-        demo.queue(max_size=20).launch(
-            debug=True, **kwargs
-        )  # Add queue with reasonable size
+        demo.queue(max_size=20).launch(debug=True, **kwargs)
 
     def _create_desktop_layout(self):
         """Create the desktop layout with sidebar and enhanced styling."""
@@ -568,11 +552,9 @@ class GradioUI:
                         clear_btn = gr.Button("Clear", variant="secondary")
                         launch_research_btn = gr.Button("Run", variant="primary")
 
-                # File upload section with better labeling
-                if self.file_upload_folder is not None:
+                if self.file_upload_folder:
                     with gr.Group():
                         gr.Markdown("📎 Upload Documents")
-                        # File input with multiple enabled with the allowed extensions
                         file_upload = gr.File(
                             label="Upload files for analysis",
                             file_types=self.allowed_extensions,
@@ -598,7 +580,6 @@ class GradioUI:
                         """
                     )
 
-            # Main chat area with improved styling
             session_state = gr.State({})
             stored_messages = gr.State([])
 
@@ -616,21 +597,19 @@ class GradioUI:
                 height=700,
             )
 
-            # Connect clear button
             clear_btn.click(
                 lambda: ([], [], {"agent": session_state.get("agent")}, []),
                 None,
                 [chatbot, stored_messages, session_state, uploaded_files_state],
             )
 
-            if self.file_upload_folder is not None:
+            if self.file_upload_folder:
                 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,
@@ -646,10 +625,8 @@ class GradioUI:
         """Create the mobile layout (simpler without sidebar)."""
         with gr.Blocks(fill_height=True) as simple_demo:
             gr.Markdown("""#OpenDeepResearch - free the AI agents!""")
-            # Add session state to store session-specific data
             session_state = gr.State({})
             stored_messages = 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,
@@ -668,8 +645,7 @@ class GradioUI:
                 scale=1,
             )
 
-            # If an upload folder is provided, enable the upload feature
-            if self.file_upload_folder is not None:
+            if self.file_upload_folder:
                 upload_status = gr.Textbox(
                     label="Upload Status", interactive=False, visible=False
                 )
@@ -707,7 +683,6 @@ class GradioUI:
         uploaded_files_state,
     ):
         """Connect the event handlers for input elements."""
-        # Connect text input submit event
         text_input.submit(
             self.log_user_message,
             [text_input],
@@ -728,7 +703,6 @@ class GradioUI:
             [text_input, launch_research_btn],
         )
 
-        # Connect button click event
         launch_research_btn.click(
             self.log_user_message,
             [text_input],
@@ -753,13 +727,8 @@ class GradioUI:
 # ------------------------ Execution ------------------------
 def main():
     """Main entry point for the application."""
-    # Initialize environment
     setup_environment()
-
-    # Ensure downloads folder exists
     os.makedirs(f"./{BROWSER_CONFIG['downloads_folder']}", exist_ok=True)
-
-    # Launch UI
     GradioUI(file_upload_folder="uploaded_files").launch()