LLMCharacter.cs

Go to the documentation of this file.

 
using System;
using System.Collections.Generic;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
using UnityEditor;
using UnityEngine;
using UnityEngine.Networking;
 
namespace LLMUnity
{
    [DefaultExecutionOrder(-2)]
    public class LLMCharacter : MonoBehaviour
    {
        [HideInInspector] public bool advancedOptions = false;
        [LocalRemote] public bool remote = false;
        [Local] public LLM llm;
        [Remote] public string host = "localhost";
        [Remote] public int port = 13333;
        [Remote] public int numRetries = 10;
        [Remote] public string APIKey;
        [LLM] public string save = "";
        [LLM] public bool saveCache = false;
        [LLM] public bool debugPrompt = false;
        [Model] public bool stream = true;
        [ModelAdvanced] public string grammar = null;
        [ModelAdvanced] public bool cachePrompt = true;
        [ModelAdvanced] public int slot = -1;
        [ModelAdvanced] public int seed = 0;
        [ModelAdvanced] public int numPredict = 256;
        [ModelAdvanced, Float(0f, 2f)] public float temperature = 0.2f;
        [ModelAdvanced, Int(-1, 100)] public int topK = 40;
        [ModelAdvanced, Float(0f, 1f)] public float topP = 0.9f;
        [ModelAdvanced, Float(0f, 1f)] public float minP = 0.05f;
        [ModelAdvanced, Float(0f, 2f)] public float repeatPenalty = 1.1f;
        [ModelAdvanced, Float(0f, 1f)] public float presencePenalty = 0f;
        [ModelAdvanced, Float(0f, 1f)] public float frequencyPenalty = 0f;
 
        [ModelAdvanced, Float(0f, 1f)] public float tfsZ = 1f;
        [ModelAdvanced, Float(0f, 1f)] public float typicalP = 1f;
        [ModelAdvanced, Int(0, 2048)] public int repeatLastN = 64;
        [ModelAdvanced] public bool penalizeNl = true;
        [ModelAdvanced] public string penaltyPrompt;
        [ModelAdvanced, Int(0, 2)] public int mirostat = 0;
        [ModelAdvanced, Float(0f, 10f)] public float mirostatTau = 5f;
        [ModelAdvanced, Float(0f, 1f)] public float mirostatEta = 0.1f;
        [ModelAdvanced, Int(0, 10)] public int nProbs = 0;
        [ModelAdvanced] public bool ignoreEos = false;
 
        public int nKeep = -1;
        public List<string> stop = new List<string>();
        public Dictionary<int, string> logitBias = null;
 
        [Chat] public string playerName = "user";
        [Chat] public string AIName = "assistant";
        [TextArea(5, 10), Chat] public string prompt = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.";
        public bool setNKeepToPrompt = true;
 
        public List<ChatMessage> chat;
        private SemaphoreSlim chatLock = new SemaphoreSlim(1, 1);
        private string chatTemplate;
        private ChatTemplate template = null;
        public string grammarString;
        private List<(string, string)> requestHeaders;
        private List<UnityWebRequest> WIPRequests = new List<UnityWebRequest>();
 
        public void Awake()
        {
            // Start the LLM server in a cross-platform way
            if (!enabled) return;
 
            requestHeaders = new List<(string, string)> { ("Content-Type", "application/json") };
            if (!remote)
            {
                AssignLLM();
                if (llm == null)
                {
                    LLMUnitySetup.LogError($"No LLM assigned or detected for LLMCharacter {name}!");
                    return;
                }
                int slotFromServer = llm.Register(this);
                if (slot == -1) slot = slotFromServer;
            }
            else
            {
                if (!String.IsNullOrEmpty(APIKey)) requestHeaders.Add(("Authorization", "Bearer " + APIKey));
            }
 
            InitGrammar();
            InitHistory();
        }
 
        void OnValidate()
        {
            AssignLLM();
            if (llm != null && llm.parallelPrompts > -1 && (slot < -1 || slot >= llm.parallelPrompts)) LLMUnitySetup.LogError($"The slot needs to be between 0 and {llm.parallelPrompts-1}, or -1 to be automatically set");
        }
 
        void Reset()
        {
            AssignLLM();
        }
 
        void AssignLLM()
        {
            if (remote || llm != null) return;
 
            LLM[] existingLLMs = FindObjectsOfType<LLM>();
            if (existingLLMs.Length == 0) return;
 
            SortBySceneAndHierarchy(existingLLMs);
            llm = existingLLMs[0];
            string msg = $"Assigning LLM {llm.name} to LLMCharacter {name}";
            if (llm.gameObject.scene != gameObject.scene) msg += $" from scene {llm.gameObject.scene}";
            LLMUnitySetup.Log(msg);
        }
 
