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 @@ -