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

perf: optimize dedup, parallelize handle searches and enrichment
This commit is contained in:
Matt Van Horn
2026-04-09 21:00:37 -07:00
committed by GitHub
8 changed files with 233 additions and 76 deletions
+14 -7
View File
@@ -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)
all_items.extend(items)
return parse_bird_response(response, query=core_topic)
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
+23 -7
View File
@@ -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))
for existing in selected
dedupe.prepared_similarity(prep, sp) for sp in selected_preps
)
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:
+102 -18
View File
@@ -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
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)
# First pass: Install* section (Firefox >= 67 format, takes priority)
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
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)
return None
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 result is not None:
return (result, name)
if name == "firefox":
result = _extract_firefox_with_source(domain, cookie_names)
if result is not None:
return result
else:
result = extractors[name](domain, cookie_names)
if result is not None:
return (result, name)
return None
+30 -2
View File
@@ -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:
if hybrid_similarity(text, item_text(existing)) >= threshold:
for existing_prep in kept_prepared:
if prepared_similarity(prep, existing_prep) >= threshold:
is_duplicate = True
break
if not is_duplicate:
kept.append(item)
kept_prepared.append(prep)
return kept
+41 -36
View File
@@ -719,44 +719,49 @@ def _retry_thin_sources(
weight=0.3,
)
for source in thin_sources:
if source in rate_limited_sources:
continue
try:
raw_items, _artifact = _retrieve_stream(
topic=topic,
subquery=retry_subquery,
source=source,
config=config,
depth=depth,
date_range=date_range,
runtime=runtime,
mock=mock,
rate_limited_sources=rate_limited_sources,
rate_limit_lock=rate_limit_lock,
web_backend=web_backend,
raw_topic=topic,
)
normalized = _normalize_score_dedupe(
source,
raw_items,
from_date,
to_date,
freshness_mode=plan.freshness_mode,
ranking_query=retry_subquery.ranking_query,
)
normalized = normalized[:settings["per_stream_limit"]]
def _retry_one_source(source: str) -> tuple[str, list[schema.SourceItem]]:
raw_items, _artifact = _retrieve_stream(
topic=topic,
subquery=retry_subquery,
source=source,
config=config,
depth=depth,
date_range=date_range,
runtime=runtime,
mock=mock,
rate_limited_sources=rate_limited_sources,
rate_limit_lock=rate_limit_lock,
web_backend=web_backend,
raw_topic=topic,
)
normalized = _normalize_score_dedupe(
source,
raw_items,
from_date,
to_date,
freshness_mode=plan.freshness_mode,
ranking_query=retry_subquery.ranking_query,
)
return source, normalized[:settings["per_stream_limit"]]
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]
retryable = [s for s in thin_sources if s not in rate_limited_sources]
if new_items:
bundle.items_by_source.setdefault(source, []).extend(new_items)
primary_label = plan.subqueries[0].label if plan.subqueries else "primary"
existing = bundle.items_by_source_and_query.get((primary_label, source), [])
bundle.items_by_source_and_query[(primary_label, source)] = existing + new_items
except Exception as exc:
print(f"[Pipeline] Retry failed for {source}: {type(exc).__name__}: {exc}", file=sys.stderr)
from concurrent.futures import ThreadPoolExecutor, as_completed
with ThreadPoolExecutor(max_workers=min(4, len(retryable) or 1)) as executor:
futures = {executor.submit(_retry_one_source, s): s for s in retryable}
for future in as_completed(futures):
source = futures[future]
try:
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(
+1 -2
View File
@@ -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[(label, source)] = existing + items
self.items_by_source_and_query.setdefault((label, source), []).extend(items)
self.items_by_source.setdefault(source, []).extend(items)
+13 -4
View File
@@ -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
for item in top_items:
from concurrent.futures import ThreadPoolExecutor, as_completed
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
+9
View File
@@ -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"])