mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 17:26:55 +08:00
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
This commit is contained in:
@@ -933,6 +933,19 @@ def save_sql_answer(session: SessionDep, record_id: int, answer: str) -> ChatRec
|
||||
return record
|
||||
|
||||
|
||||
def save_extracted_keywords(session: SessionDep, record_id: int, expanded_keywords: str,
|
||||
extracted_keywords: str = None) -> None:
|
||||
"""保存提取的关键词到 record,供后续独立接口复用。"""
|
||||
if not record_id:
|
||||
return
|
||||
values = {'expanded_keywords': expanded_keywords}
|
||||
if extracted_keywords:
|
||||
values['extracted_keywords'] = extracted_keywords
|
||||
stmt = update(ChatRecord).where(and_(ChatRecord.id == record_id)).values(**values)
|
||||
session.execute(stmt)
|
||||
session.commit()
|
||||
|
||||
|
||||
def save_analysis_answer(session: SessionDep, record_id: int, answer: str = '') -> ChatRecord:
|
||||
if not record_id:
|
||||
raise Exception("Record id cannot be None")
|
||||
|
||||
@@ -47,6 +47,7 @@ class OperationEnum(Enum):
|
||||
FILTER_CUSTOM_PROMPT = '11'
|
||||
EXECUTE_SQL = '12'
|
||||
GENERATE_PICTURE = '13'
|
||||
EXTRACT_KEYWORDS = '14'
|
||||
|
||||
|
||||
class ChatFinishStep(Enum):
|
||||
@@ -129,6 +130,8 @@ class ChatRecord(SQLModel, table=True):
|
||||
analysis_record_id: int = Field(sa_column=Column(BigInteger, nullable=True))
|
||||
predict_record_id: int = Field(sa_column=Column(BigInteger, nullable=True))
|
||||
regenerate_record_id: int = Field(sa_column=Column(BigInteger, nullable=True))
|
||||
extracted_keywords: str = Field(sa_column=Column(Text, nullable=True))
|
||||
expanded_keywords: str = Field(sa_column=Column(Text, nullable=True))
|
||||
|
||||
|
||||
class ChatRecordResult(BaseModel):
|
||||
@@ -238,6 +241,8 @@ class AiModelQuestion(BaseModel):
|
||||
question: str = None
|
||||
ai_modal_id: int = None
|
||||
ai_modal_name: str = None # Specific model name
|
||||
extracted_keywords: str = "" # LLM 提取的核心业务关键词(逗号分隔,未经扩展)
|
||||
expanded_keywords: str = "" # 经术语同义词扩展后的关键词(用于表匹配)
|
||||
engine: str = ""
|
||||
db_schema: str = ""
|
||||
sql: str = ""
|
||||
@@ -255,6 +260,10 @@ class AiModelQuestion(BaseModel):
|
||||
sample_data: str = ""
|
||||
sqlbot_name: str = "SQLBot"
|
||||
|
||||
def extract_keywords_sys_prompt(self) -> str:
|
||||
return get_sql_template().get("extract_keywords").format(lang=self.lang,
|
||||
sqlbot_name=self.sqlbot_name)
|
||||
|
||||
def sql_sys_question(self, db_type: Union[str, DB], enable_query_limit: bool = True):
|
||||
templates: dict[str, str] = {}
|
||||
_sql_template = get_sql_example_template(db_type)
|
||||
|
||||
@@ -2,6 +2,7 @@ import asyncio
|
||||
import concurrent
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import traceback
|
||||
import urllib.parse
|
||||
import warnings
|
||||
@@ -30,7 +31,7 @@ from apps.ai_model.model_factory import LLMConfig, LLMFactory, get_default_confi
|
||||
from apps.chat.curd.chat import save_question, save_sql_answer, save_sql, \
|
||||
save_error_message, save_sql_exec_data, save_chart_answer, save_chart, \
|
||||
finish_record, save_analysis_answer, save_predict_answer, save_predict_data, \
|
||||
save_select_datasource_answer, save_recommend_question_answer, \
|
||||
save_select_datasource_answer, save_recommend_question_answer, save_extracted_keywords, \
|
||||
get_old_questions, save_analysis_predict_record, rename_chat, get_chart_config, \
|
||||
get_chat_chart_data, list_generate_sql_logs, list_generate_chart_logs, start_log, end_log, \
|
||||
get_last_execute_sql_error, format_json_data, format_chart_fields, get_chat_brief_generate, get_chat_predict_data, \
|
||||
@@ -48,6 +49,7 @@ from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, ge
|
||||
from apps.system.crud.parameter_manage import get_groups
|
||||
from apps.system.crud.user import user_ws_list
|
||||
from apps.system.schemas.system_schema import AssistantOutDsSchema
|
||||
from apps.terminology.curd.terminology import expand_with_terminology
|
||||
from apps.terminology.curd.terminology import get_terminology_template
|
||||
from common.core.config import settings
|
||||
from common.core.db import engine
|
||||
@@ -378,14 +380,16 @@ class LLMService:
|
||||
calculate_oid = self.current_assistant.oid if self.current_assistant.type != 4 else self.oid
|
||||
if self.current_assistant.type == 1:
|
||||
calculate_ds_id = None
|
||||
# 使用原始提取的关键词(未扩展)进行术语匹配,避免同义词重复查询
|
||||
match_text = self.chat_question.extracted_keywords or self.chat_question.question
|
||||
if self.current_assistant and self.current_assistant.type == 1:
|
||||
self.chat_question.terminologies, term_list = get_terminology_template(_session,
|
||||
self.chat_question.question,
|
||||
match_text,
|
||||
calculate_oid,
|
||||
None, self.current_assistant.id)
|
||||
else:
|
||||
self.chat_question.terminologies, term_list = get_terminology_template(_session,
|
||||
self.chat_question.question,
|
||||
match_text,
|
||||
calculate_oid,
|
||||
calculate_ds_id)
|
||||
|
||||
@@ -447,17 +451,108 @@ class LLMService:
|
||||
OperationEnum.FILTER_SQL_EXAMPLE],
|
||||
full_message=example_list)
|
||||
|
||||
def extract_keywords(self, _session: Session) -> tuple[bool, str]:
|
||||
"""从用户问题中提取核心业务实体词。
|
||||
|
||||
使用 LLM 提取名词性业务实体,忽略时间词、动作词、修饰词。
|
||||
支持上下文历史,可根据对话上下文理解用户意图。
|
||||
提取结果可复用于表匹配、术语扩展等环节。
|
||||
|
||||
Returns:
|
||||
(is_error, result) 元组:
|
||||
- (False, keywords) 成功,关键词逗号分隔
|
||||
- (True, error_message) 错误(LLM 失败或非查数据意图)
|
||||
"""
|
||||
system_prompt = self.chat_question.extract_keywords_sys_prompt()
|
||||
|
||||
keywords_msg: List[Union[BaseMessage, dict[str, Any]]] = []
|
||||
keywords_msg.append(SystemPromptMessage(content=system_prompt))
|
||||
|
||||
# 加载上下文历史(仅取用户的历史提问,不含 AI 生成的 SQL 等回复)
|
||||
last_sql_messages: List[dict[str, Any]] = self.generate_sql_logs[-1].messages if len(
|
||||
self.generate_sql_logs) > 0 else []
|
||||
if self.chat_question.regenerate_record_id:
|
||||
_temp_log = next(
|
||||
filter(lambda obj: obj.pid == self.chat_question.regenerate_record_id, self.generate_sql_logs), None)
|
||||
last_sql_messages: List[dict[str, Any]] = _temp_log.messages if _temp_log else []
|
||||
|
||||
# 排除系统提示词和 AI 回复,只保留用户的历史提问
|
||||
last_user_messages = [
|
||||
obj for obj in last_sql_messages
|
||||
if obj.get("sqlbot_system") != True and obj.get('type') == 'human'
|
||||
]
|
||||
|
||||
if last_user_messages:
|
||||
last_rounds = get_last_conversation_rounds(last_user_messages, rounds=self.base_message_round_count_limit)
|
||||
for _msg_dict in last_rounds:
|
||||
content = _msg_dict.get('content', '')
|
||||
# 提取 <user-question> 标签内的用户原始提问,排除 error-msg 等干扰信息
|
||||
match = re.search(r'<user-question>(.*?)</user-question>', content, re.DOTALL)
|
||||
if match:
|
||||
question_text = match.group(1).strip()
|
||||
if question_text:
|
||||
keywords_msg.append(HumanMessage(content=question_text))
|
||||
|
||||
# 当前问题
|
||||
keywords_msg.append(HumanMessage(content=self.chat_question.question))
|
||||
|
||||
self.current_logs[OperationEnum.EXTRACT_KEYWORDS] = start_log(
|
||||
session=_session,
|
||||
ai_modal_id=self.chat_question.ai_modal_id,
|
||||
ai_modal_name=self.chat_question.ai_modal_name,
|
||||
operate=OperationEnum.EXTRACT_KEYWORDS,
|
||||
record_id=self.record.id,
|
||||
full_message=[{'type': msg.type, 'content': msg.content} for msg in keywords_msg]
|
||||
)
|
||||
|
||||
full_thinking_text = ''
|
||||
full_text = ''
|
||||
token_usage = {}
|
||||
res = process_stream(self.llm.stream(keywords_msg), token_usage)
|
||||
for chunk in res:
|
||||
if chunk.get('content'):
|
||||
full_text += chunk.get('content')
|
||||
if chunk.get('reasoning_content'):
|
||||
full_thinking_text += chunk.get('reasoning_content')
|
||||
|
||||
result_text = full_text.strip()
|
||||
|
||||
keywords_msg.append(AIMessage(result_text))
|
||||
self.current_logs[OperationEnum.EXTRACT_KEYWORDS] = end_log(
|
||||
session=_session,
|
||||
log=self.current_logs[OperationEnum.EXTRACT_KEYWORDS],
|
||||
full_message=[{'type': msg.type, 'content': msg.content} for msg in keywords_msg],
|
||||
reasoning_content=full_thinking_text if full_thinking_text else None,
|
||||
token_usage=token_usage
|
||||
)
|
||||
|
||||
if result_text.startswith('ERROR:NOT_DATA_QUERY'):
|
||||
# return True, '您的问题似乎不是查询数据的问题,请提出与数据查询相关的问题。'
|
||||
return False, ''
|
||||
|
||||
if result_text == 'EMPTY' or not result_text:
|
||||
self.chat_question.extracted_keywords = self.chat_question.question
|
||||
return False, ''
|
||||
|
||||
self.chat_question.extracted_keywords = result_text
|
||||
return False, result_text
|
||||
|
||||
def choose_table_schema(self, _session: Session):
|
||||
self.current_logs[OperationEnum.CHOOSE_TABLE] = start_log(session=_session,
|
||||
operate=OperationEnum.CHOOSE_TABLE,
|
||||
record_id=self.record.id,
|
||||
local_operation=True)
|
||||
|
||||
# 使用扩展后的关键词进行表匹配(包含同义词,匹配更全面)
|
||||
keywords = self.chat_question.expanded_keywords or self.chat_question.question
|
||||
|
||||
self.chat_question.db_schema, tables = self.out_ds_instance.get_db_schema(
|
||||
self.ds.id, self.chat_question.question) if self.out_ds_instance else get_table_schema(
|
||||
session=_session,
|
||||
current_user=self.current_user,
|
||||
ds=self.ds,
|
||||
question=self.chat_question.question)
|
||||
question=self.chat_question.question,
|
||||
keywords=keywords)
|
||||
|
||||
# Get sample data for all tables
|
||||
if not self.out_ds_instance:
|
||||
@@ -479,6 +574,12 @@ class LLMService:
|
||||
self.chat_question.data = orjson.dumps(data.get('data')).decode()
|
||||
analysis_msg: List[Union[BaseMessage, dict[str, Any]]] = []
|
||||
|
||||
# 从 record 加载已提取的关键词(SQL 生成阶段已持久化)
|
||||
if self.record.extracted_keywords:
|
||||
self.chat_question.extracted_keywords = self.record.extracted_keywords
|
||||
if self.record.expanded_keywords:
|
||||
self.chat_question.expanded_keywords = self.record.expanded_keywords
|
||||
|
||||
ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None
|
||||
|
||||
self.filter_terminology_template(_session, self.oid, ds_id)
|
||||
@@ -781,6 +882,22 @@ class LLMService:
|
||||
oid = self.ds.oid if isinstance(self.ds, CoreDatasource) else 1
|
||||
ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None
|
||||
|
||||
# 提取关键词(前置步骤,结果复用于术语匹配和表匹配)
|
||||
if settings.TABLE_EMBEDDING_KEYWORD_ENABLED:
|
||||
is_error, kw_result = self.extract_keywords(_session)
|
||||
if is_error:
|
||||
raise SingleMessageError(kw_result)
|
||||
if kw_result:
|
||||
self.chat_question.extracted_keywords = kw_result
|
||||
SQLBotLogUtil.info(f"提取的关键词: {kw_result}")
|
||||
# 术语同义词扩展(用于表匹配,不用于术语模板查询)
|
||||
expanded = expand_with_terminology(kw_result, _session, oid)
|
||||
self.chat_question.expanded_keywords = expanded
|
||||
SQLBotLogUtil.info(f"扩展后的关键词: {expanded}")
|
||||
# 持久化到 record,供后续独立接口(如 generate_analysis)使用
|
||||
save_extracted_keywords(_session, self.record.id,
|
||||
expanded_keywords=expanded, extracted_keywords=kw_result)
|
||||
|
||||
self.filter_terminology_template(_session, oid, ds_id)
|
||||
|
||||
self.filter_training_template(_session, oid, ds_id)
|
||||
@@ -1247,6 +1364,22 @@ class LLMService:
|
||||
oid = self.ds.oid if isinstance(self.ds, CoreDatasource) else 1
|
||||
ds_id = self.ds.id if isinstance(self.ds, CoreDatasource) else None
|
||||
|
||||
# 提取关键词(前置步骤,结果复用于术语匹配和表匹配)
|
||||
if settings.TABLE_EMBEDDING_KEYWORD_ENABLED:
|
||||
is_error, kw_result = self.extract_keywords(_session)
|
||||
if is_error:
|
||||
raise SingleMessageError(kw_result)
|
||||
if kw_result:
|
||||
self.chat_question.extracted_keywords = kw_result
|
||||
SQLBotLogUtil.info(f"提取的关键词: {kw_result}")
|
||||
# 术语同义词扩展(用于表匹配,不用于术语模板查询)
|
||||
expanded = expand_with_terminology(kw_result, _session, oid)
|
||||
self.chat_question.expanded_keywords = expanded
|
||||
SQLBotLogUtil.info(f"扩展后的关键词: {expanded}")
|
||||
# 持久化到 record,供后续独立接口(如 generate_analysis)使用
|
||||
save_extracted_keywords(_session, self.record.id,
|
||||
expanded_keywords=expanded, extracted_keywords=kw_result)
|
||||
|
||||
self.filter_terminology_template(_session, oid, ds_id)
|
||||
|
||||
self.filter_training_template(_session, oid, ds_id)
|
||||
|
||||
@@ -517,7 +517,8 @@ def get_tables_sample_data(session: SessionDep, current_user: CurrentUser, ds: C
|
||||
|
||||
|
||||
def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource, question: str,
|
||||
embedding: bool = True, table_list: list[str] = None) -> tuple[str, list]:
|
||||
embedding: bool = True, table_list: list[str] = None,
|
||||
keywords: str = None) -> tuple[str, list]:
|
||||
schema_str = ""
|
||||
table_objs = get_table_obj_by_ds(session=session, current_user=current_user, ds=ds)
|
||||
if len(table_objs) == 0:
|
||||
@@ -538,17 +539,25 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
|
||||
table_comment = ''
|
||||
if obj.table.custom_comment:
|
||||
table_comment = obj.table.custom_comment.strip()
|
||||
if not table_comment and obj.table.table_comment:
|
||||
table_comment = obj.table.table_comment.strip()
|
||||
if table_comment == '':
|
||||
schema_table += '\n[\n'
|
||||
else:
|
||||
schema_table += f", {table_comment}\n[\n"
|
||||
|
||||
fields_info = []
|
||||
if obj.fields:
|
||||
field_list = []
|
||||
for field in obj.fields:
|
||||
field_comment = ''
|
||||
if field.custom_comment:
|
||||
field_comment = field.custom_comment.strip()
|
||||
fields_info.append({
|
||||
"field_name": field.field_name,
|
||||
"field_comment": (field.field_comment or '').strip(),
|
||||
"custom_comment": (field.custom_comment or '').strip()
|
||||
})
|
||||
if field_comment == '':
|
||||
field_list.append(f"({field.field_name}:{field.field_type})")
|
||||
else:
|
||||
@@ -557,7 +566,7 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
|
||||
schema_table += '\n]\n'
|
||||
|
||||
t_obj = {"id": obj.table.id, "table_name": obj.table.table_name, "schema_table": schema_table,
|
||||
"embedding": obj.table.embedding}
|
||||
"embedding": obj.table.embedding, "table_comment": table_comment, "fields": fields_info}
|
||||
tables.append(t_obj)
|
||||
all_tables.append(t_obj)
|
||||
|
||||
@@ -565,9 +574,14 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
|
||||
if not tables:
|
||||
return schema_str, []
|
||||
|
||||
# do table embedding
|
||||
# 执行表 embedding 匹配
|
||||
if embedding and tables and settings.TABLE_EMBEDDING_ENABLED:
|
||||
tables = calc_table_embedding(tables, question)
|
||||
try:
|
||||
tables = calc_table_embedding(tables, question, session=session, oid=current_user.oid,
|
||||
keywords=keywords)
|
||||
except ValueError as e:
|
||||
# LLM 判断意图不是查数据或关键词提取失败
|
||||
return str(e), []
|
||||
# splice schema
|
||||
if tables:
|
||||
for s in tables:
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# Author: Junjun
|
||||
# Date: 2025/9/23
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
import traceback
|
||||
|
||||
@@ -10,6 +11,15 @@ from common.core.config import settings
|
||||
from common.utils.utils import SQLBotLogUtil
|
||||
|
||||
|
||||
def _parse_embedding(embedding):
|
||||
"""解析 embedding,兼容 JSON 字符串和原生列表两种格式。"""
|
||||
if isinstance(embedding, list):
|
||||
return embedding
|
||||
if isinstance(embedding, str):
|
||||
return json.loads(embedding)
|
||||
return None
|
||||
|
||||
|
||||
def get_table_embedding(tables: list[dict], question: str):
|
||||
_list = []
|
||||
for table in tables:
|
||||
@@ -40,38 +50,244 @@ def get_table_embedding(tables: list[dict], question: str):
|
||||
return _list
|
||||
|
||||
|
||||
def calc_table_embedding(tables: list[dict], question: str):
|
||||
def calc_keyword_score(
|
||||
question: str,
|
||||
table_name: str,
|
||||
table_comment: str,
|
||||
fields: list[dict]
|
||||
) -> float:
|
||||
"""计算关键词与表的匹配分数。
|
||||
|
||||
综合考虑所有关键词的匹配情况:
|
||||
- 表名完整匹配:1.0
|
||||
- 表注释完整匹配:0.9
|
||||
- 部分匹配时,按匹配质量 * 关键词覆盖率计算
|
||||
|
||||
Args:
|
||||
question: 提取后的关键词(逗号分隔)或原始问题
|
||||
table_name: 表名
|
||||
table_comment: 表注释(已合并 custom_comment 优先)
|
||||
fields: 字段列表,每个 dict 包含 field_name, field_comment, custom_comment
|
||||
|
||||
Returns:
|
||||
0.0~1.0 的匹配分数
|
||||
"""
|
||||
if not question:
|
||||
return 0.0
|
||||
|
||||
keywords = [kw.strip().lower() for kw in question.split(',') if kw.strip()]
|
||||
if not keywords:
|
||||
return 0.0
|
||||
|
||||
tn = table_name.lower()
|
||||
tc = (table_comment or '').lower()
|
||||
total_kw = len(keywords)
|
||||
|
||||
# 收集每个关键词的最佳匹配分数
|
||||
kw_scores = {} # keyword -> best_score
|
||||
|
||||
for kw in keywords:
|
||||
# 1. 表名完整匹配:1.0
|
||||
if kw == tn:
|
||||
kw_scores[kw] = 1.0
|
||||
continue
|
||||
|
||||
# 2. 表注释完整匹配:0.9
|
||||
if tc and kw == tc:
|
||||
kw_scores[kw] = max(kw_scores.get(kw, 0), 0.9)
|
||||
continue
|
||||
|
||||
# 3. 表注释子串匹配:0.7(关键词是注释的子串,如 "用户" 匹配 "用户表")
|
||||
if tc and kw in tc:
|
||||
kw_scores[kw] = max(kw_scores.get(kw, 0), 0.7)
|
||||
continue
|
||||
|
||||
# 4. 表名部分匹配:收集匹配的关键词
|
||||
tn_parts = set(re.split(r'[_\s]+', tn))
|
||||
tn_parts.discard('')
|
||||
if kw in tn_parts:
|
||||
kw_scores[kw] = max(kw_scores.get(kw, 0), 0) # 标记匹配,分数稍后计算
|
||||
|
||||
# 4. 字段注释匹配:0.5(精确)/ 0.4(子串)
|
||||
for field in fields:
|
||||
field_comment = ''
|
||||
if field.get('custom_comment'):
|
||||
field_comment = field['custom_comment'].strip().lower()
|
||||
if not field_comment and field.get('field_comment'):
|
||||
field_comment = field['field_comment'].strip().lower()
|
||||
if field_comment:
|
||||
if kw == field_comment:
|
||||
kw_scores[kw] = max(kw_scores.get(kw, 0), 0.5)
|
||||
break
|
||||
elif kw in field_comment:
|
||||
kw_scores[kw] = max(kw_scores.get(kw, 0), 0.4)
|
||||
|
||||
# 5. 字段名匹配:0.3
|
||||
if kw not in kw_scores or kw_scores[kw] < 0.3:
|
||||
for field in fields:
|
||||
fname = field.get('field_name', '').lower()
|
||||
if fname:
|
||||
fname_parts = set(re.split(r'[_\s]+', fname))
|
||||
fname_parts.discard('')
|
||||
if kw in fname_parts:
|
||||
kw_scores[kw] = max(kw_scores.get(kw, 0), 0.3)
|
||||
break
|
||||
|
||||
if not kw_scores:
|
||||
return 0.0
|
||||
|
||||
# 表名/注释完整匹配 → 直接返回 1.0 或 0.9
|
||||
if kw_scores.get(tn) == 1.0:
|
||||
return 1.0
|
||||
if tc and kw_scores.get(tc) == 0.9:
|
||||
return 0.9
|
||||
|
||||
# 表名部分匹配:按覆盖率计算(多个关键词共同覆盖表名)
|
||||
tn_parts = set(re.split(r'[_\s]+', tn))
|
||||
tn_parts.discard('')
|
||||
if tn_parts:
|
||||
matched_parts = {kw for kw in keywords if kw in tn_parts}
|
||||
if matched_parts:
|
||||
coverage = len(matched_parts) / len(tn_parts)
|
||||
tn_score = round(0.25 + 0.25 * coverage, 2)
|
||||
# 更新这些关键词的分数为表名匹配分数
|
||||
for kw in matched_parts:
|
||||
kw_scores[kw] = max(kw_scores.get(kw, 0), tn_score)
|
||||
|
||||
# 综合评分:最佳匹配分数 * 关键词覆盖率
|
||||
best_score = max(kw_scores.values())
|
||||
match_rate = len(kw_scores) / total_kw
|
||||
return round(best_score * match_rate, 2)
|
||||
|
||||
|
||||
def calc_table_embedding(tables: list[dict], question: str, session=None, oid: int = None,
|
||||
keywords: str = None):
|
||||
"""使用向量相似度和关键词匹配的融合评分计算表的相关性分数。
|
||||
|
||||
当 keywords 不为空时(已由调用方提取并扩展):
|
||||
1. 用 keywords 计算向量相似度(vec_score)
|
||||
2. 计算关键词匹配度(keyword_score)
|
||||
3. 融合评分:final_score = α * vec_score + (1-α) * keyword_score
|
||||
4. keyword_score=1.0 的表保证排在最前面
|
||||
|
||||
当 keywords 为空或 TABLE_EMBEDDING_KEYWORD_ENABLED 为 False 时,回退到纯向量匹配。
|
||||
|
||||
Args:
|
||||
tables: 表列表,每个 dict 包含 id, table_name, schema_table, embedding, table_comment, fields
|
||||
question: 原始用户问题(用于纯向量匹配回退)
|
||||
session: 数据库会话(预留)
|
||||
oid: 组织 ID(预留)
|
||||
keywords: 已提取并扩展的关键词(由调用方传入)
|
||||
"""
|
||||
if not keywords or not settings.TABLE_EMBEDDING_KEYWORD_ENABLED:
|
||||
# 回退到纯向量匹配
|
||||
return _calc_vector_only(tables, question)
|
||||
|
||||
_list = []
|
||||
for table in tables:
|
||||
_list.append(
|
||||
{"id": table.get('id'), "schema_table": table.get('schema_table'), "embedding": table.get('embedding'),
|
||||
"cosine_similarity": 0.0, "table_name": table.get('table_name')})
|
||||
_list.append({
|
||||
"id": table.get('id'),
|
||||
"schema_table": table.get('schema_table'),
|
||||
"embedding": table.get('embedding'),
|
||||
"cosine_similarity": 0.0,
|
||||
"table_name": table.get('table_name'),
|
||||
"table_comment": table.get('table_comment', ''),
|
||||
"fields": table.get('fields', [])
|
||||
})
|
||||
|
||||
if not _list:
|
||||
return _list
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
|
||||
# 步骤 1:用关键词计算向量相似度
|
||||
model = EmbeddingModelCache.get_model()
|
||||
results = [item.get('embedding') for item in _list]
|
||||
|
||||
q_embedding = model.embed_query(keywords)
|
||||
for index in range(len(results)):
|
||||
item = results[index]
|
||||
if item:
|
||||
_list[index]['cosine_similarity'] = cosine_similarity(q_embedding, _parse_embedding(item))
|
||||
|
||||
# 步骤 2 & 3:计算关键词分数并融合
|
||||
alpha = settings.TABLE_EMBEDDING_ALPHA
|
||||
for table in _list:
|
||||
vec_score = table['cosine_similarity']
|
||||
kw_score = calc_keyword_score(
|
||||
keywords,
|
||||
table.get('table_name', ''),
|
||||
table.get('table_comment', ''),
|
||||
table.get('fields', [])
|
||||
)
|
||||
table['keyword_score'] = kw_score
|
||||
table['cosine_similarity'] = alpha * vec_score + (1 - alpha) * kw_score
|
||||
|
||||
# 步骤 4:排序 - 精确匹配(keyword_score=1.0)始终排最前
|
||||
exact_matches = [t for t in _list if t.get('keyword_score') == 1.0]
|
||||
other_tables = [t for t in _list if t.get('keyword_score') != 1.0]
|
||||
|
||||
exact_matches.sort(key=lambda x: x['cosine_similarity'], reverse=True)
|
||||
other_tables.sort(key=lambda x: x['cosine_similarity'], reverse=True)
|
||||
|
||||
_list = exact_matches + other_tables
|
||||
_list = _list[:settings.TABLE_EMBEDDING_COUNT]
|
||||
|
||||
end_time = time.time()
|
||||
SQLBotLogUtil.info(f"融合评分耗时 {end_time - start_time:.3f}s")
|
||||
SQLBotLogUtil.info(json.dumps([{
|
||||
"id": ele.get('id'),
|
||||
"schema_table": ele.get('schema_table'),
|
||||
"cosine_similarity": ele.get('cosine_similarity'),
|
||||
"keyword_score": ele.get('keyword_score'),
|
||||
"table_name": ele.get('table_name')
|
||||
} for ele in _list]))
|
||||
|
||||
return _list
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
# 异常回退到纯向量匹配
|
||||
return _calc_vector_only(tables, question)
|
||||
|
||||
|
||||
def _calc_vector_only(tables: list[dict], question: str):
|
||||
"""纯向量相似度评分(原始逻辑)。"""
|
||||
_list = []
|
||||
for table in tables:
|
||||
_list.append({
|
||||
"id": table.get('id'),
|
||||
"schema_table": table.get('schema_table'),
|
||||
"embedding": table.get('embedding'),
|
||||
"cosine_similarity": 0.0,
|
||||
"table_name": table.get('table_name'),
|
||||
"table_comment": table.get('table_comment', ''),
|
||||
"fields": table.get('fields', [])
|
||||
})
|
||||
|
||||
if _list:
|
||||
try:
|
||||
# text = [s.get('schema_table') for s in _list]
|
||||
#
|
||||
model = EmbeddingModelCache.get_model()
|
||||
start_time = time.time()
|
||||
# results = model.embed_documents(text)
|
||||
# end_time = time.time()
|
||||
# SQLBotLogUtil.info(str(end_time - start_time))
|
||||
results = [item.get('embedding') for item in _list]
|
||||
|
||||
q_embedding = model.embed_query(question)
|
||||
for index in range(len(results)):
|
||||
item = results[index]
|
||||
if item:
|
||||
_list[index]['cosine_similarity'] = cosine_similarity(q_embedding, json.loads(item))
|
||||
_list[index]['cosine_similarity'] = cosine_similarity(q_embedding, _parse_embedding(item))
|
||||
|
||||
_list.sort(key=lambda x: x['cosine_similarity'], reverse=True)
|
||||
_list = _list[:settings.TABLE_EMBEDDING_COUNT]
|
||||
# print(len(_list))
|
||||
|
||||
end_time = time.time()
|
||||
SQLBotLogUtil.info(str(end_time - start_time))
|
||||
SQLBotLogUtil.info(json.dumps([{"id": ele.get('id'), "schema_table": ele.get('schema_table'),
|
||||
"cosine_similarity": ele.get('cosine_similarity'), "table_name": ele.get('table_name')}
|
||||
for ele in _list]))
|
||||
SQLBotLogUtil.info(json.dumps([{
|
||||
"id": ele.get('id'),
|
||||
"schema_table": ele.get('schema_table'),
|
||||
"cosine_similarity": ele.get('cosine_similarity'),
|
||||
"table_name": ele.get('table_name')
|
||||
} for ele in _list]))
|
||||
return _list
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
|
||||
@@ -977,10 +977,122 @@ def get_terminology_template(session: SessionDep, question: str, oid: Optional[i
|
||||
advanced_application_id: Optional[int] = None) -> tuple[str, list[dict]]:
|
||||
if not oid:
|
||||
oid = 1
|
||||
_results = select_terminology_by_word(session, question, oid, datasource, advanced_application_id)
|
||||
if _results and len(_results) > 0:
|
||||
|
||||
# 有逗号时视为关键词列表,逐个匹配(精度更高,避免整体 ILIKE 误匹配)
|
||||
# 无逗号时视为自然语言问题,整体匹配
|
||||
if ',' in question:
|
||||
parts = [kw.strip() for kw in question.split(',') if kw.strip()]
|
||||
else:
|
||||
parts = [question]
|
||||
|
||||
# select_terminology_by_word 返回 list[dict],每个 dict 结构:
|
||||
# {'words': ['销售', 'sales'], 'description': '...'}
|
||||
seen_words = set()
|
||||
_results = []
|
||||
for part in parts:
|
||||
for item in select_terminology_by_word(session, part, oid, datasource, advanced_application_id):
|
||||
item_words = set(item.get('words', [])) if isinstance(item, dict) else set()
|
||||
if item_words and not item_words.issubset(seen_words):
|
||||
seen_words.update(item_words)
|
||||
_results.append(item)
|
||||
|
||||
if _results:
|
||||
terminology = to_xml_string(_results)
|
||||
template = get_base_terminology_template().format(terminologies=terminology)
|
||||
return template, _results
|
||||
else:
|
||||
return '', []
|
||||
|
||||
|
||||
def expand_with_terminology(keywords: str, session: SessionDep, oid: int) -> str:
|
||||
"""使用术语同义词扩展关键词。
|
||||
|
||||
查询当前租户启用的术语表,将匹配的关键词扩展为同义词集合。
|
||||
pid=None 表示父术语(主词),pid=<id> 表示子术语(同义词)。
|
||||
|
||||
Args:
|
||||
keywords: 逗号分隔的关键词字符串
|
||||
session: 数据库会话
|
||||
oid: 组织 ID
|
||||
|
||||
Returns:
|
||||
扩展后的关键词字符串(逗号分隔),最多 10 个关键词
|
||||
"""
|
||||
if not keywords or not session or oid is None:
|
||||
return keywords
|
||||
|
||||
keyword_list = [kw.strip() for kw in keywords.split(',') if kw.strip()]
|
||||
if not keyword_list:
|
||||
return keywords
|
||||
|
||||
try:
|
||||
# 查询当前租户启用的术语(限制 500 条避免内存问题)
|
||||
terms = session.query(Terminology).filter(
|
||||
Terminology.oid == oid,
|
||||
Terminology.enabled == True
|
||||
).limit(500).all()
|
||||
|
||||
if not terms:
|
||||
return keywords
|
||||
|
||||
# 构建父子结构
|
||||
parent_map = {} # id -> Terminology(仅父术语)
|
||||
children_map = {} # parent_id -> [Terminology](子术语列表)
|
||||
word_to_parent_ids = {} # 小写词 -> 父术语 id 集合
|
||||
|
||||
for term in terms:
|
||||
if term.pid is None or term.pid == 0:
|
||||
parent_map[term.id] = term
|
||||
if term.word:
|
||||
w = term.word.strip().lower()
|
||||
if w not in word_to_parent_ids:
|
||||
word_to_parent_ids[w] = set()
|
||||
word_to_parent_ids[w].add(term.id)
|
||||
else:
|
||||
if term.pid not in children_map:
|
||||
children_map[term.pid] = []
|
||||
children_map[term.pid].append(term)
|
||||
if term.word:
|
||||
w = term.word.strip().lower()
|
||||
if w not in word_to_parent_ids:
|
||||
word_to_parent_ids[w] = set()
|
||||
word_to_parent_ids[w].add(term.pid)
|
||||
|
||||
# 扩展关键词
|
||||
expanded = list(keyword_list)
|
||||
seen = set(kw.lower() for kw in keyword_list)
|
||||
max_keywords = 10
|
||||
|
||||
for kw in keyword_list:
|
||||
if len(expanded) >= max_keywords:
|
||||
break
|
||||
|
||||
kw_lower = kw.lower()
|
||||
matched_parent_ids = word_to_parent_ids.get(kw_lower, set())
|
||||
|
||||
for pid in matched_parent_ids:
|
||||
if len(expanded) >= max_keywords:
|
||||
break
|
||||
|
||||
# 添加父术语
|
||||
parent = parent_map.get(pid)
|
||||
if parent and parent.word:
|
||||
parent_word = parent.word.strip()
|
||||
if parent_word.lower() not in seen:
|
||||
expanded.append(parent_word)
|
||||
seen.add(parent_word.lower())
|
||||
|
||||
# 添加子术语(同义词)
|
||||
for child in children_map.get(pid, []):
|
||||
if len(expanded) >= max_keywords:
|
||||
break
|
||||
if child.word:
|
||||
child_word = child.word.strip()
|
||||
if child_word.lower() not in seen:
|
||||
expanded.append(child_word)
|
||||
seen.add(child_word.lower())
|
||||
|
||||
return ','.join(expanded)
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
return keywords
|
||||
|
||||
Reference in New Issue
Block a user