Merge pull request #111 from ag2ai/falkor_db_graph_rag

[GraphRAG] FalkorDB graph rag integration

.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

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=[])

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

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):

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"""
+        ...