Trip planning with a FalkorDB GraphRAG agent using a Swarm
FalkorDB GraphRAG utilises a knowledge graph and can be added as a capability to agents. Together with a swarm orchestration of agents is highly effective at providing a RAG capability.
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:
Note: If you have been using
autogen
orpyautogen
, all you need to do is upgrade it using:or
as
pyautogen
,autogen
, andag2
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.
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.
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.
Create Ontology
Entities: Country, City, Attraction, Restaurant
Relationships: City in Country, Attraction in City, Restaurant in City
Establish FalkorDB and load
Remember: Change your host, port, and preferred OpenAI model if needed (gpt-4o-mini and better is recommended).
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)
Google Maps Platform
The functions necessary to query the Directions API to get travel times.
Swarm
Context Variables
Our swarm agents will have access to a couple of context variables in relation to the itinerary.
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.
Agents
Our Swarm agents and a UserProxyAgent (human) which the swarm will interact with.
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.
Run the swarm
Let’s get an itinerary for a couple of days in Rome.