swarm-assistent/AssistentMemory.Examples.cs

using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Security.Cryptography;
using System.Text;
using Microsoft.Data.Sqlite;
using Newtonsoft.Json.Linq;
using SwarmUI.Utils;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Civitai Krea2 example prompts — FTS only (no embeddings), like Danbooru tags.</summary>
public sealed partial class AssistentMemory
{
    const string ExamplesMetaFp = "civitai_examples_fp";
    const string ExamplesMetaCount = "civitai_examples_count";

    void TryIndexExamples()
    {
        try
        {
            EnsureExamplesIndex();
        }
        catch (Exception ex)
        {
            Logs.Debug($"AssistentMemory examples index: {ex.Message}");
        }
    }

    void EnsureExamplesSchema()
    {
        Exec(
            """
            CREATE TABLE IF NOT EXISTS civitai_examples (
              id INTEGER PRIMARY KEY,
              rating TEXT NOT NULL DEFAULT 'pg',
              score INTEGER NOT NULL DEFAULT 0,
              kind TEXT NOT NULL DEFAULT '',
              model_version_id INTEGER NOT NULL DEFAULT 0,
              tags TEXT NOT NULL DEFAULT '',
              prompt TEXT NOT NULL DEFAULT '',
              negative TEXT NOT NULL DEFAULT '',
              params_json TEXT NOT NULL DEFAULT '{}',
              loras_json TEXT NOT NULL DEFAULT '[]',
              search_blob TEXT NOT NULL DEFAULT ''
            );
            """);
        Exec("CREATE INDEX IF NOT EXISTS idx_civitai_examples_score ON civitai_examples(score DESC);");
        Exec(
            """
            CREATE VIRTUAL TABLE IF NOT EXISTS civitai_examples_fts USING fts5(
              tags,
              prompt,
              search_blob,
              tokenize = 'unicode61 remove_diacritics 2'
            );
            """);
        Exec(
            """
            CREATE TRIGGER IF NOT EXISTS civitai_examples_fts_ai AFTER INSERT ON civitai_examples BEGIN
              INSERT INTO civitai_examples_fts(rowid, tags, prompt, search_blob)
              VALUES (new.id, new.tags, new.prompt, new.search_blob);
            END;
            """);
        Exec(
            """
            CREATE TRIGGER IF NOT EXISTS civitai_examples_fts_ad AFTER DELETE ON civitai_examples BEGIN
              INSERT INTO civitai_examples_fts(civitai_examples_fts, rowid) VALUES('delete', old.id);
            END;
            """);
        Exec(
            """
            CREATE TRIGGER IF NOT EXISTS civitai_examples_fts_au AFTER UPDATE ON civitai_examples BEGIN
              INSERT INTO civitai_examples_fts(civitai_examples_fts, rowid) VALUES('delete', old.id);
              INSERT INTO civitai_examples_fts(rowid, tags, prompt, search_blob)
              VALUES (new.id, new.tags, new.prompt, new.search_blob);
            END;
            """);
    }

    public string FindCivitaiExamplesJsonl()
    {
        string[] candidates =
        [
            Path.Combine(_dataRoot, "Assistent", "civitai-examples.jsonl"),
            Path.Combine(_dataRoot, "civitai-examples.jsonl"),
        ];
        foreach (string path in candidates)
        {
            if (File.Exists(path))
            {
                return path;
            }
        }
        return null;
    }

    static string FileFingerprint(string path)
    {
        var info = new FileInfo(path);
        using SHA256 sha = SHA256.Create();
        using FileStream fs = File.OpenRead(path);
        byte[] hash = sha.ComputeHash(fs);
        string hex = Convert.ToHexString(hash).ToLowerInvariant();
        return $"{info.Length}:{info.LastWriteTimeUtc.Ticks}:{hex}";
    }

