swarm-assistent/SwarmAssistentExtension.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.Http;
using System.Text;
using System.Threading.Tasks;
using FreneticUtilities.FreneticExtensions;
using Newtonsoft.Json.Linq;
using SwarmUI.Accounts;
using SwarmUI.Core;
using SwarmUI.Utils;
using SwarmUI.WebAPI;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Krea 2 collaborative assistant: Ollama chat + vision + prompt/LoRA/params patches.</summary>
public 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 static readonly string[] PackNames =
    [
        "base_krea2",
        "write_prompt",
        "critique_image",
        "compose_scene",
        "fix_params",
    ];

    public override void OnPreInit()
    {
        ScriptFiles.Add("Assets/assistent.js");
        StyleSheetFiles.Add("Assets/assistent.css");
        ExtensionAuthor = "mrleo1nid";
        Description = "Collaborative Krea 2 assistant via Ollama: chat, vision, prompts, LoRA triggers, size patches.";
        License = "MIT";
        Version = "0.2.0";
        Tags = ["tabs", "ui", "llm", "ollama", "krea"];
    }

    public override void OnInit()
    {
        HttpClient ??= new HttpClient { Timeout = TimeSpan.FromMinutes(10) };
        API.RegisterAPICall(AssistentListModels, false, PermUse);
        API.RegisterAPICall(AssistentGetPacks, false, PermUse);
        API.RegisterAPICall(AssistentChat, true, PermUse);
        Logs.Init("Swarm Assistent extension loaded (Ollama proxy + Krea 2 packs)");
    }

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

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

    public string ReadPackFile(string name)
    {
        string safe = name.Replace('\\', '/').AfterLast('/').Replace("..", "");
        if (!PackNames.Contains(safe))
        {
            return null;
        }
        string path = Path.Combine(FilePath, "Prompts", $"{safe}.md");
        if (!File.Exists(path))
        {
            return null;
        }
        return File.ReadAllText(path, Encoding.UTF8);
    }

    public async Task<JObject> AssistentListModels(Session session, string baseUrl)
    {
        string root = NormalizeBaseUrl(baseUrl);
        try
        {
            using HttpResponseMessage resp = await HttpClient.GetAsync($"{root}/api/tags");
            string body = await resp.Content.ReadAsStringAsync();
            if (!resp.IsSuccessStatusCode)
            {
                return new JObject { ["error"] = $"Ollama /api/tags HTTP {(int)resp.StatusCode}: {Clip(body, 400)}" };
            }
            JObject parsed = JObject.Parse(body);
            JArray models = [];
            foreach (JToken m in parsed["models"] as JArray ?? [])
            {
                models.Add(m["name"]?.ToString() ?? "");
            }
            return new JObject { ["success"] = true, ["base_url"] = root, ["models"] = models };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"Ollama unreachable at {root}: {ex.Message}" };
        }
    }

    public async Task<JObject> AssistentGetPacks(Session session)
    {
        JObject packs = new();
        foreach (string name in PackNames)
        {
            string text = ReadPackFile(name);
            if (text is not null)
            {
                packs[name] = text;
            }
        }
        return new JObject { ["success"] = true, ["packs"] = packs, ["order"] = new JArray(PackNames) };
    }

    /// <summary>
    /// Proxy to Ollama /api/chat (non-stream).
    /// <paramref name="raw"/> must include messages (JArray) and optional context_json.
    /// </summary>
    public async Task<JObject> AssistentChat(Session session, string baseUrl, string model, string pack, bool includeBase, JObject raw)
    {
        string root = NormalizeBaseUrl(baseUrl ?? raw?["base_url"]?.ToString());
        string modelName = (model ?? raw?["model"]?.ToString() ?? "").Trim();
        if (string.IsNullOrWhiteSpace(modelName))
        {
            return new JObject { ["error"] = "model is required" };
        }
        JArray userMessages = raw?["messages"] as JArray;
        if (userMessages is null || userMessages.Count == 0)
        {
            return new JObject { ["error"] = "messages required" };
        }

        List<JObject> ollamaMessages = [];
        StringBuilder system = new();
        if (includeBase)
        {
            string basePack = ReadPackFile("base_krea2");
            if (!string.IsNullOrWhiteSpace(basePack))
            {
                system.AppendLine(basePack);
            }
        }
        string packName = (pack ?? raw?["pack"]?.ToString() ?? "write_prompt").Trim();
        if (!string.IsNullOrWhiteSpace(packName) && packName != "base_krea2")
        {
            string situational = ReadPackFile(packName);
            if (!string.IsNullOrWhiteSpace(situational))
            {
                system.AppendLine();
                system.AppendLine($"## Active mode: {packName}");
                system.AppendLine(situational);
            }
        }
        string contextJson = raw?["context_json"]?.ToString();
        if (!string.IsNullOrWhiteSpace(contextJson))
        {
            system.AppendLine();
            system.AppendLine("## Live SwarmUI context (JSON — trust this over guesses)");
            system.AppendLine("```json");
            system.AppendLine(contextJson);
            system.AppendLine("```");
        }
        if (system.Length > 0)
        {
            ollamaMessages.Add(new JObject
            {
                ["role"] = "system",
                ["content"] = system.ToString(),
            });
        }
        foreach (JToken msg in userMessages)
        {
            if (msg is not JObject mo)
            {
                continue;
            }
            JObject copy = new()
            {
                ["role"] = mo["role"]?.ToString() ?? "user",
                ["content"] = mo["content"]?.ToString() ?? "",
            };
            if (mo["images"] is JArray images && images.Count > 0)
            {
                copy["images"] = images;
            }
            ollamaMessages.Add(copy);
        }

        JObject payload = new()
        {
            ["model"] = modelName,
            ["stream"] = false,
            ["messages"] = new JArray(ollamaMessages),
        };
        try
        {
            using StringContent content = new(payload.ToString(Newtonsoft.Json.Formatting.None), Encoding.UTF8, "application/json");
            using HttpResponseMessage resp = await HttpClient.PostAsync($"{root}/api/chat", content);
            string body = await resp.Content.ReadAsStringAsync();
            if (!resp.IsSuccessStatusCode)
            {
                return new JObject { ["error"] = $"Ollama /api/chat HTTP {(int)resp.StatusCode}: {Clip(body, 800)}" };
            }
            JObject parsed = JObject.Parse(body);
            string reply = parsed["message"]?["content"]?.ToString() ?? parsed["response"]?.ToString() ?? "";
            return new JObject
            {
                ["success"] = true,
                ["reply"] = reply,
                ["model"] = modelName,
                ["pack"] = packName,
                ["raw"] = parsed,
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"Ollama chat failed: {ex.Message}" };
        }
    }
}