From 15e43e54043257544d6f7d7539bc011f5964ed5e Mon Sep 17 00:00:00 2001 From: yihong0618 Date: Fri, 28 Mar 2025 18:10:28 +0800 Subject: [PATCH 1/3] fix: make tugraph work again Signed-off-by: yihong0618 --- .../src/dbgpt_ext/datasource/conn_tugraph.py | 8 +++++-- .../dbgpt_ext/rag/summary/gdbms_db_summary.py | 24 ++++++++++++++----- 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/packages/dbgpt-ext/src/dbgpt_ext/datasource/conn_tugraph.py b/packages/dbgpt-ext/src/dbgpt_ext/datasource/conn_tugraph.py index 210a72f2ae..3f9183cd89 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/datasource/conn_tugraph.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/datasource/conn_tugraph.py @@ -158,11 +158,15 @@ def get_table_names(self) -> Iterator[str]: 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) diff --git a/packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py b/packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py index aaac9b6514..5f1ffac8d8 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py @@ -42,14 +42,17 @@ def __init__( collation=self.db.get_collation(), ) tables = self.db.get_table_names() + print(list(tables), "!!!!!!!!!!!!!!!!!!!!!!1") 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 +79,17 @@ 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") + ] + print(v_tables, e_tables) table_info_summaries = [ _parse_table_summary(conn, summary_template, table_name, "vertex") for table_name in v_tables From c7e33319982e9b3af1bf4b3c786231b349e2f873 Mon Sep 17 00:00:00 2001 From: yihong0618 Date: Fri, 28 Mar 2025 18:13:41 +0800 Subject: [PATCH 2/3] fix: debug info Signed-off-by: yihong0618 --- packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py b/packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py index 5f1ffac8d8..7ac1fea315 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py @@ -42,7 +42,6 @@ def __init__( collation=self.db.get_collation(), ) tables = self.db.get_table_names() - print(list(tables), "!!!!!!!!!!!!!!!!!!!!!!1") self.table_info_summaries = { "vertex_tables": [ self.get_table_summary(table_name.split("_")[0], "vertex") From 222d149ad1ed7463a914f04a9a180e3145124f7d Mon Sep 17 00:00:00 2001 From: yihong0618 Date: Sat, 29 Mar 2025 08:36:23 +0800 Subject: [PATCH 3/3] fix: drop useless print Signed-off-by: yihong0618 --- packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py b/packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py index 7ac1fea315..3eef4ec9ee 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/rag/summary/gdbms_db_summary.py @@ -88,7 +88,6 @@ def _parse_db_summary( for table_name in table_names if table_name.endswith("_edge") ] - print(v_tables, e_tables) table_info_summaries = [ _parse_table_summary(conn, summary_template, table_name, "vertex") for table_name in v_tables