Update app.py

app.py

@@ -256,6 +256,264 @@ def stream_to_gradio(
         )
 
 
+"""Main application for the OpenDeepResearch Gradio interface."""
+
+import mimetypes
+import os
+import re
+import shutil
+from typing import Optional
+
+from dotenv import load_dotenv
+from huggingface_hub import login
+import gradio as gr
+
+from scripts.text_inspector_tool import TextInspectorTool
+from scripts.text_web_browser import (
+    ArchiveSearchTool,
+    FinderTool,
+    FindNextTool,
+    PageDownTool,
+    PageUpTool,
+    SimpleTextBrowser,
+    VisitTool,
+)
+from scripts.visual_qa import visualizer
+from scripts.text_cleaner_tool import TextCleanerTool
+
+from smolagents import (
+    CodeAgent,
+    HfApiModel,
+    LiteLLMModel,
+    OpenAIServerModel,
+    TransformersModel,
+    GoogleSearchTool,
+    Tool,
+)
+from smolagents.agent_types import AgentText, AgentImage, AgentAudio
+from smolagents.gradio_ui import pull_messages_from_step, handle_agent_output_types
+
+
+# Constants and configurations
+AUTHORIZED_IMPORTS = [
+    "requests",
+    "zipfile",
+    "pandas",
+    "numpy",
+    "sympy",
+    "json",
+    "bs4",
+    "pubchempy",
+    "xml",
+    "yahoo_finance",
+    "Bio",
+    "sklearn",
+    "scipy",
+    "pydub",
+    "PIL",
+    "chess",
+    "PyPDF2",
+    "pptx",
+    "torch",
+    "datetime",
+    "fractions",
+    "csv",
+    "clean-text",
+]
+
+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"
+BROWSER_CONFIG = {
+    "viewport_size": 1024 * 5,
+    "downloads_folder": "downloads_folder",
+    "request_kwargs": {
+        "headers": {"User-Agent": USER_AGENT},
+        "timeout": 300,
+    },
+    "serpapi_key": os.getenv("SERPAPI_API_KEY"),
+}
+
+CUSTOM_ROLE_CONVERSIONS = {"tool-call": "assistant", "tool-response": "user"}
+
+
+ALLOWED_FILE_TYPES = [
+    "application/pdf",
+    "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+    "text/plain",
+    "text/markdown",  # Added Markdown support
+    "application/json",  # Added JSON support
+    "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
+]
+
+
+def setup_environment():
+    """Initialize environment variables and authentication."""
+    load_dotenv(override=True)
+    hf_token = os.getenv("HF_TOKEN")
+    if hf_token:  # Check if token is actually set
+        login(hf_token)
+        print("HF_TOKEN (last 10 characters):", hf_token[-10:])
+    else:
+        print("HF_TOKEN not found in environment variables.")
+
+
+class ModelManager:
+    """Manages model loading and initialization."""
+
+    @staticmethod
+    def load_model(chosen_inference: str, model_id: str, key_manager=None):
+        """Load the specified model with appropriate configuration."""
+        try:
+            if chosen_inference == "hf_api":
+                return HfApiModel(model_id=model_id)
+
+            if chosen_inference == "hf_api_provider":
+                return HfApiModel(provider="together")
+
+            if chosen_inference == "litellm":
+                return LiteLLMModel(model_id=model_id)
+
+            if chosen_inference == "openai":
+                if not key_manager:
+                    raise ValueError("Key manager required for OpenAI model")
+                return OpenAIServerModel(
+                    model_id=model_id, api_key=key_manager.get_key("openai_api_key")
+                )
+
+            if 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}")
+        except Exception as e:
+            print(f"✗ Couldn't load model: {e}")
+            raise
+
+
+class ToolRegistry:
+    """Manages tool initialization and organization."""
+
+    @staticmethod
+    def load_web_tools(model, browser, text_limit=20000):
+        """Initialize and return web-related tools."""
+        return [
+            GoogleSearchTool(provider="serper"),
+            VisitTool(browser),
+            PageUpTool(browser),
+            PageDownTool(browser),
+            FinderTool(browser),
+            FindNextTool(browser),
+            ArchiveSearchTool(browser),
+            TextInspectorTool(model, text_limit),
+        ]
+
+    @staticmethod
+    def load_image_generation_tools():
+        """Initialize and return image generation tools."""
+        try:
+            return Tool.from_space(
+                space_id="xkerser/FLUX.1-dev",
+                name="image_generator",
+                description="Generates high-quality AgentImage with text prompt (77 token limit).",
+            )
+        except Exception as e:
+            print(f"✗ Couldn't initialize image generation tool: {e}")
+            raise
+
+    @staticmethod
+    def load_clean_text_tool():
+        """Initialize and return image generation tools."""
+        try:
+            return TextCleanerTool
+        except Exception as e:
+            print(f"✗ Couldn't initialize clean text tool: {e}")
+            raise
+
+
+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/deepseek/deepseek-chat-v3-0324:free",  # 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()
+    clean_text = TextCleanerTool()  # Instantiate TextCleanerTool
+
+    # Combine all tools into a single list (not a tuple)
+    all_tools = [visualizer] + web_tools + [image_generator] + [clean_text]
+
+    # Validate tools before creating agent
+    for tool in all_tools:
+        if not isinstance(tool, Tool):
+            raise ValueError(
+                "Invalid tool type: "
+                f"{type(tool)}. All tools must be instances of Tool class."
+            )
+
+    return CodeAgent(
+        model=model,
+        tools=all_tools,  # Pass a single list containing all tools
+        max_steps=12,
+        verbosity_level=2,
+        additional_authorized_imports=AUTHORIZED_IMPORTS,
+        planning_interval=4,
+    )
+
+
+def stream_to_gradio(
+    agent,
+    task: str,
+    reset_agent_memory: bool = False,
+    additional_args: Optional[dict] = None,
+):
+    """Runs an agent with the given task and streams messages as Gradio ChatMessages."""
+    for step_log in agent.run(
+        task, stream=True, reset=reset_agent_memory, additional_args=additional_args
+    ):
+        yield from pull_messages_from_step(step_log)
+
+    # 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={"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)}"
+        )
+
+
 class GradioUI:
     """A one-line interface to launch your agent in Gradio."""
 
@@ -351,7 +609,26 @@ class GradioUI:
 
     def log_user_message(self, text_input, file_uploads_log):
         """Process user message and handle file references."""
-        message = text_input
+        # Clean the user input using the TextCleanerTool
+        from cleantext import clean
+
+        cleaned_message = clean(
+            text_input,
+            fix_unicode=True,
+            to_ascii=True,
+            lower=True,
+            no_line_breaks=False,
+            no_urls=False,
+            no_emails=False,
+            no_phone_numbers=False,
+            no_numbers=False,
+            no_digits=False,
+            no_currency_symbols=False,
+            no_punct=False,
+            lang="en",
+        )  # Can change default behaviour by instantiating an object of TextCleanerTool class.
+
+        message = cleaned_message  # Use the cleaned message
 
         if file_uploads_log:
             message += f"\nYou have been provided with these files, which might be helpful or not: {file_uploads_log}"  # Added file list
@@ -464,7 +741,7 @@ class GradioUI:
                 file_uploads_log = gr.State([])
 
             chatbot = gr.Chatbot(
-                label="open-Deep-Research",
+                label="ODR",
                 type="messages",
                 avatar_images=(
                     None,
@@ -603,3 +880,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+