LLM.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;
 
namespace LLMUnity
{
    [DefaultExecutionOrder(-1)]
    public class LLM : MonoBehaviour
    {
        [HideInInspector] public bool advancedOptions = false;
        [LocalRemote] public bool remote = false;
        [Remote] public int port = 13333;
        [LLM] public int numThreads = -1;
        [LLM] public int numGPULayers = 0;
        [LLM] public bool debug = false;
        [LLMAdvanced] public int parallelPrompts = -1;
        [LLMAdvanced] public bool dontDestroyOnLoad = true;
        [DynamicRange("minContextLength", "maxContextLength", false), Model] public int contextSize = 8192;
        [ModelAdvanced] public int batchSize = 512;
        public bool started { get; protected set; } = false;
        public bool failed { get; protected set; } = false;
        public static bool modelSetupFailed { get; protected set; } = false;
        public static bool modelSetupComplete { get; protected set; } = false;
 
        [ModelAdvanced] public string model = "";
        [ModelAdvanced] public string chatTemplate = ChatTemplate.DefaultTemplate;
        [ModelAdvanced] public string lora = "";
        [ModelAdvanced] public string loraWeights = "";
        [ModelExtras] public bool flashAttention = false;
 
        public string APIKey;
        // SSL certificate
        [SerializeField]
        private string SSLCert = "";
        public string SSLCertPath = "";
        // SSL key
        [SerializeField]
        private string SSLKey = "";
        public string SSLKeyPath = "";
 
        public int minContextLength = 0;
        public int maxContextLength = 0;
 
        IntPtr LLMObject = IntPtr.Zero;
        List<LLMCaller> clients = new List<LLMCaller>();
        LLMLib llmlib;
        StreamWrapper logStreamWrapper = null;
        Thread llmThread = null;
        List<StreamWrapper> streamWrappers = new List<StreamWrapper>();
        public LLMManager llmManager = new LLMManager();
        private readonly object startLock = new object();
        static readonly object staticLock = new object();
        public LoraManager loraManager = new LoraManager();
        string loraPre = "";
        string loraWeightsPre = "";
        public bool embeddingsOnly = false;
        public int embeddingLength = 0;
 
 
        public LLM()
        {
            LLMManager.Register(this);
        }
 
        void OnValidate()
        {
            if (lora != loraPre || loraWeights != loraWeightsPre)
            {
                loraManager.FromStrings(lora, loraWeights);
                (loraPre, loraWeightsPre) = (lora, loraWeights);
            }
        }
 
        public async void Awake()
        {
            if (!enabled) return;
#if !UNITY_EDITOR
            modelSetupFailed = !await LLMManager.Setup();
#endif
            modelSetupComplete = true;
            if (modelSetupFailed)
            {
                failed = true;
                return;
            }
            string arguments = GetLlamaccpArguments();
            if (arguments == null)
            {
                failed = true;
                return;
            }
            await Task.Run(() => StartLLMServer(arguments));
            if (!started) return;
            if (dontDestroyOnLoad) DontDestroyOnLoad(transform.root.gameObject);
        }
 
        public async Task WaitUntilReady()
        {
            while (!started) await Task.Yield();
        }
 
        public static async Task<bool> WaitUntilModelSetup(Callback<float> downloadProgressCallback = null)
        {
            if (downloadProgressCallback != null) LLMManager.downloadProgressCallbacks.Add(downloadProgressCallback);
            while (!modelSetupComplete) await Task.Yield();
            return !modelSetupFailed;
        }
 
        public static string GetLLMManagerAsset(string path)
        {
#if UNITY_EDITOR
            if (!EditorApplication.isPlaying) return GetLLMManagerAssetEditor(path);
#endif
            return GetLLMManagerAssetRuntime(path);
        }
 
        public static string GetLLMManagerAssetEditor(string path)
        {
            // empty
            if (string.IsNullOrEmpty(path)) return path;
            // LLMManager - return location the file will be stored in StreamingAssets
            ModelEntry modelEntry = LLMManager.Get(path);
            if (modelEntry != null) return modelEntry.filename;
            // StreamingAssets - return relative location within StreamingAssets
            string assetPath = LLMUnitySetup.GetAssetPath(path); // Note: this will return the full path if a full path is passed
            string basePath = LLMUnitySetup.GetAssetPath();
            if (File.Exists(assetPath))
            {
                if (LLMUnitySetup.IsSubPath(assetPath, basePath)) return LLMUnitySetup.RelativePath(assetPath, basePath);
            }
            // full path
            if (!File.Exists(assetPath))
            {
                LLMUnitySetup.LogError($"Model {path} was not found.");
            }
            else
            {
                string errorMessage = $"The model {path} was loaded locally. You can include it in the build in one of these ways:";
                errorMessage += $"\n-Copy the model inside the StreamingAssets folder and use its StreamingAssets path";
                errorMessage += $"\n-Load the model with the model manager inside the LLM GameObject and use its filename";
                LLMUnitySetup.LogWarning(errorMessage);
            }
            return path;
        }
 