        void SortBySceneAndHierarchy(LLM[] array)
        {
            for (int i = 0; i < array.Length - 1; i++)
            {
                bool swapped = false;
                for (int j = 0; j < array.Length - i - 1; j++)
                {
                    bool sameScene = array[j].gameObject.scene == array[j + 1].gameObject.scene;
                    bool swap = (
                        (!sameScene && array[j + 1].gameObject.scene == gameObject.scene) ||
                        (sameScene && array[j].transform.GetSiblingIndex() > array[j + 1].transform.GetSiblingIndex())
                    );
                    if (swap)
                    {
                        LLM temp = array[j];
                        array[j] = array[j + 1];
                        array[j + 1] = temp;
                        swapped = true;
                    }
                }
                if (!swapped) break;
            }
        }
 
        protected void InitHistory()
        {
            InitPrompt();
            _ = LoadHistory();
        }
 
        protected async Task LoadHistory()
        {
            if (save == "" || !File.Exists(GetJsonSavePath(save))) return;
            await chatLock.WaitAsync(); // Acquire the lock
            try
            {
                await Load(save);
            }
            finally
            {
                chatLock.Release(); // Release the lock
            }
        }
 
        public virtual string GetSavePath(string filename)
        {
            return Path.Combine(Application.persistentDataPath, filename).Replace('\\', '/');
        }
 
        public virtual string GetJsonSavePath(string filename)
        {
            return GetSavePath(filename + ".json");
        }
 
        public virtual string GetCacheSavePath(string filename)
        {
            return GetSavePath(filename + ".cache");
        }
 
        private void InitPrompt(bool clearChat = true)
        {
            if (chat != null)
            {
                if (clearChat) chat.Clear();
            }
            else
            {
                chat = new List<ChatMessage>();
            }
            ChatMessage promptMessage = new ChatMessage { role = "system", content = prompt };
            if (chat.Count == 0)
            {
                chat.Add(promptMessage);
            }
            else
            {
                chat[0] = promptMessage;
            }
        }
 
        public void SetPrompt(string newPrompt, bool clearChat = true)
        {
            prompt = newPrompt;
            nKeep = -1;
            InitPrompt(clearChat);
        }
 
        private bool CheckTemplate()
        {
            if (template == null)
            {
                LLMUnitySetup.LogError("Template not set!");
                return false;
            }
            return true;
        }
 
        private async Task<bool> InitNKeep()
        {
            if (setNKeepToPrompt && nKeep == -1)
            {
                if (!CheckTemplate()) return false;
                string systemPrompt = template.ComputePrompt(new List<ChatMessage>(){chat[0]}, playerName, "", false);
                List<int> tokens = await Tokenize(systemPrompt);
                if (tokens == null) return false;
                SetNKeep(tokens);
            }
            return true;
        }
 
        private void InitGrammar()
        {
            if (grammar != null && grammar != "")
            {
                grammarString = File.ReadAllText(LLMUnitySetup.GetAssetPath(grammar));
            }
        }
 
        private void SetNKeep(List<int> tokens)
        {
            // set the tokens to keep
            nKeep = tokens.Count;
        }
 
        public async Task LoadTemplate()
        {
            string llmTemplate;
            if (remote)
            {
                llmTemplate = await AskTemplate();
            }
            else
            {
                llmTemplate = llm.GetTemplate();
            }
            if (llmTemplate != chatTemplate)
            {
                chatTemplate = llmTemplate;
                template = chatTemplate == null ? null : ChatTemplate.GetTemplate(chatTemplate);
                nKeep = -1;
            }
        }
 
        public async void SetGrammar(string path)
        {
#if UNITY_EDITOR
            if (!EditorApplication.isPlaying) path = LLMUnitySetup.AddAsset(path);
#endif
            await LLMUnitySetup.AndroidExtractAsset(path, true);
            grammar = path;
            InitGrammar();
        }
 
        List<string> GetStopwords()
        {
            if (!CheckTemplate()) return null;
            List<string> stopAll = new List<string>(template.GetStop(playerName, AIName));
            if (stop != null) stopAll.AddRange(stop);
            return stopAll;
        }
 
