1use std::sync::Arc;
2
3use superfluid_abi::Status;
4use superfluid_engine::Tokenizer;
5
6use crate::codec::{BundleCodec, Channelizer, Run, TextCodec, TurnState};
7use crate::wal::{channel, role};
8
9pub struct AtemCodec {
10 inner: BundleCodec,
11 tok_start: u32,
12 tok_message: u32,
13 tok_eom: u32,
14 tok_eot: u32,
15 bos: Option<u32>,
16}
17
18impl AtemCodec {
19 #[cfg(feature = "basert")]
20 pub fn load(model_path: &std::path::Path) -> Result<AtemCodec, Status> {
21 let tok = superfluid_engine_ffi::TokenizerHandle::load(model_path).map_err(|_| Status::Fatal)?;
22 Self::from_tokenizer(Arc::new(tok))
23 }
24
25 pub fn from_tokenizer(tok: Arc<dyn Tokenizer>) -> Result<AtemCodec, Status> {
26 let inner = BundleCodec::from_tokenizer(tok);
27 let bos = inner.tokenizer().bos_token();
28 let one = |s: &str| -> Result<u32, Status> {
29 let mut t = inner.encode(s);
30 if bos.is_some() && t.first() == bos.as_ref() {
31 t.remove(0);
32 }
33 if t.len() == 1 {
34 Ok(t[0])
35 } else {
36 Err(Status::Unsupported)
37 }
38 };
39 Ok(AtemCodec {
40 tok_start: one("<|start|>")?,
41 tok_message: one("<|message|>")?,
42 tok_eom: one("<|eom|>")?,
43 tok_eot: one("<|eot|>")?,
44 bos,
45 inner,
46 })
47 }
48
49 fn enc(&self, text: &str) -> Vec<u32> {
50 let mut t = self.inner.encode(text);
51 if self.bos.is_some() && t.first() == self.bos.as_ref() {
52 t.remove(0);
53 }
54 t
55 }
56
57 fn enc_content(&self, text: &str) -> Vec<u32> {
58 let mut t = self.inner.encode_content(text);
59 if self.bos.is_some() && t.first() == self.bos.as_ref() {
60 t.remove(0);
61 }
62 t
63 }
64
65 fn enc_pieces(&self, pieces: &[(&str, bool)]) -> Vec<u32> {
66 let mut t = self.inner.encode_pieces(pieces);
67 if self.bos.is_some() && t.first() == self.bos.as_ref() {
68 t.remove(0);
69 }
70 t
71 }
72
73 fn turn(&self, header: &[(&str, bool)], body: &[(&str, bool)]) -> Vec<u32> {
74 let mut span = vec![self.tok_start];
75 span.extend(self.enc_pieces(header));
76 span.push(self.tok_message);
77 span.extend(self.enc_pieces(body));
78 span.push(self.tok_eot);
79 span
80 }
81
82 fn system_tail(&self, tool_jsons: &[String]) -> String {
83 let mut body = String::new();
84 body.push_str("\n\nReasoning strength: high.");
85 let mut namespaces: Vec<String> = Vec::new();
86 if !tool_jsons.is_empty() {
87 body.push_str("\n\n");
88 body.push_str(&render_tool_defs(tool_jsons, &mut namespaces));
89 }
90 body.push_str("\n\n# Valid recipients: \"self\"");
91 for ns in &namespaces {
92 body.push_str(&format!(", \"{ns}.*\""));
93 }
94 body.push_str(", \"user\".");
95 body
96 }
97}
98
99fn render_tool_defs(tool_jsons: &[String], namespaces: &mut Vec<String>) -> String {
100 let mut s = String::new();
101 s.push_str(
102 "In this environment you have access to a set of tools you can use to answer \
103 the user's question.\n\n",
104 );
105 s.push_str(
106 "You can invoke a function by writing a \"<atem:function_calls>\" block like \
107 the following:\n<atem:function_calls>\n<atem:invoke name=\"$FUNCTION_NAME\">\n\
108 <atem:parameter name=\"$PARAMETER_NAME\">$PARAMETER_VALUE</atem:parameter>\n...\n\
109 </atem:invoke>\n</atem:function_calls>\n\n",
110 );
111 s.push_str(
112 "String and scalar parameters should be specified as is, while lists and objects \
113 should use JSON format. Note that spaces for string values are not stripped. The \
114 output is not expected to be valid XML and is parsed with regular expressions.\n",
115 );
116 s.push_str("Here are the functions available in JSONSchema format:\n");
117 s.push_str("// Tool metadata\n");
118 let tools: Vec<serde_json::Value> = tool_jsons
119 .iter()
120 .filter_map(|t| serde_json::from_str::<serde_json::Value>(t).ok())
121 .map(|t| t.get("function").cloned().unwrap_or(t))
122 .collect();
123 for t in &tools {
124 if let Some(name) = t.get("name").and_then(|n| n.as_str()) {
125 let ns = name.split('.').next().unwrap_or(name).to_string();
126 if !namespaces.contains(&ns) {
127 namespaces.push(ns);
128 }
129 }
130 }
131 for ns in namespaces.iter() {
132 s.push_str(&format!(
133 "{{\"name\": {}, \"description\": \"\"}}\n",
134 serde_json::json!(ns)
135 ));
136 }
137 s.push_str("// Function schemas");
138 for t in &tools {
139 let name = t.get("name").cloned().unwrap_or(serde_json::json!(""));
140 let desc = t
141 .get("description")
142 .cloned()
143 .unwrap_or(serde_json::json!(""));
144 let params = t
145 .get("parameters")
146 .cloned()
147 .unwrap_or(serde_json::json!({}));
148 s.push_str(&format!(
149 "\n{{\"name\": {name}, \"description\": {desc}, \"parameters\": {params}}}"
150 ));
151 }
152 s.push_str("\n\nHere's an example of how to call a function in the tool set:\n");
153 s.push_str(
154 "(If the tool namespace is not specified, invoke the function directly as \
155 `example_function_name` rather than `example_tool_name.example_function_name`)\n\n",
156 );
157 s.push_str("to=example_tool_name.example_function_name\n\n");
158 s.push_str(
159 "<atem:function_calls>\n<atem:invoke name=\"example_tool_name.example_function_name\">\n\
160 <atem:parameter name=\"example_parameter_1\">value_1</atem:parameter>\n\
161 <atem:parameter name=\"example_parameter_2\">This is the value for the second parameter\n\
162 that can span\n\"multiple\" lines\n</atem:parameter>\n\
163 </atem:invoke>\n</atem:function_calls>",
164 );
165 s
166}
167
168pub fn render_atem_for_test(name: &str, arguments: &str) -> String {
169 render_atem(name, arguments)
170}
171
172fn render_atem(name: &str, arguments: &str) -> String {
173 render_atem_pieces(name, arguments).into_iter().map(|(s, _)| s).collect()
174}
175
176type Pieces<'a, const N: usize> = [(&'a str, bool); N];
177
178fn tool_result_pieces<'a>(name: &'a str, content: &'a str) -> (Pieces<'a, 2>, Pieces<'a, 5>) {
179 (
180 [("tool ", false), (name, true)],
181 [
182 ("<tool_output name=\"", false),
183 (name, true),
184 ("\">\n", false),
185 (content, true),
186 ("\n</tool_output>", false),
187 ],
188 )
189}
190
191fn render_atem_pieces(name: &str, arguments: &str) -> Vec<(String, bool)> {
192 let mut p: Vec<(String, bool)> = vec![
193 ("<atem:function_calls>\n<atem:invoke name=\"".to_string(), false),
194 (name.to_string(), true),
195 ("\">\n".to_string(), false),
196 ];
197 if let Ok(serde_json::Value::Object(args)) = serde_json::from_str(arguments) {
198 for (k, v) in args {
199 p.push(("<atem:parameter name=\"".to_string(), false));
200 p.push((k, true));
201 p.push(("\">".to_string(), false));
202 let value = match v {
203 serde_json::Value::Bool(true) => "true".to_string(),
204 serde_json::Value::Bool(false) => "false".to_string(),
205 serde_json::Value::Null => "null".to_string(),
206 serde_json::Value::String(v) => v,
207 serde_json::Value::Number(n) => n.to_string(),
208 other => other.to_string(),
209 };
210 p.push((value, true));
211 p.push(("</atem:parameter>\n".to_string(), false));
212 }
213 }
214 p.push(("</atem:invoke>\n</atem:function_calls>".to_string(), false));
215 p
216}
217
218impl TextCodec for AtemCodec {
219 fn encode(&self, text: &str) -> Vec<u32> {
220 self.enc(text)
221 }
222 fn encode_content(&self, text: &str) -> Vec<u32> {
223 self.enc_content(text)
224 }
225 fn encode_pieces(&self, pieces: &[(&str, bool)]) -> Vec<u32> {
226 self.enc_pieces(pieces)
227 }
228 fn chat_template_source(&self) -> Option<String> {
229 self.inner.chat_template_source()
230 }
231 fn token_bytes(&self, token: u32) -> Vec<u8> {
232 self.inner.token_bytes(token)
233 }
234 fn marker_literal(&self, token: u32) -> String {
235 self.inner.marker_literal(token)
236 }
237
238 fn stream_prologue(&self) -> Vec<u32> {
239 self.bos.into_iter().collect()
240 }
241
242 fn render_message(&self, msg_role: u32, text: &str) -> Option<Vec<u32>> {
243 match msg_role {
244 role::SYSTEM => {
245 let tail = self.system_tail(&[]);
246 Some(self.turn(&[("system", false)], &[(text, true), (&tail, false)]))
247 }
248 role::USER => Some(self.turn(&[("user", false)], &[(text, true)])),
249 role::ASSISTANT => Some(self.turn(&[("assistant to=user", false)], &[(text, true)])),
250 _ => None,
251 }
252 }
253
254 fn render_turn_with_tokens(
255 &self,
256 msg_role: u32,
257 pre_text: &str,
258 inner: &[u32],
259 post_text: &str,
260 ) -> Option<Vec<u32>> {
261 let header = match msg_role {
262 role::SYSTEM => "system",
263 role::USER => "user",
264 role::ASSISTANT => "assistant to=user",
265 _ => return None,
266 };
267 let mut span = vec![self.tok_start];
268 span.extend(self.enc(header));
269 span.push(self.tok_message);
270 if !pre_text.is_empty() {
271 span.extend(self.enc_content(pre_text));
272 }
273 span.extend_from_slice(inner);
274 if !post_text.is_empty() {
275 span.extend(self.enc_content(post_text));
276 }
277 span.push(self.tok_eot);
278 Some(span)
279 }
280
281 fn generation_prefix(&self, state: TurnState) -> Option<Vec<u32>> {
282 match state {
283 TurnState::AfterInput | TurnState::AfterFinishedTurn => {
284 let mut s = vec![self.tok_start];
285 s.extend(self.enc("assistant"));
286 Some(s)
287 }
288 TurnState::MidTurn => None,
289 }
290 }
291
292 fn channel_markers(&self) -> Vec<(u32, u32, u32)> {
293 Vec::new()
294 }
295
296 fn render_system_with_tools(
297 &self,
298 system: Option<&str>,
299 tool_jsons: &[String],
300 ) -> Option<Vec<u32>> {
301 let content = system.map(str::to_string).unwrap_or_else(|| {
302 "You are a helpful AI assistant.\nKnowledge cutoff: 2026-01-04.".to_string()
303 });
304 let tail = self.system_tail(tool_jsons);
305 Some(self.turn(&[("system", false)], &[(&content, true), (&tail, false)]))
306 }
307
308 fn render_tool_result(&self, name: &str, content: &str) -> Option<Vec<u32>> {
309 let (header, body) = tool_result_pieces(name, content);
310 Some(self.turn(&header, &body))
311 }
312
313 fn render_assistant_with_tool_calls(
314 &self,
315 _content: &str,
316 calls: &[(String, String)],
317 ) -> Option<Vec<u32>> {
318 let mut span = Vec::new();
319 for (i, (name, args)) in calls.iter().enumerate() {
320 span.push(self.tok_start);
321 span.extend(self.enc_pieces(&[("assistant to=", false), (name, true)]));
322 span.push(self.tok_message);
323 let body = render_atem_pieces(name, args);
324 let body: Vec<(&str, bool)> = body.iter().map(|(s, c)| (s.as_str(), *c)).collect();
325 span.extend(self.enc_pieces(&body));
326 span.push(if i + 1 == calls.len() {
327 self.tok_eot
328 } else {
329 self.tok_eom
330 });
331 }
332 Some(span)
333 }
334
335 fn channelizer(&self) -> Box<dyn Channelizer> {
336 Box::new(AtemChannelizer {
337 tok: self.inner.tokenizer(),
338 tok_start: self.tok_start,
339 tok_message: self.tok_message,
340 tok_eom: self.tok_eom,
341 tok_eot: self.tok_eot,
342 state: ChState::Header(String::new()),
343 })
344 }
345
346 fn tool_envelope(&self) -> crate::codec::ToolEnvelope {
347 crate::codec::ToolEnvelope::Atem
348 }
349
350 fn tool_call_delimiters(&self) -> Option<(String, String)> {
351 Some(("<atem:function_calls>".to_string(), "</atem:function_calls>".to_string()))
352 }
353
354 fn behavior_fingerprint(&self) -> Option<superfluid_fingerprint::BehaviorFingerprint> {
355 Some(superfluid_fingerprint::BehaviorFingerprint {
356 codec_id: "atem".into(),
357 codec_version: 2,
358 tokenizer_hash: self.inner.fingerprint("atem", &[]).tokenizer_hash,
359 marker_digest: superfluid_fingerprint::BehaviorFingerprint::marker_digest_of(&[
360 (self.tok_start, self.tok_message, 0),
361 (self.tok_eom, self.tok_eot, 0),
362 ]),
363 sampling_impl_version: 1,
364 block_projection_version: crate::codec::MODEL_VISIBILITY_VERSION,
365 grammar_version: 0,
366 speculation: None,
367 })
368 }
369}
370
371pub(crate) fn parse_atem_call(raw: &str) -> Option<(String, String)> {
372 let inv = raw.find("<atem:invoke name=\"")?;
373 let name_start = inv + "<atem:invoke name=\"".len();
374 let name_end = raw[name_start..].find('"')? + name_start;
375 let name = raw[name_start..name_end].to_string();
376 let mut args = serde_json::Map::new();
377 let mut rest = &raw[name_end..];
378 while let Some(p) = rest.find("<atem:parameter name=\"") {
379 let k_start = p + "<atem:parameter name=\"".len();
380 let k_len = rest[k_start..].find('"')?;
381 let key = rest[k_start..k_start + k_len].to_string();
382 let v_start = rest[k_start + k_len..].find('>')?;
383 let v_start = k_start + k_len + v_start + 1;
384 let v_len = rest[v_start..].find("</atem:parameter>")?;
385 let val_raw = &rest[v_start..v_start + v_len];
386 let value = serde_json::from_str::<serde_json::Value>(val_raw)
387 .unwrap_or_else(|_| serde_json::Value::String(val_raw.to_string()));
388 args.insert(key, value);
389 rest = &rest[v_start + v_len..];
390 }
391 Some((name, serde_json::Value::Object(args).to_string()))
392}
393
394enum ChState {
395 Header(String),
396 Body(u32),
397}
398
399struct AtemChannelizer {
400 tok: Arc<dyn Tokenizer>,
401 tok_start: u32,
402 tok_message: u32,
403 tok_eom: u32,
404 tok_eot: u32,
405 state: ChState,
406}
407
408impl AtemChannelizer {
409 fn classify(header: &str) -> u32 {
410 if header.contains("=self") {
411 return channel::REASONING;
412 }
413 match header.find("to=") {
414 Some(p) => {
415 let recipient = header[p + 3..]
416 .split_whitespace()
417 .next()
418 .unwrap_or("user");
419 if recipient == "user" {
420 channel::TEXT
421 } else {
422 channel::TOOL_CALL
423 }
424 }
425 None => channel::TEXT,
426 }
427 }
428}
429
430impl Channelizer for AtemChannelizer {
431 fn split(&mut self, tokens: &[u32]) -> Vec<Run> {
432 let mut out: Vec<Run> = Vec::new();
433 let mut span: Vec<u32> = Vec::new();
434 let mut text: Vec<u32> = Vec::new();
435 for &t in tokens {
436 match &mut self.state {
437 ChState::Header(hdr) => {
438 span.push(t);
439 if t == self.tok_message {
440 self.state = ChState::Body(Self::classify(hdr));
441 } else if t != self.tok_start {
442 hdr.push_str(&String::from_utf8_lossy(&self.tok.token_bytes(t)));
443 }
444 }
445 ChState::Body(ch) => {
446 let ch = *ch;
447 if t == self.tok_eom || t == self.tok_eot {
448 span.push(t);
449 out.push(Run {
450 span: std::mem::take(&mut span),
451 text: std::mem::take(&mut text),
452 channel: ch,
453 closes: true,
454 });
455 self.state = ChState::Header(String::new());
456 } else if t == self.tok_start {
457 out.push(Run {
458 span: std::mem::take(&mut span),
459 text: std::mem::take(&mut text),
460 channel: ch,
461 closes: false,
462 });
463 span.push(t);
464 self.state = ChState::Header(String::new());
465 } else {
466 span.push(t);
467 text.push(t);
468 }
469 }
470 }
471 }
472 if !span.is_empty() {
473 let (channel, closes) = match &self.state {
474 ChState::Header(_) => (self.current_channel(), false),
475 ChState::Body(ch) => (*ch, false),
476 };
477 out.push(Run {
478 span,
479 text,
480 channel,
481 closes,
482 });
483 }
484 out
485 }
486
487 fn current_channel(&self) -> u32 {
488 match &self.state {
489 ChState::Header(_) => channel::TEXT,
490 ChState::Body(ch) => *ch,
491 }
492 }
493}
494
495#[cfg(test)]
496mod tests {
497 use super::*;
498
499 #[test]
500 fn atem_call_parsing_covers_the_value_forms() {
501 let raw = "<atem:function_calls>\n<atem:invoke name=\"weather.get\">\n\
502 <atem:parameter name=\"city\">Paris</atem:parameter>\n\
503 <atem:parameter name=\"days\">3</atem:parameter>\n\
504 <atem:parameter name=\"detailed\">true</atem:parameter>\n\
505 <atem:parameter name=\"opts\">{\"units\":\"C\"}</atem:parameter>\n\
506 </atem:invoke>\n</atem:function_calls>";
507 let (name, args) = parse_atem_call(raw).expect("parses");
508 assert_eq!(name, "weather.get");
509 let v: serde_json::Value = serde_json::from_str(&args).unwrap();
510 assert_eq!(v["city"], "Paris");
511 assert_eq!(v["days"], 3);
512 assert_eq!(v["detailed"], true);
513 assert_eq!(v["opts"]["units"], "C");
514 }
515
516 #[test]
517 fn malformed_parameter_refuses_the_whole_call() {
518 let truncated = "<atem:invoke name=\"get\">\n\
519 <atem:parameter name=\"city\">Paris</atem:parameter>\n\
520 <atem:parameter name=\"days\">3";
521 assert!(parse_atem_call(truncated).is_none());
522 let bad_name = "<atem:invoke name=\"get\"><atem:parameter name=\"city>Paris";
523 assert!(parse_atem_call(bad_name).is_none());
524 }
525
526 #[test]
527 fn tool_defs_unwrap_the_function_wrapper() {
528 let wrapped = serde_json::json!({
529 "type": "function",
530 "function": {"name": "weather.get", "description": "d", "parameters": {}},
531 })
532 .to_string();
533 let mut namespaces = Vec::new();
534 let defs = render_tool_defs(&[wrapped], &mut namespaces);
535 assert!(defs.contains("{\"name\": \"weather.get\""));
536 assert_eq!(namespaces, vec!["weather".to_string()]);
537 let flat = serde_json::json!({"name": "ping", "description": "", "parameters": {}})
538 .to_string();
539 let mut ns2 = Vec::new();
540 let defs2 = render_tool_defs(&[flat], &mut ns2);
541 assert!(defs2.contains("{\"name\": \"ping\""));
542 }
543
544 #[test]
545 fn atem_call_parsing_refuses_garbage() {
546 assert!(parse_atem_call("not a call").is_none());
547 let (name, args) =
548 parse_atem_call("<atem:invoke name=\"ping\"></atem:invoke>").expect("zero-arg");
549 assert_eq!(name, "ping");
550 assert_eq!(args, "{}");
551 }
552
553 #[test]
554 fn tool_result_pieces_mark_the_tool_name_as_content_in_header_and_body() {
555 let (header, body) = tool_result_pieces("weather <|eot|>", "sunny");
556 let content = |p: &[(&str, bool)]| -> Vec<String> {
557 p.iter().filter(|(_, c)| *c).map(|(s, _)| s.to_string()).collect()
558 };
559 assert_eq!(content(&header), ["weather <|eot|>"]);
560 assert_eq!(content(&body), ["weather <|eot|>", "sunny"]);
561 let joined = |p: &[(&str, bool)]| -> String { p.iter().map(|(s, _)| *s).collect() };
562 assert_eq!(joined(&header), "tool weather <|eot|>");
563 assert_eq!(joined(&body), "<tool_output name=\"weather <|eot|>\">\nsunny\n</tool_output>");
564 }
565
566 #[test]
567 fn render_atem_pieces_mark_the_clients_text_as_content() {
568 let pieces = render_atem_pieces("weather.get", "{\"city\":\"Paris <|eot|>\",\"days\":3}");
569 let content: Vec<&str> = pieces.iter().filter(|(_, c)| *c).map(|(s, _)| s.as_str()).collect();
570 assert_eq!(content, ["weather.get", "city", "Paris <|eot|>", "days", "3"]);
571 let framing: String = pieces.iter().filter(|(_, c)| !*c).map(|(s, _)| s.as_str()).collect();
572 assert!(!framing.contains("Paris") && framing.contains("<atem:parameter name=\""), "{framing}");
573 let joined: String = pieces.into_iter().map(|(s, _)| s).collect();
574 assert_eq!(joined, render_atem("weather.get", "{\"city\":\"Paris <|eot|>\",\"days\":3}"));
575 }
576
577 #[test]
578 fn render_atem_round_trips_through_the_parser() {
579 let rendered = render_atem(
580 "weather.get",
581 "{\"city\":\"Paris\",\"days\":3,\"detailed\":true}",
582 );
583 let (name, args) = parse_atem_call(&rendered).expect("parses");
584 assert_eq!(name, "weather.get");
585 let v: serde_json::Value = serde_json::from_str(&args).unwrap();
586 assert_eq!(v["city"], "Paris");
587 assert_eq!(v["days"], 3);
588 assert_eq!(v["detailed"], true);
589 }
590}