        public static string GetLLMManagerAssetRuntime(string path)
        {
            // empty
            if (string.IsNullOrEmpty(path)) return path;
            // LLMManager
            string managerPath = LLMManager.GetAssetPath(path);
            if (!string.IsNullOrEmpty(managerPath) && File.Exists(managerPath)) return managerPath;
            // StreamingAssets
            string assetPath = LLMUnitySetup.GetAssetPath(path);
            if (File.Exists(assetPath)) return assetPath;
            // download path
            assetPath = LLMUnitySetup.GetDownloadAssetPath(path);
            if (File.Exists(assetPath)) return assetPath;
            // give up
            return path;
        }
 
 
        public void SetModel(string path)
        {
            model = GetLLMManagerAsset(path);
            if (!string.IsNullOrEmpty(model))
            {
                ModelEntry modelEntry = LLMManager.Get(model);
                if (modelEntry == null) modelEntry = new ModelEntry(GetLLMManagerAssetRuntime(model));
                SetTemplate(modelEntry.chatTemplate);
 
                maxContextLength = modelEntry.contextLength;
                if (contextSize > maxContextLength) contextSize = maxContextLength;
                SetEmbeddings(modelEntry.embeddingLength, modelEntry.embeddingOnly);
                if (contextSize == 0 && modelEntry.contextLength > 32768)
                {
                    LLMUnitySetup.LogWarning($"The model {path} has very large context size ({modelEntry.contextLength}), consider setting it to a smaller value (<=32768) to avoid filling up the RAM");
                }
            }
#if UNITY_EDITOR
            if (!EditorApplication.isPlaying) EditorUtility.SetDirty(this);
#endif
        }
 
        public void SetLora(string path, float weight = 1)
        {
            AssertNotStarted();
            loraManager.Clear();
            AddLora(path, weight);
        }
 
        public void AddLora(string path, float weight = 1)
        {
            AssertNotStarted();
            loraManager.Add(path, weight);
            UpdateLoras();
        }
 
        public void RemoveLora(string path)
        {
            AssertNotStarted();
            loraManager.Remove(path);
            UpdateLoras();
        }
 
        public void RemoveLoras()
        {
            AssertNotStarted();
            loraManager.Clear();
            UpdateLoras();
        }
 
        public void SetLoraWeight(string path, float weight)
        {
            loraManager.SetWeight(path, weight);
            UpdateLoras();
            if (started) ApplyLoras();
        }
 
        public void SetLoraWeights(Dictionary<string, float> loraToWeight)
        {
            foreach (KeyValuePair<string, float> entry in loraToWeight) loraManager.SetWeight(entry.Key, entry.Value);
            UpdateLoras();
            if (started) ApplyLoras();
        }
 
        public void UpdateLoras()
        {
            (lora, loraWeights) = loraManager.ToStrings();
            (loraPre, loraWeightsPre) = (lora, loraWeights);
#if UNITY_EDITOR
            if (!EditorApplication.isPlaying) EditorUtility.SetDirty(this);
#endif
        }
 
        public void SetTemplate(string templateName, bool setDirty = true)
        {
            chatTemplate = templateName;
            if (started) llmlib?.LLM_SetTemplate(LLMObject, chatTemplate);
#if UNITY_EDITOR
            if (setDirty && !EditorApplication.isPlaying) EditorUtility.SetDirty(this);
#endif
        }
 
        public void SetEmbeddings(int embeddingLength, bool embeddingsOnly)
        {
            this.embeddingsOnly = embeddingsOnly;
            this.embeddingLength = embeddingLength;
#if UNITY_EDITOR
            if (!EditorApplication.isPlaying) EditorUtility.SetDirty(this);
#endif
        }
 
 
        string ReadFileContents(string path)
        {
            if (String.IsNullOrEmpty(path)) return "";
            else if (!File.Exists(path))
            {
                LLMUnitySetup.LogError($"File {path} not found!");
                return "";
            }
            return File.ReadAllText(path);
        }
 
 
        public void SetSSLCert(string path)
        {
            SSLCertPath = path;
            SSLCert = ReadFileContents(path);
        }
 
