In this notebook, we’re building a trip planning swarm which has an objective to create an itinerary together with a customer. The end result will be an itinerary that has route times and distances calculated between activities. The following diagram outlines the key components of the Swarm, with highlights being:

FalkorDB agent using a GraphRAG database of restaurants and attractions
Structured Output agent that will enforce a strict format for the accepted itinerary
Routing agent that utilises the Google Maps API to calculate distances between activities
Swarm orchestration utilising context variables

This notebook has been updated as swarms can now accommodate any ConversableAgent.

FalkorDB’s GraphRAG-SDK is a dependency for this notebook, which can be installed with ag2 via pip:

pip install -U ag2[openai,graph-rag-falkor-db]

Note: If you have been using autogen or ag2, all you need to do is upgrade it using:
pip install -U autogen[openai,graph-rag-falkor-db]
or
pip install -U ag2[openai,graph-rag-falkor-db]
as autogen, and ag2 are aliases for the same PyPI package.

For more information, please refer to the installation guide.

Pydantic

Please ensure you have Pydantic version 2+ installed.

pip install --upgrade pydantic

Running a FalkorDB

Note: You need to have a FalkorDB graph database running. If you are running one in a Docker container, please ensure your Docker network is setup to allow access to it. In this example, we’ve set the FalkorDB host and port, please adjust them accordingly. For how to set up FalkorDB, please refer to https://docs.falkordb.com/.

Google Maps API Key

To use Google’s API to calculate travel times, you will need to have enabled the Directions API in your Google Maps Platform. You can get an API key and free quota, see here and here for more details. Once you have your API key, set your environment variable GOOGLE_MAP_API_KEY to the key

Set Configuration and OpenAI API Key

By default, FalkorDB uses OpenAI LLMs and that requires an OpenAI key in your environment variable OPENAI_API_KEY. You can utilise an OAI_CONFIG_LIST file and extract the OpenAI API key and put it in the environment, as will be shown in the following cell. Alternatively, you can load the environment variable yourself.

Learn more about configuring LLMs for agents here.

import os

import autogen

llm_config = autogen.LLMConfig.from_json(path="OAI_CONFIG_LIST", timeout=120).where(tags=["gpt-4o"])

# Put the OpenAI API key into the environment
os.environ["OPENAI_API_KEY"] = llm_config.config_list[0].api_key

Prepare the FalkorDB GraphRAG database

Using 3 sample JSON data files from our GitHub repository, we will create a specific ontology for our GraphRAG database and then populate it. Creating a specific ontology that matches with the types of queries makes for a more optimal database and is more cost efficient when populating the knowledge graph.

from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType

# 3 Files (adjust path as necessary)
input_paths = [
    "../test/agentchat/contrib/graph_rag/trip_planner_data/attractions.jsonl",
    "../test/agentchat/contrib/graph_rag/trip_planner_data/cities.jsonl",
    "../test/agentchat/contrib/graph_rag/trip_planner_data/restaurants.jsonl",
]
input_documents = [Document(doctype=DocumentType.TEXT, path_or_url=input_path) for input_path in input_paths]

Create Ontology

Entities: Country, City, Attraction, Restaurant Relationships: City in Country, Attraction in City, Restaurant in City

from graphrag_sdk import Attribute, AttributeType, Entity, Ontology, Relation

# Attraction + Restaurant + City + Country Ontology
trip_data_ontology = Ontology()

trip_data_ontology.add_entity(
    Entity(
        label="Country",
        attributes=[
            Attribute(
                name="name",
                attr_type=AttributeType.STRING,
                required=True,
                unique=True,
            ),
        ],
    )
)
trip_data_ontology.add_entity(
    Entity(
        label="City",
        attributes=[
            Attribute(
                name="name",
                attr_type=AttributeType.STRING,
                required=True,
                unique=True,
            ),
            Attribute(
                name="weather",
                attr_type=AttributeType.STRING,
                required=False,
                unique=False,
            ),
            Attribute(
                name="population",
                attr_type=AttributeType.NUMBER,
                required=False,
                unique=False,
            ),
        ],
    )
)
trip_data_ontology.add_entity(
    Entity(
        label="Restaurant",
        attributes=[
            Attribute(
                name="name",
                attr_type=AttributeType.STRING,
                required=True,
                unique=True,
            ),
            Attribute(
                name="description",
                attr_type=AttributeType.STRING,
                required=False,
                unique=False,
            ),
            Attribute(
                name="rating",
                attr_type=AttributeType.NUMBER,
                required=False,
                unique=False,
            ),
            Attribute(
                name="food_type",
                attr_type=AttributeType.STRING,
                required=False,
                unique=False,
            ),
        ],
    )
)
trip_data_ontology.add_entity(
    Entity(
        label="Attraction",
        attributes=[
            Attribute(
                name="name",
                attr_type=AttributeType.STRING,
                required=True,
                unique=True,
            ),
            Attribute(
                name="description",
                attr_type=AttributeType.STRING,
                required=False,
                unique=False,
            ),
            Attribute(
                name="type",
                attr_type=AttributeType.STRING,
                required=False,
                unique=False,
            ),
        ],
    )
)
trip_data_ontology.add_relation(
    Relation(
        label="IN_COUNTRY",
        source="City",
        target="Country",
    )
)
trip_data_ontology.add_relation(
    Relation(
        label="IN_CITY",
        source="Restaurant",
        target="City",
    )
)
trip_data_ontology.add_relation(
    Relation(
        label="IN_CITY",
        source="Attraction",
        target="City",
    )
)

