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. /// 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"; 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(); } 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}"); } } 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(); } }