swarm-assistent/AssistentMemoryApi.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using Newtonsoft.Json.Linq;
using SwarmUI.Accounts;
using SwarmUI.Utils;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Read/write routes for the vector memory list in ⚙ and the gpu-rent wanted queue badge.</summary>
public partial class SwarmAssistentExtension
{
    /// <summary>Embed model the UI should use: settings overlay wins, then persona assistant.json.</summary>
    string MemoryEmbedModel(string requested = null)
    {
        if (!string.IsNullOrWhiteSpace(requested))
        {
            return requested.Trim();
        }
        return Config?.LoadSettings()["embed_model"]?.ToString()
            ?? Config?.LoadAssistant(Config.DefaultPersonaId())["embed_model"]?.ToString()
            ?? "nomic-embed-text";
    }

    string MemoryBaseUrl(string requested = null)
    {
        return NormalizeBaseUrl(string.IsNullOrWhiteSpace(requested)
            ? Config?.LoadSettings()["base_url"]?.ToString()
            : requested);
    }

    string ResolveApiPersona(string persona, string scope)
    {
        string s = (scope ?? "").Trim().ToLowerInvariant();
        if (s is "shared" or "common" or "global")
        {
            return AssistentMemory.SharedPersona;
        }
        if (s is "personal")
        {
            return AssistentConfig.SafeId(persona) ?? Config?.DefaultPersonaId() ?? "neutral";
        }
        if (string.IsNullOrWhiteSpace(persona) || AssistentMemory.IsShared(persona))
        {
            return AssistentMemory.SharedPersona;
        }
        return AssistentMemory.NormalizePersona(persona);
    }

    public async Task<JObject> AssistentListMemory(Session session, int limit = 200, string kind = null, string persona = null, string scope = null)
    {
        await Task.CompletedTask;
        if (Memory is null)
        {
            return new JObject { ["error"] = "memory not ready" };
        }
        try
        {
            JArray all = Memory.ListAll(limit);
            string filter = (kind ?? "").Trim().ToLowerInvariant();
            string wantScope = (scope ?? "").Trim().ToLowerInvariant();
            string wantPersona = (persona ?? "").Trim();
            IEnumerable<JToken> q = all;
            if (!string.IsNullOrWhiteSpace(filter) && filter != "all")
            {
                q = q.Where(t => string.Equals(t?["kind"]?.ToString(), filter, StringComparison.OrdinalIgnoreCase));
            }
            bool personaGiven = !string.IsNullOrWhiteSpace(wantPersona);
            if (wantScope is "shared" or "common" or "global"
                || (personaGiven && AssistentMemory.IsShared(wantPersona)))
            {
                q = q.Where(t => string.Equals(t?["scope"]?.ToString(), "shared", StringComparison.OrdinalIgnoreCase));
            }
            else if (wantScope is "personal" || personaGiven)
            {
                string pid = AssistentMemory.NormalizePersona(personaGiven ? wantPersona : Config?.DefaultPersonaId());
                q = q.Where(t => string.Equals(t?["persona"]?.ToString(), pid, StringComparison.OrdinalIgnoreCase));
            }
            JArray rows = new(q);
            JArray kinds = new(all
                .Select(t => t?["kind"]?.ToString())
                .Where(s => !string.IsNullOrWhiteSpace(s))
                .Distinct(StringComparer.OrdinalIgnoreCase)
                .OrderBy(s => s, StringComparer.OrdinalIgnoreCase));
            return new JObject
            {
                ["success"] = true,
                ["memories"] = rows,
                ["kinds"] = kinds,
                ["total"] = Memory.CountAll(),
                ["embed_model"] = Memory.EmbedModel,
                ["dims"] = Memory.Dims,
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"memory list: {ex.Message}" };
        }
    }

    public async Task<JObject> AssistentUpsertMemory(Session session, string kind, string key, string text, string source = "user", string baseUrl = null, string embed_model = null, string persona = null, string scope = null)
    {
        if (Memory is null)
        {
            return new JObject { ["error"] = "memory not ready" };
        }
        kind = (kind ?? "note").Trim().ToLowerInvariant();
        key = (key ?? "").Trim();
        text = (text ?? "").Trim();
        if (string.IsNullOrWhiteSpace(key) || string.IsNullOrWhiteSpace(text))
        {
            return new JObject { ["error"] = "key and text required" };
        }
        string src = (source ?? "user").Trim().ToLowerInvariant();
        if (src == "bundled")
        {
            return new JObject { ["error"] = "bundled memories are read-only — use memory-seed/" };
        }
        string target = ResolveApiPersona(persona, scope);
        try
        {
            await Memory.UpsertTextAsync(MemoryBaseUrl(baseUrl), kind, key, text, src, null, MemoryEmbedModel(embed_model), target);
            return new JObject
            {
                ["success"] = true,
                ["kind"] = kind,
                ["key"] = key,
                ["source"] = src,
                ["scope"] = AssistentMemory.IsShared(target) ? "shared" : "personal",
                ["persona"] = AssistentMemory.IsShared(target) ? "shared" : target,
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"memory upsert: {ex.Message}" };
        }
    }

