3199 lines
134 KiB
Rust
3199 lines
134 KiB
Rust
use anyhow::{Context, Result};
|
|
use clap::{Parser, Subcommand};
|
|
use futures_util::StreamExt;
|
|
use std::path::Path;
|
|
use std::str;
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
use momentry_core::core::api_key::{ApiKeyService, ApiKeyType};
|
|
use momentry_core::core::chunk::types::{Chunk, ChunkRule, ChunkType};
|
|
use momentry_core::core::db::Database;
|
|
use momentry_core::core::time::FrameTime;
|
|
use momentry_core::ui::progress::{ProcessorType, ProgressState, ProgressUi};
|
|
use momentry_core::{
|
|
Embedder, OutputDir, PostgresDb, QdrantDb, RedisClient, VectorPayload, VideoRecord, VideoStatus,
|
|
};
|
|
|
|
fn parse_key_type(s: Option<&str>) -> ApiKeyType {
|
|
match s.map(|s| s.to_lowercase()).as_deref() {
|
|
Some("system") => ApiKeyType::System,
|
|
Some("user") => ApiKeyType::User,
|
|
Some("service") => ApiKeyType::Service,
|
|
Some("integration") => ApiKeyType::Integration,
|
|
Some("emergency") => ApiKeyType::Emergency,
|
|
_ => ApiKeyType::User,
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
|
pub enum ProcessingDecision {
|
|
Process,
|
|
SkipComplete,
|
|
ResumePartial,
|
|
ForceReprocess,
|
|
}
|
|
|
|
impl ProcessingDecision {
|
|
pub fn should_process(&self) -> bool {
|
|
matches!(
|
|
self,
|
|
ProcessingDecision::Process
|
|
| ProcessingDecision::ResumePartial
|
|
| ProcessingDecision::ForceReprocess
|
|
)
|
|
}
|
|
|
|
pub fn should_resume(&self) -> bool {
|
|
matches!(self, ProcessingDecision::ResumePartial)
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct SystemResources {
|
|
pub cpu_idle_percent: f64,
|
|
pub memory_available_mb: u64,
|
|
pub memory_total_mb: u64,
|
|
pub memory_used_percent: f64,
|
|
pub gpu_available: bool,
|
|
pub gpu_type: GpuType,
|
|
pub gpu_utilization: Option<f64>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub enum GpuType {
|
|
Nvidia,
|
|
AppleMps,
|
|
}
|
|
|
|
impl SystemResources {
|
|
pub fn check() -> Self {
|
|
let cpu_idle = Self::get_cpu_idle();
|
|
let (mem_available, mem_total) = Self::get_memory_info();
|
|
let mem_used_pct = if mem_total > 0 && mem_available <= mem_total {
|
|
((mem_total - mem_available) as f64 / mem_total as f64) * 100.0
|
|
} else if mem_total > 0 {
|
|
100.0
|
|
} else {
|
|
0.0
|
|
};
|
|
let (gpu_available, gpu_type, gpu_util) = Self::get_gpu_info();
|
|
|
|
Self {
|
|
cpu_idle_percent: cpu_idle,
|
|
memory_available_mb: mem_available,
|
|
memory_total_mb: mem_total,
|
|
memory_used_percent: mem_used_pct,
|
|
gpu_available,
|
|
gpu_type,
|
|
gpu_utilization: gpu_util,
|
|
}
|
|
}
|
|
|
|
pub fn can_parallel(&self, required_memory_mb: u64) -> bool {
|
|
const MIN_CPU_IDLE: f64 = 30.0;
|
|
const MIN_MEMORY_MB: u64 = 4096;
|
|
|
|
self.cpu_idle_percent >= MIN_CPU_IDLE
|
|
&& self.memory_available_mb >= required_memory_mb
|
|
&& self.memory_available_mb >= MIN_MEMORY_MB
|
|
}
|
|
|
|
pub fn recommend_parallel_modules(&self) -> Vec<&'static str> {
|
|
let mut recommended = Vec::new();
|
|
|
|
if self.gpu_available {
|
|
recommended.push("yolo");
|
|
}
|
|
|
|
if self.memory_available_mb >= 8192 {
|
|
recommended.push("ocr");
|
|
recommended.push("face");
|
|
recommended.push("pose");
|
|
}
|
|
|
|
recommended
|
|
}
|
|
|
|
fn get_cpu_idle() -> f64 {
|
|
use std::process::Command;
|
|
let output = Command::new("top").args(["-l", "1", "-n", "1"]).output();
|
|
match output {
|
|
Ok(o) => {
|
|
let s = String::from_utf8_lossy(&o.stdout);
|
|
if let Some(line) = s.lines().find(|l| l.contains("idle")) {
|
|
if let Some(pct) = line
|
|
.split_whitespace()
|
|
.find_map(|s| s.strip_suffix("%idle"))
|
|
{
|
|
pct.trim().parse().ok().unwrap_or(50.0)
|
|
} else {
|
|
50.0
|
|
}
|
|
} else {
|
|
50.0
|
|
}
|
|
}
|
|
Err(_) => 50.0,
|
|
}
|
|
}
|
|
|
|
fn get_memory_info() -> (u64, u64) {
|
|
use std::process::Command;
|
|
let output = Command::new("sysctl").args(["hw.memsize"]).output();
|
|
match output {
|
|
Ok(o) => {
|
|
let s = String::from_utf8_lossy(&o.stdout);
|
|
let total = s
|
|
.split_whitespace()
|
|
.nth(1)
|
|
.and_then(|v| v.parse::<u64>().ok())
|
|
.unwrap_or(0)
|
|
/ 1024
|
|
/ 1024;
|
|
|
|
let vm_stat = Command::new("vm_stat").output();
|
|
let available = match vm_stat {
|
|
Ok(v) => {
|
|
let vs = String::from_utf8_lossy(&v.stdout);
|
|
let mut free_pages: u64 = 0;
|
|
let mut inactive_pages: u64 = 0;
|
|
|
|
for line in vs.lines() {
|
|
if line.contains("Pages free:") {
|
|
free_pages = line
|
|
.split_whitespace()
|
|
.last()
|
|
.and_then(|v| v.trim_end_matches('.').parse().ok())
|
|
.unwrap_or(0);
|
|
} else if line.contains("Pages inactive:") {
|
|
inactive_pages = line
|
|
.split_whitespace()
|
|
.last()
|
|
.and_then(|v| v.trim_end_matches('.').parse().ok())
|
|
.unwrap_or(0);
|
|
}
|
|
}
|
|
|
|
// Pages * 4096 bytes / 1024 / 1024 = MB
|
|
(free_pages + inactive_pages) * 4096 / 1024 / 1024
|
|
}
|
|
Err(_) => total / 4,
|
|
};
|
|
|
|
(available, total)
|
|
}
|
|
Err(_) => (0, 0),
|
|
}
|
|
}
|
|
|
|
fn get_gpu_info() -> (bool, GpuType, Option<f64>) {
|
|
use std::process::Command;
|
|
|
|
// Check NVIDIA GPU
|
|
let nvidia_output = Command::new("nvidia-smi")
|
|
.args([
|
|
"--query-gpu=utilization.gpu",
|
|
"--format=csv,noheader,nounits",
|
|
])
|
|
.output();
|
|
|
|
if let Ok(o) = nvidia_output {
|
|
if o.status.success() {
|
|
let s = String::from_utf8_lossy(&o.stdout);
|
|
let util = s.trim().parse::<f64>().ok();
|
|
return (true, GpuType::Nvidia, util);
|
|
}
|
|
}
|
|
|
|
// Check Apple MPS (Metal Performance Shaders)
|
|
let mps_output = Command::new("system_profiler")
|
|
.args(["SPDisplaysDataType", "-detailLevel", "mini"])
|
|
.output();
|
|
|
|
if let Ok(o) = mps_output {
|
|
let s = String::from_utf8_lossy(&o.stdout);
|
|
if s.contains("Metal") || s.contains("Apple") {
|
|
return (true, GpuType::AppleMps, Some(0.0));
|
|
}
|
|
}
|
|
|
|
(false, GpuType::Nvidia, None)
|
|
}
|
|
}
|
|
|
|
impl std::fmt::Display for SystemResources {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
write!(
|
|
f,
|
|
"CPU: {:.1}% idle, Memory: {:.1}GB/{:.1}GB ({:.0}% used), GPU: {}",
|
|
self.cpu_idle_percent,
|
|
self.memory_available_mb as f64 / 1024.0,
|
|
self.memory_total_mb as f64 / 1024.0,
|
|
self.memory_used_percent,
|
|
if self.gpu_available {
|
|
format!("{:.0}% utilized", self.gpu_utilization.unwrap_or(0.0))
|
|
} else {
|
|
"N/A".to_string()
|
|
}
|
|
)
|
|
}
|
|
}
|
|
|
|
fn decide_processing(json_path: &Path, force: bool, resume: bool) -> ProcessingDecision {
|
|
if !json_path.exists() {
|
|
return ProcessingDecision::Process;
|
|
}
|
|
|
|
if force {
|
|
return ProcessingDecision::ForceReprocess;
|
|
}
|
|
|
|
if resume {
|
|
return ProcessingDecision::ResumePartial;
|
|
}
|
|
|
|
match check_json_completeness(json_path) {
|
|
JsonCompleteness::Complete => ProcessingDecision::SkipComplete,
|
|
JsonCompleteness::Partial { current, total } => {
|
|
eprintln!("\n⚠️ Found incomplete JSON file: {}", json_path.display());
|
|
eprintln!(
|
|
" Progress: {}/{} ({:.1}%)",
|
|
current,
|
|
total,
|
|
(current as f64 / total as f64) * 100.0
|
|
);
|
|
eprintln!(" Use --resume to continue from checkpoint");
|
|
eprintln!(" Use --force to reprocess from scratch");
|
|
ProcessingDecision::SkipComplete
|
|
}
|
|
JsonCompleteness::Empty => ProcessingDecision::Process,
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub enum JsonCompleteness {
|
|
Complete,
|
|
Partial { current: u32, total: u32 },
|
|
Empty,
|
|
}
|
|
|
|
fn check_json_completeness(json_path: &Path) -> JsonCompleteness {
|
|
let content = match std::fs::read_to_string(json_path) {
|
|
Ok(c) => c,
|
|
Err(_) => return JsonCompleteness::Empty,
|
|
};
|
|
|
|
if content.trim().is_empty() {
|
|
return JsonCompleteness::Empty;
|
|
}
|
|
|
|
let json: serde_json::Value = match serde_json::from_str(&content) {
|
|
Ok(v) => v,
|
|
Err(_) => return JsonCompleteness::Empty,
|
|
};
|
|
|
|
match json.get("segments") {
|
|
Some(serde_json::Value::Array(arr)) if !arr.is_empty() => JsonCompleteness::Complete,
|
|
Some(serde_json::Value::Object(obj)) => {
|
|
let current = obj.get("current").and_then(|v| v.as_u64()).unwrap_or(0) as u32;
|
|
let total = obj.get("total").and_then(|v| v.as_u64()).unwrap_or(0) as u32;
|
|
if total > 0 && current < total {
|
|
JsonCompleteness::Partial { current, total }
|
|
} else {
|
|
JsonCompleteness::Complete
|
|
}
|
|
}
|
|
_ => JsonCompleteness::Complete,
|
|
}
|
|
}
|
|
|
|
async fn process_asr_module(
|
|
asr_path: &Path,
|
|
video_path: &str,
|
|
uuid: &str,
|
|
progress_state: &Arc<Mutex<ProgressState>>,
|
|
ui: &Arc<Mutex<Option<ProgressUi>>>,
|
|
) -> anyhow::Result<()> {
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state.get_processor(ProcessorType::Asr).start(1);
|
|
}
|
|
let asr_result = momentry_core::core::processor::process_asr(
|
|
video_path,
|
|
asr_path.to_str().unwrap(),
|
|
Some(uuid),
|
|
)
|
|
.await?;
|
|
let asr_json = serde_json::to_string_pretty(&asr_result)?;
|
|
std::fs::write(asr_path, &asr_json)?;
|
|
let output_dir = OutputDir::new();
|
|
let _ = output_dir.backup_file(uuid, "asr.json");
|
|
println!(" ✓ ASR saved: {} segments", asr_result.segments.len());
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state
|
|
.get_processor(ProcessorType::Asr)
|
|
.complete(&format!("{} segments", asr_result.segments.len()));
|
|
}
|
|
if let Some(ref mut ui) = *ui.lock().unwrap() {
|
|
let _ = ui.render();
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn process_cut_module(
|
|
cut_path: &Path,
|
|
video_path: &str,
|
|
uuid: &str,
|
|
progress_state: &Arc<Mutex<ProgressState>>,
|
|
ui: &Arc<Mutex<Option<ProgressUi>>>,
|
|
) -> anyhow::Result<()> {
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state.get_processor(ProcessorType::Cut).start(1);
|
|
}
|
|
let cut_result = momentry_core::core::processor::process_cut(
|
|
video_path,
|
|
cut_path.to_str().unwrap(),
|
|
Some(uuid),
|
|
)
|
|
.await?;
|
|
let cut_json = serde_json::to_string_pretty(&cut_result)?;
|
|
std::fs::write(cut_path, &cut_json)?;
|
|
let output_dir = OutputDir::new();
|
|
let _ = output_dir.backup_file(uuid, "cut.json");
|
|
println!(" ✓ CUT saved: {} scenes", cut_result.scenes.len());
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state
|
|
.get_processor(ProcessorType::Cut)
|
|
.complete(&format!("{} scenes", cut_result.scenes.len()));
|
|
}
|
|
if let Some(ref mut ui) = *ui.lock().unwrap() {
|
|
let _ = ui.render();
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn process_asrx_module(
|
|
asrx_path: &Path,
|
|
video_path: &str,
|
|
uuid: &str,
|
|
progress_state: &Arc<Mutex<ProgressState>>,
|
|
ui: &Arc<Mutex<Option<ProgressUi>>>,
|
|
) -> anyhow::Result<()> {
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state.get_processor(ProcessorType::Asrx).start(1);
|
|
}
|
|
let asrx_result = momentry_core::core::processor::process_asrx(
|
|
video_path,
|
|
asrx_path.to_str().unwrap(),
|
|
Some(uuid),
|
|
)
|
|
.await?;
|
|
let asrx_json = serde_json::to_string_pretty(&asrx_result)?;
|
|
std::fs::write(asrx_path, &asrx_json)?;
|
|
let output_dir = OutputDir::new();
|
|
let _ = output_dir.backup_file(uuid, "asrx.json");
|
|
println!(" ✓ ASRX saved: {} segments", asrx_result.segments.len());
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state
|
|
.get_processor(ProcessorType::Asrx)
|
|
.complete(&format!("{} segments", asrx_result.segments.len()));
|
|
}
|
|
if let Some(ref mut ui) = *ui.lock().unwrap() {
|
|
let _ = ui.render();
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn process_yolo_module(
|
|
yolo_path: &Path,
|
|
video_path: &str,
|
|
uuid: &str,
|
|
progress_state: &Arc<Mutex<ProgressState>>,
|
|
ui: &Arc<Mutex<Option<ProgressUi>>>,
|
|
) -> anyhow::Result<()> {
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state.get_processor(ProcessorType::Yolo).start(1);
|
|
}
|
|
let yolo_result = momentry_core::core::processor::process_yolo(
|
|
video_path,
|
|
yolo_path.to_str().unwrap(),
|
|
Some(uuid),
|
|
)
|
|
.await?;
|
|
let yolo_json = serde_json::to_string_pretty(&yolo_result)?;
|
|
std::fs::write(yolo_path, &yolo_json)?;
|
|
let output_dir = OutputDir::new();
|
|
let _ = output_dir.backup_file(uuid, "yolo.json");
|
|
println!(" ✓ YOLO saved: {} frames", yolo_result.frame_count);
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state
|
|
.get_processor(ProcessorType::Yolo)
|
|
.complete(&format!("{} frames", yolo_result.frame_count));
|
|
}
|
|
if let Some(ref mut ui) = *ui.lock().unwrap() {
|
|
let _ = ui.render();
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn process_ocr_module(
|
|
ocr_path: &Path,
|
|
video_path: &str,
|
|
uuid: &str,
|
|
progress_state: &Arc<Mutex<ProgressState>>,
|
|
ui: &Arc<Mutex<Option<ProgressUi>>>,
|
|
) -> anyhow::Result<()> {
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state.get_processor(ProcessorType::Ocr).start(1);
|
|
}
|
|
let ocr_result = momentry_core::core::processor::process_ocr(
|
|
video_path,
|
|
ocr_path.to_str().unwrap(),
|
|
Some(uuid),
|
|
)
|
|
.await?;
|
|
let ocr_json = serde_json::to_string_pretty(&ocr_result)?;
|
|
std::fs::write(ocr_path, &ocr_json)?;
|
|
let output_dir = OutputDir::new();
|
|
let _ = output_dir.backup_file(uuid, "ocr.json");
|
|
println!(
|
|
" ✓ OCR saved: {} frames with text",
|
|
ocr_result.frames.len()
|
|
);
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state
|
|
.get_processor(ProcessorType::Ocr)
|
|
.complete(&format!("{} frames", ocr_result.frames.len()));
|
|
}
|
|
if let Some(ref mut ui) = *ui.lock().unwrap() {
|
|
let _ = ui.render();
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn process_face_module(
|
|
face_path: &Path,
|
|
video_path: &str,
|
|
uuid: &str,
|
|
progress_state: &Arc<Mutex<ProgressState>>,
|
|
ui: &Arc<Mutex<Option<ProgressUi>>>,
|
|
) -> anyhow::Result<()> {
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state.get_processor(ProcessorType::Face).start(1);
|
|
}
|
|
let face_result = momentry_core::core::processor::process_face(
|
|
video_path,
|
|
face_path.to_str().unwrap(),
|
|
Some(uuid),
|
|
)
|
|
.await?;
|
|
let face_json = serde_json::to_string_pretty(&face_result)?;
|
|
std::fs::write(face_path, &face_json)?;
|
|
let output_dir = OutputDir::new();
|
|
let _ = output_dir.backup_file(uuid, "face.json");
|
|
println!(" ✓ Face saved: {} frames", face_result.frames.len());
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state
|
|
.get_processor(ProcessorType::Face)
|
|
.complete(&format!("{} frames", face_result.frames.len()));
|
|
}
|
|
if let Some(ref mut ui) = *ui.lock().unwrap() {
|
|
let _ = ui.render();
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn process_pose_module(
|
|
pose_path: &Path,
|
|
video_path: &str,
|
|
uuid: &str,
|
|
progress_state: &Arc<Mutex<ProgressState>>,
|
|
ui: &Arc<Mutex<Option<ProgressUi>>>,
|
|
) -> anyhow::Result<()> {
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state.get_processor(ProcessorType::Pose).start(1);
|
|
}
|
|
let pose_result = momentry_core::core::processor::process_pose(
|
|
video_path,
|
|
pose_path.to_str().unwrap(),
|
|
Some(uuid),
|
|
)
|
|
.await?;
|
|
let pose_json = serde_json::to_string_pretty(&pose_result)?;
|
|
std::fs::write(pose_path, &pose_json)?;
|
|
let output_dir = OutputDir::new();
|
|
let _ = output_dir.backup_file(uuid, "pose.json");
|
|
println!(" ✓ Pose saved: {} frames", pose_result.frames.len());
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state
|
|
.get_processor(ProcessorType::Pose)
|
|
.complete(&format!("{} frames", pose_result.frames.len()));
|
|
state.stop();
|
|
}
|
|
if let Some(ref mut ui) = *ui.lock().unwrap() {
|
|
let _ = ui.render();
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn process_story_module(
|
|
story_path: &Path,
|
|
video_path: &str,
|
|
uuid: &str,
|
|
progress_state: &Arc<Mutex<ProgressState>>,
|
|
ui: &Arc<Mutex<Option<ProgressUi>>>,
|
|
) -> anyhow::Result<()> {
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state.get_processor(ProcessorType::Story).start(1);
|
|
}
|
|
let story_result = momentry_core::core::processor::process_story(
|
|
video_path,
|
|
story_path.to_str().unwrap(),
|
|
Some(uuid),
|
|
)
|
|
.await?;
|
|
let story_json = serde_json::to_string_pretty(&story_result)?;
|
|
std::fs::write(story_path, &story_json)?;
|
|
let output_dir = OutputDir::new();
|
|
let _ = output_dir.backup_file(uuid, "story.json");
|
|
println!(
|
|
" ✓ Story saved: {} parent chunks, {} child chunks",
|
|
story_result.stats.total_parent_chunks, story_result.stats.total_child_chunks
|
|
);
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state.get_processor(ProcessorType::Story).complete(&format!(
|
|
"{} parents, {} children",
|
|
story_result.stats.total_parent_chunks, story_result.stats.total_child_chunks
|
|
));
|
|
}
|
|
if let Some(ref mut ui) = *ui.lock().unwrap() {
|
|
let _ = ui.render();
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn process_caption_module(
|
|
caption_path: &Path,
|
|
video_path: &str,
|
|
uuid: &str,
|
|
progress_state: &Arc<Mutex<ProgressState>>,
|
|
ui: &Arc<Mutex<Option<ProgressUi>>>,
|
|
) -> anyhow::Result<()> {
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state.get_processor(ProcessorType::Caption).start(1);
|
|
}
|
|
let caption_result = momentry_core::core::processor::process_caption(
|
|
video_path,
|
|
caption_path.to_str().unwrap(),
|
|
Some(uuid),
|
|
)
|
|
.await?;
|
|
let caption_json = serde_json::to_string_pretty(&caption_result)?;
|
|
std::fs::write(caption_path, &caption_json)?;
|
|
let output_dir = OutputDir::new();
|
|
let _ = output_dir.backup_file(uuid, "caption.json");
|
|
println!(" ✓ Caption saved: {} frames", caption_result.total_frames);
|
|
{
|
|
let mut state = progress_state.lock().unwrap();
|
|
state
|
|
.get_processor(ProcessorType::Caption)
|
|
.complete(&format!("{} frames", caption_result.total_frames));
|
|
}
|
|
if let Some(ref mut ui) = *ui.lock().unwrap() {
|
|
let _ = ui.render();
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[derive(Parser)]
|
|
#[command(name = "momentry")]
|
|
#[command(about = "Digital asset management system with video analysis and RAG")]
|
|
struct Cli {
|
|
#[command(subcommand)]
|
|
command: Commands,
|
|
}
|
|
|
|
#[derive(Subcommand)]
|
|
enum Commands {
|
|
/// Register a video file
|
|
Register {
|
|
/// Video file path or URL
|
|
path: String,
|
|
},
|
|
/// Process video (generate all JSON files)
|
|
Process {
|
|
/// UUID or path
|
|
target: String,
|
|
/// Modules to process (comma separated: asr,cut,asrx,yolo,ocr,face,pose,story,caption)
|
|
/// If not specified, processes all modules
|
|
#[arg(short, long, value_delimiter = ',')]
|
|
modules: Option<Vec<String>>,
|
|
/// Modules to process via cloud (comma separated)
|
|
/// Example: --cloud asr,yolo
|
|
#[arg(long, value_delimiter = ',')]
|
|
cloud: Option<Vec<String>>,
|
|
/// Force reprocess even if JSON exists (skip completeness check)
|
|
#[arg(long, default_value = "false")]
|
|
force: bool,
|
|
/// Resume from last checkpoint if processing was interrupted
|
|
#[arg(long, default_value = "false")]
|
|
resume: bool,
|
|
},
|
|
/// Generate chunks and store in database
|
|
Chunk {
|
|
/// UUID
|
|
uuid: String,
|
|
},
|
|
/// Generate story for cut scenes
|
|
Story {
|
|
/// UUID
|
|
uuid: String,
|
|
},
|
|
/// Vectorize chunks
|
|
Vectorize {
|
|
/// UUID (or 'all' for all)
|
|
uuid: String,
|
|
},
|
|
/// Play video with overlays
|
|
Play {
|
|
/// Video path or UUID
|
|
target: String,
|
|
},
|
|
/// Start watching directories
|
|
Watch {
|
|
/// Directories to watch (comma separated)
|
|
directories: Option<String>,
|
|
},
|
|
/// Check system resources and recommend processing strategy
|
|
System {
|
|
/// Show detailed GPU info (NVIDIA/MPS)
|
|
#[arg(long)]
|
|
gpu: bool,
|
|
},
|
|
/// Start API server
|
|
Server {
|
|
/// Host
|
|
#[arg(long, default_value = "127.0.0.1")]
|
|
host: String,
|
|
/// Port (defaults to MOMENTRY_SERVER_PORT env var, or3002 for production)
|
|
#[arg(long)]
|
|
port: Option<u16>,
|
|
},
|
|
/// Start job worker
|
|
Worker {
|
|
/// Max concurrent processors
|
|
#[arg(long)]
|
|
max_concurrent: Option<usize>,
|
|
/// Poll interval in seconds
|
|
#[arg(long)]
|
|
poll_interval: Option<u64>,
|
|
/// Batch size
|
|
#[arg(long)]
|
|
batch_size: Option<i32>,
|
|
},
|
|
/// Query using RAG
|
|
Query {
|
|
/// Query text
|
|
query: String,
|
|
},
|
|
/// Lookup UUID from path
|
|
Lookup {
|
|
/// File path
|
|
path: String,
|
|
},
|
|
/// Resolve path from UUID
|
|
Resolve {
|
|
/// UUID
|
|
uuid: String,
|
|
},
|
|
/// Generate thumbnails for videos
|
|
Thumbnails {
|
|
/// UUID (optional, generates for all if not specified)
|
|
uuid: Option<String>,
|
|
/// Number of thumbnails per video
|
|
#[arg(short, long, default_value = "6")]
|
|
count: u32,
|
|
},
|
|
/// Show storage status report
|
|
Status {
|
|
/// UUID (optional, shows all if not specified)
|
|
uuid: Option<String>,
|
|
},
|
|
/// Manage output backups
|
|
Backup {
|
|
/// Action: list, cleanup
|
|
action: String,
|
|
/// Days to keep (for cleanup)
|
|
days: Option<u32>,
|
|
},
|
|
/// Manage API keys
|
|
ApiKey {
|
|
/// Action: create, list, validate, revoke, rotate, stats
|
|
#[arg(value_enum)]
|
|
action: ApiKeyAction,
|
|
/// Key name (for create)
|
|
name: Option<String>,
|
|
/// Key type (system, user, service, integration, emergency)
|
|
#[arg(long)]
|
|
key_type: Option<String>,
|
|
/// TTL in days (for create)
|
|
#[arg(long)]
|
|
ttl: Option<i64>,
|
|
/// API key to validate/revoke
|
|
#[arg(long)]
|
|
key: Option<String>,
|
|
},
|
|
/// Manage Gitea API tokens
|
|
Gitea {
|
|
/// Action: create, list, delete, verify
|
|
#[arg(value_enum)]
|
|
action: GiteaAction,
|
|
/// Gitea username
|
|
#[arg(long)]
|
|
username: Option<String>,
|
|
/// Gitea password (for create/list/delete)
|
|
#[arg(long)]
|
|
password: Option<String>,
|
|
/// Token name (for create/delete)
|
|
#[arg(long)]
|
|
token_name: Option<String>,
|
|
/// Token scopes (comma separated: read:repository,write:issue)
|
|
#[arg(long)]
|
|
scopes: Option<String>,
|
|
},
|
|
/// Manage n8n API keys
|
|
N8n {
|
|
/// Action: create, list, delete, verify
|
|
#[arg(value_enum)]
|
|
action: N8nAction,
|
|
/// n8n API key (for create/list/delete)
|
|
#[arg(long)]
|
|
api_key: Option<String>,
|
|
/// API key label (for create/delete)
|
|
#[arg(long)]
|
|
label: Option<String>,
|
|
/// Expiration days (for create)
|
|
#[arg(long)]
|
|
expires_in_days: Option<i64>,
|
|
},
|
|
}
|
|
|
|
#[derive(clap::ValueEnum, Clone)]
|
|
enum ApiKeyAction {
|
|
Create,
|
|
List,
|
|
Validate,
|
|
Revoke,
|
|
Rotate,
|
|
Stats,
|
|
}
|
|
|
|
#[derive(clap::ValueEnum, Clone)]
|
|
enum GiteaAction {
|
|
Create,
|
|
List,
|
|
Delete,
|
|
Verify,
|
|
}
|
|
|
|
#[derive(clap::ValueEnum, Clone)]
|
|
enum N8nAction {
|
|
Create,
|
|
List,
|
|
Delete,
|
|
Verify,
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<()> {
|
|
dotenv::dotenv().ok();
|
|
tracing_subscriber::fmt::init();
|
|
|
|
let cli = Cli::parse();
|
|
|
|
match cli.command {
|
|
Commands::Register { path } => {
|
|
println!("Registering: {}", path);
|
|
|
|
// Compute UUID
|
|
let uuid = momentry_core::uuid::compute_uuid_from_path(&path);
|
|
println!("UUID: {}", uuid);
|
|
|
|
// Run ffprobe
|
|
let probe_result = momentry_core::core::probe::probe_video(&path)?;
|
|
|
|
println!("\nVideo probe results:");
|
|
let duration = probe_result
|
|
.format
|
|
.duration
|
|
.as_ref()
|
|
.and_then(|s| s.parse::<f64>().ok())
|
|
.unwrap_or(0.0);
|
|
println!(" Duration: {}s", duration);
|
|
if let Some(size) = &probe_result.format.size {
|
|
println!(" Size: {}", size);
|
|
}
|
|
|
|
let mut width = 0u32;
|
|
let mut height = 0u32;
|
|
let mut fps = 0.0;
|
|
|
|
for stream in &probe_result.streams {
|
|
if stream.codec_type.as_deref() == Some("video") {
|
|
width = stream.width.unwrap_or(0);
|
|
height = stream.height.unwrap_or(0);
|
|
if let Some(fps_str) = &stream.r_frame_rate {
|
|
if let Some((num, den)) = fps_str.split_once('/') {
|
|
if let (Ok(n), Ok(d)) = (num.parse::<f64>(), den.parse::<f64>()) {
|
|
if d > 0.0 {
|
|
fps = n / d;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
println!(" Video: {}x{}", width, height);
|
|
if let Some(fps) = &stream.r_frame_rate {
|
|
println!(" FPS: {}", fps);
|
|
}
|
|
}
|
|
if stream.codec_type.as_deref() == Some("audio") {
|
|
println!(" Audio: {} channels", stream.channels.unwrap_or(0));
|
|
if let Some(sr) = &stream.sample_rate {
|
|
println!(" Sample Rate: {}", sr);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Save probe JSON to file
|
|
let file_manager = momentry_core::FileManager::new(std::path::PathBuf::from("."));
|
|
let json_str = serde_json::to_string_pretty(&probe_result)?;
|
|
let json_path = file_manager.save_json(&uuid, "probe", &json_str)?;
|
|
println!("\nProbe JSON saved to: {:?}", json_path);
|
|
|
|
// Store in PostgreSQL
|
|
println!("\nStoring in database...");
|
|
let db = PostgresDb::init().await?;
|
|
let file_path = Path::new(&path)
|
|
.canonicalize()
|
|
.map(|p| p.to_string_lossy().to_string())
|
|
.unwrap_or_else(|_| path.clone());
|
|
let file_name = Path::new(&path)
|
|
.file_name()
|
|
.map(|n| n.to_string_lossy().to_string())
|
|
.unwrap_or_default();
|
|
|
|
let record = VideoRecord {
|
|
id: 0,
|
|
uuid: uuid.clone(),
|
|
file_path,
|
|
file_name,
|
|
duration,
|
|
width,
|
|
height,
|
|
fps,
|
|
probe_json: Some(json_str),
|
|
storage: Default::default(),
|
|
status: VideoStatus::Pending,
|
|
user_id: None,
|
|
job_id: None,
|
|
created_at: String::new(),
|
|
};
|
|
|
|
let video_id = db.register_video(&record).await?;
|
|
println!("Video registered with ID: {}", video_id);
|
|
|
|
Ok(())
|
|
}
|
|
Commands::Process {
|
|
target,
|
|
modules,
|
|
cloud,
|
|
force,
|
|
resume,
|
|
} => {
|
|
println!("Processing: {}", target);
|
|
println!(" force: {}, resume: {}", force, resume);
|
|
|
|
// Parse selected modules
|
|
let selected_modules: Option<Vec<ProcessorType>> = modules.as_ref().map(|m| {
|
|
m.iter()
|
|
.filter_map(|name| {
|
|
let name_lower = name.to_lowercase();
|
|
match name_lower.as_str() {
|
|
"asr" => Some(ProcessorType::Asr),
|
|
"cut" => Some(ProcessorType::Cut),
|
|
"asrx" => Some(ProcessorType::Asrx),
|
|
"yolo" => Some(ProcessorType::Yolo),
|
|
"ocr" => Some(ProcessorType::Ocr),
|
|
"face" => Some(ProcessorType::Face),
|
|
"pose" => Some(ProcessorType::Pose),
|
|
"story" => Some(ProcessorType::Story),
|
|
"caption" => Some(ProcessorType::Caption),
|
|
_ => {
|
|
eprintln!("Unknown module: {}", name);
|
|
None
|
|
}
|
|
}
|
|
})
|
|
.collect()
|
|
});
|
|
|
|
// Parse cloud modules
|
|
let cloud_modules: Vec<ProcessorType> = cloud
|
|
.as_ref()
|
|
.map(|c| {
|
|
c.iter()
|
|
.filter_map(|name| {
|
|
let name_lower = name.to_lowercase();
|
|
match name_lower.as_str() {
|
|
"asr" => Some(ProcessorType::Asr),
|
|
"cut" => Some(ProcessorType::Cut),
|
|
"asrx" => Some(ProcessorType::Asrx),
|
|
"yolo" => Some(ProcessorType::Yolo),
|
|
"ocr" => Some(ProcessorType::Ocr),
|
|
"face" => Some(ProcessorType::Face),
|
|
"pose" => Some(ProcessorType::Pose),
|
|
"story" => Some(ProcessorType::Story),
|
|
"caption" => Some(ProcessorType::Caption),
|
|
_ => {
|
|
eprintln!("Unknown cloud module: {}", name);
|
|
None
|
|
}
|
|
}
|
|
})
|
|
.collect()
|
|
})
|
|
.unwrap_or_default();
|
|
|
|
if let Some(ref mods) = selected_modules {
|
|
println!(
|
|
" Modules: {}",
|
|
mods.iter()
|
|
.map(|m| m.to_string())
|
|
.collect::<Vec<_>>()
|
|
.join(", ")
|
|
);
|
|
} else {
|
|
println!(" Modules: ALL");
|
|
}
|
|
|
|
if !cloud_modules.is_empty() {
|
|
println!(
|
|
" Cloud: {}",
|
|
cloud_modules
|
|
.iter()
|
|
.map(|m| m.to_string())
|
|
.collect::<Vec<_>>()
|
|
.join(", ")
|
|
);
|
|
}
|
|
|
|
let processing_mode = if force {
|
|
"FORCE (reprocess all)"
|
|
} else if resume {
|
|
"RESUME (continue from checkpoint)"
|
|
} else {
|
|
"SMART (skip complete, resume partial)"
|
|
};
|
|
println!(" Mode: {}", processing_mode);
|
|
|
|
// Compute UUID if path is given
|
|
let uuid = if target.len() == 16 && !target.contains('/') {
|
|
target.clone()
|
|
} else {
|
|
momentry_core::uuid::compute_uuid_from_path(&target)
|
|
};
|
|
|
|
// Get video from database
|
|
let db = PostgresDb::init().await?;
|
|
let video = db
|
|
.get_video_by_uuid(&uuid)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("Video not found: {}", uuid))?;
|
|
|
|
let video_path = &video.file_path;
|
|
let video_name = video.file_name.clone();
|
|
let _file_manager = momentry_core::FileManager::new(std::path::PathBuf::from("."));
|
|
|
|
// Initialize output directory
|
|
let output_dir = OutputDir::new();
|
|
output_dir.ensure_dir()?;
|
|
println!("Output directory: {:?}", output_dir.get_base_path());
|
|
|
|
// Initialize progress UI
|
|
let progress_state = Arc::new(Mutex::new(ProgressState::new(&video_name)));
|
|
progress_state.lock().unwrap().start();
|
|
|
|
// Helper closure to check if a module should be processed
|
|
let should_process = |module: ProcessorType| -> bool {
|
|
selected_modules
|
|
.as_ref()
|
|
.map(|mods| mods.contains(&module))
|
|
.unwrap_or(true)
|
|
};
|
|
|
|
// Helper closure to check if a module should run in the cloud
|
|
let is_cloud = |module: ProcessorType| -> bool { cloud_modules.contains(&module) };
|
|
|
|
// Create UI and wrap in Arc for sharing with Redis subscriber
|
|
let ui = Arc::new(Mutex::new(ProgressUi::new(&video_name).ok()));
|
|
if let Some(ref mut ui) = *ui.lock().unwrap() {
|
|
let _ = ui.render();
|
|
}
|
|
|
|
// Spawn Redis subscriber for real-time progress updates
|
|
let redis_progress_state = progress_state.clone();
|
|
let redis_ui = ui.clone();
|
|
let redis_uuid = uuid.clone();
|
|
let redis_handle = tokio::spawn(async move {
|
|
if let Ok(redis_client) = momentry_core::core::db::RedisClient::new() {
|
|
loop {
|
|
if let Ok(mut pubsub) = redis_client.subscribe_progress(&redis_uuid).await {
|
|
let mut stream = pubsub.on_message();
|
|
while let Some(msg) = stream.next().await {
|
|
if let Ok(payload) = msg.get_payload::<String>() {
|
|
if let Ok(progress_msg) =
|
|
serde_json::from_str::<
|
|
momentry_core::core::db::ProgressMessage,
|
|
>(&payload)
|
|
{
|
|
let mut state = redis_progress_state.lock().unwrap();
|
|
state.update_from_redis(
|
|
&progress_msg.msg_type,
|
|
&progress_msg.processor,
|
|
progress_msg.data.current,
|
|
progress_msg.data.total,
|
|
progress_msg.data.message.as_deref(),
|
|
);
|
|
|
|
// Store progress in Redis Hash for HTTP API
|
|
let uuid = progress_msg.uuid.clone();
|
|
let processor = progress_msg.processor.clone();
|
|
let msg_type = progress_msg.msg_type.clone();
|
|
let current = progress_msg.data.current;
|
|
let total = progress_msg.data.total;
|
|
let message = progress_msg.data.message.clone();
|
|
|
|
tokio::spawn(async move {
|
|
if let Ok(redis_client) =
|
|
momentry_core::core::db::RedisClient::new()
|
|
{
|
|
if let Ok(mut conn) = redis_client.get_conn().await
|
|
{
|
|
let prefix = momentry_core::core::config::REDIS_KEY_PREFIX.as_str();
|
|
let key = format!(
|
|
"{}job:{}:processor:{}",
|
|
prefix, uuid, processor
|
|
);
|
|
let _: () = redis::cmd("HSET")
|
|
.arg(&key)
|
|
.arg("status")
|
|
.arg(&msg_type)
|
|
.query_async(&mut conn)
|
|
.await
|
|
.unwrap_or(());
|
|
if let Some(c) = current {
|
|
let _: () = redis::cmd("HSET")
|
|
.arg(&key)
|
|
.arg("current")
|
|
.arg(c)
|
|
.query_async(&mut conn)
|
|
.await
|
|
.unwrap_or(());
|
|
}
|
|
if let Some(t) = total {
|
|
let _: () = redis::cmd("HSET")
|
|
.arg(&key)
|
|
.arg("total")
|
|
.arg(t)
|
|
.query_async(&mut conn)
|
|
.await
|
|
.unwrap_or(());
|
|
}
|
|
if let Some(ref m) = message {
|
|
let _: () = redis::cmd("HSET")
|
|
.arg(&key)
|
|
.arg("message")
|
|
.arg(m)
|
|
.query_async(&mut conn)
|
|
.await
|
|
.unwrap_or(());
|
|
}
|
|
let _: () = redis::cmd("EXPIRE")
|
|
.arg(&key)
|
|
.arg(86400i64)
|
|
.query_async(&mut conn)
|
|
.await
|
|
.unwrap_or(());
|
|
}
|
|
}
|
|
});
|
|
|
|
// Trigger UI render on progress update
|
|
if let Some(ref mut ui) = *redis_ui.lock().unwrap() {
|
|
let _ = ui.render();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
|
}
|
|
}
|
|
});
|
|
|
|
// Process ASR (Automatic Speech Recognition)
|
|
if should_process(ProcessorType::Asr) {
|
|
let asr_path = output_dir.get_output_path(&uuid, "asr.json");
|
|
let decision = decide_processing(&asr_path, force, resume);
|
|
|
|
match decision {
|
|
ProcessingDecision::SkipComplete => {
|
|
println!("\nASR: ✓ Already complete, skipping");
|
|
}
|
|
ProcessingDecision::ForceReprocess => {
|
|
println!("\nASR: ⟳ Force reprocessing from scratch...");
|
|
std::fs::remove_file(&asr_path).ok();
|
|
if is_cloud(ProcessorType::Asr) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_asr_module(&asr_path, video_path, &uuid, &progress_state, &ui)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::ResumePartial => {
|
|
println!("\nASR: ↻ Resuming from checkpoint...");
|
|
if is_cloud(ProcessorType::Asr) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_asr_module(&asr_path, video_path, &uuid, &progress_state, &ui)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::Process => {
|
|
if is_cloud(ProcessorType::Asr) {
|
|
println!("\nASR: ☁️ Running via cloud...");
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
println!("\nASR: ⚙️ Processing...");
|
|
process_asr_module(&asr_path, video_path, &uuid, &progress_state, &ui)
|
|
.await?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Update storage status
|
|
db.update_storage_status(&uuid, "fs_json", true).await?;
|
|
|
|
// Process CUT (scene detection)
|
|
if should_process(ProcessorType::Cut) {
|
|
let cut_path = output_dir.get_output_path(&uuid, "cut.json");
|
|
let decision = decide_processing(&cut_path, force, resume);
|
|
|
|
match decision {
|
|
ProcessingDecision::SkipComplete => {
|
|
println!("\nCUT: ✓ Already complete, skipping");
|
|
}
|
|
ProcessingDecision::ForceReprocess => {
|
|
println!("\nCUT: ⟳ Force reprocessing from scratch...");
|
|
std::fs::remove_file(&cut_path).ok();
|
|
if is_cloud(ProcessorType::Cut) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_cut_module(&cut_path, video_path, &uuid, &progress_state, &ui)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::ResumePartial => {
|
|
println!("\nCUT: ↻ Resuming from checkpoint...");
|
|
if is_cloud(ProcessorType::Cut) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_cut_module(&cut_path, video_path, &uuid, &progress_state, &ui)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::Process => {
|
|
if is_cloud(ProcessorType::Cut) {
|
|
println!("\nCUT: ☁️ Running via cloud...");
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
println!("\nCUT: ⚙️ Processing...");
|
|
process_cut_module(&cut_path, video_path, &uuid, &progress_state, &ui)
|
|
.await?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Process ASRX (speaker diarization)
|
|
if should_process(ProcessorType::Asrx) {
|
|
let asrx_path = output_dir.get_output_path(&uuid, "asrx.json");
|
|
let decision = decide_processing(&asrx_path, force, resume);
|
|
|
|
match decision {
|
|
ProcessingDecision::SkipComplete => {
|
|
println!("\nASRX: ✓ Already complete, skipping");
|
|
}
|
|
ProcessingDecision::ForceReprocess => {
|
|
println!("\nASRX: ⟳ Force reprocessing from scratch...");
|
|
std::fs::remove_file(&asrx_path).ok();
|
|
if is_cloud(ProcessorType::Asrx) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_asrx_module(
|
|
&asrx_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::ResumePartial => {
|
|
println!("\nASRX: ↻ Resuming from checkpoint...");
|
|
if is_cloud(ProcessorType::Asrx) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_asrx_module(
|
|
&asrx_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::Process => {
|
|
if is_cloud(ProcessorType::Asrx) {
|
|
println!("\nASRX: ☁️ Running via cloud...");
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
println!("\nASRX: ⚙️ Processing...");
|
|
process_asrx_module(
|
|
&asrx_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Process YOLO (object detection)
|
|
if should_process(ProcessorType::Yolo) {
|
|
let yolo_path = output_dir.get_output_path(&uuid, "yolo.json");
|
|
let decision = decide_processing(&yolo_path, force, resume);
|
|
|
|
match decision {
|
|
ProcessingDecision::SkipComplete => {
|
|
println!("\nYOLO: ✓ Already complete, skipping");
|
|
}
|
|
ProcessingDecision::ForceReprocess => {
|
|
println!("\nYOLO: ⟳ Force reprocessing from scratch...");
|
|
std::fs::remove_file(&yolo_path).ok();
|
|
if is_cloud(ProcessorType::Yolo) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_yolo_module(
|
|
&yolo_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::ResumePartial => {
|
|
println!("\nYOLO: ↻ Resuming from checkpoint...");
|
|
if is_cloud(ProcessorType::Yolo) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_yolo_module(
|
|
&yolo_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::Process => {
|
|
if is_cloud(ProcessorType::Yolo) {
|
|
println!("\nYOLO: ☁️ Running via cloud...");
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
println!("\nYOLO: ⚙️ Processing...");
|
|
process_yolo_module(
|
|
&yolo_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Process OCR (text recognition)
|
|
if should_process(ProcessorType::Ocr) {
|
|
let ocr_path = output_dir.get_output_path(&uuid, "ocr.json");
|
|
let decision = decide_processing(&ocr_path, force, resume);
|
|
|
|
match decision {
|
|
ProcessingDecision::SkipComplete => {
|
|
println!("\nOCR: ✓ Already complete, skipping");
|
|
}
|
|
ProcessingDecision::ForceReprocess => {
|
|
println!("\nOCR: ⟳ Force reprocessing from scratch...");
|
|
std::fs::remove_file(&ocr_path).ok();
|
|
if is_cloud(ProcessorType::Ocr) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_ocr_module(&ocr_path, video_path, &uuid, &progress_state, &ui)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::ResumePartial => {
|
|
println!("\nOCR: ↻ Resuming from checkpoint...");
|
|
if is_cloud(ProcessorType::Ocr) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_ocr_module(&ocr_path, video_path, &uuid, &progress_state, &ui)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::Process => {
|
|
if is_cloud(ProcessorType::Ocr) {
|
|
println!("\nOCR: ☁️ Running via cloud...");
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
println!("\nOCR: ⚙️ Processing...");
|
|
process_ocr_module(&ocr_path, video_path, &uuid, &progress_state, &ui)
|
|
.await?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Process Face (face detection)
|
|
if should_process(ProcessorType::Face) {
|
|
let face_path = output_dir.get_output_path(&uuid, "face.json");
|
|
let decision = decide_processing(&face_path, force, resume);
|
|
|
|
match decision {
|
|
ProcessingDecision::SkipComplete => {
|
|
println!("\nFace: ✓ Already complete, skipping");
|
|
}
|
|
ProcessingDecision::ForceReprocess => {
|
|
println!("\nFace: ⟳ Force reprocessing from scratch...");
|
|
std::fs::remove_file(&face_path).ok();
|
|
if is_cloud(ProcessorType::Face) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_face_module(
|
|
&face_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::ResumePartial => {
|
|
println!("\nFace: ↻ Resuming from checkpoint...");
|
|
if is_cloud(ProcessorType::Face) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_face_module(
|
|
&face_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::Process => {
|
|
if is_cloud(ProcessorType::Face) {
|
|
println!("\nFace: ☁️ Running via cloud...");
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
println!("\nFace: ⚙️ Processing...");
|
|
process_face_module(
|
|
&face_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Process Pose (pose estimation)
|
|
if should_process(ProcessorType::Pose) {
|
|
let pose_path = output_dir.get_output_path(&uuid, "pose.json");
|
|
let decision = decide_processing(&pose_path, force, resume);
|
|
|
|
match decision {
|
|
ProcessingDecision::SkipComplete => {
|
|
println!("\nPose: ✓ Already complete, skipping");
|
|
}
|
|
ProcessingDecision::ForceReprocess => {
|
|
println!("\nPose: ⟳ Force reprocessing from scratch...");
|
|
std::fs::remove_file(&pose_path).ok();
|
|
if is_cloud(ProcessorType::Pose) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_pose_module(
|
|
&pose_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::ResumePartial => {
|
|
println!("\nPose: ↻ Resuming from checkpoint...");
|
|
if is_cloud(ProcessorType::Pose) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_pose_module(
|
|
&pose_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::Process => {
|
|
if is_cloud(ProcessorType::Pose) {
|
|
println!("\nPose: ☁️ Running via cloud...");
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
println!("\nPose: ⚙️ Processing...");
|
|
process_pose_module(
|
|
&pose_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Process Story (video narrative)
|
|
if should_process(ProcessorType::Story) {
|
|
let story_path = output_dir.get_output_path(&uuid, "story.json");
|
|
let decision = decide_processing(&story_path, force, resume);
|
|
|
|
match decision {
|
|
ProcessingDecision::SkipComplete => {
|
|
println!("\nStory: ✓ Already complete, skipping");
|
|
}
|
|
ProcessingDecision::ForceReprocess => {
|
|
println!("\nStory: ⟳ Force reprocessing from scratch...");
|
|
std::fs::remove_file(&story_path).ok();
|
|
if is_cloud(ProcessorType::Story) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_story_module(
|
|
&story_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::ResumePartial => {
|
|
println!("\nStory: ↻ Resuming from checkpoint...");
|
|
if is_cloud(ProcessorType::Story) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_story_module(
|
|
&story_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::Process => {
|
|
if is_cloud(ProcessorType::Story) {
|
|
println!("\nStory: ☁️ Running via cloud...");
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
println!("\nStory: ⚙️ Processing...");
|
|
process_story_module(
|
|
&story_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Process Caption (image captions)
|
|
if should_process(ProcessorType::Caption) {
|
|
let caption_path = output_dir.get_output_path(&uuid, "caption.json");
|
|
let decision = decide_processing(&caption_path, force, resume);
|
|
|
|
match decision {
|
|
ProcessingDecision::SkipComplete => {
|
|
println!("\nCaption: ✓ Already complete, skipping");
|
|
}
|
|
ProcessingDecision::ForceReprocess => {
|
|
println!("\nCaption: ⟳ Force reprocessing from scratch...");
|
|
std::fs::remove_file(&caption_path).ok();
|
|
if is_cloud(ProcessorType::Caption) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_caption_module(
|
|
&caption_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::ResumePartial => {
|
|
println!("\nCaption: ↻ Resuming from checkpoint...");
|
|
if is_cloud(ProcessorType::Caption) {
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
process_caption_module(
|
|
&caption_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
ProcessingDecision::Process => {
|
|
if is_cloud(ProcessorType::Caption) {
|
|
println!("\nCaption: ☁️ Running via cloud...");
|
|
println!(" [Cloud processing not implemented yet - run locally]");
|
|
} else {
|
|
println!("\nCaption: ⚙️ Processing...");
|
|
process_caption_module(
|
|
&caption_path,
|
|
video_path,
|
|
&uuid,
|
|
&progress_state,
|
|
&ui,
|
|
)
|
|
.await?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// TODO: Store pre_chunks and frames to database
|
|
|
|
// Stop Redis subscriber
|
|
redis_handle.abort();
|
|
|
|
println!("\n✓ Process stage completed!");
|
|
if should_process(ProcessorType::Asr) {
|
|
let path = output_dir.get_output_path(&uuid, "asr.json");
|
|
println!(" - ASR JSON: {}", path.display());
|
|
}
|
|
if should_process(ProcessorType::Cut) {
|
|
let path = output_dir.get_output_path(&uuid, "cut.json");
|
|
println!(" - CUT JSON: {}", path.display());
|
|
}
|
|
if should_process(ProcessorType::Asrx) {
|
|
let path = output_dir.get_output_path(&uuid, "asrx.json");
|
|
println!(" - ASRX JSON: {}", path.display());
|
|
}
|
|
if should_process(ProcessorType::Yolo) {
|
|
let path = output_dir.get_output_path(&uuid, "yolo.json");
|
|
println!(" - YOLO JSON: {}", path.display());
|
|
}
|
|
if should_process(ProcessorType::Ocr) {
|
|
let path = output_dir.get_output_path(&uuid, "ocr.json");
|
|
println!(" - OCR JSON: {}", path.display());
|
|
}
|
|
if should_process(ProcessorType::Face) {
|
|
let path = output_dir.get_output_path(&uuid, "face.json");
|
|
println!(" - Face JSON: {}", path.display());
|
|
}
|
|
if should_process(ProcessorType::Pose) {
|
|
let path = output_dir.get_output_path(&uuid, "pose.json");
|
|
println!(" - Pose JSON: {}", path.display());
|
|
}
|
|
if should_process(ProcessorType::Story) {
|
|
let path = output_dir.get_output_path(&uuid, "story.json");
|
|
println!(" - Story JSON: {}", path.display());
|
|
}
|
|
if should_process(ProcessorType::Caption) {
|
|
let path = output_dir.get_output_path(&uuid, "caption.json");
|
|
println!(" - Caption JSON: {}", path.display());
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
Commands::Chunk { uuid } => {
|
|
println!("Chunking: {}", uuid);
|
|
|
|
let db = PostgresDb::init().await?;
|
|
let video = db
|
|
.get_video_by_uuid(&uuid)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("Video not found: {}", uuid))?;
|
|
|
|
let file_id = video.id;
|
|
let fps = video.fps;
|
|
|
|
// ========== Read all JSON files ==========
|
|
|
|
// Read ASR JSON
|
|
let asr_path = format!("{}.asr.json", uuid);
|
|
let asr_json = std::fs::read_to_string(&asr_path)
|
|
.context("ASR file not found. Run 'process' first.")?;
|
|
let asr_result: momentry_core::core::processor::asr::AsrResult =
|
|
serde_json::from_str(&asr_json)?;
|
|
println!("Loaded ASR: {} segments", asr_result.segments.len());
|
|
|
|
// Read CUT JSON
|
|
let cut_path = format!("{}.cut.json", uuid);
|
|
let cut_json = std::fs::read_to_string(&cut_path)
|
|
.context("CUT file not found. Run 'process' first.")?;
|
|
let cut_result: momentry_core::core::processor::cut::CutResult =
|
|
serde_json::from_str(&cut_json)?;
|
|
println!("Loaded CUT: {} scenes", cut_result.scenes.len());
|
|
|
|
// Read YOLO JSON (optional)
|
|
let yolo_path = format!("{}.yolo.json", uuid);
|
|
let yolo_result = match std::fs::read_to_string(&yolo_path) {
|
|
Ok(yolo_json) => match serde_json::from_str::<
|
|
momentry_core::core::processor::yolo::YoloResult,
|
|
>(&yolo_json)
|
|
{
|
|
Ok(result) => {
|
|
println!("Loaded YOLO: {} frames", result.frames.len());
|
|
result
|
|
}
|
|
Err(e) => {
|
|
println!("Warning: Failed to parse YOLO JSON: {}. Skipping YOLO.", e);
|
|
momentry_core::core::processor::yolo::YoloResult {
|
|
frame_count: 0,
|
|
fps: 0.0,
|
|
frames: vec![],
|
|
}
|
|
}
|
|
},
|
|
Err(_) => {
|
|
println!("Warning: YOLO file not found. Skipping YOLO.");
|
|
momentry_core::core::processor::yolo::YoloResult {
|
|
frame_count: 0,
|
|
fps: 0.0,
|
|
frames: vec![],
|
|
}
|
|
}
|
|
};
|
|
|
|
// Read OCR JSON (optional)
|
|
let ocr_path = format!("{}.ocr.json", uuid);
|
|
let ocr_result = match std::fs::read_to_string(&ocr_path) {
|
|
Ok(ocr_json) => match serde_json::from_str::<
|
|
momentry_core::core::processor::ocr::OcrResult,
|
|
>(&ocr_json)
|
|
{
|
|
Ok(result) => {
|
|
println!("Loaded OCR: {} frames", result.frames.len());
|
|
result
|
|
}
|
|
Err(e) => {
|
|
println!("Warning: Failed to parse OCR JSON: {}. Skipping OCR.", e);
|
|
momentry_core::core::processor::ocr::OcrResult {
|
|
frame_count: 0,
|
|
fps: 0.0,
|
|
frames: vec![],
|
|
}
|
|
}
|
|
},
|
|
Err(_) => {
|
|
println!("Warning: OCR file not found. Skipping OCR.");
|
|
momentry_core::core::processor::ocr::OcrResult {
|
|
frame_count: 0,
|
|
fps: 0.0,
|
|
frames: vec![],
|
|
}
|
|
}
|
|
};
|
|
|
|
// Read Face JSON (optional)
|
|
let face_path = format!("{}.face.json", uuid);
|
|
let face_result = match std::fs::read_to_string(&face_path) {
|
|
Ok(face_json) => match serde_json::from_str::<
|
|
momentry_core::core::processor::face::FaceResult,
|
|
>(&face_json)
|
|
{
|
|
Ok(result) => {
|
|
println!("Loaded Face: {} frames", result.frames.len());
|
|
result
|
|
}
|
|
Err(e) => {
|
|
println!("Warning: Failed to parse Face JSON: {}. Skipping Face.", e);
|
|
momentry_core::core::processor::face::FaceResult {
|
|
frame_count: 0,
|
|
fps: 0.0,
|
|
frames: vec![],
|
|
}
|
|
}
|
|
},
|
|
Err(_) => {
|
|
println!("Warning: Face file not found. Skipping Face.");
|
|
momentry_core::core::processor::face::FaceResult {
|
|
frame_count: 0,
|
|
fps: 0.0,
|
|
frames: vec![],
|
|
}
|
|
}
|
|
};
|
|
|
|
// ========== Store pre_chunks (from ASR, CUT) ==========
|
|
|
|
println!("\nStoring pre_chunks...");
|
|
|
|
// Store ASR sentence pre_chunks
|
|
let mut asr_pre_chunk_ids = Vec::new();
|
|
for seg in asr_result.segments.iter() {
|
|
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 {
|
|
id: 0,
|
|
file_id,
|
|
source_type: "asr".to_string(),
|
|
source_file: Some(asr_path.clone()),
|
|
chunk_type: "sentence".to_string(),
|
|
start_frame,
|
|
end_frame,
|
|
fps,
|
|
raw_json: serde_json::json!({"text": seg.text}),
|
|
text_content: Some(seg.text.clone()),
|
|
processed: false,
|
|
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);
|
|
}
|
|
|
|
// Store CUT scene pre_chunks
|
|
let mut cut_pre_chunk_ids = Vec::new();
|
|
for scene in &cut_result.scenes {
|
|
let pre_chunk = momentry_core::core::db::postgres_db::PreChunk {
|
|
id: 0,
|
|
file_id,
|
|
source_type: "cut".to_string(),
|
|
source_file: Some(cut_path.clone()),
|
|
chunk_type: "cut".to_string(),
|
|
start_frame: scene.start_frame as i64,
|
|
end_frame: scene.end_frame as i64,
|
|
fps,
|
|
raw_json: serde_json::json!({
|
|
"scene_number": scene.scene_number,
|
|
}),
|
|
text_content: None,
|
|
processed: false,
|
|
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);
|
|
}
|
|
|
|
// Store time-based pre_chunks (every 10 seconds)
|
|
let duration = video.duration;
|
|
let mut time_pre_chunk_ids = Vec::new();
|
|
let mut time_start = 0.0;
|
|
while time_start < duration {
|
|
let time_end = (time_start + 10.0).min(duration);
|
|
let start_frame = FrameTime::from_seconds(time_start, fps).frames();
|
|
let end_frame = FrameTime::from_seconds(time_end, fps).frames();
|
|
|
|
let pre_chunk = momentry_core::core::db::postgres_db::PreChunk {
|
|
id: 0,
|
|
file_id,
|
|
source_type: "time".to_string(),
|
|
source_file: None,
|
|
chunk_type: "time".to_string(),
|
|
start_frame,
|
|
end_frame,
|
|
fps,
|
|
raw_json: serde_json::json!({"interval": 10.0}),
|
|
text_content: None,
|
|
processed: false,
|
|
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);
|
|
time_start = time_end;
|
|
}
|
|
|
|
println!(
|
|
"Stored pre_chunks: {} asr + {} cut + {} time",
|
|
asr_result.segments.len(),
|
|
cut_result.scenes.len(),
|
|
time_pre_chunk_ids.len()
|
|
);
|
|
|
|
// ========== Store frames (from YOLO, OCR, Face) ==========
|
|
|
|
println!("\nStoring frames...");
|
|
|
|
// Group YOLO, OCR, Face results by frame_number
|
|
let mut frame_data: std::collections::HashMap<
|
|
u64,
|
|
momentry_core::core::processor::yolo::YoloFrame,
|
|
> = std::collections::HashMap::new();
|
|
for frame in &yolo_result.frames {
|
|
frame_data.insert(frame.frame, frame.clone());
|
|
}
|
|
|
|
let mut ocr_by_frame: std::collections::HashMap<
|
|
u64,
|
|
momentry_core::core::processor::ocr::OcrFrame,
|
|
> = std::collections::HashMap::new();
|
|
for frame in &ocr_result.frames {
|
|
ocr_by_frame.insert(frame.frame, frame.clone());
|
|
}
|
|
|
|
let mut face_by_frame: std::collections::HashMap<
|
|
u64,
|
|
momentry_core::core::processor::face::FaceFrame,
|
|
> = std::collections::HashMap::new();
|
|
for frame in &face_result.frames {
|
|
face_by_frame.insert(frame.frame, frame.clone());
|
|
}
|
|
|
|
// Store frames (merge data from YOLO, OCR, Face)
|
|
let mut all_frames: Vec<u64> = frame_data
|
|
.keys()
|
|
.cloned()
|
|
.chain(ocr_by_frame.keys().cloned())
|
|
.chain(face_by_frame.keys().cloned())
|
|
.collect();
|
|
all_frames.sort();
|
|
all_frames.dedup();
|
|
|
|
for frame_num in &all_frames {
|
|
let timestamp = (*frame_num as f64) / fps;
|
|
let yolo_frame = frame_data.get(frame_num);
|
|
let ocr_frame = ocr_by_frame.get(frame_num);
|
|
let face_frame = face_by_frame.get(frame_num);
|
|
|
|
let frame = momentry_core::core::db::postgres_db::Frame {
|
|
id: 0,
|
|
file_id,
|
|
frame_number: *frame_num as i64,
|
|
timestamp,
|
|
fps,
|
|
yolo_objects: yolo_frame.map(|f| serde_json::json!(&f.objects)),
|
|
ocr_results: ocr_frame.map(|f| serde_json::json!(&f.texts)),
|
|
face_results: face_frame.map(|f| serde_json::json!(&f.faces)),
|
|
frame_path: None,
|
|
created_at: String::new(),
|
|
};
|
|
db.store_frame(&frame).await?;
|
|
}
|
|
|
|
println!("Stored {} frames", all_frames.len());
|
|
|
|
// ========== Create chunks ==========
|
|
|
|
println!("\nCreating chunks...");
|
|
|
|
// Rule 1: Direct conversion (sentence pre_chunk -> sentence chunk)
|
|
let mut sentence_chunks = Vec::new();
|
|
for (i, seg) in asr_result.segments.iter().enumerate() {
|
|
let pre_chunk_id = asr_pre_chunk_ids.get(i).copied().unwrap_or(0);
|
|
let chunk = Chunk::from_seconds(
|
|
file_id as i32,
|
|
uuid.clone(),
|
|
i as u32,
|
|
ChunkType::Sentence,
|
|
ChunkRule::Rule1,
|
|
seg.start,
|
|
seg.end,
|
|
fps,
|
|
serde_json::json!({
|
|
"text": seg.text,
|
|
}),
|
|
)
|
|
.with_text_content(seg.text.clone())
|
|
.with_pre_chunk_ids(vec![pre_chunk_id as i32]);
|
|
sentence_chunks.push(chunk);
|
|
}
|
|
|
|
// Rule 1: CUT chunks
|
|
let mut cut_chunks = Vec::new();
|
|
for (i, scene) in cut_result.scenes.iter().enumerate() {
|
|
let pre_chunk_id = cut_pre_chunk_ids.get(i).copied().unwrap_or(0);
|
|
let chunk = Chunk::from_seconds(
|
|
file_id as i32,
|
|
uuid.clone(),
|
|
i as u32,
|
|
ChunkType::Cut,
|
|
ChunkRule::Rule1,
|
|
scene.start_time,
|
|
scene.end_time,
|
|
fps,
|
|
serde_json::json!({
|
|
"scene_number": scene.scene_number,
|
|
}),
|
|
)
|
|
.with_pre_chunk_ids(vec![pre_chunk_id as i32]);
|
|
cut_chunks.push(chunk);
|
|
}
|
|
|
|
// Rule 1: Time-based chunks
|
|
let splitter = momentry_core::core::chunk::ChunkSplitter::new(10.0);
|
|
let mut time_chunks = Vec::new();
|
|
let time_chunk_list = splitter.split_time_based(&uuid, video.duration);
|
|
for (i, tc) in time_chunk_list.iter().enumerate() {
|
|
let pre_chunk_id = time_pre_chunk_ids.get(i).copied().unwrap_or(0);
|
|
let chunk = Chunk::new(
|
|
file_id as i32,
|
|
uuid.clone(),
|
|
i as u32,
|
|
ChunkType::TimeBased,
|
|
ChunkRule::Rule1,
|
|
tc.start_frame,
|
|
tc.end_frame,
|
|
fps,
|
|
serde_json::json!({"interval": 10.0}),
|
|
)
|
|
.with_pre_chunk_ids(vec![pre_chunk_id as i32]);
|
|
time_chunks.push(chunk);
|
|
}
|
|
|
|
// Store chunks
|
|
println!(
|
|
"Storing {} sentence chunks (rule_1)...",
|
|
sentence_chunks.len()
|
|
);
|
|
for chunk in &sentence_chunks {
|
|
db.store_chunk(chunk).await?;
|
|
}
|
|
|
|
println!("Storing {} cut chunks (rule_1)...", cut_chunks.len());
|
|
for chunk in &cut_chunks {
|
|
db.store_chunk(chunk).await?;
|
|
}
|
|
|
|
println!(
|
|
"Storing {} time-based chunks (rule_1)...",
|
|
time_chunks.len()
|
|
);
|
|
for chunk in &time_chunks {
|
|
db.store_chunk(chunk).await?;
|
|
}
|
|
|
|
let total_chunks = sentence_chunks.len() + cut_chunks.len() + time_chunks.len();
|
|
|
|
// Update storage status
|
|
db.update_storage_status(&uuid, "psql_chunk", true).await?;
|
|
|
|
println!("\n✓ Chunk stage completed!");
|
|
println!(
|
|
" - pre_chunks: {} (asr + cut + time)",
|
|
asr_result.segments.len() + cut_result.scenes.len() + time_pre_chunk_ids.len()
|
|
);
|
|
println!(" - frames: {}", all_frames.len());
|
|
println!(" - chunks: {} (sentence + cut + time_based)", total_chunks);
|
|
|
|
Ok(())
|
|
}
|
|
Commands::Story { uuid } => {
|
|
println!("Generating story for: {}", uuid);
|
|
|
|
let db = PostgresDb::init().await?;
|
|
let video = db
|
|
.get_video_by_uuid(&uuid)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("Video not found: {}", uuid))?;
|
|
|
|
let file_id = video.id;
|
|
let _fps = video.fps;
|
|
let duration = video.duration;
|
|
|
|
// Get all chunks
|
|
let all_chunks = db.get_chunks_by_uuid(&uuid).await?;
|
|
|
|
// Try cut chunks first, fall back to sentence chunks
|
|
let mut story_chunks: Vec<&Chunk> = all_chunks
|
|
.iter()
|
|
.filter(|c| c.chunk_type == ChunkType::Cut)
|
|
.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())
|
|
.collect();
|
|
"sentence"
|
|
} else {
|
|
"cut"
|
|
};
|
|
|
|
if story_chunks.is_empty() {
|
|
println!("No story chunks found. Run 'chunk' command first.");
|
|
return Ok(());
|
|
}
|
|
|
|
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!(
|
|
"Time: {:.2}s - {:.2}s",
|
|
story_chunk.start_time().seconds(),
|
|
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)
|
|
.await?;
|
|
|
|
// Get frames in context range
|
|
let context_frames = db
|
|
.get_frames_by_time_range(file_id, context_start, context_end)
|
|
.await?;
|
|
|
|
// Build story
|
|
let mut story = String::new();
|
|
story.push_str(&format!(
|
|
"Scene {} ({:.1}s - {:.1}s)\n\n",
|
|
i + 1,
|
|
story_chunk.start_time().seconds(),
|
|
story_chunk.end_time().seconds()
|
|
));
|
|
|
|
// Add audio/text content
|
|
let sentence_chunks: Vec<&Chunk> = context_chunks
|
|
.iter()
|
|
.filter(|c| c.chunk_type == ChunkType::Sentence)
|
|
.collect();
|
|
|
|
if !sentence_chunks.is_empty() {
|
|
story.push_str("【Speech】\n");
|
|
for sc in &sentence_chunks {
|
|
if let Some(text) = &sc.text_content {
|
|
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 !all_objects.is_empty() {
|
|
story.push_str("【Objects】\n");
|
|
let mut sorted_objects: Vec<_> = all_objects.iter().collect();
|
|
sorted_objects.sort_by(|a, b| b.1.cmp(a.1));
|
|
for (obj, count) in sorted_objects.iter().take(10) {
|
|
story.push_str(&format!(" - {} ({} frames)\n", obj, count));
|
|
}
|
|
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 !all_texts.is_empty() {
|
|
story.push_str("【Text in video】\n");
|
|
for txt in all_texts.iter().take(10) {
|
|
story.push_str(&format!(" - {}\n", txt));
|
|
}
|
|
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 face_count > 0 {
|
|
story.push_str(&format!(
|
|
"【Faces】\n - {} face(s) detected\n\n",
|
|
face_count
|
|
));
|
|
}
|
|
|
|
println!("{}", story);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
Commands::Vectorize { uuid } => {
|
|
println!("Vectorizing: {}", uuid);
|
|
|
|
let pg = PostgresDb::init()
|
|
.await
|
|
.context("Failed to init PostgreSQL")?;
|
|
let qdrant = QdrantDb::init().await.context("Failed to init Qdrant")?;
|
|
let embedder = Embedder::new("nomic-embed-text:v1.5".to_string());
|
|
|
|
let target_uuid = if uuid == "all" {
|
|
None
|
|
} else {
|
|
Some(uuid.as_str())
|
|
};
|
|
|
|
let mut stored_count = 0usize;
|
|
|
|
if let Some(target) = target_uuid {
|
|
let chunks = pg.get_chunks_by_uuid(target).await?;
|
|
let sentence_chunks: Vec<_> = chunks
|
|
.into_iter()
|
|
.filter(|c| c.chunk_type == ChunkType::Sentence)
|
|
.collect();
|
|
|
|
println!(
|
|
"Found {} sentence chunks for {}",
|
|
sentence_chunks.len(),
|
|
target
|
|
);
|
|
|
|
for chunk in sentence_chunks {
|
|
let text = chunk
|
|
.content
|
|
.get("text")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("");
|
|
|
|
if text.is_empty() {
|
|
continue;
|
|
}
|
|
|
|
print!("Embedding chunk {}... ", chunk.chunk_id);
|
|
|
|
match embedder.embed_document(text).await {
|
|
Ok(vector) => {
|
|
let vector_id = format!("{}_{}", chunk.uuid, chunk.chunk_id);
|
|
|
|
if let Err(e) =
|
|
pg.store_vector(&chunk.chunk_id, &vector, &chunk.uuid).await
|
|
{
|
|
eprintln!("store_vector error for {}: {}", chunk.chunk_id, e);
|
|
continue;
|
|
}
|
|
|
|
let qdrant_payload = VectorPayload {
|
|
uuid: chunk.uuid.clone(),
|
|
chunk_id: chunk.chunk_id.clone(),
|
|
chunk_type: "sentence".to_string(),
|
|
start_time: chunk.start_time().seconds(),
|
|
end_time: chunk.end_time().seconds(),
|
|
text: Some(text.to_string()),
|
|
};
|
|
if let Err(e) = qdrant
|
|
.upsert_vector(&chunk.chunk_id, &vector, qdrant_payload)
|
|
.await
|
|
{
|
|
eprintln!("upsert_vector error for {}: {}", chunk.chunk_id, e);
|
|
continue;
|
|
}
|
|
|
|
if let Err(e) = pg.update_vector_id(&chunk.chunk_id, &vector_id).await {
|
|
eprintln!("update_vector_id error for {}: {}", chunk.chunk_id, e);
|
|
continue;
|
|
}
|
|
|
|
stored_count += 1;
|
|
println!("done ({} dims)", vector.len());
|
|
}
|
|
Err(e) => {
|
|
println!("failed: {}", e);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Only update storage status if vectors were actually stored
|
|
if stored_count > 0 {
|
|
pg.update_storage_status(target, "pvector_chunk", true)
|
|
.await?;
|
|
pg.update_storage_status(target, "qvector_chunk", true)
|
|
.await?;
|
|
println!(
|
|
"\n✓ Vectorize stage completed for {}! ({} vectors stored)",
|
|
target, stored_count
|
|
);
|
|
} else {
|
|
println!(
|
|
"\n✗ Vectorize stage failed for {}! (0 vectors stored)",
|
|
target
|
|
);
|
|
}
|
|
} else {
|
|
println!("\n✓ Vectorize stage completed for all videos!");
|
|
}
|
|
Ok(())
|
|
}
|
|
Commands::Play { target } => {
|
|
println!("Playing: {}", target);
|
|
// TODO: Implement play
|
|
Ok(())
|
|
}
|
|
Commands::Watch { directories } => {
|
|
println!("Starting watcher: {:?}", directories);
|
|
// TODO: Implement watch
|
|
Ok(())
|
|
}
|
|
Commands::System { gpu } => {
|
|
let resources = SystemResources::check();
|
|
println!("╔══════════════════════════════════════════════════════════════╗");
|
|
println!("║ System Resources Report ║");
|
|
println!("╠══════════════════════════════════════════════════════════════╣");
|
|
println!(
|
|
"║ CPU: {:.1}% idle ║",
|
|
resources.cpu_idle_percent
|
|
);
|
|
println!(
|
|
"║ Memory: {:.1}GB / {:.1}GB available ({:.0}% used) ║",
|
|
resources.memory_available_mb as f64 / 1024.0,
|
|
resources.memory_total_mb as f64 / 1024.0,
|
|
resources.memory_used_percent
|
|
);
|
|
|
|
if resources.gpu_available {
|
|
match resources.gpu_type {
|
|
GpuType::Nvidia => {
|
|
let util = resources.gpu_utilization.unwrap_or(0.0);
|
|
println!(
|
|
"║ GPU: NVIDIA - {:.0}% utilized ║",
|
|
util
|
|
);
|
|
}
|
|
GpuType::AppleMps => {
|
|
println!(
|
|
"║ GPU: Apple MPS (Metal) - available ║"
|
|
);
|
|
}
|
|
}
|
|
} else {
|
|
println!("║ GPU: None detected ║");
|
|
}
|
|
println!("╠══════════════════════════════════════════════════════════════╣");
|
|
|
|
if resources.can_parallel(4096) {
|
|
println!("║ Mode: PARALLEL - Can run multiple modules together ║");
|
|
println!(
|
|
"║ Recommended modules: {} ║",
|
|
resources.recommend_parallel_modules().join(", ")
|
|
);
|
|
} else {
|
|
println!("║ Mode: SEQUENTIAL - Low resources, run one at a time ║");
|
|
}
|
|
println!("╚══════════════════════════════════════════════════════════════╝");
|
|
|
|
if gpu {
|
|
println!("\n=== GPU Details ===");
|
|
let output = std::process::Command::new("system_profiler")
|
|
.args(["SPDisplaysDataType", "-detailLevel", "mini"])
|
|
.output();
|
|
if let Ok(o) = output {
|
|
println!("{}", String::from_utf8_lossy(&o.stdout));
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
Commands::Server { host, port } => {
|
|
let port = port.unwrap_or_else(|| *momentry_core::core::config::SERVER_PORT);
|
|
momentry_core::api::start_server(&host, port).await?;
|
|
Ok(())
|
|
}
|
|
Commands::Worker {
|
|
max_concurrent,
|
|
poll_interval,
|
|
batch_size,
|
|
} => {
|
|
use momentry_core::worker::{JobWorker, WorkerConfig};
|
|
|
|
let mut config = WorkerConfig::default();
|
|
if let Some(max) = max_concurrent {
|
|
config.max_concurrent = max;
|
|
}
|
|
if let Some(interval) = poll_interval {
|
|
config.poll_interval_secs = interval;
|
|
}
|
|
if let Some(batch) = batch_size {
|
|
config.batch_size = batch;
|
|
}
|
|
|
|
let db = PostgresDb::init().await?;
|
|
let redis = RedisClient::new()?;
|
|
|
|
let worker = JobWorker::new(
|
|
std::sync::Arc::new(db),
|
|
std::sync::Arc::new(redis),
|
|
config.clone(),
|
|
);
|
|
|
|
println!(
|
|
"Starting worker with max_concurrent={}, poll_interval={}s",
|
|
config.max_concurrent, config.poll_interval_secs
|
|
);
|
|
|
|
worker.run().await?;
|
|
Ok(())
|
|
}
|
|
Commands::Query { query } => {
|
|
println!("Query: {}", query);
|
|
// TODO: Implement query
|
|
Ok(())
|
|
}
|
|
Commands::Lookup { path } => {
|
|
let uuid = momentry_core::uuid::compute_uuid_from_path(&path);
|
|
println!("Path: {}", path);
|
|
println!("UUID: {}", uuid);
|
|
Ok(())
|
|
}
|
|
Commands::Resolve { uuid } => {
|
|
println!("Resolving UUID: {}", uuid);
|
|
// TODO: Look up path from UUID in database
|
|
println!("(Database lookup not implemented yet)");
|
|
Ok(())
|
|
}
|
|
Commands::Thumbnails { uuid, count } => {
|
|
let db = PostgresDb::init().await?;
|
|
|
|
let videos = if let Some(ref uuid) = uuid {
|
|
vec![db
|
|
.get_video_by_uuid(uuid)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("Video not found: {}", uuid))?]
|
|
} else {
|
|
db.list_videos().await?
|
|
};
|
|
|
|
let output_dir = std::path::PathBuf::from("thumbnails");
|
|
let extractor = momentry_core::ThumbnailExtractor::new(output_dir, count);
|
|
|
|
for video in videos {
|
|
println!(
|
|
"\nGenerating thumbnails for: {} ({})",
|
|
video.file_name, video.uuid
|
|
);
|
|
|
|
match extractor.get_or_create(&video.file_path, &video.uuid) {
|
|
Ok(result) => {
|
|
println!(" Generated {} thumbnails", result.count);
|
|
}
|
|
Err(e) => {
|
|
println!(" Error: {}", e);
|
|
}
|
|
}
|
|
}
|
|
|
|
println!("\nThumbnails generated successfully!");
|
|
Ok(())
|
|
}
|
|
Commands::Status { uuid } => {
|
|
let db = PostgresDb::init().await?;
|
|
|
|
let videos = if let Some(ref u) = uuid {
|
|
vec![db
|
|
.get_video_by_uuid(u)
|
|
.await?
|
|
.ok_or_else(|| anyhow::anyhow!("Video not found: {}", u))?]
|
|
} else {
|
|
db.list_videos().await?
|
|
};
|
|
|
|
println!("\n╔══════════════════════════════════════════════════════════════════════════════════╗");
|
|
println!(
|
|
"║ 📊 Storage Status Report ║"
|
|
);
|
|
println!("╠══════════════════════════════════════════════════════════════════════════════════╣");
|
|
println!(
|
|
"║ {:32} │ {:8} │ {:8} │ {:8} │ {:8} │ {:8} │ {:8} │ {:8} ║",
|
|
"Video", "FS", "FS", "PSQL", "PObj", "MObj", "PVec", "QVec"
|
|
);
|
|
println!(
|
|
"║ {:32} │ {:8} │ {:8} │ {:8} │ {:8} │ {:8} │ {:8} │ {:8} ║",
|
|
"", "Video", "JSON", "Chunk", "Chunk", "Chunk", "Chunk", "Chunk"
|
|
);
|
|
println!(
|
|
"╠{:33}╪{:9}╪{:9}╪{:9}╪{:9}╪{:9}╪{:9}╪{:9}╣",
|
|
str::repeat("─", 32),
|
|
str::repeat("─", 8),
|
|
str::repeat("─", 8),
|
|
str::repeat("─", 8),
|
|
str::repeat("─", 8),
|
|
str::repeat("─", 8),
|
|
str::repeat("─", 8),
|
|
str::repeat("─", 8)
|
|
);
|
|
|
|
for video in videos {
|
|
let (sentence_count, time_count) =
|
|
db.get_chunk_count(&video.uuid).await.unwrap_or((0, 0));
|
|
let vector_count = db.get_vector_count(&video.uuid).await.unwrap_or(0);
|
|
let total_chunks = sentence_count + time_count;
|
|
|
|
let psql_status = if total_chunks > 0 { "✓" } else { "-" };
|
|
let pvec_status = if vector_count > 0 && total_chunks > 0 {
|
|
if vector_count >= total_chunks {
|
|
"✓"
|
|
} else {
|
|
"◐"
|
|
}
|
|
} else {
|
|
"-"
|
|
};
|
|
let qvec_status = if video.storage.qvector_chunk {
|
|
"✓"
|
|
} else {
|
|
"-"
|
|
};
|
|
|
|
let file_name = if video.file_name.len() > 30 {
|
|
format!("...{}", &video.file_name[video.file_name.len() - 27..])
|
|
} else {
|
|
video.file_name
|
|
};
|
|
|
|
println!(
|
|
"║ {:32} │ {} │ {} │ {} │ - │ - │ {} │ {} ║",
|
|
file_name,
|
|
if video.storage.fs_video { "✓" } else { "✗" },
|
|
if video.storage.fs_json { "✓" } else { "-" },
|
|
psql_status,
|
|
pvec_status,
|
|
qvec_status
|
|
);
|
|
}
|
|
|
|
println!("╠══════════════════════════════════════════════════════════════════════════════════╣");
|
|
println!(
|
|
"║ Storage Types: ║"
|
|
);
|
|
println!(
|
|
"║ FS_Video - Video file on filesystem ║"
|
|
);
|
|
println!(
|
|
"║ FS_JSON - JSON files (probe, ASR, YOLO, etc.) ║"
|
|
);
|
|
println!(
|
|
"║ PSQL_Chunk - Chunks stored in PostgreSQL ║"
|
|
);
|
|
println!(
|
|
"║ PObject - Chunks as JSON objects in PostgreSQL (future) ║"
|
|
);
|
|
println!(
|
|
"║ MObject - Chunks as JSON objects in MongoDB (future) ║"
|
|
);
|
|
println!(
|
|
"║ PVector - Vectors in PostgreSQL ║"
|
|
);
|
|
println!(
|
|
"║ QVector - Vectors in Qdrant ║"
|
|
);
|
|
println!("╚══════════════════════════════════════════════════════════════════════════════════╝");
|
|
Ok(())
|
|
}
|
|
Commands::Backup { action, days } => {
|
|
let output_dir = OutputDir::new();
|
|
output_dir.ensure_dir()?;
|
|
|
|
println!("\n📁 Backup directory: {:?}", output_dir.get_backup_dir());
|
|
|
|
match action.as_str() {
|
|
"list" => {
|
|
let backups = output_dir.list_backups()?;
|
|
println!("\n📦 Available backups:");
|
|
if backups.is_empty() {
|
|
println!(" (no backups found)");
|
|
} else {
|
|
for backup in &backups {
|
|
println!(" - {}", backup.filename);
|
|
}
|
|
}
|
|
println!("\nTotal: {} backup(s)", backups.len());
|
|
}
|
|
"cleanup" => {
|
|
let days = days.unwrap_or(30);
|
|
let deleted = output_dir.cleanup_old_backups(days)?;
|
|
println!(
|
|
"\n🗑️ Cleaned up {} old backup(s) (older than {} days)",
|
|
deleted, days
|
|
);
|
|
}
|
|
"verify" => {
|
|
println!("\n🔍 Verifying backups...");
|
|
let backups = output_dir.list_backups()?;
|
|
let mut verified = 0;
|
|
let mut failed = 0;
|
|
for backup in &backups {
|
|
match output_dir.verify_backup(&backup.path) {
|
|
Ok(true) => {
|
|
println!(" ✓ {}", backup.filename);
|
|
verified += 1;
|
|
}
|
|
Ok(false) => {
|
|
println!(" ✗ {} (missing checksum)", backup.filename);
|
|
failed += 1;
|
|
}
|
|
Err(e) => {
|
|
println!(" ✗ {} ({})", backup.filename, e);
|
|
failed += 1;
|
|
}
|
|
}
|
|
}
|
|
println!("\nVerified: {} OK, {} failed", verified, failed);
|
|
}
|
|
_ => {
|
|
println!("\n⚠️ Unknown action: {}", action);
|
|
println!("Available actions: list, cleanup, verify");
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
Commands::ApiKey {
|
|
action,
|
|
name,
|
|
key_type,
|
|
ttl,
|
|
key,
|
|
} => {
|
|
let db = PostgresDb::init().await?;
|
|
let db_url = std::env::var("DATABASE_URL")
|
|
.unwrap_or_else(|_| "postgres://accusys@localhost:5432/momentry".to_string());
|
|
|
|
let service = ApiKeyService::new(db_url);
|
|
|
|
match action {
|
|
ApiKeyAction::Create => {
|
|
let name = name.unwrap_or_else(|| "unnamed-key".to_string());
|
|
let kt = parse_key_type(key_type.as_deref());
|
|
let request = momentry_core::core::api_key::CreateApiKeyRequest {
|
|
name: name.clone(),
|
|
key_type: kt,
|
|
user_id: None,
|
|
service_name: None,
|
|
permissions: vec!["read".to_string(), "write".to_string()],
|
|
ttl_days: ttl,
|
|
};
|
|
|
|
match service.create_key(request) {
|
|
Ok(response) => {
|
|
let key_hash = service.hash_key(&response.key);
|
|
let key_type_str =
|
|
serde_json::to_string(&kt).unwrap_or_else(|_| "user".to_string());
|
|
let permissions = serde_json::json!(["read", "write"]);
|
|
|
|
let config = momentry_core::core::db::CreateApiKeyConfig::new(
|
|
&response.key_id,
|
|
&key_hash,
|
|
kt.prefix(),
|
|
&name,
|
|
&key_type_str,
|
|
)
|
|
.with_permissions(&permissions)
|
|
.with_expires_at(response.expires_at);
|
|
|
|
if let Err(e) = db.create_api_key(config).await {
|
|
eprintln!(
|
|
"\n⚠️ Key generated but failed to store in database: {}",
|
|
e
|
|
);
|
|
}
|
|
|
|
println!("\n✅ API Key created successfully!");
|
|
println!("\n┌─────────────────────────────────────────────────────────────────────────────┐");
|
|
println!("│ ⚠️ IMPORTANT: Save this key now - it will not be shown again! │");
|
|
println!("└─────────────────────────────────────────────────────────────────────────────┘");
|
|
println!("\nKey ID: {}", response.key_id);
|
|
println!("API Key: {}", response.key);
|
|
println!("Expires: {}", response.expires_at);
|
|
if !response.warning.is_empty() {
|
|
println!("\n⚠️ {}", response.warning);
|
|
}
|
|
}
|
|
Err(e) => {
|
|
eprintln!("\n❌ Failed to create API key: {}", e);
|
|
}
|
|
}
|
|
}
|
|
ApiKeyAction::List => match db.list_api_keys().await {
|
|
Ok(keys) => {
|
|
println!("\n📋 API Key List");
|
|
if keys.is_empty() {
|
|
println!(" (no API keys found)");
|
|
} else {
|
|
println!("\n┌────────────────────────────────────────────────────────────────────────────┐");
|
|
println!(
|
|
"│ {:8} │ {:20} │ {:12} │ {:8} │ {:15} │",
|
|
"Status", "Name", "Type", "Usage", "Last Used"
|
|
);
|
|
println!("├────────────────────────────────────────────────────────────────────────────┤");
|
|
for k in &keys {
|
|
let status = if k.status == "active" {
|
|
"✓ active"
|
|
} else {
|
|
&k.status
|
|
};
|
|
let last_used = k
|
|
.last_used_at
|
|
.map(|dt| dt.format("%Y-%m-%d %H:%M").to_string())
|
|
.unwrap_or_else(|| "never".to_string());
|
|
println!(
|
|
"│ {:8} │ {:20} │ {:12} │ {:8} │ {:15} │",
|
|
status,
|
|
if k.name.len() > 20 {
|
|
&k.name[..17]
|
|
} else {
|
|
&k.name
|
|
},
|
|
k.key_type,
|
|
k.usage_count,
|
|
last_used
|
|
);
|
|
}
|
|
println!("└────────────────────────────────────────────────────────────────────────────┘");
|
|
println!("\nTotal: {} key(s)", keys.len());
|
|
}
|
|
}
|
|
Err(e) => {
|
|
eprintln!("\n❌ Failed to list API keys: {}", e);
|
|
}
|
|
},
|
|
ApiKeyAction::Validate => {
|
|
let api_key =
|
|
key.ok_or_else(|| anyhow::anyhow!("--key required for validate"))?;
|
|
let key_hash = service.hash_key(&api_key);
|
|
|
|
match db.get_api_key_by_hash(&key_hash).await {
|
|
Ok(Some(record)) => {
|
|
if record.status == "active" {
|
|
db.update_api_key_usage(&record.key_id, None).await.ok();
|
|
println!("\n✅ API Key is valid");
|
|
println!("Key ID: {}", record.key_id);
|
|
println!("Name: {}", record.name);
|
|
println!("Type: {}", record.key_type);
|
|
println!("Usage: {} times", record.usage_count + 1);
|
|
if record.rotation_required {
|
|
println!(
|
|
"⚠️ Rotation required: {}",
|
|
record.rotation_reason.as_deref().unwrap_or("unknown")
|
|
);
|
|
}
|
|
} else {
|
|
println!("\n❌ API Key is {}", record.status);
|
|
}
|
|
}
|
|
Ok(None) => {
|
|
println!("\n❌ API Key is invalid or not found");
|
|
}
|
|
Err(e) => {
|
|
eprintln!("\n❌ Validation error: {}", e);
|
|
}
|
|
}
|
|
}
|
|
ApiKeyAction::Revoke => {
|
|
let key = key.ok_or_else(|| anyhow::anyhow!("--key required for revoke"))?;
|
|
let key_id = service.extract_key_id(&key);
|
|
match db.revoke_api_key(&key_id).await {
|
|
Ok(_) => {
|
|
println!("\n🔴 API Key {} revoked successfully", key_id);
|
|
}
|
|
Err(e) => {
|
|
eprintln!("\n❌ Failed to revoke API key: {}", e);
|
|
}
|
|
}
|
|
}
|
|
ApiKeyAction::Rotate => {
|
|
let key = key.ok_or_else(|| anyhow::anyhow!("--key required for rotate"))?;
|
|
let key_id = service.extract_key_id(&key);
|
|
let grace_period_end =
|
|
service.calculate_grace_period_end(parse_key_type(key_type.as_deref()));
|
|
match db
|
|
.require_api_key_rotation(
|
|
&key_id,
|
|
"manual rotation requested",
|
|
grace_period_end,
|
|
)
|
|
.await
|
|
{
|
|
Ok(_) => {
|
|
println!("\n🔄 Rotation requested for key: {}", key_id);
|
|
println!("Grace period ends: {}", grace_period_end);
|
|
}
|
|
Err(e) => {
|
|
eprintln!("\n❌ Rotation request failed: {}", e);
|
|
}
|
|
}
|
|
}
|
|
ApiKeyAction::Stats => {
|
|
match db.get_api_key_stats().await {
|
|
Ok(stats) => {
|
|
println!("\n📊 API Key Statistics");
|
|
println!("\n┌─────────────────────────────────────────┐");
|
|
println!("│ Total Keys: {:5} │", stats.total_keys);
|
|
println!(
|
|
"│ Active Keys: {:5} │",
|
|
stats.active_keys
|
|
);
|
|
println!(
|
|
"│ Expired Keys: {:5} │",
|
|
stats.expired_keys
|
|
);
|
|
println!(
|
|
"│ Rotation Required: {:4} │",
|
|
stats.rotation_required
|
|
);
|
|
println!(
|
|
"│ Anomalies (24h): {:5} │",
|
|
stats.anomalies_last_24h
|
|
);
|
|
println!("└─────────────────────────────────────────┘");
|
|
}
|
|
Err(e) => {
|
|
eprintln!("\n⚠️ Failed to get stats: {}", e);
|
|
}
|
|
}
|
|
|
|
let config = service.get_config();
|
|
println!("\n┌─────────────────────────────────────────┐");
|
|
println!("│ Anomaly Detection Thresholds │");
|
|
println!("├─────────────────────────────────────────┤");
|
|
println!(
|
|
"│ Requests/minute: {:5} │",
|
|
config.requests_per_minute_threshold
|
|
);
|
|
println!(
|
|
"│ Requests/hour: {:5} │",
|
|
config.requests_per_hour_threshold
|
|
);
|
|
println!(
|
|
"│ Error rate: {:5.1}% │",
|
|
config.error_rate_threshold * 100.0
|
|
);
|
|
println!(
|
|
"│ Unique IPs/hour: {:5} │",
|
|
config.unique_ips_per_hour_threshold
|
|
);
|
|
println!(
|
|
"│ Lockout threshold: {:5} │",
|
|
config.lockout_threshold
|
|
);
|
|
println!("└─────────────────────────────────────────┘");
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
Commands::Gitea {
|
|
action,
|
|
username,
|
|
password,
|
|
token_name,
|
|
scopes,
|
|
} => {
|
|
use momentry_core::core::api_key::gitea::{
|
|
CreateGiteaTokenRequest, GiteaClient, GiteaScope,
|
|
};
|
|
|
|
let db = PostgresDb::init().await?;
|
|
let gitea = GiteaClient::new()?;
|
|
|
|
match action {
|
|
GiteaAction::Create => {
|
|
let username = username
|
|
.ok_or_else(|| anyhow::anyhow!("--username required for create"))?;
|
|
let password = password
|
|
.ok_or_else(|| anyhow::anyhow!("--password required for create"))?;
|
|
let token_name = token_name
|
|
.ok_or_else(|| anyhow::anyhow!("--token-name required for create"))?;
|
|
|
|
let scopes_vec: Vec<GiteaScope> = scopes
|
|
.map(|s| {
|
|
s.split(',')
|
|
.filter_map(|scope| scope.trim().parse::<GiteaScope>().ok())
|
|
.collect()
|
|
})
|
|
.unwrap_or_else(|| {
|
|
vec![GiteaScope::ReadRepository, GiteaScope::WriteRepository]
|
|
});
|
|
|
|
let request = CreateGiteaTokenRequest {
|
|
username: username.clone(),
|
|
password,
|
|
token_name: token_name.clone(),
|
|
scopes: scopes_vec.clone(),
|
|
};
|
|
|
|
match gitea.create_token(&request).await {
|
|
Ok(response) => {
|
|
if let Err(e) = db
|
|
.create_gitea_token(
|
|
response.id,
|
|
&username,
|
|
&token_name,
|
|
&response.token_last_eight,
|
|
&serde_json::json!(scopes_vec
|
|
.iter()
|
|
.map(|s| s.as_str())
|
|
.collect::<Vec<_>>()),
|
|
None,
|
|
)
|
|
.await
|
|
{
|
|
eprintln!("\n⚠️ Token created but failed to store: {}", e);
|
|
}
|
|
|
|
println!("\n✅ Gitea Token created successfully!");
|
|
println!("\n┌─────────────────────────────────────────────────────────────────────────────┐");
|
|
println!("│ ⚠️ IMPORTANT: Save this token now - it will not be shown again! │");
|
|
println!("└─────────────────────────────────────────────────────────────────────────────┘");
|
|
println!("\nToken ID: {}", response.id);
|
|
println!("Token Name: {}", response.name);
|
|
println!("SHA1: {}", response.sha1);
|
|
println!("Last 8: {}", response.token_last_eight);
|
|
println!("\nAuthorization Header:");
|
|
println!(" Authorization: token {}", response.sha1);
|
|
}
|
|
Err(e) => {
|
|
eprintln!("\n❌ Failed to create Gitea token: {}", e);
|
|
}
|
|
}
|
|
}
|
|
GiteaAction::List => {
|
|
let username =
|
|
username.ok_or_else(|| anyhow::anyhow!("--username required for list"))?;
|
|
let password =
|
|
password.ok_or_else(|| anyhow::anyhow!("--password required for list"))?;
|
|
|
|
match gitea.list_tokens(&username, &password).await {
|
|
Ok(tokens) => {
|
|
println!("\n📋 Gitea Tokens for user: {}", username);
|
|
if tokens.is_empty() {
|
|
println!(" (no tokens found)");
|
|
} else {
|
|
println!("\n┌────────────────────────────────────────────────────────────────────────────┐");
|
|
println!("│ ID │ Name │ Last 8 │ Registered │");
|
|
println!("├────────────────────────────────────────────────────────────────────────────┤");
|
|
for token in &tokens {
|
|
let registered = db
|
|
.get_gitea_token_by_name(&username, &token.name)
|
|
.await
|
|
.ok()
|
|
.flatten()
|
|
.map(|_| "✓")
|
|
.unwrap_or("-");
|
|
println!(
|
|
"│ {:8} │ {:20} │ {:9} │ {:27} │",
|
|
token.id,
|
|
if token.name.len() > 20 {
|
|
&token.name[..17]
|
|
} else {
|
|
&token.name
|
|
},
|
|
token.token_last_eight,
|
|
registered
|
|
);
|
|
}
|
|
println!("└────────────────────────────────────────────────────────────────────────────┘");
|
|
println!("\nTotal: {} token(s)", tokens.len());
|
|
}
|
|
}
|
|
Err(e) => {
|
|
eprintln!("\n❌ Failed to list Gitea tokens: {}", e);
|
|
}
|
|
}
|
|
}
|
|
GiteaAction::Delete => {
|
|
let username = username
|
|
.ok_or_else(|| anyhow::anyhow!("--username required for delete"))?;
|
|
let password = password
|
|
.ok_or_else(|| anyhow::anyhow!("--password required for delete"))?;
|
|
let token_name = token_name
|
|
.ok_or_else(|| anyhow::anyhow!("--token-name required for delete"))?;
|
|
|
|
match gitea.delete_token(&username, &password, &token_name).await {
|
|
Ok(_) => {
|
|
let _ = db.delete_gitea_token(&username, &token_name).await;
|
|
println!("\n🗑️ Token '{}' deleted successfully", token_name);
|
|
}
|
|
Err(e) => {
|
|
eprintln!("\n❌ Failed to delete Gitea token: {}", e);
|
|
}
|
|
}
|
|
}
|
|
GiteaAction::Verify => {
|
|
let token_name = token_name
|
|
.ok_or_else(|| anyhow::anyhow!("--token-name required for verify"))?;
|
|
|
|
let record = db
|
|
.get_gitea_token_by_name(
|
|
&username.unwrap_or_else(|| "unknown".to_string()),
|
|
&token_name,
|
|
)
|
|
.await?;
|
|
|
|
match record {
|
|
Some(r) => {
|
|
println!("\n📋 Gitea Token: {}", r.token_name);
|
|
println!(" User: {}", r.gitea_user);
|
|
println!(" Token ID: {}", r.gitea_token_id);
|
|
println!(" Last 8: {}", r.token_last_eight);
|
|
println!(" Scopes: {}", r.scopes);
|
|
println!(" Created: {}", r.created_at);
|
|
if let Some(verified) = r.last_verified {
|
|
println!(" Last Verified: {}", verified);
|
|
} else {
|
|
println!(" Last Verified: never");
|
|
}
|
|
}
|
|
None => {
|
|
println!("\n❌ Token not found in local database");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
Commands::N8n {
|
|
action,
|
|
api_key,
|
|
label,
|
|
expires_in_days,
|
|
} => {
|
|
use momentry_core::core::api_key::n8n::{
|
|
extract_last_eight, CreateN8nApiKeyRequest, N8nClient,
|
|
};
|
|
|
|
let db = PostgresDb::init().await?;
|
|
|
|
match action {
|
|
N8nAction::Create => {
|
|
let api_key_value = api_key.ok_or_else(|| {
|
|
anyhow::anyhow!("--api-key required for create (existing n8n API key)")
|
|
})?;
|
|
let label =
|
|
label.ok_or_else(|| anyhow::anyhow!("--label required for create"))?;
|
|
|
|
let n8n = N8nClient::new(api_key_value)?;
|
|
|
|
let expires_at = expires_in_days
|
|
.map(|days| chrono::Utc::now() + chrono::Duration::days(days));
|
|
|
|
let request = CreateN8nApiKeyRequest {
|
|
label: label.clone(),
|
|
expires_at,
|
|
};
|
|
|
|
match n8n.create_api_key(&request).await {
|
|
Ok(response) => {
|
|
if let Err(e) = db
|
|
.create_n8n_api_key(
|
|
&response.id,
|
|
&label,
|
|
&extract_last_eight(&response.api_key),
|
|
None,
|
|
response.expires_at,
|
|
)
|
|
.await
|
|
{
|
|
eprintln!("\n⚠️ API key created but failed to store: {}", e);
|
|
}
|
|
|
|
println!("\n✅ n8n API Key created successfully!");
|
|
println!("\n┌─────────────────────────────────────────────────────────────────────────────┐");
|
|
println!("│ ⚠️ IMPORTANT: Save this API key now - it will not be shown again! │");
|
|
println!("└─────────────────────────────────────────────────────────────────────────────┘");
|
|
println!("\nKey ID: {}", response.id);
|
|
println!("Label: {}", response.label);
|
|
println!("API Key: {}", response.api_key);
|
|
println!("\nUsage:");
|
|
println!(" curl -H 'X-N8N-API-KEY: {}' https://n8n.momentry.ddns.net/api/v1/workflows", response.api_key);
|
|
}
|
|
Err(e) => {
|
|
eprintln!("\n❌ Failed to create n8n API key: {}", e);
|
|
}
|
|
}
|
|
}
|
|
N8nAction::List => {
|
|
let api_key_value =
|
|
api_key.ok_or_else(|| anyhow::anyhow!("--api-key required for list"))?;
|
|
|
|
let n8n = N8nClient::new(api_key_value)?;
|
|
|
|
match n8n.list_api_keys().await {
|
|
Ok(keys) => {
|
|
println!("\n📋 n8n API Keys");
|
|
if keys.is_empty() {
|
|
println!(" (no API keys found)");
|
|
} else {
|
|
println!("\n┌────────────────────────────────────────────────────────────────────────────┐");
|
|
println!("│ Label │ ID │");
|
|
println!("├────────────────────────────────────────────────────────────────────────────┤");
|
|
for key in &keys {
|
|
println!(
|
|
"│ {:27} │ {:39} │",
|
|
if key.label.len() > 27 {
|
|
&key.label[..24]
|
|
} else {
|
|
&key.label
|
|
},
|
|
key.id
|
|
);
|
|
}
|
|
println!("└────────────────────────────────────────────────────────────────────────────┘");
|
|
println!("\nTotal: {} key(s)", keys.len());
|
|
}
|
|
}
|
|
Err(e) => {
|
|
eprintln!("\n❌ Failed to list n8n API keys: {}", e);
|
|
}
|
|
}
|
|
}
|
|
N8nAction::Delete => {
|
|
let api_key_value =
|
|
api_key.ok_or_else(|| anyhow::anyhow!("--api-key required for delete"))?;
|
|
let label =
|
|
label.ok_or_else(|| anyhow::anyhow!("--label required for delete"))?;
|
|
|
|
let record = db.get_n8n_api_key_by_label(&label).await?;
|
|
if let Some(r) = record {
|
|
let n8n = N8nClient::new(api_key_value)?;
|
|
match n8n.delete_api_key(&r.n8n_key_id).await {
|
|
Ok(_) => {
|
|
let _ = db.delete_n8n_api_key(&label).await;
|
|
println!("\n🗑️ API key '{}' deleted successfully", label);
|
|
}
|
|
Err(e) => {
|
|
eprintln!("\n❌ Failed to delete n8n API key: {}", e);
|
|
}
|
|
}
|
|
} else {
|
|
println!("\n❌ API key '{}' not found in local database", label);
|
|
}
|
|
}
|
|
N8nAction::Verify => {
|
|
let label =
|
|
label.ok_or_else(|| anyhow::anyhow!("--label required for verify"))?;
|
|
|
|
let record = db.get_n8n_api_key_by_label(&label).await?;
|
|
|
|
match record {
|
|
Some(r) => {
|
|
println!("\n📋 n8n API Key: {}", r.label);
|
|
println!(" Key ID: {}", r.n8n_key_id);
|
|
println!(" Last 8: {}", r.api_key_last_eight);
|
|
println!(" Created: {}", r.created_at);
|
|
if let Some(expires) = r.expires_at {
|
|
println!(" Expires: {}", expires);
|
|
}
|
|
if let Some(verified) = r.last_verified {
|
|
println!(" Last Verified: {}", verified);
|
|
} else {
|
|
println!(" Last Verified: never");
|
|
}
|
|
}
|
|
None => {
|
|
println!("\n❌ API key not found in local database");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|