app.py

from smolagents import (
    CodeAgent,
    ToolCallingAgent,
    HfApiModel,
    DuckDuckGoSearchTool,
    LiteLLMModel,
    tool,
    load_tool,
)
import datetime
import re
import requests
from markdownify import markdownify
from requests.exceptions import RequestException
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
from huggingface_hub import login

# Handle Hugging Face authentication
# Option 1: Use token from environment variable (recommended for production)
import os

# Try to get token from environment variable
hf_token = os.environ.get("HUGGINGFACE_TOKEN")
if hf_token:
    # Use token without interactive prompt
    login(token=hf_token)
else:
    # Option 2: Skip login and use API without authentication
    print("No Hugging Face token found in environment. Running without authentication.")
    # You may encounter rate limits without authentication

# Set model ID
model_id = "Qwen/Qwen2.5-Coder-32B-Instruct"

# Web browsing tool
@tool
def visit_webpage(url: str) -> str:
    """Visits a webpage at the given URL and returns its content as a markdown string.
    
    Args:
        url: The URL of the webpage to visit.
        
    Returns:
        The content of the webpage converted to Markdown, or an error message if the request fails.
    """
    try:
        # Send a GET request to the URL
        response = requests.get(url)
        response.raise_for_status()  # Raise an exception for bad status codes
        
        # Convert the HTML content to Markdown
        markdown_content = markdownify(response.text).strip()
        
        # Remove multiple line breaks
        markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
        
        return markdown_content
    except RequestException as e:
        return f"Error fetching the webpage: {str(e)}"
    except Exception as e:
        return f"An unexpected error occurred: {str(e)}"

# Custom tool example
@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """A tool that does nothing yet 
    
    Args:
        arg1: the first argument
        arg2: the second argument
    """
    return "What magic will you build?"

# Timezone tool
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

# Initialize final answer tool
final_answer = FinalAnswerTool()

# Import image generation tool from Hub with error handling
try:
    image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
    print("Successfully loaded image generation tool")
except Exception as e:
    print(f"Failed to load image generation tool: {e}")
    # Create a simple placeholder tool
    @tool
    def image_generation_tool(prompt: str) -> str:
        """A placeholder for image generation when the real tool fails to load.
        
        Args:
            prompt: Description of the image to generate
        """
        return f"Image generation unavailable. Your prompt was: {prompt}"

# Initialize model with fallback options
try:
    # Try to use primary model
    model = HfApiModel(
        max_tokens=2096,
        temperature=0.5,
        model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
        custom_role_conversions=None,
    )
    print("Using primary model: Qwen/Qwen2.5-Coder-32B-Instruct")
except Exception as e:
    print(f"Failed to initialize primary model: {e}")
    # Fallback to alternative endpoint
    try:
        model = HfApiModel(
            max_tokens=2096,
            temperature=0.5,
            model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
            custom_role_conversions=None,
        )
        print("Using fallback endpoint model")
    except Exception as e2:
        print(f"Failed to initialize fallback model: {e2}")
        # Last resort: try a smaller model
        model = HfApiModel(
            max_tokens=2096,
            temperature=0.5,
            model_id='Qwen/Qwen2-7B-Instruct',  # Smaller model as last resort
            custom_role_conversions=None,
        )
        print("Using smaller fallback model: Qwen/Qwen2-7B-Instruct")

# Load prompt templates
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# ===== MULTI-AGENT SYSTEM SETUP =====
# 1. Create specialized web search agent
web_agent = ToolCallingAgent(
    tools=[DuckDuckGoSearchTool(), visit_webpage],
    model=model,
    max_steps=10,
    name="web_search_agent",
    description="Runs web searches for you and can visit webpages to extract information.",
)

# 2. Create time agent for timezone operations
time_agent = ToolCallingAgent(
    tools=[get_current_time_in_timezone],
    model=model,
    max_steps=3,
    name="time_agent",
    description="Handles time-related queries and timezone conversions.",
)

# 3. Create creative agent for image generation
creative_agent = ToolCallingAgent(
    tools=[image_generation_tool, my_custom_tool],
    model=model,
    max_steps=5,
    name="creative_agent",
    description="Creates images and handles other creative tasks.",
)

# 4. Create manager agent (main agent)
manager_agent = CodeAgent(
    model=model,
    tools=[final_answer],  # The manager can directly use the final_answer tool
    managed_agents=[web_agent, time_agent, creative_agent],  # Manages three specialized agents
    max_steps=6,
    verbosity_level=1,
    additional_authorized_imports=["time", "numpy", "pandas"],  # For data calculations
    grammar=None,
    planning_interval=None,
    name="manager_agent",
    description="Coordinates specialized agents to solve complex problems.",
    prompt_templates=prompt_templates
)

# Main execution block with error handling
def main():
    try:
        print("Starting Multi-Agent System UI...")
        # Initialize Gradio UI with the manager agent
        ui = GradioUI(manager_agent)
        ui.launch()
        return True
    except Exception as e:
        print(f"Error starting application: {e}")
        import traceback
        traceback.print_exc()
        return False

# Execute main function if run directly
if __name__ == "__main__":
    success = main()
    if not success:
        print("Application failed to start properly. Check logs for details.")