From 76f8abeeba420af2440ecadce4293cbc7e4d18c4 Mon Sep 17 00:00:00 2001
From: Michal Moskal <michal@moskal.me>
Date: Mon, 13 May 2024 21:29:04 +0000
Subject: [PATCH] fix crash in gctrl

---
 controllers/guidance_ctrl/src/tokenparser.rs | 14 ++++++++++----
 py/guidance                                  |  2 +-
 scripts/test-guidance.sh                     |  1 +
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/controllers/guidance_ctrl/src/tokenparser.rs b/controllers/guidance_ctrl/src/tokenparser.rs
index 0615d905..78f788da 100644
--- a/controllers/guidance_ctrl/src/tokenparser.rs
+++ b/controllers/guidance_ctrl/src/tokenparser.rs
@@ -101,14 +101,20 @@ impl TokenParser {
         infoln!("\n");
         let trie = self.token_env.tok_trie();
 
-        infoln!("post tokens: {}", trie.tokens_dbg(&arg.tokens));
+        infoln!("post tokens: bt={} {}", arg.backtrack, trie.tokens_dbg(&arg.tokens));
         arg.save_tokens(&mut self.llm_tokens);
 
         let new_bytes = trie.decode(&arg.tokens);
         self.llm_bytes.extend_from_slice(&new_bytes);
 
         // TODO maybe remove in future
-        assert!(self.llm_bytes == trie.decode(&self.llm_tokens));
+        if self.llm_bytes != trie.decode(&self.llm_tokens) {
+            panic!(
+                "llm_bytes mismatch: {:?} {:?}",
+                String::from_utf8_lossy(&self.llm_bytes),
+                String::from_utf8_lossy(&trie.decode(&self.llm_tokens))
+            );
+        }
 
         let res = self
             .parser
@@ -169,8 +175,8 @@ impl TokenParser {
             // here we remove a suffix from grm_tokens that could be possibly tokenized differently
             grm_tokens.truncate(grm_tokens.len() - chop_tokens);
 
-            if grm_tokens.len() > 0 {
-                infoln!("fixed_tokens: {}", trie.tokens_dbg(&grm_tokens));
+            if grm_tokens.len() > 0 || backtrack > 0  {
+                infoln!("fixed_tokens: {} bt={}", trie.tokens_dbg(&grm_tokens), backtrack);
                 return MidProcessResult::splice(backtrack as u32, grm_tokens);
             } else {
                 infoln!("no fixed tokens");
diff --git a/py/guidance b/py/guidance
index 1a48df37..b7e8e589 160000
--- a/py/guidance
+++ b/py/guidance
@@ -1 +1 @@
-Subproject commit 1a48df37d8fa95735012fafd303ce7ff1e345f3a
+Subproject commit b7e8e58905915a356c76a53254786e4c465eb732
diff --git a/scripts/test-guidance.sh b/scripts/test-guidance.sh
index bc0436e3..2e4caba5 100755
--- a/scripts/test-guidance.sh
+++ b/scripts/test-guidance.sh
@@ -10,3 +10,4 @@ export AZURE_GUIDANCE_URL
 
 cd $(dirname $0)/../py/guidance
 pytest --selected_model azure_guidance tests/models/test_azure_guidance.py"$@"
+pytest --selected_model azure_guidance_chat tests/models/test_azure_guidance_chat.py"$@"