rig_core/providers/gemini/
transcription.rs1use 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
16fn 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 if let Some(temp) = request.temperature {
33 generation_config.insert("temperature".to_owned(), Value::from(temp));
34 }
35 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 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#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
68pub struct Transcriptions {
69 pub provider: super::GeminiConfig,
71 pub model: String,
74}
75
76impl Transcriptions {
77 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 Ok(Encoded::new(request, Framing::Whole))
114 }
115
116 fn decoder<'id>(&self) -> Self::Decoder<'id> {
119 DocumentDecoder::new(super::completion::classify_reply, |reply| {
120 transcript_of(&reply)
121 })
122 }
123}
124
125pub 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 pub(crate) fn transcription(&self, model: impl Into<String>) -> Transcriptions {
157 Transcriptions::new(self.clone(), model)
158 }
159}
160
161#[cfg(test)]
162mod tests;