Microsoft Agent Framework をローカルLLMで試してみる その11(AIContextProviderで簡易RAG)

using System.Text.RegularExpressions;

namespace SimpleRagSample;

// 固定ドキュメント群に対する最小の全文検索。RAG の入口だけ確認したいので実装は意図的に単純にしている。
public sealed class InMemoryKnowledgeSearchService
{
    private readonly IReadOnlyList<KnowledgeDocument> _documents;

    public InMemoryKnowledgeSearchService(IReadOnlyList<KnowledgeDocument> documents)
    {
        ArgumentNullException.ThrowIfNull(documents);
        _documents = documents;
    }

    public IReadOnlyList<KnowledgeSearchResult> Search(string query, int topK = 3)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(query);

        if (topK <= 0)
        {
            throw new ArgumentOutOfRangeException(nameof(topK), topK, "topK must be greater than 0.");
        }

        string[] terms = SplitTerms(query);

        // 完全一致、タイトル一致、本文中の出現回数だけで粗く順位付けする。
        return _documents
            .Select(document => new KnowledgeSearchResult(document, CalculateScore(document, query, terms)))
            .Where(result => result.Score > 0)
            .OrderByDescending(result => result.Score)
            .ThenBy(result => result.Document.Title, StringComparer.Ordinal)
            .Take(topK)
            .ToArray();
    }

    private static string[] SplitTerms(string text)
    {
        // 日本語の簡易入力を想定し、句読点や括弧でだけ分割する。
        string[] terms = Regex.Split(text, @"[\s,、。・/\\()\[\]{}]+")
            .Select(term => term.Trim())
            .Where(term => term.Length >= 2)
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .ToArray();

        return terms.Length == 0 ? [text.Trim()] : terms;
    }

    private static double CalculateScore(KnowledgeDocument document, string query, IReadOnlyList<string> terms)
    {
        string normalizedTitle = document.Title.ToLowerInvariant();
        string normalizedContent = document.Content.ToLowerInvariant();
        string normalizedQuery = query.ToLowerInvariant();
        double score = 0;

        // 質問全体がそのまま含まれる場合は、まず強めに加点する。
        if (normalizedTitle.Contains(normalizedQuery, StringComparison.Ordinal) ||
            normalizedContent.Contains(normalizedQuery, StringComparison.Ordinal))
        {
            score += 3.0d;
        }

        foreach (string term in terms)
        {
            string normalizedTerm = term.ToLowerInvariant();

            if (normalizedTitle.Contains(normalizedTerm, StringComparison.Ordinal))
            {
                score += 2.0d;
            }

            // 本文一致は出現回数ベースにして、実装を追いやすくする。
            score += CountOccurrences(normalizedContent, normalizedTerm);
        }

        return score;
    }

    private static int CountOccurrences(string text, string term)
    {
        if (string.IsNullOrWhiteSpace(text) || string.IsNullOrWhiteSpace(term))
        {
            return 0;
        }

        int count = 0;
        int index = 0;

        while ((index = text.IndexOf(term, index, StringComparison.Ordinal)) >= 0)
        {
            count++;
            index += term.Length;
        }

        return count;
    }
}

public sealed record KnowledgeDocument(string Id, string Title, string Content);

public sealed record KnowledgeSearchResult(KnowledgeDocument Document, double Score);

using Microsoft.Agents.AI;
using Microsoft.Extensions.AI;
using System.Text;

namespace SimpleRagSample;

// 検索結果を assistant message に整形して、モデル呼び出し直前に差し込む。
public sealed class SearchAiContextProvider : AIContextProvider
{
    private readonly InMemoryKnowledgeSearchService _searchService;
    private readonly InMemoryChatHistoryProvider? _chatHistoryProvider;

    public SearchAiContextProvider(
        InMemoryKnowledgeSearchService searchService,
        InMemoryChatHistoryProvider? chatHistoryProvider = null)
    {
        _searchService = searchService ?? throw new ArgumentNullException(nameof(searchService));
        _chatHistoryProvider = chatHistoryProvider;
    }

