swarm-assistent/AssistentMemory.Heard.cs

using System;
using System.Collections.Generic;
using System.Text;
using System.Threading.Tasks;
using Microsoft.Data.Sqlite;
using Newtonsoft.Json.Linq;
using SwarmUI.Utils;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Train samples linked to the agent as retrievable &quot;heard&quot; dialogue examples.</summary>
public sealed partial class AssistentMemory
{
    public const string HeardKind = "heard";
    public const string HeardSource = "train";

    public static string FormatHeardEmbedText(JArray messages, string pack, string persona)
    {
        StringBuilder sb = new();
        if (!string.IsNullOrWhiteSpace(pack))
        {
            sb.AppendLine($"pack: {pack.Trim()}");
        }
        if (!string.IsNullOrWhiteSpace(persona))
        {
            sb.AppendLine($"persona: {persona.Trim()}");
        }
        foreach (JToken t in messages ?? [])
        {
            if (t is not JObject m)
            {
                continue;
            }
            string role = m["role"]?.ToString() ?? "user";
            string content = m["content"]?.ToString() ?? "";
            if (string.IsNullOrWhiteSpace(content))
            {
                continue;
            }
            sb.AppendLine($"{role}: {content.Trim()}");
        }
        return sb.ToString().Trim();
    }

    public void SetTrainSampleAgentLinked(string id, bool linked)
    {
        if (string.IsNullOrWhiteSpace(id))
        {
            return;
        }
        lock (_lock)
        {
            EnsureTrainingReady();
            if (!HasColumn("train_samples", "agent_linked"))
            {
                return;
            }
            using SqliteCommand cmd = _conn.CreateCommand();
            cmd.CommandText = "UPDATE train_samples SET agent_linked = $v, updated_at = $u WHERE id = $id";
            cmd.Parameters.AddWithValue("$v", linked ? 1 : 0);
            cmd.Parameters.AddWithValue("$u", DateTimeOffset.UtcNow.ToUnixTimeMilliseconds());
            cmd.Parameters.AddWithValue("$id", id.Trim());
            cmd.ExecuteNonQuery();
        }
    }

    public int CountAgentLinkedTrainSamples()
    {
        lock (_lock)
        {
            try
            {
                EnsureTrainingReady();
            }
            catch
            {
                return 0;
            }
            if (!HasColumn("train_samples", "agent_linked"))
            {
                return 0;
            }
            using SqliteCommand cmd = _conn.CreateCommand();
            cmd.CommandText = "SELECT COUNT(*) FROM train_samples WHERE agent_linked = 1";
            return Convert.ToInt32(cmd.ExecuteScalar());
        }
    }

    public async Task<bool> LinkTrainSampleToAgentAsync(string baseUrl, JObject sample, string embedModel)
    {
        if (sample is null)
        {
            return false;
        }
        string id = sample["id"]?.ToString()?.Trim();
        JArray messages = sample["messages"] as JArray ?? [];
        if (string.IsNullOrWhiteSpace(id) || messages.Count == 0)
        {
            return false;
        }
        string status = sample["status"]?.ToString() ?? "draft";
        if (!string.Equals(status, "approved", StringComparison.OrdinalIgnoreCase))
        {
            return false;
        }
        string persona = NormalizePersona(sample["persona"]?.ToString());
        string pack = sample["pack"]?.ToString()?.Trim() ?? "";
        string text = FormatHeardEmbedText(messages, pack, persona);
        if (string.IsNullOrWhiteSpace(text))
        {
            return false;
        }
        JObject meta = new()
        {
            ["train_sample_id"] = id,
            ["pack"] = pack,
            ["messages"] = messages,
            ["source_type"] = sample["source"]?.ToString() ?? "manual",
        };
        if (!string.IsNullOrWhiteSpace(sample["chat_id"]?.ToString()))
        {
            meta["chat_id"] = sample["chat_id"]?.ToString();
        }
        float[] vec = await EmbedAsync(baseUrl, embedModel, text);
        Upsert(HeardKind, id, text, HeardSource, meta, vec, persona);
        SetTrainSampleAgentLinked(id, true);
        return true;
    }

    public void UnlinkTrainSampleFromAgent(JObject sample)
    {
        if (sample is null)
        {
            return;
        }
        string id = sample["id"]?.ToString()?.Trim();
        if (string.IsNullOrWhiteSpace(id))
        {
            return;
        }
        string persona = NormalizePersona(sample["persona"]?.ToString());
        Forget(HeardKind, id, HeardSource, persona);
        if (persona != SharedPersona)
        {
            Forget(HeardKind, id, HeardSource, SharedPersona);
        }
        SetTrainSampleAgentLinked(id, false);
    }

    public JObject BuildHeardExampleFromHit(JObject hit, IEnumerable<string> personaChain = null)
    {
        if (hit is null)
        {
            return null;
        }
        string key = hit["key"]?.ToString() ?? "";
        if (string.IsNullOrWhiteSpace(key))
        {
            return null;
        }
        JObject row = Get(HeardKind, key, personaChain);
        JObject ex = new()
        {
            ["id"] = key,
            ["score"] = hit["score"],
            ["persona"] = hit["persona"],
            ["source"] = hit["source"],
        };
        JObject meta = null;
        try
        {
            string metaRaw = row?["meta_json"]?.ToString();
            if (!string.IsNullOrWhiteSpace(metaRaw))
            {
                meta = JObject.Parse(metaRaw);
            }
        }
        catch
        {
            // ignore
        }
        if (meta?["pack"] is not null)
        {
            ex["pack"] = meta["pack"];
        }
        if (meta?["messages"] is JArray msgs && msgs.Count > 0)
        {
            ex["messages"] = msgs;
        }
        else
        {
            ex["text"] = hit["text"]?.ToString() ?? row?["text"]?.ToString() ?? "";
        }
        return ex;
    }
}