| 1 | import unittest |
| 2 | |
| 3 | from lib import fusion, schema |
| 4 | |
| 5 | |
| 6 | def make_item(item_id: str, source: str, url: str, title: str, rank_score: float) -> schema.SourceItem: |
| 7 | return schema.SourceItem( |
| 8 | item_id=item_id, |
| 9 | source=source, |
| 10 | title=title, |
| 11 | body=title, |
| 12 | url=url, |
| 13 | relevance_hint=rank_score, |
| 14 | snippet=title, |
| 15 | metadata={ |
| 16 | "local_relevance": rank_score, |
| 17 | "freshness": 80, |
| 18 | "engagement_score": 5, |
| 19 | "source_quality": 0.7, |
| 20 | }, |
| 21 | ) |
| 22 | |
| 23 | |
| 24 | class FusionV3Tests(unittest.TestCase): |
| 25 | def test_weighted_rrf_merges_duplicate_urls(self): |
| 26 | plan = schema.QueryPlan( |
| 27 | intent="breaking_news", |
| 28 | freshness_mode="strict_recent", |
| 29 | cluster_mode="story", |
| 30 | raw_topic="test", |
| 31 | subqueries=[ |
| 32 | schema.SubQuery(label="primary", search_query="test", ranking_query="What happened in test?", sources=["reddit", "x"], weight=0.7), |
| 33 | schema.SubQuery(label="reaction", search_query="test reaction", ranking_query="What are the reactions to test?", sources=["x"], weight=0.3), |
| 34 | ], |
| 35 | source_weights={"reddit": 0.4, "x": 0.6}, |
| 36 | ) |
| 37 | shared = "https://example.com/shared" |
| 38 | streams = { |
| 39 | ("primary", "reddit"): [make_item("r1", "reddit", shared, "Shared item", 0.8)], |
| 40 | ("primary", "x"): [make_item("x1", "x", shared, "Shared item", 0.9)], |
| 41 | ("reaction", "x"): [make_item("x2", "x", "https://example.com/unique", "Unique item", 0.7)], |
| 42 | } |
| 43 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=10) |
| 44 | self.assertEqual(2, len(candidates)) |
| 45 | merged = next(candidate for candidate in candidates if candidate.url == shared) |
| 46 | self.assertEqual({"primary"}, set(merged.subquery_labels)) |
| 47 | self.assertEqual(2, len(merged.native_ranks)) |
| 48 | self.assertEqual({"reddit", "x"}, set(merged.sources)) |
| 49 | self.assertEqual(2, len(merged.source_items)) |
| 50 | |
| 51 | def test_diversify_pool_guarantees_min_per_qualifying_source(self): |
| 52 | """Every qualifying source (local_relevance >= 0.25) gets at least 2 |
| 53 | items in the fused pool. |
| 54 | |
| 55 | Dominant sources (x, tiktok) get high weights, so pure-RRF truncation |
| 56 | would squeeze out low-weight sources entirely. The diversity guarantee |
| 57 | must reserve at least 2 slots per qualifying active source. All sources |
| 58 | here have rank_score=0.8 (well above the 0.25 threshold), so every |
| 59 | source qualifies for reserved slots. |
| 60 | """ |
| 61 | sources = ["reddit", "hackernews", "x", "tiktok", "bluesky", "youtube"] |
| 62 | # Heavily skewed weights: x and tiktok dominate. |
| 63 | weights = { |
| 64 | "x": 3.0, |
| 65 | "tiktok": 2.5, |
| 66 | "reddit": 0.5, |
| 67 | "hackernews": 0.4, |
| 68 | "bluesky": 0.3, |
| 69 | "youtube": 0.3, |
| 70 | } |
| 71 | plan = schema.QueryPlan( |
| 72 | intent="concept", |
| 73 | freshness_mode="relaxed", |
| 74 | cluster_mode="concept", |
| 75 | raw_topic="RAG", |
| 76 | subqueries=[ |
| 77 | schema.SubQuery( |
| 78 | label="primary", |
| 79 | search_query="RAG", |
| 80 | ranking_query="What is RAG?", |
| 81 | sources=sources, |
| 82 | weight=1.0, |
| 83 | ), |
| 84 | ], |
| 85 | source_weights=weights, |
| 86 | ) |
| 87 | streams: dict[tuple[str, str], list[schema.SourceItem]] = {} |
| 88 | for src in sources: |
| 89 | items = [] |
| 90 | for rank in range(4): |
| 91 | items.append( |
| 92 | make_item( |
| 93 | item_id=f"{src}_{rank}", |
| 94 | source=src, |
| 95 | url=f"https://{src}.example.com/{rank}", |
| 96 | title=f"{src} item {rank}", |
| 97 | rank_score=0.8, |
| 98 | ) |
| 99 | ) |
| 100 | streams[("primary", src)] = items |
| 101 | |
| 102 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=12) |
| 103 | self.assertEqual(12, len(candidates)) |
| 104 | |
| 105 | source_counts: dict[str, int] = {} |
| 106 | for c in candidates: |
| 107 | source_counts[c.source] = source_counts.get(c.source, 0) + 1 |
| 108 | |
| 109 | for src in sources: |
| 110 | self.assertGreaterEqual( |
| 111 | source_counts.get(src, 0), |
| 112 | 2, |
| 113 | f"Source '{src}' has {source_counts.get(src, 0)} items, expected >= 2", |
| 114 | ) |
| 115 | |
| 116 | def test_diversify_pool_denies_slots_for_low_relevance_source(self): |
| 117 | """Sources with best local_relevance < 0.25 do not get reserved slots. |
| 118 | |
| 119 | Create two sources: 'x' with local_relevance=0.5 (qualifies) and |
| 120 | 'reddit' with local_relevance=0.1 (below threshold). With a tight |
| 121 | pool_limit, the high-relevance source gets reserved slots while |
| 122 | the low-relevance source must compete on RRF merit alone. |
| 123 | """ |
| 124 | plan = schema.QueryPlan( |
| 125 | intent="concept", |
| 126 | freshness_mode="relaxed", |
| 127 | cluster_mode="concept", |
| 128 | raw_topic="test", |
| 129 | subqueries=[ |
| 130 | schema.SubQuery( |
| 131 | label="primary", |
| 132 | search_query="test", |
| 133 | ranking_query="What is test?", |
| 134 | sources=["x", "reddit"], |
| 135 | weight=1.0, |
| 136 | ), |
| 137 | ], |
| 138 | source_weights={"x": 1.0, "reddit": 1.0}, |
| 139 | ) |
| 140 | |
| 141 | # x items: high relevance (0.5) -- qualifies for diversity reservation |
| 142 | x_items = [ |
| 143 | make_item(f"x_{i}", "x", f"https://x.example.com/{i}", f"x item {i}", 0.5) |
| 144 | for i in range(4) |
| 145 | ] |
| 146 | |
| 147 | # reddit items: low relevance (0.1) -- below threshold, no reserved slots |
| 148 | reddit_items = [ |
| 149 | make_item(f"r_{i}", "reddit", f"https://reddit.example.com/{i}", f"reddit item {i}", 0.1) |
| 150 | for i in range(4) |
| 151 | ] |
| 152 | |
| 153 | streams = { |
| 154 | ("primary", "x"): x_items, |
| 155 | ("primary", "reddit"): reddit_items, |
| 156 | } |
| 157 | |
| 158 | # pool_limit=3: x gets 2 reserved + 1 more by RRF. Reddit has no |
| 159 | # reserved slots, so it must out-score x items in the remainder. |
| 160 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=3) |
| 161 | self.assertEqual(3, len(candidates)) |
| 162 | |
| 163 | # x must have at least 2 (reserved slots) |
| 164 | x_count = sum(1 for c in candidates if c.source == "x") |
| 165 | self.assertGreaterEqual(x_count, 2, "x should have at least 2 reserved slots") |
| 166 | |
| 167 | def test_diversify_pool_no_reservation_when_all_below_threshold(self): |
| 168 | """When all sources are below the relevance threshold, no reserved slots |
| 169 | are granted. The pool is filled purely by RRF score order.""" |
| 170 | plan = schema.QueryPlan( |
| 171 | intent="concept", |
| 172 | freshness_mode="relaxed", |
| 173 | cluster_mode="concept", |
| 174 | raw_topic="test", |
| 175 | subqueries=[ |
| 176 | schema.SubQuery( |
| 177 | label="primary", |
| 178 | search_query="test", |
| 179 | ranking_query="What is test?", |
| 180 | sources=["x", "reddit", "hackernews"], |
| 181 | weight=1.0, |
| 182 | ), |
| 183 | ], |
| 184 | # Give x a much higher weight so its items get higher RRF scores |
| 185 | source_weights={"x": 3.0, "reddit": 0.3, "hackernews": 0.3}, |
| 186 | ) |
| 187 | |
| 188 | streams: dict[tuple[str, str], list[schema.SourceItem]] = {} |
| 189 | # All sources below threshold (local_relevance = 0.1) |
| 190 | for src in ["x", "reddit", "hackernews"]: |
| 191 | items = [ |
| 192 | make_item(f"{src}_{i}", src, f"https://{src}.example.com/{i}", f"{src} item {i}", 0.1) |
| 193 | for i in range(4) |
| 194 | ] |
| 195 | streams[("primary", src)] = items |
| 196 | |
| 197 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=4) |
| 198 | self.assertEqual(4, len(candidates)) |
| 199 | |
| 200 | # With no diversity reservation and x having 3x the weight, |
| 201 | # x should dominate the top slots purely on RRF score |
| 202 | source_counts: dict[str, int] = {} |
| 203 | for c in candidates: |
| 204 | source_counts[c.source] = source_counts.get(c.source, 0) + 1 |
| 205 | |
| 206 | # x has 3x weight so its RRF scores are ~3x higher than reddit/hn. |
| 207 | # All 4 x items should beat all reddit/hackernews items. |
| 208 | self.assertEqual( |
| 209 | source_counts.get("x", 0), |
| 210 | 4, |
| 211 | f"Expected x to take all 4 slots on pure RRF merit, got {source_counts}", |
| 212 | ) |
| 213 | |
| 214 | def test_diversify_pool_threshold_boundary(self): |
| 215 | """Source with best local_relevance exactly at the threshold (0.25) |
| 216 | qualifies for reserved slots.""" |
| 217 | plan = schema.QueryPlan( |
| 218 | intent="concept", |
| 219 | freshness_mode="relaxed", |
| 220 | cluster_mode="concept", |
| 221 | raw_topic="boundary", |
| 222 | subqueries=[ |
| 223 | schema.SubQuery( |
| 224 | label="primary", |
| 225 | search_query="boundary", |
| 226 | ranking_query="What is boundary?", |
| 227 | sources=["x", "reddit"], |
| 228 | weight=1.0, |
| 229 | ), |
| 230 | ], |
| 231 | # Give x much higher weight so it would dominate without reservation |
| 232 | source_weights={"x": 5.0, "reddit": 0.1}, |
| 233 | ) |
| 234 | |
| 235 | x_items = [ |
| 236 | make_item(f"x_{i}", "x", f"https://x.example.com/{i}", f"x item {i}", 0.8) |
| 237 | for i in range(6) |
| 238 | ] |
| 239 | |
| 240 | # reddit at exactly the threshold |
| 241 | reddit_items = [ |
| 242 | make_item(f"r_{i}", "reddit", f"https://reddit.example.com/{i}", f"reddit item {i}", 0.25) |
| 243 | for i in range(3) |
| 244 | ] |
| 245 | |
| 246 | streams = { |
| 247 | ("primary", "x"): x_items, |
| 248 | ("primary", "reddit"): reddit_items, |
| 249 | } |
| 250 | |
| 251 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=6) |
| 252 | self.assertEqual(6, len(candidates)) |
| 253 | |
| 254 | reddit_count = sum(1 for c in candidates if c.source == "reddit") |
| 255 | self.assertGreaterEqual( |
| 256 | reddit_count, |
| 257 | 2, |
| 258 | f"reddit (local_relevance=0.25, at threshold) should get 2 reserved slots, got {reddit_count}", |
| 259 | ) |
| 260 | |
| 261 | |
| 262 | def make_item_with_author( |
| 263 | item_id: str, source: str, url: str, title: str, rank_score: float, author: str | None = None, |
| 264 | ) -> schema.SourceItem: |
| 265 | return schema.SourceItem( |
| 266 | item_id=item_id, |
| 267 | source=source, |
| 268 | title=title, |
| 269 | body=title, |
| 270 | url=url, |
| 271 | author=author, |
| 272 | relevance_hint=rank_score, |
| 273 | snippet=title, |
| 274 | metadata={ |
| 275 | "local_relevance": rank_score, |
| 276 | "freshness": 80, |
| 277 | "engagement_score": 5, |
| 278 | "source_quality": 0.7, |
| 279 | }, |
| 280 | ) |
| 281 | |
| 282 | |
| 283 | class TestPerAuthorCap(unittest.TestCase): |
| 284 | """Per-author cap: no single author should have more than 3 items in fused pool.""" |
| 285 | |
| 286 | def _make_plan(self, sources: list[str]) -> schema.QueryPlan: |
| 287 | return schema.QueryPlan( |
| 288 | intent="breaking_news", |
| 289 | freshness_mode="strict_recent", |
| 290 | cluster_mode="story", |
| 291 | raw_topic="test", |
| 292 | subqueries=[ |
| 293 | schema.SubQuery( |
| 294 | label="primary", |
| 295 | search_query="test", |
| 296 | ranking_query="test", |
| 297 | sources=sources, |
| 298 | weight=1.0, |
| 299 | ), |
| 300 | ], |
| 301 | source_weights={s: 1.0 for s in sources}, |
| 302 | ) |
| 303 | |
| 304 | def test_author_with_8_items_capped_to_3(self): |
| 305 | """@grok scenario: 8 items from the same author, only best 3 survive.""" |
| 306 | plan = self._make_plan(["x"]) |
| 307 | items = [ |
| 308 | make_item_with_author( |
| 309 | f"x_{i}", "x", f"https://x.com/{i}", f"grok summary {i}", 0.7, author="@grok", |
| 310 | ) |
| 311 | for i in range(8) |
| 312 | ] |
| 313 | streams = {("primary", "x"): items} |
| 314 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=20) |
| 315 | grok_count = sum( |
| 316 | 1 for c in candidates |
| 317 | if any(si.author == "@grok" for si in c.source_items) |
| 318 | ) |
| 319 | self.assertLessEqual(grok_count, 3, f"@grok should be capped at 3, got {grok_count}") |
| 320 | |
| 321 | def test_author_with_3_items_all_kept(self): |
| 322 | """Author with exactly 3 items should keep all of them.""" |
| 323 | plan = self._make_plan(["x"]) |
| 324 | items = [ |
| 325 | make_item_with_author( |
| 326 | f"x_{i}", "x", f"https://x.com/{i}", f"author3 post {i}", 0.7, author="@author3", |
| 327 | ) |
| 328 | for i in range(3) |
| 329 | ] |
| 330 | streams = {("primary", "x"): items} |
| 331 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=20) |
| 332 | count = sum( |
| 333 | 1 for c in candidates |
| 334 | if any(si.author == "@author3" for si in c.source_items) |
| 335 | ) |
| 336 | self.assertEqual(count, 3) |
| 337 | |
| 338 | def test_items_without_author_not_capped(self): |
| 339 | """Items with no author field should never be dropped by the cap.""" |
| 340 | plan = self._make_plan(["reddit"]) |
| 341 | items = [ |
| 342 | make_item_with_author( |
| 343 | f"r_{i}", "reddit", f"https://reddit.com/{i}", f"post {i}", 0.7, author=None, |
| 344 | ) |
| 345 | for i in range(6) |
| 346 | ] |
| 347 | streams = {("primary", "reddit"): items} |
| 348 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=20) |
| 349 | self.assertEqual(len(candidates), 6) |
| 350 | |
| 351 | def test_multiple_authors_capped_independently(self): |
| 352 | """Two prolific authors each get capped to 3 independently.""" |
| 353 | plan = self._make_plan(["x"]) |
| 354 | items = [] |
| 355 | for i in range(5): |
| 356 | items.append(make_item_with_author( |
| 357 | f"grok_{i}", "x", f"https://x.com/grok/{i}", f"grok {i}", 0.7, author="@grok", |
| 358 | )) |
| 359 | for i in range(5): |
| 360 | items.append(make_item_with_author( |
| 361 | f"spam_{i}", "x", f"https://x.com/spam/{i}", f"spam {i}", 0.6, author="@spammer", |
| 362 | )) |
| 363 | streams = {("primary", "x"): items} |
| 364 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=20) |
| 365 | grok_count = sum(1 for c in candidates if any(si.author == "@grok" for si in c.source_items)) |
| 366 | spam_count = sum(1 for c in candidates if any(si.author == "@spammer" for si in c.source_items)) |
| 367 | self.assertLessEqual(grok_count, 3) |
| 368 | self.assertLessEqual(spam_count, 3) |
| 369 | |
| 370 | def test_cap_keeps_best_items_by_rrf_order(self): |
| 371 | """The cap should keep the first (highest-ranked) items per author.""" |
| 372 | plan = self._make_plan(["x"]) |
| 373 | # Items with decreasing relevance scores so ranking is deterministic |
| 374 | items = [ |
| 375 | make_item_with_author( |
| 376 | f"x_{i}", "x", f"https://x.com/{i}", f"post {i}", 0.9 - (i * 0.05), author="@prolific", |
| 377 | ) |
| 378 | for i in range(5) |
| 379 | ] |
| 380 | streams = {("primary", "x"): items} |
| 381 | candidates = fusion.weighted_rrf(streams, plan, pool_limit=20) |
| 382 | kept_ids = {c.item_id for c in candidates if any(si.author == "@prolific" for si in c.source_items)} |
| 383 | # The top 3 items (x_0, x_1, x_2) should be kept |
| 384 | self.assertLessEqual(len(kept_ids), 3) |
| 385 | |
| 386 | |
| 387 | class TestUrlNormalization(unittest.TestCase): |
| 388 | def test_strips_www(self): |
| 389 | from lib.fusion import _normalize_url |
| 390 | self.assertEqual( |
| 391 | _normalize_url("https://www.reddit.com/r/test"), |
| 392 | _normalize_url("https://reddit.com/r/test"), |
| 393 | ) |
| 394 | |
| 395 | def test_strips_old_prefix(self): |
| 396 | from lib.fusion import _normalize_url |
| 397 | self.assertEqual( |
| 398 | _normalize_url("https://old.reddit.com/r/test"), |
| 399 | _normalize_url("https://reddit.com/r/test"), |
| 400 | ) |
| 401 | |
| 402 | def test_strips_mobile_prefix(self): |
| 403 | from lib.fusion import _normalize_url |
| 404 | self.assertEqual( |
| 405 | _normalize_url("https://m.youtube.com/watch?v=abc"), |
| 406 | _normalize_url("https://youtube.com/watch?v=abc"), |
| 407 | ) |
| 408 | |
| 409 | def test_strips_utm_params(self): |
| 410 | from lib.fusion import _normalize_url |
| 411 | self.assertEqual( |
| 412 | _normalize_url("https://example.com/page?utm_source=twitter&id=5"), |
| 413 | _normalize_url("https://example.com/page?id=5"), |
| 414 | ) |
| 415 | |
| 416 | def test_strips_trailing_slash(self): |
| 417 | from lib.fusion import _normalize_url |
| 418 | self.assertEqual( |
| 419 | _normalize_url("https://example.com/page/"), |
| 420 | _normalize_url("https://example.com/page"), |
| 421 | ) |
| 422 | |
| 423 | def test_preserves_non_tracking_params(self): |
| 424 | from lib.fusion import _normalize_url |
| 425 | result = _normalize_url("https://example.com/page?id=5&sort=new") |
| 426 | self.assertIn("id=5", result) |
| 427 | self.assertIn("sort=new", result) |
| 428 | |
| 429 | def test_case_insensitive(self): |
| 430 | from lib.fusion import _normalize_url |
| 431 | self.assertEqual( |
| 432 | _normalize_url("https://Reddit.com/r/Test"), |
| 433 | _normalize_url("https://reddit.com/r/test"), |
| 434 | ) |
| 435 | |
| 436 | if __name__ == "__main__": |
| 437 | unittest.main() |
| 438 |