Update app.py

app.py

@@ -1,11 +1,11 @@
 #!/usr/bin/env python
 # coding=utf-8
 # Copyright 2024 The Footscray Coding Collective. All rights reserved.
+import mimetypes
 import os
 import re
 import shutil
-import datetime
-from typing import Optional, List, Dict, Tuple, Generator, Union
+from typing import Optional
 
 from dotenv import load_dotenv
 from huggingface_hub import login
@@ -38,48 +38,53 @@ from smolagents.agent_types import AgentText, AgentImage, AgentAudio
 from smolagents.gradio_ui import pull_messages_from_step, handle_agent_output_types
 
 # ------------------------ Configuration and Setup ------------------------
+# Constants and configurations
 AUTHORIZED_IMPORTS = [
-    "requests",
-    "zipfile",
-    "pandas",
-    "numpy",
-    "sympy",
-    "json",
-    "bs4",
-    "pubchempy",
+    "requests",  # Web requests (fetching data from the internet)
+    "zipfile",  # Working with ZIP archives
+    "pandas",  # Data manipulation and analysis (DataFrames)
+    "numpy",  # Numerical computing (arrays, linear algebra)
+    "sympy",  # Symbolic mathematics (algebra, calculus)
+    "json",  # JSON data serialization/deserialization
+    "bs4",  # Beautiful Soup for HTML/XML parsing
+    "pubchempy",  # Accessing PubChem chemical database
     "yaml",
-    "xml",
-    "yahoo_finance",
-    "Bio",
-    "sklearn",
-    "scipy",
-    "pydub",
-    "PIL",
-    "chess",
-    "PyPDF2",
-    "pptx",
-    "torch",
-    "datetime",
-    "fractions",
-    "csv",
-    "cleantext",
-    "os",
-    "re",
-    "collections",
-    "math",
-    "random",
-    "io",
-    "urllib.parse",
-    "typing",
-    "concurrent.futures",
-    "time",
-    "tempfile",
-    "matplotlib",
-    "seaborn",
-    "lxml",
-    "selenium",
-    "sqlite3",
-    "schedule",
+    "xml",  # XML processing
+    "yahoo_finance",  # Fetching stock data
+    "Bio",  # Bioinformatics tools (e.g., sequence analysis)
+    "sklearn",  # Scikit-learn for machine learning
+    "scipy",  # Scientific computing (stats, optimization)
+    "pydub",  # Audio manipulation
+    "PIL",  # Pillow for image processing
+    "chess",  # Chess-related functionality
+    "PyPDF2",  # PDF manipulation
+    "pptx",  # PowerPoint file manipulation
+    "torch",  # PyTorch for neural networks
+    "datetime",  # Date and time handling
+    "fractions",  # Rational number arithmetic
+    "csv",  # CSV file reading/writing
+    "cleantext",  # Text cleaning and normalization
+    "os",  # Operating system interaction (file system, etc.) VERY IMPORTANT
+    "re",  # Regular expressions for text processing
+    "collections",  # Useful data structures (e.g., defaultdict, Counter)
+    "math",  # Basic mathematical functions
+    "random",  # Random number generation
+    "io",  # Input/output streams
+    "urllib.parse",  # URL parsing and manipulation (safe URL handling)
+    "typing",  # Support for type hints (improve code clarity)
+    "concurrent.futures",  # For parallel execution
+    "time",  # Measuring time
+    "tempfile",  # Creating temporary files and directories
+    # Data Visualization (if needed) - Consider security implications carefully
+    "matplotlib",  # Plotting library (basic charts)
+    "seaborn",  # Statistical data visualization (more advanced)
+    # Web Scraping (more specific/controlled) - Consider ethical implications
+    "lxml",  # Faster XML/HTML processing (alternative to bs4)
+    "selenium",  # Automated browser control (for dynamic websites)
+    # Database interaction (if needed) - Handle credentials securely!
+    "sqlite3",  # SQLite database access
+    # Task scheduling
+    "schedule",  # Allow the agent to schedule tasks
 ]
 
 USER_AGENT = (
@@ -98,6 +103,7 @@ BROWSER_CONFIG = {
 
 CUSTOM_ROLE_CONVERSIONS = {"tool-call": "assistant", "tool-response": "user"}
 
+
 ALLOWED_FILE_TYPES = [
     "application/pdf",
     "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
@@ -113,31 +119,14 @@ 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():
-    """Initialize environment variables and authenticate with Hugging Face Hub."""
+    """Initialize environment variables and authentication."""
     load_dotenv(override=True)
-    hf_token = os.getenv("HF_TOKEN")
-    if hf_token:
-        login(hf_token)
-        print(f"HF_TOKEN (last 10 characters): {hf_token[-10:]}")
+    if os.getenv("HF_TOKEN"):  # Check if token is actually set
+        login(os.getenv("HF_TOKEN"))
+        print("HF_TOKEN (last 10 characters):", os.getenv("HF_TOKEN")[-10:])
     else:
         print("HF_TOKEN not found in environment variables.")
 
@@ -147,44 +136,36 @@ class ModelManager:
     """Manages model loading and initialization."""
 
     @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:
-            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.
-
-        Raises:
-            ValueError: If an invalid inference type is specified or if the key manager
-                is missing for OpenAI models.
-            Exception: If the model fails to load.
-        """
+    def load_model(chosen_inference: str, model_id: str, key_manager=None):
+        """Load the specified model with appropriate configuration."""
         try:
             if chosen_inference == "hf_api":
                 return HfApiModel(model_id=model_id)
-            if chosen_inference == "hf_api_provider":
+
+            elif chosen_inference == "hf_api_provider":
                 return HfApiModel(provider="together")
-            if chosen_inference == "litellm":
+
+            elif chosen_inference == "litellm":
                 return LiteLLMModel(model_id=model_id)
-            if chosen_inference == "openai":
+
+            elif 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")
                 )
-            if chosen_inference == "transformers":
+
+            elif chosen_inference == "transformers":
                 return TransformersModel(
                     model_id="HuggingFaceTB/SmolLM2-1.7B-Instruct",
                     device_map="auto",
                     max_new_tokens=1000,
                 )
-            raise ValueError(f"Invalid inference type: {chosen_inference}")
+
+            else:
+                raise ValueError(f"Invalid inference type: {chosen_inference}")
+
         except Exception as e:
             print(f"✗ Couldn't load model: {e}")
             raise
@@ -194,17 +175,8 @@ class ToolRegistry:
     """Manages tool initialization and organization."""
 
     @staticmethod
-    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: The maximum text length for the text inspector tool.
-
-        Returns:
-            A list of web-related tools.
-        """
+    def load_web_tools(model, browser, text_limit=20000):
+        """Initialize and return web-related tools."""
         return [
             GoogleSearchTool(provider="serper"),
             VisitTool(browser),
@@ -217,21 +189,20 @@ class ToolRegistry:
         ]
 
     @staticmethod
-    def load_document_tools() -> List[Tool]:
-        """Initialize and return document processing tools.
-
+    def load_document_tools():
+        """
+        Initialize and return document processing, i.e. sanitisation and indexing, tools.
         Returns:
-            List of document tools.
+            List of document tools
         """
-        return [FrontmatterGeneratorTool(), TextCleanerTool()]
+        return [
+            FrontmatterGeneratorTool(),
+            TextCleanerTool(),
+        ]
 
     @staticmethod
-    def load_image_generation_tools() -> Optional[Tool]:
-        """Initialize and return image generation tools.
-
-        Returns:
-            The image generation tool or None if initialization fails.
-        """
+    def load_image_generation_tools():
+        """Initialize and return image generation tools."""
         try:
             return Tool.from_space(
                 space_id="xkerser/FLUX.1-dev",
@@ -240,56 +211,45 @@ class ToolRegistry:
             )
         except Exception as e:
             print(f"✗ Couldn't initialize image generation tool: {e}")
-            return None
+            raise
 
 
 # ------------------------ Agent Creation and Execution ------------------------
-def create_agent() -> CodeAgent:
-    """Creates a fresh agent instance with configured tools.
-
+def create_agent():
+    """
+    Creates a fresh agent instance with properly configured tools.
     Returns:
-        CodeAgent: Configured agent ready for use.
-
+        CodeAgent: Configured agent ready for use
     Raises:
-        ValueError: If tool validation fails.
-        RuntimeError: If agent creation fails.
+        ValueError: If tool validation fails
+        RuntimeError: If agent creation fails
     """
     try:
-        # Model Configuration
-        model_id = "openrouter/deepseek/deepseek-chat-v3-0324:free"  # Define Model
         # Initialize model
         model = LiteLLMModel(
             custom_role_conversions=CUSTOM_ROLE_CONVERSIONS,
-            model_id=model_id,
+            model_id="openrouter/deepseek/deepseek-chat-v3-0324:free",
         )
 
         # Initialize tools
         text_limit = 30000
         browser = SimpleTextBrowser(**BROWSER_CONFIG)
 
-        # Create tool instances with proper error handling
+        # Collect all tools in a single list
         web_tools = ToolRegistry.load_web_tools(model, browser, text_limit)
-        doc_tools = []  # Initialize as empty list
-        image_generator = None  # Initialize as None
-
-        try:
-            doc_tools = ToolRegistry.load_document_tools()
-        except Exception as e:
-            print(f"Warning: Error loading document tools: {str(e)}")
-            print("Attempting to continue with available tools...")
-            doc_tools = []
-
+        doc_tools = ToolRegistry.load_document_tools()  # New document 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:
-            all_tools.append(image_generator)
+        # Combine all tools into a single list
+        all_tools = [visualizer] + web_tools + doc_tools + [image_generator]
 
-        # Log available tools
-        print(f"Loaded {len(all_tools)} tools successfully")
+        # Validate tools before creating agent
         for tool in all_tools:
-            print(f"- {tool.name}: {tool.description[:50]}...")
+            if not isinstance(tool, Tool):
+                raise ValueError(
+                    f"Invalid tool type: {type(tool)}. "
+                    f"All tools must be instances of Tool class."
+                )
 
         return CodeAgent(
             model=model,
@@ -299,54 +259,46 @@ def create_agent() -> CodeAgent:
             additional_authorized_imports=AUTHORIZED_IMPORTS,
             planning_interval=4,
         )
-    except Exception as e:
+    except (ValueError, RuntimeError) as e:
         print(f"Failed to create agent: {e}")
         raise RuntimeError(f"Agent creation failed: {e}")
 
 
 def stream_to_gradio(
-    agent: CodeAgent,
+    agent,
     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(
-            task, stream=True, reset=reset_agent_memory, additional_args=additional_args
-        ):
-            for message in pull_messages_from_step(step_log):
-                yield message
-
-        # Process final answer with comprehensive media output
-        final_answer = step_log  # Last log is the run's final_answer
-        final_answer = handle_agent_output_types(final_answer)
-
-        # Output handling based on type
-        if isinstance(final_answer, AgentText):
-            yield gr.ChatMessage(
-                role="assistant",
-                content=f"Final answer:\n{final_answer.to_string()}\n",
-            )
-        elif isinstance(final_answer, AgentImage):
-            yield gr.ChatMessage(
-                role="assistant",
-                content={"image": final_answer.to_string(), "type": "file"},
-            )
-        elif isinstance(final_answer, AgentAudio):
-            yield gr.ChatMessage(
-                role="assistant",
-                content={"audio": final_answer.to_string(), "type": "file"},
-            )
-        else:
-            yield gr.ChatMessage(
-                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."
+    for step_log in agent.run(
+        task, stream=True, reset=reset_agent_memory, additional_args=additional_args
+    ):
+        for message in pull_messages_from_step(step_log):
+            yield message
+
+    # Process final answer : Use a more comprehensive media output
+    final_answer = step_log  # Last log is the run's final_answer
+    final_answer = handle_agent_output_types(final_answer)
+
+    if isinstance(final_answer, AgentText):
+        yield gr.ChatMessage(
+            role="assistant",
+            content=f"**Final answer:**\n{final_answer.to_string()}\n",
+        )
+    elif isinstance(final_answer, AgentImage):
+        yield gr.ChatMessage(
+            role="assistant",
+            content={"image": final_answer.to_string(), "type": "file"},
+        )  # Send as Gradio-compatible file object:
+    elif isinstance(final_answer, AgentAudio):
         yield gr.ChatMessage(
             role="assistant",
-            content=error_message,
+            content={"audio": final_answer.to_string(), "type": "file"},
+        )  # Send as Gradio-compatible file object
+    else:
+        yield gr.ChatMessage(
+            role="assistant", content=f"**Final answer:** {str(final_answer)}"
         )
 
 
@@ -354,170 +306,132 @@ def stream_to_gradio(
 class GradioUI:
     """A one-line interface to launch your agent in Gradio."""
 
-    def __init__(self, file_upload_folder: Optional[str] = None):
+    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
 
-        if self.file_upload_folder:
-            os.makedirs(self.file_upload_folder, exist_ok=True)
+        if self.file_upload_folder is not None:
+            if not os.path.exists(file_upload_folder):
+                os.mkdir(file_upload_folder)
 
-    def interact_with_agent(
-        self,
-        prompt: str,
-        messages: List[Dict],
-        session_state: Dict,
-        uploaded_files: List[str],
-    ) -> List[Dict]:
+    def interact_with_agent(self, prompt, messages, session_state):
         """Main interaction handler with the agent."""
 
+        # Get or create session-specific agent
         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=error_message,
-                    )
-                )
-                yield messages
-                return
+            session_state["agent"] = create_agent()
 
-        session_state["request_count"] += 1
-
-        messages.append(gr.ChatMessage(role="user", content=prompt))
-        yield messages
-
-        file_message = ""
+        # Adding monitoring
         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."
-            )
+            # Log the existence of agent memory
+            has_memory = hasattr(session_state["agent"], "memory")
+            print(f"Agent has memory: {has_memory}")
+            if has_memory:
+                print(f"Memory type: {type(session_state['agent'].memory)}")
 
-        try:
-            reset_needed = session_state["request_count"] > 15
+            messages.append(gr.ChatMessage(role="user", content=prompt))
+            yield messages
 
             for msg in stream_to_gradio(
-                session_state["agent"],
-                task=prompt_with_files,
-                reset_agent_memory=reset_needed,
+                session_state["agent"], task=prompt, reset_agent_memory=False
             ):
                 messages.append(msg)
-                yield messages
-
-            if reset_needed:
-                session_state["request_count"] = 1
+                yield messages  # Yield messages after each step
+            yield messages  # Yield messages one last time
 
         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=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..."),
-            gr.Button(interactive=False),
-        )
-
-    def upload_file(self, files: List[str]) -> Tuple[str, List[str]]:
-        """Handle file uploads with validation, security, and clear feedback.
-
-        Args:
-            files: List of file paths to upload
+            print(f"Error in interaction: {str(e)}")
+            raise
 
-        Returns:
-            Tuple of (status message, list of uploaded file paths)
-        """
-        if not files:
-            return "No file uploaded", []
-
-        uploaded_files = []
-        error_message = None
-
-        for file_path in files:
-            try:
-                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)}"
-                    )
-                    return error_message, []
+    def upload_file(
+        self,
+        file,
+        file_uploads_log,
+    ):
+        """Handle file uploads with proper validation and security."""
+        if file is None:
+            return gr.Textbox("No file uploaded", visible=True), file_uploads_log
 
