From a816408cd4c41e30a13f7f68b09a09289c620d49 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Tue, 28 Oct 2025 14:16:29 +0000 Subject: [PATCH] .NET: Add Rag AIContext Provider (#1630) * Add Rag AIContext Provider * Fix issues * Improve options naming based on PR feedback. * Move Rag Provider to Data namespace * Add Raw Representation to RagSearchResult * Renaming RagProvider to TextSearchProvider --- .../Microsoft.Agents.AI/AgentJsonUtilities.cs | 2 + .../Data/TextSearchProvider.cs | 311 ++++++++++ .../Data/TextSearchProviderOptions.cs | 77 +++ .../Data/TextSearchProviderTests.cs | 542 ++++++++++++++++++ 4 files changed, 932 insertions(+) create mode 100644 dotnet/src/Microsoft.Agents.AI/Data/TextSearchProvider.cs create mode 100644 dotnet/src/Microsoft.Agents.AI/Data/TextSearchProviderOptions.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs index 4e1e0ac27f..5ef6978c00 100644 --- a/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs @@ -4,6 +4,7 @@ using System.Diagnostics.CodeAnalysis; using System.Text.Encodings.Web; using System.Text.Json; using System.Text.Json.Serialization; +using Microsoft.Agents.AI.Data; namespace Microsoft.Agents.AI; @@ -62,6 +63,7 @@ internal static partial class AgentJsonUtilities // Agent abstraction types [JsonSerializable(typeof(ChatClientAgentThread.ThreadState))] + [JsonSerializable(typeof(TextSearchProvider.TextSearchProviderState))] [ExcludeFromCodeCoverage] internal sealed partial class JsonContext : JsonSerializerContext; diff --git a/dotnet/src/Microsoft.Agents.AI/Data/TextSearchProvider.cs b/dotnet/src/Microsoft.Agents.AI/Data/TextSearchProvider.cs new file mode 100644 index 0000000000..b3b62318f6 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI/Data/TextSearchProvider.cs @@ -0,0 +1,311 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Text; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.Logging; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Agents.AI.Data; + +/// +/// A text search context provider that performs a search over external knowledge +/// and injects the formatted results into the AI invocation context, or exposes a search tool for on-demand use. +/// This provider can be used to enable Retrieval Augmented Generation (RAG) on an agent. +/// +/// +/// +/// The provider supports two behaviors controlled via : +/// +/// – Automatically performs a search prior to every AI invocation and injects results as additional instructions. +/// – Exposes a function tool that the model may invoke to retrieve contextual information when needed. +/// +/// +/// +/// When is greater than zero the provider will retain the most recent +/// user and assistant messages (up to the configured limit) across invocations and prepend them (in chronological order) +/// to the current request messages when forming the search input. This can improve search relevance by providing +/// multi-turn context to the retrieval layer without permanently altering the conversation history. +/// +/// +public sealed class TextSearchProvider : AIContextProvider +{ + private const string DefaultPluginSearchFunctionName = "Search"; + private const string DefaultPluginSearchFunctionDescription = "Allows searching for additional information to help answer the user question."; + private const string DefaultContextPrompt = "## Additional Context\nConsider the following information from source documents when responding to the user:"; + private const string DefaultCitationsPrompt = "Include citations to the source document with document name and link if document name and link is available."; + + private readonly Func>> _searchAsync; + private readonly ILogger? _logger; + private readonly AITool[] _tools; + private readonly Queue _recentMessagesText; + private readonly TextSearchProviderOptions _options; + + /// + /// Initializes a new instance of the class. + /// + /// Delegate that executes the search logic. Must not be . + /// Optional configuration options. + /// Optional logger factory. + /// Thrown when is . + public TextSearchProvider(Func>> searchAsync, TextSearchProviderOptions? options = null, ILoggerFactory? loggerFactory = null) + { + this._searchAsync = searchAsync ?? throw new ArgumentNullException(nameof(searchAsync)); + this._options = options ?? new(); + Throw.IfLessThan(this._options.RecentMessageMemoryLimit, 0); + this._logger = loggerFactory?.CreateLogger(); + this._recentMessagesText = new(); + + // Create the on-demand search tool (only used if behavior is OnDemandFunctionCalling) + this._tools = + [ + AIFunctionFactory.Create( + this.SearchAsync, + name: this._options.FunctionToolName ?? DefaultPluginSearchFunctionName, + description: this._options.FunctionToolDescription ?? DefaultPluginSearchFunctionDescription) + ]; + } + + /// + /// Initializes a new instance of the class from previously serialized state. + /// + /// Delegate that executes the search logic. Must not be . + /// A representing the serialized provider state. + /// Optional serializer options (unused - source generated context is used). + /// Optional configuration options. + /// Optional logger factory. + /// Thrown when is . + /// + /// Only overridden prompts (function name, function description, context prompt, citations prompt) are restored. + /// If a value was not persisted or matches the defaults it will fall back to the built-in defaults. + /// Custom delegates are not serialized. + /// + public TextSearchProvider(Func>> searchAsync, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, TextSearchProviderOptions? options = null, ILoggerFactory? loggerFactory = null) + { + this._searchAsync = searchAsync ?? throw new ArgumentNullException(nameof(searchAsync)); + this._options = options ?? new(); + Throw.IfLessThan(this._options.RecentMessageMemoryLimit, 0); + this._logger = loggerFactory?.CreateLogger(); + + List? restoredMessages = null; + + var state = serializedState.Deserialize(AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(TextSearchProviderState))) as TextSearchProviderState; + if (state?.RecentMessagesText is { Count: > 0 }) + { + restoredMessages = state.RecentMessagesText; + } + + // Restore recent messages respecting the limit (may truncate if limit changed afterwards). + this._recentMessagesText = restoredMessages is null ? new() : new(restoredMessages.Take(this._options.RecentMessageMemoryLimit)); + + // Create the on-demand search tool (only used if behavior is OnDemandFunctionCalling) + this._tools = + [ + AIFunctionFactory.Create( + this.SearchAsync, + name: this._options.FunctionToolName ?? DefaultPluginSearchFunctionName, + description: this._options.FunctionToolDescription ?? DefaultPluginSearchFunctionDescription) + ]; + } + + /// + public override async ValueTask InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + if (this._options.SearchTime != TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke) + { + // Expose the search tool for on-demand invocation. + return new AIContext { Tools = this._tools }; // No automatic instructions injection. + } + + // Aggregate text from memory + current request messages. + var sbInput = new StringBuilder(); + foreach (var messageText in this._recentMessagesText) + { + if (sbInput.Length > 0) + { + sbInput.Append('\n'); + } + sbInput.Append(messageText); + } + + foreach (var message in context.RequestMessages) + { + if (!string.IsNullOrWhiteSpace(message?.Text)) + { + if (sbInput.Length > 0) + { + sbInput.Append('\n'); + } + sbInput.Append(message.Text); + } + } + string input = sbInput.ToString(); + + // Search + var results = await this._searchAsync(input, cancellationToken).ConfigureAwait(false); + IList materialized = results as IList ?? results.ToList(); + if (materialized.Count == 0) + { + this._logger?.LogWarning("TextSearchProvider: No search results found."); + return new AIContext(); + } + + // Format search results + string formatted = this.FormatResults(materialized); + + this._logger?.LogInformation("TextSearchProvider: Retrieved {Count} search results.", materialized.Count); + this._logger?.LogTrace("TextSearchProvider Input:{Input}\nContext Instructions:{Instructions}", input, formatted); + + return new AIContext + { + Messages = [new ChatMessage(ChatRole.User, formatted)] + }; + } + + /// + public override ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + int limit = this._options.RecentMessageMemoryLimit; + if (limit <= 0) + { + return default; // Memory disabled. + } + + var messagesText = context.RequestMessages + .Concat(context.ResponseMessages ?? []) + .Where(m => (m.Role == ChatRole.User || m.Role == ChatRole.Assistant) && !string.IsNullOrWhiteSpace(m.Text)) + .Select(m => m.Text) + .ToList(); + if (messagesText.Count > limit) + { + // If the current request/response exceeds the limit, only keep the most recent messages from it. + messagesText = messagesText.Skip(messagesText.Count - limit).ToList(); + } + + foreach (var message in messagesText) + { + this._recentMessagesText.Enqueue(message); + } + + while (this._recentMessagesText.Count > limit) + { + this._recentMessagesText.Dequeue(); + } + + return default; + } + + /// + /// Serializes the current provider state to a containing any overridden prompts or descriptions. + /// + /// Optional serializer options (ignored, source generated context is used). + /// A with overridden values, or default if nothing was overridden. + public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) + { + // Only persist values that differ from defaults plus recent memory configuration & messages. + TextSearchProviderState state = new(); + if (this._options.RecentMessageMemoryLimit > 0 && this._recentMessagesText.Count > 0) + { + state.RecentMessagesText = this._recentMessagesText.Take(this._options.RecentMessageMemoryLimit).ToList(); + } + + return JsonSerializer.SerializeToElement(state, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(TextSearchProviderState))); + } + + /// + /// Function callable by the AI model (when enabled) to perform an ad-hoc search. + /// + /// The query text. + /// Cancellation token. + /// Formatted search results. + internal async Task SearchAsync(string userQuestion, CancellationToken cancellationToken = default) + { + var results = await this._searchAsync(userQuestion, cancellationToken).ConfigureAwait(false); + IList materialized = results as IList ?? results.ToList(); + string formatted = this.FormatResults(materialized); + + this._logger?.LogInformation("TextSearchProvider: Retrieved {Count} search results.", materialized.Count); + this._logger?.LogTrace("TextSearchProvider Input:{UserQuestion}\nContext Instructions:{Instructions}", userQuestion, formatted); + + return formatted; + } + + /// + /// Formats search results into an instructions string for model consumption. + /// + /// The results. + /// Formatted string (may be empty). + private string FormatResults(IList results) + { + if (this._options.ContextFormatter is not null) + { + return this._options.ContextFormatter(results) ?? string.Empty; + } + + if (results.Count == 0) + { + return string.Empty; // No extra context. + } + + var sb = new StringBuilder(); + sb.AppendLine(this._options.ContextPrompt ?? DefaultContextPrompt); + for (int i = 0; i < results.Count; i++) + { + var result = results[i]; + if (!string.IsNullOrWhiteSpace(result.Name)) + { + sb.AppendLine($"SourceDocName: {result.Name}"); + } + if (!string.IsNullOrWhiteSpace(result.Link)) + { + sb.AppendLine($"SourceDocLink: {result.Link}"); + } + sb.AppendLine($"Contents: {result.Value}"); + sb.AppendLine("----"); + } + sb.AppendLine(this._options.CitationsPrompt ?? DefaultCitationsPrompt); + sb.AppendLine(); + return sb.ToString(); + } + + /// + /// Represents a single retrieved text search result. + /// + public sealed class TextSearchSearchResult + { + /// + /// Gets or sets the display name of the source document (optional). + /// + public string? Name { get; set; } + + /// + /// Gets or sets a link/URL to the source document (optional). + /// + public string? Link { get; set; } + + /// + /// Gets or sets the textual content of the retrieved chunk. + /// + public string Value { get; set; } = string.Empty; + + /// + /// Gets or sets the raw representation of the search result from the data source. + /// + /// + /// If a is created to represent some underlying object from another object + /// model, this property can be used to store that original object. This can be useful for debugging or + /// for enabling the to access the underlying object model if needed. + /// + public object? RawRepresentation { get; set; } + } + + internal sealed class TextSearchProviderState + { + public List? RecentMessagesText { get; set; } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI/Data/TextSearchProviderOptions.cs b/dotnet/src/Microsoft.Agents.AI/Data/TextSearchProviderOptions.cs new file mode 100644 index 0000000000..7709df435e --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI/Data/TextSearchProviderOptions.cs @@ -0,0 +1,77 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI.Data; + +/// +/// Options controlling the behavior of . +/// +public sealed class TextSearchProviderOptions +{ + /// + /// Gets or sets a value indicating when the search should be executed. + /// + /// by default. + public TextSearchBehavior SearchTime { get; set; } = TextSearchBehavior.BeforeAIInvoke; + + /// + /// Gets or sets the name of the exposed search tool when operating in on-demand mode. + /// + /// Defaults to "Search". + public string? FunctionToolName { get; set; } + + /// + /// Gets or sets the description of the exposed search tool when operating in on-demand mode. + /// + /// Defaults to "Allows searching for additional information to help answer the user question.". + public string? FunctionToolDescription { get; set; } + + /// + /// Gets or sets the context prompt prefixed to automatically injected results. + /// + public string? ContextPrompt { get; set; } + + /// + /// Gets or sets the instruction appended after automatically injected results to request citations. + /// + public string? CitationsPrompt { get; set; } + + /// + /// Optional delegate to fully customize formatting of the result list. + /// + /// + /// If provided, and are ignored. + /// + public Func, string>? ContextFormatter { get; set; } + + /// + /// Gets or sets the number of recent conversation messages (both user and assistant) to keep in memory + /// and include when constructing the search input for searches. + /// + /// + /// The maximum number of most recent messages to retain. A value of 0 (default) disables memory and + /// only the current request's messages are used for search input. The value is a count of individual + /// messages, not turns. Only messages with role or + /// are retained. + /// + public int RecentMessageMemoryLimit { get; set; } + + /// + /// Behavior choices for the provider. + /// + public enum TextSearchBehavior + { + /// + /// Execute search prior to each invocation and inject results as instructions. + /// + BeforeAIInvoke, + + /// + /// Expose a function tool to perform search on-demand via function/tool calling. + /// + OnDemandFunctionCalling + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs new file mode 100644 index 0000000000..fa5f33ac95 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs @@ -0,0 +1,542 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Agents.AI.Data; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.Logging; +using Moq; + +namespace Microsoft.Agents.AI.UnitTests.Data; + +/// +/// Contains unit tests for . +/// +public sealed class TextSearchProviderTests +{ + private readonly Mock> _loggerMock; + private readonly Mock _loggerFactoryMock; + + public TextSearchProviderTests() + { + this._loggerMock = new(); + this._loggerFactoryMock = new(); + this._loggerFactoryMock + .Setup(f => f.CreateLogger(It.IsAny())) + .Returns(this._loggerMock.Object); + this._loggerFactoryMock + .Setup(f => f.CreateLogger(typeof(TextSearchProvider).FullName!)) + .Returns(this._loggerMock.Object); + } + + [Theory] + [InlineData(null, null, true)] + [InlineData("Custom context prompt", "Custom citations prompt", false)] + public async Task InvokingAsync_ShouldInjectFormattedResultsAsync(string? overrideContextPrompt, string? overrideCitationsPrompt, bool withLogging) + { + // Arrange + List results = + [ + new() { Name = "Doc1", Link = "http://example.com/doc1", Value = "Content of Doc1" }, + new() { Name = "Doc2", Link = "http://example.com/doc2", Value = "Content of Doc2" } + ]; + + string? capturedInput = null; + Task> SearchDelegateAsync(string input, CancellationToken ct) + { + capturedInput = input; + return Task.FromResult>(results); + } + + var options = new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, + ContextPrompt = overrideContextPrompt, + CitationsPrompt = overrideCitationsPrompt + }; + var provider = new TextSearchProvider(SearchDelegateAsync, options, withLogging ? this._loggerFactoryMock.Object : null); + + var invokingContext = new AIContextProvider.InvokingContext(new[] + { + new ChatMessage(ChatRole.User, "Sample user question?"), + new ChatMessage(ChatRole.User, "Additional part") + }); + + // Act + var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert + Assert.Equal("Sample user question?\nAdditional part", capturedInput); + Assert.Null(aiContext.Instructions); // TextSearchProvider uses a user message for context injection. + Assert.NotNull(aiContext.Messages); + Assert.Single(aiContext.Messages!); + var message = aiContext.Messages!.Single(); + Assert.Equal(ChatRole.User, message.Role); + string text = message.Text!; + + if (overrideContextPrompt is null) + { + Assert.Contains("## Additional Context", text); + Assert.Contains("Consider the following information from source documents when responding to the user:", text); + } + else + { + Assert.Contains(overrideContextPrompt, text); + } + Assert.Contains("SourceDocName: Doc1", text); + Assert.Contains("SourceDocLink: http://example.com/doc1", text); + Assert.Contains("Contents: Content of Doc1", text); + Assert.Contains("SourceDocName: Doc2", text); + Assert.Contains("SourceDocLink: http://example.com/doc2", text); + Assert.Contains("Contents: Content of Doc2", text); + if (overrideCitationsPrompt is null) + { + Assert.Contains("Include citations to the source document with document name and link if document name and link is available.", text); + } + else + { + Assert.Contains(overrideCitationsPrompt, text); + } + + if (withLogging) + { + this._loggerMock.Verify( + l => l.Log( + LogLevel.Information, + It.IsAny(), + It.Is((v, t) => v.ToString()!.Contains("TextSearchProvider: Retrieved 2 search results.")), + It.IsAny(), + It.IsAny>()), + Times.AtLeastOnce); + this._loggerMock.Verify( + l => l.Log( + LogLevel.Trace, + It.IsAny(), + It.Is((v, t) => v.ToString()!.Contains("TextSearchProvider Input:Sample user question?\nAdditional part\nContext Instructions")), + It.IsAny(), + It.IsAny>()), + Times.AtLeastOnce); + } + } + + [Theory] + [InlineData(null, null, "Search", "Allows searching for additional information to help answer the user question.")] + [InlineData("CustomSearch", "CustomDescription", "CustomSearch", "CustomDescription")] + public async Task InvokingAsync_OnDemand_ShouldExposeSearchToolAsync(string? overrideName, string? overrideDescription, string expectedName, string expectedDescription) + { + // Arrange + var options = new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.OnDemandFunctionCalling, + FunctionToolName = overrideName, + FunctionToolDescription = overrideDescription + }; + var provider = new TextSearchProvider(this.NoResultSearchAsync, options); + var invokingContext = new AIContextProvider.InvokingContext(new[] { new ChatMessage(ChatRole.User, "Q?") }); + + // Act + var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert + Assert.Null(aiContext.Messages); // No automatic injection. + Assert.NotNull(aiContext.Tools); + Assert.Single(aiContext.Tools); + var tool = aiContext.Tools.Single(); + Assert.Equal(expectedName, tool.Name); + Assert.Equal(expectedDescription, tool.Description); + } + + [Theory] + [InlineData(null, null)] + [InlineData("Custom context prompt", "Custom citations prompt")] + public async Task SearchAsync_ShouldReturnFormattedResultsAsync(string? overrideContextPrompt, string? overrideCitationsPrompt) + { + // Arrange + List results = + [ + new() { Name = "Doc1", Link = "http://example.com/doc1", Value = "Content of Doc1" }, + new() { Name = "Doc2", Link = "http://example.com/doc2", Value = "Content of Doc2" } + ]; + + Task> SearchDelegateAsync(string input, CancellationToken ct) + { + return Task.FromResult>(results); + } + + var options = new TextSearchProviderOptions + { + ContextPrompt = overrideContextPrompt, + CitationsPrompt = overrideCitationsPrompt + }; + var provider = new TextSearchProvider(SearchDelegateAsync, options); + + // Act + var formatted = await provider.SearchAsync("Sample user question?", CancellationToken.None); + + // Assert + if (overrideContextPrompt is null) + { + Assert.Contains("## Additional Context", formatted); + Assert.Contains("Consider the following information from source documents when responding to the user:", formatted); + } + else + { + Assert.Contains(overrideContextPrompt, formatted); + } + + Assert.Contains("SourceDocName: Doc1", formatted); + Assert.Contains("SourceDocLink: http://example.com/doc1", formatted); + Assert.Contains("Contents: Content of Doc1", formatted); + Assert.Contains("SourceDocName: Doc2", formatted); + Assert.Contains("SourceDocLink: http://example.com/doc2", formatted); + Assert.Contains("Contents: Content of Doc2", formatted); + if (overrideCitationsPrompt is null) + { + Assert.Contains("Include citations to the source document with document name and link if document name and link is available.", formatted); + } + else + { + Assert.Contains(overrideCitationsPrompt, formatted); + } + } + + [Fact] + public async Task InvokingAsync_ShouldUseContextFormatterWhenProvidedAsync() + { + // Arrange + List results = + [ + new() { Name = "Doc1", Link = "http://example.com/doc1", Value = "Content of Doc1" }, + new() { Name = "Doc2", Link = "http://example.com/doc2", Value = "Content of Doc2" } + ]; + + Task> SearchDelegateAsync(string input, CancellationToken ct) + { + return Task.FromResult>(results); + } + + var options = new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, + ContextFormatter = r => $"Custom formatted context with {r.Count} results." + }; + var provider = new TextSearchProvider(SearchDelegateAsync, options); + var invokingContext = new AIContextProvider.InvokingContext(new[] { new ChatMessage(ChatRole.User, "Q?") }); + + // Act + var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert + Assert.NotNull(aiContext.Messages); + Assert.Single(aiContext.Messages!); + Assert.Equal("Custom formatted context with 2 results.", aiContext.Messages![0].Text); + } + + [Fact] + public async Task InvokingAsync_WithRawRepresentations_ContextFormatterCanAccessAsync() + { + // Arrange + var payload1 = new RawPayload { Id = "R1" }; + var payload2 = new RawPayload { Id = "R2" }; + List results = + [ + new() { Name = "Doc1", Value = "Content 1", RawRepresentation = payload1 }, + new() { Name = "Doc2", Value = "Content 2", RawRepresentation = payload2 } + ]; + + Task> SearchDelegateAsync(string input, CancellationToken ct) + { + return Task.FromResult>(results); + } + + var options = new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, + ContextFormatter = r => string.Join(",", r.Select(x => ((RawPayload)x.RawRepresentation!).Id)) + }; + var provider = new TextSearchProvider(SearchDelegateAsync, options); + var invokingContext = new AIContextProvider.InvokingContext(new[] { new ChatMessage(ChatRole.User, "Q?") }); + + // Act + var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert + Assert.NotNull(aiContext.Messages); + Assert.Single(aiContext.Messages!); + Assert.Equal("R1,R2", aiContext.Messages![0].Text); + } + + [Fact] + public async Task InvokingAsync_WithNoResults_ShouldReturnEmptyContextAsync() + { + // Arrange + var options = new TextSearchProviderOptions { SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke }; + var provider = new TextSearchProvider(this.NoResultSearchAsync, options); + var invokingContext = new AIContextProvider.InvokingContext(new[] { new ChatMessage(ChatRole.User, "Q?") }); + + // Act + var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert + Assert.Null(aiContext.Messages); + Assert.Null(aiContext.Instructions); + Assert.Null(aiContext.Tools); + } + + #region Recent Message Memory Tests + + [Fact] + public async Task InvokingAsync_WithRecentMessageMemory_ShouldIncludeStoredMessagesInSearchInputAsync() + { + // Arrange + var options = new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, + RecentMessageMemoryLimit = 3 + }; + string? capturedInput = null; + Task> SearchDelegateAsync(string input, CancellationToken ct) + { + capturedInput = input; + return Task.FromResult>([]); // No results needed. + } + var provider = new TextSearchProvider(SearchDelegateAsync, options); + + // Populate memory with more messages than the limit (A,B,C,D) -> should retain B,C,D + var initialMessages = new[] + { + new ChatMessage(ChatRole.User, "A"), + new ChatMessage(ChatRole.Assistant, "B"), + new ChatMessage(ChatRole.User, "C"), + new ChatMessage(ChatRole.Assistant, "D"), + }; + await provider.InvokedAsync(new(initialMessages)); + + var invokingContext = new AIContextProvider.InvokingContext(new[] + { + new ChatMessage(ChatRole.User, "E") + }); + + // Act + await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert + Assert.Equal("B\nC\nD\nE", capturedInput); // Memory first (truncated) then current request. + } + + [Fact] + public async Task InvokingAsync_WithAccumulatedMemoryAcrossInvocations_ShouldIncludeAllUpToLimitAsync() + { + // Arrange + var options = new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, + RecentMessageMemoryLimit = 5 + }; + string? capturedInput = null; + Task> SearchDelegateAsync(string input, CancellationToken ct) + { + capturedInput = input; + return Task.FromResult>([]); + } + var provider = new TextSearchProvider(SearchDelegateAsync, options); + + // First memory update (A,B) + await provider.InvokedAsync(new(new[] + { + new ChatMessage(ChatRole.User, "A"), + new ChatMessage(ChatRole.Assistant, "B"), + })); + + // Second memory update (C,D,E) + await provider.InvokedAsync(new(new[] + { + new ChatMessage(ChatRole.User, "C"), + new ChatMessage(ChatRole.Assistant, "D"), + new ChatMessage(ChatRole.User, "E"), + })); + + var invokingContext = new AIContextProvider.InvokingContext(new[] { new ChatMessage(ChatRole.User, "F") }); + + // Act + await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert + Assert.Equal("A\nB\nC\nD\nE\nF", capturedInput); // All retained (limit 5) + current request message. + } + + #endregion + + #region Serialization Tests + + [Fact] + public void Serialize_WithNoRecentMessages_ShouldReturnEmptyState() + { + // Arrange + var options = new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, + RecentMessageMemoryLimit = 3 + }; + var provider = new TextSearchProvider(this.NoResultSearchAsync, options); + + // Act + var state = provider.Serialize(); + + // Assert + Assert.Equal(JsonValueKind.Object, state.ValueKind); + Assert.False(state.TryGetProperty("recentMessagesText", out _)); + } + + [Fact] + public async Task Serialize_WithRecentMessages_ShouldPersistMessagesUpToLimitAsync() + { + // Arrange + var options = new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, + RecentMessageMemoryLimit = 3 + }; + var provider = new TextSearchProvider(this.NoResultSearchAsync, options); + var messages = new[] + { + new ChatMessage(ChatRole.User, "M1"), + new ChatMessage(ChatRole.Assistant, "M2"), + new ChatMessage(ChatRole.User, "M3"), + }; + + // Act + await provider.InvokedAsync(new(messages)); // Populate recent memory. + var state = provider.Serialize(); + + // Assert + Assert.True(state.TryGetProperty("recentMessagesText", out var recentProperty)); + Assert.Equal(JsonValueKind.Array, recentProperty.ValueKind); + var list = recentProperty.EnumerateArray().Select(e => e.GetString()).ToList(); + Assert.Equal(3, list.Count); + Assert.Equal(["M1", "M2", "M3"], list); + } + + [Fact] + public async Task SerializeAndDeserialize_RoundtripRestoresMessagesAsync() + { + // Arrange + var options = new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, + RecentMessageMemoryLimit = 4 + }; + var provider = new TextSearchProvider(this.NoResultSearchAsync, options); + var messages = new[] + { + new ChatMessage(ChatRole.User, "A"), + new ChatMessage(ChatRole.Assistant, "B"), + new ChatMessage(ChatRole.User, "C"), + new ChatMessage(ChatRole.Assistant, "D"), + }; + await provider.InvokedAsync(new(messages)); + + // Act + var state = provider.Serialize(); + string? capturedInput = null; + Task> SearchDelegate2Async(string input, CancellationToken ct) + { + capturedInput = input; + return Task.FromResult>([]); + } + var roundTrippedProvider = new TextSearchProvider(SearchDelegate2Async, state, options: new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, + RecentMessageMemoryLimit = 4 + }); + var emptyMessages = Array.Empty(); + await roundTrippedProvider.InvokingAsync(new(emptyMessages), CancellationToken.None); // Trigger search to read memory. + + // Assert + Assert.NotNull(capturedInput); + Assert.Equal("A\nB\nC\nD", capturedInput); + } + + [Fact] + public async Task Deserialize_WithChangedLowerLimit_ShouldTruncateToNewLimitAsync() + { + // Arrange + var initialProvider = new TextSearchProvider(this.NoResultSearchAsync, new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, + RecentMessageMemoryLimit = 5 + }); + var messages = new[] + { + new ChatMessage(ChatRole.User, "L1"), + new ChatMessage(ChatRole.Assistant, "L2"), + new ChatMessage(ChatRole.User, "L3"), + new ChatMessage(ChatRole.Assistant, "L4"), + new ChatMessage(ChatRole.User, "L5"), + }; + await initialProvider.InvokedAsync(new(messages)); + var state = initialProvider.Serialize(); + + string? capturedInput = null; + Task> SearchDelegate2Async(string input, CancellationToken ct) + { + capturedInput = input; + return Task.FromResult>([]); + } + + // Act + var restoredProvider = new TextSearchProvider(SearchDelegate2Async, state, options: new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, + RecentMessageMemoryLimit = 3 // Lower limit + }); + await restoredProvider.InvokingAsync(new(Array.Empty()), CancellationToken.None); + + // Assert + Assert.NotNull(capturedInput); + Assert.Equal("L1\nL2\nL3", capturedInput); + } + + [Fact] + public async Task Deserialize_WithEmptyState_ShouldHaveNoMessagesAsync() + { + // Arrange + var emptyState = JsonSerializer.Deserialize("{}", TestJsonSerializerContext.Default.JsonElement); + + string? capturedInput = null; + Task> SearchDelegate2Async(string input, CancellationToken ct) + { + capturedInput = input; + return Task.FromResult>([]); + } + + // Act + var provider = new TextSearchProvider(SearchDelegate2Async, emptyState, options: new TextSearchProviderOptions + { + SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, + RecentMessageMemoryLimit = 3 + }); + var emptyMessages = Array.Empty(); + await provider.InvokingAsync(new(emptyMessages), CancellationToken.None); + + // Assert + Assert.NotNull(capturedInput); + Assert.Equal(string.Empty, capturedInput); // No recent messages serialized => empty input. + } + + #endregion + + private Task> NoResultSearchAsync(string input, CancellationToken ct) + { + return Task.FromResult>([]); + } + + private sealed class RawPayload + { + public string Id { get; set; } = string.Empty; + } +}