app.py

from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Import AgentText to return a composite Markdown message
from smolagents.agent_types import AgentText

# Import replicate directly to generate images using Stable Diffusion.
try:
    import replicate
except ModuleNotFoundError:
    replicate = None

@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """
    A tool that shares interesting facts about Chakwal, highlighting its flora, fauna, natural landscapes,
    and cultural heritage. It also shows the current local time in Chakwal (Asia/Karachi timezone).

    Args:
        arg1: A custom note or additional message to include in the output.
        arg2: The number of detailed facts to include about Chakwal (1 to 4).
              (1: Flora only, 2: Flora & Fauna, 3: + Natural Landscapes, 4: + Cultural Heritage)
    """
    messages = []
    
    # Introduction about Chakwal
    messages.append("Welcome to Chakwal, a city with a rich history, vibrant culture, and breathtaking natural beauty!")
    
    # Facts about Chakwal's natural and cultural highlights
    facts = [
        ("Flora", "The region boasts a diverse array of native plants, lush greenery, and seasonal wildflowers that add vibrant colors to the landscape."),
        ("Fauna", "Local wildlife includes a variety of bird species, reptiles, and small mammals that thrive in Chakwal's varied ecosystems."),
        ("Natural Landscapes", "Chakwal is surrounded by scenic hills, valleys, and water bodies, offering breathtaking views and numerous opportunities for outdoor adventures."),
        ("Cultural Heritage", "Deep-rooted traditions and lively festivals reflect Chakwal's rich cultural tapestry and historical significance.")
    ]
    
    # Determine the number of facts to display (minimum 1, maximum 4)
    num_facts = min(max(arg2, 1), len(facts))
    for i in range(num_facts):
        title, detail = facts[i]
        messages.append(f"{title}: {detail}")
    
    # Add current local time in Chakwal (using Asia/Karachi timezone)
    try:
        tz = pytz.timezone("Asia/Karachi")
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        messages.append(f"Local Time in Chakwal (Asia/Karachi): {local_time}")
    except Exception as e:
        messages.append("Error fetching local time for Chakwal.")
    
    # Include any custom note provided via arg1
    if arg1:
        messages.append(f"Note: {arg1}")
    
    return "\n".join(messages)


@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:
        tz = pytz.timezone(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)}"


# Create the image generation tool from smolagents.
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# ── NEW TOOL: prompte_chakwal ──
@tool
def prompte_chakwal(custom_note: str = "", num_results: int = 3) -> AgentText:
    """
    A tool that promotes Chakwal by retrieving real-time information and generating promotional images.
    It provides the current local time, live news about Chakwal from DuckDuckGo, and creates images that
    highlight the city's natural beauty and vibrant culture.

    The final output is a composite Markdown message that includes text details and an embedded image.

    Args:
        custom_note: A custom note to add a specific theme to the generated images.
        num_results: The maximum number of real-time search results (news/information) to include.
    """
    output_lines = []
    
    # 1. Add current local time in Chakwal (Asia/Karachi)
    try:
        tz = pytz.timezone("Asia/Karachi")
        current_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        output_lines.append(f"**Current local time in Chakwal:** {current_time}")
    except Exception as e:
        output_lines.append("Error fetching current time in Chakwal.")
    
    # 2. Fetch real-time information using DuckDuckGoSearchTool
    duck_tool = DuckDuckGoSearchTool()
    query = "Chakwal tourism investment news"
    search_results = duck_tool.run(query)
    
    if isinstance(search_results, list):
        output_lines.append("\n**Real-Time News and Information:**")
        for i, result in enumerate(search_results[:num_results]):
            title = result.get("title", "No Title")
            snippet = result.get("snippet", "No description available")
            output_lines.append(f"{i+1}. **{title}**: {snippet}")
    elif isinstance(search_results, str):
        output_lines.append("\n**Real-Time News and Information:**")
        output_lines.append(search_results)
    else:
        output_lines.append("\nNo real-time information available.")
    
    # 3. Generate a promotional image using the original smolagents image_generation_tool.
    image_prompt = "A breathtaking view of Chakwal showcasing its natural beauty, vibrant culture, and historical landmarks"
    if custom_note:
        image_prompt += f". Theme: {custom_note}"
        
    generated_image_url = image_generation_tool(image_prompt)
    output_lines.append("\n**Generated Promotional Image (smolagents):**")
    output_lines.append(f"Image URL: {generated_image_url}")
    
    # 4. Generate a promotional image using the Replicate API (Stable Diffusion) if available.
    generated_image_replicate = None
    if replicate:
        try:
            stable_diffusion = replicate.models.get("stability-ai/stable-diffusion")
            # The predict method typically returns a list of image URLs.
            result = stable_diffusion.predict(prompt=image_prompt)
            if isinstance(result, list) and len(result) > 0:
                generated_image_replicate = result[0]
            else:
                generated_image_replicate = result
            output_lines.append("\n**Generated Promotional Image (Replicate):**")
            output_lines.append(f"Image URL: {generated_image_replicate}")
        except Exception as e:
            output_lines.append(f"\nError generating image with Replicate: {e}")
    else:
        output_lines.append("\nReplicate module is not available. Install it via 'pip install replicate'.")
    
    # Choose one image to embed: Prefer the Replicate image if available, else the smolagents image.
    final_image_url = generated_image_replicate if generated_image_replicate else generated_image_url
    
    # Create a composite Markdown with text details and embed the chosen image.
    composite_markdown = "\n".join(output_lines) + f"\n\n![Chakwal Promotional Image]({final_image_url})"
    
    return AgentText(composite_markdown)


# Set up the final answer tool.
final_answer = FinalAnswerTool()

# Define the model.
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

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

# Create the CodeAgent including all our tools.
agent = CodeAgent(
    model=model,
    tools=[final_answer, my_custom_tool, get_current_time_in_timezone, prompte_chakwal, image_generation_tool],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

# Launch the Gradio UI for interactive use.
GradioUI(agent).launch()