Establish FalkorDB and load

Remember: Change your host, port, and preferred OpenAI model if needed (gpt-4o-mini and better is recommended).

from graphrag_sdk.models.openai import OpenAiGenerativeModel

from autogen.agentchat.contrib.graph_rag.falkor_graph_query_engine import FalkorGraphQueryEngine
from autogen.agentchat.contrib.graph_rag.falkor_graph_rag_capability import FalkorGraphRagCapability

# Create FalkorGraphQueryEngine
query_engine = FalkorGraphQueryEngine(
    name="trip_data",
    host="192.168.0.115",  # Change
    port=6379,  # if needed
    ontology=trip_data_ontology,
    model=OpenAiGenerativeModel("gpt-4o"),
)

# Ingest data and initialize the database
# query_engine.init_db(input_doc=input_documents)

# If you have already ingested and created the database, you can use this connect_db instead of init_db
query_engine.connect_db()

# IMPORTS
import json
import os
from typing import Any, Dict

import requests
from pydantic import BaseModel

from autogen import (
    AfterWork,
    AfterWorkOption,
    ConversableAgent,
    OnCondition,
    SwarmResult,
    UserProxyAgent,
    initiate_swarm_chat,
    register_hand_off,
)

Pydantic model for Structured Output

Utilising OpenAI’s Structured Outputs, our Structured Output agent’s responses will be constrained to this Pydantic model. The itinerary is structured as: Itinerary has Day(s) has Event(s)

class Event(BaseModel):
    type: str  # Attraction, Restaurant, Travel
    location: str
    city: str
    description: str


class Day(BaseModel):
    events: list[Event]


class Itinerary(BaseModel):
    days: list[Day]

Google Maps Platform

The functions necessary to query the Directions API to get travel times.

def _fetch_travel_time(origin: str, destination: str) -> dict:
    """Retrieves route information using Google Maps Directions API.
    API documentation at https://developers.google.com/maps/documentation/directions/get-directions
    """
    endpoint = "https://maps.googleapis.com/maps/api/directions/json"
    params = {
        "origin": origin,
        "destination": destination,
        "mode": "walking",  # driving (default), bicycling, transit
        "key": os.environ.get("GOOGLE_MAP_API_KEY"),
    }

    response = requests.get(endpoint, params=params)
    if response.status_code == 200:
        return response.json()
    else:
        return {"error": "Failed to retrieve the route information", "status_code": response.status_code}


def update_itinerary_with_travel_times(context_variables: dict) -> SwarmResult:
    """Update the complete itinerary with travel times between each event."""
    """
    Retrieves route information using Google Maps Directions API.
    API documentation at https://developers.google.com/maps/documentation/directions/get-directions
    """

    # Ensure that we have a structured itinerary, if not, back to the structured_output_agent to make it
    if context_variables.get("structured_itinerary") is None:
        return SwarmResult(
            agent="structured_output_agent",
            values="Structured itinerary not found, please create the structured output, structured_output_agent.",
        )
    elif "timed_itinerary" in context_variables:
        return SwarmResult(values="Timed itinerary already done, inform the customer that their itinerary is ready!")

    # Process the itinerary, converting it back to an object and working through each event to work out travel time and distance
    itinerary_object = Itinerary.model_validate(json.loads(context_variables["structured_itinerary"]))
    for day in itinerary_object.days:
        events = day.events
        new_events = []
        pre_event, cur_event = None, None
        event_count = len(events)
        index = 0
        while index < event_count:
            if index > 0:
                pre_event = events[index - 1]

            cur_event = events[index]
            if pre_event:
                origin = ", ".join([pre_event.location, pre_event.city])
                destination = ", ".join([cur_event.location, cur_event.city])
                maps_api_response = _fetch_travel_time(origin=origin, destination=destination)
                try:
                    leg = maps_api_response["routes"][0]["legs"][0]
                    travel_time_txt = f"{leg['duration']['text']}, ({leg['distance']['text']})"
                    new_events.append(
                        Event(
                            type="Travel",
                            location=f"walking from {pre_event.location} to {cur_event.location}",
                            city=cur_event.city,
                            description=travel_time_txt,
                        )
                    )
                except Exception:
                    print(f"Note: Unable to get travel time from {origin} to {destination}")
            new_events.append(cur_event)
            index += 1
        day.events = new_events

    context_variables["timed_itinerary"] = itinerary_object.model_dump()

    return SwarmResult(context_variables=context_variables, values="Timed itinerary added to context with travel times")

