swarm-assistent/AssistentWanted.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using Newtonsoft.Json.Linq;
using SwarmUI.Accounts;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Queue of models the assistant wants downloaded (merged into gpu-rent models.yaml on next up/capture).</summary>
public partial class SwarmAssistentExtension
{
    static readonly object WantedFileLock = new();

    string WantedModelsPath() => Path.Combine(DataRoot(), ".gpu-rent-wanted-models.yaml");

    string WantedCardsDir() => Path.Combine(DataRoot(), ".gpu-rent-wanted-cards");

    public async Task<JObject> AssistentEnqueueWanted(Session session, string kind, string url, int version_id = 0, string title = null, JObject card = null)
    {
        await Task.CompletedTask;
        kind = (kind ?? "lora").Trim().ToLowerInvariant();
        if (kind is not ("lora" or "checkpoint" or "vae" or "embedding" or "controlnet" or "upscaler" or "clip"))
        {
            kind = "lora";
        }
        url = (url ?? "").Trim();
        if (string.IsNullOrWhiteSpace(url) && version_id > 0)
        {
            url = $"https://civitai.red/models/0?modelVersionId={version_id}";
        }
        if (string.IsNullOrWhiteSpace(url))
        {
            return new JObject { ["error"] = "url or version_id required" };
        }
        if (version_id <= 0)
        {
            Match m = Regex.Match(url, @"modelVersionId=(\d+)", RegexOptions.IgnoreCase);
            if (m.Success)
            {
                version_id = int.Parse(m.Groups[1].Value);
            }
        }

        string path = WantedModelsPath();
        Directory.CreateDirectory(Path.GetDirectoryName(path) ?? DataRoot());
        lock (WantedFileLock)
        {
            Dictionary<string, List<WantedEntry>> sections = LoadWantedYaml(File.Exists(path) ? File.ReadAllText(path, Encoding.UTF8) : "");

            if (version_id > 0)
            {
                foreach (List<WantedEntry> list in sections.Values)
                {
                    if (list.Any(e => e.VersionId == version_id))
                    {
                        return new JObject { ["success"] = true, ["already"] = true, ["path"] = path, ["version_id"] = version_id };
                    }
                }
            }
            else
            {
                foreach (List<WantedEntry> list in sections.Values)
                {
                    if (list.Any(e => string.Equals(e.Url, url, StringComparison.OrdinalIgnoreCase)))
                    {
                        return new JObject { ["success"] = true, ["already"] = true, ["path"] = path };
                    }
                }
            }

            if (!sections.TryGetValue(kind, out List<WantedEntry> bucket))
            {
                bucket = [];
                sections[kind] = bucket;
            }
            bucket.Add(new WantedEntry { Url = url, Title = title, VersionId = version_id });
            File.WriteAllText(path, WriteWantedYaml(sections), Encoding.UTF8);
        }

        if (card is not null && version_id > 0)
        {
            Directory.CreateDirectory(WantedCardsDir());
            string draft = Path.Combine(WantedCardsDir(), $"{version_id}.assistent.json");
            File.WriteAllText(draft, card.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8);
        }
        return new JObject { ["success"] = true, ["path"] = path, ["version_id"] = version_id };
    }

    sealed class WantedEntry
    {
        public string Url;
        public string Title;
        public int VersionId;
    }

    static Dictionary<string, List<WantedEntry>> LoadWantedYaml(string raw)
    {
        Dictionary<string, List<WantedEntry>> sections = new(StringComparer.OrdinalIgnoreCase);
        string currentKind = null;
        WantedEntry cur = null;
        void Flush()
        {
            if (cur is null || string.IsNullOrWhiteSpace(cur.Url) || string.IsNullOrWhiteSpace(currentKind))
            {
                cur = null;
                return;
            }
            if (!sections.TryGetValue(currentKind, out List<WantedEntry> list))
            {
                list = [];
                sections[currentKind] = list;
            }
            list.Add(cur);
            cur = null;
        }
        foreach (string line in (raw ?? "").Split('\n'))
        {
            string t = line.TrimEnd();
            if (string.IsNullOrWhiteSpace(t) || t.TrimStart().StartsWith('#'))
            {
                continue;
            }
            Match kindLine = Regex.Match(t, @"^([A-Za-z0-9_-]+):\s*$");
            if (kindLine.Success && !t.TrimStart().StartsWith('-'))
            {
                Flush();
                currentKind = kindLine.Groups[1].Value.Trim().ToLowerInvariant();
                continue;
            }
            Match urlLine = Regex.Match(t, @"^\s*-\s*url:\s*[""']?(.+?)[""']?\s*$");
            if (urlLine.Success)
            {
                Flush();
                cur = new WantedEntry { Url = urlLine.Groups[1].Value.Trim() };
                continue;
            }
            if (cur is null)
            {
                continue;
            }
            Match titleLine = Regex.Match(t, @"^\s*title:\s*[""']?(.+?)[""']?\s*$");
            if (titleLine.Success)
            {
                cur.Title = titleLine.Groups[1].Value.Trim();
                continue;
            }
            Match vidLine = Regex.Match(t, @"^\s*version_id:\s*(\d+)\s*$");
            if (vidLine.Success && int.TryParse(vidLine.Groups[1].Value, out int vid))
            {
                cur.VersionId = vid;
            }
        }
        Flush();
        return sections;
    }

    static string WriteWantedYaml(Dictionary<string, List<WantedEntry>> sections)
    {
        StringBuilder sb = new();
        sb.AppendLine("# Assistent wanted queue — merged into local models.yaml on gpu-rent up/capture");
        string[] order = ["checkpoint", "lora", "vae", "embedding", "controlnet", "upscaler", "clip"];
        HashSet<string> seen = new(StringComparer.OrdinalIgnoreCase);
        foreach (string kind in order.Concat(sections.Keys.OrderBy(k => k, StringComparer.OrdinalIgnoreCase)))
        {
            if (!seen.Add(kind) || !sections.TryGetValue(kind, out List<WantedEntry> list) || list.Count == 0)
            {
                continue;
            }
            sb.AppendLine($"{kind}:");
            foreach (WantedEntry e in list)
            {
                sb.AppendLine($"  - url: \"{e.Url.Replace("\"", "%22")}\"");
                if (!string.IsNullOrWhiteSpace(e.Title))
                {
                    sb.AppendLine($"    title: \"{e.Title.Replace("\"", "'")}\"");
                }
                if (e.VersionId > 0)
                {
                    sb.AppendLine($"    version_id: {e.VersionId}");
                }
            }
        }
        return sb.ToString();
    }
}