pub use kcode_speaker_types::{FeatureVector, Key};
use kcode_speaker_types::FEATURE_COUNT;
use serde::Deserialize;
use std::fmt;
use std::sync::OnceLock;
const SINGLE_SEGMENT_LIMIT_MS: u64 = 240_000;
const MAX_SEGMENT_MS: u128 = 239_999;
const OVERLAP_MS: u128 = 5_000;
const UNIQUE_CAPACITY_MS: u128 = MAX_SEGMENT_MS - OVERLAP_MS;
pub struct SegmentPlan {
pub policy: Key,
pub source_duration_ms: u64,
pub segments: Vec<PlannedSegment>,
}
pub struct PlannedSegment {
pub ordinal: u16,
pub start_ms: u64,
pub end_ms: u64,
}
pub struct ExtractionContract {
pub schema_key: Key,
pub prompt: &'static str,
pub response_mime: &'static str,
}
pub enum ExtractionOutcome {
Scored(ScoredClip),
Unscorable { reason: String },
}
pub struct ScoredClip {
pub recording_quality: u8,
pub speakers: Vec<SpeakerSample>,
}
pub struct SpeakerSample {
pub speaker_ordinal: u16,
pub primary_language: Key,
pub closest_dialect: String,
pub usable_speech_ms: u32,
pub features: FeatureVector,
}
#[derive(Debug, PartialEq, Eq)]
pub enum ExtractError {
ZeroDuration,
SegmentCountExceedsU16 { required: u128 },
InvalidJson(String),
InvalidResponse(&'static str),
}
impl fmt::Display for ExtractError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::ZeroDuration => formatter.write_str("source duration must be positive"),
Self::SegmentCountExceedsU16 { required } => {
write!(formatter, "required segment count {required} exceeds u16")
}
Self::InvalidJson(message) => write!(formatter, "invalid response JSON: {message}"),
Self::InvalidResponse(message) => write!(formatter, "invalid response: {message}"),
}
}
}
impl std::error::Error for ExtractError {}
pub fn plan_segments(duration_ms: u64) -> Result<SegmentPlan, ExtractError> {
if duration_ms == 0 {
return Err(ExtractError::ZeroDuration);
}
if duration_ms < SINGLE_SEGMENT_LIMIT_MS {
return Ok(SegmentPlan {
policy: segment_policy(),
source_duration_ms: duration_ms,
segments: vec![PlannedSegment {
ordinal: 0,
start_ms: 0,
end_ms: duration_ms,
}],
});
}
let duration = u128::from(duration_ms);
let unique_duration = duration - OVERLAP_MS;
let required = unique_duration.div_ceil(UNIQUE_CAPACITY_MS);
if required > u128::from(u16::MAX) {
return Err(ExtractError::SegmentCountExceedsU16 { required });
}
let count = usize::try_from(required).expect("u16-bounded count fits usize");
let total_segment_duration = duration + (required - 1) * OVERLAP_MS;
let short_length = total_segment_duration / required;
let longer_count = total_segment_duration % required;
let mut segments = Vec::with_capacity(count);
let mut start = 0_u128;
for index in 0..count {
let index_u128 = u128::try_from(index).expect("u16-bounded index fits u128");
let length = short_length + if index_u128 < longer_count { 1 } else { 0 };
let end = start + length;
segments.push(PlannedSegment {
ordinal: u16::try_from(index).expect("count is bounded by u16::MAX"),
start_ms: u64::try_from(start).expect("segment start is within source duration"),
end_ms: u64::try_from(end).expect("segment end is within source duration"),
});
start = end - OVERLAP_MS;
}
debug_assert_eq!(
segments.last().map(|segment| segment.end_ms),
Some(duration_ms)
);
debug_assert!(
segments
.iter()
.all(|segment| u128::from(segment.end_ms - segment.start_ms) <= MAX_SEGMENT_MS)
);
Ok(SegmentPlan {
policy: segment_policy(),
source_duration_ms: duration_ms,
segments,
})
}
pub fn contract() -> &'static ExtractionContract {
static CONTRACT: OnceLock<ExtractionContract> = OnceLock::new();
CONTRACT.get_or_init(|| ExtractionContract {
schema_key: frozen_key("gemini-speaker-35/1"),
prompt: include_str!("assets/prompt-v1.txt"),
response_mime: "application/json",
})
}
pub fn parse(response: &str, clip_duration_ms: u64) -> Result<ExtractionOutcome, ExtractError> {
let response: WireOutcome = serde_json::from_str(response)
.map_err(|error| ExtractError::InvalidJson(error.to_string()))?;
match response {
WireOutcome::Unscorable { reason } => {
if reason.is_empty() || reason.len() > 240 {
return Err(ExtractError::InvalidResponse(
"unscorable reason must contain 1..=240 bytes",
));
}
Ok(ExtractionOutcome::Unscorable { reason })
}
WireOutcome::Scored {
recording_quality,
speakers,
} => {
if recording_quality > 100 {
return Err(ExtractError::InvalidResponse(
"recording quality must be in 0..=100",
));
}
if speakers.is_empty() {
return Err(ExtractError::InvalidResponse(
"scored response must contain at least one speaker",
));
}
let mut validated = Vec::with_capacity(speakers.len());
for (index, speaker) in speakers.into_iter().enumerate() {
if usize::from(speaker.speaker_ordinal) != index {
return Err(ExtractError::InvalidResponse(
"speaker ordinals must be unique and contiguous from zero",
));
}
if speaker.closest_dialect.is_empty() || speaker.closest_dialect.len() > 128 {
return Err(ExtractError::InvalidResponse(
"closest dialect must contain 1..=128 bytes",
));
}
if speaker.usable_speech_ms == 0
|| u64::from(speaker.usable_speech_ms) > clip_duration_ms
{
return Err(ExtractError::InvalidResponse(
"usable speech must be positive and no greater than clip duration",
));
}
let primary_language = Key::parse(&speaker.primary_language).map_err(|_| {
ExtractError::InvalidResponse("primary language is not a valid shared key")
})?;
let feature_values: [u8; FEATURE_COUNT] =
speaker.features.try_into().map_err(|_| {
ExtractError::InvalidResponse("features must contain exactly 35 integers")
})?;
let features = FeatureVector::new(feature_values).map_err(|_| {
ExtractError::InvalidResponse("feature value is outside shared validation")
})?;
validated.push(SpeakerSample {
speaker_ordinal: speaker.speaker_ordinal,
primary_language,
closest_dialect: speaker.closest_dialect,
usable_speech_ms: speaker.usable_speech_ms,
features,
});
}
Ok(ExtractionOutcome::Scored(ScoredClip {
recording_quality,
speakers: validated,
}))
}
}
}
fn segment_policy() -> Key {
frozen_key("speaker-segments/1")
}
fn frozen_key(value: &str) -> Key {
match Key::parse(value) {
Ok(key) => key,
Err(_) => panic!("invalid frozen key"),
}
}
#[derive(Deserialize)]
#[serde(tag = "status", deny_unknown_fields)]
enum WireOutcome {
#[serde(rename = "scored")]
Scored {
#[serde(rename = "recordingQuality")]
recording_quality: u8,
speakers: Vec<WireSpeaker>,
},
#[serde(rename = "unscorable")]
Unscorable { reason: String },
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct WireSpeaker {
#[serde(rename = "speakerOrdinal")]
speaker_ordinal: u16,
#[serde(rename = "primaryLanguage")]
primary_language: String,
#[serde(rename = "closestDialect")]
closest_dialect: String,
#[serde(rename = "usableSpeechMs")]
usable_speech_ms: u32,
features: Vec<u8>,
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::{Value, json};
fn base_features() -> Vec<u8> {
let mut values = vec![50; FEATURE_COUNT];
values[11] = 60;
values
}
fn speaker(ordinal: u16, usable_speech_ms: u32) -> Value {
json!({
"speakerOrdinal": ordinal,
"primaryLanguage": "eng",
"closestDialect": "General American English",
"usableSpeechMs": usable_speech_ms,
"features": base_features(),
})
}
fn scored(speakers: Vec<Value>) -> Value {
json!({
"status": "scored",
"recordingQuality": 82,
"speakers": speakers,
})
}
fn assert_rejected(value: Value, clip_duration_ms: u64) {
assert!(parse(&value.to_string(), clip_duration_ms).is_err());
}
fn assert_plan_invariants(duration_ms: u64) -> SegmentPlan {
let plan = plan_segments(duration_ms).expect("plan should succeed");
assert_eq!(plan.source_duration_ms, duration_ms);
assert_eq!(plan.policy.as_ref(), "speaker-segments/1");
assert_eq!(plan.segments.first().expect("segment").start_ms, 0);
assert_eq!(plan.segments.last().expect("segment").end_ms, duration_ms);
let lengths: Vec<u64> = plan
.segments
.iter()
.enumerate()
.map(|(index, segment)| {
assert_eq!(usize::from(segment.ordinal), index);
assert!(segment.start_ms < segment.end_ms);
let length = segment.end_ms - segment.start_ms;
assert!(length <= u64::try_from(MAX_SEGMENT_MS).unwrap());
length
})
.collect();
for (index, pair) in plan.segments.windows(2).enumerate() {
assert_eq!(pair[0].end_ms - pair[1].start_ms, 5_000);
assert_eq!(usize::from(pair[1].ordinal), index + 1);
}
let minimum = *lengths.iter().min().expect("length");
let maximum = *lengths.iter().max().expect("length");
assert!(maximum - minimum <= 1);
if plan.segments.len() > 1 {
let fewer = u128::try_from(plan.segments.len() - 1).unwrap();
let fewer_capacity = fewer * UNIQUE_CAPACITY_MS + OVERLAP_MS;
assert!(u128::from(duration_ms) > fewer_capacity);
}
plan
}
#[test]
fn rejects_zero_duration() {
assert_eq!(plan_segments(0).err(), Some(ExtractError::ZeroDuration));
}
#[test]
fn plans_below_equal_and_above_four_minutes() {
let below = assert_plan_invariants(239_999);
assert_eq!(below.segments.len(), 1);
assert_eq!(below.segments[0].end_ms, 239_999);
let equal = assert_plan_invariants(240_000);
assert_eq!(equal.segments.len(), 2);
let above = assert_plan_invariants(240_001);
assert_eq!(above.segments.len(), 2);
}
#[test]
fn plans_long_and_remainder_sources() {
for duration in [474_998, 474_999, 475_000, 1_000_003, 9_876_543] {
assert_plan_invariants(duration);
}
}
#[test]
fn plans_exact_capacity_with_minimum_count() {
let duration = 7 * 234_999 + 5_000;
let plan = assert_plan_invariants(duration);
assert_eq!(plan.segments.len(), 7);
assert!(
plan.segments
.iter()
.all(|segment| segment.end_ms - segment.start_ms == 239_999)
);
}
#[test]
fn rejects_count_that_does_not_fit_u16() {
let duration = u64::from(u16::MAX) * 234_999 + 5_001;
assert!(matches!(
plan_segments(duration),
Err(ExtractError::SegmentCountExceedsU16 { required: 65_536 })
));
}
#[test]
fn maximum_supported_count_is_valid() {
let duration = u64::from(u16::MAX) * 234_999 + 5_000;
let plan = assert_plan_invariants(duration);
assert_eq!(plan.segments.len(), usize::from(u16::MAX));
}
#[test]
fn planning_is_deterministic() {
let first = plan_segments(1_234_567).unwrap();
let second = plan_segments(1_234_567).unwrap();
assert_eq!(first.policy.as_ref(), second.policy.as_ref());
assert_eq!(first.source_duration_ms, second.source_duration_ms);
assert_eq!(first.segments.len(), second.segments.len());
for (left, right) in first.segments.iter().zip(&second.segments) {
assert_eq!(left.ordinal, right.ordinal);
assert_eq!(left.start_ms, right.start_ms);
assert_eq!(left.end_ms, right.end_ms);
}
}
#[test]
fn contract_has_frozen_identity_and_prompt() {
let value = contract();
assert_eq!(value.schema_key.as_ref(), "gemini-speaker-35/1");
assert_eq!(value.response_mime, "application/json");
assert!(
value
.prompt
.starts_with("Analyze the attached audio directly.")
);
assert!(value.prompt.contains("0 filler_preference"));
assert!(value.prompt.contains("11 perceived_age"));
assert!(value.prompt.contains("34 aspiration_intensity"));
assert!(
value
.prompt
.trim_end()
.ends_with("Return the JSON object only.")
);
assert!(std::ptr::eq(value, contract()));
}
#[test]
fn accepts_reordered_scored_fields() {
let response = format!(
r#"{{"speakers":[{}],"recordingQuality":0,"status":"scored"}}"#,
speaker(0, 1_000)
);
let outcome = parse(&response, 1_000).expect("reordered keys are valid");
let ExtractionOutcome::Scored(clip) = outcome else {
panic!("expected scored response");
};
assert_eq!(clip.recording_quality, 0);
assert_eq!(clip.speakers.len(), 1);
}
#[test]
fn accepts_multi_speaker_scored_with_overlapping_durations() {
let mut second = speaker(1, 1_000);
second["primaryLanguage"] = json!("spa");
second["closestDialect"] = json!("Mexican Spanish");
let response = scored(vec![speaker(0, 1_000), second]).to_string();
let outcome = parse(&response, 1_000).expect("overlap is valid");
let ExtractionOutcome::Scored(clip) = outcome else {
panic!("expected scored response");
};
assert_eq!(clip.recording_quality, 82);
assert_eq!(clip.speakers.len(), 2);
assert_eq!(clip.speakers[0].speaker_ordinal, 0);
assert_eq!(clip.speakers[1].speaker_ordinal, 1);
assert_eq!(clip.speakers[1].primary_language.as_ref(), "spa");
}
#[test]
fn accepts_whole_clip_unscorable() {
let outcome = parse(
r#"{"reason":"Insufficient substantive speech.","status":"unscorable"}"#,
0,
)
.expect("whole-clip abstention is valid");
let ExtractionOutcome::Unscorable { reason } = outcome else {
panic!("expected unscorable response");
};
assert_eq!(reason, "Insufficient substantive speech.");
}
#[test]
fn accepts_quality_and_feature_boundaries() {
for quality in [0, 100] {
for maximum in [false, true] {
let mut values = if maximum {
vec![100; FEATURE_COUNT]
} else {
vec![0; FEATURE_COUNT]
};
if maximum {
values[11] = 120;
}
let mut value = scored(vec![speaker(0, 1)]);
value["recordingQuality"] = json!(quality);
value["speakers"][0]["features"] = json!(values);
let outcome = parse(&value.to_string(), 1).expect("boundary is valid");
let ExtractionOutcome::Scored(clip) = outcome else {
panic!("expected scored response");
};
let actual: &[u8; FEATURE_COUNT] = clip.speakers[0].features.as_ref();
assert_eq!(actual[11], if maximum { 120 } else { 0 });
}
}
}
#[test]
fn rejects_each_feature_above_its_upper_boundary() {
for index in 0..FEATURE_COUNT {
let mut values = base_features();
values[index] = if index == 11 { 121 } else { 101 };
let mut value = scored(vec![speaker(0, 1)]);
value["speakers"][0]["features"] = json!(values);
assert_rejected(value, 1);
}
}
#[test]
fn rejects_wrong_feature_counts() {
for count in [0, 1, FEATURE_COUNT - 1, FEATURE_COUNT + 1] {
let mut value = scored(vec![speaker(0, 1)]);
value["speakers"][0]["features"] = json!(vec![50; count]);
assert_rejected(value, 1);
}
}
#[test]
fn rejects_quality_and_duration_ranges() {
let mut quality = scored(vec![speaker(0, 1)]);
quality["recordingQuality"] = json!(101);
assert_rejected(quality, 1);
let mut zero = scored(vec![speaker(0, 1)]);
zero["speakers"][0]["usableSpeechMs"] = json!(0);
assert_rejected(zero, 1);
assert_rejected(scored(vec![speaker(0, 2)]), 1);
assert_rejected(scored(vec![speaker(0, 1)]), 0);
}
#[test]
fn validates_dialect_byte_lengths() {
let mut accepted = scored(vec![speaker(0, 1)]);
accepted["speakers"][0]["closestDialect"] = json!("é".repeat(64));
assert!(parse(&accepted.to_string(), 1).is_ok());
for dialect in [
"".to_owned(),
"a".repeat(129),
format!("{}a", "é".repeat(64)),
] {
let mut value = scored(vec![speaker(0, 1)]);
value["speakers"][0]["closestDialect"] = json!(dialect);
assert_rejected(value, 1);
}
}
#[test]
fn validates_reason_byte_lengths() {
for reason in ["a".to_owned(), "a".repeat(240), "é".repeat(120)] {
let value = json!({"status": "unscorable", "reason": reason});
assert!(parse(&value.to_string(), 0).is_ok());
}
for reason in [
"".to_owned(),
"a".repeat(241),
format!("{}a", "é".repeat(120)),
] {
let value = json!({"status": "unscorable", "reason": reason});
assert_rejected(value, 0);
}
}
#[test]
fn rejects_invalid_primary_language_keys() {
let languages = ["".to_owned(), "en g".to_owned(), "a".repeat(129)];
for language in languages {
let mut value = scored(vec![speaker(0, 1)]);
value["speakers"][0]["primaryLanguage"] = json!(language);
assert_rejected(value, 1);
}
}
#[test]
fn rejects_empty_and_noncontiguous_speaker_sets() {
assert_rejected(scored(vec![]), 1);
assert_rejected(scored(vec![speaker(1, 1)]), 1);
assert_rejected(scored(vec![speaker(0, 1), speaker(0, 1)]), 1);
assert_rejected(scored(vec![speaker(0, 1), speaker(2, 1)]), 1);
}
#[test]
fn rejects_malformed_and_non_object_json() {
for response in [
"",
"{",
"null",
"[]",
"true",
"42",
r#""scored""#,
r#"{"status":"unknown"}"#,
r#"{"status":null}"#,
r#"{"status":"unscorable","reason":"x"} trailing"#,
r#"{"status":"scored","recordingQuality":NaN,"speakers":[]}"#,
r#"{"status":"scored","recordingQuality":Infinity,"speakers":[]}"#,
] {
assert!(parse(response, 1).is_err(), "{response}");
}
}
#[test]
fn rejects_missing_unknown_extra_and_mixed_top_level_fields() {
for value in [
json!({}),
json!({"status": "scored"}),
json!({"status": "scored", "recordingQuality": 50}),
json!({"status": "scored", "speakers": [speaker(0, 1)]}),
json!({"status": "unscorable"}),
json!({"status": "unscorable", "reason": "x", "extra": 1}),
json!({"status": "scored", "recordingQuality": 50, "speakers": [speaker(0, 1)], "extra": 1}),
json!({"status": "scored", "recordingQuality": 50, "speakers": [speaker(0, 1)], "reason": "x"}),
json!({"status": "unscorable", "reason": "x", "recordingQuality": 50}),
json!({"status": "unscorable", "reason": "x", "speakers": []}),
] {
assert_rejected(value, 1);
}
}
#[test]
fn rejects_duplicate_json_fields() {
for response in [
r#"{"status":"unscorable","reason":"a","reason":"b"}"#,
r#"{"status":"unscorable","status":"unscorable","reason":"a"}"#,
r#"{"status":"scored","recordingQuality":50,"recordingQuality":51,"speakers":[]}"#,
r#"{"status":"scored","recordingQuality":50,"speakers":[{"speakerOrdinal":0,"speakerOrdinal":0,"primaryLanguage":"eng","closestDialect":"x","usableSpeechMs":1,"features":[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]}]}"#,
] {
assert!(parse(response, 1).is_err(), "{response}");
}
}
#[test]
fn rejects_missing_and_extra_speaker_fields() {
for field in [
"speakerOrdinal",
"primaryLanguage",
"closestDialect",
"usableSpeechMs",
"features",
] {
let mut value = scored(vec![speaker(0, 1)]);
value["speakers"][0].as_object_mut().unwrap().remove(field);
assert_rejected(value, 1);
}
let mut extra = scored(vec![speaker(0, 1)]);
extra["speakers"][0]["confidence"] = json!(50);
assert_rejected(extra, 1);
}
#[test]
fn rejects_wrong_top_level_types() {
for quality in [json!(null), json!("50"), json!(50.5), json!(-1)] {
let mut value = scored(vec![speaker(0, 1)]);
value["recordingQuality"] = quality;
assert_rejected(value, 1);
}
for speakers in [json!(null), json!("speaker"), json!({}), json!(50)] {
let mut value = scored(vec![speaker(0, 1)]);
value["speakers"] = speakers;
assert_rejected(value, 1);
}
for reason in [json!(null), json!(50), json!([]), json!({})] {
assert_rejected(json!({"status": "unscorable", "reason": reason}), 1);
}
}
#[test]
fn rejects_wrong_speaker_field_types() {
let cases = [
("speakerOrdinal", json!(null)),
("speakerOrdinal", json!("0")),
("speakerOrdinal", json!(0.5)),
("speakerOrdinal", json!(-1)),
("primaryLanguage", json!(null)),
("primaryLanguage", json!(1)),
("closestDialect", json!(null)),
("closestDialect", json!(1)),
("usableSpeechMs", json!(null)),
("usableSpeechMs", json!("1")),
("usableSpeechMs", json!(1.5)),
("usableSpeechMs", json!(-1)),
("features", json!(null)),
("features", json!("values")),
("features", json!({})),
];
for (field, replacement) in cases {
let mut value = scored(vec![speaker(0, 1)]);
value["speakers"][0][field] = replacement;
assert_rejected(value, 1);
}
}
#[test]
fn rejects_non_integer_partial_and_abstaining_features() {
for replacement in [
json!(null),
json!("50"),
json!(50.5),
json!(-1),
json!({"value": 50, "confidence": 80}),
json!([50, 60]),
] {
let mut value = scored(vec![speaker(0, 1)]);
value["speakers"][0]["features"][4] = replacement;
assert_rejected(value, 1);
}
let raw = format!(
r#"{{"status":"scored","recordingQuality":50,"speakers":[{{"speakerOrdinal":0,"primaryLanguage":"eng","closestDialect":"x","usableSpeechMs":1,"features":[NaN,{}]}}]}}"#,
vec!["0"; FEATURE_COUNT - 1].join(",")
);
assert!(parse(&raw, 1).is_err());
}
#[test]
fn rejects_out_of_representation_numeric_values() {
let responses = [
r#"{"status":"scored","recordingQuality":256,"speakers":[]}"#.to_owned(),
format!(
r#"{{"status":"scored","recordingQuality":50,"speakers":[{{"speakerOrdinal":65536,"primaryLanguage":"eng","closestDialect":"x","usableSpeechMs":1,"features":[{}]}}]}}"#,
vec!["0"; FEATURE_COUNT].join(",")
),
format!(
r#"{{"status":"scored","recordingQuality":50,"speakers":[{{"speakerOrdinal":0,"primaryLanguage":"eng","closestDialect":"x","usableSpeechMs":4294967296,"features":[{}]}}]}}"#,
vec!["0"; FEATURE_COUNT].join(",")
),
format!(
r#"{{"status":"scored","recordingQuality":50,"speakers":[{{"speakerOrdinal":0,"primaryLanguage":"eng","closestDialect":"x","usableSpeechMs":1,"features":[256,{}]}}]}}"#,
vec!["0"; FEATURE_COUNT - 1].join(",")
),
];
for response in responses {
assert!(parse(&response, u64::MAX).is_err(), "{response}");
}
}
}