提交qiming-mcp-proxy

This commit is contained in:
Codex
2026-06-01 13:03:20 +08:00
parent 9e9486b7c2
commit afb3d9f4e6
394 changed files with 124494 additions and 0 deletions

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,392 @@
use crate::VoiceCliError;
use axum::extract::multipart::Field;
use bytes::Bytes;
use futures::TryStreamExt; // StreamExt 未使用,移除
use std::fs;
use std::path::{Path, PathBuf};
use tokio::io::AsyncWriteExt;
use tracing::{error, info, warn};
/// Service for managing audio files on disk
#[derive(Debug, Clone)]
pub struct AudioFileManager {
pub storage_dir: PathBuf,
}
impl AudioFileManager {
/// Create a new AudioFileManager
pub fn new<P: AsRef<Path>>(storage_dir: P) -> Result<Self, VoiceCliError> {
let storage_dir = storage_dir.as_ref().to_path_buf();
// Create storage directory if it doesn't exist
if !storage_dir.exists() {
fs::create_dir_all(&storage_dir).map_err(|e| {
VoiceCliError::Storage(format!(
"Failed to create audio storage directory '{}': {}",
storage_dir.display(),
e
))
})?;
}
info!(
"AudioFileManager initialized with storage directory: {}",
storage_dir.display()
);
Ok(Self { storage_dir })
}
/// Save audio data to disk and return the file path
pub async fn save_audio_file(
&self,
task_id: &str,
audio_data: &Bytes,
original_filename: &str,
) -> Result<PathBuf, VoiceCliError> {
// Extract file extension from original filename
let extension = Path::new(original_filename)
.extension()
.and_then(|ext| ext.to_str())
.unwrap_or("bin");
// Create a unique filename using task_id
let filename = format!("{}_{}.{}", task_id, uuid::Uuid::new_v4(), extension);
let file_path = self.storage_dir.join(&filename);
// Write audio data to file
tokio::fs::write(&file_path, audio_data)
.await
.map_err(|e| {
VoiceCliError::Storage(format!(
"Failed to write audio file '{}': {}",
file_path.display(),
e
))
})?;
info!(
"Saved audio file: {} ({} bytes) -> {}",
original_filename,
audio_data.len(),
file_path.display()
);
Ok(file_path)
}
/// Save audio data from multipart field stream directly to disk
pub async fn save_audio_file_streaming(
&self,
task_id: &str,
field: Field<'_>,
temp_file_name: &str,
) -> Result<String, VoiceCliError> {
// 获取原始文件名(如果有)用于日志记录
let original_filename = field
.file_name()
.map(|s| s.to_string())
.unwrap_or_else(|| "unknown".to_string());
info!(
"[Task {}] Start receiving audio file stream: {}, target temporary file name: {}",
task_id, original_filename, temp_file_name
);
let file_path = self.storage_dir.join(&temp_file_name);
// 确保存储目录存在
if let Some(parent) = file_path.parent() {
if !parent.exists() {
tokio::fs::create_dir_all(parent).await.map_err(|e| {
error!(
"[Task {}] Unable to create storage directory '{}': {}",
task_id,
parent.display(),
e
);
VoiceCliError::Storage(format!(
"无法创建存储目录 '{}': {}",
parent.display(),
e
))
})?;
}
}
// 创建文件
let file = tokio::fs::File::create(&file_path).await.map_err(|e| {
error!(
"[Task {}] Unable to create audio file '{}': {}",
task_id,
file_path.display(),
e
);
VoiceCliError::Storage(format!("无法创建音频文件 '{}': {}", file_path.display(), e))
})?;
// 创建缓冲写入器以提高性能
let mut writer = tokio::io::BufWriter::new(file);
// 将 field 转换为 StreamReader (实现 AsyncRead trait)
let mut reader = tokio_util::io::StreamReader::new(
field.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e)),
);
// 使用 tokio::io::copy 进行高效的流式复制
let total_bytes = tokio::io::copy(&mut reader, &mut writer)
.await
.map_err(|e| {
error!(
"[Task {}] Failed to stream audio file data ({} -> {}): {}",
task_id,
original_filename,
file_path.display(),
e
);
VoiceCliError::Storage(format!(
"流式复制音频文件数据失败 ({} -> {}): {}",
original_filename,
file_path.display(),
e
))
})?;
// 确保所有数据都写入磁盘
writer.flush().await.map_err(|e| {
error!(
"[Task {}] Unable to refresh data to file '{}': {}",
task_id,
file_path.display(),
e
);
VoiceCliError::Storage(format!(
"无法刷新数据到文件 '{}': {}",
file_path.display(),
e
))
})?;
info!(
"[Task {}] Successfully received and saved audio file: {} ({} bytes) -> {}",
task_id,
original_filename,
total_bytes,
file_path.display()
);
Ok(file_path.to_string_lossy().into_owned())
}
/// Delete an audio file from disk
pub async fn delete_audio_file<P: AsRef<Path>>(
&self,
file_path: P,
) -> Result<(), VoiceCliError> {
let file_path = file_path.as_ref();
if file_path.exists() {
tokio::fs::remove_file(file_path).await.map_err(|e| {
VoiceCliError::Storage(format!(
"Failed to delete audio file '{}': {}",
file_path.display(),
e
))
})?;
info!("Deleted audio file: {}", file_path.display());
} else {
warn!("Audio file not found for deletion: {}", file_path.display());
}
Ok(())
}
/// Delete multiple audio files
pub async fn delete_audio_files<P: AsRef<Path>>(
&self,
file_paths: &[P],
) -> Result<(), VoiceCliError> {
for file_path in file_paths {
if let Err(e) = self.delete_audio_file(file_path).await {
// Log error but continue with other files
error!("Failed to delete audio file: {}", e);
}
}
Ok(())
}
/// Clean up old audio files based on age
pub async fn cleanup_old_files(&self, max_age_hours: u64) -> Result<u32, VoiceCliError> {
let mut cleaned_count = 0u32;
let cutoff_time =
std::time::SystemTime::now() - std::time::Duration::from_secs(max_age_hours * 3600);
let mut entries = tokio::fs::read_dir(&self.storage_dir).await.map_err(|e| {
VoiceCliError::Storage(format!(
"Failed to read storage directory '{}': {}",
self.storage_dir.display(),
e
))
})?;
while let Some(entry) = entries
.next_entry()
.await
.map_err(|e| VoiceCliError::Storage(format!("Failed to read directory entry: {}", e)))?
{
let path = entry.path();
if path.is_file() {
if let Ok(metadata) = entry.metadata().await {
if let Ok(modified) = metadata.modified() {
if modified < cutoff_time {
if let Err(e) = self.delete_audio_file(&path).await {
error!("Failed to cleanup old file '{}': {}", path.display(), e);
} else {
cleaned_count += 1;
}
}
}
}
}
}
if cleaned_count > 0 {
info!("Cleaned up {} old audio files", cleaned_count);
}
Ok(cleaned_count)
}
/// Get the size of a file
pub async fn get_file_size<P: AsRef<Path>>(&self, file_path: P) -> Result<u64, VoiceCliError> {
let metadata = tokio::fs::metadata(file_path.as_ref()).await.map_err(|e| {
VoiceCliError::Storage(format!(
"Failed to get file metadata for '{}': {}",
file_path.as_ref().display(),
e
))
})?;
Ok(metadata.len())
}
/// Get total storage usage
pub async fn get_storage_usage(&self) -> Result<u64, VoiceCliError> {
let mut total_size = 0u64;
let mut entries = tokio::fs::read_dir(&self.storage_dir).await.map_err(|e| {
VoiceCliError::Storage(format!(
"Failed to read storage directory '{}': {}",
self.storage_dir.display(),
e
))
})?;
while let Some(entry) = entries
.next_entry()
.await
.map_err(|e| VoiceCliError::Storage(format!("Failed to read directory entry: {}", e)))?
{
if entry.path().is_file() {
if let Ok(metadata) = entry.metadata().await {
total_size += metadata.len();
}
}
}
Ok(total_size)
}
/// Check if a file exists
pub async fn file_exists<P: AsRef<Path>>(&self, file_path: P) -> bool {
tokio::fs::metadata(file_path.as_ref()).await.is_ok()
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[tokio::test]
async fn test_audio_file_manager_creation() {
let temp_dir = TempDir::new().unwrap();
let manager = AudioFileManager::new(temp_dir.path()).unwrap();
assert_eq!(manager.storage_dir, temp_dir.path());
}
#[tokio::test]
async fn test_save_and_delete_audio_file() {
let temp_dir = TempDir::new().unwrap();
let manager = AudioFileManager::new(temp_dir.path()).unwrap();
let audio_data = Bytes::from(vec![1, 2, 3, 4, 5]);
let task_id = "test-task-123";
let original_filename = "test.mp3";
// Save file
let file_path = manager
.save_audio_file(task_id, &audio_data, original_filename)
.await
.unwrap();
// Check file exists
assert!(manager.file_exists(&file_path).await);
// Check file size
let size = manager.get_file_size(&file_path).await.unwrap();
assert_eq!(size, 5);
// Delete file
manager.delete_audio_file(&file_path).await.unwrap();
// Check file no longer exists
assert!(!manager.file_exists(&file_path).await);
}
#[tokio::test]
async fn test_cleanup_old_files() {
let temp_dir = TempDir::new().unwrap();
let manager = AudioFileManager::new(temp_dir.path()).unwrap();
let audio_data = Bytes::from(vec![1, 2, 3, 4, 5]);
// Save a file
let file_path = manager
.save_audio_file("test-task", &audio_data, "test.mp3")
.await
.unwrap();
// File should exist
assert!(manager.file_exists(&file_path).await);
// Cleanup files older than 0 hours (should clean everything)
let cleaned = manager.cleanup_old_files(0).await.unwrap();
// Should have cleaned at least 1 file
assert!(cleaned >= 1);
}
#[tokio::test]
async fn test_storage_usage() {
let temp_dir = TempDir::new().unwrap();
let manager = AudioFileManager::new(temp_dir.path()).unwrap();
// Initially should be 0
let initial_usage = manager.get_storage_usage().await.unwrap();
assert_eq!(initial_usage, 0);
// Save a file
let audio_data = Bytes::from(vec![1, 2, 3, 4, 5]);
let _file_path = manager
.save_audio_file("test-task", &audio_data, "test.mp3")
.await
.unwrap();
// Usage should increase
let usage_after = manager.get_storage_usage().await.unwrap();
assert_eq!(usage_after, 5);
}
}

