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

backend/apps/datasource/embedding/table_embedding.py

@@ -6,7 +6,7 @@
 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]

backend/apps/datasource/embedding/utils.py

@@ -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)