432 lines
14 KiB
Python
432 lines
14 KiB
Python
|
"""Identity enrichment pipeline.
|
||
|
|
|
||
|
|
Adds bilingual titles and reputation signals with incremental SQLite caching.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
from typing import Callable, Optional
|
||
|
|
from urllib.request import urlopen, Request
|
||
|
|
|
||
|
|
from vlm.cache import EnrichmentCache
|
||
|
|
from vlm.config import Config
|
||
|
|
from vlm.providers import ProviderResult, TMDBProvider, DoubanProvider
|
||
|
|
from vlm.parser import normalize_title
|
||
|
|
|
||
|
|
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)
|
||
|
|
providers = _build_providers(
|
||
|
|
config,
|
||
|
|
request_timeout=request_timeout,
|
||
|
|
retries=retries,
|
||
|
|
)
|
||
|
|
|
||
|
|
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": [],
|
||
|
|
}
|
||
|
|
|
||
|
|
refresh_all = refresh_mode == "refresh_all"
|
||
|
|
|
||
|
|
for section, media_type in (("movies", "movie"), ("series", "series"), ("anime", "anime")):
|
||
|
|
records = identities_data.get(section, [])
|
||
|
|
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
|
||
|
|
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
|
||
|
|
|
||
|
|
payload, api_calls, failures = _enrich_record(
|
||
|
|
record,
|
||
|
|
media_type,
|
||
|
|
providers,
|
||
|
|
config,
|
||
|
|
request_timeout=request_timeout,
|
||
|
|
retries=retries,
|
||
|
|
)
|
||
|
|
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)
|
||
|
|
|
||
|
|
_apply_payload(record, payload)
|
||
|
|
cache.put_identity(identity_key, fingerprint, payload)
|
||
|
|
|
||
|
|
if payload.get("enriched"):
|
||
|
|
stats["enriched"] = int(stats["enriched"]) + 1
|
||
|
|
else:
|
||
|
|
stats["skipped"] = int(stats["skipped"]) + 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)
|
||
|
|
|
||
|
|
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 _build_providers(config: Config, *, request_timeout: int, retries: int) -> list:
|
||
|
|
providers = []
|
||
|
|
for name in config.enrichment_providers:
|
||
|
|
key = name.lower()
|
||
|
|
if key == "tmdb":
|
||
|
|
providers.append(
|
||
|
|
TMDBProvider(
|
||
|
|
config.tmdb_api_key,
|
||
|
|
timeout_seconds=request_timeout,
|
||
|
|
retries=retries,
|
||
|
|
min_interval_seconds=0.25,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
elif key == "douban":
|
||
|
|
providers.append(
|
||
|
|
DoubanProvider(
|
||
|
|
config.douban_api_key,
|
||
|
|
endpoint=config.douban_api_endpoint,
|
||
|
|
timeout_seconds=request_timeout,
|
||
|
|
retries=retries,
|
||
|
|
min_interval_seconds=0.4,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return providers
|
||
|
|
|
||
|
|
|
||
|
|
def _enrich_record(
|
||
|
|
record: dict,
|
||
|
|
media_type: str,
|
||
|
|
providers: list,
|
||
|
|
config: Config,
|
||
|
|
*,
|
||
|
|
request_timeout: int,
|
||
|
|
retries: int,
|
||
|
|
) -> tuple[dict, int, list[dict[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
|
||
|
|
|
||
|
|
for provider in providers:
|
||
|
|
api_calls += 1
|
||
|
|
try:
|
||
|
|
result = provider.enrich(title=title, media_type=media_type, year=year)
|
||
|
|
except Exception as exc:
|
||
|
|
failures.append(
|
||
|
|
{
|
||
|
|
"path": str(record.get("path", "")),
|
||
|
|
"title": str(title),
|
||
|
|
"provider": getattr(provider, "name", "unknown"),
|
||
|
|
"reason": str(exc),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
continue
|
||
|
|
|
||
|
|
if result:
|
||
|
|
provider_results.append(result)
|
||
|
|
|
||
|
|
merged = _merge_provider_results(provider_results)
|
||
|
|
|
||
|
|
# Optional AI fallback for missing translated titles.
|
||
|
|
if config.translation_fallback_machine:
|
||
|
|
if not merged.get("title_zh"):
|
||
|
|
translated = _translate_with_openai(
|
||
|
|
title,
|
||
|
|
target_language="Chinese (Simplified)",
|
||
|
|
api_key=config.openai_api_key,
|
||
|
|
timeout_seconds=request_timeout,
|
||
|
|
retries=retries,
|
||
|
|
)
|
||
|
|
api_calls += 1
|
||
|
|
if translated:
|
||
|
|
merged["title_zh"] = translated
|
||
|
|
merged["translation_source"] = merged.get("translation_source") or "openai"
|
||
|
|
|
||
|
|
if not merged.get("title_en"):
|
||
|
|
translated = _translate_with_openai(
|
||
|
|
title,
|
||
|
|
target_language="English",
|
||
|
|
api_key=config.openai_api_key,
|
||
|
|
timeout_seconds=request_timeout,
|
||
|
|
retries=retries,
|
||
|
|
)
|
||
|
|
api_calls += 1
|
||
|
|
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)
|
||
|
|
|
||
|
|
return merged, api_calls, failures
|
||
|
|
|
||
|
|
|
||
|
|
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 and payload[key] is not None:
|
||
|
|
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:
|
||
|
|
title_zh = payload.get("title_zh") or record.get("title")
|
||
|
|
title_en = payload.get("title_en") or record.get("title")
|
||
|
|
try:
|
||
|
|
formatted = config.naming_title_format.format(title_zh=title_zh, title_en=title_en).strip()
|
||
|
|
except Exception:
|
||
|
|
formatted = f"{title_zh} {title_en}".strip()
|
||
|
|
return " ".join(formatted.split())
|