swarm-assistent/AssistentMemory.Store.cs

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;

/// <summary>Runtime store in the same sqlite file: chats, ui-state.
/// Config overlays, sidecar cards, and ollama-roles stay on disk.</summary>
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<JObject> ListChats(bool withMessages, int limit, string query = null)
    {
        lock (_lock)
        {
            EnsureOpen();
            int take = Math.Clamp(limit, 1, MaxChatsStored);
            List<JObject> 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<JObject> 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<long?>() ?? now;
        long updated = chat["updatedAt"]?.Value<long?>() ?? 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<JObject> 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();
    }
}