LLMAgent.cs

Go to the documentation of this file.

 
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading.Tasks;
using UndreamAI.LlamaLib;
using UnityEngine;
 
namespace LLMUnity
{
    [DefaultExecutionOrder(-1)]
    public class LLMAgent : LLMClient
    {
        #region Inspector Fields
        [Tooltip("Filename for saving chat history (saved in persistentDataPath)")]
        [LLM] public string save = "";
 
        [Tooltip("Debug LLM prompts")]
        [LLM] public bool debugPrompt = false;
 
        [Tooltip("Strategy to apply when the conversation history exceeds the model's context window")]
        [LLM] public ContextOverflowStrategy overflowStrategy = ContextOverflowStrategy.Truncate;
 
        [Tooltip("Target fraction of the context window to fill after truncation or summarization (0.0–1.0)")]
        [Overflow1, Range(0.1f, 0.95f)] public float overflowTargetRatio = 0.5f;
 
        [Tooltip("Custom prompt used when asking the LLM to summarize history (leave empty for default)")]
        [Overflow2, TextArea(2, 4)] public string overflowSummarizePrompt = "";
 
        [Tooltip("Server slot to use for processing (affects caching behavior)")]
        [ModelAdvanced, SerializeField] protected int _slot = -1;
 
        [TextArea(5, 10), Chat, SerializeField]
        [Tooltip("System prompt that defines the AI's personality and behavior")]
        protected string _systemPrompt = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.";
        #endregion
 
        #region Public Properties
        public int slot
        {
            get => _slot;
            set
            {
                if (_slot != value)
                {
                    _slot = value;
                    if (llmAgent != null) llmAgent.SlotId = _slot;
                }
            }
        }
 
        public string systemPrompt
        {
            get => _systemPrompt;
            set
            {
                if (_systemPrompt != value)
                {
                    _systemPrompt = value;
                    if (llmAgent != null) llmAgent.SystemPrompt = _systemPrompt;
                }
            }
        }
 
        public UndreamAI.LlamaLib.LLMAgent llmAgent { get; protected set; }
 
        public List<ChatMessage> chat
        {
            get
            {
                if (llmAgent == null) return new List<ChatMessage>();
 
                // convert each UndreamAI.LlamaLib.ChatMessage to LLMUnity.ChatMessage
                return llmAgent.GetHistory()
                    .Select(m => new ChatMessage(m))
                    .ToList();
            }
            set
            {
                if (llmAgent != null)
                {
                    // convert LLMUnity.ChatMessage back to UndreamAI.LlamaLib.ChatMessage
                    var history = value?.Select(m => (UndreamAI.LlamaLib.ChatMessage)m).ToList()
                        ?? new List<UndreamAI.LlamaLib.ChatMessage>();
 
                    llmAgent.SetHistory(history);
                }
            }
        }
        #endregion
 
        #region Unity Lifecycle and Initialization
        public override void Awake()
        {
            if (!remote) llm?.Register(this);
            base.Awake();
        }
 
        protected override async Task SetupCallerObject()
        {
            await base.SetupCallerObject();
 
            string exceptionMessage = "";
            try
            {
                llmAgent = new UndreamAI.LlamaLib.LLMAgent(llmClient, systemPrompt);
            }
            catch (Exception ex)
            {
                exceptionMessage = ex.Message;
            }
            if (llmAgent == null || exceptionMessage != "")
            {
                string error = "LLMAgent not initialized";
                if (exceptionMessage != "") error += ", error: " + exceptionMessage;
                LLMUnitySetup.LogError(error, true);
            }
        }
 
        protected override async Task PostSetupCallerObject()
        {
            await base.PostSetupCallerObject();
            if (slot != -1) llmAgent.SlotId = slot;
            if (overflowStrategy != ContextOverflowStrategy.None)
            {
                string prompt = string.IsNullOrEmpty(overflowSummarizePrompt) ? null : overflowSummarizePrompt;
                llmAgent.SetOverflowStrategy(overflowStrategy, overflowTargetRatio, prompt);
            }
            await InitHistory();
        }
 
        protected override void OnValidate()
        {
            base.OnValidate();
 
            // Validate slot configuration
            if (llm != null && llm.parallelPrompts > -1 && (slot < -1 || slot >= llm.parallelPrompts))
            {
                LLMUnitySetup.LogError($"Slot must be between 0 and {llm.parallelPrompts - 1}, or -1 for auto-assignment");
            }
        }
 
        protected override LLMLocal GetCaller()
        {
            return llmAgent;
        }
 
        protected virtual async Task InitHistory()
        {
            await ClearHistory();
            if (!string.IsNullOrEmpty(save) && File.Exists(GetSavePath()))
            {
                await LoadHistory();
            }
        }
 
        #endregion
 