Swarm

Context Variables

Our swarm agents will have access to a couple of context variables in relation to the itinerary.

trip_context = {
    "itinerary_confirmed": False,
    "itinerary": "",
    "structured_itinerary": None,
}

Agent Functions

We have two functions/tools for our agents. One for our Planner agent to mark an itinerary as confirmed by the customer and to store the final text itinerary. This will then transfer to our Structured Output agent. Another for the Structured Output Agent to save the structured itinerary and transfer to the Route Timing agent.

def mark_itinerary_as_complete(final_itinerary: str, context_variables: Dict[str, Any]) -> SwarmResult:
    """Store and mark our itinerary as accepted by the customer."""
    context_variables["itinerary_confirmed"] = True
    context_variables["itinerary"] = final_itinerary

    # This will update the context variables and then transfer to the Structured Output agent
    return SwarmResult(
        agent="structured_output_agent", context_variables=context_variables, values="Itinerary recorded and confirmed."
    )


def create_structured_itinerary(context_variables: Dict[str, Any], structured_itinerary: str) -> SwarmResult:
    """Once a structured itinerary is created, store it and pass on to the Route Timing agent."""
    # Ensure the itinerary is confirmed, if not, back to the Planner agent to confirm it with the customer
    if not context_variables["itinerary_confirmed"]:
        return SwarmResult(
            agent="planner_agent",
            values="Itinerary not confirmed, please confirm the itinerary with the customer first.",
        )

    context_variables["structured_itinerary"] = structured_itinerary

    # This will update the context variables and then transfer to the Route Timing agent
    return SwarmResult(
        agent="route_timing_agent", context_variables=context_variables, values="Structured itinerary stored."
    )

Agents

Our Swarm agents and a UserProxyAgent (human) which the swarm will interact with.

# Planner agent, interacting with the customer and GraphRag agent, to create an itinerary
planner_agent = ConversableAgent(
    name="planner_agent",
    system_message="You are a trip planner agent. It is important to know where the customer is going, how many days, what they want to do."
    + "You will work with another agent, graphrag_agent, to get information about restaurant and attractions. "
    + "You are also working with the customer, so you must ask the customer what they want to do if you don’t have LOCATION, NUMBER OF DAYS, MEALS, and ATTRACTIONS. "
    + "When you have the customer's requirements, work with graphrag_agent to get information for an itinerary."
    + "You are responsible for creating the itinerary and for each day in the itinerary you MUST HAVE events and EACH EVENT MUST HAVE a 'type' ('Restaurant' or 'Attraction'), 'location' (name of restaurant or attraction), 'city', and 'description'. "
    + "Finally, YOU MUST ask the customer if they are happy with the itinerary before marking the itinerary as complete.",
    functions=[mark_itinerary_as_complete],
    llm_config=llm_config,
)

# FalkorDB GraphRAG agent, utilising the FalkorDB to gather data for the Planner agent
graphrag_agent = ConversableAgent(
    name="graphrag_agent",
    system_message="Return a list of restaurants and/or attractions. List them separately and provide ALL the options in the location. Do not provide travel advice.",
)

# Adding the FalkorDB capability to the agent
graph_rag_capability = FalkorGraphRagCapability(query_engine)
graph_rag_capability.add_to_agent(graphrag_agent)

# Structured Output agent, formatting the itinerary into a structured format through the response_format on the LLM Configuration
structured_llm_config = autogen.LLMConfig.from_json(path="OAI_CONFIG_LIST", timeout=120).where(tags=["gpt-4o"])

for config in structured_llm_config.config_list:
    config.response_format = Itinerary

structured_output_agent = ConversableAgent(
    name="structured_output_agent",
    system_message="You are a data formatting agent, format the provided itinerary in the context below into the provided format.",
    llm_config=structured_llm_config,
    functions=[create_structured_itinerary],
)

# Route Timing agent, adding estimated travel times to the itinerary by utilising the Google Maps Platform
route_timing_agent = ConversableAgent(
    name="route_timing_agent",
    system_message="You are a route timing agent. YOU MUST call the update_itinerary_with_travel_times tool if you do not see the exact phrase 'Timed itinerary added to context with travel times' is seen in this conversation. Only after this please tell the customer 'Your itinerary is ready!'.",
    llm_config=llm_config,
    functions=[update_itinerary_with_travel_times],
)

