Update app.py

app.py

@@ -1,8 +1,8 @@
-import mimetypes
 import os
-import re
 import shutil
-from typing import Optional
+import threading
+import tempfile
+from typing import Optional, List, Dict, Any
 
 from dotenv import load_dotenv
 from huggingface_hub import login
@@ -19,7 +19,7 @@ from scripts.text_web_browser import (
     VisitTool,
 )
 from scripts.visual_qa import visualizer
-
+from scripts.legal_document_tool import LegalDocumentTool
 from smolagents import (
     CodeAgent,
     HfApiModel,
@@ -58,6 +58,8 @@ AUTHORIZED_IMPORTS = [
     "fractions",
     "csv",
     "clean-text",
+    "langchain",
+    "llama_index",
 ]
 
 user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0"
@@ -73,19 +75,20 @@ BROWSER_CONFIG = {
 
 custom_role_conversions = {"tool-call": "assistant", "tool-response": "user"}
 
-# Multimedia file types supported:
+# Multimedia file types supported (using Gradio-compatible format)
 ALLOWED_FILE_TYPES = [
-    "application/pdf",
-    "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
-    "text/plain",
-    "image/png",
-    "image/webp",
-    "image/jpeg",  # Added JPEG support
-    "image/gif",  # Added GIF support
-    "video/mp4",
-    "audio/mpeg",  # Added MP3 support
-    "audio/wav",  # Added WAV support
-    "audio/ogg",  # Added OGG support
+    ".pdf",  # application/pdf
+    ".docx",  # application/vnd.openxmlformats-officedocument.wordprocessingml.document
+    ".txt",  # text/plain
+    ".png",  # image/png
+    ".webp",  # image/webp
+    ".jpeg",  # image/jpeg
+    ".jpg",  # image/jpeg
+    ".gif",  # image/gif
+    ".mp4",  # video/mp4
+    ".mp3",  # audio/mpeg
+    ".wav",  # audio/wav
+    ".ogg",  # audio/ogg
 ]
 
 
@@ -101,31 +104,52 @@ def setup_environment():
 
 # ------------------------ Model and Tool Management ------------------------
 class ModelManager:
-    """Manages model loading and initialization."""
+    """Manages model loading and initialization with Zhou Protocol patterns."""
+
+    _instance = None
+    _lock = threading.Lock()
+
+    @classmethod
+    def get_instance(cls):
+        """Thread-safe singleton access to model manager."""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+
+    def __init__(self):
+        """Initialize with model cache."""
+        self.model_cache = {}
+
+    def load_model(self, chosen_inference: str, model_id: str, key_manager=None):
+        """Load the specified model with appropriate configuration and caching."""
+        cache_key = f"{chosen_inference}:{model_id}"
+
+        # Return cached model if available
+        if cache_key in self.model_cache:
+            return self.model_cache[cache_key]
 
-    @staticmethod
-    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)
+                model = HfApiModel(model_id=model_id)
 
             elif chosen_inference == "hf_api_provider":
-                return HfApiModel(provider="together")
+                model = HfApiModel(provider="together")
 
             elif chosen_inference == "litellm":
-                return LiteLLMModel(model_id=model_id)
+                model = LiteLLMModel(model_id=model_id)
 
             elif chosen_inference == "openai":
                 if not key_manager:
                     raise ValueError("Key manager required for OpenAI model")
 
-                return OpenAIServerModel(
+                model = OpenAIServerModel(
                     model_id=model_id, api_key=key_manager.get_key("openai_api_key")
                 )
 
             elif chosen_inference == "transformers":
-                return TransformersModel(
+                model = TransformersModel(
                     model_id="HuggingFaceTB/SmolLM2-1.7B-Instruct",
                     device_map="auto",
                     max_new_tokens=1000,
@@ -134,13 +158,22 @@ class ModelManager:
             else:
                 raise ValueError(f"Invalid inference type: {chosen_inference}")
 
+            # Cache the model for future use
+            self.model_cache[cache_key] = model
+            return model
+
         except Exception as e:
             print(f"✗ Couldn't load model: {e}")
             raise
 
 
 class ToolRegistry:
-    """Manages tool initialization and organization."""
+    """Manages tool initialization and organization with validation."""
+
+    @staticmethod
+    def validate_tools(tools: List[Tool]) -> List[Tool]:
+        """Validate tools and filter out any None values."""
+        return [tool for tool in tools if isinstance(tool, Tool)]
 
     @staticmethod
     def load_web_tools(model, browser, text_limit=20000):
@@ -167,51 +200,208 @@ class ToolRegistry:
             )
         except Exception as e:
             print(f"✗ Couldn't initialize image generation tool: {e}")
-            raise
+            return None
+
+    @staticmethod
+    def load_legal_document_tool():
+        """Initialize and return the legal document processing tool."""
+        try:
+            # Create a simple instance with default parameters
+            return LegalDocumentTool()
+        except Exception as e:
+            print(f"✗ Couldn't initialize legal document tool: {e}")
+            # Return None instead of raising to make this tool optional
+            return None
+
+
+# ------------------------ Session Management ------------------------
+class SessionManager:
+    """Manages agent sessions with proper cleanup and lifecycle management."""
+
+    _instance = None
+    _lock = threading.Lock()
+
+    @classmethod
+    def get_instance(cls):
+        """Thread-safe singleton access."""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+
+    def __init__(self):
+        """Initialize session management structures."""
+        self.sessions = {}
+        self.temp_files = {}
+        self.last_activity = {}
+        self._cleanup_thread = None
+        self._running = False
+        self._start_cleanup_thread()
+
+    def _start_cleanup_thread(self):
+        """Start a background thread for session cleanup."""
+        if self._cleanup_thread is None:
+            self._running = True
+            self._cleanup_thread = threading.Thread(
+                target=self._cleanup_inactive_sessions, daemon=True
+            )
+            self._cleanup_thread.start()
+
+    def _cleanup_inactive_sessions(self):
+        """Periodically clean up inactive sessions."""
+        import time
+
+        # Session timeout in seconds (30 minutes)
+        SESSION_TIMEOUT = 30 * 60
+
+        while self._running:
+            current_time = time.time()
+
+            # Find inactive sessions
+            inactive_sessions = [
+                session_id
+                for session_id, last_time in self.last_activity.items()
+                if (current_time - last_time) > SESSION_TIMEOUT
+            ]
+
+            # Clean up each inactive session
+            for session_id in inactive_sessions:
+                self.cleanup_session(session_id)
+
+            # Sleep for a minute before next check
+            time.sleep(60)
+
+    def register_session(self, session_id):
+        """Register a new session."""
+        if session_id not in self.sessions:
+            self.sessions[session_id] = {}
+            self.temp_files[session_id] = []
+
+        # Update activity timestamp
+        self.last_activity[session_id] = time.time()
+
+    def update_activity(self, session_id):
+        """Update the last activity timestamp for a session."""
+        self.last_activity[session_id] = time.time()
+
+    def register_temp_file(self, session_id, file_path):
+        """Register a temporary file with a session for later cleanup."""
+        if session_id not in self.temp_files:
+            self.temp_files[session_id] = []
+
+        self.temp_files[session_id].append(file_path)
+
+    def cleanup_session(self, session_id):
+        """Clean up resources for a session."""
+        # Remove temporary files
+        if session_id in self.temp_files:
+            for file_path in self.temp_files[session_id]:
+                try:
+                    if os.path.exists(file_path):
+                        os.remove(file_path)
+                except Exception as e:
+                    print(f"Error removing temp file {file_path}: {e}")
+
+            del self.temp_files[session_id]
+
+        # Clean up session data
+        if session_id in self.sessions:
+            del self.sessions[session_id]
+
+        # Clean up activity record
+        if session_id in self.last_activity:
+            del self.last_activity[session_id]
+
+    def __del__(self):
+        """Clean up all sessions when the manager is destroyed."""
+        self._running = False
+        if self._cleanup_thread and self._cleanup_thread.is_alive():
+            self._cleanup_thread.join(timeout=1.0)
+
+        # Clean up all remaining sessions
+        for session_id in list(self.sessions.keys()):
+            self.cleanup_session(session_id)
 
 
 # ------------------------ Agent Creation and Execution ------------------------
-def create_agent():
-    """Creates a fresh agent instance with properly configured tools."""
-    # 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
-
-    # Initialize tools
-    text_limit = 30000
-    browser = SimpleTextBrowser(**BROWSER_CONFIG)
-
-    # Collect all tools in a single list
-    web_tools = ToolRegistry.load_web_tools(model, browser, text_limit)
-    image_generator = ToolRegistry.load_image_generation_tools()
-
-    # Combine all tools into a single list (not a tuple)
-    all_tools = [visualizer] + web_tools + [image_generator]
-
-    # Validate tools before creating agent
-    for tool in all_tools:
-        if not isinstance(tool, Tool):
-            raise ValueError(
-                f"Invalid tool type: {type(tool)}. All tools must be instances of Tool class."
-            )
+class AgentFactory:
+    """Factory for creating and managing agent instances with Zhou Protocol patterns."""
+
+    _instance = None
+    _lock = threading.Lock()
+
+    @classmethod
+    def get_instance(cls):
+        """Thread-safe singleton access."""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+
+    def __init__(self):
+        """Initialize with agent cache."""
+        self.agent_cache = {}
+
+    def create_agent(self, session_id: str = "default") -> CodeAgent:
+        """Creates a fresh agent instance with properly configured tools."""
+        # Return cached agent if available for this session
+        if session_id in self.agent_cache:
+            return self.agent_cache[session_id]
+
+        # 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
+
+        # Initialize tools
+        text_limit = 30000
+        browser = SimpleTextBrowser(**BROWSER_CONFIG)
+
+        # Collect all tools in a single list
+        web_tools = ToolRegistry.load_web_tools(model, browser, text_limit)
+        image_generator = ToolRegistry.load_image_generation_tools()
+        legal_tool = ToolRegistry.load_legal_document_tool()
+
+        # Combine and validate all tools
+        all_tools = [visualizer] + web_tools
+
+        # Only add tools that are properly initialized (not None)
+        if image_generator:
+            all_tools.append(image_generator)
+
+        if legal_tool:
+            all_tools.append(legal_tool)
+
+        # Final validation to ensure all tools are valid
+        all_tools = ToolRegistry.validate_tools(all_tools)
+
+        agent = CodeAgent(
+            model=model,
+            tools=all_tools,  # Pass a single list containing all tools
+            max_steps=10,
+            verbosity_level=1,
+            additional_authorized_imports=AUTHORIZED_IMPORTS,
+            planning_interval=4,
+        )
+
+        # Cache the agent for future use
+        self.agent_cache[session_id] = agent
+        return agent
 
-    return CodeAgent(
-        model=model,
-        tools=all_tools,  # Pass a single list containing all tools
-        max_steps=10,
-        verbosity_level=1,
-        additional_authorized_imports=AUTHORIZED_IMPORTS,
-        planning_interval=4,
-    )
+    def clear_agent(self, session_id: str):
+        """Remove an agent from the cache."""
+        if session_id in self.agent_cache:
+            del self.agent_cache[session_id]
 
 
 def stream_to_gradio(
     agent,
     task: str,
     reset_agent_memory: bool = False,
-    additional_args: Optional[dict] = None,
+    additional_args: Optional[Dict[str, Any]] = None,
 ):
     """Runs an agent with the given task and streams messages as Gradio ChatMessages."""
     for step_log in agent.run(
@@ -247,337 +437,296 @@ def stream_to_gradio(
 
 # ------------------------ Gradio UI Components ------------------------
 class GradioUI:
-    """A one-line interface to launch your agent in Gradio."""
+    """A Gradio-compliant interface to launch your agent with proper resource management."""
 
-    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
+    def __init__(self):
+        """Initialize the Gradio UI with proper session management."""
+        self.session_manager = SessionManager.get_instance()
+        self.agent_factory = AgentFactory.get_instance()
+        self.temp_dir = tempfile.mkdtemp(prefix="gradio_")
 
-        if self.file_upload_folder is not None:
-            if not os.path.exists(file_upload_folder):
-                os.mkdir(file_upload_folder)
+    def __del__(self):
+        """Clean up resources when the UI is destroyed."""
+        try:
+            # Clean up the temporary directory
+            if os.path.exists(self.temp_dir):
+                shutil.rmtree(self.temp_dir, ignore_errors=True)
+        except Exception as e:
+            print(f"Error cleaning up temporary directory: {e}")
+
+    @staticmethod
+    def _get_session_id(session_state):
+        """Generate or retrieve a session ID."""
+        if "session_id" not in session_state:
+            session_state["session_id"] = f"session_{id(session_state)}"
+        return session_state["session_id"]
 
     def interact_with_agent(self, prompt, messages, session_state):
         """Main interaction handler with the agent."""
+        # Get or create session ID
+        session_id = self._get_session_id(session_state)
+
+        # Register/update the session
+        self.session_manager.register_session(session_id)
+        self.session_manager.update_activity(session_id)
 
         # Get or create session-specific agent
-        if "agent" not in session_state:
-            session_state["agent"] = create_agent()
+        agent = self.agent_factory.create_agent(session_id)
 
-        # Adding monitoring
         try:
             # Log the existence of agent memory
-            has_memory = hasattr(session_state["agent"], "memory")
+            has_memory = hasattr(agent, "memory")
             print(f"Agent has memory: {has_memory}")
             if has_memory:
-                print(f"Memory type: {type(session_state['agent'].memory)}")
+                print(f"Memory type: {type(agent.memory)}")
 
             messages.append(gr.ChatMessage(role="user", content=prompt))
             yield messages
 
-            for msg in stream_to_gradio(
-                session_state["agent"], task=prompt, reset_agent_memory=False
-            ):
+            for msg in stream_to_gradio(agent, task=prompt, reset_agent_memory=False):
                 messages.append(msg)
+                self.session_manager.update_activity(session_id)
                 yield messages  # Yield messages after each step
+
             yield messages  # Yield messages one last time
 
+        except gr.Error as e:
+            # Handle Gradio-specific errors
+            messages.append(
+                gr.ChatMessage(role="assistant", content=f"Error: {str(e)}")
+            )
+            yield messages
         except Exception as e:
+            # Log the error but present a user-friendly message
             print(f"Error in interaction: {str(e)}")
-            raise
+            messages.append(
+                gr.ChatMessage(
+                    role="assistant",
+                    content="I encountered an error processing your request. Please try again with a different query.",
+                )
+            )
+            yield messages
+
+    @gr.validate_input(file="file")
+    def upload_file(self, file, file_uploads_log, session_state):
+        """Handle file uploads with Gradio-compliant temporary file handling."""
+        session_id = self._get_session_id(session_state)
+        self.session_manager.update_activity(session_id)
 
-    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
 
         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
+            # Create a temporary file with a secure random name
+            temp_file_path = ""
 
-        # Sanitize file name
-        original_name = os.path.basename(file.name)
-        sanitized_name = re.sub(
-            r"[^\w\-.]", "_", original_name
-        )  # Replace invalid chars with underscores
+            with tempfile.NamedTemporaryFile(delete=False, dir=self.temp_dir) as tmp:
+                # Copy the uploaded file to the temporary file
+                shutil.copyfileobj(open(file.name, "rb"), tmp)
+                temp_file_path = tmp.name
 
-        # 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
+            # Register the temporary file with the session manager
+            self.session_manager.register_temp_file(session_id, temp_file_path)
 
-        # 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
+            # Store the original filename for reference
+            orig_filename = os.path.basename(file.name)
 
-        # 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
+            return (
+                gr.Textbox(f"File uploaded: {orig_filename}", visible=True),
+                file_uploads_log + [(temp_file_path, orig_filename)],
+            )
 
-        if file_size_mb > max_file_size_mb:
+        except Exception as e:
+            print(f"Error handling file upload: {e}")
             return (
-                gr.Textbox(
-                    f"File size exceeds {max_file_size_mb} MB limit.", visible=True
-                ),
+                gr.Textbox(f"Error uploading file: {str(e)}", visible=True),
                 file_uploads_log,
             )
 
-        # 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)
-
-        return gr.Textbox(
-            f"File uploaded: {file_path}", visible=True
-        ), file_uploads_log + [file_path]
-
-    def log_user_message(self, text_input, file_uploads_log):
+    def log_user_message(self, text_input, file_uploads_log, session_state):
         """Process user message and handle file references."""
+        session_id = self._get_session_id(session_state)
+        self.session_manager.update_activity(session_id)
+
         message = text_input
 
-        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
+        if file_uploads_log and len(file_uploads_log) > 0:
+            # Include only the original filenames in the message
+            filenames = [f[1] for f in file_uploads_log]
+            message += f"\nYou have been provided with these files: {filenames}"
+
+            # Include the actual file paths in the additional_args
+            if "additional_args" not in session_state:
+                session_state["additional_args"] = {}
+
+            session_state["additional_args"]["file_paths"] = [
+                f[0] for f in file_uploads_log
+            ]
 
         return (
             message,
             gr.Textbox(
                 value="",
                 interactive=False,
-                placeholder="Processing...",  # Changed placeholder.
+                placeholder="Processing...",
             ),
             gr.Button(interactive=False),
+            session_state,
         )
 
-    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.
-
-        # 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"']:
-                return "Mobile"
-            elif 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()
-                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"""
-                )
-                with gr.Group():
-                    gr.Markdown("**What's on your mind mate?**", container=True)
-                    text_input = gr.Textbox(
-                        lines=3,
-                        label="Your request",
-                        container=False,
-                        placeholder="Enter your prompt here and press Shift+Enter or press the button",
+        with gr.Blocks(theme="soft", css=self._get_responsive_css()) as demo:
+            with gr.Row(equal_height=True) as main_row:
+                # Sidebar (adapts to screen size via CSS)
+                with gr.Column(scale=1, min_width=100) as sidebar:
+                    gr.Markdown(
+                        """# OpenDeepResearch
+                        ## Powered by Smolagents"""
+                    )
+                    with gr.Group():
+                        gr.Markdown("**What's on your mind?**", container=True)
+                        text_input = gr.Textbox(
+                            lines=3,
+                            label="Your request",
+                            container=False,
+                            placeholder="Enter your prompt here and press Shift+Enter or press the button",
+                        )
+                        launch_research_btn = gr.Button("Run", variant="primary")
+
+                    # Clean file upload with Gradio-compliant file_types
+                    upload_file = gr.File(
+                        label="Upload a file",
+                        file_types=ALLOWED_FILE_TYPES,
+                        type="file",
                     )
-                    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
                     )
+
+                    # Footer with proper responsive behavior
+                    with gr.Row(visible=True) as footer:
+                        gr.HTML(
+                            """
+                            <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>
+                            """
+                        )
+
+                # Main content area
+                with gr.Column(scale=4, min_width=400) as content:
+                    # Add session state to store session-specific data
+                    session_state = gr.State({})
+                    stored_messages = gr.State([])
                     file_uploads_log = gr.State([])
-                    upload_file.change(
-                        self.upload_file,
-                        [upload_file, file_uploads_log],
-                        [upload_status, file_uploads_log],
-                    )
 
-                gr.HTML("<br><br><h4><center>Powered by:</center></h4>")
-                with gr.Row():
-                    gr.HTML(
-                        """
-                    <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>
-                    """
+                    chatbot = gr.Chatbot(
+                        label="OpenDeepResearch",
+                        show_label=True,
+                        type="messages",
+                        avatar_images=(
+                            None,
+                            "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png",
+                        ),
+                        height=600,
+                        elem_id="research-chatbot",
                     )
 
-            # 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="open-Deep-Research",
-                type="messages",
-                avatar_images=(
-                    None,
-                    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png",
+            # Connect event handlers
+            text_input.submit(
+                self.log_user_message,
+                [text_input, file_uploads_log, session_state],
+                [stored_messages, text_input, launch_research_btn, session_state],
+            ).then(
+                self.interact_with_agent,
+                [stored_messages, chatbot, session_state],
+                [chatbot],
+            ).then(
+                lambda: (
+                    gr.Textbox(
+                        interactive=True,
+                        placeholder="Enter your prompt here and press the button",
+                    ),
+                    gr.Button(interactive=True),
                 ),
-                resizeable=False,
-                scale=1,
-                elem_id="my-chatbot",
+                None,
+                [text_input, launch_research_btn],
             )
 
-            self._connect_event_handlers(
-                text_input,
-                launch_research_btn,
-                file_uploads_log,
-                stored_messages,
-                chatbot,
-                session_state,
-            )
-
-            return sidebar_demo
-
-    def _create_mobile_layout(self):
-        """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_uploads_log = gr.State([])
-
-            chatbot = gr.Chatbot(
-                label="open-Deep-Research",
-                type="messages",
-                avatar_images=(
-                    None,
-                    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png",
+            launch_research_btn.click(
+                self.log_user_message,
+                [text_input, file_uploads_log, session_state],
+                [stored_messages, text_input, launch_research_btn, session_state],
+            ).then(
+                self.interact_with_agent,
+                [stored_messages, chatbot, session_state],
+                [chatbot],
+            ).then(
+                lambda: (
+                    gr.Textbox(
+                        interactive=True,
+                        placeholder="Enter your prompt here and press the button",
+                    ),
+                    gr.Button(interactive=True),
                 ),
-                resizeable=True,
-                scale=1,
+                None,
+                [text_input, launch_research_btn],
             )
 
-            # 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(
-                    self.upload_file,
-                    [upload_file, file_uploads_log],
-                    [upload_status, file_uploads_log],
-                )
-
-            text_input = gr.Textbox(
-                lines=1,
-                label="What's on your mind mate?",
-                placeholder="Chuck in a question and we'll take care of the rest",
+            upload_file.change(
+                self.upload_file,
+                [upload_file, file_uploads_log, session_state],
+                [upload_status, file_uploads_log],
             )
-            launch_research_btn = gr.Button("Run", variant="primary")
 
-            self._connect_event_handlers(
-                text_input,
-                launch_research_btn,
-                file_uploads_log,
-                stored_messages,
-                chatbot,
-                session_state,
+            # Clean up session on page unload
+            demo.load(
+                lambda: None,
+                None,
+                None,
+                _js="""
+                () => {
+                    window.addEventListener('beforeunload', function() {
+                        // Notify backend about session end (would require additional endpoint)
+                        console.log('Cleaning up session');
+                    });
+                }
+            """,
             )
 
-            return simple_demo
-
-    def _connect_event_handlers(
-        self,
-        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,
-            [text_input, file_uploads_log],
-            [stored_messages, text_input, launch_research_btn],
-        ).then(
-            self.interact_with_agent,
-            [stored_messages, chatbot, session_state],
-            [chatbot],
-        ).then(
-            lambda: (
-                gr.Textbox(
-                    interactive=True,
-                    placeholder="Enter your prompt here and press the button",
-                ),
-                gr.Button(interactive=True),
-            ),
-            None,
-            [text_input, launch_research_btn],
-        )
-
-        # Connect button click event
-        launch_research_btn.click(
-            self.log_user_message,
-            [text_input, file_uploads_log],
-            [stored_messages, text_input, launch_research_btn],
-        ).then(
-            self.interact_with_agent,
-            [stored_messages, chatbot, session_state],
-            [chatbot],
-        ).then(
-            lambda: (
-                gr.Textbox(
-                    interactive=True,
-                    placeholder="Enter your prompt here and press the button",
-                ),
-                gr.Button(interactive=True),
-            ),
-            None,
-            [text_input, launch_research_btn],
-        )
+        demo.queue(max_size=20).launch(debug=True, **kwargs)
+
+    def _get_responsive_css(self):
+        """Get CSS for responsive layout."""
+        return """
+        /* Responsive layout */
+        @media (max-width: 768px) {
+            #research-chatbot {
+                height: 400px !important;
+            }
+            
+            /* Stack columns on small screens */
+            .gradio-row {
+                flex-direction: column;
+            }
+            
+            /* Adjust column widths */
+            .gradio-column {
+                min-width: 100% !important;
+                width: 100% !important;
+            }
+        }
+        
+        /* Base styling */
+        .gradio-container {
+            max-width: 100% !important;
+        }
+        """
 
 
 # ------------------------ Execution ------------------------
@@ -586,11 +735,11 @@ def main():
     # Initialize environment
     setup_environment()
 
-    # Ensure downloads folder exists
+    # Ensure downloads folder exists for browser
     os.makedirs(f"./{BROWSER_CONFIG['downloads_folder']}", exist_ok=True)
 
     # Launch UI
-    GradioUI(file_upload_folder="uploaded_files").launch()
+    GradioUI().launch(share=True)
 
 
 if __name__ == "__main__":