Skip to content

Commit 3fd6fa8

Browse files
committed
feat: implement LLM keyword extraction and fusion scoring for table matching
- Add extract_keywords() to LLMService with conversation context support, using process_stream + start_log/end_log for logging consistency - Add expand_with_terminology() to expand keywords with synonym matching - Add fusion scoring in calc_table_embedding(): α * vec_score + (1-α) * keyword_score - Add calc_keyword_score() with multi-level matching: table name exact (1.0), comment exact (0.9), comment substring (0.7), table name partial (0.25-0.5), field comment exact (0.5), field comment substring (0.4), field name partial (0.3) - Scoring considers keyword coverage rate to differentiate tables with different numbers of matched keywords - Extracted keywords persisted to ChatRecord for reuse across endpoints - get_terminology_template() supports comma-separated keyword input with per-keyword matching for higher precision - Add TABLE_EMBEDDING_ALPHA and TABLE_EMBEDDING_KEYWORD_ENABLED config
1 parent aac200a commit 3fd6fa8

15 files changed

Lines changed: 1200 additions & 29 deletions

File tree

‎backend/alembic/env.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,14 +24,14 @@
2424

2525
# from apps.system.models.user import SQLModel # noqa
2626
# from apps.settings.models.setting_models import SQLModel
27-
#from apps.chat.models.chat_model import SQLModel
28-
from apps.terminology.models.terminology_model import SQLModel
29-
from sqlbot_xpack.custom_prompt.models.custom_prompt_model import SQLModel
27+
from apps.chat.models.chat_model import SQLModel
28+
#from apps.terminology.models.terminology_model import SQLModel
29+
#from sqlbot_xpack.custom_prompt.models.custom_prompt_model import SQLModel
3030
#from apps.data_training.models.data_training_model import SQLModel
3131
# from apps.dashboard.models.dashboard_model import SQLModel
3232
from common.core.config import settings # noqa
3333
#from apps.datasource.models.datasource import SQLModel
34-
from apps.system.models.system_model import SQLModel
34+
#from apps.system.models.system_model import SQLModel
3535

3636
target_metadata = SQLModel.metadata
3737

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
"""072_chat_record
2+
3+
Revision ID: f898322341a5
4+
Revises: a2e2ecfa5a9c
5+
Create Date: 2026-09-14 15:54:47.921806
6+
7+
"""
8+
from alembic import op
9+
import sqlalchemy as sa
10+
import sqlmodel.sql.sqltypes
11+
from sqlalchemy.dialects import postgresql
12+
13+
# revision identifiers, used by Alembic.
14+
revision = 'f898322341a5'
15+
down_revision = 'a2e2ecfa5a9c'
16+
branch_labels = None
17+
depends_on = None
18+
19+
20+
def upgrade():
21+
# ### commands auto generated by Alembic - please adjust! ###
22+
op.add_column('chat_record', sa.Column('extracted_keywords', sa.Text(), nullable=True, comment='关键词'))
23+
op.add_column('chat_record', sa.Column('expanded_keywords', sa.Text(), nullable=True, comment='扩展后的关键词'))
24+
# ### end Alembic commands ###
25+
26+
27+
def downgrade():
28+
# ### commands auto generated by Alembic - please adjust! ###
29+
op.drop_column('chat_record', 'extracted_keywords')
30+
op.drop_column('chat_record', 'expanded_keywords')
31+
# ### end Alembic commands ###

‎backend/apps/chat/curd/chat.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -933,6 +933,19 @@ def save_sql_answer(session: SessionDep, record_id: int, answer: str) -> ChatRec
933933
return record
934934

935935

936+
def save_extracted_keywords(session: SessionDep, record_id: int, expanded_keywords: str,
937+
extracted_keywords: str = None) -> None:
938+
"""保存提取的关键词到 record,供后续独立接口复用。"""
939+
if not record_id:
940+
return
941+
values = {'expanded_keywords': expanded_keywords}
942+
if extracted_keywords:
943+
values['extracted_keywords'] = extracted_keywords
944+
stmt = update(ChatRecord).where(and_(ChatRecord.id == record_id)).values(**values)
945+
session.execute(stmt)
946+
session.commit()
947+
948+
936949
def save_analysis_answer(session: SessionDep, record_id: int, answer: str = '') -> ChatRecord:
937950
if not record_id:
938951
raise Exception("Record id cannot be None")

‎backend/apps/chat/models/chat_model.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ class OperationEnum(Enum):
4747
FILTER_CUSTOM_PROMPT = '11'
4848
EXECUTE_SQL = '12'
4949
GENERATE_PICTURE = '13'
50+
EXTRACT_KEYWORDS = '14'
5051

5152

