swarm-assistent/AssistentPatch.cs

using System;
using System.Text.RegularExpressions;
using Newtonsoft.Json.Linq;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Parsing and normalization of the JSON patch the model emits inside fenced code blocks.</summary>
public partial class SwarmAssistentExtension
{
    static readonly Regex JsonFenceRe = new(@"```(?:json)?\s*([\s\S]*?)```", RegexOptions.IgnoreCase | RegexOptions.Compiled);

    static readonly string[] PatchKeys =
    [
        "prompt", "loras", "width", "height", "steps", "cfg", "seed", "sigma_shift", "sampler",
        "actions", "search_query", "civitai_query",
        "use_init_image", "clear_init_image", "init_creativity", "denoise",
        "use_mask_image", "clear_mask_image", "mask_blur", "mask_grow",
        "look_at", "vision_from", "vision_slots", "slot_to_init", "slot_to_mask",
        "snapshot_generate", "select_slot", "aspect", "images", "batch", "vary", "lock_seed",
        "creativity", "intensity", "complexity", "movement",
        "clear_prompt_images", "slot_to_prompt_image", "pack", "memories", "memory",
        "memory_query", "memory_kind", "tag_query", "user_prefs",
        "inventory_query", "skills", "persona_shelves", "controls",
    ];

    static bool HasValue(JObject obj, string key)
    {
        JToken token = obj?[key];
        return token is not null && token.Type != JTokenType.Null;
    }

    /// <summary>Maps legacy/alias patch fields onto their canonical names. Aliases are kept so older consumers still work.</summary>
    public static JObject NormalizePatch(JObject patch)
    {
        if (patch is null)
        {
            return null;
        }
        if (!HasValue(patch, "search_query") && HasValue(patch, "civitai_query"))
        {
            patch["search_query"] = patch["civitai_query"];
        }
        if (!HasValue(patch, "init_creativity") && HasValue(patch, "denoise"))
        {
            patch["init_creativity"] = patch["denoise"];
        }
        if (!HasValue(patch, "look_at"))
        {
            if (HasValue(patch, "vision_from"))
            {
                patch["look_at"] = patch["vision_from"];
            }
            else if (HasValue(patch, "vision_slots"))
            {
                patch["look_at"] = patch["vision_slots"];
            }
        }
        return patch;
    }

    static bool LooksLikeCardObject(JObject obj)
    {
        if (obj is null)
        {
            return false;
        }
        bool cardish = HasValue(obj, "kind") || HasValue(obj, "triggers") || HasValue(obj, "when") || HasValue(obj, "prompt_hint");
        bool genish = HasValue(obj, "prompt") || HasValue(obj, "loras") || HasValue(obj, "actions")
            || HasValue(obj, "width") || HasValue(obj, "height") || HasValue(obj, "steps") || HasValue(obj, "cfg")
            || HasValue(obj, "aspect") || HasValue(obj, "seed") || HasValue(obj, "search_query")
            || HasValue(obj, "civitai_query") || HasValue(obj, "look_at") || HasValue(obj, "controls");
        if (cardish && !genish && (HasValue(obj, "name") || HasValue(obj, "triggers") || HasValue(obj, "when")))
        {
            return true;
        }
        return HasValue(obj, "kind") && HasValue(obj, "name")
            && (HasValue(obj, "triggers") || HasValue(obj, "when") || HasValue(obj, "prompt_hint") || HasValue(obj, "notes"));
    }

    static JObject TryParsePatch(string reply)
    {
        if (string.IsNullOrWhiteSpace(reply))
        {
            return null;
        }
        foreach (Match match in JsonFenceRe.Matches(reply))
        {
            string raw = match.Groups[1].Value.Trim();
            try
            {
                JObject obj = JObject.Parse(raw);
                if (obj is null || LooksLikeCardObject(obj))
                {
                    continue;
                }
                if (Array.Exists(PatchKeys, k => obj[k] is not null))
                {
                    return NormalizePatch(obj);
                }
            }
            catch
            {
                // not json
            }
        }
        return null;
    }

