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
config/loader.py
| ... | ... | @@ -608,6 +608,12 @@ class AppConfigLoader: |
| 608 | 608 | rerank=RerankConfig( |
| 609 | 609 | enabled=bool(rerank_cfg.get("enabled", True)), |
| 610 | 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 | 617 | timeout_sec=float(rerank_cfg.get("timeout_sec", 15.0)), |
| 612 | 618 | weight_es=float(rerank_cfg.get("weight_es", 0.4)), |
| 613 | 619 | weight_ai=float(rerank_cfg.get("weight_ai", 0.6)), | ... | ... |
config/schema.py
| ... | ... | @@ -176,6 +176,9 @@ class RerankConfig: |
| 176 | 176 | |
| 177 | 177 | enabled: bool = True |
| 178 | 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 | 182 | timeout_sec: float = 15.0 |
| 180 | 183 | weight_es: float = 0.4 |
| 181 | 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 | 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 @@ |
| 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 | 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 | 191 | def _collect_knn_score_components( |
| 157 | 192 | matched_queries: Any, |
| 158 | 193 | fusion: RerankFusionConfig, |
| 159 | 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 | 210 | weighted_text_knn_score = text_knn_score * float(fusion.knn_text_weight) |
| 164 | 211 | weighted_image_knn_score = image_knn_score * float(fusion.knn_image_weight) |
| ... | ... | @@ -171,6 +218,14 @@ def _collect_knn_score_components( |
| 171 | 218 | return { |
| 172 | 219 | "text_knn_score": text_knn_score, |
| 173 | 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 | 229 | "weighted_text_knn_score": weighted_text_knn_score, |
| 175 | 230 | "weighted_image_knn_score": weighted_image_knn_score, |
| 176 | 231 | "primary_knn_score": primary_knn_score, |
| ... | ... | @@ -322,6 +377,10 @@ def _build_ltr_feature_block( |
| 322 | 377 | "text_support_score": float(text_components["support_text_score"]), |
| 323 | 378 | "text_knn_score": text_knn_score, |
| 324 | 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 | 384 | "knn_primary_score": float(knn_components["primary_knn_score"]), |
| 326 | 385 | "knn_support_score": float(knn_components["support_knn_score"]), |
| 327 | 386 | "has_text_match": source_score > 0.0, |
| ... | ... | @@ -433,6 +492,8 @@ def coarse_resort_hits( |
| 433 | 492 | hit["_knn_score"] = knn_score |
| 434 | 493 | hit["_text_knn_score"] = knn_components["text_knn_score"] |
| 435 | 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 | 497 | hit["_coarse_score"] = coarse_score |
| 437 | 498 | |
| 438 | 499 | if debug: |
| ... | ... | @@ -460,6 +521,12 @@ def coarse_resort_hits( |
| 460 | 521 | ), |
| 461 | 522 | "text_knn_score": knn_components["text_knn_score"], |
| 462 | 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 | 530 | "weighted_text_knn_score": knn_components["weighted_text_knn_score"], |
| 464 | 531 | "weighted_image_knn_score": knn_components["weighted_image_knn_score"], |
| 465 | 532 | "knn_primary_score": knn_components["primary_knn_score"], |
| ... | ... | @@ -557,6 +624,8 @@ def fuse_scores_and_resort( |
| 557 | 624 | hit["_knn_score"] = knn_score |
| 558 | 625 | hit["_text_knn_score"] = knn_components["text_knn_score"] |
| 559 | 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 | 629 | hit["_fused_score"] = fused |
| 561 | 630 | hit["_style_intent_selected_sku_boost"] = style_boost |
| 562 | 631 | |
| ... | ... | @@ -589,6 +658,12 @@ def fuse_scores_and_resort( |
| 589 | 658 | "text_support_score": text_components["support_text_score"], |
| 590 | 659 | "text_knn_score": knn_components["text_knn_score"], |
| 591 | 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 | 667 | "weighted_text_knn_score": knn_components["weighted_text_knn_score"], |
| 593 | 668 | "weighted_image_knn_score": knn_components["weighted_image_knn_score"], |
| 594 | 669 | "knn_primary_score": knn_components["primary_knn_score"], |
| ... | ... | @@ -744,6 +819,8 @@ def run_lightweight_rerank( |
| 744 | 819 | hit["_knn_score"] = knn_score |
| 745 | 820 | hit["_text_knn_score"] = signal_bundle["knn_components"]["text_knn_score"] |
| 746 | 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 | 824 | hit["_style_intent_selected_sku_boost"] = style_boost |
| 748 | 825 | |
| 749 | 826 | if debug: | ... | ... |
search/searcher.py
| ... | ... | @@ -236,6 +236,117 @@ class Searcher: |
| 236 | 236 | return |
| 237 | 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 | 350 | def _resolve_rerank_source_filter( |
| 240 | 351 | self, |
| 241 | 352 | doc_template: str, |
| ... | ... | @@ -573,6 +684,12 @@ class Searcher: |
| 573 | 684 | min_score=min_score, |
| 574 | 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 | 694 | # Add facets for faceted search |
| 578 | 695 | if facets: |
| ... | ... | @@ -1430,6 +1547,12 @@ class Searcher: |
| 1430 | 1547 | "es_fetch_size": es_fetch_size, |
| 1431 | 1548 | "in_rank_window": in_rank_window, |
| 1432 | 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 | 1557 | "es_response": { |
| 1435 | 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 | 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 | 226 | def test_fuse_scores_and_resort_applies_knn_dismax_weights_and_tie_breaker(): |
| 176 | 227 | hits = [ |
| 177 | 228 | { | ... | ... |
tests/test_search_rerank_window.py
| ... | ... | @@ -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 | 207 | return SearchConfig( |
| 202 | 208 | field_boosts={"title.en": 3.0}, |
| 203 | 209 | indexes=[IndexConfig(name="default", label="default", fields=["title.en"])], |
| 204 | 210 | query_config=QueryConfig(enable_text_embedding=False, enable_query_rewrite=False), |
| 205 | 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 | 218 | spu_config=SPUConfig(enabled=False), |
| 208 | 219 | es_index_name="test_products", |
| 209 | 220 | es_settings={}, |
| ... | ... | @@ -301,7 +312,11 @@ def test_config_loader_rerank_enabled_defaults_true(tmp_path: Path): |
| 301 | 312 | }, |
| 302 | 313 | "spu_config": {"enabled": False}, |
| 303 | 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 | 321 | config_path = tmp_path / "config.yaml" |
| 307 | 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 | 325 | loaded = loader.load_config(validate=False) |
| 311 | 326 | |
| 312 | 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 | 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 | 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 | 1173 | def test_searcher_rerank_rank_change_falls_back_to_coarse_rank_when_fine_disabled(monkeypatch): |
| 1032 | 1174 | es_client = _FakeESClient(total_hits=5) |
| 1033 | 1175 | config = _build_search_config(rerank_enabled=True, rerank_window=5) | ... | ... |