{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Selector Group Chat" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import asyncio\n", "from typing import Sequence\n", "\n", "from autogen_agentchat.agents import (\n", " BaseChatAgent,\n", " ChatMessage,\n", " CodingAssistantAgent,\n", " StopMessage,\n", " TextMessage,\n", " ToolUseAssistantAgent,\n", ")\n", "from autogen_agentchat.teams import SelectorGroupChat, StopMessageTermination\n", "from autogen_core.base import CancellationToken\n", "from autogen_core.components.models import OpenAIChatCompletionClient\n", "from autogen_core.components.tools import FunctionTool" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "class UserProxyAgent(BaseChatAgent):\n", " def __init__(self, name: str) -> None:\n", " super().__init__(name, \"A human user.\")\n", "\n", " async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> ChatMessage:\n", " user_input = await asyncio.get_event_loop().run_in_executor(None, input, \"Enter your response: \")\n", " if \"TERMINATE\" in user_input:\n", " return StopMessage(content=\"User has terminated the conversation.\", source=self.name)\n", " return TextMessage(content=user_input, source=self.name)" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "async def flight_search(start: str, destination: str, date: str) -> str:\n", " return \"\\n\".join(\n", " [\n", " f\"AC24 from {start} to {destination} on {date} is $500\",\n", " f\"UA23 from {start} to {destination} on {date} is $450\",\n", " f\"AL21 from {start} to {destination} on {date} is $400\",\n", " ]\n", " )\n", "\n", "\n", "async def flight_booking(flight: str, date: str) -> str:\n", " return f\"Booked flight {flight} on {date}\"" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "--------------------------------------------------------------------------- \n", "\u001b[91m[2024-10-08T20:35:30.283450]:\u001b[0m\n", "\n", "Help user plan a trip and book a flight." ] }, { "name": "stdout", "output_type": "stream", "text": [ "\n", "--------------------------------------------------------------------------- \n", "\u001b[91m[2024-10-08T20:35:48.275743], User:\u001b[0m\n", "\n", "\n", "--------------------------------------------------------------------------- \n", "\u001b[91m[2024-10-08T20:35:50.795496], TravelAssistant:\u001b[0m\n", "\n", "I'd be happy to help you plan your trip! To get started, could you please provide me with the following details:\n", "\n", "1. Your departure city and the destination city.\n", "2. Your travel dates (departure and return).\n", "3. The number of travelers and their ages (if any children are involved).\n", "4. Your budget for flights and accommodations, if you have one in mind.\n", "5. Any specific activities or attractions you're interested in at the destination.\n", "\n", "Once I have this information, I can help you find the best options!\n", "--------------------------------------------------------------------------- \n", "\u001b[91m[2024-10-08T20:35:59.701486], User:\u001b[0m\n", "\n", "Traveling to toronto from new york\n", "--------------------------------------------------------------------------- \n", "\u001b[91m[2024-10-08T20:36:02.325330], TravelAssistant:\u001b[0m\n", "\n", "Great choice! Toronto is a vibrant city with a lot to offer. Now, could you please provide the following additional details to help me assist you better?\n", "\n", "1. What are your travel dates (departure and return)?\n", "2. How many travelers will be going, and what are their ages?\n", "3. Do you have a budget for the flight and accommodations?\n", "4. Are there any specific activities or attractions you’re interested in while in Toronto?\n", "\n", "Once I have this information, I can help you find the best flights and suggestions for your trip!\n", "--------------------------------------------------------------------------- \n", "\u001b[91m[2024-10-08T20:36:20.633004], User:\u001b[0m\n", "\n", "leaving on december 7 and returning on 12\n", "--------------------------------------------------------------------------- \n", "\u001b[91m[2024-10-08T20:36:23.202871], TravelAssistant:\u001b[0m\n", "\n", "Thank you for the details! Here's what I have so far:\n", "\n", "- **Departure City:** New York\n", "- **Destination City:** Toronto\n", "- **Departure Date:** December 7\n", "- **Return Date:** December 12\n", "\n", "Now, could you please provide:\n", "\n", "1. The number of travelers and their ages.\n", "2. Your budget for flights and accommodations (if applicable).\n", "3. Any specific activities or attractions you're interested in while in Toronto.\n", "\n", "This will help me provide more tailored options for your trip!\n", "--------------------------------------------------------------------------- \n", "\u001b[91m[2024-10-08T20:36:38.096554], User:\u001b[0m\n", "\n", "just myself one adult\n", "--------------------------------------------------------------------------- \n", "\u001b[91m[2024-10-08T20:36:40.307824], FlightBroker:\u001b[0m\n", "\n", "Thanks for the information! Here's what I have:\n", "\n", "- **Departure City:** New York\n", "- **Destination City:** Toronto\n", "- **Departure Date:** December 7\n", "- **Return Date:** December 12\n", "- **Number of Travelers:** 1 Adult\n", "\n", "Could you let me know if you have a budget for flights and accommodations? Additionally, are there any specific activities or attractions you're interested in while in Toronto? This will help me provide the best options for your trip!\n", "--------------------------------------------------------------------------- \n", "\u001b[91m[2024-10-08T20:36:45.875280], User:\u001b[0m\n", "\n", "that's it\n", "--------------------------------------------------------------------------- \n", "\u001b[91m[2024-10-08T20:36:50.925624], FlightBroker:\u001b[0m\n", "\n", "Your flights have been successfully booked! Here are the details:\n", "\n", "- **Departure:** New York to Toronto\n", " - **Flight:** AL21\n", " - **Date:** December 7, 2023\n", "\n", "- **Return:** Toronto to New York\n", " - **Flight:** AL21\n", " - **Date:** December 12, 2023\n", "\n", "If you need help with accommodations, activities, or anything else for your trip, feel free to let me know! \n", "\n", "TERMINATE" ] }, { "data": { "text/plain": [ "TeamRunResult(messages=[TextMessage(source='user', content='Help user plan a trip and book a flight.'), TextMessage(source='User', content=''), TextMessage(source='TravelAssistant', content=\"I'd be happy to help you plan your trip! To get started, could you please provide me with the following details:\\n\\n1. Your departure city and the destination city.\\n2. Your travel dates (departure and return).\\n3. The number of travelers and their ages (if any children are involved).\\n4. Your budget for flights and accommodations, if you have one in mind.\\n5. Any specific activities or attractions you're interested in at the destination.\\n\\nOnce I have this information, I can help you find the best options!\"), TextMessage(source='User', content='Traveling to toronto from new york'), TextMessage(source='TravelAssistant', content='Great choice! Toronto is a vibrant city with a lot to offer. Now, could you please provide the following additional details to help me assist you better?\\n\\n1. What are your travel dates (departure and return)?\\n2. How many travelers will be going, and what are their ages?\\n3. Do you have a budget for the flight and accommodations?\\n4. Are there any specific activities or attractions you’re interested in while in Toronto?\\n\\nOnce I have this information, I can help you find the best flights and suggestions for your trip!'), TextMessage(source='User', content='leaving on december 7 and returning on 12'), TextMessage(source='TravelAssistant', content=\"Thank you for the details! Here's what I have so far:\\n\\n- **Departure City:** New York\\n- **Destination City:** Toronto\\n- **Departure Date:** December 7\\n- **Return Date:** December 12\\n\\nNow, could you please provide:\\n\\n1. The number of travelers and their ages.\\n2. Your budget for flights and accommodations (if applicable).\\n3. Any specific activities or attractions you're interested in while in Toronto.\\n\\nThis will help me provide more tailored options for your trip!\"), TextMessage(source='User', content='just myself one adult'), TextMessage(source='FlightBroker', content=\"Thanks for the information! Here's what I have:\\n\\n- **Departure City:** New York\\n- **Destination City:** Toronto\\n- **Departure Date:** December 7\\n- **Return Date:** December 12\\n- **Number of Travelers:** 1 Adult\\n\\nCould you let me know if you have a budget for flights and accommodations? Additionally, are there any specific activities or attractions you're interested in while in Toronto? This will help me provide the best options for your trip!\"), TextMessage(source='User', content=\"that's it\"), StopMessage(source='FlightBroker', content='Your flights have been successfully booked! Here are the details:\\n\\n- **Departure:** New York to Toronto\\n - **Flight:** AL21\\n - **Date:** December 7, 2023\\n\\n- **Return:** Toronto to New York\\n - **Flight:** AL21\\n - **Date:** December 12, 2023\\n\\nIf you need help with accommodations, activities, or anything else for your trip, feel free to let me know! \\n\\nTERMINATE'), StopMessage(source='StopMessageTermination', content='Stop message received')])" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "user_proxy = UserProxyAgent(\"User\")\n", "flight_broker = ToolUseAssistantAgent(\n", " \"FlightBroker\",\n", " description=\"An assistant for booking flights\",\n", " model_client=OpenAIChatCompletionClient(model=\"gpt-4o-mini\"),\n", " registered_tools=[\n", " FunctionTool(flight_search, description=\"Search for flights\"),\n", " FunctionTool(flight_booking, description=\"Book a flight\"),\n", " ],\n", ")\n", "travel_assistant = CodingAssistantAgent(\n", " \"TravelAssistant\",\n", " description=\"A travel assistant\",\n", " model_client=OpenAIChatCompletionClient(model=\"gpt-4o-mini\"),\n", " system_message=\"You are a travel assistant.\",\n", ")\n", "team = SelectorGroupChat(\n", " [user_proxy, flight_broker, travel_assistant], model_client=OpenAIChatCompletionClient(model=\"gpt-4o-mini\")\n", ")\n", "await team.run(\"Help user plan a trip and book a flight.\", termination_condition=StopMessageTermination())" ] } ], "metadata": { "kernelspec": { "display_name": ".venv", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.6" } }, "nbformat": 4, "nbformat_minor": 2 }