swarm-assistent/SwarmAssistentExtension.cs

using System;
using System.IO;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using Newtonsoft.Json.Linq;
using SwarmUI.Accounts;
using SwarmUI.Core;
using SwarmUI.Utils;
using SwarmUI.WebAPI;
using System.Net.Http;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Krea 2 collaborative assistant: Ollama chat + vision + prompt/LoRA/params patches + Generate/Civitai.</summary>
public partial class SwarmAssistentExtension : Extension
{
    public static PermInfo PermUse = Permissions.Register(new(
        "swarm_assistent_use",
        "[Swarm Assistent] Use",
        "Allows using the Swarm Assistent chat (Ollama proxy).",
        PermissionDefault.USER,
        Permissions.GroupUser));

    public static HttpClient HttpClient;

    public AssistentConfig Config;
    public AssistentMemory Memory;

    public override void OnPreInit()
    {
        ScriptFiles.Add("Assets/assistent.bundle.js");
        StyleSheetFiles.Add("Assets/assistent.css");
        ExtensionAuthor = "mrleo1nid";
        Description = "Collaborative Krea 2 assistant: Ollama chat, persona presets, vector memory, model cards, Generate loop.";
        License = "MIT";
        Version = "0.15.3";
        Tags = ["tabs", "ui", "llm", "ollama", "krea", "inpaint", "memory", "training", "heard", "qlora"];
    }

    public override void OnInit()
    {
        HttpClient ??= new HttpClient { Timeout = TimeSpan.FromMinutes(10) };
        Config = new AssistentConfig(FilePath, DataRoot());
        Memory = new AssistentMemory(DataRoot(), HttpClient, Config.LoadAssistant(Config.DefaultPersonaId())["embed_model"]?.ToString() ?? "nomic-embed-text");
        API.RegisterAPICall(AssistentListModels, false, PermUse);
        API.RegisterAPICall(AssistentListPersonas, false, PermUse);
        API.RegisterAPICall(AssistentGetConfig, false, PermUse);
        API.RegisterAPICall(AssistentSaveSettings, true, PermUse);
        API.RegisterAPICall(AssistentListInventory, false, PermUse);
        API.RegisterAPICall(AssistentChat, true, PermUse);
        API.RegisterAPICall(AssistentChatWS, true, PermUse);
        API.RegisterAPICall(AssistentListChats, false, PermUse);
        API.RegisterAPICall(AssistentGetChat, false, PermUse);
        API.RegisterAPICall(AssistentSaveChat, true, PermUse);
        API.RegisterAPICall(AssistentDeleteChat, true, PermUse);
        API.RegisterAPICall(AssistentGetUiState, false, PermUse);
        API.RegisterAPICall(AssistentSaveUiState, true, PermUse);
        API.RegisterAPICall(AssistentParkLlm, true, PermUse);
        API.RegisterAPICall(AssistentWarmLlm, true, PermUse);
        API.RegisterAPICall(AssistentListMemory, false, PermUse);
        API.RegisterAPICall(AssistentUpsertMemory, true, PermUse);
        API.RegisterAPICall(AssistentForgetMemory, true, PermUse);
        API.RegisterAPICall(AssistentSearchMemory, false, PermUse);
        API.RegisterAPICall(AssistentGetMemory, false, PermUse);
        API.RegisterAPICall(AssistentLookupTags, false, PermUse);
        API.RegisterAPICall(AssistentLookupExamples, false, PermUse);
        API.RegisterAPICall(AssistentSaveControls, true, PermUse);
        API.RegisterAPICall(AssistentGetPersonaShelves, false, PermUse);
        API.RegisterAPICall(AssistentClonePersona, true, PermUse);
        API.RegisterAPICall(AssistentSavePersona, true, PermUse);
        API.RegisterAPICall(AssistentDeletePersona, true, PermUse);
        API.RegisterAPICall(AssistentExportPersona, false, PermUse);
        API.RegisterAPICall(AssistentImportPersona, true, PermUse);
        API.RegisterAPICall(AssistentSaveKnobs, true, PermUse);
        API.RegisterAPICall(AssistentListUserPrefs, false, PermUse);
        API.RegisterAPICall(AssistentUpsertUserPref, true, PermUse);
        API.RegisterAPICall(AssistentForgetUserPref, true, PermUse);
        API.RegisterAPICall(AssistentClearUserPrefs, true, PermUse);
        API.RegisterAPICall(AssistentClearMemory, true, PermUse);
        API.RegisterAPICall(AssistentListTrainSamples, false, PermUse);
        API.RegisterAPICall(AssistentUpsertTrainSample, true, PermUse);
        API.RegisterAPICall(AssistentDeleteTrainSample, true, PermUse);
        API.RegisterAPICall(AssistentBuildDatasetFromChats, false, PermUse);
        API.RegisterAPICall(AssistentImportDataset, true, PermUse);
        API.RegisterAPICall(AssistentExportDataset, false, PermUse);
        API.RegisterAPICall(AssistentCreateOllamaModel, true, PermUse);
        API.RegisterAPICall(AssistentSearchHfDatasets, false, PermUse);
        API.RegisterAPICall(AssistentCheckHfDataset, false, PermUse);
        API.RegisterAPICall(AssistentPreviewHfDataset, false, PermUse);
        API.RegisterAPICall(AssistentImportHfDataset, true, PermUse);
        API.RegisterAPICall(AssistentStartTrainJob, true, PermUse);
        API.RegisterAPICall(AssistentCancelTrainJob, true, PermUse);
        API.RegisterAPICall(AssistentGetTrainJob, false, PermUse);
        API.RegisterAPICall(AssistentTrainWS, true, PermUse);
        API.RegisterAPICall(AssistentSaveRunnerSettings, true, PermUse);
        API.RegisterAPICall(AssistentGetRunnerSettings, false, PermUse);
        API.RegisterAPICall(AssistentGetDatasetAgentSettings, false, PermUse);
        API.RegisterAPICall(AssistentSaveDatasetAgentSettings, true, PermUse);
        API.RegisterAPICall(AssistentLinkTrainSampleToAgent, true, PermUse);
        API.RegisterAPICall(AssistentUnlinkTrainSampleFromAgent, true, PermUse);
        API.RegisterAPICall(AssistentSyncDatasetToAgent, true, PermUse);
        Logs.Init("Swarm Assistent extension loaded (0.15.0 persona packs / dreamer)");
    }

