.NET: Add serialization support for ContextualFunctionProvider (#2601)

* Initial plan

* Add serialization support for ContextualFunctionProvider

- Added constructor overload that accepts JsonElement serializedState
- Implemented Serialize method to persist recent messages
- Added ContextualFunctionProviderState class
- Registered state class with AgentJsonUtilities JsonContext
- Added comprehensive unit tests for serialization/deserialization

Co-authored-by: westey-m <164392973+westey-m@users.noreply.github.com>

* Fix constructor parameter ordering and make default more explicit

- Reordered constructor parameters to maintain consistency (options before jsonSerializerOptions)
- Made default JsonElement explicit with default(JsonElement) in constructor chaining
- All tests continue to pass

Co-authored-by: westey-m <164392973+westey-m@users.noreply.github.com>

* Improve XML documentation comments for clarity

- Updated jsonSerializerOptions parameter comments to clarify actual usage
- Made it clear that AgentJsonUtilities.DefaultOptions is used
- All tests continue to pass

Co-authored-by: westey-m <164392973+westey-m@users.noreply.github.com>

* Change NumberOfRecentMessagesInContext to 2 to demonstrate truncation

- Updated test to use limit of 2 messages instead of 3
- Now test shows truncation in action (3 messages added, only 2 kept)
- Updated assertion to expect 2 messages in serialized state
- All tests pass

Co-authored-by: westey-m <164392973+westey-m@users.noreply.github.com>

* Apply suggestion from @westey-m

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: westey-m <164392973+westey-m@users.noreply.github.com>
This commit is contained in:
Copilot
2025-12-03 12:18:19 +00:00
committed by GitHub
co-authored by westey-m copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
parent 9008bb8cb7
commit 7965285808
3 changed files with 246 additions and 2 deletions
@@ -69,6 +69,7 @@ internal static partial class AgentJsonUtilities
[JsonSerializable(typeof(ChatClientAgentThread.ThreadState))]
[JsonSerializable(typeof(TextSearchProvider.TextSearchProviderState))]
[JsonSerializable(typeof(ChatHistoryMemoryProvider.ChatHistoryMemoryProviderState))]
[JsonSerializable(typeof(Functions.ContextualFunctionProvider.ContextualFunctionProviderState))]
[ExcludeFromCodeCoverage]
internal sealed partial class JsonContext : JsonSerializerContext;
@@ -4,6 +4,7 @@ using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
@@ -60,6 +61,30 @@ public sealed class ContextualFunctionProvider : AIContextProvider
int maxNumberOfFunctions,
ContextualFunctionProviderOptions? options = null,
ILoggerFactory? loggerFactory = null)
: this(vectorStore, vectorDimensions, functions, maxNumberOfFunctions, default(JsonElement), options, null, loggerFactory)
{
}
/// <summary>
/// Initializes a new instance of the <see cref="ContextualFunctionProvider"/> class.
/// </summary>
/// <param name="vectorStore">An instance of a vector store.</param>
/// <param name="vectorDimensions">The number of dimensions to use for the memory embeddings.</param>
/// <param name="functions">The functions to vectorize and store for searching related functions.</param>
/// <param name="maxNumberOfFunctions">The maximum number of relevant functions to retrieve from the vector store.</param>
/// <param name="serializedState">A <see cref="JsonElement"/> representing the serialized provider state.</param>
/// <param name="options">Further optional settings for configuring the provider.</param>
/// <param name="jsonSerializerOptions">Optional serializer options. If not provided, <see cref="AgentJsonUtilities.DefaultOptions"/> will be used.</param>
/// <param name="loggerFactory">The logger factory to use for logging. If not provided, no logging will be performed.</param>
public ContextualFunctionProvider(
VectorStore vectorStore,
int vectorDimensions,
IEnumerable<AIFunction> functions,
int maxNumberOfFunctions,
JsonElement serializedState,
ContextualFunctionProviderOptions? options = null,
JsonSerializerOptions? jsonSerializerOptions = null,
ILoggerFactory? loggerFactory = null)
{
Throw.IfNull(vectorStore);
Throw.IfLessThan(vectorDimensions, 1, "Vector dimensions must be greater than 0");
@@ -81,6 +106,21 @@ public sealed class ContextualFunctionProvider : AIContextProvider
EmbeddingValueProvider = this._options.EmbeddingValueProvider,
}
);
// Restore recent messages from serialized state if provided
if (serializedState.ValueKind is not JsonValueKind.Null and not JsonValueKind.Undefined)
{
JsonSerializerOptions jso = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions;
ContextualFunctionProviderState? state = serializedState.Deserialize(jso.GetTypeInfo(typeof(ContextualFunctionProviderState))) as ContextualFunctionProviderState;
if (state?.RecentMessages is { Count: > 0 })
{
// Restore recent messages respecting the limit (may truncate if limit changed afterwards).
foreach (ChatMessage message in state.RecentMessages.Take(this._options.NumberOfRecentMessagesInContext))
{
this._recentMessages.Enqueue(message);
}
}
}
}
/// <inheritdoc />
@@ -141,6 +181,22 @@ public sealed class ContextualFunctionProvider : AIContextProvider
return default;
}
/// <summary>
/// Serializes the current provider state to a <see cref="JsonElement"/> containing the recent messages.
/// </summary>
/// <param name="jsonSerializerOptions">Optional serializer options. This parameter is not used; <see cref="AgentJsonUtilities.DefaultOptions"/> is always used for serialization.</param>
/// <returns>A <see cref="JsonElement"/> with the recent messages, or default if there are no recent messages.</returns>
public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
{
ContextualFunctionProviderState state = new();
if (this._options.NumberOfRecentMessagesInContext > 0 && !this._recentMessages.IsEmpty)
{
state.RecentMessages = this._recentMessages.Take(this._options.NumberOfRecentMessagesInContext).ToList();
}
return JsonSerializer.SerializeToElement(state, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ContextualFunctionProviderState)));
}
/// <summary>
/// Builds the context from chat messages.
/// </summary>
@@ -166,4 +222,9 @@ public sealed class ContextualFunctionProvider : AIContextProvider
.Where(m => !string.IsNullOrWhiteSpace(m?.Text))
.Select(m => m.Text));
}
internal sealed class ContextualFunctionProviderState
{
public List<ChatMessage>? RecentMessages { get; set; }
}
}
@@ -3,6 +3,7 @@
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Agents.AI.Functions;
@@ -305,7 +306,7 @@ public sealed class ContextualFunctionProviderTests
}
[Fact]
public async Task InvokedAsync_ShouldNotAddMessages_WhenExceptionIsPresent_Async()
public void Serialize_WithNoRecentMessages_ShouldReturnEmptyState()
{
// Arrange
var functions = new List<AIFunction> { CreateFunction("f1") };
@@ -320,7 +321,6 @@ public sealed class ContextualFunctionProviderTests
functions: functions,
maxNumberOfFunctions: 5,
options: options);
var message1 = new ChatMessage() { Contents = [new TextContent("msg1")] };
var message2 = new ChatMessage() { Contents = [new TextContent("msg2")] };
var message3 = new ChatMessage() { Contents = [new TextContent("msg3")] };
@@ -343,6 +343,188 @@ public sealed class ContextualFunctionProviderTests
var expected = string.Join(Environment.NewLine, ["msg1", "msg2", "new message"]);
this._collectionMock.Verify(c => c.SearchAsync(expected, It.IsAny<int>(), null, It.IsAny<CancellationToken>()), Times.Once);
}
[Fact]
public async Task InvokedAsync_ShouldNotAddMessages_WhenExceptionIsPresent_Async()
{
// Arrange
var functions = new List<AIFunction> { CreateFunction("f1") };
var options = new ContextualFunctionProviderOptions
{
NumberOfRecentMessagesInContext = 3
};
var provider = new ContextualFunctionProvider(
vectorStore: this._vectorStoreMock.Object,
vectorDimensions: 1536,
functions: functions,
maxNumberOfFunctions: 5,
options: options);
// Act
JsonElement state = provider.Serialize();
// Assert
Assert.Equal(JsonValueKind.Object, state.ValueKind);
Assert.False(state.TryGetProperty("recentMessages", out _));
}
[Fact]
public async Task Serialize_WithRecentMessages_ShouldPersistMessagesUpToLimitAsync()
{
// Arrange
var functions = new List<AIFunction> { CreateFunction("f1") };
var options = new ContextualFunctionProviderOptions
{
NumberOfRecentMessagesInContext = 2
};
var provider = new ContextualFunctionProvider(
vectorStore: this._vectorStoreMock.Object,
vectorDimensions: 1536,
functions: functions,
maxNumberOfFunctions: 5,
options: options);
var messages = new[]
{
new ChatMessage() { Contents = [new TextContent("M1")] },
new ChatMessage() { Contents = [new TextContent("M2")] },
new ChatMessage() { Contents = [new TextContent("M3")] }
};
// Act
await provider.InvokedAsync(new AIContextProvider.InvokedContext(messages, aiContextProviderMessages: null));
JsonElement state = provider.Serialize();
// Assert
Assert.True(state.TryGetProperty("recentMessages", out JsonElement recentProperty));
Assert.Equal(JsonValueKind.Array, recentProperty.ValueKind);
int count = recentProperty.GetArrayLength();
Assert.Equal(2, count);
}
[Fact]
public async Task SerializeAndDeserialize_RoundtripRestoresMessagesAsync()
{
// Arrange
var functions = new List<AIFunction> { CreateFunction("f1") };
var options = new ContextualFunctionProviderOptions
{
NumberOfRecentMessagesInContext = 4
};
var provider = new ContextualFunctionProvider(
vectorStore: this._vectorStoreMock.Object,
vectorDimensions: 1536,
functions: functions,
maxNumberOfFunctions: 5,
options: options);
var messages = new[]
{
new ChatMessage() { Contents = [new TextContent("A")] },
new ChatMessage() { Contents = [new TextContent("B")] },
new ChatMessage() { Contents = [new TextContent("C")] },
new ChatMessage() { Contents = [new TextContent("D")] }
};
await provider.InvokedAsync(new AIContextProvider.InvokedContext(messages, aiContextProviderMessages: null));
// Act
JsonElement state = provider.Serialize();
var roundTrippedProvider = new ContextualFunctionProvider(
vectorStore: this._vectorStoreMock.Object,
vectorDimensions: 1536,
functions: functions,
maxNumberOfFunctions: 5,
serializedState: state,
options: new ContextualFunctionProviderOptions
{
NumberOfRecentMessagesInContext = 4
});
// Trigger search to verify messages are used
var invokingContext = new AIContextProvider.InvokingContext(Array.Empty<ChatMessage>());
await roundTrippedProvider.InvokingAsync(invokingContext);
// Assert
string expected = string.Join(Environment.NewLine, ["A", "B", "C", "D"]);
this._collectionMock.Verify(c => c.SearchAsync(expected, It.IsAny<int>(), null, It.IsAny<CancellationToken>()), Times.Once);
}
[Fact]
public async Task Deserialize_WithChangedLowerLimit_ShouldTruncateToNewLimitAsync()
{
// Arrange
var functions = new List<AIFunction> { CreateFunction("f1") };
var initialProvider = new ContextualFunctionProvider(
vectorStore: this._vectorStoreMock.Object,
vectorDimensions: 1536,
functions: functions,
maxNumberOfFunctions: 5,
options: new ContextualFunctionProviderOptions
{
NumberOfRecentMessagesInContext = 5
});
var messages = new[]
{
new ChatMessage() { Contents = [new TextContent("L1")] },
new ChatMessage() { Contents = [new TextContent("L2")] },
new ChatMessage() { Contents = [new TextContent("L3")] },
new ChatMessage() { Contents = [new TextContent("L4")] },
new ChatMessage() { Contents = [new TextContent("L5")] }
};
await initialProvider.InvokedAsync(new AIContextProvider.InvokedContext(messages, aiContextProviderMessages: null));
JsonElement state = initialProvider.Serialize();
// Act
var restoredProvider = new ContextualFunctionProvider(
vectorStore: this._vectorStoreMock.Object,
vectorDimensions: 1536,
functions: functions,
maxNumberOfFunctions: 5,
serializedState: state,
options: new ContextualFunctionProviderOptions
{
NumberOfRecentMessagesInContext = 3 // Lower limit
});
var invokingContext = new AIContextProvider.InvokingContext(Array.Empty<ChatMessage>());
await restoredProvider.InvokingAsync(invokingContext);
// Assert
string expected = string.Join(Environment.NewLine, ["L1", "L2", "L3"]);
this._collectionMock.Verify(c => c.SearchAsync(expected, It.IsAny<int>(), null, It.IsAny<CancellationToken>()), Times.Once);
}
[Fact]
public async Task Deserialize_WithEmptyState_ShouldHaveNoMessagesAsync()
{
// Arrange
var functions = new List<AIFunction> { CreateFunction("f1") };
JsonElement emptyState = JsonSerializer.Deserialize("{}", TestJsonSerializerContext.Default.JsonElement);
// Act
var provider = new ContextualFunctionProvider(
vectorStore: this._vectorStoreMock.Object,
vectorDimensions: 1536,
functions: functions,
maxNumberOfFunctions: 5,
serializedState: emptyState,
options: new ContextualFunctionProviderOptions
{
NumberOfRecentMessagesInContext = 3
});
var invokingContext = new AIContextProvider.InvokingContext(Array.Empty<ChatMessage>());
await provider.InvokingAsync(invokingContext);
// Assert
this._collectionMock.Verify(c => c.SearchAsync(string.Empty, It.IsAny<int>(), null, It.IsAny<CancellationToken>()), Times.Once);
}
private static AIFunction CreateFunction(string name, string description = "")
{