6f0df5a774
Stop tracking personal workflow artifacts at repo root, add CI and MIT license, align README and agent skills with artifacts/ defaults, and enable Ruff in dev/CI so releases are verifiable without local-only runs. Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com>
573 lines
18 KiB
Python
573 lines
18 KiB
Python
"""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")
|