1use serde::Serialize;
2use serde_json::{Map, Value, json};
3use std::{collections::BTreeSet, error::Error, fmt};
4
5pub use kcode_speaker_v3_gemini_protocol::{
6 FEATURE_PACKETS, FeaturePacket, GEMINI_FEATURE_PROMPT_ONE, GEMINI_FEATURE_PROMPT_ONE_REVISION,
7 GEMINI_FEATURE_PROMPT_REVISIONS, GEMINI_FEATURE_PROMPT_THREE,
8 GEMINI_FEATURE_PROMPT_THREE_REVISION, GEMINI_FEATURE_PROMPT_TWO,
9 GEMINI_FEATURE_PROMPT_TWO_REVISION, GEMINI_TRANSCRIPT_PROMPT,
10 GEMINI_TRANSCRIPT_PROMPT_REVISION, GeminiRequestPart, gemini_feature_cached_prefix,
11 gemini_feature_suffix, gemini_transcript_request,
12};
13pub use kcode_speaker_v3_schema::{
14 FEATURE_NAMES, FeatureVector24, LocalSpeakerLabel, StructuredAnalysis, StructuredSpeaker,
15 ValidationError as AnalysisValidationError, VocalGenderPresentation,
16};
17pub use kcode_speaker_v3_terra_labels_protocol::{
18 RECORD_SPEAKER_LABELS_TOOL_DESCRIPTION, RECORD_SPEAKER_LABELS_TOOL_NAME,
19 TERRA_SPEAKER_LABELS_PROMPT, TERRA_SPEAKER_LABELS_PROMPT_REVISION,
20};
21
22use kcode_speaker_v3_gemini_protocol::{
23 GeminiProtocolError, extract_gemini_text as extract_gemini_text_leaf,
24};
25use kcode_speaker_v3_terra_labels_protocol::{
26 TerraLabelsProtocolError, TerraSpeakerLabelsInput as TerraSpeakerLabelsInputLeaf,
27 decode_record_speaker_labels_arguments as decode_record_speaker_labels_arguments_leaf,
28 record_speaker_labels_tool as record_speaker_labels_tool_leaf,
29};
30
31pub const GPT_STRUCTURING_PROMPT_REVISION: &str = "speaker-v3-gpt-structure-r2";
32pub const GPT_STRUCTURING_PROMPT: &str = include_str!("gpt-structuring-prompt.txt");
33
34pub const RECORD_SPEAKER_ANALYSIS_TOOL_NAME: &str = "record_speaker_analysis";
35pub const RECORD_SPEAKER_ANALYSIS_TOOL_DESCRIPTION: &str = "Record the exact transcript and complete structured 24-feature analysis for every local speaker.";
36
37#[derive(Debug, Clone, PartialEq, Eq)]
38pub enum ProtocolError {
39 Blank(&'static str),
40 InvalidGeminiResponse(String),
41 GeminiTextCandidateCount(usize),
42 InvalidSpeakerLabelsArguments(String),
43 DuplicateSpeakerLabel(LocalSpeakerLabel),
44 InvalidFinalArguments(String),
45 InvalidStructuredAnalysis(AnalysisValidationError),
46}
47
48impl fmt::Display for ProtocolError {
49 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
50 match self {
51 Self::Blank(field) => write!(formatter, "{field} is blank"),
52 Self::InvalidGeminiResponse(message) => {
53 write!(formatter, "invalid Gemini response: {message}")
54 }
55 Self::GeminiTextCandidateCount(count) => {
56 write!(
57 formatter,
58 "Gemini returned {count} nonblank textual candidates"
59 )
60 }
61 Self::InvalidSpeakerLabelsArguments(message) => {
62 write!(formatter, "invalid speaker-label arguments: {message}")
63 }
64 Self::DuplicateSpeakerLabel(label) => {
65 write!(formatter, "duplicate speaker label: {label}")
66 }
67 Self::InvalidFinalArguments(message) => {
68 write!(formatter, "invalid final analysis arguments: {message}")
69 }
70 Self::InvalidStructuredAnalysis(error) => {
71 write!(formatter, "invalid structured analysis: {error}")
72 }
73 }
74 }
75}
76
77impl Error for ProtocolError {
78 fn source(&self) -> Option<&(dyn Error + 'static)> {
79 match self {
80 Self::InvalidStructuredAnalysis(error) => Some(error),
81 _ => None,
82 }
83 }
84}
85
86#[derive(Debug, Clone, PartialEq, Eq)]
87pub struct ToolDefinition {
88 pub name: &'static str,
89 pub description: &'static str,
90 pub input_schema: Value,
91}
92
93#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
94#[serde(transparent)]
95pub struct TerraSpeakerLabelsInput {
96 inner: TerraSpeakerLabelsInputLeaf,
97}
98
99impl TerraSpeakerLabelsInput {
100 pub fn new(transcript: String) -> Result<Self, ProtocolError> {
101 TerraSpeakerLabelsInputLeaf::new(transcript)
102 .map(|inner| Self { inner })
103 .map_err(map_terra_labels_error)
104 }
105
106 pub fn instruction(&self) -> &'static str {
107 self.inner.instruction()
108 }
109
110 pub fn transcript(&self) -> &str {
111 self.inner.transcript()
112 }
113
114 pub fn render(&self) -> String {
115 self.inner.render()
116 }
117}
118
119#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
120pub struct SpeakerFeatureEvidence {
121 speaker: LocalSpeakerLabel,
122 feature_packets: [String; 3],
123}
124
125impl SpeakerFeatureEvidence {
126 pub fn new(
127 speaker: LocalSpeakerLabel,
128 packet_one: String,
129 packet_two: String,
130 packet_three: String,
131 ) -> Result<Self, ProtocolError> {
132 require_nonblank(&packet_one, "feature_packet_one")?;
133 require_nonblank(&packet_two, "feature_packet_two")?;
134 require_nonblank(&packet_three, "feature_packet_three")?;
135 Ok(Self {
136 speaker,
137 feature_packets: [packet_one, packet_two, packet_three],
138 })
139 }
140
141 pub fn speaker(&self) -> LocalSpeakerLabel {
142 self.speaker
143 }
144
145 pub fn packet(&self, packet: FeaturePacket) -> &str {
146 &self.feature_packets[packet.index()]
147 }
148
149 pub fn packets(&self) -> [&str; 3] {
150 self.feature_packets.each_ref().map(String::as_str)
151 }
152}
153
154#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
155pub struct TerraFinalInput {
156 transcript: String,
157 speakers: Vec<SpeakerFeatureEvidence>,
158}
159
160impl TerraFinalInput {
161 pub fn new(
162 transcript: String,
163 speakers: Vec<SpeakerFeatureEvidence>,
164 ) -> Result<Self, ProtocolError> {
165 require_nonblank(&transcript, "transcript")?;
166 let mut labels = BTreeSet::new();
167 for speaker in &speakers {
168 if !labels.insert(speaker.speaker()) {
169 return Err(ProtocolError::DuplicateSpeakerLabel(speaker.speaker()));
170 }
171 }
172 Ok(Self {
173 transcript,
174 speakers,
175 })
176 }
177
178 pub fn instruction(&self) -> &'static str {
179 GPT_STRUCTURING_PROMPT
180 }
181
182 pub fn transcript(&self) -> &str {
183 &self.transcript
184 }
185
186 pub fn speakers(&self) -> &[SpeakerFeatureEvidence] {
187 &self.speakers
188 }
189
190 pub fn render(&self) -> String {
191 format!(
192 "{}Input:\n{}",
193 self.instruction(),
194 serde_json::to_string(self).expect("serializing strings cannot fail")
195 )
196 }
197}
198
199pub fn extract_gemini_text(response: &Value) -> Result<String, ProtocolError> {
200 extract_gemini_text_leaf(response).map_err(|error| match error {
201 GeminiProtocolError::InvalidResponse(message) => {
202 ProtocolError::InvalidGeminiResponse(message)
203 }
204 GeminiProtocolError::TextCandidateCount(count) => {
205 ProtocolError::GeminiTextCandidateCount(count)
206 }
207 })
208}
209
210pub fn record_speaker_labels_tool() -> ToolDefinition {
211 let tool = record_speaker_labels_tool_leaf();
212 ToolDefinition {
213 name: tool.name,
214 description: tool.description,
215 input_schema: tool.input_schema,
216 }
217}
218
219pub fn decode_record_speaker_labels_arguments(
220 arguments: &Value,
221) -> Result<Vec<LocalSpeakerLabel>, ProtocolError> {
222 decode_record_speaker_labels_arguments_leaf(arguments).map_err(map_terra_labels_error)
223}
224
225pub fn record_speaker_analysis_tool() -> ToolDefinition {
226 let mut feature_properties = Map::new();
227 for name in FEATURE_NAMES {
228 feature_properties.insert(name.into(), feature_property(name));
229 }
230 ToolDefinition {
231 name: RECORD_SPEAKER_ANALYSIS_TOOL_NAME,
232 description: RECORD_SPEAKER_ANALYSIS_TOOL_DESCRIPTION,
233 input_schema: json!({
234 "type": "object",
235 "additionalProperties": false,
236 "required": ["transcript", "speakers"],
237 "properties": {
238 "transcript": {
239 "type": "string"
240 },
241 "speakers": {
242 "type": "array",
243 "items": {
244 "type": "object",
245 "additionalProperties": false,
246 "required": [
247 "speaker",
248 "language",
249 "features",
250 "features_usable_for_training"
251 ],
252 "properties": {
253 "speaker": {
254 "type": "string",
255 "pattern": "^Speaker [1-9][0-9]*$"
256 },
257 "language": {
258 "type": "string"
259 },
260 "features": {
261 "type": "object",
262 "additionalProperties": false,
263 "required": FEATURE_NAMES,
264 "properties": feature_properties
265 },
266 "features_usable_for_training": {
267 "type": "boolean"
268 }
269 }
270 }
271 }
272 }
273 }),
274 }
275}
276
277pub fn decode_record_speaker_analysis_arguments(
278 arguments: &Value,
279) -> Result<StructuredAnalysis, ProtocolError> {
280 validate_final_shape(arguments)?;
281 let analysis: StructuredAnalysis = serde_json::from_value(arguments.clone())
282 .map_err(|error| ProtocolError::InvalidFinalArguments(error.to_string()))?;
283 analysis
284 .validate()
285 .map_err(ProtocolError::InvalidStructuredAnalysis)?;
286 Ok(analysis)
287}
288
289fn map_terra_labels_error(error: TerraLabelsProtocolError) -> ProtocolError {
290 match error {
291 TerraLabelsProtocolError::Blank(field) => ProtocolError::Blank(field),
292 TerraLabelsProtocolError::InvalidArguments(message) => {
293 ProtocolError::InvalidSpeakerLabelsArguments(message)
294 }
295 TerraLabelsProtocolError::DuplicateSpeakerLabel(label) => {
296 ProtocolError::DuplicateSpeakerLabel(label)
297 }
298 }
299}
300
301fn require_nonblank(value: &str, field: &'static str) -> Result<(), ProtocolError> {
302 if value.trim().is_empty() {
303 return Err(ProtocolError::Blank(field));
304 }
305 Ok(())
306}
307
308fn parse_canonical_label(value: &str) -> Result<LocalSpeakerLabel, String> {
309 let label = value
310 .parse::<LocalSpeakerLabel>()
311 .map_err(|error| error.to_string())?;
312 if label.to_string() != value {
313 return Err(format!("noncanonical speaker label: {value}"));
314 }
315 Ok(label)
316}
317
318fn require_exact_keys(object: &Map<String, Value>, expected: &[&str]) -> Result<(), String> {
319 if object.len() != expected.len() || expected.iter().any(|key| !object.contains_key(*key)) {
320 let found = object.keys().cloned().collect::<Vec<_>>().join(", ");
321 return Err(format!(
322 "object keys must be exactly [{}], found [{found}]",
323 expected.join(", ")
324 ));
325 }
326 Ok(())
327}
328
329fn feature_property(name: &str) -> Value {
330 match name {
331 "dominant_rhotic_realization" | "dominant_lateral_realization" => {
332 json!({ "type": ["string", "null"] })
333 }
334 "vocal_gender_presentation" => json!({
335 "type": ["string", "null"],
336 "enum": [
337 "strongly_feminine",
338 "feminine",
339 "androgynous",
340 "masculine",
341 "strongly_masculine",
342 null
343 ]
344 }),
345 _ => json!({ "type": ["number", "null"] }),
346 }
347}
348
349fn validate_final_shape(arguments: &Value) -> Result<(), ProtocolError> {
350 let object = arguments
351 .as_object()
352 .ok_or_else(|| ProtocolError::InvalidFinalArguments("arguments is not an object".into()))?;
353 require_exact_keys(object, &["transcript", "speakers"])
354 .map_err(ProtocolError::InvalidFinalArguments)?;
355 let speakers = object
356 .get("speakers")
357 .and_then(Value::as_array)
358 .ok_or_else(|| ProtocolError::InvalidFinalArguments("speakers is not an array".into()))?;
359 for (speaker_index, speaker) in speakers.iter().enumerate() {
360 let speaker_object = speaker.as_object().ok_or_else(|| {
361 ProtocolError::InvalidFinalArguments(format!(
362 "speaker {speaker_index} is not an object"
363 ))
364 })?;
365 require_exact_keys(
366 speaker_object,
367 &[
368 "speaker",
369 "language",
370 "features",
371 "features_usable_for_training",
372 ],
373 )
374 .map_err(|message| {
375 ProtocolError::InvalidFinalArguments(format!("speaker {speaker_index}: {message}"))
376 })?;
377 let raw_label = speaker_object
378 .get("speaker")
379 .and_then(Value::as_str)
380 .ok_or_else(|| {
381 ProtocolError::InvalidFinalArguments(format!(
382 "speaker {speaker_index} label is not a string"
383 ))
384 })?;
385 parse_canonical_label(raw_label).map_err(|message| {
386 ProtocolError::InvalidFinalArguments(format!("speaker {speaker_index}: {message}"))
387 })?;
388 let features = speaker_object
389 .get("features")
390 .and_then(Value::as_object)
391 .ok_or_else(|| {
392 ProtocolError::InvalidFinalArguments(format!(
393 "speaker {speaker_index} features is not an object"
394 ))
395 })?;
396 require_exact_keys(features, &FEATURE_NAMES).map_err(|message| {
397 ProtocolError::InvalidFinalArguments(format!(
398 "speaker {speaker_index} features: {message}"
399 ))
400 })?;
401 }
402 Ok(())
403}
404
405#[cfg(test)]
406mod tests {
407 use super::*;
408 use std::time::Instant;
409
410 fn label(number: u32) -> LocalSpeakerLabel {
411 LocalSpeakerLabel::new(number).unwrap()
412 }
413
414 fn null_features() -> Map<String, Value> {
415 FEATURE_NAMES
416 .into_iter()
417 .map(|name| (name.into(), Value::Null))
418 .collect()
419 }
420
421 fn final_arguments() -> Value {
422 json!({
423 "transcript": "[high] Speaker 1: exact words",
424 "speakers": [{
425 "speaker": "Speaker 1",
426 "language": "English",
427 "features": null_features(),
428 "features_usable_for_training": true
429 }]
430 })
431 }
432
433 #[test]
434 fn compatibility_facade_preserves_protocol_errors_and_inputs() {
435 assert_eq!(
436 extract_gemini_text(&json!({})),
437 Err(ProtocolError::InvalidGeminiResponse(
438 "candidates is not an array".into()
439 ))
440 );
441 assert_eq!(
442 extract_gemini_text(&json!({ "candidates": [] })),
443 Err(ProtocolError::GeminiTextCandidateCount(0))
444 );
445 let duplicate = json!({ "speakers": ["Speaker 1", "Speaker 1"] });
446 assert_eq!(
447 decode_record_speaker_labels_arguments(&duplicate),
448 Err(ProtocolError::DuplicateSpeakerLabel(label(1)))
449 );
450 let input = TerraSpeakerLabelsInput::new("x\n\"Input:\\n\" {{raw}}".into()).unwrap();
451 let serialized = input
452 .render()
453 .strip_prefix(TERRA_SPEAKER_LABELS_PROMPT)
454 .unwrap()
455 .strip_prefix("Input:\n")
456 .unwrap()
457 .to_owned();
458 let decoded: Value = serde_json::from_str(&serialized).unwrap();
459 assert_eq!(decoded["transcript"], input.transcript());
460 }
461
462 #[test]
463 fn evidence_and_final_input_preserve_adversarial_text_and_order() {
464 let first = SpeakerFeatureEvidence::new(
465 label(2),
466 "one\n\"x\"".into(),
467 "two {{raw}}".into(),
468 "three Input:\n".into(),
469 )
470 .unwrap();
471 let second =
472 SpeakerFeatureEvidence::new(label(1), "four".into(), "five".into(), "six".into())
473 .unwrap();
474 let input = TerraFinalInput::new(
475 "transcript\n\"quoted\"".into(),
476 vec![first.clone(), second.clone()],
477 )
478 .unwrap();
479 assert_eq!(input.speakers()[0].speaker(), label(2));
480 assert_eq!(input.speakers()[1].speaker(), label(1));
481 let serialized = input
482 .render()
483 .strip_prefix(GPT_STRUCTURING_PROMPT)
484 .unwrap()
485 .strip_prefix("Input:\n")
486 .unwrap()
487 .to_owned();
488 let decoded: Value = serde_json::from_str(&serialized).unwrap();
489 assert_eq!(decoded["transcript"], input.transcript());
490 assert_eq!(decoded["speakers"][0]["speaker"], "Speaker 2");
491 assert_eq!(
492 decoded["speakers"][0]["feature_packets"][0],
493 first.packet(FeaturePacket::One)
494 );
495 assert!(TerraFinalInput::new("x".into(), vec![first.clone(), first]).is_err());
496 }
497
498 #[test]
499 fn final_tool_schema_is_closed_complete_and_nullable() {
500 let tool = record_speaker_analysis_tool();
501 assert_eq!(tool.name, "record_speaker_analysis");
502 assert_eq!(tool.input_schema["additionalProperties"], false);
503 let speaker = &tool.input_schema["properties"]["speakers"]["items"];
504 assert_eq!(speaker["additionalProperties"], false);
505 let features = &speaker["properties"]["features"];
506 assert_eq!(features["additionalProperties"], false);
507 assert_eq!(features["required"].as_array().unwrap().len(), 24);
508 assert_eq!(features["properties"].as_object().unwrap().len(), 24);
509 for name in FEATURE_NAMES {
510 assert!(features["properties"][name].to_string().contains("null"));
511 }
512 }
513
514 #[test]
515 fn final_decode_requires_complete_closed_valid_analysis() {
516 let arguments = final_arguments();
517 let analysis = decode_record_speaker_analysis_arguments(&arguments).unwrap();
518 assert_eq!(analysis.transcript, "[high] Speaker 1: exact words");
519 assert_eq!(analysis.speakers.len(), 1);
520 assert!(analysis.speakers[0].features_usable_for_training);
521 let mut extra = arguments.clone();
522 extra
523 .as_object_mut()
524 .unwrap()
525 .insert("extra".into(), json!(true));
526 assert!(decode_record_speaker_analysis_arguments(&extra).is_err());
527 let mut missing = arguments.clone();
528 missing["speakers"][0]["features"]
529 .as_object_mut()
530 .unwrap()
531 .remove("median_f0_hz");
532 assert!(decode_record_speaker_analysis_arguments(&missing).is_err());
533 let mut nested_extra = arguments.clone();
534 nested_extra["speakers"][0]["features"]
535 .as_object_mut()
536 .unwrap()
537 .insert("extra".into(), Value::Null);
538 assert!(decode_record_speaker_analysis_arguments(&nested_extra).is_err());
539 let mut noncanonical = arguments;
540 noncanonical["speakers"][0]["speaker"] = json!("Speaker 01");
541 assert!(decode_record_speaker_analysis_arguments(&noncanonical).is_err());
542 }
543
544 #[test]
545 fn reference_scale_canary_completes_local_work() {
546 let started = Instant::now();
547 let speakers = (1..=1000)
548 .map(|number| {
549 SpeakerFeatureEvidence::new(
550 label(number),
551 "a".repeat(1024),
552 "b".repeat(1024),
553 "c".repeat(1024),
554 )
555 .unwrap()
556 })
557 .collect();
558 let input = TerraFinalInput::new("x".repeat(1_048_576), speakers).unwrap();
559 assert!(input.render().len() > 4_000_000);
560 assert!(started.elapsed().as_secs() < 10);
561 }
562}