use serde::Serialize;
use serde_json::{Map, Value, json};
use std::{collections::BTreeSet, error::Error, fmt};
pub use kcode_speaker_v3_gemini_protocol::{
FEATURE_PACKETS, FeaturePacket, GEMINI_FEATURE_PROMPT_ONE, GEMINI_FEATURE_PROMPT_ONE_REVISION,
GEMINI_FEATURE_PROMPT_REVISIONS, GEMINI_FEATURE_PROMPT_THREE,
GEMINI_FEATURE_PROMPT_THREE_REVISION, GEMINI_FEATURE_PROMPT_TWO,
GEMINI_FEATURE_PROMPT_TWO_REVISION, GEMINI_TRANSCRIPT_PROMPT,
GEMINI_TRANSCRIPT_PROMPT_REVISION, GeminiRequestPart, gemini_feature_cached_prefix,
gemini_feature_suffix, gemini_transcript_request,
};
pub use kcode_speaker_v3_schema::{
FEATURE_NAMES, FeatureVector24, LocalSpeakerLabel, StructuredAnalysis, StructuredSpeaker,
ValidationError as AnalysisValidationError, VocalGenderPresentation,
};
pub use kcode_speaker_v3_terra_labels_protocol::{
RECORD_SPEAKER_LABELS_TOOL_DESCRIPTION, RECORD_SPEAKER_LABELS_TOOL_NAME,
TERRA_SPEAKER_LABELS_PROMPT, TERRA_SPEAKER_LABELS_PROMPT_REVISION,
};
use kcode_speaker_v3_gemini_protocol::{
GeminiProtocolError, extract_gemini_text as extract_gemini_text_leaf,
};
use kcode_speaker_v3_terra_labels_protocol::{
TerraLabelsProtocolError, TerraSpeakerLabelsInput as TerraSpeakerLabelsInputLeaf,
decode_record_speaker_labels_arguments as decode_record_speaker_labels_arguments_leaf,
record_speaker_labels_tool as record_speaker_labels_tool_leaf,
};
pub const GPT_STRUCTURING_PROMPT_REVISION: &str = "speaker-v3-gpt-structure-r2";
pub const GPT_STRUCTURING_PROMPT: &str = include_str!("gpt-structuring-prompt.txt");
pub const RECORD_SPEAKER_ANALYSIS_TOOL_NAME: &str = "record_speaker_analysis";
pub const RECORD_SPEAKER_ANALYSIS_TOOL_DESCRIPTION: &str = "Record the exact transcript and complete structured 24-feature analysis for every local speaker.";
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProtocolError {
Blank(&'static str),
InvalidGeminiResponse(String),
GeminiTextCandidateCount(usize),
InvalidSpeakerLabelsArguments(String),
DuplicateSpeakerLabel(LocalSpeakerLabel),
InvalidFinalArguments(String),
InvalidStructuredAnalysis(AnalysisValidationError),
}
impl fmt::Display for ProtocolError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Blank(field) => write!(formatter, "{field} is blank"),
Self::InvalidGeminiResponse(message) => {
write!(formatter, "invalid Gemini response: {message}")
}
Self::GeminiTextCandidateCount(count) => {
write!(
formatter,
"Gemini returned {count} nonblank textual candidates"
)
}
Self::InvalidSpeakerLabelsArguments(message) => {
write!(formatter, "invalid speaker-label arguments: {message}")
}
Self::DuplicateSpeakerLabel(label) => {
write!(formatter, "duplicate speaker label: {label}")
}
Self::InvalidFinalArguments(message) => {
write!(formatter, "invalid final analysis arguments: {message}")
}
Self::InvalidStructuredAnalysis(error) => {
write!(formatter, "invalid structured analysis: {error}")
}
}
}
}
impl Error for ProtocolError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
Self::InvalidStructuredAnalysis(error) => Some(error),
_ => None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolDefinition {
pub name: &'static str,
pub description: &'static str,
pub input_schema: Value,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(transparent)]
pub struct TerraSpeakerLabelsInput {
inner: TerraSpeakerLabelsInputLeaf,
}
impl TerraSpeakerLabelsInput {
pub fn new(transcript: String) -> Result<Self, ProtocolError> {
TerraSpeakerLabelsInputLeaf::new(transcript)
.map(|inner| Self { inner })
.map_err(map_terra_labels_error)
}
pub fn instruction(&self) -> &'static str {
self.inner.instruction()
}
pub fn transcript(&self) -> &str {
self.inner.transcript()
}
pub fn render(&self) -> String {
self.inner.render()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct SpeakerFeatureEvidence {
speaker: LocalSpeakerLabel,
feature_packets: [String; 3],
}
impl SpeakerFeatureEvidence {
pub fn new(
speaker: LocalSpeakerLabel,
packet_one: String,
packet_two: String,
packet_three: String,
) -> Result<Self, ProtocolError> {
require_nonblank(&packet_one, "feature_packet_one")?;
require_nonblank(&packet_two, "feature_packet_two")?;
require_nonblank(&packet_three, "feature_packet_three")?;
Ok(Self {
speaker,
feature_packets: [packet_one, packet_two, packet_three],
})
}
pub fn speaker(&self) -> LocalSpeakerLabel {
self.speaker
}
pub fn packet(&self, packet: FeaturePacket) -> &str {
&self.feature_packets[packet.index()]
}
pub fn packets(&self) -> [&str; 3] {
self.feature_packets.each_ref().map(String::as_str)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct TerraFinalInput {
transcript: String,
speakers: Vec<SpeakerFeatureEvidence>,
}
impl TerraFinalInput {
pub fn new(
transcript: String,
speakers: Vec<SpeakerFeatureEvidence>,
) -> Result<Self, ProtocolError> {
require_nonblank(&transcript, "transcript")?;
let mut labels = BTreeSet::new();
for speaker in &speakers {
if !labels.insert(speaker.speaker()) {
return Err(ProtocolError::DuplicateSpeakerLabel(speaker.speaker()));
}
}
Ok(Self {
transcript,
speakers,
})
}
pub fn instruction(&self) -> &'static str {
GPT_STRUCTURING_PROMPT
}
pub fn transcript(&self) -> &str {
&self.transcript
}
pub fn speakers(&self) -> &[SpeakerFeatureEvidence] {
&self.speakers
}
pub fn render(&self) -> String {
format!(
"{}Input:\n{}",
self.instruction(),
serde_json::to_string(self).expect("serializing strings cannot fail")
)
}
}
pub fn extract_gemini_text(response: &Value) -> Result<String, ProtocolError> {
extract_gemini_text_leaf(response).map_err(|error| match error {
GeminiProtocolError::InvalidResponse(message) => {
ProtocolError::InvalidGeminiResponse(message)
}
GeminiProtocolError::TextCandidateCount(count) => {
ProtocolError::GeminiTextCandidateCount(count)
}
})
}
pub fn record_speaker_labels_tool() -> ToolDefinition {
let tool = record_speaker_labels_tool_leaf();
ToolDefinition {
name: tool.name,
description: tool.description,
input_schema: tool.input_schema,
}
}
pub fn decode_record_speaker_labels_arguments(
arguments: &Value,
) -> Result<Vec<LocalSpeakerLabel>, ProtocolError> {
decode_record_speaker_labels_arguments_leaf(arguments).map_err(map_terra_labels_error)
}
pub fn record_speaker_analysis_tool() -> ToolDefinition {
let mut feature_properties = Map::new();
for name in FEATURE_NAMES {
feature_properties.insert(name.into(), feature_property(name));
}
ToolDefinition {
name: RECORD_SPEAKER_ANALYSIS_TOOL_NAME,
description: RECORD_SPEAKER_ANALYSIS_TOOL_DESCRIPTION,
input_schema: json!({
"type": "object",
"additionalProperties": false,
"required": ["transcript", "speakers"],
"properties": {
"transcript": {
"type": "string"
},
"speakers": {
"type": "array",
"items": {
"type": "object",
"additionalProperties": false,
"required": [
"speaker",
"language",
"features",
"features_usable_for_training"
],
"properties": {
"speaker": {
"type": "string",
"pattern": "^Speaker [1-9][0-9]*$"
},
"language": {
"type": "string"
},
"features": {
"type": "object",
"additionalProperties": false,
"required": FEATURE_NAMES,
"properties": feature_properties
},
"features_usable_for_training": {
"type": "boolean"
}
}
}
}
}
}),
}
}
pub fn decode_record_speaker_analysis_arguments(
arguments: &Value,
) -> Result<StructuredAnalysis, ProtocolError> {
validate_final_shape(arguments)?;
let analysis: StructuredAnalysis = serde_json::from_value(arguments.clone())
.map_err(|error| ProtocolError::InvalidFinalArguments(error.to_string()))?;
analysis
.validate()
.map_err(ProtocolError::InvalidStructuredAnalysis)?;
Ok(analysis)
}
fn map_terra_labels_error(error: TerraLabelsProtocolError) -> ProtocolError {
match error {
TerraLabelsProtocolError::Blank(field) => ProtocolError::Blank(field),
TerraLabelsProtocolError::InvalidArguments(message) => {
ProtocolError::InvalidSpeakerLabelsArguments(message)
}
TerraLabelsProtocolError::DuplicateSpeakerLabel(label) => {
ProtocolError::DuplicateSpeakerLabel(label)
}
}
}
fn require_nonblank(value: &str, field: &'static str) -> Result<(), ProtocolError> {
if value.trim().is_empty() {
return Err(ProtocolError::Blank(field));
}
Ok(())
}
fn parse_canonical_label(value: &str) -> Result<LocalSpeakerLabel, String> {
let label = value
.parse::<LocalSpeakerLabel>()
.map_err(|error| error.to_string())?;
if label.to_string() != value {
return Err(format!("noncanonical speaker label: {value}"));
}
Ok(label)
}
fn require_exact_keys(object: &Map<String, Value>, expected: &[&str]) -> Result<(), String> {
if object.len() != expected.len() || expected.iter().any(|key| !object.contains_key(*key)) {
let found = object.keys().cloned().collect::<Vec<_>>().join(", ");
return Err(format!(
"object keys must be exactly [{}], found [{found}]",
expected.join(", ")
));
}
Ok(())
}
fn feature_property(name: &str) -> Value {
match name {
"dominant_rhotic_realization" | "dominant_lateral_realization" => {
json!({ "type": ["string", "null"] })
}
"vocal_gender_presentation" => json!({
"type": ["string", "null"],
"enum": [
"strongly_feminine",
"feminine",
"androgynous",
"masculine",
"strongly_masculine",
null
]
}),
_ => json!({ "type": ["number", "null"] }),
}
}
fn validate_final_shape(arguments: &Value) -> Result<(), ProtocolError> {
let object = arguments
.as_object()
.ok_or_else(|| ProtocolError::InvalidFinalArguments("arguments is not an object".into()))?;
require_exact_keys(object, &["transcript", "speakers"])
.map_err(ProtocolError::InvalidFinalArguments)?;
let speakers = object
.get("speakers")
.and_then(Value::as_array)
.ok_or_else(|| ProtocolError::InvalidFinalArguments("speakers is not an array".into()))?;
for (speaker_index, speaker) in speakers.iter().enumerate() {
let speaker_object = speaker.as_object().ok_or_else(|| {
ProtocolError::InvalidFinalArguments(format!(
"speaker {speaker_index} is not an object"
))
})?;
require_exact_keys(
speaker_object,
&[
"speaker",
"language",
"features",
"features_usable_for_training",
],
)
.map_err(|message| {
ProtocolError::InvalidFinalArguments(format!("speaker {speaker_index}: {message}"))
})?;
let raw_label = speaker_object
.get("speaker")
.and_then(Value::as_str)
.ok_or_else(|| {
ProtocolError::InvalidFinalArguments(format!(
"speaker {speaker_index} label is not a string"
))
})?;
parse_canonical_label(raw_label).map_err(|message| {
ProtocolError::InvalidFinalArguments(format!("speaker {speaker_index}: {message}"))
})?;
let features = speaker_object
.get("features")
.and_then(Value::as_object)
.ok_or_else(|| {
ProtocolError::InvalidFinalArguments(format!(
"speaker {speaker_index} features is not an object"
))
})?;
require_exact_keys(features, &FEATURE_NAMES).map_err(|message| {
ProtocolError::InvalidFinalArguments(format!(
"speaker {speaker_index} features: {message}"
))
})?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Instant;
fn label(number: u32) -> LocalSpeakerLabel {
LocalSpeakerLabel::new(number).unwrap()
}
fn null_features() -> Map<String, Value> {
FEATURE_NAMES
.into_iter()
.map(|name| (name.into(), Value::Null))
.collect()
}
fn final_arguments() -> Value {
json!({
"transcript": "[high] Speaker 1: exact words",
"speakers": [{
"speaker": "Speaker 1",
"language": "English",
"features": null_features(),
"features_usable_for_training": true
}]
})
}
#[test]
fn compatibility_facade_preserves_protocol_errors_and_inputs() {
assert_eq!(
extract_gemini_text(&json!({})),
Err(ProtocolError::InvalidGeminiResponse(
"candidates is not an array".into()
))
);
assert_eq!(
extract_gemini_text(&json!({ "candidates": [] })),
Err(ProtocolError::GeminiTextCandidateCount(0))
);
let duplicate = json!({ "speakers": ["Speaker 1", "Speaker 1"] });
assert_eq!(
decode_record_speaker_labels_arguments(&duplicate),
Err(ProtocolError::DuplicateSpeakerLabel(label(1)))
);
let input = TerraSpeakerLabelsInput::new("x\n\"Input:\\n\" {{raw}}".into()).unwrap();
let serialized = input
.render()
.strip_prefix(TERRA_SPEAKER_LABELS_PROMPT)
.unwrap()
.strip_prefix("Input:\n")
.unwrap()
.to_owned();
let decoded: Value = serde_json::from_str(&serialized).unwrap();
assert_eq!(decoded["transcript"], input.transcript());
}
#[test]
fn evidence_and_final_input_preserve_adversarial_text_and_order() {
let first = SpeakerFeatureEvidence::new(
label(2),
"one\n\"x\"".into(),
"two {{raw}}".into(),
"three Input:\n".into(),
)
.unwrap();
let second =
SpeakerFeatureEvidence::new(label(1), "four".into(), "five".into(), "six".into())
.unwrap();
let input = TerraFinalInput::new(
"transcript\n\"quoted\"".into(),
vec![first.clone(), second.clone()],
)
.unwrap();
assert_eq!(input.speakers()[0].speaker(), label(2));
assert_eq!(input.speakers()[1].speaker(), label(1));
let serialized = input
.render()
.strip_prefix(GPT_STRUCTURING_PROMPT)
.unwrap()
.strip_prefix("Input:\n")
.unwrap()
.to_owned();
let decoded: Value = serde_json::from_str(&serialized).unwrap();
assert_eq!(decoded["transcript"], input.transcript());
assert_eq!(decoded["speakers"][0]["speaker"], "Speaker 2");
assert_eq!(
decoded["speakers"][0]["feature_packets"][0],
first.packet(FeaturePacket::One)
);
assert!(TerraFinalInput::new("x".into(), vec![first.clone(), first]).is_err());
}
#[test]
fn final_tool_schema_is_closed_complete_and_nullable() {
let tool = record_speaker_analysis_tool();
assert_eq!(tool.name, "record_speaker_analysis");
assert_eq!(tool.input_schema["additionalProperties"], false);
let speaker = &tool.input_schema["properties"]["speakers"]["items"];
assert_eq!(speaker["additionalProperties"], false);
let features = &speaker["properties"]["features"];
assert_eq!(features["additionalProperties"], false);
assert_eq!(features["required"].as_array().unwrap().len(), 24);
assert_eq!(features["properties"].as_object().unwrap().len(), 24);
for name in FEATURE_NAMES {
assert!(features["properties"][name].to_string().contains("null"));
}
}
#[test]
fn final_decode_requires_complete_closed_valid_analysis() {
let arguments = final_arguments();
let analysis = decode_record_speaker_analysis_arguments(&arguments).unwrap();
assert_eq!(analysis.transcript, "[high] Speaker 1: exact words");
assert_eq!(analysis.speakers.len(), 1);
assert!(analysis.speakers[0].features_usable_for_training);
let mut extra = arguments.clone();
extra
.as_object_mut()
.unwrap()
.insert("extra".into(), json!(true));
assert!(decode_record_speaker_analysis_arguments(&extra).is_err());
let mut missing = arguments.clone();
missing["speakers"][0]["features"]
.as_object_mut()
.unwrap()
.remove("median_f0_hz");
assert!(decode_record_speaker_analysis_arguments(&missing).is_err());
let mut nested_extra = arguments.clone();
nested_extra["speakers"][0]["features"]
.as_object_mut()
.unwrap()
.insert("extra".into(), Value::Null);
assert!(decode_record_speaker_analysis_arguments(&nested_extra).is_err());
let mut noncanonical = arguments;
noncanonical["speakers"][0]["speaker"] = json!("Speaker 01");
assert!(decode_record_speaker_analysis_arguments(&noncanonical).is_err());
}
#[test]
fn reference_scale_canary_completes_local_work() {
let started = Instant::now();
let speakers = (1..=1000)
.map(|number| {
SpeakerFeatureEvidence::new(
label(number),
"a".repeat(1024),
"b".repeat(1024),
"c".repeat(1024),
)
.unwrap()
})
.collect();
let input = TerraFinalInput::new("x".repeat(1_048_576), speakers).unwrap();
assert!(input.render().len() > 4_000_000);
assert!(started.elapsed().as_secs() < 10);
}
}