AI-Studio/app/MindWork AI Studio/Tools/PluginSystem/Assistants/PluginAssistants.cs

using AIStudio.Tools.PluginSystem.Assistants.DataModel;
using AIStudio.Tools.PluginSystem.Assistants.DataModel.Layout;
using Lua;
using System.Text;

namespace AIStudio.Tools.PluginSystem.Assistants;

public sealed class PluginAssistants(bool isInternal, LuaState state, PluginType type) : PluginBase(isInternal, state, type)
{
    private static string TB(string fallbackEn) => I18N.I.T(fallbackEn, typeof(PluginAssistants).Namespace, nameof(PluginAssistants));

    private static readonly ILogger<PluginAssistants> LOGGER = Program.LOGGER_FACTORY.CreateLogger<PluginAssistants>();

    public AssistantForm? RootComponent { get; private set; }
    public string AssistantTitle { get; private set; } = string.Empty;
    public string AssistantDescription { get; private set; } = string.Empty;
    public string SystemPrompt { get; private set; } = string.Empty;
    public string SubmitText { get; private set; } = string.Empty;
    public bool AllowProfiles { get; private set; } = true;
    public bool HasEmbeddedProfileSelection { get; private set; }
    public bool HasCustomPromptBuilder => this.buildPromptFunction is not null;
    public const int TEXT_AREA_MAX_VALUE = 524288;

    private LuaFunction? buildPromptFunction;

    public void TryLoad()
    {
        if(!this.TryProcessAssistant(out var issue))
            this.pluginIssues.Add(issue);
    }

    /// <summary>
    /// Tries to parse the assistant table into our internal assistant render tree data model. It follows this process:
    /// <list type="number">
    /// <item><description>ASSISTANT ? Title/Description ? UI</description></item>
    /// <item><description>UI: Root element ? required Children ? Components</description></item>
    /// <item><description>Components: Type ? Props ? Children (recursively)</description></item>
    /// </list>
    /// </summary>
    /// <param name="message">The error message, when parameters from the table could not be read.</param>
    /// <returns>True, when the assistant could be read successfully indicating the data model is populated.</returns>
    private bool TryProcessAssistant(out string message)
    {
        message = string.Empty;
        this.HasEmbeddedProfileSelection = false;
        this.buildPromptFunction = null;

        this.RegisterLuaHelpers();
        
        // Ensure that the main ASSISTANT table exists and is a valid Lua table:
        if (!this.state.Environment["ASSISTANT"].TryRead<LuaTable>(out var assistantTable))
        {
            message = TB("The ASSISTANT lua table does not exist or is not a valid table.");
            return false;
        }
        
        if (!assistantTable.TryGetValue("Title", out var assistantTitleValue) ||
            !assistantTitleValue.TryRead<string>(out var assistantTitle))
        {
            message = TB("The provided ASSISTANT lua table does not contain a valid title.");
            return false;
        }

        if (!assistantTable.TryGetValue("Description", out var assistantDescriptionValue) ||
            !assistantDescriptionValue.TryRead<string>(out var assistantDescription))
        {
            message = TB("The provided ASSISTANT lua table does not contain a valid description.");
            return false;
        }
        
        if (!assistantTable.TryGetValue("SystemPrompt", out var assistantSystemPromptValue) ||
            !assistantSystemPromptValue.TryRead<string>(out var assistantSystemPrompt))
        {
            message = TB("The provided ASSISTANT lua table does not contain a valid system prompt.");
            return false;
        }
        
        if (!assistantTable.TryGetValue("SubmitText", out var assistantSubmitTextValue) ||
            !assistantSubmitTextValue.TryRead<string>(out var assistantSubmitText))
        {
            message = TB("The ASSISTANT table does not contain a valid system prompt.");
            return false;
        }
        
        if (!assistantTable.TryGetValue("AllowProfiles", out var assistantAllowProfilesValue) ||
            !assistantAllowProfilesValue.TryRead<bool>(out var assistantAllowProfiles))
        {
            message = TB("The provided ASSISTANT lua table does not contain the boolean flag to control the allowance of profiles.");
            return false;
        }

        if (assistantTable.TryGetValue("BuildPrompt", out var buildPromptValue))
        {
            if (buildPromptValue.TryRead<LuaFunction>(out var buildPrompt))
                this.buildPromptFunction = buildPrompt;
            else
                message = TB("ASSISTANT.BuildPrompt exists but is not a Lua function or has invalid syntax.");
        }

        this.AssistantTitle = assistantTitle;
        this.AssistantDescription = assistantDescription;
        this.SystemPrompt = assistantSystemPrompt;
        this.SubmitText = assistantSubmitText;
        this.AllowProfiles = assistantAllowProfiles;

        // Ensure that the UI table exists nested in the ASSISTANT table and is a valid Lua table:
        if (!assistantTable.TryGetValue("UI", out var uiVal) || !uiVal.TryRead<LuaTable>(out var uiTable))
        {
            message = TB("The provided ASSISTANT lua table does not contain a valid UI table.");
            return false;
        }
        
        if (!this.TryReadRenderTree(uiTable, out var rootComponent))
        {
            message = TB("Failed to parse the UI render tree from the ASSISTANT lua table.");
            return false;
        }

        this.RootComponent = (AssistantForm)rootComponent;
        return true;
    }

