From 5dedd61778671e2b9b346849eb803e338dbe8f98 Mon Sep 17 00:00:00 2001 From: TheMaddax Date: Tue, 28 Jan 2025 12:26:11 -0600 Subject: [PATCH] Updated to use LazyStore for persistent storage of API key and Visual and Subtitle prompts --- package.json | 1 + pnpm-lock.yaml | 10 ++ src-tauri/Cargo.lock | 17 +++ src-tauri/Cargo.toml | 1 + src-tauri/capabilities/default.json | 23 +++- src-tauri/src/lib.rs | 201 ++++++++++++++++++---------- src-tauri/src/main.rs | 27 ++-- src/App.tsx | 52 +++++-- 8 files changed, 225 insertions(+), 107 deletions(-) diff --git a/package.json b/package.json index 0184288..5e5e901 100644 --- a/package.json +++ b/package.json @@ -14,6 +14,7 @@ "@tauri-apps/plugin-dialog": "2", "@tauri-apps/plugin-fs": "2", "@tauri-apps/plugin-opener": "^2", + "@tauri-apps/plugin-store": "~2", "react": "^18.3.1", "react-dom": "^18.3.1" }, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 6a4eb4e..f1222f2 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -20,6 +20,9 @@ importers: '@tauri-apps/plugin-opener': specifier: ^2 version: 2.2.2 + '@tauri-apps/plugin-store': + specifier: ~2 + version: 2.2.0 react: specifier: ^18.3.1 version: 18.3.1 @@ -498,6 +501,9 @@ packages: '@tauri-apps/plugin-opener@2.2.2': resolution: {integrity: sha512-E/XIHKqGV+FT8PDdkfMETmgPUxcR79Rk8USuzbadD/ZdvsKCfQR5q+6rpZC9zEnG2wzi9lVQM4D3xwrtGGIB8A==} + '@tauri-apps/plugin-store@2.2.0': + resolution: {integrity: sha512-hJTRtuJis4w5fW1dkcgftsYxKXK0+DbAqurZ3CURHG5WkAyyZgbxpeYctw12bbzF9ZbZREXZklPq8mocCC3Sgg==} + '@types/babel__core@7.20.5': resolution: {integrity: sha512-qoQprZvz5wQFJwMDqeseRXWv3rqMvhgpbXFfVyWhbx9X47POIA6i/+dXefEmZKoAgOaTdaIgNSMqMIU61yRyzA==} @@ -1413,6 +1419,10 @@ snapshots: dependencies: '@tauri-apps/api': 2.1.1 + '@tauri-apps/plugin-store@2.2.0': + dependencies: + '@tauri-apps/api': 2.1.1 + '@types/babel__core@7.20.5': dependencies: '@babel/parser': 7.26.3 diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 7cfc69d..2460659 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -3898,6 +3898,7 @@ dependencies = [ "tauri-plugin-dialog", "tauri-plugin-fs", "tauri-plugin-shell", + "tauri-plugin-store", "tempfile", "thiserror 1.0.69", "tokio", @@ -4241,6 +4242,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "tauri-plugin-store" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c0c08fae6995909f5e9a0da6038273b750221319f2c0f3b526d6de1cde21505" +dependencies = [ + "dunce", + "serde", + "serde_json", + "tauri", + "tauri-plugin", + "thiserror 2.0.9", + "tokio", + "tracing", +] + [[package]] name = "tauri-runtime" version = "2.2.0" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 705f8c5..b68aa74 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -26,6 +26,7 @@ reqwest = { version = "0.11", features = ["json", "multipart"] } base64 = "0.21" image = "0.24" tempfile = "3.9" +tauri-plugin-store = "2" [features] custom-protocol = ["tauri/custom-protocol"] diff --git a/src-tauri/capabilities/default.json b/src-tauri/capabilities/default.json index 8d2c9ac..adddd13 100644 --- a/src-tauri/capabilities/default.json +++ b/src-tauri/capabilities/default.json @@ -1,22 +1,32 @@ { "identifier": "default", "description": "Default capability for the app", - "windows": ["main"], + "windows": [ + "main" + ], "permissions": [ { "identifier": "fs:allow-read", "description": "Allows reading video and subtitle files", - "paths": ["**"] + "paths": [ + "**" + ] }, { "identifier": "fs:allow-write", "description": "Allows writing output video files", - "paths": ["**"] + "paths": [ + "**" + ] }, { "identifier": "fs:scope", "description": "Allow access to app directory", - "paths": ["$APP", "$APPDATA", "$APPLOCALDATA"] + "paths": [ + "$APP", + "$APPDATA", + "$APPLOCALDATA" + ] }, { "identifier": "dialog:allow-open", @@ -39,6 +49,7 @@ { "identifier": "core:window:allow-set-shadow", "description": "Allows setting window shadow" - } + }, + "store:default" ] -} +} \ No newline at end of file diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index a361b48..0cd7f0f 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -1,12 +1,12 @@ +use base64::Engine; use serde::{Deserialize, Serialize}; +use std::fs; +use std::io::BufRead; use std::process::{Command, Stdio}; use std::sync::Mutex; -use std::io::BufRead; -use tauri::{Runtime, State, Emitter}; +use tauri::{Emitter, Runtime, State}; use tauri_plugin_shell::ShellExt; use tempfile::TempDir; -use std::fs; -use base64::Engine; #[derive(Debug, thiserror::Error)] pub enum Error { @@ -77,9 +77,7 @@ pub struct ThumbnailState { impl Default for ThumbnailState { fn default() -> Self { - Self { - temp_dir: None, - } + Self { temp_dir: None } } } @@ -105,7 +103,10 @@ pub async fn check_ffmpeg() -> Result<(), Error> { .output() .map_err(|e| { println!("FFmpeg check failed: {}", e); - Error::FFmpegError(format!("FFmpeg is not installed or not found in PATH: {}", e)) + Error::FFmpegError(format!( + "FFmpeg is not installed or not found in PATH: {}", + e + )) })?; if !output.status.success() { @@ -118,16 +119,21 @@ pub async fn check_ffmpeg() -> Result<(), Error> { pub async fn get_video_duration(video_path: &str) -> Result { let output = Command::new("/opt/homebrew/bin/ffprobe") .args([ - "-v", "error", - "-show_entries", "format=duration", - "-of", "default=noprint_wrappers=1:nokey=1", - video_path + "-v", + "error", + "-show_entries", + "format=duration", + "-of", + "default=noprint_wrappers=1:nokey=1", + video_path, ]) .output() .map_err(|e| Error::FFmpegError(format!("Failed to run ffprobe: {}", e)))?; if !output.status.success() { - return Err(Error::FFmpegError("Failed to get video duration".to_string())); + return Err(Error::FFmpegError( + "Failed to get video duration".to_string(), + )); } String::from_utf8_lossy(&output.stdout) @@ -139,12 +145,16 @@ pub async fn get_video_duration(video_path: &str) -> Result { async fn extract_frame(video_path: &str, time: u32, output_path: &str) -> Result<(), Error> { let output = Command::new("/opt/homebrew/bin/ffmpeg") .args([ - "-ss", &time.to_string(), - "-i", video_path, - "-vframes", "1", - "-q:v", "2", + "-ss", + &time.to_string(), + "-i", + video_path, + "-vframes", + "1", + "-q:v", + "2", "-y", - output_path + output_path, ]) .stdout(Stdio::piped()) .stderr(Stdio::piped()) @@ -153,11 +163,17 @@ async fn extract_frame(video_path: &str, time: u32, output_path: &str) -> Result if !output.status.success() { let error = String::from_utf8_lossy(&output.stderr); - return Err(Error::FFmpegError(format!("Frame extraction failed: {}", error))); + return Err(Error::FFmpegError(format!( + "Frame extraction failed: {}", + error + ))); } if !std::path::Path::new(output_path).exists() { - return Err(Error::FFmpegError(format!("Output file was not created: {}", output_path))); + return Err(Error::FFmpegError(format!( + "Output file was not created: {}", + output_path + ))); } Ok(()) @@ -167,7 +183,7 @@ async fn query_gemini( app_handle: &tauri::AppHandle, api_key: &str, prompt: &str, - image_path: Option<&str> + image_path: Option<&str>, ) -> Result { println!("\n=== Gemini API Request ==="); println!("Prompt: {}", prompt); @@ -175,21 +191,25 @@ async fn query_gemini( println!("Image Path: {}", path); } - let sidecar = app_handle.shell().sidecar("gemini-query") + let sidecar = app_handle + .shell() + .sidecar("gemini-query") .map_err(|e| Error::GeminiError(format!("Failed to get sidecar: {}", e)))?; - + let mut command = sidecar.arg(api_key).arg(prompt); - + if let Some(path) = image_path { command = command.arg(path); } println!("\nExecuting sidecar command"); - let output = command.output().await + let output = command + .output() + .await .map_err(|e| Error::GeminiError(format!("Failed to execute Gemini query: {}", e)))?; let response_text = String::from_utf8_lossy(&output.stdout); - + // Try to parse the response as JSON match serde_json::from_str::(&response_text) { Ok(response) => { @@ -199,35 +219,47 @@ async fn query_gemini( println!("Response length: {} characters", result.len()); Ok(result) } else { - let error = response["error"].as_str().unwrap_or("Unknown error").to_string(); + let error = response["error"] + .as_str() + .unwrap_or("Unknown error") + .to_string(); println!("\n=== Gemini API Error ==="); println!("Error: {}", error); Err(Error::GeminiError(error)) } - }, + } Err(e) => { println!("\n=== Gemini API Parse Error ==="); println!("Raw output: {}", response_text); println!("Parse error: {}", e); - Err(Error::GeminiError(format!("Failed to parse response: {}", e))) + Err(Error::GeminiError(format!( + "Failed to parse response: {}", + e + ))) } } } -pub async fn test_gemini_api_impl(app_handle: &tauri::AppHandle, api_key: &str) -> Result<(), Error> { +pub async fn test_gemini_api_impl( + app_handle: &tauri::AppHandle, + api_key: &str, +) -> Result<(), Error> { query_gemini(app_handle, api_key, "Test connection", None).await?; Ok(()) } pub async fn extract_thumbnails_impl( - video_path: &str, + video_path: &str, interval: u32, - thumbnail_state: State<'_, Mutex> + thumbnail_state: State<'_, Mutex>, ) -> Result, Error> { // Create a new temporary directory let temp_dir = TempDir::new().map_err(|e| Error::IoError(e))?; - println!("Extracting thumbnails to temporary directory: {:?}", temp_dir.path()); - + println!( + "Extracting thumbnails to temporary directory: {:?}", + temp_dir.path() + ); + let duration = get_video_duration(video_path).await? as u32; let mut thumbnails = Vec::new(); @@ -241,7 +273,7 @@ pub async fn extract_thumbnails_impl( path: output_path.to_str().unwrap().to_string(), time, }); - }, + } Err(e) => { println!("Failed to extract frame at time {}: {}", time, e); continue; @@ -250,7 +282,9 @@ pub async fn extract_thumbnails_impl( } if thumbnails.is_empty() { - return Err(Error::FFmpegError("Failed to extract any thumbnails".to_string())); + return Err(Error::FFmpegError( + "Failed to extract any thumbnails".to_string(), + )); } // Store the TempDir in the state only after we're done using it @@ -260,7 +294,7 @@ pub async fn extract_thumbnails_impl( pub async fn process_transcription_impl( app_handle: &tauri::AppHandle, - args: TranscriptionArgs + args: TranscriptionArgs, ) -> Result { println!("\n=== Starting Transcription Process ==="); println!("Video Path: {}", args.video_path); @@ -270,27 +304,34 @@ pub async fn process_transcription_impl( println!("Number of selected frames: {}", args.selected_frames.len()); let temp_dir = TempDir::new().map_err(|e| Error::IoError(e))?; - + // Process each frame and collect descriptions let mut frame_descriptions = Vec::new(); for (index, frame) in args.selected_frames.iter().enumerate() { println!("\nProcessing frame {} at time {}", index + 1, frame.time); - + // Save base64 image to temp file let frame_path = temp_dir.path().join(format!("frame_{}.jpg", frame.time)); - let image_data = base64::engine::general_purpose::STANDARD.decode(&frame.base64_image) + let image_data = base64::engine::general_purpose::STANDARD + .decode(&frame.base64_image) .map_err(|e| Error::ImageError(format!("Failed to decode base64 image: {}", e)))?; - fs::write(&frame_path, image_data) - .map_err(|e| Error::IoError(e))?; + fs::write(&frame_path, image_data).map_err(|e| Error::IoError(e))?; let frame_prompt = format!("{}\n\nTimestamp: {}", args.visual_prompt, frame.time); - let description = query_gemini(app_handle, &args.api_key, &frame_prompt, Some(frame_path.to_str().unwrap())).await?; + let description = query_gemini( + app_handle, + &args.api_key, + &frame_prompt, + Some(frame_path.to_str().unwrap()), + ) + .await?; frame_descriptions.push((frame.time, description)); } // Format scene descriptions with timestamps in MM:SS format - let scene_descriptions = frame_descriptions.iter() + let scene_descriptions = frame_descriptions + .iter() .map(|(time, desc)| { let minutes = time / 60; let seconds = time % 60; @@ -302,14 +343,16 @@ pub async fn process_transcription_impl( // Process subtitle content if provided let narrative = if let Some(subtitle_path) = &args.subtitle_path { println!("\nReading and parsing subtitle content..."); - let subtitle_content = fs::read_to_string(subtitle_path) - .map_err(|e| Error::IoError(e))?; - println!("Subtitle content length: {} characters", subtitle_content.len()); + let subtitle_content = fs::read_to_string(subtitle_path).map_err(|e| Error::IoError(e))?; + println!( + "Subtitle content length: {} characters", + subtitle_content.len() + ); // Parse subtitle timestamps and content, handling both SRT and ASS formats let mut subtitle_lines = Vec::new(); let lines: Vec<&str> = subtitle_content.lines().collect(); - + // Check if it's an ASS file if subtitle_path.to_lowercase().ends_with(".ass") { let mut in_events = false; @@ -339,15 +382,18 @@ pub async fn process_transcription_impl( if timestamp_parts.len() == 2 { let start_time = timestamp_parts[0].trim(); let mut text = lines[i + 2].trim().to_string(); - + // Collect multi-line subtitle text let mut j = i + 3; - while j < lines.len() && !lines[j].trim().is_empty() && !lines[j].contains("-->") { + while j < lines.len() + && !lines[j].trim().is_empty() + && !lines[j].contains("-->") + { text.push_str(" "); text.push_str(lines[j].trim()); j += 1; } - + subtitle_lines.push((start_time.to_string(), text)); i = j; } @@ -360,17 +406,17 @@ pub async fn process_transcription_impl( subtitle_lines.sort_by(|a, b| a.0.cmp(&b.0)); // Format subtitle content with timestamps - let formatted_subtitles = subtitle_lines.iter() + let formatted_subtitles = subtitle_lines + .iter() .map(|(time, text)| format!("[{}] {}", time, text)) .collect::>() .join("\n"); - let narrative_prompt = format!("{}\n\nScene Descriptions:\n{}\n\nSubtitle Content with Timestamps:\n{}", - args.subtitle_prompt, - scene_descriptions, - formatted_subtitles + let narrative_prompt = format!( + "{}\n\nScene Descriptions:\n{}\n\nSubtitle Content with Timestamps:\n{}", + args.subtitle_prompt, scene_descriptions, formatted_subtitles ); - + // Process the subtitle content with scene descriptions query_gemini(app_handle, &args.api_key, &narrative_prompt, None).await? } else { @@ -405,10 +451,10 @@ pub async fn merge_video_subtitle_impl( app_handle: tauri::AppHandle, ) -> Result { let output_path = format!("{}-subbed.mp4", args.video_path.trim_end_matches(".mp4")); - + let total_duration = get_video_duration(&args.video_path).await?; println!("Video duration: {} seconds", total_duration); - + let mut filters = Vec::new(); if args.add_black_bar { filters.push("pad=iw:ih+170:0:0:black".to_string()); @@ -416,17 +462,24 @@ pub async fn merge_video_subtitle_impl( if args.resize_to_720p { filters.push("scale=-1:720".to_string()); } - - let subtitle_path = args.subtitle_path.replace("'", "'\\''").replace(",", "\\,").replace(" ", "\\ "); + + let subtitle_path = args + .subtitle_path + .replace("'", "'\\''") + .replace(",", "\\,") + .replace(" ", "\\ "); filters.push(format!("subtitles='{}'", subtitle_path)); filters.push("format=nv12|qsv".to_string()); - + let filter_chain = filters.join(","); println!("Filter chain: {}", filter_chain); - + check_ffmpeg().await?; - - println!("Starting FFmpeg process with paths: video={}, subtitle={}", args.video_path, args.subtitle_path); + + println!( + "Starting FFmpeg process with paths: video={}, subtitle={}", + args.video_path, args.subtitle_path + ); let mut command = Command::new("/opt/homebrew/bin/ffmpeg"); command .arg("-y") @@ -462,7 +515,7 @@ pub async fn merge_video_subtitle_impl( if bytes == 0 { break; } - + if line.starts_with("out_time=") { if let Some(time_str) = line.split('=').nth(1) { let parts: Vec<&str> = time_str.trim().split(':').collect(); @@ -470,14 +523,17 @@ pub async fn merge_video_subtitle_impl( if let (Ok(hours), Ok(minutes), Ok(seconds)) = ( parts[0].parse::(), parts[1].parse::(), - parts[2].parse::() + parts[2].parse::(), ) { let current_seconds = hours * 3600.0 + minutes * 60.0 + seconds; - let progress = ((current_seconds / total_duration) * 100.0).min(100.0).max(0.0); - + let progress = ((current_seconds / total_duration) * 100.0) + .min(100.0) + .max(0.0); + if progress - last_progress >= 1.0 || progress == 100.0 { *progress_state.0.lock().unwrap() = progress; - if let Err(e) = app_handle.emit("progress", ProgressUpdate { progress }) { + if let Err(e) = app_handle.emit("progress", ProgressUpdate { progress }) + { println!("Failed to emit progress: {}", e); } println!("Progress: {:.1}%", progress); @@ -498,10 +554,7 @@ pub async fn merge_video_subtitle_impl( } println!("FFmpeg process completed successfully"); - let _ = Command::new("open") - .arg("-R") - .arg(&output_path) - .spawn(); + let _ = Command::new("open").arg("-R").arg(&output_path).spawn(); Ok(output_path) } diff --git a/src-tauri/src/main.rs b/src-tauri/src/main.rs index 82d388f..c2721e5 100644 --- a/src-tauri/src/main.rs +++ b/src-tauri/src/main.rs @@ -1,15 +1,15 @@ // Prevents additional console window on Windows in release #![cfg_attr(not(debug_assertions), windows_subsystem = "windows")] -use subtitle_merge_lib::{ - ProgressState, MergeArgs, Error, merge_video_subtitle_impl, check_ffmpeg, - test_gemini_api_impl, extract_thumbnails_impl, process_transcription_impl, - TranscriptionArgs, get_video_duration, ThumbnailState -}; -use std::fs; use base64::Engine; -use tauri::Runtime; +use std::fs; use std::sync::Mutex; +use subtitle_merge_lib::{ + check_ffmpeg, extract_thumbnails_impl, get_video_duration, merge_video_subtitle_impl, + process_transcription_impl, test_gemini_api_impl, Error, MergeArgs, ProgressState, + ThumbnailState, TranscriptionArgs, +}; +use tauri::Runtime; #[tauri::command] async fn check_dependencies() -> Result<(), Error> { @@ -33,7 +33,7 @@ async fn merge_video_subtitle( #[tauri::command] async fn test_gemini_api( app_handle: tauri::AppHandle, - api_key: String + api_key: String, ) -> Result<(), Error> { test_gemini_api_impl(&app_handle, &api_key).await } @@ -55,7 +55,7 @@ async fn extract_thumbnails( #[tauri::command] async fn process_transcription( app_handle: tauri::AppHandle, - args: TranscriptionArgs + args: TranscriptionArgs, ) -> Result { process_transcription_impl(&app_handle, args).await } @@ -79,7 +79,7 @@ async fn read_thumbnail( if let Some(temp_dir) = state.get_temp_dir() { let temp_path = temp_dir.path(); let requested_path = std::path::Path::new(&path); - + // Check if the requested path is within our temp directory if !requested_path.starts_with(temp_path) { return Err(Error::IoError(std::io::Error::new( @@ -87,7 +87,7 @@ async fn read_thumbnail( "Attempted to access file outside of temporary directory", ))); } - + let data = fs::read(&path).map_err(|e| Error::IoError(e))?; Ok(base64::engine::general_purpose::STANDARD.encode(&data)) } else { @@ -100,14 +100,13 @@ async fn read_thumbnail( fn main() { tauri::Builder::default() + .plugin(tauri_plugin_store::Builder::new().build()) .plugin(tauri_plugin_shell::init()) .plugin(tauri_plugin_dialog::init()) .plugin(tauri_plugin_fs::init()) .manage(ProgressState::default()) .manage(Mutex::new(ThumbnailState::default())) - .setup(|_app| { - Ok(()) - }) + .setup(|_app| Ok(())) .invoke_handler(tauri::generate_handler![ merge_video_subtitle, check_dependencies, diff --git a/src/App.tsx b/src/App.tsx index dce8659..9cae7e1 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -1,5 +1,6 @@ import { useState, useEffect } from 'react'; import { invoke } from '@tauri-apps/api/core'; +import { LazyStore } from '@tauri-apps/plugin-store'; import { Settings, DEFAULT_VISUAL_PROMPT, DEFAULT_SUBTITLE_PROMPT } from './components/Settings'; import { TranscriptionTask } from './components/TranscriptionTask'; import { MergeTask } from './components/MergeTask'; @@ -12,23 +13,48 @@ function App() { const [error, setError] = useState(null); const [isSettingsOpen, setIsSettingsOpen] = useState(false); const [apiKey, setApiKey] = useState(''); - const [visualPrompt, setVisualPrompt] = useState(() => { - const saved = localStorage.getItem('visualPrompt'); - return saved || DEFAULT_VISUAL_PROMPT; - }); - const [subtitlePrompt, setSubtitlePrompt] = useState(() => { - const saved = localStorage.getItem('subtitlePrompt'); - return saved || DEFAULT_SUBTITLE_PROMPT; - }); + const [store] = useState(() => new LazyStore('settings.json')); - // Persist prompts to localStorage when they change + // Load API key from store on mount useEffect(() => { - localStorage.setItem('visualPrompt', visualPrompt); - }, [visualPrompt]); + store.get('apiKey').then((savedKey) => { + if (savedKey) { + setApiKey(savedKey); + } + }); + }, [store]); + + // Save API key to store when it changes + useEffect(() => { + if (apiKey) { + store.set('apiKey', apiKey); + } + }, [apiKey, store]); + const [visualPrompt, setVisualPrompt] = useState(DEFAULT_VISUAL_PROMPT); + const [subtitlePrompt, setSubtitlePrompt] = useState(DEFAULT_SUBTITLE_PROMPT); + + // Load prompts from store on mount + useEffect(() => { + store.get('visualPrompt').then((saved) => { + if (saved) { + setVisualPrompt(saved); + } + }); + store.get('subtitlePrompt').then((saved) => { + if (saved) { + setSubtitlePrompt(saved); + } + }); + }, [store]); + + // Save prompts to store when they change + useEffect(() => { + store.set('visualPrompt', visualPrompt); + }, [visualPrompt, store]); useEffect(() => { - localStorage.setItem('subtitlePrompt', subtitlePrompt); - }, [subtitlePrompt]); + store.set('subtitlePrompt', subtitlePrompt); + }, [subtitlePrompt, store]); useEffect(() => { invoke('check_dependencies')