    protected override ValueTask<AIContext> ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default)
    {
        if (context.Session is null)
        {
            return new ValueTask<AIContext>(new AIContext());
        }

        // 直前に保存した最新入力を優先し、見つからなければ履歴へフォールバックする。
        string userInput = GetLatestUserMessage(context.Session);
        if (string.IsNullOrWhiteSpace(userInput))
        {
            return new ValueTask<AIContext>(new AIContext());
        }

        IReadOnlyList<KnowledgeSearchResult> results = _searchService.Search(userInput, topK: 3);
        if (results.Count == 0)
        {
            return new ValueTask<AIContext>(new AIContext
            {
                Instructions = "検索結果が見つからない場合は、その旨を明示して一般的な説明に留めてください。"
            });
        }

        // 最小サンプルなので、検索結果はそのまま読みやすい箇条書きで埋め込む。
        StringBuilder builder = new();
        builder.AppendLine("以下はローカル検索で抽出した参考コンテキストです。質問に関連する場合は優先して参照してください。");

        foreach (KnowledgeSearchResult result in results)
        {
            builder.AppendLine($"- タイトル: {result.Document.Title}");
            builder.AppendLine($"  スコア: {result.Score:F2}");
            builder.AppendLine($"  内容: {result.Document.Content}");
        }
        Console.WriteLine( builder.ToString() );

        AIContext aiContext = new()
        {
            Instructions = "参考コンテキストに書かれている内容を優先し、根拠が足りない場合は断定しないでください。",
            Messages = [new ChatMessage(ChatRole.Assistant, builder.ToString())]
        };

        return new ValueTask<AIContext>(aiContext);
    }

    private string GetLatestUserMessage(AgentSession session)
    {
        string? pendingUserInput = PendingUserInputState.GetPendingUserInput(session);
        if (!string.IsNullOrWhiteSpace(pendingUserInput))
        {
            return pendingUserInput;
        }

        // pending input がない場面では、会話履歴の最後の user message を使う。
        if (_chatHistoryProvider is null)
        {
            return string.Empty;
        }

        ChatMessage? userMessage = _chatHistoryProvider
            .GetMessages(session)
            .LastOrDefault(message => message.Role == ChatRole.User);

        return ExtractText(userMessage);
    }

    private static string ExtractText(ChatMessage? message)
    {
        if (message is null)
        {
            return string.Empty;
        }

        return string.Concat(message.Contents.OfType<TextContent>().Select(content => content.Text)).Trim();
    }
}

using Microsoft.Agents.AI;
using System.Text.Json.Serialization;

namespace SimpleRagSample;

// AIContextProvider から最新質問を確実に読めるよう、実行中だけセッション状態へ退避する。
public static class PendingUserInputState
{
    public const string StateKey = "SimpleRagSample.PendingUserInput";

    private static readonly ProviderSessionState<State> SessionState = new(
        _ => new State(string.Empty),
        StateKey);

    public static void SetPendingUserInput(AgentSession session, string userInput)
    {
        ArgumentNullException.ThrowIfNull(session);
        ArgumentException.ThrowIfNullOrWhiteSpace(userInput);

        SessionState.SaveState(session, new State(userInput));
    }

    public static string? GetPendingUserInput(AgentSession session)
    {
        ArgumentNullException.ThrowIfNull(session);

        // 値が空文字なら「未設定」とみなして null を返す。
        State state = SessionState.GetOrInitializeState(session);
        return string.IsNullOrWhiteSpace(state.UserInput) ? null : state.UserInput;
    }

    public static void ClearPendingUserInput(AgentSession session)
    {
        ArgumentNullException.ThrowIfNull(session);
        SessionState.SaveState(session, new State(string.Empty));
    }

    public sealed class State
    {
        [JsonConstructor]
        public State(string userInput)
        {
            UserInput = userInput;
        }

        [JsonPropertyName("userInput")]
        public string UserInput { get; }
    }
}

