| from smolagents import CodeAgent, DuckDuckGoSearchTool, LiteLLMModel, OpenAIServerModel, tool |
| import datetime |
| from superhero_party_theme_tool import SuperheroPartyThemeTool |
| from model import model |
|
|
| @tool |
| def suggest_menu(occasion: str) -> str: |
| """ |
| Suggests a menu based on the occasion. |
| Args: |
| occasion: The type of occasion for the party. Allowed values are: |
| - "casual": Menu for casual party. |
| - "formal": Menu for formal party. |
| - "superhero": Menu for superhero party. |
| - "custom": Custom menu. |
| """ |
| 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." |
|
|
|
|
| def main_v1() -> None: |
| |
| music_agent = CodeAgent(tools=[DuckDuckGoSearchTool()], model=model) |
| music_agent.run("Search for the best music recommendations for a party at the Wayne's mansion.") |
|
|
| |
| menu_agent = CodeAgent(tools=[suggest_menu], model=model) |
| menu_agent.run("Prepare a formal menu for the party.") |
|
|
| def main_v2() -> None: |
| agent = CodeAgent(tools=[], model=model, additional_authorized_imports=['datetime']) |
|
|
| agent.run( |
| """ |
| Alfred needs to prepare for the party. Here are the tasks: |
| 1. Prepare the drinks - 30 minutes |
| 2. Decorate the mansion - 60 minutes |
| 3. Set up the menu - 45 minutes |
| 4. Prepare the music and playlist - 45 minutes |
| |
| If we start right now, at what time will the party be ready? |
| """ |
| ) |
|
|
| def main() -> None: |
| agent = CodeAgent(tools=[SuperheroPartyThemeTool()], model=model) |
|
|
| agent.run( |
| """ |
| whats the theme for classic heroes party |
| """ |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
