gradbot 0.6.4

Core library for gradbot voice AI multiplexing
Documentation 
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
148
149
150
151
152
153
154

use anyhow::{Context, Result};

pub struct TtsClient(gradium::Client);
pub struct TtsStreamSender(gradium::tts::TtsStreamSender);
pub struct TtsStreamReceiver(gradium::tts::TtsStreamReceiver);

/// TTS output message with turn index.
/// The turn_idx is used to skip stale audio after an interruption:
/// when the user interrupts, the turn_idx is incremented, and any
/// TtsOut messages with an older turn_idx are discarded.
#[derive(Debug)]
pub enum TtsOut {
    Audio {
        pcm: Vec<f32>,
        start_s: f64,
        stop_s: f64,
        turn_idx: u64,
        /// Set by multiplex when this is the last audio before an interruption.
        /// The TTS receiver always sets this to false.
        interrupted: bool,
    },
    Text {
        text: String,
        start_s: f64,
        stop_s: f64,
        turn_idx: u64,
    },
    /// Signals that all output for this turn has been sent.
    /// The stop_s is the timestamp of the last audio sent, used to
    /// set the listening start time for silence detection.
    TurnComplete { turn_idx: u64, stop_s: f64 },
}

impl TtsClient {
    pub fn new(gradium_api_key: Option<&str>, base_url: &str) -> Result<Self> {
        let api_key = match gradium_api_key {
            Some(key) => key.to_string(),
            None => std::env::var("GRADIUM_API_KEY")
                .map_err(|_| anyhow::anyhow!("GRADIUM_API_KEY environment variable not set"))?,
        };
        let client = gradium::Client::new(&api_key)
            .with_base_url(base_url)
            .context("TTS: failed to initialize Gradium client")?;
        Ok(Self(client))
    }

    pub async fn tts_stream(
        &self,
        voice_id: Option<String>,
        padding_bonus: f64,
        rewrite_rules: Option<String>,
        extra_config: Option<&str>,
    ) -> Result<(TtsStreamSender, TtsStreamReceiver)> {
        let json_config = {
            let mut config = serde_json::Map::new();
            if padding_bonus != 0.0 {
                config.insert("padding_bonus".into(), serde_json::json!(padding_bonus));
            }
            if let Some(rules) = rewrite_rules {
                config.insert("rewrite_rules".into(), serde_json::json!(rules));
            }
            if let Some(extra) = extra_config
                && let Ok(serde_json::Value::Object(map)) = serde_json::from_str(extra)
            {
                config.extend(map);
            }
            if config.is_empty() {
                None
            } else {
                Some(serde_json::Value::Object(config).to_string())
            }
        };

        let setup = gradium::protocol::tts::Setup {
            model_name: "default".to_string(),
            voice_id,
            voice: None,
            output_format: gradium::protocol::AudioFormat::Pcm,
            client_req_id: None,
            close_ws_on_eos: None,
            pronunciation_id: None,
            json_config,
        };
        let stream = self
            .0
            .tts_stream(setup)
            .await
            .context("TTS: failed to connect")?;
        let (tx, rx) = stream.split();
        Ok((TtsStreamSender(tx), TtsStreamReceiver(rx)))
    }
}

impl TtsStreamSender {
    pub async fn send_text(&mut self, text: &str) -> Result<()> {
        self.0
            .send_text(text)
            .await
            .context("TTS: failed to send text")?;
        Ok(())
    }

    pub async fn send_end_of_stream(&mut self) -> Result<()> {
        self.0
            .send_eos()
            .await
            .context("TTS: failed to send end-of-stream")?;
        Ok(())
    }
}

impl TtsStreamReceiver {
    pub async fn next_message(&mut self, turn_idx: u64) -> Result<Option<TtsOut>> {
        use gradium::protocol::tts::Response;
        while let Some(msg) = self.0.next_message().await? {
            match msg {
                Response::Audio(a) => {
                    let pcm = a.raw_audio()?;
                    let pcm = pcm
                        .chunks_exact(2)
                        .map(|b| i16::from_le_bytes([b[0], b[1]]) as f32 / 32768.0)
                        .collect::<Vec<f32>>();
                    return Ok(Some(TtsOut::Audio {
                        pcm,
                        start_s: a.start_s,
                        stop_s: a.stop_s,
                        turn_idx,
                        interrupted: false,
                    }));
                }
                Response::Text(text) => {
                    return Ok(Some(TtsOut::Text {
                        text: text.text,
                        start_s: text.start_s,
                        stop_s: text.stop_s,
                        turn_idx,
                    }));
                }
                Response::Error { code, message, .. } => {
                    tracing::error!(?code, ?message, "TTS API returned error");
                    anyhow::bail!("TTS Error {code:?}: {message}")
                }
                Response::Ready(_) => {
                    tracing::debug!("TTS Ready received");
                }
                Response::EndOfStream { client_req_id: _ } => {
                    tracing::debug!("TTS EndOfStream received");
                }
            }
        }
        tracing::debug!("TTS stream ended (no more messages)");
        Ok(None)
    }
}