loopd: the self-test checks the model's tags against the tokenizer

Co-Authored-By: ornith-1.5-35b-a3b <agent@wntrmute.dev>

crates/loopd/src/selftest.rs

@@ -1,6 +1,9 @@
 //! The startup self-test: before `loopd` serves anyone, it checks that the server is the one its
-//! config describes, that a tool call comes back parsed, and that a second turn reuses the first
-//! turn's cache. Any check failing means `loopd` does not start.
+//! config describes, that the model's tags tokenize as recorded, that a tool call comes back
+//! parsed, and that a second turn reuses the first turn's cache. Any check failing means `loopd`
+//! does not start.
+
+use bxcore::sanitize::{TAGS, Tag, neutralise};
 
 use crate::llama::{
     ChatMessage, ChatRequest, Client, FinishReason, InferError, ToolSchema,
@@ -21,6 +24,45 @@ pub enum SelfTestError {
     },
     Infer(InferError),
     Hash,
+    /// The tag with the recorded id `id` is `count` tokens, not one.
+    TagNotOneToken {
+        id: u32,
+        count: usize,
+    },
+    /// The tag with the recorded id `id` is one token, but the token `got`.
+    TagOtherId {
+        id: u32,
+        got: u32,
+    },
+    /// The neutralised form of the tag with the recorded id `id` is `count` tokens, not more
+    /// than one.
+    NeutralisedTooFew {
+        id: u32,
+        count: usize,
+    },
+    /// `/tokenize` failed for the tag with the recorded id `id`, or for its neutralised form.
+    Tokenize {
+        id: u32,
+        error: InferError,
+    },
+}
+
+impl SelfTestError {
+    /// The runbook pointer that ends the line `loopd` prints for this failure. A tag failure has
+    /// its own entry; every other failure shares the self-test's.
+    pub fn pointer(&self) -> &'static str {
+        match self {
+            SelfTestError::TagNotOneToken { .. }
+            | SelfTestError::TagOtherId { .. }
+            | SelfTestError::NeutralisedTooFew { .. }
+            | SelfTestError::Tokenize { .. } => "see docs/runbook.md#loopd-tags-unexpected",
+            SelfTestError::Mismatch { .. }
+            | SelfTestError::ToolCall(_)
+            | SelfTestError::CacheMiss { .. }
+            | SelfTestError::Infer(_)
+            | SelfTestError::Hash => "see docs/runbook.md#loopd-selftest-failed",
+        }
+    }
 }
 
 impl std::fmt::Display for SelfTestError {
@@ -40,6 +82,19 @@ impl std::fmt::Display for SelfTestError {
             ),
             SelfTestError::Infer(e) => write!(f, "{e}"),
             SelfTestError::Hash => write!(f, "the chat template could not be hashed"),
+            SelfTestError::TagNotOneToken { id, count } => {
+                write!(f, "the tag with id {id} is {count} tokens, not one")
+            }
+            SelfTestError::TagOtherId { id, got } => {
+                write!(f, "the tag with id {id} is one token, but the token {got}")
+            }
+            SelfTestError::NeutralisedTooFew { id, count } => write!(
+                f,
+                "the neutralised form of the tag with id {id} does not split into several tokens (the server gave {count})"
+            ),
+            SelfTestError::Tokenize { id, error } => {
+                write!(f, "could not tokenize the tag with id {id}: {error}")
+            }
         }
     }
 }
@@ -54,11 +109,14 @@ impl From<InferError> for SelfTestError {
 
 const SYSTEM: &str = "You are Boxmaker, a careful personal agent.";
 
-/// Runs the three checks in order. `on_step` is told the name of each check as it starts.
+/// Runs the four checks in order. `on_step` is told the name of each check as it starts.
 pub fn run(client: &Client, on_step: &mut dyn FnMut(&str)) -> Result<(), SelfTestError> {
     on_step("server matches config");
     server_matches(client)?;
 
+    on_step("tags tokenize as recorded");
+    tags_tokenize(client)?;
+
     on_step("tool call round trip");
     tool_call_round_trip(client)?;
 
@@ -103,7 +161,45 @@ fn server_matches(client: &Client) -> Result<(), SelfTestError> {
     Ok(())
 }
 
-/// Check 2: a tool call comes back parsed, with a `read_file` tool and a string `path`.
+/// Check 2: each of the model's tags (`bxcore::sanitize::TAGS`) is one token, the one recorded
+/// with it, and its neutralised form is more than one (M5b spec, section 8). Like check 1 it sends
+/// no prompt and uses no slot, so a server whose tokenizer is not the recorded one is sent no
+/// prompt either. A tag missing from the list cannot be found: the server does not list them.
+fn tags_tokenize(client: &Client) -> Result<(), SelfTestError> {
+    for tag in TAGS {
+        check_tag(client, tag)?;
+    }
+    Ok(())
+}
+
+/// Two `/tokenize` requests: `tag.text`, then `neutralise(tag.text)`. An error names the tag by
+/// its id, never by its text, which the operator's terminal and log would show whole.
+fn check_tag(client: &Client, tag: &Tag) -> Result<(), SelfTestError> {
+    let id = tag.id;
+    let tokenize = |text: &str| {
+        client
+            .token_ids(text)
+            .map_err(|error| SelfTestError::Tokenize { id, error })
+    };
+    let ids = tokenize(tag.text)?;
+    match ids.as_slice() {
+        [got] if *got == id => {}
+        [got] => return Err(SelfTestError::TagOtherId { id, got: *got }),
+        _ => {
+            return Err(SelfTestError::TagNotOneToken {
+                id,
+                count: ids.len(),
+            });
+        }
+    }
+    let count = tokenize(&neutralise(tag.text))?.len();
+    if count <= 1 {
+        return Err(SelfTestError::NeutralisedTooFew { id, count });
+    }
+    Ok(())
+}
+
+/// Check 3: a tool call comes back parsed, with a `read_file` tool and a string `path`.
 fn tool_call_round_trip(client: &Client) -> Result<(), SelfTestError> {
     let req = ChatRequest {
         slot: client.config().slots.main,
@@ -159,7 +255,7 @@ fn tool_call_round_trip(client: &Client) -> Result<(), SelfTestError> {
     Ok(())
 }
 
-/// Check 3: a second turn that extends the first reuses the first turn's cache.
+/// Check 4: a second turn that extends the first reuses the first turn's cache.
 fn cache_hit(client: &Client) -> Result<(), SelfTestError> {
     let first = ChatRequest {
         slot: client.config().slots.main,