swarm-assistent/AssistentKnowledge.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Newtonsoft.Json.Linq;
using SwarmUI.Utils;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Unified knowledge hub: books FTS + legacy examples shim + catalog for personas.</summary>
public partial class SwarmAssistentExtension
{
    public List<string> ResolveAttachedBooks(string personaId)
        => Config.ResolveKnowledgeAttach(personaId);

    public JObject BuildKnowledgeCatalog(string personaId)
    {
        string pid = AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId();
        List<string> attached = ResolveAttachedBooks(pid);
        JArray books = Memory?.ListBooksOnDisk() ?? [];
        HashSet<string> attachSet = attached.ToHashSet(StringComparer.OrdinalIgnoreCase);
        JArray catalog = [];
        foreach (JToken t in books)
        {
            if (t is not JObject b)
            {
                continue;
            }
            string id = b["id"]?.ToString() ?? "";
            catalog.Add(new JObject
            {
                ["id"] = id,
                ["title"] = b["title"] ?? id,
                ["description"] = b["description"] ?? "",
                ["content_kind"] = b["content_kind"] ?? "",
                ["language"] = b["language"] ?? "",
                ["attached"] = attachSet.Contains(id),
                ["indexed"] = b["indexed"]?.Value<bool?>() ?? false,
            });
        }
        return new JObject
        {
            ["persona"] = pid,
            ["attach"] = new JArray(attached),
            ["books"] = catalog,
            ["books_indexed"] = Memory?.BookRowCount() ?? 0,
            ["examples_indexed"] = Memory?.ExampleCount() ?? 0,
            ["tags_indexed"] = Memory?.TagCount() ?? 0,
        };
    }

    public string BuildKnowledgeSystemLayer(string personaId)
    {
        JObject cat = BuildKnowledgeCatalog(personaId);
        JArray attach = cat["attach"] as JArray ?? [];
        if (attach.Count == 0)
        {
            return "";
        }
        StringBuilder sb = new();
        sb.AppendLine("## Knowledge books (FTS reference — read-only)");
        sb.AppendLine("Attached corpora for this persona. Search via ask:[\"knowledge\"] + knowledge_query (or legacy ask:[\"examples\"] + example_query).");
        foreach (JToken t in attach)
        {
            string id = t?.ToString()?.Trim();
            if (string.IsNullOrWhiteSpace(id))
            {
                continue;
            }
            JObject meta = (cat["books"] as JArray)?.OfType<JObject>()
                .FirstOrDefault(b => string.Equals(b["id"]?.ToString(), id, StringComparison.OrdinalIgnoreCase));
            string title = meta?["title"]?.ToString() ?? id;
            string kind = meta?["content_kind"]?.ToString() ?? "";
            string desc = meta?["description"]?.ToString() ?? "";
            sb.AppendLine($"- **{id}** ({title}){ (string.IsNullOrWhiteSpace(kind) ? "" : $" · {kind}")}");
            if (!string.IsNullOrWhiteSpace(desc))
            {
                sb.AppendLine($"  {desc}");
            }
        }
        return sb.ToString().TrimEnd();
    }

    public JArray SearchKnowledge(string personaId, string query, int limit = 8, string rating = null)
    {
        query = (query ?? "").Trim();
        if (query.Length < 1 || Memory is null)
        {
            return [];
        }
        List<string> attached = ResolveAttachedBooks(personaId);
        JArray hits = [];
        HashSet<string> seen = [];
        if (attached.Count > 0)
        {
            foreach (JToken t in Memory.LookupBooks(query, limit, attached, rating))
            {
                string key = t?["id"]?.ToString() ?? t?.ToString();
                if (!string.IsNullOrWhiteSpace(key) && seen.Add(key))
                {
                    hits.Add(t);
                }
            }
        }
        if (hits.Count < limit && attached.Any(b => string.Equals(b, "civitai-krea2", StringComparison.OrdinalIgnoreCase)))
        {
            int remain = limit - hits.Count;
            foreach (JToken t in Memory.LookupExamples(query, remain, rating))
            {
                string key = "legacy:" + (t?["id"]?.ToString() ?? "");
                if (!seen.Add(key))
                {
                    continue;
                }
                JObject row = t as JObject ?? new JObject();
                hits.Add(new JObject
                {
                    ["id"] = row["id"],
                    ["book"] = "civitai-krea2",
                    ["source"] = "examples-legacy",
                    ["title"] = "",
                    ["tags"] = row["tags"],
                    ["text"] = row["prompt"],
                    ["body"] = row["prompt"],
                    ["rating"] = row["rating"],
                    ["params"] = row["params"],
                    ["loras"] = row["loras"],
                    ["note"] = row["note"],
                });
            }
        }
        return hits;
    }

    public static JArray CivitaiResultsShim(JArray knowledgeResults)
    {
        JArray outRows = [];
        foreach (JToken t in knowledgeResults ?? [])
        {
            if (t is not JObject row)
            {
                continue;
            }
            string book = row["book"]?.ToString() ?? "";
            string source = row["source"]?.ToString() ?? "";
            if (!string.Equals(book, "civitai-krea2", StringComparison.OrdinalIgnoreCase)
                && source is not "examples-legacy")
            {
                continue;
            }
            if (row["prompt"] is not null)
            {
                outRows.Add(row);
                continue;
            }
            outRows.Add(new JObject
            {
                ["id"] = row["id"],
                ["rating"] = row["rating"],
                ["tags"] = row["tags"],
                ["prompt"] = row["text"] ?? row["body"],
                ["negative"] = row["meta"]?["negative"] ?? "",
                ["params"] = row["params"] ?? row["meta"]?["params"],
                ["loras"] = row["loras"] ?? row["meta"]?["loras"],
                ["note"] = row["note"] ?? "EXAMPLE from Civitai — remix, do not copy 1:1",
            });
        }
        return outRows;
    }

    JObject BuildKnowledgeResponse(string personaId, JArray hops, JArray results)
        => new()
        {
            ["catalog"] = BuildKnowledgeCatalog(personaId),
            ["hops"] = hops ?? new JArray(),
            ["results"] = results ?? new JArray(),
        };

    static void MergeKnowledgeResults(JArray target, JArray rows, bool legacyExamples = false)
    {
        if (target is null || rows is null)
        {
            return;
        }
        HashSet<string> seen = target
            .Select(t => t?["id"]?.ToString() ?? t?.ToString())
            .Where(s => !string.IsNullOrWhiteSpace(s))
            .ToHashSet(StringComparer.OrdinalIgnoreCase);
        foreach (JToken t in rows)
        {
            JObject row = t as JObject ?? new JObject { ["text"] = t?.ToString() };
            if (legacyExamples)
            {
                row = new JObject
                {
                    ["id"] = row["id"],
                    ["book"] = "civitai-krea2",
                    ["source"] = "examples-legacy",
                    ["tags"] = row["tags"],
                    ["text"] = row["prompt"],
                    ["body"] = row["prompt"],
                    ["rating"] = row["rating"],
                    ["note"] = row["note"],
                };
            }
            string key = row["id"]?.ToString() ?? row["text"]?.ToString();
            if (string.IsNullOrWhiteSpace(key) || !seen.Add(key))
            {
                continue;
            }
            target.Add(row);
        }
    }
}