Merge pull request #141 from ag2ai/graphrag_demo

GraphRAG demo with SwarmAgents

.pre-commit-config.yaml

@@ -47,7 +47,8 @@ repos:
               website/docs/tutorial/code-executors.ipynb |
               website/docs/topics/code-execution/custom-executor.ipynb |
               website/docs/topics/non-openai-models/cloud-gemini.ipynb |
-              notebook/.*
+              notebook/.* |
+              test/agentchat/contrib/graph_rag/trip_planner_data/.*
             )$
   # See https://jaredkhan.com/blog/mypy-pre-commit
   - repo: local

autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py

@@ -3,9 +3,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import os
-from dataclasses import dataclass, field
+import warnings
 from typing import List
 
+from falkordb import FalkorDB, Graph
 from graphrag_sdk import KnowledgeGraph, Source
 from graphrag_sdk.model_config import KnowledgeGraphModelConfig
 from graphrag_sdk.models import GenerativeModel
@@ -35,6 +36,8 @@ def __init__(
         Initialize a FalkorDB knowledge graph.
         Please also refer to https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/kg.py
 
+        TODO: Fix LLM API cost calculation for FalkorDB useages.
+
         Args:
             name (str): Knowledge graph name.
             host (str): FalkorDB hostname.
@@ -53,8 +56,38 @@ def __init__(
         self.model = model
         self.model_config = KnowledgeGraphModelConfig.with_model(model)
         self.ontology = ontology
+        self.knowledge_graph = None
+        self.falkordb = FalkorDB(host=self.host, port=self.port, username=self.username, password=self.password)
+
+    def connect_db(self):
+        """
+        Connect to an existing knowledge graph.
+        """
+        if self.name in self.falkordb.list_graphs():
+            try:
+                self.ontology = self._load_ontology_from_db(self.name)
+            except Exception:
+                warnings.warn("Graph Ontology is not loaded.")
+
+            if self.ontology is None:
+                raise ValueError(f"Ontology of the knowledge graph '{self.name}' can't be None.")
+
+            self.knowledge_graph = KnowledgeGraph(
+                name=self.name,
+                host=self.host,
+                port=self.port,
+                username=self.username,
+                password=self.password,
+                model_config=self.model_config,
+                ontology=self.ontology,
+            )
+
+            # Establishing a chat session will maintain the history
+            self._chat_session = self.knowledge_graph.chat_session()
+        else:
+            raise ValueError(f"Knowledge graph '{self.name}' does not exist")
 
-    def init_db(self, input_doc: List[Document] | None):
+    def init_db(self, input_doc: List[Document]):
         """
         Build the knowledge graph with input documents.
         """
@@ -81,10 +114,16 @@ def init_db(self, input_doc: List[Document] | None):
                 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)
 
+            # Establishing a chat session will maintain the history
+            self._chat_session = self.knowledge_graph.chat_session()
+
+            # Save Ontology to graph for future access.
+            self._save_ontology_to_db(self.name, self.ontology)
+        else:
+            raise ValueError("No input documents could be loaded.")
+
     def add_records(self, new_records: List) -> bool:
         raise NotImplementedError("This method is not supported by FalkorDB SDK yet.")
 
@@ -101,11 +140,26 @@ def query(self, question: str, n_results: int = 1, **kwargs) -> GraphStoreQueryR
         Returns: FalkorGraphQueryResult
         """
         if self.knowledge_graph is None:
-            raise ValueError("Knowledge graph is not created.")
+            raise ValueError("Knowledge graph has not been selected or 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=[])
+
+    def __get_ontology_storage_graph(self, graph_name: str) -> Graph:
+        ontology_table_name = graph_name + "_ontology"
+        return self.falkordb.select_graph(ontology_table_name)
+
+    def _save_ontology_to_db(self, graph_name: str, ontology: Ontology):
+        """
+        Save graph ontology to a separate table with {graph_name}_ontology
+        """
+        graph = self.__get_ontology_storage_graph(graph_name)
+        ontology.save_to_graph(graph)
+
+    def _load_ontology_from_db(self, graph_name: str) -> Ontology:
+        graph = self.__get_ontology_storage_graph(graph_name)
+        return Ontology.from_graph(graph)

autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py

@@ -55,6 +55,8 @@ def _reply_using_falkordb_query(
         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.
 
+        The agent's system message will be incorporated into the query, if it's not blank.
+
         If no results are found, a default message is returned: "I'm sorry, I don't have an answer for that."
 
         Args:
@@ -66,16 +68,38 @@ def _reply_using_falkordb_query(
         Returns:
             A tuple containing a boolean indicating success and the assistant's reply.
         """
-        question = self._get_last_question(messages[-1])
+        # question = self._get_last_question(messages[-1])
+        question = self._messages_summary(messages, recipient.system_message)
         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
+    def _messages_summary(self, messages: Union[Dict, str], system_message: str) -> str:
+        """Summarize the messages in the conversation history. Excluding any message with 'tool_calls' and 'tool_responses'
+        Includes the 'name' (if it exists) and the 'content', with a new line between each one, like:
+        customer:
+        <content>
+
+        agent:
+        <content>
+        """
+
+        if isinstance(messages, str):
+            if system_message:
+                summary = f"IMPORTANT: {system_message}\nContext:\n\n{messages}"
+            else:
+                return messages
+
+        elif isinstance(messages, List):
+            summary = ""
+            for message in messages:
+                if "content" in message and "tool_calls" not in message and "tool_responses" not in message:
+                    summary += f"{message.get('name', '')}: {message.get('content','')}\n\n"
+
+            if system_message:
+                summary = f"IMPORTANT: {system_message}\nContext:\n\n{summary}"
+
+            return summary
+
+        else:
+            raise ValueError("Invalid messages format. Must be a list of messages or a string.")