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 bf67bd5571..9673cafea0 100644 --- a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py +++ b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py @@ -7,17 +7,14 @@ 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.openai import OpenAiGenerativeModel +from graphrag_sdk.ontology import Ontology from .document import Document from .graph_query_engine import GraphStoreQueryResult -@dataclass -class FalkorGraphQueryResult(GraphStoreQueryResult): - messages: list = field(default_factory=list) - - class FalkorGraphQueryEngine: """ This is a wrapper for FalkorDB KnowledgeGraph. @@ -31,7 +28,7 @@ def __init__( username: str | None = None, password: str | None = None, model: str = "gpt-4o", - schema: Schema | None = None, + ontology: Ontology | None = None, ): """ Initialize a FalkorDB knowledge graph. @@ -44,10 +41,22 @@ def __init__( username (str|None): FalkorDB username. password (str|None): FalkorDB password. model (str): OpenAI model to use for FalkorDB to build and retrieve from the graph. - schema: FalkorDB 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, FalkorDB will auto generate a schema from the input docs. + ontology: FalkorDB 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, FalkorDB will auto generate an ontology from the input docs. """ - self.knowledge_graph = KnowledgeGraph(name, host, port, username, password, model, schema) + openai_model = OpenAiGenerativeModel(model) + self.knowledge_graph = KnowledgeGraph( + name=name, + host=host, + port=port, + username=username, + password=password, + model_config=KnowledgeGraphModelConfig.with_model(openai_model), + ontology=ontology, + ) + + # Establish a chat session, this will maintain the history + self._chat_session = self.knowledge_graph.chat_session() def init_db(self, input_doc: List[Document] | None): """ @@ -64,7 +73,7 @@ def init_db(self, input_doc: List[Document] | None): def add_records(self, new_records: List) -> bool: 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 knowledge graph with a question and optional message history. @@ -76,6 +85,9 @@ 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) + 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 index 842edc975d..7e8bf9cc90 100644 --- a/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py +++ b/autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py @@ -6,7 +6,8 @@ from autogen import Agent, ConversableAgent, UserProxyAgent -from .falkor_graph_query_engine import FalkorGraphQueryEngine, FalkorGraphQueryResult +from .falkor_graph_query_engine import FalkorGraphQueryEngine +from .graph_query_engine import GraphStoreQueryResult from .graph_rag_capability import GraphRagCapability @@ -24,9 +25,6 @@ def __init__(self, query_engine: FalkorGraphQueryEngine): """ self.query_engine = query_engine - # Graph DB query history. - self._history = [] - def add_to_agent(self, agent: UserProxyAgent): """ Add FalkorDB GraphRAG capability to a UserProxyAgent. @@ -69,9 +67,7 @@ def _reply_using_falkordb_query( A tuple containing a boolean indicating success and the assistant's reply. """ question = self._get_last_question(messages[-1]) - result: FalkorGraphQueryResult = self.query_engine.query(question, messages=self._history) - - self._history = result.messages + 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." diff --git a/notebook/agentchat_graph_rag_falkordb.ipynb b/notebook/agentchat_graph_rag_falkordb.ipynb index 023a893d50..8af8cc52ca 100644 --- a/notebook/agentchat_graph_rag_falkordb.ipynb +++ b/notebook/agentchat_graph_rag_falkordb.ipynb @@ -22,7 +22,7 @@ "metadata": {}, "outputs": [], "source": [ - "%pip install graphrag_sdk==0.1.3b0" + "%pip install graphrag_sdk==0.3.3" ] }, { @@ -51,9 +51,18 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 1, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/usr/local/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + } + ], "source": [ "import os\n", "\n", @@ -100,7 +109,7 @@ "metadata": {}, "outputs": [], "source": [ - "from graphrag_sdk.schema import Schema\n", + "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", @@ -108,16 +117,20 @@ "# 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", + "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=\"192.168.0.1\", # Change\n", " port=6379, # if needed\n", - " schema=movie_schema,\n", + " ontology=movie_ontology,\n", ")\n", "\n", "input_documents = [Document(doctype=DocumentType.TEXT, path_or_url=input_path)]\n", @@ -136,7 +149,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 3, "metadata": {}, "outputs": [ { @@ -150,17 +163,27 @@ "--------------------------------------------------------------------------------\n", "\u001b[33mmatrix_agent\u001b[0m (to user_proxy):\n", "\n", - "A few actors who have played in 'The Matrix' are Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving.\n", + "Keanu Reeves, Laurence Fishburne, and Carrie-Anne Moss are a few actors who've played in 'The Matrix'.\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33muser_proxy\u001b[0m (to matrix_agent):\n", + "\n", + "Anyone else?\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33mmatrix_agent\u001b[0m (to user_proxy):\n", + "\n", + "Hugo Weaving, Lilly Wachowski, and Lana Wachowski are other individuals associated with 'The Matrix'.\n", "\n", "--------------------------------------------------------------------------------\n", "\u001b[33muser_proxy\u001b[0m (to matrix_agent):\n", "\n", - "Who else stared in The Matrix?\n", + "Anyone else?\n", "\n", "--------------------------------------------------------------------------------\n", "\u001b[33mmatrix_agent\u001b[0m (to user_proxy):\n", "\n", - "I'm sorry, I don't have an answer for that.\n", + "No, there are no other actors associated with 'The Matrix' besides those already mentioned.\n", "\n", "--------------------------------------------------------------------------------\n" ] @@ -168,10 +191,10 @@ { "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 Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving.\", 'role': 'user', 'name': 'matrix_agent'}, {'content': 'Who else stared in The Matrix?', 'role': 'assistant', 'name': 'user_proxy'}, {'content': \"I'm sorry, I don't have an answer for that.\", 'role': 'user', 'name': 'matrix_agent'}], summary=\"I'm sorry, I don't have an answer for that.\", cost={'usage_including_cached_inference': {'total_cost': 0}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=['Who else stared in The Matrix?', 'exit'])" + "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, and Carrie-Anne Moss are a few actors who've played in 'The Matrix'.\", 'role': 'user', 'name': 'matrix_agent'}, {'content': 'Anyone else?', 'role': 'assistant', 'name': 'user_proxy'}, {'content': \"Hugo Weaving, Lilly Wachowski, and Lana Wachowski are other individuals associated with 'The Matrix'.\", 'role': 'user', 'name': 'matrix_agent'}, {'content': 'Anyone else?', 'role': 'assistant', 'name': 'user_proxy'}, {'content': \"No, there are no other actors associated with 'The Matrix' besides those already mentioned.\", 'role': 'user', 'name': 'matrix_agent'}], summary=\"No, there are no other actors associated with 'The Matrix' besides those already mentioned.\", cost={'usage_including_cached_inference': {'total_cost': 0}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=['Anyone else?', 'Anyone else?', 'exit'])" ] }, - "execution_count": 5, + "execution_count": 3, "metadata": {}, "output_type": "execute_result" } diff --git a/setup.py b/setup.py index a77fbce05c..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==0.1.3b0", + "graphrag_sdk==0.3.3", ] if current_os in ["Windows", "Darwin"]: