From a5666b3120cc565cbc9111e6328c4d13b8a7f4ce Mon Sep 17 00:00:00 2001 From: lipengfei <15566300566@163.com> Date: Fri, 26 Apr 2024 09:48:40 +0800 Subject: [PATCH] feat: Added support for TuGraph graph database (#1451) Co-authored-by: aries_ckt <916701291@qq.com> --- .mypy.ini | 3 + .../app/scene/chat_db/professional_qa/chat.py | 18 ++- dbgpt/app/static/404.html | 2 +- dbgpt/app/static/404/index.html | 2 +- .../_buildManifest.js | 0 .../_ssgManifest.js | 0 ...98b291d43c.js => _app-f113d80168f7f0f2.js} | 34 ++--- dbgpt/app/static/agent/index.html | 2 +- dbgpt/app/static/app/index.html | 2 +- dbgpt/app/static/chat/index.html | 2 +- dbgpt/app/static/database/index.html | 2 +- dbgpt/app/static/flow/canvas/index.html | 2 +- dbgpt/app/static/flow/index.html | 2 +- dbgpt/app/static/icons/tugraph.png | Bin 0 -> 252992 bytes dbgpt/app/static/index.html | 2 +- dbgpt/app/static/knowledge/chunk/index.html | 2 +- dbgpt/app/static/knowledge/index.html | 2 +- dbgpt/app/static/models/index.html | 2 +- dbgpt/app/static/prompt/index.html | 2 +- dbgpt/datasource/base.py | 5 + dbgpt/datasource/conn_tugraph.py | 127 +++++++++++++++++ dbgpt/datasource/manages/connector_manager.py | 1 + dbgpt/rag/knowledge/datasource.py | 7 +- dbgpt/rag/summary/db_summary_client.py | 17 ++- dbgpt/rag/summary/gdbms_db_summary.py | 134 ++++++++++++++++++ dbgpt/storage/schema.py | 1 + setup.py | 1 + .../datasource/test_conn_tugraph.py | 42 ++++++ web/public/icons/tugraph.png | Bin 0 -> 252992 bytes web/utils/constants.ts | 1 + 30 files changed, 379 insertions(+), 38 deletions(-) rename dbgpt/app/static/_next/static/{GRS1EW__ol-4Jl6ARii0d => BaAlw8DGDQ-bWcVWS80mT}/_buildManifest.js (100%) rename dbgpt/app/static/_next/static/{GRS1EW__ol-4Jl6ARii0d => BaAlw8DGDQ-bWcVWS80mT}/_ssgManifest.js (100%) rename dbgpt/app/static/_next/static/chunks/pages/{_app-d2565d98b291d43c.js => _app-f113d80168f7f0f2.js} (77%) create mode 100644 dbgpt/app/static/icons/tugraph.png create mode 100644 dbgpt/datasource/conn_tugraph.py create mode 100644 dbgpt/rag/summary/gdbms_db_summary.py create mode 100644 tests/intetration_tests/datasource/test_conn_tugraph.py create mode 100644 web/public/icons/tugraph.png diff --git a/.mypy.ini b/.mypy.ini index d9dfd8233..91f6d0c19 100644 --- a/.mypy.ini +++ b/.mypy.ini @@ -88,6 +88,9 @@ ignore_missing_imports = True [mypy-fastchat.protocol.api_protocol] ignore_missing_imports = True +[mypy-neo4j.*] +ignore_missing_imports = True + # Agent [mypy-seaborn.*] ignore_missing_imports = True diff --git a/dbgpt/app/scene/chat_db/professional_qa/chat.py b/dbgpt/app/scene/chat_db/professional_qa/chat.py index 06411319c..52ed56abc 100644 --- a/dbgpt/app/scene/chat_db/professional_qa/chat.py +++ b/dbgpt/app/scene/chat_db/professional_qa/chat.py @@ -31,12 +31,18 @@ class ChatWithDbQA(BaseChat): if self.db_name: self.database = CFG.local_db_manager.get_connector(self.db_name) self.tables = self.database.get_table_names() - - self.top_k = ( - CFG.KNOWLEDGE_SEARCH_TOP_SIZE - if len(self.tables) > CFG.KNOWLEDGE_SEARCH_TOP_SIZE - else len(self.tables) - ) + if self.database.is_graph_type(): + # When the current graph database retrieves source data from ChatDB, the topk uses the sum of node table and edge table. + self.top_k = len(self.tables["vertex_tables"]) + len( + self.tables["edge_tables"] + ) + else: + print(self.database.db_type) + self.top_k = ( + CFG.KNOWLEDGE_SEARCH_TOP_SIZE + if len(self.tables) > CFG.KNOWLEDGE_SEARCH_TOP_SIZE + else len(self.tables) + ) @trace() async def generate_input_values(self) -> Dict: diff --git a/dbgpt/app/static/404.html b/dbgpt/app/static/404.html index 167212bd7..184e3b581 100644 --- a/dbgpt/app/static/404.html +++ b/dbgpt/app/static/404.html @@ -1 +1 @@ -