1use std::sync::Arc;
4
5use superfluid_engine::Tokenizer;
6
7use crate::codec::{Channelizer, Run};
8use crate::wal::channel;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub struct HarmonyIds {
12 pub start: u32,
13 pub channel: u32,
14 pub message: u32,
15 pub end: u32,
16 pub ret: u32,
17 pub call: u32,
18 pub constrain: Option<u32>,
19}
20
21impl HarmonyIds {
22 pub fn discover(specials: &[(String, u32)]) -> Option<HarmonyIds> {
23 let find = |name: &str| specials.iter().find(|(s, _)| s == name).map(|(_, i)| *i);
24 Some(HarmonyIds {
25 start: find("<|start|>")?,
26 channel: find("<|channel|>")?,
27 message: find("<|message|>")?,
28 end: find("<|end|>")?,
29 ret: find("<|return|>")?,
30 call: find("<|call|>")?,
31 constrain: find("<|constrain|>"),
32 })
33 }
34
35 fn is_marker(&self, t: u32) -> bool {
36 t == self.start
37 || t == self.channel
38 || t == self.message
39 || t == self.end
40 || t == self.ret
41 || t == self.call
42 || Some(t) == self.constrain
43 }
44}
45
46fn classify(header: &str) -> u32 {
47 let h = header.trim_start();
48 if h.contains("to=") {
49 channel::TOOL_CALL
50 } else if h.starts_with("analysis") {
51 channel::REASONING
52 } else {
53 channel::TEXT
54 }
55}
56
57enum State {
58 Header,
59 Body(u32),
60}
61
62pub struct HarmonyChannelizer {
63 tok: Arc<dyn Tokenizer>,
64 ids: HarmonyIds,
65 state: State,
66 header: String,
67 header_tokens: Vec<u32>,
68 header_cut: Option<usize>,
69 recipient: Vec<u32>,
70}
71
72impl HarmonyChannelizer {
73 pub fn new(tok: Arc<dyn Tokenizer>, ids: HarmonyIds) -> HarmonyChannelizer {
74 HarmonyChannelizer {
75 tok,
76 ids,
77 state: State::Header,
78 header: String::new(),
79 header_tokens: Vec::new(),
80 header_cut: None,
81 recipient: Vec::new(),
82 }
83 }
84
85 fn reset_header(&mut self) {
86 self.header.clear();
87 self.header_tokens.clear();
88 self.header_cut = None;
89 }
90
91 fn recipient_tokens(&self) -> Vec<u32> {
92 let end = self.header_cut.unwrap_or(self.header_tokens.len());
93 let Some(pos) = self.header.find("to=") else {
94 return Vec::new();
95 };
96 let mut at = 0;
97 for (i, &t) in self.header_tokens[..end].iter().enumerate() {
98 at += String::from_utf8_lossy(&self.tok.token_bytes(t)).len();
99 if at > pos {
100 return self.header_tokens[i..end].to_vec();
101 }
102 }
103 Vec::new()
104 }
105}
106
107impl Channelizer for HarmonyChannelizer {
108 fn split(&mut self, tokens: &[u32]) -> Vec<Run> {
109 let mut out: Vec<Run> = Vec::new();
110 let mut span: Vec<u32> = Vec::new();
111 let mut text: Vec<u32> = Vec::new();
112 for &t in tokens {
113 match self.state {
114 State::Header => {
115 span.push(t);
116 if t == self.ids.message {
117 let mut header = self.header.clone();
118 if !self.recipient.is_empty() {
119 header.push(' ');
120 header.push_str(&String::from_utf8_lossy(&self.recipient.iter().flat_map(|&r| self.tok.token_bytes(r)).collect::<Vec<u8>>()));
121 }
122 let ch = classify(&header);
123 if ch == channel::TOOL_CALL {
124 let cut = self.header_cut.unwrap_or(self.header_tokens.len());
125 text.extend(self.header_tokens[..cut].iter().copied());
126 text.extend(self.recipient.iter().copied());
127 }
128 self.reset_header();
129 self.recipient.clear();
130 self.state = State::Body(ch);
131 } else if t == self.ids.start {
132 self.reset_header();
133 self.recipient.clear();
134 } else if t == self.ids.channel {
135 let carried = self.recipient_tokens();
136 self.reset_header();
137 if !carried.is_empty() {
138 self.recipient = carried;
139 }
140 } else if Some(t) == self.ids.constrain {
141 self.header_cut.get_or_insert(self.header_tokens.len());
142 } else if self.ids.is_marker(t) {
143 } else {
144 self.header.push_str(&String::from_utf8_lossy(&self.tok.token_bytes(t)));
145 self.header_tokens.push(t);
146 }
147 }
148 State::Body(ch) => {
149 if t == self.ids.end || t == self.ids.ret || t == self.ids.call {
150 span.push(t);
151 out.push(Run {
152 span: std::mem::take(&mut span),
153 text: std::mem::take(&mut text),
154 channel: ch,
155 closes: true,
156 });
157 self.state = State::Header;
158 } else if t == self.ids.start || t == self.ids.channel {
159 if !span.is_empty() {
160 out.push(Run {
161 span: std::mem::take(&mut span),
162 text: std::mem::take(&mut text),
163 channel: ch,
164 closes: false,
165 });
166 }
167 span.push(t);
168 self.reset_header();
169 self.state = State::Header;
170 } else if self.ids.is_marker(t) {
171 span.push(t);
172 } else {
173 span.push(t);
174 text.push(t);
175 }
176 }
177 }
178 }
179 if !span.is_empty() {
180 out.push(Run { span, text, channel: self.current_channel(), closes: false });
181 }
182 out
183 }
184
185 fn current_channel(&self) -> u32 {
186 match self.state {
187 State::Header => channel::TEXT,
188 State::Body(ch) => ch,
189 }
190 }
191
192 fn prime(&mut self, channel: u32) {
193 self.state = State::Body(channel);
194 }
195}
196
197pub fn parse_harmony_tool_call(raw: &str) -> Option<(String, String)> {
198 let (name, after) = recipient(raw)?;
199 let body = after.get(after.find('{')?..)?;
200 let args = crate::codec::repair_json_object(body)?;
201 Some((name, args))
202}
203
204fn recipient(raw: &str) -> Option<(String, &str)> {
205 let p = raw.find("to=")?;
206 let rest = &raw[p + 3..];
207 let len = rest.find(|c: char| c.is_whitespace() || c == '<' || c == '{').unwrap_or(rest.len());
208 let recipient = &rest[..len];
209 let name = recipient.strip_prefix("functions.").unwrap_or(recipient);
210 if name.is_empty() {
211 return None;
212 }
213 Some((name.to_string(), &rest[len..]))
214}
215
216pub fn named_call(raw: &str) -> Option<(String, String)> {
217 let (name, after) = recipient(raw)?;
218 let args = after.find('{').map_or("", |i| &after[i..]);
219 Some((name, args.to_string()))
220}
221
222fn tag_token(id: u32) -> serde_json::Value {
223 serde_json::json!({ "type": "token", "token": id })
224}
225
226fn tag_text(s: &str) -> serde_json::Value {
227 serde_json::json!({ "type": "const_string", "value": s })
228}
229
230fn markers(ids: HarmonyIds) -> Vec<u32> {
231 let mut markers = vec![ids.start, ids.channel, ids.message, ids.end, ids.ret, ids.call];
232 markers.extend(ids.constrain);
233 markers
234}
235
236fn body_then_next_message(ids: HarmonyIds) -> [serde_json::Value; 3] {
237 [
238 serde_json::json!({ "type": "tag", "begin": tag_token(ids.message),
239 "content": { "type": "any_tokens", "exclude_tokens": markers(ids) },
240 "end": tag_token(ids.end) }),
241 tag_token(ids.start),
242 tag_text("assistant"),
243 ]
244}
245
246fn optional_analysis(ids: HarmonyIds) -> serde_json::Value {
247 let [body, start, assistant] = body_then_next_message(ids);
248 serde_json::json!({ "type": "optional", "content": { "type": "sequence", "elements": [
249 tag_token(ids.channel), tag_text("analysis"), body, start, assistant,
250 ]}})
251}
252
253fn optional_content_type(ids: HarmonyIds) -> serde_json::Value {
254 let written = match ids.constrain {
255 Some(c) => serde_json::json!({ "type": "or", "elements": [
256 { "type": "sequence", "elements": [
257 { "type": "optional", "content": tag_text(" ") }, tag_token(c), tag_text("json"),
258 ]},
259 tag_text(" json"),
260 ]}),
261 None => tag_text(" json"),
262 };
263 serde_json::json!({ "type": "optional", "content": written })
264}
265
266fn structural(elements: Vec<serde_json::Value>) -> String {
267 serde_json::json!({ "type": "structural_tag", "format": { "type": "sequence", "elements": elements } }).to_string()
268}
269
270pub fn forced_call_tag(ids: HarmonyIds, tool_jsons: &[String]) -> Option<String> {
271 use serde_json::json;
272 let mut calls: Vec<serde_json::Value> = declared_functions(tool_jsons)?
273 .into_iter()
274 .map(|(name, params)| {
275 json!({ "type": "sequence", "elements": [
276 tag_token(ids.channel),
277 tag_text(&format!("commentary to=functions.{name}")),
278 optional_content_type(ids),
279 tag_token(ids.message),
280 { "type": "json_schema", "json_schema": params },
281 ]})
282 })
283 .collect();
284 let call = if calls.len() == 1 { calls.remove(0) } else { json!({ "type": "or", "elements": calls }) };
285 Some(structural(vec![optional_analysis(ids), call]))
286}
287
288fn declared_functions(tool_jsons: &[String]) -> Option<Vec<(String, serde_json::Value)>> {
289 if tool_jsons.is_empty() {
290 return None;
291 }
292 tool_jsons
293 .iter()
294 .map(|t| {
295 let v: serde_json::Value = serde_json::from_str(t).ok()?;
296 let f = v.get("function").unwrap_or(&v);
297 let name = f.get("name").and_then(|n| n.as_str())?.to_string();
298 let params = f.get("parameters").cloned().unwrap_or_else(|| serde_json::json!({ "type": "object" }));
299 Some((name, params))
300 })
301 .collect()
302}
303
304pub const FREE_MESSAGES: usize = 3;
305
306pub fn auto_call_tag(ids: HarmonyIds, tool_jsons: &[String]) -> Option<String> {
307 use serde_json::json;
308 let channel_of_call = || json!({ "type": "or", "elements": [tag_text("commentary"), tag_text("analysis")] });
309 let mut last = vec![json!({ "type": "sequence", "elements": [
310 tag_token(ids.channel), tag_text("final"), optional_content_type(ids), tag_token(ids.message),
311 { "type": "any_tokens", "exclude_tokens": markers(ids) },
312 ]})];
313 for (name, params) in declared_functions(tool_jsons)? {
314 let to = tag_text(&format!(" to=functions.{name}"));
315 last.push(json!({ "type": "sequence", "elements": [
316 { "type": "or", "elements": [
317 { "type": "sequence", "elements": [tag_token(ids.channel), channel_of_call(), to.clone()] },
318 { "type": "sequence", "elements": [to, tag_token(ids.channel), channel_of_call()] },
319 ]},
320 optional_content_type(ids),
321 tag_token(ids.message),
322 { "type": "json_schema", "json_schema": params },
323 { "type": "star", "content": tag_text(" ") },
324 ]}));
325 }
326 let free_message = || {
327 let [body, start, assistant] = body_then_next_message(ids);
328 json!({ "type": "sequence", "elements": [
329 tag_token(ids.channel),
330 { "type": "or", "elements": [tag_text("analysis"), tag_text("commentary")] },
331 optional_content_type(ids),
332 body, start, assistant,
333 ]})
334 };
335 let mut before = json!({ "type": "optional", "content": free_message() });
336 for _ in 1..FREE_MESSAGES {
337 before = json!({ "type": "optional", "content": { "type": "sequence", "elements": [free_message(), before] } });
338 }
339 Some(structural(vec![before, json!({ "type": "or", "elements": last })]))
340}
341
342pub fn final_json_tag(ids: HarmonyIds, json_schema: &str) -> Option<String> {
343 let schema: serde_json::Value = serde_json::from_str(json_schema).ok()?;
344 Some(structural(vec![
345 optional_analysis(ids),
346 tag_token(ids.channel),
347 tag_text("final"),
348 optional_content_type(ids),
349 tag_token(ids.message),
350 serde_json::json!({ "type": "json_schema", "json_schema": schema }),
351 ]))
352}
353
354#[cfg(test)]
355mod tests {
356 use super::*;
357
358 struct Tok;
359 const START: u32 = 1000;
360 const CHANNEL: u32 = 1001;
361 const MESSAGE: u32 = 1002;
362 const END: u32 = 1003;
363 const RETURN: u32 = 1004;
364 const CALL: u32 = 1005;
365 const CONSTRAIN: u32 = 1006;
366
367 fn ids() -> HarmonyIds {
368 HarmonyIds { start: START, channel: CHANNEL, message: MESSAGE, end: END, ret: RETURN, call: CALL, constrain: Some(CONSTRAIN) }
369 }
370
371 impl Tokenizer for Tok {
372 fn encode(&self, text: &str) -> Vec<u32> {
373 text.bytes().map(u32::from).collect()
374 }
375 fn token_bytes(&self, token: u32) -> Vec<u8> {
376 if token < 256 {
377 vec![token as u8]
378 } else {
379 Vec::new()
380 }
381 }
382 fn vocab_size(&self) -> u32 {
383 1100
384 }
385 fn special_tokens(&self) -> Vec<(String, u32)> {
386 vec![
387 ("<|start|>".into(), START),
388 ("<|channel|>".into(), CHANNEL),
389 ("<|message|>".into(), MESSAGE),
390 ("<|end|>".into(), END),
391 ("<|return|>".into(), RETURN),
392 ("<|call|>".into(), CALL),
393 ("<|constrain|>".into(), CONSTRAIN),
394 ]
395 }
396 fn bos_token(&self) -> Option<u32> {
397 None
398 }
399 fn eos_token(&self) -> u32 {
400 RETURN
401 }
402 fn chat_template_jinja(&self) -> String {
403 String::new()
404 }
405 }
406
407 fn bytes(s: &str) -> Vec<u32> {
408 s.bytes().map(u32::from).collect()
409 }
410
411 fn stream(parts: &[&[u32]]) -> Vec<u32> {
412 parts.iter().flat_map(|p| p.iter().copied()).collect()
413 }
414
415 fn decode(text: &[u32]) -> String {
416 String::from_utf8(text.iter().flat_map(|&t| Tok.token_bytes(t)).collect()).unwrap()
417 }
418
419 #[test]
420 fn discovers_the_markers_from_the_specials() {
421 assert_eq!(HarmonyIds::discover(&Tok.special_tokens()), Some(ids()));
422 let mut fewer = Tok.special_tokens();
423 fewer.retain(|(s, _)| s != "<|return|>");
424 assert_eq!(HarmonyIds::discover(&fewer), None, "every marker the format needs");
425 fewer = Tok.special_tokens();
426 fewer.retain(|(s, _)| s != "<|constrain|>");
427 assert!(HarmonyIds::discover(&fewer).is_some(), "the constraint marker is optional");
428 }
429
430 #[test]
431 fn routes_analysis_and_final() {
432 let toks = stream(&[
433 &[CHANNEL],
434 &bytes("analysis"),
435 &[MESSAGE],
436 &bytes("think hard"),
437 &[END, START],
438 &bytes("assistant"),
439 &[CHANNEL],
440 &bytes("final"),
441 &[MESSAGE],
442 &bytes("Hello!"),
443 &[RETURN],
444 ]);
445 let mut ch = HarmonyChannelizer::new(Arc::new(Tok), ids());
446 let runs = ch.split(&toks);
447 let joined: Vec<u32> = runs.iter().flat_map(|r| r.span.iter().copied()).collect();
448 assert_eq!(joined, toks, "spans partition the input");
449 let routed: Vec<(u32, String, bool)> =
450 runs.iter().map(|r| (r.channel, decode(&r.text), r.closes)).collect();
451 assert_eq!(
452 routed,
453 vec![
454 (channel::REASONING, "think hard".into(), true),
455 (channel::TEXT, "Hello!".into(), true),
456 ]
457 );
458 }
459
460 #[test]
461 fn routes_a_tool_call_with_its_header() {
462 let toks = stream(&[
463 &[CHANNEL],
464 &bytes("commentary to=functions.get_weather "),
465 &[CONSTRAIN],
466 &bytes("json"),
467 &[MESSAGE],
468 &bytes(r#"{"city":"Paris"}"#),
469 &[CALL],
470 ]);
471 let mut ch = HarmonyChannelizer::new(Arc::new(Tok), ids());
472 let runs = ch.split(&toks);
473 assert_eq!(runs.len(), 1);
474 assert_eq!(runs[0].channel, channel::TOOL_CALL);
475 assert!(runs[0].closes);
476 let raw = decode(&runs[0].text);
477 assert_eq!(raw, r#"commentary to=functions.get_weather {"city":"Paris"}"#, "the type is dropped");
478 assert_eq!(parse_harmony_tool_call(&raw), Some(("get_weather".into(), r#"{"city":"Paris"}"#.into())));
479 }
480
481 #[test]
482 fn a_recipient_glued_to_the_constraint_marker_keeps_its_name() {
483 let toks = stream(&[
484 &[CHANNEL],
485 &bytes("commentary to=functions.bash"),
486 &[CONSTRAIN],
487 &bytes("json"),
488 &[MESSAGE],
489 &bytes(r#"{"command":"ls -la"}"#),
490 &[CALL],
491 ]);
492 let mut ch = HarmonyChannelizer::new(Arc::new(Tok), ids());
493 let runs = ch.split(&toks);
494 assert_eq!(runs.len(), 1);
495 assert_eq!(runs[0].channel, channel::TOOL_CALL);
496 let raw = decode(&runs[0].text);
497 assert_eq!(raw, r#"commentary to=functions.bash{"command":"ls -la"}"#);
498 assert_eq!(parse_harmony_tool_call(&raw), Some(("bash".into(), r#"{"command":"ls -la"}"#.into())));
499 assert_eq!(parse_harmony_tool_call(r#"commentary to=functions.bash json{"command":"ls"}"#), Some(("bash".into(), r#"{"command":"ls"}"#.into())));
500 }
501
502 #[test]
503 fn a_recipient_before_the_channel_still_makes_a_call() {
504 let calls = |toks: &[u32]| -> Vec<(u32, String, bool)> {
505 let mut ch = HarmonyChannelizer::new(Arc::new(Tok), ids());
506 ch.split(toks).into_iter().filter(|r| !r.text.is_empty()).map(|r| (r.channel, decode(&r.text), r.closes)).collect()
507 };
508 let role_header = stream(&[&[START], &bytes("assistant to=functions.bash"), &[CHANNEL], &bytes("commentary json"), &[MESSAGE], &bytes(r#"{"command":"ls"}"#), &[CALL]]);
509 let got = calls(&role_header);
510 assert_eq!(got.len(), 1);
511 assert_eq!((got[0].0, got[0].2), (channel::TOOL_CALL, true), "{got:?}");
512 assert_eq!(parse_harmony_tool_call(&got[0].1), Some(("bash".into(), r#"{"command":"ls"}"#.into())), "{got:?}");
513
514 let restarted = stream(&[&[CHANNEL], &bytes("commentary to=functions.read"), &[CHANNEL], &bytes("analysis code"), &[MESSAGE], &bytes(r#"{"path":"a"}"#), &[CALL]]);
515 let got = calls(&restarted);
516 assert_eq!(got[0].0, channel::TOOL_CALL, "{got:?}");
517 assert_eq!(parse_harmony_tool_call(&got[0].1), Some(("read".into(), r#"{"path":"a"}"#.into())), "{got:?}");
518
519 let mut ch = HarmonyChannelizer::new(Arc::new(Tok), ids());
520 let sliced: Vec<Run> = role_header.iter().flat_map(|&t| ch.split(&[t])).collect();
521 let text: Vec<u32> = sliced.iter().filter(|r| r.channel == channel::TOOL_CALL).flat_map(|r| r.text.clone()).collect();
522 assert_eq!(parse_harmony_tool_call(&decode(&text)).map(|(n, _)| n), Some("bash".into()));
523
524 let answer = stream(&[&[START], &bytes("assistant"), &[CHANNEL], &bytes("final"), &[MESSAGE], &bytes("hi"), &[RETURN]]);
525 assert_eq!(calls(&answer), vec![(channel::TEXT, "hi".into(), true)]);
526 let forgotten = stream(&[&[START], &bytes("assistant to=functions.f"), &[START], &bytes("assistant"), &[CHANNEL], &bytes("final"), &[MESSAGE], &bytes("x"), &[RETURN]]);
527 assert_eq!(calls(&forgotten), vec![(channel::TEXT, "x".into(), true)]);
528 }
529
530 #[test]
531 fn preambles_are_text_and_builtin_tools_are_calls() {
532 let toks = stream(&[
533 &[CHANNEL],
534 &bytes("commentary"),
535 &[MESSAGE],
536 &bytes("Let me check."),
537 &[END, START],
538 &bytes("assistant"),
539 &[CHANNEL],
540 &bytes("analysis to=browser.search "),
541 &[CONSTRAIN],
542 &bytes("json"),
543 &[MESSAGE],
544 &bytes(r#"{"query":"weather"}"#),
545 &[CALL],
546 ]);
547 let mut ch = HarmonyChannelizer::new(Arc::new(Tok), ids());
548 let runs = ch.split(&toks);
549 assert_eq!(runs[0].channel, channel::TEXT);
550 assert_eq!(decode(&runs[0].text), "Let me check.");
551 assert_eq!(runs[1].channel, channel::TOOL_CALL);
552 assert_eq!(parse_harmony_tool_call(&decode(&runs[1].text)), Some(("browser.search".into(), r#"{"query":"weather"}"#.into())));
553 }
554
555 #[test]
556 fn slicing_anywhere_routes_like_the_whole() {
557 let toks = stream(&[
558 &[CHANNEL],
559 &bytes("analysis"),
560 &[MESSAGE],
561 &bytes("hmm"),
562 &[END, START],
563 &bytes("assistant"),
564 &[CHANNEL],
565 &bytes("commentary to=functions.f "),
566 &[CONSTRAIN],
567 &bytes("json"),
568 &[MESSAGE],
569 &bytes(r#"{"a":1}"#),
570 &[CALL],
571 ]);
572 let whole: Vec<(u32, Vec<u32>)> = {
573 let mut ch = HarmonyChannelizer::new(Arc::new(Tok), ids());
574 merge(ch.split(&toks))
575 };
576 for cut in 1..toks.len() {
577 let mut ch = HarmonyChannelizer::new(Arc::new(Tok), ids());
578 let mut runs = ch.split(&toks[..cut]);
579 runs.extend(ch.split(&toks[cut..]));
580 assert_eq!(merge(runs), whole, "cut at {cut}");
581 }
582 }
583
584 fn merge(runs: Vec<Run>) -> Vec<(u32, Vec<u32>)> {
585 let mut out: Vec<(u32, Vec<u32>)> = Vec::new();
586 for r in runs {
587 if r.text.is_empty() {
588 continue;
589 }
590 match out.last_mut() {
591 Some((ch, text)) if *ch == r.channel => text.extend(r.text),
592 _ => out.push((r.channel, r.text)),
593 }
594 }
595 out
596 }
597
598 #[test]
599 fn a_stop_mid_message_closes_nothing_and_a_new_header_recovers() {
600 let toks = stream(&[&[CHANNEL], &bytes("final"), &[MESSAGE], &bytes("partial"), &[START], &bytes("assistant"), &[CHANNEL], &bytes("final"), &[MESSAGE], &bytes("again")]);
601 let mut ch = HarmonyChannelizer::new(Arc::new(Tok), ids());
602 let runs = ch.split(&toks);
603 let texts: Vec<(String, bool)> = runs.iter().filter(|r| !r.text.is_empty()).map(|r| (decode(&r.text), r.closes)).collect();
604 assert_eq!(texts, vec![("partial".into(), false), ("again".into(), false)]);
605 assert_eq!(ch.current_channel(), channel::TEXT);
606 }
607
608 #[test]
609 fn the_forced_call_tag_matches_markers_by_id_and_bodies_by_schema() {
610 use serde_json::json;
611 let weather = r#"{"type":"function","function":{"name":"get_weather","parameters":{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}}}"#.to_string();
612 let now = r#"{"name":"now"}"#.to_string();
613 let tag: serde_json::Value = serde_json::from_str(&forced_call_tag(ids(), &[weather.clone(), now]).expect("a tag")).unwrap();
614 assert_eq!(tag["type"], "structural_tag");
615 let seq = tag["format"]["elements"].as_array().expect("a sequence");
616 assert_eq!(seq.len(), 2, "an optional analysis, then the call");
617
618 assert_eq!(seq[0]["type"], "optional");
619 let analysis = seq[0]["content"]["elements"].as_array().expect("the analysis message");
620 assert_eq!(analysis[0], json!({"type": "token", "token": CHANNEL}), "markers are matched by id");
621 assert_eq!(analysis[1], json!({"type": "const_string", "value": "analysis"}));
622 let body = &analysis[2];
623 assert_eq!(body["type"], "tag", "the unbounded body is closed by its end marker");
624 assert_eq!(body["begin"], json!({"type": "token", "token": MESSAGE}));
625 assert_eq!(body["end"], json!({"type": "token", "token": END}));
626 assert_eq!(body["content"]["type"], "any_tokens");
627 let excluded: Vec<u32> = serde_json::from_value(body["content"]["exclude_tokens"].clone()).unwrap();
628 for marker in [START, CHANNEL, MESSAGE, END, RETURN, CALL, CONSTRAIN] {
629 assert!(excluded.contains(&marker), "the analysis body cannot contain marker {marker}");
630 }
631 assert_eq!(&analysis[3..], &[json!({"type": "token", "token": START}), json!({"type": "const_string", "value": "assistant"})]);
632
633 let calls = seq[1]["elements"].as_array().expect("one alternative per tool");
634 assert_eq!(seq[1]["type"], "or");
635 assert_eq!(calls.len(), 2);
636 let first = calls[0]["elements"].as_array().unwrap();
637 assert_eq!(first[0], json!({"type": "token", "token": CHANNEL}));
638 assert_eq!(first[1], json!({"type": "const_string", "value": "commentary to=functions.get_weather"}));
639 assert_eq!(first[2]["type"], "optional", "the content type may be written or left out");
640 assert_eq!(first[3], json!({"type": "token", "token": MESSAGE}));
641 assert_eq!(first[4]["json_schema"]["required"], json!(["city"]), "the body is the tool's own parameters");
642 assert_eq!(first.len(), 5, "the tag ends with the body; the turn's stop token ends the call");
643 assert_eq!(calls[1]["elements"][4]["json_schema"], json!({"type": "object"}), "no parameters: any object");
644
645 let one: serde_json::Value = serde_json::from_str(&forced_call_tag(ids(), &[weather]).unwrap()).unwrap();
646 assert_eq!(one["format"]["elements"][1]["type"], "sequence", "one tool needs no alternation");
647 assert!(forced_call_tag(ids(), &[]).is_none(), "no tools, no tag");
648 assert!(forced_call_tag(ids(), &[r#"{"parameters":{}}"#.to_string()]).is_none(), "a nameless tool, no tag");
649 }
650
651 #[test]
652 fn the_answer_tag_holds_the_schema_in_the_final_channel() {
653 use serde_json::json;
654 let schema = r#"{"type":"object","properties":{"ok":{"type":"boolean"}},"required":["ok"]}"#;
655 let tag: serde_json::Value = serde_json::from_str(&final_json_tag(ids(), schema).expect("a tag")).unwrap();
656 let seq = tag["format"]["elements"].as_array().expect("a sequence");
657 assert_eq!(seq[0]["type"], "optional", "the model may reason first");
658 assert_eq!(seq[0]["content"]["elements"][1], json!({"type": "const_string", "value": "analysis"}));
659 assert_eq!(&seq[1..3], &[json!({"type": "token", "token": CHANNEL}), json!({"type": "const_string", "value": "final"})]);
660 assert_eq!(seq[3]["type"], "optional", "the content type may be written or left out");
661 assert_eq!(seq[4], json!({"type": "token", "token": MESSAGE}));
662 assert_eq!(seq[5], json!({"type": "json_schema", "json_schema": serde_json::from_str::<serde_json::Value>(schema).unwrap()}));
663 assert_eq!(seq.len(), 6, "the value is the last thing; `<|return|>` is the stop that ends it");
664 assert!(final_json_tag(ids(), "{not json").is_none());
665 }
666
667 #[test]
668 fn the_parser_repairs_a_truncated_body_and_refuses_the_rest() {
669 assert_eq!(parse_harmony_tool_call(r#"commentary to=functions.f json{"a":1,"b":[1,2"#), Some(("f".into(), r#"{"a":1,"b":[1,2]}"#.into())));
670 assert_eq!(parse_harmony_tool_call("final no recipient here {}"), None);
671 assert_eq!(parse_harmony_tool_call("commentary to=functions.f json"), None, "no body");
672 assert_eq!(parse_harmony_tool_call("commentary to= json{}"), None, "no name");
673 }
674
675 #[test]
676 fn a_call_whose_arguments_do_not_parse_is_still_named_by_its_header() {
677 let raw = r#"commentary to=functions.write {"path":"server.py","content":"main()\n","}"#;
678 assert_eq!(parse_harmony_tool_call(raw), None, "the arguments are not JSON");
679 assert_eq!(named_call(raw), Some(("write".into(), r#"{"path":"server.py","content":"main()\n","}"#.into())));
680 assert_eq!(named_call(r#"commentary to=functions.bash json{"command":"ls"#), Some(("bash".into(), r#"{"command":"ls"#.into())), "the type is not the arguments");
681 assert_eq!(named_call("commentary to=functions.f json"), Some(("f".into(), String::new())), "no body, no arguments");
682 assert_eq!(named_call("final no recipient here {}"), None);
683 assert_eq!(named_call("commentary to= json{}"), None, "no name");
684 }
685
686 #[test]
687 fn the_auto_tag_holds_only_a_calls_arguments() {
688 use serde_json::json;
689 let weather = r#"{"type":"function","function":{"name":"get_weather","parameters":{"type":"object","properties":{"city":{"type":"string"}},"required":["city"]}}}"#.to_string();
690 let now = r#"{"name":"now"}"#.to_string();
691 let tag: serde_json::Value = serde_json::from_str(&auto_call_tag(ids(), &[weather, now]).expect("a tag")).unwrap();
692 let seq = tag["format"]["elements"].as_array().expect("a sequence");
693 assert_eq!(seq.len(), 2, "the messages before the last, then the last");
694
695 let none = json!(null);
696 let mut depth = 0;
697 let mut free = &seq[0];
698 while free["type"] == "optional" {
699 depth += 1;
700 let inner = &free["content"];
701 let message = if inner["elements"][0]["type"] == "sequence" { &inner["elements"][0] } else { inner };
702 let parts = message["elements"].as_array().expect("a message");
703 assert_eq!(parts[0], json!({"type": "token", "token": CHANNEL}));
704 assert_eq!(parts[1], json!({"type": "or", "elements": [
705 {"type": "const_string", "value": "analysis"}, {"type": "const_string", "value": "commentary"}]}));
706 assert_eq!(parts[3]["type"], "tag", "an unbounded body closed by its end marker");
707 assert_eq!(parts[3]["end"], json!({"type": "token", "token": END}));
708 assert_eq!(&parts[4..], &[json!({"type": "token", "token": START}), json!({"type": "const_string", "value": "assistant"})]);
709 free = if inner["elements"][0]["type"] == "sequence" { &inner["elements"][1] } else { &none };
710 }
711 assert_eq!(depth, FREE_MESSAGES);
712
713 assert_eq!(seq[1]["type"], "or");
714 let last = seq[1]["elements"].as_array().unwrap();
715 assert_eq!(last.len(), 3, "the answer, then one call per tool");
716 let answer = last[0]["elements"].as_array().unwrap();
717 assert_eq!(&answer[..2], &[json!({"type": "token", "token": CHANNEL}), json!({"type": "const_string", "value": "final"})]);
718 assert_eq!(answer.last().unwrap()["type"], "any_tokens", "the answer runs to the stop token");
719 for (call, params) in last[1..].iter().zip([json!(["city"]), json!(null)]) {
720 let parts = call["elements"].as_array().unwrap();
721 assert_eq!(parts[0]["type"], "or", "the recipient after the channel, or before it");
722 let body = parts.iter().find(|p| p["type"] == "json_schema").expect("the arguments");
723 assert_eq!(body["json_schema"]["required"], params);
724 assert_eq!(parts.last().unwrap()["type"], "star", "every branch of the last message runs to the stop token");
725 }
726 assert!(auto_call_tag(ids(), &[]).is_none(), "no tools, no tag");
727 assert!(auto_call_tag(ids(), &[r#"{"parameters":{}}"#.to_string()]).is_none(), "a nameless tool, no tag");
728 }
729}