Update app.py

app.py

@@ -1,23 +1,27 @@
-import argparse
 import json
+import mimetypes
 import os
+import re
+import shutil
 import threading
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from datetime import datetime
-from pathlib import Path
-from typing import List, Optional
+from typing import Optional
 
-import datasets
-import pandas as pd
+import gradio as gr
 from dotenv import load_dotenv
 from huggingface_hub import login
-import gradio as gr
-
-from scripts.reformulator import prepare_response
-from scripts.run_agents import (
-    get_single_file_description,
-    get_zip_description,
+from smolagents import (
+    CodeAgent,
+    HfApiModel,
+    Tool,
 )
+from smolagents.agent_types import (
+    AgentAudio,
+    AgentImage,
+    AgentText,
+    handle_agent_output_types,
+)
+from smolagents.gradio_ui import pull_messages_from_step
+
 from scripts.text_inspector_tool import TextInspectorTool
 from scripts.text_web_browser import (
     ArchiveSearchTool,
@@ -29,19 +33,6 @@ from scripts.text_web_browser import (
     VisitTool,
 )
 from scripts.visual_qa import visualizer
-from tqdm import tqdm
-
-from smolagents import (
-    CodeAgent,
-    HfApiModel,
-    LiteLLMModel,
-    Model,
-    ToolCallingAgent,
-)
-from smolagents.agent_types import AgentText, AgentImage, AgentAudio
-from smolagents.gradio_ui import pull_messages_from_step, handle_agent_output_types
-
-from smolagents import Tool
 
 
 class GoogleSearchTool(Tool):
@@ -67,7 +58,9 @@ class GoogleSearchTool(Tool):
         import requests
 
         if self.serpapi_key is None:
-            raise ValueError("Missing SerpAPI key. Make sure you have 'SERPER_API_KEY' in your env variables.")
+            raise ValueError(
+                "Missing SerpAPI key. Make sure you have 'SERPER_API_KEY' in your env variables."
+            )
 
         params = {
             "engine": "google",
@@ -76,16 +69,19 @@ class GoogleSearchTool(Tool):
             "google_domain": "google.com",
         }
 
-        headers = {
-        'X-API-KEY': self.serpapi_key,
-        'Content-Type': 'application/json'
-        }
+        headers = {"X-API-KEY": self.serpapi_key, "Content-Type": "application/json"}
 
         if filter_year is not None:
-            params["tbs"] = f"cdr:1,cd_min:01/01/{filter_year},cd_max:12/31/{filter_year}"
-
-        response = requests.request("POST", "https://google.serper.dev/search", headers=headers, data=json.dumps(params))
-
+            params["tbs"] = (
+                f"cdr:1,cd_min:01/01/{filter_year},cd_max:12/31/{filter_year}"
+            )
+
+        response = requests.request(
+            "POST",
+            "https://google.serper.dev/search",
+            headers=headers,
+            data=json.dumps(params),
+        )
 
         if response.status_code == 200:
             results = response.json()
@@ -99,9 +95,13 @@ class GoogleSearchTool(Tool):
                     f"No results found for query: '{query}' with filtering on year={filter_year}. Use a less restrictive query or do not filter on year."
                 )
             else:
-                raise Exception(f"No results found for query: '{query}'. Use a less restrictive query.")
+                raise Exception(
+                    f"No results found for query: '{query}'. Use a less restrictive query."
+                )
         if len(results["organic"]) == 0:
-            year_filter_message = f" with filter year={filter_year}" if filter_year is not None else ""
+            year_filter_message = (
+                f" with filter year={filter_year}" if filter_year is not None else ""
+            )
             return f"No results found for '{query}'{year_filter_message}. Try with a more general query, or remove the year filter."
 
         web_snippets = []
@@ -121,11 +121,14 @@ class GoogleSearchTool(Tool):
 
                 redacted_version = f"{idx}. [{page['title']}]({page['link']}){date_published}{source}\n{snippet}"
 
-                redacted_version = redacted_version.replace("Your browser can't play this video.", "")
+                redacted_version = redacted_version.replace(
+                    "Your browser can't play this video.", ""
+                )
                 web_snippets.append(redacted_version)
 
         return "## Search Results\n" + "\n\n".join(web_snippets)
 
+
 # web_search = GoogleSearchTool()
 
 # print(web_search(query="Donald Trump news"))
@@ -198,6 +201,7 @@ WEB_TOOLS = [
     TextInspectorTool(model, text_limit),
 ]
 
+
 # Agent creation in a factory function
 def create_agent():
     """Creates a fresh agent instance for each session"""
@@ -210,8 +214,10 @@ def create_agent():
         planning_interval=4,
     )
 