    /// <summary>Load search.jsonl into FTS (no embeddings). No-op if fingerprint matches.</summary>
    public int EnsureExamplesIndex()
    {
        string path = FindCivitaiExamplesJsonl();
        if (string.IsNullOrWhiteSpace(path) || !File.Exists(path))
        {
            return 0;
        }
        string fp = FileFingerprint(path);
        lock (_lock)
        {
            EnsureOpen();
            EnsureExamplesSchema();
            if (string.Equals(GetMeta(ExamplesMetaFp), fp, StringComparison.Ordinal))
            {
                using SqliteCommand c = _conn.CreateCommand();
                c.CommandText = "SELECT COUNT(*) FROM civitai_examples";
                return Convert.ToInt32(c.ExecuteScalar());
            }
        }

        List<(long id, string rating, int score, string kind, int mvid, string tags, string prompt, string negative, string paramsJson, string lorasJson, string blob)> rows = [];
        foreach (string line in File.ReadLines(path, Encoding.UTF8))
        {
            if (string.IsNullOrWhiteSpace(line))
            {
                continue;
            }
            JObject o;
            try
            {
                o = JObject.Parse(line);
            }
            catch
            {
                continue;
            }
            long id = o["id"]?.Value<long?>() ?? 0;
            if (id <= 0)
            {
                continue;
            }
            string prompt = o["prompt"]?.ToString() ?? "";
            if (string.IsNullOrWhiteSpace(prompt))
            {
                continue;
            }
            string tagsJoined = "";
            if (o["tags"] is JArray tagArr)
            {
                tagsJoined = string.Join(", ", tagArr.Select(t => t?.ToString()?.Trim()).Where(t => !string.IsNullOrWhiteSpace(t)));
            }
            string rating = o["rating"]?.ToString() ?? "pg";
            int score = o["score"]?.Value<int?>() ?? 0;
            string kind = o["kind"]?.ToString() ?? "";
            int mvid = o["modelVersionId"]?.Value<int?>() ?? 0;
            string negative = o["negative"]?.ToString() ?? "";
            string paramsJson = (o["params"] as JObject)?.ToString(Newtonsoft.Json.Formatting.None) ?? "{}";
            string lorasJson = (o["loras"] as JArray)?.ToString(Newtonsoft.Json.Formatting.None) ?? "[]";
            string blob = $"{tagsJoined}\n{prompt}\n{rating}\n{kind}";
            rows.Add((id, rating, score, kind, mvid, tagsJoined, prompt, negative, paramsJson, lorasJson, blob));
        }

        lock (_lock)
        {
            EnsureOpen();
            EnsureExamplesSchema();
            using SqliteTransaction tx = _conn.BeginTransaction();
            using (SqliteCommand del = _conn.CreateCommand())
            {
                del.Transaction = tx;
                del.CommandText = "DELETE FROM civitai_examples";
                del.ExecuteNonQuery();
            }
            try
            {
                using (SqliteCommand delFts = _conn.CreateCommand())
                {
                    delFts.Transaction = tx;
                    delFts.CommandText = "DELETE FROM civitai_examples_fts";
                    delFts.ExecuteNonQuery();
                }
            }
            catch
            {
                // FTS missing
            }
            using (SqliteCommand ins = _conn.CreateCommand())
            {
                ins.Transaction = tx;
                ins.CommandText =
                    """
                    INSERT OR REPLACE INTO civitai_examples(
                      id, rating, score, kind, model_version_id, tags, prompt, negative, params_json, loras_json, search_blob)
                    VALUES($id,$r,$s,$k,$m,$t,$p,$n,$pj,$lj,$b)
                    """;
                var pid = ins.Parameters.Add("$id", SqliteType.Integer);
                var pr = ins.Parameters.Add("$r", SqliteType.Text);
                var ps = ins.Parameters.Add("$s", SqliteType.Integer);
                var pk = ins.Parameters.Add("$k", SqliteType.Text);
                var pm = ins.Parameters.Add("$m", SqliteType.Integer);
                var pt = ins.Parameters.Add("$t", SqliteType.Text);
                var pp = ins.Parameters.Add("$p", SqliteType.Text);
                var pn = ins.Parameters.Add("$n", SqliteType.Text);
                var ppj = ins.Parameters.Add("$pj", SqliteType.Text);
                var plj = ins.Parameters.Add("$lj", SqliteType.Text);
                var pb = ins.Parameters.Add("$b", SqliteType.Text);
                foreach (var row in rows)
                {
                    pid.Value = row.id;
                    pr.Value = row.rating;
                    ps.Value = row.score;
                    pk.Value = row.kind;
                    pm.Value = row.mvid;
                    pt.Value = row.tags;
                    pp.Value = row.prompt;
                    pn.Value = row.negative;
                    ppj.Value = row.paramsJson;
                    plj.Value = row.lorasJson;
                    pb.Value = row.blob;
                    ins.ExecuteNonQuery();
                }
            }
            tx.Commit();
            SetMeta(ExamplesMetaFp, fp);
            SetMeta(ExamplesMetaCount, rows.Count.ToString(CultureInfo.InvariantCulture));
            Logs.Info($"AssistentMemory: indexed {rows.Count} civitai examples from {Path.GetFileName(path)}");
            return rows.Count;
        }
    }

    public int ExampleCount()
    {
        lock (_lock)
        {
            EnsureOpen();
            if (!TableExists("civitai_examples"))
            {
                return 0;
            }
            using SqliteCommand c = _conn.CreateCommand();
            c.CommandText = "SELECT COUNT(*) FROM civitai_examples";
            return Convert.ToInt32(c.ExecuteScalar());
        }
    }

