From 5691c9135087e9d0977c85d112761c50b0e4eecd Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Fri, 7 Feb 2025 10:36:49 +0100
Subject: [PATCH 1/2] Add `loop_controls` feature to `minijinja`

---
 router/Cargo.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/router/Cargo.toml b/router/Cargo.toml
index e4d0179a426..9326258daa2 100644
--- a/router/Cargo.toml
+++ b/router/Cargo.toml
@@ -48,7 +48,7 @@ ngrok = { version = "0.13.1", features = ["axum"], optional = true }
 init-tracing-opentelemetry = { version = "0.14.1", features = [
   "opentelemetry-otlp",
 ] }
-minijinja = { workspace = true }
+minijinja = { workspace = true, features = ["loop_controls"] }
 minijinja-contrib = { workspace = true }
 futures-util = "0.3.30"
 regex = "1.10.3"

From 20603881e377b08a143a233d3460687497370ee3 Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Fri, 7 Feb 2025 13:35:26 +0100
Subject: [PATCH 2/2] Add `test_chat_template_loop_controls` to test `break`

---
 router/src/infer/chat_template.rs | 66 +++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)

diff --git a/router/src/infer/chat_template.rs b/router/src/infer/chat_template.rs
index 8303ee76829..2fef2dcb87f 100644
--- a/router/src/infer/chat_template.rs
+++ b/router/src/infer/chat_template.rs
@@ -186,6 +186,72 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_chat_template_loop_controls() {
+        // some chat templates as e.g. CohereForAI/c4ai-command-r7b-12-202 contain `break`
+        // statements in their chat templates, so the feature `loop_controls` has been included
+        // in `minijinja`
+        let env = Environment::new();
+
+        let source = r#"
+        {% set user_count = 0 %}
+        {% for message in messages %}
+            {% if message['role'] == 'user' %}
+                {{'### User:\n' + message['content']+'\n\n'}}
+                {% set user_count = user_count + 1 %}
+                {% if user_count >= 2 %}
+                    {% break %}
+                {% endif %}
+            {% elif message['role'] == 'assistant' %}
+                {{'### Assistant:\n'  + message['content']}}
+            {% endif %}
+        {% endfor %}
+        {% if add_generation_prompt %}
+            {{ '### Assistant:\n' }}
+        {% endif %}"#;
+
+        // trim all the whitespace
+        let source = source
+            .lines()
+            .map(|line| line.trim())
+            .collect::<Vec<&str>>()
+            .join("");
+
+        let tmpl = env.template_from_str(&source);
+
+        let chat_template_inputs = ChatTemplateInputs {
+            messages: vec![
+                TextMessage {
+                    role: "user".to_string(),
+                    content: "Hi!".to_string(),
+                },
+                TextMessage {
+                    role: "assistant".to_string(),
+                    content: "Hello how can I help?".to_string(),
+                },
+                TextMessage {
+                    role: "user".to_string(),
+                    content: "What is Deep Learning?".to_string(),
+                },
+                TextMessage {
+                    role: "assistant".to_string(),
+                    content: "magic!".to_string(),
+                },
+            ],
+            bos_token: Some("[BOS]"),
+            eos_token: Some("[EOS]"),
+            add_generation_prompt: true,
+            ..Default::default()
+        };
+
+        let result = tmpl.unwrap().render(chat_template_inputs).unwrap();
+
+        assert_eq!(
+            result,
+            "### User:\nHi!\n\n### Assistant:\nHello how can I help?### User:\nWhat is Deep Learning?\n\n### Assistant:\n"
+        );
+    }
+
     #[test]
     fn test_chat_template_invalid_with_raise() {
         let mut env = Environment::new();