Skip to content

Commit

Permalink
Merge pull request #178 from ag2ai/falkor_db_exception
Browse files Browse the repository at this point in the history 
Falkor DB Query Engine: Add delete method and handle exceptions
  • Loading branch information
@AgentGenie
AgentGenie authored Dec 12, 2024
2 parents 4fe28fc + 359e327 commit a10979c
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 35 deletions.
36 changes: 25 additions & 11 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 @@ -49,6 +49,7 @@ def __init__(
If None, FalkorDB will auto generate an ontology from the input docs.
"""
self.name = name
self.ontology_table_name = name + "_ontology"
self.host = host
self.port = port
self.username = username
Expand All @@ -65,7 +66,7 @@ def connect_db(self):
"""
if self.name in self.falkordb.list_graphs():
try:
self.ontology = self._load_ontology_from_db(self.name)
self.ontology = self._load_ontology_from_db()
except Exception:
warnings.warn("Graph Ontology is not loaded.")

Expand Down Expand Up @@ -103,6 +104,8 @@ def init_db(self, input_doc: List[Document]):
sources=sources,
model=self.model,
)
# Save Ontology to graph for future access.
self._save_ontology_to_db(self.ontology)

self.knowledge_graph = KnowledgeGraph(
name=self.name,
Expand All @@ -118,9 +121,6 @@ def init_db(self, input_doc: List[Document]):

# 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.")

Expand Down Expand Up @@ -149,17 +149,31 @@ def query(self, question: str, n_results: int = 1, **kwargs) -> GraphStoreQueryR

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 delete(self) -> bool:
"""
Delete graph and its data from database.
"""
all_graphs = self.falkordb.list_graphs()
if self.name in all_graphs:
self.falkordb.select_graph(self.name).delete()
if self.ontology_table_name in all_graphs:
self.falkordb.select_graph(self.ontology_table_name).delete()
return True

def __get_ontology_storage_graph(self) -> Graph:
return self.falkordb.select_graph(self.ontology_table_name)

def _save_ontology_to_db(self, graph_name: str, ontology: Ontology):
def _save_ontology_to_db(self, ontology: Ontology):
"""
Save graph ontology to a separate table with {graph_name}_ontology
"""
graph = self.__get_ontology_storage_graph(graph_name)
if self.ontology_table_name in self.falkordb.list_graphs():
raise ValueError("Knowledge graph {} is already created.".format(self.name))
graph = self.__get_ontology_storage_graph()
ontology.save_to_graph(graph)

def _load_ontology_from_db(self, graph_name: str) -> Ontology:
graph = self.__get_ontology_storage_graph(graph_name)
def _load_ontology_from_db(self) -> Ontology:
if self.ontology_table_name not in self.falkordb.list_graphs():
raise ValueError("Knowledge graph {} has not been created.".format(self.name))
graph = self.__get_ontology_storage_graph()
return Ontology.from_graph(graph)
60 changes: 36 additions & 24 deletions notebook/agentchat_graph_rag_falkordb.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,18 +48,9 @@
},
{
"cell_type": "code",
"execution_count": 1,
"execution_count": null,
"metadata": {},
"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"
]
}
],
"outputs": [],
"source": [
"import os\n",
"\n",
Expand Down Expand Up @@ -108,7 +99,7 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": 3,
"metadata": {},
"outputs": [
{
Expand All @@ -127,23 +118,23 @@
"--------------------------------------------------------------------------------\n",
"\u001b[33muser_proxy\u001b[0m (to matrix_agent):\n",
"\n",
"Who else acted in The Matrix?\n",
"Any additional actors?\n",
"\n",
"--------------------------------------------------------------------------------\n",
"\u001b[33mmatrix_agent\u001b[0m (to user_proxy):\n",
"\n",
"Based on the provided information, there is no additional data about other actors who acted in 'The Matrix' outside of Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving.\n",
"I'm sorry, but I don't have information about any additional actors from 'The Matrix' beyond Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving.\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': \"Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving are a few actors who've played in 'The Matrix'.\", 'role': 'user', 'name': 'matrix_agent'}, {'content': 'Who else acted in The Matrix?', 'role': 'assistant', 'name': 'user_proxy'}, {'content': \"Based on the provided information, there is no additional data about other actors who acted in 'The Matrix' outside of Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving.\", 'role': 'user', 'name': 'matrix_agent'}], summary=\"Based on the provided information, there is no additional data about other actors who acted in 'The Matrix' outside of Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving.\", cost={'usage_including_cached_inference': {'total_cost': 0}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=['Who else acted 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, Carrie-Anne Moss, and Hugo Weaving are a few actors who've played in 'The Matrix'.\", 'role': 'user', 'name': 'matrix_agent'}, {'content': 'Any additional actors?', 'role': 'assistant', 'name': 'user_proxy'}, {'content': \"I'm sorry, but I don't have information about any additional actors from 'The Matrix' beyond Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving.\", 'role': 'user', 'name': 'matrix_agent'}], summary=\"I'm sorry, but I don't have information about any additional actors from 'The Matrix' beyond Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving.\", cost={'usage_including_cached_inference': {'total_cost': 0}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=['Any additional actors?', 'exit'])"
]
},
"execution_count": 2,
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
Expand All @@ -161,7 +152,7 @@
"# Create FalkorGraphQueryEngine\n",
"query_engine = FalkorGraphQueryEngine(\n",
" name=\"The_Matrix_Auto\",\n",
" host=\"172.18.0.3\", # Change\n",
" host=\"172.17.0.4\", # Change\n",
" port=6379, # if needed\n",
")\n",
"\n",
Expand Down Expand Up @@ -197,7 +188,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -219,7 +210,7 @@
"\n",
"query_engine = FalkorGraphQueryEngine(\n",
" name=\"IMDB\",\n",
" host=\"172.18.0.3\", # Change\n",
" host=\"172.17.0.4\", # Change\n",
" port=6379, # if needed\n",
" ontology=movie_ontology,\n",
")\n",
Expand All @@ -240,7 +231,7 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 6,
"metadata": {},
"outputs": [
{
Expand All @@ -254,7 +245,7 @@
"--------------------------------------------------------------------------------\n",
"\u001b[33mmatrix_agent\u001b[0m (to user_proxy):\n",
"\n",
"Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving are some of the actors who have played in \"The Matrix\".\n",
"Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving have played in 'The Matrix'.\n",
"\n",
"--------------------------------------------------------------------------------\n",
"\u001b[33muser_proxy\u001b[0m (to matrix_agent):\n",
Expand All @@ -264,7 +255,7 @@
"--------------------------------------------------------------------------------\n",
"\u001b[33mmatrix_agent\u001b[0m (to user_proxy):\n",
"\n",
"No other actors are listed as having acted in \"The Matrix\" beyond Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving.\n",
"Based on the provided information, Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving are the actors listed for 'The Matrix'. No additional actors are mentioned in the data.\n",
"\n",
"--------------------------------------------------------------------------------\n",
"\u001b[33muser_proxy\u001b[0m (to matrix_agent):\n",
Expand All @@ -282,10 +273,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': 'Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving are some of the actors who have played in \"The Matrix\".', 'role': 'user', 'name': 'matrix_agent'}, {'content': 'Who else acted in The Matrix?', 'role': 'assistant', 'name': 'user_proxy'}, {'content': 'No other actors are listed as having acted in \"The Matrix\" beyond Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving.', 'role': 'user', 'name': 'matrix_agent'}, {'content': 'Where was The Matrix filmed?', 'role': 'assistant', 'name': 'user_proxy'}, {'content': 'I am sorry, I could not find the answer to your question', 'role': 'user', 'name': 'matrix_agent'}], summary='I am sorry, I could not find the answer to your question', cost={'usage_including_cached_inference': {'total_cost': 0}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=['Who else acted in The Matrix?', 'Where was The Matrix filmed?', '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, Carrie-Anne Moss, and Hugo Weaving have played in 'The Matrix'.\", 'role': 'user', 'name': 'matrix_agent'}, {'content': 'Who else acted in The Matrix?', 'role': 'assistant', 'name': 'user_proxy'}, {'content': \"Based on the provided information, Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving are the actors listed for 'The Matrix'. No additional actors are mentioned in the data.\", 'role': 'user', 'name': 'matrix_agent'}, {'content': 'Where was The Matrix filmed?', 'role': 'assistant', 'name': 'user_proxy'}, {'content': 'I am sorry, I could not find the answer to your question', 'role': 'user', 'name': 'matrix_agent'}], summary='I am sorry, I could not find the answer to your question', cost={'usage_including_cached_inference': {'total_cost': 0}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=['Who else acted in The Matrix?', 'Where was The Matrix filmed?', 'exit'])"
]
},
"execution_count": 4,
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
Expand Down Expand Up @@ -326,6 +317,27 @@
"- There is inbuilt awareness of answers it has provided previously\n",
"- It can't answer any questions outside of the actors and movies entities as we didn't include anything else in the ontology"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"True"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Remove graph from database:\n",
"query_engine.delete()"
]
}
],
"metadata": {
Expand Down

0 comments on commit a10979c

Please sign in to comment.