use kcode_speaker_v3_schema::LocalSpeakerLabel;
use serde_json::Value;
use std::{error::Error, fmt};
pub const GEMINI_TRANSCRIPT_PROMPT_REVISION: &str = "speaker-v3-gemini-transcript-r1";
pub const GEMINI_FEATURE_PROMPT_ONE_REVISION: &str = "speaker-v3-gemini-feature-1-r2";
pub const GEMINI_FEATURE_PROMPT_TWO_REVISION: &str = "speaker-v3-gemini-feature-2-r2";
pub const GEMINI_FEATURE_PROMPT_THREE_REVISION: &str = "speaker-v3-gemini-feature-3-r2";
pub const GEMINI_FEATURE_PROMPT_REVISIONS: [&str; 3] = [
GEMINI_FEATURE_PROMPT_ONE_REVISION,
GEMINI_FEATURE_PROMPT_TWO_REVISION,
GEMINI_FEATURE_PROMPT_THREE_REVISION,
];
pub const GEMINI_TRANSCRIPT_PROMPT: &str = include_str!("gemini-transcript-prompt.txt");
pub const GEMINI_FEATURE_PROMPT_ONE: &str = include_str!("gemini-feature-1-prompt.txt");
pub const GEMINI_FEATURE_PROMPT_TWO: &str = include_str!("gemini-feature-2-prompt.txt");
pub const GEMINI_FEATURE_PROMPT_THREE: &str = include_str!("gemini-feature-3-prompt.txt");
pub const FEATURE_PACKETS: [FeaturePacket; 3] =
[FeaturePacket::One, FeaturePacket::Two, FeaturePacket::Three];
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GeminiRequestPart<'a> {
Audio {
media_type: &'static str,
bytes: &'a [u8],
},
Text(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum FeaturePacket {
One,
Two,
Three,
}
impl FeaturePacket {
pub fn index(self) -> usize {
match self {
Self::One => 0,
Self::Two => 1,
Self::Three => 2,
}
}
pub fn prompt(self) -> &'static str {
match self {
Self::One => GEMINI_FEATURE_PROMPT_ONE,
Self::Two => GEMINI_FEATURE_PROMPT_TWO,
Self::Three => GEMINI_FEATURE_PROMPT_THREE,
}
}
pub fn revision(self) -> &'static str {
match self {
Self::One => GEMINI_FEATURE_PROMPT_ONE_REVISION,
Self::Two => GEMINI_FEATURE_PROMPT_TWO_REVISION,
Self::Three => GEMINI_FEATURE_PROMPT_THREE_REVISION,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GeminiProtocolError {
InvalidResponse(String),
TextCandidateCount(usize),
}
impl fmt::Display for GeminiProtocolError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::InvalidResponse(message) => {
write!(formatter, "invalid Gemini response: {message}")
}
Self::TextCandidateCount(count) => {
write!(
formatter,
"Gemini returned {count} nonblank textual candidates"
)
}
}
}
}
impl Error for GeminiProtocolError {}
pub fn gemini_transcript_request(audio: &[u8]) -> [GeminiRequestPart<'_>; 2] {
[
GeminiRequestPart::Audio {
media_type: "audio/ogg",
bytes: audio,
},
GeminiRequestPart::Text(GEMINI_TRANSCRIPT_PROMPT.to_owned()),
]
}
pub fn gemini_feature_cached_prefix<'a>(
audio: &'a [u8],
transcript: &str,
) -> [GeminiRequestPart<'a>; 2] {
let (shared_prefix, _) = feature_prompt_parts(FeaturePacket::One);
[
GeminiRequestPart::Audio {
media_type: "audio/ogg",
bytes: audio,
},
GeminiRequestPart::Text(format!("{shared_prefix}{transcript}")),
]
}
pub fn gemini_feature_suffix(packet: FeaturePacket, target: LocalSpeakerLabel) -> String {
let (_, suffix) = feature_prompt_parts(packet);
suffix.replace("{{TARGET_SPEAKER}}", &target.to_string())
}
pub fn extract_gemini_text(response: &Value) -> Result<String, GeminiProtocolError> {
let candidates = response
.get("candidates")
.and_then(Value::as_array)
.ok_or_else(|| GeminiProtocolError::InvalidResponse("candidates is not an array".into()))?;
let mut textual_candidates = Vec::new();
for (candidate_index, candidate) in candidates.iter().enumerate() {
let parts = candidate
.get("content")
.and_then(|content| content.get("parts"))
.and_then(Value::as_array)
.ok_or_else(|| {
GeminiProtocolError::InvalidResponse(format!(
"candidate {candidate_index} has no content parts array"
))
})?;
let mut text = String::new();
for (part_index, part) in parts.iter().enumerate() {
if let Some(value) = part.get("text") {
let value = value.as_str().ok_or_else(|| {
GeminiProtocolError::InvalidResponse(format!(
"candidate {candidate_index} part {part_index} text is not a string"
))
})?;
text.push_str(value);
}
}
if !text.trim().is_empty() {
textual_candidates.push(text);
}
}
if textual_candidates.len() != 1 {
return Err(GeminiProtocolError::TextCandidateCount(
textual_candidates.len(),
));
}
Ok(textual_candidates.pop().expect("length checked"))
}
fn feature_prompt_parts(packet: FeaturePacket) -> (&'static str, &'static str) {
packet
.prompt()
.split_once("{{TRANSCRIPT}}")
.expect("frozen feature prompt contains transcript placeholder")
}
#[cfg(test)]
mod tests {
use super::*;
use kcode_speaker_v3_schema::FEATURE_NAMES;
use serde_json::json;
use std::time::Instant;
const EXPECTED_PROMPT_FINGERPRINTS: [u64; 4] = [
15056256181481631076,
14503216141444052946,
5759944296588955155,
17635457290013820104,
];
fn label(number: u32) -> LocalSpeakerLabel {
LocalSpeakerLabel::new(number).unwrap()
}
fn candidate(parts: Vec<Value>) -> Value {
json!({ "content": { "parts": parts } })
}
fn fingerprint(value: &str) -> u64 {
value.bytes().fold(0xcbf29ce484222325, |hash, byte| {
(hash ^ u64::from(byte)).wrapping_mul(0x100000001b3)
})
}
#[test]
fn prompt_bytes_revisions_and_feature_coverage_are_frozen() {
let actual = [
fingerprint(GEMINI_TRANSCRIPT_PROMPT),
fingerprint(GEMINI_FEATURE_PROMPT_ONE),
fingerprint(GEMINI_FEATURE_PROMPT_TWO),
fingerprint(GEMINI_FEATURE_PROMPT_THREE),
];
assert_eq!(actual, EXPECTED_PROMPT_FINGERPRINTS);
assert_eq!(
GEMINI_FEATURE_PROMPT_REVISIONS,
[
"speaker-v3-gemini-feature-1-r2",
"speaker-v3-gemini-feature-2-r2",
"speaker-v3-gemini-feature-3-r2",
]
);
assert!(GEMINI_TRANSCRIPT_PROMPT.contains("[high] Speaker N:"));
for (packet, names) in FEATURE_PACKETS.into_iter().zip([
&FEATURE_NAMES[..8],
&FEATURE_NAMES[8..16],
&FEATURE_NAMES[16..],
]) {
assert!(packet.prompt().contains("{{TRANSCRIPT}}"));
assert!(packet.prompt().contains("{{TARGET_SPEAKER}}"));
for name in names {
assert!(packet.prompt().contains(name));
}
}
}
#[test]
fn requests_keep_audio_prefix_shared_text_and_target_last() {
let audio = b"OggS bytes";
assert_eq!(
gemini_transcript_request(audio),
[
GeminiRequestPart::Audio {
media_type: "audio/ogg",
bytes: audio,
},
GeminiRequestPart::Text(GEMINI_TRANSCRIPT_PROMPT.into()),
]
);
let transcript = "literal {{TARGET_SPEAKER}}\ntranscript";
let prefix = gemini_feature_cached_prefix(audio, transcript);
assert_eq!(
prefix[0],
GeminiRequestPart::Audio {
media_type: "audio/ogg",
bytes: audio,
}
);
let GeminiRequestPart::Text(prefix_text) = &prefix[1] else {
panic!()
};
let (shared, _) = feature_prompt_parts(FeaturePacket::One);
assert_eq!(prefix_text, &format!("{shared}{transcript}"));
for packet in FEATURE_PACKETS {
let suffix = gemini_feature_suffix(packet, label(12));
assert!(!suffix.contains("{{TARGET_SPEAKER}}"));
assert!(suffix.ends_with("Speaker 12\n"));
assert_eq!(
packet.revision(),
GEMINI_FEATURE_PROMPT_REVISIONS[packet.index()]
);
assert_eq!(
feature_prompt_parts(packet).0,
feature_prompt_parts(FeaturePacket::One).0
);
}
}
#[test]
fn extraction_preserves_one_textual_candidate() {
let response = serde_json::json!({
"candidates": [
candidate(vec![
serde_json::json!({"text": "a\n"}),
serde_json::json!({"inlineData": {}}),
serde_json::json!({"text": " b"})
]),
candidate(vec![serde_json::json!({"text": " "})])
]
});
assert_eq!(extract_gemini_text(&response).unwrap(), "a\n b");
let ambiguous = serde_json::json!({
"candidates": [
candidate(vec![serde_json::json!({"text": "one"})]),
candidate(vec![serde_json::json!({"text": "two"})])
]
});
assert_eq!(
extract_gemini_text(&ambiguous),
Err(GeminiProtocolError::TextCandidateCount(2))
);
assert_eq!(
extract_gemini_text(&serde_json::json!({})),
Err(GeminiProtocolError::InvalidResponse(
"candidates is not an array".into()
))
);
}
#[test]
fn one_mebibyte_extraction_completes_within_reference_envelope() {
let started = Instant::now();
let text = "x".repeat(1_048_576);
let response = serde_json::json!({
"candidates": [candidate(vec![serde_json::json!({"text": text})])]
});
assert_eq!(extract_gemini_text(&response).unwrap().len(), 1_048_576);
assert!(started.elapsed().as_secs() < 10);
}
}