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.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.</summary>
public sealed class AssistentMemory : IDisposable
{
    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")
    {
        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;

    void EnsureOpen()
    {
        if (_conn is not null)
        {
            return;
        }
        _conn = new SqliteConnection($"Data Source={_dbPath}");
        _conn.Open();
        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,
                  text TEXT NOT NULL,
                  source TEXT NOT NULL DEFAULT 'user',
                  meta_json TEXT,
                  embedding BLOB,
                  updated INTEGER NOT NULL,
                  UNIQUE(kind, key, source)
                );
                CREATE INDEX IF NOT EXISTS idx_memories_kind ON memories(kind);
                """;
            cmd.ExecuteNonQuery();
        }
        _embedModel = GetMeta("embed_model") ?? _embedModel;
        _ = int.TryParse(GetMeta("dims"), out _dims);
        _ = int.TryParse(GetMeta("seed_version"), out _seedVersion);
    }

    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)));
    }

    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;
    }

    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);
        if (!needReseed)
        {
            int bundledCount;
            lock (_lock)
            {
                using SqliteCommand cmd = _conn.CreateCommand();
                cmd.CommandText = "SELECT COUNT(*) FROM memories WHERE source = 'bundled'";
                bundledCount = Convert.ToInt32(cmd.ExecuteScalar());
            }
            if (bundledCount > 0)
            {
                return;
            }
        }

        List<JObject> docs = config.LoadMemorySeedDocs();
        if (docs.Count == 0)
        {
            return;
        }

        // Probe embed
        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}");
            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();
            if (string.IsNullOrWhiteSpace(key) || string.IsNullOrWhiteSpace(text))
            {
                continue;
            }
            try
            {
                float[] vec = await EmbedAsync(baseUrl, wantModel, text);
                Upsert(kind, key, text, "bundled", doc["tags"], vec);
            }
            catch (Exception ex)
            {
                Logs.Debug($"AssistentMemory seed item {kind}/{key}: {ex.Message}");
            }
        }
    }

    public void Upsert(string kind, string key, string text, string source, JToken meta, float[] embedding)
    {
        lock (_lock)
        {
            EnsureOpen();
            kind = (kind ?? "note").Trim().ToLowerInvariant();
            key = (key ?? "").Trim();
            text = (text ?? "").Trim();
            source = string.IsNullOrWhiteSpace(source) ? "user" : source.Trim();
            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, text, source, meta_json, embedding, updated)
                VALUES($kind, $key, $text, $source, $meta, $emb, $upd)
                ON CONFLICT(kind, key, source) 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("$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();
        }
    }

    public void Forget(string kind, string key, string source = null)
    {
        lock (_lock)
        {
            EnsureOpen();
            using SqliteCommand cmd = _conn.CreateCommand();
            if (string.IsNullOrWhiteSpace(source))
            {
                cmd.CommandText = "DELETE FROM memories WHERE kind = $kind AND key = $key AND source != 'bundled'";
            }
            else
            {
                cmd.CommandText = "DELETE FROM memories WHERE kind = $kind AND key = $key AND source = $source";
                cmd.Parameters.AddWithValue("$source", source);
            }
            cmd.Parameters.AddWithValue("$kind", (kind ?? "").Trim().ToLowerInvariant());
            cmd.Parameters.AddWithValue("$key", (key ?? "").Trim());
            cmd.ExecuteNonQuery();
        }
    }

    public async Task<JArray> RetrieveAsync(string baseUrl, string query, int topK = 10, string modelOverride = null)
    {
        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 [];
        }

        List<(float score, JObject row)> scored = [];
        lock (_lock)
        {
            EnsureOpen();
            using SqliteCommand cmd = _conn.CreateCommand();
            cmd.CommandText = "SELECT kind, key, text, source, meta_json, embedding FROM memories WHERE embedding IS NOT NULL";
            using SqliteDataReader reader = cmd.ExecuteReader();
            while (reader.Read())
            {
                float[] emb = BytesToFloats(reader.IsDBNull(5) ? null : (byte[])reader.GetValue(5));
                float score = Cosine(q, emb);
                if (float.IsNegativeInfinity(score))
                {
                    continue;
                }
                scored.Add((score, new JObject
                {
                    ["kind"] = reader.GetString(0),
                    ["key"] = reader.GetString(1),
                    ["text"] = reader.GetString(2),
                    ["source"] = reader.GetString(3),
                    ["score"] = Math.Round(score, 4),
                }));
            }
        }
        return new JArray(scored.OrderByDescending(s => s.score).Take(Math.Clamp(topK, 1, 30)).Select(s => s.row));
    }

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

    public async Task ReembedAllAsync(string baseUrl, string newModel)
    {
        List<(string kind, string key, string text, string source, string meta)> rows = [];
        lock (_lock)
        {
            EnsureOpen();
            using SqliteCommand cmd = _conn.CreateCommand();
            cmd.CommandText = "SELECT kind, key, text, source, meta_json 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)));
            }
        }
        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);
            }
            catch (Exception ex)
            {
                Logs.Debug($"AssistentMemory reembed {row.kind}/{row.key}: {ex.Message}");
            }
        }
    }

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