-                file_size_mb = os.path.getsize(file_path) / (1024 * 1024)
-                max_file_size_mb = 50
+        try:
+            mime_type, _ = mimetypes.guess_type(file.name)
+        except Exception as e:
+            return gr.Textbox(f"Error: {e}", visible=True), file_uploads_log
+
+        if mime_type not in ALLOWED_FILE_TYPES:
+            return gr.Textbox("File type disallowed", visible=True), file_uploads_log
+
+        # Sanitize file name
+        original_name = os.path.basename(file.name)
+        sanitized_name = re.sub(
+            r"[^\w\-.]", "_", original_name
+        )  # Replace invalid chars with underscores
+
+        # Ensure the extension correlates to the mime type
+        type_to_ext = {}
+        for ext, t in mimetypes.types_map.items():
+            if t not in type_to_ext:
+                type_to_ext[t] = ext
+
+        # Build sanitized filename with proper extension
+        name_parts = sanitized_name.split(".")[:-1]
+        extension = type_to_ext.get(mime_type, "")
+        sanitized_name = "".join(name_parts) + extension
+
+        # Limit File Size, and Throw Error
+        max_file_size_mb = 50  # Define the limit
+        file_size_mb = os.path.getsize(file.name) / (1024 * 1024)  # Size in MB
+
+        if file_size_mb > max_file_size_mb:
+            return (
+                gr.Textbox(
+                    f"File size exceeds {max_file_size_mb} MB limit.", visible=True
+                ),
+                file_uploads_log,
+            )
 
