From 6f58e064734c89fb715d59667f7a2821e7674898 Mon Sep 17 00:00:00 2001
From: Antonis Makropoulos <benuix@gmail.com>
Date: Tue, 21 Jan 2025 11:37:02 +0200
Subject: [PATCH] add warm-up for provided prompt

---
 Runtime/LLMCharacter.cs | 67 ++++++++++++++++++++++++++++-------------
 1 file changed, 46 insertions(+), 21 deletions(-)

diff --git a/Runtime/LLMCharacter.cs b/Runtime/LLMCharacter.cs
index e22e8604..fbf97ac7 100644
--- a/Runtime/LLMCharacter.cs
+++ b/Runtime/LLMCharacter.cs
@@ -416,6 +416,24 @@ protected virtual async Task<string> CompletionRequest(string json, Callback<str
             return result;
         }
 
+        protected async Task<ChatRequest> PromptWithQuery(string query)
+        {
+            ChatRequest result = default;
+            await chatLock.WaitAsync();
+            try
+            {
+                AddPlayerMessage(query);
+                string prompt = template.ComputePrompt(chat, playerName, AIName);
+                result = GenerateRequest(prompt);
+                chat.RemoveAt(chat.Count - 1);
+            }
+            finally
+            {
+                chatLock.Release();
+            }
+            return result;
+        }
+
         /// <summary>
         /// Chat functionality of the LLM.
         /// It calls the LLM completion based on the provided query including the previous chat history.
@@ -436,20 +454,7 @@ public virtual async Task<string> Chat(string query, Callback<string> callback =
             if (!CheckTemplate()) return null;
             if (!await InitNKeep()) return null;
 
-            string json;
-            await chatLock.WaitAsync();
-            try
-            {
-                AddPlayerMessage(query);
-                string prompt = template.ComputePrompt(chat, playerName, AIName);
-                json = JsonUtility.ToJson(GenerateRequest(prompt));
-                chat.RemoveAt(chat.Count - 1);
-            }
-            finally
-            {
-                chatLock.Release();
-            }
-
+            string json = JsonUtility.ToJson(await PromptWithQuery(query));
             string result = await CompletionRequest(json, callback);
 
             if (addToHistory && result != null)
@@ -494,23 +499,43 @@ public virtual async Task<string> Complete(string prompt, Callback<string> callb
         }
 
         /// <summary>
-        /// Allow to warm-up a model by processing the prompt.
+        /// Allow to warm-up a model by processing the system prompt.
         /// The prompt processing will be cached (if cachePrompt=true) allowing for faster initialisation.
-        /// The function allows callback for when the prompt is processed and the response received.
-        ///
-        /// The function calls the Chat function with a predefined query without adding it to history.
+        /// The function allows a callback function for when the prompt is processed and the response received.
         /// </summary>
         /// <param name="completionCallback">callback function called when the full response has been received</param>
-        /// <param name="query">user prompt used during the initialisation (not added to history)</param>
         /// <returns>the LLM response</returns>
         public virtual async Task Warmup(EmptyCallback completionCallback = null)
+        {
+            await Warmup(null, completionCallback);
+        }
+
+        /// <summary>
+        /// Allow to warm-up a model by processing the provided prompt without adding it to history.
+        /// The prompt processing will be cached (if cachePrompt=true) allowing for faster initialisation.
+        /// The function allows a callback function for when the prompt is processed and the response received.
+        ///
+        /// </summary>
+        /// <param name="query">user prompt used during the initialisation (not added to history)</param>
+        /// <param name="completionCallback">callback function called when the full response has been received</param>
+        /// <returns>the LLM response</returns>
+        public virtual async Task Warmup(string query, EmptyCallback completionCallback = null)
         {
             await LoadTemplate();
             if (!CheckTemplate()) return;
             if (!await InitNKeep()) return;
 
-            string prompt = template.ComputePrompt(chat, playerName, AIName);
-            ChatRequest request = GenerateRequest(prompt);
+            ChatRequest request;
+            if (String.IsNullOrEmpty(query))
+            {
+                string prompt = template.ComputePrompt(chat, playerName, AIName);
+                request = GenerateRequest(prompt);
+            }
+            else
+            {
+                request = await PromptWithQuery(query);
+            }
+
             request.n_predict = 0;
             string json = JsonUtility.ToJson(request);
             await CompletionRequest(json);