github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-07-22 11:51:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: make tugraph work again (#2551)

Browse Source 
Signed-off-by: yihong0618 <zouzou0208@gmail.com>
This commit is contained in:
yihong 2025-03-31 19:46:44 +08:00 committed by GitHub
parent 3bfbdc3fc7
commit beafca5c6c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 22 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
packages/dbgpt-ext/src/dbgpt_ext/datasource/conn_tugraph.py
View File

@ -158,11 +158,15 @@ class TuGraphConnector(BaseConnector):
with self._driver.session(database=self._graph) as session:
# Run the query to get vertex labels
raw_vertex_labels = session.run("CALL db.vertexLabels()").data()
vertex_labels = [table_name["label"] for table_name in raw_vertex_labels]
vertex_labels = [
table_name["label"] + "_vertex" for table_name in raw_vertex_labels
]
# Run the query to get edge labels
raw_edge_labels = session.run("CALL db.edgeLabels()").data()
edge_labels = [table_name["label"] for table_name in raw_edge_labels]
edge_labels = [
table_name["label"] + "_edge" for table_name in raw_edge_labels
]
return iter(vertex_labels + edge_labels)

22
packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py
View File

@ -44,12 +44,14 @@ class GdbmsSummary(DBSummary):
tables = self.db.get_table_names()
self.table_info_summaries = {
"vertex_tables": [
self.get_table_summary(table_name, "vertex")
for table_name in tables["vertex_tables"]
self.get_table_summary(table_name.split("_")[0], "vertex")
for table_name in tables
if table_name.endswith("_vertex")
],
"edge_tables": [
self.get_table_summary(table_name, "edge")
for table_name in tables["edge_tables"]
self.get_table_summary(table_name.split("_")[0], "edge")
for table_name in tables
if table_name.endswith("_edge")
],
}
@ -76,8 +78,16 @@ def _parse_db_summary(
table_info_summaries = None
if isinstance(conn, TuGraphConnector):
table_names = conn.get_table_names()
v_tables = table_names.get("vertex_tables", []) # type: ignore
e_tables = table_names.get("edge_tables", []) # type: ignore
v_tables = [
table_name.split("_")[0]
for table_name in table_names
if table_name.endswith("_vertex")
]
e_tables = [
table_name.split("_")[0]
for table_name in table_names
if table_name.endswith("_edge")
]
table_info_summaries = [
_parse_table_summary(conn, summary_template, table_name, "vertex")
for table_name in v_tables