swarm-assistent/AssistentPatch.cs

using System;
using System.Collections.Generic;
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);

    string[] PatchKeys => _patchKeys ??= Config?.LoadPatchKeys() ?? [];

    static string[] _patchKeys;

    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"];
            }
        }
        if (ActionsContain(patch, "generate") || GenerateFlagOn(patch))
        {
            patch["generate"] = true;
        }
        if (patch["ask"] is JValue askVal && askVal.Type == JTokenType.String)
        {
            string one = askVal.ToString()?.Trim();
            if (!string.IsNullOrWhiteSpace(one))
            {
                patch["ask"] = new JArray(one);
            }
            else
            {
                patch.Remove("ask");
            }
        }
        return patch;
    }

    JObject TryParsePatch(string reply)
    {
        if (string.IsNullOrWhiteSpace(reply))
        {
            return null;
        }
        JObject lastAny = null;
        JObject lastTerminal = 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)
                {
                    continue;
                }
                if (Array.Exists(PatchKeys, k => obj[k] is not null))
                {
                    JObject normalized = NormalizePatch(obj);
                    lastAny = normalized;
                    if (FenceIsTerminalPatch(normalized))
                    {
                        lastTerminal = normalized;
                    }
                }
            }
            catch
            {
                // not json
            }
        }
        if (lastTerminal is not null || lastAny is not null)
        {
            return lastTerminal ?? lastAny;
        }
        int brace = reply.LastIndexOf('{');
        if (brace < 0)
        {
            return null;
        }
        try
        {
            JObject obj = JObject.Parse(reply.Substring(brace).Trim());
            if (obj is not null && 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 that is "done enough" to act on,
    /// cut everything after it. Do NOT stop on weak fences (pack/creativity/notes-only) — models
    /// often emit a tiny JSON first then the real prompt fence; aborting early cuts the prompt.
    /// </summary>
    static bool TryTruncateAtCompleteFence(string reply, out string truncated)
    {
        truncated = reply ?? "";
        if (string.IsNullOrWhiteSpace(reply))
        {
            return false;
        }
        MatchCollection matches = JsonFenceRe.Matches(reply);
        for (int i = 0; i < matches.Count; i++)
        {
            Match match = matches[i];
            string raw = match.Groups[1].Value.Trim();
            try
            {
                JObject obj = NormalizePatch(JObject.Parse(raw));
                if (obj is null || !FenceIsTerminalPatch(obj))
                {
                    continue;
                }
                truncated = reply.Substring(0, match.Index + match.Length).TrimEnd();
                return true;
            }
            catch
            {
                // incomplete / invalid json inside fence
            }
        }
        int brace = reply.LastIndexOf('{');
        if (brace >= 0)
        {
            try
            {
                JObject obj = NormalizePatch(JObject.Parse(reply.Substring(brace).Trim()));
                if (obj is not null && FenceIsTerminalPatch(obj))
                {
                    truncated = reply.TrimEnd();
                    return true;
                }
            }
            catch
            {
                // incomplete unfenced json
            }
        }
        return false;
    }

    /// <summary>
    /// True when a closed fence is worth aborting the Ollama stream (real deliverable or ask hop).
    /// </summary>
    static bool FenceIsTerminalPatch(JObject obj)
    {
        if (obj is null)
        {
            return false;
        }
        if (GenerateFlagOn(obj))
        {
            return true;
        }
        if (HasAsk(obj))
        {
            return true;
        }
        if (obj["variants"] is JArray variants && variants.Count > 0)
        {
            return true;
        }
        if (HasValue(obj, "look_at") || HasValue(obj, "vision_from") || HasValue(obj, "vision_slots"))
        {
            return true;
        }
        if (ActionsContain(obj, "generate"))
        {
            return true;
        }
        string prompt = obj["prompt"]?.ToString()?.Trim() ?? "";
        if (prompt.Length >= 48)
        {
            return true;
        }
        if (HasValue(obj, "loras") || HasValue(obj, "aspect") || HasValue(obj, "steps")
            || HasValue(obj, "width") || HasValue(obj, "height") || HasValue(obj, "cfg")
            || HasValue(obj, "seed") || HasValue(obj, "controls")
            || HasValue(obj, "checkpoint") || HasValue(obj, "negative"))
        {
            return true;
        }
        // Weak: pack / creativity / intensity / empty actions / notes-only → keep streaming
        return false;
    }

    static bool HasAsk(JObject patch)
    {
        if (patch?["ask"] is JArray asks)
        {
            foreach (JToken t in asks)
            {
                if (!string.IsNullOrWhiteSpace(t?.ToString()))
                {
                    return true;
                }
            }
            return false;
        }
        return !string.IsNullOrWhiteSpace(patch?["ask"]?.ToString());
    }

    static bool AskContains(JObject patch, string name)
    {
        if (patch is null || string.IsNullOrWhiteSpace(name))
        {
            return false;
        }
        if (patch["ask"] is JArray asks)
        {
            foreach (JToken t in asks)
            {
                if (string.Equals(t?.ToString()?.Trim(), name, StringComparison.OrdinalIgnoreCase))
                {
                    return true;
                }
            }
            return false;
        }
        return string.Equals(patch["ask"]?.ToString()?.Trim(), name, StringComparison.OrdinalIgnoreCase);
    }

    static bool GenerateFlagOn(JObject obj)
    {
        if (obj is null)
        {
            return false;
        }
        if (obj["generate"]?.Type == JTokenType.Boolean && obj["generate"].Value<bool>())
        {
            return true;
        }
        if (obj["generate"]?.Type == JTokenType.Integer && obj["generate"].Value<int>() != 0)
        {
            return true;
        }
        string raw = obj["generate"]?.ToString()?.Trim();
        if (!string.IsNullOrWhiteSpace(raw)
            && (raw.Equals("true", StringComparison.OrdinalIgnoreCase)
                || raw == "1"
                || raw.Equals("yes", StringComparison.OrdinalIgnoreCase)
                || raw.Equals("on", StringComparison.OrdinalIgnoreCase)))
        {
            return true;
        }
        return ActionsContain(obj, "generate");
    }

    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;
    }

    /// <summary>Server tool hops: settings dump, inventory, or Civitai example FTS.</summary>
    static string NextToolHop(JObject patch, HashSet<string> skip = null)
    {
        if (patch is null)
        {
            return null;
        }
        bool Skip(string tool) => skip is not null && skip.Contains(tool);
        if (AskContains(patch, "settings") && !Skip("ask_settings"))
        {
            return "ask_settings";
        }
        if (AskContains(patch, "inventory") && !Skip("ask_inventory"))
        {
            return "ask_inventory";
        }
        if ((AskContains(patch, "knowledge") || ActionsContain(patch, "lookup_knowledge")) && !Skip("ask_knowledge"))
        {
            return "ask_knowledge";
        }
        if ((AskContains(patch, "examples") || ActionsContain(patch, "lookup_examples")) && !Skip("ask_examples"))
        {
            return "ask_examples";
        }
        return null;
    }
}