-                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, []
+        # Save the uploaded file to the specified folder
+        file_path = os.path.join(self.file_upload_folder, sanitized_name)
+        shutil.copy(file.name, file_path)
 
-                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}")
+        return gr.Textbox(
+            f"File uploaded: {file_path}", visible=True
+        ), file_uploads_log + [file_path]
 
-            except Exception as e:
-                error_message = f"❌ Upload error: {str(e)}"
-                return error_message, []
+    def log_user_message(self, text_input, file_uploads_log):
+        """Process user message and handle file references."""
+        message = text_input
 
-        if error_message:
-            return error_message, []
+        if len(file_uploads_log) > 0:
+            message += f"\nYou have been provided with these files, which might be helpful or not: {file_uploads_log}"  # Added file list
 
         return (
-            f"✓ Files uploaded successfully: {', '.join([os.path.basename(f) for f in uploaded_files])}",
-            uploaded_files,
+            message,
+            gr.Textbox(
+                value="",
+                interactive=False,
+                placeholder="Processing...",  # Changed placeholder.
+            ),
+            gr.Button(interactive=False),
         )
 
-    def detect_device(self, request: gr.Request) -> str:
+    def detect_device(self, request: gr.Request):
         """Detect whether the user is on mobile or desktop device."""
         if not request:
-            return "Desktop"
+            return "Unknown device"  # Handle case where request is none.
 
+        # 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,6 +439,7 @@ class GradioUI:
             if platform in ['"windows"', '"macos"', '"linux"']:
                 return "Desktop"
 
+        # Default case if no clear indicators
         return "Desktop"
 
     def launch(self, **kwargs):
@@ -534,118 +449,101 @@ class GradioUI:
             @gr.render()
             def layout(request: gr.Request):
                 device = self.detect_device(request)
-                print(f"Device detected: {device}")
+                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)
+        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 and enhanced styling."""
-        with gr.Column(visible=False) as sidebar_demo:
+        """Create the desktop layout with sidebar."""
+        with gr.Blocks(fill_height=True) as sidebar_demo:
             with gr.Sidebar():
                 gr.Markdown(
-                    """# 🔍 OpenDeepResearch
-                    ### Smolagents + Document Tools
-                    """
+                    """#OpenDeepResearch - 3theSmolagents!
+                Model_id: google/gemini-2.0-flash-001"""
                 )
                 with gr.Group():
-                    gr.Markdown("What can I help you with today?", container=True)
+                    gr.Markdown("**What's on your mind mate?**", container=True)
                     text_input = gr.Textbox(
-                        lines=4,
+                        lines=3,
                         label="Your request",
                         container=False,
-                        placeholder="Enter your question or task here...",
-                        show_label=False,
+                        placeholder="Enter your prompt here and press Shift+Enter or press the button",
+                    )
+                    launch_research_btn = gr.Button("Run", variant="primary")
+
+                # 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
+                    )
+                    file_uploads_log = gr.State([])
+                    upload_file.change(
+                        self.upload_file,
+                        [upload_file, file_uploads_log],
+                        [upload_status, file_uploads_log],
                     )
 
-                    with gr.Row():
-                        clear_btn = gr.Button("Clear", variant="secondary")
-                        launch_research_btn = gr.Button("Run", variant="primary")
-
-                if self.file_upload_folder:
-                    with gr.Group():
-                        gr.Markdown("📎 Upload Documents")
-                        file_upload = gr.File(
-                            label="Upload files for analysis",
-                            file_types=self.allowed_extensions,
-                            file_count="multiple",
-                        )
-
-                        upload_status = gr.Textbox(
-                            label="Upload Status", interactive=False, visible=False
-                        )
-                        uploaded_files_state = gr.State([])
-
-                gr.HTML("<br><hr><h4><center>Powered by:</center></h4>")
+                gr.HTML("<br><br><h4><center>Powered by:</center></h4>")
                 with gr.Row():
                     gr.HTML(
                         """
