| 1 | """Weighted reciprocal rank fusion for per-(subquery, source) streams.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from collections.abc import Iterable |
| 6 | |
| 7 | from urllib.parse import parse_qs, urlencode, urlparse, urlunparse |
| 8 | |
| 9 | from . import schema |
| 10 | |
| 11 | # Standard RRF smoothing constant (Cormack et al. 2009) |
| 12 | RRF_K = 60 |
| 13 | |
| 14 | |
| 15 | def _candidate_sort_key(c: schema.Candidate) -> tuple: |
| 16 | # Out-of-window evidence sorts strictly below anything in the window. A |
| 17 | # "last 30 days" brief that ranks a nine-month-old video at #1 breaks its |
| 18 | # own contract, however relevant that video is; it still appears, just |
| 19 | # never above in-window evidence. |
| 20 | return ( |
| 21 | 1 if schema.candidate_out_of_window(c) else 0, |
| 22 | -c.rrf_score, |
| 23 | -c.local_relevance, |
| 24 | -c.freshness, |
| 25 | schema.candidate_source_label(c), |
| 26 | c.title, |
| 27 | ) |
| 28 | |
| 29 | |
| 30 | def _normalize_url(url: str) -> str: |
| 31 | """Normalize URL for dedup: lowercase host, strip www/old/m prefixes, remove tracking params.""" |
| 32 | parsed = urlparse(url.strip()) |
| 33 | netloc = parsed.netloc.lower() |
| 34 | for prefix in ("www.", "old.", "m."): |
| 35 | if netloc.startswith(prefix): |
| 36 | netloc = netloc[len(prefix):] |
| 37 | # Strip tracking params |
| 38 | params = parse_qs(parsed.query) |
| 39 | clean_params = {k: v for k, v in params.items() if not k.lower().startswith("utm_")} |
| 40 | query = urlencode(clean_params, doseq=True) |
| 41 | return urlunparse((parsed.scheme.lower(), netloc, parsed.path.rstrip("/"), "", query, "")) |
| 42 | |
| 43 | |
| 44 | def candidate_key(item: schema.SourceItem) -> str: |
| 45 | if item.url: |
| 46 | return _normalize_url(item.url) |
| 47 | return f"{item.source}:{item.item_id}" |
| 48 | |
| 49 | |
| 50 | # Enrichment that one copy of a thread may carry and another may lack. When |
| 51 | # the same URL arrives from two subquery streams (each stream enriches its own |
| 52 | # top-N), the copy that won a comment slot must survive de-duplication. |
| 53 | _ENRICHMENT_KEYS = ( |
| 54 | "top_comments", |
| 55 | "comment_insights", |
| 56 | "transcript_highlights", |
| 57 | "transcript_snippet", |
| 58 | "transcript", |
| 59 | ) |
| 60 | |
| 61 | |
| 62 | def merge_source_items(existing: schema.SourceItem, incoming: schema.SourceItem) -> schema.SourceItem: |
| 63 | """Fold ``incoming``'s enrichment and counts into ``existing`` (same thread). |
| 64 | |
| 65 | Keeps the richer value per field: an enrichment list the existing copy |
| 66 | lacks (or a longer one), the larger numeric engagement counters, and the |
| 67 | longer body/snippet. Mutates and returns ``existing``. |
| 68 | """ |
| 69 | for key in _ENRICHMENT_KEYS: |
| 70 | theirs = incoming.metadata.get(key) |
| 71 | if not theirs: |
| 72 | continue |
| 73 | mine = existing.metadata.get(key) |
| 74 | if not mine or (isinstance(theirs, list) and isinstance(mine, list) and len(theirs) > len(mine)): |
| 75 | existing.metadata[key] = theirs |
| 76 | for field_name, value in (incoming.engagement or {}).items(): |
| 77 | if isinstance(value, bool) or not isinstance(value, (int, float)): |
| 78 | if existing.engagement.get(field_name) is None and value is not None: |
| 79 | existing.engagement[field_name] = value |
| 80 | continue |
| 81 | current = existing.engagement.get(field_name) |
| 82 | if not isinstance(current, (int, float)) or isinstance(current, bool) or value > current: |
| 83 | existing.engagement[field_name] = value |
| 84 | if len(incoming.body or "") > len(existing.body or ""): |
| 85 | existing.body = incoming.body |
| 86 | if len(incoming.snippet or "") > len(existing.snippet or ""): |
| 87 | existing.snippet = incoming.snippet |
| 88 | return existing |
| 89 | |
| 90 | |
| 91 | def collapse_duplicate_urls(items: list[schema.SourceItem]) -> list[schema.SourceItem]: |
| 92 | """Collapse same-source, same-URL copies, keeping order and merging enrichment. |
| 93 | |
| 94 | Per-stream item ids (``R1``, ``X1``) collide across subqueries, so identity |
| 95 | is the normalized URL, not the id. The first occurrence stays in place and |
| 96 | absorbs later copies via :func:`merge_source_items`. |
| 97 | """ |
| 98 | first_by_key: dict[tuple[str, str], schema.SourceItem] = {} |
| 99 | kept: list[schema.SourceItem] = [] |
| 100 | for item in items: |
| 101 | key = (item.source, candidate_key(item)) |
| 102 | existing = first_by_key.get(key) |
| 103 | if existing is None: |
| 104 | first_by_key[key] = item |
| 105 | kept.append(item) |
| 106 | else: |
| 107 | merge_source_items(existing, item) |
| 108 | return kept |
| 109 | |
| 110 | |
| 111 | _DIVERSITY_RELEVANCE_THRESHOLD = 0.25 |
| 112 | |
| 113 | # Reddit engagement reservation: pool slots held for the highest-engagement |
| 114 | # entity-grounded in-window Reddit candidates, scaled by pool size (quick 15 |
| 115 | # -> 2, default 40 -> 3, deep 60 -> 4). The fused order is RRF-first and each |
| 116 | # stream is relevance-first, so the month's most-discussed on-topic thread can |
| 117 | # otherwise lose its slot to a one-upvote post with better title overlap. |
| 118 | _REDDIT_RESERVE_BY_POOL = ((15, 2), (40, 3)) |
| 119 | _REDDIT_RESERVE_MAX = 4 |
| 120 | |
| 121 | |
| 122 | def _reddit_reserve_for(pool_limit: int) -> int: |
| 123 | for ceiling, reserve in _REDDIT_RESERVE_BY_POOL: |
| 124 | if pool_limit <= ceiling: |
| 125 | return reserve |
| 126 | return _REDDIT_RESERVE_MAX |
| 127 | |
| 128 | |
| 129 | def relevance_floor_for_entity(entity: str) -> float: |
| 130 | """Relevance a Reddit thread must clear to earn a reservation or keeper slot. |
| 131 | |
| 132 | Grounding keys on the entity's head token (see ``rerank._entity_grounded``). |
| 133 | A generic head ("ai", "x", "new") matches almost anything, so the floor |
| 134 | rises to the diversity threshold; a distinctive head keeps the shared |
| 135 | ``RELEVANCE_FLOOR``. |
| 136 | """ |
| 137 | from . import relevance |
| 138 | |
| 139 | head = (entity or "").lower().split()[:1] |
| 140 | if not head: |
| 141 | return relevance.RELEVANCE_FLOOR |
| 142 | token = head[0] |
| 143 | generic = ( |
| 144 | len(token) <= 2 |
| 145 | or token in relevance.STOPWORDS |
| 146 | or token in relevance.LOW_SIGNAL_QUERY_TOKENS |
| 147 | ) |
| 148 | return _DIVERSITY_RELEVANCE_THRESHOLD if generic else relevance.RELEVANCE_FLOOR |
| 149 | |
| 150 | |
| 151 | def raw_engagement(item: schema.SourceItem) -> float: |
| 152 | """Upvotes plus comments as a plain number, for engagement-first ordering.""" |
| 153 | eng = item.engagement or {} |
| 154 | total = 0.0 |
| 155 | for key in ("score", "num_comments"): |
| 156 | value = eng.get(key) |
| 157 | if isinstance(value, (int, float)) and not isinstance(value, bool): |
| 158 | total += float(value) |
| 159 | return total |
| 160 | |
| 161 | |
| 162 | def reddit_thread_qualifies( |
| 163 | item: schema.SourceItem, |
| 164 | entity: str, |
| 165 | floor: float, |
| 166 | relevance: float | None = None, |
| 167 | ) -> bool: |
| 168 | """On-topic enough for an engagement slot: clears the floor and names the entity. |
| 169 | |
| 170 | ``relevance`` overrides the item's own ``local_relevance`` (a fused |
| 171 | candidate carries the max across its copies). |
| 172 | """ |
| 173 | from . import rerank |
| 174 | |
| 175 | score = relevance if relevance is not None else (item.local_relevance or 0.0) |
| 176 | if score < floor: |
| 177 | return False |
| 178 | if raw_engagement(item) <= 0: |
| 179 | return False |
| 180 | if not entity: |
| 181 | return True |
| 182 | return rerank._entity_grounded(f"{item.title or ''} {item.body or ''}", entity) |
| 183 | |
| 184 | # Per-author cap: no single author/handle should dominate the pool. |
| 185 | _MAX_ITEMS_PER_AUTHOR = 3 |
| 186 | |
| 187 | # Raised cap for the subject of the topic (a handle in the run's |
| 188 | # resolved_handles). On a person or company topic the subject is what the user |
| 189 | # asked about, so the flat cap discards exactly the evidence the run worked |
| 190 | # hardest to retrieve -- the measured 'Peter Steinberger steipete' baseline |
| 191 | # recovered 8 subject-authored posts and would have kept 3. Still bounded: a |
| 192 | # prolific subject must not crowd out commentary about them, which is the other |
| 193 | # half of the answer a user wants. |
| 194 | _MAX_ITEMS_PER_FIRST_PARTY_AUTHOR = 8 |
| 195 | |
| 196 | |
| 197 | def _extract_author(candidate: schema.Candidate) -> str | None: |
| 198 | """Return a normalized author key from a candidate's source items.""" |
| 199 | for item in candidate.source_items: |
| 200 | # Jobs use the ATS provider as author, not an individual contributor. |
| 201 | if item.source != "jobs" and item.author: |
| 202 | return item.author.strip().lower() |
| 203 | return None |
| 204 | |
| 205 | |
| 206 | def _apply_per_author_cap( |
| 207 | candidates: list[schema.Candidate], |
| 208 | max_per_author: int = _MAX_ITEMS_PER_AUTHOR, |
| 209 | first_party_handles: Iterable[str] | None = None, |
| 210 | max_per_first_party_author: int = _MAX_ITEMS_PER_FIRST_PARTY_AUTHOR, |
| 211 | ) -> list[schema.Candidate]: |
| 212 | """Keep at most *max_per_author* items from any single author. |
| 213 | |
| 214 | Authors named in *first_party_handles* -- the subject of the topic -- get |
| 215 | the higher *max_per_first_party_author* allowance instead, because their |
| 216 | own posts are the point of the query rather than one voice among many. |
| 217 | |
| 218 | Candidates are assumed to already be sorted by quality (rrf_score etc.), |
| 219 | so the first N encountered per author are the best ones. |
| 220 | """ |
| 221 | first_party = { |
| 222 | h.strip().lstrip("@").lower() |
| 223 | for h in (first_party_handles or ()) |
| 224 | if h and h.strip() |
| 225 | } |
| 226 | author_counts: dict[str, int] = {} |
| 227 | result: list[schema.Candidate] = [] |
| 228 | for c in candidates: |
| 229 | author = _extract_author(c) |
| 230 | if author is None: |
| 231 | result.append(c) |
| 232 | continue |
| 233 | limit = ( |
| 234 | max_per_first_party_author |
| 235 | if author.strip().lstrip("@").lower() in first_party |
| 236 | else max_per_author |
| 237 | ) |
| 238 | count = author_counts.get(author, 0) |
| 239 | if count < limit: |
| 240 | result.append(c) |
| 241 | author_counts[author] = count + 1 |
| 242 | return result |
| 243 | |
| 244 | |
| 245 | def _reddit_engagement_reservation( |
| 246 | fused: list[schema.Candidate], |
| 247 | reserve: int, |
| 248 | entity: str, |
| 249 | ) -> list[schema.Candidate]: |
| 250 | """The *reserve* highest-engagement Reddit candidates that are in-window |
| 251 | and on-topic, in engagement order.""" |
| 252 | if reserve <= 0: |
| 253 | return [] |
| 254 | floor = relevance_floor_for_entity(entity) |
| 255 | eligible = [] |
| 256 | for c in fused: |
| 257 | if c.source != "reddit" or schema.candidate_out_of_window(c): |
| 258 | continue |
| 259 | reddit_items = [it for it in c.source_items if it.source == "reddit"] |
| 260 | if not reddit_items: |
| 261 | continue |
| 262 | best = max(reddit_items, key=raw_engagement) |
| 263 | relevance = max(c.local_relevance or 0.0, best.local_relevance or 0.0) |
| 264 | if not reddit_thread_qualifies(best, entity, floor, relevance=relevance): |
| 265 | continue |
| 266 | eligible.append((raw_engagement(best), c)) |
| 267 | eligible.sort(key=lambda pair: -pair[0]) |
| 268 | return [c for _, c in eligible[:reserve]] |
| 269 | |
| 270 | |
| 271 | def _diversify_pool( |
| 272 | fused: list[schema.Candidate], |
| 273 | pool_limit: int, |
| 274 | min_per_source: int = 2, |
| 275 | entity: str = "", |
| 276 | ) -> list[schema.Candidate]: |
| 277 | """Ensure at least *min_per_source* items per qualifying source survive truncation. |
| 278 | |
| 279 | Sources only qualify for reserved slots if their best item exceeds |
| 280 | the relevance threshold. Low-relevance sources compete on merit only. |
| 281 | Reddit additionally gets an engagement reservation (see |
| 282 | ``_reddit_reserve_for``) for its most-discussed on-topic threads. |
| 283 | """ |
| 284 | max_relevance: dict[str, float] = {} |
| 285 | for c in fused: |
| 286 | current = max_relevance.get(c.source, 0.0) |
| 287 | if c.local_relevance > current: |
| 288 | max_relevance[c.source] = c.local_relevance |
| 289 | |
| 290 | protected = _reddit_engagement_reservation(fused, _reddit_reserve_for(pool_limit), entity) |
| 291 | protected_ids = {c.candidate_id for c in protected} |
| 292 | pool: list[schema.Candidate] = list(protected) |
| 293 | seen = set(protected_ids) |
| 294 | reserved: dict[str, list[schema.Candidate]] = {} |
| 295 | remainder: list[schema.Candidate] = [] |
| 296 | for c in fused: |
| 297 | if c.candidate_id in seen: |
| 298 | continue |
| 299 | qualifies = max_relevance.get(c.source, 0.0) >= _DIVERSITY_RELEVANCE_THRESHOLD |
| 300 | bucket = reserved.setdefault(c.source, []) |
| 301 | if qualifies and len(bucket) < min_per_source: |
| 302 | bucket.append(c) |
| 303 | else: |
| 304 | remainder.append(c) |
| 305 | pool.extend(c for per_source in reserved.values() for c in per_source) |
| 306 | seen = {c.candidate_id for c in pool} |
| 307 | for c in remainder: |
| 308 | if len(pool) >= pool_limit: |
| 309 | break |
| 310 | if c.candidate_id not in seen: |
| 311 | pool.append(c) |
| 312 | pool.sort(key=_candidate_sort_key) |
| 313 | if len(pool) > pool_limit: |
| 314 | # The per-source buckets can overfill a small pool. The Reddit |
| 315 | # reservation is low-RRF by construction, so a plain slice would cut |
| 316 | # exactly the threads it exists to keep: trim unprotected candidates |
| 317 | # from the sorted tail instead. |
| 318 | keep_unprotected = pool_limit - len(protected_ids) |
| 319 | trimmed: list[schema.Candidate] = [] |
| 320 | for c in pool: |
| 321 | if c.candidate_id in protected_ids: |
| 322 | trimmed.append(c) |
| 323 | elif keep_unprotected > 0: |
| 324 | trimmed.append(c) |
| 325 | keep_unprotected -= 1 |
| 326 | pool = trimmed |
| 327 | return pool[:pool_limit] |
| 328 | |
| 329 | |
| 330 | def weighted_rrf( |
| 331 | streams: dict[tuple[str, str], list[schema.SourceItem]], |
| 332 | plan: schema.QueryPlan, |
| 333 | *, |
| 334 | pool_limit: int, |
| 335 | range_from: str | None = None, |
| 336 | range_to: str | None = None, |
| 337 | first_party_handles: Iterable[str] | None = None, |
| 338 | ) -> list[schema.Candidate]: |
| 339 | """Fuse ranked lists into a single candidate pool. |
| 340 | |
| 341 | When ``range_from`` and ``range_to`` are provided, they are stored in each |
| 342 | candidate's metadata so ``candidate_out_of_window`` can compare the actual |
| 343 | date against the run window (instead of relying solely on adapter-provided |
| 344 | ``date_confidence``). ``first_party_handles`` raises the per-author cap |
| 345 | for the topic's subject so their own posts are not flattened to the |
| 346 | incidental-account allowance. |
| 347 | """ |
| 348 | subqueries = {subquery.label: subquery for subquery in plan.subqueries} |
| 349 | candidates: dict[str, schema.Candidate] = {} |
| 350 | # Track source items already attached to each candidate, keyed by |
| 351 | # (source, normalized URL): per-stream ids collide across subqueries, and |
| 352 | # a repeat copy may carry enrichment the first one lacks. |
| 353 | seen_source_items: dict[str, dict[tuple[str, str], schema.SourceItem]] = {} |
| 354 | |
| 355 | for (label, source), items in streams.items(): |
| 356 | subquery = subqueries[label] |
| 357 | weight = subquery.weight * plan.source_weights.get(source, 1.0) |
| 358 | for rank, item in enumerate(items, start=1): |
| 359 | key = candidate_key(item) |
| 360 | score = weight / (RRF_K + rank) |
| 361 | item_local_relevance = item.local_relevance if item.local_relevance is not None else float(item.metadata.get("local_relevance", item.relevance_hint)) |
| 362 | item_freshness = item.freshness if item.freshness is not None else int(item.metadata.get("freshness", 0)) |
| 363 | item_source_quality = item.source_quality if item.source_quality is not None else float(item.metadata.get("source_quality", 0.6)) |
| 364 | if key not in candidates: |
| 365 | candidate_metadata: dict = { |
| 366 | "provenance": [ |
| 367 | { |
| 368 | "source": source, |
| 369 | "subquery_label": label, |
| 370 | "native_rank": rank, |
| 371 | "item_id": item.item_id, |
| 372 | } |
| 373 | ] |
| 374 | } |
| 375 | if range_from: |
| 376 | candidate_metadata["range_from"] = range_from |
| 377 | if range_to: |
| 378 | candidate_metadata["range_to"] = range_to |
| 379 | candidates[key] = schema.Candidate( |
| 380 | candidate_id=key, |
| 381 | item_id=item.item_id, |
| 382 | source=item.source, |
| 383 | title=item.title, |
| 384 | url=item.url, |
| 385 | snippet=item.snippet, |
| 386 | subquery_labels=[label], |
| 387 | native_ranks={f"{label}:{source}": rank}, |
| 388 | local_relevance=item_local_relevance, |
| 389 | freshness=item_freshness, |
| 390 | engagement=item.engagement_score if item.engagement_score is not None else item.metadata.get("engagement_score"), |
| 391 | source_quality=item_source_quality, |
| 392 | rrf_score=score, |
| 393 | sources=[item.source], |
| 394 | source_items=[item], |
| 395 | metadata=candidate_metadata, |
| 396 | ) |
| 397 | seen_source_items[key] = {(item.source, candidate_key(item)): item} |
| 398 | continue |
| 399 | |
| 400 | candidate = candidates[key] |
| 401 | candidate.rrf_score += score |
| 402 | previous_primary_score = (candidate.local_relevance * 100.0) + candidate.freshness + (candidate.source_quality * 10.0) |
| 403 | incoming_primary_score = (item_local_relevance * 100.0) + item_freshness + (item_source_quality * 10.0) |
| 404 | candidate.local_relevance = max( |
| 405 | candidate.local_relevance, |
| 406 | item_local_relevance, |
| 407 | ) |
| 408 | candidate.freshness = max(candidate.freshness, item_freshness) |
| 409 | item_eng = item.engagement_score if item.engagement_score is not None else item.metadata.get("engagement_score") |
| 410 | if candidate.engagement is None: |
| 411 | candidate.engagement = item_eng |
| 412 | elif item_eng is not None: |
| 413 | candidate.engagement = max(candidate.engagement, item_eng) |
| 414 | candidate.source_quality = max( |
| 415 | candidate.source_quality, |
| 416 | item_source_quality, |
| 417 | ) |
| 418 | candidate.native_ranks[f"{label}:{source}"] = rank |
| 419 | if label not in candidate.subquery_labels: |
| 420 | candidate.subquery_labels.append(label) |
| 421 | if item.source not in candidate.sources: |
| 422 | candidate.sources.append(item.source) |
| 423 | source_item_key = (item.source, candidate_key(item)) |
| 424 | existing_item = seen_source_items[key].get(source_item_key) |
| 425 | if existing_item is None: |
| 426 | seen_source_items[key][source_item_key] = item |
| 427 | candidate.source_items.append(item) |
| 428 | else: |
| 429 | merge_source_items(existing_item, item) |
| 430 | candidate.metadata.setdefault("provenance", []).append( |
| 431 | { |
| 432 | "source": source, |
| 433 | "subquery_label": label, |
| 434 | "native_rank": rank, |
| 435 | "item_id": item.item_id, |
| 436 | } |
| 437 | ) |
| 438 | if incoming_primary_score > previous_primary_score: |
| 439 | candidate.item_id = item.item_id |
| 440 | candidate.source = item.source |
| 441 | candidate.title = item.title |
| 442 | candidate.snippet = item.snippet |
| 443 | if len(candidate.snippet.split()) < len(item.snippet.split()): |
| 444 | candidate.snippet = item.snippet |
| 445 | |
| 446 | fused = sorted(candidates.values(), key=_candidate_sort_key) |
| 447 | fused = _apply_per_author_cap(fused, first_party_handles=first_party_handles) |
| 448 | from . import rerank |
| 449 | |
| 450 | entity = rerank._primary_entity(plan.raw_topic or "") if plan.raw_topic else "" |
| 451 | return _diversify_pool(fused, pool_limit, entity=entity) |
| 452 |