Commit 317c5d2c61e172bd952b54da7eb9075e5a278e17
1 parent
0a440fb3
feat(search): 引入 exact vector rescore 为 topN 补全精确向量分,解决 rerank 阶段部分文档缺失 text/image knn 分数的问题
背景与问题
- 现有粗排/重排依赖 `knn_query` 和 `image_knn_query` 分数,但这两路分数来自 ANN 召回,并非所有进入 rerank_window (160) 的文档都同时命中文本和图片向量召回,导致部分文档得分为 0,影响融合公式的稳定性。
- 简单扩大 ANN 的 k 无法保证 lexical 召回带来的文档也包含两路向量分;二次查询或拉回向量本地计算均有额外开销且实现复杂。
解决方案
采用 ES rescore 机制,在第一次搜索的 `window_size` 内对每个文档执行精确的向量 script_score,并将分数以 named query 形式附加到 `matched_queries` 中,供后续 coarse/rerank 优先使用。
**设计决策**:
- **只补分,不改排序**:rescore 使用 `score_mode: total` 且 `rescore_query_weight: 0.0`,原始 `_score` 保持不变,避免干扰现有排序逻辑,风险最小。
- **精确分数命名**:`exact_text_knn_query` 和 `exact_image_knn_query`,便于客户端识别和回退。
- **可配置**:通过 `exact_knn_rescore_enabled` 开关和 `exact_knn_rescore_window` 控制窗口大小,默认 160。
技术实现细节
1. 配置扩展 (`config/config.yaml`, `config/loader.py`)
```yaml
exact_knn_rescore_enabled: true
exact_knn_rescore_window: 160
```
新增配置项并注入到 `RerankConfig`。
2. Searcher 构建 rescore 查询 (`search/searcher.py`)
- 在 `_build_es_search_request` 中,当 `enable_rerank=True` 且配置开启时,构造 rescore 对象:
- `window_size` = `exact_knn_rescore_window`
- `query` 为一个 `bool` 查询,内嵌两个 `script_score` 子查询,分别计算文本和图片向量的点积相似度:
```painless
// exact_text_knn_query
(dotProduct(params.query_vector, 'title_embedding') + 1.0) / 2.0
// exact_image_knn_query
(dotProduct(params.image_query_vector, 'image_embedding.vector') + 1.0) / 2.0
```
- 每个 `script_score` 都设置 `_name` 为对应的 named query。
- 注意:当前实现的脚本分数**尚未乘以 knn_text_boost / knn_image_boost**,保持与原始 ANN 分数尺度对齐的后续待办。
3. RerankClient 优先读取 exact 分数 (`search/rerank_client.py`)
- 在 `_extract_coarse_signals` 中,从文档的 `matched_queries` 里读取 `exact_text_knn_query` 和 `exact_image_knn_query` 分数。
- 若存在且值有效,则用作 `text_knn_score` / `image_knn_score`,并标记 `text_knn_source='exact_text_knn_query'`。
- 若不存在,则回退到原有的 `knn_query` / `image_knn_query` (ANN 分数)。
- 同时保留原始 ANN 分数到 `approx_text_knn_score` / `approx_image_knn_score` 供调试对比。
4. 调试信息增强
- `debug_info.per_result[*].ranking_funnel.coarse_rank.signals` 中输出 exact 分数、回退分数及来源标记,便于线上观察覆盖率和数值分布。
验证结果
- 通过单元测试 `tests/test_rerank_client.py` 和 `tests/test_search_rerank_window.py`,验证 exact 优先级、配置解析及 ES 请求体结构。
- 线上真实查询采样(6 个 query,top160)显示:
- **exact 覆盖率达到 100%**(文本和图片均有分),解决了原 ANN 部分缺失的问题。
- 但 exact 分数与原始 ANN 分数存在量级差异(ANN/exact 中位数比值约 4.1 倍),原因是 exact 脚本未乘 boost 因子。
- 当前排名影响:粗排 top10 重叠度最低降至 1/10,最大排名漂移超过 100。
后续计划
1. 对齐 exact 分与 ANN 分的尺度:在 script_score 中乘以 `knn_text_boost` / `knn_image_boost`,并对长查询额外乘 1.4。
2. 重新评估 top10 重叠度和漂移,若收敛则可将 coarse 融合公式整体迁移至 ES rescore 阶段。
3. 当前版本保持“只补分不改排序”的安全策略,已解决核心的分数缺失问题。
涉及文件
- `config/config.yaml`
- `config/loader.py`
- `search/searcher.py`
- `search/rerank_client.py`
- `tests/test_rerank_client.py`
- `tests/test_search_rerank_window.py`
Showing
10 changed files
with
461 additions
and
5 deletions
Show diff stats
config/config.yaml
| @@ -324,6 +324,8 @@ fine_rank: | @@ -324,6 +324,8 @@ fine_rank: | ||
| 324 | rerank: | 324 | rerank: |
| 325 | enabled: true | 325 | enabled: true |
| 326 | rerank_window: 160 | 326 | rerank_window: 160 |
| 327 | + exact_knn_rescore_enabled: true | ||
| 328 | + exact_knn_rescore_window: 160 | ||
| 327 | timeout_sec: 15.0 | 329 | timeout_sec: 15.0 |
| 328 | weight_es: 0.4 | 330 | weight_es: 0.4 |
| 329 | weight_ai: 0.6 | 331 | weight_ai: 0.6 |
config/loader.py
| @@ -608,6 +608,12 @@ class AppConfigLoader: | @@ -608,6 +608,12 @@ class AppConfigLoader: | ||
| 608 | rerank=RerankConfig( | 608 | rerank=RerankConfig( |
| 609 | enabled=bool(rerank_cfg.get("enabled", True)), | 609 | enabled=bool(rerank_cfg.get("enabled", True)), |
| 610 | rerank_window=int(rerank_cfg.get("rerank_window", 384)), | 610 | rerank_window=int(rerank_cfg.get("rerank_window", 384)), |
| 611 | + exact_knn_rescore_enabled=bool( | ||
| 612 | + rerank_cfg.get("exact_knn_rescore_enabled", False) | ||
| 613 | + ), | ||
| 614 | + exact_knn_rescore_window=int( | ||
| 615 | + rerank_cfg.get("exact_knn_rescore_window", 0) | ||
| 616 | + ), | ||
| 611 | timeout_sec=float(rerank_cfg.get("timeout_sec", 15.0)), | 617 | timeout_sec=float(rerank_cfg.get("timeout_sec", 15.0)), |
| 612 | weight_es=float(rerank_cfg.get("weight_es", 0.4)), | 618 | weight_es=float(rerank_cfg.get("weight_es", 0.4)), |
| 613 | weight_ai=float(rerank_cfg.get("weight_ai", 0.6)), | 619 | weight_ai=float(rerank_cfg.get("weight_ai", 0.6)), |
config/schema.py
| @@ -176,6 +176,9 @@ class RerankConfig: | @@ -176,6 +176,9 @@ class RerankConfig: | ||
| 176 | 176 | ||
| 177 | enabled: bool = True | 177 | enabled: bool = True |
| 178 | rerank_window: int = 384 | 178 | rerank_window: int = 384 |
| 179 | + exact_knn_rescore_enabled: bool = False | ||
| 180 | + #: topN exact vector scoring window; <=0 means "follow rerank_window" | ||
| 181 | + exact_knn_rescore_window: int = 0 | ||
| 179 | timeout_sec: float = 15.0 | 182 | timeout_sec: float = 15.0 |
| 180 | weight_es: float = 0.4 | 183 | weight_es: float = 0.4 |
| 181 | weight_ai: float = 0.6 | 184 | weight_ai: float = 0.6 |
docs/issues/a deleted
docs/issues/issue-2026-04-12-test-env.md
| @@ -14,3 +14,30 @@ nohup bash scripts/start_embedding_service.sh > log.start_embedding_service.0412 | @@ -14,3 +14,30 @@ nohup bash scripts/start_embedding_service.sh > log.start_embedding_service.0412 | ||
| 14 | 看他陪的文本是用的哪套方案、哪个模型,跟他对齐(我指的是当前的测试分支) | 14 | 看他陪的文本是用的哪套方案、哪个模型,跟他对齐(我指的是当前的测试分支) |
| 15 | 15 | ||
| 16 | 16 | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | + | ||
| 22 | +我在这个机器上部署了一个测试环境: | ||
| 23 | +120.76.41.98 端口22 用户名和密码: | ||
| 24 | +tw twtw@123 (有sudo权限) | ||
| 25 | +cd /home/tw/saas-search | ||
| 26 | +$ git branch | ||
| 27 | + masters RETURN) | ||
| 28 | +* test/small-gpu-es9 | ||
| 29 | + | ||
| 30 | +我希望差异只是: | ||
| 31 | +1. es配置不同(测试环境要连接到那台机器的一个docer的es 19200端口)、redis配置不同 | ||
| 32 | +2. reranker关闭、不要启动reranker服务 | ||
| 33 | + | ||
| 34 | +其余没什么不同。 | ||
| 35 | + | ||
| 36 | +但是启动有问题,现在翻译报错。 | ||
| 37 | +这体现了当前项目移植性比较差,我希望你检查一下失败原因,然后先到本地(本机 即当前目录master分支)优化好、提升移植性之后,那边更新,保持测试分支跟master只有少量的、配置层面的不同,让后到测试机器把翻译启动起来,最后包括整个服务都要启动起来。 | ||
| 38 | + | ||
| 39 | + | ||
| 40 | + | ||
| 41 | + | ||
| 42 | + | ||
| 43 | + |
| @@ -0,0 +1,25 @@ | @@ -0,0 +1,25 @@ | ||
| 1 | +需求: | ||
| 2 | +目前160条结果(rerank_window: 160)会进入重排,重排中 文本和图片向量的相关性,都会作为融合公式的因子之一(粗排和reranker都有): | ||
| 3 | +knn_score | ||
| 4 | +text_knn | ||
| 5 | +image_knn | ||
| 6 | +text_factor | ||
| 7 | +knn_factor | ||
| 8 | +但是文本向量召回和图片向量召回,是使用 KNN 索引召回的方式,并不是所有结果都有这两个得分,这两项得分都有为0的。 | ||
| 9 | +为了解决这个问题,有一个方法是对最终能进入重排的 160 条,看其中还有哪些分别缺失文本和图片向量召回的得分,再通过某种方式让 ES 去算,或者从 ES 把向量拉回来,自己算,或者在召回的时候请求 ES 的时候,就通过某种设定,确保前面的若干条都带有这两个分数,不知道有哪些方法,我感觉这些方法都不太好,请你思考一下 | ||
| 10 | + | ||
| 11 | +考虑的一个方案: | ||
| 12 | +想在“第一次 ES 搜索”里,只对 topN 补向量精算,考虑 rescore 或 retriever.rescorer的方案(官方明确支持多段 rescore/支持 score_mode: multiply,甚至示例里就有 function_score/script_score 放进 rescore 的写法。) | ||
| 13 | +这意味着你完全可以: | ||
| 14 | +初检仍然用现在的 lexical + text knn + image knn 召回候选 | ||
| 15 | +对 window_size=160 做 rescore | ||
| 16 | +用 exact script_score 给 top160 补 text/image vector 分 | ||
| 17 | +顺手把你现在本地 coarse 融合迁回 ES | ||
| 18 | + | ||
| 19 | +export ES_AUTH="saas:4hOaLaf41y2VuI8y" | ||
| 20 | +export ES="http://127.0.0.1:9200" | ||
| 21 | +"index":"search_products_tenant_163" | ||
| 22 | + | ||
| 23 | +有个细节暴露出来了:dotProduct() 这类向量函数在 script_score 评分上下文能用,但在 script_fields 取字段上下文里不认。所以如果我们要把 exact 分顺手回传给 rerank,用 script_fields 的话得自己写数组循环,不能直接调向量内建函数。 | ||
| 24 | + | ||
| 25 | +重排打分公式需要的base_query base_query_trans_zh knn_query image_knn_query还能不能拿到?请你考虑,尽量想想如何得到这些打分,如果实在拿不到去想替代的办法比如简化打分公式。 |
search/rerank_client.py
| @@ -153,12 +153,59 @@ def _extract_named_query_score(matched_queries: Any, name: str) -> float: | @@ -153,12 +153,59 @@ def _extract_named_query_score(matched_queries: Any, name: str) -> float: | ||
| 153 | return 0.0 | 153 | return 0.0 |
| 154 | 154 | ||
| 155 | 155 | ||
| 156 | +def _resolve_named_query_score( | ||
| 157 | + matched_queries: Any, | ||
| 158 | + *, | ||
| 159 | + preferred_names: List[str], | ||
| 160 | + fallback_names: List[str], | ||
| 161 | +) -> Tuple[float, Optional[str], float, Optional[str]]: | ||
| 162 | + preferred_score = 0.0 | ||
| 163 | + preferred_name: Optional[str] = None | ||
| 164 | + for name in preferred_names: | ||
| 165 | + score = _extract_named_query_score(matched_queries, name) | ||
| 166 | + if score > 0.0: | ||
| 167 | + preferred_score = score | ||
| 168 | + preferred_name = name | ||
| 169 | + break | ||
| 170 | + | ||
| 171 | + fallback_score = 0.0 | ||
| 172 | + fallback_name: Optional[str] = None | ||
| 173 | + for name in fallback_names: | ||
| 174 | + score = _extract_named_query_score(matched_queries, name) | ||
| 175 | + if score > 0.0: | ||
| 176 | + fallback_score = score | ||
| 177 | + fallback_name = name | ||
| 178 | + break | ||
| 179 | + | ||
| 180 | + if preferred_name is None and preferred_names: | ||
| 181 | + preferred_name = preferred_names[0] | ||
| 182 | + preferred_score = _extract_named_query_score(matched_queries, preferred_name) | ||
| 183 | + if fallback_name is None and fallback_names: | ||
| 184 | + fallback_name = fallback_names[0] | ||
| 185 | + fallback_score = _extract_named_query_score(matched_queries, fallback_name) | ||
| 186 | + if preferred_score > 0.0: | ||
| 187 | + return preferred_score, preferred_name, fallback_score, fallback_name | ||
| 188 | + return fallback_score, fallback_name, preferred_score, preferred_name | ||
| 189 | + | ||
| 190 | + | ||
| 156 | def _collect_knn_score_components( | 191 | def _collect_knn_score_components( |
| 157 | matched_queries: Any, | 192 | matched_queries: Any, |
| 158 | fusion: RerankFusionConfig, | 193 | fusion: RerankFusionConfig, |
| 159 | ) -> Dict[str, float]: | 194 | ) -> Dict[str, float]: |
| 160 | - text_knn_score = _extract_named_query_score(matched_queries, "knn_query") | ||
| 161 | - image_knn_score = _extract_named_query_score(matched_queries, "image_knn_query") | 195 | + text_knn_score, text_knn_source, _, _ = _resolve_named_query_score( |
| 196 | + matched_queries, | ||
| 197 | + preferred_names=["exact_text_knn_query"], | ||
| 198 | + fallback_names=["knn_query"], | ||
| 199 | + ) | ||
| 200 | + image_knn_score, image_knn_source, _, _ = _resolve_named_query_score( | ||
| 201 | + matched_queries, | ||
| 202 | + preferred_names=["exact_image_knn_query"], | ||
| 203 | + fallback_names=["image_knn_query"], | ||
| 204 | + ) | ||
| 205 | + exact_text_knn_score = _extract_named_query_score(matched_queries, "exact_text_knn_query") | ||
| 206 | + exact_image_knn_score = _extract_named_query_score(matched_queries, "exact_image_knn_query") | ||
| 207 | + approx_text_knn_score = _extract_named_query_score(matched_queries, "knn_query") | ||
| 208 | + approx_image_knn_score = _extract_named_query_score(matched_queries, "image_knn_query") | ||
| 162 | 209 | ||
| 163 | weighted_text_knn_score = text_knn_score * float(fusion.knn_text_weight) | 210 | weighted_text_knn_score = text_knn_score * float(fusion.knn_text_weight) |
| 164 | weighted_image_knn_score = image_knn_score * float(fusion.knn_image_weight) | 211 | weighted_image_knn_score = image_knn_score * float(fusion.knn_image_weight) |
| @@ -171,6 +218,14 @@ def _collect_knn_score_components( | @@ -171,6 +218,14 @@ def _collect_knn_score_components( | ||
| 171 | return { | 218 | return { |
| 172 | "text_knn_score": text_knn_score, | 219 | "text_knn_score": text_knn_score, |
| 173 | "image_knn_score": image_knn_score, | 220 | "image_knn_score": image_knn_score, |
| 221 | + "exact_text_knn_score": exact_text_knn_score, | ||
| 222 | + "exact_image_knn_score": exact_image_knn_score, | ||
| 223 | + "approx_text_knn_score": approx_text_knn_score, | ||
| 224 | + "approx_image_knn_score": approx_image_knn_score, | ||
| 225 | + "text_knn_source": text_knn_source, | ||
| 226 | + "image_knn_source": image_knn_source, | ||
| 227 | + "approx_text_knn_source": "knn_query", | ||
| 228 | + "approx_image_knn_source": "image_knn_query", | ||
| 174 | "weighted_text_knn_score": weighted_text_knn_score, | 229 | "weighted_text_knn_score": weighted_text_knn_score, |
| 175 | "weighted_image_knn_score": weighted_image_knn_score, | 230 | "weighted_image_knn_score": weighted_image_knn_score, |
| 176 | "primary_knn_score": primary_knn_score, | 231 | "primary_knn_score": primary_knn_score, |
| @@ -322,6 +377,10 @@ def _build_ltr_feature_block( | @@ -322,6 +377,10 @@ def _build_ltr_feature_block( | ||
| 322 | "text_support_score": float(text_components["support_text_score"]), | 377 | "text_support_score": float(text_components["support_text_score"]), |
| 323 | "text_knn_score": text_knn_score, | 378 | "text_knn_score": text_knn_score, |
| 324 | "image_knn_score": image_knn_score, | 379 | "image_knn_score": image_knn_score, |
| 380 | + "exact_text_knn_score": float(knn_components["exact_text_knn_score"]), | ||
| 381 | + "exact_image_knn_score": float(knn_components["exact_image_knn_score"]), | ||
| 382 | + "approx_text_knn_score": float(knn_components["approx_text_knn_score"]), | ||
| 383 | + "approx_image_knn_score": float(knn_components["approx_image_knn_score"]), | ||
| 325 | "knn_primary_score": float(knn_components["primary_knn_score"]), | 384 | "knn_primary_score": float(knn_components["primary_knn_score"]), |
| 326 | "knn_support_score": float(knn_components["support_knn_score"]), | 385 | "knn_support_score": float(knn_components["support_knn_score"]), |
| 327 | "has_text_match": source_score > 0.0, | 386 | "has_text_match": source_score > 0.0, |
| @@ -433,6 +492,8 @@ def coarse_resort_hits( | @@ -433,6 +492,8 @@ def coarse_resort_hits( | ||
| 433 | hit["_knn_score"] = knn_score | 492 | hit["_knn_score"] = knn_score |
| 434 | hit["_text_knn_score"] = knn_components["text_knn_score"] | 493 | hit["_text_knn_score"] = knn_components["text_knn_score"] |
| 435 | hit["_image_knn_score"] = knn_components["image_knn_score"] | 494 | hit["_image_knn_score"] = knn_components["image_knn_score"] |
| 495 | + hit["_exact_text_knn_score"] = knn_components["exact_text_knn_score"] | ||
| 496 | + hit["_exact_image_knn_score"] = knn_components["exact_image_knn_score"] | ||
| 436 | hit["_coarse_score"] = coarse_score | 497 | hit["_coarse_score"] = coarse_score |
| 437 | 498 | ||
| 438 | if debug: | 499 | if debug: |
| @@ -460,6 +521,12 @@ def coarse_resort_hits( | @@ -460,6 +521,12 @@ def coarse_resort_hits( | ||
| 460 | ), | 521 | ), |
| 461 | "text_knn_score": knn_components["text_knn_score"], | 522 | "text_knn_score": knn_components["text_knn_score"], |
| 462 | "image_knn_score": knn_components["image_knn_score"], | 523 | "image_knn_score": knn_components["image_knn_score"], |
| 524 | + "exact_text_knn_score": knn_components["exact_text_knn_score"], | ||
| 525 | + "exact_image_knn_score": knn_components["exact_image_knn_score"], | ||
| 526 | + "approx_text_knn_score": knn_components["approx_text_knn_score"], | ||
| 527 | + "approx_image_knn_score": knn_components["approx_image_knn_score"], | ||
| 528 | + "text_knn_source": knn_components["text_knn_source"], | ||
| 529 | + "image_knn_source": knn_components["image_knn_source"], | ||
| 463 | "weighted_text_knn_score": knn_components["weighted_text_knn_score"], | 530 | "weighted_text_knn_score": knn_components["weighted_text_knn_score"], |
| 464 | "weighted_image_knn_score": knn_components["weighted_image_knn_score"], | 531 | "weighted_image_knn_score": knn_components["weighted_image_knn_score"], |
| 465 | "knn_primary_score": knn_components["primary_knn_score"], | 532 | "knn_primary_score": knn_components["primary_knn_score"], |
| @@ -557,6 +624,8 @@ def fuse_scores_and_resort( | @@ -557,6 +624,8 @@ def fuse_scores_and_resort( | ||
| 557 | hit["_knn_score"] = knn_score | 624 | hit["_knn_score"] = knn_score |
| 558 | hit["_text_knn_score"] = knn_components["text_knn_score"] | 625 | hit["_text_knn_score"] = knn_components["text_knn_score"] |
| 559 | hit["_image_knn_score"] = knn_components["image_knn_score"] | 626 | hit["_image_knn_score"] = knn_components["image_knn_score"] |
| 627 | + hit["_exact_text_knn_score"] = knn_components["exact_text_knn_score"] | ||
| 628 | + hit["_exact_image_knn_score"] = knn_components["exact_image_knn_score"] | ||
| 560 | hit["_fused_score"] = fused | 629 | hit["_fused_score"] = fused |
| 561 | hit["_style_intent_selected_sku_boost"] = style_boost | 630 | hit["_style_intent_selected_sku_boost"] = style_boost |
| 562 | 631 | ||
| @@ -589,6 +658,12 @@ def fuse_scores_and_resort( | @@ -589,6 +658,12 @@ def fuse_scores_and_resort( | ||
| 589 | "text_support_score": text_components["support_text_score"], | 658 | "text_support_score": text_components["support_text_score"], |
| 590 | "text_knn_score": knn_components["text_knn_score"], | 659 | "text_knn_score": knn_components["text_knn_score"], |
| 591 | "image_knn_score": knn_components["image_knn_score"], | 660 | "image_knn_score": knn_components["image_knn_score"], |
| 661 | + "exact_text_knn_score": knn_components["exact_text_knn_score"], | ||
| 662 | + "exact_image_knn_score": knn_components["exact_image_knn_score"], | ||
| 663 | + "approx_text_knn_score": knn_components["approx_text_knn_score"], | ||
| 664 | + "approx_image_knn_score": knn_components["approx_image_knn_score"], | ||
| 665 | + "text_knn_source": knn_components["text_knn_source"], | ||
| 666 | + "image_knn_source": knn_components["image_knn_source"], | ||
| 592 | "weighted_text_knn_score": knn_components["weighted_text_knn_score"], | 667 | "weighted_text_knn_score": knn_components["weighted_text_knn_score"], |
| 593 | "weighted_image_knn_score": knn_components["weighted_image_knn_score"], | 668 | "weighted_image_knn_score": knn_components["weighted_image_knn_score"], |
| 594 | "knn_primary_score": knn_components["primary_knn_score"], | 669 | "knn_primary_score": knn_components["primary_knn_score"], |
| @@ -744,6 +819,8 @@ def run_lightweight_rerank( | @@ -744,6 +819,8 @@ def run_lightweight_rerank( | ||
| 744 | hit["_knn_score"] = knn_score | 819 | hit["_knn_score"] = knn_score |
| 745 | hit["_text_knn_score"] = signal_bundle["knn_components"]["text_knn_score"] | 820 | hit["_text_knn_score"] = signal_bundle["knn_components"]["text_knn_score"] |
| 746 | hit["_image_knn_score"] = signal_bundle["knn_components"]["image_knn_score"] | 821 | hit["_image_knn_score"] = signal_bundle["knn_components"]["image_knn_score"] |
| 822 | + hit["_exact_text_knn_score"] = signal_bundle["knn_components"]["exact_text_knn_score"] | ||
| 823 | + hit["_exact_image_knn_score"] = signal_bundle["knn_components"]["exact_image_knn_score"] | ||
| 747 | hit["_style_intent_selected_sku_boost"] = style_boost | 824 | hit["_style_intent_selected_sku_boost"] = style_boost |
| 748 | 825 | ||
| 749 | if debug: | 826 | if debug: |
search/searcher.py
| @@ -236,6 +236,117 @@ class Searcher: | @@ -236,6 +236,117 @@ class Searcher: | ||
| 236 | return | 236 | return |
| 237 | es_query["_source"] = {"includes": self.source_fields} | 237 | es_query["_source"] = {"includes": self.source_fields} |
| 238 | 238 | ||
| 239 | + def _resolve_exact_knn_rescore_window(self) -> int: | ||
| 240 | + configured = int(self.config.rerank.exact_knn_rescore_window) | ||
| 241 | + if configured > 0: | ||
| 242 | + return configured | ||
| 243 | + return int(self.config.rerank.rerank_window) | ||
| 244 | + | ||
| 245 | + @staticmethod | ||
| 246 | + def _vector_to_list(vector: Any) -> List[float]: | ||
| 247 | + if vector is None: | ||
| 248 | + return [] | ||
| 249 | + if hasattr(vector, "tolist"): | ||
| 250 | + values = vector.tolist() | ||
| 251 | + else: | ||
| 252 | + values = list(vector) | ||
| 253 | + return [float(v) for v in values] | ||
| 254 | + | ||
| 255 | + def _build_exact_knn_rescore( | ||
| 256 | + self, | ||
| 257 | + *, | ||
| 258 | + query_vector: Any, | ||
| 259 | + image_query_vector: Any, | ||
| 260 | + ) -> Optional[Dict[str, Any]]: | ||
| 261 | + clauses: List[Dict[str, Any]] = [] | ||
| 262 | + | ||
| 263 | + if query_vector is not None and self.text_embedding_field: | ||
| 264 | + clauses.append( | ||
| 265 | + { | ||
| 266 | + "script_score": { | ||
| 267 | + "_name": "exact_text_knn_query", | ||
| 268 | + "query": {"exists": {"field": self.text_embedding_field}}, | ||
| 269 | + "script": { | ||
| 270 | + # Keep exact score on the same [0, 1]-ish scale as KNN dot_product recall. | ||
| 271 | + "source": ( | ||
| 272 | + f"(dotProduct(params.query_vector, '{self.text_embedding_field}') + 1.0) / 2.0" | ||
| 273 | + ), | ||
| 274 | + "params": {"query_vector": self._vector_to_list(query_vector)}, | ||
| 275 | + }, | ||
| 276 | + } | ||
| 277 | + } | ||
| 278 | + ) | ||
| 279 | + | ||
| 280 | + if image_query_vector is not None and self.image_embedding_field: | ||
| 281 | + nested_path, _, _ = str(self.image_embedding_field).rpartition(".") | ||
| 282 | + if nested_path: | ||
| 283 | + clauses.append( | ||
| 284 | + { | ||
| 285 | + "nested": { | ||
| 286 | + "path": nested_path, | ||
| 287 | + "_name": "exact_image_knn_query", | ||
| 288 | + "score_mode": "max", | ||
| 289 | + "query": { | ||
| 290 | + "script_score": { | ||
| 291 | + "query": {"exists": {"field": self.image_embedding_field}}, | ||
| 292 | + "script": { | ||
| 293 | + # Keep exact score on the same [0, 1]-ish scale as KNN dot_product recall. | ||
| 294 | + "source": ( | ||
| 295 | + f"(dotProduct(params.query_vector, '{self.image_embedding_field}') + 1.0) / 2.0" | ||
| 296 | + ), | ||
| 297 | + "params": { | ||
| 298 | + "query_vector": self._vector_to_list(image_query_vector), | ||
| 299 | + }, | ||
| 300 | + }, | ||
| 301 | + } | ||
| 302 | + }, | ||
| 303 | + } | ||
| 304 | + } | ||
| 305 | + ) | ||
| 306 | + | ||
| 307 | + if not clauses: | ||
| 308 | + return None | ||
| 309 | + | ||
| 310 | + return { | ||
| 311 | + "window_size": self._resolve_exact_knn_rescore_window(), | ||
| 312 | + "query": { | ||
| 313 | + # Phase 1: only compute exact vector scores and expose them in matched_queries. | ||
| 314 | + "score_mode": "total", | ||
| 315 | + "query_weight": 1.0, | ||
| 316 | + "rescore_query_weight": 0.0, | ||
| 317 | + "rescore_query": { | ||
| 318 | + "bool": { | ||
| 319 | + "should": clauses, | ||
| 320 | + "minimum_should_match": 1, | ||
| 321 | + } | ||
| 322 | + }, | ||
| 323 | + }, | ||
| 324 | + } | ||
| 325 | + | ||
| 326 | + def _attach_exact_knn_rescore( | ||
| 327 | + self, | ||
| 328 | + es_query: Dict[str, Any], | ||
| 329 | + *, | ||
| 330 | + in_rank_window: bool, | ||
| 331 | + query_vector: Any, | ||
| 332 | + image_query_vector: Any, | ||
| 333 | + ) -> None: | ||
| 334 | + if not in_rank_window or not self.config.rerank.exact_knn_rescore_enabled: | ||
| 335 | + return | ||
| 336 | + rescore = self._build_exact_knn_rescore( | ||
| 337 | + query_vector=query_vector, | ||
| 338 | + image_query_vector=image_query_vector, | ||
| 339 | + ) | ||
| 340 | + if not rescore: | ||
| 341 | + return | ||
| 342 | + existing = es_query.get("rescore") | ||
| 343 | + if existing is None: | ||
| 344 | + es_query["rescore"] = rescore | ||
| 345 | + elif isinstance(existing, list): | ||
| 346 | + es_query["rescore"] = [*existing, rescore] | ||
| 347 | + else: | ||
| 348 | + es_query["rescore"] = [existing, rescore] | ||
| 349 | + | ||
| 239 | def _resolve_rerank_source_filter( | 350 | def _resolve_rerank_source_filter( |
| 240 | self, | 351 | self, |
| 241 | doc_template: str, | 352 | doc_template: str, |
| @@ -573,6 +684,12 @@ class Searcher: | @@ -573,6 +684,12 @@ class Searcher: | ||
| 573 | min_score=min_score, | 684 | min_score=min_score, |
| 574 | parsed_query=parsed_query, | 685 | parsed_query=parsed_query, |
| 575 | ) | 686 | ) |
| 687 | + self._attach_exact_knn_rescore( | ||
| 688 | + es_query, | ||
| 689 | + in_rank_window=in_rank_window, | ||
| 690 | + query_vector=parsed_query.query_vector if enable_embedding else None, | ||
| 691 | + image_query_vector=image_query_vector, | ||
| 692 | + ) | ||
| 576 | 693 | ||
| 577 | # Add facets for faceted search | 694 | # Add facets for faceted search |
| 578 | if facets: | 695 | if facets: |
| @@ -1430,6 +1547,12 @@ class Searcher: | @@ -1430,6 +1547,12 @@ class Searcher: | ||
| 1430 | "es_fetch_size": es_fetch_size, | 1547 | "es_fetch_size": es_fetch_size, |
| 1431 | "in_rank_window": in_rank_window, | 1548 | "in_rank_window": in_rank_window, |
| 1432 | "include_named_queries_score": bool(in_rank_window), | 1549 | "include_named_queries_score": bool(in_rank_window), |
| 1550 | + "exact_knn_rescore_enabled": bool(rc.exact_knn_rescore_enabled and in_rank_window), | ||
| 1551 | + "exact_knn_rescore_window": ( | ||
| 1552 | + self._resolve_exact_knn_rescore_window() | ||
| 1553 | + if rc.exact_knn_rescore_enabled and in_rank_window | ||
| 1554 | + else None | ||
| 1555 | + ), | ||
| 1433 | }, | 1556 | }, |
| 1434 | "es_response": { | 1557 | "es_response": { |
| 1435 | "took_ms": es_response.get('took', 0), | 1558 | "took_ms": es_response.get('took', 0), |
tests/test_rerank_client.py
| @@ -172,6 +172,57 @@ def test_fuse_scores_and_resort_uses_max_of_text_and_image_knn_scores(): | @@ -172,6 +172,57 @@ def test_fuse_scores_and_resort_uses_max_of_text_and_image_knn_scores(): | ||
| 172 | assert isclose(debug[0]["image_knn_score"], 0.7, rel_tol=1e-9) | 172 | assert isclose(debug[0]["image_knn_score"], 0.7, rel_tol=1e-9) |
| 173 | 173 | ||
| 174 | 174 | ||
| 175 | +def test_fuse_scores_and_resort_prefers_exact_knn_scores_over_ann_scores(): | ||
| 176 | + hits = [ | ||
| 177 | + { | ||
| 178 | + "_id": "exact-mm-hit", | ||
| 179 | + "_score": 1.0, | ||
| 180 | + "matched_queries": { | ||
| 181 | + "base_query": 1.5, | ||
| 182 | + "knn_query": 0.2, | ||
| 183 | + "image_knn_query": 0.7, | ||
| 184 | + "exact_text_knn_query": 0.9, | ||
| 185 | + "exact_image_knn_query": 0.1, | ||
| 186 | + }, | ||
| 187 | + } | ||
| 188 | + ] | ||
| 189 | + | ||
| 190 | + debug = fuse_scores_and_resort(hits, [0.8], debug=True) | ||
| 191 | + | ||
| 192 | + assert isclose(hits[0]["_knn_score"], 0.9, rel_tol=1e-9) | ||
| 193 | + assert isclose(debug[0]["text_knn_score"], 0.9, rel_tol=1e-9) | ||
| 194 | + assert isclose(debug[0]["image_knn_score"], 0.1, rel_tol=1e-9) | ||
| 195 | + assert isclose(debug[0]["exact_text_knn_score"], 0.9, rel_tol=1e-9) | ||
| 196 | + assert isclose(debug[0]["exact_image_knn_score"], 0.1, rel_tol=1e-9) | ||
| 197 | + assert isclose(debug[0]["approx_text_knn_score"], 0.2, rel_tol=1e-9) | ||
| 198 | + assert isclose(debug[0]["approx_image_knn_score"], 0.7, rel_tol=1e-9) | ||
| 199 | + assert debug[0]["text_knn_source"] == "exact_text_knn_query" | ||
| 200 | + assert debug[0]["image_knn_source"] == "exact_image_knn_query" | ||
| 201 | + | ||
| 202 | + | ||
| 203 | +def test_fuse_scores_and_resort_falls_back_to_ann_when_exact_knn_missing(): | ||
| 204 | + hits = [ | ||
| 205 | + { | ||
| 206 | + "_id": "ann-only-hit", | ||
| 207 | + "_score": 1.0, | ||
| 208 | + "matched_queries": { | ||
| 209 | + "base_query": 1.5, | ||
| 210 | + "knn_query": 0.4, | ||
| 211 | + "image_knn_query": 0.5, | ||
| 212 | + }, | ||
| 213 | + } | ||
| 214 | + ] | ||
| 215 | + | ||
| 216 | + debug = fuse_scores_and_resort(hits, [0.8], debug=True) | ||
| 217 | + | ||
| 218 | + assert isclose(debug[0]["text_knn_score"], 0.4, rel_tol=1e-9) | ||
| 219 | + assert isclose(debug[0]["image_knn_score"], 0.5, rel_tol=1e-9) | ||
| 220 | + assert isclose(debug[0]["approx_text_knn_score"], 0.4, rel_tol=1e-9) | ||
| 221 | + assert isclose(debug[0]["approx_image_knn_score"], 0.5, rel_tol=1e-9) | ||
| 222 | + assert debug[0]["text_knn_source"] == "knn_query" | ||
| 223 | + assert debug[0]["image_knn_source"] == "image_knn_query" | ||
| 224 | + | ||
| 225 | + | ||
| 175 | def test_fuse_scores_and_resort_applies_knn_dismax_weights_and_tie_breaker(): | 226 | def test_fuse_scores_and_resort_applies_knn_dismax_weights_and_tie_breaker(): |
| 176 | hits = [ | 227 | hits = [ |
| 177 | { | 228 | { |
tests/test_search_rerank_window.py
| @@ -197,13 +197,24 @@ class _FakeESClient: | @@ -197,13 +197,24 @@ class _FakeESClient: | ||
| 197 | } | 197 | } |
| 198 | 198 | ||
| 199 | 199 | ||
| 200 | -def _build_search_config(*, rerank_enabled: bool = True, rerank_window: int = 384): | 200 | +def _build_search_config( |
| 201 | + *, | ||
| 202 | + rerank_enabled: bool = True, | ||
| 203 | + rerank_window: int = 384, | ||
| 204 | + exact_knn_rescore_enabled: bool = False, | ||
| 205 | + exact_knn_rescore_window: int = 0, | ||
| 206 | +): | ||
| 201 | return SearchConfig( | 207 | return SearchConfig( |
| 202 | field_boosts={"title.en": 3.0}, | 208 | field_boosts={"title.en": 3.0}, |
| 203 | indexes=[IndexConfig(name="default", label="default", fields=["title.en"])], | 209 | indexes=[IndexConfig(name="default", label="default", fields=["title.en"])], |
| 204 | query_config=QueryConfig(enable_text_embedding=False, enable_query_rewrite=False), | 210 | query_config=QueryConfig(enable_text_embedding=False, enable_query_rewrite=False), |
| 205 | function_score=FunctionScoreConfig(), | 211 | function_score=FunctionScoreConfig(), |
| 206 | - rerank=RerankConfig(enabled=rerank_enabled, rerank_window=rerank_window), | 212 | + rerank=RerankConfig( |
| 213 | + enabled=rerank_enabled, | ||
| 214 | + rerank_window=rerank_window, | ||
| 215 | + exact_knn_rescore_enabled=exact_knn_rescore_enabled, | ||
| 216 | + exact_knn_rescore_window=exact_knn_rescore_window, | ||
| 217 | + ), | ||
| 207 | spu_config=SPUConfig(enabled=False), | 218 | spu_config=SPUConfig(enabled=False), |
| 208 | es_index_name="test_products", | 219 | es_index_name="test_products", |
| 209 | es_settings={}, | 220 | es_settings={}, |
| @@ -301,7 +312,11 @@ def test_config_loader_rerank_enabled_defaults_true(tmp_path: Path): | @@ -301,7 +312,11 @@ def test_config_loader_rerank_enabled_defaults_true(tmp_path: Path): | ||
| 301 | }, | 312 | }, |
| 302 | "spu_config": {"enabled": False}, | 313 | "spu_config": {"enabled": False}, |
| 303 | "function_score": {"score_mode": "sum", "boost_mode": "multiply", "functions": []}, | 314 | "function_score": {"score_mode": "sum", "boost_mode": "multiply", "functions": []}, |
| 304 | - "rerank": {"rerank_window": 384}, | 315 | + "rerank": { |
| 316 | + "rerank_window": 384, | ||
| 317 | + "exact_knn_rescore_enabled": True, | ||
| 318 | + "exact_knn_rescore_window": 160, | ||
| 319 | + }, | ||
| 305 | } | 320 | } |
| 306 | config_path = tmp_path / "config.yaml" | 321 | config_path = tmp_path / "config.yaml" |
| 307 | config_path.write_text(yaml.safe_dump(config_data), encoding="utf-8") | 322 | config_path.write_text(yaml.safe_dump(config_data), encoding="utf-8") |
| @@ -310,6 +325,8 @@ def test_config_loader_rerank_enabled_defaults_true(tmp_path: Path): | @@ -310,6 +325,8 @@ def test_config_loader_rerank_enabled_defaults_true(tmp_path: Path): | ||
| 310 | loaded = loader.load_config(validate=False) | 325 | loaded = loader.load_config(validate=False) |
| 311 | 326 | ||
| 312 | assert loaded.rerank.enabled is True | 327 | assert loaded.rerank.enabled is True |
| 328 | + assert loaded.rerank.exact_knn_rescore_enabled is True | ||
| 329 | + assert loaded.rerank.exact_knn_rescore_window == 160 | ||
| 313 | 330 | ||
| 314 | 331 | ||
| 315 | def test_config_loader_parses_named_rerank_instances(tmp_path: Path): | 332 | def test_config_loader_parses_named_rerank_instances(tmp_path: Path): |
| @@ -1028,6 +1045,131 @@ def test_searcher_debug_info_uses_initial_es_max_score_for_normalization(monkeyp | @@ -1028,6 +1045,131 @@ def test_searcher_debug_info_uses_initial_es_max_score_for_normalization(monkeyp | ||
| 1028 | assert result.debug_info["per_result"][1]["es_score_normalized"] == 2.0 / 3.0 | 1045 | assert result.debug_info["per_result"][1]["es_score_normalized"] == 2.0 / 3.0 |
| 1029 | 1046 | ||
| 1030 | 1047 | ||
| 1048 | +def test_searcher_attaches_exact_knn_rescore_for_rank_window(monkeypatch): | ||
| 1049 | + class _VectorQueryParser: | ||
| 1050 | + def parse(self, query: str, tenant_id: str, generate_vector: bool, context: Any, target_languages: Any = None): | ||
| 1051 | + return _FakeParsedQuery( | ||
| 1052 | + original_query=query, | ||
| 1053 | + query_normalized=query, | ||
| 1054 | + rewritten_query=query, | ||
| 1055 | + translations={}, | ||
| 1056 | + query_vector=np.array([0.1, 0.2, 0.3], dtype=np.float32), | ||
| 1057 | + image_query_vector=np.array([0.4, 0.5, 0.6], dtype=np.float32), | ||
| 1058 | + ) | ||
| 1059 | + | ||
| 1060 | + es_client = _FakeESClient(total_hits=5) | ||
| 1061 | + base = _build_search_config( | ||
| 1062 | + rerank_enabled=True, | ||
| 1063 | + rerank_window=5, | ||
| 1064 | + exact_knn_rescore_enabled=True, | ||
| 1065 | + exact_knn_rescore_window=3, | ||
| 1066 | + ) | ||
| 1067 | + config = SearchConfig( | ||
| 1068 | + field_boosts=base.field_boosts, | ||
| 1069 | + indexes=base.indexes, | ||
| 1070 | + query_config=QueryConfig( | ||
| 1071 | + enable_text_embedding=True, | ||
| 1072 | + enable_query_rewrite=False, | ||
| 1073 | + text_embedding_field="title_embedding", | ||
| 1074 | + image_embedding_field="image_embedding.vector", | ||
| 1075 | + ), | ||
| 1076 | + function_score=base.function_score, | ||
| 1077 | + coarse_rank=base.coarse_rank, | ||
| 1078 | + fine_rank=FineRankConfig(enabled=False, input_window=5, output_window=5), | ||
| 1079 | + rerank=base.rerank, | ||
| 1080 | + spu_config=base.spu_config, | ||
| 1081 | + es_index_name=base.es_index_name, | ||
| 1082 | + es_settings=base.es_settings, | ||
| 1083 | + ) | ||
| 1084 | + searcher = _build_searcher(config, es_client) | ||
| 1085 | + searcher.query_parser = _VectorQueryParser() | ||
| 1086 | + context = create_request_context(reqid="exact-rescore", uid="u-exact") | ||
| 1087 | + | ||
| 1088 | + monkeypatch.setattr( | ||
| 1089 | + "search.searcher.get_tenant_config_loader", | ||
| 1090 | + lambda: SimpleNamespace(get_tenant_config=lambda tenant_id: {"index_languages": ["en"]}), | ||
| 1091 | + ) | ||
| 1092 | + | ||
| 1093 | + searcher.search( | ||
| 1094 | + query="dress", | ||
| 1095 | + tenant_id="162", | ||
| 1096 | + from_=0, | ||
| 1097 | + size=2, | ||
| 1098 | + context=context, | ||
| 1099 | + enable_rerank=False, | ||
| 1100 | + debug=True, | ||
| 1101 | + ) | ||
| 1102 | + | ||
| 1103 | + body = es_client.calls[0]["body"] | ||
| 1104 | + assert body["rescore"]["window_size"] == 3 | ||
| 1105 | + assert body["rescore"]["query"]["score_mode"] == "total" | ||
| 1106 | + assert body["rescore"]["query"]["rescore_query_weight"] == 0.0 | ||
| 1107 | + should = body["rescore"]["query"]["rescore_query"]["bool"]["should"] | ||
| 1108 | + names = [] | ||
| 1109 | + for clause in should: | ||
| 1110 | + if "script_score" in clause: | ||
| 1111 | + names.append(clause["script_score"]["_name"]) | ||
| 1112 | + elif "nested" in clause: | ||
| 1113 | + names.append(clause["nested"]["_name"]) | ||
| 1114 | + assert names == ["exact_text_knn_query", "exact_image_knn_query"] | ||
| 1115 | + | ||
| 1116 | + | ||
| 1117 | +def test_searcher_skips_exact_knn_rescore_outside_rank_window(monkeypatch): | ||
| 1118 | + class _VectorQueryParser: | ||
| 1119 | + def parse(self, query: str, tenant_id: str, generate_vector: bool, context: Any, target_languages: Any = None): | ||
| 1120 | + return _FakeParsedQuery( | ||
| 1121 | + original_query=query, | ||
| 1122 | + query_normalized=query, | ||
| 1123 | + rewritten_query=query, | ||
| 1124 | + translations={}, | ||
| 1125 | + query_vector=np.array([0.1, 0.2, 0.3], dtype=np.float32), | ||
| 1126 | + ) | ||
| 1127 | + | ||
| 1128 | + es_client = _FakeESClient(total_hits=20) | ||
| 1129 | + base = _build_search_config( | ||
| 1130 | + rerank_enabled=True, | ||
| 1131 | + rerank_window=5, | ||
| 1132 | + exact_knn_rescore_enabled=True, | ||
| 1133 | + exact_knn_rescore_window=4, | ||
| 1134 | + ) | ||
| 1135 | + config = SearchConfig( | ||
| 1136 | + field_boosts=base.field_boosts, | ||
| 1137 | + indexes=base.indexes, | ||
| 1138 | + query_config=QueryConfig( | ||
| 1139 | + enable_text_embedding=True, | ||
| 1140 | + enable_query_rewrite=False, | ||
| 1141 | + text_embedding_field="title_embedding", | ||
| 1142 | + ), | ||
| 1143 | + function_score=base.function_score, | ||
| 1144 | + coarse_rank=base.coarse_rank, | ||
| 1145 | + fine_rank=FineRankConfig(enabled=False, input_window=5, output_window=5), | ||
| 1146 | + rerank=base.rerank, | ||
| 1147 | + spu_config=base.spu_config, | ||
| 1148 | + es_index_name=base.es_index_name, | ||
| 1149 | + es_settings=base.es_settings, | ||
| 1150 | + ) | ||
| 1151 | + searcher = _build_searcher(config, es_client) | ||
| 1152 | + searcher.query_parser = _VectorQueryParser() | ||
| 1153 | + context = create_request_context(reqid="exact-rescore-off", uid="u-exact-off") | ||
| 1154 | + | ||
| 1155 | + monkeypatch.setattr( | ||
| 1156 | + "search.searcher.get_tenant_config_loader", | ||
| 1157 | + lambda: SimpleNamespace(get_tenant_config=lambda tenant_id: {"index_languages": ["en"]}), | ||
| 1158 | + ) | ||
| 1159 | + | ||
| 1160 | + searcher.search( | ||
| 1161 | + query="dress", | ||
| 1162 | + tenant_id="162", | ||
| 1163 | + from_=5, | ||
| 1164 | + size=2, | ||
| 1165 | + context=context, | ||
| 1166 | + enable_rerank=False, | ||
| 1167 | + ) | ||
| 1168 | + | ||
| 1169 | + body = es_client.calls[0]["body"] | ||
| 1170 | + assert "rescore" not in body | ||
| 1171 | + | ||
| 1172 | + | ||
| 1031 | def test_searcher_rerank_rank_change_falls_back_to_coarse_rank_when_fine_disabled(monkeypatch): | 1173 | def test_searcher_rerank_rank_change_falls_back_to_coarse_rank_when_fine_disabled(monkeypatch): |
| 1032 | es_client = _FakeESClient(total_hits=5) | 1174 | es_client = _FakeESClient(total_hits=5) |
| 1033 | config = _build_search_config(rerank_enabled=True, rerank_window=5) | 1175 | config = _build_search_config(rerank_enabled=True, rerank_window=5) |