-                        <div style="display: flex; align-items: center; justify-content: center; gap: 8px; font-family: system-ui, -apple-system, sans-serif;">
-                        <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png"
-                            style="width: 32px; height: 32px; object-fit: contain;" alt="logo">
-                        <a target="_blank" href="https://github.com/huggingface/smolagents">
-                            <b>huggingface/smolagents</b>
-                        </a>
-                        </div>
-                        """
+                    <div style="display: flex; align-items: center; gap: 8px; font-family: system-ui, -apple-system, sans-serif;">
+                    <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png"
+                         style="width: 32px; height: 32px; object-fit: contain;" alt="logo">
+                    <a target="_blank" href="https://github.com/huggingface/smolagents">
+                        <b>huggingface/smolagents</b>
+                    </a>
+                    </div>
+                    """
                     )
 
-            session_state = gr.State({})
+            # Add session state to store session-specific data
+            session_state = gr.State({})  # Initialize empty state for each session
             stored_messages = gr.State([])
+            if "file_uploads_log" not in locals():
+                file_uploads_log = gr.State([])
 
             chatbot = gr.Chatbot(
-                label="OpenDeepResearch Assistant",
+                label="open-Deep-Research",
                 type="messages",
                 avatar_images=(
                     None,
                     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png",
                 ),
-                resizeable=True,
-                show_copy_button=True,
+                resizeable=False,
                 scale=1,
                 elem_id="my-chatbot",
-                height=700,
-            )
-
-            clear_btn.click(
-                lambda: ([], [], {"agent": session_state.get("agent")}, []),
-                outputs=[chatbot, stored_messages, session_state, uploaded_files_state],
-            )
-
-            file_upload.change(
-                self.upload_file,
-                inputs=[file_upload],
-                outputs=[upload_status, uploaded_files_state],
             )
 
             self._connect_event_handlers(
                 text_input,
                 launch_research_btn,
+                file_uploads_log,
                 stored_messages,
                 chatbot,
                 session_state,
-                uploaded_files_state,
             )
 
             return sidebar_demo
 
     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!")
