Merge pull request #198 from iliaal/perf/pipeline-optimizations

perf: optimize dedup, parallelize handle searches and enrichment

scripts/lib/bird_x.py

@@ -312,10 +312,9 @@ def search_handles(
     Returns:
         List of raw item dicts (same format as parse_bird_response output).
     """
-    all_items = []
     core_topic = _extract_core_subject(topic) if topic else None
 
-    for handle in handles:
+    def _search_one_handle(handle: str) -> List[Dict[str, Any]]:
         handle = handle.lstrip("@")
         if core_topic:
             query = f"from:{handle} {core_topic} since:{from_date}"
@@ -350,24 +349,32 @@ def search_handles(
                     proc.kill()
                 proc.wait(timeout=5)
                 _log(f"Handle search timed out for @{handle}")
-                continue
+                return []
 
             if proc.returncode != 0:
                 _log(f"Handle search failed for @{handle}: {(stderr or '').strip()}")
-                continue
+                return []
 
             output = (stdout or "").strip()
             if not output:
-                continue
+                return []
 
             response = json.loads(output)
-            items = parse_bird_response(response, query=core_topic)
+            return parse_bird_response(response, query=core_topic)
-            all_items.extend(items)
 
         except json.JSONDecodeError:
             _log(f"Invalid JSON from handle search for @{handle}")
         except (OSError, subprocess.SubprocessError) as e:
             _log(f"Handle search error for @{handle}: {e}")
+        return []
+
+    from concurrent.futures import ThreadPoolExecutor, as_completed
+
+    all_items: List[Dict[str, Any]] = []
+    with ThreadPoolExecutor(max_workers=min(5, len(handles))) as executor:
+        futures = {executor.submit(_search_one_handle, h): h for h in handles}
+        for future in as_completed(futures):
+            all_items.extend(future.result())
 
     return all_items
 

scripts/lib/cluster.py

@@ -57,28 +57,34 @@ def _entity_overlap(entities_a: set[str], entities_b: set[str]) -> float:
 
 def _mmr_representatives(
     candidates: list[schema.Candidate],
+    text_cache: dict[str, dedupe._PreparedText],
     limit: int = 3,
     diversity_lambda: float = 0.75,
 ) -> list[str]:
     selected: list[schema.Candidate] = []
+    remaining_set = {c.candidate_id for c in candidates}
     remaining = list(candidates)
     while remaining and len(selected) < limit:
         if not selected:
             best = max(remaining, key=lambda candidate: candidate.final_score)
             selected.append(best)
-            remaining.remove(best)
+            remaining_set.discard(best.candidate_id)
+            remaining = [c for c in remaining if c.candidate_id in remaining_set]
             continue
 
+        selected_preps = [text_cache[c.candidate_id] for c in selected]
+
         def score(candidate: schema.Candidate) -> float:
+            prep = text_cache[candidate.candidate_id]
             diversity_penalty = max(
-                dedupe.hybrid_similarity(_candidate_text(candidate), _candidate_text(existing))
+                dedupe.prepared_similarity(prep, sp) for sp in selected_preps
-                for existing in selected
             )
             return (diversity_lambda * candidate.final_score) - ((1 - diversity_lambda) * diversity_penalty * 100)
 
         best = max(remaining, key=score)
         selected.append(best)
-        remaining.remove(best)
+        remaining_set.discard(best.candidate_id)
+        remaining = [c for c in remaining if c.candidate_id in remaining_set]
     return [candidate.candidate_id for candidate in selected]
 
 
@@ -105,15 +111,21 @@ def cluster_candidates(
             )
         return clusters
 
+    text_cache: dict[str, dedupe._PreparedText] = {
+        c.candidate_id: dedupe._PreparedText(_candidate_text(c))
+        for c in candidates
+    }
+
     groups: list[list[schema.Candidate]] = []
     # Lower threshold for breaking_news: related articles share fewer exact
     # words but cover the same event.
     threshold = 0.42 if plan.intent == "breaking_news" else 0.48
     for candidate in candidates:
         assigned = False
+        cand_prep = text_cache[candidate.candidate_id]
         for group in groups:
             leader = group[0]
-            similarity = dedupe.hybrid_similarity(_candidate_text(candidate), _candidate_text(leader))
+            similarity = dedupe.prepared_similarity(cand_prep, text_cache[leader.candidate_id])
             if similarity >= threshold:
                 group.append(candidate)
                 assigned = True
@@ -125,7 +137,7 @@ def cluster_candidates(
     for index, group in enumerate(groups, start=1):
         group.sort(key=lambda candidate: candidate.final_score, reverse=True)
         cluster_id = f"cluster-{index}"
-        representatives = _mmr_representatives(group)
+        representatives = _mmr_representatives(group, text_cache)
         for candidate in group:
             candidate.cluster_id = cluster_id
         clusters.append(
@@ -225,7 +237,11 @@ def _merge_entity_clusters(
         # Pick representatives from combined pool
         combined_candidates = [candidate_map[cid] for cid in combined_cids if cid in candidate_map]
         combined_candidates.sort(key=lambda c: c.final_score, reverse=True)
-        reps = _mmr_representatives(combined_candidates)
+        merge_text_cache = {
+            c.candidate_id: dedupe._PreparedText(_candidate_text(c))
+            for c in combined_candidates
+        }
+        reps = _mmr_representatives(combined_candidates, merge_text_cache)
 
         cluster_id = cl.cluster_id
         for cid in combined_cids:

scripts/lib/cookie_extract.py

@@ -7,6 +7,7 @@ Only uses Python stdlib — no external dependencies.
 """
 
 import configparser
+import functools
 import logging
 import platform
 import shutil
@@ -18,6 +19,40 @@ from typing import Dict, List, Optional
 logger = logging.getLogger(__name__)
 
 
+@functools.lru_cache(maxsize=1)
+def _is_wsl() -> bool:
+    """Detect if running under Windows Subsystem for Linux.
+
+    Cached after the first call since /proc/version doesn't change at runtime.
+    """
+    try:
+        return "microsoft" in Path("/proc/version").read_text().lower()
+    except OSError:
+        return False
+
+
+def _get_wsl_firefox_profiles_dir() -> Optional[Path]:
+    """Find Firefox profiles directory on the Windows host from WSL.
+
+    Scans /mnt/c/Users/*/AppData/Roaming/Mozilla/Firefox for real user
+    directories (skips Public, Default, etc.).
+    """
+    mnt_users = Path("/mnt/c/Users")
+    if not mnt_users.is_dir():
+        return None
+    skip = {"Public", "Default", "Default User", "All Users"}
+    try:
+        for user_dir in sorted(mnt_users.iterdir()):
+            if user_dir.name in skip or not user_dir.is_dir():
+                continue
+            ff_dir = user_dir / "AppData" / "Roaming" / "Mozilla" / "Firefox"
+            if ff_dir.is_dir():
+                return ff_dir
+    except OSError:
+        pass
+    return None
+
+
 def _get_firefox_profiles_dir() -> Optional[Path]:
     """Return the Firefox profiles directory for the current platform, or None."""
     system = platform.system()
@@ -45,12 +80,7 @@ def _find_default_profile(profiles_dir: Path) -> Optional[Path]:
             config = configparser.ConfigParser()
             config.read(str(ini_path), encoding="utf-8")
 
-            # First pass: look for Default=1
+            # First pass: Install* section (Firefox >= 67 format, takes priority)
-            for section in config.sections():
-                if config.has_option(section, "Default") and config.get(section, "Default") == "1":
-                    return _resolve_profile_path(profiles_dir, config, section)
-
-            # Second pass: first Install* section with Default key (Firefox >= 67 format)
             for section in config.sections():
                 if section.startswith("Install") and config.has_option(section, "Default"):
                     raw = config.get(section, "Default")
@@ -58,6 +88,11 @@ def _find_default_profile(profiles_dir: Path) -> Optional[Path]:
                     if candidate.is_dir():
                         return candidate
 
+            # Second pass: Profile section with Default=1
+            for section in config.sections():
+                if section.startswith("Profile") and config.has_option(section, "Default") and config.get(section, "Default") == "1":
+                    return _resolve_profile_path(profiles_dir, config, section)
+
             # Third pass: first Profile section that exists on disk
             for section in config.sections():
                 if section.startswith("Profile"):
@@ -153,6 +188,15 @@ def _query_cookies_db(
                 pass
 
 
+def _try_firefox_dir(profiles_dir: Path, domain: str, cookie_names: List[str]) -> Optional[Dict[str, str]]:
+    """Try to extract cookies from a Firefox profiles directory."""
+    profile_path = _find_default_profile(profiles_dir)
+    if profile_path is None:
+        logger.debug("No Firefox profile found in %s", profiles_dir)
+        return None
+    return _query_cookies_db(profile_path / "cookies.sqlite", domain, cookie_names)
+
+
 def extract_firefox_cookies(
     domain: str, cookie_names: List[str]
 ) -> Optional[Dict[str, str]]:
@@ -161,6 +205,10 @@ def extract_firefox_cookies(
     Finds the default Firefox profile, copies cookies.sqlite to a temp file
     (to avoid lock conflicts), and queries for the requested cookies.
 
+    On WSL2, falls back to Windows Firefox if native Linux Firefox has no
+    matching cookies. Windows Firefox cookies are unencrypted, so this works
+    without DPAPI or any Windows-side helpers.
+
     Args:
         domain: The cookie domain to match (e.g. ".x.com"). Matched with LIKE %domain.
         cookie_names: List of cookie names to extract (e.g. ["auth_token", "ct0"]).
@@ -169,17 +217,20 @@ def extract_firefox_cookies(
         Dict of {cookie_name: cookie_value} or None if extraction fails.
     """
     profiles_dir = _get_firefox_profiles_dir()
+    if profiles_dir is not None:
+        result = _try_firefox_dir(profiles_dir, domain, cookie_names)
+        if result is not None:
+            return result
+
+    if platform.system() == "Linux" and _is_wsl():
+        wsl_dir = _get_wsl_firefox_profiles_dir()
+        if wsl_dir is not None:
+            logger.debug("Trying Windows Firefox via WSL: %s", wsl_dir)
+            return _try_firefox_dir(wsl_dir, domain, cookie_names)
+
     if profiles_dir is None:
         logger.debug("Firefox profiles directory not found")
-        return None
+    return None
-
-    profile_path = _find_default_profile(profiles_dir)
-    if profile_path is None:
-        logger.debug("No Firefox profile found in %s", profiles_dir)
-        return None
-
-    db_path = profile_path / "cookies.sqlite"
-    return _query_cookies_db(db_path, domain, cookie_names)
 
 
 def extract_chrome_cookies(
@@ -248,6 +299,31 @@ def extract_cookies(
     return cookies
 
 
+def _extract_firefox_with_source(
+    domain: str, cookie_names: List[str]
+) -> Optional[tuple[Dict[str, str], str]]:
+    """Extract Firefox cookies and report whether they came from native or WSL.
+
+    Returns (cookies, "firefox") for native Linux/macOS Firefox, or
+    (cookies, "firefox-wsl") for Windows Firefox accessed via WSL2.
+    """
+    profiles_dir = _get_firefox_profiles_dir()
+    if profiles_dir is not None:
+        result = _try_firefox_dir(profiles_dir, domain, cookie_names)
+        if result is not None:
+            return (result, "firefox")
+
+    if platform.system() == "Linux" and _is_wsl():
+        wsl_dir = _get_wsl_firefox_profiles_dir()
+        if wsl_dir is not None:
+            logger.debug("Trying Windows Firefox via WSL: %s", wsl_dir)
+            result = _try_firefox_dir(wsl_dir, domain, cookie_names)
+            if result is not None:
+                return (result, "firefox-wsl")
+
+    return None
+
+
 def extract_cookies_with_source(
     browser: str, domain: str, cookie_names: list[str]
 ) -> Optional[tuple[dict[str, str], str]]:
@@ -263,6 +339,7 @@ def extract_cookies_with_source(
 
     Returns:
         Tuple of ({cookie_name: cookie_value}, browser_name) or None.
+        browser_name is "firefox-wsl" when cookies came from Windows Firefox via WSL2.
     """
     extractors = {
         "firefox": extract_firefox_cookies,
@@ -271,6 +348,8 @@ def extract_cookies_with_source(
     }
 
     if browser != "auto":
+        if browser == "firefox":
+            return _extract_firefox_with_source(domain, cookie_names)
         extractor = extractors.get(browser)
         if extractor is None:
             logger.warning("Unknown browser: %s", browser)
@@ -288,8 +367,13 @@ def extract_cookies_with_source(
         order = ["firefox"]
 
     for name in order:
-        result = extractors[name](domain, cookie_names)
+        if name == "firefox":
-        if result is not None:
+            result = _extract_firefox_with_source(domain, cookie_names)
-            return (result, name)
+            if result is not None:
+                return result
+        else:
+            result = extractors[name](domain, cookie_names)
+            if result is not None:
+                return (result, name)
 
     return None

scripts/lib/dedupe.py

@@ -76,6 +76,31 @@ def hybrid_similarity(text_a: str, text_b: str) -> float:
     )
 
 
+def _tokenize(normalized: str) -> frozenset[str]:
+    return frozenset(
+        tok for tok in normalized.split()
+        if len(tok) > 1 and tok not in STOPWORDS
+    )
+
+
+class _PreparedText:
+    """Pre-computed text representations for fast repeated similarity checks."""
+
+    __slots__ = ("ngrams", "tokens")
+
+    def __init__(self, raw: str) -> None:
+        norm = normalize_text(raw)
+        self.ngrams = get_ngrams(norm) if norm else set()
+        self.tokens = _tokenize(norm)
+
+
+def prepared_similarity(a: _PreparedText, b: _PreparedText) -> float:
+    return max(
+        jaccard_similarity(a.ngrams, b.ngrams),
+        jaccard_similarity(a.tokens, b.tokens),
+    )
+
+
 def item_text(item: schema.SourceItem) -> str:
     parts = [item.title, item.body, item.author or "", item.container or ""]
     return " ".join(part for part in parts if part).strip()
@@ -84,16 +109,19 @@ def item_text(item: schema.SourceItem) -> str:
 def dedupe_items(items: list[schema.SourceItem], threshold: float = 0.7) -> list[schema.SourceItem]:
     """Remove near-duplicates while keeping earlier, better-scored items."""
     kept: list[schema.SourceItem] = []
+    kept_prepared: list[_PreparedText] = []
     for item in items:
         text = item_text(item)
         if not text:
             kept.append(item)
             continue
+        prep = _PreparedText(text)
         is_duplicate = False
-        for existing in kept:
+        for existing_prep in kept_prepared:
-            if hybrid_similarity(text, item_text(existing)) >= threshold:
+            if prepared_similarity(prep, existing_prep) >= threshold:
                 is_duplicate = True
                 break
         if not is_duplicate:
             kept.append(item)
+            kept_prepared.append(prep)
     return kept

scripts/lib/pipeline.py

@@ -719,44 +719,49 @@ def _retry_thin_sources(
         weight=0.3,
     )
 
-    for source in thin_sources:
+    def _retry_one_source(source: str) -> tuple[str, list[schema.SourceItem]]:
-        if source in rate_limited_sources:
+        raw_items, _artifact = _retrieve_stream(
-            continue
+            topic=topic,
-        try:
+            subquery=retry_subquery,
-            raw_items, _artifact = _retrieve_stream(
+            source=source,
-                topic=topic,
+            config=config,
-                subquery=retry_subquery,
+            depth=depth,
-                source=source,
+            date_range=date_range,
-                config=config,
+            runtime=runtime,
-                depth=depth,
+            mock=mock,
-                date_range=date_range,
+            rate_limited_sources=rate_limited_sources,
-                runtime=runtime,
+            rate_limit_lock=rate_limit_lock,
-                mock=mock,
+            web_backend=web_backend,
-                rate_limited_sources=rate_limited_sources,
+            raw_topic=topic,
-                rate_limit_lock=rate_limit_lock,
+        )
-                web_backend=web_backend,
+        normalized = _normalize_score_dedupe(
-                raw_topic=topic,
+            source,
-            )
+            raw_items,
-            normalized = _normalize_score_dedupe(
+            from_date,
-                source,
+            to_date,
-                raw_items,
+            freshness_mode=plan.freshness_mode,
-                from_date,
+            ranking_query=retry_subquery.ranking_query,
-                to_date,
+        )
-                freshness_mode=plan.freshness_mode,
+        return source, normalized[:settings["per_stream_limit"]]
-                ranking_query=retry_subquery.ranking_query,
-            )
-            normalized = normalized[:settings["per_stream_limit"]]
 
-            existing_urls = {item.url for item in bundle.items_by_source.get(source, []) if item.url}
+    retryable = [s for s in thin_sources if s not in rate_limited_sources]
-            new_items = [item for item in normalized if item.url not in existing_urls]
 
-            if new_items:
+    from concurrent.futures import ThreadPoolExecutor, as_completed
-                bundle.items_by_source.setdefault(source, []).extend(new_items)
+    with ThreadPoolExecutor(max_workers=min(4, len(retryable) or 1)) as executor:
-                primary_label = plan.subqueries[0].label if plan.subqueries else "primary"
+        futures = {executor.submit(_retry_one_source, s): s for s in retryable}
-                existing = bundle.items_by_source_and_query.get((primary_label, source), [])
+        for future in as_completed(futures):
-                bundle.items_by_source_and_query[(primary_label, source)] = existing + new_items
+            source = futures[future]
-        except Exception as exc:
+            try:
-            print(f"[Pipeline] Retry failed for {source}: {type(exc).__name__}: {exc}", file=sys.stderr)
+                source, normalized = future.result()
+                existing_urls = {item.url for item in bundle.items_by_source.get(source, []) if item.url}
+                new_items = [item for item in normalized if item.url not in existing_urls]
+
+                if new_items:
+                    bundle.items_by_source.setdefault(source, []).extend(new_items)
+                    primary_label = plan.subqueries[0].label if plan.subqueries else "primary"
+                    bundle.items_by_source_and_query.setdefault((primary_label, source), []).extend(new_items)
+            except Exception as exc:
+                print(f"[Pipeline] Retry failed for {source}: {type(exc).__name__}: {exc}", file=sys.stderr)
 
 
 def _retrieve_stream(

scripts/lib/schema.py

@@ -168,8 +168,7 @@ class RetrievalBundle:
 
     def add_items(self, label: str, source: str, items: list[SourceItem]) -> None:
         """Atomically append items to both items_by_source_and_query and items_by_source."""
-        existing = self.items_by_source_and_query.get((label, source), [])
+        self.items_by_source_and_query.setdefault((label, source), []).extend(items)
-        self.items_by_source_and_query[(label, source)] = existing + items
         self.items_by_source.setdefault(source, []).extend(items)
 
 

scripts/lib/youtube_yt.py

@@ -699,18 +699,27 @@ def enrich_with_comments(
     top_items = ranked[:max_videos]
     _log(f"Enriching comments for {len(top_items)} YouTube videos")
 
-    enriched_count = 0
+    from concurrent.futures import ThreadPoolExecutor, as_completed
-    for item in top_items:
+
+    def _enrich_one(item: dict) -> bool:
         video_id = item.get("video_id", "")
         if not video_id:
-            continue
+            return False
         try:
             comments = _fetch_video_comments(video_id, token, max_comments)
             if comments:
                 item["top_comments"] = comments
-                enriched_count += 1
+                return True
         except Exception as exc:
             _log(f"Comment enrichment failed for {video_id}: {exc}")
+        return False
+
+    enriched_count = 0
+    with ThreadPoolExecutor(max_workers=min(4, len(top_items))) as executor:
+        futures = {executor.submit(_enrich_one, item): item for item in top_items}
+        for future in as_completed(futures):
+            if future.result():
+                enriched_count += 1
 
     _log(f"Enriched {enriched_count}/{len(top_items)} videos with comments")
     return items

tests/test_cookie_extract.py

@@ -153,6 +153,9 @@ class TestExtractFirefoxCookies:
         with patch(
             "scripts.lib.cookie_extract._get_firefox_profiles_dir",
             return_value=None,
+        ), patch(
+            "scripts.lib.cookie_extract._is_wsl",
+            return_value=False,
         ):
             result = extract_firefox_cookies(".x.com", ["auth_token"])
 
@@ -167,6 +170,9 @@ class TestExtractFirefoxCookies:
         with patch(
             "scripts.lib.cookie_extract._get_firefox_profiles_dir",
             return_value=profiles_dir,
+        ), patch(
+            "scripts.lib.cookie_extract._is_wsl",
+            return_value=False,
         ):
             result = extract_firefox_cookies(".x.com", ["auth_token", "ct0"])
 
@@ -185,6 +191,9 @@ class TestExtractFirefoxCookies:
         with patch(
             "scripts.lib.cookie_extract._get_firefox_profiles_dir",
             return_value=profiles_dir,
+        ), patch(
+            "scripts.lib.cookie_extract._is_wsl",
+            return_value=False,
         ):
             result = extract_firefox_cookies(".x.com", ["auth_token", "ct0"])