Merge branch 'graphrag_demo' of https://github.com/ag2ai/ag2 into gra…

…phrag_demo
ag2ai · Dec 4, 2024 · ebc2eda · ebc2eda
2 parents b61c223 + bf81c00
commit ebc2eda
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 33 deletions.
diff --git a/notebook/agentchat_swarm_graphrag_trip_planner.ipynb b/notebook/agentchat_swarm_graphrag_trip_planner.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -18,7 +18,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -200,36 +200,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "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"
-    }
-   ],
+   "outputs": [],
    "source": [
     "from autogen import ConversableAgent, UserProxyAgent\n",
     "from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType\n",
@@ -238,8 +211,8 @@
     "\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/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",

diff --git a/test/agentchat/contrib/graph_rag/trip_planner_data/restaurants.json b/test/agentchat/contrib/graph_rag/trip_planner_data/restaurants.json
@@ -86,5 +86,13 @@
     "city": "Rome",
     "rating": 4.7,
     "food_type": "Italian"
+  },
+  {
+    "id": 12,
+    "name": "Osteria delle Commari",
+    "description": "Local home style Roman restaurant near Vatican Museums.",
+    "city": "Rome",
+    "rating": 4.3,
+    "food_type": "Italian"
   }
 ]