klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-08 17:26:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: implement LLM keyword extraction and fusion scoring for table matching

Browse Source 
- 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:
ulleo
2026-09-14 17:51:45 +08:00
committed by ulleo
parent aac200aef5
commit 3fd6fa867d
15 changed files with 1200 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/datasource/crud/datasource.py
+18 -4
View File
@@ -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:
backend/apps/datasource/embedding/table_embedding.py
+230 -14
View File
@@ -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()