        public void SetSSLKey(string path)
        {
            SSLKeyPath = path;
            SSLKey = ReadFileContents(path);
        }
 
        public string GetTemplate()
        {
            return chatTemplate;
        }
 
        protected virtual string GetLlamaccpArguments()
        {
            // Start the LLM server in a cross-platform way
            if ((SSLCert != "" && SSLKey == "") || (SSLCert == "" && SSLKey != ""))
            {
                LLMUnitySetup.LogError($"Both SSL certificate and key need to be provided!");
                return null;
            }
 
            if (model == "")
            {
                LLMUnitySetup.LogError("No model file provided!");
                return null;
            }
            string modelPath = GetLLMManagerAssetRuntime(model);
            if (!File.Exists(modelPath))
            {
                LLMUnitySetup.LogError($"File {modelPath} not found!");
                return null;
            }
 
            loraManager.FromStrings(lora, loraWeights);
            string loraArgument = "";
            foreach (string lora in loraManager.GetLoras())
            {
                string loraPath = GetLLMManagerAssetRuntime(lora);
                if (!File.Exists(loraPath))
                {
                    LLMUnitySetup.LogError($"File {loraPath} not found!");
                    return null;
                }
                loraArgument += $" --lora \"{loraPath}\"";
            }
 
            int numThreadsToUse = numThreads;
            if (Application.platform == RuntimePlatform.Android && numThreads <= 0) numThreadsToUse = LLMUnitySetup.AndroidGetNumBigCores();
 
            int slots = GetNumClients();
            string arguments = $"-m \"{modelPath}\" -c {contextSize} -b {batchSize} --log-disable -np {slots}";
            if (embeddingsOnly) arguments += " --embedding";
            if (numThreadsToUse > 0) arguments += $" -t {numThreadsToUse}";
            arguments += loraArgument;
            arguments += $" -ngl {numGPULayers}";
            if (LLMUnitySetup.FullLlamaLib && flashAttention) arguments += $" --flash-attn";
            if (remote)
            {
                arguments += $" --port {port} --host 0.0.0.0";
                if (!String.IsNullOrEmpty(APIKey)) arguments += $" --api-key {APIKey}";
            }
 
            // the following is the equivalent for running from command line
            string serverCommand;
            if (Application.platform == RuntimePlatform.WindowsEditor || Application.platform == RuntimePlatform.WindowsPlayer) serverCommand = "undreamai_server.exe";
            else serverCommand = "./undreamai_server";
            serverCommand += " " + arguments;
            serverCommand += $" --template \"{chatTemplate}\"";
            if (remote && SSLCert != "" && SSLKey != "") serverCommand += $" --ssl-cert-file {SSLCertPath} --ssl-key-file {SSLKeyPath}";
            LLMUnitySetup.Log($"Deploy server command: {serverCommand}");
            return arguments;
        }
 
        private void SetupLogging()
        {
            logStreamWrapper = ConstructStreamWrapper(LLMUnitySetup.LogWarning, true);
            llmlib?.Logging(logStreamWrapper.GetStringWrapper());
        }
 
        private void StopLogging()
        {
            if (logStreamWrapper == null) return;
            llmlib?.StopLogging();
            DestroyStreamWrapper(logStreamWrapper);
        }
 
        private void StartLLMServer(string arguments)
        {
            started = false;
            failed = false;
            bool useGPU = numGPULayers > 0;
 
            foreach (string arch in LLMLib.PossibleArchitectures(useGPU))
            {
                string error;
                try
                {
                    InitLib(arch);
                    InitService(arguments);
                    LLMUnitySetup.Log($"Using architecture: {arch}");
                    break;
                }
                catch (LLMException e)
                {
                    error = e.Message;
                    Destroy();
                }
                catch (DestroyException)
                {
                    break;
                }
                catch (Exception e)
                {
                    error = $"{e.GetType()}: {e.Message}";
                }
                LLMUnitySetup.Log($"Tried architecture: {arch}, error: " + error);
            }
            if (llmlib == null)
            {
                LLMUnitySetup.LogError("LLM service couldn't be created");
                failed = true;
                return;
            }
            CallWithLock(StartService);
            LLMUnitySetup.Log("LLM service created");
        }
 
        private void InitLib(string arch)
        {
            llmlib = new LLMLib(arch);
            CheckLLMStatus(false);
        }
 
