feat: Execute with custom sql dataset

Author: fit2cloud-chenyw

backend/apps/chat/models/chat_model.py

@@ -9,6 +9,7 @@
 from apps.template.filter.generator import get_permissions_template
 from apps.template.generate_analysis.generator import get_analysis_template
 from apps.template.generate_chart.generator import get_chart_template
+from apps.template.generate_dynamic.generator import get_dynamic_template
 from apps.template.generate_guess_question.generator import get_guess_question_template
 from apps.template.generate_predict.generator import get_predict_template
 from apps.template.generate_sql.generator import get_sql_template
@@ -107,6 +108,7 @@ class AiModelQuestion(BaseModel):
     data: str = ""
     lang: str = "简体中文"
     filter: str = []
+    sub_query: Optional[list[dict]] = None
 
     def sql_sys_question(self):
         return get_sql_template()['system'].format(engine=self.engine, schema=self.db_schema, question=self.question, lang=self.lang)
@@ -151,7 +153,12 @@ def filter_sys_question(self):
 
     def filter_user_question(self):
         return get_permissions_template()['user'].format(sql=self.sql, filter=self.filter)
-
+    
+    def dynamic_sys_question(self): 
+        return get_dynamic_template()['system'].format(lang=self.lang, engine=self.engine)
+    
+    def dynamic_user_question(self): 
+        return get_dynamic_template()['user'].format(sql=self.sql, sub_query=self.sub_query)
 
 class ChatQuestion(AiModelQuestion):
     question: str = Body(description='用户提问')

backend/apps/chat/task/llm.py

@@ -516,6 +516,41 @@ def generate_sql(self):
                                                            [{'type': msg.type, 'content': msg.content} for msg in
                                                             self.sql_message]).decode())
 
+    def generate_with_sub_sql(self, sql, sub_mappings: list):
+        sub_query = json.dumps(sub_mappings, ensure_ascii=False)
+        self.chat_question.sql = sql
+        self.chat_question.sub_query = sub_query
+        msg: List[Union[BaseMessage, dict[str, Any]]] = []
+        msg.append(SystemMessage(content=self.chat_question.dynamic_sys_question()))
+        msg.append(HumanMessage(content=self.chat_question.dynamic_user_question()))
+        full_thinking_text = ''
+        full_dynamic_text = ''
+        res = self.llm.stream(msg)
+        token_usage = {}
+        for chunk in res:
+            SQLBotLogUtil.info(chunk)
+            reasoning_content_chunk = ''
+            if 'reasoning_content' in chunk.additional_kwargs:
+                reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
+            if reasoning_content_chunk is None:
+                reasoning_content_chunk = ''
+            full_thinking_text += reasoning_content_chunk
+            full_dynamic_text += chunk.content
+            get_token_usage(chunk, token_usage)
+
+        SQLBotLogUtil.info(full_dynamic_text)
+        return full_dynamic_text
+    
+    def generate_assistant_dynamic_sql(self, sql, tables: List):
+        ds: AssistantOutDsSchema = self.ds
+        sub_query = []
+        for table in ds.tables:
+            if table.name in tables and table.sql:
+                sub_query.append({"table": table.name, "query": table.sql})
+        if not sub_query:
+            return None
+        return self.generate_with_sub_sql(sql=sql, sub_mappings=sub_query)
+        
     def build_table_filter(self, sql: str, filters: list):
         filter = json.dumps(filters, ensure_ascii=False)
         self.chat_question.sql = sql
@@ -635,27 +670,23 @@ def generate_chart(self):
                                                          full_message=orjson.dumps(
                                                              [{'type': msg.type, 'content': msg.content} for msg in
                                                               self.chart_message]).decode())
-
-    def check_save_sql(self, res: str) -> str:
-
+    def check_sql(self, res: str) -> tuple[any]:
         json_str = extract_nested_json(res)
-        data = orjson.loads(json_str)
-
+        data: dict = orjson.loads(json_str)
         sql = ''
         message = ''
-        error = False
-
         if data['success']:
             sql = data['sql']
         else:
             message = data['message']
