add class tool SuperheroPartyThemeTool

app.py

@@ -3,6 +3,7 @@ import os
 from dotenv import load_dotenv
 from smolagents import CodeAgent, DuckDuckGoSearchTool, LiteLLMModel, OpenAIServerModel, tool
 import datetime
+from superhero_party_theme_tool import SuperheroPartyThemeTool
 
 load_dotenv()  # reads MODEL_PROVIDER and OPENAI_API_KEY from .env
 
@@ -61,7 +62,7 @@ def main_v1() -> None:
     menu_agent = CodeAgent(tools=[suggest_menu], model=model)
     menu_agent.run("Prepare a formal menu for the party.")
 
-def main() -> None:
+def main_v2() -> None:
     agent = CodeAgent(tools=[], model=model, additional_authorized_imports=['datetime'])
 
     agent.run(
@@ -76,6 +77,15 @@ def main() -> None:
         """
     )
 
+def main() -> None:
+    agent = CodeAgent(tools=[SuperheroPartyThemeTool()], model=model)
+
+    agent.run(
+        """
+        whats the theme for classic heroes party
+        """
+    )
+
 
 if __name__ == "__main__":
     main()

superhero_party_theme_tool.py

@@ -0,0 +1,24 @@
+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'.")