        void CallWithLock(EmptyCallback fn)
        {
            lock (startLock)
            {
                if (llmlib == null) throw new DestroyException();
                fn();
            }
        }
 
        private void InitService(string arguments)
        {
            lock (staticLock)
            {
                if (debug) CallWithLock(SetupLogging);
                CallWithLock(() => { LLMObject = llmlib.LLM_Construct(arguments); });
                CallWithLock(() => llmlib.LLM_SetTemplate(LLMObject, chatTemplate));
                if (remote)
                {
                    if (SSLCert != "" && SSLKey != "")
                    {
                        LLMUnitySetup.Log("Using SSL");
                        CallWithLock(() => llmlib.LLM_SetSSL(LLMObject, SSLCert, SSLKey));
                    }
                    CallWithLock(() => llmlib.LLM_StartServer(LLMObject));
                }
                CallWithLock(() => CheckLLMStatus(false));
            }
        }
 
        private void StartService()
        {
            llmThread = new Thread(() => llmlib.LLM_Start(LLMObject));
            llmThread.Start();
            while (!llmlib.LLM_Started(LLMObject)) {}
            ApplyLoras();
            started = true;
        }
 
        public int Register(LLMCaller llmCaller)
        {
            clients.Add(llmCaller);
            int index = clients.IndexOf(llmCaller);
            if (parallelPrompts != -1) return index % parallelPrompts;
            return index;
        }
 
        protected int GetNumClients()
        {
            return Math.Max(parallelPrompts == -1 ? clients.Count : parallelPrompts, 1);
        }
 
        public delegate void LLMStatusCallback(IntPtr LLMObject, IntPtr stringWrapper);
        public delegate void LLMNoInputReplyCallback(IntPtr LLMObject, IntPtr stringWrapper);
        public delegate void LLMReplyCallback(IntPtr LLMObject, string json_data, IntPtr stringWrapper);
 
        StreamWrapper ConstructStreamWrapper(Callback<string> streamCallback = null, bool clearOnUpdate = false)
        {
            StreamWrapper streamWrapper = new StreamWrapper(llmlib, streamCallback, clearOnUpdate);
            streamWrappers.Add(streamWrapper);
            return streamWrapper;
        }
 
        void DestroyStreamWrapper(StreamWrapper streamWrapper)
        {
            streamWrappers.Remove(streamWrapper);
            streamWrapper.Destroy();
        }
 
        public void Update()
        {
            foreach (StreamWrapper streamWrapper in streamWrappers) streamWrapper.Update();
        }
 
        void AssertStarted()
        {
            string error = null;
            if (failed) error = "LLM service couldn't be created";
            else if (!started) error = "LLM service not started";
            if (error != null)
            {
                LLMUnitySetup.LogError(error);
                throw new Exception(error);
            }
        }
 
        void AssertNotStarted()
        {
            if (started)
            {
                string error = "This method can't be called when the LLM has started";
                LLMUnitySetup.LogError(error);
                throw new Exception(error);
            }
        }
 
        void CheckLLMStatus(bool log = true)
        {
            if (llmlib == null) { return; }
            IntPtr stringWrapper = llmlib.StringWrapper_Construct();
            int status = llmlib.LLM_Status(LLMObject, stringWrapper);
            string result = llmlib.GetStringWrapperResult(stringWrapper);
            llmlib.StringWrapper_Delete(stringWrapper);
            string message = $"LLM {status}: {result}";
            if (status > 0)
            {
                if (log) LLMUnitySetup.LogError(message);
                throw new LLMException(message, status);
            }
            else if (status < 0)
            {
                if (log) LLMUnitySetup.LogWarning(message);
            }
        }
 
        async Task<string> LLMNoInputReply(LLMNoInputReplyCallback callback)
        {
            AssertStarted();
            IntPtr stringWrapper = llmlib.StringWrapper_Construct();
            await Task.Run(() => callback(LLMObject, stringWrapper));
            string result = llmlib?.GetStringWrapperResult(stringWrapper);
            llmlib?.StringWrapper_Delete(stringWrapper);
            CheckLLMStatus();
            return result;
        }
 
        async Task<string> LLMReply(LLMReplyCallback callback, string json)
        {
            AssertStarted();
            IntPtr stringWrapper = llmlib.StringWrapper_Construct();
            await Task.Run(() => callback(LLMObject, json, stringWrapper));
            string result = llmlib?.GetStringWrapperResult(stringWrapper);
            llmlib?.StringWrapper_Delete(stringWrapper);
            CheckLLMStatus();
            return result;
        }
 
