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

- improve table matching
This commit is contained in:
ulleo
2026-09-14 18:52:15 +08:00
committed by ulleo
parent 3fd6fa867d
commit e685a749c6
2 changed files with 99 additions and 43 deletions
@@ -6,7 +6,7 @@ import time
import traceback
from apps.ai_model.embedding import EmbeddingModelCache
from apps.datasource.embedding.utils import cosine_similarity
from apps.datasource.embedding.utils import batch_cosine_similarity, cosine_similarity
from common.core.config import settings
from common.utils.utils import SQLBotLogUtil
@@ -83,6 +83,27 @@ def calc_keyword_score(
tc = (table_comment or '').lower()
total_kw = len(keywords)
# 表名只分割一次,后续复用
tn_parts = set(re.split(r'[_\s]+', tn))
tn_parts.discard('')
# 预处理字段:只 strip/lower 一次
field_comments = []
field_name_parts = []
for field in fields:
fc = ''
if field.get('custom_comment'):
fc = field['custom_comment'].strip().lower()
if not fc and field.get('field_comment'):
fc = field['field_comment'].strip().lower()
field_comments.append(fc)
fname = field.get('field_name', '').lower()
if fname:
field_name_parts.append(set(re.split(r'[_\s]+', fname)) - {''})
else:
field_name_parts.append(set())
# 收集每个关键词的最佳匹配分数
kw_scores = {} # keyword -> best_score
@@ -103,35 +124,24 @@ def calc_keyword_score(
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:
# 5. 字段注释匹配:0.5(精确)/ 0.4(子串)
for fc in field_comments:
if fc:
if kw == fc:
kw_scores[kw] = max(kw_scores.get(kw, 0), 0.5)
break
elif kw in field_comment:
elif kw in fc:
kw_scores[kw] = max(kw_scores.get(kw, 0), 0.4)
# 5. 字段名匹配:0.3
# 6. 字段名匹配: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
for fparts in field_name_parts:
if kw in fparts:
kw_scores[kw] = max(kw_scores.get(kw, 0), 0.3)
break
if not kw_scores:
return 0.0
@@ -141,10 +151,6 @@ def calc_keyword_score(
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:
@@ -201,15 +207,26 @@ def calc_table_embedding(tables: list[dict], question: str, session=None, oid: i
try:
start_time = time.time()
# 步骤 1:用关键词计算向量相似度
# 步骤 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))
# 预解析所有 embedding,跳过空的
parsed_embeddings = []
valid_indices = []
for i, item in enumerate(_list):
emb = item.get('embedding')
if emb:
parsed = _parse_embedding(emb)
if parsed:
parsed_embeddings.append(parsed)
valid_indices.append(i)
# 一次性批量计算所有余弦相似度
if parsed_embeddings:
similarities = batch_cosine_similarity(q_embedding, parsed_embeddings)
for idx, sim in zip(valid_indices, similarities):
_list[idx]['cosine_similarity'] = float(sim)
# 步骤 2 & 3:计算关键词分数并融合
alpha = settings.TABLE_EMBEDDING_ALPHA
@@ -269,13 +286,24 @@ def _calc_vector_only(tables: list[dict], question: str):
try:
model = EmbeddingModelCache.get_model()
start_time = time.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, _parse_embedding(item))
# 预解析所有 embedding,跳过空的
parsed_embeddings = []
valid_indices = []
for i, item in enumerate(_list):
emb = item.get('embedding')
if emb:
parsed = _parse_embedding(emb)
if parsed:
parsed_embeddings.append(parsed)
valid_indices.append(i)
# 一次性批量计算所有余弦相似度
if parsed_embeddings:
similarities = batch_cosine_similarity(q_embedding, parsed_embeddings)
for idx, sim in zip(valid_indices, similarities):
_list[idx]['cosine_similarity'] = float(sim)
_list.sort(key=lambda x: x['cosine_similarity'], reverse=True)
_list = _list[:settings.TABLE_EMBEDDING_COUNT]
+32 -4
View File
@@ -2,17 +2,45 @@
# Date: 2025/9/23
import math
import numpy as np
def cosine_similarity(vec_a, vec_b):
if len(vec_a) != len(vec_b):
raise ValueError("The vector dimension must be the same")
dot_product = sum(a * b for a, b in zip(vec_a, vec_b))
a = np.asarray(vec_a, dtype=np.float64)
b = np.asarray(vec_b, dtype=np.float64)
norm_a = math.sqrt(sum(a * a for a in vec_a))
norm_b = math.sqrt(sum(b * b for b in vec_b))
dot_product = np.dot(a, b)
norm_a = np.linalg.norm(a)
norm_b = np.linalg.norm(b)
if norm_a == 0 or norm_b == 0:
return 0.0
return dot_product / (norm_a * norm_b)
return float(dot_product / (norm_a * norm_b))
def batch_cosine_similarity(query_vec, table_vecs):
"""批量计算一个 query 向量与多个 table 向量的余弦相似度。
Args:
query_vec: query embedding (list or array)
table_vecs: 二维数组,每行是一个 table embedding
Returns:
一维 numpy array,每个元素是对应 table 与 query 的余弦相似度
"""
q = np.asarray(query_vec, dtype=np.float64)
mat = np.asarray(table_vecs, dtype=np.float64)
q_norm = np.linalg.norm(q)
if q_norm == 0:
return np.zeros(len(table_vecs))
# 按行计算范数,避免除零
norms = np.linalg.norm(mat, axis=1)
norms[norms == 0] = 1.0 # 零向量相似度为 0,除以 1 保持结果为 0
return mat @ q / (norms * q_norm)