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 = ferrin_schema::json::parse(&output).unwrap_or_else(|_| json!({}));
450                let mut function_response = JsonObject::new();
451                function_response.insert("id".to_owned(), JsonValue::from(call_id));
452                if let Some(name) = name {
453                    function_response.insert("name".to_owned(), JsonValue::from(name));
454                }
455                function_response.insert("response".to_owned(), response);
456                json!({"toolResponse": {"functionResponses": [function_response]}})
457            }
458            ConversationItem::AudioMessage { .. } => JsonValue::Null,
459            #[allow(unreachable_patterns, reason = "ConversationItem is non-exhaustive")]
460            _ => return Err(ProviderError::unsupported("realtime conversation item")),
461        },
462        RealtimeClientEvent::InputAudioClear
463        | RealtimeClientEvent::ResponseCreate { .. }
464        | RealtimeClientEvent::ResponseCancel
465        | RealtimeClientEvent::ConversationItemTruncate { .. } => JsonValue::Null,
466        #[allow(unreachable_patterns, reason = "RealtimeClientEvent is non-exhaustive")]
467        _ => return Err(ProviderError::unsupported("realtime client event")),
468    })
469}
470
471/// Live API realtime model.
472#[derive(Debug)]
473pub struct GoogleRealtimeModel {
474    config: SharedConfig,
475    provider: ProviderId,
476    model_id: ModelId,
477    state: Mutex<MapperState>,
478}
479
480impl GoogleRealtimeModel {
481    /// Creates the model.
482    #[must_use]
483    pub fn new(config: SharedConfig, model_id: impl Into<ModelId>) -> Self {
484        Self {
485            provider: config.provider_id(FAMILY),
486            config,
487            model_id: model_id.into(),
488            state: Mutex::new(MapperState::default()),
489        }
490    }
491
492    /// WebSocket URL of the constrained bidi endpoint (without the token).
493    #[must_use]
494    pub fn session_url(&self) -> Url {
495        self.config.websocket_url(WEBSOCKET_SERVICE_PATH)
496    }
497
498    fn lock(&self) -> std::sync::MutexGuard<'_, MapperState> {
499        self.state
500            .lock()
501            .unwrap_or_else(std::sync::PoisonError::into_inner)
502    }
503}
504
505impl RealtimeModel for GoogleRealtimeModel {
506    fn provider(&self) -> &ProviderId {
507        &self.provider
508    }
509
510    fn model_id(&self) -> &ModelId {
511        &self.model_id
512    }
513
514    #[tracing::instrument(skip_all, fields(model = %self.model_id))]
515    async fn do_create_client_secret(
516        &self,
517        options: ClientSecretOptions,
518    ) -> Result<ClientSecret, ProviderError> {
519        let api_key = self.config.api_key()?;
520        let now = Utc::now();
521        let window = i64::try_from(
522            options
523                .expires_after_seconds
524                .unwrap_or(DEFAULT_EXPIRES_AFTER_SECONDS),
525        )
526        .unwrap_or(i64::MAX / 4);
527        let new_session_expire_time = now + Duration::seconds(window);
528        let expire_time = new_session_expire_time + Duration::minutes(30);
529        let body = json!({
530            "uses": 0,
531            "expireTime": expire_time.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
532            "newSessionExpireTime": new_session_expire_time.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
533            "bidiGenerateContentSetup": build_session_config(
534                options.session_config.as_ref(),
535                self.model_id.as_str(),
536            )?,
537        });
538        let mut url = self.config.origin_url(AUTH_TOKENS_PATH);
539        url.query_pairs_mut()
540            .append_pair("key", api_key.expose_secret());
541        let mut headers = self.config.headers.clone();
542        headers = headers.with_user_agent_suffix([crate::config::USER_AGENT]);
543        let handlers = ResponseHandlers::new(
544            json_response_handler::<AuthTokenResponse>(),
545            failed_response_handler(),
546        );
547        let response = post_json(
548            self.config.transport.as_ref(),
549            url,
550            headers,
551            &body,
552            &handlers,
553            CancellationToken::new(),
554        )
555        .await?;
556        let expires_at = response
557            .value
558            .expire_time
559            .as_deref()
560            .and_then(|time| chrono::DateTime::parse_from_rfc3339(time).ok())
561            .and_then(|time| u64::try_from(time.timestamp()).ok());
562        Ok(ClientSecret {
563            token: response.value.name,
564            url: self.session_url(),
565            expires_at,
566        })
567    }
568
569    fn websocket_config(&self, token: &str, url: &Url) -> WebSocketConfig {
570        let mut url = url.clone();
571        url.query_pairs_mut().append_pair("access_token", token);
572        WebSocketConfig {
573            url,
574            protocols: Vec::new(),
575        }
576    }
577
578    fn parse_server_event(
579        &self,
580        raw: JsonValue,
581    ) -> Result<Vec<RealtimeServerEvent>, ProviderError> {
582        let mut state = self.lock();
583        Ok(parse_event(&mut state, &raw))
584    }
585
586    async fn serialize_client_event(
587        &self,
588        event: RealtimeClientEvent,
589    ) -> Result<JsonValue, ProviderError> {
590        let mut state = self.lock();
591        serialize_event(&mut state, event, self.model_id.as_str())
592    }
593
594    fn build_session_config(
595        &self,
596        config: &RealtimeSessionConfig,
597    ) -> Result<JsonValue, ProviderError> {
598        build_session_config(Some(config), self.model_id.as_str())
599    }
600}
601
602/// Factory for Live API models and session tokens.
603#[derive(Debug, Clone)]
604pub struct GoogleRealtimeFactory {
605    config: SharedConfig,
606    provider: ProviderId,
607}
608
609impl GoogleRealtimeFactory {
610    /// Creates the factory.
611    #[must_use]
612    pub fn new(config: SharedConfig) -> Self {
613        Self {
614            provider: config.provider_id(FAMILY),
615            config,
616        }
617    }
618
619    /// Creates a realtime model.
620    #[must_use]
621    pub fn realtime_model(&self, model_id: &str) -> GoogleRealtimeModel {
622        GoogleRealtimeModel::new(self.config.clone(), model_id)
623    }
624}
625
626impl RealtimeFactory for GoogleRealtimeFactory {
627    fn provider(&self) -> &ProviderId {
628        &self.provider
629    }
630
631    fn model(&self, model_id: &str) -> Result<RealtimeModelRef, NoSuchModelError> {
632        Ok(self.realtime_model(model_id).into())
633    }
634
635    async fn get_token(&self, options: GetTokenOptions) -> Result<ClientSecret, ProviderError> {
636        self.realtime_model(options.model.as_str())
637            .do_create_client_secret(ClientSecretOptions {
638                expires_after_seconds: options.expires_after_seconds,
639                session_config: options.session_config,
640            })
641            .await
642    }
643}
644
645/// Headers of the auth token request: configuration headers only (the key
646/// travels in the query string).
647#[must_use]
648pub fn token_request_headers(config: &GoogleConfig) -> Headers {
649    config
650        .headers
651        .clone()
652        .with_user_agent_suffix([crate::config::USER_AGENT])
653}
ferrin_google/realtime.rs

ferrin_google/
realtime.rs