        public async Task<string> Tokenize(string json)
        {
            AssertStarted();
            LLMReplyCallback callback = (IntPtr LLMObject, string jsonData, IntPtr strWrapper) =>
            {
                llmlib.LLM_Tokenize(LLMObject, jsonData, strWrapper);
            };
            return await LLMReply(callback, json);
        }
 
        public async Task<string> Detokenize(string json)
        {
            AssertStarted();
            LLMReplyCallback callback = (IntPtr LLMObject, string jsonData, IntPtr strWrapper) =>
            {
                llmlib.LLM_Detokenize(LLMObject, jsonData, strWrapper);
            };
            return await LLMReply(callback, json);
        }
 
        public async Task<string> Embeddings(string json)
        {
            AssertStarted();
            LLMReplyCallback callback = (IntPtr LLMObject, string jsonData, IntPtr strWrapper) =>
            {
                llmlib.LLM_Embeddings(LLMObject, jsonData, strWrapper);
            };
            return await LLMReply(callback, json);
        }
 
        public void ApplyLoras()
        {
            LoraWeightRequestList loraWeightRequest = new LoraWeightRequestList();
            loraWeightRequest.loraWeights = new List<LoraWeightRequest>();
            float[] weights = loraManager.GetWeights();
            if (weights.Length == 0) return;
            for (int i = 0; i < weights.Length; i++)
            {
                loraWeightRequest.loraWeights.Add(new LoraWeightRequest() { id = i, scale = weights[i] });
            }
 
            string json = JsonUtility.ToJson(loraWeightRequest);
            int startIndex = json.IndexOf("[");
            int endIndex = json.LastIndexOf("]") + 1;
            json = json.Substring(startIndex, endIndex - startIndex);
 
            IntPtr stringWrapper = llmlib.StringWrapper_Construct();
            llmlib.LLM_LoraWeight(LLMObject, json, stringWrapper);
            llmlib.StringWrapper_Delete(stringWrapper);
        }
 
        public async Task<List<LoraWeightResult>> ListLoras()
        {
            AssertStarted();
            LLMNoInputReplyCallback callback = (IntPtr LLMObject, IntPtr strWrapper) =>
            {
                llmlib.LLM_LoraList(LLMObject, strWrapper);
            };
            string json = await LLMNoInputReply(callback);
            if (String.IsNullOrEmpty(json)) return null;
            LoraWeightResultList loraRequest = JsonUtility.FromJson<LoraWeightResultList>("{\"loraWeights\": " + json + "}");
            return loraRequest.loraWeights;
        }
 
        public async Task<string> Slot(string json)
        {
            AssertStarted();
            LLMReplyCallback callback = (IntPtr LLMObject, string jsonData, IntPtr strWrapper) =>
            {
                llmlib.LLM_Slot(LLMObject, jsonData, strWrapper);
            };
            return await LLMReply(callback, json);
        }
 
        public async Task<string> Completion(string json, Callback<string> streamCallback = null)
        {
            AssertStarted();
            if (streamCallback == null) streamCallback = (string s) => {};
            StreamWrapper streamWrapper = ConstructStreamWrapper(streamCallback);
            await Task.Run(() => llmlib.LLM_Completion(LLMObject, json, streamWrapper.GetStringWrapper()));
            if (!started) return null;
            streamWrapper.Update();
            string result = streamWrapper.GetString();
            DestroyStreamWrapper(streamWrapper);
            CheckLLMStatus();
            return result;
        }
 
        public void CancelRequest(int id_slot)
        {
            AssertStarted();
            llmlib?.LLM_Cancel(LLMObject, id_slot);
            CheckLLMStatus();
        }
 
        public void Destroy()
        {
            lock (staticLock)
                lock (startLock)
                {
                    try
                    {
                        if (llmlib != null)
                        {
                            if (LLMObject != IntPtr.Zero)
                            {
                                llmlib.LLM_Stop(LLMObject);
                                if (remote) llmlib.LLM_StopServer(LLMObject);
                                StopLogging();
                                llmThread?.Join();
                                llmlib.LLM_Delete(LLMObject);
                                LLMObject = IntPtr.Zero;
                            }
                            llmlib.Destroy();
                            llmlib = null;
                        }
                        started = false;
                        failed = false;
                    }
                    catch (Exception e)
                    {
                        LLMUnitySetup.LogError(e.Message);
                    }
                }
        }
 
        public void OnDestroy()
        {
            Destroy();
            LLMManager.Unregister(this);
        }
    }
}