| 1 | import unittest |
| 2 | |
| 3 | from lib import cluster, schema |
| 4 | |
| 5 | |
| 6 | def make_candidate(candidate_id: str, source: str, title: str, snippet: str, score: float) -> schema.Candidate: |
| 7 | return schema.Candidate( |
| 8 | candidate_id=candidate_id, |
| 9 | item_id=candidate_id, |
| 10 | source=source, |
| 11 | title=title, |
| 12 | url=f"https://example.com/{candidate_id}", |
| 13 | snippet=snippet, |
| 14 | subquery_labels=["primary"], |
| 15 | native_ranks={"primary:reddit": 1}, |
| 16 | local_relevance=0.8, |
| 17 | freshness=80, |
| 18 | engagement=10, |
| 19 | source_quality=0.7, |
| 20 | rrf_score=0.02, |
| 21 | rerank_score=score, |
| 22 | final_score=score, |
| 23 | ) |
| 24 | |
| 25 | |
| 26 | class ClusterV3Tests(unittest.TestCase): |
| 27 | def test_singleton_clusters_for_non_clustered_plan(self): |
| 28 | plan = schema.QueryPlan( |
| 29 | intent="how_to", |
| 30 | freshness_mode="balanced_recent", |
| 31 | cluster_mode="none", |
| 32 | raw_topic="docker setup", |
| 33 | subqueries=[schema.SubQuery(label="primary", search_query="docker setup", ranking_query="How do I set up Docker?", sources=["reddit"])], |
| 34 | source_weights={"reddit": 1.0}, |
| 35 | ) |
| 36 | candidates = [ |
| 37 | make_candidate("c1", "reddit", "Docker setup guide", "Step by step setup", 80), |
| 38 | make_candidate("c2", "youtube", "Docker install video", "Video walkthrough", 75), |
| 39 | ] |
| 40 | clusters = cluster.cluster_candidates(candidates, plan) |
| 41 | self.assertEqual(2, len(clusters)) |
| 42 | self.assertEqual(["c1"], clusters[0].representative_ids) |
| 43 | self.assertEqual(["c2"], clusters[1].representative_ids) |
| 44 | |
| 45 | def test_breaking_news_clusters_related_items(self): |
| 46 | plan = schema.QueryPlan( |
| 47 | intent="breaking_news", |
| 48 | freshness_mode="strict_recent", |
| 49 | cluster_mode="story", |
| 50 | raw_topic="model launch", |
| 51 | subqueries=[schema.SubQuery(label="primary", search_query="model launch", ranking_query="What happened in the model launch?", sources=["reddit", "x"])], |
| 52 | source_weights={"reddit": 0.5, "x": 0.5}, |
| 53 | ) |
| 54 | candidates = [ |
| 55 | make_candidate("c1", "reddit", "Open model launch reactions", "People are reacting to the open model launch today.", 88), |
| 56 | make_candidate("c2", "x", "Open model launch update", "People are reacting to the open model launch today on X.", 84), |
| 57 | make_candidate("c3", "youtube", "Different topic", "A separate discussion about hardware benchmarks.", 70), |
| 58 | ] |
| 59 | clusters = cluster.cluster_candidates(candidates, plan) |
| 60 | self.assertEqual(2, len(clusters)) |
| 61 | self.assertEqual(2, len(clusters[0].candidate_ids)) |
| 62 | self.assertIn("c1", clusters[0].candidate_ids) |
| 63 | self.assertIn("c2", clusters[0].candidate_ids) |
| 64 | |
| 65 | |
| 66 | class TestCrossSourceMerging(unittest.TestCase): |
| 67 | """Test the entity-based second pass that merges same-story clusters across sources.""" |
| 68 | |
| 69 | def _plan(self, intent="breaking_news"): |
| 70 | return schema.QueryPlan( |
| 71 | intent=intent, |
| 72 | freshness_mode="strict_recent", |
| 73 | cluster_mode="story", |
| 74 | raw_topic="test", |
| 75 | subqueries=[schema.SubQuery(label="primary", search_query="test", ranking_query="test", sources=["reddit", "x", "tiktok"])], |
| 76 | source_weights={"reddit": 0.5, "x": 0.5, "tiktok": 0.5}, |
| 77 | ) |
| 78 | |
| 79 | def test_same_story_different_phrasing_merges(self): |
| 80 | """Wireless Festival example: same event, different wording, different sources.""" |
| 81 | candidates = [ |
| 82 | make_candidate("c1", "reddit", "Kanye West to headline all three nights of Wireless Festival 2026", "Big announcement for Wireless.", 80), |
| 83 | make_candidate("c2", "x", "BREAKING: Kanye West is making his massive UK comeback at Wireless Festival this July", "Ye returns to UK.", 75), |
| 84 | make_candidate("c3", "youtube", "Kanye West BULLY Album Review - Knox Hill Reacts", "Full album reaction and breakdown.", 70), |
| 85 | ] |
| 86 | clusters = cluster.cluster_candidates(candidates, self._plan()) |
| 87 | # c1 and c2 should merge (Kanye + Wireless + Festival overlap), c3 should stay separate |
| 88 | self.assertEqual(2, len(clusters)) |
| 89 | wireless_cluster = next(cl for cl in clusters if len(cl.candidate_ids) == 2) |
| 90 | self.assertIn("c1", wireless_cluster.candidate_ids) |
| 91 | self.assertIn("c2", wireless_cluster.candidate_ids) |
| 92 | self.assertEqual(sorted(["reddit", "x"]), wireless_cluster.sources) |
| 93 | # Multi-source cluster should not have "single-source" uncertainty |
| 94 | self.assertNotEqual("single-source", wireless_cluster.uncertainty) |
| 95 | |
| 96 | def test_different_stories_dont_merge(self): |
| 97 | """Different topics should stay separate even with some entity overlap (e.g., 'Kanye').""" |
| 98 | candidates = [ |
| 99 | make_candidate("c1", "reddit", "Kanye West BULLY Album First Impressions Thread", "What do you think of BULLY?", 80), |
| 100 | make_candidate("c2", "x", "Kanye West apology for antisemitism in Wall Street Journal ad", "Full page WSJ ad.", 75), |
| 101 | make_candidate("c3", "tiktok", "Kanye West Wireless Festival ticket prices breakdown", "How much for Wireless tickets?", 70), |
| 102 | ] |
| 103 | clusters = cluster.cluster_candidates(candidates, self._plan()) |
| 104 | # These are 3 different stories, should remain as 3 clusters |
| 105 | self.assertEqual(3, len(clusters)) |
| 106 | |
| 107 | def test_same_source_clusters_dont_merge(self): |
| 108 | """Two single-source clusters from the same source should not merge via entity pass.""" |
| 109 | candidates = [ |
| 110 | make_candidate("c1", "reddit", "Kanye West Wireless Festival headline announcement", "Three nights!", 80), |
| 111 | make_candidate("c2", "reddit", "Kanye West returning to Wireless Festival confirmed", "UK comeback.", 70), |
| 112 | ] |
| 113 | clusters = cluster.cluster_candidates(candidates, self._plan()) |
| 114 | # The initial greedy pass may or may not merge these (depends on token similarity). |
| 115 | # But if they end up as separate clusters, the entity pass should NOT merge them |
| 116 | # since they're both from reddit. |
| 117 | for cl in clusters: |
| 118 | self.assertTrue(len(cl.sources) >= 1) # basic sanity |
| 119 | |
| 120 | |
| 121 | class TestPolymarketIsolation(unittest.TestCase): |
| 122 | """Polymarket clusters must not merge with non-Polymarket clusters via entity overlap.""" |
| 123 | |
| 124 | def _plan(self): |
| 125 | return schema.QueryPlan( |
| 126 | intent="breaking_news", |
| 127 | freshness_mode="strict_recent", |
| 128 | cluster_mode="story", |
| 129 | raw_topic="test", |
| 130 | subqueries=[schema.SubQuery(label="primary", search_query="test", ranking_query="test", sources=["reddit", "x", "polymarket"])], |
| 131 | source_weights={"reddit": 0.5, "x": 0.5, "polymarket": 0.5}, |
| 132 | ) |
| 133 | |
| 134 | def test_polymarket_does_not_merge_into_news_cluster(self): |
| 135 | """A Polymarket prediction about Sam Altman should not merge into a news cluster about Sam Altman.""" |
| 136 | candidates = [ |
| 137 | make_candidate("c1", "reddit", "Sam Altman personal rivalry with Elon Musk escalates", "The feud between Sam Altman and Elon Musk continues.", 80), |
| 138 | make_candidate("c2", "polymarket", "Sam Altman equity stake in OpenAI valued at $500M", "Will Sam Altman receive equity in OpenAI restructuring?", 75), |
| 139 | ] |
| 140 | clusters = cluster.cluster_candidates(candidates, self._plan()) |
| 141 | self.assertEqual(2, len(clusters), "Polymarket and news clusters should remain separate") |
| 142 | # Each cluster should have exactly one candidate |
| 143 | for cl in clusters: |
| 144 | self.assertEqual(1, len(cl.candidate_ids)) |
| 145 | |
| 146 | def test_two_polymarket_clusters_not_blocked_by_poly_guard(self): |
| 147 | """Two Polymarket items about the same topic are not blocked by the Polymarket guard. |
| 148 | |
| 149 | Note: same-source clusters are still blocked by the existing same-source |
| 150 | guard, so we verify the poly guard specifically by checking that two |
| 151 | polymarket items with high text similarity merge via the greedy pass. |
| 152 | """ |
| 153 | candidates = [ |
| 154 | make_candidate("c1", "polymarket", "Sam Altman equity stake in OpenAI restructuring", "Will Sam Altman get equity in the OpenAI restructuring deal?", 80), |
| 155 | make_candidate("c2", "polymarket", "Sam Altman equity stake in OpenAI restructuring odds", "Will Sam Altman get equity in the OpenAI restructuring deal? Current odds.", 75), |
| 156 | ] |
| 157 | clusters = cluster.cluster_candidates(candidates, self._plan()) |
| 158 | # High text similarity means greedy pass merges them |
| 159 | self.assertEqual(1, len(clusters)) |
| 160 | self.assertEqual(2, len(clusters[0].candidate_ids)) |
| 161 | |
| 162 | def test_neither_polymarket_still_merges(self): |
| 163 | """Non-Polymarket clusters with entity overlap should still merge (existing behavior).""" |
| 164 | candidates = [ |
| 165 | make_candidate("c1", "reddit", "Sam Altman OpenAI restructuring announcement details", "Sam Altman announces major OpenAI restructuring.", 80), |
| 166 | make_candidate("c2", "x", "Sam Altman reveals OpenAI restructuring plan for 2026", "Major OpenAI restructuring coming says Sam Altman.", 75), |
| 167 | ] |
| 168 | clusters = cluster.cluster_candidates(candidates, self._plan()) |
| 169 | self.assertEqual(1, len(clusters)) |
| 170 | self.assertEqual(2, len(clusters[0].candidate_ids)) |
| 171 | |
| 172 | |
| 173 | class TestStaleClusterDemotion(unittest.TestCase): |
| 174 | """Stale candidates must never become cluster representatives or titles.""" |
| 175 | |
| 176 | def _plan(self): |
| 177 | return schema.QueryPlan( |
| 178 | intent="breaking_news", |
| 179 | freshness_mode="strict_recent", |
| 180 | cluster_mode="story", |
| 181 | raw_topic="test", |
| 182 | subqueries=[schema.SubQuery(label="primary", search_query="test", ranking_query="test", sources=["reddit", "x"])], |
| 183 | source_weights={"reddit": 0.5, "x": 0.5}, |
| 184 | ) |
| 185 | |
| 186 | def _candidate_with_date( |
| 187 | self, candidate_id: str, source: str, title: str, score: float, |
| 188 | published_at: str, date_confidence: str, range_from: str, range_to: str, |
| 189 | ) -> schema.Candidate: |
| 190 | item = schema.SourceItem( |
| 191 | item_id=candidate_id, |
| 192 | source=source, |
| 193 | title=title, |
| 194 | body=title, |
| 195 | url=f"https://example.com/{candidate_id}", |
| 196 | published_at=published_at, |
| 197 | date_confidence=date_confidence, |
| 198 | ) |
| 199 | return schema.Candidate( |
| 200 | candidate_id=candidate_id, |
| 201 | item_id=candidate_id, |
| 202 | source=source, |
| 203 | title=title, |
| 204 | url=f"https://example.com/{candidate_id}", |
| 205 | snippet=title, |
| 206 | subquery_labels=["primary"], |
| 207 | native_ranks={"primary:reddit": 1}, |
| 208 | local_relevance=0.8, |
| 209 | freshness=80, |
| 210 | engagement=10, |
| 211 | source_quality=0.7, |
| 212 | rrf_score=0.02, |
| 213 | rerank_score=score, |
| 214 | final_score=score, |
| 215 | source_items=[item], |
| 216 | metadata={"range_from": range_from, "range_to": range_to}, |
| 217 | ) |
| 218 | |
| 219 | def test_stale_candidate_not_cluster_representative(self): |
| 220 | """A stale item with higher final_score must not lead a cluster over a fresh item. |
| 221 | |
| 222 | This guards against the issue where a 2025-10 video ranked #1 in a |
| 223 | 2026-07 brief because clustering re-sorted by final_score alone. |
| 224 | """ |
| 225 | range_from = "2026-06-15" |
| 226 | range_to = "2026-07-15" |
| 227 | stale = self._candidate_with_date( |
| 228 | "stale", "reddit", "Model launch reactions discussion", |
| 229 | score=95.0, |
| 230 | published_at="2025-10-15", |
| 231 | date_confidence="low", |
| 232 | range_from=range_from, |
| 233 | range_to=range_to, |
| 234 | ) |
| 235 | fresh = self._candidate_with_date( |
| 236 | "fresh", "x", "Model launch reactions update", |
| 237 | score=50.0, |
| 238 | published_at="2026-07-10", |
| 239 | date_confidence="high", |
| 240 | range_from=range_from, |
| 241 | range_to=range_to, |
| 242 | ) |
| 243 | candidates = [stale, fresh] |
| 244 | clusters = cluster.cluster_candidates(candidates, self._plan()) |
| 245 | |
| 246 | self.assertEqual(1, len(clusters)) |
| 247 | self.assertEqual("fresh", clusters[0].representative_ids[0]) |
| 248 | self.assertEqual(fresh.title, clusters[0].title) |
| 249 | |
| 250 | |
| 251 | class TestClusterUncertainty(unittest.TestCase): |
| 252 | def test_single_source_returns_single_source(self): |
| 253 | candidates = [make_candidate("c1", "reddit", "Title", "Body", 80)] |
| 254 | result = cluster._cluster_uncertainty(candidates) |
| 255 | self.assertEqual("single-source", result) |
| 256 | |
| 257 | def test_multi_source_high_score_returns_none(self): |
| 258 | candidates = [ |
| 259 | make_candidate("c1", "reddit", "Title", "Body", 80), |
| 260 | make_candidate("c2", "x", "Title2", "Body2", 70), |
| 261 | ] |
| 262 | result = cluster._cluster_uncertainty(candidates) |
| 263 | self.assertIsNone(result) |
| 264 | |
| 265 | def test_multi_source_low_score_returns_thin_evidence(self): |
| 266 | candidates = [ |
| 267 | make_candidate("c1", "reddit", "Title", "Body", 30), |
| 268 | make_candidate("c2", "x", "Title2", "Body2", 40), |
| 269 | ] |
| 270 | result = cluster._cluster_uncertainty(candidates) |
| 271 | self.assertEqual("thin-evidence", result) |
| 272 | |
| 273 | if __name__ == "__main__": |
| 274 | unittest.main() |
| 275 |