kcode_speaker_v3_gemini_protocol/
lib.rs

1use kcode_speaker_v3_schema::LocalSpeakerLabel;
2use serde_json::Value;
3use std::{error::Error, fmt};
4
5pub const GEMINI_TRANSCRIPT_PROMPT_REVISION: &str = "speaker-v3-gemini-transcript-r1";
6pub const GEMINI_FEATURE_PROMPT_ONE_REVISION: &str = "speaker-v3-gemini-feature-1-r2";
7pub const GEMINI_FEATURE_PROMPT_TWO_REVISION: &str = "speaker-v3-gemini-feature-2-r2";
8pub const GEMINI_FEATURE_PROMPT_THREE_REVISION: &str = "speaker-v3-gemini-feature-3-r2";
9
10pub const GEMINI_FEATURE_PROMPT_REVISIONS: [&str; 3] = [
11    GEMINI_FEATURE_PROMPT_ONE_REVISION,
12    GEMINI_FEATURE_PROMPT_TWO_REVISION,
13    GEMINI_FEATURE_PROMPT_THREE_REVISION,
14];
15
16pub const GEMINI_TRANSCRIPT_PROMPT: &str = include_str!("gemini-transcript-prompt.txt");
17pub const GEMINI_FEATURE_PROMPT_ONE: &str = include_str!("gemini-feature-1-prompt.txt");
18pub const GEMINI_FEATURE_PROMPT_TWO: &str = include_str!("gemini-feature-2-prompt.txt");
19pub const GEMINI_FEATURE_PROMPT_THREE: &str = include_str!("gemini-feature-3-prompt.txt");
20
21pub const FEATURE_PACKETS: [FeaturePacket; 3] =
22    [FeaturePacket::One, FeaturePacket::Two, FeaturePacket::Three];
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub enum GeminiRequestPart<'a> {
26    Audio {
27        media_type: &'static str,
28        bytes: &'a [u8],
29    },
30    Text(String),
31}
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
34pub enum FeaturePacket {
35    One,
36    Two,
37    Three,
38}
39
40impl FeaturePacket {
41    pub fn index(self) -> usize {
42        match self {
43            Self::One => 0,
44            Self::Two => 1,
45            Self::Three => 2,
46        }
47    }
48
49    pub fn prompt(self) -> &'static str {
50        match self {
51            Self::One => GEMINI_FEATURE_PROMPT_ONE,
52            Self::Two => GEMINI_FEATURE_PROMPT_TWO,
53            Self::Three => GEMINI_FEATURE_PROMPT_THREE,
54        }
55    }
56
57    pub fn revision(self) -> &'static str {
58        match self {
59            Self::One => GEMINI_FEATURE_PROMPT_ONE_REVISION,
60            Self::Two => GEMINI_FEATURE_PROMPT_TWO_REVISION,
61            Self::Three => GEMINI_FEATURE_PROMPT_THREE_REVISION,
62        }
63    }
64}
65
66#[derive(Debug, Clone, PartialEq, Eq)]
67pub enum GeminiProtocolError {
68    InvalidResponse(String),
69    TextCandidateCount(usize),
70}
71
72impl fmt::Display for GeminiProtocolError {
73    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
74        match self {
75            Self::InvalidResponse(message) => {
76                write!(formatter, "invalid Gemini response: {message}")
77            }
78            Self::TextCandidateCount(count) => {
79                write!(
80                    formatter,
81                    "Gemini returned {count} nonblank textual candidates"
82                )
83            }
84        }
85    }
86}
87
88impl Error for GeminiProtocolError {}
89
90pub fn gemini_transcript_request(audio: &[u8]) -> [GeminiRequestPart<'_>; 2] {
91    [
92        GeminiRequestPart::Audio {
93            media_type: "audio/ogg",
94            bytes: audio,
95        },
96        GeminiRequestPart::Text(GEMINI_TRANSCRIPT_PROMPT.to_owned()),
97    ]
98}
99
100pub fn gemini_feature_cached_prefix<'a>(
101    audio: &'a [u8],
102    transcript: &str,
103) -> [GeminiRequestPart<'a>; 2] {
104    let (shared_prefix, _) = feature_prompt_parts(FeaturePacket::One);
105    [
106        GeminiRequestPart::Audio {
107            media_type: "audio/ogg",
108            bytes: audio,
109        },
110        GeminiRequestPart::Text(format!("{shared_prefix}{transcript}")),
111    ]
112}
113
114pub fn gemini_feature_suffix(packet: FeaturePacket, target: LocalSpeakerLabel) -> String {
115    let (_, suffix) = feature_prompt_parts(packet);
116    suffix.replace("{{TARGET_SPEAKER}}", &target.to_string())
117}
118
119pub fn extract_gemini_text(response: &Value) -> Result<String, GeminiProtocolError> {
120    let candidates = response
121        .get("candidates")
122        .and_then(Value::as_array)
123        .ok_or_else(|| GeminiProtocolError::InvalidResponse("candidates is not an array".into()))?;
124    let mut textual_candidates = Vec::new();
125    for (candidate_index, candidate) in candidates.iter().enumerate() {
126        let parts = candidate
127            .get("content")
128            .and_then(|content| content.get("parts"))
129            .and_then(Value::as_array)
130            .ok_or_else(|| {
131                GeminiProtocolError::InvalidResponse(format!(
132                    "candidate {candidate_index} has no content parts array"
133                ))
134            })?;
135        let mut text = String::new();
136        for (part_index, part) in parts.iter().enumerate() {
137            if let Some(value) = part.get("text") {
138                let value = value.as_str().ok_or_else(|| {
139                    GeminiProtocolError::InvalidResponse(format!(
140                        "candidate {candidate_index} part {part_index} text is not a string"
141                    ))
142                })?;
143                text.push_str(value);
144            }
145        }
146        if !text.trim().is_empty() {
147            textual_candidates.push(text);
148        }
149    }
150    if textual_candidates.len() != 1 {
151        return Err(GeminiProtocolError::TextCandidateCount(
152            textual_candidates.len(),
153        ));
154    }
155    Ok(textual_candidates.pop().expect("length checked"))
156}
157
158fn feature_prompt_parts(packet: FeaturePacket) -> (&'static str, &'static str) {
159    packet
160        .prompt()
161        .split_once("{{TRANSCRIPT}}")
162        .expect("frozen feature prompt contains transcript placeholder")
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168    use kcode_speaker_v3_schema::FEATURE_NAMES;
169    use serde_json::json;
170    use std::time::Instant;
171
172    const EXPECTED_PROMPT_FINGERPRINTS: [u64; 4] = [
173        15056256181481631076,
174        14503216141444052946,
175        5759944296588955155,
176        17635457290013820104,
177    ];
178
179    fn label(number: u32) -> LocalSpeakerLabel {
180        LocalSpeakerLabel::new(number).unwrap()
181    }
182
183    fn candidate(parts: Vec<Value>) -> Value {
184        json!({ "content": { "parts": parts } })
185    }
186
187    fn fingerprint(value: &str) -> u64 {
188        value.bytes().fold(0xcbf29ce484222325, |hash, byte| {
189            (hash ^ u64::from(byte)).wrapping_mul(0x100000001b3)
190        })
191    }
192
193    #[test]
194    fn prompt_bytes_revisions_and_feature_coverage_are_frozen() {
195        let actual = [
196            fingerprint(GEMINI_TRANSCRIPT_PROMPT),
197            fingerprint(GEMINI_FEATURE_PROMPT_ONE),
198            fingerprint(GEMINI_FEATURE_PROMPT_TWO),
199            fingerprint(GEMINI_FEATURE_PROMPT_THREE),
200        ];
201        assert_eq!(actual, EXPECTED_PROMPT_FINGERPRINTS);
202        assert_eq!(
203            GEMINI_FEATURE_PROMPT_REVISIONS,
204            [
205                "speaker-v3-gemini-feature-1-r2",
206                "speaker-v3-gemini-feature-2-r2",
207                "speaker-v3-gemini-feature-3-r2",
208            ]
209        );
210        assert!(GEMINI_TRANSCRIPT_PROMPT.contains("[high] Speaker N:"));
211        for (packet, names) in FEATURE_PACKETS.into_iter().zip([
212            &FEATURE_NAMES[..8],
213            &FEATURE_NAMES[8..16],
214            &FEATURE_NAMES[16..],
215        ]) {
216            assert!(packet.prompt().contains("{{TRANSCRIPT}}"));
217            assert!(packet.prompt().contains("{{TARGET_SPEAKER}}"));
218            for name in names {
219                assert!(packet.prompt().contains(name));
220            }
221        }
222    }
223
224    #[test]
225    fn requests_keep_audio_prefix_shared_text_and_target_last() {
226        let audio = b"OggS bytes";
227        assert_eq!(
228            gemini_transcript_request(audio),
229            [
230                GeminiRequestPart::Audio {
231                    media_type: "audio/ogg",
232                    bytes: audio,
233                },
234                GeminiRequestPart::Text(GEMINI_TRANSCRIPT_PROMPT.into()),
235            ]
236        );
237        let transcript = "literal {{TARGET_SPEAKER}}\ntranscript";
238        let prefix = gemini_feature_cached_prefix(audio, transcript);
239        assert_eq!(
240            prefix[0],
241            GeminiRequestPart::Audio {
242                media_type: "audio/ogg",
243                bytes: audio,
244            }
245        );
246        let GeminiRequestPart::Text(prefix_text) = &prefix[1] else {
247            panic!()
248        };
249        let (shared, _) = feature_prompt_parts(FeaturePacket::One);
250        assert_eq!(prefix_text, &format!("{shared}{transcript}"));
251        for packet in FEATURE_PACKETS {
252            let suffix = gemini_feature_suffix(packet, label(12));
253            assert!(!suffix.contains("{{TARGET_SPEAKER}}"));
254            assert!(suffix.ends_with("Speaker 12\n"));
255            assert_eq!(
256                packet.revision(),
257                GEMINI_FEATURE_PROMPT_REVISIONS[packet.index()]
258            );
259            assert_eq!(
260                feature_prompt_parts(packet).0,
261                feature_prompt_parts(FeaturePacket::One).0
262            );
263        }
264    }
265
266    #[test]
267    fn extraction_preserves_one_textual_candidate() {
268        let response = serde_json::json!({
269            "candidates": [
270                candidate(vec![
271                    serde_json::json!({"text": "a\n"}),
272                    serde_json::json!({"inlineData": {}}),
273                    serde_json::json!({"text": " b"})
274                ]),
275                candidate(vec![serde_json::json!({"text": "   "})])
276            ]
277        });
278        assert_eq!(extract_gemini_text(&response).unwrap(), "a\n b");
279        let ambiguous = serde_json::json!({
280            "candidates": [
281                candidate(vec![serde_json::json!({"text": "one"})]),
282                candidate(vec![serde_json::json!({"text": "two"})])
283            ]
284        });
285        assert_eq!(
286            extract_gemini_text(&ambiguous),
287            Err(GeminiProtocolError::TextCandidateCount(2))
288        );
289        assert_eq!(
290            extract_gemini_text(&serde_json::json!({})),
291            Err(GeminiProtocolError::InvalidResponse(
292                "candidates is not an array".into()
293            ))
294        );
295    }
296
297    #[test]
298    fn one_mebibyte_extraction_completes_within_reference_envelope() {
299        let started = Instant::now();
300        let text = "x".repeat(1_048_576);
301        let response = serde_json::json!({
302            "candidates": [candidate(vec![serde_json::json!({"text": text})])]
303        });
304        assert_eq!(extract_gemini_text(&response).unwrap().len(), 1_048_576);
305        assert!(started.elapsed().as_secs() < 10);
306    }
307}
kcode_speaker_v3_gemini_protocol/lib.rs

kcode_speaker_v3_gemini_protocol/
lib.rs