    public async Task<string?> TryBuildPromptAsync(LuaTable input, CancellationToken cancellationToken = default)
    {
        if (this.buildPromptFunction is null)
            return null;

        try
        {
            cancellationToken.ThrowIfCancellationRequested();
            var results = await this.state.CallAsync(this.buildPromptFunction, [input]);
            if (results.Length == 0)
                return string.Empty;

            if (results[0].TryRead<string>(out var prompt))
                return prompt;

            LOGGER.LogWarning("ASSISTANT.BuildPrompt returned a non-string value.");
            return string.Empty;
        }
        catch (Exception e)
        {
            LOGGER.LogError(e, "ASSISTANT.BuildPrompt failed to execute.");
            return string.Empty;
        }
    }

    public async Task<LuaTable?> TryInvokeButtonActionAsync(AssistantButton button, LuaTable input, CancellationToken cancellationToken = default)
    {
        return await this.TryInvokeComponentCallbackAsync(button.Action, AssistantComponentType.BUTTON, button.Name, input, cancellationToken);
    }

    public async Task<LuaTable?> TryInvokeSwitchChangedAsync(AssistantSwitch switchComponent, LuaTable input, CancellationToken cancellationToken = default)
    {
        return await this.TryInvokeComponentCallbackAsync(switchComponent.OnChanged, AssistantComponentType.SWITCH, switchComponent.Name, input, cancellationToken);
    }

    private async Task<LuaTable?> TryInvokeComponentCallbackAsync(LuaFunction? callback, AssistantComponentType componentType, string componentName, LuaTable input, CancellationToken cancellationToken = default)
    {
        if (callback is null)
            return null;

        try
        {
            cancellationToken.ThrowIfCancellationRequested();
            var results = await this.state.CallAsync(callback, [input]);
            if (results.Length == 0)
                return null;

            if (results[0].Type is LuaValueType.Nil)
                return null;

            if (results[0].TryRead<LuaTable>(out var updateTable))
                return updateTable;

            LOGGER.LogWarning("Assistant plugin '{PluginName}' {ComponentType} '{ComponentName}' callback returned a non-table value. The result is ignored.", this.Name, componentType, componentName);
            return null;
        }
        catch (Exception e)
        {
            LOGGER.LogError(e, "Assistant plugin '{PluginName}' {ComponentType} '{ComponentName}' callback failed to execute.", this.Name, componentType, componentName);
            return null;
        }
    }

