Skip to content

Commit

Permalink
GraphRAG Demo: Data Setup
Browse files Browse the repository at this point in the history
  • Loading branch information
@AgentGenie
AgentGenie committed Dec 3, 2024
1 parent 321e833 commit a2408d3
Show file tree
Hide file tree
Showing 5 changed files with 659 additions and 1 deletion.
3 changes: 2 additions & 1 deletion .pre-commit-config.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,8 @@ repos:
website/docs/tutorial/code-executors.ipynb |
website/docs/topics/code-execution/custom-executor.ipynb |
website/docs/topics/non-openai-models/cloud-gemini.ipynb |
notebook/.*
notebook/.* |
test/agentchat/contrib/graph_rag/trip_planner_data/.*
)$
# See https://jaredkhan.com/blog/mypy-pre-commit
- repo: local
Expand Down
353 changes: 353 additions & 0 deletions notebook/agentchat_swarm_graphrag_trip_planner.ipynb
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,353 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"\n",
"import autogen\n",
"\n",
"config_list = autogen.config_list_from_json(env_or_file=\"../OAI_CONFIG_LIST\")\n",
"\n",
"# Put the OpenAI API key into the environment\n",
"os.environ[\"OPENAI_API_KEY\"] = config_list[0][\"api_key\"]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from graphrag_sdk import (\n",
" Attribute,\n",
" AttributeType,\n",
" Entity,\n",
" Ontology,\n",
" Relation,\n",
")\n",
"\n",
"# Attraction + Restaurant + City + Country Ontology\n",
"trip_data_ontology = Ontology()\n",
"\n",
"\n",
"# Manually created Ontology by adding entities and relations\n",
"trip_data_ontology.add_entity(\n",
" Entity(\n",
" label=\"Country\",\n",
" attributes=[\n",
" Attribute(\n",
" name=\"name\",\n",
" attr_type=AttributeType.STRING,\n",
" required=True,\n",
" unique=True,\n",
" ),\n",
" ],\n",
" )\n",
")\n",
"trip_data_ontology.add_entity(\n",
" Entity(\n",
" label=\"City\",\n",
" attributes=[\n",
" Attribute(\n",
" name=\"name\",\n",
" attr_type=AttributeType.STRING,\n",
" required=True,\n",
" unique=True,\n",
" ),\n",
" Attribute(\n",
" name=\"weather\",\n",
" attr_type=AttributeType.STRING,\n",
" required=False,\n",
" unique=False,\n",
" ),\n",
" Attribute(\n",
" name=\"population\",\n",
" attr_type=AttributeType.NUMBER,\n",
" required=False,\n",
" unique=False,\n",
" ),\n",
" ],\n",
" )\n",
")\n",
"trip_data_ontology.add_entity(\n",
" Entity(\n",
" label=\"Restaurant\",\n",
" attributes=[\n",
" Attribute(\n",
" name=\"name\",\n",
" attr_type=AttributeType.STRING,\n",
" required=True,\n",
" unique=True,\n",
" ),\n",
" Attribute(\n",
" name=\"description\",\n",
" attr_type=AttributeType.STRING,\n",
" required=False,\n",
" unique=False,\n",
" ),\n",
" Attribute(\n",
" name=\"rating\",\n",
" attr_type=AttributeType.NUMBER,\n",
" required=False,\n",
" unique=False,\n",
" ),\n",
" Attribute(\n",
" name=\"food_type\",\n",
" attr_type=AttributeType.STRING,\n",
" required=False,\n",
" unique=False,\n",
" ),\n",
" ],\n",
" )\n",
")\n",
"trip_data_ontology.add_relation(\n",
" Relation(\n",
" label=\"IN_COUNTRY\",\n",
" source=\"City\",\n",
" target=\"Country\",\n",
" )\n",
")\n",
"trip_data_ontology.add_relation(\n",
" Relation(\n",
" label=\"IN_CITY\",\n",
" source=\"Restaurant\",\n",
" target=\"City\",\n",
" )\n",
")\n",
"\n",
"trip_data_ontology.add_entity(\n",
" Entity(\n",
" label=\"Country\",\n",
" attributes=[\n",
" Attribute(\n",
" name=\"name\",\n",
" attr_type=AttributeType.STRING,\n",
" required=True,\n",
" unique=True,\n",
" ),\n",
" ],\n",
" )\n",
")\n",
"trip_data_ontology.add_entity(\n",
" Entity(\n",
" label=\"City\",\n",
" attributes=[\n",
" Attribute(\n",
" name=\"name\",\n",
" attr_type=AttributeType.STRING,\n",
" required=True,\n",
" unique=True,\n",
" ),\n",
" Attribute(\n",
" name=\"weather\",\n",
" attr_type=AttributeType.STRING,\n",
" required=False,\n",
" unique=False,\n",
" ),\n",
" Attribute(\n",
" name=\"population\",\n",
" attr_type=AttributeType.NUMBER,\n",
" required=False,\n",
" unique=False,\n",
" ),\n",
" ],\n",
" )\n",
")\n",
"trip_data_ontology.add_entity(\n",
" Entity(\n",
" label=\"Attraction\",\n",
" attributes=[\n",
" Attribute(\n",
" name=\"name\",\n",
" attr_type=AttributeType.STRING,\n",
" required=True,\n",
" unique=True,\n",
" ),\n",
" Attribute(\n",
" name=\"description\",\n",
" attr_type=AttributeType.STRING,\n",
" required=False,\n",
" unique=False,\n",
" ),\n",
" Attribute(\n",
" name=\"type\",\n",
" attr_type=AttributeType.STRING,\n",
" required=False,\n",
" unique=False,\n",
" ),\n",
" ],\n",
" )\n",
")\n",
"trip_data_ontology.add_relation(\n",
" Relation(\n",
" label=\"IN_COUNTRY\",\n",
" source=\"City\",\n",
" target=\"Country\",\n",
" )\n",
")\n",
"trip_data_ontology.add_relation(\n",
" Relation(\n",
" label=\"IN_CITY\",\n",
" source=\"Attraction\",\n",
" target=\"City\",\n",
" )\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[33muser_proxy\u001b[0m (to graph_rag_agent):\n",
"\n",
"Find 3 restaurant in Milan.\n",
"\n",
"--------------------------------------------------------------------------------\n",
"\u001b[33mgraph_rag_agent\u001b[0m (to user_proxy):\n",
"\n",
"Bros' Burger in Milan is known for its gourmet burgers and offers a relaxed vibe.\n",
"\n",
"--------------------------------------------------------------------------------\n"
]
},
{
"data": {
"text/plain": [
"ChatResult(chat_id=None, chat_history=[{'content': 'Find 3 restaurant in Milan.', 'role': 'assistant', 'name': 'user_proxy'}, {'content': \"Bros' Burger in Milan is known for its gourmet burgers and offers a relaxed vibe.\", 'role': 'user', 'name': 'graph_rag_agent'}], summary=\"Bros' Burger in Milan is known for its gourmet burgers and offers a relaxed vibe.\", cost={'usage_including_cached_inference': {'total_cost': 0}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=['exit'])"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from autogen import ConversableAgent, UserProxyAgent\n",
"from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType\n",
"from autogen.agentchat.contrib.graph_rag.falkor_graph_query_engine import FalkorGraphQueryEngine\n",
"from autogen.agentchat.contrib.graph_rag.falkor_graph_rag_capability import FalkorGraphRagCapability\n",
"\n",
"# Auto generate graph schema from unstructured data\n",
"input_paths = [\n",
" \"../test/agentchat/contrib/graph_rag/trip_planner_data/attractions.json\",\n",
" \"../test/agentchat/contrib/graph_rag/trip_planner_data/cities.json\",\n",
" \"../test/agentchat/contrib/graph_rag/trip_planner_data/restaurants.json\",\n",
"]\n",
"input_documents = [Document(doctype=DocumentType.TEXT, path_or_url=input_path) for input_path in input_paths]\n",
"\n",
"# Create FalkorGraphQueryEngine\n",
"query_engine = FalkorGraphQueryEngine(\n",
" name=\"trip_data\",\n",
" host=\"172.18.0.3\", # Change\n",
" port=6379, # if needed\n",
" ontology=trip_data_ontology,\n",
")\n",
"\n",
"# Ingest data and initialize the database\n",
"query_engine.init_db(input_doc=input_documents)\n",
"\n",
"# Create a ConversableAgent (no LLM configuration)\n",
"graph_rag_agent = ConversableAgent(\n",
" name=\"graph_rag_agent\",\n",
" human_input_mode=\"NEVER\",\n",
")\n",
"\n",
"# Associate the capability with the agent\n",
"graph_rag_capability = FalkorGraphRagCapability(query_engine)\n",
"graph_rag_capability.add_to_agent(graph_rag_agent)\n",
"\n",
"# Create a user proxy agent to converse with our RAG agent\n",
"user_proxy = UserProxyAgent(\n",
" name=\"user_proxy\",\n",
" human_input_mode=\"ALWAYS\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[33muser_proxy\u001b[0m (to graph_rag_agent):\n",
"\n",
"Find 3 restaurants in Rome.\n",
"\n",
"--------------------------------------------------------------------------------\n",
"\u001b[33mgraph_rag_agent\u001b[0m (to user_proxy):\n",
"\n",
"Three restaurants in Rome are Trattoria da Enzo, known for its traditional Roman dishes and welcoming atmosphere; Il Pagliaccio, an elegant Michelin-starred restaurant offering contemporary Italian cuisine; and Tonnarello, a casual spot for pasta, meatballs, and other simple Roman dishes with patio seating and acoustic guitar music.\n",
"\n",
"--------------------------------------------------------------------------------\n",
"\u001b[33muser_proxy\u001b[0m (to graph_rag_agent):\n",
"\n",
"Find 3 attractions in Rome.\n",
"\n",
"--------------------------------------------------------------------------------\n",
"\u001b[33mgraph_rag_agent\u001b[0m (to user_proxy):\n",
"\n",
"Three attractions in Rome are the Colosseum, an ancient amphitheater known for gladiatorial contests and public spectacles; the Vatican Museums, a complex of museums and galleries showcasing works of art collected by Popes over centuries; and the Trevi Fountain, a Baroque fountain famous for its stunning sculptures and the tradition of tossing coins.\n",
"\n",
"--------------------------------------------------------------------------------\n",
"\u001b[33muser_proxy\u001b[0m (to graph_rag_agent):\n",
"\n",
"How many attractions in Rome are in the database?\n",
"\n",
"--------------------------------------------------------------------------------\n",
"\u001b[33mgraph_rag_agent\u001b[0m (to user_proxy):\n",
"\n",
"There are 4 attractions in Rome listed in the database.\n",
"\n",
"--------------------------------------------------------------------------------\n"
]
},
{
"data": {
"text/plain": [
"ChatResult(chat_id=None, chat_history=[{'content': 'Find 3 restaurants in Rome.', 'role': 'assistant', 'name': 'user_proxy'}, {'content': 'Three restaurants in Rome are Trattoria da Enzo, known for its traditional Roman dishes and welcoming atmosphere; Il Pagliaccio, an elegant Michelin-starred restaurant offering contemporary Italian cuisine; and Tonnarello, a casual spot for pasta, meatballs, and other simple Roman dishes with patio seating and acoustic guitar music.', 'role': 'user', 'name': 'graph_rag_agent'}, {'content': 'Find 3 attractions in Rome.', 'role': 'assistant', 'name': 'user_proxy'}, {'content': 'Three attractions in Rome are the Colosseum, an ancient amphitheater known for gladiatorial contests and public spectacles; the Vatican Museums, a complex of museums and galleries showcasing works of art collected by Popes over centuries; and the Trevi Fountain, a Baroque fountain famous for its stunning sculptures and the tradition of tossing coins.', 'role': 'user', 'name': 'graph_rag_agent'}, {'content': 'How many attractions in Rome are in the database?', 'role': 'assistant', 'name': 'user_proxy'}, {'content': 'There are 4 attractions in Rome listed in the database.', 'role': 'user', 'name': 'graph_rag_agent'}], summary='There are 4 attractions in Rome listed in the database.', cost={'usage_including_cached_inference': {'total_cost': 0}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=['Find 3 attractions in Rome.', 'How many attractions in Rome are in the database?', 'exit'])"
]
},
"execution_count": 12,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"user_proxy.initiate_chat(graph_rag_agent, message=\"Find 3 restaurants in Rome.\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.11.10"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Loading

0 comments on commit a2408d3

Please sign in to comment.