From 77940ac73f4efd0d8735466504ce12a66a1eb061 Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Fri, 31 Jan 2025 21:52:55 +0100
Subject: [PATCH] Fix `test_chat_template_valid_with_strftime_now`

---
 router/src/infer/chat_template.rs | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/router/src/infer/chat_template.rs b/router/src/infer/chat_template.rs
index 00938aba..8303ee76 100644
--- a/router/src/infer/chat_template.rs
+++ b/router/src/infer/chat_template.rs
@@ -326,13 +326,15 @@ mod tests {
         let source = r#"
         {% set today = strftime_now("%Y-%m-%d") %}
         {% set default_system_message = "The current date is " + today + "." %}
-        {{ default_system_message }}
         {{ bos_token }}
         {% if messages[0]['role'] == 'system' %}
-            {% set loop_messages = messages[1:] %}
+            { set system_message = messages[0]['content'] %}
+            {%- set loop_messages = messages[1:] %}
         {% else %}
-            {% set loop_messages = messages %}
+            {%- set system_message = default_system_message %}
+            {%- set loop_messages = messages %}
         {% endif %}
+        {{ '[SYSTEM_PROMPT]' + system_message + '[/SYSTEM_PROMPT]' }}
         {% for message in loop_messages %}
             {% if message['role'] == 'user' %}
                 {{ '[INST]' + message['content'] + '[/INST]' }}
@@ -380,7 +382,7 @@ mod tests {
 
         let current_date = Local::now().format("%Y-%m-%d").to_string();
         let result = tmpl.unwrap().render(chat_template_inputs).unwrap();
-        assert_eq!(result, format!("The current date is {}.[BOS][INST] Hi! [/INST]Hello how can I help?[EOS][INST] What is Deep Learning? [/INST]magic![EOS]", current_date));
+        assert_eq!(result, format!("[BOS][SYSTEM_PROMPT]The current date is {}.[/SYSTEM_PROMPT][INST]Hi![/INST]Hello how can I help?[EOS][INST]What is Deep Learning?[/INST]magic![EOS]", current_date));
     }
 
     #[test]