| 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 TestClusterUncertainty(unittest.TestCase): |
| 174 | def test_single_source_returns_single_source(self): |
| 175 | candidates = [make_candidate("c1", "reddit", "Title", "Body", 80)] |
| 176 | result = cluster._cluster_uncertainty(candidates) |
| 177 | self.assertEqual("single-source", result) |
| 178 | |
| 179 | def test_multi_source_high_score_returns_none(self): |
| 180 | candidates = [ |
| 181 | make_candidate("c1", "reddit", "Title", "Body", 80), |
| 182 | make_candidate("c2", "x", "Title2", "Body2", 70), |
| 183 | ] |
| 184 | result = cluster._cluster_uncertainty(candidates) |
| 185 | self.assertIsNone(result) |
| 186 | |
| 187 | def test_multi_source_low_score_returns_thin_evidence(self): |
| 188 | candidates = [ |
| 189 | make_candidate("c1", "reddit", "Title", "Body", 30), |
| 190 | make_candidate("c2", "x", "Title2", "Body2", 40), |
| 191 | ] |
| 192 | result = cluster._cluster_uncertainty(candidates) |
| 193 | self.assertEqual("thin-evidence", result) |
| 194 | |
| 195 | if __name__ == "__main__": |
| 196 | unittest.main() |
| 197 |