diff --git a/.github/workflows/contrib-graph-rag-tests.yml b/.github/workflows/contrib-graph-rag-tests.yml
index 2df2d3c93a..080eea3b9f 100644
--- a/.github/workflows/contrib-graph-rag-tests.yml
+++ b/.github/workflows/contrib-graph-rag-tests.yml
@@ -43,7 +43,7 @@ jobs:
run: |
python -m pip install --upgrade pip wheel
pip install pytest
- - name: Install Falkor DB SDK when on linux
+ - name: Install FalkorDB SDK when on linux
run: |
pip install -e .[graph_rag_falkor_db]
- name: Set AUTOGEN_USE_DOCKER based on OS
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..114e6d8ed3 100644
--- a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py
+++ b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py
@@ -1,27 +1,24 @@
# 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
-from graphrag_sdk.schema import Schema
+from graphrag_sdk.model_config import KnowledgeGraphModelConfig
+from graphrag_sdk.models import GenerativeModel
+from graphrag_sdk.models.openai import OpenAiGenerativeModel
+from graphrag_sdk.ontology import Ontology
from .document import Document
from .graph_query_engine import GraphStoreQueryResult
-class FalkorGraphQueryResult(GraphStoreQueryResult):
- messages: list = field(default_factory=list)
-
-
class FalkorGraphQueryEngine:
"""
- This is a wrapper for Falkor DB KnowledgeGraph.
+ This is a wrapper for FalkorDB KnowledgeGraph.
"""
def __init__(
@@ -31,11 +28,11 @@ def __init__(
port: int = 6379,
username: str | None = None,
password: str | None = None,
- model: str = "gpt-4-1106-preview",
- schema: Schema | None = None,
+ model: GenerativeModel = OpenAiGenerativeModel("gpt-4o"),
+ ontology: Ontology | None = None,
):
"""
- Initialize a Falkor DB knowledge graph.
+ Initialize a FalkorDB knowledge graph.
Please also refer to https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/kg.py
Args:
@@ -44,11 +41,18 @@ def __init__(
port (int): FalkorDB port number.
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
- If None, Falkor DB will auto generate a schema from the input docs.
+ model (GenerativeModel): LLM model to use for FalkorDB to build and retrieve from the graph, default to use OAI gpt-4o.
+ ontology: FalkorDB knowledge graph schema/ontology, https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/ontology.py
+ If None, FalkorDB will auto generate an ontology from the input docs.
"""
- self.knowledge_graph = KnowledgeGraph(name, host, port, username, password, model, schema)
+ self.name = name
+ self.host = host
+ self.port = port
+ self.username = username
+ self.password = password
+ self.model = model
+ self.model_config = KnowledgeGraphModelConfig.with_model(model)
+ self.ontology = ontology
def init_db(self, input_doc: List[Document] | None):
"""
@@ -60,14 +64,33 @@ def init_db(self, input_doc: List[Document] | None):
sources.append(Source(doc.path_or_url))
if sources:
+ # Auto generate graph ontology if not created by user.
+ if self.ontology is None:
+ self.ontology = Ontology.from_sources(
+ sources=sources,
+ model=self.model,
+ )
+
+ self.knowledge_graph = KnowledgeGraph(
+ name=self.name,
+ host=self.host,
+ port=self.port,
+ username=self.username,
+ password=self.password,
+ model_config=KnowledgeGraphModelConfig.with_model(self.model),
+ ontology=self.ontology,
+ )
+
+ # Establish a chat session, this will maintain the history
+ self._chat_session = self.knowledge_graph.chat_session()
self.knowledge_graph.process_sources(sources)
def add_records(self, new_records: List) -> bool:
- raise NotImplementedError("This method is not supported by Falkor DB SDK yet.")
+ raise NotImplementedError("This method is not supported by FalkorDB SDK yet.")
- def query(self, question: str, n_results: int = 1, **kwargs) -> FalkorGraphQueryResult:
+ def query(self, question: str, n_results: int = 1, **kwargs) -> GraphStoreQueryResult:
"""
- Query the knowledage graph with a question and optional message history.
+ Query the knowledge graph with a question and optional message history.
Args:
question: a human input question.
@@ -77,6 +100,12 @@ def query(self, question: str, n_results: int = 1, **kwargs) -> FalkorGraphQuery
Returns: FalkorGraphQueryResult
"""
- messages = kwargs.pop("messages", [])
- answer, messages = self.knowledge_graph.ask(question, messages)
- return FalkorGraphQueryResult(answer=answer, results=[], messages=messages)
+ if self.knowledge_graph is None:
+ raise ValueError("Knowledge graph is not created.")
+
+ response = self._chat_session.send_message(question)
+
+ # History will be considered when querying by setting the last_answer
+ self._chat_session.last_answer = response["response"]
+
+ return GraphStoreQueryResult(answer=response["response"], results=[])
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..7e8bf9cc90
--- /dev/null
+++ b/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py
@@ -0,0 +1,81 @@
+# 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
+from .graph_query_engine import GraphStoreQueryResult
+from .graph_rag_capability import GraphRagCapability
+
+
+class FalkorGraphRagCapability(GraphRagCapability):
+ """
+ The FalkorDB GraphRAG capability integrate FalkorDB with 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 GraphRAG capability with a graph query engine
+ """
+ self.query_engine = query_engine
+
+ def add_to_agent(self, agent: UserProxyAgent):
+ """
+ Add FalkorDB GraphRAG 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(
+ "Agents with GraphRAG capabilities do not use an LLM configuration. Please set your llm_config to None or False."
+ )
+
+ # Register method to generate the reply using a FalkorDB query
+ # All other reply methods will be removed
+ 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 and return the message. Internally, it utilises OpenAI to generate a reply based on the given messages.
+ The history with FalkorDB is also logged and updated.
+
+ If no results are found, a default message is returned: "I'm sorry, I don't have an answer for that."
+
+ 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: GraphStoreQueryResult = self.query_engine.query(question)
+
+ 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]):
+ """Retrieves the last message from the conversation history."""
+ if isinstance(message, str):
+ return message
+ if isinstance(message, Dict):
+ if "content" in message:
+ return message["content"]
+ return None
diff --git a/autogen/agentchat/contrib/graph_rag/graph_query_engine.py b/autogen/agentchat/contrib/graph_rag/graph_query_engine.py
index 7f533fe98b..b15866f2db 100644
--- a/autogen/agentchat/contrib/graph_rag/graph_query_engine.py
+++ b/autogen/agentchat/contrib/graph_rag/graph_query_engine.py
@@ -26,7 +26,7 @@ class GraphStoreQueryResult:
class GraphQueryEngine(Protocol):
"""An abstract base class that represents a graph query engine on top of a underlying graph database.
- This interface defines the basic methods for graph rag.
+ This interface defines the basic methods for graph-based RAG.
"""
def init_db(self, input_doc: List[Document] | None = None):
diff --git a/autogen/agentchat/contrib/graph_rag/graph_rag_capability.py b/autogen/agentchat/contrib/graph_rag/graph_rag_capability.py
index 2819bcee8a..7e2ade403d 100644
--- a/autogen/agentchat/contrib/graph_rag/graph_rag_capability.py
+++ b/autogen/agentchat/contrib/graph_rag/graph_rag_capability.py
@@ -12,9 +12,9 @@
class GraphRagCapability(AgentCapability):
"""
- A graph rag capability uses a graph query engine to give a conversable agent the graph rag ability.
+ A graph-based RAG capability uses a graph query engine to give a conversable agent the graph-based RAG ability.
- An agent class with graph rag capability could
+ An agent class with graph-based RAG capability could
1. create a graph in the underlying database with input documents.
2. retrieved relevant information based on messages received by the agent.
3. generate answers from retrieved information and send messages back.
@@ -55,8 +55,10 @@ class GraphRagCapability(AgentCapability):
def __init__(self, query_engine: GraphQueryEngine):
"""
- initialize graph rag capability with a graph query engine
+ Initialize graph-based RAG capability with a graph query engine
"""
...
- def add_to_agent(self, agent: ConversableAgent): ...
+ def add_to_agent(self, agent: ConversableAgent):
+ """Add the capability to an agent"""
+ ...
diff --git a/notebook/agentchat_graph_rag_falkordb.ipynb b/notebook/agentchat_graph_rag_falkordb.ipynb
new file mode 100644
index 0000000000..9d85924569
--- /dev/null
+++ b/notebook/agentchat_graph_rag_falkordb.ipynb
@@ -0,0 +1,426 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Using FalkorGraphRagCapability with agents for GraphRAG Question & Answering\n",
+ "\n",
+ "AG2 provides GraphRAG integration using agent capabilities. This is an example to integrate FalkorDB (a Knowledge Graph database)."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Install FalkorDB SDK"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 6,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Defaulting to user installation because normal site-packages is not writeable\n",
+ "Requirement already satisfied: graphrag_sdk==0.3.3 in /home/autogen/.local/lib/python3.11/site-packages (0.3.3)\n",
+ "Requirement already satisfied: backoff<3.0.0,>=2.2.1 in /usr/local/lib/python3.11/site-packages (from graphrag_sdk==0.3.3) (2.2.1)\n",
+ "Requirement already satisfied: bs4<0.0.3,>=0.0.2 in /home/autogen/.local/lib/python3.11/site-packages (from graphrag_sdk==0.3.3) (0.0.2)\n",
+ "Requirement already satisfied: falkordb<2.0.0,>=1.0.9 in /home/autogen/.local/lib/python3.11/site-packages (from graphrag_sdk==0.3.3) (1.0.10)\n",
+ "Requirement already satisfied: fix-busted-json<0.0.19,>=0.0.18 in /home/autogen/.local/lib/python3.11/site-packages (from graphrag_sdk==0.3.3) (0.0.18)\n",
+ "Requirement already satisfied: ipykernel<7.0.0,>=6.29.5 in /usr/local/lib/python3.11/site-packages (from graphrag_sdk==0.3.3) (6.29.5)\n",
+ "Requirement already satisfied: pypdf<5.0.0,>=4.2.0 in /home/autogen/.local/lib/python3.11/site-packages (from graphrag_sdk==0.3.3) (4.3.1)\n",
+ "Requirement already satisfied: python-abc<0.3.0,>=0.2.0 in /home/autogen/.local/lib/python3.11/site-packages (from graphrag_sdk==0.3.3) (0.2.0)\n",
+ "Requirement already satisfied: python-dotenv<2.0.0,>=1.0.1 in /usr/local/lib/python3.11/site-packages (from graphrag_sdk==0.3.3) (1.0.1)\n",
+ "Requirement already satisfied: ratelimit<3.0.0,>=2.2.1 in /home/autogen/.local/lib/python3.11/site-packages (from graphrag_sdk==0.3.3) (2.2.1)\n",
+ "Requirement already satisfied: requests<3.0.0,>=2.32.3 in /usr/local/lib/python3.11/site-packages (from graphrag_sdk==0.3.3) (2.32.3)\n",
+ "Requirement already satisfied: typing-extensions<5.0.0,>=4.12.1 in /usr/local/lib/python3.11/site-packages (from graphrag_sdk==0.3.3) (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.3.3) (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.9->graphrag_sdk==0.3.3) (5.2.0)\n",
+ "Requirement already satisfied: comm>=0.1.1 in /usr/local/lib/python3.11/site-packages (from ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (0.2.2)\n",
+ "Requirement already satisfied: debugpy>=1.6.5 in /usr/local/lib/python3.11/site-packages (from ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (1.8.9)\n",
+ "Requirement already satisfied: ipython>=7.23.1 in /usr/local/lib/python3.11/site-packages (from ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (8.29.0)\n",
+ "Requirement already satisfied: jupyter-client>=6.1.12 in /usr/local/lib/python3.11/site-packages (from ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (8.6.3)\n",
+ "Requirement already satisfied: jupyter-core!=5.0.*,>=4.12 in /usr/local/lib/python3.11/site-packages (from ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (5.7.2)\n",
+ "Requirement already satisfied: matplotlib-inline>=0.1 in /usr/local/lib/python3.11/site-packages (from ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (0.1.7)\n",
+ "Requirement already satisfied: nest-asyncio in /usr/local/lib/python3.11/site-packages (from ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (1.6.0)\n",
+ "Requirement already satisfied: packaging in /usr/local/lib/python3.11/site-packages (from ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (24.2)\n",
+ "Requirement already satisfied: psutil in /usr/local/lib/python3.11/site-packages (from ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (6.1.0)\n",
+ "Requirement already satisfied: pyzmq>=24 in /usr/local/lib/python3.11/site-packages (from ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (26.2.0)\n",
+ "Requirement already satisfied: tornado>=6.1 in /usr/local/lib/python3.11/site-packages (from ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (6.4.2)\n",
+ "Requirement already satisfied: traitlets>=5.4.0 in /usr/local/lib/python3.11/site-packages (from ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (5.14.3)\n",
+ "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/site-packages (from requests<3.0.0,>=2.32.3->graphrag_sdk==0.3.3) (3.4.0)\n",
+ "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.11/site-packages (from requests<3.0.0,>=2.32.3->graphrag_sdk==0.3.3) (3.10)\n",
+ "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/site-packages (from requests<3.0.0,>=2.32.3->graphrag_sdk==0.3.3) (2.2.3)\n",
+ "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.11/site-packages (from requests<3.0.0,>=2.32.3->graphrag_sdk==0.3.3) (2024.8.30)\n",
+ "Requirement already satisfied: decorator in /usr/local/lib/python3.11/site-packages (from ipython>=7.23.1->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (5.1.1)\n",
+ "Requirement already satisfied: jedi>=0.16 in /usr/local/lib/python3.11/site-packages (from ipython>=7.23.1->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (0.19.2)\n",
+ "Requirement already satisfied: prompt-toolkit<3.1.0,>=3.0.41 in /usr/local/lib/python3.11/site-packages (from ipython>=7.23.1->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (3.0.48)\n",
+ "Requirement already satisfied: pygments>=2.4.0 in /usr/local/lib/python3.11/site-packages (from ipython>=7.23.1->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (2.18.0)\n",
+ "Requirement already satisfied: stack-data in /usr/local/lib/python3.11/site-packages (from ipython>=7.23.1->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (0.6.3)\n",
+ "Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.11/site-packages (from ipython>=7.23.1->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (4.9.0)\n",
+ "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.11/site-packages (from jupyter-client>=6.1.12->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (2.9.0.post0)\n",
+ "Requirement already satisfied: platformdirs>=2.5 in /usr/local/lib/python3.11/site-packages (from jupyter-core!=5.0.*,>=4.12->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (4.3.6)\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.3.3) (2.6)\n",
+ "Requirement already satisfied: parso<0.9.0,>=0.8.4 in /usr/local/lib/python3.11/site-packages (from jedi>=0.16->ipython>=7.23.1->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (0.8.4)\n",
+ "Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.11/site-packages (from pexpect>4.3->ipython>=7.23.1->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (0.7.0)\n",
+ "Requirement already satisfied: wcwidth in /usr/local/lib/python3.11/site-packages (from prompt-toolkit<3.1.0,>=3.0.41->ipython>=7.23.1->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (0.2.13)\n",
+ "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.11/site-packages (from python-dateutil>=2.8.2->jupyter-client>=6.1.12->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (1.16.0)\n",
+ "Requirement already satisfied: executing>=1.2.0 in /usr/local/lib/python3.11/site-packages (from stack-data->ipython>=7.23.1->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (2.1.0)\n",
+ "Requirement already satisfied: asttokens>=2.1.0 in /usr/local/lib/python3.11/site-packages (from stack-data->ipython>=7.23.1->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (2.4.1)\n",
+ "Requirement already satisfied: pure-eval in /usr/local/lib/python3.11/site-packages (from stack-data->ipython>=7.23.1->ipykernel<7.0.0,>=6.29.5->graphrag_sdk==0.3.3) (0.2.3)\n",
+ "Note: you may need to restart the kernel to use updated packages.\n",
+ "Defaulting to user installation because normal site-packages is not writeable\n",
+ "Obtaining file:///workspaces/ag2\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 ag2==0.4) (1.55.0)\n",
+ "Requirement already satisfied: diskcache in /usr/local/lib/python3.11/site-packages (from ag2==0.4) (5.6.3)\n",
+ "Requirement already satisfied: termcolor in /usr/local/lib/python3.11/site-packages (from ag2==0.4) (2.5.0)\n",
+ "Requirement already satisfied: flaml in /usr/local/lib/python3.11/site-packages (from ag2==0.4) (2.3.2)\n",
+ "Requirement already satisfied: numpy<2.0.0,>=1.24.0 in /usr/local/lib/python3.11/site-packages (from ag2==0.4) (1.26.4)\n",
+ "Requirement already satisfied: python-dotenv in /usr/local/lib/python3.11/site-packages (from ag2==0.4) (1.0.1)\n",
+ "Requirement already satisfied: tiktoken in /usr/local/lib/python3.11/site-packages (from ag2==0.4) (0.8.0)\n",
+ "Requirement already satisfied: pydantic!=2.6.0,<3,>=1.10 in /usr/local/lib/python3.11/site-packages (from ag2==0.4) (1.10.9)\n",
+ "Requirement already satisfied: docker in /usr/local/lib/python3.11/site-packages (from ag2==0.4) (7.1.0)\n",
+ "Requirement already satisfied: packaging in /usr/local/lib/python3.11/site-packages (from ag2==0.4) (24.2)\n",
+ "Requirement already satisfied: anyio<5,>=3.5.0 in /usr/local/lib/python3.11/site-packages (from openai>=1.3->ag2==0.4) (4.6.2.post1)\n",
+ "Requirement already satisfied: distro<2,>=1.7.0 in /usr/local/lib/python3.11/site-packages (from openai>=1.3->ag2==0.4) (1.9.0)\n",
+ "Requirement already satisfied: httpx<1,>=0.23.0 in /usr/local/lib/python3.11/site-packages (from openai>=1.3->ag2==0.4) (0.27.2)\n",
+ "Requirement already satisfied: jiter<1,>=0.4.0 in /usr/local/lib/python3.11/site-packages (from openai>=1.3->ag2==0.4) (0.7.1)\n",
+ "Requirement already satisfied: sniffio in /usr/local/lib/python3.11/site-packages (from openai>=1.3->ag2==0.4) (1.3.1)\n",
+ "Requirement already satisfied: tqdm>4 in /usr/local/lib/python3.11/site-packages (from openai>=1.3->ag2==0.4) (4.67.1)\n",
+ "Requirement already satisfied: typing-extensions<5,>=4.11 in /usr/local/lib/python3.11/site-packages (from openai>=1.3->ag2==0.4) (4.12.2)\n",
+ "Requirement already satisfied: requests>=2.26.0 in /usr/local/lib/python3.11/site-packages (from docker->ag2==0.4) (2.32.3)\n",
+ "Requirement already satisfied: urllib3>=1.26.0 in /usr/local/lib/python3.11/site-packages (from docker->ag2==0.4) (2.2.3)\n",
+ "Requirement already satisfied: regex>=2022.1.18 in /usr/local/lib/python3.11/site-packages (from tiktoken->ag2==0.4) (2024.11.6)\n",
+ "Requirement already satisfied: idna>=2.8 in /usr/local/lib/python3.11/site-packages (from anyio<5,>=3.5.0->openai>=1.3->ag2==0.4) (3.10)\n",
+ "Requirement already satisfied: certifi in /usr/local/lib/python3.11/site-packages (from httpx<1,>=0.23.0->openai>=1.3->ag2==0.4) (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->ag2==0.4) (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>=1.3->ag2==0.4) (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->ag2==0.4) (3.4.0)\n",
+ "Building wheels for collected packages: ag2\n",
+ " Building editable for ag2 (pyproject.toml) ... \u001b[?25ldone\n",
+ "\u001b[?25h Created wheel for ag2: filename=ag2-0.4-0.editable-py3-none-any.whl size=15779 sha256=6352d50605c2b333aabeb86bee921ade8e34d3429ed67f175db6d0d23dcacd05\n",
+ " Stored in directory: /tmp/pip-ephem-wheel-cache-62aq_s_d/wheels/81/96/a6/c21a93c37cecb7b27b25b264f04fa27be94acd8d0f5e637527\n",
+ "Successfully built ag2\n",
+ "Installing collected packages: ag2\n",
+ " Attempting uninstall: ag2\n",
+ " Found existing installation: ag2 0.4\n",
+ " Uninstalling ag2-0.4:\n",
+ " Successfully uninstalled ag2-0.4\n",
+ "Successfully installed ag2-0.4\n",
+ "Note: you may need to restart the kernel to use updated packages.\n"
+ ]
+ }
+ ],
+ "source": [
+ "%pip install graphrag_sdk==0.3.3\n",
+ "%pip install -e /workspaces/ag2"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Set Configuration and OpenAI API Key\n",
+ "\n",
+ "In order to use FalkorDB you need to have an OpenAI key in your environment variable `OPENAI_API_KEY`.\n",
+ "\n",
+ "You can utilise an OAI_CONFIG_LIST file and extract the OpenAI API key and put it in the environment, as will be shown in the following cell.\n",
+ "\n",
+ "Alternatively, you can load the environment variable yourself.\n",
+ "\n",
+ "````{=mdx}\n",
+ ":::tip\n",
+ "Learn more about configuring LLMs for agents [here](/docs/topics/llm_configuration).\n",
+ ":::\n",
+ "````"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "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": "markdown",
+ "metadata": {},
+ "source": [
+ "````{=mdx}\n",
+ ":::important\n",
+ "The default model for loading graph data and answering questions using FalkorDB's SDK is OpenAI's GPT 4o and this can be changed by setting the `model` parameter on the FalkorGraphQueryEngine.\n",
+ ":::\n",
+ "````"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Create a Knowledge Graph with Your Own Data\n",
+ "\n",
+ "**Note:** You need to have a FalkorDB graph database running. If you are running one in a Docker container, please ensure your Docker network is setup to allow access to it.\n",
+ "\n",
+ "In this example, the FalkorDB endpoint is set to host=\"172.18.0.3\" and port=6379, please adjust accordingly. For how to set up FalkorDB, please refer to https://docs.falkordb.com/\n",
+ "\n",
+ "\n",
+ "Below, we have some sample data from IMDB on the movie 'The Matrix'.\n",
+ "\n",
+ "We then initialise the database with that text document, creating the graph in FalkorDB."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### A Simple Example\n",
+ "In this example, the graph ontology is auto-generated."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "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",
+ "Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving all played in 'The Matrix'.\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': \"Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving all played in 'The Matrix'.\", 'role': 'user', 'name': 'matrix_agent'}], summary=\"Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving all played in 'The Matrix'.\", cost={'usage_including_cached_inference': {'total_cost': 0}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=['exit'])"
+ ]
+ },
+ "execution_count": 4,
+ "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_path = \"../test/agentchat/contrib/graph_rag/the_matrix.txt\"\n",
+ "input_documents = [Document(doctype=DocumentType.TEXT, path_or_url=input_path)]\n",
+ "\n",
+ "# Create FalkorGraphQueryEngine\n",
+ "query_engine = FalkorGraphQueryEngine(\n",
+ " name=\"The_Matrix_Auto\",\n",
+ " host=\"172.18.0.3\", # Change\n",
+ " port=6379, # if needed\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=\"matrix_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",
+ ")\n",
+ "\n",
+ "user_proxy.initiate_chat(graph_rag_agent, message=\"Name a few actors who've played in 'The Matrix'\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Revisit the exmple with self defined graph ontology.\n",
+ "We define the entities and attributes that we want in the graph database (`Actor` with `name` and `Movie` with `title`) and the relationships (Actors `ACTED` in Movies). This allows your RAG agent to answer questions about actors in the movie."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from graphrag_sdk import Attribute, AttributeType, Entity, Ontology, Relation\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",
+ "input_path = \"../test/agentchat/contrib/graph_rag/the_matrix.txt\"\n",
+ "\n",
+ "movie_ontology = Ontology()\n",
+ "movie_ontology.add_entity(\n",
+ " Entity(label=\"Actor\", attributes=[Attribute(name=\"name\", attr_type=AttributeType.STRING, unique=True)])\n",
+ ")\n",
+ "movie_ontology.add_entity(\n",
+ " Entity(label=\"Movie\", attributes=[Attribute(name=\"title\", attr_type=AttributeType.STRING, unique=True)])\n",
+ ")\n",
+ "movie_ontology.add_relation(Relation(label=\"ACTED\", source=\"Actor\", target=\"Movie\"))\n",
+ "\n",
+ "query_engine = FalkorGraphQueryEngine(\n",
+ " name=\"IMDB\",\n",
+ " host=\"172.18.0.3\", # Change\n",
+ " port=6379, # if needed\n",
+ " ontology=movie_ontology,\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": [
+ "## Add capability to a ConversableAgent and query them\n",
+ "\n",
+ "With FalkorDB setup, we can now associate the GraphRAG capability with a ConversableAgent and have a chat with it."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\u001b[33muser_proxy\u001b[0m (to matrix_agent):\n",
+ "\n"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Name a few actors who've played in 'The Matrix'\n",
+ "\n",
+ "--------------------------------------------------------------------------------\n",
+ "\u001b[33mmatrix_agent\u001b[0m (to user_proxy):\n",
+ "\n",
+ "Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving have played in 'The Matrix'.\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",
+ "The actors Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving are listed. There are no additional actors provided in the context related to 'The Matrix'.\n",
+ "\n",
+ "--------------------------------------------------------------------------------\n",
+ "\u001b[33muser_proxy\u001b[0m (to matrix_agent):\n",
+ "\n",
+ "List additional people associated with the movie.\n",
+ "\n",
+ "--------------------------------------------------------------------------------\n",
+ "\u001b[33mmatrix_agent\u001b[0m (to user_proxy):\n",
+ "\n",
+ "I am sorry, I could not find the answer to your question\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': \"Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving have played in 'The Matrix'.\", 'role': 'user', 'name': 'matrix_agent'}, {'content': 'List additional actors.', 'role': 'assistant', 'name': 'user_proxy'}, {'content': \"The actors Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving are listed. There are no additional actors provided in the context related to 'The Matrix'.\", 'role': 'user', 'name': 'matrix_agent'}, {'content': 'List additional people associated with the movie.', 'role': 'assistant', 'name': 'user_proxy'}, {'content': 'I am sorry, I could not find the answer to your question', 'role': 'user', 'name': 'matrix_agent'}], summary='I am sorry, I could not find the answer to your question', cost={'usage_including_cached_inference': {'total_cost': 0}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=['List additional actors.', 'List additional people associated with the movie.', 'exit'])"
+ ]
+ },
+ "execution_count": 3,
+ "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",
+ "# Create a ConversableAgent (no LLM configuration)\n",
+ "graph_rag_agent = ConversableAgent(\n",
+ " name=\"matrix_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",
+ " 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'\")\n",
+ "\n",
+ "# You will be prompted, as a human in the loop, after the response - feel free to ask more questions."
+ ]
+ }
+ ],
+ "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 e551daa664..79336ab8a5 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.3.3",
]
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..0e23414b43 100644
--- a/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py
+++ b/test/agentchat/contrib/graph_rag/test_falkor_graph_rag.py
@@ -8,14 +8,10 @@
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 import Attribute, AttributeType, Entity, Ontology, Relation
try:
- from autogen.agentchat.contrib.graph_rag.document import (
- Document,
- DocumentType,
- )
+ from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType
from autogen.agentchat.contrib.graph_rag.falkor_graph_query_engine import (
FalkorGraphQueryEngine,
GraphStoreQueryResult,
@@ -34,18 +30,27 @@
)
def test_falkor_db_query_engine():
"""
- Test Falkor DB Query Engine.
- 1. create a test Falkor DB Query Engine with a schema.
+ Test FalkorDB Query Engine.
+ 1. create a test FalkorDB Query Engine with a schema.
2. Initialize it with an input txt file.
3. Query it with a question and verify the result contains the critical information.
"""
# Arrange
- test_schema = Ontology()
- 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)
+ movie_ontology = Ontology()
+ movie_ontology.add_entity(
+ Entity(label="Actor", attributes=[Attribute(name="name", attr_type=AttributeType.STRING, unique=True)])
+ )
+ movie_ontology.add_entity(
+ Entity(label="Movie", attributes=[Attribute(name="title", attr_type=AttributeType.STRING, unique=True)])
+ )
+ movie_ontology.add_relation(Relation(label="ACTED", source="Actor", target="Movie"))
- query_engine = FalkorGraphQueryEngine(schema=test_schema)
+ query_engine = FalkorGraphQueryEngine(
+ name="IMDB",
+ # host="192.168.0.115", # Change
+ # port=6379, # if needed
+ ontology=movie_ontology,
+ )
source_file = "test/agentchat/contrib/graph_rag/the_matrix.txt"
input_docs = [Document(doctype=DocumentType.TEXT, path_or_url=source_file)]
@@ -59,9 +64,3 @@ def test_falkor_db_query_engine():
# Assert
assert query_result.answer.find("Keanu Reeves") >= 0
- for message in query_result.messages:
- if isinstance(message, dict) and "role" in message and message["role"] == "user":
- assert "content" in message
- assert message["content"] is question
- return
- pytest.fail("Question not found in message history.")