    /// <summary>
    /// Parses the root <c>FORM</c> component and start to parse its required children (main ui components)
    /// </summary>
    /// <param name="uiTable">The <c>LuaTable</c> containing all UI components</param>
    /// <param name="root">Outputs the root <c>FORM</c> component, if the parsing is successful. </param>
    /// <returns>True, when the UI table could be read successfully.</returns>
    private bool TryReadRenderTree(LuaTable uiTable, out IAssistantComponent root)
    {
        root = null!;
        
        if (!uiTable.TryGetValue("Type", out var typeVal)
            || !typeVal.TryRead<string>(out var typeText)
            || !Enum.TryParse<AssistantComponentType>(typeText, true, out var type)
            || type != AssistantComponentType.FORM)
        {
            LOGGER.LogWarning("UI table of the ASSISTANT table has no valid Form type.");
            return false;
        }

        if (!uiTable.TryGetValue("Children", out var childrenVal) ||
            !childrenVal.TryRead<LuaTable>(out var childrenTable))
        {
            LOGGER.LogWarning("Form has no valid Children table.");
            return false;
        }

        var children = new List<IAssistantComponent>();
        var count = childrenTable.ArrayLength;
        for (var idx = 1; idx <= count; idx++)
        {
            var childVal = childrenTable[idx];
            if (!childVal.TryRead<LuaTable>(out var childTable))
            {
                LOGGER.LogWarning($"Child #{idx} is not a table.");
                continue;
            }

            if (!this.TryReadComponentTable(idx, childTable, out var comp))
            {
                LOGGER.LogWarning($"Child #{idx} could not be parsed.");
                continue;
            }

            children.Add(comp);
        }

        root = AssistantComponentFactory.CreateComponent(AssistantComponentType.FORM, new Dictionary<string, object>(), children);
        return true;
    }

    /// <summary>
    /// Parses the components' table containing all members and properties.
    /// Recursively calls itself, if the component has a children table
    /// </summary>
    /// <param name="idx">Current index inside the <c>FORM</c> children</param>
    /// <param name="componentTable">The <c>LuaTable</c> containing all component properties</param>
    /// <param name="component">Outputs the component if the parsing is successful</param>
    /// <returns>True, when the component table could be read successfully.</returns>
    private bool TryReadComponentTable(int idx, LuaTable componentTable, out IAssistantComponent component)
    {
        component = null!;

        if (!componentTable.TryGetValue("Type", out var typeVal)
            || !typeVal.TryRead<string>(out var typeText)
            || !Enum.TryParse<AssistantComponentType>(typeText, true, out var type))
        {
            LOGGER.LogWarning($"Component #{idx} missing valid Type.");
            return false;
        }
        
        if (type == AssistantComponentType.PROFILE_SELECTION)
            this.HasEmbeddedProfileSelection = true;

        Dictionary<string, object> props = new();
        if (componentTable.TryGetValue("Props", out var propsVal)
            && propsVal.TryRead<LuaTable>(out var propsTable))
        {
            if (!this.TryReadComponentProps(type, propsTable, out props))
                LOGGER.LogWarning($"Component #{idx} Props could not be fully read.");
        }
        
        var children = new List<IAssistantComponent>();
        if (componentTable.TryGetValue("Children", out var childVal)
            && childVal.TryRead<LuaTable>(out var childTable))
        {
            var cnt = childTable.ArrayLength;
            for (var i = 1; i <= cnt; i++)
            {
                var cv = childTable[i];
                if (cv.TryRead<LuaTable>(out var ct)
                    && this.TryReadComponentTable(i, ct, out var childComp))
                {
                    children.Add(childComp);
                }
            }
        }

        component = AssistantComponentFactory.CreateComponent(type, props, children);

        if (component is AssistantTextArea textArea)
        {
            if (!string.IsNullOrWhiteSpace(textArea.AdornmentIcon) && !string.IsNullOrWhiteSpace(textArea.AdornmentText))
                LOGGER.LogWarning($"Assistant plugin '{this.Name}' TEXT_AREA '{textArea.Name}' defines both '[\"AdornmentIcon\"]' and '[\"AdornmentText\"]', thus both will be ignored by the renderer. You`re only allowed to use either one of them.");

            if (textArea.MaxLength == 0)
            {
                LOGGER.LogWarning($"Assistant plugin '{this.Name}' TEXT_AREA '{textArea.Name}' defines a MaxLength of `0`. This is not applicable, if you want a readonly Textfield, set the [\"ReadOnly\"] field to `true`. MAXLENGTH IS SET TO DEFAULT {TEXT_AREA_MAX_VALUE}.");
                textArea.MaxLength = TEXT_AREA_MAX_VALUE;
            }
            
            if (textArea.MaxLength != 0 && textArea.MaxLength != TEXT_AREA_MAX_VALUE)
                textArea.Counter = textArea.MaxLength;
            
            if (textArea.Counter != null)
                textArea.IsImmediate = true;
        }

        if (component is AssistantButtonGroup buttonGroup)
        {
            var invalidChildren = buttonGroup.Children.Where(child => child.Type != AssistantComponentType.BUTTON).ToList();
            if (invalidChildren.Count > 0)
            {
                LOGGER.LogWarning("Assistant plugin '{PluginName}' BUTTON_GROUP contains non-BUTTON children. Only BUTTON children are supported and invalid children are ignored.", this.Name);
                buttonGroup.Children = buttonGroup.Children.Where(child => child.Type == AssistantComponentType.BUTTON).ToList();
            }
        }

        if (component is AssistantGrid grid)
        {
            var invalidChildren = grid.Children.Where(child => child.Type != AssistantComponentType.LAYOUT_ITEM).ToList();
            if (invalidChildren.Count > 0)
            {
                LOGGER.LogWarning("Assistant plugin '{PluginName}' LAYOUT_GRID contains non-LAYOUT_ITEM children. Only LAYOUT_ITEM children are supported and invalid children are ignored.", this.Name);
                grid.Children = grid.Children.Where(child => child.Type == AssistantComponentType.LAYOUT_ITEM).ToList();
            }
        }

        return true;
    }

