diff --git a/src/api/trace_agent_api.rs b/src/api/trace_agent_api.rs index 90cb8ad..c6abb78 100644 --- a/src/api/trace_agent_api.rs +++ b/src/api/trace_agent_api.rs @@ -3,7 +3,7 @@ use axum::{ extract::{Path, Query, State}, http::{header, StatusCode}, response::{IntoResponse, Json, Response}, - routing::{get, post}, + routing::{delete, get, post}, Router, }; use serde::{Deserialize, Serialize}; @@ -39,6 +39,8 @@ pub fn trace_agent_routes() -> Router { get(get_cooccurrence), ) .route("/api/v1/file/:file_uuid/tkg/rebuild", post(rebuild_tkg)) + .route("/api/v1/file/:file_uuid/tkg", get(get_tkg_operations)) + .route("/api/v1/file/:file_uuid/tkg", delete(delete_tkg)) .route("/api/v1/file/:file_uuid/rule2", post(ingest_rule2)) .route( "/api/v1/file/:file_uuid/representative-frame", @@ -980,58 +982,41 @@ struct TkgRebuildResponse { error: Option, } +#[derive(Deserialize)] +struct RebuildParams { + force: Option, +} + async fn rebuild_tkg( State(state): State, Path(file_uuid): Path, + Query(params): Query, ) -> Json { use crate::core::chunk::rule2_ingest::ingest_rule2; + use crate::core::tkg::TkgService; use tracing::info; - let redis = crate::core::db::RedisClient::new().ok(); - let result = crate::core::processor::tkg::build_tkg(&state.db, &file_uuid, &OUTPUT_DIR, redis.map(Arc::new)).await; + let db = state.db.clone(); + let tkg_service = TkgService::new(state.db); + let force = params.force.unwrap_or(false); + let result = tkg_service.rebuild(&file_uuid, &OUTPUT_DIR, force).await; match result { Ok(r) => { - let total_edges = r.speaker_face_edges - + r.mutual_gaze_edges - + r.face_face_edges - + r.co_occurrence_edges - + r.has_appearance_edges - + r.wears_edges; - - if total_edges > 0 { - info!( - "[TKG] {} relationship edges found, triggering Rule 2 ingestion...", - total_edges - ); - match ingest_rule2(state.db.pool(), &file_uuid, None, None).await { - Ok(count) => info!("[TKG] Rule 2 created {} relationship chunks", count), - Err(e) => info!("[TKG] Rule 2 ingestion failed: {}", e), - } + // Always trigger Rule 2 (even with 0 edges) + info!( + "[TKG] Rebuild completed for {}: {} nodes, {} edges", + file_uuid, r.total_nodes(), r.total_edges() + ); + match ingest_rule2(db.pool(), &file_uuid, None, None).await { + Ok(count) => info!("[TKG] Rule 2 created {} relationship chunks", count), + Err(e) => info!("[TKG] Rule 2 ingestion failed: {}", e), } Json(TkgRebuildResponse { success: true, file_uuid, - result: Some(serde_json::json!({ - "face_track_nodes": r.face_track_nodes, - "gaze_track_nodes": r.gaze_track_nodes, - "lip_track_nodes": r.lip_track_nodes, - "text_region_nodes": r.text_region_nodes, - "appearance_trace_nodes": r.appearance_trace_nodes, - "accessory_nodes": r.accessory_nodes, - "object_nodes": r.object_nodes, - "hand_nodes": r.hand_nodes, - "speaker_nodes": r.speaker_nodes, - "co_occurrence_edges": r.co_occurrence_edges, - "speaker_face_edges": r.speaker_face_edges, - "face_face_edges": r.face_face_edges, - "mutual_gaze_edges": r.mutual_gaze_edges, - "lip_sync_edges": r.lip_sync_edges, - "has_appearance_edges": r.has_appearance_edges, - "wears_edges": r.wears_edges, - "hand_object_edges": r.hand_object_edges, - })), + result: Some(r.to_json()), error: None, }) } @@ -1044,6 +1029,28 @@ async fn rebuild_tkg( } } +async fn get_tkg_operations( + State(state): State, + Path(file_uuid): Path, +) -> Json> { + use crate::core::tkg::TkgService; + let tkg_service = TkgService::new(state.db); + let operations = tkg_service.get_operations(&file_uuid).await.unwrap_or_default(); + Json(operations) +} + +async fn delete_tkg( + State(state): State, + Path(file_uuid): Path, +) -> Json { + use crate::core::tkg::TkgService; + let tkg_service = TkgService::new(state.db); + match tkg_service.delete_tkg_and_logs(&file_uuid).await { + Ok(_) => Json(serde_json::json!({"success": true, "message": "TKG deleted successfully"})), + Err(e) => Json(serde_json::json!({"success": false, "error": e.to_string()})), + } +} + // ── Representative Frame (JSON) ─────────────────────────────────── use crate::core::processor::tkg; diff --git a/src/core/db/postgres_db.rs b/src/core/db/postgres_db.rs index b3129d7..7ae38b3 100644 --- a/src/core/db/postgres_db.rs +++ b/src/core/db/postgres_db.rs @@ -1324,6 +1324,29 @@ impl PostgresDb { .execute(pool) .await?; + // ── TKG Operation Log ── + sqlx::query("CREATE TABLE IF NOT EXISTS tkg_operation_log ( + id BIGSERIAL PRIMARY KEY, + file_uuid VARCHAR(32) NOT NULL, + operation VARCHAR(20) NOT NULL, + node_type VARCHAR(50), + edge_type VARCHAR(50), + nodes_created INT DEFAULT 0, + nodes_updated INT DEFAULT 0, + nodes_deleted INT DEFAULT 0, + edges_created INT DEFAULT 0, + edges_updated INT DEFAULT 0, + edges_deleted INT DEFAULT 0, + status VARCHAR(20) DEFAULT 'pending', + error_message TEXT, + started_at TIMESTAMPTZ DEFAULT NOW(), + completed_at TIMESTAMPTZ, + properties JSONB + )").execute(pool).await?; + sqlx::query("CREATE INDEX IF NOT EXISTS idx_tkg_op_file_uuid ON tkg_operation_log(file_uuid)") + .execute(pool) + .await?; + // ── Functions & Triggers ── sqlx::query( "CREATE OR REPLACE FUNCTION update_search_vector() RETURNS TRIGGER AS $func$ diff --git a/src/core/mod.rs b/src/core/mod.rs index 7566f82..45e0f53 100644 --- a/src/core/mod.rs +++ b/src/core/mod.rs @@ -23,4 +23,5 @@ pub mod text; pub mod thumbnail; pub mod time; pub mod tmdb; +pub mod tkg; pub mod vision; diff --git a/src/core/tkg/log.rs b/src/core/tkg/log.rs new file mode 100644 index 0000000..9fa53cc --- /dev/null +++ b/src/core/tkg/log.rs @@ -0,0 +1,119 @@ +use crate::core::db::schema; +use crate::core::db::PostgresDb; +use anyhow::Result; +use sqlx::PgPool; + +pub struct TkgLogger { + pool: PgPool, + pub log_id: Option, +} + +impl TkgLogger { + pub fn new(pool: PgPool) -> Self { + Self { pool, log_id: None } + } + + pub async fn start_operation(pool: &PgPool, file_uuid: &str, operation: &str) -> Result { + let table = schema::table_name("tkg_operation_log"); + let id = sqlx::query_scalar::<_, i64>(&format!( + "INSERT INTO {table} (file_uuid, operation, status, started_at) \ + VALUES ($1, $2, 'pending', NOW()) RETURNING id" + )) + .bind(file_uuid) + .bind(operation) + .fetch_one(pool) + .await?; + tracing::info!("[TKG-Log] Started operation {} for {}", operation, file_uuid); + Ok(id) + } + + pub async fn update_progress( + &self, + node_type: &str, + nodes_created: i32, + edges_created: i32, + ) -> Result<()> { + if let Some(log_id) = self.log_id { + let table = schema::table_name("tkg_operation_log"); + sqlx::query(&format!( + "UPDATE {table} \ + SET nodes_created = nodes_created + $1, \ + edges_created = edges_created + $2, \ + node_type = COALESCE(node_type, $3) \ + WHERE id = $4" + )) + .bind(nodes_created) + .bind(edges_created) + .bind(node_type) + .bind(log_id) + .execute(&self.pool) + .await?; + } + Ok(()) + } + + pub async fn complete_operation(&self, error: Option<&str>) -> Result<()> { + if let Some(log_id) = self.log_id { + let table = schema::table_name("tkg_operation_log"); + let status = if error.is_some() { "failed" } else { "completed" }; + sqlx::query(&format!( + "UPDATE {table} \ + SET status = $1, \ + error_message = $2, \ + completed_at = NOW() \ + WHERE id = $3" + )) + .bind(status) + .bind(error) + .bind(log_id) + .execute(&self.pool) + .await?; + tracing::info!("[TKG-Log] Operation {} completed with status: {}", log_id, status); + } + Ok(()) + } + + pub async fn get_operations(pool: &PgPool, file_uuid: &str) -> Result> { + let table = schema::table_name("tkg_operation_log"); + let rows = sqlx::query_as::<_, (i64, String, String, Option, Option, i32, i32, i32, i32, i32, i32, String, Option, String, Option, Option)>(&format!( + "SELECT id, file_uuid, operation, node_type, edge_type, \ + nodes_created, nodes_updated, nodes_deleted, \ + edges_created, edges_updated, edges_deleted, \ + status, error_message, started_at::text, completed_at::text, properties \ + FROM {table} WHERE file_uuid = $1 ORDER BY started_at DESC" + )) + .bind(file_uuid) + .fetch_all(pool) + .await?; + + Ok(rows.into_iter().map(|r| crate::core::tkg::models::TkgOperationLog { + id: r.0, + file_uuid: r.1, + operation: r.2, + node_type: r.3, + edge_type: r.4, + nodes_created: r.5, + nodes_updated: r.6, + nodes_deleted: r.7, + edges_created: r.8, + edges_updated: r.9, + edges_deleted: r.10, + status: r.11, + error_message: r.12, + started_at: r.13, + completed_at: r.14, + properties: r.15, + }).collect()) + } + + pub async fn delete_operations(pool: &PgPool, file_uuid: &str) -> Result { + let table = schema::table_name("tkg_operation_log"); + let result = sqlx::query(&format!( + "DELETE FROM {table} WHERE file_uuid = $1" + )) + .bind(file_uuid) + .execute(pool) + .await?; + Ok(result.rows_affected() as i64) + } +} diff --git a/src/core/tkg/mod.rs b/src/core/tkg/mod.rs new file mode 100644 index 0000000..2ac0b81 --- /dev/null +++ b/src/core/tkg/mod.rs @@ -0,0 +1,7 @@ +pub mod service; +pub mod log; +pub mod models; + +pub use service::TkgService; +pub use log::TkgLogger; +pub use models::*; diff --git a/src/core/tkg/models.rs b/src/core/tkg/models.rs new file mode 100644 index 0000000..cb1ce0b --- /dev/null +++ b/src/core/tkg/models.rs @@ -0,0 +1,94 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TkgOperationLog { + pub id: i64, + pub file_uuid: String, + pub operation: String, + pub node_type: Option, + pub edge_type: Option, + pub nodes_created: i32, + pub nodes_updated: i32, + pub nodes_deleted: i32, + pub edges_created: i32, + pub edges_updated: i32, + pub edges_deleted: i32, + pub status: String, + pub error_message: Option, + pub started_at: String, + pub completed_at: Option, + pub properties: Option, +} + +#[derive(Debug, Clone, Default)] +pub struct TkgNodeStats { + pub created: i32, + pub updated: i32, + pub deleted: i32, +} + +#[derive(Debug, Clone, Default)] +pub struct TkgEdgeStats { + pub created: i32, + pub updated: i32, + pub deleted: i32, +} + +#[derive(Debug, Clone, Default)] +pub struct TkgBuildStats { + pub face_track_nodes: i32, + pub gaze_track_nodes: i32, + pub lip_track_nodes: i32, + pub text_region_nodes: i32, + pub appearance_trace_nodes: i32, + pub accessory_nodes: i32, + pub object_nodes: i32, + pub hand_nodes: i32, + pub speaker_nodes: i32, + pub co_occurrence_edges: i32, + pub speaker_face_edges: i32, + pub face_face_edges: i32, + pub mutual_gaze_edges: i32, + pub lip_sync_edges: i32, + pub has_appearance_edges: i32, + pub wears_edges: i32, + pub hand_object_edges: i32, +} + +impl TkgBuildStats { + pub fn total_nodes(&self) -> i32 { + self.face_track_nodes + self.gaze_track_nodes + self.lip_track_nodes + + self.text_region_nodes + self.appearance_trace_nodes + self.accessory_nodes + + self.object_nodes + self.hand_nodes + self.speaker_nodes + } + + pub fn total_edges(&self) -> i32 { + self.co_occurrence_edges + self.speaker_face_edges + self.face_face_edges + + self.mutual_gaze_edges + self.lip_sync_edges + self.has_appearance_edges + + self.wears_edges + self.hand_object_edges + } + + pub fn to_json(&self) -> serde_json::Value { + serde_json::json!({ + "face_track_nodes": self.face_track_nodes, + "gaze_track_nodes": self.gaze_track_nodes, + "lip_track_nodes": self.lip_track_nodes, + "text_region_nodes": self.text_region_nodes, + "appearance_trace_nodes": self.appearance_trace_nodes, + "accessory_nodes": self.accessory_nodes, + "object_nodes": self.object_nodes, + "hand_nodes": self.hand_nodes, + "speaker_nodes": self.speaker_nodes, + "co_occurrence_edges": self.co_occurrence_edges, + "speaker_face_edges": self.speaker_face_edges, + "face_face_edges": self.face_face_edges, + "mutual_gaze_edges": self.mutual_gaze_edges, + "lip_sync_edges": self.lip_sync_edges, + "has_appearance_edges": self.has_appearance_edges, + "wears_edges": self.wears_edges, + "hand_object_edges": self.hand_object_edges, + "total_nodes": self.total_nodes(), + "total_edges": self.total_edges(), + }) + } +} diff --git a/src/core/tkg/service.rs b/src/core/tkg/service.rs new file mode 100644 index 0000000..0b3e9e9 --- /dev/null +++ b/src/core/tkg/service.rs @@ -0,0 +1,140 @@ +use crate::core::db::PostgresDb; +use crate::core::db::schema; +use crate::core::tkg::log::TkgLogger; +use crate::core::tkg::models::{TkgBuildStats, TkgOperationLog}; +use anyhow::Result; +use sqlx::PgPool; +use std::sync::Arc; + +pub struct TkgService { + db: Arc, +} + +impl TkgService { + pub fn new(db: Arc) -> Self { + Self { db } + } + + pub async fn build(&self, file_uuid: &str, output_dir: &str) -> Result { + let log_id = TkgLogger::start_operation(self.db.pool(), file_uuid, "build").await?; + let mut logger = TkgLogger::new(self.db.pool().clone()); + logger.log_id = Some(log_id); + + let redis = crate::core::db::RedisClient::new().ok().map(|r| std::sync::Arc::new(r)); + let result = crate::core::processor::tkg::build_tkg(&self.db, file_uuid, output_dir, redis).await; + + match result { + Ok(r) => { + let stats = TkgBuildStats { + face_track_nodes: r.face_track_nodes as i32, + gaze_track_nodes: r.gaze_track_nodes as i32, + lip_track_nodes: r.lip_track_nodes as i32, + text_region_nodes: r.text_region_nodes as i32, + appearance_trace_nodes: r.appearance_trace_nodes as i32, + accessory_nodes: r.accessory_nodes as i32, + object_nodes: r.object_nodes as i32, + hand_nodes: r.hand_nodes as i32, + speaker_nodes: r.speaker_nodes as i32, + co_occurrence_edges: r.co_occurrence_edges as i32, + speaker_face_edges: r.speaker_face_edges as i32, + face_face_edges: r.face_face_edges as i32, + mutual_gaze_edges: r.mutual_gaze_edges as i32, + lip_sync_edges: r.lip_sync_edges as i32, + has_appearance_edges: r.has_appearance_edges as i32, + wears_edges: r.wears_edges as i32, + hand_object_edges: r.hand_object_edges as i32, + }; + logger.complete_operation(None).await?; + tracing::info!("[TKG-Service] Build completed for {}: {} nodes, {} edges", + file_uuid, stats.total_nodes(), stats.total_edges()); + Ok(stats) + } + Err(e) => { + logger.complete_operation(Some(&e.to_string())).await?; + tracing::error!("[TKG-Service] Build failed for {}: {}", file_uuid, e); + Err(e) + } + } + } + + pub async fn rebuild(&self, file_uuid: &str, output_dir: &str, force: bool) -> Result { + let operation = if force { "rebuild_force" } else { "rebuild" }; + let log_id = TkgLogger::start_operation(self.db.pool(), file_uuid, operation).await?; + + if force { + self.delete_internal(file_uuid).await?; + } + + let result = self.build(file_uuid, output_dir).await; + + // Update the original log entry + if let Some(id) = Some(log_id) { + let table = schema::table_name("tkg_operation_log"); + let status = if result.is_ok() { "completed" } else { "failed" }; + let error = result.as_ref().err().map(|e| e.to_string()); + sqlx::query(&format!( + "UPDATE {table} SET status = $1, error_message = $2, completed_at = NOW() WHERE id = $3" + )) + .bind(status) + .bind(error) + .bind(id) + .execute(self.db.pool()) + .await?; + } + + result + } + + pub async fn delete(&self, file_uuid: &str) -> Result<()> { + let log_id = TkgLogger::start_operation(self.db.pool(), file_uuid, "delete").await?; + + let result = self.delete_internal(file_uuid).await; + + match result { + Ok(_) => { + TkgLogger::new(self.db.pool().clone()).complete_operation(None).await?; + Ok(()) + } + Err(e) => { + TkgLogger::new(self.db.pool().clone()).complete_operation(Some(&e.to_string())).await?; + Err(e) + } + } + } + + async fn delete_internal(&self, file_uuid: &str) -> Result<()> { + let nodes_table = schema::table_name("tkg_nodes"); + let edges_table = schema::table_name("tkg_edges"); + + // Delete edges first (foreign key constraint) + let edges_deleted = sqlx::query(&format!( + "DELETE FROM {edges_table} WHERE file_uuid = $1" + )) + .bind(file_uuid) + .execute(self.db.pool()) + .await?; + + // Delete nodes + let nodes_deleted = sqlx::query(&format!( + "DELETE FROM {nodes_table} WHERE file_uuid = $1" + )) + .bind(file_uuid) + .execute(self.db.pool()) + .await?; + + tracing::info!("[TKG-Service] Deleted {} nodes and {} edges for {}", + nodes_deleted.rows_affected(), edges_deleted.rows_affected(), file_uuid); + + Ok(()) + } + + pub async fn get_operations(&self, file_uuid: &str) -> Result> { + TkgLogger::get_operations(self.db.pool(), file_uuid).await + } + + pub async fn delete_tkg_and_logs(&self, file_uuid: &str) -> Result<()> { + self.delete_internal(file_uuid).await?; + TkgLogger::delete_operations(self.db.pool(), file_uuid).await?; + Ok(()) + } +}