deveshka's picture
added custom tool for party planner
3fffc75 verified
from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool,Tool,VisitWebpageTool
import datetime
import requests
import os
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
@tool
def suggest_menu(occasion: str) -> str:
"""
Suggests a menu based on the occasion.
Args:
occasion: The type of occasion for the party.
"""
if occasion == "casual":
return "Pizza, snacks, and drinks."
elif occasion == "formal":
return "3-course dinner with wine and dessert."
elif occasion == "superhero":
return "Buffet with high-energy and healthy food."
else:
return "Custom menu for the butler."
@tool
def catering_service_tool(query: str) -> str:
"""
This tool returns the highest-rated catering service in Gotham City.
Args:
query: A search term for finding catering services.
"""
# Example list of catering services and their ratings
services = {
"Gotham Catering Co.": 4.9,
"Wayne Manor Catering": 4.8,
"Gotham City Events": 4.7,
}
# Find the highest rated catering service (simulating search query filtering)
best_service = max(services, key=services.get)
return best_service
class SuperheroPartyThemeTool(Tool):
name = "superhero_party_theme_generator"
description = """
This tool suggests creative superhero-themed party ideas based on a category.
It returns a unique party theme idea."""
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):
themes = {
"classic heroes": "Justice League Gala: Guests come dressed as their favorite DC heroes with themed cocktails like 'The Kryptonite Punch'.",
"villain masquerade": "Gotham Rogues' Ball: A mysterious masquerade where guests dress as classic Batman villains.",
"futuristic Gotham": "Neo-Gotham Night: A cyberpunk-style party inspired by Batman Beyond, with neon decorations and futuristic gadgets."
}
return themes.get(category.lower(), "Themed party idea not found. Try 'classic heroes', 'villain masquerade', or 'futuristic Gotham'.")
final_answer = FinalAnswerTool()
# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud'
model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)
# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
with open("prompts.yaml", 'r') as stream:
prompt_templates = yaml.safe_load(stream)
agent = CodeAgent(
model=model,
tools=[final_answer, DuckDuckGoSearchTool(),VisitWebpageTool(),
suggest_menu,
catering_service_tool,
SuperheroPartyThemeTool()
], ## add your tools here (don't remove final answer)
max_steps=6,
verbosity_level=1,
grammar=None,
planning_interval=None,
name="Party planner",
description="Search for the best music recommendations for a party at the Wayne's mansion.",
prompt_templates=prompt_templates,
additional_authorized_imports=['datetime']
)
GradioUI(agent).launch()