"""Identity enrichment pipeline. Adds bilingual titles and reputation signals with incremental SQLite caching. """ from __future__ import annotations import hashlib import json from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Callable, Optional from urllib.request import Request, urlopen from vlm.cache import EnrichmentCache from vlm.config import Config from vlm.parser import normalize_title from vlm.providers import ProviderResult, TMDBAuthError, TMDBProvider, TMDBProviderError from vlm.utils import sanitize_path_component RefreshMode = str ProgressCallback = Callable[[int, int, dict[str, int]], None] def enrich_identities_data( identities_data: dict, config: Config, *, refresh_mode: RefreshMode = "incremental", request_timeout: int = 6, retries: int = 2, logger=None, progress_callback: Optional[ProgressCallback] = None, ) -> tuple[dict, dict[str, int | list[dict[str, str]]]]: """Enrich parsed identities in memory and return updated data + stats. refresh_mode: - incremental: default, uses fingerprint cache checks - refresh_changed_only: semantic alias of incremental mode - refresh_all: bypasses cache and re-fetches all records """ cache = EnrichmentCache(config.enrichment_cache_db) total_records = ( len(identities_data.get("movies", [])) + len(identities_data.get("series", [])) + len(identities_data.get("anime", [])) ) stats: dict[str, int | list[dict[str, str]]] = { "total": total_records, "processed": 0, "enriched": 0, "cache_hits": 0, "skipped": 0, "needs_review": 0, "failed": 0, "api_calls": 0, "failed_items": [], "skip_reasons": {}, } refresh_all = refresh_mode == "refresh_all" max_workers = max(1, int(config.enrichment_max_concurrency)) for section, media_type in (("movies", "movie"), ("series", "series"), ("anime", "anime")): records = identities_data.get(section, []) pending_jobs: list[tuple[dict, str, str, str]] = [] for record in records: title = record.get("title") or _fallback_title_from_filename(record.get("filename")) if not title: stats["skipped"] = int(stats["skipped"]) + 1 _increment_skip_reason(stats, "invalid_input") stats["processed"] = int(stats["processed"]) + 1 _emit_progress(stats, progress_callback) continue if "title" not in record: record["title"] = title fingerprint = _fingerprint(record, media_type) identity_key = _identity_key(record, media_type) cached = None if not refresh_all: cached = cache.get_identity(identity_key, fingerprint) if cached: _apply_payload(record, cached) stats["cache_hits"] = int(stats["cache_hits"]) + 1 if record.get("needs_review"): stats["needs_review"] = int(stats["needs_review"]) + 1 stats["processed"] = int(stats["processed"]) + 1 _emit_progress(stats, progress_callback) continue pending_jobs.append((record, identity_key, fingerprint, media_type)) if not pending_jobs: continue if max_workers <= 1: for record, identity_key, fingerprint, pending_media_type in pending_jobs: payload, api_calls, failures, skip_reason = _enrich_record_with_fresh_providers( record, pending_media_type, config, request_timeout=request_timeout, retries=retries, ) _apply_payload(record, payload) cache.put_identity(identity_key, fingerprint, payload) _update_stats_after_enrich( stats, payload, failures, api_calls, record, progress_callback, skip_reason ) else: future_map = {} with ThreadPoolExecutor(max_workers=max_workers) as pool: for record, identity_key, fingerprint, pending_media_type in pending_jobs: future = pool.submit( _enrich_record_with_fresh_providers, record, pending_media_type, config, request_timeout, retries, ) future_map[future] = (record, identity_key, fingerprint) for future in as_completed(future_map): record, identity_key, fingerprint = future_map[future] payload, api_calls, failures, skip_reason = future.result() _apply_payload(record, payload) cache.put_identity(identity_key, fingerprint, payload) _update_stats_after_enrich( stats, payload, failures, api_calls, record, progress_callback, skip_reason ) metadata = identities_data.setdefault("metadata", {}) metadata["enriched"] = True metadata["enrichment_policy"] = "incremental" if refresh_mode != "refresh_all" else "full" metadata["enrichment_refresh_mode"] = refresh_mode if logger: logger.info( "Enrichment completed: total=%s enriched=%s cache_hits=%s skipped=%s failed=%s api_calls=%s", stats["total"], stats["enriched"], stats["cache_hits"], stats["skipped"], stats["failed"], stats["api_calls"], ) return identities_data, stats def _emit_progress(stats: dict[str, int | list[dict[str, str]]], progress_callback: Optional[ProgressCallback]) -> None: if not progress_callback: return progress_callback( int(stats["processed"]), int(stats["total"]), { "cache_hits": int(stats["cache_hits"]), "api_calls": int(stats["api_calls"]), "failed": int(stats["failed"]), }, ) def _update_stats_after_enrich( stats: dict[str, int | list[dict[str, str]]], payload: dict, failures: list[dict[str, str]], api_calls: int, record: dict, progress_callback: Optional[ProgressCallback], skip_reason: str, ) -> None: """Update stats and emit progress after enriching a single record.""" stats["api_calls"] = int(stats["api_calls"]) + api_calls if failures: failed_items = stats["failed_items"] assert isinstance(failed_items, list) failed_items.extend(failures) stats["failed"] = int(stats["failed"]) + len(failures) if payload.get("enriched"): stats["enriched"] = int(stats["enriched"]) + 1 else: stats["skipped"] = int(stats["skipped"]) + 1 _increment_skip_reason(stats, skip_reason or "no_match") if record.get("needs_review"): stats["needs_review"] = int(stats["needs_review"]) + 1 stats["processed"] = int(stats["processed"]) + 1 _emit_progress(stats, progress_callback) def _build_providers(config: Config, *, request_timeout: int, retries: int) -> list: providers = [] unsupported: list[str] = [] for name in config.enrichment_providers: key = name.lower() if key == "tmdb": providers.append( TMDBProvider( config.tmdb_api_key, bearer_token=config.tmdb_bearer_token, language=config.tmdb_language, region=config.tmdb_region, include_adult=config.tmdb_include_adult, timeout_seconds=request_timeout, retries=retries, min_interval_seconds=0.25, ) ) else: unsupported.append(name) if unsupported: raise ValueError( f"Unsupported enrichment providers: {unsupported}. Supported providers: ['tmdb']" ) return providers def _enrich_record_with_fresh_providers( record: dict, media_type: str, config: Config, request_timeout: int, retries: int, ) -> tuple[dict, int, list[dict[str, str]], str]: providers = _build_providers( config, request_timeout=request_timeout, retries=retries, ) return _enrich_record( record, media_type, providers, config, request_timeout=request_timeout, retries=retries, ) def _enrich_record( record: dict, media_type: str, providers: list, config: Config, *, request_timeout: int, retries: int, ) -> tuple[dict, int, list[dict[str, str]], str]: title = record.get("title") year = record.get("year") if media_type == "movie" else None provider_results: list[ProviderResult] = [] failures: list[dict[str, str]] = [] api_calls = 0 configured_provider_count = 0 for provider in providers: if not _provider_is_configured(provider, config): continue configured_provider_count += 1 try: result = provider.enrich(title=title, media_type=media_type, year=year) except TMDBAuthError as exc: raise RuntimeError(str(exc)) from exc except TMDBProviderError as exc: failures.append( { "path": str(record.get("path", "")), "title": str(title), "provider": provider.name, "reason": str(exc), } ) api_calls += provider.last_request_count continue except Exception as exc: failures.append( { "path": str(record.get("path", "")), "title": str(title), "provider": provider.name, "reason": str(exc), } ) api_calls += provider.last_request_count continue api_calls += provider.last_request_count if result: provider_results.append(result) merged = _merge_provider_results(provider_results) # Optional AI fallback for missing translated titles. if config.translation_fallback_machine and config.openai_api_key: if not merged.get("title_zh"): api_calls += 1 translated = _translate_with_openai( title, target_language="Chinese (Simplified)", api_key=config.openai_api_key, timeout_seconds=request_timeout, retries=retries, ) if translated: merged["title_zh"] = translated merged["translation_source"] = merged.get("translation_source") or "openai" if not merged.get("title_en"): api_calls += 1 translated = _translate_with_openai( title, target_language="English", api_key=config.openai_api_key, timeout_seconds=request_timeout, retries=retries, ) if translated: merged["title_en"] = translated merged["translation_source"] = merged.get("translation_source") or "openai" confidence = _enrichment_confidence(merged) merged["enrichment_confidence"] = confidence review_status = record.get("review_status", "pending") needs_review = bool(record.get("needs_review", False)) if confidence < config.enrichment_min_match_score: needs_review = True score = merged.get("reputation_score") votes = merged.get("reputation_votes") or 0 if ( score is not None and votes >= config.reputation_min_votes and score < config.reputation_low_score_threshold ): needs_review = True merged["review_status"] = review_status merged["needs_review"] = needs_review merged["enriched"] = bool(provider_results or merged.get("translation_source")) merged["display_title"] = _build_display_title(record, merged, config) skip_reason = _determine_skip_reason( provider_results=provider_results, failures=failures, configured_provider_count=configured_provider_count, api_calls=api_calls, ) return merged, api_calls, failures, skip_reason def _determine_skip_reason( *, provider_results: list[ProviderResult], failures: list[dict[str, str]], configured_provider_count: int, api_calls: int, ) -> str: if provider_results: return "" if configured_provider_count == 0 and api_calls == 0: return "no_key" if failures: for failure in failures: reason = str(failure.get("reason", "")).lower() if "rate limit" in reason or "(429)" in reason: return "rate_limited" if "authentication failed" in reason or "(401/403)" in reason: return "auth_error" return "provider_error" return "no_match" def _increment_skip_reason(stats: dict[str, int | list[dict[str, str]]], reason: str) -> None: if not reason: return current = stats.get("skip_reasons") if not isinstance(current, dict): current = {} stats["skip_reasons"] = current current[reason] = int(current.get(reason, 0)) + 1 def _provider_is_configured(provider: object, config: Config) -> bool: provider_name = provider.name.lower() if provider_name == "tmdb": return bool(config.tmdb_bearer_token or config.tmdb_api_key) return True def _merge_provider_results(results: list[ProviderResult]) -> dict: payload: dict = { "canonical_id": None, "title_zh": None, "title_en": None, "translation_source": None, "reputation_score": None, "reputation_votes": None, "reputation_source": None, "provider_metadata": {}, } if not results: return payload first = results[0] payload["canonical_id"] = first.canonical_id payload["title_zh"] = first.title_zh payload["title_en"] = first.title_en payload["translation_source"] = first.translation_source total_weight = 0 weighted_score = 0.0 source_names = [] for result in results: source_names.append(result.provider) payload["provider_metadata"][result.provider] = json.dumps(result.raw_metadata, ensure_ascii=False) if not payload["title_zh"] and result.title_zh: payload["title_zh"] = result.title_zh payload["translation_source"] = result.translation_source or result.provider if not payload["title_en"] and result.title_en: payload["title_en"] = result.title_en if result.reputation_score is None: continue votes = result.reputation_votes if result.reputation_votes and result.reputation_votes > 0 else 1 weighted_score += result.reputation_score * votes total_weight += votes if total_weight > 0: payload["reputation_score"] = round(weighted_score / total_weight, 3) payload["reputation_votes"] = total_weight payload["reputation_source"] = "+".join(sorted(set(source_names))) return payload def _enrichment_confidence(payload: dict) -> float: score = 0.0 if payload.get("canonical_id"): score += 0.4 if payload.get("title_zh"): score += 0.2 if payload.get("title_en"): score += 0.2 if payload.get("reputation_score") is not None: score += 0.2 return round(score, 3) def _translate_with_openai( text: str, *, target_language: str, api_key: Optional[str], timeout_seconds: int, retries: int, ) -> Optional[str]: if not api_key: return None body = { "model": "gpt-4o-mini", "input": ( f"Translate the movie or TV title into {target_language}. " "Return only the translated title without explanations." f"\nTitle: {text}" ), } request = Request( "https://api.openai.com/v1/responses", data=json.dumps(body).encode("utf-8"), headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", }, method="POST", ) for _ in range(max(retries + 1, 1)): try: with urlopen(request, timeout=timeout_seconds) as response: payload = json.loads(response.read().decode("utf-8")) output_text = payload.get("output_text") if isinstance(output_text, str) and output_text.strip(): return output_text.strip() except Exception: continue return None def _apply_payload(record: dict, payload: dict) -> None: for key in ( "canonical_id", "title_zh", "title_en", "translation_source", "reputation_score", "reputation_votes", "reputation_source", "review_status", "enrichment_confidence", "provider_metadata", "display_title", ): if key in payload: record[key] = payload[key] if "needs_review" in payload: record["needs_review"] = payload["needs_review"] def _identity_key(record: dict, media_type: str) -> str: return f"{media_type}:{record.get('path', '')}" def _fingerprint(record: dict, media_type: str) -> str: fields = [ media_type, str(record.get("path", "")), str(record.get("title", "")), str(record.get("year", "")), str(record.get("season", "")), json.dumps(record.get("episodes", [])), ] digest = hashlib.sha256("|".join(fields).encode("utf-8")).hexdigest() return digest def _fallback_title_from_filename(filename: Optional[str]) -> Optional[str]: if not filename: return None base = filename.rsplit(".", 1)[0] return normalize_title(base) def _build_display_title(record: dict, payload: dict, config: Config) -> str: fallback_title = record.get("title") title_zh = payload.get("title_zh") or fallback_title title_en = payload.get("title_en") or fallback_title if title_zh and title_en and title_zh == title_en: title_en = "" try: formatted = config.naming_title_format.format( title_zh=title_zh or "", title_en=title_en or "", ).strip() except Exception: formatted = f"{title_zh or ''} {title_en or ''}".strip() return sanitize_path_component(" ".join(formatted.split()), fallback=fallback_title or "untitled")