AG2 provides GraphRAG integration using agent capabilities. This is an
example to integrate FalkorDB (a Knowledge Graph database).
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.
Set Configuration and OpenAI API Key
By default, in order to use FalkorDB you need to have 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
config_list = autogen.config_list_from_json(env_or_file="OAI_CONFIG_LIST")
# Put the OpenAI API key into the environment
os.environ["OPENAI_API_KEY"] = config_list[0]["api_key"]
Create a Knowledge Graph with Your Own Data
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, the FalkorDB endpoint is set to host=“172.18.0.3” and
port=6379, please adjust accordingly. For how to set up FalkorDB, please
refer to https://docs.falkordb.com/
Below, we have some sample data from IMDB on the movie ‘The Matrix’. See
the contents of the file
here.
We then initialise the database with that text document, creating the
graph in FalkorDB.
A Simple Example
In this example, the graph ontology is auto-generated. This allows you
to load data without specifying the specific types of entities and
relationships that will make up the database (however, this may not be
optimal and not cost efficient).
from autogen import ConversableAgent, UserProxyAgent
from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType
from autogen.agentchat.contrib.graph_rag.falkor_graph_query_engine import FalkorGraphQueryEngine
from autogen.agentchat.contrib.graph_rag.falkor_graph_rag_capability import FalkorGraphRagCapability
# Auto generate graph schema from unstructured data
input_path = "../test/agentchat/contrib/graph_rag/the_matrix.txt"
input_documents = [Document(doctype=DocumentType.TEXT, path_or_url=input_path)]
# Create FalkorGraphQueryEngine
query_engine = FalkorGraphQueryEngine(
name="The_Matrix_Auto",
host="172.17.0.4", # Change
port=6379, # if needed
)
# Ingest data and initialize the database
query_engine.init_db(input_doc=input_documents)
# Create a ConversableAgent (no LLM configuration)
graph_rag_agent = ConversableAgent(
name="matrix_agent",
human_input_mode="NEVER",
)
# Associate the capability with the agent
graph_rag_capability = FalkorGraphRagCapability(query_engine)
graph_rag_capability.add_to_agent(graph_rag_agent)
# Create a user proxy agent to converse with our RAG agent
user_proxy = UserProxyAgent(
name="user_proxy",
human_input_mode="ALWAYS",
)
user_proxy.initiate_chat(graph_rag_agent, message="Name a few actors who've played in 'The Matrix'")
Revisit the example with self defined graph ontology
If you want to be more specific on the entities and attributes in the
graph database you can create an Ontology before loading your data, in
the case below we define these entities (Actor with name and Movie
with title) and relationships (Actors ACTED in Movies). This allows
the RAG agent to answer questions about actors in the movie.
from graphrag_sdk import Attribute, AttributeType, Entity, Ontology, Relation
from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType
from autogen.agentchat.contrib.graph_rag.falkor_graph_query_engine import FalkorGraphQueryEngine
input_path = "../test/agentchat/contrib/graph_rag/the_matrix.txt"
movie_ontology = Ontology()
movie_ontology.add_entity(
Entity(label="Actor", attributes=[Attribute(name="name", attr_type=AttributeType.STRING, unique=True)])
)
movie_ontology.add_entity(
Entity(label="Movie", attributes=[Attribute(name="title", attr_type=AttributeType.STRING, unique=True)])
)
movie_ontology.add_relation(Relation(label="ACTED", source="Actor", target="Movie"))
query_engine = FalkorGraphQueryEngine(
name="IMDB",
host="172.17.0.4", # Change
port=6379, # if needed
ontology=movie_ontology,
)
input_documents = [Document(doctype=DocumentType.TEXT, path_or_url=input_path)]
query_engine.init_db(input_doc=input_documents)
Add capability to a ConversableAgent and query them
With FalkorDB setup with an ontology and data, we can now associate the
GraphRAG capability with a ConversableAgent and have a chat with it.
from autogen import ConversableAgent, UserProxyAgent
from autogen.agentchat.contrib.graph_rag.falkor_graph_rag_capability import FalkorGraphRagCapability
# Create a ConversableAgent (no LLM configuration)
graph_rag_agent = ConversableAgent(
name="matrix_agent",
human_input_mode="NEVER",
)
# Associate the capability with the agent
graph_rag_capability = FalkorGraphRagCapability(query_engine)
graph_rag_capability.add_to_agent(graph_rag_agent)
# Create a user proxy agent to converse with our RAG agent
user_proxy = UserProxyAgent(
name="user_proxy",
code_execution_config=False,
is_termination_msg=lambda msg: "TERMINATE" in msg["content"],
human_input_mode="ALWAYS",
)
user_proxy.initiate_chat(graph_rag_agent, message="Name a few actors who've played in 'The Matrix'")
# You will be prompted, as a human in the loop, after the response - feel free to ask more questions.
From the output we can see that: - We can ask about actors and movies
because we defined the ontology (and they were also automatically
created in the prior example) - There is inbuilt awareness of answers it
has provided previously - It can’t answer any questions outside of the
actors and movies entities as we didn’t include anything else in the
ontology
# Remove graph from database:
query_engine.delete()