    private bool TryReadComponentProps(AssistantComponentType type, LuaTable propsTable, out Dictionary<string, object> props)
    {
        props = new Dictionary<string, object>();

        if (!ComponentPropSpecs.SPECS.TryGetValue(type, out var spec))
        {
            LOGGER.LogWarning($"No PropSpec defined for component type {type}");
            return false;
        }

        foreach (var key in spec.Required)
        {
            if (!propsTable.TryGetValue(key, out var luaVal))
            {
                LOGGER.LogWarning($"Component {type} missing required prop '{key}'.");
                return false;
            }
            if (!this.TryConvertComponentPropValue(type, key, luaVal, out var dotNetVal))
            {
                LOGGER.LogWarning($"Component {type}: prop '{key}' has wrong type.");
                return false;
            }
            props[key] = dotNetVal;
        }

        foreach (var key in spec.Optional)
        {
            if (!propsTable.TryGetValue(key, out var luaVal))
                continue;

            if (!this.TryConvertComponentPropValue(type, key, luaVal, out var dotNetVal))
            {
                LOGGER.LogWarning($"Component {type}: optional prop '{key}' has wrong type, skipping.");
                continue;
            }
            props[key] = dotNetVal;
        }

        return true;
    }

    private bool TryConvertComponentPropValue(AssistantComponentType type, string key, LuaValue val, out object result)
    {
        if (type == AssistantComponentType.BUTTON && key == "Action" && val.TryRead<LuaFunction>(out var action))
        {
            result = action;
            return true;
        }
        
        if (type == AssistantComponentType.SWITCH && key == "OnChanged" && val.TryRead<LuaFunction>(out var onChanged))
        {
            result = onChanged;
            return true;
        }

        return this.TryConvertLuaValue(val, out result);
    }
    
    private bool TryConvertLuaValue(LuaValue val, out object result)
    {
        if (val.TryRead<string>(out var s))
        {
            result = s;
            return true;
        }
        if (val.TryRead<bool>(out var b))
        {
            result = b;
            return true;
        }
        if (val.TryRead<double>(out var d))
        {
            result = d;
            return true;
        }
        
        if (val.TryRead<LuaTable>(out var table) && this.TryParseDropdownItem(table, out var item))
        {
            result = item;
            return true;
        }

        if (val.TryRead<LuaTable>(out var listTable) && this.TryParseDropdownItemList(listTable, out var itemList))
        {
            result = itemList;
            return true;
        }

        if (val.TryRead<LuaTable>(out var listItemListTable) && this.TryParseListItemList(listItemListTable, out var listItemList))
        {
            result = listItemList;
            return true;
        }
        
        result = null!;
        return false;
    }
    
