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 0bb1f51a05b3..d0670de102be 100644 --- a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py +++ b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py @@ -1,5 +1,5 @@ import os -from dataclasses import field +from dataclasses import dataclass, field from typing import List from graphrag_sdk import KnowledgeGraph, Source @@ -9,6 +9,7 @@ from .graph_query_engine import GraphStoreQueryResult +@dataclass class FalkorGraphQueryResult(GraphStoreQueryResult): messages: list = field(default_factory=list) @@ -30,7 +31,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. @@ -39,7 +40,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 000000000000..1860047500bc --- /dev/null +++ b/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py @@ -0,0 +1,63 @@ +from typing import Dict, Union + +from autogen import 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 a hook for processing the last message. + agent.register_hook(hookable_method="process_last_received_message", hook=self.process_last_received_message) + + # Append extra info to the system message. + agent.update_system_message( + agent.system_message + "\nYou've been given the special ability to use graph rag to retrieve information." + ) + + def process_last_received_message(self, message: Union[Dict, str]): + """ + Query FalkorDB before return the message. + The history with FalkorDB is also logged and updated. + """ + question = self._get_last_question(message) + result: FalkorGraphQueryResult = self.query_engine.query(question, self._history) + self._history = result.messages + return result.answer + + 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 000000000000..ffbab4cd0314 --- /dev/null +++ b/notebook/agentchat_graph_rag_falkordb.ipynb @@ -0,0 +1,290 @@ +{ + "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": 17, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "1222214.22s - 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-ai/.local/lib/python3.11/site-packages (0.1.3b0)\n", + "Requirement already satisfied: bs4<0.0.3,>=0.0.2 in /home/autogen-ai/.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-ai/.local/lib/python3.11/site-packages (from graphrag_sdk==0.1.3b0) (1.0.8)\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.49.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.35)\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-ai/.local/lib/python3.11/site-packages (from falkordb<2.0.0,>=1.0.4->graphrag_sdk==0.1.3b0) (5.1.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.0)\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.5.0)\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.66.5)\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.5)\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", + "\n", + "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.3.1\u001b[0m\n", + "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n", + "Note: you may need to restart the kernel to use updated packages.\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "1222219.78s - 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", + "Obtaining file:///workspaces/autogen\n", + " Installing build dependencies ... \u001b[?25ldone\n", + "\u001b[?25h Checking if build backend supports build_editable ... \u001b[?25ldone\n", + "\u001b[?25h Getting requirements to build editable ... \u001b[?25ldone\n", + "\u001b[?25h Preparing editable metadata (pyproject.toml) ... \u001b[?25ldone\n", + "\u001b[?25hRequirement already satisfied: openai>=1.3 in /usr/local/lib/python3.11/site-packages (from autogen==0.3.2) (1.49.0)\n", + "Requirement already satisfied: diskcache in /usr/local/lib/python3.11/site-packages (from autogen==0.3.2) (5.6.3)\n", + "Requirement already satisfied: termcolor in /usr/local/lib/python3.11/site-packages (from autogen==0.3.2) (2.4.0)\n", + "Requirement already satisfied: flaml in /home/autogen-ai/.local/lib/python3.11/site-packages (from autogen==0.3.2) (2.2.0)\n", + "Requirement already satisfied: numpy<2,>=1.17.0 in /usr/local/lib/python3.11/site-packages (from autogen==0.3.2) (1.26.4)\n", + "Requirement already satisfied: python-dotenv in /usr/local/lib/python3.11/site-packages (from autogen==0.3.2) (1.0.1)\n", + "Requirement already satisfied: tiktoken in /usr/local/lib/python3.11/site-packages (from autogen==0.3.2) (0.7.0)\n", + "Requirement already satisfied: pydantic!=2.6.0,<3,>=1.10 in /usr/local/lib/python3.11/site-packages (from autogen==0.3.2) (1.10.9)\n", + "Requirement already satisfied: docker in /usr/local/lib/python3.11/site-packages (from autogen==0.3.2) (7.1.0)\n", + "Requirement already satisfied: packaging in /usr/local/lib/python3.11/site-packages (from autogen==0.3.2) (24.1)\n", + "Requirement already satisfied: anyio<5,>=3.5.0 in /usr/local/lib/python3.11/site-packages (from openai>=1.3->autogen==0.3.2) (4.6.0)\n", + "Requirement already satisfied: distro<2,>=1.7.0 in /usr/local/lib/python3.11/site-packages (from openai>=1.3->autogen==0.3.2) (1.9.0)\n", + "Requirement already satisfied: httpx<1,>=0.23.0 in /usr/local/lib/python3.11/site-packages (from openai>=1.3->autogen==0.3.2) (0.27.2)\n", + "Requirement already satisfied: jiter<1,>=0.4.0 in /usr/local/lib/python3.11/site-packages (from openai>=1.3->autogen==0.3.2) (0.5.0)\n", + "Requirement already satisfied: sniffio in /usr/local/lib/python3.11/site-packages (from openai>=1.3->autogen==0.3.2) (1.3.1)\n", + "Requirement already satisfied: tqdm>4 in /usr/local/lib/python3.11/site-packages (from openai>=1.3->autogen==0.3.2) (4.66.5)\n", + "Requirement already satisfied: typing-extensions<5,>=4.11 in /usr/local/lib/python3.11/site-packages (from openai>=1.3->autogen==0.3.2) (4.12.2)\n", + "Requirement already satisfied: requests>=2.26.0 in /usr/local/lib/python3.11/site-packages (from docker->autogen==0.3.2) (2.32.3)\n", + "Requirement already satisfied: urllib3>=1.26.0 in /usr/local/lib/python3.11/site-packages (from docker->autogen==0.3.2) (2.2.3)\n", + "Requirement already satisfied: regex>=2022.1.18 in /usr/local/lib/python3.11/site-packages (from tiktoken->autogen==0.3.2) (2024.9.11)\n", + "Requirement already satisfied: idna>=2.8 in /usr/local/lib/python3.11/site-packages (from anyio<5,>=3.5.0->openai>=1.3->autogen==0.3.2) (3.10)\n", + "Requirement already satisfied: certifi in /usr/local/lib/python3.11/site-packages (from httpx<1,>=0.23.0->openai>=1.3->autogen==0.3.2) (2024.8.30)\n", + "Requirement already satisfied: httpcore==1.* in /usr/local/lib/python3.11/site-packages (from httpx<1,>=0.23.0->openai>=1.3->autogen==0.3.2) (1.0.5)\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>=1.3->autogen==0.3.2) (0.14.0)\n", + "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/site-packages (from requests>=2.26.0->docker->autogen==0.3.2) (3.3.2)\n", + "Building wheels for collected packages: autogen\n", + " Building editable for autogen (pyproject.toml) ... \u001b[?25ldone\n", + "\u001b[?25h Created wheel for autogen: filename=autogen-0.3.2-0.editable-py3-none-any.whl size=16216 sha256=6a1d2928c2581b5d66f17c8cafed5444dab21fcb5a52b4ba34fe43abbc4055ec\n", + " Stored in directory: /tmp/pip-ephem-wheel-cache-9wh6yuuq/wheels/d4/86/1f/a538740a449c67aa3ff8343698c29d70fc94236f70bde92144\n", + "Successfully built autogen\n", + "Installing collected packages: autogen\n", + " Attempting uninstall: autogen\n", + " Found existing installation: autogen 0.3.0\n", + " Uninstalling autogen-0.3.0:\n", + " Successfully uninstalled autogen-0.3.0\n", + "Successfully installed autogen-0.3.2\n", + "\n", + "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.3.1\u001b[0m\n", + "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n", + "Note: you may need to restart the kernel to use updated packages.\n" + ] + } + ], + "source": [ + "%pip install graphrag_sdk==0.1.3b0\n", + "\n", + "# For debug only,\n", + "# %pip install -e /workspaces/autogen" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Set OpenAI API" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "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. \n", + "In this example, Falker DB endpint is already set at host=\"172.18.0.3\" and port=6379.\n", + "For how to set up Falkor DB, please refer to https://docs.falkordb.com/" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "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.3\",\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": 18, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\u001b[33muser_proxy\u001b[0m (to user_proxy):\n", + "\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Name a few actors who've played in 'The Matrix'\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[31m\n", + ">>>>>>>> NO HUMAN INPUT RECEIVED.\u001b[0m\n", + "\u001b[31m\n", + ">>>>>>>> USING AUTO REPLY...\u001b[0m\n", + "\u001b[33muser_proxy\u001b[0m (to user_proxy):\n", + "\n", + "\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33muser_proxy\u001b[0m (to user_proxy):\n", + "\n", + "List additional actors\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': \"A few actors who have played in 'The Matrix' are:\\n\\n- Keanu Reeves\\n- Laurence Fishburne\\n- Carrie-Anne Moss\\n- Hugo Weaving\", 'role': 'user', 'name': 'user_proxy'}, {'content': '', 'role': 'assistant', 'name': 'user_proxy'}, {'content': None, 'role': 'user', 'name': 'user_proxy'}, {'content': 'List additional actors', 'role': 'assistant', 'name': 'user_proxy'}, {'content': 'The additional actors in the knowledge graph, ordered alphabetically by name, are:\\n\\n- Carrie-Anne Moss\\n- Hugo Weaving\\n- Keanu Reeves\\n- Lana Wachowski\\n- Laurence Fishburne\\n- Lilly Wachowski', 'role': 'user', 'name': 'user_proxy'}], summary='The additional actors in the knowledge graph, ordered alphabetically by name, are:\\n\\n- Carrie-Anne Moss\\n- Hugo Weaving\\n- Keanu Reeves\\n- Lana Wachowski\\n- Laurence Fishburne\\n- Lilly Wachowski', cost={'usage_including_cached_inference': {'total_cost': 0}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=['', 'List additional actors', 'exit'])" + ] + }, + "execution_count": 18, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from autogen import UserProxyAgent\n", + "from autogen.agentchat.contrib.graph_rag.falkor_graph_rag_capability import FalkorGraphRagCapability\n", + "\n", + "graph_rag_agent = 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", + "graph_rag_capability = FalkorGraphRagCapability(query_engine)\n", + "graph_rag_capability.add_to_agent(graph_rag_agent)\n", + "\n", + "graph_rag_agent.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 64004e772435..e22974b0f46e 100644 --- a/setup.py +++ b/setup.py @@ -60,7 +60,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"]: