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.api.js");
        ScriptFiles.Add("Assets/assistent.patch.js");
        ScriptFiles.Add("Assets/assistent.persist.js");
        ScriptFiles.Add("Assets/assistent.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.9.0";
        Tags = ["tabs", "ui", "llm", "ollama", "krea", "inpaint", "memory"];
    }

    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(AssistentGetPacks, false, PermUse);
        API.RegisterAPICall(AssistentListPersonas, false, PermUse);
        API.RegisterAPICall(AssistentGetConfig, false, PermUse);
        API.RegisterAPICall(AssistentGetSettings, false, PermUse);
        API.RegisterAPICall(AssistentSaveSettings, true, PermUse);
        API.RegisterAPICall(AssistentListInventory, false, PermUse);
        API.RegisterAPICall(AssistentGetCard, false, PermUse);
        API.RegisterAPICall(AssistentSaveCard, true, PermUse);
        API.RegisterAPICall(AssistentEnqueueWanted, true, PermUse);
        API.RegisterAPICall(AssistentGetCardMeta, false, PermUse);
        API.RegisterAPICall(AssistentSearchCivitai, false, PermUse);
        API.RegisterAPICall(AssistentGetTaste, false, PermUse);
        API.RegisterAPICall(AssistentSaveTaste, true, 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(AssistentListWanted, 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);
        Logs.Init("Swarm Assistent extension loaded (persona shelves + controls + overlay clone)");
    }

    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 string ReadPackFile(string name)
    {
        return Config?.LoadPackPrompt(Config.DefaultPersonaId(), name);
    }

    public async Task<JObject> AssistentGetPacks(Session session, string persona = null)
    {
        await Task.CompletedTask;
        string pid = AssistentConfig.SafeId(persona) ?? Config.DefaultPersonaId();
        JObject packs = new();
        JArray order = [];
        foreach (var p in Config.ListPacks(pid))
        {
            string text = Config.LoadPackPrompt(pid, p.id);
            if (text is not null)
            {
                packs[p.id] = text;
            }
            order.Add(p.id);
        }
        return new JObject { ["success"] = true, ["packs"] = packs, ["order"] = order, ["persona"] = pid };
    }

    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> AssistentGetSettings(Session session)
    {
        await Task.CompletedTask;
        return new JObject { ["success"] = true, ["settings"] = Config.LoadSettings() };
    }

    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),
                ["source"] = p.source,
            });
        }
        return new JObject
        {
            ["success"] = true,
            ["default"] = Config.DefaultPersonaId(),
            ["personas"] = list,
        };
    }

    public async Task<JObject> AssistentGetTaste(Session session)
    {
        await Task.CompletedTask;
        try
        {
            return new JObject { ["success"] = true, ["taste"] = Memory.GetKvObject(AssistentMemory.KvTaste) };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"taste: {ex.Message}" };
        }
    }

    public async Task<JObject> AssistentSaveTaste(Session session, JObject taste)
    {
        await Task.CompletedTask;
        if (taste is null)
        {
            return new JObject { ["error"] = "taste required" };
        }
        if (taste["updated"] == null)
        {
            taste["updated"] = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
        }
        try
        {
            Memory.SetKvObject(AssistentMemory.KvTaste, taste);
            return new JObject { ["success"] = true, ["path"] = "Assistent/memory/assistent.sqlite" };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"taste save: {ex.Message}" };
        }
    }
}