+
 document_inspection_tool = TextInspectorTool(model, 20000)
 
+
 def stream_to_gradio(
     agent,
     task: str,
@@ -219,7 +225,9 @@ def stream_to_gradio(
     additional_args: Optional[dict] = None,
 ):
     """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages."""
-    for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
+    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,
         ):
@@ -244,14 +252,15 @@ def stream_to_gradio(
             content={"path": final_answer.to_string(), "mime_type": "audio/wav"},
         )
     else:
-        yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(final_answer)}")
+        yield gr.ChatMessage(
+            role="assistant", content=f"**Final answer:** {str(final_answer)}"
+        )
 
 
 class GradioUI:
     """A one-line interface to launch your agent in Gradio"""
 
     def __init__(self, file_upload_folder: str | None = None):
-        
         self.file_upload_folder = file_upload_folder
         if self.file_upload_folder is not None:
             if not os.path.exists(file_upload_folder):
@@ -259,21 +268,23 @@ class GradioUI:
 
     def interact_with_agent(self, prompt, messages, session_state):
         # Get or create session-specific agent
-        if 'agent' not in session_state:
-            session_state['agent'] = create_agent()
+        if "agent" not in session_state:
+            session_state["agent"] = create_agent()
 
         # Adding monitoring
         try:
             # log the existence of agent memory
-            has_memory = hasattr(session_state['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)}")
-                
+
             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(
+                session_state["agent"], task=prompt, reset_agent_memory=False
+            ):
                 messages.append(msg)
                 yield messages
             yield messages
@@ -322,10 +333,14 @@ class GradioUI:
         sanitized_name = "".join(sanitized_name)
 
         # Save the uploaded file to the specified folder
-        file_path = os.path.join(self.file_upload_folder, os.path.basename(sanitized_name))
+        file_path = os.path.join(
+            self.file_upload_folder, os.path.basename(sanitized_name)
+        )
         shutil.copy(file.name, file_path)
 
-        return gr.Textbox(f"File uploaded: {file_path}", visible=True), file_uploads_log + [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):
         return (
@@ -335,49 +350,55 @@ class GradioUI:
                 if len(file_uploads_log) > 0
                 else ""
             ),
-            gr.Textbox(value="", interactive=False, placeholder="Please wait while Steps are getting populated"),
-            gr.Button(interactive=False)
+            gr.Textbox(
+                value="",
+                interactive=False,
+                placeholder="Please wait while Steps are getting populated",
+            ),
+            gr.Button(interactive=False),
         )
 
     def detect_device(self, request: gr.Request):
         # Check whether the user device is a mobile or a computer
-    
+
         if not request:
             return "Unknown device"
         # Method 1: Check sec-ch-ua-mobile header
-        is_mobile_header = request.headers.get('sec-ch-ua-mobile')
+        is_mobile_header = request.headers.get("sec-ch-ua-mobile")
         if is_mobile_header:
-            return "Mobile" if '?1' in is_mobile_header else "Desktop"
-    
+            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']
-    
+        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()
+        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):
+        return "Desktop"
 
+    def launch(self, **kwargs):
         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 
+                # Render layout with sidebar
                 if device == "Desktop":
