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:
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
+612
View File
@@ -0,0 +1,612 @@
"""表匹配融合评分测试。
测试覆盖:
1. calc_keyword_score - 纯函数,无外部依赖
2. expand_with_terminology - 需要 mock 数据库
3. calc_table_embedding - 集成测试,mock 所有依赖
"""
import json
from unittest.mock import MagicMock, patch
import pytest
from apps.datasource.embedding.table_embedding import (
calc_keyword_score,
calc_table_embedding,
)
from apps.terminology.curd.terminology import expand_with_terminology
# mock 目标路径
EMBEDDING_CACHE_PATCH = 'apps.datasource.embedding.table_embedding.EmbeddingModelCache'
SETTINGS_PATCH = 'apps.datasource.embedding.table_embedding.settings'
TERMINOLOGY_PATCH = 'apps.terminology.curd.terminology'
# =============================================================================
# calc_keyword_score 测试(纯函数,无需 mock)
# =============================================================================
class TestCalcKeywordScore:
"""关键词匹配分数计算测试。"""
def test_full_table_name_match(self):
"""表名完整匹配关键词 -> 1.0"""
score = calc_keyword_score(
question="sales_order",
table_name="sales_order",
table_comment="销售订单",
fields=[]
)
assert score == 1.0
def test_full_table_name_match_case_insensitive(self):
"""表名匹配应忽略大小写。"""
score = calc_keyword_score(
question="Sales_Order",
table_name="sales_order",
table_comment="",
fields=[]
)
assert score == 1.0
def test_full_table_name_match_among_multiple_keywords(self):
"""多个关键词时,表名完整匹配仍返回 1.0。"""
score = calc_keyword_score(
question="sales_order,本月,销售额",
table_name="sales_order",
table_comment="",
fields=[]
)
assert score == 1.0
def test_table_comment_full_match(self):
"""表注释完整匹配关键词 -> 0.9"""
score = calc_keyword_score(
question="销售订单",
table_name="so_header",
table_comment="销售订单",
fields=[]
)
assert score == 0.9
def test_table_comment_full_match_custom_comment_priority(self):
"""自定义备注应作为 table_comment 使用(调用方负责合并)。"""
# 调用方传入已合并的注释,这里模拟该场景
score = calc_keyword_score(
question="销售订单",
table_name="so_header",
table_comment="自定义备注",
fields=[]
)
# "销售订单" != "自定义备注",不匹配注释
assert score == 0.0
def test_partial_table_name_match_single_keyword(self):
"""关键词匹配表名部分 -> 0.25~0.5(按覆盖比例)。"""
# "sales" 匹配 "sales_order" 的一部分(2 个部分:sales, order)
# 覆盖率 = 1/2 = 0.5 -> 0.25 + 0.25*0.5 = 0.375
score = calc_keyword_score(
question="sales",
table_name="sales_order",
table_comment="",
fields=[]
)
assert score == 0.38 # 0.25 + 0.25 * 0.5 = 0.375,四舍五入为 0.38
def test_partial_table_name_match_full_coverage(self):
"""表名所有部分都被关键词覆盖 -> 0.5。"""
# "sales" 和 "order" 覆盖了 "sales_order" 的所有部分
score = calc_keyword_score(
question="sales,order",
table_name="sales_order",
table_comment="",
fields=[]
)
assert score == 0.5
def test_field_comment_match(self):
"""关键词匹配字段注释 -> 0.5"""
fields = [
{"field_name": "amount", "field_comment": "销售额", "custom_comment": ""},
{"field_name": "qty", "field_comment": "数量", "custom_comment": ""},
]
score = calc_keyword_score(
question="销售额",
table_name="order_detail",
table_comment="",
fields=fields
)
assert score == 0.5
def test_field_custom_comment_priority(self):
"""自定义注释应优先于字段注释。"""
fields = [
{"field_name": "amount", "field_comment": "原始注释", "custom_comment": "销售额"},
]
score = calc_keyword_score(
question="销售额",
table_name="order_detail",
table_comment="",
fields=fields
)
assert score == 0.5
def test_field_name_match(self):
"""关键词匹配字段名部分 -> 0.3"""
fields = [
{"field_name": "order_amount", "field_comment": "", "custom_comment": ""},
]
score = calc_keyword_score(
question="amount",
table_name="transactions",
table_comment="",
fields=fields
)
assert score == 0.3
def test_field_name_match_no_duplicate(self):
"""同一词根匹配多个字段时只计算一次。"""
fields = [
{"field_name": "order_amount", "field_comment": "", "custom_comment": ""},
{"field_name": "total_amount", "field_comment": "", "custom_comment": ""},
]
# "amount" 匹配了两个字段,但仍应返回 0.3(不会更高)
score = calc_keyword_score(
question="amount",
table_name="transactions",
table_comment="",
fields=fields
)
assert score == 0.3
def test_no_match(self):
"""无匹配 -> 0.0"""
fields = [
{"field_name": "id", "field_comment": "主键", "custom_comment": ""},
]
score = calc_keyword_score(
question="销售额",
table_name="users",
table_comment="用户表",
fields=fields
)
assert score == 0.0
def test_empty_question(self):
"""空问题 -> 0.0"""
score = calc_keyword_score(
question="",
table_name="sales_order",
table_comment="",
fields=[]
)
assert score == 0.0
def test_empty_keywords_after_split(self):
"""问题只有逗号/空格 -> 0.0"""
score = calc_keyword_score(
question=",, ,",
table_name="sales_order",
table_comment="",
fields=[]
)
assert score == 0.0
def test_priority_table_name_over_comment(self):
"""表名匹配应优先于注释匹配。"""
score = calc_keyword_score(
question="sales_order,销售订单",
table_name="sales_order",
table_comment="销售订单",
fields=[]
)
assert score == 1.0
def test_priority_comment_over_field_comment(self):
"""注释匹配(0.9)应优先于字段注释匹配(0.5)。"""
fields = [
{"field_name": "amt", "field_comment": "金额", "custom_comment": ""},
]
score = calc_keyword_score(
question="销售订单",
table_name="so_detail",
table_comment="销售订单",
fields=fields
)
assert score == 0.9
def test_priority_field_comment_over_field_name(self):
"""字段注释匹配(0.5)应优先于字段名匹配(0.3)。"""
fields = [
{"field_name": "销售额_column", "field_comment": "", "custom_comment": ""},
{"field_name": "amt", "field_comment": "销售额", "custom_comment": ""},
]
# 字段注释匹配 "销售额" -> 0.5
score = calc_keyword_score(
question="销售额",
table_name="detail",
table_comment="",
fields=fields
)
assert score == 0.5
def test_multiple_keywords_best_score_wins(self):
"""多个关键词时,返回最佳匹配分数。"""
fields = [
{"field_name": "qty", "field_comment": "数量", "custom_comment": ""},
]
# "sales_order" 匹配表名 -> 1.0
# "数量" 匹配字段注释 -> 0.5
# 最佳为 1.0
score = calc_keyword_score(
question="sales_order,数量",
table_name="sales_order",
table_comment="",
fields=fields
)
assert score == 1.0
def test_underscore_table_name_parts(self):
"""带下划线的表名应正确分割。"""
# "sales_order_detail" -> 部分: {"sales", "order", "detail"}
# "sales" 覆盖 1/3 -> 0.25 + 0.25*(1/3) ≈ 0.33
score = calc_keyword_score(
question="sales",
table_name="sales_order_detail",
table_comment="",
fields=[]
)
assert 0.3 <= score <= 0.35
def test_field_name_with_underscores(self):
"""带下划线的字段名应正确分割。"""
fields = [
{"field_name": "customer_order_count", "field_comment": "", "custom_comment": ""},
]
score = calc_keyword_score(
question="order",
table_name="stats",
table_comment="",
fields=fields
)
assert score == 0.3
def test_real_scenario_sales_query(self):
"""真实场景:'查询销售额' 应匹配包含 '销售额' 字段的表。"""
fields = [
{"field_name": "id", "field_comment": "主键", "custom_comment": ""},
{"field_name": "order_amount", "field_comment": "订单金额", "custom_comment": ""},
{"field_name": "sales_amount", "field_comment": "销售额", "custom_comment": ""},
]
score = calc_keyword_score(
question="销售额",
table_name="daily_report",
table_comment="日报表",
fields=fields
)
assert score == 0.5 # 字段注释匹配
def test_real_scenario_table_name_query(self):
"""真实场景:'查 sales_order 表最近的订单' -> 提取 'sales_order'。"""
score = calc_keyword_score(
question="sales_order,订单",
table_name="sales_order",
table_comment="销售订单",
fields=[]
)
assert score == 1.0 # 表名完整匹配
# =============================================================================
# expand_with_terminology 测试(mock 数据库)
# =============================================================================
class TestExpandWithTerminology:
"""术语同义词扩展测试。"""
def _make_term(self, id, word, pid=None, enabled=True, oid=1):
"""创建 mock 术语对象的辅助方法。"""
term = MagicMock()
term.id = id
term.word = word
term.pid = pid
term.enabled = enabled
term.oid = oid
return term
def test_expand_hit(self):
"""匹配术语的关键词应扩展同义词。"""
session = MagicMock()
session.query.return_value.filter.return_value.limit.return_value.all.return_value = [
self._make_term(1, "销售", pid=None),
self._make_term(2, "sales", pid=1),
self._make_term(3, "selling", pid=1),
]
result = expand_with_terminology("销售", session, oid=1)
keywords = [kw.strip() for kw in result.split(',')]
assert "销售" in keywords
assert "sales" in keywords
assert "selling" in keywords
def test_no_hit(self):
"""未匹配任何术语的关键词应保持不变。"""
session = MagicMock()
session.query.return_value.filter.return_value.limit.return_value.all.return_value = [
self._make_term(1, "客户", pid=None),
]
result = expand_with_terminology("销售额", session, oid=1)
assert result == "销售额"
def test_max_10_keywords(self):
"""扩展后的关键词总数不应超过 10 个。"""
session = MagicMock()
# 创建一个有多个同义词的术语
terms = [self._make_term(1, "销售", pid=None)]
for i in range(15):
terms.append(self._make_term(100 + i, f"synonym_{i}", pid=1))
session.query.return_value.filter.return_value.limit.return_value.all.return_value = terms
result = expand_with_terminology("销售", session, oid=1)
keywords = [kw.strip() for kw in result.split(',') if kw.strip()]
assert len(keywords) <= 10
def test_empty_terminology(self):
"""空术语表应返回原始关键词。"""
session = MagicMock()
session.query.return_value.filter.return_value.limit.return_value.all.return_value = []
result = expand_with_terminology("销售额,部门", session, oid=1)
assert result == "销售额,部门"
def test_none_session(self):
"""session 为 None 时应返回原始关键词。"""
result = expand_with_terminology("销售额", None, oid=1)
assert result == "销售额"
def test_none_oid(self):
"""oid 为 None 时应返回原始关键词。"""
session = MagicMock()
result = expand_with_terminology("销售额", session, oid=None)
assert result == "销售额"
def test_oid_zero_should_still_query(self):
"""oid=0 不应被视为假值 — 它是有效的组织 ID。"""
session = MagicMock()
session.query.return_value.filter.return_value.limit.return_value.all.return_value = [
self._make_term(1, "销售", pid=None, oid=0),
self._make_term(2, "revenue", pid=1, oid=0),
]
result = expand_with_terminology("销售", session, oid=0)
keywords = [kw.strip() for kw in result.split(',')]
assert "销售" in keywords
assert "revenue" in keywords
def test_empty_keywords(self):
"""空关键词字符串应原样返回。"""
session = MagicMock()
result = expand_with_terminology("", session, oid=1)
assert result == ""
def test_case_insensitive_match(self):
"""术语匹配应忽略大小写。"""
session = MagicMock()
session.query.return_value.filter.return_value.limit.return_value.all.return_value = [
self._make_term(1, "Sales", pid=None),
self._make_term(2, "revenue", pid=1),
]
result = expand_with_terminology("sales", session, oid=1)
keywords = [kw.strip() for kw in result.split(',')]
# "sales"(输入)匹配 "Sales"(术语),忽略大小写
# "Sales" 不会被添加,因为 "sales" 已在列表中(避免重复)
# "revenue"(子术语同义词)被添加
assert "sales" in keywords
assert "revenue" in keywords
assert len(keywords) == 2
def test_multiple_keywords_expansion(self):
"""多个关键词应各自独立扩展。"""
session = MagicMock()
session.query.return_value.filter.return_value.limit.return_value.all.return_value = [
self._make_term(1, "销售", pid=None),
self._make_term(2, "sales", pid=1),
self._make_term(3, "部门", pid=None),
self._make_term(4, "department", pid=3),
]
result = expand_with_terminology("销售,部门", session, oid=1)
keywords = [kw.strip() for kw in result.split(',')]
assert "sales" in keywords
assert "department" in keywords
def test_child_term_word_also_matches(self):
"""如果关键词匹配子术语的词,应扩展父术语和兄弟术语。"""
session = MagicMock()
session.query.return_value.filter.return_value.limit.return_value.all.return_value = [
self._make_term(1, "销售", pid=None),
self._make_term(2, "sales", pid=1),
self._make_term(3, "selling", pid=1),
]
# "sales" 匹配子术语,应扩展包含父术语 "销售" 和兄弟术语 "selling"
result = expand_with_terminology("sales", session, oid=1)
keywords = [kw.strip() for kw in result.split(',')]
assert "sales" in keywords
assert "销售" in keywords
assert "selling" in keywords
# =============================================================================
# calc_table_embedding 集成测试(mock embedding)
# =============================================================================
class TestCalcTableEmbedding:
"""融合评分管线集成测试。"""
def _make_table(self, id, table_name, schema_table, embedding, table_comment="", fields=None):
return {
"id": id,
"table_name": table_name,
"schema_table": schema_table,
"embedding": json.dumps(embedding),
"table_comment": table_comment,
"fields": fields or []
}
@patch('apps.datasource.embedding.table_embedding.settings')
@patch('apps.datasource.embedding.table_embedding.EmbeddingModelCache')
def test_exact_table_name_match_ranks_first(self, mock_embed_cache, mock_settings):
"""表名完整匹配关键词时,应排第一。"""
mock_settings.TABLE_EMBEDDING_KEYWORD_ENABLED = True
mock_settings.TABLE_EMBEDDING_ALPHA = 0.4
mock_settings.TABLE_EMBEDDING_COUNT = 10
mock_model = MagicMock()
mock_model.embed_query.return_value = [1.0, 0.0, 0.0]
mock_embed_cache.get_model.return_value = mock_model
tables = [
self._make_table(1, "customer", "# Table: customer", [0.0, 1.0, 0.0], "客户表",
[{"field_name": "name", "field_comment": "姓名", "custom_comment": ""}]),
self._make_table(2, "product", "# Table: product", [0.0, 0.0, 1.0], "产品表",
[{"field_name": "name", "field_comment": "产品名", "custom_comment": ""}]),
self._make_table(3, "sales_order", "# Table: sales_order", [0.9, 0.1, 0.0], "销售订单",
[{"field_name": "amount", "field_comment": "金额", "custom_comment": ""}]),
]
# 直接传入已提取的关键词
result = calc_table_embedding(tables, "查 sales_order 表最近的订单",
keywords="sales_order,订单")
# sales_order 应排第一,因为 keyword_score=1.0
assert result[0]['table_name'] == "sales_order"
assert result[0].get('keyword_score') == 1.0
@patch('apps.datasource.embedding.table_embedding.settings')
@patch('apps.datasource.embedding.table_embedding.EmbeddingModelCache')
def test_keyword_disabled_uses_vector_only(self, mock_embed_cache, mock_settings):
"""关键词匹配关闭时,应使用纯向量匹配。"""
mock_settings.TABLE_EMBEDDING_KEYWORD_ENABLED = False
mock_settings.TABLE_EMBEDDING_COUNT = 10
mock_model = MagicMock()
mock_model.embed_query.return_value = [1.0, 0.0, 0.0]
mock_embed_cache.get_model.return_value = mock_model
tables = [
self._make_table(1, "customer", "# Table: customer", [0.0, 1.0, 0.0]),
self._make_table(2, "sales_order", "# Table: sales_order", [0.9, 0.1, 0.0]),
]
result = calc_table_embedding(tables, "查 sales_order 表", keywords="sales_order,订单")
# 关键词匹配关闭,即使传入了 keywords 也应使用纯向量匹配
assert result[0]['table_name'] == "sales_order"
@patch('apps.datasource.embedding.table_embedding.settings')
@patch('apps.datasource.embedding.table_embedding.EmbeddingModelCache')
def test_no_keywords_falls_back_to_vector(self, mock_embed_cache, mock_settings):
"""未传入 keywords 时,应回退到纯向量匹配。"""
mock_settings.TABLE_EMBEDDING_KEYWORD_ENABLED = True
mock_settings.TABLE_EMBEDDING_COUNT = 10
mock_model = MagicMock()
mock_model.embed_query.return_value = [1.0, 0.0, 0.0]
mock_embed_cache.get_model.return_value = mock_model
tables = [
self._make_table(1, "customer", "# Table: customer", [0.0, 1.0, 0.0]),
self._make_table(2, "sales_order", "# Table: sales_order", [0.9, 0.1, 0.0]),
]
# 不传 keywords
result = calc_table_embedding(tables, "查 sales_order 表")
# 应回退到纯向量匹配
assert result[0]['table_name'] == "sales_order"
@patch('apps.datasource.embedding.table_embedding.settings')
@patch('apps.datasource.embedding.table_embedding.EmbeddingModelCache')
def test_multiple_exact_matches_all_rank_first(self, mock_embed_cache, mock_settings):
"""多个 keyword_score=1.0 的表都应排在其他表之前。"""
mock_settings.TABLE_EMBEDDING_KEYWORD_ENABLED = True
mock_settings.TABLE_EMBEDDING_ALPHA = 0.4
mock_settings.TABLE_EMBEDDING_COUNT = 10
mock_model = MagicMock()
mock_model.embed_query.return_value = [1.0, 0.0, 0.0]
mock_embed_cache.get_model.return_value = mock_model
tables = [
self._make_table(1, "product", "# Table: product", [0.0, 0.0, 1.0], "产品",
[{"field_name": "name", "field_comment": "名称", "custom_comment": ""}]),
self._make_table(2, "sales_order", "# Table: sales_order", [0.9, 0.1, 0.0], "",
[{"field_name": "amount", "field_comment": "", "custom_comment": ""}]),
self._make_table(3, "customer", "# Table: customer", [0.8, 0.2, 0.0], "",
[{"field_name": "name", "field_comment": "", "custom_comment": ""}]),
]
result = calc_table_embedding(tables, "查 sales_order 和 customer",
keywords="sales_order,customer")
# 前两个应是精确匹配(它们之间的顺序无所谓)
exact_match_names = {result[0]['table_name'], result[1]['table_name']}
assert exact_match_names == {"sales_order", "customer"}
# product 应排最后
assert result[2]['table_name'] == "product"
@patch('apps.datasource.embedding.table_embedding.settings')
@patch('apps.datasource.embedding.table_embedding.EmbeddingModelCache')
def test_fusion_score_formula(self, mock_embed_cache, mock_settings):
"""验证融合公式:final = α * vec + (1-α) * keyword。"""
mock_settings.TABLE_EMBEDDING_KEYWORD_ENABLED = True
mock_settings.TABLE_EMBEDDING_ALPHA = 0.4
mock_settings.TABLE_EMBEDDING_COUNT = 10
mock_model = MagicMock()
mock_model.embed_query.return_value = [1.0, 0.0]
mock_embed_cache.get_model.return_value = mock_model
# 表的 vec_score=0.8(通过余弦相似度),keyword_score 来自部分匹配
tables = [
self._make_table(1, "sales_order", "# Table: sales_order", [0.8, 0.6], "",
[{"field_name": "id", "field_comment": "", "custom_comment": ""}]),
]
result = calc_table_embedding(tables, "sales", keywords="sales")
# "sales" 部分匹配 "sales_order"(2 个部分中的 1 个)
# keyword_score = 0.25 + 0.25 * 0.5 = 0.375
# vec_score = cosine([1,0], [0.8, 0.6]) = 0.8 / (1 * 1) = 0.8
# final = 0.4 * 0.8 + 0.6 * 0.375 = 0.32 + 0.225 = 0.545
assert len(result) == 1
score = result[0]['cosine_similarity']
assert 0.5 <= score <= 0.6 # 允许小的浮点误差
@patch('apps.datasource.embedding.table_embedding.settings')
@patch('apps.datasource.embedding.table_embedding.EmbeddingModelCache')
def test_fallback_on_unexpected_error(self, mock_embed_cache, mock_settings):
"""遇到意外错误时,应回退到纯向量评分。"""
mock_settings.TABLE_EMBEDDING_KEYWORD_ENABLED = True
mock_settings.TABLE_EMBEDDING_ALPHA = 0.4
mock_settings.TABLE_EMBEDDING_COUNT = 10
# 让 embed_query 抛出意外错误
mock_model = MagicMock()
mock_model.embed_query.side_effect = RuntimeError("模型加载失败")
mock_embed_cache.get_model.return_value = mock_model
tables = [
self._make_table(1, "t1", "# Table: t1", [0.8, 0.6]),
self._make_table(2, "t2", "# Table: t2", [0.3, 0.7]),
]
# 不应抛出异常,应回退到纯向量匹配
result = calc_table_embedding(tables, "test", keywords="test")
assert len(result) == 2
if __name__ == '__main__':
pytest.main([__file__, '-v'])