[GraphRAG] FalkorDB graph rag integration

autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py

@@ -1,11 +1,9 @@
 # Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
 #
 # SPDX-License-Identifier: Apache-2.0
-#
-# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
-# SPDX-License-Identifier: MIT
+
 import os
-from dataclasses import field
+from dataclasses import dataclass, field
 from typing import List
 
 from graphrag_sdk import KnowledgeGraph, Source
@@ -15,6 +13,7 @@
 from .graph_query_engine import GraphStoreQueryResult
 
 
+@dataclass
 class FalkorGraphQueryResult(GraphStoreQueryResult):
     messages: list = field(default_factory=list)
 
@@ -36,7 +35,7 @@ def __init__(
     ):
         """
         Initialize a Falkor DB knowledge graph.
-        Please also refer to https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/kg.py
+        Please also refer to https://github.com/FalkorDB/GraphRAG-SDK/blob/2-move-away-from-sql-to-json-ontology-detection/graphrag_sdk/kg.py
 
         Args:
             name (str): Knowledge graph name.
@@ -45,7 +44,7 @@ def __init__(
             username (str|None): FalkorDB username.
             password (str|None): FalkorDB password.
             model (str): OpenAI model to use for Falkor DB to build and retrieve from the graph.
-            schema: Falkor DB knowledge graph schema (ontology), https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/schema/schema.py
+            schema: Falkor DB knowledge graph schema (ontology), https://github.com/FalkorDB/GraphRAG-SDK/blob/2-move-away-from-sql-to-json-ontology-detection/graphrag_sdk/schema/schema.py
                     If None, Falkor DB will auto generate a schema from the input docs.
         """
         self.knowledge_graph = KnowledgeGraph(name, host, port, username, password, model, schema)

autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from autogen import Agent, ConversableAgent, UserProxyAgent
+
+from .falkor_graph_query_engine import FalkorGraphQueryEngine, FalkorGraphQueryResult
+from .graph_rag_capability import GraphRagCapability
+
+
+class FalkorGraphRagCapability(GraphRagCapability):
+    """
+    The Falkor graph rag capability integrate FalkorDB graphrag_sdk version: 0.1.3b0.
+    Ref: https://github.com/FalkorDB/GraphRAG-SDK/tree/2-move-away-from-sql-to-json-ontology-detection
+
+    For usage, please refer to example notebook/agentchat_graph_rag_falkordb.ipynb
+    """
+
+    def __init__(self, query_engine: FalkorGraphQueryEngine):
+        """
+        initialize graph rag capability with a graph query engine
+        """
+        self.query_engine = query_engine
+
+        # Graph DB query history.
+        self._history = []
+
+    def add_to_agent(self, agent: UserProxyAgent):
+        """
+        Add FalkorDB graph RAG capability to a UserProxyAgent.
+        The restriction to a UserProxyAgent to make sure the returned message does not contain information retrieved from the graph DB instead of any LLMs.
+        """
+        self.graph_rag_agent = agent
+
+        # Validate the agent config
+        if agent.llm_config not in (None, False):
+            raise Exception(
+                "Graph rag capability limits the query to graph DB, llm_config must be a dict or False or None."
+            )
+
+        # Register method to generate reply
+        agent.register_reply(
+            [ConversableAgent, None], self._reply_using_falkordb_query, position=0, remove_other_reply_funcs=True
+        )
+
+    def _reply_using_falkordb_query(
+        self,
+        recipient: ConversableAgent,
+        messages: Optional[List[Dict]] = None,
+        sender: Optional[Agent] = None,
+        config: Optional[Any] = None,
+    ) -> Tuple[bool, Union[str, Dict, None]]:
+        """
+        Query FalkorDB before return the message. Internally, it invokes the OpenAI assistant to generate a reply based on the given messages.
+        The history with FalkorDB is also logged and updated.
+
+        Args:
+            recipient: The agent instance that will receive the message.
+            messages: A list of messages in the conversation history with the sender.
+            sender: The agent instance that sent the message.
+            config: Optional configuration for message processing.
+
+        Returns:
+            A tuple containing a boolean indicating success and the assistant's reply.
+        """
+        question = self._get_last_question(messages[-1])
+        result: FalkorGraphQueryResult = self.query_engine.query(question, messages=self._history)
+        self._history = result.messages
+        return True, result.answer if result.answer else "I'm sorry, I don't have an answer for that."
+
+    def _get_last_question(self, message: Union[Dict, str]):
+        if isinstance(message, str):
+            return message
+        if isinstance(message, Dict):
+            if "content" in message:
+                return message["content"]
+        return None