-                    with gr.Blocks(fill_height=True,) as sidebar_demo:
+                    with gr.Blocks(
+                        fill_height=True,
+                    ):
+                        file_uploads_log = gr.State([])
                         with gr.Sidebar():
                             gr.Markdown("""# open Deep Research - free the AI agents!
                             
@@ -388,30 +409,42 @@ class GradioUI:
                 You can try a simplified version here (uses `Qwen-Coder-32B` instead of `o1`, so much less powerful than the original open-Deep-Research).<br><br>""")
                             with gr.Group():
                                 gr.Markdown("**Your request**", 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")
+                                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"
+                                )
 
                             # 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_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],
                                 )
-                                    
+
                             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>""")                        
+                        </div>""")
 
                         # Add session state to store session-specific data
-                        session_state = gr.State({})  # Initialize empty state for each session
+                        session_state = gr.State(
+                            {}
+                        )  # Initialize empty state for each session
                         stored_messages = gr.State([])
-                        file_uploads_log = gr.State([])
                         chatbot = gr.Chatbot(
                             label="open-Deep-Research",
                             type="messages",
@@ -421,35 +454,55 @@ class GradioUI:
                             ),
                             resizeable=False,
                             scale=1,
-                            elem_id="my-chatbot"
+                            elem_id="my-chatbot",
                         )
 
                         text_input.submit(
                             self.log_user_message,
                             [text_input, file_uploads_log],
                             [stored_messages, text_input, launch_research_btn],
-                        ).then(self.interact_with_agent,
+                        ).then(
+                            self.interact_with_agent,
                             # Include session_state in function calls
                             [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])
+                            [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],
+                        )
                         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,
+                        ).then(
+                            self.interact_with_agent,
                             # Include session_state in function calls
                             [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])
-                        
+                            [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],
+                        )
+
                 # Render simple layout
                 else:
-                    with gr.Blocks(fill_height=True,) as simple_demo:
+                    with gr.Blocks(
+                        fill_height=True,
+                    ):
                         gr.Markdown("""# open Deep Research - free the AI agents!
             _Built with [smolagents](https://github.com/huggingface/smolagents)_
             
@@ -459,7 +512,9 @@ class GradioUI:
             
             You can try a simplified version below (uses `Qwen-Coder-32B` instead of `o1`, so much less powerful than the original open-Deep-Research)👇""")
                         # Add session state to store session-specific data
-                        session_state = gr.State({})  # Initialize empty state for each session
+                        session_state = gr.State(
+                            {}
+                        )  # Initialize empty state for each session
                         stored_messages = gr.State([])
                         file_uploads_log = gr.State([])
                         chatbot = gr.Chatbot(
@@ -475,38 +530,66 @@ class GradioUI:
                         # If an upload folder is provided, enable the upload feature
                         if self.file_upload_folder is not None:
                             upload_file = gr.File(label="Upload a file")
-                            upload_status = gr.Textbox(label="Upload Status", interactive=False, visible=False)
+                            upload_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="Your request", placeholder="Enter your prompt here and press the button")
-                        launch_research_btn = gr.Button("Run", variant="primary",)
-                        
+                        text_input = gr.Textbox(
+                            lines=1,
+                            label="Your request",
+                            placeholder="Enter your prompt here and press the button",
+                        )
+                        launch_research_btn = gr.Button(
+                            "Run",
+                            variant="primary",
+                        )
+
                         text_input.submit(
                             self.log_user_message,
                             [text_input, file_uploads_log],
                             [stored_messages, text_input, launch_research_btn],
-                        ).then(self.interact_with_agent,
+                        ).then(
+                            self.interact_with_agent,
                             # Include session_state in function calls
                             [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])
+                            [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],
+                        )
                         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,
+                        ).then(
+                            self.interact_with_agent,
                             # Include session_state in function calls
                             [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])
-            
+                            [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.launch(debug=True, **kwargs)
 
+
 GradioUI().launch()