| 1 | """Weighted reciprocal rank fusion for per-(subquery, source) streams.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from urllib.parse import parse_qs, urlencode, urlparse, urlunparse |
| 6 | |
| 7 | from . import schema |
| 8 | |
| 9 | # Standard RRF smoothing constant (Cormack et al. 2009) |
| 10 | RRF_K = 60 |
| 11 | |
| 12 | |
| 13 | def _candidate_sort_key(c: schema.Candidate) -> tuple: |
| 14 | return (-c.rrf_score, -c.local_relevance, -c.freshness, schema.candidate_source_label(c), c.title) |
| 15 | |
| 16 | |
| 17 | def _normalize_url(url: str) -> str: |
| 18 | """Normalize URL for dedup: lowercase, strip www/old/m prefixes, remove tracking params.""" |
| 19 | parsed = urlparse(url.strip().lower()) |
| 20 | netloc = parsed.netloc |
| 21 | for prefix in ("www.", "old.", "m."): |
| 22 | if netloc.startswith(prefix): |
| 23 | netloc = netloc[len(prefix):] |
| 24 | # Strip tracking params |
| 25 | params = parse_qs(parsed.query) |
| 26 | clean_params = {k: v for k, v in params.items() if not k.startswith("utm_")} |
| 27 | query = urlencode(clean_params, doseq=True) |
| 28 | return urlunparse((parsed.scheme, netloc, parsed.path.rstrip("/"), "", query, "")) |
| 29 | |
| 30 | |
| 31 | def candidate_key(item: schema.SourceItem) -> str: |
| 32 | if item.url: |
| 33 | return _normalize_url(item.url) |
| 34 | return f"{item.source}:{item.item_id}" |
| 35 | |
| 36 | |
| 37 | _DIVERSITY_RELEVANCE_THRESHOLD = 0.25 |
| 38 | |
| 39 | # Per-author cap: no single author/handle should dominate the pool. |
| 40 | _MAX_ITEMS_PER_AUTHOR = 3 |
| 41 | |
| 42 | |
| 43 | def _extract_author(candidate: schema.Candidate) -> str | None: |
| 44 | """Return a normalized author key from a candidate's source items.""" |
| 45 | for item in candidate.source_items: |
| 46 | if item.author: |
| 47 | return item.author.strip().lower() |
| 48 | return None |
| 49 | |
| 50 | |
| 51 | def _apply_per_author_cap( |
| 52 | candidates: list[schema.Candidate], |
| 53 | max_per_author: int = _MAX_ITEMS_PER_AUTHOR, |
| 54 | ) -> list[schema.Candidate]: |
| 55 | """Keep at most *max_per_author* items from any single author. |
| 56 | |
| 57 | Candidates are assumed to already be sorted by quality (rrf_score etc.), |
| 58 | so the first N encountered per author are the best ones. |
| 59 | """ |
| 60 | author_counts: dict[str, int] = {} |
| 61 | result: list[schema.Candidate] = [] |
| 62 | for c in candidates: |
| 63 | author = _extract_author(c) |
| 64 | if author is None: |
| 65 | result.append(c) |
| 66 | continue |
| 67 | count = author_counts.get(author, 0) |
| 68 | if count < max_per_author: |
| 69 | result.append(c) |
| 70 | author_counts[author] = count + 1 |
| 71 | return result |
| 72 | |
| 73 | |
| 74 | def _diversify_pool( |
| 75 | fused: list[schema.Candidate], |
| 76 | pool_limit: int, |
| 77 | min_per_source: int = 2, |
| 78 | ) -> list[schema.Candidate]: |
| 79 | """Ensure at least *min_per_source* items per qualifying source survive truncation. |
| 80 | |
| 81 | Sources only qualify for reserved slots if their best item exceeds |
| 82 | the relevance threshold. Low-relevance sources compete on merit only. |
| 83 | """ |
| 84 | max_relevance: dict[str, float] = {} |
| 85 | for c in fused: |
| 86 | current = max_relevance.get(c.source, 0.0) |
| 87 | if c.local_relevance > current: |
| 88 | max_relevance[c.source] = c.local_relevance |
| 89 | |
| 90 | reserved: dict[str, list[schema.Candidate]] = {} |
| 91 | remainder: list[schema.Candidate] = [] |
| 92 | for c in fused: |
| 93 | qualifies = max_relevance.get(c.source, 0.0) >= _DIVERSITY_RELEVANCE_THRESHOLD |
| 94 | bucket = reserved.setdefault(c.source, []) |
| 95 | if qualifies and len(bucket) < min_per_source: |
| 96 | bucket.append(c) |
| 97 | else: |
| 98 | remainder.append(c) |
| 99 | pool = [c for per_source in reserved.values() for c in per_source] |
| 100 | seen = {c.candidate_id for c in pool} |
| 101 | for c in remainder: |
| 102 | if len(pool) >= pool_limit: |
| 103 | break |
| 104 | if c.candidate_id not in seen: |
| 105 | pool.append(c) |
| 106 | pool.sort(key=_candidate_sort_key) |
| 107 | return pool[:pool_limit] |
| 108 | |
| 109 | |
| 110 | def weighted_rrf( |
| 111 | streams: dict[tuple[str, str], list[schema.SourceItem]], |
| 112 | plan: schema.QueryPlan, |
| 113 | *, |
| 114 | pool_limit: int, |
| 115 | ) -> list[schema.Candidate]: |
| 116 | """Fuse ranked lists into a single candidate pool.""" |
| 117 | subqueries = {subquery.label: subquery for subquery in plan.subqueries} |
| 118 | candidates: dict[str, schema.Candidate] = {} |
| 119 | # Track (source, item_id) pairs already attached to each candidate for O(1) dedup. |
| 120 | seen_source_items: dict[str, set[tuple[str, str]]] = {} |
| 121 | |
| 122 | for (label, source), items in streams.items(): |
| 123 | subquery = subqueries[label] |
| 124 | weight = subquery.weight * plan.source_weights.get(source, 1.0) |
| 125 | for rank, item in enumerate(items, start=1): |
| 126 | key = candidate_key(item) |
| 127 | score = weight / (RRF_K + rank) |
| 128 | item_local_relevance = item.local_relevance if item.local_relevance is not None else float(item.metadata.get("local_relevance", item.relevance_hint)) |
| 129 | item_freshness = item.freshness if item.freshness is not None else int(item.metadata.get("freshness", 0)) |
| 130 | item_source_quality = item.source_quality if item.source_quality is not None else float(item.metadata.get("source_quality", 0.6)) |
| 131 | if key not in candidates: |
| 132 | candidates[key] = schema.Candidate( |
| 133 | candidate_id=key, |
| 134 | item_id=item.item_id, |
| 135 | source=item.source, |
| 136 | title=item.title, |
| 137 | url=item.url, |
| 138 | snippet=item.snippet, |
| 139 | subquery_labels=[label], |
| 140 | native_ranks={f"{label}:{source}": rank}, |
| 141 | local_relevance=item_local_relevance, |
| 142 | freshness=item_freshness, |
| 143 | engagement=item.engagement_score if item.engagement_score is not None else item.metadata.get("engagement_score"), |
| 144 | source_quality=item_source_quality, |
| 145 | rrf_score=score, |
| 146 | sources=[item.source], |
| 147 | source_items=[item], |
| 148 | metadata={ |
| 149 | "provenance": [ |
| 150 | { |
| 151 | "source": source, |
| 152 | "subquery_label": label, |
| 153 | "native_rank": rank, |
| 154 | "item_id": item.item_id, |
| 155 | } |
| 156 | ] |
| 157 | }, |
| 158 | ) |
| 159 | seen_source_items[key] = {(item.source, item.item_id)} |
| 160 | continue |
| 161 | |
| 162 | candidate = candidates[key] |
| 163 | candidate.rrf_score += score |
| 164 | previous_primary_score = (candidate.local_relevance * 100.0) + candidate.freshness + (candidate.source_quality * 10.0) |
| 165 | incoming_primary_score = (item_local_relevance * 100.0) + item_freshness + (item_source_quality * 10.0) |
| 166 | candidate.local_relevance = max( |
| 167 | candidate.local_relevance, |
| 168 | item_local_relevance, |
| 169 | ) |
| 170 | candidate.freshness = max(candidate.freshness, item_freshness) |
| 171 | item_eng = item.engagement_score if item.engagement_score is not None else item.metadata.get("engagement_score") |
| 172 | if candidate.engagement is None: |
| 173 | candidate.engagement = item_eng |
| 174 | elif item_eng is not None: |
| 175 | candidate.engagement = max(candidate.engagement, item_eng) |
| 176 | candidate.source_quality = max( |
| 177 | candidate.source_quality, |
| 178 | item_source_quality, |
| 179 | ) |
| 180 | candidate.native_ranks[f"{label}:{source}"] = rank |
| 181 | if label not in candidate.subquery_labels: |
| 182 | candidate.subquery_labels.append(label) |
| 183 | if item.source not in candidate.sources: |
| 184 | candidate.sources.append(item.source) |
| 185 | source_item_key = (item.source, item.item_id) |
| 186 | if source_item_key not in seen_source_items[key]: |
| 187 | seen_source_items[key].add(source_item_key) |
| 188 | candidate.source_items.append(item) |
| 189 | candidate.metadata.setdefault("provenance", []).append( |
| 190 | { |
| 191 | "source": source, |
| 192 | "subquery_label": label, |
| 193 | "native_rank": rank, |
| 194 | "item_id": item.item_id, |
| 195 | } |
| 196 | ) |
| 197 | if incoming_primary_score > previous_primary_score: |
| 198 | candidate.item_id = item.item_id |
| 199 | candidate.source = item.source |
| 200 | candidate.title = item.title |
| 201 | candidate.snippet = item.snippet |
| 202 | if len(candidate.snippet.split()) < len(item.snippet.split()): |
| 203 | candidate.snippet = item.snippet |
| 204 | |
| 205 | fused = sorted(candidates.values(), key=_candidate_sort_key) |
| 206 | fused = _apply_per_author_cap(fused) |
| 207 | return _diversify_pool(fused, pool_limit) |
| 208 |