using Microsoft.Agents.AI;
using Microsoft.Extensions.AI;
using OpenAI;
using SimpleRagSample;
using System.ClientModel;

// AIContextProvider が確実に最新の質問を参照できるよう、実行中だけ一時状態に保持する。
static async Task<AgentResponse> RunWithSearchContextAsync(AIAgent agent, AgentSession session, string userInput)
{
    PendingUserInputState.SetPendingUserInput(session, userInput);

    try
    {
        return await agent.RunAsync(userInput, session);
    }
    finally
    {
        PendingUserInputState.ClearPendingUserInput(session);
    }
}

const string lmStudioUrl = "http://localhost:1234/v1";
const string modelName = "openai/gpt-oss-20b";
const string dummyApiKey = "sk-dummy";

// 既存サンプルと同じく LM Studio の OpenAI 互換 API を使う。
var clientOptions = new OpenAIClientOptions
{
    Endpoint = new Uri(lmStudioUrl)
};

var openAIClient = new OpenAIClient(
    new ApiKeyCredential(dummyApiKey),
    clientOptions);

IChatClient chatClient = openAIClient
    .GetChatClient(modelName)
    .AsIChatClient();

var chatHistoryProvider = new InMemoryChatHistoryProvider();

// 検索対象は固定データだけに絞り、RAG の流れそのものを追いやすくする。
var searchService = new InMemoryKnowledgeSearchService(
    [
        new KnowledgeDocument(
            "maf-overview",
            "Microsoft Agent Framework の基本構成",
            "Microsoft Agent Framework では AIAgent を中心に ChatHistoryProvider と AIContextProvider を組み合わせて拡張できます。"),
        new KnowledgeDocument(
            "simple-rag-flow",
            "最小 RAG の流れ",
            "最小構成の RAG ではユーザー質問を検索し、上位ドキュメントを AIContextProvider から assistant message として注入して最終回答を生成します。"),
        new KnowledgeDocument(
            "search-strategy",
            "検索サービスの考え方",
            "この確認用サンプルでは固定ドキュメントに対するキーワード検索だけを実装します。タイトル一致を少し強く評価し、複雑な前処理や加工は行いません。"),
        new KnowledgeDocument(
            "context-provider-role",
            "AIContextProvider の役割",
            "AIContextProvider はモデル呼び出し直前に補助文脈を差し込む拡張ポイントです。検索結果をここで渡すと、モデルは検索内容を根拠として回答できます。"),
        new KnowledgeDocument(
            "pending-input",
            "最新質問の受け渡し",
            "このサンプルでは PendingUserInputState を使って最新の質問を一時保存し、AIContextProvider 側から確実に取得できるようにしています。")
    ]);

// 検索結果を AIContext に詰める責務は AIContextProvider 側に寄せる。
AIContextProvider contextProvider = new SearchAiContextProvider(searchService, chatHistoryProvider);

AIAgent agent = chatClient.AsAIAgent(new ChatClientAgentOptions
{
    Name = "SimpleRagAgent",
    Description = "検索と AIContextProvider ベースの最小 RAG を確認するサンプルです。",
    ChatOptions = new ChatOptions
    {
        Instructions = "提供された参考コンテキストが質問に関連する場合はそれを優先して、日本語で簡潔に回答してください。根拠が不足する場合は不足していると明示してください。",
        MaxOutputTokens = 1200
    },
    ChatHistoryProvider = chatHistoryProvider,
    AIContextProviders = [contextProvider]
});

// サンプル実行は 1 問だけにし、最小の検索 + 文脈注入ループを確認しやすくする。
AgentSession session = await agent.CreateSessionAsync();

string userInput = "Microsoft Agent Framework で AIContextProvider を使った最小 RAG はどう構成しますか?";
AgentResponse response = await RunWithSearchContextAsync(agent, session, userInput);

Console.WriteLine("[UserInput]");
Console.WriteLine(userInput);
Console.WriteLine();
Console.WriteLine("[Response]");
Console.WriteLine(response);