View File

@@ -0,0 +1,267 @@
use bytes::Bytes;
use infer::{self, Type};
use std::io::Cursor;
use std::path::Path;
use symphonia::core::formats::{FormatReader, Track};
use symphonia::core::io::MediaSourceStream;
use symphonia::core::meta::MetadataOptions;
use symphonia::core::probe::Hint;
use symphonia::core::probe::ProbeResult;
use symphonia::default::get_probe;
use tracing::{error, info, warn};
use crate::error::VoiceCliError;
use crate::models::request::{AudioFormat, AudioFormatResult, AudioMetadata, DetectionMethod};
/// Service for intelligent audio format detection using Symphonia
pub struct AudioFormatDetector;
impl AudioFormatDetector {
/// Detect audio format using infer library (magic number detection)
pub fn detect_format_from_path(path: &Path) -> anyhow::Result<Option<Type>> {
let kind = infer::get_from_path(path)
.map_err(|e| anyhow::anyhow!("Failed to read file for format detection: {}", e))?;
Ok(kind)
}
/// Detect audio format using Symphonia probe with fallback to infer (magic number detection) and filename extension
pub fn detect_format(
audio_data: &Bytes,
filename: Option<&str>,
) -> Result<AudioFormatResult, VoiceCliError> {
info!(
"Starting audio format detection for {} byte audio data",
audio_data.len()
);
// Try Symphonia probe first (primary method)
if let Ok(result) = Self::symphonia_probe(audio_data, filename) {
info!(
"Successfully detected format using Symphonia probe: {:?}",
result.format
);
return Ok(result);
}
// Fallback to filename extension if provided
if let Some(filename) = filename {
let format = AudioFormat::from_filename(filename);
if format.is_supported() {
info!("Format detected from filename extension: {:?}", format);
return Ok(AudioFormatResult {
format,
confidence: 0.5, // Lower confidence for filename-based detection
metadata: None,
detection_method: DetectionMethod::FileExtension,
});
}
}
// All detection methods failed
error!("All audio format detection methods failed");
Err(VoiceCliError::UnsupportedFormat(
"Unable to detect audio format using any available method".to_string(),
))
}
/// Primary detection method using Symphonia
fn symphonia_probe(
audio_data: &Bytes,
filename: Option<&str>,
) -> Result<AudioFormatResult, VoiceCliError> {
// Create a cursor from copied audio data to avoid lifetime issues
let data_copy = audio_data.to_vec();
let cursor = Cursor::new(data_copy);
let media_source = MediaSourceStream::new(Box::new(cursor), Default::default());
// Create a hint based on filename if available
let mut hint = Hint::new();
if let Some(filename) = filename {
if let Some(extension) = std::path::Path::new(filename)
.extension()
.and_then(|ext| ext.to_str())
{
hint.with_extension(extension);
}
}
// Get the default probe
let probe = get_probe();
// Attempt to probe the media source
let probe_result = probe
.format(
&hint,
media_source,
&FormatOptions::default(),
&MetadataOptions::default(),
)
.map_err(|e| {
warn!("Symphonia probe failed: {}", e);
VoiceCliError::UnsupportedFormat(format!("Symphonia probe error: {}", e))
})?;
// Extract format information
let format_info = Self::extract_format_info(&probe_result)?;
Ok(format_info)
}
/// Extract format and metadata information from Symphonia probe result
fn extract_format_info(probe_result: &ProbeResult) -> Result<AudioFormatResult, VoiceCliError> {
let reader = &probe_result.format;
let tracks = reader.tracks();
// Find the first audio track (any track with codec parameters)
let track = tracks
.iter()
.find(|t| t.codec_params.codec != symphonia::core::codecs::CODEC_TYPE_NULL)
.ok_or_else(|| {
VoiceCliError::UnsupportedFormat("No audio tracks found in file".to_string())
})?;
// Convert codec type to our AudioFormat
let format = AudioFormat::from_symphonia_codec(track.codec_params.codec);
// Extract metadata
let metadata = Self::extract_metadata(track, reader);
// Calculate confidence based on detection success
let confidence = if format != AudioFormat::Unknown {
0.95
} else {
0.0
};
if format == AudioFormat::Unknown {
return Err(VoiceCliError::UnsupportedFormat(format!(
"Unsupported codec type: {:?}",
track.codec_params.codec
)));
}
Ok(AudioFormatResult {
format,
confidence,
metadata: Some(metadata),
detection_method: DetectionMethod::SymphoniaProbe,
})
}
/// Extract detailed audio metadata from track and format reader
fn extract_metadata(track: &Track, _reader: &Box<dyn FormatReader>) -> AudioMetadata {
let codec_params = &track.codec_params;
// Extract basic parameters
let sample_rate = codec_params.sample_rate;
let channels = codec_params.channels.map(|ch| ch.count() as u8);
let bit_depth = codec_params.bits_per_sample.map(|bits| bits as u8);
// Calculate duration if time base and n_frames are available
let duration = if let (Some(time_base), Some(n_frames)) =
(codec_params.time_base, codec_params.n_frames)
{
let duration_secs = n_frames as f64 * time_base.numer as f64 / time_base.denom as f64;
Some(std::time::Duration::from_secs_f64(duration_secs))
} else {
None
};
// Estimate bitrate if possible
let bitrate = if let (Some(sample_rate), Some(channels), Some(bit_depth)) =
(sample_rate, channels, bit_depth)
{
Some(sample_rate * channels as u32 * bit_depth as u32)
} else {
None
};
// Create codec info string
let codec_info = format!("Codec: {:?}", codec_params.codec);
AudioMetadata {
duration,
sample_rate,
channels,
bit_depth,
bitrate,
codec_info,
}
}
/// Validate that the detected format is supported for transcription
pub fn validate_format_support(format_result: &AudioFormatResult) -> Result<(), VoiceCliError> {
if !format_result.format.is_supported() {
return Err(VoiceCliError::UnsupportedFormat(format!(
"Format {} is not supported for transcription",
format_result.format.to_string()
)));
}
// Check confidence threshold
if format_result.confidence < 0.3 {
warn!(
"Low confidence format detection: {}",
format_result.confidence
);
}
Ok(())
}
}
// Import FormatOptions for compilation
use symphonia::core::formats::FormatOptions;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_format_detection_with_filename() {
let test_data = Bytes::from(vec![0u8; 1024]); // Dummy data
// Test filename-based fallback
let result = AudioFormatDetector::detect_format(&test_data, Some("test.mp3"));
match result {
Ok(format_result) => {
assert_eq!(
format_result.detection_method,
DetectionMethod::FileExtension
);
assert_eq!(format_result.format, AudioFormat::Mp3);
}
Err(_) => {
// Expected for dummy data, but we should get filename-based detection
}
}
}
#[test]
fn test_unsupported_format() {
let test_data = Bytes::from(vec![0u8; 1024]);
let result = AudioFormatDetector::detect_format(&test_data, Some("test.xyz"));
assert!(result.is_err());
}
#[test]
fn test_format_validation() {
let format_result = AudioFormatResult {
format: AudioFormat::Mp3,
confidence: 0.9,
metadata: None,
detection_method: DetectionMethod::SymphoniaProbe,
};
assert!(AudioFormatDetector::validate_format_support(&format_result).is_ok());
let unsupported_result = AudioFormatResult {
format: AudioFormat::Unknown,
confidence: 0.9,
metadata: None,
detection_method: DetectionMethod::SymphoniaProbe,
};
assert!(AudioFormatDetector::validate_format_support(&unsupported_result).is_err());
}
}

View File

