Skip to main content

rig_core/providers/gemini/
transcription.rs

1use std::path::Path;
2
3use base64::{Engine, prelude::BASE64_STANDARD};
4use serde_json::{Map, Value, json};
5
6use super::completion::usage_of;
7use crate::error::{EncodeError, ProviderError};
8use crate::json_utils::Lenient;
9use crate::operation::Transcription;
10use crate::transcription;
11use crate::wire::{Body, Descriptor, DocumentDecoder, Encoded, Framing, Mode, Wire};
12
13const TRANSCRIPTION_PREAMBLE: &str =
14    "Translate the provided audio exactly. Do not add additional information.";
15
16/// Encode audio as inline base64 with a transcription system instruction.
17/// Reject invalid generation parameters or JSON serialization failures.
18fn transcription_body(
19    request: transcription::TranscriptionRequest,
20) -> Result<Vec<u8>, EncodeError> {
21    let mut generation_config = match request.additional_params {
22        None | Some(Value::Null) => Map::new(),
23        Some(Value::Object(config)) => config,
24        Some(other) => {
25            return Err(EncodeError::request(format!(
26                "Gemini transcription `additional_params` should be an object, got {other}"
27            )));
28        }
29    };
30    // A temperature named on the request outranks one carried inside
31    // `additional_params`.
32    if let Some(temp) = request.temperature {
33        generation_config.insert("temperature".to_owned(), Value::from(temp));
34    }
35    // The request supplies no explicit MIME type, so infer it from the filename.
36    let mime_type = mime_guess::from_path(Path::new(&request.filename))
37        .first()
38        .map_or_else(|| "audio/mpeg".to_string(), |mime| mime.to_string());
39    let mut body = json!({
40        "contents": [{
41            "parts": [{ "inlineData": { "mimeType": mime_type, "data": null }, "thought": false }],
42            "role": "user",
43        }],
44        "generationConfig": generation_config,
45        "safetySettings": null,
46        "toolConfig": null,
47        "systemInstruction": {
48            "parts": [{ "text": TRANSCRIPTION_PREAMBLE, "thought": false }],
49            "role": "model",
50        },
51    });
52    // The audio moves in: `json!` would copy it.
53    let data = body.pointer_mut("/contents/0/parts/0/inlineData/data");
54    *data.ok_or_else(|| EncodeError::request("the body has no slot for the audio"))? =
55        Value::String(BASE64_STANDARD.encode(request.data));
56    tracing::trace!(
57        target: "rig::transcription",
58        "Sending completion request to Gemini API {}",
59        serde_json::to_string_pretty(&body)?
60    );
61    Ok(serde_json::to_vec(&body)?)
62}
63
64/// The transcription wire: `POST /v1beta/models/{model}:generateContent`.
65///
66/// Both [`Mode`]s send inline audio in JSON and read a whole response.
67#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
68pub struct Transcriptions {
69    /// The provider this wire speaks to.
70    pub provider: super::GeminiConfig,
71    /// The model transcribing, for example
72    /// [`GEMINI_2_0_FLASH`](super::completion::GEMINI_2_0_FLASH).
73    pub model: String,
74}
75
76impl Transcriptions {
77    /// The transcription wire for `model`.
78    pub fn new(provider: super::GeminiConfig, model: impl Into<String>) -> Self {
79        Self {
80            provider,
81            model: model.into(),
82        }
83    }
84}
85
86impl Wire for Transcriptions {
87    type Op = Transcription;
88    type Payload = crate::wire::Encoded;
89    type Frame = crate::wire::WireFrame;
90    type Event = Value;
91    type Decoder<'id> = DocumentDecoder<Transcription, Value>;
92    type Reassembler = crate::wire::document::Unreassembled;
93
94    fn describe(&self) -> Descriptor<'_> {
95        Descriptor::new(super::PROVIDER_NAME).model(self.model.as_str())
96    }
97
98    fn encode(
99        &self,
100        request: transcription::TranscriptionRequest,
101        _mode: Mode,
102    ) -> Result<Encoded, EncodeError> {
103        let body = transcription_body(request)?;
104        let request = http::Request::post(format!(
105            "{}/v1beta/models/{}:generateContent?key={}",
106            self.provider.base_url,
107            self.model,
108            self.provider.api_key.expose()
109        ))
110        .header(http::header::CONTENT_TYPE, "application/json")
111        .body(Body::Bytes(body))?;
112        // Gemini reports no transport request-id header.
113        Ok(Encoded::new(request, Framing::Whole))
114    }
115
116    /// The first candidate's visible text; missing candidates or text are
117    /// response errors.
118    fn decoder<'id>(&self) -> Self::Decoder<'id> {
119        DocumentDecoder::new(super::completion::classify_reply, |reply| {
120            transcript_of(&reply)
121        })
122    }
123}
124
125/// The first candidate's visible text parts, those not marked `thought`, as
126/// a transcript of `reply`. Errors when there is no candidate or no visible
127/// text part.
128pub fn transcript_of(reply: &Value) -> Result<transcription::TranscriptionResponse, ProviderError> {
129    let candidate = reply
130        .arr("candidates")
131        .first()
132        .ok_or_else(|| ProviderError::Response("No response candidates in response".into()))?;
133    let parts: Vec<&str> = candidate
134        .get("content")
135        .map(|content| content.arr("parts"))
136        .unwrap_or_default()
137        .iter()
138        .filter(|part| part.bool("thought") != Some(true))
139        .filter_map(|part| part.str("text"))
140        .collect();
141    if parts.is_empty() {
142        return Err(ProviderError::Response(
143            "Response content contains no text".to_string(),
144        ));
145    }
146    Ok(transcription::TranscriptionResponse {
147        model: reply.str("modelVersion").map(str::to_owned),
148        response_id: Some(reply.str("responseId").unwrap_or_default().to_owned()),
149        usage: reply.get("usageMetadata").map(usage_of).unwrap_or_default(),
150        ..transcription::TranscriptionResponse::new(parts.concat())
151    })
152}
153
154impl super::GeminiConfig {
155    /// The audio transcription wire.
156    pub(crate) fn transcription(&self, model: impl Into<String>) -> Transcriptions {
157        Transcriptions::new(self.clone(), model)
158    }
159}
160
161#[cfg(test)]
162mod tests;