.NET: Add AIContextProvider support (#691)

* Add AIContextProvider support * Address feedback. * Address PR comments. * Switch to valuetask and remove parallel calls for AIContextProvider * Remove Model from ModelInvokingAsync method name * Remove agent thread id again and remove it from context provider interface * Add AIContextProvider serialization support to AgentThread and update sample to show this feature * Address PR comments * Improve memory sample * Update sample comment. * Remove AggregateAIContextProvider for now since it makes too many assumptions. We can include it later as a sample if needed. * Update AIContextProviders to have an Invoked method instead of MessagesAddingAsync. * Remove unused using. * Address PR comments. * Address PR comment. * Update comment. * Update comment * Address PR comments.
2026-06-16 21:04:09 +08:00 · 2025-09-16 10:54:18 +01:00
parent 89cb94b5c2
commit 66fe1c957c
25 changed files with 981 additions and 32 deletions
@@ -114,7 +114,17 @@ public sealed class ChatClientAgent : AIAgent

        this._logger.LogAgentChatClientInvokingAgent(nameof(RunAsync), this.Id, agentName, this._chatClientType);

-        ChatResponse chatResponse = await this.ChatClient.GetResponseAsync(threadMessages, chatOptions, cancellationToken).ConfigureAwait(false);
+        // Call the IChatClient and notify the AIContextProvider of any failures.
+        ChatResponse chatResponse;
+        try
+        {
+            chatResponse = await this.ChatClient.GetResponseAsync(threadMessages, chatOptions, cancellationToken).ConfigureAwait(false);
+        }
+        catch (Exception ex)
+        {
+            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, cancellationToken).ConfigureAwait(false);
+            throw;
+        }

        this._logger.LogAgentChatClientInvokedAgent(nameof(RunAsync), this.Id, agentName, this._chatClientType, inputMessages.Count);

@@ -122,19 +132,17 @@ public sealed class ChatClientAgent : AIAgent
        // so let's update it and set the conversation id for the service thread case.
        this.UpdateThreadWithTypeAndConversationId(safeThread, chatResponse.ConversationId);

-        // Only notify the thread of new messages if the chatResponse was successful to avoid inconsistent messages state in the thread.
-        await NotifyThreadOfNewMessagesAsync(safeThread, inputMessages, cancellationToken).ConfigureAwait(false);
-
        // Ensure that the author name is set for each message in the response.
        foreach (ChatMessage chatResponseMessage in chatResponse.Messages)
        {
            chatResponseMessage.AuthorName ??= agentName;
        }

-        // Convert the chat response messages to a valid IReadOnlyCollection for notification signatures below.
-        var chatResponseMessages = chatResponse.Messages as IReadOnlyCollection<ChatMessage> ?? [.. chatResponse.Messages];
+        // Only notify the thread of new messages if the chatResponse was successful to avoid inconsistent message state in the thread.
+        await NotifyThreadOfNewMessagesAsync(safeThread, inputMessages.Concat(chatResponse.Messages), cancellationToken).ConfigureAwait(false);

-        await NotifyThreadOfNewMessagesAsync(safeThread, chatResponseMessages, cancellationToken).ConfigureAwait(false);
+        // Notify the AIContextProvider of all new messages.
+        await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);

        return new(chatResponse) { AgentId = this.Id };
    }
@@ -156,15 +164,34 @@ public sealed class ChatClientAgent : AIAgent

        this._logger.LogAgentChatClientInvokingAgent(nameof(RunStreamingAsync), this.Id, loggingAgentName, this._chatClientType);

-        // Using the enumerator to ensure we consider the case where no updates are returned for notification.
-        var responseUpdatesEnumerator = this.ChatClient.GetStreamingResponseAsync(threadMessages, chatOptions, cancellationToken).GetAsyncEnumerator(cancellationToken);
+        List<ChatResponseUpdate> responseUpdates = [];
+
+        IAsyncEnumerator<ChatResponseUpdate> responseUpdatesEnumerator;
+
+        try
+        {
+            // Using the enumerator to ensure we consider the case where no updates are returned for notification.
+            responseUpdatesEnumerator = this.ChatClient.GetStreamingResponseAsync(threadMessages, chatOptions, cancellationToken).GetAsyncEnumerator(cancellationToken);
+        }
+        catch (Exception ex)
+        {
+            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, cancellationToken).ConfigureAwait(false);
+            throw;
+        }

        this._logger.LogAgentChatClientInvokedStreamingAgent(nameof(RunStreamingAsync), this.Id, loggingAgentName, this._chatClientType);

-        List<ChatResponseUpdate> responseUpdates = [];
-
-        // Ensure we start the streaming request
-        var hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false);
+        bool hasUpdates;
+        try
+        {
+            // Ensure we start the streaming request
+            hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false);
+        }
+        catch (Exception ex)
+        {
+            await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, cancellationToken).ConfigureAwait(false);
+            throw;
+        }

        while (hasUpdates)
        {
@@ -176,20 +203,28 @@ public sealed class ChatClientAgent : AIAgent
                yield return new(update) { AgentId = this.Id };
            }

-            hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false);
+            try
+            {
+                hasUpdates = await responseUpdatesEnumerator.MoveNextAsync().ConfigureAwait(false);
+            }
+            catch (Exception ex)
+            {
+                await NotifyAIContextProviderOfFailureAsync(safeThread, ex, inputMessages, cancellationToken).ConfigureAwait(false);
+                throw;
+            }
        }

        var chatResponse = responseUpdates.ToChatResponse();
