diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index e86fc11f..c1f431d7 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/backend/apps/chat/models/chat_model.py @@ -9,6 +9,7 @@ from sqlmodel import SQLModel, Field 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 @@ class AiModelQuestion(BaseModel): 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='用户提问') diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 64992e52..3311c9f8 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -516,6 +516,41 @@ class LLMService: [{'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 @@ class LLMService: 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 @@ class LLMService: # 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: diff --git a/backend/apps/template/generate_dynamic/__init__.py b/backend/apps/template/generate_dynamic/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/apps/template/generate_dynamic/generator.py b/backend/apps/template/generate_dynamic/generator.py new file mode 100644 index 00000000..4c69e0e4 --- /dev/null +++ b/backend/apps/template/generate_dynamic/generator.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'] diff --git a/backend/template.yaml b/backend/template.yaml index b87fbac8..76f65a77 100644 --- a/backend/template.yaml +++ b/backend/template.yaml @@ -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}