    int CfgInt(string key, int fallback)
    {
        try
        {
            return Config?.LoadAssistant(Config.DefaultPersonaId())[key]?.Value<int?>() ?? fallback;
        }
        catch
        {
            return fallback;
        }
    }

    static string Clip(string text, int max)
    {
        if (string.IsNullOrEmpty(text) || text.Length <= max)
        {
            return text ?? "";
        }
        return text[..max] + "…";
    }

    static string CollapseWs(string text)
    {
        if (string.IsNullOrWhiteSpace(text))
        {
            return "";
        }
        return Regex.Replace(text.Trim(), @"\s+", " ");
    }

    public static string NormalizeBaseUrl(string raw)
    {
        string url = (raw ?? "").Trim();
        if (string.IsNullOrWhiteSpace(url))
        {
            url = "http://127.0.0.1:11434";
        }
        return url.TrimEnd('/');
    }

    static string DataRoot()
    {
        if (Directory.Exists("/mnt/swarm_data"))
        {
            return "/mnt/swarm_data";
        }
        try
        {
            string models = Path.GetFullPath(Path.Combine(Environment.CurrentDirectory, "Models"));
            if (Directory.Exists(models))
            {
                return Path.GetDirectoryName(models) ?? Environment.CurrentDirectory;
            }
        }
        catch
        {
            // ignore
        }
        return Environment.CurrentDirectory;
    }

    public async Task<JObject> AssistentGetConfig(Session session, string persona = null)
    {
        await Task.CompletedTask;
        string pid = AssistentConfig.SafeId(persona) ?? Config.DefaultPersonaId();
        return Config.BuildMergedConfigPayload(pid);
    }

    public async Task<JObject> AssistentSaveSettings(Session session, JObject settings)
    {
        await Task.CompletedTask;
        if (settings is null)
        {
            return new JObject { ["error"] = "settings required" };
        }
        string prevEmbed = Config.LoadSettings()["embed_model"]?.ToString();
        Config.SaveSettings(settings);
        string nextEmbed = settings["embed_model"]?.ToString();
        if (!string.IsNullOrWhiteSpace(nextEmbed) && !string.Equals(prevEmbed, nextEmbed, StringComparison.OrdinalIgnoreCase))
        {
            try
            {
                await Memory.ReembedAllAsync(NormalizeBaseUrl(settings["base_url"]?.ToString()), nextEmbed);
            }
            catch (Exception ex)
            {
                Logs.Debug($"AssistentSaveSettings reembed: {ex.Message}");
            }
        }
        return new JObject { ["success"] = true, ["path"] = Path.Combine(Config.OverlayRoot, "settings.json") };
    }

    public async Task<JObject> AssistentListPersonas(Session session)
    {
        await Task.CompletedTask;
        var catalog = Config.ListPersonaCatalog();
        JArray list = [];
        foreach (var p in catalog)
        {
            list.Add(new JObject
            {
                ["id"] = p.id,
                ["title"] = p.title,
                ["accent"] = p.accent,
                ["prompt"] = Config.RenderIdentityBlock(p.id, includeAllShelves: true),
                ["source"] = p.source,
            });
        }
        return new JObject
        {
            ["success"] = true,
            ["default"] = Config.DefaultPersonaId(),
            ["personas"] = list,
        };
    }
}