Add FTS lookup for seeded Civitai Krea2 prompt examples.
Index Assistent/civitai-examples.jsonl without embeddings and expose AssistentLookupExamples plus ask:examples hop for remix references. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Leonid Pershin
Cursor
@@ -0,0 +1,383 @@
|
||||
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));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user