-        var chatResponseMessages = chatResponse.Messages as IReadOnlyCollection<ChatMessage> ?? [.. chatResponse.Messages];

        // We can derive the type of supported thread from whether we have a conversation id,
        // so let's update it and set the conversation id for the service thread case.
        this.UpdateThreadWithTypeAndConversationId(safeThread, chatResponse.ConversationId);

        // To avoid inconsistent state we only notify the thread of the input messages if no error occurs after the initial request.
-        await NotifyThreadOfNewMessagesAsync(safeThread, inputMessages, cancellationToken).ConfigureAwait(false);
+        await NotifyThreadOfNewMessagesAsync(safeThread, inputMessages.Concat(chatResponse.Messages), cancellationToken).ConfigureAwait(false);

-        await NotifyThreadOfNewMessagesAsync(safeThread, chatResponseMessages, cancellationToken).ConfigureAwait(false);
+        // Notify the AIContextProvider of all new messages.
+        await NotifyAIContextProviderOfSuccessAsync(safeThread, inputMessages, chatResponse.Messages, cancellationToken).ConfigureAwait(false);
    }

    /// <inheritdoc/>
@@ -204,12 +239,40 @@ public sealed class ChatClientAgent : AIAgent
    /// <inheritdoc/>
    public override AgentThread GetNewThread()
    {
-        var thread = new AgentThread { MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke() };
+        var thread = new AgentThread
+        {
+            MessageStore = this._agentOptions?.ChatMessageStoreFactory?.Invoke(),
+            AIContextProvider = this._agentOptions?.AIContextProviderFactory?.Invoke()
+        };
        return thread;
    }

    #region Private

+    /// <summary>
+    /// Notify the <see cref="AIContextProvider"/> when an agent run succeeded, if there is an <see cref="AIContextProvider"/>.
+    /// </summary>
+    private static async Task NotifyAIContextProviderOfSuccessAsync(AgentThread thread, IEnumerable<ChatMessage> inputMessages, IEnumerable<ChatMessage> responseMessages, CancellationToken cancellationToken)
+    {
+        if (thread.AIContextProvider is not null)
+        {
+            await thread.AIContextProvider.InvokedAsync(new(inputMessages) { ResponseMessages = responseMessages },
+                cancellationToken).ConfigureAwait(false);
+        }
+    }
+
+    /// <summary>
+    /// Notify the <see cref="AIContextProvider"/> of any failure during an agent run, if there is an <see cref="AIContextProvider"/>.
+    /// </summary>
+    private static async Task NotifyAIContextProviderOfFailureAsync(AgentThread thread, Exception ex, IEnumerable<ChatMessage> inputMessages, CancellationToken cancellationToken)
+    {
+        if (thread.AIContextProvider is not null)
+        {
+            await thread.AIContextProvider.InvokedAsync(new(inputMessages) { InvokeException = ex },
+                cancellationToken).ConfigureAwait(false);
+        }
+    }
+
    /// <summary>
    /// Configures and returns chat options by merging the provided run options with the agent's default chat options.
    /// </summary>
@@ -350,6 +413,34 @@ public sealed class ChatClientAgent : AIAgent
            threadMessages.AddRange(await thread.MessageStore.GetMessagesAsync(cancellationToken).ConfigureAwait(false));
        }

+        // If we have an AIContextProvider, we should get context from it, and update our
+        // messages and options with the additional context.
+        if (thread.AIContextProvider is not null)
+        {
+            var invokingContext = new AIContextProvider.InvokingContext(inputMessages);
+            var aiContext = await thread.AIContextProvider.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false);
+            if (aiContext.Messages is { Count: > 0 })
+            {
+                threadMessages.AddRange(aiContext.Messages);
+            }
+
+            if (aiContext.Tools is { Count: > 0 })
+            {
+                chatOptions ??= new();
+                chatOptions.Tools ??= [];
+                foreach (AITool tool in aiContext.Tools)
+                {
+                    chatOptions.Tools.Add(tool);
+                }
+            }
+
+            if (aiContext.Instructions is not null)
+            {
+                chatOptions ??= new();
+                chatOptions.Instructions = string.IsNullOrWhiteSpace(chatOptions.Instructions) ? aiContext.Instructions : $"{chatOptions.Instructions}\n{aiContext.Instructions}";
+            }
+        }
+
        // Add the input messages to the end of thread messages.
        threadMessages.AddRange(inputMessages);

@@ -80,6 +80,13 @@ public class ChatClientAgentOptions
    /// </summary>
    public Func<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; }
+
    /// <summary>
    /// Gets or sets a value indicating whether to use the provided <see cref="IChatClient"/> instance as is,
    /// without applying any default decorators.
@@ -105,6 +112,7 @@ public class ChatClientAgentOptions
            Instructions = this.Instructions,
            Description = this.Description,
            ChatOptions = this.ChatOptions?.Clone(),
-            ChatMessageStoreFactory = this.ChatMessageStoreFactory
+            ChatMessageStoreFactory = this.ChatMessageStoreFactory,
+            AIContextProviderFactory = this.AIContextProviderFactory,
        };
 }