use tauri::{command, Emitter}; use serde::{Deserialize, Serialize}; use std::collections::HashSet; use std::fs::File; use std::io::{BufRead, BufReader, Write}; use std::path::Path; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DataCleaningProgress { pub current: usize, pub total: usize, pub percentage: f64, pub status: String, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DataCleaningResult { pub success: bool, pub message: String, pub original_count: usize, pub removed_count: usize, pub final_count: usize, pub output_file: String, } #[derive(Debug, Serialize, Deserialize)] pub struct StorageObject { pub uri: String, // 其他字段可以忽略,我们只关心 uri } #[derive(Debug, Serialize, Deserialize)] pub struct ProcessedData { pub json_data: String, } /// AI检索图片/数据清洗 - JSONL格式处理 #[command] pub async fn clean_jsonl_data( all_data_file: String, remove_data_file: String, output_file: String, window: tauri::Window, ) -> Result { // 验证文件存在 if !Path::new(&all_data_file).exists() { return Err(format!("全部数据文件不存在: {}", all_data_file)); } if !Path::new(&remove_data_file).exists() { return Err(format!("要去除的数据文件不存在: {}", remove_data_file)); } // 发送开始处理的进度 let _ = window.emit("data-cleaning-progress", DataCleaningProgress { current: 0, total: 0, percentage: 0.0, status: "开始处理...".to_string(), }); // 第一步:读取要去除的数据文件,提取所有URI let _ = window.emit("data-cleaning-progress", DataCleaningProgress { current: 0, total: 0, percentage: 10.0, status: "读取要去除的数据文件...".to_string(), }); let remove_uris = match extract_remove_uris(&remove_data_file).await { Ok(uris) => uris, Err(e) => return Err(format!("读取要去除的数据文件失败: {}", e)), }; let _ = window.emit("data-cleaning-progress", DataCleaningProgress { current: 0, total: 0, percentage: 30.0, status: format!("找到 {} 个要去除的URI", remove_uris.len()), }); // 第二步:处理全部数据文件 let _ = window.emit("data-cleaning-progress", DataCleaningProgress { current: 0, total: 0, percentage: 40.0, status: "开始处理全部数据文件...".to_string(), }); let result = match process_all_data_file(&all_data_file, &remove_uris, &output_file, &window).await { Ok(result) => result, Err(e) => return Err(format!("处理全部数据文件失败: {}", e)), }; let _ = window.emit("data-cleaning-progress", DataCleaningProgress { current: result.final_count, total: result.original_count, percentage: 100.0, status: "处理完成".to_string(), }); Ok(result) } /// 从要去除的数据文件中提取所有URI async fn extract_remove_uris(file_path: &str) -> Result, String> { let file = File::open(file_path) .map_err(|e| format!("无法打开文件: {}", e))?; let reader = BufReader::new(file); let mut uris = HashSet::new(); for line in reader.lines() { let line = line.map_err(|e| format!("读取行失败: {}", e))?; if line.trim().is_empty() { continue; } // 解析JSON let processed_data: ProcessedData = serde_json::from_str(&line) .map_err(|e| format!("解析JSON失败: {} - 行内容: {}", e, line))?; // 从json_data字段中解析内部JSON let inner_json: serde_json::Value = serde_json::from_str(&processed_data.json_data) .map_err(|e| format!("解析内部JSON失败: {}", e))?; // 提取uri字段 if let Some(uri) = inner_json.get("uri").and_then(|v| v.as_str()) { uris.insert(uri.to_string()); } } Ok(uris) } /// 处理全部数据文件,过滤掉要去除的URI async fn process_all_data_file( input_file: &str, remove_uris: &HashSet, output_file: &str, window: &tauri::Window, ) -> Result { let input_file_handle = File::open(input_file) .map_err(|e| format!("无法打开输入文件: {}", e))?; let reader = BufReader::new(input_file_handle); let mut output_file_handle = File::create(output_file) .map_err(|e| format!("无法创建输出文件: {}", e))?; let mut original_count = 0; let mut removed_count = 0; let mut final_count = 0; // 首先计算总行数用于进度显示 let total_lines = count_lines(input_file)?; for (line_number, line) in reader.lines().enumerate() { let line = line.map_err(|e| format!("读取行失败: {}", e))?; if line.trim().is_empty() { continue; } original_count += 1; // 解析JSON let storage_object: StorageObject = match serde_json::from_str(&line) { Ok(obj) => obj, Err(e) => { // 如果解析失败,跳过这一行但记录警告 eprintln!("警告: 解析JSON失败,跳过行 {}: {} - 内容: {}", line_number + 1, e, line); continue; } }; // 检查URI是否在要去除的列表中 if remove_uris.contains(&storage_object.uri) { removed_count += 1; } else { // 写入输出文件 writeln!(output_file_handle, "{}", line) .map_err(|e| format!("写入输出文件失败: {}", e))?; final_count += 1; } // 每处理100行发送一次进度更新 if line_number % 100 == 0 { let percentage = 40.0 + (line_number as f64 / total_lines as f64) * 50.0; let _ = window.emit("data-cleaning-progress", DataCleaningProgress { current: line_number, total: total_lines, percentage, status: format!("已处理 {} 行,已去除 {} 个重复项", line_number, removed_count), }); } } // 确保文件被正确写入 output_file_handle.flush() .map_err(|e| format!("刷新输出文件失败: {}", e))?; Ok(DataCleaningResult { success: true, message: "数据清洗完成".to_string(), original_count, removed_count, final_count, output_file: output_file.to_string(), }) } /// 计算文件行数 fn count_lines(file_path: &str) -> Result { let file = File::open(file_path) .map_err(|e| format!("无法打开文件计算行数: {}", e))?; let reader = BufReader::new(file); let count = reader.lines().count(); Ok(count) }