|
| 1 | +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. |
| 2 | +# SPDX-License-Identifier: AGPL-3.0 |
| 3 | +"""Keyword/dense score fusion for find/search. |
| 4 | +
|
| 5 | +The local FTS5 sidecar provides search-time BM25 recall for exact tokens (code |
| 6 | +names, acronyms, tickers, version strings) that dense retrieval handles poorly. |
| 7 | +This module merges dense ``MatchedContext`` results with keyword candidates |
| 8 | +using Reciprocal Rank Fusion (robust, no score calibration) or a weighted blend |
| 9 | +of normalized scores. |
| 10 | +""" |
| 11 | + |
| 12 | +from __future__ import annotations |
| 13 | + |
| 14 | +from dataclasses import replace |
| 15 | +from typing import Any, Awaitable, Callable, Dict, List, Optional, Sequence, Tuple |
| 16 | + |
| 17 | +from openviking_cli.retrieve.types import MatchedContext |
| 18 | + |
| 19 | +ReadAbstract = Callable[[str], Awaitable[str]] |
| 20 | + |
| 21 | + |
| 22 | +class HybridKeywordRecaller: |
| 23 | + """Fuse keyword-sidecar recall into a dense retrieval result list.""" |
| 24 | + |
| 25 | + def __init__( |
| 26 | + self, |
| 27 | + keyword_fs: Any, |
| 28 | + hybrid_config: Any, |
| 29 | + keyword_config: Optional[Any] = None, |
| 30 | + ): |
| 31 | + self._keyword_fs = keyword_fs |
| 32 | + self._config = hybrid_config |
| 33 | + self._keyword_config = keyword_config |
| 34 | + |
| 35 | + def enabled(self, ctx: Any = None) -> bool: |
| 36 | + if self._keyword_fs is None or self._config is None: |
| 37 | + return False |
| 38 | + if not getattr(self._config, "enabled", False): |
| 39 | + return False |
| 40 | + account_id = getattr(ctx, "account_id", None) or "default" |
| 41 | + try: |
| 42 | + return self._keyword_fs.is_ready(account_id) |
| 43 | + except Exception: |
| 44 | + return False |
| 45 | + |
| 46 | + async def enhance( |
| 47 | + self, |
| 48 | + query: str, |
| 49 | + dense: Sequence[MatchedContext], |
| 50 | + scope_uris: Sequence[str], |
| 51 | + ctx: Any, |
| 52 | + limit: int, |
| 53 | + exclude_uri: str = "", |
| 54 | + read_abstract: Optional[ReadAbstract] = None, |
| 55 | + ) -> List[MatchedContext]: |
| 56 | + """Merge keyword candidates into ``dense`` and return the fused top ``limit``.""" |
| 57 | + if not self.enabled(ctx) or not query: |
| 58 | + return list(dense) |
| 59 | + candidates = await self._recall(query, scope_uris, exclude_uri, ctx, limit) |
| 60 | + if not candidates: |
| 61 | + return list(dense) |
| 62 | + fused = self._fuse(dense, candidates, limit) |
| 63 | + # Enrich keyword-only hits with best-effort abstract text. |
| 64 | + dense_uris = {m.uri for m in dense} |
| 65 | + for mc in fused: |
| 66 | + if mc.uri in dense_uris or not mc.abstract: |
| 67 | + if read_abstract is not None and not mc.abstract: |
| 68 | + try: |
| 69 | + mc.abstract = await read_abstract(mc.uri) |
| 70 | + except Exception: |
| 71 | + pass |
| 72 | + return fused |
| 73 | + |
| 74 | + async def _recall( |
| 75 | + self, |
| 76 | + query: str, |
| 77 | + scope_uris: Sequence[str], |
| 78 | + exclude_uri: str, |
| 79 | + ctx: Any, |
| 80 | + limit: int, |
| 81 | + ) -> List[Tuple[str, float]]: |
| 82 | + account_id = getattr(ctx, "account_id", None) or "default" |
| 83 | + collected: Dict[str, float] = {} |
| 84 | + scopes = list(scope_uris) or [""] |
| 85 | + for scope in scopes: |
| 86 | + try: |
| 87 | + hits = self._keyword_fs.lookup( |
| 88 | + account_id=account_id, |
| 89 | + query=query, |
| 90 | + scope_uri=scope, |
| 91 | + exclude_uri=exclude_uri, |
| 92 | + limit=max(limit * 3, 30), |
| 93 | + ) |
| 94 | + except Exception: |
| 95 | + continue |
| 96 | + for uri, score in hits: |
| 97 | + # Keep the best (lowest) bm25 score for a URI across scopes. |
| 98 | + if uri not in collected or score < collected[uri]: |
| 99 | + collected[uri] = score |
| 100 | + ranked = sorted(collected.items(), key=lambda x: x[1]) |
| 101 | + return ranked[: max(limit * 3, 30)] |
| 102 | + |
| 103 | + def _fuse( |
| 104 | + self, |
| 105 | + dense: Sequence[MatchedContext], |
| 106 | + candidates: Sequence[Tuple[str, float]], |
| 107 | + limit: int, |
| 108 | + ) -> List[MatchedContext]: |
| 109 | + fusion = getattr(self._config, "fusion", "rrf") |
| 110 | + if fusion == "weighted": |
| 111 | + return self._fuse_weighted(dense, candidates, limit) |
| 112 | + return self._fuse_rrf(dense, candidates, limit) |
| 113 | + |
| 114 | + def _fuse_rrf( |
| 115 | + self, |
| 116 | + dense: Sequence[MatchedContext], |
| 117 | + candidates: Sequence[Tuple[str, float]], |
| 118 | + limit: int, |
| 119 | + ) -> List[MatchedContext]: |
| 120 | + k = float(getattr(self._config, "rrf_k", 60.0) or 60.0) |
| 121 | + dense_rank = {m.uri: i for i, m in enumerate(dense)} |
| 122 | + kw_rank = {uri: i for i, (uri, _s) in enumerate(candidates)} |
| 123 | + scores: Dict[str, float] = {} |
| 124 | + for uri in dense_rank: |
| 125 | + scores[uri] = 1.0 / (k + dense_rank[uri] + 1) |
| 126 | + for uri in kw_rank: |
| 127 | + scores[uri] = scores.get(uri, 0.0) + 1.0 / (k + kw_rank[uri] + 1) |
| 128 | + |
| 129 | + by_uri = {m.uri: m for m in dense} |
| 130 | + ordered_uris = sorted(scores.keys(), key=lambda u: (-scores[u], dense_rank.get(u, 10**9))) |
| 131 | + out: List[MatchedContext] = [] |
| 132 | + for uri in ordered_uris: |
| 133 | + mc = by_uri.get(uri) |
| 134 | + if mc is None: |
| 135 | + mc = self._make_keyword_context(uri, candidates) |
| 136 | + out.append(replace(mc, score=scores[uri])) |
| 137 | + if len(out) >= limit: |
| 138 | + break |
| 139 | + # Always keep the dense ordering for ties handled above; append leftover |
| 140 | + # dense results only if there is still headroom. |
| 141 | + if len(out) < limit: |
| 142 | + seen = {m.uri for m in out} |
| 143 | + for m in dense: |
| 144 | + if m.uri not in seen: |
| 145 | + out.append(m) |
| 146 | + seen.add(m.uri) |
| 147 | + if len(out) >= limit: |
| 148 | + break |
| 149 | + return out |
| 150 | + |
| 151 | + def _fuse_weighted( |
| 152 | + self, |
| 153 | + dense: Sequence[MatchedContext], |
| 154 | + candidates: Sequence[Tuple[str, float]], |
| 155 | + limit: int, |
| 156 | + ) -> List[MatchedContext]: |
| 157 | + w = float(getattr(self._config, "keyword_weight", 0.3) or 0.3) |
| 158 | + raw_scores = {uri: score for uri, score in candidates} |
| 159 | + if raw_scores: |
| 160 | + lo = min(raw_scores.values()) |
| 161 | + hi = max(raw_scores.values()) |
| 162 | + else: |
| 163 | + lo = hi = 0.0 |
| 164 | + |
| 165 | + def norm(uri: str) -> float: |
| 166 | + if hi > lo: |
| 167 | + return (raw_scores[uri] - lo) / (hi - lo) |
| 168 | + return 0.5 |
| 169 | + |
| 170 | + by_uri = {m.uri: m for m in dense} |
| 171 | + scores: Dict[str, float] = {} |
| 172 | + for m in dense: |
| 173 | + kw = norm(m.uri) if m.uri in raw_scores else 0.0 |
| 174 | + scores[m.uri] = (1 - w) * m.score + w * kw |
| 175 | + for uri in raw_scores: |
| 176 | + if uri not in scores: |
| 177 | + scores[uri] = w * norm(uri) |
| 178 | + |
| 179 | + ordered = sorted(scores.keys(), key=lambda u: -scores[u]) |
| 180 | + out: List[MatchedContext] = [] |
| 181 | + for uri in ordered: |
| 182 | + mc = by_uri.get(uri) |
| 183 | + if mc is None: |
| 184 | + mc = self._make_keyword_context(uri, candidates) |
| 185 | + out.append(replace(mc, score=scores[uri])) |
| 186 | + if len(out) >= limit: |
| 187 | + break |
| 188 | + return out |
| 189 | + |
| 190 | + def _make_keyword_context(self, uri: str, candidates: Sequence[Tuple[str, float]]) -> MatchedContext: |
| 191 | + from openviking.core.context import ContextType |
| 192 | + from openviking.core.namespace import context_type_for_uri |
| 193 | + |
| 194 | + score = 0.0 |
| 195 | + for u, s in candidates: |
| 196 | + if u == uri: |
| 197 | + score = s |
| 198 | + break |
| 199 | + ctype = context_type_for_uri(uri) |
| 200 | + try: |
| 201 | + context_type = ContextType(ctype) |
| 202 | + except ValueError: |
| 203 | + context_type = ContextType.RESOURCE |
| 204 | + return MatchedContext( |
| 205 | + uri=uri, |
| 206 | + context_type=context_type, |
| 207 | + level=2, |
| 208 | + abstract="", |
| 209 | + category="", |
| 210 | + score=score, |
| 211 | + match_reason="keyword", |
| 212 | + ) |
0 commit comments