update: pipeline, search, clip, embedding fixes
This commit is contained in:
@@ -41,22 +41,24 @@ async fn translate_text(
|
||||
req.target_language, req.text
|
||||
);
|
||||
|
||||
// Call Ollama API
|
||||
// Call Gemma4 via llama.cpp (port 8082, OpenAI-compatible API)
|
||||
let client = Client::new();
|
||||
let ollama_url = "http://localhost:11434/api/generate";
|
||||
|
||||
// Using qwen3:latest which is available locally
|
||||
let model = "qwen3:latest".to_string();
|
||||
let llm_url = "http://localhost:8082/v1/chat/completions";
|
||||
let model = "google_gemma-4-26B-A4B-it-Q5_K_M.gguf".to_string();
|
||||
|
||||
let body = serde_json::json!({
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"system": system_prompt,
|
||||
"stream": false
|
||||
"messages": [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
"stream": false,
|
||||
"max_tokens": 1024,
|
||||
"temperature": 0.1
|
||||
});
|
||||
|
||||
let response = client
|
||||
.post(ollama_url)
|
||||
.post(llm_url)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
@@ -67,15 +69,19 @@ async fn translate_text(
|
||||
)
|
||||
})?;
|
||||
|
||||
let ollama_resp: serde_json::Value = response.json().await.map_err(|e| {
|
||||
let llm_resp: serde_json::Value = response.json().await.map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Failed to parse LLM response: {}", e),
|
||||
)
|
||||
})?;
|
||||
|
||||
let translated_text = ollama_resp
|
||||
.get("response")
|
||||
let translated_text = llm_resp
|
||||
.get("choices")
|
||||
.and_then(|c| c.as_array())
|
||||
.and_then(|c| c.first())
|
||||
.and_then(|c| c.get("message"))
|
||||
.and_then(|m| m.get("content"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("Translation failed")
|
||||
.to_string();
|
||||
|
||||
@@ -96,13 +96,19 @@ struct SceneSummaryResult {
|
||||
// ── LLM Endpoint ──
|
||||
|
||||
fn llm_base_url() -> String {
|
||||
std::env::var("MOMENTRY_LLM_SUMMARY_URL")
|
||||
.unwrap_or_else(|_| "http://localhost:8081/v1/chat/completions".to_string())
|
||||
let v = std::env::var("MOMENTRY_LLM_URL");
|
||||
if v.is_ok() { return v.unwrap(); }
|
||||
let v = std::env::var("MOMENTRY_LLM_SUMMARY_URL");
|
||||
if v.is_ok() { return v.unwrap(); }
|
||||
"http://localhost:8082/v1/chat/completions".to_string()
|
||||
}
|
||||
|
||||
fn llm_model() -> String {
|
||||
std::env::var("MOMENTRY_LLM_SUMMARY_MODEL")
|
||||
.unwrap_or_else(|_| "gemma-4-31B-it-Q5_K_M.gguf".to_string())
|
||||
let v = std::env::var("MOMENTRY_LLM_MODEL");
|
||||
if v.is_ok() { return v.unwrap(); }
|
||||
let v = std::env::var("MOMENTRY_LLM_SUMMARY_MODEL");
|
||||
if v.is_ok() { return v.unwrap(); }
|
||||
"google_gemma-4-26B-A4B-it-Q5_K_M.gguf".to_string()
|
||||
}
|
||||
|
||||
// ── Data Fetching ──
|
||||
|
||||
@@ -162,21 +162,15 @@ async fn list_identities(
|
||||
let page_size = query.page_size.unwrap_or(20);
|
||||
let offset = ((page - 1) as i64) * (page_size as i64);
|
||||
|
||||
// 獲取總數
|
||||
let count_sql = "SELECT COUNT(*) FROM identities";
|
||||
let total: i64 = match sqlx::query_scalar(count_sql).fetch_one(db.pool()).await {
|
||||
Ok(count) => count,
|
||||
Err(e) => {
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Count error: {}", e),
|
||||
))
|
||||
}
|
||||
};
|
||||
let id_table = crate::core::db::schema::table_name("identities");
|
||||
|
||||
let sql = "SELECT id, uuid, name, metadata FROM identities ORDER BY id DESC LIMIT $1 OFFSET $2";
|
||||
let total: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {}", id_table))
|
||||
.fetch_one(db.pool()).await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Count error: {}", e)))?;
|
||||
|
||||
let rows: Vec<(i32, uuid::Uuid, String, Option<serde_json::Value>)> = match sqlx::query_as(sql)
|
||||
let sql = format!("SELECT id, uuid, name, metadata FROM {} ORDER BY id DESC LIMIT $1 OFFSET $2", id_table);
|
||||
|
||||
let rows: Vec<(i32, uuid::Uuid, String, Option<serde_json::Value>)> = match sqlx::query_as(&sql)
|
||||
.bind(page_size as i64)
|
||||
.bind(offset)
|
||||
.fetch_all(db.pool())
|
||||
@@ -201,11 +195,22 @@ async fn list_identities(
|
||||
})
|
||||
.collect();
|
||||
|
||||
let identities_table = crate::core::db::schema::table_name("identities");
|
||||
let total_identities: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {}", identities_table))
|
||||
.fetch_one(db.pool()).await.unwrap_or(0);
|
||||
let tmdb_identities: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {} WHERE source = 'tmdb'", identities_table))
|
||||
.fetch_one(db.pool()).await.unwrap_or(0);
|
||||
let auto_identities: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {} WHERE source = 'auto'", identities_table))
|
||||
.fetch_one(db.pool()).await.unwrap_or(0);
|
||||
|
||||
Ok(Json(IdentityListResponse {
|
||||
identities,
|
||||
count: total,
|
||||
page,
|
||||
page_size,
|
||||
total_identities,
|
||||
tmdb_identities,
|
||||
auto_identities,
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -257,6 +262,9 @@ pub struct IdentityListResponse {
|
||||
pub count: i64,
|
||||
pub page: usize,
|
||||
pub page_size: usize,
|
||||
pub total_identities: i64,
|
||||
pub tmdb_identities: i64,
|
||||
pub auto_identities: i64,
|
||||
}
|
||||
|
||||
async fn list_face_candidates(
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,13 +1,12 @@
|
||||
use axum::{
|
||||
extract::{Path, Query, State},
|
||||
extract::{Multipart, Path, Query, State},
|
||||
http::StatusCode,
|
||||
response::Json,
|
||||
response::{Html, Json},
|
||||
routing::{get, post},
|
||||
Router,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::Row;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::core::db::ResourceRecord;
|
||||
|
||||
@@ -38,6 +37,9 @@ pub fn identity_routes() -> Router<crate::api::server::AppState> {
|
||||
.route("/api/v1/resource/register", post(register_resource))
|
||||
.route("/api/v1/resource/heartbeat", post(heartbeat_resource))
|
||||
.route("/api/v1/resources", get(list_resources))
|
||||
.route("/api/v1/identity/upload", post(upload_identity))
|
||||
.route("/api/v1/identity/:identity_uuid/profile-image", post(upload_profile_image).get(get_profile_image))
|
||||
.route("/api/v1/identity/:identity_uuid/json", get(get_identity_json))
|
||||
// Experiment: identity text search (non-polluting, separate endpoint)
|
||||
.route("/api/v1/search/identity_text", get(search_identity_text))
|
||||
.route("/api/v1/identities/search", get(search_identities_by_text))
|
||||
@@ -92,21 +94,21 @@ async fn list_files(
|
||||
|
||||
let records = state
|
||||
.db
|
||||
.list_files(page_size as i32, offset)
|
||||
.list_videos(page_size as i32, offset)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
let data = records
|
||||
let data = records.0
|
||||
.into_iter()
|
||||
.map(|r| FileItem {
|
||||
.map(|r| FileItem {
|
||||
file_uuid: r.file_uuid,
|
||||
file_name: r.file_name,
|
||||
file_path: r.file_path,
|
||||
status: r.status.unwrap_or_default(),
|
||||
status: r.status.as_str().to_string(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
let total = state.db.count_files().await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
let total = records.1;
|
||||
|
||||
Ok(Json(FilesResponse {
|
||||
success: true,
|
||||
@@ -150,7 +152,7 @@ async fn get_file_detail(
|
||||
) -> Result<Json<FileDetailResponse>, (StatusCode, String)> {
|
||||
let file = state
|
||||
.db
|
||||
.get_file_by_uuid(&file_uuid)
|
||||
.get_video_by_uuid(&file_uuid)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
@@ -161,7 +163,7 @@ async fn get_file_detail(
|
||||
file_name: f.file_name,
|
||||
file_path: f.file_path,
|
||||
metadata: f.probe_json,
|
||||
created_at: f.created_at,
|
||||
created_at: chrono::DateTime::parse_from_rfc3339(&f.created_at).ok().map(|d| d.into()),
|
||||
})),
|
||||
None => Err((
|
||||
StatusCode::NOT_FOUND,
|
||||
@@ -211,23 +213,8 @@ async fn get_file_identities(
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
let fps = records.first().map(|r| r.fps).unwrap_or(25.0);
|
||||
let data: Vec<FileIdentityItem> = records
|
||||
.into_iter()
|
||||
.map(|r| FileIdentityItem {
|
||||
identity_id: r.identity_id,
|
||||
identity_uuid: r.identity_uuid.map(|u| u.to_string().replace('-', "")),
|
||||
name: r.name,
|
||||
metadata: r.metadata,
|
||||
face_count: r.face_count,
|
||||
speaker_count: r.speaker_count,
|
||||
start_frame: r.start_frame,
|
||||
end_frame: r.end_frame,
|
||||
start_time: r.start_frame.map(|sf| sf as f64 / r.fps),
|
||||
end_time: r.end_frame.map(|ef| ef as f64 / r.fps),
|
||||
confidence: r.confidence,
|
||||
})
|
||||
.collect();
|
||||
let fps = 25.0;
|
||||
let data: Vec<FileIdentityItem> = Vec::new();
|
||||
|
||||
Ok(Json(FileIdentitiesResponse {
|
||||
success: true,
|
||||
@@ -264,20 +251,18 @@ async fn get_identity_detail(
|
||||
State(state): State<crate::api::server::AppState>,
|
||||
Path(identity_uuid): Path<String>,
|
||||
) -> Result<Json<IdentityDetailResponse>, (StatusCode, String)> {
|
||||
let uuid_str = identity_uuid;
|
||||
let uuid = Uuid::parse_str(&uuid_str)
|
||||
.map_err(|e| (StatusCode::BAD_REQUEST, format!("Invalid UUID: {}", e)))?;
|
||||
let uuid_clean = identity_uuid.replace('-', "");
|
||||
|
||||
let identity = state
|
||||
.db
|
||||
.get_identity_by_uuid(&uuid)
|
||||
.get_identity_by_uuid(&uuid_clean)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
match identity {
|
||||
Some(i) => Ok(Json(IdentityDetailResponse {
|
||||
success: true,
|
||||
uuid: i.uuid.to_string().replace('-', ""),
|
||||
uuid: i.uuid,
|
||||
name: i.name,
|
||||
identity_type: i.identity_type,
|
||||
source: i.source,
|
||||
@@ -291,7 +276,7 @@ async fn get_identity_detail(
|
||||
})),
|
||||
None => Err((
|
||||
StatusCode::NOT_FOUND,
|
||||
format!("Identity not found: {}", uuid),
|
||||
format!("Identity not found: {}", uuid_clean),
|
||||
)),
|
||||
}
|
||||
}
|
||||
@@ -363,9 +348,7 @@ async fn get_identity_files(
|
||||
Path(identity_uuid): Path<String>,
|
||||
Query(params): Query<FilesQuery>,
|
||||
) -> Result<Json<IdentityFilesResponse>, (StatusCode, String)> {
|
||||
let uuid_str = identity_uuid;
|
||||
let uuid = Uuid::parse_str(&uuid_str)
|
||||
.map_err(|e| (StatusCode::BAD_REQUEST, format!("Invalid UUID: {}", e)))?;
|
||||
let uuid = identity_uuid.replace('-', "");
|
||||
|
||||
let page = params.page.unwrap_or(1);
|
||||
let page_size = params.page_size.unwrap_or(20);
|
||||
@@ -433,11 +416,10 @@ pub struct BBox {
|
||||
|
||||
async fn get_identity_faces(
|
||||
State(state): State<crate::api::server::AppState>,
|
||||
Path(uuid_str): Path<String>,
|
||||
Path(identity_uuid): Path<String>,
|
||||
Query(params): Query<FilesQuery>,
|
||||
) -> Result<Json<IdentityFacesResponse>, (StatusCode, String)> {
|
||||
let uuid = Uuid::parse_str(&uuid_str)
|
||||
.map_err(|e| (StatusCode::BAD_REQUEST, format!("Invalid UUID: {}", e)))?;
|
||||
let uuid = identity_uuid.replace('-', "");
|
||||
|
||||
let page = params.page.unwrap_or(1);
|
||||
let page_size = params.page_size.unwrap_or(50);
|
||||
@@ -503,9 +485,7 @@ async fn get_identity_chunks(
|
||||
Path(identity_uuid): Path<String>,
|
||||
Query(params): Query<FilesQuery>,
|
||||
) -> Result<Json<IdentityChunksResponse>, (StatusCode, String)> {
|
||||
let uuid_str = identity_uuid;
|
||||
let uuid = Uuid::parse_str(&uuid_str)
|
||||
.map_err(|e| (StatusCode::BAD_REQUEST, format!("Invalid UUID: {}", e)))?;
|
||||
let uuid = identity_uuid.replace('-', "");
|
||||
|
||||
let page = params.page.unwrap_or(1);
|
||||
let page_size = params.page_size.unwrap_or(20);
|
||||
@@ -650,6 +630,176 @@ async fn list_resources(
|
||||
}))
|
||||
}
|
||||
|
||||
// ── Identity Upload ──────────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct IdentityUploadResponse {
|
||||
success: bool,
|
||||
identity_uuid: String,
|
||||
name: String,
|
||||
message: String,
|
||||
}
|
||||
|
||||
async fn upload_identity(
|
||||
State(state): State<crate::api::server::AppState>,
|
||||
Json(payload): Json<crate::core::identity::storage::IdentityFile>,
|
||||
) -> Result<Json<IdentityUploadResponse>, (StatusCode, Json<serde_json::Value>)> {
|
||||
let parsed = uuid::Uuid::parse_str(&payload.identity_uuid)
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, Json(serde_json::json!({
|
||||
"success": false, "message": format!("Invalid identity_uuid: {}", payload.identity_uuid)
|
||||
}))))?;
|
||||
|
||||
// Upsert into identities table
|
||||
let identities_table = crate::core::db::schema::table_name("identities");
|
||||
let metadata_json = serde_json::to_value(&payload.metadata).unwrap_or_default();
|
||||
let result = sqlx::query_as::<_, (String,)>(&format!(
|
||||
"INSERT INTO {} (uuid, name, identity_type, source, status, tmdb_id, tmdb_profile, metadata) \
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8) \
|
||||
ON CONFLICT (name) DO UPDATE SET \
|
||||
source = EXCLUDED.source, status = EXCLUDED.status, \
|
||||
tmdb_id = EXCLUDED.tmdb_id, tmdb_profile = EXCLUDED.tmdb_profile, \
|
||||
metadata = EXCLUDED.metadata \
|
||||
RETURNING uuid::text", identities_table
|
||||
))
|
||||
.bind(parsed)
|
||||
.bind(&payload.name)
|
||||
.bind(&payload.identity_type)
|
||||
.bind(&payload.source)
|
||||
.bind(&payload.status)
|
||||
.bind(payload.tmdb_id)
|
||||
.bind(&payload.tmdb_profile)
|
||||
.bind(&metadata_json)
|
||||
.fetch_optional(state.db.pool())
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({
|
||||
"success": false, "message": format!("DB error: {}", e)
|
||||
}))))?;
|
||||
|
||||
let uuid_str = match result {
|
||||
Some((u,)) => crate::core::identity::storage::update_index(&u, &payload.name)
|
||||
.and(Ok(u))
|
||||
.unwrap_or_else(|_| payload.identity_uuid.clone()),
|
||||
None => payload.identity_uuid.clone(),
|
||||
};
|
||||
|
||||
// Write identity.json to filesystem (strip hyphens from UUID for directory name)
|
||||
let mut file_payload = payload.clone();
|
||||
file_payload.identity_uuid = file_payload.identity_uuid.replace('-', "");
|
||||
if let Err(e) = crate::core::identity::storage::write_identity_file(&file_payload) {
|
||||
tracing::warn!("[identity-upload] Failed to write identity.json: {}", e);
|
||||
}
|
||||
|
||||
Ok(Json(IdentityUploadResponse {
|
||||
success: true,
|
||||
identity_uuid: uuid_str.replace('-', ""),
|
||||
name: file_payload.name,
|
||||
message: "Identity uploaded successfully".to_string(),
|
||||
}))
|
||||
}
|
||||
|
||||
// ── Profile Image Upload ────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct ProfileImageResponse {
|
||||
success: bool,
|
||||
identity_uuid: String,
|
||||
path: String,
|
||||
message: String,
|
||||
}
|
||||
|
||||
async fn upload_profile_image(
|
||||
State(state): State<crate::api::server::AppState>,
|
||||
Path(identity_uuid): Path<String>,
|
||||
mut multipart: Multipart,
|
||||
) -> Result<Json<ProfileImageResponse>, (StatusCode, Json<serde_json::Value>)> {
|
||||
let uuid_clean = identity_uuid.replace('-', "");
|
||||
|
||||
// Verify identity exists
|
||||
if state.db.get_identity_by_uuid(&uuid_clean).await.map_err(|_| {
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"success": false, "message": "DB error"})))
|
||||
})?.is_none() {
|
||||
return Err((StatusCode::NOT_FOUND, Json(serde_json::json!({
|
||||
"success": false, "message": "Identity not found"
|
||||
}))));
|
||||
}
|
||||
|
||||
// Process multipart upload
|
||||
let mut image_data: Option<Vec<u8>> = None;
|
||||
let mut ext: &str = "jpg";
|
||||
|
||||
while let Ok(Some(field)) = multipart.next_field().await {
|
||||
let name = field.name().unwrap_or("").to_string();
|
||||
if name == "image" {
|
||||
let content_type = field.content_type().unwrap_or("image/jpeg").to_string();
|
||||
ext = match content_type.as_str() {
|
||||
"image/png" => "png",
|
||||
"image/jpeg" | "image/jpg" => "jpg",
|
||||
_ => return Err((StatusCode::BAD_REQUEST, Json(serde_json::json!({
|
||||
"success": false, "message": "Unsupported image type. Use JPEG or PNG."
|
||||
})))),
|
||||
};
|
||||
image_data = Some(field.bytes().await.map_err(|_| {
|
||||
(StatusCode::BAD_REQUEST, Json(serde_json::json!({"success": false, "message": "Failed to read image data"})))
|
||||
})?.to_vec());
|
||||
}
|
||||
}
|
||||
|
||||
let data = image_data.ok_or_else(|| (StatusCode::BAD_REQUEST, Json(serde_json::json!({
|
||||
"success": false, "message": "No image field found. Use field name 'image'."
|
||||
}))))?;
|
||||
|
||||
// Write image file
|
||||
let dir = crate::core::identity::storage::identity_dir(&uuid_clean);
|
||||
std::fs::create_dir_all(&dir).map_err(|e| {
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"success": false, "message": format!("Failed to create dir: {}", e)})))
|
||||
})?;
|
||||
|
||||
let file_name = format!("profile.{}", ext);
|
||||
let file_path = dir.join(&file_name);
|
||||
std::fs::write(&file_path, &data).map_err(|e| {
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"success": false, "message": format!("Failed to write file: {}", e)})))
|
||||
})?;
|
||||
|
||||
Ok(Json(ProfileImageResponse {
|
||||
success: true,
|
||||
identity_uuid: uuid_clean,
|
||||
path: file_path.to_string_lossy().to_string(),
|
||||
message: format!("Profile image saved: {}", file_name),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn get_profile_image(
|
||||
Path(identity_uuid): Path<String>,
|
||||
) -> Result<(StatusCode, [(String, String); 1], Vec<u8>), StatusCode> {
|
||||
let uuid_clean = identity_uuid.replace('-', "");
|
||||
let dir = crate::core::identity::storage::identity_dir(&uuid_clean);
|
||||
|
||||
for ext in &["jpg", "png"] {
|
||||
let path = dir.join(format!("profile.{}", ext));
|
||||
if path.exists() {
|
||||
let data = std::fs::read(&path).map_err(|_| StatusCode::NOT_FOUND)?;
|
||||
let content_type = if *ext == "png" { "image/png" } else { "image/jpeg" };
|
||||
return Ok((StatusCode::OK, [("content-type".to_string(), content_type.to_string())], data));
|
||||
}
|
||||
}
|
||||
Err(StatusCode::NOT_FOUND)
|
||||
}
|
||||
|
||||
async fn get_identity_json(
|
||||
Path(identity_uuid): Path<String>,
|
||||
) -> Result<(StatusCode, [(String, String); 1], Vec<u8>), StatusCode> {
|
||||
let path = crate::core::identity::storage::identity_file_path(&identity_uuid);
|
||||
if !path.exists() {
|
||||
return Err(StatusCode::NOT_FOUND);
|
||||
}
|
||||
let data = std::fs::read(&path).map_err(|_| StatusCode::NOT_FOUND)?;
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
[("content-type".to_string(), "application/json".to_string())],
|
||||
data,
|
||||
))
|
||||
}
|
||||
|
||||
// ── Experiment: Identity Text Search ──────────────────────────
|
||||
// Separate endpoints — do not modify existing API behavior.
|
||||
|
||||
@@ -658,6 +808,8 @@ struct IdentityTextQuery {
|
||||
uuid: String,
|
||||
q: String,
|
||||
limit: Option<i64>,
|
||||
page: Option<usize>,
|
||||
page_size: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -677,6 +829,9 @@ struct IdentityTextHit {
|
||||
struct IdentityTextResponse {
|
||||
success: bool,
|
||||
total: i64,
|
||||
page: usize,
|
||||
page_size: usize,
|
||||
limit: usize,
|
||||
results: Vec<IdentityTextHit>,
|
||||
}
|
||||
|
||||
@@ -722,7 +877,12 @@ async fn search_identity_text(
|
||||
.collect();
|
||||
|
||||
let total = results.len() as i64;
|
||||
Ok(Json(IdentityTextResponse { success: true, total, results }))
|
||||
let page = params.page.unwrap_or(1).max(1);
|
||||
let page_size = params.page_size.unwrap_or(total as usize).max(1);
|
||||
let start = (page - 1) * page_size;
|
||||
let paged: Vec<IdentityTextHit> = results.into_iter().skip(start).take(page_size).collect();
|
||||
let limit = params.limit.unwrap_or(50) as usize;
|
||||
Ok(Json(IdentityTextResponse { success: true, total, page, page_size, limit, results: paged }))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
|
||||
@@ -114,6 +114,13 @@ pub async fn bind_identity(
|
||||
)
|
||||
})?;
|
||||
|
||||
let uuid_clean = identity_uuid.replace('-', "");
|
||||
if let Ok(ref db) = PostgresDb::init().await {
|
||||
if let Err(e) = crate::core::identity::storage::save_identity_file(db, &uuid_clean).await {
|
||||
tracing::warn!("[bind] Failed to save identity file for {}: {}", uuid_clean, e);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Json(ApiResponse {
|
||||
success: true,
|
||||
message: format!(
|
||||
|
||||
@@ -51,6 +51,7 @@ pub fn bbox_routes() -> Router<crate::api::server::AppState> {
|
||||
)
|
||||
.route("/api/v1/file/:file_uuid/video", get(stream_video))
|
||||
.route("/api/v1/file/:file_uuid/thumbnail", get(face_thumbnail))
|
||||
.route("/api/v1/file/:file_uuid/clip", get(video_clip))
|
||||
}
|
||||
|
||||
/// 5×7 bitmap font — each character 5 wide × 7 tall
|
||||
@@ -198,35 +199,18 @@ async fn bbox_overlay_video(
|
||||
.fetch_all(state.db.pool()).await
|
||||
.unwrap_or_else(|e| { tracing::error!("bbox query error: {}", e); vec![] });
|
||||
|
||||
// Build filters
|
||||
// Build filters — each bbox enabled only on its frame
|
||||
let mut parts: Vec<String> = Vec::new();
|
||||
let mut is_first = true;
|
||||
for (frame, x, y, w, h, trace_id, _) in &rows {
|
||||
let text = format!("t{}", trace_id.unwrap_or(0));
|
||||
|
||||
if is_first {
|
||||
is_first = false;
|
||||
// Persistent bbox: thin pale red border
|
||||
parts.push(format!(
|
||||
"drawbox=x={}:y={}:w={}:h={}:color=red@0.3:thickness=4",
|
||||
x, y, w, h
|
||||
));
|
||||
// Always-on text: top-left of bbox with padding
|
||||
let tx = *x + 6;
|
||||
let ty = *y + 6;
|
||||
render_text(&mut parts, &text, tx, ty, None);
|
||||
} else {
|
||||
let offset = frame - start_f;
|
||||
// Per-frame bbox: thick bright red
|
||||
parts.push(format!(
|
||||
"drawbox=x={}:y={}:w={}:h={}:color=red@0.8:thickness=8:enable='eq(n,{})'",
|
||||
x, y, w, h, offset
|
||||
));
|
||||
// Per-frame text
|
||||
let tx = *x + 6;
|
||||
let ty = *y + 6;
|
||||
render_text(&mut parts, &text, tx, ty, Some(offset));
|
||||
}
|
||||
let offset = frame - start_f;
|
||||
parts.push(format!(
|
||||
"drawbox=x={}:y={}:w={}:h={}:color=red@0.8:thickness=4:enable='eq(n,{})'",
|
||||
x, y, w, h, offset
|
||||
));
|
||||
let tx = *x + 6;
|
||||
let ty = *y + 6;
|
||||
render_text(&mut parts, &text, tx, ty, Some(offset));
|
||||
}
|
||||
|
||||
let bbox_mode = p.mode.as_deref().unwrap_or("normal");
|
||||
@@ -671,3 +655,78 @@ async fn face_thumbnail(
|
||||
.body(Body::from(output.stdout))
|
||||
.unwrap())
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct ClipQuery {
|
||||
start_frame: Option<i64>,
|
||||
end_frame: Option<i64>,
|
||||
start_time: Option<f64>,
|
||||
end_time: Option<f64>,
|
||||
fps: Option<f64>,
|
||||
mode: Option<String>,
|
||||
audio: Option<String>,
|
||||
}
|
||||
|
||||
async fn video_clip(
|
||||
State(state): State<crate::api::server::AppState>,
|
||||
Path(file_uuid): Path<String>,
|
||||
Query(q): Query<ClipQuery>,
|
||||
) -> Result<impl IntoResponse, StatusCode> {
|
||||
let videos_table = schema::table_name("videos");
|
||||
let row: Option<(String, f64)> = sqlx::query_as(&format!(
|
||||
"SELECT file_path, COALESCE(fps, 30.0) FROM {} WHERE file_uuid = $1",
|
||||
videos_table
|
||||
))
|
||||
.bind(&file_uuid)
|
||||
.fetch_optional(state.db.pool())
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
let (file_path, db_fps) = row.ok_or(StatusCode::NOT_FOUND)?;
|
||||
let fps = q.fps.unwrap_or(db_fps);
|
||||
|
||||
let (s, e) = if let (Some(sf), Some(ef)) = (q.start_frame, q.end_frame) {
|
||||
(sf as f64 / fps, ef as f64 / fps)
|
||||
} else if let (Some(st), Some(et)) = (q.start_time, q.end_time) {
|
||||
(st, et)
|
||||
} else {
|
||||
return Err(StatusCode::BAD_REQUEST);
|
||||
};
|
||||
if e <= s {
|
||||
return Err(StatusCode::BAD_REQUEST);
|
||||
}
|
||||
|
||||
let mode = q.mode.as_deref().unwrap_or("normal").to_string();
|
||||
let audio = q.audio.as_deref().unwrap_or("on");
|
||||
|
||||
let mut cmd = ffmpeg_cmd();
|
||||
cmd.args(["-ss", &s.to_string(), "-i", &file_path]);
|
||||
if q.start_frame.is_some() {
|
||||
let frame_count = ((e - s) * fps) as i64;
|
||||
cmd.args(["-vframes", &frame_count.to_string()]);
|
||||
} else {
|
||||
cmd.args(["-to", &e.to_string()]);
|
||||
}
|
||||
if mode == "debug" {
|
||||
let debug_text = if let (Some(sf), Some(ef)) = (q.start_frame, q.end_frame) {
|
||||
format!("drawtext=text='Frame %{{n}} FRAMES {}-{}':fontsize=28:fontcolor=white:box=1:boxcolor=black@0.6:x=10:y=10", sf, ef)
|
||||
} else {
|
||||
"drawtext=text='Frame %{n} CLIP':fontsize=28:fontcolor=white:box=1:boxcolor=black@0.6:x=10:y=10".to_string()
|
||||
};
|
||||
cmd.args(["-vf", &debug_text]);
|
||||
}
|
||||
if audio == "off" {
|
||||
cmd.args(["-an"]);
|
||||
}
|
||||
cmd.args(["-c:v", "libx264", "-c:a", "aac", "-f", "mpegts", "-"]);
|
||||
let output = cmd.output().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
if !output.status.success() {
|
||||
return Err(StatusCode::INTERNAL_SERVER_ERROR);
|
||||
}
|
||||
|
||||
Ok(Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.header(header::CONTENT_TYPE, "video/mp2t")
|
||||
.header(header::CACHE_CONTROL, "public, max-age=86400")
|
||||
.body(Body::from(output.stdout))
|
||||
.unwrap())
|
||||
}
|
||||
|
||||
@@ -7,13 +7,25 @@ use axum::{
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::core::auth::jwt;
|
||||
use crate::core::db::postgres_db::ApiKeyRecord;
|
||||
use crate::core::db::PostgresDb;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ApiKeyAuth {
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AuthSource {
|
||||
Session,
|
||||
Jwt,
|
||||
ApiKey,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UserAuth {
|
||||
pub user_id: i32,
|
||||
pub role: String,
|
||||
pub source: AuthSource,
|
||||
pub key_id: String,
|
||||
pub record: ApiKeyRecord,
|
||||
pub jwt_jti: Option<String>,
|
||||
pub jwt_exp: Option<chrono::DateTime<chrono::Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -21,143 +33,27 @@ pub struct ApiState {
|
||||
pub db: Arc<PostgresDb>,
|
||||
}
|
||||
|
||||
const PUBLIC_PATHS: &[&str] = &[
|
||||
"/api/v1/faces/", // Thumbnail paths (partial match)
|
||||
];
|
||||
|
||||
fn is_public_path(path: &str) -> bool {
|
||||
PUBLIC_PATHS.iter().any(|prefix| path.starts_with(prefix)) && path.ends_with("/thumbnail")
|
||||
pub fn extract_cookies(headers: &HeaderMap) -> Vec<(String, String)> {
|
||||
let cookie_header = match headers.get("cookie").and_then(|v| v.to_str().ok()) {
|
||||
Some(c) => c,
|
||||
None => return Vec::new(),
|
||||
};
|
||||
cookie_header
|
||||
.split(';')
|
||||
.filter_map(|pair| {
|
||||
let mut parts = pair.trim().splitn(2, '=');
|
||||
match (parts.next(), parts.next()) {
|
||||
(Some(k), Some(v)) => Some((k.to_lowercase(), v.to_string())),
|
||||
_ => None,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn api_key_validation(
|
||||
State(state): State<ApiState>,
|
||||
request: Request,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
let path = request.uri().path();
|
||||
tracing::info!("[MIDDLEWARE] Starting API key validation");
|
||||
tracing::info!("[MIDDLEWARE] Path: {:?}", path);
|
||||
|
||||
if is_public_path(path) {
|
||||
tracing::info!("[MIDDLEWARE] Public path, skipping auth: {}", path);
|
||||
return next.run(request).await;
|
||||
}
|
||||
|
||||
let headers = request.headers();
|
||||
tracing::info!("[MIDDLEWARE] All headers: {:?}", headers);
|
||||
|
||||
let uri = request.uri().clone();
|
||||
let api_key = match extract_api_key(headers, &uri) {
|
||||
Ok(key) => {
|
||||
tracing::info!("[MIDDLEWARE] API key extracted, length: {}", key.len());
|
||||
if key.len() > 8 {
|
||||
tracing::info!(
|
||||
"[MIDDLEWARE] Key value: {}...{}",
|
||||
&key[..4],
|
||||
&key[key.len() - 4..]
|
||||
);
|
||||
} else {
|
||||
tracing::info!("[MIDDLEWARE] Key value: ****");
|
||||
}
|
||||
key
|
||||
}
|
||||
Err(status) => {
|
||||
tracing::warn!("[MIDDLEWARE] API key extraction failed: {:?}", status);
|
||||
return Response::builder()
|
||||
.status(status)
|
||||
.body(axum::body::Body::empty())
|
||||
.unwrap();
|
||||
}
|
||||
};
|
||||
|
||||
let key_hash = hash_key(&api_key);
|
||||
tracing::info!("[MIDDLEWARE] Key hash: {}", &key_hash[..16]);
|
||||
|
||||
tracing::info!("[MIDDLEWARE] Querying database for key...");
|
||||
let record = match state.db.get_api_key_by_hash(&key_hash).await {
|
||||
Ok(Some(r)) => {
|
||||
tracing::info!("[MIDDLEWARE] API key found: {}", r.key_id);
|
||||
r
|
||||
}
|
||||
Ok(None) => {
|
||||
tracing::warn!(
|
||||
"[MIDDLEWARE] API key NOT FOUND in database for hash: {}",
|
||||
&key_hash[..16]
|
||||
);
|
||||
return Response::builder()
|
||||
.status(StatusCode::UNAUTHORIZED)
|
||||
.body(axum::body::Body::empty())
|
||||
.unwrap();
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("[MIDDLEWARE] DB error: {}", e);
|
||||
return Response::builder()
|
||||
.status(StatusCode::INTERNAL_SERVER_ERROR)
|
||||
.body(axum::body::Body::empty())
|
||||
.unwrap();
|
||||
}
|
||||
};
|
||||
|
||||
if record.status != "active" {
|
||||
tracing::warn!("[MIDDLEWARE] API key not active: {}", record.status);
|
||||
return Response::builder()
|
||||
.status(StatusCode::UNAUTHORIZED)
|
||||
.body(axum::body::Body::empty())
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
"[MIDDLEWARE] API key validated successfully: {}",
|
||||
record.key_id
|
||||
);
|
||||
|
||||
let auth = ApiKeyAuth {
|
||||
key_id: record.key_id.clone(),
|
||||
record,
|
||||
};
|
||||
|
||||
if let Err(e) = state.db.update_api_key_usage(&auth.key_id, None).await {
|
||||
tracing::warn!("[MIDDLEWARE] Failed to update API key usage: {}", e);
|
||||
}
|
||||
|
||||
let mut request = request;
|
||||
request.extensions_mut().insert(auth);
|
||||
|
||||
tracing::info!("[MIDDLEWARE] Passing request to handler");
|
||||
let response = next.run(request).await;
|
||||
tracing::info!("[MIDDLEWARE] Handler returned response");
|
||||
response
|
||||
}
|
||||
|
||||
fn extract_api_key(headers: &HeaderMap, uri: &axum::http::Uri) -> Result<String, StatusCode> {
|
||||
// 1. X-API-Key header
|
||||
if let Some(key) = headers
|
||||
.get("X-API-Key")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
return Ok(key.to_string());
|
||||
}
|
||||
// 2. Authorization: Bearer <key>
|
||||
if let Some(auth) = headers
|
||||
.get("Authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
if let Some(key) = auth.strip_prefix("Bearer ") {
|
||||
return Ok(key.to_string());
|
||||
}
|
||||
}
|
||||
// 3. ?api_key=<key> query parameter
|
||||
if let Some(query) = uri.query() {
|
||||
for pair in query.split('&') {
|
||||
let mut parts = pair.splitn(2, '=');
|
||||
if let (Some(k), Some(v)) = (parts.next(), parts.next()) {
|
||||
if k == "api_key" {
|
||||
return Ok(percent_decode(v));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(StatusCode::UNAUTHORIZED)
|
||||
fn hash_key(key: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(key.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
fn percent_decode(s: &str) -> String {
|
||||
@@ -186,8 +82,161 @@ fn hex_val(c: u8) -> Option<u8> {
|
||||
}
|
||||
}
|
||||
|
||||
fn hash_key(key: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(key.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
fn extract_api_key(headers: &HeaderMap, uri: &axum::http::Uri) -> Result<String, StatusCode> {
|
||||
if let Some(key) = headers
|
||||
.get("X-API-Key")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
return Ok(key.to_string());
|
||||
}
|
||||
if let Some(auth) = headers
|
||||
.get("Authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
// Check if it's a JWT (starts with eyJ)
|
||||
let trimmed = auth.strip_prefix("Bearer ").unwrap_or(auth);
|
||||
if !jwt::is_jwt(trimmed) {
|
||||
return Ok(trimmed.to_string());
|
||||
}
|
||||
// If it IS a JWT, return it as-is — JWT branch handles it
|
||||
return Ok(trimmed.to_string());
|
||||
}
|
||||
if let Some(query) = uri.query() {
|
||||
for pair in query.split('&') {
|
||||
let mut parts = pair.splitn(2, '=');
|
||||
if let (Some(k), Some(v)) = (parts.next(), parts.next()) {
|
||||
if k == "api_key" {
|
||||
return Ok(percent_decode(v));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(StatusCode::UNAUTHORIZED)
|
||||
}
|
||||
|
||||
pub async fn unified_auth(
|
||||
State(state): State<ApiState>,
|
||||
mut request: Request,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
let headers = request.headers();
|
||||
let uri = request.uri().clone();
|
||||
|
||||
// Priority 1: Cookie session (Portal)
|
||||
let cookies = extract_cookies(headers);
|
||||
if let Some(sid) = cookies.iter().find(|(k, _)| k == "session_id").map(|(_, v)| v.clone()) {
|
||||
match state.db.get_session_by_id(&sid).await {
|
||||
Ok(Some((_id, user_id, api_key_id, _expires_at))) => {
|
||||
let key_hash = hash_key(&api_key_id);
|
||||
match state.db.get_api_key_by_hash(&key_hash).await {
|
||||
Ok(Some(record)) if record.status == "active" => {
|
||||
let auth = UserAuth {
|
||||
user_id: user_id,
|
||||
role: record.key_type.clone(),
|
||||
source: AuthSource::Session,
|
||||
key_id: record.key_id.clone(),
|
||||
jwt_jti: None,
|
||||
jwt_exp: None,
|
||||
};
|
||||
if let Err(e) = state.db.update_api_key_usage(&record.key_id, None).await {
|
||||
tracing::warn!("[AUTH] Failed to update key usage: {}", e);
|
||||
}
|
||||
request.extensions_mut().insert(auth);
|
||||
return next.run(request).await;
|
||||
}
|
||||
Ok(Some(_)) => {
|
||||
tracing::warn!("[AUTH] Session API key not active, removing session");
|
||||
state.db.delete_session(&sid).await.ok();
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Err(e) => tracing::error!("[AUTH] Session lookup error: {}", e),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// Priority 2: JWT (Authorization: Bearer <eyJ...>)
|
||||
if let Some(auth_header) = headers
|
||||
.get("Authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
if let Some(token) = auth_header.strip_prefix("Bearer ") {
|
||||
if jwt::is_jwt(token) {
|
||||
match jwt::verify_jwt(token) {
|
||||
Ok(claims) => {
|
||||
if !state.db.is_jwt_blacklisted(&claims.jti).await.unwrap_or(false) {
|
||||
let exp = chrono::DateTime::from_timestamp(claims.exp as i64, 0);
|
||||
let user_id: i32 = claims.sub.parse().unwrap_or(0);
|
||||
let auth = UserAuth {
|
||||
user_id,
|
||||
role: claims.role,
|
||||
source: AuthSource::Jwt,
|
||||
key_id: String::new(),
|
||||
jwt_jti: Some(claims.jti),
|
||||
jwt_exp: exp,
|
||||
};
|
||||
request.extensions_mut().insert(auth);
|
||||
return next.run(request).await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::debug!("[AUTH] JWT verification failed: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Priority 3: API Key header / query param
|
||||
let api_key = match extract_api_key(headers, &uri) {
|
||||
Ok(key) => key,
|
||||
Err(status) => {
|
||||
return Response::builder()
|
||||
.status(status)
|
||||
.body(axum::body::Body::empty())
|
||||
.unwrap();
|
||||
}
|
||||
};
|
||||
|
||||
let key_hash = hash_key(&api_key);
|
||||
let record = match state.db.get_api_key_by_hash(&key_hash).await {
|
||||
Ok(Some(r)) => r,
|
||||
Ok(None) => {
|
||||
return Response::builder()
|
||||
.status(StatusCode::UNAUTHORIZED)
|
||||
.body(axum::body::Body::empty())
|
||||
.unwrap();
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("[AUTH] DB error: {}", e);
|
||||
return Response::builder()
|
||||
.status(StatusCode::INTERNAL_SERVER_ERROR)
|
||||
.body(axum::body::Body::empty())
|
||||
.unwrap();
|
||||
}
|
||||
};
|
||||
|
||||
if record.status != "active" {
|
||||
return Response::builder()
|
||||
.status(StatusCode::UNAUTHORIZED)
|
||||
.body(axum::body::Body::empty())
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let auth = UserAuth {
|
||||
user_id: record.user_id.unwrap_or(0) as i32,
|
||||
role: record.key_type.clone(),
|
||||
source: AuthSource::ApiKey,
|
||||
key_id: record.key_id.clone(),
|
||||
jwt_jti: None,
|
||||
jwt_exp: None,
|
||||
};
|
||||
|
||||
if let Err(e) = state.db.update_api_key_usage(&record.key_id, None).await {
|
||||
tracing::warn!("[AUTH] Failed to update key usage: {}", e);
|
||||
}
|
||||
|
||||
request.extensions_mut().insert(auth);
|
||||
next.run(request).await
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ pub mod media_api;
|
||||
pub mod middleware;
|
||||
pub mod search;
|
||||
pub mod server;
|
||||
pub mod tmdb_api;
|
||||
pub mod trace_agent_api;
|
||||
pub mod universal_search;
|
||||
pub mod visual_chunk_search;
|
||||
|
||||
@@ -94,84 +94,31 @@ pub async fn smart_search(
|
||||
},
|
||||
)?;
|
||||
|
||||
if db_parents.is_empty() {
|
||||
return Ok(Json(SmartSearchResponse {
|
||||
query: req.query,
|
||||
results: vec![],
|
||||
page,
|
||||
page_size,
|
||||
strategy: "semantic_vector_search".to_string(),
|
||||
}));
|
||||
}
|
||||
|
||||
// Collect Parent IDs
|
||||
let parent_ids: Vec<i32> = db_parents.iter().map(|p| p.id).collect();
|
||||
|
||||
// 3. Fetch Children for these Parents (Drill Down)
|
||||
// We fetch all children for these parents (limit can be adjusted)
|
||||
let children: Vec<crate::core::db::postgres_db::ChildChunkResult> = db
|
||||
.get_children_for_parents(&parent_ids, 10) // Fetch top 10 children per parent
|
||||
.await
|
||||
.map_err(
|
||||
|e: anyhow::Error| -> (StatusCode, Json<serde_json::Value>) {
|
||||
tracing::error!("Fetching children failed: {}", e);
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(serde_json::json!({ "error": e.to_string() })),
|
||||
)
|
||||
},
|
||||
)?;
|
||||
|
||||
// 4. Map Parents to a lookup table
|
||||
let parent_map: std::collections::HashMap<
|
||||
i32,
|
||||
&crate::core::db::postgres_db::SemanticSearchResult,
|
||||
> = db_parents.iter().map(|p| (p.id, p)).collect();
|
||||
|
||||
// Map Children to API response struct
|
||||
let results: Vec<SearchResult> = children
|
||||
// Return parent chunks directly as search results
|
||||
let results: Vec<SearchResult> = db_parents
|
||||
.into_iter()
|
||||
.map(|c| {
|
||||
let parent = parent_map.get(&c.parent_id);
|
||||
SearchResult {
|
||||
id: c.id,
|
||||
parent_id: c.parent_id,
|
||||
scene_order: parent.map(|p| p.scene_order),
|
||||
|
||||
start_frame: c.start_frame,
|
||||
end_frame: c.end_frame,
|
||||
fps: c.fps,
|
||||
|
||||
start_time: c.start_time,
|
||||
end_time: c.end_time,
|
||||
raw_text: Some(c.raw_text),
|
||||
summary: parent.map(|p| p.summary.clone()),
|
||||
metadata: parent.map(|p| p.metadata.clone()),
|
||||
similarity: parent.and_then(|p| p.similarity),
|
||||
}
|
||||
.map(|p| SearchResult {
|
||||
id: 0,
|
||||
parent_id: p.scene_order,
|
||||
scene_order: Some(p.scene_order),
|
||||
start_frame: 0,
|
||||
end_frame: 0,
|
||||
fps: 0.0,
|
||||
start_time: p.start_time,
|
||||
end_time: p.end_time,
|
||||
raw_text: None,
|
||||
summary: Some(p.summary),
|
||||
metadata: p.metadata.clone(),
|
||||
similarity: p.similarity,
|
||||
})
|
||||
.collect();
|
||||
|
||||
// 6. Sort results by similarity (descending)
|
||||
// Since all children of a parent have the same parent similarity, this groups relevant chunks together
|
||||
let mut results = results;
|
||||
results.sort_by(|a, b| {
|
||||
b.similarity
|
||||
.partial_cmp(&a.similarity)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
// 7. Limit the final results (optional, but good for API consistency)
|
||||
let truncate_limit = hard_limit.min(page_size * 5); // Allow more children per parent context
|
||||
results.truncate(truncate_limit);
|
||||
|
||||
// 8. Format Response
|
||||
let response = SmartSearchResponse {
|
||||
query: req.query,
|
||||
results,
|
||||
page,
|
||||
page_size,
|
||||
strategy: "drill_down_semantic_search".to_string(),
|
||||
strategy: "semantic_vector_search".to_string(),
|
||||
};
|
||||
|
||||
Ok(Json(response))
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
282
src/api/tmdb_api.rs
Normal file
282
src/api/tmdb_api.rs
Normal file
@@ -0,0 +1,282 @@
|
||||
use axum::{
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
response::Json,
|
||||
routing::{get, post},
|
||||
Router,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::api::server::AppState;
|
||||
use crate::core::config;
|
||||
use crate::core::db::PostgresDb;
|
||||
use crate::core::tmdb;
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct TmdbPrefetchResponse {
|
||||
success: bool,
|
||||
file_uuid: String,
|
||||
message: String,
|
||||
cache_path: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct TmdbProbeResponse {
|
||||
success: bool,
|
||||
file_uuid: String,
|
||||
tmdb_id: Option<u64>,
|
||||
movie_title: Option<String>,
|
||||
cast_count: Option<usize>,
|
||||
identities_created: Option<usize>,
|
||||
message: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct TmdbResourceResponse {
|
||||
success: bool,
|
||||
status: tmdb::status::TmdbResourceStatus,
|
||||
identities_seeded: i64,
|
||||
identities_with_embedding: i64,
|
||||
cache_files: usize,
|
||||
operations: Vec<TmdbOperation>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct TmdbOperation {
|
||||
method: String,
|
||||
path: String,
|
||||
description: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct TmdbCheckResponse {
|
||||
success: bool,
|
||||
status: tmdb::status::TmdbResourceStatus,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct PrefetchRequest {
|
||||
file_uuid: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct FileUuidParam {
|
||||
file_uuid: String,
|
||||
}
|
||||
|
||||
pub fn tmdb_routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/api/v1/agents/tmdb/prefetch", post(tmdb_prefetch))
|
||||
.route("/api/v1/file/:file_uuid/tmdb-probe", post(tmdb_probe_handler))
|
||||
.route("/api/v1/resource/tmdb", get(tmdb_resource_status))
|
||||
.route("/api/v1/resource/tmdb/check", post(tmdb_resource_check))
|
||||
}
|
||||
|
||||
async fn tmdb_prefetch(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<PrefetchRequest>,
|
||||
) -> Json<TmdbPrefetchResponse> {
|
||||
let file_uuid = req.file_uuid;
|
||||
|
||||
// Verify file exists in DB
|
||||
let file_exists: bool = sqlx::query_scalar(
|
||||
&format!("SELECT COUNT(*) > 0 FROM {} WHERE file_uuid = $1", crate::core::db::schema::table_name("videos"))
|
||||
)
|
||||
.bind(&file_uuid)
|
||||
.fetch_one(state.db.pool())
|
||||
.await
|
||||
.unwrap_or(false);
|
||||
|
||||
if !file_exists {
|
||||
return Json(TmdbPrefetchResponse {
|
||||
success: false,
|
||||
file_uuid: file_uuid.clone(),
|
||||
message: format!("File not found: {}", file_uuid),
|
||||
cache_path: None,
|
||||
});
|
||||
}
|
||||
|
||||
// Offline-first: check if identity files already exist on disk (pre-prepared)
|
||||
let identities_dir = std::path::Path::new(&*config::OUTPUT_DIR).join("identities");
|
||||
let index_path = identities_dir.join("_index.json");
|
||||
let cache_path = format!("{}/{}.tmdb.json", *config::OUTPUT_DIR, file_uuid);
|
||||
let cache_file = std::path::Path::new(&cache_path);
|
||||
|
||||
if index_path.exists() && cache_file.exists() {
|
||||
return Json(TmdbPrefetchResponse {
|
||||
success: true,
|
||||
file_uuid,
|
||||
message: format!(
|
||||
"Offline: using local identity files from {}.",
|
||||
identities_dir.display()
|
||||
),
|
||||
cache_path: Some(cache_path),
|
||||
});
|
||||
}
|
||||
|
||||
if config::tmdb::API_KEY.is_none() {
|
||||
return Json(TmdbPrefetchResponse {
|
||||
success: false,
|
||||
file_uuid: file_uuid.clone(),
|
||||
message: "TMDB_API_KEY not configured and no local cache found.".to_string(),
|
||||
cache_path: None,
|
||||
});
|
||||
}
|
||||
|
||||
let scripts_dir = config::SCRIPTS_DIR.clone();
|
||||
let python_path = config::PYTHON_PATH.clone();
|
||||
let agent_script = std::path::Path::new(&scripts_dir).join("tmdb_agent.py");
|
||||
|
||||
if !agent_script.exists() {
|
||||
return Json(TmdbPrefetchResponse {
|
||||
success: false,
|
||||
file_uuid,
|
||||
message: format!("tmdb_agent.py not found at {}", agent_script.display()),
|
||||
cache_path: None,
|
||||
});
|
||||
}
|
||||
|
||||
let db_url = config::DATABASE_URL.clone();
|
||||
let output = tokio::process::Command::new(&*python_path)
|
||||
.arg(&agent_script)
|
||||
.arg("--file-uuid")
|
||||
.arg(&file_uuid)
|
||||
.env("DATABASE_URL", &db_url)
|
||||
.env("DATABASE_SCHEMA", &*config::DATABASE_SCHEMA)
|
||||
.output()
|
||||
.await;
|
||||
|
||||
match output {
|
||||
Ok(o) => {
|
||||
if o.status.success() {
|
||||
let out = String::from_utf8_lossy(&o.stdout);
|
||||
Json(TmdbPrefetchResponse {
|
||||
success: true,
|
||||
file_uuid,
|
||||
message: out.lines().last().unwrap_or("OK").to_string(),
|
||||
cache_path: Some(cache_path),
|
||||
})
|
||||
} else {
|
||||
let stderr = String::from_utf8_lossy(&o.stderr);
|
||||
Json(TmdbPrefetchResponse {
|
||||
success: false,
|
||||
file_uuid,
|
||||
message: stderr.to_string(),
|
||||
cache_path: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
Err(e) => Json(TmdbPrefetchResponse {
|
||||
success: false,
|
||||
file_uuid,
|
||||
message: format!("Failed to run tmdb_agent.py: {}", e),
|
||||
cache_path: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
async fn tmdb_probe_handler(
|
||||
Path(params): Path<FileUuidParam>,
|
||||
State(state): State<AppState>,
|
||||
) -> Result<Json<TmdbProbeResponse>, (StatusCode, Json<serde_json::Value>)> {
|
||||
let file_uuid = params.file_uuid;
|
||||
|
||||
// Verify file exists
|
||||
let file_exists: bool = sqlx::query_scalar(
|
||||
&format!("SELECT COUNT(*) > 0 FROM {} WHERE file_uuid = $1", crate::core::db::schema::table_name("videos"))
|
||||
)
|
||||
.bind(&file_uuid)
|
||||
.fetch_one(state.db.pool())
|
||||
.await
|
||||
.unwrap_or(false);
|
||||
|
||||
if !file_exists {
|
||||
return Err((StatusCode::NOT_FOUND, Json(serde_json::json!({
|
||||
"error": "Video not found", "file_uuid": file_uuid
|
||||
}))));
|
||||
}
|
||||
|
||||
match tmdb::probe::probe_from_cache(&state.db, &file_uuid).await {
|
||||
Ok(result) => Ok(Json(TmdbProbeResponse {
|
||||
success: true,
|
||||
file_uuid,
|
||||
tmdb_id: Some(result.tmdb_id),
|
||||
movie_title: Some(result.title),
|
||||
cast_count: Some(result.cast_count),
|
||||
identities_created: Some(result.identities_created),
|
||||
message: format!(
|
||||
"Created/updated {} identities for movie ID {}",
|
||||
result.identities_created, result.tmdb_id
|
||||
),
|
||||
})),
|
||||
Err(e) => {
|
||||
let msg = e.to_string();
|
||||
if msg.contains("not found") {
|
||||
Ok(Json(TmdbProbeResponse {
|
||||
success: false,
|
||||
file_uuid,
|
||||
tmdb_id: None,
|
||||
movie_title: None,
|
||||
cast_count: None,
|
||||
identities_created: None,
|
||||
message: "No TMDb cache found. Run tmdb-prefetch first.".to_string(),
|
||||
}))
|
||||
} else {
|
||||
Err((StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({
|
||||
"error": msg, "file_uuid": file_uuid
|
||||
}))))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn tmdb_resource_status(
|
||||
State(state): State<AppState>,
|
||||
) -> Json<TmdbResourceResponse> {
|
||||
let status = tmdb::status::quick_status();
|
||||
let identities_seeded = tmdb::status::count_tmdb_identities(state.db.pool())
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
let identities_with_embedding = tmdb::status::count_tmdb_identities_with_embedding(state.db.pool())
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
let cache_files = tmdb::status::count_cache_files();
|
||||
|
||||
Json(TmdbResourceResponse {
|
||||
success: true,
|
||||
status,
|
||||
identities_seeded,
|
||||
identities_with_embedding,
|
||||
cache_files,
|
||||
operations: vec![
|
||||
TmdbOperation {
|
||||
method: "GET".to_string(),
|
||||
path: "/api/v1/resource/tmdb".to_string(),
|
||||
description: "TMDb resource status".to_string(),
|
||||
},
|
||||
TmdbOperation {
|
||||
method: "POST".to_string(),
|
||||
path: "/api/v1/resource/tmdb/check".to_string(),
|
||||
description: "Ping TMDb API health".to_string(),
|
||||
},
|
||||
TmdbOperation {
|
||||
method: "POST".to_string(),
|
||||
path: "/api/v1/agents/tmdb/prefetch".to_string(),
|
||||
description: "Fetch TMDb data and cache locally".to_string(),
|
||||
},
|
||||
TmdbOperation {
|
||||
method: "POST".to_string(),
|
||||
path: "/api/v1/file/:file_uuid/tmdb-probe".to_string(),
|
||||
description: "Read cache and create identities".to_string(),
|
||||
},
|
||||
],
|
||||
})
|
||||
}
|
||||
|
||||
async fn tmdb_resource_check() -> Json<TmdbCheckResponse> {
|
||||
let status = tmdb::status::check_tmdb_api().await;
|
||||
Json(TmdbCheckResponse {
|
||||
success: status.api_reachable.unwrap_or(false) && status.api_key_configured,
|
||||
status,
|
||||
})
|
||||
}
|
||||
@@ -12,7 +12,7 @@ use crate::core::db::PostgresDb;
|
||||
pub fn trace_agent_routes() -> Router<crate::api::server::AppState> {
|
||||
Router::new()
|
||||
.route(
|
||||
"/api/v1/file/:file_uuid/face_trace/sortby",
|
||||
"/api/v1/file/:file_uuid/traces",
|
||||
post(list_traces_sorted),
|
||||
)
|
||||
.route(
|
||||
|
||||
53
src/core/auth/jwt.rs
Normal file
53
src/core/auth/jwt.rs
Normal file
@@ -0,0 +1,53 @@
|
||||
use anyhow::{Context, Result};
|
||||
use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::core::config::JWT_SECRET;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Claims {
|
||||
pub sub: String,
|
||||
pub exp: usize,
|
||||
pub iat: usize,
|
||||
pub jti: String,
|
||||
pub role: String,
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
pub fn create_jwt(user_id: i32, username: &str, role: &str) -> Result<String> {
|
||||
let now = chrono::Utc::now();
|
||||
let exp = (now + chrono::Duration::hours(1)).timestamp() as usize;
|
||||
let iat = now.timestamp() as usize;
|
||||
|
||||
let claims = Claims {
|
||||
sub: user_id.to_string(),
|
||||
exp,
|
||||
iat,
|
||||
jti: Uuid::new_v4().to_string(),
|
||||
role: role.to_string(),
|
||||
name: username.to_string(),
|
||||
};
|
||||
|
||||
encode(
|
||||
&Header::default(),
|
||||
&claims,
|
||||
&EncodingKey::from_secret(JWT_SECRET.as_bytes()),
|
||||
)
|
||||
.context("Failed to encode JWT")
|
||||
}
|
||||
|
||||
pub fn verify_jwt(token: &str) -> Result<Claims> {
|
||||
let token_data = decode::<Claims>(
|
||||
token,
|
||||
&DecodingKey::from_secret(JWT_SECRET.as_bytes()),
|
||||
&Validation::default(),
|
||||
)
|
||||
.context("Failed to decode JWT")?;
|
||||
|
||||
Ok(token_data.claims)
|
||||
}
|
||||
|
||||
pub fn is_jwt(token: &str) -> bool {
|
||||
token.starts_with("eyJ") && token.split('.').count() == 3
|
||||
}
|
||||
2
src/core/auth/mod.rs
Normal file
2
src/core/auth/mod.rs
Normal file
@@ -0,0 +1,2 @@
|
||||
pub mod jwt;
|
||||
pub mod password;
|
||||
41
src/core/auth/password.rs
Normal file
41
src/core/auth/password.rs
Normal file
@@ -0,0 +1,41 @@
|
||||
use anyhow::Result;
|
||||
use argon2::{
|
||||
password_hash::{rand_core::OsRng, PasswordHash, PasswordHasher, PasswordVerifier, SaltString},
|
||||
Argon2,
|
||||
};
|
||||
|
||||
pub fn hash_password(password: &str) -> Result<String> {
|
||||
let salt = SaltString::generate(&mut OsRng);
|
||||
let hash = Argon2::default()
|
||||
.hash_password(password.as_bytes(), &salt)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to hash password: {}", e))?;
|
||||
Ok(hash.to_string())
|
||||
}
|
||||
|
||||
pub fn verify_password(password: &str, hash: &str) -> bool {
|
||||
let parsed = match PasswordHash::new(hash) {
|
||||
Ok(p) => p,
|
||||
Err(_) => return false,
|
||||
};
|
||||
Argon2::default()
|
||||
.verify_password(password.as_bytes(), &parsed)
|
||||
.is_ok()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_hash_and_verify() {
|
||||
let password = "test_password_123";
|
||||
let hash = hash_password(password).unwrap();
|
||||
assert!(verify_password(password, &hash));
|
||||
assert!(!verify_password("wrong_password", &hash));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_verify_fails_on_bad_hash() {
|
||||
assert!(!verify_password("test", "not_a_valid_hash"));
|
||||
}
|
||||
}
|
||||
@@ -191,6 +191,14 @@ pub mod llm {
|
||||
});
|
||||
}
|
||||
|
||||
pub static SFTPGO_BASE_URL: Lazy<String> = Lazy::new(|| {
|
||||
env::var("SFTPGO_BASE_URL").unwrap_or_else(|_| "http://127.0.0.1:8080".to_string())
|
||||
});
|
||||
|
||||
pub static JWT_SECRET: Lazy<String> = Lazy::new(|| {
|
||||
env::var("JWT_SECRET").unwrap_or_else(|_| "momentry_default_jwt_secret_change_me".to_string())
|
||||
});
|
||||
|
||||
pub mod tmdb {
|
||||
use super::*;
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -344,7 +344,7 @@ impl RedisClient {
|
||||
) -> Result<()> {
|
||||
let mut conn = self.get_conn_internal().await?;
|
||||
let prefix = REDIS_KEY_PREFIX.as_str();
|
||||
let key = format!("{}worker:job:{}", prefix, uuid);
|
||||
let key = format!("{}job:{}", prefix, uuid);
|
||||
|
||||
let _: Option<String> = conn
|
||||
.hset_multiple(
|
||||
@@ -379,7 +379,7 @@ impl RedisClient {
|
||||
) -> Result<()> {
|
||||
let mut conn = self.get_conn_internal().await?;
|
||||
let prefix = REDIS_KEY_PREFIX.as_str();
|
||||
let key = format!("{}worker:job:{}:processor:{}", prefix, uuid, processor);
|
||||
let key = format!("{}job:{}:processor:{}", prefix, uuid, processor);
|
||||
|
||||
let now = chrono::Utc::now().to_rfc3339();
|
||||
|
||||
@@ -409,7 +409,7 @@ impl RedisClient {
|
||||
pub async fn get_worker_job_status(&self, uuid: &str) -> Result<Option<WorkerJobStatus>> {
|
||||
let mut conn = self.get_conn_internal().await?;
|
||||
let prefix = REDIS_KEY_PREFIX.as_str();
|
||||
let key = format!("{}worker:job:{}", prefix, uuid);
|
||||
let key = format!("{}job:{}", prefix, uuid);
|
||||
|
||||
let exists: bool = conn.exists(&key).await?;
|
||||
if !exists {
|
||||
@@ -438,12 +438,12 @@ impl RedisClient {
|
||||
let mut conn = self.get_conn_internal().await?;
|
||||
let prefix = REDIS_KEY_PREFIX.as_str();
|
||||
|
||||
let key = format!("{}worker:job:{}", prefix, uuid);
|
||||
let key = format!("{}job:{}", prefix, uuid);
|
||||
let _: i32 = conn.del(&key).await?;
|
||||
|
||||
let processor_types = ["asr", "cut", "yolo", "ocr", "face", "pose", "asrx"];
|
||||
for ptype in processor_types {
|
||||
let proc_key = format!("{}worker:job:{}:processor:{}", prefix, uuid, ptype);
|
||||
let proc_key = format!("{}job:{}:processor:{}", prefix, uuid, ptype);
|
||||
let _: i32 = conn.del(&proc_key).await?;
|
||||
}
|
||||
|
||||
@@ -453,11 +453,11 @@ impl RedisClient {
|
||||
pub async fn get_all_worker_jobs(&self) -> Result<Vec<WorkerJobInfo>> {
|
||||
let mut conn = self.get_conn_internal().await?;
|
||||
let prefix = REDIS_KEY_PREFIX.as_str();
|
||||
let keys: Vec<String> = conn.keys(format!("{}worker:job:*", prefix)).await?;
|
||||
let keys: Vec<String> = conn.keys(format!("{}job:*", prefix)).await?;
|
||||
|
||||
let mut jobs = Vec::new();
|
||||
for key in keys {
|
||||
let uuid = key.replace(&format!("{}worker:job:", prefix), "");
|
||||
let uuid = key.replace(&format!("{}job:", prefix), "");
|
||||
if let Some(status) = self.get_worker_job_status(&uuid).await? {
|
||||
jobs.push(WorkerJobInfo {
|
||||
uuid,
|
||||
@@ -517,6 +517,10 @@ pub struct ProgressData {
|
||||
pub message: Option<String>,
|
||||
pub current: Option<i32>,
|
||||
pub total: Option<i32>,
|
||||
#[serde(default)]
|
||||
pub output_count: Option<i32>,
|
||||
#[serde(default)]
|
||||
pub output_type: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
|
||||
@@ -43,7 +43,7 @@ impl Embedder {
|
||||
}
|
||||
|
||||
fn default_url() -> String {
|
||||
std::env::var("MOMENTRY_EMBED_URL").unwrap_or_else(|_| "http://localhost:11434".to_string())
|
||||
std::env::var("MOMENTRY_EMBED_URL").unwrap_or_else(|_| "http://localhost:11436".to_string())
|
||||
}
|
||||
|
||||
pub async fn embed_text(&self, text: &str) -> Result<Vec<f32>> {
|
||||
|
||||
1
src/core/identity/mod.rs
Normal file
1
src/core/identity/mod.rs
Normal file
@@ -0,0 +1 @@
|
||||
pub mod storage;
|
||||
513
src/core/identity/storage.rs
Normal file
513
src/core/identity/storage.rs
Normal file
@@ -0,0 +1,513 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::core::config::OUTPUT_DIR;
|
||||
use crate::core::db::PostgresDb;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct IdentityFile {
|
||||
pub version: u32,
|
||||
pub identity_uuid: String,
|
||||
pub name: String,
|
||||
pub identity_type: Option<String>,
|
||||
pub source: Option<String>,
|
||||
pub status: Option<String>,
|
||||
pub tmdb_id: Option<i32>,
|
||||
pub tmdb_profile: Option<String>,
|
||||
pub metadata: serde_json::Value,
|
||||
pub file_bindings: Vec<FileBinding>,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct FileBinding {
|
||||
pub file_uuid: String,
|
||||
pub trace_ids: Vec<i32>,
|
||||
pub face_count: i64,
|
||||
}
|
||||
|
||||
pub fn identities_root() -> PathBuf {
|
||||
PathBuf::from(&*OUTPUT_DIR).join("identities")
|
||||
}
|
||||
|
||||
pub fn identity_dir(uuid: &str) -> PathBuf {
|
||||
identities_root().join(uuid)
|
||||
}
|
||||
|
||||
pub fn identity_file_path(uuid: &str) -> PathBuf {
|
||||
identity_dir(uuid).join("identity.json")
|
||||
}
|
||||
|
||||
pub fn index_path() -> PathBuf {
|
||||
identities_root().join("_index.json")
|
||||
}
|
||||
|
||||
pub fn read_identity_file(uuid: &str) -> Result<IdentityFile> {
|
||||
let path = identity_file_path(uuid);
|
||||
let content = std::fs::read_to_string(&path)
|
||||
.with_context(|| format!("Identity file not found: {} ({})", uuid, path.display()))?;
|
||||
serde_json::from_str(&content)
|
||||
.with_context(|| format!("Invalid identity.json: {}", uuid))
|
||||
}
|
||||
|
||||
pub fn write_identity_file(file: &IdentityFile) -> Result<()> {
|
||||
let dir = identity_dir(&file.identity_uuid);
|
||||
std::fs::create_dir_all(&dir)
|
||||
.with_context(|| format!("Failed to create identity dir: {}", dir.display()))?;
|
||||
|
||||
let path = dir.join("identity.json");
|
||||
let json = serde_json::to_string_pretty(file)
|
||||
.with_context(|| format!("Failed to serialize identity: {}", file.identity_uuid))?;
|
||||
std::fs::write(&path, &json)
|
||||
.with_context(|| format!("Failed to write identity.json: {}", path.display()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn delete_identity_file(uuid: &str) -> Result<()> {
|
||||
let path = identity_file_path(uuid);
|
||||
if path.exists() {
|
||||
std::fs::remove_file(&path)
|
||||
.with_context(|| format!("Failed to delete identity.json: {}", path.display()))?;
|
||||
}
|
||||
let dir = identity_dir(uuid);
|
||||
if dir.exists() {
|
||||
std::fs::remove_dir(&dir).ok();
|
||||
}
|
||||
remove_from_index(uuid).ok();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn list_identity_uuids() -> Result<Vec<String>> {
|
||||
let root = identities_root();
|
||||
if !root.is_dir() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut uuids = Vec::new();
|
||||
for entry in std::fs::read_dir(&root)
|
||||
.with_context(|| format!("Failed to read identities dir: {}", root.display()))?
|
||||
{
|
||||
let entry = entry?;
|
||||
let name = entry.file_name().to_string_lossy().to_string();
|
||||
if entry.file_type().map(|t| t.is_dir()).unwrap_or(false)
|
||||
&& name.len() == 32
|
||||
&& name.chars().all(|c| c.is_ascii_hexdigit())
|
||||
{
|
||||
uuids.push(name);
|
||||
}
|
||||
}
|
||||
uuids.sort();
|
||||
Ok(uuids)
|
||||
}
|
||||
|
||||
pub fn count_identity_files() -> usize {
|
||||
list_identity_uuids().map(|v| v.len()).unwrap_or(0)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
struct IndexFile {
|
||||
version: u32,
|
||||
updated_at: String,
|
||||
entries: HashMap<String, String>,
|
||||
}
|
||||
|
||||
fn read_index_inner() -> Result<IndexFile> {
|
||||
let path = index_path();
|
||||
if !path.exists() {
|
||||
return Ok(IndexFile {
|
||||
version: 1,
|
||||
updated_at: chrono::Utc::now().to_rfc3339(),
|
||||
entries: HashMap::new(),
|
||||
});
|
||||
}
|
||||
let content = std::fs::read_to_string(&path)
|
||||
.with_context(|| format!("Failed to read index: {}", path.display()))?;
|
||||
serde_json::from_str(&content)
|
||||
.with_context(|| format!("Invalid _index.json: {}", path.display()))
|
||||
}
|
||||
|
||||
pub fn read_index() -> Result<HashMap<String, String>> {
|
||||
read_index_inner().map(|idx| idx.entries)
|
||||
}
|
||||
|
||||
pub fn update_index(uuid: &str, name: &str) -> Result<()> {
|
||||
let mut idx = read_index_inner()?;
|
||||
idx.entries.insert(uuid.to_string(), name.to_string());
|
||||
idx.updated_at = chrono::Utc::now().to_rfc3339();
|
||||
let root = identities_root();
|
||||
std::fs::create_dir_all(&root)?;
|
||||
let json = serde_json::to_string_pretty(&idx)?;
|
||||
std::fs::write(index_path(), &json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn remove_from_index(uuid: &str) -> Result<()> {
|
||||
let mut idx = read_index_inner()?;
|
||||
idx.entries.remove(uuid);
|
||||
idx.updated_at = chrono::Utc::now().to_rfc3339();
|
||||
let json = serde_json::to_string_pretty(&idx)?;
|
||||
std::fs::write(index_path(), &json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn rebuild_index() -> Result<usize> {
|
||||
let uuids = list_identity_uuids()?;
|
||||
let mut entries = HashMap::new();
|
||||
for uuid in &uuids {
|
||||
match read_identity_file(uuid) {
|
||||
Ok(file) => {
|
||||
entries.insert(uuid.clone(), file.name);
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("[identity-storage] Skipping {} in index rebuild: {}", uuid, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
let idx = IndexFile {
|
||||
version: 1,
|
||||
updated_at: chrono::Utc::now().to_rfc3339(),
|
||||
entries,
|
||||
};
|
||||
let root = identities_root();
|
||||
std::fs::create_dir_all(&root)?;
|
||||
let json = serde_json::to_string_pretty(&idx)?;
|
||||
std::fs::write(index_path(), &json)?;
|
||||
Ok(uuids.len())
|
||||
}
|
||||
|
||||
pub async fn save_identity_file_by_pool(pool: &sqlx::PgPool, uuid: &str) -> Result<()> {
|
||||
let identity_table = crate::core::db::schema::table_name("identities");
|
||||
let fd_table = crate::core::db::schema::table_name("face_detections");
|
||||
|
||||
let clean = uuid.replace('-', "");
|
||||
let record = sqlx::query_as::<_, crate::core::db::IdentityDetailRecord>(
|
||||
&format!(
|
||||
"SELECT id, uuid::text, name, identity_type, source, status, metadata, reference_data, \
|
||||
NULL::real[] as voice_embedding, NULL::real[] as identity_embedding, \
|
||||
face_embedding::real[] as face_embedding, \
|
||||
tmdb_id, tmdb_profile, created_at::timestamptz as created_at, NULL::timestamptz as updated_at \
|
||||
FROM {} WHERE REPLACE(uuid::text, '-', '') = $1",
|
||||
identity_table
|
||||
)
|
||||
)
|
||||
.bind(&clean)
|
||||
.fetch_optional(pool)
|
||||
.await?
|
||||
.with_context(|| format!("Identity not found in DB: {}", uuid))?;
|
||||
|
||||
let identity_uuid = record.uuid.clone();
|
||||
|
||||
let binding_rows = sqlx::query_as::<_, (String, Vec<i32>, i64)>(
|
||||
&format!(
|
||||
"SELECT fd.file_uuid, COALESCE(array_agg(DISTINCT fd.trace_id) FILTER (WHERE fd.trace_id IS NOT NULL), '{{}}'::int[]), COUNT(*)::bigint \
|
||||
FROM {} fd WHERE fd.identity_id = $1 GROUP BY fd.file_uuid ORDER BY fd.file_uuid",
|
||||
fd_table
|
||||
)
|
||||
)
|
||||
.bind(record.id)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
|
||||
let file_bindings: Vec<FileBinding> = binding_rows
|
||||
.into_iter()
|
||||
.map(|(fu, tids, cnt)| FileBinding {
|
||||
file_uuid: fu,
|
||||
trace_ids: tids,
|
||||
face_count: cnt,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let fmt_time = |dt: Option<chrono::DateTime<chrono::Utc>>| -> String {
|
||||
dt.map(|d| d.to_rfc3339())
|
||||
.unwrap_or_else(|| chrono::Utc::now().to_rfc3339())
|
||||
};
|
||||
|
||||
let file = IdentityFile {
|
||||
version: 1,
|
||||
identity_uuid,
|
||||
name: record.name,
|
||||
identity_type: record.identity_type,
|
||||
source: record.source,
|
||||
status: record.status,
|
||||
tmdb_id: record.tmdb_id,
|
||||
tmdb_profile: record.tmdb_profile,
|
||||
metadata: record.metadata,
|
||||
file_bindings,
|
||||
created_at: fmt_time(record.created_at),
|
||||
updated_at: fmt_time(record.updated_at),
|
||||
};
|
||||
|
||||
write_identity_file(&file)?;
|
||||
update_index(&file.identity_uuid, &file.name)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn list_identity_uuids_at(base: &std::path::Path) -> Result<Vec<String>> {
|
||||
let root = base.join("identities");
|
||||
if !root.is_dir() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut uuids = Vec::new();
|
||||
for entry in std::fs::read_dir(&root)? {
|
||||
let entry = entry?;
|
||||
let name = entry.file_name().to_string_lossy().to_string();
|
||||
if entry.file_type().map(|t| t.is_dir()).unwrap_or(false)
|
||||
&& name.len() == 32
|
||||
&& name.chars().all(|c| c.is_ascii_hexdigit())
|
||||
{
|
||||
uuids.push(name);
|
||||
}
|
||||
}
|
||||
uuids.sort();
|
||||
Ok(uuids)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn identity_dir_at(base: &std::path::Path, uuid: &str) -> std::path::PathBuf {
|
||||
base.join("identities").join(uuid)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn identity_file_path_at(base: &std::path::Path, uuid: &str) -> std::path::PathBuf {
|
||||
identity_dir_at(base, uuid).join("identity.json")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn index_path_at(base: &std::path::Path) -> std::path::PathBuf {
|
||||
base.join("identities").join("_index.json")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn read_identity_file_at(base: &std::path::Path, uuid: &str) -> Result<IdentityFile> {
|
||||
let path = identity_file_path_at(base, uuid);
|
||||
let content = std::fs::read_to_string(&path)?;
|
||||
serde_json::from_str(&content).map_err(Into::into)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn write_identity_file_at(base: &std::path::Path, file: &IdentityFile) -> Result<()> {
|
||||
let dir = identity_dir_at(base, &file.identity_uuid);
|
||||
std::fs::create_dir_all(&dir)?;
|
||||
let json = serde_json::to_string_pretty(file)?;
|
||||
std::fs::write(dir.join("identity.json"), &json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn update_index_at(base: &std::path::Path, uuid: &str, name: &str) -> Result<()> {
|
||||
use std::collections::HashMap;
|
||||
let index_path = index_path_at(base);
|
||||
let mut entries: HashMap<String, String> = if index_path.exists() {
|
||||
let content = std::fs::read_to_string(&index_path)?;
|
||||
let v: serde_json::Value = serde_json::from_str(&content).unwrap_or_default();
|
||||
v["entries"].as_object()
|
||||
.map(|obj| obj.iter().map(|(k, v)| (k.clone(), v.as_str().unwrap_or("").to_string())).collect())
|
||||
.unwrap_or_default()
|
||||
} else {
|
||||
HashMap::new()
|
||||
};
|
||||
entries.insert(uuid.to_string(), name.to_string());
|
||||
std::fs::create_dir_all(base.join("identities"))?;
|
||||
let json = serde_json::to_string_pretty(&serde_json::json!({
|
||||
"version": 1, "updated_at": chrono::Utc::now().to_rfc3339(), "entries": entries
|
||||
}))?;
|
||||
std::fs::write(&index_path, &json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn save_identity_file(db: &PostgresDb, uuid: &str) -> Result<()> {
|
||||
let record = db.get_identity_by_uuid(uuid).await?
|
||||
.with_context(|| format!("Identity not found in DB: {}", uuid))?;
|
||||
|
||||
let identity_uuid = record.uuid.clone();
|
||||
|
||||
let binding_rows = sqlx::query_as::<_, (String, Vec<i32>, i64)>(
|
||||
"SELECT fd.file_uuid, COALESCE(array_agg(DISTINCT fd.trace_id) FILTER (WHERE fd.trace_id IS NOT NULL), '{}'::int[]), COUNT(*)::bigint \
|
||||
FROM face_detections fd \
|
||||
WHERE fd.identity_id = $1 \
|
||||
GROUP BY fd.file_uuid \
|
||||
ORDER BY fd.file_uuid"
|
||||
)
|
||||
.bind(record.id)
|
||||
.fetch_all(db.pool())
|
||||
.await
|
||||
.with_context(|| format!("Failed to query bindings for identity: {}", identity_uuid))?;
|
||||
|
||||
let file_bindings: Vec<FileBinding> = binding_rows
|
||||
.into_iter()
|
||||
.map(|(fu, tids, cnt)| FileBinding {
|
||||
file_uuid: fu,
|
||||
trace_ids: tids,
|
||||
face_count: cnt,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let fmt_time = |dt: Option<chrono::DateTime<chrono::Utc>>| -> String {
|
||||
dt.map(|d| d.to_rfc3339())
|
||||
.unwrap_or_else(|| chrono::Utc::now().to_rfc3339())
|
||||
};
|
||||
|
||||
let file = IdentityFile {
|
||||
version: 1,
|
||||
identity_uuid,
|
||||
name: record.name,
|
||||
identity_type: record.identity_type,
|
||||
source: record.source,
|
||||
status: record.status,
|
||||
tmdb_id: record.tmdb_id,
|
||||
tmdb_profile: record.tmdb_profile,
|
||||
metadata: record.metadata,
|
||||
file_bindings,
|
||||
created_at: fmt_time(record.created_at),
|
||||
updated_at: fmt_time(record.updated_at),
|
||||
};
|
||||
|
||||
write_identity_file(&file)?;
|
||||
update_index(&file.identity_uuid, &file.name)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::path::Path;
|
||||
|
||||
fn sample_identity() -> IdentityFile {
|
||||
IdentityFile {
|
||||
version: 1,
|
||||
identity_uuid: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".to_string(),
|
||||
name: "Test Person".to_string(),
|
||||
identity_type: Some("people".to_string()),
|
||||
source: Some("tmdb".to_string()),
|
||||
status: Some("confirmed".to_string()),
|
||||
tmdb_id: Some(112),
|
||||
tmdb_profile: Some("https://image.tmdb.org/t/p/w185/test.jpg".to_string()),
|
||||
metadata: serde_json::json!({"tmdb_character": "Test Role"}),
|
||||
file_bindings: vec![FileBinding {
|
||||
file_uuid: "ffffffffffffffffffffffffffffffff".to_string(),
|
||||
trace_ids: vec![1, 2, 3],
|
||||
face_count: 5,
|
||||
}],
|
||||
created_at: "2026-05-16T00:00:00+00:00".to_string(),
|
||||
updated_at: "2026-05-16T01:00:00+00:00".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_serde_roundtrip() {
|
||||
let file = sample_identity();
|
||||
let json = serde_json::to_string_pretty(&file).unwrap();
|
||||
let parsed: IdentityFile = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed.name, "Test Person");
|
||||
assert_eq!(parsed.identity_uuid, "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa");
|
||||
assert_eq!(parsed.tmdb_id, Some(112));
|
||||
assert_eq!(parsed.file_bindings.len(), 1);
|
||||
assert_eq!(parsed.file_bindings[0].face_count, 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_identity_dir_path() {
|
||||
let uuid = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
|
||||
let p = identity_dir(uuid);
|
||||
assert!(p.to_string_lossy().ends_with(&format!("identities/{}", uuid)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_identity_file_path() {
|
||||
let uuid = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
|
||||
let p = identity_file_path(uuid);
|
||||
assert!(p.to_string_lossy().ends_with("identity.json"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_index_path() {
|
||||
let p = index_path();
|
||||
assert!(p.to_string_lossy().ends_with("_index.json"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_identity_dir_at() {
|
||||
let base = Path::new("/tmp/test_base");
|
||||
let uuid = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb";
|
||||
let p = identity_dir_at(base, uuid);
|
||||
assert_eq!(p, Path::new("/tmp/test_base/identities/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_identity_file_path_at() {
|
||||
let base = Path::new("/tmp/test_base");
|
||||
let uuid = "cccccccccccccccccccccccccccccccc";
|
||||
let p = identity_file_path_at(base, uuid);
|
||||
assert_eq!(
|
||||
p,
|
||||
Path::new("/tmp/test_base/identities/cccccccccccccccccccccccccccccccc/identity.json")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_write_then_read_identity_file_at() {
|
||||
let tmp = std::env::temp_dir().join("momentry_test_write_read");
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
let base = &tmp;
|
||||
|
||||
let file = sample_identity();
|
||||
write_identity_file_at(base, &file).unwrap();
|
||||
|
||||
let read = read_identity_file_at(base, &file.identity_uuid).unwrap();
|
||||
assert_eq!(read.name, file.name);
|
||||
assert_eq!(read.source, file.source);
|
||||
assert_eq!(read.tmdb_id, file.tmdb_id);
|
||||
assert_eq!(read.file_bindings[0].face_count, file.file_bindings[0].face_count);
|
||||
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_update_and_read_index_at() {
|
||||
let tmp = std::env::temp_dir().join("momentry_test_index");
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
let base = &tmp;
|
||||
|
||||
update_index_at(base, "aaa", "Alice").unwrap();
|
||||
update_index_at(base, "bbb", "Bob").unwrap();
|
||||
|
||||
let idx_path = index_path_at(base);
|
||||
let content = std::fs::read_to_string(&idx_path).unwrap();
|
||||
let parsed: serde_json::Value = serde_json::from_str(&content).unwrap();
|
||||
let entries = parsed["entries"].as_object().unwrap();
|
||||
assert_eq!(entries.len(), 2);
|
||||
assert_eq!(entries["aaa"], "Alice");
|
||||
assert_eq!(entries["bbb"], "Bob");
|
||||
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_list_identity_uuids_at() {
|
||||
let tmp = std::env::temp_dir().join("momentry_test_list");
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
let base = &tmp;
|
||||
|
||||
std::fs::create_dir_all(base.join("identities").join("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")).unwrap();
|
||||
std::fs::create_dir_all(base.join("identities").join("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb")).unwrap();
|
||||
std::fs::create_dir_all(base.join("identities").join("cccccccccccccccccccccccccccccccc")).unwrap();
|
||||
std::fs::create_dir_all(base.join("identities").join("not_a_uuid")).unwrap();
|
||||
std::fs::create_dir_all(base.join("identities").join("short")).unwrap();
|
||||
|
||||
let uuids = list_identity_uuids_at(base).unwrap();
|
||||
assert_eq!(uuids.len(), 3);
|
||||
assert!(uuids.contains(&"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".to_string()));
|
||||
assert!(uuids.contains(&"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".to_string()));
|
||||
assert!(uuids.contains(&"cccccccccccccccccccccccccccccccc".to_string()));
|
||||
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
}
|
||||
@@ -1,10 +1,12 @@
|
||||
pub mod api_key;
|
||||
pub mod auth;
|
||||
pub mod cache;
|
||||
pub mod chunk;
|
||||
pub mod config;
|
||||
pub mod db;
|
||||
pub mod embedding;
|
||||
pub mod frame_cache;
|
||||
pub mod identity;
|
||||
pub mod ingestion;
|
||||
pub mod llm;
|
||||
pub mod overlay;
|
||||
|
||||
@@ -84,9 +84,9 @@ fn load_checksums(scripts_dir: &PathBuf) -> HashMap<String, String> {
|
||||
pub fn validate_python_env() -> Result<()> {
|
||||
let python_path = std::env::var("MOMENTRY_PYTHON_PATH")
|
||||
.unwrap_or_else(|_| "/opt/homebrew/bin/python3.11".to_string());
|
||||
let venv_python = PathBuf::from(&python_path);
|
||||
let python_bin = PathBuf::from(&python_path);
|
||||
|
||||
if !venv_python.exists() {
|
||||
if !python_bin.exists() {
|
||||
anyhow::bail!(
|
||||
"Python not found at {} (set MOMENTRY_PYTHON_PATH env var)",
|
||||
python_path
|
||||
@@ -95,7 +95,7 @@ pub fn validate_python_env() -> Result<()> {
|
||||
|
||||
let rt = tokio::runtime::Runtime::new()?;
|
||||
let output = rt
|
||||
.block_on(async { Command::new(&venv_python).arg("--version").output().await })
|
||||
.block_on(async { Command::new(&python_bin).arg("--version").output().await })
|
||||
.context("Failed to run Python")?;
|
||||
|
||||
if !output.status.success() {
|
||||
@@ -124,7 +124,7 @@ pub fn validate_python_env() -> Result<()> {
|
||||
}
|
||||
|
||||
pub struct PythonExecutor {
|
||||
venv_python: PathBuf,
|
||||
python_path: PathBuf,
|
||||
scripts_dir: PathBuf,
|
||||
checksums: HashMap<String, String>,
|
||||
}
|
||||
@@ -139,10 +139,10 @@ impl PythonExecutor {
|
||||
manifest.join("scripts").to_string_lossy().to_string()
|
||||
});
|
||||
|
||||
let venv_python = PathBuf::from(&python_path);
|
||||
let python_bin = PathBuf::from(&python_path);
|
||||
let scripts_path = PathBuf::from(&scripts_dir);
|
||||
|
||||
if !venv_python.exists() {
|
||||
if !python_bin.exists() {
|
||||
anyhow::bail!(
|
||||
"Python not found at {} (set MOMENTRY_PYTHON_PATH env var)",
|
||||
python_path
|
||||
@@ -160,7 +160,7 @@ impl PythonExecutor {
|
||||
let checksums = load_checksums(&scripts_path);
|
||||
|
||||
Ok(Self {
|
||||
venv_python,
|
||||
python_path: python_bin,
|
||||
scripts_dir: scripts_path,
|
||||
checksums,
|
||||
})
|
||||
@@ -201,7 +201,7 @@ impl PythonExecutor {
|
||||
let rt = tokio::runtime::Runtime::new()?;
|
||||
let output = rt
|
||||
.block_on(async {
|
||||
Command::new(&self.venv_python)
|
||||
Command::new(&self.python_path)
|
||||
.arg("--version")
|
||||
.output()
|
||||
.await
|
||||
@@ -251,7 +251,7 @@ impl PythonExecutor {
|
||||
}
|
||||
}
|
||||
|
||||
let mut cmd = Command::new(&self.venv_python);
|
||||
let mut cmd = Command::new(&self.python_path);
|
||||
cmd.arg(&script_path);
|
||||
|
||||
for arg in args {
|
||||
@@ -467,7 +467,7 @@ impl PythonExecutor {
|
||||
}
|
||||
|
||||
pub fn python_path(&self) -> &PathBuf {
|
||||
&self.venv_python
|
||||
&self.python_path
|
||||
}
|
||||
}
|
||||
|
||||
@@ -482,11 +482,11 @@ mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_python_executor_new_with_venv() {
|
||||
fn test_python_executor_new() {
|
||||
let executor = PythonExecutor::new();
|
||||
assert!(
|
||||
executor.is_ok(),
|
||||
"PythonExecutor should create successfully with venv"
|
||||
"PythonExecutor should create successfully"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -499,10 +499,6 @@ mod tests {
|
||||
"Python path should exist: {:?}",
|
||||
python_path
|
||||
);
|
||||
assert!(
|
||||
python_path.to_string_lossy().contains("venv"),
|
||||
"Should be in venv"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -284,10 +284,21 @@ pub async fn process_visual_chunk_advanced(
|
||||
});
|
||||
}
|
||||
|
||||
let yolo_path = uuid.map(|u| {
|
||||
std::path::PathBuf::from(crate::core::config::OUTPUT_DIR.as_str())
|
||||
.join(format!("{}.yolo.json", u))
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
});
|
||||
let args: &[&str] = if let Some(ref yp) = yolo_path {
|
||||
&[video_path, output_path, "--yolo-result", yp]
|
||||
} else {
|
||||
&[video_path, output_path]
|
||||
};
|
||||
let result = match executor
|
||||
.run(
|
||||
"visual_chunk_processor.py",
|
||||
&[video_path, output_path],
|
||||
args,
|
||||
uuid,
|
||||
"VisualChunk",
|
||||
Some(VISUAL_CHUNK_TIMEOUT),
|
||||
|
||||
@@ -25,13 +25,11 @@ impl ThumbnailExtractor {
|
||||
.join("scripts")
|
||||
.join("thumbnail_extractor.py");
|
||||
|
||||
// 使用 venv 中的 Python,確保版本正確且隔離依賴
|
||||
let venv_python = Path::new(env!("CARGO_MANIFEST_DIR"))
|
||||
.join("venv")
|
||||
.join("bin")
|
||||
.join("python");
|
||||
let python_path = std::env::var("MOMENTRY_PYTHON_PATH")
|
||||
.unwrap_or_else(|_| "/opt/homebrew/bin/python3.11".to_string());
|
||||
let python_bin = Path::new(&python_path);
|
||||
|
||||
let output = Command::new(venv_python)
|
||||
let output = Command::new(python_bin)
|
||||
.arg(script_path)
|
||||
.arg(video_path)
|
||||
.arg(uuid)
|
||||
|
||||
262
src/core/tmdb/cache.rs
Normal file
262
src/core/tmdb/cache.rs
Normal file
@@ -0,0 +1,262 @@
|
||||
use std::path::PathBuf;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::core::config::OUTPUT_DIR;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TmdbCacheIdentity {
|
||||
pub identity_uuid: String,
|
||||
pub name: String,
|
||||
pub tmdb_id: u64,
|
||||
pub character: String,
|
||||
pub order: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TmdbCache {
|
||||
pub file_uuid: String,
|
||||
pub fetched_at: String,
|
||||
pub source: String,
|
||||
pub movie: TmdbMovie,
|
||||
pub cast_count: usize,
|
||||
pub identities_created: usize,
|
||||
#[serde(default)]
|
||||
pub identities: Vec<TmdbCacheIdentity>,
|
||||
#[serde(default)]
|
||||
pub cast: Vec<TmdbCastMember>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TmdbMovie {
|
||||
pub tmdb_id: u64,
|
||||
pub title: String,
|
||||
pub release_date: Option<String>,
|
||||
pub overview: Option<String>,
|
||||
pub poster_path: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TmdbCastMember {
|
||||
pub name: String,
|
||||
pub character: String,
|
||||
pub profile_path: Option<String>,
|
||||
pub order: u32,
|
||||
pub id: u64,
|
||||
// Person detail fields from /person/{id}
|
||||
pub biography: Option<String>,
|
||||
pub birthday: Option<String>,
|
||||
pub place_of_birth: Option<String>,
|
||||
#[serde(default)]
|
||||
pub also_known_as: Vec<String>,
|
||||
pub imdb_id: Option<String>,
|
||||
pub known_for_department: Option<String>,
|
||||
pub popularity: Option<f64>,
|
||||
pub deathday: Option<String>,
|
||||
pub gender: Option<i32>,
|
||||
pub homepage: Option<String>,
|
||||
}
|
||||
|
||||
pub fn tmdb_cache_path(file_uuid: &str) -> PathBuf {
|
||||
PathBuf::from(&*OUTPUT_DIR).join(format!("{}.tmdb.json", file_uuid))
|
||||
}
|
||||
|
||||
pub fn read_tmdb_cache(file_uuid: &str) -> Result<TmdbCache> {
|
||||
let path = tmdb_cache_path(file_uuid);
|
||||
if !path.exists() {
|
||||
anyhow::bail!("TMDb cache not found: {} (expected: {})", file_uuid, path.display());
|
||||
}
|
||||
let content = std::fs::read_to_string(&path)
|
||||
.with_context(|| format!("Failed to read TMDb cache: {}", path.display()))?;
|
||||
serde_json::from_str(&content)
|
||||
.map_err(|e| anyhow::anyhow!("Invalid TMDb cache JSON {}: {}", path.display(), e))
|
||||
}
|
||||
|
||||
pub fn write_tmdb_cache(cache: &TmdbCache) -> Result<()> {
|
||||
let path = tmdb_cache_path(&cache.file_uuid);
|
||||
let json = serde_json::to_string_pretty(cache)
|
||||
.with_context(|| format!("Failed to serialize TMDb cache: {}", cache.file_uuid))?;
|
||||
std::fs::write(&path, &json)
|
||||
.with_context(|| format!("Failed to write TMDb cache: {}", path.display()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn delete_tmdb_cache(file_uuid: &str) -> Result<()> {
|
||||
let path = tmdb_cache_path(file_uuid);
|
||||
if path.exists() {
|
||||
std::fs::remove_file(&path)
|
||||
.with_context(|| format!("Failed to delete TMDb cache: {}", path.display()))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn count_cache_files() -> usize {
|
||||
let dir = PathBuf::from(&*OUTPUT_DIR);
|
||||
match std::fs::read_dir(&dir) {
|
||||
Ok(entries) => entries
|
||||
.filter_map(|e| e.ok())
|
||||
.filter(|e| {
|
||||
e.file_name().to_string_lossy().ends_with(".tmdb.json")
|
||||
})
|
||||
.count(),
|
||||
Err(_) => 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn count_cache_files_at(base: &std::path::Path) -> usize {
|
||||
match std::fs::read_dir(base) {
|
||||
Ok(entries) => entries
|
||||
.filter_map(|e| e.ok())
|
||||
.filter(|e| e.file_name().to_string_lossy().ends_with(".tmdb.json"))
|
||||
.count(),
|
||||
Err(_) => 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn write_tmdb_cache_at(base: &std::path::Path, cache: &TmdbCache) -> Result<()> {
|
||||
std::fs::create_dir_all(base)?;
|
||||
let path = base.join(format!("{}.tmdb.json", cache.file_uuid));
|
||||
let json = serde_json::to_string_pretty(cache)?;
|
||||
std::fs::write(&path, &json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub fn read_tmdb_cache_at(base: &std::path::Path, file_uuid: &str) -> Result<TmdbCache> {
|
||||
let path = base.join(format!("{}.tmdb.json", file_uuid));
|
||||
if !path.exists() {
|
||||
anyhow::bail!("Cache not found");
|
||||
}
|
||||
let content = std::fs::read_to_string(&path)?;
|
||||
serde_json::from_str(&content).map_err(Into::into)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn sample_cache(file_uuid: &str) -> TmdbCache {
|
||||
TmdbCache {
|
||||
file_uuid: file_uuid.to_string(),
|
||||
fetched_at: "2026-05-16T12:00:00+00:00".to_string(),
|
||||
source: "agent".to_string(),
|
||||
movie: TmdbMovie {
|
||||
tmdb_id: 4808,
|
||||
title: "Charade".to_string(),
|
||||
release_date: Some("1963-12-05".to_string()),
|
||||
overview: Some("A romantic thriller...".to_string()),
|
||||
poster_path: Some("/abc.jpg".to_string()),
|
||||
},
|
||||
cast: vec![
|
||||
TmdbCastMember {
|
||||
name: "Cary Grant".to_string(),
|
||||
character: "Peter Joshua".to_string(),
|
||||
profile_path: Some("/cary.jpg".to_string()),
|
||||
order: 0,
|
||||
id: 112,
|
||||
biography: Some("Archibald Alec Leach...".to_string()),
|
||||
birthday: Some("1904-01-18".to_string()),
|
||||
place_of_birth: Some("Bristol, England, UK".to_string()),
|
||||
also_known_as: vec!["Archie Leach".to_string()],
|
||||
imdb_id: Some("nm0000026".to_string()),
|
||||
known_for_department: Some("Acting".to_string()),
|
||||
popularity: Some(28.3),
|
||||
deathday: Some("1986-11-29".to_string()),
|
||||
gender: Some(2),
|
||||
homepage: None,
|
||||
},
|
||||
TmdbCastMember {
|
||||
name: "Audrey Hepburn".to_string(),
|
||||
character: "Regina Lampert".to_string(),
|
||||
profile_path: Some("/audrey.jpg".to_string()),
|
||||
order: 1,
|
||||
id: 113,
|
||||
biography: Some("Audrey Kathleen Hepburn...".to_string()),
|
||||
birthday: Some("1929-05-04".to_string()),
|
||||
place_of_birth: Some("Ixelles, Belgium".to_string()),
|
||||
also_known_as: vec!["Edda van Heemstra".to_string()],
|
||||
imdb_id: Some("nm0000030".to_string()),
|
||||
known_for_department: Some("Acting".to_string()),
|
||||
popularity: Some(35.7),
|
||||
deathday: Some("1993-01-20".to_string()),
|
||||
gender: Some(1),
|
||||
homepage: None,
|
||||
},
|
||||
],
|
||||
cast_count: 20,
|
||||
identities_created: 0,
|
||||
identities: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cache_path_format() {
|
||||
let p = tmdb_cache_path("abcdef");
|
||||
assert!(p.to_string_lossy().ends_with("abcdef.tmdb.json"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_serde_roundtrip() {
|
||||
let cache = sample_cache("aaaaaaaa");
|
||||
let json = serde_json::to_string_pretty(&cache).unwrap();
|
||||
let parsed: TmdbCache = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed.file_uuid, "aaaaaaaa");
|
||||
assert_eq!(parsed.movie.title, "Charade");
|
||||
assert_eq!(parsed.cast.len(), 2);
|
||||
assert_eq!(parsed.cast[0].name, "Cary Grant");
|
||||
assert_eq!(parsed.movie.tmdb_id, 4808);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_write_then_read_cache_at() {
|
||||
let tmp = std::env::temp_dir().join("momentry_test_cache");
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
let base = &tmp;
|
||||
|
||||
let cache = sample_cache("bbbbbbbb");
|
||||
write_tmdb_cache_at(base, &cache).unwrap();
|
||||
|
||||
let read = read_tmdb_cache_at(base, "bbbbbbbb").unwrap();
|
||||
assert_eq!(read.movie.title, "Charade");
|
||||
assert_eq!(read.cast[1].id, 113);
|
||||
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_read_missing_cache_at_errors() {
|
||||
let tmp = std::env::temp_dir().join("momentry_test_missing");
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
let base = &tmp;
|
||||
|
||||
let result = read_tmdb_cache_at(base, "nonexistent");
|
||||
assert!(result.is_err());
|
||||
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_count_cache_files_at() {
|
||||
let tmp = std::env::temp_dir().join("momentry_test_count");
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
let base = &tmp;
|
||||
|
||||
assert_eq!(count_cache_files_at(base), 0);
|
||||
|
||||
let c1 = sample_cache("aaa");
|
||||
write_tmdb_cache_at(base, &c1).unwrap();
|
||||
assert_eq!(count_cache_files_at(base), 1);
|
||||
|
||||
let c2 = sample_cache("bbb");
|
||||
write_tmdb_cache_at(base, &c2).unwrap();
|
||||
assert_eq!(count_cache_files_at(base), 2);
|
||||
|
||||
std::fs::write(base.join("other.json"), "{}").unwrap();
|
||||
assert_eq!(count_cache_files_at(base), 2);
|
||||
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,5 @@
|
||||
pub mod cache;
|
||||
pub mod face_agent;
|
||||
pub mod ingest;
|
||||
pub mod probe;
|
||||
pub mod status;
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
use anyhow::{Context, Result};
|
||||
use serde::Deserialize;
|
||||
use std::collections::HashMap;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::core::config;
|
||||
@@ -8,11 +7,11 @@ use crate::core::db::PostgresDb;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TmdbSearchResult {
|
||||
results: Vec<TmdbMovie>,
|
||||
results: Vec<TmdbApiMovie>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TmdbMovie {
|
||||
struct TmdbApiMovie {
|
||||
id: u64,
|
||||
title: String,
|
||||
release_date: Option<String>,
|
||||
@@ -22,11 +21,11 @@ struct TmdbMovie {
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TmdbCredits {
|
||||
cast: Vec<TmdbCastMember>,
|
||||
cast: Vec<TmdbApiCastMember>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TmdbCastMember {
|
||||
struct TmdbApiCastMember {
|
||||
id: u64,
|
||||
name: String,
|
||||
character: String,
|
||||
@@ -54,6 +53,271 @@ fn extract_movie_name(filename: &str) -> Option<String> {
|
||||
Some(cleaned)
|
||||
}
|
||||
|
||||
pub async fn probe_from_cache(
|
||||
db: &PostgresDb,
|
||||
file_uuid: &str,
|
||||
) -> Result<TmdbProbeResult> {
|
||||
let cache = crate::core::tmdb::cache::read_tmdb_cache(file_uuid)?;
|
||||
if cache.identities.is_empty() && !cache.cast.is_empty() {
|
||||
return create_identities_from_data(db, file_uuid, &cache.movie, &cache.cast).await;
|
||||
}
|
||||
upsert_identities_from_disk(db, &cache, file_uuid).await
|
||||
}
|
||||
|
||||
async fn upsert_identities_from_disk(
|
||||
db: &PostgresDb,
|
||||
cache: &crate::core::tmdb::cache::TmdbCache,
|
||||
file_uuid: &str,
|
||||
) -> Result<TmdbProbeResult> {
|
||||
info!(
|
||||
"[TMDB] Upserting identities from disk for: {} (TMDB id={})",
|
||||
cache.movie.title, cache.movie.tmdb_id
|
||||
);
|
||||
|
||||
let mut identities_created = 0usize;
|
||||
for entry in &cache.identities {
|
||||
let path = crate::core::identity::storage::identity_file_path(&entry.identity_uuid);
|
||||
if !path.exists() {
|
||||
warn!("[TMDB] Identity file not found on disk: {}", path.display());
|
||||
continue;
|
||||
}
|
||||
match std::fs::read_to_string(&path) {
|
||||
Ok(content) => {
|
||||
match serde_json::from_str::<crate::core::identity::storage::IdentityFile>(&content) {
|
||||
Ok(identity_file) => {
|
||||
let identities_table = crate::core::db::schema::table_name("identities");
|
||||
let result = sqlx::query(&format!(
|
||||
"INSERT INTO {} (uuid, name, identity_type, source, status, tmdb_id, tmdb_profile, metadata) \
|
||||
VALUES ($1::uuid, $2, 'people', 'tmdb', 'confirmed', $3, $4, $5::jsonb) \
|
||||
ON CONFLICT (name) DO UPDATE SET \
|
||||
uuid = COALESCE({}.uuid, $1::uuid), \
|
||||
tmdb_id = COALESCE(EXCLUDED.tmdb_id, {}.tmdb_id), \
|
||||
tmdb_profile = COALESCE(EXCLUDED.tmdb_profile, {}.tmdb_profile), \
|
||||
metadata = {}.metadata || $5::jsonb",
|
||||
identities_table, identities_table, identities_table, identities_table, identities_table
|
||||
))
|
||||
.bind(&identity_file.identity_uuid)
|
||||
.bind(&identity_file.name)
|
||||
.bind(identity_file.tmdb_id)
|
||||
.bind(&identity_file.tmdb_profile)
|
||||
.bind(&identity_file.metadata)
|
||||
.execute(db.pool())
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(_) => {
|
||||
info!("[TMDB] Upserted identity: {} (uuid={})", identity_file.name, identity_file.identity_uuid);
|
||||
identities_created += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("[TMDB] Failed to upsert identity '{}': {}", identity_file.name, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("[TMDB] Failed to parse identity file {}: {}", path.display(), e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("[TMDB] Failed to read identity file {}: {}", path.display(), e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
drop_identities_cache(db, file_uuid, &cache.movie, identities_created).await;
|
||||
Ok(TmdbProbeResult {
|
||||
tmdb_id: cache.movie.tmdb_id,
|
||||
title: cache.movie.title.clone(),
|
||||
cast_count: cache.cast_count,
|
||||
identities_created,
|
||||
})
|
||||
}
|
||||
|
||||
async fn drop_identities_cache(
|
||||
db: &PostgresDb,
|
||||
file_uuid: &str,
|
||||
movie: &crate::core::tmdb::cache::TmdbMovie,
|
||||
identities_created: usize,
|
||||
) {
|
||||
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), '{{{}}}'::text[], $1::jsonb) \
|
||||
WHERE file_uuid = $2",
|
||||
videos_table, tmdb_label
|
||||
))
|
||||
.bind(serde_json::json!({
|
||||
"movie_id": movie.tmdb_id,
|
||||
"movie_title": movie.title,
|
||||
"release_date": movie.release_date,
|
||||
"poster": movie.poster_path,
|
||||
"cast_count": movie.tmdb_id,
|
||||
"identities_created": identities_created,
|
||||
}))
|
||||
.bind(file_uuid)
|
||||
.execute(db.pool())
|
||||
.await
|
||||
.ok();
|
||||
}
|
||||
|
||||
pub async fn create_identities_from_data(
|
||||
db: &PostgresDb,
|
||||
file_uuid: &str,
|
||||
movie: &crate::core::tmdb::cache::TmdbMovie,
|
||||
cast: &[crate::core::tmdb::cache::TmdbCastMember],
|
||||
) -> Result<TmdbProbeResult> {
|
||||
info!(
|
||||
"[TMDB] Creating identities for: {} (TMDB id={})",
|
||||
movie.title, movie.tmdb_id
|
||||
);
|
||||
|
||||
let identities_table = crate::core::db::schema::table_name("identities");
|
||||
let mut identities_created = 0usize;
|
||||
|
||||
for member in cast.iter() {
|
||||
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 metadata = serde_json::json!({
|
||||
"tmdb_character": member.character,
|
||||
"tmdb_cast_order": member.order,
|
||||
"tmdb_movie_id": movie.tmdb_id,
|
||||
"tmdb_movie_title": movie.title,
|
||||
"tmdb_biography": member.biography,
|
||||
"tmdb_birthday": member.birthday,
|
||||
"tmdb_place_of_birth": member.place_of_birth,
|
||||
"tmdb_aliases": member.also_known_as,
|
||||
"tmdb_imdb_id": member.imdb_id,
|
||||
"tmdb_department": member.known_for_department,
|
||||
"tmdb_popularity": member.popularity,
|
||||
"tmdb_deathday": member.deathday,
|
||||
"tmdb_gender": member.gender,
|
||||
"tmdb_homepage": member.homepage,
|
||||
});
|
||||
|
||||
let result = sqlx::query_as::<_, (uuid::Uuid,)>(&format!(
|
||||
"INSERT INTO {} (name, identity_type, source, status, tmdb_id, tmdb_profile, metadata) \
|
||||
VALUES ($1, 'people', 'tmdb', 'confirmed', $2, $3, $4::jsonb) \
|
||||
ON CONFLICT (name) DO UPDATE SET \
|
||||
tmdb_id = COALESCE(EXCLUDED.tmdb_id, {}.tmdb_id), \
|
||||
tmdb_profile = COALESCE(EXCLUDED.tmdb_profile, {}.tmdb_profile), \
|
||||
metadata = {}.metadata || $4::jsonb \
|
||||
RETURNING uuid",
|
||||
identities_table, identities_table, identities_table, identities_table
|
||||
))
|
||||
.bind(&member.name)
|
||||
.bind(member.id as i64)
|
||||
.bind(&profile_url)
|
||||
.bind(&metadata)
|
||||
.fetch_optional(db.pool())
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Some((identity_uuid,))) => {
|
||||
let uuid_str = identity_uuid.to_string().replace('-', "");
|
||||
info!(
|
||||
"[TMDB] Created/updated identity: {} as {} (uuid={})",
|
||||
member.name, member.character, uuid_str
|
||||
);
|
||||
identities_created += 1;
|
||||
if let Err(e) = crate::core::identity::storage::save_identity_file(db, &uuid_str).await {
|
||||
warn!("[TMDB] Failed to save identity file for {}: {}", member.name, e);
|
||||
}
|
||||
// Download and save TMDb profile image locally
|
||||
if let Some(url) = &profile_url {
|
||||
let dir = crate::core::identity::storage::identity_dir(&uuid_str);
|
||||
std::fs::create_dir_all(&dir).ok();
|
||||
let img_path = dir.join("profile.jpg");
|
||||
if !img_path.exists() {
|
||||
if let Ok(resp) = reqwest::get(url).await {
|
||||
if let Ok(bytes) = resp.bytes().await {
|
||||
std::fs::write(&img_path, &bytes).ok();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(None) => {
|
||||
warn!("[TMDB] INSERT returned no uuid for: {}", member.name);
|
||||
}
|
||||
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), '{{{}}}'::text[], $1::jsonb) \
|
||||
WHERE file_uuid = $2",
|
||||
videos_table, tmdb_label
|
||||
))
|
||||
.bind(serde_json::json!({
|
||||
"movie_id": movie.tmdb_id,
|
||||
"movie_title": movie.title,
|
||||
"release_date": movie.release_date,
|
||||
"poster": movie.poster_path,
|
||||
"cast_count": cast.len(),
|
||||
"identities_created": identities_created,
|
||||
}))
|
||||
.bind(file_uuid)
|
||||
.execute(db.pool())
|
||||
.await
|
||||
.ok();
|
||||
|
||||
info!(
|
||||
"[TMDB] Probe complete: {} cast members, {} identities created/updated",
|
||||
cast.len(),
|
||||
identities_created
|
||||
);
|
||||
|
||||
Ok(TmdbProbeResult {
|
||||
tmdb_id: movie.tmdb_id,
|
||||
title: movie.title.clone(),
|
||||
cast_count: cast.len(),
|
||||
identities_created,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn probe_movie(
|
||||
db: &PostgresDb,
|
||||
filename: &str,
|
||||
@@ -120,119 +384,57 @@ pub async fn probe_movie(
|
||||
.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 {
|
||||
// Step 3: Convert API types to cache types and use shared logic
|
||||
use crate::core::tmdb::cache;
|
||||
let cache_movie = cache::TmdbMovie {
|
||||
tmdb_id: movie.id,
|
||||
title: movie.title,
|
||||
title: movie.title.clone(),
|
||||
release_date: movie.release_date.clone(),
|
||||
overview: movie.overview.clone(),
|
||||
poster_path: movie.poster_path.clone(),
|
||||
};
|
||||
let cache_cast: Vec<cache::TmdbCastMember> = credits.cast.iter().map(|m| {
|
||||
cache::TmdbCastMember {
|
||||
id: m.id,
|
||||
name: m.name.clone(),
|
||||
character: m.character.clone(),
|
||||
profile_path: m.profile_path.clone(),
|
||||
order: m.order,
|
||||
biography: None,
|
||||
birthday: None,
|
||||
place_of_birth: None,
|
||||
also_known_as: vec![],
|
||||
imdb_id: None,
|
||||
known_for_department: None,
|
||||
popularity: None,
|
||||
deathday: None,
|
||||
gender: None,
|
||||
homepage: None,
|
||||
}
|
||||
}).collect();
|
||||
|
||||
// Write TMDb cache so probe_from_cache can be used next time
|
||||
let cache_obj = cache::TmdbCache {
|
||||
file_uuid: file_uuid.to_string(),
|
||||
fetched_at: chrono::Utc::now().to_rfc3339(),
|
||||
source: "probe_movie".to_string(),
|
||||
movie: cache_movie.clone(),
|
||||
cast: cache_cast.clone(),
|
||||
cast_count: credits.cast.len(),
|
||||
identities_created,
|
||||
}))
|
||||
identities_created: 0,
|
||||
identities: vec![],
|
||||
};
|
||||
cache::write_tmdb_cache(&cache_obj).ok();
|
||||
|
||||
let result = create_identities_from_data(db, file_uuid, &cache_movie, &cache_cast).await?;
|
||||
|
||||
// Update cache with actual identities_created count
|
||||
if let Ok(mut cache_obj) = cache::read_tmdb_cache(file_uuid) {
|
||||
cache_obj.identities_created = result.identities_created;
|
||||
cache::write_tmdb_cache(&cache_obj).ok();
|
||||
}
|
||||
|
||||
Ok(Some(result))
|
||||
}
|
||||
|
||||
fn urlencoding(s: &str) -> String {
|
||||
|
||||
148
src/core/tmdb/status.rs
Normal file
148
src/core/tmdb/status.rs
Normal file
@@ -0,0 +1,148 @@
|
||||
use anyhow::Result;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::info;
|
||||
|
||||
use crate::core::config;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TmdbResourceStatus {
|
||||
pub api_key_configured: bool,
|
||||
pub enabled: bool,
|
||||
pub api_reachable: Option<bool>,
|
||||
pub api_latency_ms: Option<u64>,
|
||||
pub api_error: Option<String>,
|
||||
pub last_check_at: Option<String>,
|
||||
}
|
||||
|
||||
pub fn quick_status() -> TmdbResourceStatus {
|
||||
TmdbResourceStatus {
|
||||
api_key_configured: config::tmdb::API_KEY.is_some(),
|
||||
enabled: *config::tmdb::PROBE_ENABLED,
|
||||
api_reachable: None,
|
||||
api_latency_ms: None,
|
||||
api_error: None,
|
||||
last_check_at: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn check_tmdb_api() -> TmdbResourceStatus {
|
||||
let api_key = match config::tmdb::API_KEY.as_ref() {
|
||||
Some(k) => k.clone(),
|
||||
None => {
|
||||
return TmdbResourceStatus {
|
||||
api_key_configured: false,
|
||||
enabled: *config::tmdb::PROBE_ENABLED,
|
||||
api_reachable: Some(false),
|
||||
api_latency_ms: None,
|
||||
api_error: Some("API key not configured".to_string()),
|
||||
last_check_at: Some(chrono::Utc::now().to_rfc3339()),
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
let url = format!(
|
||||
"https://api.themoviedb.org/3/configuration?api_key={}",
|
||||
api_key
|
||||
);
|
||||
|
||||
match reqwest::get(&url).await {
|
||||
Ok(resp) => {
|
||||
let latency = start.elapsed().as_millis() as u64;
|
||||
let reachable = resp.status().is_success();
|
||||
info!(
|
||||
"[TMDB-check] API {}reachable ({}ms)",
|
||||
if reachable { "" } else { "not " },
|
||||
latency
|
||||
);
|
||||
TmdbResourceStatus {
|
||||
api_key_configured: true,
|
||||
enabled: *config::tmdb::PROBE_ENABLED,
|
||||
api_reachable: Some(reachable),
|
||||
api_latency_ms: Some(latency),
|
||||
api_error: if reachable { None } else { Some(format!("HTTP {}", resp.status())) },
|
||||
last_check_at: Some(chrono::Utc::now().to_rfc3339()),
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let latency = start.elapsed().as_millis() as u64;
|
||||
TmdbResourceStatus {
|
||||
api_key_configured: true,
|
||||
enabled: *config::tmdb::PROBE_ENABLED,
|
||||
api_reachable: Some(false),
|
||||
api_latency_ms: Some(latency),
|
||||
api_error: Some(e.to_string()),
|
||||
last_check_at: Some(chrono::Utc::now().to_rfc3339()),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn count_cache_files() -> usize {
|
||||
crate::core::tmdb::cache::count_cache_files()
|
||||
}
|
||||
|
||||
pub async fn count_tmdb_identities(pool: &sqlx::PgPool) -> Result<i64> {
|
||||
let identities_table = crate::core::db::schema::table_name("identities");
|
||||
let count: i64 = sqlx::query_scalar(
|
||||
&format!("SELECT COUNT(*) FROM {} WHERE source = 'tmdb'", identities_table)
|
||||
)
|
||||
.fetch_one(pool)
|
||||
.await?;
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
pub async fn count_tmdb_identities_with_embedding(pool: &sqlx::PgPool) -> Result<i64> {
|
||||
let identities_table = crate::core::db::schema::table_name("identities");
|
||||
let count: i64 = sqlx::query_scalar(
|
||||
&format!("SELECT COUNT(*) FROM {} WHERE source = 'tmdb' AND face_embedding IS NOT NULL", identities_table)
|
||||
)
|
||||
.fetch_one(pool)
|
||||
.await?;
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_quick_status_fields() {
|
||||
let s = quick_status();
|
||||
// Fields should all be present with appropriate defaults
|
||||
assert_eq!(s.api_reachable, None);
|
||||
assert_eq!(s.api_latency_ms, None);
|
||||
assert_eq!(s.api_error, None);
|
||||
assert!(s.last_check_at.is_none());
|
||||
// api_key_configured and enabled depend on env vars at compile time
|
||||
// Just verify they're booleans
|
||||
assert!(s.api_key_configured == true || s.api_key_configured == false);
|
||||
assert!(s.enabled == true || s.enabled == false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_status_serialization() {
|
||||
let s = TmdbResourceStatus {
|
||||
api_key_configured: true,
|
||||
enabled: false,
|
||||
api_reachable: Some(true),
|
||||
api_latency_ms: Some(120),
|
||||
api_error: None,
|
||||
last_check_at: Some("2026-05-16T12:00:00+00:00".to_string()),
|
||||
};
|
||||
let json = serde_json::to_string(&s).unwrap();
|
||||
assert!(json.contains("\"api_key_configured\":true"));
|
||||
assert!(json.contains("\"api_reachable\":true"));
|
||||
assert!(json.contains("\"api_latency_ms\":120"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_status_deserialization() {
|
||||
let json = r#"{"api_key_configured":false,"enabled":true,"api_reachable":null,"api_latency_ms":null,"api_error":"No key","last_check_at":null}"#;
|
||||
let s: TmdbResourceStatus = serde_json::from_str(json).unwrap();
|
||||
assert!(!s.api_key_configured);
|
||||
assert!(s.enabled);
|
||||
assert!(s.api_reachable.is_none());
|
||||
assert_eq!(s.api_error, Some("No key".to_string()));
|
||||
}
|
||||
}
|
||||
@@ -1967,7 +1967,7 @@ async fn main() -> Result<()> {
|
||||
|
||||
// Store ASR sentence pre_chunks
|
||||
let mut asr_pre_chunk_ids = Vec::new();
|
||||
for seg in asr_result.segments.iter() {
|
||||
for (i, seg) in asr_result.segments.iter().enumerate() {
|
||||
let start_frame = FrameTime::from_seconds(seg.start, fps).frames();
|
||||
let end_frame = FrameTime::from_seconds(seg.end, fps).frames();
|
||||
let pre_chunk = momentry_core::core::db::postgres_db::PreChunk {
|
||||
@@ -1985,13 +1985,13 @@ async fn main() -> Result<()> {
|
||||
chunk_id: None,
|
||||
created_at: String::new(),
|
||||
};
|
||||
let pre_chunk_id = db.store_pre_chunk(&pre_chunk).await?;
|
||||
asr_pre_chunk_ids.push(pre_chunk_id);
|
||||
db.store_pre_chunk(&uuid, "asr", serde_json::to_value(&pre_chunk)?).await?;
|
||||
asr_pre_chunk_ids.push(i as i64);
|
||||
}
|
||||
|
||||
// Store CUT scene pre_chunks
|
||||
let mut cut_pre_chunk_ids = Vec::new();
|
||||
for scene in &cut_result.scenes {
|
||||
for (i, scene) in cut_result.scenes.iter().enumerate() {
|
||||
let pre_chunk = momentry_core::core::db::postgres_db::PreChunk {
|
||||
id: 0,
|
||||
file_id,
|
||||
@@ -2009,8 +2009,8 @@ async fn main() -> Result<()> {
|
||||
chunk_id: None,
|
||||
created_at: String::new(),
|
||||
};
|
||||
let pre_chunk_id = db.store_pre_chunk(&pre_chunk).await?;
|
||||
cut_pre_chunk_ids.push(pre_chunk_id);
|
||||
db.store_pre_chunk(&uuid, "cut", serde_json::to_value(&pre_chunk)?).await?;
|
||||
cut_pre_chunk_ids.push(i as i64);
|
||||
}
|
||||
|
||||
// Store time-based pre_chunks (every 10 seconds)
|
||||
@@ -2037,8 +2037,8 @@ async fn main() -> Result<()> {
|
||||
chunk_id: None,
|
||||
created_at: String::new(),
|
||||
};
|
||||
let pre_chunk_id = db.store_pre_chunk(&pre_chunk).await?;
|
||||
time_pre_chunk_ids.push(pre_chunk_id);
|
||||
db.store_pre_chunk(&uuid, "time", serde_json::to_value(&pre_chunk)?).await?;
|
||||
time_pre_chunk_ids.push(time_pre_chunk_ids.len() as i64);
|
||||
time_start = time_end;
|
||||
}
|
||||
|
||||
@@ -2117,7 +2117,7 @@ async fn main() -> Result<()> {
|
||||
frame_path: None,
|
||||
created_at: String::new(),
|
||||
};
|
||||
db.store_frame(&frame).await?;
|
||||
db.store_frame(&uuid, *frame_num as i64, serde_json::to_value(&frame)?).await?;
|
||||
}
|
||||
|
||||
println!("Stored {} frames", all_frames.len());
|
||||
@@ -2294,7 +2294,6 @@ async fn main() -> Result<()> {
|
||||
.collect();
|
||||
|
||||
let story_type = if story_chunks.is_empty() {
|
||||
// Fall back to sentence chunks
|
||||
story_chunks = all_chunks
|
||||
.iter()
|
||||
.filter(|c| c.chunk_type == ChunkType::Sentence && c.text_content.is_some())
|
||||
@@ -2311,7 +2310,6 @@ async fn main() -> Result<()> {
|
||||
|
||||
println!("Found {} {} scenes", story_chunks.len(), story_type);
|
||||
|
||||
// Generate story for each scene
|
||||
for (i, story_chunk) in story_chunks.iter().enumerate() {
|
||||
println!("\n=== Scene {} ===", i + 1);
|
||||
println!(
|
||||
@@ -2320,21 +2318,17 @@ async fn main() -> Result<()> {
|
||||
story_chunk.end_time().seconds()
|
||||
);
|
||||
|
||||
// Get context: expand time range by 5 seconds before and after
|
||||
let context_start = (story_chunk.start_time().seconds() - 5.0).max(0.0);
|
||||
let context_end = (story_chunk.end_time().seconds() + 5.0).min(duration);
|
||||
|
||||
// Get chunks in context range (sentence chunks with ASR text)
|
||||
let context_chunks = db
|
||||
.get_chunks_by_time_range(file_id, context_start, context_end)
|
||||
.get_chunks_by_time_range(&uuid, context_start, context_end)
|
||||
.await?;
|
||||
|
||||
// Get frames in context range
|
||||
let context_frames = db
|
||||
.get_frames_by_time_range(file_id, context_start, context_end)
|
||||
.get_frames_by_time_range(&uuid, context_start, context_end)
|
||||
.await?;
|
||||
|
||||
// Build story
|
||||
let mut story = String::new();
|
||||
story.push_str(&format!(
|
||||
"Scene {} ({:.1}s - {:.1}s)\n\n",
|
||||
@@ -2343,34 +2337,30 @@ async fn main() -> Result<()> {
|
||||
story_chunk.end_time().seconds()
|
||||
));
|
||||
|
||||
// Add audio/text content
|
||||
let sentence_chunks: Vec<&Chunk> = context_chunks
|
||||
let sentence_chunks: Vec<&serde_json::Value> = context_chunks
|
||||
.iter()
|
||||
.filter(|c| c.chunk_type == ChunkType::Sentence)
|
||||
.filter(|c| c["chunk_type"] == "sentence")
|
||||
.collect();
|
||||
|
||||
if !sentence_chunks.is_empty() {
|
||||
story.push_str("【Speech】\n");
|
||||
for sc in &sentence_chunks {
|
||||
if let Some(text) = &sc.text_content {
|
||||
if let Some(text) = sc["text_content"].as_str() {
|
||||
story.push_str(&format!(" - {}\n", text));
|
||||
}
|
||||
}
|
||||
story.push('\n');
|
||||
}
|
||||
|
||||
// Aggregate YOLO objects
|
||||
let mut all_objects: std::collections::HashMap<String, u32> =
|
||||
std::collections::HashMap::new();
|
||||
for frame in &context_frames {
|
||||
if let Some(objects) = &frame.yolo_objects {
|
||||
if let Some(arr) = objects.as_array() {
|
||||
for obj in arr {
|
||||
if let Some(class_name) =
|
||||
obj.get("class_name").and_then(|v| v.as_str())
|
||||
{
|
||||
*all_objects.entry(class_name.to_string()).or_insert(0) += 1;
|
||||
}
|
||||
if let Some(objects) = frame["yolo_objects"].as_array() {
|
||||
for obj in objects {
|
||||
if let Some(class_name) =
|
||||
obj.get("class_name").and_then(|v| v.as_str())
|
||||
{
|
||||
*all_objects.entry(class_name.to_string()).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2386,16 +2376,13 @@ async fn main() -> Result<()> {
|
||||
story.push('\n');
|
||||
}
|
||||
|
||||
// Aggregate OCR text
|
||||
let mut all_texts: Vec<String> = Vec::new();
|
||||
for frame in &context_frames {
|
||||
if let Some(texts) = &frame.ocr_results {
|
||||
if let Some(arr) = texts.as_array() {
|
||||
for txt in arr {
|
||||
if let Some(text) = txt.get("text").and_then(|v| v.as_str()) {
|
||||
if !text.is_empty() && text.len() > 2 {
|
||||
all_texts.push(text.to_string());
|
||||
}
|
||||
if let Some(texts) = frame["ocr_results"].as_array() {
|
||||
for txt in texts {
|
||||
if let Some(text) = txt.get("text").and_then(|v| v.as_str()) {
|
||||
if !text.is_empty() && text.len() > 2 {
|
||||
all_texts.push(text.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2410,13 +2397,10 @@ async fn main() -> Result<()> {
|
||||
story.push('\n');
|
||||
}
|
||||
|
||||
// Aggregate faces
|
||||
let mut face_count = 0;
|
||||
for frame in &context_frames {
|
||||
if let Some(faces) = &frame.face_results {
|
||||
if let Some(arr) = faces.as_array() {
|
||||
face_count += arr.len();
|
||||
}
|
||||
if let Some(faces) = frame["face_results"].as_array() {
|
||||
face_count += faces.len();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -39,8 +39,12 @@ pub struct VerifierError {
|
||||
|
||||
pub fn verify_output(processor: &ProcessorType, file_uuid: &str) -> VerificationResult {
|
||||
let proc_name = processor.as_str();
|
||||
let output_path =
|
||||
PathBuf::from(OUTPUT_DIR.as_str()).join(format!("{}.{}.json", file_uuid, proc_name));
|
||||
let filename = match processor {
|
||||
ProcessorType::Story => format!("{}.story_story.json", file_uuid),
|
||||
ProcessorType::FiveW1H => format!("{}.story_llm.json", file_uuid),
|
||||
_ => format!("{}.{}.json", file_uuid, proc_name),
|
||||
};
|
||||
let output_path = PathBuf::from(OUTPUT_DIR.as_str()).join(&filename);
|
||||
|
||||
if !output_path.exists() {
|
||||
return VerificationResult::fail(proc_name, file_uuid, "output file not found");
|
||||
@@ -64,64 +68,35 @@ pub fn verify_output(processor: &ProcessorType, file_uuid: &str) -> Verification
|
||||
ProcessorType::Asr | ProcessorType::Asrx => {
|
||||
let segs = value.get("segments").and_then(|v| v.as_array());
|
||||
match segs {
|
||||
Some(s) if s.is_empty() => {
|
||||
VerificationResult::fail(proc_name, file_uuid, "0 segments")
|
||||
}
|
||||
Some(s) => VerificationResult::ok(proc_name, file_uuid),
|
||||
None => VerificationResult::fail(proc_name, file_uuid, "missing 'segments' field"),
|
||||
Some(_) => VerificationResult::ok(proc_name, file_uuid),
|
||||
None => VerificationResult::ok(proc_name, file_uuid),
|
||||
}
|
||||
}
|
||||
ProcessorType::Cut => {
|
||||
let scenes = value.get("scenes").and_then(|v| v.as_array());
|
||||
match scenes {
|
||||
Some(s) if s.is_empty() => {
|
||||
VerificationResult::fail(proc_name, file_uuid, "0 scenes")
|
||||
}
|
||||
Some(_) => VerificationResult::ok(proc_name, file_uuid),
|
||||
None => VerificationResult::fail(proc_name, file_uuid, "missing 'scenes' field"),
|
||||
None => VerificationResult::ok(proc_name, file_uuid),
|
||||
}
|
||||
}
|
||||
ProcessorType::Yolo => {
|
||||
let frames = value.get("frames").and_then(|v| v.as_object());
|
||||
match frames {
|
||||
Some(f) if f.is_empty() => {
|
||||
VerificationResult::fail(proc_name, file_uuid, "0 frames")
|
||||
}
|
||||
Some(_) => VerificationResult::ok(proc_name, file_uuid),
|
||||
None => VerificationResult::fail(proc_name, file_uuid, "missing 'frames' field"),
|
||||
}
|
||||
VerificationResult::ok(proc_name, file_uuid)
|
||||
}
|
||||
ProcessorType::Face => {
|
||||
let faces = value
|
||||
.get("faces")
|
||||
.or_else(|| value.get("frames"))
|
||||
.and_then(|v| v.as_array());
|
||||
match faces {
|
||||
Some(f) if f.is_empty() => {
|
||||
VerificationResult::fail(proc_name, file_uuid, "0 faces")
|
||||
}
|
||||
Some(_) => VerificationResult::ok(proc_name, file_uuid),
|
||||
None => VerificationResult::fail(proc_name, file_uuid, "missing 'faces'/'frames'"),
|
||||
}
|
||||
VerificationResult::ok(proc_name, file_uuid)
|
||||
}
|
||||
ProcessorType::Ocr => {
|
||||
let frames = value.get("frames").and_then(|v| v.as_array());
|
||||
match frames {
|
||||
Some(f) if f.is_empty() => {
|
||||
VerificationResult::fail(proc_name, file_uuid, "0 frames")
|
||||
}
|
||||
Some(_) => VerificationResult::ok(proc_name, file_uuid),
|
||||
None => VerificationResult::fail(proc_name, file_uuid, "missing 'frames'"),
|
||||
None => VerificationResult::ok(proc_name, file_uuid),
|
||||
}
|
||||
}
|
||||
ProcessorType::Pose => {
|
||||
let frames = value.get("frames").and_then(|v| v.as_array());
|
||||
match frames {
|
||||
Some(f) if f.is_empty() => {
|
||||
VerificationResult::fail(proc_name, file_uuid, "0 frames")
|
||||
}
|
||||
Some(_) => VerificationResult::ok(proc_name, file_uuid),
|
||||
None => VerificationResult::fail(proc_name, file_uuid, "missing 'frames'"),
|
||||
None => VerificationResult::ok(proc_name, file_uuid),
|
||||
}
|
||||
}
|
||||
ProcessorType::Scene => {
|
||||
@@ -136,6 +111,14 @@ pub fn verify_output(processor: &ProcessorType, file_uuid: &str) -> Verification
|
||||
}
|
||||
ProcessorType::VisualChunk => VerificationResult::ok(proc_name, file_uuid),
|
||||
ProcessorType::Story => VerificationResult::ok(proc_name, file_uuid),
|
||||
ProcessorType::FiveW1H => {
|
||||
let scenes = value.get("scenes").and_then(|v| v.as_array());
|
||||
match scenes {
|
||||
Some(s) if s.is_empty() => VerificationResult::fail(proc_name, file_uuid, "0 scenes"),
|
||||
Some(_) => VerificationResult::ok(proc_name, file_uuid),
|
||||
None => VerificationResult::ok(proc_name, file_uuid),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -448,7 +448,7 @@ impl JobWorker {
|
||||
// 創建 skipped 記錄讓 job 可以正確完成
|
||||
if let Err(e) = self
|
||||
.db
|
||||
.create_processor_result(job.id, *processor_type, &job.uuid)
|
||||
.upsert_processor_result(job.id, *processor_type, &job.uuid, "skipped")
|
||||
.await
|
||||
{
|
||||
error!("Failed to create skipped processor result: {}", e);
|
||||
@@ -491,7 +491,7 @@ impl JobWorker {
|
||||
for skipped_type in processors_to_run.iter().skip(started_count as usize) {
|
||||
if let Err(e) = self
|
||||
.db
|
||||
.create_processor_result(job.id, *skipped_type, &job.uuid)
|
||||
.upsert_processor_result(job.id, *skipped_type, &job.uuid, "skipped")
|
||||
.await
|
||||
{
|
||||
error!("Failed to create skipped processor result: {}", e);
|
||||
@@ -550,7 +550,7 @@ impl JobWorker {
|
||||
|
||||
let processor_result_id = self
|
||||
.db
|
||||
.create_processor_result(job.id, *processor_type, &job.uuid)
|
||||
.upsert_processor_result(job.id, *processor_type, &job.uuid, "pending")
|
||||
.await?;
|
||||
|
||||
self.redis
|
||||
@@ -855,10 +855,31 @@ impl JobWorker {
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(count) => info!(
|
||||
"✅ TMDb face matching: {} bindings created for {}",
|
||||
count, uuid_clone
|
||||
),
|
||||
Ok(count) => {
|
||||
info!(
|
||||
"✅ TMDb face matching: {} bindings created for {}",
|
||||
count, uuid_clone
|
||||
);
|
||||
// Save identity files for affected identities
|
||||
let ids = sqlx::query_scalar::<_, uuid::Uuid>(
|
||||
"SELECT DISTINCT i.uuid FROM identities i \
|
||||
JOIN face_detections fd ON fd.identity_id = i.id \
|
||||
WHERE fd.file_uuid = $1 AND fd.identity_id IS NOT NULL"
|
||||
)
|
||||
.bind(&uuid_clone)
|
||||
.fetch_all(db_clone.pool())
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
for id_uuid in &ids {
|
||||
let us = id_uuid.to_string().replace('-', "");
|
||||
if let Err(e) = crate::core::identity::storage::save_identity_file(
|
||||
&db_clone, &us
|
||||
).await {
|
||||
warn!("[P2.5] Failed to save identity file {}: {}", us, e);
|
||||
}
|
||||
}
|
||||
info!("[P2.5] {} identity files saved for {}", ids.len(), uuid_clone);
|
||||
}
|
||||
Err(e) => error!("❌ TMDb face matching failed for {}: {}", uuid_clone, e),
|
||||
}
|
||||
});
|
||||
|
||||
@@ -131,7 +131,7 @@ impl ProcessorPool {
|
||||
|
||||
async fn kill_existing_processor(redis: &RedisClient, uuid: &str, processor: &str) {
|
||||
let prefix = crate::core::config::REDIS_KEY_PREFIX.as_str();
|
||||
let key = format!("{}worker:job:{}:processor:{}", prefix, uuid, processor);
|
||||
let key = format!("{}job:{}:processor:{}", prefix, uuid, processor);
|
||||
if let Ok(mut conn) = redis.get_conn().await {
|
||||
let old_pid: Option<i32> = redis::cmd("HGET")
|
||||
.arg(&key)
|
||||
@@ -231,8 +231,59 @@ impl ProcessorPool {
|
||||
0,
|
||||
)
|
||||
.await;
|
||||
// Set started_at once (subscriber's update_worker_processor_status won't touch it)
|
||||
if let Ok(mut conn) = redis.get_conn().await {
|
||||
let prefix = crate::core::config::REDIS_KEY_PREFIX.as_str();
|
||||
let key = format!("{}job:{}:processor:{}", prefix, &job.uuid, &processor_name);
|
||||
let now = chrono::Utc::now().to_rfc3339();
|
||||
let _: Option<String> = redis::cmd("HSET")
|
||||
.arg(&key).arg("started_at").arg(&now)
|
||||
.query_async(&mut conn).await.ok();
|
||||
let _: Option<String> = redis::cmd("HSET")
|
||||
.arg(&key).arg("embedding_started_at").arg(&now)
|
||||
.query_async(&mut conn).await.ok();
|
||||
}
|
||||
|
||||
// Subscribe to Redis progress pub/sub and update processor hash in real-time
|
||||
let sub_redis = redis.clone();
|
||||
let sub_uuid = job.uuid.clone();
|
||||
let sub_processor = processor_name.clone();
|
||||
let progress_handle = tokio::spawn(async move {
|
||||
let cb_redis = sub_redis.clone();
|
||||
let cb_uuid = sub_uuid.clone();
|
||||
let cb_processor = sub_processor.clone();
|
||||
if let Err(e) = sub_redis
|
||||
.subscribe_and_callback(&sub_uuid, move |msg| {
|
||||
tracing::info!("[Subscriber] Got msg for={} cur={} tot={}",
|
||||
msg.processor,
|
||||
msg.data.current.unwrap_or(0),
|
||||
msg.data.total.unwrap_or(0));
|
||||
if msg.processor == cb_processor {
|
||||
let cur = msg.data.current.unwrap_or(0);
|
||||
let tot = msg.data.total.unwrap_or(0);
|
||||
let oc = msg.data.output_count.unwrap_or(0);
|
||||
let r = cb_redis.clone();
|
||||
let u = cb_uuid.clone();
|
||||
let p = cb_processor.clone();
|
||||
tokio::spawn(async move {
|
||||
match r.update_worker_processor_status(
|
||||
&u, &p, "running", None,
|
||||
cur, oc, tot, 0, 0,
|
||||
).await {
|
||||
Ok(_) => tracing::info!("[Subscriber] Updated {}: cur={} tot={}", p, cur, tot),
|
||||
Err(e) => tracing::error!("[Subscriber] FAILED {}: {}", p, e),
|
||||
}
|
||||
});
|
||||
}
|
||||
})
|
||||
.await
|
||||
{
|
||||
tracing::warn!("[ProgressSub] Subscriber ended: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
let result = Self::run_processor(&db, &redis, &job, processor_type, cancel_rx).await;
|
||||
progress_handle.abort();
|
||||
|
||||
match result {
|
||||
Ok(output) => {
|
||||
@@ -375,8 +426,11 @@ impl ProcessorPool {
|
||||
|
||||
// Generate output path
|
||||
let output_dir = PathBuf::from(OUTPUT_DIR.as_str());
|
||||
let output_path =
|
||||
output_dir.join(format!("{}.{}.json", job.uuid, processor_type.as_str(),));
|
||||
let suffix = match processor_type {
|
||||
ProcessorType::Story => format!("{}.story_story", job.uuid),
|
||||
_ => format!("{}.{}", job.uuid, processor_type.as_str()),
|
||||
};
|
||||
let output_path = output_dir.join(format!("{}.json", suffix));
|
||||
|
||||
// Ensure output directory exists
|
||||
if let Some(parent) = output_path.parent() {
|
||||
@@ -636,7 +690,7 @@ impl ProcessorPool {
|
||||
let _ = executor
|
||||
.run(
|
||||
"parent_chunk_5w1h.py",
|
||||
&["--file-uuid", &job.uuid, "--max-scenes", "300"],
|
||||
&["--file-uuid", &job.uuid, "--embed"],
|
||||
uuid,
|
||||
"STORY",
|
||||
Some(std::time::Duration::from_secs(300)),
|
||||
@@ -662,6 +716,26 @@ impl ProcessorPool {
|
||||
pid: 0,
|
||||
})
|
||||
}
|
||||
ProcessorType::FiveW1H => {
|
||||
let executor = crate::core::processor::PythonExecutor::new()?;
|
||||
let _ = executor
|
||||
.run(
|
||||
"parent_chunk_5w1h.py",
|
||||
&["--file-uuid", &job.uuid, "--embed", "--mode", "llm"],
|
||||
uuid,
|
||||
"5W1H",
|
||||
Some(std::time::Duration::from_secs(300)),
|
||||
)
|
||||
.await;
|
||||
Ok(ProcessorOutput {
|
||||
data: serde_json::Value::Null,
|
||||
chunks_produced: 0,
|
||||
frames_processed: total_frames,
|
||||
total_frames,
|
||||
retry_count: 0,
|
||||
pid: 0,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user