# Our customer will be a human in the loop
customer = UserProxyAgent(name="customer")

Hand offs and After works

In conjunction with the agent’s associated functions, we establish rules that govern the swarm orchestration through hand offs and After works. For more details on the swarm orchestration, see the documentation.

register_hand_off(
    agent=planner_agent,
    hand_to=[
        OnCondition(
            graphrag_agent,
            "Need information on the restaurants and attractions for a location. DO NOT call more than once at a time.",
        ),  # Get info from FalkorDB GraphRAG
        OnCondition(structured_output_agent, "Itinerary is confirmed by the customer"),
        AfterWork(AfterWorkOption.REVERT_TO_USER),  # Revert to the customer for more information on their plans
    ],
)


# Back to the Planner when information has been retrieved
register_hand_off(agent=graphrag_agent, hand_to=[AfterWork(planner_agent)])

# Once we have formatted our itinerary, we can hand off to the route timing agent to add in the travel timings
register_hand_off(agent=structured_output_agent, hand_to=[AfterWork(route_timing_agent)])

# Finally, once the route timing agent has finished, we can terminate the swarm
register_hand_off(agent=route_timing_agent, hand_to=[AfterWork(AfterWorkOption.TERMINATE)])

Run the swarm

Let’s get an itinerary for a couple of days in Rome.

# Start the conversation

chat_result, context_variables, last_agent = initiate_swarm_chat(
    initial_agent=planner_agent,
    agents=[planner_agent, graphrag_agent, structured_output_agent, route_timing_agent],
    user_agent=customer,
    context_variables=trip_context,
    messages="I want to go to Rome for a couple of days. Can you help me plan my trip?",
    after_work=AfterWorkOption.TERMINATE,
    max_rounds=100,
)

Bonus itinerary output

def print_itinerary(itinerary_data):
    header = "█             █\n █           █ \n  █  █████  █  \n   ██     ██   \n  █         █  \n █  ███████  █ \n █ ██ ███ ██ █ \n   █████████   \n\n ██   ███ ███  \n█  █ █       █ \n████ █ ██  ██  \n█  █ █  █ █    \n█  █  ██  ████ \n"
    width = 80
    icons = {"Travel": "🚶", "Restaurant": "🍽️", "Attraction": "🏛️"}

    for line in header.split("\n"):
        print(line.center(width))
    print(f"Itinerary for {itinerary_data['days'][0]['events'][0]['city']}".center(width))
    print("=" * width)

    for day_num, day in enumerate(itinerary_data["days"], 1):
        print(f"\nDay {day_num}".center(width))
        print("-" * width)

        for event in day["events"]:
            event_type = event["type"]
            print(f"\n  {icons[event_type]} {event['location']}")
            if event_type != "Travel":
                words = event["description"].split()
                line = "    "
                for word in words:
                    if len(line) + len(word) + 1 <= 76:
                        line += word + " "
                    else:
                        print(line)
                        line = "    " + word + " "
                if line.strip():
                    print(line)
            else:
                print(f"    {event['description']}")
        print("\n" + "-" * width)


if "timed_itinerary" in context_variables:
    print_itinerary(context_variables["timed_itinerary"])
else:
    print("No itinerary available to print.")

Edit this page

Use Cases

Trip planning with a FalkorDB GraphRAG agent using a Swarm

Pydantic

Running a FalkorDB

Google Maps API Key

Set Configuration and OpenAI API Key

Prepare the FalkorDB GraphRAG database

Create Ontology

Establish FalkorDB and load

Pydantic model for Structured Output

Google Maps Platform

Swarm

Context Variables

Agent Functions

Agents

Hand offs and After works

Run the swarm

Bonus itinerary output

Use Cases

​Pydantic

​Running a FalkorDB

​Google Maps API Key

​Set Configuration and OpenAI API Key

​Prepare the FalkorDB GraphRAG database

​Create Ontology

​Establish FalkorDB and load

​Pydantic model for Structured Output

​Google Maps Platform

​Swarm

​Context Variables

​Agent Functions

​Agents

​Hand offs and After works

​Run the swarm

​Bonus itinerary output

Pydantic

Running a FalkorDB

Google Maps API Key

Set Configuration and OpenAI API Key

Prepare the FalkorDB GraphRAG database

Create Ontology

Establish FalkorDB and load

Pydantic model for Structured Output

Google Maps Platform

Swarm

Context Variables

Agent Functions

Agents

Hand offs and After works

Run the swarm

Bonus itinerary output