AI-Studio/app/MindWork AI Studio/Dialogs/AssistantPluginRevisionDialog.razor.cs

using System.Text;
using AIStudio.Agents.AssistantAudit;
using AIStudio.Components;
using AIStudio.Provider;
using AIStudio.Tools.PluginSystem;
using AIStudio.Tools.PluginSystem.Assistants;
using AIStudio.Tools.Services;
using Microsoft.AspNetCore.Components;

namespace AIStudio.Dialogs;

public partial class AssistantPluginRevisionDialog : MSGComponentBase
{
    private const string PLUGIN_FILE_NAME = "plugin.lua";
    private static readonly ILogger LOGGER = Program.LOGGER_FACTORY.CreateLogger(nameof(AssistantPluginRevisionDialog));

    [CascadingParameter]
    private IMudDialogInstance MudDialog { get; set; } = null!;

    [Inject]
    private AssistantPluginGenerationService AssistantPluginGenerationService { get; init; } = null!;

    [Inject]
    private PluginInstallService PluginInstallService { get; init; } = null!;

    [Inject]
    private AssistantPluginAuditService AssistantPluginAuditService { get; init; } = null!;

    [Parameter]
    public Guid PluginId { get; set; }

    [Parameter]
    public string PluginLocalPath { get; set; } = string.Empty;

    [Parameter]
    public string TestContext { get; set; } = string.Empty;

    private IAvailablePlugin? availablePlugin;
    private PluginAssistants? assistantPlugin;
    private AIStudio.Settings.Provider providerSettings = AIStudio.Settings.Provider.NONE;
    private string pluginFile = string.Empty;
    private string currentLua = string.Empty;
    private string changeRequest = string.Empty;
    private string revisedLua = string.Empty;
    private string revisedPluginName = string.Empty;
    private string issue = string.Empty;
    private AssistantPluginCheckResult? revisionCheckResult;
    private bool isLoading = true;
    private bool isGenerating;
    private bool isApplying;
    private bool isAuditing;

    private bool CanGenerate => this.assistantPlugin is not null &&
                                !this.isLoading &&
                                !this.isGenerating &&
                                !this.isApplying &&
                                !string.IsNullOrWhiteSpace(this.changeRequest);

    private bool CanApply => this.availablePlugin is not null &&
                             this.assistantPlugin is not null &&
                             !this.isGenerating &&
                             !this.isApplying &&
                             !this.isAuditing &&
                             this.revisionCheckResult?.Success is true &&
                             !string.IsNullOrWhiteSpace(this.revisedLua);

    protected override async Task OnInitializedAsync()
    {
        try
        {
            this.providerSettings = this.SettingsManager.GetPreselectedProvider(Tools.Components.META_ASSISTANT);
            this.availablePlugin = PluginFactory.AvailablePlugins
                .OfType<IAvailablePlugin>()
                .FirstOrDefault(x => x.Id == this.PluginId && AreSamePath(x.LocalPath, this.PluginLocalPath));

            this.assistantPlugin = PluginFactory.RunningPlugins
                .OfType<PluginAssistants>()
                .FirstOrDefault(x => x.Id == this.PluginId && AreSamePath(x.PluginPath, this.PluginLocalPath));

            if (this.availablePlugin is null || this.assistantPlugin is null)
            {
                this.issue = T("The assistant plugin could not be resolved.");
                return;
            }

            if (!CanReviseAssistantPlugin(this.availablePlugin, this.assistantPlugin))
            {
                this.issue = T("Only locally managed assistant plugins can be revised with AI.");
                return;
            }

            this.pluginFile = Path.Join(this.availablePlugin.LocalPath, PLUGIN_FILE_NAME);
            if (!File.Exists(this.pluginFile))
            {
                this.issue = T("The plugin.lua file could not be found.");
                return;
            }

            this.currentLua = await File.ReadAllTextAsync(this.pluginFile, Encoding.UTF8);
        }
        catch (Exception e)
        {
            this.issue = string.Format(T("The assistant plugin could not be loaded: {0}"), e.Message);
        }
        finally
        {
            this.isLoading = false;
        }

        await base.OnInitializedAsync();
    }