        ChatRequest GenerateRequest(string prompt)
        {
            // setup the request struct
            ChatRequest chatRequest = new ChatRequest();
            if (debugPrompt) LLMUnitySetup.Log(prompt);
            chatRequest.prompt = prompt;
            chatRequest.id_slot = slot;
            chatRequest.temperature = temperature;
            chatRequest.top_k = topK;
            chatRequest.top_p = topP;
            chatRequest.min_p = minP;
            chatRequest.n_predict = numPredict;
            chatRequest.n_keep = nKeep;
            chatRequest.stream = stream;
            chatRequest.stop = GetStopwords();
            chatRequest.tfs_z = tfsZ;
            chatRequest.typical_p = typicalP;
            chatRequest.repeat_penalty = repeatPenalty;
            chatRequest.repeat_last_n = repeatLastN;
            chatRequest.penalize_nl = penalizeNl;
            chatRequest.presence_penalty = presencePenalty;
            chatRequest.frequency_penalty = frequencyPenalty;
            chatRequest.penalty_prompt = (penaltyPrompt != null && penaltyPrompt != "") ? penaltyPrompt : null;
            chatRequest.mirostat = mirostat;
            chatRequest.mirostat_tau = mirostatTau;
            chatRequest.mirostat_eta = mirostatEta;
            chatRequest.grammar = grammarString;
            chatRequest.seed = seed;
            chatRequest.ignore_eos = ignoreEos;
            chatRequest.logit_bias = logitBias;
            chatRequest.n_probs = nProbs;
            chatRequest.cache_prompt = cachePrompt;
            return chatRequest;
        }
 
        public void AddMessage(string role, string content)
        {
            // add the question / answer to the chat list, update prompt
            chat.Add(new ChatMessage { role = role, content = content });
        }
 
        public void AddPlayerMessage(string content)
        {
            AddMessage(playerName, content);
        }
 
        public void AddAIMessage(string content)
        {
            AddMessage(AIName, content);
        }
 
        protected string ChatContent(ChatResult result)
        {
            // get content from a chat result received from the endpoint
            return result.content.Trim();
        }
 
        protected string MultiChatContent(MultiChatResult result)
        {
            // get content from a chat result received from the endpoint
            string response = "";
            foreach (ChatResult resultPart in result.data)
            {
                response += resultPart.content;
            }
            return response.Trim();
        }
 
        async Task<string> CompletionRequest(string json, Callback<string> callback = null)
        {
            string result = "";
            if (stream)
            {
                result = await PostRequest<MultiChatResult, string>(json, "completion", MultiChatContent, callback);
            }
            else
            {
                result = await PostRequest<ChatResult, string>(json, "completion", ChatContent, callback);
            }
            return result;
        }
 
        protected string TemplateContent(TemplateResult result)
        {
            // get content from a char result received from the endpoint in open AI format
            return result.template;
        }
 
        protected List<int> TokenizeContent(TokenizeResult result)
        {
            // get the tokens from a tokenize result received from the endpoint
            return result.tokens;
        }
 
        protected string DetokenizeContent(TokenizeRequest result)
        {
            // get content from a chat result received from the endpoint
            return result.content;
        }
 
        protected List<float> EmbeddingsContent(EmbeddingsResult result)
        {
            // get content from a chat result received from the endpoint
            return result.embedding;
        }
 
        protected string SlotContent(SlotResult result)
        {
            // get the tokens from a tokenize result received from the endpoint
            return result.filename;
        }
 
        public async Task<string> Chat(string query, Callback<string> callback = null, EmptyCallback completionCallback = null, bool addToHistory = true)
        {
            // handle a chat message by the user
            // call the callback function while the answer is received
            // call the completionCallback function when the answer is fully received
            await LoadTemplate();
            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 result = await CompletionRequest(json, callback);
 
            if (addToHistory && result != null)
            {
                await chatLock.WaitAsync();
                try
                {
                    AddPlayerMessage(query);
                    AddAIMessage(result);
                }
                finally
                {
                    chatLock.Release();
                }
                if (save != "") _ = Save(save);
            }
 
            completionCallback?.Invoke();
            return result;
        }
 
        public async Task<string> Complete(string prompt, Callback<string> callback = null, EmptyCallback completionCallback = null)
        {
            // handle a completion request by the user
            // call the callback function while the answer is received
            // call the completionCallback function when the answer is fully received
            await LoadTemplate();
 
            string json = JsonUtility.ToJson(GenerateRequest(prompt));
            string result = await CompletionRequest(json, callback);
            completionCallback?.Invoke();
            return result;
        }
 
