swarm-assistent/AssistentTraining.Agent.cs

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

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Link training dataset samples to the live agent as retrievable &quot;heard&quot; examples.</summary>
public partial class SwarmAssistentExtension
{
    string MemoryEmbedForTraining(string personaId = null)
    {
        string pid = AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId();
        return Config.LoadSettings()["embed_model"]?.ToString()
            ?? Config.LoadAssistant(pid)["embed_model"]?.ToString()
            ?? "nomic-embed-text";
    }

    string MemoryBaseForTraining(JObject raw = null)
        => MemoryBaseUrl(raw?["base_url"]?.ToString());

    public async Task<JObject> AssistentGetDatasetAgentSettings(Session session)
    {
        await Task.CompletedTask;
        try
        {
            JObject settings = Config.LoadTrainingAgent();
            return new JObject
            {
                ["success"] = true,
                ["settings"] = settings,
                ["linked"] = Memory.CountAgentLinkedTrainSamples(),
                ["approved"] = Memory.CountTrainSamples("approved"),
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = ex.Message };
        }
    }

    public async Task<JObject> AssistentSaveDatasetAgentSettings(Session session, JObject settings)
    {
        await Task.CompletedTask;
        if (settings is null)
        {
            return new JObject { ["error"] = "settings required" };
        }
        Config.SaveTrainingAgent(settings);
        return new JObject { ["success"] = true, ["settings"] = Config.LoadTrainingAgent() };
    }

    public async Task<JObject> AssistentLinkTrainSampleToAgent(Session session, string id, JObject raw = null)
    {
        if (string.IsNullOrWhiteSpace(id))
        {
            return new JObject { ["error"] = "id required" };
        }
        JObject sample = Memory.GetTrainSample(id.Trim());
        if (sample is null)
        {
            return new JObject { ["error"] = "sample not found" };
        }
        if (!string.Equals(sample["status"]?.ToString(), "approved", StringComparison.OrdinalIgnoreCase))
        {
            return new JObject { ["error"] = "only approved samples can be linked to the agent" };
        }
        try
        {
            string embed = MemoryEmbedForTraining(sample["persona"]?.ToString());
            bool ok = await Memory.LinkTrainSampleToAgentAsync(MemoryBaseForTraining(raw), sample, embed);
            return new JObject { ["success"] = ok, ["linked"] = Memory.CountAgentLinkedTrainSamples() };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = ex.Message };
        }
    }

    public async Task<JObject> AssistentUnlinkTrainSampleFromAgent(Session session, string id)
    {
        await Task.CompletedTask;
        if (string.IsNullOrWhiteSpace(id))
        {
            return new JObject { ["error"] = "id required" };
        }
        JObject sample = Memory.GetTrainSample(id.Trim());
        if (sample is null)
        {
            return new JObject { ["error"] = "sample not found" };
        }
        Memory.UnlinkTrainSampleFromAgent(sample);
        return new JObject { ["success"] = true, ["linked"] = Memory.CountAgentLinkedTrainSamples() };
    }

    public async Task<JObject> AssistentSyncDatasetToAgent(Session session, JObject raw = null)
    {
        bool approvedOnly = raw?["approved_only"]?.Value<bool?>() ?? true;
        bool relink = raw?["relink"]?.Value<bool?>() ?? false;
        string persona = raw?["persona"]?.ToString();
        string status = approvedOnly ? "approved" : "all";
        List<JObject> samples = Memory.ListTrainSamples(status, persona, null, 2000);
        string baseUrl = MemoryBaseForTraining(raw);
        int linked = 0;
        int skipped = 0;
        List<string> errors = [];
        foreach (JObject sample in samples)
        {
            bool already = sample["agent_linked"]?.Value<bool?>() ?? false;
            if (already && !relink)
            {
                skipped++;
                continue;
            }
            if (!string.Equals(sample["status"]?.ToString(), "approved", StringComparison.OrdinalIgnoreCase))
            {
                skipped++;
                continue;
            }
            try
            {
                string embed = MemoryEmbedForTraining(sample["persona"]?.ToString());
                if (await Memory.LinkTrainSampleToAgentAsync(baseUrl, sample, embed))
                {
                    linked++;
                }
                else
                {
                    skipped++;
                }
            }
            catch (Exception ex)
            {
                errors.Add($"{sample["id"]}: {ex.Message}");
            }
        }
        return new JObject
        {
            ["success"] = true,
            ["linked_now"] = linked,
            ["skipped"] = skipped,
            ["total_linked"] = Memory.CountAgentLinkedTrainSamples(),
            ["errors"] = new JArray(errors.Take(8)),
        };
    }

    async Task TryAutoLinkTrainSample(Session session, JObject sample, JObject raw = null)
    {
        if (sample is null || Memory is null)
        {
            return;
        }
        JObject agent = Config.LoadTrainingAgent();
        if (agent["enabled"]?.Value<bool?>() == false)
        {
            return;
        }
        if (agent["auto_link_on_approve"]?.Value<bool?>() == false)
        {
            return;
        }
        string status = sample["status"]?.ToString() ?? "";
        if (!string.Equals(status, "approved", StringComparison.OrdinalIgnoreCase))
        {
            Memory.UnlinkTrainSampleFromAgent(sample);
            return;
        }
        try
        {
            string embed = MemoryEmbedForTraining(sample["persona"]?.ToString());
            await Memory.LinkTrainSampleToAgentAsync(MemoryBaseForTraining(raw), sample, embed);
        }
        catch (Exception ex)
        {
            Logs.Debug($"TryAutoLinkTrainSample: {ex.Message}");
        }
    }
}