.NET: Update AgentThread, MessageStores and Context Providers to deserialize via constructor. (#779)

* Update AgentThread, MessageStores and Context Providers to deserialize via constructor.

* Fix pr comment.

* Add additional validation for AgentThread deserialization

* Update WorkflowMessageStore desreialize to improve error checking.

* Reduce allocations in InMemoryChatMessageStore

---------

Co-authored-by: Chris <66376200+crickman@users.noreply.github.com>
This commit is contained in:
westey
2025-09-17 11:08:13 +01:00
committed by GitHub
Unverified
parent 3050c8c8bd
commit 4247fc26dd
22 changed files with 234 additions and 195 deletions
@@ -40,7 +40,7 @@ await File.WriteAllTextAsync(tempFilePath, JsonSerializer.Serialize(serializedTh
JsonElement reloadedSerializedThread = JsonSerializer.Deserialize<JsonElement>(await File.ReadAllTextAsync(tempFilePath));
// Deserialize the thread state after loading from storage.
AgentThread resumedThread = await agent.DeserializeThreadAsync(reloadedSerializedThread);
AgentThread resumedThread = agent.DeserializeThread(reloadedSerializedThread);
// Run the agent again with the resumed thread.
Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
@@ -38,12 +38,12 @@ AIAgent agent = new AzureOpenAIClient(
{
Name = JokerName,
Instructions = JokerInstructions,
ChatMessageStoreFactory = () =>
ChatMessageStoreFactory = (jsonElement, jso) =>
{
// Create a new chat message store for this agent that stores the messages in a vector store.
// Each thread must get its own copy of the VectorChatMessageStore, since the store
// also contains the id that the thread is stored under.
return new VectorChatMessageStore(vectorStore);
return new VectorChatMessageStore(vectorStore, jsonElement, jso);
}
});
@@ -65,7 +65,7 @@ Console.WriteLine(JsonSerializer.Serialize(serializedThread, new JsonSerializerO
// and loaded again later.
// Deserialize the thread state after loading from storage.
AgentThread resumedThread = await agent.DeserializeThreadAsync(serializedThread);
AgentThread resumedThread = agent.DeserializeThread(serializedThread);
// Run the agent with the thread that stores conversation history in the vector store a second time.
Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedThread));
@@ -75,10 +75,21 @@ namespace SampleApp
/// <summary>
/// A sample implementation of <see cref="IChatMessageStore"/> that stores chat messages in a vector store.
/// </summary>
/// <param name="vectorStore">The vector store to store the messages in.</param>
internal sealed class VectorChatMessageStore(VectorStore vectorStore) : IChatMessageStore
internal sealed class VectorChatMessageStore : IChatMessageStore
{
private string? _threadId;
private readonly VectorStore _vectorStore;
public VectorChatMessageStore(VectorStore vectorStore, JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null)
{
this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
if (serializedStoreState.ValueKind == JsonValueKind.String)
{
// Here we can deserialize the thread id so that we can access the same messages as before the suspension.
this._threadId = JsonSerializer.Deserialize<string>(serializedStoreState);
}
}
public string? ThreadId => this._threadId;
@@ -86,7 +97,7 @@ namespace SampleApp
{
this._threadId ??= Guid.NewGuid().ToString();
var collection = vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
await collection.EnsureCollectionExistsAsync(cancellationToken);
await collection.UpsertAsync(messages.Select(x => new ChatHistoryItem()
@@ -101,7 +112,7 @@ namespace SampleApp
public async Task<IEnumerable<ChatMessage>> GetMessagesAsync(CancellationToken cancellationToken)
{
var collection = vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
var collection = this._vectorStore.GetCollection<string, ChatHistoryItem>("ChatHistory");
await collection.EnsureCollectionExistsAsync(cancellationToken);
var records = await collection
@@ -124,13 +135,6 @@ namespace SampleApp
return new ValueTask<JsonElement?>(JsonSerializer.SerializeToElement(this._threadId));
}
public ValueTask DeserializeStateAsync(JsonElement? serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
{
// Here we can deserialize the thread id so that we can access the same messages as before the suspension.
this._threadId = JsonSerializer.Deserialize<string>((JsonElement)serializedStoreState!);
return new ValueTask();
}
/// <summary>
/// The data structure used to store chat history items in the vector store.
/// </summary>
@@ -38,7 +38,7 @@ ChatClient chatClient = new AzureOpenAIClient(
AIAgent agent = chatClient.CreateAIAgent(new ChatClientAgentOptions()
{
Instructions = "You are a friendly assistant. Always address the user by their name.",
AIContextProviderFactory = () => new SampleApp.UserInfoMemory(chatClient.AsIChatClient())
AIContextProviderFactory = (jse, jso) => new UserInfoMemory(chatClient.AsIChatClient(), jse, jso)
});
// Create a new thread for the conversation.
@@ -57,7 +57,7 @@ var threadElement = await thread.SerializeAsync();
Console.WriteLine("\n>> Use deserialized thread with previously created memories\n");
// Later we can deserialize the thread and continue the conversation with the previous memory component state.
var deserializedThread = await agent.DeserializeThreadAsync(threadElement);
var deserializedThread = agent.DeserializeThread(threadElement);
Console.WriteLine(await agent.RunAsync("What is my name and age?", deserializedThread));
Console.WriteLine("\n>> Read memories from memory component\n");
@@ -89,16 +89,30 @@ namespace SampleApp
/// <summary>
/// Sample memory component that can remember a user's name and age.
/// </summary>
internal sealed class UserInfoMemory(IChatClient chatClient, UserInfo? userInfo = null) : AIContextProvider
internal sealed class UserInfoMemory : AIContextProvider
{
public UserInfo UserInfo { get; set; } = userInfo ?? new();
private readonly IChatClient _chatClient;
public UserInfoMemory(IChatClient chatClient, UserInfo? userInfo = null)
{
this._chatClient = chatClient;
this.UserInfo = userInfo ?? new UserInfo();
}
public UserInfoMemory(IChatClient chatClient, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null)
{
this._chatClient = chatClient;
this.UserInfo = JsonSerializer.Deserialize<UserInfo>(serializedState, jsonSerializerOptions) ?? new UserInfo();
}
public UserInfo UserInfo { get; set; }
public override async ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default)
{
// Try and extract the user name and age from the message if we don't have it already and it's a user message.
if ((this.UserInfo.UserName == null || this.UserInfo.UserAge == null) && context.RequestMessages.Any(x => x.Role == ChatRole.User))
{
var result = await chatClient.GetResponseAsync<UserInfo>(
var result = await this._chatClient.GetResponseAsync<UserInfo>(
context.RequestMessages,
new ChatOptions()
{
@@ -77,8 +77,7 @@ internal class AIAgentHostExecutor : Executor
JsonElement? threadValue = await context.ReadStateAsync<JsonElement?>(ThreadStateKey).ConfigureAwait(false);
if (threadValue.HasValue)
{
this._thread = await this._agent.DeserializeThreadAsync(threadValue.Value, cancellationToken: cancellation)
.ConfigureAwait(false);
this._thread = this._agent.DeserializeThread(threadValue.Value, cancellationToken: cancellation);
}
JsonElement? messagesValue = await context.ReadStateAsync<JsonElement?>(PendingMessagesStateKey).ConfigureAwait(false);
@@ -6,6 +6,7 @@ using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
@@ -47,6 +48,9 @@ internal class WorkflowHostAgent : AIAgent
public override AgentThread GetNewThread() => new WorkflowThread(this.Id, this.Name, this.GenerateNewId());
public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
=> new WorkflowThread(serializedThread, jsonSerializerOptions);
private async
IAsyncEnumerable<AgentRunResponseUpdate> InvokeStageAsync(
WorkflowThread conversation,
@@ -1,5 +1,6 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Text.Json;
using System.Threading;
@@ -14,6 +15,30 @@ internal class WorkflowMessageStore : IChatMessageStore
private int _bookmark = 0;
private readonly List<ChatMessage> _chatMessages = new();
public WorkflowMessageStore()
{
}
public WorkflowMessageStore(JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null)
{
if (serializedStoreState.ValueKind != JsonValueKind.Object)
{
throw new ArgumentException("The provided JsonElement must be a json object", nameof(serializedStoreState));
}
StoreState? state =
JsonSerializer.Deserialize(
serializedStoreState,
AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))) as StoreState;
if (state?.Messages is not null)
{
this._chatMessages.AddRange(state.Messages);
}
this._bookmark = state?.Bookmark ?? 0;
}
internal class StoreState
{
public int Bookmark { get; set; }
@@ -50,31 +75,6 @@ internal class WorkflowMessageStore : IChatMessageStore
this._bookmark = this._chatMessages.Count;
}
public ValueTask DeserializeStateAsync(JsonElement? serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
{
if (serializedStoreState is null)
{
return default;
}
object? maybeState =
JsonSerializer.Deserialize(
serializedStoreState.Value,
AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState)));
if (maybeState is not StoreState state)
{
throw new JsonException("Invalid state format for WorkflowMessageStore.");
}
this._chatMessages.Clear();
this._chatMessages.AddRange(state.Messages);
this._bookmark = state.Bookmark;
return default;
}
public ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
{
StoreState state = new()
@@ -25,6 +25,11 @@ internal class WorkflowThread : AgentThread
this._workflowName = workflowName;
}
public WorkflowThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null)
{
throw new NotImplementedException("Pending Checkpointing work.");
}
public string RunId { get; }
public int Halts { get; } = 0;
@@ -35,11 +40,6 @@ internal class WorkflowThread : AgentThread
throw new NotImplementedException("Pending Checkpointing work.");
}
protected override Task DeserializeAsync(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
{
throw new NotImplementedException("Pending Checkpointing work.");
}
public AgentRunResponseUpdate CreateUpdate(params AIContent[] parts)
{
Throw.IfNullOrEmpty(parts);
@@ -90,12 +90,8 @@ public abstract class AIAgent
/// <param name="jsonSerializerOptions">Optional <see cref="JsonSerializerOptions"/> to use for deserializing the thread state.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>The deserialized <see cref="AgentThread"/> instance.</returns>
public async ValueTask<AgentThread> DeserializeThreadAsync(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
{
var thread = this.GetNewThread();
await thread.DeserializeAsync(serializedThread, jsonSerializerOptions, cancellationToken).ConfigureAwait(false);
return thread;
}
public virtual AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
=> new(serializedThread, jsonSerializerOptions);
/// <summary>
/// Run the agent with no message assuming that all required instructions are already provided to the agent or on the thread.
@@ -26,6 +26,46 @@ public class AgentThread
{
}
/// <summary>
/// Initializes a new instance of the <see cref="AgentThread"/> class from serialized state.
/// </summary>
/// <param name="serializedThreadState">A <see cref="JsonElement"/> representing the serialized state of the thread.</param>
/// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
/// <param name="chatMessageStoreFactory">An optional factory function to create a custom <see cref="IChatMessageStore"/>.</param>
/// <param name="aiContextProviderFactory">An optional factory function to create a custom <see cref="AIContextProvider"/>.</param>
public AgentThread(
JsonElement serializedThreadState,
JsonSerializerOptions? jsonSerializerOptions = null,
Func<JsonElement, JsonSerializerOptions?, IChatMessageStore>? chatMessageStoreFactory = null,
Func<JsonElement, JsonSerializerOptions?, AIContextProvider>? aiContextProviderFactory = null)
{
if (serializedThreadState.ValueKind != JsonValueKind.Object)
{
throw new ArgumentException("The serialized thread state must be a JSON object.", nameof(serializedThreadState));
}
var state = JsonSerializer.Deserialize(
serializedThreadState,
AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ThreadState))) as ThreadState;
this.AIContextProvider = aiContextProviderFactory?.Invoke(state?.AIContextProviderState ?? default, jsonSerializerOptions);
if (state?.ConversationId is string threadId)
{
this.ConversationId = threadId;
// Since we have an ID, we should not have a chat message store and we can return here.
return;
}
this._messageStore = chatMessageStoreFactory?.Invoke(state?.StoreState ?? default, jsonSerializerOptions);
if (this._messageStore is null)
{
// If we didn't get a custom store, create an in-memory one.
this._messageStore = new InMemoryChatMessageStore(state?.StoreState ?? default, jsonSerializerOptions);
}
}
/// <summary>
/// Gets or sets the ID of the underlying service thread to support cases where the chat history is stored by the agent service.
/// </summary>
@@ -175,47 +215,6 @@ public class AgentThread
}
}
/// <summary>
/// Deserializes the state contained in the provided <see cref="JsonElement"/> into the properties on this thread.
/// </summary>
/// <param name="serializedThread">A <see cref="JsonElement"/> representing the state of the thread.</param>
/// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A <see cref="ValueTask"/> that completes when the state has been deserialized.</returns>
protected internal virtual async Task DeserializeAsync(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
{
var state = JsonSerializer.Deserialize(
serializedThread,
AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ThreadState))) as ThreadState;
if (state?.ConversationId is string threadId)
{
this.ConversationId = threadId;
// Since we have an ID, we should not have a chat message store and we can return here.
return;
}
if (state?.AIContextProviderState.HasValue is true && this.AIContextProvider is not null)
{
await this.AIContextProvider.DeserializeAsync(state.AIContextProviderState.Value, jsonSerializerOptions, cancellationToken).ConfigureAwait(false);
}
// If we don't have any IChatMessageStore state return here.
if (state?.StoreState is null || state?.StoreState.Value.ValueKind is JsonValueKind.Undefined or JsonValueKind.Null)
{
return;
}
if (this._messageStore is null)
{
// If we don't have a chat message store yet, create an in-memory one.
this._messageStore = new InMemoryChatMessageStore();
}
await this._messageStore.DeserializeStateAsync(state!.StoreState.Value, jsonSerializerOptions, cancellationToken).ConfigureAwait(false);
}
internal sealed class ThreadState
{
public string? ConversationId { get; set; }
@@ -2,6 +2,7 @@
using System;
using System.Collections.Generic;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Shared.Diagnostics;
@@ -52,6 +53,10 @@ public class DelegatingAIAgent : AIAgent
/// <inheritdoc />
public override AgentThread GetNewThread() => this.InnerAgent.GetNewThread();
/// <inheritdoc />
public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
=> this.InnerAgent.DeserializeThread(serializedThread, jsonSerializerOptions, cancellationToken);
/// <inheritdoc />
public override Task<AgentRunResponse> RunAsync(
IEnumerable<ChatMessage> messages,
@@ -44,28 +44,11 @@ public interface IChatMessageStore
/// <returns>An async task.</returns>
Task AddMessagesAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default);
/// <summary>
/// Deserializes the state contained in the provided <see cref="JsonElement"/> into the properties on this store.
/// </summary>
/// <param name="serializedStoreState">A <see cref="JsonElement"/> representing the state of the store.</param>
/// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A <see cref="ValueTask"/> that completes when the state has been deserialized.</returns>
/// <remarks>
/// This method, together with <see cref="SerializeStateAsync(JsonSerializerOptions?, CancellationToken)"/> can be used to save and load messages from a persistent store
/// if this store only has messages in memory.
/// </remarks>
ValueTask DeserializeStateAsync(JsonElement? serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default);
/// <summary>
/// Serializes the current object's state to a <see cref="JsonElement"/> using the specified serialization options.
/// </summary>
/// <param name="jsonSerializerOptions">The JSON serialization options to use.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>A <see cref="JsonElement"/> representation of the object's state.</returns>
/// <remarks>
/// This method, together with <see cref="DeserializeStateAsync(JsonElement?, JsonSerializerOptions?, CancellationToken)"/> can be used to save and load messages from a persistent store
/// if this store only has messages in memory.
/// </remarks>
ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default);
}
@@ -17,12 +17,23 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
{
private readonly IChatReducer? _chatReducer;
private readonly ChatReducerTriggerEvent _reducerTriggerEvent;
private List<ChatMessage> _messages = new();
private List<ChatMessage> _messages;
/// <summary>
/// Initializes a new instance of the <see cref="InMemoryChatMessageStore"/> class.
/// </summary>
public InMemoryChatMessageStore()
{
this._messages = new();
}
/// <summary>
/// Initializes a new instance of the <see cref="InMemoryChatMessageStore"/> class, with an existing state from a serialized JSON element.
/// </summary>
/// <param name="serializedStoreState">A <see cref="JsonElement"/> representing the serialized state of the store.</param>
/// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
public InMemoryChatMessageStore(JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null)
: this(null, serializedStoreState, jsonSerializerOptions, ChatReducerTriggerEvent.BeforeMessagesRetrieval)
{
}
@@ -32,9 +43,40 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
/// <param name="chatReducer">An optional <see cref="IChatReducer"/> instance used to process or reduce chat messages. If null, no reduction logic will be applied.</param>
/// <param name="reducerTriggerEvent">The event that should trigger the reducer invocation.</param>
public InMemoryChatMessageStore(IChatReducer chatReducer, ChatReducerTriggerEvent reducerTriggerEvent = ChatReducerTriggerEvent.BeforeMessagesRetrieval)
: this(chatReducer, default, null, reducerTriggerEvent)
{
this._chatReducer = Throw.IfNull(chatReducer);
Throw.IfNull(chatReducer);
}
/// <summary>
/// Initializes a new instance of the <see cref="InMemoryChatMessageStore"/> class, with an existing state from a serialized JSON element.
/// </summary>
/// <param name="chatReducer">An optional <see cref="IChatReducer"/> instance used to process or reduce chat messages. If null, no reduction logic will be applied.</param>
/// <param name="serializedStoreState">A <see cref="JsonElement"/> representing the serialized state of the store.</param>
/// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
/// <param name="reducerTriggerEvent">The event that should trigger the reducer invocation.</param>
public InMemoryChatMessageStore(IChatReducer? chatReducer, JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, ChatReducerTriggerEvent reducerTriggerEvent = ChatReducerTriggerEvent.BeforeMessagesRetrieval)
{
this._chatReducer = chatReducer;
this._reducerTriggerEvent = reducerTriggerEvent;
if (serializedStoreState.ValueKind != JsonValueKind.Object)
{
this._messages = new();
return;
}
var state = JsonSerializer.Deserialize(
serializedStoreState,
AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))) as StoreState;
if (state?.Messages is { Count: > 0 } messages)
{
this._messages = messages;
return;
}
this._messages = new();
}
/// <summary>
@@ -84,26 +126,6 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
return this._messages;
}
/// <inheritdoc />
public ValueTask DeserializeStateAsync(JsonElement? serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
{
if (serializedStoreState is null)
{
return default;
}
var state = JsonSerializer.Deserialize(
serializedStoreState.Value,
AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))) as StoreState;
if (state?.Messages is { Count: > 0 } messages)
{
this._messages.AddRange(messages);
}
return default;
}
/// <inheritdoc />
public ValueTask<JsonElement?> SerializeStateAsync(JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
{
@@ -157,7 +179,7 @@ public sealed class InMemoryChatMessageStore : IList<ChatMessage>, IChatMessageS
internal sealed class StoreState
{
public IList<ChatMessage> Messages { get; set; } = new List<ChatMessage>();
public List<ChatMessage> Messages { get; set; } = new List<ChatMessage>();
}
/// <summary>
@@ -39,7 +39,7 @@ internal sealed class AgentActor(
if (threadResult.Value is { } threadJson)
{
// Deserialize the thread state if it exists
this._thread = await agent.DeserializeThreadAsync(threadJson, cancellationToken: cancellationToken).ConfigureAwait(false);
this._thread = agent.DeserializeThread(threadJson, cancellationToken: cancellationToken);
hasExistingThread = true;
}
}
@@ -36,6 +36,10 @@ public sealed class AgentProxy : AIAgent
/// <inheritdoc/>
public override AgentThread GetNewThread() => new AgentProxyThread();
/// <inheritdoc/>
public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
=> new AgentProxyThread(serializedThread, jsonSerializerOptions);
/// <summary>
/// Gets a thread by its <see cref="AgentThread.ConversationId"/>.
/// </summary>
@@ -1,6 +1,7 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Text.Json;
using System.Text.RegularExpressions;
using Microsoft.Shared.Diagnostics;
@@ -55,6 +56,16 @@ internal sealed partial class AgentProxyThread : AgentThread
{
}
/// <summary>
/// Initializes a new instance of the <see cref="AgentProxyThread"/> class from serialized state.
/// </summary>
/// <param name="serializedThreadState">A <see cref="JsonElement"/> representing the serialized state of the thread.</param>
/// <param name="jsonSerializerOptions">Optional settings for customizing the JSON deserialization process.</param>
public AgentProxyThread(JsonElement serializedThreadState, JsonSerializerOptions? jsonSerializerOptions = null)
: base(serializedThreadState, jsonSerializerOptions)
{
}
internal static string CreateId() => Guid.NewGuid().ToString("N");
/// <summary>
@@ -1,5 +1,6 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Text.Json;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.AI.Agents;
using Microsoft.Extensions.Logging;
@@ -83,6 +84,10 @@ public class OpenAIChatClientAgent : AIAgent
public sealed override AgentThread GetNewThread()
=> this._chatClientAgent.GetNewThread();
/// <inheritdoc/>
public sealed override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
=> this._chatClientAgent.DeserializeThread(serializedThread, jsonSerializerOptions, cancellationToken);
/// <inheritdoc/>
public sealed override Task<AgentRunResponse> RunAsync(
IEnumerable<Microsoft.Extensions.AI.ChatMessage> messages,
@@ -4,6 +4,7 @@ using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
@@ -241,12 +242,16 @@ public sealed class ChatClientAgent : AIAgent
{
var thread = new AgentThread
{
MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke(),
AIContextProvider = this._agentOptions?.AIContextProviderFactory?.Invoke()
MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke(default, null),
AIContextProvider = this._agentOptions?.AIContextProviderFactory?.Invoke(default, null)
};
return thread;
}
/// <inheritdoc/>
public override AgentThread DeserializeThread(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default)
=> new(serializedThread, jsonSerializerOptions, this._agentOptions?.ChatMessageStoreFactory, this._agentOptions?.AIContextProviderFactory);
#region Private
/// <summary>
@@ -493,7 +498,7 @@ public sealed class ChatClientAgent : AIAgent
// If the service doesn't use service side thread storage (i.e. we got no id back from invocation), and
// the thread has no MessageStore yet, and we have a custom messages store, we should update the thread
// with the custom MessageStore so that it has somewhere to store the chat history.
thread.MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke();
thread.MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke(default, null);
}
}
@@ -2,6 +2,7 @@
using System;
using System.Collections.Generic;
using System.Text.Json;
namespace Microsoft.Extensions.AI.Agents;
@@ -78,14 +79,14 @@ public class ChatClientAgentOptions
/// Gets or sets a factory function to create an instance of <see cref="IChatMessageStore"/>
/// which will be used to store chat messages for this agent.
/// </summary>
public Func<IChatMessageStore>? ChatMessageStoreFactory { get; set; }
public Func<JsonElement, JsonSerializerOptions?, IChatMessageStore>? ChatMessageStoreFactory { get; set; }
/// <summary>
/// Gets or sets a factory function to create an instance of <see cref="AIContextProvider"/>
/// which will be used to create a context provider for each new thread, and can then
/// provide additional context for each agent run.
/// </summary>
public Func<AIContextProvider>? AIContextProviderFactory { get; set; }
public Func<JsonElement, JsonSerializerOptions?, AIContextProvider>? AIContextProviderFactory { get; set; }
/// <summary>
/// Gets or sets a value indicating whether to use the provided <see cref="IChatClient"/> instance as is,
@@ -135,40 +135,39 @@ public class AgentThreadTests
#region Deserialize Tests
[Fact]
public async Task VerifyDeserializeWithMessagesAsync()
public async Task VerifyDeserializeConstructorWithMessagesAsync()
{
// Arrange
var chatMessageStore = new InMemoryChatMessageStore();
var json = JsonSerializer.Deserialize("""
{
"storeState": { "messages": [{"authorName": "testAuthor"}] }
}
""", TestJsonSerializerContext.Default.JsonElement);
var thread = new AgentThread { MessageStore = chatMessageStore };
// Act.
await thread.DeserializeAsync(json);
var thread = new AgentThread(json);
// Assert
Assert.Null(thread.ConversationId);
Assert.Single(chatMessageStore);
Assert.Equal("testAuthor", chatMessageStore[0].AuthorName);
var messageStore = thread.MessageStore as InMemoryChatMessageStore;
Assert.NotNull(messageStore);
Assert.Single(messageStore);
Assert.Equal("testAuthor", messageStore[0].AuthorName);
}
[Fact]
public async Task VerifyDeserializeWithIdAsync()
public async Task VerifyDeserializeConstructorWithIdAsync()
{
// Arrange
var json = JsonSerializer.Deserialize("""
{
"conversationId": "TestConvId"
"ConversationId": "TestConvId"
}
""", TestJsonSerializerContext.Default.JsonElement);
var thread = new AgentThread();
// Act
await thread.DeserializeAsync(json);
var thread = new AgentThread(json);
// Assert
Assert.Equal("TestConvId", thread.ConversationId);
@@ -176,34 +175,33 @@ public class AgentThreadTests
}
[Fact]
public async Task VerifyDeserializeWithAIContextProviderAsync()
public async Task VerifyDeserializeConstructorWithAIContextProviderAsync()
{
// Arrange
var json = JsonSerializer.Deserialize("""
{
"ConversationId": "TestConvId",
"aiContextProviderState": ["CP1"]
}
""", TestJsonSerializerContext.Default.JsonElement);
Mock<AIContextProvider> mockProvider = new();
var thread = new AgentThread() { AIContextProvider = mockProvider.Object };
// Act
await thread.DeserializeAsync(json);
var thread = new AgentThread(json, aiContextProviderFactory: (_, _) => mockProvider.Object);
// Assert
Assert.Null(thread.MessageStore);
mockProvider.Verify(m => m.DeserializeAsync(It.Is<JsonElement>(e => e.ValueKind == JsonValueKind.Array && e.GetArrayLength() == 1), It.IsAny<JsonSerializerOptions?>(), It.IsAny<CancellationToken>()), Times.Once);
Assert.Same(thread.AIContextProvider, mockProvider.Object);
}
[Fact]
public async Task DeserializeWithInvalidJsonThrowsAsync()
public async Task DeserializeContructorWithInvalidJsonThrowsAsync()
{
// Arrange
var invalidJson = JsonSerializer.Deserialize("[42]", TestJsonSerializerContext.Default.JsonElement);
var thread = new AgentThread();
// Act & Assert
await Assert.ThrowsAsync<JsonException>(() => thread.DeserializeAsync(invalidJson));
Assert.Throws<ArgumentException>(() => new AgentThread(invalidJson));
}
#endregion Deserialize Tests
@@ -89,19 +89,17 @@ public class InMemoryChatMessageStoreTests
}
[Fact]
public async Task DeserializeWithEmptyElementAsync()
public async Task DeserializeConstructorWithEmptyElementAsync()
{
var newStore = new InMemoryChatMessageStore();
var emptyObject = JsonSerializer.Deserialize<JsonElement>("{}", TestJsonSerializerContext.Default.JsonElement);
await newStore.DeserializeStateAsync(emptyObject);
var newStore = new InMemoryChatMessageStore(emptyObject);
Assert.Empty(newStore);
}
[Fact]
public async Task SerializeAndDeserializeRoundtripsAsync()
public async Task SerializeAndDeserializeConstructorRoundtripsAsync()
{
var store = new InMemoryChatMessageStore
{
@@ -110,9 +108,7 @@ public class InMemoryChatMessageStoreTests
};
var jsonElement = await store.SerializeStateAsync();
var newStore = new InMemoryChatMessageStore();
await newStore.DeserializeStateAsync(jsonElement);
var newStore = new InMemoryChatMessageStore(jsonElement.Value);
Assert.Equal(2, newStore.Count);
Assert.Equal("A", newStore[0].Text);
@@ -141,57 +137,49 @@ public class InMemoryChatMessageStoreTests
}
[Fact]
public async Task DeserializeStateAsync_WithNullSerializedState_DoesNothingAsync()
public void DeserializeContructor_WithNullSerializedState_CreatesEmptyStore()
{
// Arrange
var store = new InMemoryChatMessageStore();
store.Add(new ChatMessage(ChatRole.User, "Existing message"));
// Act
await store.DeserializeStateAsync(null);
var store = new InMemoryChatMessageStore(new JsonElement());
// Assert - Should still have the existing message
Assert.Single(store);
Assert.Equal("Existing message", store[0].Text);
// Assert
Assert.Empty(store);
}
[Fact]
public async Task DeserializeStateAsync_WithEmptyMessages_DoesNotAddMessagesAsync()
public async Task DeserializeContructor_WithEmptyMessages_DoesNotAddMessagesAsync()
{
// Arrange
var store = new InMemoryChatMessageStore();
var stateWithEmptyMessages = JsonSerializer.SerializeToElement(
new Dictionary<string, object> { ["Messages"] = new List<ChatMessage>() },
TestJsonSerializerContext.Default.IDictionaryStringObject);
// Act
await store.DeserializeStateAsync(stateWithEmptyMessages);
var store = new InMemoryChatMessageStore(stateWithEmptyMessages);
// Assert
Assert.Empty(store);
}
[Fact]
public async Task DeserializeStateAsync_WithNullMessages_DoesNotAddMessagesAsync()
public async Task DeserializeConstructor_WithNullMessages_DoesNotAddMessagesAsync()
{
// Arrange
var store = new InMemoryChatMessageStore();
var stateWithNullMessages = JsonSerializer.SerializeToElement(
new Dictionary<string, object> { ["Messages"] = null! },
TestJsonSerializerContext.Default.DictionaryStringObject);
// Act
await store.DeserializeStateAsync(stateWithNullMessages);
var store = new InMemoryChatMessageStore(stateWithNullMessages);
// Assert
Assert.Empty(store);
}
[Fact]
public async Task DeserializeStateAsync_WithValidMessages_AddsMessagesAsync()
public async Task DeserializeConstructor_WithValidMessages_AddsMessagesAsync()
{
// Arrange
var store = new InMemoryChatMessageStore();
var messages = new List<ChatMessage>
{
new(ChatRole.User, "User message"),
@@ -203,7 +191,7 @@ public class InMemoryChatMessageStoreTests
TestJsonSerializerContext.Default.DictionaryStringObject);
// Act
await store.DeserializeStateAsync(serializedState);
var store = new InMemoryChatMessageStore(serializedState);
// Assert
Assert.Equal(2, store.Count);
@@ -1,6 +1,7 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Collections.Generic;
using System.Text.Json;
using Moq;
namespace Microsoft.Extensions.AI.Agents.UnitTests.ChatCompletion;
@@ -164,8 +165,8 @@ public class ChatClientAgentOptionsTests
const string Name = "Test name";
const string Description = "Test description";
var tools = new List<AITool> { AIFunctionFactory.Create(() => "test") };
static IChatMessageStore ChatMessageStoreFactory() => new Mock<IChatMessageStore>().Object;
static AIContextProvider AIContextProviderFactory() => new Mock<AIContextProvider>().Object;
static IChatMessageStore ChatMessageStoreFactory(JsonElement jse, JsonSerializerOptions? jso) => new Mock<IChatMessageStore>().Object;
static AIContextProvider AIContextProviderFactory(JsonElement jse, JsonSerializerOptions? jso) => new Mock<AIContextProvider>().Object;
var original = new ChatClientAgentOptions(Instructions, Name, Description, tools)
{
@@ -461,7 +461,7 @@ public class ChatClientAgentTests
.Setup(p => p.InvokedAsync(It.IsAny<AIContextProvider.InvokedContext>(), It.IsAny<CancellationToken>()))
.Returns(new ValueTask());
ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "base instructions", AIContextProviderFactory = () => mockProvider.Object, ChatOptions = new() { Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "base instructions", AIContextProviderFactory = (_, _) => mockProvider.Object, ChatOptions = new() { Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
// Act
await agent.RunAsync(requestMessages);
@@ -506,7 +506,7 @@ public class ChatClientAgentTests
.Setup(p => p.InvokedAsync(It.IsAny<AIContextProvider.InvokedContext>(), It.IsAny<CancellationToken>()))
.Returns(new ValueTask());
ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "base instructions", AIContextProviderFactory = () => mockProvider.Object, ChatOptions = new() { Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "base instructions", AIContextProviderFactory = (_, _) => mockProvider.Object, ChatOptions = new() { Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
// Act
await Assert.ThrowsAsync<InvalidOperationException>(() => agent.RunAsync(requestMessages));
@@ -548,7 +548,7 @@ public class ChatClientAgentTests
.Setup(p => p.InvokingAsync(It.IsAny<AIContextProvider.InvokingContext>(), It.IsAny<CancellationToken>()))
.ReturnsAsync(new AIContext());
ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "base instructions", AIContextProviderFactory = () => mockProvider.Object, ChatOptions = new() { Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
ChatClientAgent agent = new(mockService.Object, options: new() { Instructions = "base instructions", AIContextProviderFactory = (_, _) => mockProvider.Object, ChatOptions = new() { Tools = [AIFunctionFactory.Create(() => { }, "base function")] } });
// Act
await agent.RunAsync([new(ChatRole.User, "user message")]);
@@ -1713,7 +1713,7 @@ public class ChatClientAgentTests
var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
{
Instructions = "Test instructions",
ChatMessageStoreFactory = () =>
ChatMessageStoreFactory = (_, _) =>
{
factoryCalled = true;
return mockStore.Object;
@@ -1738,7 +1738,7 @@ public class ChatClientAgentTests
var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions
{
Instructions = "Test instructions",
AIContextProviderFactory = () =>
AIContextProviderFactory = (_, _) =>
{
factoryCalled = true;
return mockContextProvider.Object;