-            error = True
-
-        if error:
             raise Exception(message)
+        
         if sql.strip() == '':
             raise Exception("SQL query is empty")
-
+        return sql, data.get('tables')
+        
+    def check_save_sql(self, res: str) -> str:
+        sql, *_ = self.check_sql(res=res)
         save_sql(session=self.session, sql=sql, record_id=self.record.id)
 
         self.chat_question.sql = sql
@@ -816,34 +847,27 @@ def run_task(self, in_chat: bool = True):
 
             # filter sql
             SQLBotLogUtil.info(full_sql_text)
+            use_dynamic_ds: bool = self.current_assistant and self.current_assistant.type == 1
 
             # todo row permission
-            if (not self.current_assistant and is_normal_user(self.current_user)) or (
-                    self.current_assistant and self.current_assistant.type == 1):
-                sql_json_str = extract_nested_json(full_sql_text)
-                data = orjson.loads(sql_json_str)
-
-                sql = ''
-                message = ''
-                error = False
-                if data['success']:
-                    sql = data['sql']
-                else:
-                    message = data['message']
-                    error = True
-                if error:
-                    raise Exception(message)
-                if sql.strip() == '':
-                    raise Exception("SQL query is empty")
+            if (not self.current_assistant and is_normal_user(self.current_user)) or use_dynamic_ds:
+                sql, tables = self.check_sql(res=full_sql_text)
 
                 if self.current_assistant:
-                    sql_result = self.generate_assistant_filter(data.get('sql'), data.get('tables'))
+                    dynamic_sql_result = self.generate_assistant_dynamic_sql(sql, tables)
+                    if dynamic_sql_result:
+                        SQLBotLogUtil.info(dynamic_sql_result)
+                        sql, *_ = self.check_sql(res=dynamic_sql_result)
+                        
+                    sql_result = self.generate_assistant_filter(sql, tables)
                 else:
-                    sql_result = self.generate_filter(data.get('sql'), data.get('tables'))  # maybe no sql and tables
+                    sql_result = self.generate_filter(sql, tables)  # maybe no sql and tables
 
                 if sql_result:
                     SQLBotLogUtil.info(sql_result)
                     sql = self.check_save_sql(res=sql_result)
+                elif dynamic_sql_result:
+                    sql = self.check_save_sql(res=dynamic_sql_result)
                 else:
                     sql = self.check_save_sql(res=full_sql_text)
             else:

backend/apps/template/generate_dynamic/__init__.py

@@ -0,0 +1,6 @@
+from apps.template.template import get_base_template
+
+
+def get_dynamic_template():
+    template = get_base_template()
+    return template['template']['dynamic_sql']

backend/apps/template/generate_dynamic/generator.py

@@ -214,14 +214,13 @@ template:
       
       ### 说明：
       提供给你一句SQL和一组子查询映射表，你需要将给定的SQL查询中的表名替换为对应的子查询。请严格保持原始SQL的结构不变，只替换表引用部分，生成符合{engine}数据库引擎规范的新SQL语句。
-      - 原始SQL（标记为`sql`）
-      - 子查询映射表（标记为`sub_query`，格式为`[{"原表名": "子查询SQL"},...]`）
+      - 子查询映射表标记为sub_query，格式为[{{"table":"表名","query":"子查询语句"}},...]
       你必须遵守以下规则:
       - 生成的SQL必须符合{engine}的规范。
       - 不要替换原来SQL中的过滤条件。
       - 完全匹配表名（注意大小写敏感）。
       - 根据子查询语句以及{engine}数据库引擎规范决定是否需要给子查询添加括号包围
-      - 若子查询包含别名，保留原表名作为别名
+      - 若原始SQL中原表名有别名则保留原有别名，否则保留原表名作为别名
       - 生成SQL时，必须避免关键字冲突。
       - 生成的SQL使用JSON格式返回：
       {{"success":true,"sql":"生成的SQL语句"}}
@@ -235,5 +234,5 @@ template:
       ### sql:
       {sql}
       
-      ### 过滤条件:
+      ### 子查询映射表:
       {sub_query}