.NET: Add ChatClient decorator for calling AIContextProviders#4097
.NET: Add ChatClient decorator for calling AIContextProviders#4097westey-m wants to merge 2 commits intomicrosoft:mainfrom
Conversation
There was a problem hiding this comment.
Pull request overview
Adds decorator-based middleware to invoke AIContextProvider/MessageAIContextProvider pipelines, enabling context enrichment before chat/agent execution and provider notifications after completion, aligning with issue #3957’s request for an IChatClient-level decorator.
Changes:
- Introduces
AIContextProviderChatClientand aChatClientBuilder.Use(...)extension to invokeAIContextProviders aroundIChatClientcalls. - Introduces
MessageAIContextProviderAgentto invokeMessageAIContextProviders aroundAIAgentruns (including streaming). - Adds unit tests and updates the Step14 middleware sample/docs to demonstrate the new pipeline.
Reviewed changes
Copilot reviewed 5 out of 7 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClient.cs | New delegating chat client that invokes AIContextProvider pipeline for non-streaming and streaming calls. |
| dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/AIContextProviderChatClientBuilderExtensions.cs | Adds ChatClientBuilder.Use(params AIContextProvider[]) extension to register the decorator. |
| dotnet/src/Microsoft.Agents.AI/AIContextProviderDecorators/MessageAIContextProviderAgent.cs | New delegating agent that invokes MessageAIContextProvider pipeline around agent runs and streaming. |
| dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs | Unit tests covering enrichment, sequencing, success/failure notifications, and builder integration for chat client decorator. |
| dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/MessageAIContextProviderAgentTests.cs | Unit tests covering enrichment, sequencing, and success/failure notifications for agent decorator (incl. streaming). |
| dotnet/samples/GettingStarted/Agents/Agent_Step14_Middleware/Program.cs | Sample updates to demonstrate using AIContextProvider in the IChatClient pipeline. |
| dotnet/samples/GettingStarted/Agents/Agent_Step14_Middleware/README.md | Documentation update listing the new chat-client-level AIContextProvider middleware example. |
| IAsyncEnumerator<ChatResponseUpdate> enumerator; | ||
| try | ||
| { | ||
| enumerator = base.GetStreamingResponseAsync(enrichedMessages, enrichedOptions, cancellationToken).GetAsyncEnumerator(cancellationToken); | ||
| } |
There was a problem hiding this comment.
The streaming implementation creates an IAsyncEnumerator via GetAsyncEnumerator(...) but never disposes it. Consider using await using / try-finally to ensure DisposeAsync is called so underlying streaming resources are released promptly.
| // AIContextProvider at the chat client level. Unlike the agent-level MessageAIContextProvider, | ||
| // this operates within the IChatClient pipeline and can also enrich tools and instructions. | ||
| // It must be used within the context of a running AIAgent (uses AIAgent.CurrentRunContext). | ||
| Console.WriteLine("\n\n=== Example 6: AIContextProvider on chat client pipeline ==="); | ||
|
|
||
| var chatClientProviderAgent = azureOpenAIClient.AsIChatClient() | ||
| .AsBuilder() | ||
| .Use(new DateTimeContextProvider()) | ||
| .BuildAIAgent( | ||
| instructions: "You are an AI assistant that helps people find information."); |
There was a problem hiding this comment.
This example/comment says the chat-client-level AIContextProvider pipeline can enrich tools and instructions, but it wires up DateTimeContextProvider, which is a MessageAIContextProvider (defined later in this file) and only injects messages. Either adjust the wording to reflect message-only enrichment for this provider, or add a dedicated AIContextProvider implementation here that demonstrates tools/instructions enrichment so the sample matches the text.
Motivation and Context
#3957
Description
Contribution Checklist