app.py

from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
import os
import sys
import uuid
import base64
from io import BytesIO
from typing import Any, Union, List
from tools.final_answer import FinalAnswerTool
from tools.web_search import DuckDuckGoSearchTool
from tools.visit_webpage import VisitWebpageTool

from Gradio_UI import GradioUI

#############################################################
# TOOL 1: TIME ZONE 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 using pytz library
        tz = pytz.timezone(timezone)
        
        # Get current time in that timezone and format it as a readable string
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        
        # Return formatted response with the timezone and current time
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        # Handle any errors that might occur (invalid timezone, etc.)
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

#############################################################
# TOOL 2: ENHANCED IMAGE GENERATION TOOL
#############################################################
@tool
def generate_image_from_text(prompt: str) -> str:
    """A tool that generates an image based on a text description and saves it to a file.
    Args:
        prompt: A detailed text description of the image you want to generate.
    """
    try:
        # Create images directory if it doesn't exist
        os.makedirs("uploads/images", exist_ok=True)
        
        # Generate a unique filename
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        unique_id = str(uuid.uuid4())[:8]
        filename = f"uploads/images/{timestamp}_{unique_id}.jpg"
        
        # Call the image generation tool to get the PIL image
        pil_image = image_generation_tool(prompt)
        
        # Save the PIL image to file
        pil_image.save(filename)
        
        # Return only the filename - the display formatting is now handled by the final_answer tool
        return filename
    except Exception as e:
        # Handle any errors that occur during image generation
        return f"Error generating image: {str(e)}"

#############################################################
# TOOL 3: WEB SEARCH TOOL - Using the existing DuckDuckGoSearchTool class
#############################################################
# Initialize the DuckDuckGo search tool from the tools directory
web_search_tool = DuckDuckGoSearchTool()

@tool
def search_web(query: str) -> str:
    """A tool that searches the web using DuckDuckGo for information.
    Args:
        query: The search query to find information on the web.
    """
    try:
        # Execute the search query using DuckDuckGo
        search_results = web_search_tool(query)
        
        # Format and return the search results
        return f"Search results for '{query}':\n\n{search_results}"
    except Exception as e:
        # Handle any errors that occur during the search
        return f"Error searching the web: {str(e)}"

#############################################################
# CUSTOM FINAL ANSWER TOOL
#############################################################
# Initialize the base final answer tool
final_answer_base = FinalAnswerTool()

@tool
def final_answer(answer: Any) -> Any:
    """Provides a final answer to the given problem.
    Args:
        answer: The final answer to the problem
    """
    # Check if the answer is a list of image paths
    if isinstance(answer, list) and all(isinstance(item, str) and 'uploads/images/' in item for item in answer):
        # Build HTML to display all images
        html_output = "<div style='display: flex; flex-direction: column; gap: 20px;'>"
        
        # Get the HuggingFace Space name from environment variables
        space_name = os.environ.get("SPACE_ID", "unknown-space")
        
        for i, image_path in enumerate(answer):
            # Construct a URL that works in Hugging Face Spaces
            file_url = f"https://huggingface.co/spaces/{space_name}/resolve/main/{image_path}"
            
            html_output += f"""
            <div style="margin-bottom: 20px; border: 1px solid #ddd; border-radius: 8px; padding: 15px; background-color: #f9f9f9;">
              <h3 style="margin-top: 0;">Image {i+1}</h3>
              <img src="{image_path}" alt="Generated image {i+1}" style="max-width:100%; height:auto; border-radius:4px;">
              <div style="margin-top:10px;">
                <a href="{file_url}" target="_blank" style="display:inline-block; padding:8px 16px; background-color:#4CAF50; color:white; text-decoration:none; border-radius:4px; margin-right:10px;">View Full Size</a>
                <code style="display:block; margin-top:10px; padding:8px; background-color:#f0f0f0; border-radius:4px; font-size:12px; word-break:break-all;">{file_url}</code>
              </div>
            </div>
            """
        
        html_output += "</div>"
        return final_answer_base(html_output)
    
    # Otherwise, just pass the answer to the regular final_answer tool
    return final_answer_base(answer)

#############################################################
# VISIT WEBPAGE TOOL - From tools directory
#############################################################
# Initialize the visit webpage tool
visit_webpage = VisitWebpageTool()

#############################################################
# TOOL INITIALIZATION AND VERIFICATION
#############################################################
# Ensure all tools are properly loaded and print status
print("Initializing tools...")
try:
    # Import the image generation tool from Hugging Face Hub
    # This tool will be used by the generate_image_from_text function
    image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
    print("✓ Image generation tool loaded successfully")
except Exception as e:
    print(f"✗ Failed to load image generation tool: {str(e)}")
    # Provide a fallback if the image generation tool fails to load
    def image_generation_tool(prompt):
        from PIL import Image, ImageDraw, ImageFont
        import numpy as np
        
        # Create a blank image with text saying this is a fallback
        img = Image.new('RGB', (512, 512), color=(245, 245, 245))
        d = ImageDraw.Draw(img)
        d.text((10, 10), f"Fallback image for: {prompt}", fill=(0, 0, 0))
        d.text((10, 30), "Image generation tool failed to load", fill=(255, 0, 0))
        return img

print("✓ Time zone tool initialized")
print("✓ Web search tool initialized")
print("✓ Visit webpage tool initialized")
print("✓ Final answer tool initialized")

#############################################################
# MODEL CONFIGURATION
#############################################################
print("Configuring model...")
# Use the alternative endpoint since the primary one seems to have authentication issues
model = HfApiModel(
    max_tokens=2096,           # Maximum number of tokens in the response
    temperature=0.5,           # Controls randomness: lower = more deterministic
    model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',  # Alternative endpoint
    custom_role_conversions=None,
)
print("✓ Model configured successfully")

#############################################################
# LOAD PROMPT TEMPLATES
#############################################################
print("Loading prompt templates...")
# Create an upload directory if it doesn't exist
os.makedirs("uploads", exist_ok=True)
os.makedirs("uploads/images", exist_ok=True)

# Load prompt templates from YAML file for consistent agent responses
try:
    with open("prompts.yaml", 'r') as stream:
        prompt_templates = yaml.safe_load(stream)
    print("✓ Prompt templates loaded successfully")
except Exception as e:
    print(f"✗ Failed to load prompt templates: {str(e)}")
    # Provide default empty templates if loading fails
    prompt_templates = {}
    
#############################################################
# AGENT CONFIGURATION
#############################################################
print("Configuring agent...")
agent = CodeAgent(
    model=model,
    tools=[
        get_current_time_in_timezone,  # Tool 1: Time zone tool
        generate_image_from_text,      # Tool 2: Image generation tool
        search_web,                    # Tool 3: Web search tool
        visit_webpage,                 # Tool 4: Visit webpage tool (added from tools directory)
        final_answer                   # Custom final answer tool
    ],
    max_steps=6,                       # Maximum number of reasoning steps
    verbosity_level=1,                 # Level of detail in agent's output
    grammar=None,                      # No specific grammar constraints
    planning_interval=None,            # No specific planning interval
    name=None,                         # No custom agent name
    description=None,                  # No custom agent description
    prompt_templates=prompt_templates  # Using loaded prompt templates
)
print("✓ Agent configured successfully")

#############################################################
# LAUNCH THE GRADIO UI
#############################################################
print("Launching Gradio UI...")
# Start the Gradio interface with our configured agent and file upload directory
GradioUI(agent, file_upload_folder="uploads").launch()