+        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_upload = gr.File(
-                label="Upload files for analysis",
-                file_types=self.allowed_extensions,
-                file_count="multiple",
-            )
-            uploaded_files_state = gr.State([])
+            file_uploads_log = gr.State([])
 
             chatbot = gr.Chatbot(
-                label="OpenDeepResearch",
+                label="open-Deep-Research",
                 type="messages",
                 avatar_images=(
                     None,
@@ -655,14 +553,16 @@ class GradioUI:
                 scale=1,
             )
 
-            if self.file_upload_folder:
+            # 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
                 )
-                file_upload.change(
+                upload_file.change(
                     self.upload_file,
-                    [file_upload],
-                    [upload_status, uploaded_files_state],
+                    [upload_file, file_uploads_log],
+                    [upload_status, file_uploads_log],
                 )
 
             text_input = gr.Textbox(
@@ -675,78 +575,79 @@ 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
 
     def _connect_event_handlers(
         self,
-        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,
+        launch_research_btn,
+        file_uploads_log,
+        stored_messages,
+        chatbot,
+        session_state,
+    ):
+        """Connect the event handlers for input elements."""
+        # Connect text input submit event
         text_input.submit(
             self.log_user_message,
-            inputs=[text_input],
-            outputs=[text_input, text_input, launch_btn],
+            [text_input, file_uploads_log],
+            [stored_messages, text_input, launch_research_btn],
         ).then(
             self.interact_with_agent,
-            inputs=[text_input, stored_messages, session_state, uploaded_files_state],
-            outputs=[chatbot],
+            [stored_messages, chatbot, session_state],
+            [chatbot],
         ).then(
             lambda: (
                 gr.Textbox(
-                    value="", interactive=True, placeholder="Enter your prompt here..."
+                    interactive=True,
+                    placeholder="Enter your prompt here and press the button",
                 ),
                 gr.Button(interactive=True),
             ),
-            outputs=[text_input, launch_btn],
+            None,
+            [text_input, launch_research_btn],
         )
 
-        # Button click handler (same flow)
-        launch_btn.click(
+        # Connect button click event
+        launch_research_btn.click(
             self.log_user_message,
-            inputs=[text_input],
-            outputs=[text_input, text_input, launch_btn],
+            [text_input, file_uploads_log],
+            [stored_messages, text_input, launch_research_btn],
         ).then(
             self.interact_with_agent,
-            inputs=[text_input, stored_messages, session_state, uploaded_files_state],
-            outputs=[chatbot],
+            [stored_messages, chatbot, session_state],
+            [chatbot],
         ).then(
             lambda: (
                 gr.Textbox(
-                    value="", interactive=True, placeholder="Enter your prompt here..."
+                    interactive=True,
+                    placeholder="Enter your prompt here and press the button",
                 ),
                 gr.Button(interactive=True),
             ),
-            outputs=[text_input, launch_btn],
+            None,
+            [text_input, launch_research_btn],
         )
 
 
 # ------------------------ Execution ------------------------
 def main():
     """Main entry point for the application."""
+    # Initialize environment
     setup_environment()
-    os.makedirs(f"./{BROWSER_CONFIG['downloads_folder']}", exist_ok=True)
 
-    # Set up file upload directory
-    upload_folder = os.path.join(os.getcwd(), "uploads")
-    os.makedirs(upload_folder, exist_ok=True)
+    # Ensure downloads folder exists
+    os.makedirs(f"./{BROWSER_CONFIG['downloads_folder']}", exist_ok=True)
 
-    # Initialize and launch the UI
-    ui = GradioUI(file_upload_folder=upload_folder)
-    ui.launch(share=True)
+    # Launch UI
+    GradioUI(file_upload_folder="uploaded_files").launch()
 
 
 if __name__ == "__main__":
-    main()
+    main()