Merge pull request #198 from iliaal/perf/pipeline-optimizations
perf: optimize dedup, parallelize handle searches and enrichment
This commit is contained in:
+14
-7
@@ -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
@@ -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:
|
||||
|
||||
@@ -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,18 +217,21 @@ 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)
|
||||
|
||||
|
||||
def extract_chrome_cookies(
|
||||
domain: str, cookie_names: List[str]
|
||||
@@ -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,6 +367,11 @@ def extract_cookies_with_source(
|
||||
order = ["firefox"]
|
||||
|
||||
for name in order:
|
||||
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)
|
||||
|
||||
+30
-2
@@ -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
|
||||
|
||||
+12
-7
@@ -719,10 +719,7 @@ def _retry_thin_sources(
|
||||
weight=0.3,
|
||||
)
|
||||
|
||||
for source in thin_sources:
|
||||
if source in rate_limited_sources:
|
||||
continue
|
||||
try:
|
||||
def _retry_one_source(source: str) -> tuple[str, list[schema.SourceItem]]:
|
||||
raw_items, _artifact = _retrieve_stream(
|
||||
topic=topic,
|
||||
subquery=retry_subquery,
|
||||
@@ -745,16 +742,24 @@ def _retry_thin_sources(
|
||||
freshness_mode=plan.freshness_mode,
|
||||
ranking_query=retry_subquery.ranking_query,
|
||||
)
|
||||
normalized = normalized[:settings["per_stream_limit"]]
|
||||
return source, normalized[:settings["per_stream_limit"]]
|
||||
|
||||
retryable = [s for s in thin_sources if s not in rate_limited_sources]
|
||||
|
||||
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"
|
||||
existing = bundle.items_by_source_and_query.get((primary_label, source), [])
|
||||
bundle.items_by_source_and_query[(primary_label, source)] = existing + new_items
|
||||
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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"])
|
||||
|
||||
|
||||
Reference in New Issue
Block a user