    private bool TryParseDropdownItem(LuaTable table, out AssistantDropdownItem item)
    {
        item = new AssistantDropdownItem();

        if (!table.TryGetValue("Value", out var valueVal) || !valueVal.TryRead<string>(out var value))
            return false;

        if (!table.TryGetValue("Display", out var displayVal) || !displayVal.TryRead<string>(out var display))
            return false;

        item.Value = value;
        item.Display = display;
        return true;
    }

    
    private bool TryParseDropdownItemList(LuaTable table, out List<AssistantDropdownItem> items)
    {
        items = new List<AssistantDropdownItem>();

        var length = table.ArrayLength;
        for (var i = 1; i <= length; i++)
        {
            var value = table[i];

            if (value.TryRead<LuaTable>(out var subTable) && this.TryParseDropdownItem(subTable, out var item))
            {
                items.Add(item);
            }
            else
            {
                items = null!;
                return false;
            }
        }

        return true;
    }

    private bool TryParseListItem(LuaTable table, out AssistantListItem item)
    {
        item = new AssistantListItem();

        if (!table.TryGetValue("Text", out var textVal) || !textVal.TryRead<string>(out var text))
            return false;
        
        if (!table.TryGetValue("Type", out var typeVal) || !typeVal.TryRead<string>(out var type))
            return false;
        
        table.TryGetValue("Icon", out var iconVal);
        iconVal.TryRead<string>(out var icon);
        icon ??= string.Empty;

        table.TryGetValue("IconColor", out var iconColorVal);
        iconColorVal.TryRead<string>(out var iconColor);
        iconColor ??= string.Empty;

        
        item.Text = text;
        item.Type = type;
        item.Icon = icon;
        item.IconColor = iconColor;

        if (table.TryGetValue("Href", out var hrefVal) && hrefVal.TryRead<string>(out var href))
        {
            item.Href = href;
        }

        return true;
    }
    
    private bool TryParseListItemList(LuaTable table, out List<AssistantListItem> items)
    {
        items = new List<AssistantListItem>();

        var length = table.ArrayLength;
        for (var i = 1; i <= length; i++)
        {
            var value = table[i];

            if (value.TryRead<LuaTable>(out var subTable) && this.TryParseListItem(subTable, out var item))
            {
                items.Add(item);
            }
            else
            {
                items = null!;
                return false;
            }
        }

        return true;
    }

    private void RegisterLuaHelpers()
    {
        
        this.state.Environment["LogInfo"] = new LuaFunction((context, _) =>
        {
            if (context.ArgumentCount == 0) return new(0);
            
            var message = context.GetArgument<string>(0);
            LOGGER.LogInformation($"[Lua] [Assistants] [{this.Name}]: {message}");
            return new(0);
        });
        
        this.state.Environment["LogDebug"] = new LuaFunction((context, _) =>
        {
            if (context.ArgumentCount == 0) return new(0);
            
            var message = context.GetArgument<string>(0);
            LOGGER.LogDebug($"[Lua] [Assistants] [{this.Name}]: {message}");
            return new(0);
        });
        
        this.state.Environment["LogWarning"] = new LuaFunction((context, _) =>
        {
            if (context.ArgumentCount == 0) return new(0);
            
            var message = context.GetArgument<string>(0);
            LOGGER.LogWarning($"[Lua] [Assistants] [{this.Name}]: {message}");
            return new(0);
        });
        
        this.state.Environment["LogError"] = new LuaFunction((context, _) =>
        {
            if (context.ArgumentCount == 0) return new(0);
            
            var message = context.GetArgument<string>(0);
            LOGGER.LogError($"[Lua] [Assistants] [{this.Name}]: {message}");
            return new(0);
        });

        this.state.Environment["DateTime"] = new LuaFunction((context, _) =>
        {
            var format = context.ArgumentCount > 0 ? context.GetArgument<string>(0) : "yyyy-MM-dd HH:mm:ss";
            var now = DateTime.Now;
            var formattedDate = now.ToString(format);
            
            var table = new LuaTable
            {
                ["year"] = now.Year,
                ["month"] = now.Month,
                ["day"] = now.Day,
                ["hour"] = now.Hour,
                ["minute"] = now.Minute,
                ["second"] = now.Second,
                ["millisecond"] = now.Millisecond,
                ["formatted"] = formattedDate,
            };
            return new(context.Return(table));
        });
        
        this.state.Environment["Timestamp"] = new LuaFunction((context, _) =>
        {
            var timestamp = DateTime.UtcNow.ToString("o");
            return new(context.Return(timestamp));
        });
    }
}