5253
class ChatFinishStep(Enum):
@@ -129,6 +130,8 @@ class ChatRecord(SQLModel, table=True):
129130
analysis_record_id: int = Field(sa_column=Column(BigInteger, nullable=True))
130131
predict_record_id: int = Field(sa_column=Column(BigInteger, nullable=True))
131132
regenerate_record_id: int = Field(sa_column=Column(BigInteger, nullable=True))
133+
extracted_keywords: str = Field(sa_column=Column(Text, nullable=True))
134+
expanded_keywords: str = Field(sa_column=Column(Text, nullable=True))
132135

133136

134137
class ChatRecordResult(BaseModel):
@@ -238,6 +241,8 @@ class AiModelQuestion(BaseModel):
238241
question: str = None
239242
ai_modal_id: int = None
240243
ai_modal_name: str = None # Specific model name
244+
extracted_keywords: str = "" # LLM 提取的核心业务关键词(逗号分隔,未经扩展)
245+
expanded_keywords: str = "" # 经术语同义词扩展后的关键词(用于表匹配)
241246
engine: str = ""
242247
db_schema: str = ""
243248
sql: str = ""
@@ -255,6 +260,10 @@ class AiModelQuestion(BaseModel):
255260
sample_data: str = ""
256261
sqlbot_name: str = "SQLBot"
257262

263+
def extract_keywords_sys_prompt(self) -> str:
264+
return get_sql_template().get("extract_keywords").format(lang=self.lang,
265+
sqlbot_name=self.sqlbot_name)
266+
258267
def sql_sys_question(self, db_type: Union[str, DB], enable_query_limit: bool = True):
259268
templates: dict[str, str] = {}
260269
_sql_template = get_sql_example_template(db_type)

‎backend/apps/chat/task/llm.py‎

Lines changed: 137 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import concurrent
33
import json
44
import os
5+
import re
56
import traceback
67
import urllib.parse
78
import warnings
@@ -30,7 +31,7 @@
3031
from 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, \
@@ -48,6 +49,7 @@
4849
from apps.system.crud.parameter_manage import get_groups
4950
from apps.system.crud.user import user_ws_list
5051
from apps.system.schemas.system_schema import AssistantOutDsSchema
52+
from apps.terminology.curd.terminology import expand_with_terminology
5153
from apps.terminology.curd.terminology import get_terminology_template
5254
from common.core.config import settings
5355
from 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)

‎backend/apps/datasource/crud/datasource.py‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -517,7 +517,8 @@ def get_tables_sample_data(session: SessionDep, current_user: CurrentUser, ds: C
517517

518518

519519
def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource, question: str,
520-
embedding: bool = True, table_list: list[str] = None) -> tuple[str, list]:
520+
embedding: bool = True, table_list: list[str] = None,
521+
keywords: str = None) -> tuple[str, list]:
521522
schema_str = ""
522523
table_objs = get_table_obj_by_ds(session=session, current_user=current_user, ds=ds)
523524
if len(table_objs) == 0:
@@ -538,17 +539,25 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
538539
table_comment = ''
539540
if obj.table.custom_comment:
540541
table_comment = obj.table.custom_comment.strip()
542+
if not table_comment and obj.table.table_comment:
543+
table_comment = obj.table.table_comment.strip()
541544
if table_comment == '':
542545
schema_table += '\n[\n'
543546
else:
544547
schema_table += f", {table_comment}\n[\n"
545548

549+
fields_info = []
546550
if obj.fields:
547551
field_list = []
548552
for field in obj.fields:
549553
field_comment = ''
550554
if field.custom_comment:
551555
field_comment = field.custom_comment.strip()
556+
fields_info.append({
557+
"field_name": field.field_name,
558+
"field_comment": (field.field_comment or '').strip(),
559+
"custom_comment": (field.custom_comment or '').strip()
560+
})
552561
if field_comment == '':
553562
field_list.append(f"({field.field_name}:{field.field_type})")
554563
else:
@@ -557,17 +566,22 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
557566
schema_table += '\n]\n'
558567

559568
t_obj = {"id": obj.table.id, "table_name": obj.table.table_name, "schema_table": schema_table,
560-
"embedding": obj.table.embedding}
569+
"embedding": obj.table.embedding, "table_comment": table_comment, "fields": fields_info}
561570
tables.append(t_obj)
562571
all_tables.append(t_obj)
563572

564573
# 如果没有符合过滤条件的表,直接返回
565574
if not tables:
566575
return schema_str, []
567576

568-
# do table embedding
577+
# 执行表 embedding 匹配
569578
if embedding and tables and settings.TABLE_EMBEDDING_ENABLED:
570-
tables = calc_table_embedding(tables, question)
579+
try:
580+
tables = calc_table_embedding(tables, question, session=session, oid=current_user.oid,
581+
keywords=keywords)
582+
except ValueError as e:
583+
# LLM 判断意图不是查数据或关键词提取失败
584+
return str(e), []
571585
# splice schema
572586
if tables:
573587
for s in tables:

0 commit comments

Comments
 (0)