notebook/agentchat_graph_rag_falkordb.ipynb

@@ -0,0 +1,242 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Using FalkorGraphRagCapability with UserProxyAgent for Graph RAG Question Answering\n",
+    "\n",
+    "AutoGen provides graph rag integration with Agent Capability. This is an example to integrate FalkorDB (a Knowledge Graph Database)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Install Falkor DB SDK"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2254.85s - pydevd: Sending message related to process being replaced timed-out after 5 seconds\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Defaulting to user installation because normal site-packages is not writeable\n",
+      "Requirement already satisfied: graphrag_sdk==0.1.3b0 in /home/autogen/.local/lib/python3.11/site-packages (0.1.3b0)\n",
+      "Requirement already satisfied: bs4<0.0.3,>=0.0.2 in /home/autogen/.local/lib/python3.11/site-packages (from graphrag_sdk==0.1.3b0) (0.0.2)\n",
+      "Requirement already satisfied: falkordb<2.0.0,>=1.0.4 in /home/autogen/.local/lib/python3.11/site-packages (from graphrag_sdk==0.1.3b0) (1.0.10)\n",
+      "Requirement already satisfied: openai<2.0.0,>=1.30.3 in /usr/local/lib/python3.11/site-packages (from graphrag_sdk==0.1.3b0) (1.55.0)\n",
+      "Requirement already satisfied: sqlalchemy<3.0.0,>=2.0.30 in /usr/local/lib/python3.11/site-packages (from graphrag_sdk==0.1.3b0) (2.0.36)\n",
+      "Requirement already satisfied: typing-extensions<5.0.0,>=4.12.1 in /usr/local/lib/python3.11/site-packages (from graphrag_sdk==0.1.3b0) (4.12.2)\n",
+      "Requirement already satisfied: beautifulsoup4 in /usr/local/lib/python3.11/site-packages (from bs4<0.0.3,>=0.0.2->graphrag_sdk==0.1.3b0) (4.12.3)\n",
+      "Requirement already satisfied: redis<6.0.0,>=5.0.1 in /home/autogen/.local/lib/python3.11/site-packages (from falkordb<2.0.0,>=1.0.4->graphrag_sdk==0.1.3b0) (5.2.0)\n",
+      "Requirement already satisfied: anyio<5,>=3.5.0 in /usr/local/lib/python3.11/site-packages (from openai<2.0.0,>=1.30.3->graphrag_sdk==0.1.3b0) (4.6.2.post1)\n",
+      "Requirement already satisfied: distro<2,>=1.7.0 in /usr/local/lib/python3.11/site-packages (from openai<2.0.0,>=1.30.3->graphrag_sdk==0.1.3b0) (1.9.0)\n",
+      "Requirement already satisfied: httpx<1,>=0.23.0 in /usr/local/lib/python3.11/site-packages (from openai<2.0.0,>=1.30.3->graphrag_sdk==0.1.3b0) (0.27.2)\n",
+      "Requirement already satisfied: jiter<1,>=0.4.0 in /usr/local/lib/python3.11/site-packages (from openai<2.0.0,>=1.30.3->graphrag_sdk==0.1.3b0) (0.7.1)\n",
+      "Requirement already satisfied: pydantic<3,>=1.9.0 in /usr/local/lib/python3.11/site-packages (from openai<2.0.0,>=1.30.3->graphrag_sdk==0.1.3b0) (1.10.9)\n",
+      "Requirement already satisfied: sniffio in /usr/local/lib/python3.11/site-packages (from openai<2.0.0,>=1.30.3->graphrag_sdk==0.1.3b0) (1.3.1)\n",
+      "Requirement already satisfied: tqdm>4 in /usr/local/lib/python3.11/site-packages (from openai<2.0.0,>=1.30.3->graphrag_sdk==0.1.3b0) (4.67.1)\n",
+      "Requirement already satisfied: greenlet!=0.4.17 in /usr/local/lib/python3.11/site-packages (from sqlalchemy<3.0.0,>=2.0.30->graphrag_sdk==0.1.3b0) (3.1.1)\n",
+      "Requirement already satisfied: idna>=2.8 in /usr/local/lib/python3.11/site-packages (from anyio<5,>=3.5.0->openai<2.0.0,>=1.30.3->graphrag_sdk==0.1.3b0) (3.10)\n",
+      "Requirement already satisfied: certifi in /usr/local/lib/python3.11/site-packages (from httpx<1,>=0.23.0->openai<2.0.0,>=1.30.3->graphrag_sdk==0.1.3b0) (2024.8.30)\n",
+      "Requirement already satisfied: httpcore==1.* in /usr/local/lib/python3.11/site-packages (from httpx<1,>=0.23.0->openai<2.0.0,>=1.30.3->graphrag_sdk==0.1.3b0) (1.0.7)\n",
+      "Requirement already satisfied: h11<0.15,>=0.13 in /usr/local/lib/python3.11/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->openai<2.0.0,>=1.30.3->graphrag_sdk==0.1.3b0) (0.14.0)\n",
+      "Requirement already satisfied: soupsieve>1.2 in /usr/local/lib/python3.11/site-packages (from beautifulsoup4->bs4<0.0.3,>=0.0.2->graphrag_sdk==0.1.3b0) (2.6)\n",
+      "Note: you may need to restart the kernel to use updated packages.\n"
+     ]
+    }
+   ],
+   "source": [
+    "%pip install graphrag_sdk==0.1.3b0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Set OpenAI API"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "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",
+    "os.environ[\"OPENAI_API_KEY\"] = config_list[0][\"api_key\"]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create Knowledge Graph with Your Own Data\n",
+    "\n",
+    "Note that, you would need to have a Falkor DB running ready. If you use docker, please set up docker network properly.\n",
+    "\n",
+    "In this example, Falker DB endpint is already set at host=\"172.18.0.2\" and port=6379.\n",
+    "For how to set up Falkor DB, please refer to https://docs.falkordb.com/"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from graphrag_sdk.schema import Schema\n",
+    "\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",
+    "\n",
+    "# Auto generate graph schema from unstructured data\n",
+    "input_path = \"../test/agentchat/contrib/graph_rag/the_matrix.txt\"\n",
+    "\n",
+    "movie_schema = Schema()\n",
+    "actor = movie_schema.add_entity(\"Actor\").add_attribute(\"name\", str, unique=True)\n",
+    "movie = movie_schema.add_entity(\"Movie\").add_attribute(\"title\", str, unique=True)\n",
+    "movie_schema.add_relation(\"ACTED\", actor, movie)\n",
+    "\n",
+    "query_engine = FalkorGraphQueryEngine(\n",
+    "    name=\"IMDB\",\n",
+    "    host=\"172.18.0.2\",\n",
+    "    port=6379,\n",
+    "    schema=movie_schema,\n",
+    ")\n",
+    "\n",
+    "input_documents = [Document(doctype=DocumentType.TEXT, path_or_url=input_path)]\n",
+    "\n",
+    "query_engine.init_db(input_doc=input_documents)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create a UserProxyAgent for FalkorDB and Answer Questions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[33muser_proxy\u001b[0m (to matrix_agent):\n",
+      "\n",
+      "Name a few actors who've played in 'The Matrix'\n",
+      "\n",
+      "--------------------------------------------------------------------------------\n",
+      "\u001b[33mmatrix_agent\u001b[0m (to user_proxy):\n",
+      "\n",
+      "The actors who've played in 'The Matrix' include:\n",
+      "\n",
+      "- Keanu Reeves\n",
+      "- Laurence Fishburne\n",
+      "- Carrie-Anne Moss\n",
+      "- Hugo Weaving\n",
+      "\n",
+      "--------------------------------------------------------------------------------\n",
+      "\u001b[33muser_proxy\u001b[0m (to matrix_agent):\n",
+      "\n",
+      "List additional actors.\n",
+      "\n",
+      "--------------------------------------------------------------------------------\n",
+      "\u001b[33mmatrix_agent\u001b[0m (to user_proxy):\n",
+      "\n",
+      "It appears that the actors I initially provided exhaust the list available in our current knowledge graph. There are no additional actors listed for 'The Matrix' beyond Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving. If the film had more actors, they are not present in our graph at this time.\n",
+      "\n",
+      "--------------------------------------------------------------------------------\n",
+      "\u001b[33muser_proxy\u001b[0m (to matrix_agent):\n",
+      "\n",
+      "Who is Lana Wachowski?\n",
+      "\n",
+      "--------------------------------------------------------------------------------\n",
+      "\u001b[33mmatrix_agent\u001b[0m (to user_proxy):\n",
+      "\n",
+      "I'm sorry, I don't have an answer for that.\n",
+      "\n",
+      "--------------------------------------------------------------------------------\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "ChatResult(chat_id=None, chat_history=[{'content': \"Name a few actors who've played in 'The Matrix'\", 'role': 'assistant', 'name': 'user_proxy'}, {'content': \"The actors who've played in 'The Matrix' include:\\n\\n- Keanu Reeves\\n- Laurence Fishburne\\n- Carrie-Anne Moss\\n- Hugo Weaving\", 'role': 'user', 'name': 'matrix_agent'}, {'content': 'List additional actors.', 'role': 'assistant', 'name': 'user_proxy'}, {'content': \"It appears that the actors I initially provided exhaust the list available in our current knowledge graph. There are no additional actors listed for 'The Matrix' beyond Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving. If the film had more actors, they are not present in our graph at this time.\", 'role': 'user', 'name': 'matrix_agent'}, {'content': 'Who is Lana Wachowski?', 'role': 'assistant', 'name': 'user_proxy'}, {'content': \"I'm sorry, I don't have an answer for that.\", 'role': 'user', 'name': 'matrix_agent'}], summary=\"I'm sorry, I don't have an answer for that.\", cost={'usage_including_cached_inference': {'total_cost': 0}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=['List additional actors.', 'Who is Lana Wachowski?', 'exit'])"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from autogen import ConversableAgent, UserProxyAgent\n",
+    "from autogen.agentchat.contrib.graph_rag.falkor_graph_rag_capability import FalkorGraphRagCapability\n",
+    "\n",
+    "graph_rag_agent = ConversableAgent(\n",
+    "    name=\"matrix_agent\",\n",
+    "    human_input_mode=\"NEVER\",\n",
+    ")\n",
+    "\n",
+    "graph_rag_capability = FalkorGraphRagCapability(query_engine)\n",
+    "graph_rag_capability.add_to_agent(graph_rag_agent)\n",
+    "\n",
+    "user_proxy = UserProxyAgent(\n",
+    "    name=\"user_proxy\",\n",
+    "    code_execution_config=False,\n",
+    "    is_termination_msg=lambda msg: \"TERMINATE\" in msg[\"content\"],\n",
+    "    human_input_mode=\"ALWAYS\",\n",
+    ")\n",
+    "\n",
+    "user_proxy.initiate_chat(graph_rag_agent, message=\"Name a few actors who've played in 'The Matrix'\")"
+   ]
+  }
+ ],
+ "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
+}

setup.py

@@ -61,7 +61,7 @@
 retrieve_chat_pgvector = [*retrieve_chat, "pgvector>=0.2.5"]
 
 graph_rag_falkor_db = [
-    "graphrag_sdk",
+    "graphrag_sdk==0.1.3b0",
 ]
 
 if current_os in ["Windows", "Darwin"]:

test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py

@@ -8,8 +8,7 @@
 
 import pytest
 from conftest import reason, skip_openai  # noqa: E402
-from graphrag_sdk import KnowledgeGraph, Source
-from graphrag_sdk.ontology import Ontology
+from graphrag_sdk.schema import Schema
 
 try:
     from autogen.agentchat.contrib.graph_rag.document import (
@@ -40,7 +39,7 @@ def test_falkor_db_query_engine():
     3. Query it with a question and verify the result contains the critical information.
     """
     # Arrange
-    test_schema = Ontology()
+    test_schema = Schema()
     actor = test_schema.add_entity("Actor").add_attribute("name", str, unique=True)
     movie = test_schema.add_entity("Movie").add_attribute("title", str, unique=True)
     test_schema.add_relation("ACTED", actor, movie)