from smolagents import Tool
import random

class SuperheroPartyThemeTool(Tool):
    name = "superhero_party_theme"
    description = "Extract the party category it can be classic heroes, villain masquerade, futuristic Gotham and return the theme"

       
    inputs = {
        "category": {
            "type": "string",
            "description": "The type of superhero party (e.g., 'classic heroes', 'villain masquerade', 'futuristic Gotham').",
        }
    }

    output_type = "string"

    def forward(self, category: str) -> str:
        themes = {
            "classic heroes": "Superman, Batman, Wonder Woman, The Flash, Aquaman, Cyborg, Green Lantern",
            "villain masquerade": "Joker, Lex Luthor, Catwoman, The Riddler, The Penguin, The Joker, The Riddler",
            "futuristic Gotham": "Batman, Catwoman, The Joker, The Riddler, The Penguin, The Batman, The Catwoman, The Joker, The Riddler, The Penguin",
        }
        return themes.get(category.lower(), "Themed party idea not found. Try 'classic heroes', 'villain masquerade', or 'futuristic Gotham'.")