using System; using System.Collections.Generic; using System.IO; using System.Linq; using System.Text; using Microsoft.Data.Sqlite; using Newtonsoft.Json.Linq; using SwarmUI.Utils; namespace Mrleo1nid.SwarmAssistent; /// Runtime store in the same sqlite file: chats, ui-state, taste. /// Config overlays, sidecar cards, and ollama-roles stay on disk. public sealed partial class AssistentMemory { public const int MaxChatsStored = 200; public const string KvUiState = "ui_state"; public const string KvTaste = "taste"; void EnsureStoreSchema() { Exec( """ CREATE TABLE IF NOT EXISTS kv ( key TEXT PRIMARY KEY, value TEXT NOT NULL, updated INTEGER NOT NULL ); CREATE TABLE IF NOT EXISTS chats ( id TEXT PRIMARY KEY, title TEXT NOT NULL DEFAULT '', created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, params_json TEXT, messages_json TEXT NOT NULL DEFAULT '[]', messages_count INTEGER NOT NULL DEFAULT 0, body TEXT NOT NULL DEFAULT '' ); CREATE INDEX IF NOT EXISTS idx_chats_updated ON chats(updated_at DESC); """); EnsureChatsFts(); MigrateJsonStoreOnce(); } void EnsureChatsFts() { try { Exec( """ CREATE VIRTUAL TABLE IF NOT EXISTS chats_fts USING fts5( title, body, tokenize = 'unicode61 remove_diacritics 2' ); """); Exec( """ CREATE TRIGGER IF NOT EXISTS chats_fts_ai AFTER INSERT ON chats BEGIN INSERT INTO chats_fts(rowid, title, body) VALUES (new.rowid, new.title, new.body); END; """); Exec( """ CREATE TRIGGER IF NOT EXISTS chats_fts_ad AFTER DELETE ON chats BEGIN INSERT INTO chats_fts(chats_fts, rowid) VALUES('delete', old.rowid); END; """); Exec( """ CREATE TRIGGER IF NOT EXISTS chats_fts_au AFTER UPDATE ON chats BEGIN INSERT INTO chats_fts(chats_fts, rowid) VALUES('delete', old.rowid); INSERT INTO chats_fts(rowid, title, body) VALUES (new.rowid, new.title, new.body); END; """); int n = 0, fts = 0; using (SqliteCommand c = _conn.CreateCommand()) { c.CommandText = "SELECT COUNT(*) FROM chats"; n = Convert.ToInt32(c.ExecuteScalar()); } using (SqliteCommand c = _conn.CreateCommand()) { c.CommandText = "SELECT COUNT(*) FROM chats_fts"; fts = Convert.ToInt32(c.ExecuteScalar()); } if (n != fts) { Exec("DELETE FROM chats_fts"); Exec("INSERT INTO chats_fts(rowid, title, body) SELECT rowid, title, body FROM chats"); } } catch (Exception ex) { Logs.Debug($"AssistentMemory chats FTS5 unavailable: {ex.Message}"); } } void MigrateJsonStoreOnce() { if (GetMeta("json_store_migrated") == "1") { return; } string root = Path.Combine(_dataRoot, "Assistent"); int chats = 0; string chatsDir = Path.Combine(root, "chats"); if (Directory.Exists(chatsDir)) { foreach (string file in Directory.EnumerateFiles(chatsDir, "*.json")) { try { JObject chat = JObject.Parse(File.ReadAllText(file, Encoding.UTF8)); string id = (chat["id"]?.ToString() ?? Path.GetFileNameWithoutExtension(file) ?? "").Trim(); if (string.IsNullOrWhiteSpace(id) || GetChatUnlocked(id) is not null) { continue; } UpsertChatUnlocked(chat, id); chats++; } catch (Exception ex) { Logs.Debug($"AssistentMemory migrate chat {file}: {ex.Message}"); } } } ImportKvFile(Path.Combine(root, "ui-state.json"), KvUiState); ImportKvFile(Path.Combine(root, "taste.json"), KvTaste); SetMeta("json_store_migrated", "1"); TryArchiveMigratedJson(root, chatsDir); if (chats > 0) { Logs.Debug($"AssistentMemory: migrated {chats} chats from JSON into sqlite"); } } void ImportKvFile(string path, string key) { if (!File.Exists(path) || !string.IsNullOrEmpty(GetKvUnlocked(key))) { return; } try { SetKvUnlocked(key, File.ReadAllText(path, Encoding.UTF8)); } catch (Exception ex) { Logs.Debug($"AssistentMemory migrate {key}: {ex.Message}"); } } void TryArchiveMigratedJson(string root, string chatsDir) { try { string dest = Path.Combine(root, "_migrated_json"); Directory.CreateDirectory(dest); MoveIfExists(Path.Combine(root, "ui-state.json"), Path.Combine(dest, "ui-state.json")); MoveIfExists(Path.Combine(root, "taste.json"), Path.Combine(dest, "taste.json")); if (!Directory.Exists(chatsDir)) { return; } string chatsDest = Path.Combine(dest, "chats"); Directory.CreateDirectory(chatsDest); foreach (string file in Directory.EnumerateFiles(chatsDir, "*.json")) { MoveIfExists(file, Path.Combine(chatsDest, Path.GetFileName(file))); } } catch (Exception ex) { Logs.Debug($"AssistentMemory archive json: {ex.Message}"); } } static void MoveIfExists(string src, string dest) { if (!File.Exists(src)) { return; } if (File.Exists(dest)) { File.Delete(src); return; } Directory.CreateDirectory(Path.GetDirectoryName(dest)!); File.Move(src, dest); } public string GetKv(string key) { lock (_lock) { EnsureOpen(); return GetKvUnlocked(key); } } public JObject GetKvObject(string key) { string raw = GetKv(key); if (string.IsNullOrWhiteSpace(raw)) { return null; } try { return JObject.Parse(raw); } catch (Exception ex) { Logs.Debug($"AssistentMemory kv {key}: {ex.Message}"); return null; } } public void SetKvObject(string key, JObject value) { SetKv(key, value?.ToString(Newtonsoft.Json.Formatting.None) ?? ""); } public void SetKv(string key, string value) { lock (_lock) { EnsureOpen(); SetKvUnlocked(key, value); } } string GetKvUnlocked(string key) { using SqliteCommand cmd = _conn.CreateCommand(); cmd.CommandText = "SELECT value FROM kv WHERE key = $k"; cmd.Parameters.AddWithValue("$k", key ?? ""); return cmd.ExecuteScalar()?.ToString(); } void SetKvUnlocked(string key, string value) { using SqliteCommand cmd = _conn.CreateCommand(); cmd.CommandText = """ INSERT INTO kv(key, value, updated) VALUES($k, $v, $u) ON CONFLICT(key) DO UPDATE SET value = excluded.value, updated = excluded.updated """; cmd.Parameters.AddWithValue("$k", key ?? ""); cmd.Parameters.AddWithValue("$v", value ?? ""); cmd.Parameters.AddWithValue("$u", DateTimeOffset.UtcNow.ToUnixTimeMilliseconds()); cmd.ExecuteNonQuery(); } public List ListChats(bool withMessages, int limit, string query = null) { lock (_lock) { EnsureOpen(); int take = Math.Clamp(limit, 1, MaxChatsStored); List list = []; string match = BuildFtsMatch(query); bool searched = false; if (!string.IsNullOrWhiteSpace(match) && TableExists("chats_fts")) { try { using SqliteCommand cmd = _conn.CreateCommand(); cmd.CommandText = """ SELECT id, title, created_at, updated_at, params_json, messages_json, messages_count FROM chats WHERE rowid IN (SELECT rowid FROM chats_fts WHERE chats_fts MATCH $q) ORDER BY updated_at DESC LIMIT $lim """; cmd.Parameters.AddWithValue("$q", match); cmd.Parameters.AddWithValue("$lim", take); ReadChats(cmd, withMessages, list); searched = true; } catch (Exception ex) { Logs.Debug($"AssistentMemory chats FTS: {ex.Message}"); } } if (!searched && !string.IsNullOrWhiteSpace(query)) { using SqliteCommand cmd = _conn.CreateCommand(); cmd.CommandText = """ SELECT id, title, created_at, updated_at, params_json, messages_json, messages_count FROM chats WHERE title LIKE $like ESCAPE '\' OR body LIKE $like ESCAPE '\' ORDER BY updated_at DESC LIMIT $lim """; cmd.Parameters.AddWithValue("$like", "%" + EscapeLike(query.Trim()) + "%"); cmd.Parameters.AddWithValue("$lim", take); ReadChats(cmd, withMessages, list); searched = true; } if (!searched) { using SqliteCommand cmd = _conn.CreateCommand(); cmd.CommandText = """ SELECT id, title, created_at, updated_at, params_json, messages_json, messages_count FROM chats ORDER BY updated_at DESC LIMIT $lim """; cmd.Parameters.AddWithValue("$lim", take); ReadChats(cmd, withMessages, list); } return list; } } public int CountChats() { lock (_lock) { EnsureOpen(); using SqliteCommand cmd = _conn.CreateCommand(); cmd.CommandText = "SELECT COUNT(*) FROM chats"; return Convert.ToInt32(cmd.ExecuteScalar()); } } public JObject GetChat(string id) { lock (_lock) { EnsureOpen(); return GetChatUnlocked(id); } } public void SaveChat(JObject chat) { lock (_lock) { EnsureOpen(); string id = chat?["id"]?.ToString()?.Trim(); if (string.IsNullOrWhiteSpace(id)) { throw new ArgumentException("chat id required"); } UpsertChatUnlocked(chat, id); PruneChatsUnlocked(); } } public bool DeleteChat(string id) { lock (_lock) { EnsureOpen(); using SqliteCommand cmd = _conn.CreateCommand(); cmd.CommandText = "DELETE FROM chats WHERE id = $id"; cmd.Parameters.AddWithValue("$id", (id ?? "").Trim()); return cmd.ExecuteNonQuery() > 0; } } JObject GetChatUnlocked(string id) { using SqliteCommand cmd = _conn.CreateCommand(); cmd.CommandText = """ SELECT id, title, created_at, updated_at, params_json, messages_json, messages_count FROM chats WHERE id = $id LIMIT 1 """; cmd.Parameters.AddWithValue("$id", (id ?? "").Trim()); List list = []; ReadChats(cmd, withMessages: true, list); return list.Count > 0 ? list[0] : null; } void UpsertChatUnlocked(JObject chat, string id) { JArray messages = chat["messages"] as JArray ?? []; string title = chat["title"]?.ToString() ?? "Новый чат"; long now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(); long created = chat["createdAt"]?.Value() ?? now; long updated = chat["updatedAt"]?.Value() ?? now; string paramsJson = chat["params"] is JObject p ? p.ToString(Newtonsoft.Json.Formatting.None) : null; string messagesJson = messages.ToString(Newtonsoft.Json.Formatting.None); using SqliteCommand cmd = _conn.CreateCommand(); cmd.CommandText = """ INSERT INTO chats(id, title, created_at, updated_at, params_json, messages_json, messages_count, body) VALUES($id, $title, $created, $updated, $params, $messages, $count, $body) ON CONFLICT(id) DO UPDATE SET title = excluded.title, created_at = excluded.created_at, updated_at = excluded.updated_at, params_json = excluded.params_json, messages_json = excluded.messages_json, messages_count = excluded.messages_count, body = excluded.body """; cmd.Parameters.AddWithValue("$id", id); cmd.Parameters.AddWithValue("$title", title); cmd.Parameters.AddWithValue("$created", created); cmd.Parameters.AddWithValue("$updated", updated); cmd.Parameters.AddWithValue("$params", (object)paramsJson ?? DBNull.Value); cmd.Parameters.AddWithValue("$messages", messagesJson); cmd.Parameters.AddWithValue("$count", messages.Count); cmd.Parameters.AddWithValue("$body", ChatFtsBody(title, messages)); cmd.ExecuteNonQuery(); } void PruneChatsUnlocked() { using SqliteCommand cmd = _conn.CreateCommand(); cmd.CommandText = """ DELETE FROM chats WHERE rowid IN ( SELECT rowid FROM chats ORDER BY updated_at DESC LIMIT -1 OFFSET $keep ) """; cmd.Parameters.AddWithValue("$keep", MaxChatsStored); cmd.ExecuteNonQuery(); } static void ReadChats(SqliteCommand cmd, bool withMessages, List list) { using SqliteDataReader reader = cmd.ExecuteReader(); while (reader.Read()) { list.Add(ReadChatRow(reader, withMessages)); } } static JObject ReadChatRow(SqliteDataReader reader, bool withMessages) { JObject chat = new() { ["id"] = reader.GetString(0), ["title"] = reader.IsDBNull(1) ? "Новый чат" : reader.GetString(1), ["createdAt"] = reader.GetInt64(2), ["updatedAt"] = reader.GetInt64(3), ["messages_count"] = reader.GetInt32(6), }; if (!reader.IsDBNull(4)) { try { chat["params"] = JObject.Parse(reader.GetString(4)); } catch { chat["params"] = null; } } if (withMessages) { try { chat["messages"] = reader.IsDBNull(5) ? new JArray() : JArray.Parse(reader.GetString(5)); } catch { chat["messages"] = new JArray(); } } return chat; } static string ChatFtsBody(string title, JArray messages) { StringBuilder sb = new(); sb.Append(title).Append('\n'); foreach (JToken token in messages ?? []) { if (token is JObject mo) { sb.Append(mo["content"]?.ToString()).Append('\n'); } } return sb.ToString(); } }