@@ -0,0 +1,314 @@
use crate::VoiceCliError;
use crate::models::AudioFormat;
use crate::models::request::ProcessedAudio;
use bytes::Bytes;
use std::io::Write;
use std::path::PathBuf;
use tempfile::NamedTempFile;
use tracing::{debug, info, warn};
#[derive(Debug)]
pub struct AudioProcessor {
temp_dir: PathBuf,
}
impl AudioProcessor {
pub fn new(temp_dir: Option<PathBuf>) -> Self {
let temp_dir = temp_dir.unwrap_or_else(|| std::env::temp_dir().join("voice-cli"));
// Ensure temp directory exists
if let Err(e) = std::fs::create_dir_all(&temp_dir) {
warn!("Failed to create temp directory {:?}: {}", temp_dir, e);
}
Self { temp_dir }
}
/// Process audio data and convert to whisper-compatible format if needed
pub async fn process_audio(
&self,
audio_data: Bytes,
filename: Option<&str>,
) -> Result<ProcessedAudio, VoiceCliError> {
debug!(
"Processing audio data: {} bytes, filename: {:?}",
audio_data.len(),
filename
);
// Detect audio format
let format = self.detect_audio_format(&audio_data, filename)?;
debug!("Detected audio format: {:?}", format);
// For WAV format, validate the content since it doesn't need conversion
if format == AudioFormat::Wav {
self.validate_whisper_format(&audio_data)?;
}
// Check if conversion is needed
if !format.needs_conversion() {
debug!("Audio format is already compatible, no conversion needed");
return Ok(ProcessedAudio {
data: audio_data,
converted: false,
original_format: Some(format.to_string().to_string()),
});
}
// Convert to Whisper-compatible format
info!("Converting audio from {} to WAV format", format.to_string());
let converted_data = self.convert_to_whisper_format(audio_data, format).await?;
Ok(ProcessedAudio {
data: converted_data,
converted: true,
original_format: Some(format.to_string().to_string()),
})
}
/// Detect audio format from data and filename
fn detect_audio_format(
&self,
audio_data: &Bytes,
filename: Option<&str>,
) -> Result<AudioFormat, VoiceCliError> {
// Use AudioFormatDetector for enhanced detection
use crate::services::AudioFormatDetector;
match AudioFormatDetector::detect_format(audio_data, filename) {
Ok(format_result) => {
// Validate format support
AudioFormatDetector::validate_format_support(&format_result)?;
Ok(format_result.format)
}
Err(e) => Err(e),
}
}
/// Convert audio to Whisper-compatible format (16kHz, mono, 16-bit PCM WAV)
async fn convert_to_whisper_format(
&self,
audio_data: Bytes,
source_format: AudioFormat,
) -> Result<Bytes, VoiceCliError> {
debug!("Converting {} to WAV format", source_format.to_string());
// Create temporary files for input and output
let input_file = self.create_temp_file(&audio_data, &source_format)?;
let output_file = NamedTempFile::new_in(&self.temp_dir).map_err(|e| {
VoiceCliError::AudioProcessing(format!("Failed to create temp output file: {}", e))
})?;
let input_path = input_file.path();
let output_path = output_file.path();
// Try to use rs-voice-toolkit for conversion
match self
.convert_with_rs_voice_toolkit(input_path, output_path)
.await
{
Ok(_) => {
// Read converted file
let converted_data = std::fs::read(output_path).map_err(|e| {
VoiceCliError::AudioProcessing(format!("Failed to read converted file: {}", e))
})?;
info!(
"Successfully converted audio: {} -> {} bytes",
audio_data.len(),
converted_data.len()
);
Ok(Bytes::from(converted_data))
}
Err(toolkit_error) => {
warn!(
"rs-voice-toolkit conversion failed: {}, trying fallback",
toolkit_error
);
return Err(toolkit_error);
}
}
}
/// Convert using rs-voice-toolkit
async fn convert_with_rs_voice_toolkit(
&self,
input_path: &std::path::Path,
output_path: &std::path::Path,
) -> Result<(), VoiceCliError> {
debug!(
"Converting audio using voice-toolkit: {:?} -> {:?}",
input_path, output_path
);
// Use voice_toolkit::audio::ensure_whisper_compatible for real audio conversion
match voice_toolkit::audio::ensure_whisper_compatible(
input_path,
Some(output_path.to_path_buf()),
) {
Ok(compatible_wav) => {
debug!(
"Successfully converted audio to whisper-compatible format: {:?}",
compatible_wav.path
);
// Verify the output file exists and is at the expected location
if compatible_wav.path != output_path {
// If the output is in a different location, move it to the expected location
if let Err(e) = std::fs::rename(&compatible_wav.path, output_path) {
warn!("Failed to move converted file to expected location: {}", e);
// Try copying instead
std::fs::copy(&compatible_wav.path, output_path).map_err(|e| {
VoiceCliError::AudioProcessing(format!(
"Failed to copy converted file: {}",
e
))
})?;
// Clean up the original if copy succeeded
let _ = std::fs::remove_file(&compatible_wav.path);
}
}
info!("Audio conversion completed successfully");
Ok(())
}
Err(e) => {
warn!("voice-toolkit conversion failed: {}", e);
Err(VoiceCliError::AudioProcessing(format!(
"Audio conversion failed: {}",
e
)))
}
}
}
/// Create temporary file with audio data
fn create_temp_file(
&self,
audio_data: &Bytes,
format: &AudioFormat,
) -> Result<NamedTempFile, VoiceCliError> {
let extension = format.to_string();
let mut temp_file =
NamedTempFile::with_suffix_in(&format!(".{}", extension), &self.temp_dir).map_err(
|e| VoiceCliError::AudioProcessing(format!("Failed to create temp file: {}", e)),
)?;
temp_file.write_all(audio_data).map_err(|e| {
VoiceCliError::AudioProcessing(format!("Failed to write temp file: {}", e))
})?;
temp_file.flush().map_err(|e| {
VoiceCliError::AudioProcessing(format!("Failed to flush temp file: {}", e))
})?;
Ok(temp_file)
}
/// Validate that the processed audio is in the correct format for Whisper
pub fn validate_whisper_format(&self, audio_data: &Bytes) -> Result<(), VoiceCliError> {
// Basic WAV header validation
if audio_data.len() < 44 {
return Err(VoiceCliError::AudioProcessing(
"Audio file too small to be valid WAV".to_string(),
));
}
let header = &audio_data[0..44];
// Check RIFF header
if &header[0..4] != b"RIFF" {
return Err(VoiceCliError::AudioProcessing(
"Invalid WAV file: missing RIFF header".to_string(),
));
}
// Check WAVE format
if &header[8..12] != b"WAVE" {
return Err(VoiceCliError::AudioProcessing(
"Invalid WAV file: missing WAVE format".to_string(),
));
}
// Check fmt chunk
if &header[12..16] != b"fmt " {
return Err(VoiceCliError::AudioProcessing(
"Invalid WAV file: missing fmt chunk".to_string(),
));
}
// Extract audio parameters
let sample_rate = u32::from_le_bytes([header[24], header[25], header[26], header[27]]);
let channels = u16::from_le_bytes([header[22], header[23]]);
let bits_per_sample = u16::from_le_bytes([header[34], header[35]]);
debug!(
"WAV format - Sample rate: {}Hz, Channels: {}, Bits: {}",
sample_rate, channels, bits_per_sample
);
// Whisper prefers 16kHz, mono, 16-bit, but can handle other formats too
// We'll just warn for non-optimal formats rather than error
if sample_rate != 16000 {
warn!(
"Non-optimal sample rate: {}Hz (Whisper prefers 16kHz)",
sample_rate
);
}
if channels != 1 {
warn!(
"Non-optimal channel count: {} (Whisper prefers mono)",
channels
);
}
if bits_per_sample != 16 {
warn!(
"Non-optimal bit depth: {} (Whisper prefers 16-bit)",
bits_per_sample
);
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[tokio::test]
async fn test_audio_processor_creation() {
let temp_dir = TempDir::new().unwrap();
let processor = AudioProcessor::new(Some(temp_dir.path().to_path_buf()));
// Test with mock WAV data (basic WAV header)
let wav_header = b"RIFF\x24\x00\x00\x00WAVE";
let audio_data = Bytes::from_static(wav_header);
let format = processor.detect_audio_format(&audio_data, Some("test.wav"));
assert!(matches!(format, Ok(AudioFormat::Wav)));
}
#[test]
fn test_audio_format_detection() {
let processor = AudioProcessor::new(None);
// Test WAV detection from filename
let wav_data = Bytes::from_static(b"RIFF\x24\x00\x00\x00WAVE");
assert!(matches!(
processor.detect_audio_format(&wav_data, Some("test.wav")),
Ok(AudioFormat::Wav)
));
// Test MP3 detection from filename
let mp3_data = Bytes::from_static(b"\xFF\xFB\x90\x00");
assert!(matches!(
processor.detect_audio_format(&mp3_data, Some("test.mp3")),
Ok(AudioFormat::Mp3)
));
}
}

View File

@@ -0,0 +1,367 @@
use crate::VoiceCliError;
use serde::{Deserialize, Serialize};
use std::fs::Metadata;
use std::path::Path;
use tokio::{fs, task};
use tracing::{debug, info, warn};
/// 音视频元数据信息
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AudioVideoMetadata {
// 基础信息
pub format: String, // 文件格式 (mp3, wav, mp4, etc.)
pub container_format: String, // 容器格式
pub duration_seconds: f64, // 时长(秒)
pub file_size_bytes: u64, // 文件大小
// 音频信息
pub audio_codec: String, // 音频编码器
pub sample_rate: u32, // 采样率 (Hz)
pub channels: u8, // 声道数
pub audio_bitrate: u32, // 音频码率 (kbps)
// 视频信息(如果是视频文件)
pub has_video: bool, // 是否包含视频
pub video_codec: Option<String>, // 视频编码器
pub width: Option<u32>, // 视频宽度
pub height: Option<u32>, // 视频高度
pub video_bitrate: Option<u32>, // 视频码率 (kbps)
pub frame_rate: Option<f64>, // 帧率
// 其他元数据
pub bitrate: u32, // 总码率 (kbps)
pub creation_time: Option<String>, // 创建时间
}
impl Default for AudioVideoMetadata {
fn default() -> Self {
Self {
format: "unknown".to_string(),
container_format: "unknown".to_string(),
duration_seconds: 0.0,
file_size_bytes: 0,
audio_codec: "unknown".to_string(),
sample_rate: 0,
channels: 1,
audio_bitrate: 0,
has_video: false,
video_codec: None,
width: None,
height: None,
video_bitrate: None,
frame_rate: None,
bitrate: 0,
creation_time: None,
}
}
}
/// 音视频元数据提取器
pub struct MetadataExtractor;
impl MetadataExtractor {
/// 从文件路径提取音视频元数据
pub async fn extract_metadata(file_path: &Path) -> Result<AudioVideoMetadata, VoiceCliError> {
info!("Start extracting audio and video metadata: {:?}", file_path);
// 首先获取文件基本信息
let file_metadata = fs::metadata(file_path)
.await
.map_err(|e| VoiceCliError::Storage(format!("无法访问文件: {}", e)))?;
let _file_size = file_metadata.len();
let file_extension = file_path
.extension()
.and_then(|ext| ext.to_str())
.unwrap_or("unknown")
.to_lowercase();
// 尝试使用 FFmpeg 提取详细元数据
if let Ok(ffmpeg_metadata) = Self::extract_with_ffmpeg(file_path).await {
info!("FFmpeg metadata extraction successful");
return Ok(ffmpeg_metadata);
}
// 如果 FFmpeg 不可用或失败,使用基础方法
warn!("FFmpeg is unavailable or failed, use base metadata extraction method");
Self::extract_basic_metadata(file_path, &file_metadata, &file_extension).await
}
/// 使用 FFmpeg 提取详细元数据
async fn extract_with_ffmpeg(file_path: &Path) -> Result<AudioVideoMetadata, VoiceCliError> {
use ffmpeg_sidecar::command::FfmpegCommand;
debug!("Extract metadata using FFmpeg: {:?}", file_path);
let file_path_buf = file_path.to_path_buf();
let metadata =
task::spawn_blocking(move || -> Result<AudioVideoMetadata, VoiceCliError> {
let mut metadata = AudioVideoMetadata::default();
let file_path_str = file_path_buf.to_string_lossy().to_string();
// 使用 FfmpegCommand 获取文件信息
let mut child = FfmpegCommand::new()
.arg("-i")
.arg(&file_path_str)
.arg("-hide_banner")
.spawn()
.map_err(|e| VoiceCliError::Storage(format!("FFmpeg 执行失败: {}", e)))?;
// 等待命令完成(在阻塞线程中执行)
let _exit_status = child
.wait()
.map_err(|e| VoiceCliError::Storage(format!("FFmpeg 执行失败: {}", e)))?;
// 使用传统方法获取输出(因为 ffmpeg-sidecar 主要用于处理媒体流,不是元数据提取)
let output = std::process::Command::new("ffmpeg")
.args(["-i", &file_path_str, "-hide_banner", "-f", "null", "-"])
.output()
.map_err(|e| VoiceCliError::Storage(format!("FFmpeg 执行失败: {}", e)))?;
// 解析 stderr 输出中的元数据信息
let stderr_output = String::from_utf8_lossy(&output.stderr);
// 解析输出中的元数据信息
for line in stderr_output.lines() {
if line.contains("Duration:") {
// 解析时长: Duration: 00:00:01.60, start: 0.000000, bitrate: 705 kb/s
if let Some(duration_part) = line.split("Duration: ").nth(1) {
if let Some(duration_str) = duration_part.split(',').next() {
let parts: Vec<&str> = duration_str.split(':').collect();
if parts.len() == 3 {
let hours: f64 = parts[0].parse().unwrap_or(0.0);
let minutes: f64 = parts[1].parse().unwrap_or(0.0);
let seconds: f64 = parts[2].parse().unwrap_or(0.0);
metadata.duration_seconds =
hours * 3600.0 + minutes * 60.0 + seconds;
}
}
}
}
if line.contains("Audio:") {
// 解析音频信息: Stream #0:0: Audio: pcm_f32le, 22050 Hz, mono, fltp, 705 kb/s
let audio_info = line.split("Audio: ").nth(1).unwrap_or("");
let parts: Vec<&str> = audio_info.split(',').collect();
if let Some(codec) = parts.first() {
metadata.audio_codec = codec.trim().to_string();
}
for part in parts {
if part.contains("Hz") {
if let Some(rate_str) = part.split("Hz").next() {
metadata.sample_rate = rate_str.trim().parse().unwrap_or(0);
}
}
if part.contains("mono") {
metadata.channels = 1;
}
if part.contains("stereo") {
metadata.channels = 2;
}
if part.contains("kb/s") {
if let Some(bitrate_str) = part.split("kb/s").next() {
metadata.audio_bitrate =
bitrate_str.trim().parse().unwrap_or(0);
}
}
}
}
if line.contains("Video:") {
// 解析视频信息: Stream #0:1: Video: h264, yuv420p, 1280x720, 24 fps, 1992 kb/s
metadata.has_video = true;
let video_info = line.split("Video: ").nth(1).unwrap_or("");
let parts: Vec<&str> = video_info.split(',').collect();
if let Some(codec) = parts.first() {
metadata.video_codec = Some(codec.trim().to_string());
}
for part in parts {
if part.contains('x') {
let resolution_parts: Vec<&str> = part.trim().split('x').collect();
if resolution_parts.len() == 2 {
metadata.width = resolution_parts[0].trim().parse().ok();
metadata.height = resolution_parts[1].trim().parse().ok();
}
}
if part.contains("fps") {
if let Some(fps_str) = part.split("fps").next() {
metadata.frame_rate = fps_str.trim().parse().ok();
}
}
if part.contains("kb/s") {
if let Some(bitrate_str) = part.split("kb/s").next() {
metadata.video_bitrate =
Some(bitrate_str.trim().parse().unwrap_or(0));
}
}
}
}
}
// 获取文件大小
if let Ok(file_meta) = std::fs::metadata(&file_path_buf) {
metadata.file_size_bytes = file_meta.len();
}
// 如果没有从输出中获取到码率,计算总码率
if metadata.bitrate == 0
&& metadata.duration_seconds > 0.0
&& metadata.file_size_bytes > 0
{
let total_bits = metadata.file_size_bytes as f64 * 8.0;
metadata.bitrate = (total_bits / metadata.duration_seconds / 1000.0) as u32;
}
// 根据文件扩展名设置格式
if let Some(extension) = file_path_buf.extension().and_then(|ext| ext.to_str()) {
metadata.format = extension.to_lowercase();
metadata.container_format = extension.to_lowercase();
}
Ok(metadata)
})
.await
.map_err(|e| VoiceCliError::Storage(format!("FFmpeg 阻塞任务失败: {}", e)))??;
info!("FFmpeg metadata extraction completed: {:?}", metadata);
Ok(metadata)
}
/// 基础元数据提取(不依赖 FFmpeg
async fn extract_basic_metadata(
file_path: &Path,
file_metadata: &Metadata,
file_extension: &str,
) -> Result<AudioVideoMetadata, VoiceCliError> {
debug!("Extract metadata using basic method: {:?}", file_path);
let mut metadata = AudioVideoMetadata {
file_size_bytes: file_metadata.len(),
format: file_extension.to_string(),
container_format: file_extension.to_string(),
..Default::default()
};
// 尝试使用现有的 AudioFormatDetector 获取音频信息
if let Ok(Some(format_type)) =
crate::services::AudioFormatDetector::detect_format_from_path(file_path)
{
metadata.format = format_type.extension().to_string();
metadata.audio_codec = format_type.mime_type().to_string();
}
// 判断是否为视频文件
metadata.has_video = Self::is_video_format(file_extension);
// 如果是视频文件,设置默认值
if metadata.has_video {
metadata.video_codec = Some("unknown".to_string());
}
// 计算码率
if metadata.duration_seconds > 0.0 && metadata.file_size_bytes > 0 {
let total_bits = metadata.file_size_bytes as f64 * 8.0;
metadata.bitrate = (total_bits / metadata.duration_seconds / 1000.0) as u32;
}
info!("Basic metadata extraction completed: {:?}", metadata);
Ok(metadata)
}
/// 判断是否为视频格式
fn is_video_format(extension: &str) -> bool {
matches!(
extension,
"mp4" | "avi" | "mkv" | "mov" | "wmv" | "flv" | "webm" | "m4v" | "3gp" | "mpg" | "mpeg"
)
}
/// 获取文件格式描述
pub fn get_format_description(metadata: &AudioVideoMetadata) -> String {
if metadata.has_video {
format!(
"视频文件 - 格式: {}, 分辨率: {}x{}, 时长: {:.2}s",
metadata.format,
metadata.width.unwrap_or(0),
metadata.height.unwrap_or(0),
metadata.duration_seconds
)
} else {
format!(
"音频文件 - 格式: {}, 采样率: {}Hz, 声道: {}, 时长: {:.2}s",
metadata.format, metadata.sample_rate, metadata.channels, metadata.duration_seconds
)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
use tempfile::NamedTempFile;
#[tokio::test]
async fn test_extract_basic_metadata() {
// 创建一个临时文件进行测试
let mut temp_file = NamedTempFile::new().unwrap();
temp_file.write_all(b"dummy audio data").unwrap();
temp_file.flush().unwrap();
let metadata = MetadataExtractor::extract_metadata(temp_file.path()).await;
// 验证基本结构
match metadata {
Ok(meta) => {
assert!(!meta.format.is_empty());
assert_eq!(meta.file_size_bytes, 16); // "dummy audio data" 的长度
}
Err(e) => {
println!("Failed to extract: {}", e);
// 对于测试环境FFmpeg 可能不可用,这是可以接受的
}
}
}
#[test]
fn test_is_video_format() {
assert!(MetadataExtractor::is_video_format("mp4"));
assert!(MetadataExtractor::is_video_format("avi"));
assert!(!MetadataExtractor::is_video_format("mp3"));
assert!(!MetadataExtractor::is_video_format("wav"));
}
#[test]
fn test_get_format_description() {
let audio_meta = AudioVideoMetadata {
format: "mp3".to_string(),
sample_rate: 44100,
channels: 2,
duration_seconds: 180.5,
has_video: false,
..Default::default()
};
let video_meta = AudioVideoMetadata {
format: "mp4".to_string(),
width: Some(1920),
height: Some(1080),
duration_seconds: 120.0,
has_video: true,
..Default::default()
};
let audio_desc = MetadataExtractor::get_format_description(&audio_meta);
let video_desc = MetadataExtractor::get_format_description(&video_meta);
assert!(audio_desc.contains("音频文件"));
assert!(audio_desc.contains("44100Hz"));
assert!(video_desc.contains("视频文件"));
assert!(video_desc.contains("1920x1080"));
}
}

View File

@@ -0,0 +1,24 @@
pub mod apalis_manager;
pub mod audio_file_manager;
pub mod audio_format_detector;
pub mod audio_processor;
pub mod metadata_extractor;
pub mod model_service;
pub mod transcription_engine;
pub mod tts_service;
pub mod tts_task_manager;
// 重新导出核心服务
pub use apalis_manager::{
ApalisManager, LockFreeApalisManager, StepContext, TaskStatusUpdate, TranscriptionTask,
init_global_apalis_manager, init_global_lock_free_apalis_manager,
transcription_pipeline_worker,
};
pub use audio_file_manager::AudioFileManager;
pub use audio_format_detector::AudioFormatDetector;
pub use audio_processor::AudioProcessor;
pub use metadata_extractor::{AudioVideoMetadata, MetadataExtractor};
pub use model_service::ModelService;
pub use transcription_engine::TranscriptionEngine;
pub use tts_service::TtsService;
pub use tts_task_manager::{TtsTaskManager, TtsTaskStats};

View File

@@ -0,0 +1,524 @@
use crate::VoiceCliError;
use crate::models::{Config, DownloadStatus, ModelDownloadStatus, ModelInfo};
use reqwest::Client;
use std::path::{Path, PathBuf};
use tokio::fs;
use tokio::io::AsyncWriteExt;
use tracing::{debug, info, warn};
#[derive(Debug)]
pub struct ModelService {
config: Config,
client: Client,
models_dir: PathBuf,
}
impl ModelService {
pub fn new(config: Config) -> Self {
Self {
models_dir: config.models_dir_path(),
config,
client: Client::new(),
}
}
/// Get the default model name from configuration
pub fn default_model(&self) -> &str {
&self.config.whisper.default_model
}
/// Get the worker timeout from configuration
pub fn worker_timeout(&self) -> u64 {
self.config.whisper.workers.worker_timeout as u64
}
/// Ensure a model is available (download if necessary)
pub async fn ensure_model(&self, model_name: &str) -> Result<(), VoiceCliError> {
if self.is_model_downloaded(model_name).await? {
debug!("Model '{}' already exists", model_name);
return Ok(());
}
if self.config.whisper.auto_download {
info!("Auto-downloading model: {}", model_name);
self.download_model(model_name).await?;
} else {
return Err(VoiceCliError::ModelNotFound(format!(
"Model '{}' not found and auto_download is disabled",
model_name
)));
}
Ok(())
}
/// Download a whisper model from the official repository
pub async fn download_model(&self, model_name: &str) -> Result<(), VoiceCliError> {
if !self
.config
.whisper
.supported_models
.contains(&model_name.to_string())
{
return Err(VoiceCliError::InvalidModelName(format!(
"Model '{}' is not supported",
model_name
)));
}
// Create models directory if it doesn't exist
fs::create_dir_all(&self.models_dir).await?;
let model_path = self.get_model_path(model_name)?;
if model_path.exists() {
info!("Model '{}' already exists at {:?}", model_name, model_path);
return Ok(());
}
info!(
"Downloading model '{}' from whisper.cpp repository...",
model_name
);
// Download from Hugging Face (official whisper.cpp models)
let download_url = self.get_model_download_url(model_name)?;
debug!("Download URL: {}", download_url);
// Download with progress tracking
let response = self
.client
.get(&download_url)
.send()
.await
.map_err(|e| VoiceCliError::Model(format!("Failed to start download: {}", e)))?;
if !response.status().is_success() {
return Err(VoiceCliError::Model(format!(
"Failed to download model: HTTP {}",
response.status()
)));
}
let total_size = response.content_length().unwrap_or(0);
info!("Downloading {} ({} bytes)...", model_name, total_size);
// Create temporary file
let temp_path = model_path.with_extension("tmp");
let mut file = fs::File::create(&temp_path)
.await
.map_err(|e| VoiceCliError::Model(format!("Failed to create file: {}", e)))?;
let mut downloaded = 0u64;
let mut stream = response.bytes_stream();
use futures::StreamExt;
while let Some(chunk) = stream.next().await {
let chunk =
chunk.map_err(|e| VoiceCliError::Model(format!("Download error: {}", e)))?;
file.write_all(&chunk)
.await
.map_err(|e| VoiceCliError::Model(format!("Failed to write file: {}", e)))?;
downloaded += chunk.len() as u64;
if total_size > 0 {
let progress = (downloaded as f32 / total_size as f32) * 100.0;
if downloaded % (1024 * 1024) == 0 {
// Log every MB
debug!(
"Downloaded {:.1}% ({} / {} bytes)",
progress, downloaded, total_size
);
}
}
}
file.flush()
.await
.map_err(|e| VoiceCliError::Model(format!("Failed to flush file: {}", e)))?;
// Move temporary file to final location
fs::rename(&temp_path, &model_path)
.await
.map_err(|e| VoiceCliError::Model(format!("Failed to finalize download: {}", e)))?;
info!(
"Successfully downloaded model '{}' to {:?}",
model_name, model_path
);
// Basic validation: just check file exists and has reasonable size
let metadata = fs::metadata(&model_path)
.await
.map_err(|e| VoiceCliError::Model(format!("Failed to check downloaded file: {}", e)))?;
if metadata.len() < 1024 {
// Clean up the invalid file
let _ = fs::remove_file(&model_path).await;
return Err(VoiceCliError::Model(format!(
"Downloaded model '{}' is too small ({} bytes), likely corrupted",
model_name,
metadata.len()
)));
}
info!(
"Model '{}' downloaded successfully - {} bytes",
model_name,
metadata.len()
);
Ok(())
}
/// Get the download URL for a specific model
fn get_model_download_url(&self, model_name: &str) -> Result<String, VoiceCliError> {
// Whisper.cpp models are hosted on Hugging Face under ggerganov organization
let base_url = "https://huggingface.co/ggerganov/whisper.cpp/resolve/main";
let model_filename = format!("ggml-{}.bin", model_name);
Ok(format!("{}/{}", base_url, model_filename))
}
/// Get the local path for a model file
pub fn get_model_path(&self, model_name: &str) -> Result<PathBuf, VoiceCliError> {
let filename = format!("ggml-{}.bin", model_name);
Ok(self.models_dir.join(filename))
}
/// Check if a model is downloaded locally
pub async fn is_model_downloaded(&self, model_name: &str) -> Result<bool, VoiceCliError> {
let model_path = self.get_model_path(model_name)?;
Ok(model_path.exists())
}
/// List all downloaded models
pub async fn list_downloaded_models(&self) -> Result<Vec<String>, VoiceCliError> {
if !self.models_dir.exists() {
return Ok(Vec::new());
}
let mut models = Vec::new();
let mut entries = fs::read_dir(&self.models_dir).await?;
while let Some(entry) = entries.next_entry().await? {
let path = entry.path();
if let Some(filename) = path.file_name().and_then(|n| n.to_str()) {
// Parse model name from filename (ggml-{model_name}.bin)
if filename.starts_with("ggml-") && filename.ends_with(".bin") {
let model_name = &filename[5..filename.len() - 4]; // Remove "ggml-" and ".bin"
if self
.config
.whisper
.supported_models
.contains(&model_name.to_string())
{
models.push(model_name.to_string());
}
}
}
}
models.sort();
Ok(models)
}
/// Get information about a downloaded model
pub async fn get_model_info(&self, model_name: &str) -> Result<ModelInfo, VoiceCliError> {
let model_path = self.get_model_path(model_name)?;
if !model_path.exists() {
return Err(VoiceCliError::ModelNotFound(format!(
"Model '{}' not found",
model_name
)));
}
let metadata = fs::metadata(&model_path)
.await
.map_err(|e| VoiceCliError::Model(format!("Failed to get model info: {}", e)))?;
let size = Self::format_size(metadata.len());
// TODO: Get actual memory usage if model is loaded
// This is a placeholder implementation - real memory tracking would require
// integration with the transcription service to monitor loaded models
let memory_usage = "Not tracked".to_string();
let status = if self.is_model_valid(&model_path).await? {
"Valid"
} else {
"Invalid"
}
.to_string();
Ok(ModelInfo {
size,
memory_usage,
status,
})
}
/// Validate a downloaded model
pub async fn validate_model(&self, model_name: &str) -> Result<(), VoiceCliError> {
let model_path = self.get_model_path(model_name)?;
if !model_path.exists() {
return Err(VoiceCliError::ModelNotFound(format!(
"Model '{}' not found",
model_name
)));
}
if !self.is_model_valid(&model_path).await? {
return Err(VoiceCliError::Model(format!(
"Model '{}' validation failed",
model_name
)));
}
debug!("Model '{}' validation passed", model_name);
Ok(())
}
/// Check if a model file is valid
async fn is_model_valid(&self, model_path: &Path) -> Result<bool, VoiceCliError> {
let metadata = fs::metadata(model_path)
.await
.map_err(|e| VoiceCliError::Model(format!("Failed to read model file: {}", e)))?;
// Basic validation: check if file is not empty and has reasonable size
if metadata.len() < 1024 {
warn!("Model file too small: {} bytes", metadata.len());
return Ok(false);
}
// Check if file size is reasonable for the model type
if let Some(expected_size) = self.get_expected_model_size(
&model_path
.file_stem()
.and_then(|s| s.to_str())
.map(|s| s.strip_prefix("ggml-").unwrap_or(s))
.unwrap_or("unknown"),
) {
let actual_size = metadata.len();
let size_diff_percent = if actual_size > expected_size {
((actual_size as f64 - expected_size as f64) / expected_size as f64) * 100.0
} else {
((expected_size as f64 - actual_size as f64) / expected_size as f64) * 100.0
};
// Allow 20% size difference to accommodate different versions
if size_diff_percent > 20.0 {
warn!(
"Model file size differs significantly from expected: actual={} bytes, expected={} bytes, diff={:.1}%",
actual_size, expected_size, size_diff_percent
);
// Don't fail validation, just warn - the file might still be valid
}
}
// File exists and has reasonable size - assume it's valid
// Let whisper.cpp handle format validation during actual loading
debug!("Model file appears to be valid: {} bytes", metadata.len());
Ok(true)
}
/// Remove a downloaded model
pub async fn remove_model(&self, model_name: &str) -> Result<(), VoiceCliError> {
let model_path = self.get_model_path(model_name)?;
if !model_path.exists() {
return Ok(()); // Already removed
}
fs::remove_file(&model_path)
.await
.map_err(|e| VoiceCliError::Model(format!("Failed to remove model: {}", e)))?;
info!("Removed model '{}' from {:?}", model_name, model_path);
Ok(())
}
/// Get download status for a model
pub async fn get_download_status(
&self,
model_name: &str,
) -> Result<ModelDownloadStatus, VoiceCliError> {
let status = if self.is_model_downloaded(model_name).await? {
DownloadStatus::Exists
} else {
DownloadStatus::NotStarted
};
Ok(ModelDownloadStatus {
model_name: model_name.to_string(),
status,
progress: None,
message: None,
})
}
/// List models that are currently loaded in memory
pub async fn list_loaded_models(&self) -> Result<Vec<String>, VoiceCliError> {
// TODO: This should track actually loaded models in transcription service
// For now, return empty list as this is not a core business feature
// Real implementation would require:
// 1. Integration with voice-toolkit to track loaded models
// 2. Memory usage monitoring of loaded model instances
// 3. Reference counting for multiple concurrent uses
Ok(Vec::new())
}
/// Format file size in human-readable format
fn format_size(size: u64) -> String {
const UNITS: &[&str] = &["B", "KB", "MB", "GB"];
let mut size = size as f64;
let mut unit_index = 0;
while size >= 1024.0 && unit_index < UNITS.len() - 1 {
size /= 1024.0;
unit_index += 1;
}
format!("{:.1} {}", size, UNITS[unit_index])
}
/// Get the expected model size for download progress
pub fn get_expected_model_size(&self, model_name: &str) -> Option<u64> {
// Approximate sizes for whisper models (in bytes)
match model_name {
"tiny" | "tiny.en" => Some(39 * 1024 * 1024), // ~39MB
"base" | "base.en" => Some(142 * 1024 * 1024), // ~142MB
"small" | "small.en" => Some(244 * 1024 * 1024), // ~244MB
"medium" | "medium.en" => Some(769 * 1024 * 1024), // ~769MB
"large-v1" | "large-v2" | "large-v3" => Some(1550 * 1024 * 1024), // ~1.5GB
_ => None,
}
}
/// Diagnose a corrupted model file and provide suggestions
pub async fn diagnose_model(&self, model_name: &str) -> Result<String, VoiceCliError> {
let model_path = self.get_model_path(model_name)?;
if !model_path.exists() {
return Ok(format!(
"Model '{}' file does not exist at {:?}",
model_name, model_path
));
}
let metadata = fs::metadata(&model_path)
.await
.map_err(|e| VoiceCliError::Model(format!("Failed to read model metadata: {}", e)))?;
let mut diagnosis = Vec::new();
// Check file size
let actual_size = metadata.len();
diagnosis.push(format!(
"File size: {} bytes ({})",
actual_size,
Self::format_size(actual_size)
));
if let Some(expected_size) = self.get_expected_model_size(model_name) {
let size_diff = if actual_size > expected_size {
actual_size - expected_size
} else {
expected_size - actual_size
};
let size_diff_percent = (size_diff as f64 / expected_size as f64) * 100.0;
diagnosis.push(format!(
"Expected size: {} bytes ({})",
expected_size,
Self::format_size(expected_size)
));
diagnosis.push(format!("Size difference: {:.1}%", size_diff_percent));
if size_diff_percent > 20.0 {
diagnosis.push(
"⚠️ File size differs significantly from expected - may be corrupted"
.to_string(),
);
} else {
diagnosis.push("✅ File size is within expected range".to_string());
}
}
// Basic file accessibility check
match fs::File::open(&model_path).await {
Ok(_) => {
diagnosis.push("✅ File is readable".to_string());
}
Err(e) => {
diagnosis.push(format!("❌ File is not readable: {}", e));
}
}
// Check if file is completely empty or too small
if actual_size == 0 {
diagnosis.push("❌ File is empty".to_string());
} else if actual_size < 1024 {
diagnosis.push("❌ File is too small to be a valid model".to_string());
} else {
diagnosis.push("✅ File has reasonable size".to_string());
}
Ok(diagnosis.join("\n"))
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[tokio::test]
async fn test_model_service_creation() {
let config = Config::default();
let service = ModelService::new(config);
assert!(!service.models_dir.as_os_str().is_empty());
}
#[tokio::test]
async fn test_model_path_generation() {
let temp_dir = TempDir::new().unwrap();
let mut config = Config::default();
config.whisper.models_dir = temp_dir.path().to_string_lossy().to_string();
let service = ModelService::new(config);
let path = service.get_model_path("base").unwrap();
assert!(path.to_string_lossy().contains("ggml-base.bin"));
}
#[tokio::test]
async fn test_list_downloaded_models_empty() {
let temp_dir = TempDir::new().unwrap();
let mut config = Config::default();
config.whisper.models_dir = temp_dir.path().to_string_lossy().to_string();
let service = ModelService::new(config);
let models = service.list_downloaded_models().await.unwrap();
assert!(models.is_empty());
}
#[test]
fn test_format_size() {
assert_eq!(ModelService::format_size(1024), "1.0 KB");
assert_eq!(ModelService::format_size(1024 * 1024), "1.0 MB");
assert_eq!(ModelService::format_size(1536 * 1024 * 1024), "1.5 GB");
}
#[test]
fn test_get_expected_model_size() {
let service = ModelService::new(Config::default());
assert!(service.get_expected_model_size("base").is_some());
assert!(service.get_expected_model_size("unknown").is_none());
}
}

View File

@@ -0,0 +1,155 @@
use crate::VoiceCliError;
use crate::services::ModelService;
use dashmap::DashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
// Reuse an already-loaded WhisperTranscriber to avoid reloading the model
use voice_toolkit::stt::{self, TranscriptionResult, WhisperConfig, WhisperTranscriber};
/// Shared transcription engine to unify model resolution, audio conversion and transcription
pub struct TranscriptionEngine {
model_service: Arc<ModelService>,
// Cache transcribers per model to avoid reloading model/VRAM each time
// Using DashMap for better concurrent performance
transcribers: DashMap<String, Arc<WhisperTranscriber>>,
}
impl std::fmt::Debug for TranscriptionEngine {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TranscriptionEngine")
.field("model_service", &self.model_service)
.field("transcribers_count", &self.transcribers.len())
.finish()
}
}
impl TranscriptionEngine {
/// Create a new transcription engine
pub fn new(model_service: Arc<ModelService>) -> Self {
Self {
model_service,
transcribers: DashMap::new(),
}
}
async fn get_or_create_transcriber(
&self,
model_name: &str,
) -> Result<Arc<WhisperTranscriber>, VoiceCliError> {
// Fast path: try get from cache
if let Some(existing) = self.transcribers.get(model_name) {
return Ok(existing.clone());
}
// Resolve model path
let model_path = self.model_service.get_model_path(model_name)?;
if !model_path.exists() {
return Err(VoiceCliError::ModelNotFound(model_name.to_string()));
}
// Create transcriber (assume construction might be CPU-heavy)
let created_res = tokio::task::spawn_blocking(move || {
let cfg = WhisperConfig::new(model_path);
WhisperTranscriber::new(cfg)
})
.await
.map_err(|e| VoiceCliError::Model(format!("Transcriber create join error: {}", e)))?;
let created = created_res.map_err(|e| VoiceCliError::Model(e.to_string()))?;
let transcriber = Arc::new(created);
// Insert into cache using DashMap's atomic operations
// Use entry API to handle race conditions where another thread might have inserted the same key
match self.transcribers.entry(model_name.to_string()) {
dashmap::mapref::entry::Entry::Occupied(entry) => {
// Another thread already inserted this transcriber, use the existing one
Ok(entry.get().clone())
}
dashmap::mapref::entry::Entry::Vacant(entry) => {
// We're the first to insert, use our transcriber
entry.insert(transcriber.clone());
Ok(transcriber)
}
}
}
/// Transcribe an audio file that is already Whisper-compatible (wav, correct params)
pub async fn transcribe_compatible_audio(
&self,
model_name: &str,
audio_path: &Path,
timeout_secs: u64,
) -> Result<TranscriptionResult, VoiceCliError> {
let transcriber = self.get_or_create_transcriber(model_name).await?;
let audio_path = audio_path.to_path_buf();
let timeout_duration = std::time::Duration::from_secs(timeout_secs);
let result = tokio::time::timeout(
timeout_duration,
// Use spawn_blocking for CPU-intensive Whisper transcription
// This moves the blocking operation to a separate thread pool
tokio::task::spawn_blocking(move || -> Result<TranscriptionResult, String> {
// Create a new runtime within the blocking thread
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|e| {
format!("Failed to create runtime for Whisper transcription: {}", e)
})?;
rt.block_on(async {
stt::transcribe_file_with_transcriber(&transcriber, &audio_path)
.await
.map_err(|e| e.to_string())
})
}),
)
.await
.map_err(|_| VoiceCliError::transcription_timeout(timeout_secs))?
.map_err(|e| {
if e.is_panic() {
VoiceCliError::transcription_failed("Whisper transcription panicked")
} else if e.is_cancelled() {
VoiceCliError::transcription_failed("Whisper transcription was cancelled")
} else {
VoiceCliError::transcription_failed(format!(
"Whisper transcription join error: {}",
e
))
}
})?;
Ok(result.map_err(|e| VoiceCliError::TranscriptionFailed(e.to_string()))?)
}
/// Get the default model name from configuration
pub fn default_model(&self) -> &str {
self.model_service.default_model()
}
/// Get the worker timeout from configuration
pub fn worker_timeout(&self) -> u64 {
self.model_service.worker_timeout()
}
/// Transcribe an input audio file, converting to Whisper-compatible format if necessary
pub async fn transcribe_with_conversion(
&self,
model_name: &str,
input_audio_path: &Path,
timeout_secs: u64,
) -> Result<TranscriptionResult, VoiceCliError> {
// Convert to Whisper-compatible format in blocking thread
let input_path = input_audio_path.to_path_buf();
let compatible = tokio::task::spawn_blocking(move || {
voice_toolkit::audio::ensure_whisper_compatible(&input_path, None::<PathBuf>)
})
.await
.map_err(|e| VoiceCliError::AudioConversionFailed(format!("Task join error: {}", e)))?
.map_err(|e| VoiceCliError::AudioConversionFailed(e.to_string()))?;
self.transcribe_compatible_audio(model_name, &compatible.path, timeout_secs)
.await
}
}

View File

@@ -0,0 +1,339 @@
use crate::VoiceCliError;
use crate::models::{TtsAsyncRequest, TtsSyncRequest, TtsTaskResponse};
use std::path::{Path, PathBuf};
use std::process::Command;
use tempfile::NamedTempFile;
use tracing::{debug, error, info};
use uuid::Uuid;
/// TTS服务 - 处理文本到语音转换
#[derive(Debug)]
pub struct TtsService {
python_path: PathBuf,
script_path: PathBuf,
model_path: Option<PathBuf>,
}
impl TtsService {
/// 创建新的TTS服务实例
pub fn new(
python_path: Option<PathBuf>,
model_path: Option<PathBuf>,
) -> Result<Self, VoiceCliError> {
let python_path = python_path.unwrap_or_else(|| {
// 尝试在多个位置查找虚拟环境中的 Python
let possible_venv_paths = vec![
// 当前目录下的虚拟环境
if cfg!(windows) {
PathBuf::from(".venv/Scripts/python.exe")
} else {
PathBuf::from(".venv/bin/python")
},
// crate 目录下的虚拟环境
if cfg!(windows) {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(".venv/Scripts/python.exe")
} else {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(".venv/bin/python")
},
];
// 查找第一个存在的 Python 解释器
for venv_python in possible_venv_paths {
if venv_python.exists() {
return venv_python;
}
}
// 回退到系统 Python
if let Ok(_output) = Command::new("python3").arg("--version").output() {
PathBuf::from("python3")
} else if let Ok(_output) = Command::new("python").arg("--version").output() {
PathBuf::from("python")
} else {
PathBuf::from("python3") // 默认使用python3
}
});
// 获取脚本路径(首先尝试当前目录,然后尝试 crate 目录)
let current_dir = std::env::current_dir()
.map_err(|e| VoiceCliError::Config(format!("获取当前目录失败: {}", e)))?;
let script_path = current_dir.join("tts_service.py");
let final_script_path = if script_path.exists() {
script_path
} else {
// 尝试在 crate 目录中查找
let crate_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
let crate_script_path = crate_path.join("tts_service.py");
if crate_script_path.exists() {
crate_script_path
} else {
return Err(VoiceCliError::Config(format!(
"TTS脚本不存在: 在 {:?}{:?} 中都未找到",
script_path, crate_script_path
)));
}
};
info!("Use TTS script: {:?}", final_script_path);
info!(
"Initialize TTS service - Python: {:?}, script: {:?}",
python_path, final_script_path
);
Ok(Self {
python_path,
script_path: final_script_path,
model_path,
})
}
/// 同步TTS合成
pub async fn synthesize_sync(&self, request: TtsSyncRequest) -> Result<PathBuf, VoiceCliError> {
let start_time = std::time::Instant::now();
// 验证输入
if request.text.trim().is_empty() {
return Err(VoiceCliError::InvalidInput("文本不能为空".to_string()));
}
if let Some(speed) = request.speed {
if !(0.5..=2.0).contains(&speed) {
return Err(VoiceCliError::InvalidInput(
"语速必须在0.5-2.0之间".to_string(),
));
}
}
if let Some(pitch) = request.pitch {
if !(-20..=20).contains(&pitch) {
return Err(VoiceCliError::InvalidInput(
"音调必须在-20到20之间".to_string(),
));
}
}
if let Some(volume) = request.volume {
if !(0.5..=2.0).contains(&volume) {
return Err(VoiceCliError::InvalidInput(
"音量必须在0.5-2.0之间".to_string(),
));
}
}
// 创建临时输出文件
let output_format = request.format.as_deref().unwrap_or("mp3");
let temp_file = NamedTempFile::new()
.map_err(|e| VoiceCliError::Io(format!("创建临时文件失败: {}", e)))?;
let output_path = temp_file.into_temp_path();
let output_path_str = output_path
.to_str()
.ok_or_else(|| VoiceCliError::Io("临时文件路径无效".to_string()))?;
info!(
"Start TTS synthesis - text length: {}, format: {}",
request.text.len(),
output_format
);
// 使用 uv run 来确保在正确的虚拟环境中运行
let mut cmd = Command::new("uv");
cmd.arg("run")
.arg(&self.script_path)
.arg(&request.text)
.arg("--output")
.arg(output_path_str)
.arg("--speed")
.arg(request.speed.unwrap_or(1.0).to_string())
.arg("--pitch")
.arg(request.pitch.unwrap_or(0).to_string())
.arg("--volume")
.arg(request.volume.unwrap_or(1.0).to_string())
.arg("--format")
.arg(output_format)
// 设置工作目录为脚本所在的目录
.current_dir(
self.script_path
.parent()
.unwrap_or(&std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))),
);
// 添加模型参数
if let Some(model) = &request.model {
cmd.arg("--model").arg(model);
}
if let Some(ref model_path) = self.model_path {
cmd.env("TTS_MODEL_PATH", model_path.to_string_lossy().as_ref());
}
cmd.env(
"TTS_PYTHON_PATH",
self.python_path.to_string_lossy().as_ref(),
);
debug!("Execute TTS command: {:?}", cmd);
// 执行命令
let output = cmd
.output()
.map_err(|e| VoiceCliError::TtsError(format!("执行TTS命令失败: {}", e)))?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
let stdout = String::from_utf8_lossy(&output.stdout);
error!(
"TTS synthesis failed - stderr: {}, stdout: {}",
stderr, stdout
);
return Err(VoiceCliError::TtsError(format!("TTS合成失败: {}", stderr)));
}
// 验证输出文件
if !output_path.exists() {
return Err(VoiceCliError::TtsError(
"TTS合成失败输出文件未创建".to_string(),
));
}
let file_size = output_path.metadata().map(|m| m.len()).unwrap_or(0);
if file_size == 0 {
return Err(VoiceCliError::TtsError(
"TTS合成失败输出文件为空".to_string(),
));
}
let processing_time = start_time.elapsed();
info!(
"TTS synthesis completed - file size: {} bytes, time taken: {:?}",
file_size, processing_time
);
// 将临时文件持久化到正式位置
let final_output_path = self
.persist_output_file(&output_path, output_format)
.await?;
Ok(final_output_path)
}
/// 创建异步TTS任务
pub async fn create_async_task(
&self,
request: TtsAsyncRequest,
) -> Result<TtsTaskResponse, VoiceCliError> {
// 验证输入
if request.text.trim().is_empty() {
return Err(VoiceCliError::InvalidInput("文本不能为空".to_string()));
}
// 预估处理时间(基于文本长度)
let estimated_duration = self.estimate_processing_time(&request.text);
info!(
"Create a TTS asynchronous task - text length: {}, estimated time: {}s",
request.text.len(),
estimated_duration
);
// TODO: 将任务提交到TTS任务管理器
// 这里暂时返回模拟的任务ID实际实现需要集成TtsTaskManager
let task_id = Uuid::new_v4().to_string();
Ok(TtsTaskResponse {
task_id: task_id.clone(),
message: "TTS任务已提交".to_string(),
estimated_duration: Some(estimated_duration),
})
}
/// 预估处理时间
fn estimate_processing_time(&self, text: &str) -> u32 {
// 简单的预估:基于文本长度
// 假设每秒处理10个字符
let chars_per_second = 10;
let estimated_seconds = (text.len() as f32 / chars_per_second as f32).ceil() as u32;
// 最少3秒最多300秒5分钟
estimated_seconds.max(3).min(300)
}
/// 持久化输出文件
async fn persist_output_file(
&self,
temp_path: &Path,
format: &str,
) -> Result<PathBuf, VoiceCliError> {
// 创建输出目录
let output_dir = PathBuf::from("./data/tts");
tokio::fs::create_dir_all(&output_dir)
.await
.map_err(|e| VoiceCliError::Io(format!("创建输出目录失败: {}", e)))?;
// 生成唯一文件名
let filename = format!("tts_{}.{}", Uuid::new_v4(), format);
let final_path = output_dir.join(filename);
// 复制文件
tokio::fs::copy(temp_path, &final_path)
.await
.map_err(|e| VoiceCliError::Io(format!("复制文件失败: {}", e)))?;
Ok(final_path)
}
/// 清理资源
pub async fn cleanup(&self) -> Result<(), VoiceCliError> {
// 清理临时文件等
info!("TTS service cleanup completed");
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_estimate_processing_time() {
let service = TtsService::new(None, None).unwrap();
// 测试短文本
let short_time = service.estimate_processing_time("Hello");
assert!(short_time >= 3);
// 测试长文本
let long_text = "A".repeat(1000);
let long_time = service.estimate_processing_time(&long_text);
assert!(long_time > 50);
// 测试最大限制
let very_long_text = "A".repeat(10000);
let max_time = service.estimate_processing_time(&very_long_text);
assert_eq!(max_time, 300);
}
#[tokio::test]
async fn test_create_async_task() {
let service = TtsService::new(None, None).unwrap();
let request = TtsAsyncRequest {
text: "Hello, world!".to_string(),
model: None,
speed: Some(1.0),
pitch: Some(0),
volume: Some(1.0),
format: Some("mp3".to_string()),
priority: None,
};
let response = service.create_async_task(request).await.unwrap();
assert!(!response.task_id.is_empty());
assert_eq!(response.message, "TTS任务已提交");
assert!(response.estimated_duration.unwrap() >= 3);
}
}

View File

@@ -0,0 +1,387 @@
use crate::VoiceCliError;
use crate::models::{
TtsAsyncRequest, TtsProcessingStage, TtsProgressDetails, TtsTaskError, TtsTaskStatus,
};
use apalis_sql::sqlite::SqliteStorage;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::Row;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::info;
use uuid::Uuid;
/// TTS任务
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TtsTask {
pub task_id: String,
pub text: String,
pub model: Option<String>,
pub speed: f32,
pub pitch: i32,
pub volume: f32,
pub format: String,
pub created_at: DateTime<Utc>,
pub priority: u32,
}
/// TTS任务状态
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TtsTaskState {
pub task_id: String,
pub status: TtsTaskStatus,
pub updated_at: DateTime<Utc>,
}
/// TTS任务管理器
pub struct TtsTaskManager {
storage: Arc<RwLock<SqliteStorage<TtsTask>>>,
max_concurrent_tasks: usize,
}
impl TtsTaskManager {
/// 创建新的TTS任务管理器
pub async fn new(
database_url: &str,
max_concurrent_tasks: usize,
) -> Result<Self, VoiceCliError> {
info!("Initialize TTS Task Manager - Database: {}", database_url);
// 创建SQLite存储
let pool = sqlx::sqlite::SqlitePoolOptions::new()
.max_connections(5)
.connect(database_url)
.await
.map_err(|e| VoiceCliError::Storage(format!("连接SQLite失败: {}", e)))?;
let storage = Arc::new(RwLock::new(SqliteStorage::new(pool)));
// 创建任务表
Self::create_tables_if_not_exists(&storage).await?;
Ok(Self {
storage,
max_concurrent_tasks,
})
}
/// 创建必要的表
async fn create_tables_if_not_exists(
storage: &Arc<RwLock<SqliteStorage<TtsTask>>>,
) -> Result<(), VoiceCliError> {
let guard = storage.read().await;
let pool = guard.pool();
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS tts_tasks (
id INTEGER PRIMARY KEY AUTOINCREMENT,
task_id TEXT NOT NULL UNIQUE,
text TEXT NOT NULL,
model TEXT,
speed REAL NOT NULL,
pitch INTEGER NOT NULL,
volume REAL NOT NULL,
format TEXT NOT NULL,
created_at TEXT NOT NULL,
priority INTEGER NOT NULL,
status TEXT NOT NULL,
updated_at TEXT NOT NULL,
result_path TEXT,
file_size INTEGER,
duration_seconds REAL,
error_message TEXT,
retry_count INTEGER DEFAULT 0
)
"#,
)
.execute(pool)
.await
.map_err(|e| VoiceCliError::Storage(format!("创建TTS任务表失败: {}", e)))?;
info!("TTS task list created successfully");
Ok(())
}
/// 提交TTS任务
pub async fn submit_task(&self, request: TtsAsyncRequest) -> Result<String, VoiceCliError> {
let task_id = Uuid::new_v4().to_string();
let created_at = Utc::now();
let task = TtsTask {
task_id: task_id.clone(),
text: request.text.clone(),
model: request.model.clone(),
speed: request.speed.unwrap_or(1.0),
pitch: request.pitch.unwrap_or(0),
volume: request.volume.unwrap_or(1.0),
format: request.format.unwrap_or_else(|| "mp3".to_string()),
created_at,
priority: request.priority.map_or(2, |p| match p {
crate::models::tts::TaskPriority::Low => 1,
crate::models::tts::TaskPriority::Normal => 2,
crate::models::tts::TaskPriority::High => 3,
}),
};
// 保存任务到数据库
let guard = self.storage.read().await;
let pool = guard.pool();
sqlx::query(
r#"
INSERT INTO tts_tasks (
task_id, text, model, speed, pitch, volume, format,
created_at, priority, status, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&task.task_id)
.bind(&task.text)
.bind(&task.model)
.bind(task.speed)
.bind(task.pitch)
.bind(task.volume)
.bind(&task.format)
.bind(task.created_at)
.bind(task.priority)
.bind("pending")
.bind(task.created_at)
.execute(pool)
.await
.map_err(|e| VoiceCliError::Storage(format!("保存TTS任务失败: {}", e)))?;
info!("TTS task has been submitted - ID: {}", task_id);
Ok(task_id)
}
/// 获取任务状态
pub async fn get_task_status(
&self,
task_id: &str,
) -> Result<Option<TtsTaskStatus>, VoiceCliError> {
let guard = self.storage.read().await;
let pool = guard.pool();
let row = sqlx::query(
"SELECT status, updated_at, result_path, file_size, duration_seconds, error_message, retry_count FROM tts_tasks WHERE task_id = ?"
)
.bind(task_id)
.fetch_optional(pool)
.await
.map_err(|e| VoiceCliError::Storage(format!("查询任务状态失败: {}", e)))?;
match row {
Some(row) => {
let status_str: String = row.get("status");
let updated_at: DateTime<Utc> = row.get("updated_at");
let result_path: Option<String> = row.get("result_path");
let file_size: Option<i64> = row.get("file_size");
let duration_seconds: Option<f64> = row.get("duration_seconds");
let error_message: Option<String> = row.get("error_message");
let retry_count: i32 = row.get("retry_count");
let status = match status_str.as_str() {
"pending" => TtsTaskStatus::Pending {
queued_at: updated_at,
},
"processing" => TtsTaskStatus::Processing {
stage: TtsProcessingStage::VoiceSynthesis,
started_at: updated_at,
progress_details: Some(TtsProgressDetails {
current_stage: TtsProcessingStage::VoiceSynthesis,
stage_progress: Some(0.5),
estimated_remaining: Some(chrono::Duration::seconds(30)),
text_length: 100,
processed_chars: 50,
}),
},
"completed" => {
if let (Some(path), Some(size), Some(duration)) =
(result_path, file_size, duration_seconds)
{
TtsTaskStatus::Completed {
completed_at: updated_at,
processing_time: updated_at.signed_duration_since(updated_at), // 这里应该用创建时间
audio_file_path: path,
file_size: size as u64,
duration_seconds: duration as f32,
}
} else {
return Err(VoiceCliError::Storage(
"完成的任务缺少结果信息".to_string(),
));
}
}
"failed" => {
let error = error_message.unwrap_or_else(|| "未知错误".to_string());
TtsTaskStatus::Failed {
error: TtsTaskError::SynthesisFailed {
model: "default".to_string(),
message: error,
is_recoverable: retry_count < 3,
},
failed_at: updated_at,
retry_count: retry_count as u32,
is_recoverable: retry_count < 3,
}
}
"cancelled" => TtsTaskStatus::Cancelled {
cancelled_at: updated_at,
reason: None,
},
_ => {
return Err(VoiceCliError::Storage(format!(
"未知的任务状态: {}",
status_str
)));
}
};
Ok(Some(status))
}
None => Ok(None),
}
}
/// 更新任务状态
pub async fn update_task_status(
&self,
task_id: &str,
status: TtsTaskStatus,
) -> Result<(), VoiceCliError> {
let guard = self.storage.read().await;
let pool = guard.pool();
let updated_at = Utc::now();
let (status_str, result_path, file_size, duration_seconds, error_message) = match status {
TtsTaskStatus::Pending { .. } => ("pending", None, None, None, None),
TtsTaskStatus::Processing { .. } => ("processing", None, None, None, None),
TtsTaskStatus::Completed {
audio_file_path,
file_size,
duration_seconds,
..
} => (
"completed",
Some(audio_file_path),
Some(file_size as i64),
Some(duration_seconds as f64),
None,
),
TtsTaskStatus::Failed { error, .. } => {
("failed", None, None, None, Some(error.to_string()))
}
TtsTaskStatus::Cancelled { .. } => ("cancelled", None, None, None, None),
};
sqlx::query(
"UPDATE tts_tasks SET status = ?, updated_at = ?, result_path = ?, file_size = ?, duration_seconds = ?, error_message = ? WHERE task_id = ?"
)
.bind(status_str)
.bind(updated_at)
.bind(result_path)
.bind(file_size)
.bind(duration_seconds)
.bind(error_message)
.bind(task_id)
.execute(pool)
.await
.map_err(|e| VoiceCliError::Storage(format!("更新任务状态失败: {}", e)))?;
Ok(())
}
/// 启动任务处理器
pub async fn start_worker(&self) -> Result<(), VoiceCliError> {
info!(
"Start the TTS task processor, the maximum number of concurrent tasks: {}",
self.max_concurrent_tasks
);
// TODO: 实现实际的任务处理逻辑
// 这里应该启动一个后台worker来处理TTS任务队列
Ok(())
}
/// 获取任务统计
pub async fn get_stats(&self) -> Result<TtsTaskStats, VoiceCliError> {
let guard = self.storage.read().await;
let pool = guard.pool();
let row = sqlx::query(
r#"
SELECT
COUNT(*) as total,
SUM(CASE WHEN status = 'pending' THEN 1 ELSE 0 END) as pending,
SUM(CASE WHEN status = 'processing' THEN 1 ELSE 0 END) as processing,
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) as completed,
SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END) as failed,
SUM(CASE WHEN status = 'cancelled' THEN 1 ELSE 0 END) as cancelled
FROM tts_tasks
"#,
)
.fetch_one(pool)
.await
.map_err(|e| VoiceCliError::Storage(format!("获取任务统计失败: {}", e)))?;
Ok(TtsTaskStats {
total_tasks: row.get("total"),
pending_tasks: row.get("pending"),
processing_tasks: row.get("processing"),
completed_tasks: row.get("completed"),
failed_tasks: row.get("failed"),
cancelled_tasks: row.get("cancelled"),
})
}
}
/// TTS任务统计
#[derive(Debug, Clone)]
pub struct TtsTaskStats {
pub total_tasks: i64,
pub pending_tasks: i64,
pub processing_tasks: i64,
pub completed_tasks: i64,
pub failed_tasks: i64,
pub cancelled_tasks: i64,
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::NamedTempFile;
#[tokio::test]
async fn test_tts_task_manager_creation() {
let temp_file = NamedTempFile::new().unwrap();
let db_path = temp_file.path().to_string_lossy().to_string();
let db_url = format!("sqlite://{}", db_path);
let manager = TtsTaskManager::new(&db_url, 2).await.unwrap();
assert_eq!(manager.max_concurrent_tasks, 2);
}
#[tokio::test]
async fn test_task_submission() {
let temp_file = NamedTempFile::new().unwrap();
let db_path = temp_file.path().to_string_lossy().to_string();
let db_url = format!("sqlite://{}", db_path);
let manager = TtsTaskManager::new(&db_url, 2).await.unwrap();
let request = TtsAsyncRequest {
text: "Hello, world!".to_string(),
model: None,
speed: Some(1.0),
pitch: Some(0),
volume: Some(1.0),
format: Some("mp3".to_string()),
priority: None,
};
let task_id = manager.submit_task(request).await.unwrap();
assert!(!task_id.is_empty());
}
}