Building an AI Travel Planner with AutoGen: A Step-by-Step Guide

# user_agent.py

from autogen_agentchat.agents import UserProxyAgent
from autogen_agentchat.message import AssistantMessage

class UserAgent(UserProxyAgent):
    pass  # Inherits functionality from UserProxyAgent

# flight_agent.py

import asyncio
from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models import OpenAIChatCompletionClient

async def search_flights(departure_city: str, destination_city: str, departure_date: str, return_date: str):
    # Mock implementation of flight search
    await asyncio.sleep(1)  # Simulate network delay
    return f"Found flights from {departure_city} to {destination_city} departing on {departure_date} and returning on {return_date}."

flight_agent = AssistantAgent(
    name="FlightAgent",
    model_client=OpenAIChatCompletionClient(
        model="gpt-4",
        # api_key will be loaded from environment variable
    ),
    instructions="""
You are an AI agent specialized in booking flights. Assist in finding flights based on user preferences.
""",
    tools=[search_flights],
)

# hotel_agent.py

import asyncio
from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models import OpenAIChatCompletionClient

async def search_hotels(destination_city: str, check_in_date: str, check_out_date: str):
    # Mock implementation of hotel search
    await asyncio.sleep(1)  # Simulate network delay
    return f"Found hotels in {destination_city} from {check_in_date} to {check_out_date}."

hotel_agent = AssistantAgent(
    name="HotelAgent",
    model_client=OpenAIChatCompletionClient(
        model="gpt-4",
    ),
    instructions="""
You are an AI agent specialized in booking accommodations. Assist in finding hotels based on user preferences.
""",
    tools=[search_hotels],
)

# activity_agent.py

import asyncio
from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models import OpenAIChatCompletionClient

async def suggest_activities(destination_city: str, interests: str):
    # Mock implementation of activity suggestions
    await asyncio.sleep(1)  # Simulate processing time
    return f"Suggested activities in {destination_city} based on your interests ({interests}): Visit the museum, explore downtown, enjoy local cuisine."

activity_agent = AssistantAgent(
    name="ActivityAgent",
    model_client=OpenAIChatCompletionClient(
        model="gpt-4",
    ),
    instructions="""
You are an AI agent specialized in suggesting activities and attractions. Provide recommendations based on user interests.
""",
    tools=[suggest_activities],
)

# main.py

import asyncio
import os
from dotenv import load_dotenv
from autogen_agentchat.agents import UserProxyAgent
from autogen_agentchat.teams import SequentialTeam
from autogen_agentchat.task import Console
from autogen_ext.models import OpenAIChatCompletionClient

# Import agents
from flight_agent import flight_agent
from hotel_agent import hotel_agent
from activity_agent import activity_agent

# Load environment variables
load_dotenv()
openai_api_key = os.getenv("OPENAI_API_KEY")

# Ensure API key is set
if not openai_api_key:
    raise ValueError("OPENAI_API_KEY is not set in the environment variables.")

# Set the API key for model clients
flight_agent.model_client.api_key = openai_api_key
hotel_agent.model_client.api_key = openai_api_key
activity_agent.model_client.api_key = openai_api_key

async def main():
    # Create the user agent
    user_agent = UserProxyAgent(
        name="UserAgent",
    )

    # Define the travel planning team
    travel_team = SequentialTeam(
        agents=[
            flight_agent,
            hotel_agent,
            activity_agent,
        ],
        user_agent=user_agent,
    )

    # Initial user message
    user_message = input("You: ")

    # Run the team
    stream = travel_team.run_stream(task=user_message)
    await Console(stream)

if __name__ == "__main__":
    asyncio.run(main())