Updated to use LazyStore for persistent storage of API key and Visual and Subtitle prompts

This commit is contained in:
TheMaddax 2025-01-28 12:26:11 -06:00
parent 436c8132d3
commit 5dedd61778
8 changed files with 225 additions and 107 deletions

View file

@ -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"
},

10
pnpm-lock.yaml generated
View file

@ -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

17
src-tauri/Cargo.lock generated
View file

@ -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"

View file

@ -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"]

View file

@ -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"
]
}
}

View file

@ -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<f32, Error> {
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<f32, Error> {
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<R: Runtime>(
app_handle: &tauri::AppHandle<R>,
api_key: &str,
prompt: &str,
image_path: Option<&str>
image_path: Option<&str>,
) -> Result<String, Error> {
println!("\n=== Gemini API Request ===");
println!("Prompt: {}", prompt);
@ -175,21 +191,25 @@ async fn query_gemini<R: Runtime>(
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::<serde_json::Value>(&response_text) {
Ok(response) => {
@ -199,35 +219,47 @@ async fn query_gemini<R: Runtime>(
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<R: Runtime>(app_handle: &tauri::AppHandle<R>, api_key: &str) -> Result<(), Error> {
pub async fn test_gemini_api_impl<R: Runtime>(
app_handle: &tauri::AppHandle<R>,
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<ThumbnailState>>
thumbnail_state: State<'_, Mutex<ThumbnailState>>,
) -> Result<Vec<Thumbnail>, 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<R: Runtime>(
app_handle: &tauri::AppHandle<R>,
args: TranscriptionArgs
args: TranscriptionArgs,
) -> Result<TranscriptionResult, Error> {
println!("\n=== Starting Transcription Process ===");
println!("Video Path: {}", args.video_path);
@ -270,27 +304,34 @@ pub async fn process_transcription_impl<R: Runtime>(
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<R: Runtime>(
// 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<R: Runtime>(
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<R: Runtime>(
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::<Vec<_>>()
.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<R: Runtime>(
app_handle: tauri::AppHandle<R>,
) -> Result<String, Error> {
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<R: Runtime>(
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<R: Runtime>(
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<R: Runtime>(
if let (Ok(hours), Ok(minutes), Ok(seconds)) = (
parts[0].parse::<f32>(),
parts[1].parse::<f32>(),
parts[2].parse::<f32>()
parts[2].parse::<f32>(),
) {
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<R: Runtime>(
}
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)
}

View file

@ -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<R: Runtime>(
#[tauri::command]
async fn test_gemini_api<R: Runtime>(
app_handle: tauri::AppHandle<R>,
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<R: Runtime>(
app_handle: tauri::AppHandle<R>,
args: TranscriptionArgs
args: TranscriptionArgs,
) -> Result<subtitle_merge_lib::TranscriptionResult, Error> {
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,

View file

@ -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<string | null>(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<string>('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<string>('visualPrompt').then((saved) => {
if (saved) {
setVisualPrompt(saved);
}
});
store.get<string>('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')