rustvani 0.2.5

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

use std::sync::Arc;

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

use super::segment::PendingAudioSegment;
use super::storage::{AudioStorage, RecordedSegment};

// ---------------------------------------------------------------------------
// 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 — collects all segments then writes one mixed recording
// ---------------------------------------------------------------------------

async fn drain_task(
    session_id: Uuid,
    mut rx: mpsc::Receiver<PendingAudioSegment>,
    storage: Arc<dyn AudioStorage>,
) {
    let mut segments: Vec<RecordedSegment> = Vec::new();

    while let Some(seg) = rx.recv().await {
        if seg.pcm.is_empty() {
            continue;
        }
        segments.push(RecordedSegment {
            speaker:      seg.speaker,
            pcm:          seg.pcm,
            sample_rate:  seg.sample_rate,
            num_channels: seg.num_channels,
            started_at:   seg.started_at,
        });
    }

    if let Err(e) = storage.finalize_recording(session_id, &segments).await {
        log::error!("AudioCapture: finalize_recording 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() }
}