rustvani 0.2.3

Voice AI framework for Rust — real-time speech pipelines with STT, LLM, TTS, and Dhara conversation flows
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

use std::sync::Arc;

use chrono::Utc;
use tokio::sync::mpsc;
use uuid::Uuid;

use super::encoder::{encode_pcm_to_wav, pcm_duration_ms};
use super::segment::{AudioSegmentMeta, PendingAudioSegment};
use super::storage::AudioStorage;

// ---------------------------------------------------------------------------
// Public trait
// ---------------------------------------------------------------------------

pub trait AudioCaptureCollector: Send + Sync {
    /// Enqueue a completed audio segment for encoding and storage.
    /// Sync and non-blocking — drops if the channel is full.
    fn record_segment(&self, segment: PendingAudioSegment);
    fn session_id(&self) -> Uuid;
}

// ---------------------------------------------------------------------------
// SessionAudioCapture — real implementation
// ---------------------------------------------------------------------------

pub struct SessionAudioCapture {
    session_id: Uuid,
    tx: mpsc::Sender<PendingAudioSegment>,
}

impl SessionAudioCapture {
    /// `channel_capacity` — how many pending segments can be buffered.
    /// 64 is a safe default (each segment is ~160KB for a 5s mono 16kHz turn).
    pub fn new(
        session_id: Uuid,
        storage: Arc<dyn AudioStorage>,
        channel_capacity: usize,
    ) -> (Arc<Self>, tokio::task::JoinHandle<()>) {
        let (tx, rx) = mpsc::channel(channel_capacity);
        let collector = Arc::new(Self { session_id, tx });
        let handle = tokio::spawn(drain_task(session_id, rx, storage));
        (collector, handle)
    }
}

impl AudioCaptureCollector for SessionAudioCapture {
    fn record_segment(&self, segment: PendingAudioSegment) {
        if let Err(e) = self.tx.try_send(segment) {
            match e {
                mpsc::error::TrySendError::Full(_) => {
                    log::warn!("AudioCaptureCollector: channel full, dropping segment");
                }
                mpsc::error::TrySendError::Closed(_) => {}
            }
        }
    }

    fn session_id(&self) -> Uuid {
        self.session_id
    }
}

// ---------------------------------------------------------------------------
// Drain task — encodes PCM and writes to storage off the hot path
// ---------------------------------------------------------------------------

async fn drain_task(
    session_id: Uuid,
    mut rx: mpsc::Receiver<PendingAudioSegment>,
    storage: Arc<dyn AudioStorage>,
) {
    while let Some(seg) = rx.recv().await {
        process_segment(session_id, seg, &*storage).await;
    }
}

async fn process_segment(
    session_id: Uuid,
    seg: PendingAudioSegment,
    storage: &dyn AudioStorage,
) {
    if seg.pcm.is_empty() {
        return;
    }

    let ended_at = Utc::now();
    let duration_ms = pcm_duration_ms(seg.pcm.len(), seg.sample_rate, seg.num_channels);
    let byte_size_before_encode = seg.pcm.len() as u64;

    // Encode in a blocking thread so we don't block the tokio runtime.
    let pcm = seg.pcm;
    let sr = seg.sample_rate;
    let ch = seg.num_channels;
    let wav_bytes = match tokio::task::spawn_blocking(move || encode_pcm_to_wav(&pcm, sr, ch)).await {
        Ok(Ok(bytes)) => bytes,
        Ok(Err(e)) => {
            log::error!("AudioCapture: encode failed: {e}");
            return;
        }
        Err(e) => {
            log::error!("AudioCapture: encode task panicked: {e}");
            return;
        }
    };

    let wav_size = wav_bytes.len() as u64;

    let url = match storage.store_segment(session_id, seg.segment_id, seg.speaker, &wav_bytes).await {
        Ok(u) => u,
        Err(e) => {
            log::error!("AudioCapture: store_segment failed: {e}");
            return;
        }
    };

    let meta = AudioSegmentMeta {
        segment_id:   seg.segment_id,
        session_id,
        turn_id:      seg.turn_id,
        speaker:      seg.speaker.to_string(),
        audio_url:    url,
        format:       "wav".to_string(),
        sample_rate:  sr,
        num_channels: ch,
        duration_ms,
        byte_size:    wav_size,
        interrupted:  seg.interrupted,
        started_at:   seg.started_at,
        ended_at,
    };

    let _ = byte_size_before_encode; // acknowledged
    if let Err(e) = storage.save_metadata(session_id, &meta).await {
        log::error!("AudioCapture: save_metadata failed: {e}");
    }
}

// ---------------------------------------------------------------------------
// NoopAudioCaptureCollector — zero-cost default
// ---------------------------------------------------------------------------

pub struct NoopAudioCaptureCollector;

impl AudioCaptureCollector for NoopAudioCaptureCollector {
    fn record_segment(&self, _segment: PendingAudioSegment) {}
    fn session_id(&self) -> Uuid { Uuid::nil() }
}