[GraphRAG] FalkorDB graph rag integration

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)

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