From 8702b64e1239da23d871fe1d869337b81a75f598 Mon Sep 17 00:00:00 2001 From: Leonid Pershin Date: Sat, 22 Aug 2026 01:03:40 +0300 Subject: [PATCH] Ship Assistent disk persist, park LLM, and memory UI cleanup. Split the extension into partials, persist chats on the data volume, park/warm the chat model around Generate, and drop dual raw/persona dump paths. Co-authored-by: Cursor --- Assets/assistent.js | 90 ++++++- AssistentChatPipeline.cs | 182 +++++++++++--- AssistentMemory.Tags.cs | 373 +++++++++++++++++++++++++++++ AssistentMemory.cs | 433 ++++++++++++++++++++++++++++++---- AssistentMemoryApi.cs | 69 ++++++ AssistentPatch.cs | 63 ++++- AssistentVram.cs | 2 +- Config/_base/assistant.json | 11 + Config/_base/core/core.md | 11 +- Config/_base/skills/memory.md | 28 ++- README.md | 9 +- SwarmAssistentExtension.cs | 35 ++- 12 files changed, 1192 insertions(+), 114 deletions(-) create mode 100644 AssistentMemory.Tags.cs diff --git a/Assets/assistent.js b/Assets/assistent.js index 39e31a3..5b110a4 100644 --- a/Assets/assistent.js +++ b/Assets/assistent.js @@ -155,6 +155,8 @@ activeChatId: null, restoringChat: false, chatsPanelOpen: false, + chatsQuery: '', + chatsSearchHits: null, slashIndex: 0, llmParked: false, memoryRows: [], @@ -1769,6 +1771,23 @@ } } + function chatMatchesQuery(chat, q) { + if (!q) { + return true; + } + const title = String(chat?.title || '').toLowerCase(); + if (title.includes(q)) { + return true; + } + const msgs = chat?.messages || []; + for (const m of msgs) { + if (String(m?.content || '').toLowerCase().includes(q)) { + return true; + } + } + return false; + } + function renderChatsList() { const root = $('sa_chats_list'); if (!root) { @@ -1776,19 +1795,28 @@ } root.innerHTML = ''; syncHistoryBadge(); - const chats = (state.chats || []) + const q = (state.chatsQuery || '').trim().toLowerCase(); + let chats = (state.chats || []) .slice() .sort((a, b) => (b.updatedAt || 0) - (a.updatedAt || 0)) - .filter((c) => (c.messages || []).length > 0); + .filter((c) => (c.messages || []).length > 0 || (c.messages_count || 0) > 0); + if (q) { + const local = chats.filter((c) => chatMatchesQuery(c, q)); + const seen = new Set(local.map((c) => c.id)); + const extra = (state.chatsSearchHits || []).filter((h) => h && h.id && !seen.has(h.id)); + chats = local.concat(extra); + } if (!chats.length) { - root.innerHTML = '
Пока пусто. Напиши что-нибудь в чат — он появится здесь. Переключение вернёт и сообщения, и параметры Generate.
'; + root.innerHTML = q + ? '
Ничего не нашлось.
' + : '
Пока пусто. Напиши что-нибудь в чат — он появится здесь. Переключение вернёт и сообщения, и параметры Generate.
'; return; } for (const c of chats) { const row = document.createElement('div'); row.className = 'sa-chat-row' + (c.id === state.activeChatId ? ' sa-chat-row-active' : ''); row.dataset.id = c.id; - const n = (c.messages || []).length; + const n = (c.messages || []).length || Number(c.messages_count) || 0; const bits = []; if (c.params?.width && c.params?.height) { bits.push(`${c.params.width}×${c.params.height}`); @@ -1818,6 +1846,11 @@ btn?.setAttribute('aria-expanded', state.chatsPanelOpen ? 'true' : 'false'); if (state.chatsPanelOpen) { saveActiveChatToStore(); + const search = $('sa_chats_search'); + if (search) { + search.value = state.chatsQuery || ''; + search.focus(); + } renderChatsList(); } } @@ -1878,7 +1911,23 @@ return; } saveActiveChatToStore({ dropEmpty: true }); - const chat = findChat(id); + let chat = findChat(id); + if (!chat || !(chat.messages || []).length) { + try { + const full = await diskPersist()?.getChat?.(id); + if (full) { + const idx = (state.chats || []).findIndex((c) => c.id === id); + if (idx >= 0) { + state.chats[idx] = full; + } else { + state.chats.unshift(full); + } + chat = full; + } + } catch (e) { + console.warn('Assistent: getChat failed', id, e); + } + } if (!chat) { setStatus('Чат не найден'); return; @@ -3346,7 +3395,9 @@ if (!state.busy) { stopBusyUi(src ? 'Generate готов' : 'Generate завершён (новое изображение не найдено)'); } - if (state.view === 'chat') { + // Auto-critique loads the model itself on the next request — don't pay for it twice. + const paneVisible = !!document.getElementById('swarm_assistent_root')?.offsetParent; + if (state.view === 'chat' && paneVisible && !$('sa_auto_critique')?.checked) { warmLlm(); } if (src) { @@ -5140,6 +5191,9 @@ ? `Карточка Assistent сохранена · ${data.path}` : `Черновик + wanted · ${data.path}`); refreshInventory(() => renderCardsList()); + if (enqueue || !data.installed) { + refreshWantedQueue(); + } }, 0, (err) => setCardStatus(String(err || 'Ошибка сохранения')), @@ -5164,6 +5218,7 @@ }, (data) => { setCardStatus(data.already ? 'Уже в wanted' : `Wanted → ${data.path}`); + refreshWantedQueue(); }, 0, (err) => setCardStatus(String(err || 'Ошибка enqueue')), @@ -6275,6 +6330,7 @@ splitter.classList.remove('sa-dragging'); document.body.style.cursor = ''; document.body.style.userSelect = ''; + saveUiStateToDisk(); }); } @@ -6394,6 +6450,28 @@ switchToChat(id); } }); + let chatsSearchTimer = null; + $('sa_chats_search')?.addEventListener('input', () => { + const q = ($('sa_chats_search')?.value || '').trim(); + state.chatsQuery = q; + if (!q) { + state.chatsSearchHits = null; + renderChatsList(); + return; + } + renderChatsList(); + clearTimeout(chatsSearchTimer); + chatsSearchTimer = setTimeout(async () => { + try { + const hits = await diskPersist()?.searchChats?.(q); + if ((state.chatsQuery || '') !== q) { + return; + } + state.chatsSearchHits = Array.isArray(hits) ? hits : []; + renderChatsList(); + } catch (e) { /* ignore */ } + }, 220); + }); $('sa_tab_chat')?.addEventListener('click', () => setView('chat')); $('sa_tab_cards')?.addEventListener('click', () => setView('cards')); diff --git a/AssistentChatPipeline.cs b/AssistentChatPipeline.cs index 271d033..9372ba7 100644 --- a/AssistentChatPipeline.cs +++ b/AssistentChatPipeline.cs @@ -13,6 +13,7 @@ namespace Mrleo1nid.SwarmAssistent; public partial class SwarmAssistentExtension { const int MaxCivitaiHopsFallback = 2; + const int MaxToolHopsFallback = 4; List BuildOllamaMessages(string packName, bool includeBase, string contextJson, JArray userMessages, string extraSystem = null, string personaId = null, IEnumerable skillIds = null) { @@ -139,12 +140,12 @@ public partial class SwarmAssistentExtension Logs.Debug($"Assistent memory seed: {ex.Message}"); } - string retrieveQuery = BuildRetrieveQuery(userMessages, contextJson); + string retrieveQuery = BuildRetrieveQuery(userMessages, contextJson, packName); JArray hits = []; try { - int topK = Config.LoadAssistant(pid)["memory_top_k"]?.Value() ?? 10; - hits = await Memory.RetrieveAsync(root, retrieveQuery, topK, embed, Config.PersonaExtendsChain(pid)); + AssistentMemory.RetrieveOptions opt = MemoryRetrieveOptions(pid); + hits = await Memory.RetrieveAsync(root, retrieveQuery, opt.TopK, embed, Config.PersonaExtendsChain(pid), opt); } catch (Exception ex) { @@ -156,7 +157,9 @@ public partial class SwarmAssistentExtension JArray civitaiResults = []; string reply = ""; JObject lastRaw = null; - int maxHops = CfgInt("max_civitai_hops", MaxCivitaiHopsFallback); + int maxHops = Math.Max(CfgInt("max_civitai_hops", MaxCivitaiHopsFallback), CfgInt("max_tool_hops", MaxToolHopsFallback)); + HashSet hopDone = new(StringComparer.OrdinalIgnoreCase); + var chain = Config.PersonaExtendsChain(pid); for (int hop = 0; hop < maxHops; hop++) { if (onHopStart is not null) @@ -166,44 +169,37 @@ public partial class SwarmAssistentExtension (reply, lastRaw) = await CallOllamaChat(root, modelName, messages, stream: onDelta is not null, onDelta, pid); JObject patch = TryParsePatch(reply); await ApplyMemoryActions(root, patch, embed, pid); - if (hop + 1 >= maxHops || !WantsCivitaiSearch(patch)) + if (hop + 1 >= maxHops) { break; } - string query = ExtractSearchQuery(patch); - if (string.IsNullOrWhiteSpace(query)) + string tool = NextToolHop(patch); + if (string.IsNullOrWhiteSpace(tool)) { break; } - JObject search = await AssistentSearchCivitai(session, query, 8); - if (search["error"] is not null) + (string follow, JArray civitaiHop) = await RunToolHop(session, root, embed, pid, chain, patch, tool, hopDone); + if (follow is null) { - messages.Add(new JObject { ["role"] = "assistant", ["content"] = reply }); - messages.Add(new JObject - { - ["role"] = "user", - ["content"] = $"Civitai search failed: {search["error"]}. Continue without download — use only available_loras from context.", - }); - continue; + break; + } + if (civitaiHop is { Count: > 0 }) + { + civitaiResults = civitaiHop; } - civitaiResults = search["results"] as JArray ?? []; messages.Add(new JObject { ["role"] = "assistant", ["content"] = reply }); - messages.Add(new JObject - { - ["role"] = "user", - ["content"] = - "Civitai search results (JSON). Prefer `krea_likely: true`. Do NOT download yourself — the UI shows Confirm cards. " + - "Pick useful LoRAs from results or available_loras, emit a normal patch (prompt/loras). " + - "Omit search_civitai from actions unless you need a different query.\n```json\n" + - civitaiResults.ToString(Newtonsoft.Json.Formatting.None) + "\n```", - }); + messages.Add(new JObject { ["role"] = "user", ["content"] = follow }); } return (reply, lastRaw, civitaiResults); } - static string BuildRetrieveQuery(JArray userMessages, string contextJson) + static string BuildRetrieveQuery(JArray userMessages, string contextJson, string packName = null) { StringBuilder sb = new(); + if (!string.IsNullOrWhiteSpace(packName)) + { + sb.Append(packName).Append(' '); + } if (!string.IsNullOrWhiteSpace(contextJson)) { try @@ -216,7 +212,7 @@ public partial class SwarmAssistentExtension } if (ctx["enabled_loras"] is JArray en) { - foreach (JToken t in en.Take(8)) + foreach (JToken t in en.Take(12)) { string n = t?["name"]?.ToString() ?? t?.ToString(); if (!string.IsNullOrWhiteSpace(n)) @@ -229,23 +225,149 @@ public partial class SwarmAssistentExtension { sb.Append("krea ").Append(ctx["krea_profile"]).Append(' '); } + string aspect = ctx["aspect"]?.ToString(); + if (!string.IsNullOrWhiteSpace(aspect)) + { + sb.Append(aspect).Append(' '); + } + string prompt = ctx["prompt"]?.ToString(); + if (!string.IsNullOrWhiteSpace(prompt)) + { + sb.Append(prompt.Length > 400 ? prompt[..400] : prompt).Append(' '); + } } catch { // ignore } } - foreach (JToken msg in (userMessages ?? []).Reverse().Take(2)) + foreach (JToken msg in (userMessages ?? []).Reverse().Take(3)) { if (msg is JObject mo && string.Equals(mo["role"]?.ToString(), "user", StringComparison.OrdinalIgnoreCase)) { - sb.Append(mo["content"]?.ToString()).Append(' '); + string c = mo["content"]?.ToString() ?? ""; + sb.Append(c.Length > 500 ? c[..500] : c).Append(' '); } } string q = CollapseWs(sb.ToString()); return string.IsNullOrWhiteSpace(q) ? "krea2 prompting" : q; } + AssistentMemory.RetrieveOptions MemoryRetrieveOptions(string pid) + { + JObject a = Config.LoadAssistant(pid) ?? new JObject(); + AssistentMemory.RetrieveOptions opt = new() + { + TopK = a["memory_top_k"]?.Value() ?? 10, + MinScore = a["memory_min_score"]?.Value() ?? 0.32f, + ApplyQuotas = true, + }; + if (a["memory_quotas"] is JObject quotas) + { + Dictionary d = new(StringComparer.OrdinalIgnoreCase); + foreach (JProperty p in quotas.Properties()) + { + d[p.Name] = p.Value?.Value() ?? 2; + } + opt.Quotas = d; + } + return opt; + } + + async Task<(string follow, JArray civitai)> RunToolHop( + Session session, + string root, + string embed, + string pid, + IEnumerable chain, + JObject patch, + string tool, + HashSet hopDone) + { + if (tool == "memory_get") + { + JArray got = []; + foreach (JToken t in patch["memories"] as JArray ?? []) + { + if (t is not JObject mo) + { + continue; + } + string kind = mo["kind"]?.ToString() ?? "note"; + string key = mo["key"]?.ToString() ?? ""; + if (string.IsNullOrWhiteSpace(key)) + { + continue; + } + string sig = $"get:{kind}:{key}"; + if (!hopDone.Add(sig)) + { + continue; + } + JObject row = Memory.Get(kind, key, chain); + got.Add(row ?? new JObject { ["kind"] = kind, ["key"] = key, ["missing"] = true }); + } + if (got.Count == 0) + { + return (null, null); + } + return ( + "memory_get results (JSON). Use these facts; omit memory_get unless you need a different key.\n```json\n" + + got.ToString(Newtonsoft.Json.Formatting.None) + "\n```", + null); + } + if (tool == "memory_search") + { + string q = ExtractMemoryQuery(patch); + if (string.IsNullOrWhiteSpace(q) || !hopDone.Add("search:" + q)) + { + return (null, null); + } + string kind = patch["memory_kind"]?.ToString(); + int topK = Config.LoadAssistant(pid)["memory_top_k"]?.Value() ?? 10; + JArray rows = await Memory.SearchAsync(root, q, kind, topK, embed, chain); + return ( + "memory_search results (JSON, hybrid FTS+vector). Omit memory_search unless you need a different query.\n```json\n" + + rows.ToString(Newtonsoft.Json.Formatting.None) + "\n```", + null); + } + if (tool == "lookup_tags") + { + string q = ExtractTagQuery(patch); + if (string.IsNullOrWhiteSpace(q) || !hopDone.Add("tags:" + q)) + { + return (null, null); + } + int lim = Config.LoadAssistant(pid)["tag_lookup_limit"]?.Value() ?? 20; + JArray tags = Memory.LookupTags(q, lim); + return ( + "lookup_tags results from Danbooru csv (canonical name, aliases, post_count). Krea prompts stay natural prose — use this to check spelling/aliases, do not dump tag soup.\n```json\n" + + tags.ToString(Newtonsoft.Json.Formatting.None) + "\n```", + null); + } + if (tool == "civitai") + { + string query = ExtractSearchQuery(patch); + if (string.IsNullOrWhiteSpace(query) || !hopDone.Add("civitai:" + query)) + { + return (null, null); + } + JObject search = await AssistentSearchCivitai(session, query, 8); + if (search["error"] is not null) + { + return ($"Civitai search failed: {search["error"]}. Continue without download — use only available_loras from context.", null); + } + JArray civitaiResults = search["results"] as JArray ?? []; + return ( + "Civitai search results (JSON). Prefer `krea_likely: true`. Do NOT download yourself — the UI shows Confirm cards. " + + "Pick useful LoRAs from results or available_loras, emit a normal patch (prompt/loras). " + + "Omit search_civitai from actions unless you need a different query.\n```json\n" + + civitaiResults.ToString(Newtonsoft.Json.Formatting.None) + "\n```", + civitaiResults); + } + return (null, null); + } + static string InjectMemoryHits(string contextJson, JArray hits, JObject exact = null) { JObject ctx; diff --git a/AssistentMemory.Tags.cs b/AssistentMemory.Tags.cs new file mode 100644 index 0000000..ea05b41 --- /dev/null +++ b/AssistentMemory.Tags.cs @@ -0,0 +1,373 @@ +using System; +using System.Collections.Generic; +using System.Globalization; +using System.IO; +using System.Linq; +using System.Text; +using Microsoft.Data.Sqlite; +using Newtonsoft.Json.Linq; +using SwarmUI.Utils; + +namespace Mrleo1nid.SwarmAssistent; + +public sealed partial class AssistentMemory +{ + static readonly Dictionary TagCategories = new() + { + [0] = "general", + [1] = "artist", + [3] = "copyright", + [4] = "character", + [5] = "meta", + }; + + void TryIndexTags() + { + try + { + EnsureTagsIndex(); + } + catch (Exception ex) + { + Logs.Debug($"AssistentMemory tags index: {ex.Message}"); + } + } + + void EnsureTagsSchema() + { + Exec( + """ + CREATE TABLE IF NOT EXISTS tags ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL COLLATE NOCASE UNIQUE, + category INTEGER NOT NULL DEFAULT 0, + post_count INTEGER NOT NULL DEFAULT 0, + aliases TEXT NOT NULL DEFAULT '' + ); + """); + Exec("CREATE INDEX IF NOT EXISTS idx_tags_count ON tags(post_count DESC);"); + Exec( + """ + CREATE VIRTUAL TABLE IF NOT EXISTS tags_fts USING fts5( + name, + aliases, + tokenize = 'unicode61 remove_diacritics 2' + ); + """); + Exec( + """ + CREATE TRIGGER IF NOT EXISTS tags_fts_ai AFTER INSERT ON tags BEGIN + INSERT INTO tags_fts(rowid, name, aliases) VALUES (new.id, new.name, new.aliases); + END; + """); + Exec( + """ + CREATE TRIGGER IF NOT EXISTS tags_fts_ad AFTER DELETE ON tags BEGIN + INSERT INTO tags_fts(tags_fts, rowid) VALUES('delete', old.id); + END; + """); + Exec( + """ + CREATE TRIGGER IF NOT EXISTS tags_fts_au AFTER UPDATE ON tags BEGIN + INSERT INTO tags_fts(tags_fts, rowid) VALUES('delete', old.id); + INSERT INTO tags_fts(rowid, name, aliases) VALUES (new.id, new.name, new.aliases); + END; + """); + } + + public string FindAutocompleteCsv() + { + string dir = Path.Combine(_dataRoot, "Data", "Autocompletions"); + if (!Directory.Exists(dir)) + { + return null; + } + string preferred = Path.Combine(dir, "danbooru.csv"); + if (File.Exists(preferred)) + { + return preferred; + } + string[] csvs = Directory.GetFiles(dir, "*.csv"); + if (csvs.Length == 0) + { + return null; + } + return csvs.OrderByDescending(f => new FileInfo(f).Length).First(); + } + + static string CsvFingerprint(string csvPath) + { + var info = new FileInfo(csvPath); + string sha = ""; + string meta = csvPath + ".gpu-rent-meta.json"; + if (File.Exists(meta)) + { + try + { + JObject o = JObject.Parse(File.ReadAllText(meta)); + sha = o["github_blob_sha"]?.ToString() ?? ""; + } + catch + { + // ignore + } + } + return $"{info.Length}:{info.LastWriteTimeUtc.Ticks}:{sha}"; + } + + /// Load SwarmUI Autocompletions csv into FTS (no embeddings). No-op if fingerprint matches. + public int EnsureTagsIndex() + { + string csv = FindAutocompleteCsv(); + if (string.IsNullOrWhiteSpace(csv) || !File.Exists(csv)) + { + return 0; + } + string fp = CsvFingerprint(csv); + lock (_lock) + { + EnsureOpen(); + if (string.Equals(GetMeta("tags_csv_fp"), fp, StringComparison.Ordinal)) + { + using SqliteCommand c = _conn.CreateCommand(); + c.CommandText = "SELECT COUNT(*) FROM tags"; + return Convert.ToInt32(c.ExecuteScalar()); + } + } + + List<(string name, int cat, int count, string aliases)> rows = []; + foreach (string line in File.ReadLines(csv, Encoding.UTF8)) + { + if (string.IsNullOrWhiteSpace(line) || line.StartsWith('#')) + { + continue; + } + List cols = ParseCsvLine(line); + if (cols.Count < 1) + { + continue; + } + string name = cols[0].Trim(); + if (string.IsNullOrWhiteSpace(name)) + { + continue; + } + int cat = 0, count = 0; + if (cols.Count > 1) + { + _ = int.TryParse(cols[1].Trim(), NumberStyles.Integer, CultureInfo.InvariantCulture, out cat); + } + if (cols.Count > 2) + { + _ = int.TryParse(cols[2].Trim(), NumberStyles.Integer, CultureInfo.InvariantCulture, out count); + } + string aliases = cols.Count > 3 ? cols[3].Trim() : ""; + rows.Add((name, cat, count, aliases)); + } + + lock (_lock) + { + EnsureOpen(); + using SqliteTransaction tx = _conn.BeginTransaction(); + using (SqliteCommand del = _conn.CreateCommand()) + { + del.Transaction = tx; + del.CommandText = "DELETE FROM tags"; + del.ExecuteNonQuery(); + } + try + { + using (SqliteCommand delFts = _conn.CreateCommand()) + { + delFts.Transaction = tx; + delFts.CommandText = "DELETE FROM tags_fts"; + delFts.ExecuteNonQuery(); + } + } + catch + { + // FTS table missing + } + using (SqliteCommand ins = _conn.CreateCommand()) + { + ins.Transaction = tx; + ins.CommandText = "INSERT OR REPLACE INTO tags(name, category, post_count, aliases) VALUES($n,$c,$p,$a)"; + var pn = ins.Parameters.Add("$n", SqliteType.Text); + var pc = ins.Parameters.Add("$c", SqliteType.Integer); + var pp = ins.Parameters.Add("$p", SqliteType.Integer); + var pa = ins.Parameters.Add("$a", SqliteType.Text); + foreach (var row in rows) + { + pn.Value = row.name; + pc.Value = row.cat; + pp.Value = row.count; + pa.Value = row.aliases ?? ""; + ins.ExecuteNonQuery(); + } + } + tx.Commit(); + SetMeta("tags_csv_fp", fp); + SetMeta("tags_count", rows.Count.ToString(CultureInfo.InvariantCulture)); + Logs.Debug($"AssistentMemory indexed {rows.Count} tags from {Path.GetFileName(csv)}"); + return rows.Count; + } + } + + public int TagCount() + { + lock (_lock) + { + EnsureOpen(); + if (!TableExists("tags")) + { + return 0; + } + using SqliteCommand cmd = _conn.CreateCommand(); + cmd.CommandText = "SELECT COUNT(*) FROM tags"; + return Convert.ToInt32(cmd.ExecuteScalar()); + } + } + + /// Prefix + FTS lookup over the Danbooru csv. No embeddings. + public JArray LookupTags(string query, int limit = 20) + { + query = (query ?? "").Trim(); + if (query.Length < 1) + { + return []; + } + TryIndexTags(); + int cap = Math.Clamp(limit, 1, 40); + string underscored = query.Replace(' ', '_'); + Dictionary byName = new(StringComparer.OrdinalIgnoreCase); + + void Add(SqliteDataReader reader) + { + string name = reader.GetString(0); + if (byName.ContainsKey(name)) + { + return; + } + int cat = reader.GetInt32(1); + TagCategories.TryGetValue(cat, out string catName); + byName[name] = new JObject + { + ["name"] = name, + ["category"] = cat, + ["category_name"] = catName ?? "general", + ["post_count"] = reader.GetInt32(2), + ["aliases"] = reader.IsDBNull(3) ? "" : reader.GetString(3), + }; + } + + lock (_lock) + { + EnsureOpen(); + if (!TableExists("tags")) + { + return []; + } + try + { + using SqliteCommand cmd = _conn.CreateCommand(); + cmd.CommandText = + """ + SELECT name, category, post_count, aliases FROM tags + WHERE name LIKE $p ESCAPE '\' OR (aliases != '' AND aliases LIKE $a) + ORDER BY post_count DESC LIMIT $lim + """; + cmd.Parameters.AddWithValue("$p", EscapeLike(underscored) + "%"); + cmd.Parameters.AddWithValue("$a", "%" + EscapeLike(query) + "%"); + cmd.Parameters.AddWithValue("$lim", cap); + using SqliteDataReader reader = cmd.ExecuteReader(); + while (reader.Read()) + { + Add(reader); + } + } + catch (Exception ex) + { + Logs.Debug($"AssistentMemory tag prefix: {ex.Message}"); + } + + string match = BuildFtsMatch(query); + if (!string.IsNullOrWhiteSpace(match) && TableExists("tags_fts") && byName.Count < cap) + { + try + { + using SqliteCommand cmd = _conn.CreateCommand(); + cmd.CommandText = + """ + SELECT t.name, t.category, t.post_count, t.aliases + FROM tags t + WHERE t.id IN (SELECT rowid FROM tags_fts WHERE tags_fts MATCH $q) + ORDER BY t.post_count DESC LIMIT $lim + """; + cmd.Parameters.AddWithValue("$q", match); + cmd.Parameters.AddWithValue("$lim", cap); + using SqliteDataReader reader = cmd.ExecuteReader(); + while (reader.Read()) + { + Add(reader); + } + } + catch (Exception ex) + { + Logs.Debug($"AssistentMemory tag FTS: {ex.Message}"); + } + } + } + + return new JArray(byName.Values.OrderByDescending(t => t["post_count"]?.Value() ?? 0).Take(cap)); + } + + static string EscapeLike(string s) + { + return (s ?? "").Replace(@"\", @"\\").Replace("%", @"\%").Replace("_", @"\_"); + } + + static List ParseCsvLine(string line) + { + List cols = []; + StringBuilder cur = new(); + bool quoted = false; + for (int i = 0; i < line.Length; i++) + { + char c = line[i]; + if (quoted) + { + if (c == '"') + { + if (i + 1 < line.Length && line[i + 1] == '"') + { + cur.Append('"'); + i++; + } + else + { + quoted = false; + } + } + else + { + cur.Append(c); + } + } + else if (c == '"') + { + quoted = true; + } + else if (c == ',') + { + cols.Add(cur.ToString()); + cur.Clear(); + } + else + { + cur.Append(c); + } + } + cols.Add(cur.ToString()); + return cols; + } +} diff --git a/AssistentMemory.cs b/AssistentMemory.cs index 2f1c0e5..27d91d4 100644 --- a/AssistentMemory.cs +++ b/AssistentMemory.cs @@ -4,6 +4,7 @@ using System.IO; using System.Linq; using System.Net.Http; using System.Text; +using System.Text.RegularExpressions; using System.Threading.Tasks; using Microsoft.Data.Sqlite; using Newtonsoft.Json.Linq; @@ -14,10 +15,31 @@ namespace Mrleo1nid.SwarmAssistent; /// Local SQLite vector memory with Ollama /api/embed. /// Two layers: shared (persona='') is visible to every personality; personal (persona=id) /// is not written back to shared. On retrieve, personal overwrites shared on the same kind+key. -public sealed class AssistentMemory : IDisposable +public sealed partial class AssistentMemory : IDisposable { public const string SharedPersona = ""; + public sealed class RetrieveOptions + { + public int TopK { get; set; } = 10; + public float MinScore { get; set; } = 0.32f; + public string KindFilter { get; set; } + public IReadOnlyDictionary Quotas { get; set; } + public bool ApplyQuotas { get; set; } = true; + } + + static readonly Dictionary DefaultQuotas = new(StringComparer.OrdinalIgnoreCase) + { + ["card"] = 3, + ["lora"] = 3, + ["pitfall"] = 3, + ["path"] = 2, + ["note"] = 4, + ["model"] = 2, + ["aspect"] = 1, + }; + + readonly string _dataRoot; readonly string _dbPath; readonly HttpClient _http; readonly object _lock = new(); @@ -28,7 +50,8 @@ public sealed class AssistentMemory : IDisposable public AssistentMemory(string dataRoot, HttpClient http, string defaultEmbedModel = "nomic-embed-text") { - string dir = Path.Combine(dataRoot ?? ".", "Assistent", "memory"); + _dataRoot = string.IsNullOrWhiteSpace(dataRoot) ? "." : dataRoot; + string dir = Path.Combine(_dataRoot, "Assistent", "memory"); Directory.CreateDirectory(dir); _dbPath = Path.Combine(dir, "assistent.sqlite"); _http = http; @@ -63,6 +86,8 @@ public sealed class AssistentMemory : IDisposable } _conn = new SqliteConnection($"Data Source={_dbPath}"); _conn.Open(); + TryPragma("journal_mode=WAL"); + TryPragma("busy_timeout=5000"); using (SqliteCommand cmd = _conn.CreateCommand()) { cmd.CommandText = @@ -96,6 +121,23 @@ public sealed class AssistentMemory : IDisposable """; idx.ExecuteNonQuery(); } + EnsureFts(); + try + { + EnsureTagsSchema(); + } + catch (Exception ex) + { + Logs.Debug($"AssistentMemory tags schema: {ex.Message}"); + } + try + { + EnsureStoreSchema(); + } + catch (Exception ex) + { + Logs.Debug($"AssistentMemory store schema: {ex.Message}"); + } _embedModel = GetMeta("embed_model") ?? _embedModel; _ = int.TryParse(GetMeta("dims"), out _dims); _ = int.TryParse(GetMeta("seed_version"), out _seedVersion); @@ -154,6 +196,89 @@ public sealed class AssistentMemory : IDisposable Logs.Debug("AssistentMemory: migrated sqlite to shared+personal persona column (existing rows → shared)"); } + void TryPragma(string pragma) + { + try + { + using SqliteCommand cmd = _conn.CreateCommand(); + cmd.CommandText = "PRAGMA " + pragma; + cmd.ExecuteNonQuery(); + } + catch (Exception ex) + { + Logs.Debug($"AssistentMemory PRAGMA {pragma}: {ex.Message}"); + } + } + + bool TableExists(string name) + { + using SqliteCommand cmd = _conn.CreateCommand(); + cmd.CommandText = "SELECT 1 FROM sqlite_master WHERE type IN ('table','view') AND name = $n LIMIT 1"; + cmd.Parameters.AddWithValue("$n", name); + return cmd.ExecuteScalar() is not null; + } + + void Exec(string sql) + { + using SqliteCommand cmd = _conn.CreateCommand(); + cmd.CommandText = sql; + cmd.ExecuteNonQuery(); + } + + void EnsureFts() + { + try + { + Exec( + """ + CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts USING fts5( + key, + text, + tokenize = 'unicode61 remove_diacritics 2' + ); + """); + Exec( + """ + CREATE TRIGGER IF NOT EXISTS memories_fts_ai AFTER INSERT ON memories BEGIN + INSERT INTO memories_fts(rowid, key, text) VALUES (new.id, new.key, new.text); + END; + """); + Exec( + """ + CREATE TRIGGER IF NOT EXISTS memories_fts_ad AFTER DELETE ON memories BEGIN + INSERT INTO memories_fts(memories_fts, rowid) VALUES('delete', old.id); + END; + """); + Exec( + """ + CREATE TRIGGER IF NOT EXISTS memories_fts_au AFTER UPDATE ON memories BEGIN + INSERT INTO memories_fts(memories_fts, rowid) VALUES('delete', old.id); + INSERT INTO memories_fts(rowid, key, text) VALUES (new.id, new.key, new.text); + END; + """); + int mem = 0, fts = 0; + using (SqliteCommand c = _conn.CreateCommand()) + { + c.CommandText = "SELECT COUNT(*) FROM memories"; + mem = Convert.ToInt32(c.ExecuteScalar()); + } + using (SqliteCommand c = _conn.CreateCommand()) + { + c.CommandText = "SELECT COUNT(*) FROM memories_fts"; + fts = Convert.ToInt32(c.ExecuteScalar()); + } + if (mem != fts) + { + Exec("DELETE FROM memories_fts"); + Exec("INSERT INTO memories_fts(rowid, key, text) SELECT id, key, text FROM memories"); + } + } + catch (Exception ex) + { + Logs.Debug($"AssistentMemory FTS5 unavailable, cosine-only: {ex.Message}"); + } + } + string GetMeta(string key) { using SqliteCommand cmd = _conn.CreateCommand(); @@ -268,6 +393,7 @@ public sealed class AssistentMemory : IDisposable List docs = config.LoadMemorySeedDocs(); if (docs.Count == 0) { + TryIndexTags(); return; } @@ -287,6 +413,7 @@ public sealed class AssistentMemory : IDisposable } if (!missing) { + TryIndexTags(); return; } } @@ -299,6 +426,7 @@ public sealed class AssistentMemory : IDisposable catch (Exception ex) { Logs.Debug($"AssistentMemory seed defer (embed unavailable): {ex.Message}"); + TryIndexTags(); return; } @@ -356,6 +484,7 @@ public sealed class AssistentMemory : IDisposable Logs.Debug($"AssistentMemory seed item {kind}/{key}: {ex.Message}"); } } + TryIndexTags(); } public void Upsert(string kind, string key, string text, string source, JToken meta, float[] embedding, string persona = null) @@ -428,29 +557,8 @@ public sealed class AssistentMemory : IDisposable } } - /// Retrieve shared + the given persona chain. Personal overwrites shared (and parent personas) on kind+key. - public async Task RetrieveAsync(string baseUrl, string query, int topK = 10, string modelOverride = null, IEnumerable personaChain = null) + static Dictionary PersonaRankMap(IEnumerable personaChain) { - if (string.IsNullOrWhiteSpace(query)) - { - return []; - } - lock (_lock) - { - EnsureOpen(); - } - string model = string.IsNullOrWhiteSpace(modelOverride) ? _embedModel : modelOverride; - float[] q; - try - { - q = await EmbedAsync(baseUrl, model, query); - } - catch (Exception ex) - { - Logs.Debug($"AssistentMemory retrieve embed: {ex.Message}"); - return []; - } - Dictionary rank = new(StringComparer.OrdinalIgnoreCase) { [SharedPersona] = 0, @@ -465,46 +573,219 @@ public sealed class AssistentMemory : IDisposable } rank[p] = i++; } + return rank; + } - List<(float score, int personaRank, int sourceRank, JObject row)> scored = []; + static string BuildFtsMatch(string query) + { + if (string.IsNullOrWhiteSpace(query)) + { + return null; + } + HashSet tokens = new(StringComparer.OrdinalIgnoreCase); + foreach (Match m in Regex.Matches(query, @"[\p{L}\p{N}_-]{2,}")) + { + string t = m.Value.Replace("\"", "").Trim('_', '-'); + if (t.Length >= 2) + { + tokens.Add(t); + } + string spaced = t.Replace('_', ' ').Replace('-', ' '); + if (!string.Equals(spaced, t, StringComparison.OrdinalIgnoreCase) && spaced.Length >= 2) + { + foreach (string p in spaced.Split(' ', StringSplitOptions.RemoveEmptyEntries)) + { + if (p.Length >= 2) + { + tokens.Add(p); + } + } + } + if (tokens.Count >= 12) + { + break; + } + } + if (tokens.Count == 0) + { + return null; + } + return string.Join(" OR ", tokens.Select(t => $"\"{t}\"")); + } + + Dictionary FtsRowRanks(string match, int limit) + { + Dictionary ranks = []; + if (string.IsNullOrWhiteSpace(match) || !TableExists("memories_fts")) + { + return ranks; + } + try + { + using SqliteCommand cmd = _conn.CreateCommand(); + cmd.CommandText = "SELECT rowid FROM memories_fts WHERE memories_fts MATCH $q ORDER BY rank LIMIT $lim"; + cmd.Parameters.AddWithValue("$q", match); + cmd.Parameters.AddWithValue("$lim", limit); + using SqliteDataReader reader = cmd.ExecuteReader(); + int r = 0; + while (reader.Read()) + { + ranks[reader.GetInt64(0)] = r++; + } + } + catch (Exception ex) + { + Logs.Debug($"AssistentMemory FTS match: {ex.Message}"); + } + return ranks; + } + + static float KeyBoost(string query, string key) + { + if (string.IsNullOrWhiteSpace(query)) + { + return 0; + } + string q = query.Replace('_', ' '); + string k = (key ?? "").Replace('_', ' '); + if (k.Length >= 3 && q.Contains(k, StringComparison.OrdinalIgnoreCase)) + { + return 0.22f; + } + if (k.Length >= 3 && k.Contains(q.Trim(), StringComparison.OrdinalIgnoreCase) && q.Trim().Length >= 4) + { + return 0.12f; + } + return 0; + } + + static JArray ApplyQuotas(IEnumerable ordered, int topK, IReadOnlyDictionary quotas) + { + Dictionary used = new(StringComparer.OrdinalIgnoreCase); + List picked = []; + List overflow = []; + foreach (JObject row in ordered) + { + string kind = row["kind"]?.ToString() ?? "note"; + int cap = 2; + if (quotas is not null && quotas.TryGetValue(kind, out int q)) + { + cap = q; + } + else if (DefaultQuotas.TryGetValue(kind, out int d)) + { + cap = d; + } + used.TryGetValue(kind, out int n); + if (n < cap) + { + picked.Add(row); + used[kind] = n + 1; + } + else + { + overflow.Add(row); + } + if (picked.Count >= topK) + { + return new JArray(picked); + } + } + foreach (JObject row in overflow) + { + if (picked.Count >= topK) + { + break; + } + picked.Add(row); + } + return new JArray(picked); + } + + /// Retrieve shared + the given persona chain. Hybrid FTS + cosine; personal overwrites shared on kind+key. + public async Task RetrieveAsync(string baseUrl, string query, int topK = 10, string modelOverride = null, IEnumerable personaChain = null, RetrieveOptions options = null) + { + options ??= new RetrieveOptions { TopK = topK }; + if (options.TopK <= 0) + { + options.TopK = topK; + } + if (string.IsNullOrWhiteSpace(query)) + { + return []; + } lock (_lock) { EnsureOpen(); + } + string model = string.IsNullOrWhiteSpace(modelOverride) ? _embedModel : modelOverride; + float[] qvec = null; + try + { + qvec = await EmbedAsync(baseUrl, model, query); + } + catch (Exception ex) + { + Logs.Debug($"AssistentMemory retrieve embed: {ex.Message}"); + } + + Dictionary rank = PersonaRankMap(personaChain); + string kindFilter = string.IsNullOrWhiteSpace(options.KindFilter) ? null : options.KindFilter.Trim().ToLowerInvariant(); + string ftsMatch = BuildFtsMatch(query); + + List<(float hybrid, float cosine, int personaRank, int sourceRank, bool fts, JObject row)> scored = []; + lock (_lock) + { + EnsureOpen(); + Dictionary ftsRanks = FtsRowRanks(ftsMatch, Math.Max(40, options.TopK * 4)); using SqliteCommand cmd = _conn.CreateCommand(); - cmd.CommandText = "SELECT kind, key, text, source, meta_json, embedding, persona FROM memories WHERE embedding IS NOT NULL"; + cmd.CommandText = "SELECT id, kind, key, text, source, embedding, persona FROM memories"; using SqliteDataReader reader = cmd.ExecuteReader(); while (reader.Read()) { + long id = reader.GetInt64(0); + string kind = reader.GetString(1); + if (kindFilter is not null && !string.Equals(kind, kindFilter, StringComparison.OrdinalIgnoreCase)) + { + continue; + } string persona = reader.IsDBNull(6) ? SharedPersona : reader.GetString(6) ?? SharedPersona; if (!rank.TryGetValue(persona, out int personaRank)) { continue; } - float[] emb = BytesToFloats(reader.IsDBNull(5) ? null : (byte[])reader.GetValue(5)); - float score = Cosine(q, emb); - if (float.IsNegativeInfinity(score)) + string key = reader.GetString(2); + string text = reader.GetString(3); + string source = reader.GetString(4); + float cosine = 0; + if (qvec is not null && !reader.IsDBNull(5)) { - continue; + cosine = Cosine(qvec, BytesToFloats((byte[])reader.GetValue(5))); + if (float.IsNegativeInfinity(cosine)) + { + cosine = 0; + } } - string kind = reader.GetString(0); - string key = reader.GetString(1); - string source = reader.GetString(3); + bool ftsHit = ftsRanks.TryGetValue(id, out int ftsRank); + float ftsBoost = ftsHit ? 0.28f * (1f - (ftsRank / 50f)) : 0; + float hybrid = cosine + ftsBoost + KeyBoost(query, key); bool shared = persona == SharedPersona; - scored.Add((score, personaRank, SourceRank(source), new JObject + scored.Add((hybrid, cosine, personaRank, SourceRank(source), ftsHit, new JObject { ["kind"] = kind, ["key"] = key, - ["text"] = reader.GetString(2), + ["text"] = text, ["source"] = source, ["scope"] = shared ? "shared" : "personal", ["persona"] = shared ? "shared" : persona, - ["score"] = Math.Round(score, 4), + ["score"] = Math.Round(hybrid, 4), + ["cosine"] = Math.Round(cosine, 4), + ["fts"] = ftsHit, })); } } - // Personal (and later parents) overwrite shared on the same kind+key; user beats bundled. - Dictionary best = new(StringComparer.OrdinalIgnoreCase); + Dictionary best = new(StringComparer.OrdinalIgnoreCase); foreach (var item in scored) { string id = $"{item.row["kind"]}\n{item.row["key"]}"; @@ -512,7 +793,7 @@ public sealed class AssistentMemory : IDisposable { if (item.personaRank < cur.personaRank || (item.personaRank == cur.personaRank && item.sourceRank < cur.sourceRank) - || (item.personaRank == cur.personaRank && item.sourceRank == cur.sourceRank && item.score <= cur.score)) + || (item.personaRank == cur.personaRank && item.sourceRank == cur.sourceRank && item.hybrid <= cur.hybrid)) { continue; } @@ -520,7 +801,79 @@ public sealed class AssistentMemory : IDisposable best[id] = item; } - return new JArray(best.Values.OrderByDescending(s => s.score).Take(Math.Clamp(topK, 1, 30)).Select(s => s.row)); + float min = options.MinScore; + IEnumerable ordered = best.Values + .Where(s => s.fts || s.hybrid >= min || s.cosine >= min) + .OrderByDescending(s => s.hybrid) + .Select(s => s.row); + + int k = Math.Clamp(options.TopK, 1, 30); + if (options.ApplyQuotas) + { + return ApplyQuotas(ordered, k, options.Quotas ?? DefaultQuotas); + } + return new JArray(ordered.Take(k)); + } + + /// Exact kind+key read with the same personal-over-shared overlay as retrieve. + public JObject Get(string kind, string key, IEnumerable personaChain = null) + { + kind = (kind ?? "note").Trim().ToLowerInvariant(); + key = (key ?? "").Trim(); + if (string.IsNullOrWhiteSpace(key)) + { + return null; + } + Dictionary rank = PersonaRankMap(personaChain); + JObject best = null; + int bestPersona = -1, bestSource = -1; + lock (_lock) + { + EnsureOpen(); + using SqliteCommand cmd = _conn.CreateCommand(); + cmd.CommandText = "SELECT kind, key, text, source, persona, updated FROM memories WHERE kind = $kind AND key = $key"; + cmd.Parameters.AddWithValue("$kind", kind); + cmd.Parameters.AddWithValue("$key", key); + using SqliteDataReader reader = cmd.ExecuteReader(); + while (reader.Read()) + { + string persona = reader.IsDBNull(4) ? SharedPersona : reader.GetString(4) ?? SharedPersona; + if (!rank.TryGetValue(persona, out int personaRank)) + { + continue; + } + int src = SourceRank(reader.GetString(3)); + if (best is not null && (personaRank < bestPersona || (personaRank == bestPersona && src <= bestSource))) + { + continue; + } + bestPersona = personaRank; + bestSource = src; + bool shared = persona == SharedPersona; + best = new JObject + { + ["kind"] = reader.GetString(0), + ["key"] = reader.GetString(1), + ["text"] = reader.GetString(2), + ["source"] = reader.GetString(3), + ["scope"] = shared ? "shared" : "personal", + ["persona"] = shared ? "shared" : persona, + ["updated"] = reader.IsDBNull(5) ? 0 : reader.GetInt64(5), + }; + } + } + return best; + } + + public async Task SearchAsync(string baseUrl, string query, string kind, int topK, string modelOverride, IEnumerable personaChain) + { + return await RetrieveAsync(baseUrl, query, topK, modelOverride, personaChain, new RetrieveOptions + { + TopK = Math.Clamp(topK, 1, 30), + MinScore = 0.18f, + KindFilter = kind, + ApplyQuotas = false, + }); } public JArray ListAll(int limit = 200) diff --git a/AssistentMemoryApi.cs b/AssistentMemoryApi.cs index 536e476..db66a72 100644 --- a/AssistentMemoryApi.cs +++ b/AssistentMemoryApi.cs @@ -173,6 +173,75 @@ public partial class SwarmAssistentExtension } } + public async Task AssistentSearchMemory(Session session, string query, string kind = null, int limit = 10, string persona = null, string baseUrl = null, string embed_model = null) + { + if (Memory is null) + { + return new JObject { ["error"] = "memory not ready" }; + } + query = (query ?? "").Trim(); + if (query.Length < 2) + { + return new JObject { ["error"] = "query required" }; + } + string pid = AssistentConfig.SafeId(persona) ?? Config.DefaultPersonaId(); + try + { + JArray rows = await Memory.SearchAsync(MemoryBaseUrl(baseUrl), query, kind, limit, MemoryEmbedModel(embed_model), Config.PersonaExtendsChain(pid)); + return new JObject { ["success"] = true, ["query"] = query, ["kind"] = kind ?? "", ["memories"] = rows }; + } + catch (Exception ex) + { + return new JObject { ["error"] = $"memory search: {ex.Message}" }; + } + } + + public async Task AssistentGetMemory(Session session, string kind, string key, string persona = null) + { + await Task.CompletedTask; + if (Memory is null) + { + return new JObject { ["error"] = "memory not ready" }; + } + string pid = AssistentConfig.SafeId(persona) ?? Config.DefaultPersonaId(); + JObject row = Memory.Get(kind, key, Config.PersonaExtendsChain(pid)); + if (row is null) + { + return new JObject { ["success"] = true, ["missing"] = true, ["kind"] = kind, ["key"] = key }; + } + return new JObject { ["success"] = true, ["memory"] = row }; + } + + public async Task AssistentLookupTags(Session session, string query, int limit = 20) + { + await Task.CompletedTask; + if (Memory is null) + { + return new JObject { ["error"] = "memory not ready" }; + } + query = (query ?? "").Trim(); + if (query.Length < 1) + { + return new JObject { ["error"] = "query required" }; + } + try + { + JArray tags = Memory.LookupTags(query, limit); + return new JObject + { + ["success"] = true, + ["query"] = query, + ["tags"] = tags, + ["indexed"] = Memory.TagCount(), + ["csv"] = Memory.FindAutocompleteCsv() ?? "", + }; + } + catch (Exception ex) + { + return new JObject { ["error"] = $"tag lookup: {ex.Message}" }; + } + } + /// The gpu-rent wanted queue (models pending the next up) — count + entries. public async Task AssistentListWanted(Session session) { diff --git a/AssistentPatch.cs b/AssistentPatch.cs index 92ec4a8..435e0af 100644 --- a/AssistentPatch.cs +++ b/AssistentPatch.cs @@ -19,6 +19,7 @@ public partial class SwarmAssistentExtension "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", ]; static bool HasValue(JObject obj, string key) @@ -91,8 +92,66 @@ public partial class SwarmAssistentExtension return string.IsNullOrWhiteSpace(q) ? null : q; } - static bool WantsCivitaiSearch(JObject patch) + static bool ActionsContain(JObject patch, string action) { - return !string.IsNullOrWhiteSpace(ExtractSearchQuery(patch)); + 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, "search_civitai") || !string.IsNullOrWhiteSpace(ExtractSearchQuery(patch))) + { + return "civitai"; + } + return null; + } + + static bool WantsCivitaiSearch(JObject patch) => NextToolHop(patch) == "civitai"; } diff --git a/AssistentVram.cs b/AssistentVram.cs index 9843c10..2b28c8b 100644 --- a/AssistentVram.cs +++ b/AssistentVram.cs @@ -57,7 +57,7 @@ public partial class SwarmAssistentExtension } /// Single-token chat so the model is resident again by the time the user types. - public async Task AssistentWarmLlm(Session session, string baseUrl, string model, string persona = null) + public async Task AssistentWarmLlm(Session session, string baseUrl, string model) { string root = NormalizeBaseUrl(baseUrl); string name = (model ?? "").Trim(); diff --git a/Config/_base/assistant.json b/Config/_base/assistant.json index 690eb02..4afe903 100644 --- a/Config/_base/assistant.json +++ b/Config/_base/assistant.json @@ -10,6 +10,17 @@ "default_persona": "neutral", "embed_model": "nomic-embed-text", "memory_top_k": 10, + "memory_min_score": 0.32, + "max_tool_hops": 4, + "tag_lookup_limit": 20, + "memory_quotas": { + "card": 3, + "lora": 3, + "pitfall": 3, + "path": 2, + "note": 4, + "model": 2 + }, "seed_version": 2, "gate": { "architecture": "krea2", diff --git a/Config/_base/core/core.md b/Config/_base/core/core.md index 8baae90..15ab6c0 100644 --- a/Config/_base/core/core.md +++ b/Config/_base/core/core.md @@ -11,7 +11,7 @@ When instructions conflict, apply this order (highest wins): 3. **Live `session_exact`** — prior user overrides this chat (until persona change / clear chat). 4. **Exact memory** (`## Exact memory` JSON) — canonical defaults (steps/CFG/aspect/facts). Persona overlays are already merged into it. 5. **Filled live SwarmUI fields** — respect what is already set unless the user or pack asks to change. -6. **`memory_hits` (vector RAG)** — notes, pitfalls, LoRA blurbs. Shared hits apply to every persona; personal hits are this persona only and overwrite shared on the same kind+key. Never override exact numbers or the user’s param request. +6. **`memory_hits` (hybrid FTS + vector RAG)** — notes, pitfalls, LoRA blurbs. Shared hits apply to every persona; personal hits overwrite shared on the same kind+key. Never override exact numbers or the user’s param request. For a missing exact row use `memory_get`; for a second search use `memory_search`; for Danbooru spelling/aliases use `lookup_tags` (do not dump tag soup into Krea prompts). 7. Guesses — last resort only. Exact = encyclopedia of defaults. RAG = soft notes. Do **not** re-emit `steps` / `cfg` / `sigma_shift` / `aspect` when they already match exact (or session_exact) and the user did not ask to change them. @@ -25,7 +25,7 @@ A JSON block named "Live SwarmUI context" is attached. Treat it as ground truth - Use only LoRAs listed in `available_loras` / `enabled_loras` (exact `name`), or Civitai search candidates. - Prefer listed `trigger_phrase` / `triggers` — **never invent** trigger words. - `session_exact` / `recommended_params` — session overrides and defaults (Exact KV is in the system block above). -- `memory_hits` are retrieved notes (LoRA tips, pitfalls). Each hit has `scope` (`shared`|`personal`). Trust them over guesses, but **not** over Exact or the user. +- `memory_hits` are retrieved notes (hybrid FTS+vector). Each hit has `scope` (`shared`|`personal`). Trust them over guesses, but **not** over Exact or the user. - `has_vision_image` — if false, do not invent what the image looks like; emit `look_at` first when you need to see it. - `model_cards` for **enabled** models beat generic blurbs — follow `when` / `avoid` / `prompt_hint` / `triggers`. - `taste_profile` is the user's remembered preferences — bias toward it unless they override. @@ -77,6 +77,9 @@ A JSON block named "Live SwarmUI context" is attached. Treat it as ground truth "pack": null, "actions": ["generate"], "search_query": null, + "memory_query": null, + "memory_kind": null, + "tag_query": null, "memories": [{"kind": "lora", "key": "name", "text": "fact", "scope": "personal"}], "notes": "one-line why" } @@ -91,7 +94,7 @@ A JSON block named "Live SwarmUI context" is attached. Treat it as ground truth - `vary: true` — new random seed. `lock_seed: true` — reuse current seed. - `pack` — switch active prompt pack for a follow-up hop. - Do not invent model or LoRA filenames. -- Memory: `actions` may include `memory_upsert` or `memory_forget` with `memories: [{kind,key,text,scope}]`. Default scope is personal (this persona). `"scope":"shared"` is visible to all personas; personal never copies into shared. +- Memory: `memory_upsert` / `memory_forget` with `memories: [{kind,key,text,scope}]`. Default scope is personal. Tools: `memory_get` + kind/key, `memory_search` + `memory_query`, `lookup_tags` + `tag_query` (Danbooru csv — spelling only, not prompt soup). ### Actions (auto-safe) @@ -99,5 +102,7 @@ A JSON block named "Live SwarmUI context" is attached. Treat it as ground truth - `"search_civitai"` — Civitai search; user Confirms downloads. - `"interrupt"` — stop generation. - `"memory_upsert"` / `"memory_forget"` — write or delete vector memory (personal by default; `scope: "shared"` for the common store). +- `"memory_get"` / `"memory_search"` — hop: exact row or hybrid search. +- `"lookup_tags"` — hop: Danbooru csv (aliases/counts). Do not emit tag soup for Krea. - `look_at: ["generate", "ref1"]` — vision hop. - Pure Q&A with no change: omit the JSON patch. diff --git a/Config/_base/skills/memory.md b/Config/_base/skills/memory.md index e69c78c..bd67c2c 100644 --- a/Config/_base/skills/memory.md +++ b/Config/_base/skills/memory.md @@ -1,26 +1,34 @@ # Skill: memory -You have two memory layers: +You have three memory tools: -1. **Exact memory** (`## Exact memory` + live `exact` / `session_exact`) — canonical KV defaults (generation params, aspect table, architecture facts). Always prefer Exact over RAG for numbers and defaults. -2. **Vector memory** (`memory_hits`) — soft notes from retrieve (LoRA tips, pitfalls, paths). Hits are **shared + this persona**. `scope: "personal"` overwrites `scope: "shared"` on the same `kind`+`key`. Other personas never see your personal rows. +1. **Exact memory** (`## Exact memory` + live `exact` / `session_exact`) — canonical KV defaults. Always prefer Exact over RAG for numbers. +2. **Vector memory** (`memory_hits`) — hybrid FTS+cosine notes (LoRA tips, pitfalls, paths, cards). Shared + this persona; personal overwrites shared on the same `kind`+`key`. +3. **Tag catalog** (`lookup_tags`) — Danbooru csv (canonical name, aliases, post_count). **Not** RAG. Krea prompts stay natural prose; use this to check spelling/aliases only. ## Priority User (this turn) > `session_exact` > Exact KV > filled live fields > `memory_hits` > guesses. Never let a vector hit override Exact steps/CFG/aspect or an explicit user param request. +## Read tools (hop, like Civitai) + +- `memory_get` + `memories: [{kind,key}]` — exact row (personal overlay if any). +- `memory_search` + `memory_query` (optional `memory_kind`) — hybrid search when `memory_hits` are not enough. +- `lookup_tags` + `tag_query` — csv lookup. Do **not** paste tag soup into the prompt. + +Omit the tool action on the follow-up turn once you have results. + ## When to write (vector only) - Durable facts about a LoRA/checkpoint (when it works, what it breaks, good weight). - Bad paths / pitfalls you discovered this session. - Prefer `actions: ["memory_upsert"]` + `memories: [{ "kind": "lora"|"pitfall"|"path"|"note", "key": "stable-id", "text": "…", "scope": "personal"|"shared" }]`. -- Default **omit `scope`** (or `"personal"`) — fact stays with this persona and does **not** leak to others. -- Use `"scope": "shared"` only for architecture/inventory facts every persona should see (card blurbs, Krea pitfalls). +- Default **omit `scope`** (or `"personal"`). `"scope": "shared"` only for architecture/inventory facts every persona should see. ## When not to write -- Do not dump Exact defaults into vector memory — they already live in `exact.json`. -- Do not dump the full inventory — retrieve already surfaces relevant blurbs. -- Do not store the user's taste profile (that is `taste_profile` / taste.json). -- Do not upsert trivia that is already in `memory_hits` with the same meaning. -- `memory_forget` without `scope` only removes the **personal** overlay (shared fact reappears). Use `"scope": "shared"` to delete a shared row. +- Do not dump Exact defaults or the full inventory into vector memory. +- Do not store the user's taste profile (`taste.json`). +- Do not upsert trivia already in `memory_hits`. +- Do not upsert Danbooru tags — the csv catalog already has them. +- `memory_forget` without `scope` only removes the personal overlay. diff --git a/README.md b/README.md index 58bc092..d3d5fc2 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ SwarmUI extension for **collaborative Krea 2** prompting via **Ollama**: chat + board (Generate | Refs tabs), LoRA chips, **persona presets** (`Config/personas/`), **vector memory**, model cards with Civitai fetch, img2img/inpaint, slash commands, auto Generate. -**Version 0.8.1** — Chats live on disk (`Assistent/chats/`), the chat model is **parked out of VRAM** before every Generate, memory + wanted queue are editable in ⚙, Ollama health sits in the chat header. Vector memory is shared + personal: personal never leaks into shared; shared is visible to every persona; personal overwrites the same kind+key. +**Version 0.8.2** — Hybrid memory (FTS5 + cosine, kind quotas, min_score) plus model hops `memory_get` / `memory_search` / `lookup_tags`. Danbooru csv is a shared FTS catalog (no embeddings); Krea prompts stay prose. ## Layout @@ -48,7 +48,8 @@ Two layers in `memory/assistent.sqlite` (`persona` column; empty = shared): - **Shared** — `Config/_base/memory-seed/`, model cards, `scope: "shared"` upserts. Visible to every persona. - **Personal** — `Config/personas//memory-seed/` and chat upserts (default). Never copied into shared. Other personas do not retrieve it. -- Retrieve = shared ∪ this persona (and `extends` parents). Same `kind`+`key`: personal overwrites parent overwrites shared. Forget without `scope` only drops the personal overlay. +- Retrieve = shared ∪ this persona (and `extends` parents). Hybrid **FTS5 + cosine**, kind quotas (e.g. 3 cards / 3 pitfalls / 4 notes), `min_score`. Same `kind`+`key`: personal overwrites parent overwrites shared. +- Tools: `memory_get`, `memory_search`, `lookup_tags` (Danbooru csv in `Data/Autocompletions`, FTS, **no embeddings**). - SQLite + Ollama `/api/embed` (default `nomic-embed-text`, pick in ⚙) - Soft notes only — Exact and the user beat RAG for params - ⚙ → **Память** lists every row (scope · source · date) with a per-row forget; bundled rows are read-only because reseed brings them back @@ -133,8 +134,10 @@ Restart / rebuild SwarmUI after clone. gpu-rent: `seed-extensions` + restart. | `AssistentEnqueueWanted` / `AssistentListWanted` | Wanted YAML queue (write / read + count) | | `AssistentGetTaste` / `AssistentSaveTaste` | Persistent taste profile | | `AssistentSearchCivitai` | Civitai LoRA search | -| `AssistentChat` / `AssistentChatWS` | Chat (+ memory retrieve + Civitai hop) | +| `AssistentChat` / `AssistentChatWS` | Chat (+ hybrid memory + Civitai/tag hops) | | `AssistentListMemory` / `AssistentUpsertMemory` / `AssistentForgetMemory` | Vector store (optional `scope` / `persona`) | +| `AssistentSearchMemory` / `AssistentGetMemory` | Hybrid search / exact kind+key | +| `AssistentLookupTags` | Danbooru csv FTS (no embeddings) | | `AssistentListChats` / `AssistentGetChat` / `AssistentSaveChat` / `AssistentDeleteChat` | `Assistent/chats/.json` | | `AssistentGetUiState` / `AssistentSaveUiState` | `Assistent/ui-state.json` | | `AssistentParkLlm` / `AssistentWarmLlm` | Unload / reload the chat model in VRAM | diff --git a/SwarmAssistentExtension.cs b/SwarmAssistentExtension.cs index 6db5004..b514a85 100644 --- a/SwarmAssistentExtension.cs +++ b/SwarmAssistentExtension.cs @@ -37,7 +37,7 @@ public partial class SwarmAssistentExtension : Extension ExtensionAuthor = "mrleo1nid"; Description = "Collaborative Krea 2 assistant: Ollama chat, persona presets, vector memory, model cards, Generate loop."; License = "MIT"; - Version = "0.8.1"; + Version = "0.8.3"; Tags = ["tabs", "ui", "llm", "ollama", "krea", "inpaint", "memory"]; } @@ -73,8 +73,11 @@ public partial class SwarmAssistentExtension : Extension API.RegisterAPICall(AssistentListMemory, false, PermUse); API.RegisterAPICall(AssistentUpsertMemory, true, PermUse); API.RegisterAPICall(AssistentForgetMemory, true, PermUse); + API.RegisterAPICall(AssistentSearchMemory, false, PermUse); + API.RegisterAPICall(AssistentGetMemory, false, PermUse); + API.RegisterAPICall(AssistentLookupTags, false, PermUse); API.RegisterAPICall(AssistentListWanted, false, PermUse); - Logs.Init("Swarm Assistent extension loaded (disk chats + park LLM + memory UI)"); + Logs.Init("Swarm Assistent extension loaded (sqlite chats/kv + park LLM + memory UI)"); } int CfgInt(string key, int fallback) @@ -138,10 +141,6 @@ public partial class SwarmAssistentExtension : Extension return Environment.CurrentDirectory; } - string PersonasOverlayJsonPath() => Path.Combine(DataRoot(), "Assistent", "personas.json"); - - string TasteJsonPath() => Path.Combine(DataRoot(), "Assistent", "taste.json"); - public string ReadPackFile(string name) { return Config?.LoadPackPrompt(Config.DefaultPersonaId(), name); @@ -229,19 +228,13 @@ public partial class SwarmAssistentExtension : Extension public async Task AssistentGetTaste(Session session) { await Task.CompletedTask; - string path = TasteJsonPath(); - if (!File.Exists(path)) - { - return new JObject { ["success"] = true, ["taste"] = null }; - } try { - JObject taste = JObject.Parse(File.ReadAllText(path, Encoding.UTF8)); - return new JObject { ["success"] = true, ["taste"] = taste }; + return new JObject { ["success"] = true, ["taste"] = Memory.GetKvObject(AssistentMemory.KvTaste) }; } catch (Exception ex) { - return new JObject { ["error"] = $"taste.json: {ex.Message}" }; + return new JObject { ["error"] = $"taste: {ex.Message}" }; } } @@ -256,10 +249,14 @@ public partial class SwarmAssistentExtension : Extension { taste["updated"] = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(); } - string dir = Path.Combine(DataRoot(), "Assistent"); - Directory.CreateDirectory(dir); - string path = TasteJsonPath(); - File.WriteAllText(path, taste.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8); - return new JObject { ["success"] = true, ["path"] = path }; + try + { + Memory.SetKvObject(AssistentMemory.KvTaste, taste); + return new JObject { ["success"] = true, ["path"] = "Assistent/memory/assistent.sqlite" }; + } + catch (Exception ex) + { + return new JObject { ["error"] = $"taste save: {ex.Message}" }; + } } }