| 1 | """Candidate clustering and representative selection.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from . import dedupe, entity_extract, schema |
| 6 | |
| 7 | |
| 8 | def _cluster_sort_key(candidate: schema.Candidate) -> tuple: |
| 9 | """Sort key that partitions stale candidates below fresh ones. |
| 10 | |
| 11 | Stale items (all dated source_items outside the window) must never lead |
| 12 | cluster representatives or render as the cluster title. |
| 13 | """ |
| 14 | return ( |
| 15 | 1 if schema.candidate_out_of_window(candidate) else 0, |
| 16 | -candidate.final_score, |
| 17 | ) |
| 18 | |
| 19 | CLUSTERABLE_INTENTS = {"breaking_news", "opinion", "comparison", "prediction"} |
| 20 | |
| 21 | def _candidate_text(candidate: schema.Candidate) -> str: |
| 22 | return " ".join(part for part in [candidate.title, candidate.snippet] if part).strip() |
| 23 | |
| 24 | def _mmr_representatives( |
| 25 | candidates: list[schema.Candidate], |
| 26 | text_cache: dict[str, dedupe._PreparedText], |
| 27 | limit: int = 3, |
| 28 | diversity_lambda: float = 0.75, |
| 29 | ) -> list[str]: |
| 30 | selected: list[schema.Candidate] = [] |
| 31 | remaining_set = {c.candidate_id for c in candidates} |
| 32 | remaining = list(candidates) |
| 33 | while remaining and len(selected) < limit: |
| 34 | if not selected: |
| 35 | best = min(remaining, key=_cluster_sort_key) |
| 36 | selected.append(best) |
| 37 | remaining_set.discard(best.candidate_id) |
| 38 | remaining = [c for c in remaining if c.candidate_id in remaining_set] |
| 39 | continue |
| 40 | |
| 41 | selected_preps = [text_cache[c.candidate_id] for c in selected] |
| 42 | |
| 43 | def score(candidate: schema.Candidate) -> tuple: |
| 44 | prep = text_cache[candidate.candidate_id] |
| 45 | diversity_penalty = max( |
| 46 | dedupe.prepared_similarity(prep, sp) for sp in selected_preps |
| 47 | ) |
| 48 | base_score = (diversity_lambda * candidate.final_score) - ((1 - diversity_lambda) * diversity_penalty * 100) |
| 49 | return ( |
| 50 | 0 if schema.candidate_out_of_window(candidate) else 1, |
| 51 | base_score, |
| 52 | ) |
| 53 | |
| 54 | best = max(remaining, key=score) |
| 55 | selected.append(best) |
| 56 | remaining_set.discard(best.candidate_id) |
| 57 | remaining = [c for c in remaining if c.candidate_id in remaining_set] |
| 58 | return [candidate.candidate_id for candidate in selected] |
| 59 | |
| 60 | |
| 61 | def cluster_candidates( |
| 62 | candidates: list[schema.Candidate], |
| 63 | plan: schema.QueryPlan, |
| 64 | ) -> list[schema.Cluster]: |
| 65 | """Greedy clustering around high-ranked leaders.""" |
| 66 | if plan.intent not in CLUSTERABLE_INTENTS or plan.cluster_mode == "none": |
| 67 | clusters = [] |
| 68 | for index, candidate in enumerate(candidates, start=1): |
| 69 | cluster_id = f"cluster-{index}" |
| 70 | candidate.cluster_id = cluster_id |
| 71 | clusters.append( |
| 72 | schema.Cluster( |
| 73 | cluster_id=cluster_id, |
| 74 | title=candidate.title, |
| 75 | candidate_ids=[candidate.candidate_id], |
| 76 | representative_ids=[candidate.candidate_id], |
| 77 | sources=sorted(schema.candidate_sources(candidate)), |
| 78 | score=candidate.final_score, |
| 79 | uncertainty=None, |
| 80 | ) |
| 81 | ) |
| 82 | return clusters |
| 83 | |
| 84 | text_cache: dict[str, dedupe._PreparedText] = { |
| 85 | c.candidate_id: dedupe._PreparedText(_candidate_text(c)) |
| 86 | for c in candidates |
| 87 | } |
| 88 | |
| 89 | groups: list[list[schema.Candidate]] = [] |
| 90 | # Lower threshold for breaking_news: related articles share fewer exact |
| 91 | # words but cover the same event. |
| 92 | threshold = 0.42 if plan.intent == "breaking_news" else 0.48 |
| 93 | for candidate in candidates: |
| 94 | assigned = False |
| 95 | cand_prep = text_cache[candidate.candidate_id] |
| 96 | for group in groups: |
| 97 | leader = group[0] |
| 98 | similarity = dedupe.prepared_similarity(cand_prep, text_cache[leader.candidate_id]) |
| 99 | if similarity >= threshold: |
| 100 | group.append(candidate) |
| 101 | assigned = True |
| 102 | break |
| 103 | if not assigned: |
| 104 | groups.append([candidate]) |
| 105 | |
| 106 | clusters: list[schema.Cluster] = [] |
| 107 | for index, group in enumerate(groups, start=1): |
| 108 | group.sort(key=_cluster_sort_key) |
| 109 | cluster_id = f"cluster-{index}" |
| 110 | representatives = _mmr_representatives(group, text_cache) |
| 111 | for candidate in group: |
| 112 | candidate.cluster_id = cluster_id |
| 113 | clusters.append( |
| 114 | schema.Cluster( |
| 115 | cluster_id=cluster_id, |
| 116 | title=group[0].title, |
| 117 | candidate_ids=[candidate.candidate_id for candidate in group], |
| 118 | representative_ids=representatives, |
| 119 | sources=sorted({source for candidate in group for source in schema.candidate_sources(candidate)}), |
| 120 | score=max(candidate.final_score for candidate in group), |
| 121 | uncertainty=_cluster_uncertainty(group), |
| 122 | ) |
| 123 | ) |
| 124 | |
| 125 | # Second pass: merge small clusters that share entities across sources. |
| 126 | clusters = _merge_entity_clusters( |
| 127 | clusters, |
| 128 | candidates, |
| 129 | min_shared_entities=2 if "discover-mode" in plan.notes else 1, |
| 130 | ) |
| 131 | |
| 132 | return sorted(clusters, key=lambda cluster: cluster.score, reverse=True) |
| 133 | |
| 134 | |
| 135 | def _merge_entity_clusters( |
| 136 | clusters: list[schema.Cluster], |
| 137 | all_candidates: list[schema.Candidate], |
| 138 | *, |
| 139 | min_shared_entities: int = 1, |
| 140 | ) -> list[schema.Cluster]: |
| 141 | """Merge small clusters that cover the same story across different sources. |
| 142 | |
| 143 | The initial greedy pass uses text similarity which misses cross-source |
| 144 | matches where phrasing differs. This second pass looks at entity overlap |
| 145 | (proper nouns, names, numbers) to catch cases like: |
| 146 | - Reddit: "Kanye West to headline all three nights of Wireless Festival 2026" |
| 147 | - X: "BREAKING: Kanye West (Ye) is making his massive UK comeback!" |
| 148 | """ |
| 149 | if len(clusters) < 2: |
| 150 | return clusters |
| 151 | |
| 152 | candidate_map = {c.candidate_id: c for c in all_candidates} |
| 153 | |
| 154 | # Build entity sets per cluster |
| 155 | cluster_entities: list[set[str]] = [] |
| 156 | for cl in clusters: |
| 157 | entities: set[str] = set() |
| 158 | for cid in cl.candidate_ids: |
| 159 | cand = candidate_map.get(cid) |
| 160 | if cand: |
| 161 | entities |= entity_extract.extract_text_entities(_candidate_text(cand)) |
| 162 | cluster_entities.append(entities) |
| 163 | |
| 164 | # Only merge clusters with <= 3 items (don't merge already-large clusters) |
| 165 | merged_into: dict[int, int] = {} # index -> merge target index |
| 166 | for i in range(len(clusters)): |
| 167 | if i in merged_into or len(clusters[i].candidate_ids) > 3: |
| 168 | continue |
| 169 | for j in range(i + 1, len(clusters)): |
| 170 | if j in merged_into or len(clusters[j].candidate_ids) > 3: |
| 171 | continue |
| 172 | # Require different sources to merge (same-source should already be grouped) |
| 173 | sources_i = set(clusters[i].sources) |
| 174 | sources_j = set(clusters[j].sources) |
| 175 | if sources_i == sources_j and len(sources_i) == 1: |
| 176 | continue |
| 177 | # Prevent Polymarket clusters from merging with non-Polymarket |
| 178 | # clusters. Prediction markets about "Sam Altman equity" should not |
| 179 | # merge into a news cluster about "Sam Altman rivalry" just because |
| 180 | # both mention the same entity. |
| 181 | poly_i = "polymarket" in sources_i |
| 182 | poly_j = "polymarket" in sources_j |
| 183 | if poly_i != poly_j: |
| 184 | continue |
| 185 | |
| 186 | shared_entities = cluster_entities[i] & cluster_entities[j] |
| 187 | overlap = entity_extract.entity_overlap(cluster_entities[i], cluster_entities[j]) |
| 188 | if len(shared_entities) >= min_shared_entities and overlap >= 0.45: |
| 189 | merged_into[j] = i |
| 190 | |
| 191 | if not merged_into: |
| 192 | return clusters |
| 193 | |
| 194 | # Build merged cluster list |
| 195 | result: list[schema.Cluster] = [] |
| 196 | for i, cl in enumerate(clusters): |
| 197 | if i in merged_into: |
| 198 | continue |
| 199 | # Collect all clusters merged into this one |
| 200 | merge_sources = [i] + [j for j, target in merged_into.items() if target == i] |
| 201 | if len(merge_sources) == 1: |
| 202 | result.append(cl) |
| 203 | continue |
| 204 | |
| 205 | # Combine candidates from all merged clusters |
| 206 | combined_cids: list[str] = [] |
| 207 | combined_sources: set[str] = set() |
| 208 | best_score = 0.0 |
| 209 | for idx in merge_sources: |
| 210 | combined_cids.extend(clusters[idx].candidate_ids) |
| 211 | combined_sources.update(clusters[idx].sources) |
| 212 | best_score = max(best_score, clusters[idx].score) |
| 213 | |
| 214 | # Pick representatives from combined pool |
| 215 | combined_candidates = [candidate_map[cid] for cid in combined_cids if cid in candidate_map] |
| 216 | combined_candidates.sort(key=_cluster_sort_key) |
| 217 | merge_text_cache = { |
| 218 | c.candidate_id: dedupe._PreparedText(_candidate_text(c)) |
| 219 | for c in combined_candidates |
| 220 | } |
| 221 | reps = _mmr_representatives(combined_candidates, merge_text_cache) |
| 222 | |
| 223 | cluster_id = cl.cluster_id |
| 224 | for cid in combined_cids: |
| 225 | cand = candidate_map.get(cid) |
| 226 | if cand: |
| 227 | cand.cluster_id = cluster_id |
| 228 | |
| 229 | result.append(schema.Cluster( |
| 230 | cluster_id=cluster_id, |
| 231 | title=combined_candidates[0].title if combined_candidates else cl.title, |
| 232 | candidate_ids=combined_cids, |
| 233 | representative_ids=reps, |
| 234 | sources=sorted(combined_sources), |
| 235 | score=best_score, |
| 236 | uncertainty=_cluster_uncertainty(combined_candidates), |
| 237 | )) |
| 238 | |
| 239 | return result |
| 240 | |
| 241 | |
| 242 | def _cluster_uncertainty(group: list[schema.Candidate]) -> str | None: |
| 243 | sources = {source for candidate in group for source in schema.candidate_sources(candidate)} |
| 244 | if len(sources) == 1: |
| 245 | return "single-source" |
| 246 | if max(candidate.final_score for candidate in group) < 55: |
| 247 | return "thin-evidence" |
| 248 | return None |
| 249 |