MyAgent/MyAgent/Agents/HeavyMemoryRecallAgent.cs

using System.Text;
using System.Text.Json;
using MyAgent.Storage;

namespace MyAgent.Agents;

public sealed class HeavyMemoryRecallAgent(HeavyMemoryStore store, IChatModelClient llm)
{
    public async Task<HeavyMemoryRecallResult> RecallIfNeededAsync(MessageBatch batch, bool requested, CancellationToken ct = default)
    {
        if (!requested) return new HeavyMemoryRecallResult { ShouldRecall = false, Note = "planner_weight_below_threshold" };

        var candidates = await store.SearchAsync(batch.GroupId, batch.MergedText, 5, ct);
        if (candidates.Count == 0)
        {
            return new HeavyMemoryRecallResult
            {
                ShouldRecall = true,
                NeedClarification = true,
                ClarifyingQuestion = "我知道有重记忆，但没按这句话直接找到对应片段。你能补一个关键词、人大概是谁，或者大概时间吗？",
                Note = "no_candidates"
            };
        }

        var decision = await AskModelAsync(batch.MergedText, candidates, ct);
        if (!decision.ShouldUse)
        {
            return new HeavyMemoryRecallResult
            {
                ShouldRecall = true,
                NeedClarification = true,
                ClarifyingQuestion = string.IsNullOrWhiteSpace(decision.ClarifyingQuestion)
                    ? "我可以查重记忆，但这个指代有点宽。你想回忆的是哪个话题/时间段？"
                    : decision.ClarifyingQuestion,
                Note = "model_requested_clarification"
            };
        }

        var selected = candidates
            .Where(x => decision.SelectedIds.Count == 0 || decision.SelectedIds.Contains(x.Id, StringComparer.OrdinalIgnoreCase))
            .Take(3)
            .ToList();
        if (selected.Count == 0) selected = candidates.Take(3).ToList();

        return new HeavyMemoryRecallResult
        {
            ShouldRecall = true,
            NeedClarification = false,
            Entries = selected,
            LightMemoryText = HeavyMemoryStore.FormatRecallText(batch.MergedText, selected),
            Note = "recalled"
        };
    }

    private async Task<RecallDecision> AskModelAsync(string query, IReadOnlyList<HeavyMemoryEntry> candidates, CancellationToken ct)
    {
        var brief = new StringBuilder();
        foreach (var e in candidates)
        {
            brief.AppendLine($"id={e.Id}");
            brief.AppendLine($"topic={e.Topic}");
            brief.AppendLine($"time={e.StartAt:yyyy-MM-dd HH:mm}~{e.EndAt:yyyy-MM-dd HH:mm}");
            brief.AppendLine($"keywords={string.Join("/", e.Keywords)}");
            brief.AppendLine($"summary={e.Summary}");
            brief.AppendLine();
        }

        var system = """
你是重记忆召回判定器。重记忆平时不访问，只有用户明显在要求回忆过往对话/决定/曾经说过的内容时才用。
任务：从候选重记忆中判断是否足以回答。若指代模糊，输出一个很小的澄清问题，不要强行召回。
只输出 JSON：{"shouldUse":true/false,"selectedIds":["..."],"clarifyingQuestion":"..."}
""";
        var prompt = $$"""
用户当前问题：{{query}}

候选重记忆：
{{brief}}
""";
        var result = await llm.CompleteMemoryAsync([
            new ChatMessage { Role = "system", Content = system },
            new ChatMessage { Role = "user", Content = prompt }
        ], 0.1, 400, ct);

        if (result.IsFallback || string.IsNullOrWhiteSpace(result.Content))
            return new RecallDecision { ShouldUse = candidates.Count > 0, SelectedIds = candidates.Take(3).Select(x => x.Id).ToList() };

        try
        {
            var json = ExtractJson(result.Content);
            using var doc = JsonDocument.Parse(json);
            var root = doc.RootElement;
            return new RecallDecision
            {
                ShouldUse = root.TryGetProperty("shouldUse", out var s) && s.GetBoolean(),
                SelectedIds = root.TryGetProperty("selectedIds", out var ids) && ids.ValueKind == JsonValueKind.Array
                    ? ids.EnumerateArray().Select(x => x.GetString() ?? "").Where(x => !string.IsNullOrWhiteSpace(x)).ToList()
                    : [],
                ClarifyingQuestion = root.TryGetProperty("clarifyingQuestion", out var q) ? q.GetString() ?? "" : ""
            };
        }
        catch
        {
            return new RecallDecision { ShouldUse = candidates.Count > 0, SelectedIds = candidates.Take(3).Select(x => x.Id).ToList() };
        }
    }

    private static string ExtractJson(string text)
    {
        var start = text.IndexOf('{');
        var end = text.LastIndexOf('}');
        return start >= 0 && end > start ? text[start..(end + 1)] : text;
    }

    private sealed class RecallDecision
    {
        public bool ShouldUse { get; set; }
        public List<string> SelectedIds { get; set; } = [];
        public string ClarifyingQuestion { get; set; } = "";
    }
}

public sealed class HeavyMemoryRecallResult
{
    public bool ShouldRecall { get; set; }
    public bool NeedClarification { get; set; }
    public string ClarifyingQuestion { get; set; } = "";
    public IReadOnlyList<HeavyMemoryEntry> Entries { get; set; } = [];
    public string LightMemoryText { get; set; } = "";
    public string Note { get; set; } = "";
}