    /// <summary>
    /// If the reply already contains a closed fenced patch/card, cut everything after it.
    /// Models often keep writing («Готово!», second aspect, …) and the stream never feels done.
    /// </summary>
    static bool TryTruncateAtCompleteFence(string reply, out string truncated)
    {
        truncated = reply ?? "";
        if (string.IsNullOrWhiteSpace(reply))
        {
            return false;
        }
        MatchCollection matches = JsonFenceRe.Matches(reply);
        if (matches.Count == 0)
        {
            return false;
        }
        for (int i = 0; i < matches.Count; i++)
        {
            Match match = matches[i];
            string raw = match.Groups[1].Value.Trim();
            try
            {
                JObject obj = JObject.Parse(raw);
                if (obj is null)
                {
                    continue;
                }
                bool usable = LooksLikeCardObject(obj)
                    || Array.Exists(PatchKeys, k => obj[k] is not null);
                if (!usable)
                {
                    continue;
                }
                truncated = reply.Substring(0, match.Index + match.Length).TrimEnd();
                // Only treat as complete if the fence actually closed (regex requires ```).
                return true;
            }
            catch
            {
                // incomplete / invalid json inside fence
            }
        }
        return false;
    }

    static string ExtractSearchQuery(JObject patch)
    {
        if (patch is null)
        {
            return null;
        }
        string q = (patch["search_query"] ?? patch["civitai_query"])?.ToString()?.Trim();
        return string.IsNullOrWhiteSpace(q) ? null : q;
    }

    static bool ActionsContain(JObject patch, string action)
    {
        if (patch?["actions"] is not JArray acts)
        {
            return false;
        }
        foreach (JToken a in acts)
        {
            if (string.Equals(a?.ToString(), action, StringComparison.OrdinalIgnoreCase))
            {
                return true;
            }
        }
        return false;
    }

    static string ExtractMemoryQuery(JObject patch)
    {
        string q = patch?["memory_query"]?.ToString()?.Trim();
        if (!string.IsNullOrWhiteSpace(q))
        {
            return q;
        }
        return ActionsContain(patch, "memory_search") ? ExtractSearchQuery(patch) : null;
    }

    static string ExtractTagQuery(JObject patch)
    {
        string q = patch?["tag_query"]?.ToString()?.Trim();
        if (!string.IsNullOrWhiteSpace(q))
        {
            return q;
        }
        return ActionsContain(patch, "lookup_tags") ? ExtractSearchQuery(patch) : null;
    }

    static string NextToolHop(JObject patch)
    {
        if (patch is null)
        {
            return null;
        }
        if (ActionsContain(patch, "memory_get"))
        {
            return "memory_get";
        }
        if (ActionsContain(patch, "memory_search") || !string.IsNullOrWhiteSpace(patch["memory_query"]?.ToString()))
        {
            return "memory_search";
        }
        if (ActionsContain(patch, "lookup_tags") || !string.IsNullOrWhiteSpace(patch["tag_query"]?.ToString()))
        {
            return "lookup_tags";
        }
        if (ActionsContain(patch, "list_inventory") || !string.IsNullOrWhiteSpace(patch["inventory_query"]?.ToString()))
        {
            return "list_inventory";
        }
        if (ActionsContain(patch, "skill_load"))
        {
            return "skill_load";
        }
        if (ActionsContain(patch, "persona_read"))
        {
            return "persona_read";
        }
        if (ActionsContain(patch, "search_civitai") || !string.IsNullOrWhiteSpace(ExtractSearchQuery(patch)))
        {
            return "civitai";
        }
        return null;
    }

    static bool WantsCivitaiSearch(JObject patch)
        => ActionsContain(patch, "search_civitai") || !string.IsNullOrWhiteSpace(ExtractSearchQuery(patch));
}