app.py

"""
OpenDeepResearch Agent Creation Module

This module provides functions to create various agent configurations
for different use cases.
"""

import os
from dotenv import load_dotenv
from huggingface_hub import login

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.frontmatter_tool import FrontmatterGeneratorTool
from scripts.text_cleaner_tool import TextCleanerTool

from smolagents import (
    CodeAgent,
    GradioUI,
    LiteLLMModel,
    GoogleSearchTool,
    Tool,
    FinalAnswerTool,
)

# Constants
AUTHORIZED_IMPORTS = [
    "requests",
    "zipfile",
    "pandas",
    "numpy",
    "sympy",
    "json",
    "bs4",
    "pubchempy",
    "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",
]

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"}


def setup_environment():
    """Initialize environment variables and authentication."""
    load_dotenv(override=True)
    if os.getenv("HF_TOKEN"):
        login(os.getenv("HF_TOKEN"))
        print("HF_TOKEN authenticated successfully")
    else:
        print("HF_TOKEN not found in environment variables")

    # Ensure download folder exists
    os.makedirs(f"./{BROWSER_CONFIG['downloads_folder']}", exist_ok=True)

    # Ensure upload folder exists
    os.makedirs("uploaded_files", exist_ok=True)


def create_agent(model_id="openrouter/google/gemini-2.0-flash-001"):
    """
    Creates an agent with all necessary tools for research and document processing.

    Args:
        model_id: Model ID to use for the agent

    Returns:
        CodeAgent: Fully configured agent
    """
    # Setup environment first
    setup_environment()

    # Initialize model
    model = LiteLLMModel(
        custom_role_conversions=CUSTOM_ROLE_CONVERSIONS,
        model_id=model_id,
    )

    # Initialize tools
    text_limit = 30000
    browser = SimpleTextBrowser(**BROWSER_CONFIG)

    # Create tool instances
    web_tools = [
        GoogleSearchTool(provider="serper"),
        VisitTool(browser),
        PageUpTool(browser),
        PageDownTool(browser),
        FinderTool(browser),
        FindNextTool(browser),
        ArchiveSearchTool(browser),
        TextInspectorTool(model, text_limit),
    ]

    # Add FinalAnswerTool explicitly
    final_answer_tool = FinalAnswerTool()

    # Load document tools
    try:
        doc_tools = [
            FrontmatterGeneratorTool(),
            TextCleanerTool(),
        ]
    except AssertionError as e:
        print(f"Warning: Error loading document tools: {str(e)}")
        doc_tools = []

    # Load image generation tool
    try:
        image_generator = Tool.from_space(
            space_id="xkerser/FLUX.1-dev",
            name="image_generator",
            description="Generates high-quality images using the FLUX.1-dev model based on text prompts.",
        )
    except Exception as e:
        print(f"Warning: Image generation tool unavailable: {str(e)}")
        image_generator = None

    # Combine all tools
    all_tools = [
        tool
        for tool in (
            [visualizer]
            + [final_answer_tool]
            + web_tools
            + doc_tools
            + ([image_generator] if image_generator else [])
        )
        if tool is not None
    ]

    # Log available tools
    print(f"Loaded {len(all_tools)} tools successfully")

    # Create and return the agent
    return CodeAgent(
        model=model,
        tools=all_tools,
        max_steps=15,
        verbosity_level=2,
        additional_authorized_imports=AUTHORIZED_IMPORTS,
        planning_interval=4,
    )


if __name__ == "__main__":
    # Create an agent instance
    agent_instance = create_agent()

    # Create Gradio UI with the agent instance and file upload folder
    demo = GradioUI(agent_instance, file_upload_folder="uploaded_files")

    # Launch the UI
    demo.launch()