Skip to content

Commit

Permalink
Support for latest GraphRAG 0.3.3
Browse files Browse the repository at this point in the history 
Signed-off-by: Mark Sze <[email protected]>
  • Loading branch information
@marklysze
marklysze committed Nov 29, 2024
1 parent 210cd24 commit 71ad707
Show file tree
Hide file tree
Showing 4 changed files with 68 additions and 37 deletions.
40 changes: 26 additions & 14 deletions autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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):
"""
Expand All @@ -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.
Expand All @@ -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=[])
10 changes: 3 additions & 7 deletions autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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.
Expand Down Expand Up @@ -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."

Expand Down
53 changes: 38 additions & 15 deletions notebook/agentchat_graph_rag_falkordb.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
"metadata": {},
"outputs": [],
"source": [
"%pip install graphrag_sdk==0.1.3b0"
"%pip install graphrag_sdk==0.3.3"
]
},
{
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -100,24 +109,28 @@
"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",
"\n",
"# 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",
Expand All @@ -136,7 +149,7 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": 3,
"metadata": {},
"outputs": [
{
Expand All @@ -150,28 +163,38 @@
"--------------------------------------------------------------------------------\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"
]
},
{
"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"
}
Expand Down
2 changes: 1 addition & 1 deletion setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"]:
Expand Down

0 comments on commit 71ad707

Please sign in to comment.