feat: trace quality agent selection report, identity clustering runner_v2 DB write, age/gender CoreML selection, updated experiment config UUID
This commit is contained in:
Warren
@@ -0,0 +1,246 @@
|
||||
use anyhow::{Context, Result};
|
||||
use serde::Deserialize;
|
||||
use std::collections::HashMap;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::core::config;
|
||||
use crate::core::db::PostgresDb;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TmdbSearchResult {
|
||||
results: Vec<TmdbMovie>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TmdbMovie {
|
||||
id: u64,
|
||||
title: String,
|
||||
release_date: Option<String>,
|
||||
overview: Option<String>,
|
||||
poster_path: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TmdbCredits {
|
||||
cast: Vec<TmdbCastMember>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TmdbCastMember {
|
||||
id: u64,
|
||||
name: String,
|
||||
character: String,
|
||||
profile_path: Option<String>,
|
||||
order: u32,
|
||||
}
|
||||
|
||||
pub struct TmdbProbeResult {
|
||||
pub tmdb_id: u64,
|
||||
pub title: String,
|
||||
pub cast_count: usize,
|
||||
pub identities_created: usize,
|
||||
}
|
||||
|
||||
fn extract_movie_name(filename: &str) -> Option<String> {
|
||||
let name = std::path::Path::new(filename)
|
||||
.file_stem()
|
||||
.and_then(|s| s.to_str())?;
|
||||
|
||||
let cleaned = name.replace(['.', '_'], " ").trim().to_string();
|
||||
|
||||
if cleaned.is_empty() || cleaned.len() < 3 {
|
||||
return None;
|
||||
}
|
||||
Some(cleaned)
|
||||
}
|
||||
|
||||
pub async fn probe_movie(
|
||||
db: &PostgresDb,
|
||||
filename: &str,
|
||||
file_uuid: &str,
|
||||
) -> Result<Option<TmdbProbeResult>> {
|
||||
let api_key = match config::tmdb::API_KEY.as_ref() {
|
||||
Some(k) => k.clone(),
|
||||
None => {
|
||||
info!("[TMDB] No API key configured, skipping TMDb probe");
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
let movie_name = match extract_movie_name(filename) {
|
||||
Some(n) => n,
|
||||
None => {
|
||||
info!("[TMDB] Could not extract movie name from: {}", filename);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
info!("[TMDB] Searching for movie: {}", movie_name);
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
// Step 1: Search movie
|
||||
let search_url = format!(
|
||||
"https://api.themoviedb.org/3/search/movie?api_key={}&query={}",
|
||||
api_key,
|
||||
urlencoding(&movie_name)
|
||||
);
|
||||
|
||||
let search_resp = client
|
||||
.get(&search_url)
|
||||
.send()
|
||||
.await
|
||||
.context("TMDb search request failed")?
|
||||
.json::<TmdbSearchResult>()
|
||||
.await
|
||||
.context("Failed to parse TMDb search response")?;
|
||||
|
||||
let movie = match search_resp.results.into_iter().next() {
|
||||
Some(m) => m,
|
||||
None => {
|
||||
info!("[TMDB] No movie found for: {}", movie_name);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
info!("[TMDB] Matched: {} (TMDB id={})", movie.title, movie.id);
|
||||
|
||||
// Step 2: Fetch cast
|
||||
let credits_url = format!(
|
||||
"https://api.themoviedb.org/3/movie/{}/credits?api_key={}",
|
||||
movie.id, api_key
|
||||
);
|
||||
|
||||
let credits = client
|
||||
.get(&credits_url)
|
||||
.send()
|
||||
.await
|
||||
.context("TMDb credits request failed")?
|
||||
.json::<TmdbCredits>()
|
||||
.await
|
||||
.context("Failed to parse TMDb credits response")?;
|
||||
|
||||
// Step 3: Create identities for top cast
|
||||
let identities_table = crate::core::db::schema::table_name("identities");
|
||||
let mut identities_created = 0usize;
|
||||
|
||||
for member in credits.cast.iter().take(20) {
|
||||
if member.name.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let profile_url = member
|
||||
.profile_path
|
||||
.as_ref()
|
||||
.map(|p| format!("https://image.tmdb.org/t/p/w185{}", p));
|
||||
|
||||
let result = sqlx::query(&format!(
|
||||
"INSERT INTO {} (name, identity_type, source, status, tmdb_id, tmdb_profile, metadata) \
|
||||
VALUES ($1, 'people', 'tmdb', 'confirmed', $2, $3, \
|
||||
jsonb_build_object('tmdb_character', $4, 'tmdb_cast_order', $5, 'tmdb_movie_id', $6, 'tmdb_movie_title', $7)) \
|
||||
ON CONFLICT (name) DO UPDATE SET \
|
||||
tmdb_id = COALESCE(EXCLUDED.tmdb_id, {}.tmdb_id), \
|
||||
tmdb_profile = COALESCE(EXCLUDED.tmdb_profile, {}.tmdb_profile), \
|
||||
metadata = {}.metadata || jsonb_build_object('tmdb_movie_id', $6, 'tmdb_movie_title', $7) \
|
||||
RETURNING id",
|
||||
identities_table, identities_table, identities_table, identities_table
|
||||
))
|
||||
.bind(&member.name)
|
||||
.bind(member.id as i64)
|
||||
.bind(&profile_url)
|
||||
.bind(&member.character)
|
||||
.bind(member.order as i32)
|
||||
.bind(movie.id as i64)
|
||||
.bind(&movie.title)
|
||||
.execute(db.pool())
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(_) => {
|
||||
info!(
|
||||
"[TMDB] Created/updated identity: {} as {}",
|
||||
member.name, member.character
|
||||
);
|
||||
identities_created += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("[TMDB] Failed to create identity '{}': {}", member.name, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Step 4: Trigger background embedding extraction
|
||||
if identities_created > 0 {
|
||||
let scripts_dir = std::env::var("MOMENTRY_SCRIPTS_DIR")
|
||||
.unwrap_or_else(|_| "/Users/accusys/momentry_core_0.1/scripts".to_string());
|
||||
let python_path = std::env::var("MOMENTRY_PYTHON_PATH")
|
||||
.unwrap_or_else(|_| "/opt/homebrew/bin/python3.11".to_string());
|
||||
let schema = crate::core::config::DATABASE_SCHEMA.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let output = tokio::process::Command::new(&python_path)
|
||||
.arg(&format!("{}/tmdb_embed_extractor.py", scripts_dir))
|
||||
.arg("--schema")
|
||||
.arg(&schema)
|
||||
.output()
|
||||
.await;
|
||||
|
||||
match output {
|
||||
Ok(o) => {
|
||||
if !o.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&o.stderr);
|
||||
warn!("[TMDB] Embed extraction script failed: {}", stderr);
|
||||
} else {
|
||||
info!("[TMDB] Background face embedding extraction complete");
|
||||
}
|
||||
}
|
||||
Err(e) => warn!("[TMDB] Failed to run embed extraction script: {}", e),
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Step 5: Store tmdb_id on the video record for later use
|
||||
let videos_table = crate::core::db::schema::table_name("videos");
|
||||
let tmdb_label = "tmdb";
|
||||
let _ = sqlx::query(&format!(
|
||||
"UPDATE {} SET birth_registration = \
|
||||
jsonb_set(COALESCE(birth_registration, '{{}}'::jsonb), '{{{}}}', $1::jsonb) \
|
||||
WHERE file_uuid = $2",
|
||||
videos_table, tmdb_label
|
||||
))
|
||||
.bind(serde_json::json!({
|
||||
"movie_id": movie.id,
|
||||
"movie_title": movie.title,
|
||||
"release_date": movie.release_date,
|
||||
"poster": movie.poster_path,
|
||||
"cast_count": credits.cast.len(),
|
||||
"identities_created": identities_created,
|
||||
}))
|
||||
.bind(file_uuid)
|
||||
.execute(db.pool())
|
||||
.await
|
||||
.ok();
|
||||
|
||||
info!(
|
||||
"[TMDB] Probe complete: {} cast members, {} identities created/updated",
|
||||
credits.cast.len(),
|
||||
identities_created
|
||||
);
|
||||
|
||||
Ok(Some(TmdbProbeResult {
|
||||
tmdb_id: movie.id,
|
||||
title: movie.title,
|
||||
cast_count: credits.cast.len(),
|
||||
identities_created,
|
||||
}))
|
||||
}
|
||||
|
||||
fn urlencoding(s: &str) -> String {
|
||||
s.chars()
|
||||
.map(|c| match c {
|
||||
'A'..='Z' | 'a'..='z' | '0'..='9' | '-' | '_' | '.' | '~' => c.to_string(),
|
||||
' ' => '+'.to_string(),
|
||||
_ => format!("%{:02X}", c as u8),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
Reference in New Issue
Block a user