    /// <summary>FTS / LIKE lookup over Civitai example prompts. No embeddings.</summary>
    public JArray LookupExamples(string query, int limit = 5, string rating = null)
    {
        query = (query ?? "").Trim();
        if (query.Length < 1)
        {
            return [];
        }
        TryIndexExamples();
        int cap = Math.Clamp(limit, 1, 20);
        string ratingFilter = string.IsNullOrWhiteSpace(rating) ? null : rating.Trim().ToLowerInvariant();
        List<JObject> hits = [];
        HashSet<long> seen = [];

        void Add(SqliteDataReader reader)
        {
            long id = reader.GetInt64(0);
            if (!seen.Add(id))
            {
                return;
            }
            string tags = reader.IsDBNull(4) ? "" : reader.GetString(4);
            string prompt = reader.IsDBNull(5) ? "" : reader.GetString(5);
            string negative = reader.IsDBNull(6) ? "" : reader.GetString(6);
            string paramsJson = reader.IsDBNull(7) ? "{}" : reader.GetString(7);
            string lorasJson = reader.IsDBNull(8) ? "[]" : reader.GetString(8);
            JToken paramsTok = null;
            JToken lorasTok = null;
            try { paramsTok = JToken.Parse(paramsJson); } catch { paramsTok = new JObject(); }
            try { lorasTok = JToken.Parse(lorasJson); } catch { lorasTok = new JArray(); }
            hits.Add(new JObject
            {
                ["id"] = id,
                ["rating"] = reader.IsDBNull(1) ? "pg" : reader.GetString(1),
                ["score"] = reader.GetInt32(2),
                ["kind"] = reader.IsDBNull(3) ? "" : reader.GetString(3),
                ["tags"] = tags,
                ["prompt"] = prompt.Length > 1200 ? prompt[..1200] + "…" : prompt,
                ["negative"] = negative.Length > 400 ? negative[..400] + "…" : negative,
                ["params"] = paramsTok,
                ["loras"] = lorasTok,
                ["note"] = "EXAMPLE from Civitai — remix, do not copy 1:1",
            });
        }

        lock (_lock)
        {
            EnsureOpen();
            if (!TableExists("civitai_examples"))
            {
                return [];
            }

            string match = BuildFtsMatch(query);
            if (!string.IsNullOrWhiteSpace(match) && TableExists("civitai_examples_fts"))
            {
                try
                {
                    using SqliteCommand cmd = _conn.CreateCommand();
                    string sql =
                        """
                        SELECT e.id, e.rating, e.score, e.kind, e.tags, e.prompt, e.negative, e.params_json, e.loras_json
                        FROM civitai_examples e
                        WHERE e.id IN (SELECT rowid FROM civitai_examples_fts WHERE civitai_examples_fts MATCH $q)
                        """;
                    if (ratingFilter is not null)
                    {
                        sql += " AND lower(e.rating) = $r";
                    }
                    sql += " ORDER BY e.score DESC LIMIT $lim";
                    cmd.CommandText = sql;
                    cmd.Parameters.AddWithValue("$q", match);
                    if (ratingFilter is not null)
                    {
                        cmd.Parameters.AddWithValue("$r", ratingFilter);
                    }
                    cmd.Parameters.AddWithValue("$lim", cap);
                    using SqliteDataReader reader = cmd.ExecuteReader();
                    while (reader.Read())
                    {
                        Add(reader);
                    }
                }
                catch (Exception ex)
                {
                    Logs.Debug($"AssistentMemory examples FTS: {ex.Message}");
                }
            }

            if (hits.Count < cap)
            {
                try
                {
                    using SqliteCommand cmd = _conn.CreateCommand();
                    string sql =
                        """
                        SELECT id, rating, score, kind, tags, prompt, negative, params_json, loras_json
                        FROM civitai_examples
                        WHERE (tags LIKE $p ESCAPE '\' OR prompt LIKE $p ESCAPE '\' OR search_blob LIKE $p ESCAPE '\')
                        """;
                    if (ratingFilter is not null)
                    {
                        sql += " AND lower(rating) = $r";
                    }
                    sql += " ORDER BY score DESC LIMIT $lim";
                    cmd.CommandText = sql;
                    cmd.Parameters.AddWithValue("$p", "%" + EscapeLike(query) + "%");
                    if (ratingFilter is not null)
                    {
                        cmd.Parameters.AddWithValue("$r", ratingFilter);
                    }
                    cmd.Parameters.AddWithValue("$lim", cap);
                    using SqliteDataReader reader = cmd.ExecuteReader();
                    while (reader.Read() && hits.Count < cap)
                    {
                        Add(reader);
                    }
                }
                catch (Exception ex)
                {
                    Logs.Debug($"AssistentMemory examples LIKE: {ex.Message}");
                }
            }
        }

        return new JArray(hits.Take(cap));
    }
}