add incremental enrich controls with progress and retry limits
This commit is contained in:
@@ -9,6 +9,8 @@ A Python-based CLI tool for managing personal video collections with a safety-fi
|
||||
- **Comprehensive Analysis**: Detect episode gaps and duplicate files
|
||||
- **Rich Metadata**: Extract video resolution, codec, duration, and bitrate
|
||||
- **Flexible Organization**: Customizable directory structure and naming templates
|
||||
- **Metadata Enrichment**: Add bilingual titles and reputation signals (TMDB + optional Douban + optional AI fallback)
|
||||
- **Incremental Performance**: SQLite-backed cache avoids repeated metadata lookups
|
||||
- **State Tracking**: Track file status throughout the workflow
|
||||
- **Detailed Reporting**: Generate inventory, completeness, and duplicate reports
|
||||
|
||||
@@ -72,7 +74,26 @@ This creates `identities.json` with parsed information.
|
||||
- **Series**: Title, season, and episode numbers (e.g., "Breaking Bad S01E01")
|
||||
- **Confidence scores**: Indicates parsing reliability
|
||||
|
||||
### 4. Analyze Your Library
|
||||
### 4. Enrich Titles and Reputation (Optional but Recommended)
|
||||
|
||||
Add translation and reputation metadata to `identities.json`:
|
||||
|
||||
```bash
|
||||
vlm enrich
|
||||
```
|
||||
|
||||
This updates `identities.json` in place and adds fields like:
|
||||
- `title_zh`, `title_en`, `display_title`
|
||||
- `reputation_score`, `reputation_votes`, `reputation_source`
|
||||
- `review_status`, `enrichment_confidence`
|
||||
|
||||
To refresh all records instead of using incremental cache:
|
||||
|
||||
```bash
|
||||
vlm enrich --refresh-all
|
||||
```
|
||||
|
||||
### 5. Analyze Your Library
|
||||
|
||||
Detect episode gaps and duplicates:
|
||||
|
||||
@@ -84,7 +105,7 @@ This creates `analysis.json` with:
|
||||
- Series with missing episodes
|
||||
- Duplicate files with quality comparison
|
||||
|
||||
### 5. Generate Execution Plan
|
||||
### 6. Generate Execution Plan
|
||||
|
||||
Create a reviewable plan of file operations:
|
||||
|
||||
@@ -96,7 +117,7 @@ This creates `plan.json` with proposed operations (move, rename, quarantine).
|
||||
|
||||
**Review the plan** by opening `plan.json` in your editor. You can edit it if needed.
|
||||
|
||||
### 6. Execute (Dry-Run First)
|
||||
### 7. Execute (Dry-Run First)
|
||||
|
||||
Preview what will happen without making changes:
|
||||
|
||||
@@ -112,7 +133,7 @@ vlm execute --confirm
|
||||
|
||||
**Important**: This creates a rollback log in `~/.vlm/rollback/` for reverting changes.
|
||||
|
||||
### 7. Rollback (If Needed)
|
||||
### 8. Rollback (If Needed)
|
||||
|
||||
If you need to undo the operations:
|
||||
|
||||
@@ -139,27 +160,31 @@ vlm scan
|
||||
vlm parse
|
||||
# Output: identities.json with parsed titles and episodes
|
||||
|
||||
# 4. Analyze for gaps and duplicates
|
||||
# 4. Enrich identities (translation + reputation)
|
||||
vlm enrich
|
||||
# Output: identities.json updated in place (incremental cache enabled)
|
||||
|
||||
# 5. Analyze for gaps and duplicates
|
||||
vlm analyze
|
||||
# Output: analysis.json with 5 series with gaps, 12 duplicate groups
|
||||
|
||||
# 5. Generate execution plan
|
||||
# 6. Generate execution plan
|
||||
vlm plan
|
||||
# Output: plan.json with 456 operations proposed
|
||||
|
||||
# 6. Review the plan
|
||||
# 7. Review the plan
|
||||
cat plan.json | less
|
||||
# or open in your editor
|
||||
|
||||
# 7. Dry-run to preview
|
||||
# 8. Dry-run to preview
|
||||
vlm execute
|
||||
# Shows what would happen without making changes
|
||||
|
||||
# 8. Execute with confirmation
|
||||
# 9. Execute with confirmation
|
||||
vlm execute --confirm
|
||||
# Actually performs the file operations
|
||||
|
||||
# 9. If needed, rollback
|
||||
# 10. If needed, rollback
|
||||
vlm rollback
|
||||
```
|
||||
|
||||
@@ -204,6 +229,25 @@ vlm parse
|
||||
vlm parse --input my_inventory.csv --output my_identities.json
|
||||
```
|
||||
|
||||
### Enrichment
|
||||
|
||||
```bash
|
||||
# Enrich identities in place (default: identities.json)
|
||||
vlm enrich
|
||||
|
||||
# Enrich custom file and write to another file
|
||||
vlm enrich --input my_identities.json --output enriched_identities.json
|
||||
|
||||
# Refresh changed records only (explicit incremental mode)
|
||||
vlm enrich --refresh-changed-only
|
||||
|
||||
# Force full refresh (ignore cache for all records)
|
||||
vlm enrich --refresh-all
|
||||
|
||||
# Tune request behavior
|
||||
vlm enrich --timeout 6 --retries 2
|
||||
```
|
||||
|
||||
### Analysis
|
||||
|
||||
```bash
|
||||
@@ -337,6 +381,30 @@ categories:
|
||||
movie: [movie, movies, films]
|
||||
series: [series, tv, shows]
|
||||
anime: [anime]
|
||||
|
||||
# Enrichment settings
|
||||
enrichment:
|
||||
enabled: true
|
||||
incremental: true
|
||||
refresh_mode: "manual"
|
||||
providers: [tmdb, douban]
|
||||
cache_db: "~/.vlm/enrichment_cache.db"
|
||||
max_concurrency: 6
|
||||
min_match_score: 0.75
|
||||
douban_endpoint: null
|
||||
translation:
|
||||
mode: "bidirectional"
|
||||
fallback_machine: true
|
||||
api_keys:
|
||||
tmdb: null
|
||||
douban: null
|
||||
openai: null
|
||||
reputation:
|
||||
min_votes: 50
|
||||
low_score_threshold: 6.0
|
||||
policy: "flag_for_review"
|
||||
naming:
|
||||
title_format: "{title_zh} {title_en}"
|
||||
```
|
||||
|
||||
### Template Variables
|
||||
@@ -414,6 +482,7 @@ vlm report summary
|
||||
# 1. Scan and parse
|
||||
vlm scan
|
||||
vlm parse
|
||||
vlm enrich
|
||||
|
||||
# 2. Analyze completeness
|
||||
vlm analyze
|
||||
@@ -428,6 +497,7 @@ vlm report completeness
|
||||
# 1. Scan and parse
|
||||
vlm scan
|
||||
vlm parse
|
||||
vlm enrich
|
||||
|
||||
# 2. Analyze for duplicates
|
||||
vlm analyze
|
||||
@@ -452,6 +522,7 @@ vlm execute
|
||||
# 2. Scan and parse
|
||||
vlm scan
|
||||
vlm parse
|
||||
vlm enrich
|
||||
|
||||
# 3. Generate plan
|
||||
vlm plan
|
||||
@@ -580,6 +651,9 @@ src/vlm/
|
||||
├── cli.py # Click-based CLI interface
|
||||
├── scanner.py # File discovery and metadata extraction
|
||||
├── parser.py # Filename parsing (titles, years, episodes)
|
||||
├── enrichment.py # Title/reputation enrichment pipeline
|
||||
├── cache.py # SQLite cache for incremental enrichment
|
||||
├── providers/ # External metadata providers (TMDB/Douban)
|
||||
├── analysis.py # Completeness and duplicate detection
|
||||
├── planner.py # Execution plan generation
|
||||
├── executor.py # File operations and rollback
|
||||
@@ -593,6 +667,7 @@ src/vlm/
|
||||
tests/
|
||||
├── test_scanner.py # Scanner tests
|
||||
├── test_parser.py # Parser tests
|
||||
├── test_enrichment.py # Enrichment tests
|
||||
├── test_analysis.py # Analysis tests
|
||||
├── test_planner.py # Planner tests
|
||||
└── ... # More test files
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
"""SQLite cache utilities for enrichment."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class EnrichmentCache:
|
||||
"""Simple SQLite-backed cache for incremental enrichment."""
|
||||
|
||||
def __init__(self, db_path: Path) -> None:
|
||||
self.db_path = db_path
|
||||
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._ensure_schema()
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
return sqlite3.connect(self.db_path)
|
||||
|
||||
def _ensure_schema(self) -> None:
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS identity_enrichment (
|
||||
identity_key TEXT PRIMARY KEY,
|
||||
fingerprint TEXT NOT NULL,
|
||||
payload_json TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
def get_identity(self, identity_key: str, fingerprint: str) -> Optional[dict]:
|
||||
with self._connect() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT payload_json FROM identity_enrichment WHERE identity_key = ? AND fingerprint = ?",
|
||||
(identity_key, fingerprint),
|
||||
).fetchone()
|
||||
|
||||
if not row:
|
||||
return None
|
||||
return json.loads(row[0])
|
||||
|
||||
def put_identity(self, identity_key: str, fingerprint: str, payload: dict) -> None:
|
||||
payload_json = json.dumps(payload, ensure_ascii=False)
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO identity_enrichment(identity_key, fingerprint, payload_json)
|
||||
VALUES (?, ?, ?)
|
||||
ON CONFLICT(identity_key) DO UPDATE SET
|
||||
fingerprint=excluded.fingerprint,
|
||||
payload_json=excluded.payload_json,
|
||||
updated_at=CURRENT_TIMESTAMP
|
||||
""",
|
||||
(identity_key, fingerprint, payload_json),
|
||||
)
|
||||
+289
-9
@@ -116,8 +116,31 @@ def main(ctx, config: Path, log_level: Optional[str]):
|
||||
default=Path('inventory.csv'),
|
||||
help='Output file for inventory (default: inventory.csv)'
|
||||
)
|
||||
@click.option(
|
||||
'--metadata/--no-metadata',
|
||||
default=True,
|
||||
help='Extract video metadata via ffprobe (default: enabled)'
|
||||
)
|
||||
@click.option(
|
||||
'--reuse-from',
|
||||
type=click.Path(exists=True, path_type=Path),
|
||||
default=None,
|
||||
help='Reuse metadata cache from an existing inventory CSV (default: output file if it exists)'
|
||||
)
|
||||
@click.option(
|
||||
'--force-refresh-metadata',
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help='Ignore metadata cache and re-run ffprobe for all files'
|
||||
)
|
||||
@pass_context
|
||||
def scan(ctx: CLIContext, output: Path):
|
||||
def scan(
|
||||
ctx: CLIContext,
|
||||
output: Path,
|
||||
metadata: bool,
|
||||
reuse_from: Optional[Path],
|
||||
force_refresh_metadata: bool
|
||||
):
|
||||
"""Scan library and generate inventory.
|
||||
|
||||
Discovers all video files in the library and records their metadata.
|
||||
@@ -127,8 +150,11 @@ def scan(ctx: CLIContext, output: Path):
|
||||
|
||||
vlm scan # Save to inventory.csv
|
||||
vlm scan --output my_library.csv # Save to custom file
|
||||
vlm scan --no-metadata # Faster scan without ffprobe
|
||||
vlm scan --reuse-from old.csv # Reuse prior metadata cache
|
||||
vlm scan --force-refresh-metadata # Re-run ffprobe for all files
|
||||
"""
|
||||
from vlm.scanner import scan_library, save_inventory_csv
|
||||
from vlm.scanner import scan_library, save_inventory_csv, load_inventory_csv
|
||||
|
||||
config = ctx.config
|
||||
logger = ctx.logger
|
||||
@@ -139,8 +165,59 @@ def scan(ctx: CLIContext, output: Path):
|
||||
click.echo("This may take a while for large libraries...")
|
||||
click.echo()
|
||||
|
||||
progress_state: dict[str, Optional[click.ProgressBar]] = {"bar": None}
|
||||
progress_position = {"current": 0}
|
||||
|
||||
def _scan_progress(processed: int, total: int) -> None:
|
||||
if total <= 0:
|
||||
return
|
||||
if progress_state["bar"] is None:
|
||||
bar = click.progressbar(
|
||||
length=total,
|
||||
label="Scanning files",
|
||||
show_pos=True,
|
||||
)
|
||||
progress_state["bar"] = bar.__enter__()
|
||||
|
||||
step = processed - progress_position["current"]
|
||||
if step > 0 and progress_state["bar"] is not None:
|
||||
progress_state["bar"].update(step)
|
||||
progress_position["current"] = processed
|
||||
|
||||
metadata_cache = None
|
||||
cache_source = None
|
||||
if not force_refresh_metadata:
|
||||
cache_source = reuse_from if reuse_from is not None else (output if output.exists() else None)
|
||||
|
||||
if metadata and force_refresh_metadata:
|
||||
click.echo("Forcing metadata refresh for all files (cache disabled).")
|
||||
click.echo()
|
||||
|
||||
if metadata and cache_source is not None:
|
||||
click.echo(f"Loading metadata cache from: {cache_source}")
|
||||
try:
|
||||
cached_files = load_inventory_csv(cache_source)
|
||||
metadata_cache = {str(vf.path): vf for vf in cached_files}
|
||||
click.echo(f"Loaded metadata cache entries: {len(metadata_cache)}")
|
||||
click.echo()
|
||||
except Exception as e:
|
||||
click.echo(f"Warning: could not load metadata cache: {e}")
|
||||
click.echo("Continuing without cache.")
|
||||
click.echo()
|
||||
|
||||
# Perform the scan
|
||||
video_files = scan_library(config.library_root, config)
|
||||
try:
|
||||
video_files = scan_library(
|
||||
config.library_root,
|
||||
config,
|
||||
progress_callback=_scan_progress,
|
||||
include_video_metadata=metadata,
|
||||
metadata_cache=metadata_cache
|
||||
)
|
||||
finally:
|
||||
if progress_state["bar"] is not None:
|
||||
progress_state["bar"].__exit__(None, None, None)
|
||||
click.echo()
|
||||
|
||||
# Display summary
|
||||
click.echo(f"Scan complete!")
|
||||
@@ -369,6 +446,187 @@ def parse(ctx: CLIContext, input: Path, output: Path):
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
@main.command()
|
||||
@click.option(
|
||||
'--input',
|
||||
type=click.Path(exists=True, path_type=Path),
|
||||
default=Path('identities.json'),
|
||||
help='Path to identities JSON file (default: identities.json)'
|
||||
)
|
||||
@click.option(
|
||||
'--output',
|
||||
type=click.Path(path_type=Path),
|
||||
default=None,
|
||||
help='Output path (default: overwrite input file)'
|
||||
)
|
||||
@click.option(
|
||||
'--refresh-changed-only',
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help='Refresh only records whose identity fingerprint changed (incremental behavior)'
|
||||
)
|
||||
@click.option(
|
||||
'--refresh-all',
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help='Bypass cache and re-enrich all records'
|
||||
)
|
||||
@click.option(
|
||||
'--timeout',
|
||||
type=int,
|
||||
default=6,
|
||||
show_default=True,
|
||||
help='Per-request timeout in seconds'
|
||||
)
|
||||
@click.option(
|
||||
'--retries',
|
||||
type=int,
|
||||
default=2,
|
||||
show_default=True,
|
||||
help='Retry attempts for provider/API requests'
|
||||
)
|
||||
@pass_context
|
||||
def enrich(
|
||||
ctx: CLIContext,
|
||||
input: Path,
|
||||
output: Optional[Path],
|
||||
refresh_changed_only: bool,
|
||||
refresh_all: bool,
|
||||
timeout: int,
|
||||
retries: int,
|
||||
):
|
||||
"""Enrich identities with translation and reputation metadata.
|
||||
|
||||
Applies incremental cache-backed enrichment to parsed identities and writes
|
||||
results back into identities JSON.
|
||||
"""
|
||||
import json
|
||||
from vlm.enrichment import enrich_identities_data
|
||||
|
||||
config = ctx.config
|
||||
logger = ctx.logger
|
||||
|
||||
if output is None:
|
||||
output = input
|
||||
|
||||
try:
|
||||
click.echo(f"Enriching identities from: {input}")
|
||||
click.echo(f"Output file: {output}")
|
||||
click.echo()
|
||||
|
||||
with open(input, 'r', encoding='utf-8') as jsonfile:
|
||||
identities_data = json.load(jsonfile)
|
||||
|
||||
if refresh_all and refresh_changed_only:
|
||||
click.echo("Error: --refresh-all and --refresh-changed-only are mutually exclusive.", err=True)
|
||||
sys.exit(1)
|
||||
|
||||
if timeout < 1:
|
||||
click.echo("Error: --timeout must be >= 1", err=True)
|
||||
sys.exit(1)
|
||||
|
||||
if retries < 0:
|
||||
click.echo("Error: --retries must be >= 0", err=True)
|
||||
sys.exit(1)
|
||||
|
||||
refresh_mode = "incremental"
|
||||
if refresh_all:
|
||||
refresh_mode = "refresh_all"
|
||||
elif refresh_changed_only:
|
||||
refresh_mode = "refresh_changed_only"
|
||||
|
||||
total = (
|
||||
len(identities_data.get('movies', []))
|
||||
+ len(identities_data.get('series', []))
|
||||
+ len(identities_data.get('anime', []))
|
||||
)
|
||||
|
||||
progress_state: dict[str, Optional[click.ProgressBar]] = {"bar": None}
|
||||
progress_position = {"current": 0}
|
||||
|
||||
def _enrich_progress(processed: int, total_count: int, _metrics: dict[str, int]) -> None:
|
||||
if total_count <= 0:
|
||||
return
|
||||
|
||||
if progress_state["bar"] is None:
|
||||
bar = click.progressbar(
|
||||
length=total_count,
|
||||
label="Enriching records",
|
||||
show_pos=True,
|
||||
)
|
||||
progress_state["bar"] = bar.__enter__()
|
||||
|
||||
step = processed - progress_position["current"]
|
||||
if step > 0 and progress_state["bar"] is not None:
|
||||
progress_state["bar"].update(step)
|
||||
progress_position["current"] = processed
|
||||
|
||||
try:
|
||||
if total == 0:
|
||||
click.echo("No movie/series/anime records found to enrich.")
|
||||
enriched_data, stats = enrich_identities_data(
|
||||
identities_data,
|
||||
config,
|
||||
refresh_mode=refresh_mode,
|
||||
request_timeout=timeout,
|
||||
retries=retries,
|
||||
logger=logger,
|
||||
progress_callback=_enrich_progress,
|
||||
)
|
||||
finally:
|
||||
if progress_state["bar"] is not None:
|
||||
progress_state["bar"].__exit__(None, None, None)
|
||||
click.echo()
|
||||
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(output, 'w', encoding='utf-8') as jsonfile:
|
||||
json.dump(enriched_data, jsonfile, indent=2, ensure_ascii=False)
|
||||
|
||||
total_records = int(stats['total']) if stats['total'] else 0
|
||||
cache_hits = int(stats['cache_hits'])
|
||||
hit_rate = (cache_hits / total_records * 100.0) if total_records else 0.0
|
||||
|
||||
click.echo("Enrichment complete!")
|
||||
click.echo(f" Total records: {total_records}")
|
||||
click.echo(f" Refresh mode: {refresh_mode}")
|
||||
click.echo(f" Enriched now: {stats['enriched']}")
|
||||
click.echo(f" Cache hits: {stats['cache_hits']}")
|
||||
click.echo(f" Cache hit rate: {hit_rate:.1f}%")
|
||||
click.echo(f" API calls: {stats['api_calls']}")
|
||||
click.echo(f" Failed requests: {stats['failed']}")
|
||||
click.echo(f" Skipped: {stats['skipped']}")
|
||||
click.echo(f" Needs review: {stats['needs_review']}")
|
||||
failed_items = stats.get('failed_items', [])
|
||||
if isinstance(failed_items, list) and failed_items:
|
||||
click.echo(" Failure sample:")
|
||||
for item in failed_items[:3]:
|
||||
click.echo(
|
||||
f" - [{item.get('provider', 'unknown')}] {item.get('title', '')}: {item.get('reason', '')}"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Enrich completed: total=%s enriched=%s cache_hits=%s failed=%s skipped=%s mode=%s",
|
||||
stats['total'],
|
||||
stats['enriched'],
|
||||
stats['cache_hits'],
|
||||
stats['failed'],
|
||||
stats['skipped'],
|
||||
refresh_mode,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
click.echo(f"Error: Input file not found: {input}", err=True)
|
||||
logger.error(f"Input file not found: {input}")
|
||||
sys.exit(1)
|
||||
except json.JSONDecodeError as e:
|
||||
click.echo(f"Error: Failed to parse JSON file: {e}", err=True)
|
||||
logger.error(f"JSON parsing failed during enrich: {e}", exc_info=True)
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
click.echo(f"Error during enrichment: {e}", err=True)
|
||||
logger.error(f"Enrich failed: {e}", exc_info=True)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
@main.command()
|
||||
@click.option(
|
||||
'--input',
|
||||
@@ -632,6 +890,7 @@ def plan(ctx: CLIContext, input: Path, output: Path):
|
||||
|
||||
# Process movies
|
||||
for m in movies_data:
|
||||
is_approved = m.get('review_status') == 'approved'
|
||||
video_file = VideoFile(
|
||||
path=Path(m['path']),
|
||||
filename=m['filename'],
|
||||
@@ -645,17 +904,28 @@ def plan(ctx: CLIContext, input: Path, output: Path):
|
||||
)
|
||||
|
||||
movie_identity = MovieIdentity(
|
||||
title=m['title'],
|
||||
title=m.get('display_title', m['title']),
|
||||
year=m.get('year'),
|
||||
confidence=m['confidence'],
|
||||
needs_review=m['needs_review'],
|
||||
original_filename=m['filename']
|
||||
needs_review=(m['needs_review'] and not is_approved),
|
||||
original_filename=m['filename'],
|
||||
canonical_id=m.get('canonical_id'),
|
||||
title_zh=m.get('title_zh'),
|
||||
title_en=m.get('title_en'),
|
||||
translation_source=m.get('translation_source'),
|
||||
reputation_score=m.get('reputation_score'),
|
||||
reputation_votes=m.get('reputation_votes'),
|
||||
reputation_source=m.get('reputation_source'),
|
||||
review_status=m.get('review_status', 'pending'),
|
||||
enrichment_confidence=m.get('enrichment_confidence'),
|
||||
provider_metadata=m.get('provider_metadata', {})
|
||||
)
|
||||
|
||||
identities_list.append((video_file, movie_identity))
|
||||
|
||||
# Process series
|
||||
for s in series_data:
|
||||
is_approved = s.get('review_status') == 'approved'
|
||||
video_file = VideoFile(
|
||||
path=Path(s['path']),
|
||||
filename=s['filename'],
|
||||
@@ -669,12 +939,22 @@ def plan(ctx: CLIContext, input: Path, output: Path):
|
||||
)
|
||||
|
||||
series_identity = SeriesIdentity(
|
||||
title=s['title'],
|
||||
title=s.get('display_title', s['title']),
|
||||
season=s.get('season'),
|
||||
episodes=s.get('episodes', []),
|
||||
confidence=s['confidence'],
|
||||
needs_review=s['needs_review'],
|
||||
original_filename=s['filename']
|
||||
needs_review=(s['needs_review'] and not is_approved),
|
||||
original_filename=s['filename'],
|
||||
canonical_id=s.get('canonical_id'),
|
||||
title_zh=s.get('title_zh'),
|
||||
title_en=s.get('title_en'),
|
||||
translation_source=s.get('translation_source'),
|
||||
reputation_score=s.get('reputation_score'),
|
||||
reputation_votes=s.get('reputation_votes'),
|
||||
reputation_source=s.get('reputation_source'),
|
||||
review_status=s.get('review_status', 'pending'),
|
||||
enrichment_confidence=s.get('enrichment_confidence'),
|
||||
provider_metadata=s.get('provider_metadata', {})
|
||||
)
|
||||
|
||||
identities_list.append((video_file, series_identity))
|
||||
|
||||
+138
-118
@@ -8,19 +8,8 @@ import yaml
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
"""Configuration for Video Library Manager.
|
||||
|
||||
Attributes:
|
||||
library_root: Root directory of the video library
|
||||
video_extensions: List of video file extensions to recognize
|
||||
movie_template: Directory template for movies (e.g., "movie/{title} ({year})/")
|
||||
series_template: Directory template for series (e.g., "series/{title}/Season {season:02d}/")
|
||||
movie_filename_template: Filename template for movies (e.g., "{title} ({year}){ext}")
|
||||
series_filename_template: Filename template for series (e.g., "S{season:02d}E{episode:02d}{ext}")
|
||||
log_level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
|
||||
quarantine_dir: Quarantine directory name relative to category root (e.g., ".quarantine")
|
||||
categories: Mapping of category names to directory name lists for file categorization
|
||||
"""
|
||||
"""Configuration for Video Library Manager."""
|
||||
|
||||
library_root: Path
|
||||
video_extensions: list[str] = field(default_factory=lambda: [
|
||||
".mp4", ".mkv", ".avi", ".mov", ".wmv", ".flv", ".webm", ".m4v"
|
||||
@@ -32,67 +21,75 @@ class Config:
|
||||
log_level: str = "INFO"
|
||||
quarantine_dir: str = ".quarantine"
|
||||
categories: dict[str, list[str]] = field(default_factory=lambda: {
|
||||
"movie": ["movie"],
|
||||
"series": ["series"],
|
||||
"movie": ["movie", "movies"],
|
||||
"series": ["series", "tv", "shows"],
|
||||
"anime": ["anime"]
|
||||
})
|
||||
|
||||
# Enrichment settings
|
||||
enrichment_enabled: bool = True
|
||||
enrichment_incremental: bool = True
|
||||
enrichment_refresh_mode: str = "manual"
|
||||
enrichment_providers: list[str] = field(default_factory=lambda: ["tmdb", "douban"])
|
||||
enrichment_cache_db: Path = field(default_factory=lambda: Path.home() / ".vlm" / "enrichment_cache.db")
|
||||
enrichment_max_concurrency: int = 6
|
||||
enrichment_min_match_score: float = 0.75
|
||||
translation_mode: str = "bidirectional"
|
||||
translation_fallback_machine: bool = True
|
||||
tmdb_api_key: Optional[str] = None
|
||||
douban_api_key: Optional[str] = None
|
||||
douban_api_endpoint: Optional[str] = None
|
||||
openai_api_key: Optional[str] = None
|
||||
reputation_min_votes: int = 50
|
||||
reputation_low_score_threshold: float = 6.0
|
||||
reputation_policy: str = "flag_for_review"
|
||||
naming_title_format: str = "{title_zh} {title_en}"
|
||||
|
||||
|
||||
def load_config(path: Path) -> Config:
|
||||
"""Load configuration from YAML file.
|
||||
|
||||
Args:
|
||||
path: Path to configuration file
|
||||
|
||||
Returns:
|
||||
Config object with loaded settings
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If config file doesn't exist (caller should handle by creating default)
|
||||
yaml.YAMLError: If YAML syntax is invalid (caller should handle by using defaults)
|
||||
"""
|
||||
"""Load configuration from YAML file."""
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Configuration file not found: {path}")
|
||||
|
||||
|
||||
try:
|
||||
with open(path, 'r', encoding='utf-8') as f:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
data = yaml.safe_load(f)
|
||||
except yaml.YAMLError as e:
|
||||
raise yaml.YAMLError(f"Invalid YAML syntax in configuration file: {e}")
|
||||
|
||||
|
||||
if data is None:
|
||||
data = {}
|
||||
|
||||
# Extract library_root (required field)
|
||||
library_root_str = data.get('library_root')
|
||||
|
||||
library_root_str = data.get("library_root")
|
||||
if not library_root_str:
|
||||
raise ValueError("Configuration must specify 'library_root'")
|
||||
|
||||
|
||||
library_root = Path(library_root_str).expanduser()
|
||||
|
||||
# Extract optional fields with defaults
|
||||
video_extensions = data.get('video_extensions', [
|
||||
|
||||
video_extensions = data.get("video_extensions", [
|
||||
".mp4", ".mkv", ".avi", ".mov", ".wmv", ".flv", ".webm", ".m4v"
|
||||
])
|
||||
|
||||
# Extract templates
|
||||
templates = data.get('templates', {})
|
||||
movie_template = templates.get('movie_dir', "movie/{title} ({year})/")
|
||||
series_template = templates.get('series_dir', "series/{title}/Season {season:02d}/")
|
||||
movie_filename_template = templates.get('movie_filename', "{title} ({year}){ext}")
|
||||
series_filename_template = templates.get('series_filename', "S{season:02d}E{episode:02d}{ext}")
|
||||
|
||||
# Extract other settings
|
||||
quarantine_dir = data.get('quarantine_dir', '.quarantine')
|
||||
log_level = data.get('log_level', 'INFO')
|
||||
|
||||
# Extract categories configuration
|
||||
categories = data.get('categories', {
|
||||
"movie": ["movie"],
|
||||
"series": ["series"],
|
||||
templates = data.get("templates", {})
|
||||
movie_template = templates.get("movie_dir", "movie/{title} ({year})/")
|
||||
series_template = templates.get("series_dir", "series/{title}/Season {season:02d}/")
|
||||
movie_filename_template = templates.get("movie_filename", "{title} ({year}){ext}")
|
||||
series_filename_template = templates.get("series_filename", "S{season:02d}E{episode:02d}{ext}")
|
||||
|
||||
quarantine_dir = data.get("quarantine_dir", ".quarantine")
|
||||
log_level = data.get("log_level", "INFO")
|
||||
categories = data.get("categories", {
|
||||
"movie": ["movie", "movies"],
|
||||
"series": ["series", "tv", "shows"],
|
||||
"anime": ["anime"]
|
||||
})
|
||||
|
||||
enrichment = data.get("enrichment", {})
|
||||
translation = enrichment.get("translation", {})
|
||||
api_keys = enrichment.get("api_keys", {})
|
||||
reputation = enrichment.get("reputation", {})
|
||||
naming = enrichment.get("naming", {})
|
||||
|
||||
return Config(
|
||||
library_root=library_root,
|
||||
video_extensions=video_extensions,
|
||||
@@ -102,79 +99,94 @@ def load_config(path: Path) -> Config:
|
||||
series_filename_template=series_filename_template,
|
||||
log_level=log_level,
|
||||
quarantine_dir=quarantine_dir,
|
||||
categories=categories
|
||||
categories=categories,
|
||||
enrichment_enabled=enrichment.get("enabled", True),
|
||||
enrichment_incremental=enrichment.get("incremental", True),
|
||||
enrichment_refresh_mode=enrichment.get("refresh_mode", "manual"),
|
||||
enrichment_providers=enrichment.get("providers", ["tmdb", "douban"]),
|
||||
enrichment_cache_db=Path(
|
||||
enrichment.get("cache_db", str(Path.home() / ".vlm" / "enrichment_cache.db"))
|
||||
).expanduser(),
|
||||
enrichment_max_concurrency=enrichment.get("max_concurrency", 6),
|
||||
enrichment_min_match_score=enrichment.get("min_match_score", 0.75),
|
||||
translation_mode=translation.get("mode", "bidirectional"),
|
||||
translation_fallback_machine=translation.get("fallback_machine", True),
|
||||
tmdb_api_key=api_keys.get("tmdb"),
|
||||
douban_api_key=api_keys.get("douban"),
|
||||
douban_api_endpoint=enrichment.get("douban_endpoint"),
|
||||
openai_api_key=api_keys.get("openai"),
|
||||
reputation_min_votes=reputation.get("min_votes", 50),
|
||||
reputation_low_score_threshold=reputation.get("low_score_threshold", 6.0),
|
||||
reputation_policy=reputation.get("policy", "flag_for_review"),
|
||||
naming_title_format=naming.get("title_format", "{title_zh} {title_en}"),
|
||||
)
|
||||
|
||||
|
||||
def create_default_config(path: Path) -> Config:
|
||||
"""Create a default configuration file and return the Config object.
|
||||
|
||||
Args:
|
||||
path: Path where configuration file should be created
|
||||
|
||||
Returns:
|
||||
Config object with default settings
|
||||
"""
|
||||
# Create default config object
|
||||
"""Create a default configuration file and return the Config object."""
|
||||
default_config = Config(
|
||||
library_root=Path.home() / "Videos",
|
||||
video_extensions=[".mp4", ".mkv", ".avi", ".mov", ".wmv", ".flv", ".webm", ".m4v"],
|
||||
movie_template="movie/{title} ({year})/",
|
||||
series_template="series/{title}/Season {season:02d}/",
|
||||
movie_filename_template="{title} ({year}){ext}",
|
||||
series_filename_template="S{season:02d}E{episode:02d}{ext}",
|
||||
log_level="INFO",
|
||||
quarantine_dir=".quarantine",
|
||||
categories={
|
||||
"movie": ["movie"],
|
||||
"series": ["series"],
|
||||
"anime": ["anime"]
|
||||
}
|
||||
)
|
||||
|
||||
# Create YAML content
|
||||
|
||||
yaml_content = {
|
||||
'library_root': str(default_config.library_root),
|
||||
'video_extensions': default_config.video_extensions,
|
||||
'templates': {
|
||||
'movie_dir': default_config.movie_template,
|
||||
'series_dir': default_config.series_template,
|
||||
'movie_filename': default_config.movie_filename_template,
|
||||
'series_filename': default_config.series_filename_template
|
||||
"library_root": str(default_config.library_root),
|
||||
"video_extensions": default_config.video_extensions,
|
||||
"templates": {
|
||||
"movie_dir": default_config.movie_template,
|
||||
"series_dir": default_config.series_template,
|
||||
"movie_filename": default_config.movie_filename_template,
|
||||
"series_filename": default_config.series_filename_template,
|
||||
},
|
||||
"quarantine_dir": default_config.quarantine_dir,
|
||||
"log_level": default_config.log_level,
|
||||
"categories": default_config.categories,
|
||||
"enrichment": {
|
||||
"enabled": default_config.enrichment_enabled,
|
||||
"incremental": default_config.enrichment_incremental,
|
||||
"refresh_mode": default_config.enrichment_refresh_mode,
|
||||
"providers": default_config.enrichment_providers,
|
||||
"cache_db": str(default_config.enrichment_cache_db),
|
||||
"max_concurrency": default_config.enrichment_max_concurrency,
|
||||
"min_match_score": default_config.enrichment_min_match_score,
|
||||
"translation": {
|
||||
"mode": default_config.translation_mode,
|
||||
"fallback_machine": default_config.translation_fallback_machine,
|
||||
},
|
||||
"api_keys": {
|
||||
"tmdb": default_config.tmdb_api_key,
|
||||
"douban": default_config.douban_api_key,
|
||||
"openai": default_config.openai_api_key,
|
||||
},
|
||||
"douban_endpoint": default_config.douban_api_endpoint,
|
||||
"reputation": {
|
||||
"min_votes": default_config.reputation_min_votes,
|
||||
"low_score_threshold": default_config.reputation_low_score_threshold,
|
||||
"policy": default_config.reputation_policy,
|
||||
},
|
||||
"naming": {
|
||||
"title_format": default_config.naming_title_format,
|
||||
},
|
||||
},
|
||||
'quarantine_dir': default_config.quarantine_dir,
|
||||
'log_level': default_config.log_level,
|
||||
'categories': default_config.categories
|
||||
}
|
||||
|
||||
# Ensure parent directory exists
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Write configuration file
|
||||
with open(path, 'w', encoding='utf-8') as f:
|
||||
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
yaml.dump(yaml_content, f, default_flow_style=False, sort_keys=False)
|
||||
|
||||
|
||||
return default_config
|
||||
|
||||
|
||||
def validate_config(config: Config) -> list[str]:
|
||||
"""Validate configuration and return list of error messages.
|
||||
|
||||
Args:
|
||||
config: Configuration object to validate
|
||||
|
||||
Returns:
|
||||
List of error messages (empty if valid)
|
||||
"""
|
||||
"""Validate configuration and return list of error messages."""
|
||||
errors = []
|
||||
|
||||
# Validate library_root
|
||||
|
||||
if not isinstance(config.library_root, Path):
|
||||
errors.append("library_root must be a Path object")
|
||||
elif not str(config.library_root) or str(config.library_root) == ".":
|
||||
errors.append("library_root cannot be empty")
|
||||
|
||||
# Validate video_extensions
|
||||
|
||||
if not config.video_extensions:
|
||||
errors.append("video_extensions cannot be empty")
|
||||
elif not isinstance(config.video_extensions, list):
|
||||
@@ -184,31 +196,29 @@ def validate_config(config: Config) -> list[str]:
|
||||
if not isinstance(ext, str):
|
||||
errors.append(f"video_extensions must contain strings, found: {type(ext)}")
|
||||
break
|
||||
if not ext.startswith('.'):
|
||||
if not ext.startswith("."):
|
||||
errors.append(f"video extension must start with '.': {ext}")
|
||||
|
||||
# Validate templates
|
||||
|
||||
if not config.movie_template:
|
||||
errors.append("movie_template cannot be empty")
|
||||
elif not isinstance(config.movie_template, str):
|
||||
errors.append("movie_template must be a string")
|
||||
|
||||
|
||||
if not config.series_template:
|
||||
errors.append("series_template cannot be empty")
|
||||
elif not isinstance(config.series_template, str):
|
||||
errors.append("series_template must be a string")
|
||||
|
||||
|
||||
if not config.movie_filename_template:
|
||||
errors.append("movie_filename_template cannot be empty")
|
||||
elif not isinstance(config.movie_filename_template, str):
|
||||
errors.append("movie_filename_template must be a string")
|
||||
|
||||
|
||||
if not config.series_filename_template:
|
||||
errors.append("series_filename_template cannot be empty")
|
||||
elif not isinstance(config.series_filename_template, str):
|
||||
errors.append("series_filename_template must be a string")
|
||||
|
||||
# Validate log_level
|
||||
|
||||
valid_log_levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]
|
||||
if not config.log_level:
|
||||
errors.append("log_level cannot be empty")
|
||||
@@ -216,28 +226,24 @@ def validate_config(config: Config) -> list[str]:
|
||||
errors.append("log_level must be a string")
|
||||
elif config.log_level.upper() not in valid_log_levels:
|
||||
errors.append(f"log_level must be one of {valid_log_levels}, got: {config.log_level}")
|
||||
|
||||
# Validate quarantine_dir
|
||||
|
||||
if not config.quarantine_dir:
|
||||
errors.append("quarantine_dir cannot be empty")
|
||||
elif not isinstance(config.quarantine_dir, str):
|
||||
errors.append("quarantine_dir must be a string")
|
||||
elif config.quarantine_dir.startswith('/') or config.quarantine_dir.startswith('\\'):
|
||||
elif config.quarantine_dir.startswith("/") or config.quarantine_dir.startswith("\\"):
|
||||
errors.append("quarantine_dir must be relative to category root, not absolute")
|
||||
|
||||
# Validate categories
|
||||
if not config.categories:
|
||||
errors.append("categories cannot be empty")
|
||||
elif not isinstance(config.categories, dict):
|
||||
errors.append("categories must be a dictionary")
|
||||
else:
|
||||
# Check required category keys exist
|
||||
required_categories = {"movie", "series", "anime"}
|
||||
missing = required_categories - set(config.categories.keys())
|
||||
if missing:
|
||||
errors.append(f"categories must include keys: {sorted(missing)}")
|
||||
|
||||
# Validate each category's directory list and check for duplicates
|
||||
seen_dirs = {}
|
||||
for category, dir_list in config.categories.items():
|
||||
if not isinstance(dir_list, list):
|
||||
@@ -257,7 +263,6 @@ def validate_config(config: Config) -> list[str]:
|
||||
errors.append(f"categories['{category}'] contains empty directory name")
|
||||
break
|
||||
|
||||
# Check for duplicates (case-insensitive)
|
||||
dir_lower = dir_name.lower()
|
||||
if dir_lower in seen_dirs:
|
||||
errors.append(
|
||||
@@ -267,4 +272,19 @@ def validate_config(config: Config) -> list[str]:
|
||||
else:
|
||||
seen_dirs[dir_lower] = category
|
||||
|
||||
if not isinstance(config.enrichment_cache_db, Path):
|
||||
errors.append("enrichment_cache_db must be a Path object")
|
||||
if not isinstance(config.enrichment_providers, list) or not config.enrichment_providers:
|
||||
errors.append("enrichment_providers must be a non-empty list")
|
||||
if config.enrichment_max_concurrency < 1:
|
||||
errors.append("enrichment_max_concurrency must be >= 1")
|
||||
if not (0.0 <= config.enrichment_min_match_score <= 1.0):
|
||||
errors.append("enrichment_min_match_score must be between 0.0 and 1.0")
|
||||
if config.enrichment_refresh_mode not in {"manual"}:
|
||||
errors.append("enrichment_refresh_mode must be 'manual'")
|
||||
if config.reputation_min_votes < 0:
|
||||
errors.append("reputation_min_votes must be >= 0")
|
||||
if not (0.0 <= config.reputation_low_score_threshold <= 10.0):
|
||||
errors.append("reputation_low_score_threshold must be between 0.0 and 10.0")
|
||||
|
||||
return errors
|
||||
|
||||
@@ -0,0 +1,431 @@
|
||||
"""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())
|
||||
+21
-1
@@ -4,7 +4,7 @@ This module defines the core data structures used throughout the application
|
||||
for representing video files and their parsed identities.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
@@ -54,6 +54,16 @@ class MovieIdentity:
|
||||
confidence: float
|
||||
needs_review: bool
|
||||
original_filename: str
|
||||
canonical_id: Optional[str] = None
|
||||
title_zh: Optional[str] = None
|
||||
title_en: Optional[str] = None
|
||||
translation_source: Optional[str] = None
|
||||
reputation_score: Optional[float] = None
|
||||
reputation_votes: Optional[int] = None
|
||||
reputation_source: Optional[str] = None
|
||||
review_status: str = "pending"
|
||||
enrichment_confidence: Optional[float] = None
|
||||
provider_metadata: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -74,6 +84,16 @@ class SeriesIdentity:
|
||||
confidence: float
|
||||
needs_review: bool
|
||||
original_filename: str
|
||||
canonical_id: Optional[str] = None
|
||||
title_zh: Optional[str] = None
|
||||
title_en: Optional[str] = None
|
||||
translation_source: Optional[str] = None
|
||||
reputation_score: Optional[float] = None
|
||||
reputation_votes: Optional[int] = None
|
||||
reputation_source: Optional[str] = None
|
||||
review_status: str = "pending"
|
||||
enrichment_confidence: Optional[float] = None
|
||||
provider_metadata: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
+24
-2
@@ -136,7 +136,18 @@ def _create_movie_operation(
|
||||
Returns:
|
||||
FileOperation for organizing the movie
|
||||
"""
|
||||
# If movie needs review (no year), generate no-op
|
||||
# Explicitly blocked by manual review workflow.
|
||||
if identity.review_status == "rejected":
|
||||
return FileOperation(
|
||||
operation_type="no-op",
|
||||
source_path=video_file.path,
|
||||
destination_path=None,
|
||||
reason="Movie rejected during manual review",
|
||||
has_conflict=False,
|
||||
conflict_reason=None
|
||||
)
|
||||
|
||||
# If movie needs review (no year or low-confidence enrichment), generate no-op
|
||||
if identity.needs_review or identity.year is None:
|
||||
return FileOperation(
|
||||
operation_type="no-op",
|
||||
@@ -214,7 +225,18 @@ def _create_series_operation(
|
||||
Returns:
|
||||
FileOperation for organizing the series episode
|
||||
"""
|
||||
# If series needs review (no season or no episodes), generate no-op (v1 constraint)
|
||||
# Explicitly blocked by manual review workflow.
|
||||
if identity.review_status == "rejected":
|
||||
return FileOperation(
|
||||
operation_type="no-op",
|
||||
source_path=video_file.path,
|
||||
destination_path=None,
|
||||
reason="Series rejected during manual review",
|
||||
has_conflict=False,
|
||||
conflict_reason=None
|
||||
)
|
||||
|
||||
# If series needs review (no season or no episodes), generate no-op
|
||||
if identity.needs_review or identity.season is None or len(identity.episodes) == 0:
|
||||
return FileOperation(
|
||||
operation_type="no-op",
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
"""Provider implementations for enrichment."""
|
||||
|
||||
from vlm.providers.base import EnrichmentProvider, ProviderResult
|
||||
from vlm.providers.tmdb import TMDBProvider
|
||||
from vlm.providers.douban import DoubanProvider
|
||||
|
||||
__all__ = [
|
||||
"EnrichmentProvider",
|
||||
"ProviderResult",
|
||||
"TMDBProvider",
|
||||
"DoubanProvider",
|
||||
]
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Provider interfaces for enrichment sources."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Protocol, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProviderResult:
|
||||
"""Normalized provider output used by enrichment pipeline."""
|
||||
|
||||
provider: str
|
||||
canonical_id: Optional[str] = None
|
||||
title_zh: Optional[str] = None
|
||||
title_en: Optional[str] = None
|
||||
translation_source: Optional[str] = None
|
||||
reputation_score: Optional[float] = None
|
||||
reputation_votes: Optional[int] = None
|
||||
reputation_source: Optional[str] = None
|
||||
match_score: Optional[float] = None
|
||||
raw_metadata: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
|
||||
class EnrichmentProvider(Protocol):
|
||||
"""Protocol for title/score providers."""
|
||||
|
||||
name: str
|
||||
|
||||
def enrich(self, *, title: str, media_type: str, year: Optional[int] = None) -> Optional[ProviderResult]:
|
||||
"""Return normalized metadata for a single identity."""
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Douban provider implementation.
|
||||
|
||||
This provider is optional. If no endpoint is configured, it silently degrades.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import Optional
|
||||
from urllib.parse import urlencode
|
||||
from urllib.request import urlopen, Request
|
||||
|
||||
from vlm.providers.base import ProviderResult
|
||||
|
||||
|
||||
class DoubanProvider:
|
||||
"""Fetch reputation data from a configurable Douban-compatible API."""
|
||||
|
||||
name = "douban"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: Optional[str],
|
||||
endpoint: Optional[str] = None,
|
||||
timeout_seconds: int = 6,
|
||||
retries: int = 2,
|
||||
min_interval_seconds: float = 0.4,
|
||||
) -> None:
|
||||
self.api_key = api_key
|
||||
self.endpoint = endpoint
|
||||
self.timeout_seconds = timeout_seconds
|
||||
self.retries = retries
|
||||
self.min_interval_seconds = min_interval_seconds
|
||||
self._last_request_at = 0.0
|
||||
|
||||
def enrich(self, *, title: str, media_type: str, year: Optional[int] = None) -> Optional[ProviderResult]:
|
||||
if not self.endpoint:
|
||||
return None
|
||||
|
||||
params = {
|
||||
"q": title,
|
||||
"type": media_type,
|
||||
}
|
||||
if year is not None:
|
||||
params["year"] = year
|
||||
if self.api_key:
|
||||
params["api_key"] = self.api_key
|
||||
|
||||
data = self._get_json(self.endpoint, params)
|
||||
if not data:
|
||||
return None
|
||||
|
||||
items = data.get("items") or data.get("subjects") or []
|
||||
if not items:
|
||||
return None
|
||||
|
||||
item = items[0]
|
||||
score = item.get("rating") or item.get("score")
|
||||
votes = item.get("vote_count") or item.get("ratings_count")
|
||||
title_zh = item.get("title")
|
||||
title_en = item.get("original_title")
|
||||
|
||||
return ProviderResult(
|
||||
provider=self.name,
|
||||
canonical_id=f"douban:{item.get('id', 'unknown')}",
|
||||
title_zh=title_zh,
|
||||
title_en=title_en,
|
||||
translation_source=self.name if title_zh or title_en else None,
|
||||
reputation_score=float(score) if score is not None else None,
|
||||
reputation_votes=int(votes) if votes is not None else None,
|
||||
reputation_source=self.name,
|
||||
raw_metadata={"id": str(item.get("id", ""))},
|
||||
)
|
||||
|
||||
def _wait_for_rate_limit(self) -> None:
|
||||
if self.min_interval_seconds <= 0:
|
||||
return
|
||||
now = time.monotonic()
|
||||
elapsed = now - self._last_request_at
|
||||
if elapsed < self.min_interval_seconds:
|
||||
time.sleep(self.min_interval_seconds - elapsed)
|
||||
|
||||
def _get_json(self, url: str, params: dict) -> Optional[dict]:
|
||||
full_url = f"{url}?{urlencode(params)}"
|
||||
request = Request(full_url, headers={"Accept": "application/json"})
|
||||
for _ in range(max(self.retries + 1, 1)):
|
||||
self._wait_for_rate_limit()
|
||||
try:
|
||||
with urlopen(request, timeout=self.timeout_seconds) as response:
|
||||
payload = response.read().decode("utf-8")
|
||||
self._last_request_at = time.monotonic()
|
||||
return json.loads(payload)
|
||||
except Exception:
|
||||
self._last_request_at = time.monotonic()
|
||||
continue
|
||||
return None
|
||||
@@ -0,0 +1,108 @@
|
||||
"""TMDB provider implementation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import Optional
|
||||
from urllib.parse import urlencode
|
||||
from urllib.request import urlopen, Request
|
||||
|
||||
from vlm.providers.base import ProviderResult
|
||||
|
||||
|
||||
class TMDBProvider:
|
||||
"""Fetch translations and reputation data from TMDB."""
|
||||
|
||||
name = "tmdb"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: Optional[str],
|
||||
language: str = "zh-CN",
|
||||
timeout_seconds: int = 6,
|
||||
retries: int = 2,
|
||||
min_interval_seconds: float = 0.25,
|
||||
) -> None:
|
||||
self.api_key = api_key
|
||||
self.language = language
|
||||
self.base_url = "https://api.themoviedb.org/3"
|
||||
self.timeout_seconds = timeout_seconds
|
||||
self.retries = retries
|
||||
self.min_interval_seconds = min_interval_seconds
|
||||
self._last_request_at = 0.0
|
||||
|
||||
def enrich(self, *, title: str, media_type: str, year: Optional[int] = None) -> Optional[ProviderResult]:
|
||||
if not self.api_key:
|
||||
return None
|
||||
|
||||
search_type = "tv" if media_type in {"series", "anime", "tv"} else "movie"
|
||||
query_params = {
|
||||
"api_key": self.api_key,
|
||||
"query": title,
|
||||
"language": self.language,
|
||||
}
|
||||
if year and search_type == "movie":
|
||||
query_params["year"] = year
|
||||
|
||||
search_data = self._get_json(f"{self.base_url}/search/{search_type}", query_params)
|
||||
if not search_data:
|
||||
return None
|
||||
|
||||
results = search_data.get("results", [])
|
||||
if not results:
|
||||
return None
|
||||
|
||||
candidate = results[0]
|
||||
tmdb_id = candidate.get("id")
|
||||
if tmdb_id is None:
|
||||
return None
|
||||
|
||||
details = self._get_json(
|
||||
f"{self.base_url}/{search_type}/{tmdb_id}",
|
||||
{"api_key": self.api_key, "language": self.language},
|
||||
)
|
||||
if not details:
|
||||
details = candidate
|
||||
|
||||
original_title = details.get("original_title") or details.get("original_name")
|
||||
localized_title = details.get("title") or details.get("name")
|
||||
vote_average = details.get("vote_average")
|
||||
vote_count = details.get("vote_count")
|
||||
|
||||
return ProviderResult(
|
||||
provider=self.name,
|
||||
canonical_id=f"tmdb:{search_type}:{tmdb_id}",
|
||||
title_zh=localized_title,
|
||||
title_en=original_title,
|
||||
translation_source=self.name,
|
||||
reputation_score=float(vote_average) if vote_average is not None else None,
|
||||
reputation_votes=int(vote_count) if vote_count is not None else None,
|
||||
reputation_source=self.name,
|
||||
match_score=float(candidate.get("popularity", 0.0)) if candidate.get("popularity") is not None else None,
|
||||
raw_metadata={"media_type": search_type, "id": str(tmdb_id)},
|
||||
)
|
||||
|
||||
def _wait_for_rate_limit(self) -> None:
|
||||
if self.min_interval_seconds <= 0:
|
||||
return
|
||||
now = time.monotonic()
|
||||
elapsed = now - self._last_request_at
|
||||
if elapsed < self.min_interval_seconds:
|
||||
time.sleep(self.min_interval_seconds - elapsed)
|
||||
|
||||
def _get_json(self, url: str, params: dict) -> Optional[dict]:
|
||||
full_url = f"{url}?{urlencode(params)}"
|
||||
request = Request(full_url, headers={"Accept": "application/json"})
|
||||
|
||||
for _ in range(max(self.retries + 1, 1)):
|
||||
self._wait_for_rate_limit()
|
||||
try:
|
||||
with urlopen(request, timeout=self.timeout_seconds) as response:
|
||||
payload = response.read().decode("utf-8")
|
||||
self._last_request_at = time.monotonic()
|
||||
return json.loads(payload)
|
||||
except Exception:
|
||||
self._last_request_at = time.monotonic()
|
||||
continue
|
||||
return None
|
||||
@@ -0,0 +1,136 @@
|
||||
"""Tests for CLI enrich command."""
|
||||
|
||||
import json
|
||||
|
||||
from click.testing import CliRunner
|
||||
|
||||
from vlm.cli import main
|
||||
|
||||
|
||||
def test_enrich_in_place_updates_identities(tmp_path):
|
||||
"""`vlm enrich` should update identities file in place by default."""
|
||||
library_root = tmp_path / "library"
|
||||
library_root.mkdir(parents=True)
|
||||
|
||||
config_file = tmp_path / "config.yaml"
|
||||
config_file.write_text(
|
||||
f"""
|
||||
library_root: {library_root}
|
||||
categories:
|
||||
movie: [movie, movies]
|
||||
series: [series, tv, shows]
|
||||
anime: [anime]
|
||||
enrichment:
|
||||
cache_db: {tmp_path / 'cache.db'}
|
||||
""".strip()
|
||||
)
|
||||
|
||||
identities_file = tmp_path / "identities.json"
|
||||
identities_file.write_text(json.dumps({
|
||||
"metadata": {},
|
||||
"movies": [
|
||||
{
|
||||
"path": "/library/movie/Test.2024.mkv",
|
||||
"filename": "Test.2024.mkv",
|
||||
"category": "movie",
|
||||
"title": "Test",
|
||||
"year": 2024,
|
||||
"confidence": 0.9,
|
||||
"needs_review": False,
|
||||
}
|
||||
],
|
||||
"series": [],
|
||||
"anime": [],
|
||||
"other": [],
|
||||
}, ensure_ascii=False))
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(main, ["--config", str(config_file), "enrich", "--input", str(identities_file)])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Enrichment complete!" in result.output
|
||||
|
||||
updated = json.loads(identities_file.read_text(encoding="utf-8"))
|
||||
assert updated["metadata"]["enriched"] is True
|
||||
|
||||
|
||||
def test_enrich_refresh_all_option_runs_successfully(tmp_path):
|
||||
"""`vlm enrich --refresh-all` should execute successfully."""
|
||||
library_root = tmp_path / "library"
|
||||
library_root.mkdir(parents=True)
|
||||
|
||||
config_file = tmp_path / "config.yaml"
|
||||
config_file.write_text(
|
||||
f"""
|
||||
library_root: {library_root}
|
||||
categories:
|
||||
movie: [movie, movies]
|
||||
series: [series, tv, shows]
|
||||
anime: [anime]
|
||||
enrichment:
|
||||
cache_db: {tmp_path / 'cache.db'}
|
||||
""".strip()
|
||||
)
|
||||
|
||||
identities_file = tmp_path / "identities.json"
|
||||
identities_file.write_text(json.dumps({
|
||||
"metadata": {},
|
||||
"movies": [],
|
||||
"series": [],
|
||||
"anime": [],
|
||||
"other": [],
|
||||
}))
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["--config", str(config_file), "enrich", "--input", str(identities_file), "--refresh-all"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Enrichment complete!" in result.output
|
||||
|
||||
|
||||
def test_enrich_rejects_conflicting_refresh_flags(tmp_path):
|
||||
"""Conflicting refresh flags should fail with a clear error."""
|
||||
library_root = tmp_path / "library"
|
||||
library_root.mkdir(parents=True)
|
||||
|
||||
config_file = tmp_path / "config.yaml"
|
||||
config_file.write_text(
|
||||
f"""
|
||||
library_root: {library_root}
|
||||
categories:
|
||||
movie: [movie, movies]
|
||||
series: [series, tv, shows]
|
||||
anime: [anime]
|
||||
enrichment:
|
||||
cache_db: {tmp_path / 'cache.db'}
|
||||
""".strip()
|
||||
)
|
||||
|
||||
identities_file = tmp_path / "identities.json"
|
||||
identities_file.write_text(json.dumps({
|
||||
"metadata": {},
|
||||
"movies": [],
|
||||
"series": [],
|
||||
"anime": [],
|
||||
"other": [],
|
||||
}))
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"--config",
|
||||
str(config_file),
|
||||
"enrich",
|
||||
"--input",
|
||||
str(identities_file),
|
||||
"--refresh-all",
|
||||
"--refresh-changed-only",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "mutually exclusive" in result.output
|
||||
@@ -40,8 +40,8 @@ class TestConfig:
|
||||
config = Config(library_root=Path("/test"))
|
||||
|
||||
assert config.categories == {
|
||||
"movie": ["movie"],
|
||||
"series": ["series"],
|
||||
"movie": ["movie", "movies"],
|
||||
"series": ["series", "tv", "shows"],
|
||||
"anime": ["anime"]
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
"""Unit tests for enrichment pipeline."""
|
||||
|
||||
from vlm.config import Config
|
||||
from vlm.enrichment import enrich_identities_data
|
||||
from vlm.providers.base import ProviderResult
|
||||
|
||||
|
||||
class DummyProvider:
|
||||
name = "dummy"
|
||||
|
||||
def __init__(self):
|
||||
self.calls = 0
|
||||
|
||||
def enrich(self, *, title: str, media_type: str, year=None):
|
||||
self.calls += 1
|
||||
return ProviderResult(
|
||||
provider="dummy",
|
||||
canonical_id=f"dummy:{title}",
|
||||
title_zh="测试中文名",
|
||||
title_en="Test English Title",
|
||||
translation_source="dummy",
|
||||
reputation_score=8.2,
|
||||
reputation_votes=100,
|
||||
reputation_source="dummy",
|
||||
)
|
||||
|
||||
|
||||
def test_enrich_incremental_cache_hit(tmp_path, monkeypatch):
|
||||
"""Second enrichment run should hit cache when fingerprint is unchanged."""
|
||||
config = Config(
|
||||
library_root=tmp_path,
|
||||
enrichment_cache_db=tmp_path / "cache.db",
|
||||
enrichment_providers=["dummy"],
|
||||
translation_fallback_machine=False,
|
||||
)
|
||||
|
||||
provider = DummyProvider()
|
||||
monkeypatch.setattr(
|
||||
"vlm.enrichment._build_providers",
|
||||
lambda _config, request_timeout, retries: [provider],
|
||||
)
|
||||
|
||||
identities = {
|
||||
"metadata": {},
|
||||
"movies": [
|
||||
{
|
||||
"path": "/library/movie/The.Matrix.1999.mkv",
|
||||
"filename": "The.Matrix.1999.mkv",
|
||||
"category": "movie",
|
||||
"title": "The Matrix",
|
||||
"year": 1999,
|
||||
"confidence": 0.9,
|
||||
"needs_review": False,
|
||||
}
|
||||
],
|
||||
"series": [],
|
||||
"anime": [],
|
||||
"other": [],
|
||||
}
|
||||
|
||||
_, stats_first = enrich_identities_data(identities, config)
|
||||
assert stats_first["enriched"] == 1
|
||||
assert stats_first["cache_hits"] == 0
|
||||
assert stats_first["api_calls"] == 1
|
||||
assert provider.calls == 1
|
||||
|
||||
_, stats_second = enrich_identities_data(identities, config)
|
||||
assert stats_second["cache_hits"] == 1
|
||||
assert stats_second["api_calls"] == 0
|
||||
assert provider.calls == 1
|
||||
|
||||
movie = identities["movies"][0]
|
||||
assert movie["title_zh"] == "测试中文名"
|
||||
assert movie["title_en"] == "Test English Title"
|
||||
assert movie["display_title"] == "测试中文名 Test English Title"
|
||||
|
||||
|
||||
def test_enrich_flags_low_reputation_for_review(tmp_path, monkeypatch):
|
||||
"""Low reputation should set needs_review when vote count is high enough."""
|
||||
|
||||
class LowScoreProvider(DummyProvider):
|
||||
def enrich(self, *, title: str, media_type: str, year=None):
|
||||
self.calls += 1
|
||||
return ProviderResult(
|
||||
provider="dummy",
|
||||
canonical_id=f"dummy:{title}",
|
||||
title_zh="低分作品",
|
||||
title_en="Low Score",
|
||||
translation_source="dummy",
|
||||
reputation_score=4.5,
|
||||
reputation_votes=500,
|
||||
reputation_source="dummy",
|
||||
)
|
||||
|
||||
config = Config(
|
||||
library_root=tmp_path,
|
||||
enrichment_cache_db=tmp_path / "cache.db",
|
||||
enrichment_providers=["dummy"],
|
||||
translation_fallback_machine=False,
|
||||
reputation_low_score_threshold=6.0,
|
||||
reputation_min_votes=50,
|
||||
)
|
||||
|
||||
provider = LowScoreProvider()
|
||||
monkeypatch.setattr(
|
||||
"vlm.enrichment._build_providers",
|
||||
lambda _config, request_timeout, retries: [provider],
|
||||
)
|
||||
|
||||
identities = {
|
||||
"metadata": {},
|
||||
"movies": [
|
||||
{
|
||||
"path": "/library/movie/Unknown.2020.mkv",
|
||||
"filename": "Unknown.2020.mkv",
|
||||
"category": "movie",
|
||||
"title": "Unknown",
|
||||
"year": 2020,
|
||||
"confidence": 0.9,
|
||||
"needs_review": False,
|
||||
}
|
||||
],
|
||||
"series": [],
|
||||
"anime": [],
|
||||
"other": [],
|
||||
}
|
||||
|
||||
enrich_identities_data(identities, config)
|
||||
|
||||
assert identities["movies"][0]["needs_review"] is True
|
||||
|
||||
|
||||
def test_enrich_refresh_all_bypasses_cache(tmp_path, monkeypatch):
|
||||
"""refresh_all should not use cache and should invoke provider again."""
|
||||
config = Config(
|
||||
library_root=tmp_path,
|
||||
enrichment_cache_db=tmp_path / "cache.db",
|
||||
enrichment_providers=["dummy"],
|
||||
translation_fallback_machine=False,
|
||||
)
|
||||
|
||||
provider = DummyProvider()
|
||||
monkeypatch.setattr(
|
||||
"vlm.enrichment._build_providers",
|
||||
lambda _config, request_timeout, retries: [provider],
|
||||
)
|
||||
|
||||
identities = {
|
||||
"metadata": {},
|
||||
"movies": [
|
||||
{
|
||||
"path": "/library/movie/Test.2020.mkv",
|
||||
"filename": "Test.2020.mkv",
|
||||
"category": "movie",
|
||||
"title": "Test",
|
||||
"year": 2020,
|
||||
"confidence": 0.9,
|
||||
"needs_review": False,
|
||||
}
|
||||
],
|
||||
"series": [],
|
||||
"anime": [],
|
||||
"other": [],
|
||||
}
|
||||
|
||||
enrich_identities_data(identities, config, refresh_mode="incremental")
|
||||
assert provider.calls == 1
|
||||
|
||||
_, stats = enrich_identities_data(identities, config, refresh_mode="refresh_all")
|
||||
assert provider.calls == 2
|
||||
assert stats["cache_hits"] == 0
|
||||
@@ -921,3 +921,52 @@ def test_plan_json_includes_all_required_fields(config, tmp_path):
|
||||
required_op_fields = ["operation_type", "source_path", "destination_path", "reason", "has_conflict", "conflict_reason"]
|
||||
for field in required_op_fields:
|
||||
assert field in operation, f"Missing required operation field: {field}"
|
||||
|
||||
|
||||
def test_movie_rejected_by_review_generates_noop(config):
|
||||
"""Rejected movie should not generate move/rename operation."""
|
||||
video_file = VideoFile(
|
||||
path=Path("/mnt/nas/videos/movie/Movie.2020.mkv"),
|
||||
filename="Movie.2020.mkv",
|
||||
size_bytes=1000000,
|
||||
modified_timestamp=datetime.now(),
|
||||
category="movie"
|
||||
)
|
||||
|
||||
identity = MovieIdentity(
|
||||
title="Movie",
|
||||
year=2020,
|
||||
confidence=0.9,
|
||||
needs_review=False,
|
||||
original_filename="Movie.2020.mkv",
|
||||
review_status="rejected"
|
||||
)
|
||||
|
||||
plan = generate_plan([(video_file, identity)], config)
|
||||
assert plan.operations[0].operation_type == "no-op"
|
||||
assert "rejected" in plan.operations[0].reason.lower()
|
||||
|
||||
|
||||
def test_series_rejected_by_review_generates_noop(config):
|
||||
"""Rejected series should not generate move/rename operation."""
|
||||
video_file = VideoFile(
|
||||
path=Path("/mnt/nas/videos/series/Show.S01E01.mkv"),
|
||||
filename="Show.S01E01.mkv",
|
||||
size_bytes=1000000,
|
||||
modified_timestamp=datetime.now(),
|
||||
category="series"
|
||||
)
|
||||
|
||||
identity = SeriesIdentity(
|
||||
title="Show",
|
||||
season=1,
|
||||
episodes=[1],
|
||||
confidence=0.9,
|
||||
needs_review=False,
|
||||
original_filename="Show.S01E01.mkv",
|
||||
review_status="rejected"
|
||||
)
|
||||
|
||||
plan = generate_plan([(video_file, identity)], config)
|
||||
assert plan.operations[0].operation_type == "no-op"
|
||||
assert "rejected" in plan.operations[0].reason.lower()
|
||||
|
||||
Reference in New Issue
Block a user