22import concurrent
33import json
44import os
5+ import re
56import traceback
67import urllib .parse
78import warnings
3031from apps .chat .curd .chat import save_question , save_sql_answer , save_sql , \
3132 save_error_message , save_sql_exec_data , save_chart_answer , save_chart , \
3233 finish_record , save_analysis_answer , save_predict_answer , save_predict_data , \
33- save_select_datasource_answer , save_recommend_question_answer , \
34+ save_select_datasource_answer , save_recommend_question_answer , save_extracted_keywords , \
3435 get_old_questions , save_analysis_predict_record , rename_chat , get_chart_config , \
3536 get_chat_chart_data , list_generate_sql_logs , list_generate_chart_logs , start_log , end_log , \
3637 get_last_execute_sql_error , format_json_data , format_chart_fields , get_chat_brief_generate , get_chat_predict_data , \
4849from apps .system .crud .parameter_manage import get_groups
4950from apps .system .crud .user import user_ws_list
5051from apps .system .schemas .system_schema import AssistantOutDsSchema
52+ from apps .terminology .curd .terminology import expand_with_terminology
5153from apps .terminology .curd .terminology import get_terminology_template
5254from common .core .config import settings
5355from common .core .db import engine
@@ -378,14 +380,16 @@ def filter_terminology_template(self, _session: Session, oid: int = None, ds_id:
378380 calculate_oid = self .current_assistant .oid if self .current_assistant .type != 4 else self .oid
379381 if self .current_assistant .type == 1 :
380382 calculate_ds_id = None
383+ # 使用原始提取的关键词(未扩展)进行术语匹配,避免同义词重复查询
384+ match_text = self .chat_question .extracted_keywords or self .chat_question .question
381385 if self .current_assistant and self .current_assistant .type == 1 :
382386 self .chat_question .terminologies , term_list = get_terminology_template (_session ,
383- self . chat_question . question ,
387+ match_text ,
384388 calculate_oid ,
385389 None , self .current_assistant .id )
386390 else :
387391 self .chat_question .terminologies , term_list = get_terminology_template (_session ,
388- self . chat_question . question ,
392+ match_text ,
389393 calculate_oid ,
390394 calculate_ds_id )
391395
@@ -447,17 +451,108 @@ def filter_training_template(self, _session: Session, oid: int = None, ds_id: in
447451 OperationEnum .FILTER_SQL_EXAMPLE ],
448452 full_message = example_list )
449453
454+ def extract_keywords (self , _session : Session ) -> tuple [bool , str ]:
455+ """从用户问题中提取核心业务实体词。
456+
457+ 使用 LLM 提取名词性业务实体,忽略时间词、动作词、修饰词。
458+ 支持上下文历史,可根据对话上下文理解用户意图。
459+ 提取结果可复用于表匹配、术语扩展等环节。
460+
461+ Returns:
462+ (is_error, result) 元组:
463+ - (False, keywords) 成功,关键词逗号分隔
464+ - (True, error_message) 错误(LLM 失败或非查数据意图)
465+ """
466+ system_prompt = self .chat_question .extract_keywords_sys_prompt ()
467+
468+ keywords_msg : List [Union [BaseMessage , dict [str , Any ]]] = []
469+ keywords_msg .append (SystemPromptMessage (content = system_prompt ))
470+
471+ # 加载上下文历史(仅取用户的历史提问,不含 AI 生成的 SQL 等回复)
472+ last_sql_messages : List [dict [str , Any ]] = self .generate_sql_logs [- 1 ].messages if len (
473+ self .generate_sql_logs ) > 0 else []
474+ if self .chat_question .regenerate_record_id :
475+ _temp_log = next (
476+ filter (lambda obj : obj .pid == self .chat_question .regenerate_record_id , self .generate_sql_logs ), None )
477+ last_sql_messages : List [dict [str , Any ]] = _temp_log .messages if _temp_log else []
478+
479+ # 排除系统提示词和 AI 回复,只保留用户的历史提问
480+ last_user_messages = [
481+ obj for obj in last_sql_messages
482+ if obj .get ("sqlbot_system" ) != True and obj .get ('type' ) == 'human'
483+ ]
484+
485+ if last_user_messages :
486+ last_rounds = get_last_conversation_rounds (last_user_messages , rounds = self .base_message_round_count_limit )
487+ for _msg_dict in last_rounds :
488+ content = _msg_dict .get ('content' , '' )
489+ # 提取 <user-question> 标签内的用户原始提问,排除 error-msg 等干扰信息
490+ match = re .search (r'<user-question>(.*?)</user-question>' , content , re .DOTALL )
491+ if match :
492+ question_text = match .group (1 ).strip ()
493+ if question_text :
494+ keywords_msg .append (HumanMessage (content = question_text ))
495+
496+ # 当前问题
497+ keywords_msg .append (HumanMessage (content = self .chat_question .question ))
498+
499+ self .current_logs [OperationEnum .EXTRACT_KEYWORDS ] = start_log (
500+ session = _session ,
501+ ai_modal_id = self .chat_question .ai_modal_id ,
502+ ai_modal_name = self .chat_question .ai_modal_name ,
503+ operate = OperationEnum .EXTRACT_KEYWORDS ,
504+ record_id = self .record .id ,
505+ full_message = [{'type' : msg .type , 'content' : msg .content } for msg in keywords_msg ]
506+ )
507+
508+ full_thinking_text = ''
509+ full_text = ''
510+ token_usage = {}
511+ res = process_stream (self .llm .stream (keywords_msg ), token_usage )
512+ for chunk in res :
513+ if chunk .get ('content' ):
514+ full_text += chunk .get ('content' )
515+ if chunk .get ('reasoning_content' ):
516+ full_thinking_text += chunk .get ('reasoning_content' )
517+
518+ result_text = full_text .strip ()
519+
520+ keywords_msg .append (AIMessage (result_text ))
521+ self .current_logs [OperationEnum .EXTRACT_KEYWORDS ] = end_log (
522+ session = _session ,
523+ log = self .current_logs [OperationEnum .EXTRACT_KEYWORDS ],
524+ full_message = [{'type' : msg .type , 'content' : msg .content } for msg in keywords_msg ],
525+ reasoning_content = full_thinking_text if full_thinking_text else None ,
526+ token_usage = token_usage
527+ )
528+
529+ if result_text .startswith ('ERROR:NOT_DATA_QUERY' ):
530+ # return True, '您的问题似乎不是查询数据的问题,请提出与数据查询相关的问题。'
531+ return False , ''
532+
533+ if result_text == 'EMPTY' or not result_text :
534+ self .chat_question .extracted_keywords = self .chat_question .question
535+ return False , ''
536+
537+ self .chat_question .extracted_keywords = result_text
538+ return False , result_text
539+
450540 def choose_table_schema (self , _session : Session ):
451541 self .current_logs [OperationEnum .CHOOSE_TABLE ] = start_log (session = _session ,
452542 operate = OperationEnum .CHOOSE_TABLE ,
453543 record_id = self .record .id ,
454544 local_operation = True )
545+
546+ # 使用扩展后的关键词进行表匹配(包含同义词,匹配更全面)
547+ keywords = self .chat_question .expanded_keywords or self .chat_question .question
548+
455549 self .chat_question .db_schema , tables = self .out_ds_instance .get_db_schema (
456550 self .ds .id , self .chat_question .question ) if self .out_ds_instance else get_table_schema (
457551 session = _session ,
458552 current_user = self .current_user ,
459553 ds = self .ds ,
460- question = self .chat_question .question )
554+ question = self .chat_question .question ,
555+ keywords = keywords )
461556
462557 # Get sample data for all tables
463558 if not self .out_ds_instance :
@@ -479,6 +574,12 @@ def generate_analysis(self, _session: Session):
479574 self .chat_question .data = orjson .dumps (data .get ('data' )).decode ()
480575 analysis_msg : List [Union [BaseMessage , dict [str , Any ]]] = []
481576
577+ # 从 record 加载已提取的关键词(SQL 生成阶段已持久化)
578+ if self .record .extracted_keywords :
579+ self .chat_question .extracted_keywords = self .record .extracted_keywords
580+ if self .record .expanded_keywords :
581+ self .chat_question .expanded_keywords = self .record .expanded_keywords
582+
482583 ds_id = self .ds .id if isinstance (self .ds , CoreDatasource ) else None
483584
484585 self .filter_terminology_template (_session , self .oid , ds_id )
@@ -781,6 +882,22 @@ def select_datasource(self, _session: Session):
781882 oid = self .ds .oid if isinstance (self .ds , CoreDatasource ) else 1
782883 ds_id = self .ds .id if isinstance (self .ds , CoreDatasource ) else None
783884
885+ # 提取关键词(前置步骤,结果复用于术语匹配和表匹配)
886+ if settings .TABLE_EMBEDDING_KEYWORD_ENABLED :
887+ is_error , kw_result = self .extract_keywords (_session )
888+ if is_error :
889+ raise SingleMessageError (kw_result )
890+ if kw_result :
891+ self .chat_question .extracted_keywords = kw_result
892+ SQLBotLogUtil .info (f"提取的关键词: { kw_result } " )
893+ # 术语同义词扩展(用于表匹配,不用于术语模板查询)
894+ expanded = expand_with_terminology (kw_result , _session , oid )
895+ self .chat_question .expanded_keywords = expanded
896+ SQLBotLogUtil .info (f"扩展后的关键词: { expanded } " )
897+ # 持久化到 record,供后续独立接口(如 generate_analysis)使用
898+ save_extracted_keywords (_session , self .record .id ,
899+ expanded_keywords = expanded , extracted_keywords = kw_result )
900+
784901 self .filter_terminology_template (_session , oid , ds_id )
785902
786903 self .filter_training_template (_session , oid , ds_id )
@@ -1247,6 +1364,22 @@ def run_task(self, in_chat: bool = True, stream: bool = True,
12471364 oid = self .ds .oid if isinstance (self .ds , CoreDatasource ) else 1
12481365 ds_id = self .ds .id if isinstance (self .ds , CoreDatasource ) else None
12491366
1367+ # 提取关键词(前置步骤,结果复用于术语匹配和表匹配)
1368+ if settings .TABLE_EMBEDDING_KEYWORD_ENABLED :
1369+ is_error , kw_result = self .extract_keywords (_session )
1370+ if is_error :
1371+ raise SingleMessageError (kw_result )
1372+ if kw_result :
1373+ self .chat_question .extracted_keywords = kw_result
1374+ SQLBotLogUtil .info (f"提取的关键词: { kw_result } " )
1375+ # 术语同义词扩展(用于表匹配,不用于术语模板查询)
1376+ expanded = expand_with_terminology (kw_result , _session , oid )
1377+ self .chat_question .expanded_keywords = expanded
1378+ SQLBotLogUtil .info (f"扩展后的关键词: { expanded } " )
1379+ # 持久化到 record,供后续独立接口(如 generate_analysis)使用
1380+ save_extracted_keywords (_session , self .record .id ,
1381+ expanded_keywords = expanded , extracted_keywords = kw_result )
1382+
12501383 self .filter_terminology_template (_session , oid , ds_id )
12511384
12521385 self .filter_training_template (_session , oid , ds_id )
0 commit comments