Skip to content

Commit

Permalink
fix(GraphRAG): create same graph name error. (#2215)
Browse files Browse the repository at this point in the history
  • Loading branch information
Aries-ckt authored Dec 19, 2024
1 parent 342a89c commit d4031f0
Show file tree
Hide file tree
Showing 5 changed files with 31 additions and 0 deletions.
13 changes: 13 additions & 0 deletions dbgpt/datasource/conn_tugraph.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,19 @@ def create_graph(self, graph_name: str) -> bool:

return not exists

def is_exist(self, graph_name: str) -> bool:
"""Check a new graph in the database if it doesn't already exist."""
try:
with self._driver.session(database="default") as session:
graph_list = session.run("CALL dbms.graph.listGraphs()").data()
exists = any(item["graph_name"] == graph_name for item in graph_list)
except Exception as e:
raise Exception(
f"Failed to check graph exist'{graph_name}': {str(e)}"
) from e

return exists

def delete_graph(self, graph_name: str) -> None:
"""Delete a graph in the database if it exists."""
with self._driver.session(database="default") as session:
Expand Down
4 changes: 4 additions & 0 deletions dbgpt/storage/graph_store/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,3 +40,7 @@ def __init__(self, config: GraphStoreConfig):
@abstractmethod
def get_config(self) -> GraphStoreConfig:
"""Get the graph store config."""

def is_exist(self, name) -> bool:
"""Check Graph Name is Exist."""
raise NotImplementedError
4 changes: 4 additions & 0 deletions dbgpt/storage/graph_store/tugraph_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,10 @@ def get_config(self) -> TuGraphStoreConfig:
"""Get the TuGraph store config."""
return self._config

def is_exist(self, name) -> bool:
"""Check Graph Name is Exist."""
return self.conn.is_exist(name)

def _add_vertex_index(self, field_name):
"""Add an index to the vertex table."""
# TODO: Not used in the current implementation.
Expand Down
2 changes: 2 additions & 0 deletions dbgpt/storage/knowledge_graph/community_summary.py
Original file line number Diff line number Diff line change
Expand Up @@ -316,6 +316,8 @@ def get_config(self) -> BuiltinKnowledgeGraphConfig:

async def aload_document(self, chunks: List[Chunk]) -> List[str]:
"""Extract and persist graph from the document file."""
if not self.vector_name_exists():
self._graph_store_apdater.create_graph(self.get_config().name)
await self._aload_document_graph(chunks)
await self._aload_triplet_graph(chunks)
await self._community_store.build_communities(
Expand Down
8 changes: 8 additions & 0 deletions dbgpt/storage/knowledge_graph/knowledge_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,8 @@ async def process_chunk(chunk: Chunk):
return chunk.chunk_id

# wait async tasks completed
if not self.vector_name_exists():
self._graph_store_apdater.create_graph(self.get_config().name)
tasks = [process_chunk(chunk) for chunk in chunks]
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
Expand All @@ -187,6 +189,8 @@ async def aload_document(self, chunks: List[Chunk]) -> List[str]: # type: ignor
Return:
List[str]: chunk ids.
"""
if not self.vector_name_exists():
self._graph_store_apdater.create_graph(self.get_config().name)
for chunk in chunks:
triplets = await self._triplet_extractor.extract(chunk.content)
for triplet in triplets:
Expand Down Expand Up @@ -279,3 +283,7 @@ def delete_by_ids(self, ids: str) -> List[str]:
"""Delete by ids."""
self._graph_store_apdater.delete_document(chunk_id=ids)
return []

def vector_name_exists(self) -> bool:
"""Whether name exists."""
return self._graph_store_apdater.graph_store.is_exist(self._config.name)

0 comments on commit d4031f0

Please sign in to comment.