        public async Task Warmup(EmptyCallback completionCallback = null)
        {
            await LoadTemplate();
            if (!CheckTemplate()) return;
            if (!await InitNKeep()) return;
 
            string prompt = template.ComputePrompt(chat, playerName, AIName);
            ChatRequest request = GenerateRequest(prompt);
            request.n_predict = 0;
            string json = JsonUtility.ToJson(request);
            await CompletionRequest(json);
            completionCallback?.Invoke();
        }
 
        public async Task<string> AskTemplate()
        {
            return await PostRequest<TemplateResult, string>("{}", "template", TemplateContent);
        }
 
        public async Task<List<int>> Tokenize(string query, Callback<List<int>> callback = null)
        {
            // handle the tokenization of a message by the user
            TokenizeRequest tokenizeRequest = new TokenizeRequest();
            tokenizeRequest.content = query;
            string json = JsonUtility.ToJson(tokenizeRequest);
            return await PostRequest<TokenizeResult, List<int>>(json, "tokenize", TokenizeContent, callback);
        }
 
        public async Task<string> Detokenize(List<int> tokens, Callback<string> callback = null)
        {
            // handle the detokenization of a message by the user
            TokenizeResult tokenizeRequest = new TokenizeResult();
            tokenizeRequest.tokens = tokens;
            string json = JsonUtility.ToJson(tokenizeRequest);
            return await PostRequest<TokenizeRequest, string>(json, "detokenize", DetokenizeContent, callback);
        }
 
        public async Task<List<float>> Embeddings(string query, Callback<List<float>> callback = null)
        {
            // handle the tokenization of a message by the user
            TokenizeRequest tokenizeRequest = new TokenizeRequest();
            tokenizeRequest.content = query;
            string json = JsonUtility.ToJson(tokenizeRequest);
            return await PostRequest<EmbeddingsResult, List<float>>(json, "embeddings", EmbeddingsContent, callback);
        }
 
        protected async Task<string> Slot(string filepath, string action)
        {
            SlotRequest slotRequest = new SlotRequest();
            slotRequest.id_slot = slot;
            slotRequest.filepath = filepath;
            slotRequest.action = action;
            string json = JsonUtility.ToJson(slotRequest);
            return await PostRequest<SlotResult, string>(json, "slots", SlotContent);
        }
 
        public virtual async Task<string> Save(string filename)
        {
            string filepath = GetJsonSavePath(filename);
            string dirname = Path.GetDirectoryName(filepath);
            if (!Directory.Exists(dirname)) Directory.CreateDirectory(dirname);
            string json = JsonUtility.ToJson(new ChatListWrapper { chat = chat.GetRange(1, chat.Count - 1) });
            File.WriteAllText(filepath, json);
 
            string cachepath = GetCacheSavePath(filename);
            if (remote || !saveCache) return null;
            string result = await Slot(cachepath, "save");
            return result;
        }
 
        public virtual async Task<string> Load(string filename)
        {
            string filepath = GetJsonSavePath(filename);
            if (!File.Exists(filepath))
            {
                LLMUnitySetup.LogError($"File {filepath} does not exist.");
                return null;
            }
            string json = File.ReadAllText(filepath);
            List<ChatMessage> chatHistory = JsonUtility.FromJson<ChatListWrapper>(json).chat;
            InitPrompt(true);
            chat.AddRange(chatHistory);
            LLMUnitySetup.Log($"Loaded {filepath}");
 
            string cachepath = GetCacheSavePath(filename);
            if (remote || !saveCache || !File.Exists(GetSavePath(cachepath))) return null;
            string result = await Slot(cachepath, "restore");
            return result;
        }
 
        protected Ret ConvertContent<Res, Ret>(string response, ContentCallback<Res, Ret> getContent = null)
        {
            // template function to convert the json received and get the content
            if (response == null) return default;
            response = response.Trim();
            if (response.StartsWith("data: "))
            {
                string responseArray = "";
                foreach (string responsePart in response.Replace("\n\n", "").Split("data: "))
                {
                    if (responsePart == "") continue;
                    if (responseArray != "") responseArray += ",\n";
                    responseArray += responsePart;
                }
                response = $"{{\"data\": [{responseArray}]}}";
            }
            return getContent(JsonUtility.FromJson<Res>(response));
        }
 
        protected void CancelRequestsLocal()
        {
            if (slot >= 0) llm.CancelRequest(slot);
        }
 
        protected void CancelRequestsRemote()
        {
            foreach (UnityWebRequest request in WIPRequests)
            {
                request.Abort();
            }
            WIPRequests.Clear();
        }
 