    public async Task<JObject> AssistentForgetMemory(Session session, string kind, string key, string source = null, string persona = null, string scope = null)
    {
        await Task.CompletedTask;
        if (Memory is null)
        {
            return new JObject { ["error"] = "memory not ready" };
        }
        if (string.IsNullOrWhiteSpace(kind) || string.IsNullOrWhiteSpace(key))
        {
            return new JObject { ["error"] = "kind and key required" };
        }
        if (string.Equals((source ?? "").Trim(), "bundled", StringComparison.OrdinalIgnoreCase))
        {
            return new JObject { ["error"] = "bundled memories come back on reseed — edit memory-seed/ instead" };
        }
        string target = ResolveApiPersona(persona, scope);
        try
        {
            Memory.Forget(kind, key, string.IsNullOrWhiteSpace(source) ? null : source.Trim(), target);
            return new JObject
            {
                ["success"] = true,
                ["kind"] = kind.Trim().ToLowerInvariant(),
                ["key"] = key.Trim(),
                ["scope"] = AssistentMemory.IsShared(target) ? "shared" : "personal",
                ["persona"] = AssistentMemory.IsShared(target) ? "shared" : target,
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"memory forget: {ex.Message}" };
        }
    }

    public async Task<JObject> AssistentSearchMemory(Session session, string query, string kind = null, int limit = 10, string persona = null, string baseUrl = null, string embed_model = null)
    {
        if (Memory is null)
        {
            return new JObject { ["error"] = "memory not ready" };
        }
        query = (query ?? "").Trim();
        if (query.Length < 2)
        {
            return new JObject { ["error"] = "query required" };
        }
        string pid = AssistentConfig.SafeId(persona) ?? Config.DefaultPersonaId();
        try
        {
            JArray rows = await Memory.SearchAsync(MemoryBaseUrl(baseUrl), query, kind, limit, MemoryEmbedModel(embed_model), Config.PersonaExtendsChain(pid));
            return new JObject { ["success"] = true, ["query"] = query, ["kind"] = kind ?? "", ["memories"] = rows };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"memory search: {ex.Message}" };
        }
    }

    public async Task<JObject> AssistentGetMemory(Session session, string kind, string key, string persona = null)
    {
        await Task.CompletedTask;
        if (Memory is null)
        {
            return new JObject { ["error"] = "memory not ready" };
        }
        string pid = AssistentConfig.SafeId(persona) ?? Config.DefaultPersonaId();
        JObject row = Memory.Get(kind, key, Config.PersonaExtendsChain(pid));
        if (row is null)
        {
            return new JObject { ["success"] = true, ["missing"] = true, ["kind"] = kind, ["key"] = key };
        }
        return new JObject { ["success"] = true, ["memory"] = row };
    }

    public async Task<JObject> AssistentLookupTags(Session session, string query, int limit = 20)
    {
        await Task.CompletedTask;
        if (Memory is null)
        {
            return new JObject { ["error"] = "memory not ready" };
        }
        query = (query ?? "").Trim();
        if (query.Length < 1)
        {
            return new JObject { ["error"] = "query required" };
        }
        try
        {
            JArray tags = Memory.LookupTags(query, limit);
            return new JObject
            {
                ["success"] = true,
                ["query"] = query,
                ["tags"] = tags,
                ["indexed"] = Memory.TagCount(),
                ["csv"] = Memory.FindAutocompleteCsv() ?? "",
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"tag lookup: {ex.Message}" };
        }
    }

    public async Task<JObject> AssistentClearMemory(Session session, string scope = null, string kind = null, string persona = null)
    {
        await Task.CompletedTask;
        if (Memory is null)
        {
            return new JObject { ["error"] = "memory not ready" };
        }
        try
        {
            string targetPersona = persona;
            string wantScope = (scope ?? "").Trim().ToLowerInvariant();
            if (wantScope is "personal" && string.IsNullOrWhiteSpace(targetPersona))
            {
                targetPersona = Config?.DefaultPersonaId() ?? "neutral";
            }
            int n = Memory.ClearCraftMemory(scope, kind, targetPersona);
            return new JObject
            {
                ["success"] = true,
                ["deleted"] = n,
                ["scope"] = scope ?? "all",
                ["kind"] = kind ?? "all",
                ["persona"] = AssistentMemory.IsShared(AssistentMemory.NormalizePersona(targetPersona))
                    ? "shared"
                    : AssistentMemory.NormalizePersona(targetPersona),
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"memory clear: {ex.Message}" };
        }
    }

    /// <summary>The gpu-rent wanted queue (models pending the next <c>up</c>) — count + entries.</summary>
    public async Task<JObject> AssistentListWanted(Session session)
    {
        await Task.CompletedTask;
        string path = WantedModelsPath();
        JArray items = [];
        if (!File.Exists(path))
        {
            return new JObject { ["success"] = true, ["count"] = 0, ["items"] = items, ["path"] = path };
        }
        try
        {
            Dictionary<string, List<WantedEntry>> sections = LoadWantedYaml(File.ReadAllText(path, Encoding.UTF8));
            foreach ((string kind, List<WantedEntry> list) in sections.OrderBy(p => p.Key, StringComparer.OrdinalIgnoreCase))
            {
                foreach (WantedEntry entry in list)
                {
                    items.Add(new JObject
                    {
                        ["kind"] = kind,
                        ["url"] = entry.Url,
                        ["title"] = entry.Title,
                        ["version_id"] = entry.VersionId,
                    });
                }
            }
            return new JObject
            {
                ["success"] = true,
                ["count"] = items.Count,
                ["items"] = items,
                ["path"] = path,
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"wanted queue: {ex.Message}" };
        }
    }
}