        #region File Path Management
        public virtual string GetSavePath()
        {
            if (string.IsNullOrEmpty(save))
            {
                LLMUnitySetup.LogError("No save path specified");
                return null;
            }
 
            return Path.Combine(Application.persistentDataPath, save).Replace('\\', '/');
        }
 
        #endregion
 
        #region Chat Management
        public virtual async Task ClearHistory()
        {
            await CheckCaller(checkConnection: false);
            llmAgent.ClearHistory();
        }
 
        public virtual async Task AddUserMessage(string content)
        {
            await CheckCaller();
            llmAgent.AddUserMessage(content);
        }
 
        public virtual async Task AddAssistantMessage(string content)
        {
            await CheckCaller();
            llmAgent.AddAssistantMessage(content);
        }
 
        public string GetSummary()
        {
            return llmAgent?.GetSummary() ?? string.Empty;
        }
 
        public void SetSummary(string summary)
        {
            llmAgent?.SetSummary(summary ?? string.Empty);
        }
 
        #endregion
 
        #region Chat Functionality
        [Serializable]
        public class CompletionResponseJson
        {
            public string prompt;
            public string content;
        }
        public virtual async Task<string> Chat(string query, Action<string> callback = null,
            Action completionCallback = null, bool addToHistory = true)
        {
            await CheckCaller();
            string result = "";
            try
            {
                LlamaLib.CharArrayCallback wrappedCallback = null;
                if (callback != null)
                {
#if ENABLE_IL2CPP
                    // For IL2CPP: wrap to IntPtr callback, then wrap for main thread
                    Action<string> mainThreadCallback = Utils.WrapActionForMainThread(callback, this);
                    wrappedCallback = IL2CPP_Completion.CreateCallback(mainThreadCallback);
#else
                    // For Mono: direct callback wrapping
                    wrappedCallback = Utils.WrapCallbackForAsync(callback, this);
#endif
                }
 
                SetCompletionParameters();
                result = await llmAgent.ChatAsync(query, addToHistory, wrappedCallback, false, debugPrompt);
                if (this == null) return null;
                if (addToHistory && result != null && save != "") _ = SaveHistory();
                if (this != null) completionCallback?.Invoke();
            }
            catch (Exception ex)
            {
                LLMUnitySetup.LogError(ex.Message, true);
            }
            return result;
        }
 
        public virtual async Task Warmup(Action completionCallback = null)
        {
            await Warmup(null, completionCallback);
        }
 
        public virtual async Task Warmup(string query, Action completionCallback = null)
        {
            int originalNumPredict = numPredict;
            try
            {
                // Set to generate no tokens for warmup
                numPredict = 0;
                await Chat(query, null, completionCallback, false);
            }
            finally
            {
                // Restore original setting
                numPredict = originalNumPredict;
                SetCompletionParameters();
            }
        }
 
        #endregion
 
        #region Persistence
        public virtual async Task SaveHistory()
        {
            if (string.IsNullOrEmpty(save))
            {
                LLMUnitySetup.LogError("No save path specified");
                return;
            }
            await CheckCaller();
 
            // Save chat history
            string jsonPath = GetSavePath();
            string directory = Path.GetDirectoryName(jsonPath);
 
            if (!Directory.Exists(directory))
            {
                Directory.CreateDirectory(directory);
            }
 
            try
            {
                llmAgent.SaveHistory(jsonPath);
                LLMUnitySetup.Log($"Saved chat history to: {jsonPath}");
            }
            catch (Exception ex)
            {
                LLMUnitySetup.LogError($"Failed to save chat history to '{jsonPath}': {ex.Message}", true);
            }
        }
 
        public virtual async Task LoadHistory()
        {
            if (string.IsNullOrEmpty(save))
            {
                LLMUnitySetup.LogError("No save path specified");
                return;
            }
            await CheckCaller();
 
            // Load chat history
            string jsonPath = GetSavePath();
            if (!File.Exists(jsonPath))
            {
                LLMUnitySetup.LogError($"Chat history file not found: {jsonPath}");
            }
 
            try
            {
                llmAgent.LoadHistory(jsonPath);
                LLMUnitySetup.Log($"Loaded chat history from: {jsonPath}");
            }
            catch (Exception ex)
            {
                LLMUnitySetup.LogError($"Failed to load chat history from '{jsonPath}': {ex.Message}", true);
            }
        }
 
        #endregion
 
        #region Request Management
        public void CancelRequests()
        {
            llmAgent?.Cancel();
        }
 
        #endregion
    }
 
    public class ChatMessage : UndreamAI.LlamaLib.ChatMessage
    {
        public ChatMessage(string role, string content) : base(role, content) {}
        public ChatMessage(UndreamAI.LlamaLib.ChatMessage other) : base(other.role, other.content) {}
    }
}