        // <summary>
        public void CancelRequests()
        {
            if (remote) CancelRequestsRemote();
            else CancelRequestsLocal();
        }
 
        protected async Task<Ret> PostRequestLocal<Res, Ret>(string json, string endpoint, ContentCallback<Res, Ret> getContent, Callback<Ret> callback = null)
        {
            // send a post request to the server and call the relevant callbacks to convert the received content and handle it
            // this function has streaming functionality i.e. handles the answer while it is being received
            string callResult = null;
            bool callbackCalled = false;
            while (!llm.failed && !llm.started) await Task.Yield();
            switch (endpoint)
            {
                case "tokenize":
                    callResult = await llm.Tokenize(json);
                    break;
                case "detokenize":
                    callResult = await llm.Detokenize(json);
                    break;
                case "embeddings":
                    callResult = await llm.Embeddings(json);
                    break;
                case "slots":
                    callResult = await llm.Slot(json);
                    break;
                case "completion":
                    Callback<string> callbackString = null;
                    if (stream && callback != null)
                    {
                        if (typeof(Ret) == typeof(string))
                        {
                            callbackString = (strArg) =>
                            {
                                callback(ConvertContent(strArg, getContent));
                            };
                        }
                        else
                        {
                            LLMUnitySetup.LogError($"wrong callback type, should be string");
                        }
                        callbackCalled = true;
                    }
                    callResult = await llm.Completion(json, callbackString);
                    break;
                default:
                    LLMUnitySetup.LogError($"Unknown endpoint {endpoint}");
                    break;
            }
 
            Ret result = ConvertContent(callResult, getContent);
            if (!callbackCalled) callback?.Invoke(result);
            return result;
        }
 
        protected async Task<Ret> PostRequestRemote<Res, Ret>(string json, string endpoint, ContentCallback<Res, Ret> getContent, Callback<Ret> callback = null)
        {
            // send a post request to the server and call the relevant callbacks to convert the received content and handle it
            // this function has streaming functionality i.e. handles the answer while it is being received
            if (endpoint == "slots")
            {
                LLMUnitySetup.LogError("Saving and loading is not currently supported in remote setting");
                return default;
            }
 
            Ret result = default;
            byte[] jsonToSend = new System.Text.UTF8Encoding().GetBytes(json);
            UnityWebRequest request = null;
            string error = null;
            int tryNr = numRetries;
 
            while (tryNr != 0)
            {
                using (request = UnityWebRequest.Put($"{host}:{port}/{endpoint}", jsonToSend))
                {
                    WIPRequests.Add(request);
 
                    request.method = "POST";
                    if (requestHeaders != null)
                    {
                        for (int i = 0; i < requestHeaders.Count; i++)
                            request.SetRequestHeader(requestHeaders[i].Item1, requestHeaders[i].Item2);
                    }
 
                    // Start the request asynchronously
                    var asyncOperation = request.SendWebRequest();
                    float lastProgress = 0f;
                    // Continue updating progress until the request is completed
                    while (!asyncOperation.isDone)
                    {
                        float currentProgress = request.downloadProgress;
                        // Check if progress has changed
                        if (currentProgress != lastProgress && callback != null)
                        {
                            try
                            {
                                callback?.Invoke(ConvertContent(request.downloadHandler.text, getContent));
                            }
                            catch (Exception) {}
                            lastProgress = currentProgress;
                        }
                        // Wait for the next frame
                        await Task.Yield();
                    }
                    WIPRequests.Remove(request);
                    if (request.result == UnityWebRequest.Result.Success)
                    {
                        result = ConvertContent(request.downloadHandler.text, getContent);
                        error = null;
                        break;
                    }
                    else
                    {
                        result = default;
                        error = request.error;
                        if (request.responseCode == (int)System.Net.HttpStatusCode.Unauthorized) break;
                    }
                }
                tryNr--;
                if (tryNr > 0) await Task.Delay(200 * (numRetries - tryNr));
            }
 
            if (error != null) LLMUnitySetup.LogError(error);
            callback?.Invoke(result);
            return result;
        }
 
        protected async Task<Ret> PostRequest<Res, Ret>(string json, string endpoint, ContentCallback<Res, Ret> getContent, Callback<Ret> callback = null)
        {
            if (remote) return await PostRequestRemote(json, endpoint, getContent, callback);
            return await PostRequestLocal(json, endpoint, getContent, callback);
        }
    }
 
    [Serializable]
    public class ChatListWrapper
    {
        public List<ChatMessage> chat;
    }
}