From 55cc2b3b25dcd6409a949399a6a5b5a6b3d9257a Mon Sep 17 00:00:00 2001
From: Antonis Makropoulos <benuix@gmail.com>
Date: Wed, 13 Mar 2024 13:33:05 +0200
Subject: [PATCH] adapt tests to template prefix changes

---
 Tests/Runtime/TestLLM.cs              |  2 +-
 Tests/Runtime/TestLLMChatTemplates.cs | 17 ++++++++---------
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/Tests/Runtime/TestLLM.cs b/Tests/Runtime/TestLLM.cs
index 33b14aeb..6ee56796 100644
--- a/Tests/Runtime/TestLLM.cs
+++ b/Tests/Runtime/TestLLM.cs
@@ -122,7 +122,7 @@ public void TestAlive()
         public void TestInitParameters(int nkeep, int chats)
         {
             Assert.That(llm.nKeep == nkeep);
-            Assert.That(llm.template.GetStop().Length > 0);
+            Assert.That(llm.template.GetStop(llm.playerName, llm.AIName).Length > 0);
             Assert.That(llm.GetChat().Count == chats);
         }
 
diff --git a/Tests/Runtime/TestLLMChatTemplates.cs b/Tests/Runtime/TestLLMChatTemplates.cs
index 8fa14d53..dc2f360e 100644
--- a/Tests/Runtime/TestLLMChatTemplates.cs
+++ b/Tests/Runtime/TestLLMChatTemplates.cs
@@ -2,7 +2,6 @@
 using System.Collections.Generic;
 using NUnit.Framework.Internal;
 using NUnit.Framework;
-using UnityEngine;
 
 namespace LLMUnityTests
 {
@@ -22,7 +21,7 @@ public class TestChatTemplate
         public void TestChatML()
         {
             Assert.AreEqual(
-                new ChatMLTemplate().ComputePrompt(messages),
+                new ChatMLTemplate().ComputePrompt(messages, "assistant"),
                 "<|im_start|>system\nyou are a bot<|im_end|>\n<|im_start|>user\nHello, how are you?<|im_end|>\n<|im_start|>assistant\nI'm doing great. How can I help you today?<|im_end|>\n<|im_start|>user\nI'd like to show off how chat templating works!<|im_end|>\n<|im_start|>assistant\nchat template is awesome<|im_end|>\n<|im_start|>user\ndo you think so?<|im_end|>\n<|im_start|>assistant\n"
             );
         }
@@ -31,7 +30,7 @@ public void TestChatML()
         public void TestMistralInstruct()
         {
             Assert.AreEqual(
-                new MistralInstructTemplate().ComputePrompt(messages),
+                new MistralInstructTemplate().ComputePrompt(messages, "assistant"),
                 "<s>[INST] you are a bot\n\nHello, how are you? [/INST]I'm doing great. How can I help you today?</s>[INST] I'd like to show off how chat templating works! [/INST]chat template is awesome</s>[INST] do you think so? [/INST]"
             );
         }
@@ -40,7 +39,7 @@ public void TestMistralInstruct()
         public void TestMistralChat()
         {
             Assert.AreEqual(
-                new MistralChatTemplate().ComputePrompt(messages),
+                new MistralChatTemplate().ComputePrompt(messages, "assistant"),
                 "<s>[INST] you are a bot\n\n### user: Hello, how are you? [/INST]### assistant: I'm doing great. How can I help you today?</s>[INST] ### user: I'd like to show off how chat templating works! [/INST]### assistant: chat template is awesome</s>[INST] ### user: do you think so? [/INST]### assistant:"
             );
         }
@@ -49,7 +48,7 @@ public void TestMistralChat()
         public void TestLLama2()
         {
             Assert.AreEqual(
-                new LLama2Template().ComputePrompt(messages),
+                new LLama2Template().ComputePrompt(messages, "assistant"),
                 "<s>[INST] <<SYS>>\nyou are a bot\n<</SYS>> Hello, how are you? [/INST]I'm doing great. How can I help you today? </s><s>[INST] I'd like to show off how chat templating works! [/INST]chat template is awesome </s><s>[INST] do you think so? [/INST]"
             );
         }
@@ -58,7 +57,7 @@ public void TestLLama2()
         public void TestLLama2Chat()
         {
             Assert.AreEqual(
-                new LLama2ChatTemplate().ComputePrompt(messages),
+                new LLama2ChatTemplate().ComputePrompt(messages, "assistant"),
                 "<s>[INST] <<SYS>>\nyou are a bot\n<</SYS>> ### user: Hello, how are you? [/INST]### assistant: I'm doing great. How can I help you today? </s><s>[INST] ### user: I'd like to show off how chat templating works! [/INST]### assistant: chat template is awesome </s><s>[INST] ### user: do you think so? [/INST]### assistant:"
             );
         }
@@ -67,7 +66,7 @@ public void TestLLama2Chat()
         public void TestAlpaca()
         {
             Assert.AreEqual(
-                new AlpacaTemplate().ComputePrompt(messages),
+                new AlpacaTemplate().ComputePrompt(messages, "assistant"),
                 "you are a bot\n\n### user: Hello, how are you?\n### assistant: I'm doing great. How can I help you today?\n### user: I'd like to show off how chat templating works!\n### assistant: chat template is awesome\n### user: do you think so?\n### assistant:"
             );
         }
@@ -76,7 +75,7 @@ public void TestAlpaca()
         public void TestPhi2()
         {
             Assert.AreEqual(
-                new Phi2Template().ComputePrompt(messages),
+                new Phi2Template().ComputePrompt(messages, "assistant"),
                 "you are a bot\n\nuser: Hello, how are you?\nassistant: I'm doing great. How can I help you today?\nuser: I'd like to show off how chat templating works!\nassistant: chat template is awesome\nuser: do you think so?\nassistant:"
             );
         }
@@ -85,7 +84,7 @@ public void TestPhi2()
         public void TestZephyr()
         {
             Assert.AreEqual(
-                new ZephyrTemplate().ComputePrompt(messages),
+                new ZephyrTemplate().ComputePrompt(messages, "assistant"),
                 "<|system|>\nyou are a bot</s>\n<|user|>\nHello, how are you?</s>\n<|assistant|>\nI'm doing great. How can I help you today?</s>\n<|user|>\nI'd like to show off how chat templating works!</s>\n<|assistant|>\nchat template is awesome</s>\n<|user|>\ndo you think so?</s>\n<|assistant|>\n"
             );
         }