ferrin_google/
realtime.rs

1//! Live API realtime model: ephemeral auth tokens, session setup and the
2//! mapping between Live API messages and the specification events.
3
4use std::sync::Mutex;
5
6use base64::Engine;
7use bytes::Bytes;
8use chrono::Duration;
9use chrono::Utc;
10use ferrin_provider_util::http::ResponseHandlers;
11use ferrin_provider_util::http::json_response_handler;
12use ferrin_provider_util::http::post_json;
13use ferrin_spec::Headers;
14use ferrin_spec::JsonObject;
15use ferrin_spec::JsonValue;
16use ferrin_spec::ModelId;
17use ferrin_spec::ProviderId;
18use ferrin_spec::RealtimeModelRef;
19use ferrin_spec::error::NoSuchModelError;
20use ferrin_spec::error::ProviderError;
21use ferrin_spec::realtime_model::ClientSecret;
22use ferrin_spec::realtime_model::ClientSecretOptions;
23use ferrin_spec::realtime_model::ConversationItem;
24use ferrin_spec::realtime_model::GetTokenOptions;
25use ferrin_spec::realtime_model::RealtimeClientEvent;
26use ferrin_spec::realtime_model::RealtimeFactory;
27use ferrin_spec::realtime_model::RealtimeModel;
28use ferrin_spec::realtime_model::RealtimeServerEvent;
29use ferrin_spec::realtime_model::RealtimeSessionConfig;
30use ferrin_spec::realtime_model::WebSocketConfig;
31use secrecy::ExposeSecret;
32use serde::Deserialize;
33use serde_json::json;
34use tokio_util::sync::CancellationToken;
35use url::Url;
36
37use crate::config::AUTH_TOKENS_PATH;
38use crate::config::CANONICAL_OPTIONS_KEY;
39use crate::config::GoogleConfig;
40use crate::config::SharedConfig;
41use crate::error::failed_response_handler;
42use crate::json_schema::convert_json_schema_to_openapi_schema;
43
44/// Provider id family.
45pub const FAMILY: &str = "realtime";
46
47/// Service path of the constrained (token-authenticated) bidi endpoint.
48pub const WEBSOCKET_SERVICE_PATH: &str =
49    "google.ai.generativelanguage.v1alpha.GenerativeService.BidiGenerateContentConstrained";
50
51/// Default window in which a token may open a session.
52pub const DEFAULT_EXPIRES_AFTER_SECONDS: u64 = 60;
53
54/// Default input audio sample rate advertised in audio blobs.
55pub const DEFAULT_INPUT_AUDIO_RATE: u32 = 16_000;
56
57#[derive(Debug, Deserialize)]
58struct AuthTokenResponse {
59    name: String,
60    #[serde(default, rename = "expireTime")]
61    expire_time: Option<String>,
62}
63
64/// Builds the `setup` message (`bidiGenerateContentSetup`) for a session.
65///
66/// # Errors
67///
68/// Returns [`ProviderError::UnsupportedFunctionality`] when a tool schema
69/// cannot be converted.
70pub fn build_session_config(
71    config: Option<&RealtimeSessionConfig>,
72    model_id: &str,
73) -> Result<JsonValue, ProviderError> {
74    let mut setup = JsonObject::new();
75    setup.insert(
76        "model".to_owned(),
77        JsonValue::from(GoogleConfig::model_path(model_id)),
78    );
79    let mut generation = JsonObject::new();
80    let modalities: Vec<JsonValue> =
81        match config.and_then(|config| config.output_modalities.as_ref()) {
82            Some(modalities) => modalities
83                .iter()
84                .map(|modality| {
85                    let text = serde_json::to_value(modality)
86                        .ok()
87                        .and_then(|value| value.as_str().map(str::to_ascii_uppercase))
88                        .unwrap_or_else(|| "AUDIO".to_owned());
89                    JsonValue::from(text)
90                })
91                .collect(),
92            None => vec![JsonValue::from("AUDIO")],
93        };
94    generation.insert(
95        "responseModalities".to_owned(),
96        JsonValue::Array(modalities),
97    );
98    if let Some(voice) = config.and_then(|config| config.voice.as_deref()) {
99        generation.insert(
100            "speechConfig".to_owned(),
101            json!({"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": voice}}}),
102        );
103    }
104    setup.insert("generationConfig".to_owned(), JsonValue::Object(generation));
105    let Some(config) = config else {
106        return Ok(JsonValue::Object(setup));
107    };
108    if let Some(instructions) = &config.instructions {
109        setup.insert(
110            "systemInstruction".to_owned(),
111            json!({"parts": [{"text": instructions}]}),
112        );
113    }
114    if !config.tools.is_empty() {
115        let mut declarations = Vec::new();
116        for tool in &config.tools {
117            let mut declaration = JsonObject::new();
118            declaration.insert("name".to_owned(), JsonValue::from(tool.name.as_str()));
119            if let Some(description) = &tool.description {
120                declaration.insert(
121                    "description".to_owned(),
122                    JsonValue::from(description.as_str()),
123                );
124            }
125            if let Some(parameters) = convert_json_schema_to_openapi_schema(&tool.parameters)? {
126                declaration.insert("parameters".to_owned(), parameters);
127            }
128            declarations.push(JsonValue::Object(declaration));
129        }
130        setup.insert(
131            "tools".to_owned(),
132            json!([{"functionDeclarations": declarations}]),
133        );
134    }
135    if config.input_audio_transcription.is_some() {
136        setup.insert("inputAudioTranscription".to_owned(), json!({}));
137    }
138    if config.output_audio_transcription.is_some() {
139        setup.insert("outputAudioTranscription".to_owned(), json!({}));
140    }
141    if let Some(provider_options) = &config.provider_options {
142        let mut google: Option<JsonObject> = None;
143        for (key, value) in provider_options {
144            if key == CANONICAL_OPTIONS_KEY {
145                if let JsonValue::Object(object) = value {
146                    google = Some(object.clone());
147                }
148            } else {
149                setup.insert(key.clone(), value.clone());
150            }
151        }
152        if let Some(translation) = google
153            .as_ref()
154            .and_then(|google| google.get("translationConfig"))
155        {
156            let target = match setup.get_mut("generationConfig") {
157                Some(JsonValue::Object(generation)) => generation,
158                _ => {
159                    setup.insert("generationConfig".to_owned(), json!({}));
160                    match setup.get_mut("generationConfig") {
161                        Some(JsonValue::Object(generation)) => generation,
162                        _ => return Ok(JsonValue::Object(setup)),
163                    }
164                }
165            };
166            target.insert("translationConfig".to_owned(), translation.clone());
167        }
168    }
169    Ok(JsonValue::Object(setup))
170}
171
172/// Turn tracking of the stateful event mapper.
173#[derive(Debug, Default)]
174struct MapperState {
175    turn_counter: u64,
176    has_audio: bool,
177    has_text: bool,
178    has_transcript: bool,
179    turn_closed: bool,
180    input_audio_rate: Option<u32>,
181}
182
183impl MapperState {
184    fn response_id(&self) -> String {
185        format!("google-resp-{}", self.turn_counter)
186    }
187
188    fn item_id(&self) -> String {
189        format!("google-item-{}", self.turn_counter)
190    }
191
192    fn input_id(&self) -> String {
193        format!("google-input-{}", self.turn_counter)
194    }
195
196    fn begin_turn_if_closed(&mut self) {
197        if self.turn_closed {
198            self.turn_counter += 1;
199            self.has_audio = false;
200            self.has_text = false;
201            self.has_transcript = false;
202            self.turn_closed = false;
203        }
204    }
205}
206
207fn custom(raw_type: &str, raw: &JsonValue) -> RealtimeServerEvent {
208    RealtimeServerEvent::Custom {
209        raw_type: raw_type.to_owned(),
210        raw: raw.clone(),
211    }
212}
213
214fn decode_audio(data: &str) -> Bytes {
215    base64::engine::general_purpose::STANDARD
216        .decode(data)
217        .map(Bytes::from)
218        .unwrap_or_default()
219}
220
221fn parse_server_content(
222    state: &mut MapperState,
223    content: &JsonValue,
224    raw: &JsonValue,
225) -> Vec<RealtimeServerEvent> {
226    let mut events = Vec::new();
227    if content.get("interrupted").and_then(JsonValue::as_bool) == Some(true) {
228        events.push(RealtimeServerEvent::SpeechStarted {
229            item_id: None,
230            raw: raw.clone(),
231        });
232    }
233    if let Some(parts) = content
234        .get("modelTurn")
235        .and_then(|turn| turn.get("parts"))
236        .and_then(JsonValue::as_array)
237    {
238        state.begin_turn_if_closed();
239        for part in parts {
240            if let Some(data) = part
241                .get("inlineData")
242                .and_then(|inline| inline.get("data"))
243                .and_then(JsonValue::as_str)
244                .filter(|data| !data.is_empty())
245            {
246                state.has_audio = true;
247                events.push(RealtimeServerEvent::AudioDelta {
248                    response_id: state.response_id(),
249                    item_id: state.item_id(),
250                    delta: decode_audio(data),
251                    raw: raw.clone(),
252                });
253            }
254            if let Some(text) = part
255                .get("text")
256                .and_then(JsonValue::as_str)
257                .filter(|text| !text.is_empty())
258            {
259                state.has_text = true;
260                events.push(RealtimeServerEvent::TextDelta {
261                    response_id: state.response_id(),
262                    item_id: state.item_id(),
263                    delta: text.to_owned(),
264                    raw: raw.clone(),
265                });
266            }
267        }
268    }
269    if let Some(text) = content
270        .get("outputTranscription")
271        .and_then(|transcription| transcription.get("text"))
272        .and_then(JsonValue::as_str)
273        .filter(|text| !text.is_empty())
274    {
275        state.has_transcript = true;
276        events.push(RealtimeServerEvent::AudioTranscriptDelta {
277            response_id: state.response_id(),
278            item_id: state.item_id(),
279            delta: text.to_owned(),
280            raw: raw.clone(),
281        });
282    }
283    if let Some(text) = content
284        .get("inputTranscription")
285        .and_then(|transcription| transcription.get("text"))
286        .and_then(JsonValue::as_str)
287        .filter(|text| !text.is_empty())
288    {
289        events.push(RealtimeServerEvent::InputTranscriptionCompleted {
290            item_id: state.input_id(),
291            transcript: text.to_owned(),
292            raw: raw.clone(),
293        });
294    }
295    if content
296        .get("generationComplete")
297        .and_then(JsonValue::as_bool)
298        == Some(true)
299    {
300        events.push(custom("generationComplete", raw));
301    }
302    if content.get("turnComplete").and_then(JsonValue::as_bool) == Some(true) {
303        if state.has_audio {
304            events.push(RealtimeServerEvent::AudioDone {
305                response_id: state.response_id(),
306                item_id: state.item_id(),
307                raw: raw.clone(),
308            });
309        }
310        if state.has_text {
311            events.push(RealtimeServerEvent::TextDone {
312                response_id: state.response_id(),
313                item_id: state.item_id(),
314                text: None,
315                raw: raw.clone(),
316            });
317        }
318        if state.has_transcript {
319            events.push(RealtimeServerEvent::AudioTranscriptDone {
320                response_id: state.response_id(),
321                item_id: state.item_id(),
322                transcript: None,
323                raw: raw.clone(),
324            });
325        }
326        events.push(RealtimeServerEvent::ResponseDone {
327            response_id: state.response_id(),
328            status: "completed".to_owned(),
329            raw: raw.clone(),
330        });
331        state.turn_closed = true;
332    }
333    if events.is_empty() {
334        events.push(custom("serverContent", raw));
335    }
336    events
337}
338
339fn parse_event(state: &mut MapperState, raw: &JsonValue) -> Vec<RealtimeServerEvent> {
340    let Some(object) = raw.as_object() else {
341        return vec![custom("unknown", raw)];
342    };
343    if object.contains_key("setupComplete") {
344        return vec![RealtimeServerEvent::SessionCreated {
345            session_id: None,
346            raw: raw.clone(),
347        }];
348    }
349    if let Some(tool_call) = object.get("toolCall") {
350        state.begin_turn_if_closed();
351        let mut events = Vec::new();
352        for call in tool_call
353            .get("functionCalls")
354            .and_then(JsonValue::as_array)
355            .into_iter()
356            .flatten()
357        {
358            let args = call
359                .get("args")
360                .cloned()
361                .unwrap_or_else(|| JsonValue::Object(JsonObject::new()))
362                .to_string();
363            let call_id = call
364                .get("id")
365                .and_then(JsonValue::as_str)
366                .unwrap_or_default()
367                .to_owned();
368            let name = call
369                .get("name")
370                .and_then(JsonValue::as_str)
371                .unwrap_or_default()
372                .to_owned();
373            events.push(RealtimeServerEvent::FunctionCallArgumentsDelta {
374                response_id: state.response_id(),
375                item_id: state.item_id(),
376                call_id: call_id.clone(),
377                delta: args.clone(),
378                raw: raw.clone(),
379            });
380            events.push(RealtimeServerEvent::FunctionCallArgumentsDone {
381                response_id: state.response_id(),
382                item_id: state.item_id(),
383                call_id,
384                name,
385                arguments: args,
386                raw: raw.clone(),
387            });
388        }
389        return events;
390    }
391    for key in ["toolCallCancellation", "goAway", "sessionResumptionUpdate"] {
392        if object.contains_key(key) {
393            return vec![custom(key, raw)];
394        }
395    }
396    if let Some(content) = object.get("serverContent") {
397        return parse_server_content(state, content, raw);
398    }
399    if let Some(text) = object
400        .get("inputTranscription")
401        .and_then(|transcription| transcription.get("text"))
402        .and_then(JsonValue::as_str)
403    {
404        return vec![RealtimeServerEvent::InputTranscriptionCompleted {
405            item_id: state.input_id(),
406            transcript: text.to_owned(),
407            raw: raw.clone(),
408        }];
409    }
410    let raw_type = object.keys().next().map_or("unknown", String::as_str);
411    vec![custom(raw_type, raw)]
412}
413
414fn serialize_event(
415    state: &mut MapperState,
416    event: RealtimeClientEvent,
417    model_id: &str,
418) -> Result<JsonValue, ProviderError> {
419    Ok(match event {
420        RealtimeClientEvent::SessionUpdate { config } => {
421            if let Some(rate) = config
422                .input_audio_format
423                .as_ref()
424                .and_then(|format| format.rate)
425            {
426                state.input_audio_rate = Some(rate);
427            }
428            json!({"setup": build_session_config(Some(&config), model_id)?})
429        }
430        RealtimeClientEvent::InputAudioAppend { audio } => {
431            let rate = state.input_audio_rate.unwrap_or(DEFAULT_INPUT_AUDIO_RATE);
432            json!({"realtimeInput": {"audio": {
433                "data": base64::engine::general_purpose::STANDARD.encode(&audio),
434                "mimeType": format!("audio/pcm;rate={rate}"),
435            }}})
436        }
437        RealtimeClientEvent::InputAudioCommit => {
438            json!({"realtimeInput": {"audioStreamEnd": true}})
439        }
440        RealtimeClientEvent::ConversationItemCreate { item } => match item {
441            ConversationItem::TextMessage { text, .. } => {
442                json!({"realtimeInput": {"text": text}})
443            }
444            ConversationItem::FunctionCallOutput {
445                call_id,
446                name,
447                output,
448            } => {
449                let response = serde_json::from_str::<JsonValue>(&output)
450                    .ok()
451                    .filter(JsonValue::is_object)
452                    .unwrap_or_else(|| JsonValue::Object(JsonObject::new()));
453                let mut function_response = JsonObject::new();
454                function_response.insert("id".to_owned(), JsonValue::from(call_id));
455                if let Some(name) = name {
456                    function_response.insert("name".to_owned(), JsonValue::from(name));
457                }
458                function_response.insert("response".to_owned(), response);
459                json!({"toolResponse": {"functionResponses": [function_response]}})
460            }
461            ConversationItem::AudioMessage { .. } => {
462                return Err(ProviderError::unsupported(
463                    "realtime conversation item: audio message",
464                ));
465            }
466            #[allow(unreachable_patterns, reason = "ConversationItem is non-exhaustive")]
467            _ => return Err(ProviderError::unsupported("realtime conversation item")),
468        },
469        RealtimeClientEvent::InputAudioClear => {
470            return Err(ProviderError::unsupported(
471                "realtime client event: input audio clear",
472            ));
473        }
474        RealtimeClientEvent::ResponseCreate { .. } => {
475            return Err(ProviderError::unsupported(
476                "realtime client event: response create",
477            ));
478        }
479        RealtimeClientEvent::ResponseCancel => {
480            return Err(ProviderError::unsupported(
481                "realtime client event: response cancel",
482            ));
483        }
484        RealtimeClientEvent::ConversationItemTruncate { .. } => {
485            return Err(ProviderError::unsupported(
486                "realtime client event: conversation item truncate",
487            ));
488        }
489        #[allow(unreachable_patterns, reason = "RealtimeClientEvent is non-exhaustive")]
490        _ => return Err(ProviderError::unsupported("realtime client event")),
491    })
492}
493
494/// Live API realtime model.
495#[derive(Debug)]
496pub struct GoogleRealtimeModel {
497    config: SharedConfig,
498    provider: ProviderId,
499    model_id: ModelId,
500    state: Mutex<MapperState>,
501}
502
503impl GoogleRealtimeModel {
504    /// Creates the model.
505    #[must_use]
506    pub fn new(config: SharedConfig, model_id: impl Into<ModelId>) -> Self {
507        Self {
508            provider: config.provider_id(FAMILY),
509            config,
510            model_id: model_id.into(),
511            state: Mutex::new(MapperState::default()),
512        }
513    }
514
515    /// WebSocket URL of the constrained bidi endpoint (without the token).
516    #[must_use]
517    pub fn session_url(&self) -> Url {
518        self.config.websocket_url(WEBSOCKET_SERVICE_PATH)
519    }
520
521    fn lock(&self) -> std::sync::MutexGuard<'_, MapperState> {
522        self.state
523            .lock()
524            .unwrap_or_else(std::sync::PoisonError::into_inner)
525    }
526}
527
528impl RealtimeModel for GoogleRealtimeModel {
529    fn provider(&self) -> &ProviderId {
530        &self.provider
531    }
532
533    fn model_id(&self) -> &ModelId {
534        &self.model_id
535    }
536
537    #[tracing::instrument(skip_all, fields(model = %self.model_id))]
538    async fn do_create_client_secret(
539        &self,
540        options: ClientSecretOptions,
541    ) -> Result<ClientSecret, ProviderError> {
542        let api_key = self.config.api_key()?;
543        let now = Utc::now();
544        let window = i64::try_from(
545            options
546                .expires_after_seconds
547                .unwrap_or(DEFAULT_EXPIRES_AFTER_SECONDS),
548        )
549        .unwrap_or(i64::MAX / 4);
550        let new_session_expire_time = now + Duration::seconds(window);
551        let expire_time = new_session_expire_time + Duration::minutes(30);
552        let body = json!({
553            "uses": 0,
554            "expireTime": expire_time.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
555            "newSessionExpireTime": new_session_expire_time.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
556            "bidiGenerateContentSetup": build_session_config(
557                options.session_config.as_ref(),
558                self.model_id.as_str(),
559            )?,
560        });
561        let mut url = self.config.origin_url(AUTH_TOKENS_PATH);
562        url.query_pairs_mut()
563            .append_pair("key", api_key.expose_secret());
564        let mut headers = self.config.headers.clone();
565        headers = headers.with_user_agent_suffix([crate::config::USER_AGENT]);
566        let handlers = ResponseHandlers::new(
567            json_response_handler::<AuthTokenResponse>(),
568            failed_response_handler(),
569        );
570        let response = post_json(
571            self.config.transport.as_ref(),
572            url,
573            headers,
574            &body,
575            &handlers,
576            CancellationToken::new(),
577        )
578        .await?;
579        let expires_at = response
580            .value
581            .expire_time
582            .as_deref()
583            .and_then(|time| chrono::DateTime::parse_from_rfc3339(time).ok())
584            .and_then(|time| u64::try_from(time.timestamp()).ok());
585        Ok(ClientSecret {
586            token: response.value.name,
587            url: self.session_url(),
588            expires_at,
589        })
590    }
591
592    fn websocket_config(&self, token: &str, url: &Url) -> WebSocketConfig {
593        let mut url = url.clone();
594        url.query_pairs_mut().append_pair("access_token", token);
595        WebSocketConfig {
596            url,
597            protocols: Vec::new(),
598        }
599    }
600
601    fn parse_server_event(
602        &self,
603        raw: JsonValue,
604    ) -> Result<Vec<RealtimeServerEvent>, ProviderError> {
605        let mut state = self.lock();
606        Ok(parse_event(&mut state, &raw))
607    }
608
609    async fn serialize_client_event(
610        &self,
611        event: RealtimeClientEvent,
612    ) -> Result<JsonValue, ProviderError> {
613        let mut state = self.lock();
614        serialize_event(&mut state, event, self.model_id.as_str())
615    }
616
617    fn build_session_config(
618        &self,
619        config: &RealtimeSessionConfig,
620    ) -> Result<JsonValue, ProviderError> {
621        build_session_config(Some(config), self.model_id.as_str())
622    }
623}
624
625/// Factory for Live API models and session tokens.
626#[derive(Debug, Clone)]
627pub struct GoogleRealtimeFactory {
628    config: SharedConfig,
629    provider: ProviderId,
630}
631
632impl GoogleRealtimeFactory {
633    /// Creates the factory.
634    #[must_use]
635    pub fn new(config: SharedConfig) -> Self {
636        Self {
637            provider: config.provider_id(FAMILY),
638            config,
639        }
640    }
641
642    /// Creates a realtime model.
643    #[must_use]
644    pub fn realtime_model(&self, model_id: &str) -> GoogleRealtimeModel {
645        GoogleRealtimeModel::new(self.config.clone(), model_id)
646    }
647}
648
649impl RealtimeFactory for GoogleRealtimeFactory {
650    fn provider(&self) -> &ProviderId {
651        &self.provider
652    }
653
654    fn model(&self, model_id: &str) -> Result<RealtimeModelRef, NoSuchModelError> {
655        Ok(self.realtime_model(model_id).into())
656    }
657
658    async fn get_token(&self, options: GetTokenOptions) -> Result<ClientSecret, ProviderError> {
659        self.realtime_model(options.model.as_str())
660            .do_create_client_secret(ClientSecretOptions {
661                expires_after_seconds: options.expires_after_seconds,
662                session_config: options.session_config,
663            })
664            .await
665    }
666}
667
668/// Headers of the auth token request: configuration headers only (the key
669/// travels in the query string).
670#[must_use]
671pub fn token_request_headers(config: &GoogleConfig) -> Headers {
672    config
673        .headers
674        .clone()
675        .with_user_agent_suffix([crate::config::USER_AGENT])
676}
ferrin_google/realtime.rs

ferrin_google/
realtime.rs
