github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-07-23 12:21:08 +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: with self._driver.session(database=self._graph) as session:
# Run the query to get vertex labels # Run the query to get vertex labels
raw_vertex_labels = session.run("CALL db.vertexLabels()").data() 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 # Run the query to get edge labels
raw_edge_labels = session.run("CALL db.edgeLabels()").data() 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) 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() tables = self.db.get_table_names()
self.table_info_summaries = { self.table_info_summaries = {
"vertex_tables": [ "vertex_tables": [
self.get_table_summary(table_name, "vertex") self.get_table_summary(table_name.split("_")[0], "vertex")
for table_name in tables["vertex_tables"] for table_name in tables
if table_name.endswith("_vertex")
], ],
"edge_tables": [ "edge_tables": [
self.get_table_summary(table_name, "edge") self.get_table_summary(table_name.split("_")[0], "edge")
for table_name in tables["edge_tables"] for table_name in tables
if table_name.endswith("_edge")
], ],
} }
@ -76,8 +78,16 @@ def _parse_db_summary(
table_info_summaries = None table_info_summaries = None
if isinstance(conn, TuGraphConnector): if isinstance(conn, TuGraphConnector):
table_names = conn.get_table_names() table_names = conn.get_table_names()
v_tables = table_names.get("vertex_tables", []) # type: ignore v_tables = [
e_tables = table_names.get("edge_tables", []) # type: ignore 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 = [ table_info_summaries = [
_parse_table_summary(conn, summary_template, table_name, "vertex") _parse_table_summary(conn, summary_template, table_name, "vertex")
for table_name in v_tables for table_name in v_tables