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