    private async Task GenerateRevisionAsync()
    {
        if (!this.CanGenerate || this.assistantPlugin is null)
            return;

        this.isGenerating = true;
        this.issue = string.Empty;
        this.revisedLua = string.Empty;
        this.revisionCheckResult = null;
        await this.InvokeAsync(this.StateHasChanged);

        try
        {
            var draft = await this.AssistantPluginGenerationService.GenerateRevisionAsync(
                this.assistantPlugin,
                this.currentLua,
                this.changeRequest,
                this.providerSettings,
                this.TestContext,
                CancellationToken.None);

            if (!draft.Success)
            {
                this.issue = draft.Issue;
                return;
            }

            this.revisedLua = draft.Lua;
            this.revisedPluginName = draft.PluginName;
            if (this.availablePlugin is null)
                return;

            this.revisionCheckResult = await this.PluginInstallService.CheckInstalledAssistantUpdateAsync(this.availablePlugin, this.revisedLua, CancellationToken.None);
            if (this.revisionCheckResult.Success)
                return;

            this.issue = this.revisionCheckResult.Issue;
        }
        finally
        {
            this.isGenerating = false;
            await this.InvokeAsync(this.StateHasChanged);
        }
    }

    private async Task ApplyRevisionAsync()
    {
        if (!this.CanApply || this.availablePlugin is null)
            return;

        this.isApplying = true;
        this.issue = string.Empty;
        await this.InvokeAsync(this.StateHasChanged);

        try
        {
            var result = await this.PluginInstallService.UpdateInstalledAssistantAsync(this.availablePlugin, this.revisedLua, CancellationToken.None);
            if (!result.Success)
            {
                LOGGER.LogError($"Failed to revise assistant plugin '{result.PluginName}' ({result.PluginId}) in '{result.PluginDirectory}' with issue '{result.Issue}'.");
                this.issue = result.Issue;
                return;
            }

            PluginAssistantAudit? audit = null;
            if (this.SettingsManager.ConfigurationData.AssistantPluginAudit.AutomaticallyAuditAssistants)
                audit = await this.TryRunAuditAsync(result.PluginId);

            this.MudDialog.Close(DialogResult.Ok(new AssistantPluginRevisionDialogResult(result.PluginId, result.PluginName, audit)));
        }
        finally
        {
            this.isApplying = false;
            if (!string.IsNullOrWhiteSpace(this.issue))
                await this.InvokeAsync(this.StateHasChanged);
        }
    }

    private async Task<PluginAssistantAudit?> TryRunAuditAsync(Guid pluginId)
    {
        var updatedPlugin = PluginFactory.RunningPlugins.OfType<PluginAssistants>().FirstOrDefault(x => x.Id == pluginId);
        if (updatedPlugin is null)
            return null;

        this.isAuditing = true;
        await this.InvokeAsync(this.StateHasChanged);
        try
        {
            var audit = await this.AssistantPluginAuditService.RunAuditAsync(updatedPlugin);
            if (audit.Level is AssistantAuditLevel.UNKNOWN)
                return audit;

            UpsertAudit(this.SettingsManager.ConfigurationData.AssistantPluginAudits, audit);
            await this.SettingsManager.StoreSettings();
            return audit;
        }
        finally
        {
            this.isAuditing = false;
        }
    }

    private string? ValidatingProvider(AIStudio.Settings.Provider provider)
    {
        if (provider.UsedLLMProvider == LLMProviders.NONE)
            return T("Please select a provider.");

        return null;
    }

    private void Cancel() => this.MudDialog.Cancel();

    private static bool CanReviseAssistantPlugin(IAvailablePlugin availablePlugin, PluginAssistants assistantPlugin) =>
        availablePlugin is { IsInternal: false, IsManagedByConfigServer: false, Type: PluginType.ASSISTANT } &&
        !string.IsNullOrWhiteSpace(availablePlugin.LocalPath) &&
        assistantPlugin is { IsInternal: false, IsManagedByConfigServer: false };

    private static void UpsertAudit(IList<PluginAssistantAudit> audits, PluginAssistantAudit audit)
    {
        var existingIndex = audits.ToList().FindIndex(x => x.PluginId == audit.PluginId);
        if (existingIndex >= 0)
            audits[existingIndex] = audit;
        else
            audits.Add(audit);
    }

    private static bool AreSamePath(string left, string right)
    {
        if (string.IsNullOrWhiteSpace(left) || string.IsNullOrWhiteSpace(right))
            return false;

        var comparison = OperatingSystem.IsWindows()
            ? StringComparison.OrdinalIgnoreCase
            : StringComparison.Ordinal;

        return string.Equals(
            Path.GetFullPath(left).TrimEnd(Path.DirectorySeparatorChar, Path.AltDirectorySeparatorChar),
            Path.GetFullPath(right).TrimEnd(Path.DirectorySeparatorChar, Path.AltDirectorySeparatorChar),
            comparison);
    }
}