swarm-assistent/AssistentMemory.cs

using System;
using System.Collections.Generic;
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;
using SwarmUI.Utils;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>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.</summary>
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<string, int> Quotas { get; set; }
        public bool ApplyQuotas { get; set; } = true;
    }

    static readonly Dictionary<string, int> DefaultQuotas = new(StringComparer.OrdinalIgnoreCase)
    {
        ["card"] = 3,
        ["lora"] = 3,
        ["pitfall"] = 3,
        ["path"] = 2,
        ["note"] = 4,
        ["model"] = 2,
        ["aspect"] = 1,
        ["heard"] = 3,
    };

    public static Dictionary<string, int> CopyDefaultQuotas()
        => new(DefaultQuotas, StringComparer.OrdinalIgnoreCase);

    readonly string _dataRoot;
    readonly string _dbPath;
    readonly HttpClient _http;
    readonly object _lock = new();
    SqliteConnection _conn;
    string _embedModel;
    int _dims;
    int _seedVersion;

    public AssistentMemory(string dataRoot, HttpClient http, string defaultEmbedModel = "nomic-embed-text")
    {
        _dataRoot = string.IsNullOrWhiteSpace(dataRoot) ? "." : dataRoot;
        string dir = Path.Combine(_dataRoot, "Assistent", "memory");
        Directory.CreateDirectory(dir);
        _dbPath = Path.Combine(dir, "assistent.sqlite");
        _http = http;
        _embedModel = string.IsNullOrWhiteSpace(defaultEmbedModel) ? "nomic-embed-text" : defaultEmbedModel.Trim();
    }

    public string EmbedModel => _embedModel;
    public int Dims => _dims;
    public int SeedVersion => _seedVersion;

    public static string NormalizePersona(string raw)
    {
        if (string.IsNullOrWhiteSpace(raw))
        {
            return SharedPersona;
        }
        string s = raw.Trim();
        if (s is "_" or "*" or "shared" or "common" or "global" or "_shared")
        {
            return SharedPersona;
        }
        return AssistentConfig.SafeId(s) ?? SharedPersona;
    }

    public static bool IsShared(string persona) => string.IsNullOrEmpty(NormalizePersona(persona));

    void EnsureOpen()
    {
        if (_conn is not null)
        {
            return;
        }
        AssistentSqliteBootstrap.EnsureInitialized();
        _conn = new SqliteConnection($"Data Source={_dbPath}");
        _conn.Open();
        TryPragma("journal_mode=WAL");
        TryPragma("busy_timeout=5000");
        using (SqliteCommand cmd = _conn.CreateCommand())
        {
            cmd.CommandText =
                """
                CREATE TABLE IF NOT EXISTS meta (
                  key TEXT PRIMARY KEY,
                  value TEXT NOT NULL
                );
                CREATE TABLE IF NOT EXISTS memories (
                  id INTEGER PRIMARY KEY AUTOINCREMENT,
                  kind TEXT NOT NULL,
                  key TEXT NOT NULL,
                  persona TEXT NOT NULL DEFAULT '',
                  text TEXT NOT NULL,
                  source TEXT NOT NULL DEFAULT 'user',
                  meta_json TEXT,
                  embedding BLOB,
                  updated INTEGER NOT NULL,
                  UNIQUE(kind, key, source, persona)
                );
                """;
            cmd.ExecuteNonQuery();
        }
        MigratePersonaColumn();
        using (SqliteCommand idx = _conn.CreateCommand())
        {
            idx.CommandText =
                """
                CREATE INDEX IF NOT EXISTS idx_memories_kind ON memories(kind);
                CREATE INDEX IF NOT EXISTS idx_memories_persona ON memories(persona);
                """;
            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}");
        }
        try
        {
            EnsureTrainingSchema();
        }
        catch (Exception ex)
        {
            Logs.Debug($"AssistentMemory training schema: {ex.Message}");
        }
        try
        {
            EnsureUserPrefsSchema();
        }
        catch (Exception ex)
        {
            Logs.Debug($"AssistentMemory user_prefs schema: {ex.Message}");
        }
        _embedModel = GetMeta("embed_model") ?? _embedModel;
        _ = int.TryParse(GetMeta("dims"), out _dims);
        _ = int.TryParse(GetMeta("seed_version"), out _seedVersion);
    }

    bool HasColumn(string table, string column)
    {
        using SqliteCommand cmd = _conn.CreateCommand();
        cmd.CommandText = $"PRAGMA table_info({table})";
        using SqliteDataReader reader = cmd.ExecuteReader();
        while (reader.Read())
        {
            if (string.Equals(reader.GetString(1), column, StringComparison.OrdinalIgnoreCase))
            {
                return true;
            }
        }
        return false;
    }

    void MigratePersonaColumn()
    {
        if (HasColumn("memories", "persona"))
        {
            return;
        }
        using SqliteTransaction tx = _conn.BeginTransaction();
        using (SqliteCommand cmd = _conn.CreateCommand())
        {
            cmd.Transaction = tx;
            cmd.CommandText =
                """
                CREATE TABLE memories_v2 (
                  id INTEGER PRIMARY KEY AUTOINCREMENT,
                  kind TEXT NOT NULL,
                  key TEXT NOT NULL,
                  persona TEXT NOT NULL DEFAULT '',
                  text TEXT NOT NULL,
                  source TEXT NOT NULL DEFAULT 'user',
                  meta_json TEXT,
                  embedding BLOB,
                  updated INTEGER NOT NULL,
                  UNIQUE(kind, key, source, persona)
                );
                INSERT INTO memories_v2 (kind, key, persona, text, source, meta_json, embedding, updated)
                  SELECT kind, key, '', text, source, meta_json, embedding, updated FROM memories;
                DROP TABLE memories;
                ALTER TABLE memories_v2 RENAME TO memories;
                CREATE INDEX IF NOT EXISTS idx_memories_kind ON memories(kind);
                CREATE INDEX IF NOT EXISTS idx_memories_persona ON memories(persona);
                """;
            cmd.ExecuteNonQuery();
        }
        tx.Commit();
        SetMeta("schema_version", "2");
        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.ExecuteScalar();
        }
        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();
        cmd.CommandText = "SELECT value FROM meta WHERE key = $k";
        cmd.Parameters.AddWithValue("$k", key);
        return cmd.ExecuteScalar()?.ToString();
    }

    void SetMeta(string key, string value)
    {
        using SqliteCommand cmd = _conn.CreateCommand();
        cmd.CommandText = "INSERT INTO meta(key, value) VALUES($k, $v) ON CONFLICT(key) DO UPDATE SET value = excluded.value";
        cmd.Parameters.AddWithValue("$k", key);
        cmd.Parameters.AddWithValue("$v", value ?? "");
        cmd.ExecuteNonQuery();
    }

    static byte[] FloatsToBytes(float[] v)
    {
        byte[] bytes = new byte[v.Length * 4];
        Buffer.BlockCopy(v, 0, bytes, 0, bytes.Length);
        return bytes;
    }

    static float[] BytesToFloats(byte[] bytes)
    {
        if (bytes is null || bytes.Length < 4 || bytes.Length % 4 != 0)
        {
            return Array.Empty<float>();
        }
        float[] v = new float[bytes.Length / 4];
        Buffer.BlockCopy(bytes, 0, v, 0, bytes.Length);
        return v;
    }

    static float Cosine(float[] a, float[] b)
    {
        if (a.Length == 0 || a.Length != b.Length)
        {
            return float.NegativeInfinity;
        }
        double dot = 0, na = 0, nb = 0;
        for (int i = 0; i < a.Length; i++)
        {
            dot += a[i] * b[i];
            na += a[i] * a[i];
            nb += b[i] * b[i];
        }
        if (na <= 0 || nb <= 0)
        {
            return float.NegativeInfinity;
        }
        return (float)(dot / (Math.Sqrt(na) * Math.Sqrt(nb)));
    }

    static int SourceRank(string source)
    {
        return string.Equals(source, "user", StringComparison.OrdinalIgnoreCase) ? 1 : 0;
    }

    public async Task<float[]> EmbedAsync(string baseUrl, string model, string text, string keepAlive = "60m")
    {
        string root = (baseUrl ?? "http://127.0.0.1:11434").TrimEnd('/');
        string m = string.IsNullOrWhiteSpace(model) ? _embedModel : model.Trim();
        JObject payload = new()
        {
            ["model"] = m,
            ["input"] = text ?? "",
            ["keep_alive"] = keepAlive,
        };
        using StringContent content = new(payload.ToString(Newtonsoft.Json.Formatting.None), Encoding.UTF8, "application/json");
        using HttpResponseMessage resp = await _http.PostAsync($"{root}/api/embed", content);
        string body = await resp.Content.ReadAsStringAsync();
        if (!resp.IsSuccessStatusCode)
        {
            throw new Exception($"Ollama /api/embed HTTP {(int)resp.StatusCode}: {body[..Math.Min(body.Length, 200)]}");
        }
        JObject parsed = JObject.Parse(body);
        JArray embeddings = parsed["embeddings"] as JArray;
        JToken first = embeddings?.FirstOrDefault() ?? parsed["embedding"];
        if (first is not JArray vec)
        {
            throw new Exception("Ollama /api/embed: no embeddings in response");
        }
        float[] floats = vec.Select(t => t.Value<float>()).ToArray();
        return floats;
    }

    bool BundledExists(string kind, string key, string persona)
    {
        using SqliteCommand cmd = _conn.CreateCommand();
        cmd.CommandText = "SELECT 1 FROM memories WHERE source = 'bundled' AND kind = $kind AND key = $key AND persona = $persona LIMIT 1";
        cmd.Parameters.AddWithValue("$kind", kind);
        cmd.Parameters.AddWithValue("$key", key);
        cmd.Parameters.AddWithValue("$persona", persona ?? SharedPersona);
        return cmd.ExecuteScalar() is not null;
    }

    public async Task EnsureSeedAsync(string baseUrl, AssistentConfig config, string modelOverride = null)
    {
        lock (_lock)
        {
            EnsureOpen();
        }
        JObject assistant = config.LoadAssistant(config.DefaultPersonaId());
        int wantVersion = assistant["seed_version"]?.Value<int?>() ?? 1;
        string wantModel = string.IsNullOrWhiteSpace(modelOverride)
            ? (assistant["embed_model"]?.ToString() ?? _embedModel)
            : modelOverride.Trim();

        bool needReseed = _seedVersion != wantVersion || !string.Equals(_embedModel, wantModel, StringComparison.OrdinalIgnoreCase);
        List<JObject> docs = config.LoadMemorySeedDocs();
        if (docs.Count == 0)
        {
            TryIndexTags();
            TryIndexExamples();
            return;
        }

        if (!needReseed)
        {
            bool missing;
            lock (_lock)
            {
                EnsureOpen();
                missing = docs.Any(d =>
                {
                    string kind = (d["kind"]?.ToString() ?? "note").Trim().ToLowerInvariant();
                    string key = (d["key"]?.ToString() ?? "").Trim();
                    string persona = NormalizePersona(d["persona"]?.ToString());
                    return !string.IsNullOrWhiteSpace(key) && !BundledExists(kind, key, persona);
                });
            }
            if (!missing)
            {
                TryIndexTags();
                TryIndexExamples();
                return;
            }
        }

        float[] probe;
        try
        {
            probe = await EmbedAsync(baseUrl, wantModel, docs[0]["text"]?.ToString() ?? "seed");
        }
        catch (Exception ex)
        {
            Logs.Debug($"AssistentMemory seed defer (embed unavailable): {ex.Message}");
            TryIndexTags();
            TryIndexExamples();
            return;
        }

        lock (_lock)
        {
            EnsureOpen();
            if (!string.Equals(_embedModel, wantModel, StringComparison.OrdinalIgnoreCase) || (_dims > 0 && _dims != probe.Length))
            {
                using SqliteCommand clear = _conn.CreateCommand();
                clear.CommandText = "DELETE FROM memories";
                clear.ExecuteNonQuery();
            }
            else if (needReseed)
            {
                using SqliteCommand clearBundled = _conn.CreateCommand();
                clearBundled.CommandText = "DELETE FROM memories WHERE source = 'bundled'";
                clearBundled.ExecuteNonQuery();
            }
            _embedModel = wantModel;
            _dims = probe.Length;
            _seedVersion = wantVersion;
            SetMeta("embed_model", _embedModel);
            SetMeta("dims", _dims.ToString());
            SetMeta("seed_version", _seedVersion.ToString());
        }

        foreach (JObject doc in docs)
        {
            string kind = (doc["kind"]?.ToString() ?? "note").Trim();
            string key = (doc["key"]?.ToString() ?? "").Trim();
            string text = (doc["text"]?.ToString() ?? "").Trim();
            string persona = NormalizePersona(doc["persona"]?.ToString());
            if (string.IsNullOrWhiteSpace(key) || string.IsNullOrWhiteSpace(text))
            {
                continue;
            }
            if (!needReseed)
            {
                lock (_lock)
                {
                    EnsureOpen();
                    if (BundledExists(kind.ToLowerInvariant(), key, persona))
                    {
                        continue;
                    }
                }
            }
            try
            {
                float[] vec = await EmbedAsync(baseUrl, wantModel, text);
                Upsert(kind, key, text, "bundled", doc["tags"], vec, persona);
            }
            catch (Exception ex)
            {
                Logs.Debug($"AssistentMemory seed item {kind}/{key}: {ex.Message}");
            }
        }
        TryIndexTags();
        TryIndexExamples();
    }

    public void Upsert(string kind, string key, string text, string source, JToken meta, float[] embedding, string persona = null)
    {
        lock (_lock)
        {
            EnsureOpen();
            kind = (kind ?? "note").Trim().ToLowerInvariant();
            key = (key ?? "").Trim();
            text = (text ?? "").Trim();
            source = string.IsNullOrWhiteSpace(source) ? "user" : source.Trim();
            persona = NormalizePersona(persona);
            if (string.IsNullOrWhiteSpace(key) || string.IsNullOrWhiteSpace(text))
            {
                return;
            }
            if (embedding is { Length: > 0 })
            {
                if (_dims <= 0)
                {
                    _dims = embedding.Length;
                    SetMeta("dims", _dims.ToString());
                }
            }
            using SqliteCommand cmd = _conn.CreateCommand();
            cmd.CommandText =
                """
                INSERT INTO memories(kind, key, persona, text, source, meta_json, embedding, updated)
                VALUES($kind, $key, $persona, $text, $source, $meta, $emb, $upd)
                ON CONFLICT(kind, key, source, persona) DO UPDATE SET
                  text = excluded.text,
                  meta_json = excluded.meta_json,
                  embedding = excluded.embedding,
                  updated = excluded.updated
                """;
            cmd.Parameters.AddWithValue("$kind", kind);
            cmd.Parameters.AddWithValue("$key", key);
            cmd.Parameters.AddWithValue("$persona", persona);
            cmd.Parameters.AddWithValue("$text", text);
            cmd.Parameters.AddWithValue("$source", source);
            cmd.Parameters.AddWithValue("$meta", meta?.ToString(Newtonsoft.Json.Formatting.None) ?? "");
            cmd.Parameters.AddWithValue("$emb", embedding is null ? (object)DBNull.Value : FloatsToBytes(embedding));
            cmd.Parameters.AddWithValue("$upd", DateTimeOffset.UtcNow.ToUnixTimeSeconds());
            cmd.ExecuteNonQuery();
        }
    }

    /// <summary>Delete a row in one layer. Default (persona=current) never touches shared.
    /// Pass SharedPersona to forget a shared fact. Bundled rows are kept unless source is set.</summary>
    public void Forget(string kind, string key, string source = null, string persona = null)
    {
        lock (_lock)
        {
            EnsureOpen();
            persona = NormalizePersona(persona);
            using SqliteCommand cmd = _conn.CreateCommand();
            if (string.IsNullOrWhiteSpace(source))
            {
                cmd.CommandText = "DELETE FROM memories WHERE kind = $kind AND key = $key AND persona = $persona AND source != 'bundled'";
            }
            else
            {
                cmd.CommandText = "DELETE FROM memories WHERE kind = $kind AND key = $key AND persona = $persona AND source = $source";
                cmd.Parameters.AddWithValue("$source", source);
            }
            cmd.Parameters.AddWithValue("$kind", (kind ?? "").Trim().ToLowerInvariant());
            cmd.Parameters.AddWithValue("$key", (key ?? "").Trim());
            cmd.Parameters.AddWithValue("$persona", persona);
            cmd.ExecuteNonQuery();
        }
    }

    static Dictionary<string, int> PersonaRankMap(IEnumerable<string> personaChain)
    {
        Dictionary<string, int> rank = new(StringComparer.OrdinalIgnoreCase)
        {
            [SharedPersona] = 0,
        };
        int i = 1;
        foreach (string id in personaChain ?? [])
        {
            string p = NormalizePersona(id);
            if (p == SharedPersona)
            {
                continue;
            }
            rank[p] = i++;
        }
        return rank;
    }

    static string BuildFtsMatch(string query)
    {
        if (string.IsNullOrWhiteSpace(query))
        {
            return null;
        }
        HashSet<string> 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<long, int> FtsRowRanks(string match, int limit)
    {
        Dictionary<long, int> 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<JObject> ordered, int topK, IReadOnlyDictionary<string, int> quotas)
    {
        Dictionary<string, int> used = new(StringComparer.OrdinalIgnoreCase);
        List<JObject> picked = [];
        List<JObject> 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);
    }

    /// <summary>Retrieve shared + the given persona chain. Hybrid FTS + cosine; personal overwrites shared on kind+key.</summary>
    public async Task<JArray> RetrieveAsync(string baseUrl, string query, int topK = 10, string modelOverride = null, IEnumerable<string> 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<string, int> 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<long, int> ftsRanks = FtsRowRanks(ftsMatch, Math.Max(40, options.TopK * 4));
            using SqliteCommand cmd = _conn.CreateCommand();
            List<string> personaKeys = rank.Keys.ToList();
            StringBuilder sql = new("SELECT id, kind, key, text, source, embedding, persona FROM memories WHERE persona IN (");
            for (int i = 0; i < personaKeys.Count; i++)
            {
                if (i > 0)
                {
                    sql.Append(',');
                }
                string pName = "$p" + i;
                sql.Append(pName);
                cmd.Parameters.AddWithValue(pName, personaKeys[i]);
            }
            sql.Append(')');
            if (kindFilter is not null)
            {
                sql.Append(" AND kind = $kind");
                cmd.Parameters.AddWithValue("$kind", kindFilter);
            }
            cmd.CommandText = sql.ToString();
            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;
                }
                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))
                {
                    cosine = Cosine(qvec, BytesToFloats((byte[])reader.GetValue(5)));
                    if (float.IsNegativeInfinity(cosine))
                    {
                        cosine = 0;
                    }
                }
                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((hybrid, cosine, personaRank, SourceRank(source), ftsHit, new JObject
                {
                    ["kind"] = kind,
                    ["key"] = key,
                    ["text"] = text,
                    ["source"] = source,
                    ["scope"] = shared ? "shared" : "personal",
                    ["persona"] = shared ? "shared" : persona,
                    ["score"] = Math.Round(hybrid, 4),
                    ["cosine"] = Math.Round(cosine, 4),
                    ["fts"] = ftsHit,
                }));
            }
        }

        Dictionary<string, (float hybrid, float cosine, int personaRank, int sourceRank, bool fts, JObject row)> best = new(StringComparer.OrdinalIgnoreCase);
        foreach (var item in scored)
        {
            string id = $"{item.row["kind"]}\n{item.row["key"]}";
            if (best.TryGetValue(id, out var cur))
            {
                if (item.personaRank < cur.personaRank
                    || (item.personaRank == cur.personaRank && item.sourceRank < cur.sourceRank)
                    || (item.personaRank == cur.personaRank && item.sourceRank == cur.sourceRank && item.hybrid <= cur.hybrid))
                {
                    continue;
                }
            }
            best[id] = item;
        }

        float min = options.MinScore;
        IEnumerable<JObject> 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));
    }

    /// <summary>Exact kind+key read with the same personal-over-shared overlay as retrieve.</summary>
    public JObject Get(string kind, string key, IEnumerable<string> personaChain = null)
    {
        kind = (kind ?? "note").Trim().ToLowerInvariant();
        key = (key ?? "").Trim();
        if (string.IsNullOrWhiteSpace(key))
        {
            return null;
        }
        Dictionary<string, int> 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, meta_json 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),
                    ["meta_json"] = reader.IsDBNull(6) ? null : reader.GetString(6),
                };
            }
        }
        return best;
    }

    public async Task<JArray> SearchAsync(string baseUrl, string query, string kind, int topK, string modelOverride, IEnumerable<string> 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)
    {
        lock (_lock)
        {
            EnsureOpen();
            List<JObject> rows = [];
            using SqliteCommand cmd = _conn.CreateCommand();
            cmd.CommandText = "SELECT kind, key, text, source, persona, updated FROM memories ORDER BY updated DESC LIMIT $lim";
            cmd.Parameters.AddWithValue("$lim", Math.Clamp(limit, 1, 2000));
            using SqliteDataReader reader = cmd.ExecuteReader();
            while (reader.Read())
            {
                string persona = reader.IsDBNull(4) ? SharedPersona : reader.GetString(4) ?? SharedPersona;
                bool shared = string.IsNullOrEmpty(persona);
                rows.Add(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 new JArray(rows);
        }
    }

    public int CountAll()
    {
        lock (_lock)
        {
            EnsureOpen();
            using SqliteCommand cmd = _conn.CreateCommand();
            cmd.CommandText = "SELECT COUNT(*) FROM memories";
            return Convert.ToInt32(cmd.ExecuteScalar());
        }
    }

    public async Task UpsertTextAsync(string baseUrl, string kind, string key, string text, string source = "user", JToken meta = null, string modelOverride = null, string persona = null)
    {
        string model = string.IsNullOrWhiteSpace(modelOverride) ? _embedModel : modelOverride;
        float[] vec = await EmbedAsync(baseUrl, model, text);
        Upsert(kind, key, text, source, meta, vec, persona);
    }

    public async Task ReembedAllAsync(string baseUrl, string newModel)
    {
        List<(string kind, string key, string text, string source, string meta, string persona)> rows = [];
        lock (_lock)
        {
            EnsureOpen();
            using SqliteCommand cmd = _conn.CreateCommand();
            cmd.CommandText = "SELECT kind, key, text, source, meta_json, persona FROM memories";
            using SqliteDataReader reader = cmd.ExecuteReader();
            while (reader.Read())
            {
                rows.Add((
                    reader.GetString(0),
                    reader.GetString(1),
                    reader.GetString(2),
                    reader.GetString(3),
                    reader.IsDBNull(4) ? "" : reader.GetString(4),
                    reader.FieldCount > 5 && !reader.IsDBNull(5) ? reader.GetString(5) : SharedPersona));
            }
        }
        if (rows.Count == 0)
        {
            _embedModel = newModel;
            lock (_lock)
            {
                EnsureOpen();
                SetMeta("embed_model", _embedModel);
            }
            return;
        }
        float[] first = await EmbedAsync(baseUrl, newModel, rows[0].text);
        lock (_lock)
        {
            EnsureOpen();
            _embedModel = newModel;
            _dims = first.Length;
            SetMeta("embed_model", _embedModel);
            SetMeta("dims", _dims.ToString());
        }
        foreach (var row in rows)
        {
            try
            {
                float[] vec = await EmbedAsync(baseUrl, newModel, row.text);
                JToken meta = null;
                if (!string.IsNullOrWhiteSpace(row.meta))
                {
                    try { meta = JToken.Parse(row.meta); } catch { /* ignore */ }
                }
                Upsert(row.kind, row.key, row.text, row.source, meta, vec, row.persona);
            }
            catch (Exception ex)
            {
                Logs.Debug($"AssistentMemory reembed {row.kind}/{row.key}: {ex.Message}");
            }
        }
    }

    public void Dispose()
    {
        lock (_lock)
        {
            _conn?.Dispose();
            _conn = null;
        }
    }
}