diff --git a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py index 64453b5307..b2940c99a6 100644 --- a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py +++ b/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) diff --git a/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py b/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py new file mode 100644 index 0000000000..bf6ed5c9a0 --- /dev/null +++ b/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 diff --git a/notebook/agentchat_graph_rag_falkordb.ipynb b/notebook/agentchat_graph_rag_falkordb.ipynb new file mode 100644 index 0000000000..be72873b98 --- /dev/null +++ b/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 +} diff --git a/setup.py b/setup.py index 16abda3298..c2b73cc2a0 100644 --- a/setup.py +++ b/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"]: diff --git a/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py b/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py index 8c22edb846..6f54126e3e 100644 --- a/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py +++ b/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)