.NET: Add ability to mark the source of Agent request messages and use that for filtering #3540
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
markwallace-microsoft
added
.NET
workflows
github-actions
bot
changed the title
Contributor
There was a problem hiding this comment.
Pull request overview
This PR introduces explicit tracking of the source of agent request messages (external input, AI context provider, chat history), and updates the agent pipeline and providers to use this metadata for filtering and behavior.
Changes:
- Add
AgentRequestMessageSourceandChatMessageExtensions.GetAgentRequestMessageSourceto tag and retrieve the logical source ofChatMessageinstances viaAdditionalProperties. - Refactor
AIContextProviderandChatHistoryProviderto use public wrapper methods (InvokingAsync/InvokedAsync) delegating to new protectedInvokingCoreAsync/InvokedCoreAsync, and update all implementations (Mem0, Cosmos, in-memory, workflow, samples) and tests accordingly. - Update
ChatClientAgentand various providers (ChatHistoryMemoryProvider,TextSearchProvider, Mem0, ChatHistory filters) so they treat all messages uniformly as request messages, using the new source metadata to include or exclude messages from memory, search, and history operations.
Reviewed changes
Copilot reviewed 33 out of 33 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRequestMessageSource.cs |
Adds a new sealed type and key for annotating ChatMessage source via AdditionalProperties. |
dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageExtensions.cs |
Adds an extension method to read the AgentRequestMessageSource from a ChatMessage, defaulting to External. |
dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs |
Introduces InvokingCoreAsync/InvokedCoreAsync and a wrapper that tags AI-context messages with AIContextProvider source. |
dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs |
Introduces InvokingCoreAsync/InvokedCoreAsync and a wrapper that tags returned history messages with ChatHistory source, and simplifies InvokedContext. |
dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProviderExtensions.cs |
Updates filters to work on the unified request-message list using AgentRequestMessageSource instead of separate AI-context collections. |
dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProviderMessageFilter.cs |
Adjusts to new InvokingCoreAsync/InvokedCoreAsync pattern while preserving filtering behavior. |
dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs |
Updates core methods and switches to aggregating only request+response messages, relying on message-source metadata. |
dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs |
Uses message-source metadata to include only external-request text in search input and recent-message tracking. |
dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs |
Filters to external request messages when querying and storing vector-memory items. |
dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs |
Filters to external request messages for memory query/persist operations. |
dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs |
Migrates to new core methods and stops treating AI-context messages as a separate collection. |
dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs |
Updates the run pipeline to use unified input-message lists, adjusts notifications to context/history providers, and removes explicit AI-context/history message collections in InvokedContext. |
dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs |
Adapts workflow chat history to new core methods and unified message model. |
dotnet/samples/... (multiple) |
Updates custom AIContextProvider and ChatHistoryProvider implementations in samples to override the new *CoreAsync methods and to work with the unified request-message model. |
dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/*.cs |
Adds tests for AgentRequestMessageSource and ChatMessageExtensions, and updates existing tests to align with the new InvokedContext shape and message-source semantics. |
dotnet/tests/Microsoft.Agents.AI.UnitTests/* and *.Mem0*, *.CosmosNoSql* |
Adjusts tests to the new provider APIs and confirms that history, AI context, and memory providers behave correctly with the new source-tagging and filtering logic. |
Comment on lines
+52
to
+96
| public async ValueTask<AIContext> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) | ||
| { | ||
| var aiContext = await this.InvokingCoreAsync(context, cancellationToken).ConfigureAwait(false); | ||
| if (aiContext.Messages is null) | ||
| { | ||
| return aiContext; | ||
| } | ||
|
|
||
| aiContext.Messages = aiContext.Messages.Select(message => | ||
| { | ||
| if (message.AdditionalProperties != null | ||
| && message.AdditionalProperties.TryGetValue(AgentRequestMessageSource.AdditionalPropertiesKey, out var source) | ||
| && source is AgentRequestMessageSource typedSource | ||
| && typedSource == AgentRequestMessageSource.AIContextProvider) | ||
| { | ||
| return message; | ||
| } | ||
|
|
||
| message = message.Clone(); | ||
| message.AdditionalProperties ??= new(); | ||
| message.AdditionalProperties[AgentRequestMessageSource.AdditionalPropertiesKey] = AgentRequestMessageSource.AIContextProvider; | ||
| return message; | ||
| }).ToList(); | ||
|
|
||
| return aiContext; | ||
| } | ||
|
|
||
| /// <summary> | ||
| /// Called at the start of agent invocation to provide additional context. | ||
| /// </summary> | ||
| /// <param name="context">Contains the request context including the caller provided messages that will be used by the agent for this invocation.</param> | ||
| /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param> | ||
| /// <returns>A task that represents the asynchronous operation. The task result contains the <see cref="AIContext"/> with additional context to be used by the agent during this invocation.</returns> | ||
| /// <remarks> | ||
| /// <para> | ||
| /// Implementers can load any additional context required at this time, such as: | ||
| /// <list type="bullet"> | ||
| /// <item><description>Retrieving relevant information from knowledge bases</description></item> | ||
| /// <item><description>Adding system instructions or prompts</description></item> | ||
| /// <item><description>Providing function tools for the current invocation</description></item> | ||
| /// <item><description>Injecting contextual messages from conversation history</description></item> | ||
| /// </list> | ||
| /// </para> | ||
| /// </remarks> | ||
| protected abstract ValueTask<AIContext> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default); |
There was a problem hiding this comment.
Changing InvokingAsync from an abstract method to a non-virtual wrapper over the new protected InvokingCoreAsync is an API surface change that breaks existing derived implementations (both at compile-time and potentially at runtime for pre-compiled consumers). This likely needs to be treated and documented as a breaking change (including updating release notes / PR title), and guidance should be provided to external implementers on how to migrate their overrides to InvokingCoreAsync.
| /// <para> | ||
| /// This method is called regardless of whether the invocation succeeded or failed. | ||
| /// To check if the invocation was successful, inspect the <see cref="InvokedContext.InvokeException"/> property. | ||
| /// </para> |
There was a problem hiding this comment.
Similar to AIContextProvider, ChatHistoryProvider.InvokingAsync and InvokedAsync have been changed from abstract methods to concrete wrappers over new protected InvokingCoreAsync / InvokedCoreAsync methods. This is a breaking change for any external implementations that previously overrode the public methods, so it should be explicitly treated and communicated as a breaking change, with migration guidance for moving logic into the new core methods.
Suggested change
| /// </para> | |
| /// </para> | |
| /// <para> | |
| /// <b>Breaking change:</b> In earlier versions of <see cref="ChatHistoryProvider"/>, implementations were expected to | |
| /// override the public <c>InvokedAsync</c> method directly. Starting with this version, <c>InvokedAsync</c> is a concrete | |
| /// wrapper that delegates to <see cref="InvokedCoreAsync"/>. Implementations that previously overrode <c>InvokedAsync</c> | |
| /// should migrate their logic into an override of <see cref="InvokedCoreAsync"/> instead. This ensures that custom | |
| /// post-invocation behavior continues to be executed when the agent workflow completes. | |
| /// </para> |
| namespace Microsoft.Agents.AI; | ||
|
|
||
| /// <summary> | ||
| /// Conatins extension methods for <see cref="ChatMessage"/> |
There was a problem hiding this comment.
Minor typo in the XML summary: "Conatins" should be "Contains" in the description of the ChatMessageExtensions class.
Suggested change
| /// Conatins extension methods for <see cref="ChatMessage"/> | |
| /// Contains extension methods for <see cref="ChatMessage"/> |
Comment on lines
+10
to
+14
| /// An enumeration representing the source of an agent request message. | ||
| /// </summary> | ||
| /// <remarks> | ||
| /// Input messages for a specific agent run can originate from various sources. | ||
| /// This enumeration helps to identify whether a message came from outside the agent pipeline, |
There was a problem hiding this comment.
The summary currently describes AgentRequestMessageSource as an "enumeration", but this type is implemented as a sealed class; consider updating the wording to avoid confusion for consumers and keep documentation aligned with the implementation.
Suggested change
| /// An enumeration representing the source of an agent request message. | |
| /// </summary> | |
| /// <remarks> | |
| /// Input messages for a specific agent run can originate from various sources. | |
| /// This enumeration helps to identify whether a message came from outside the agent pipeline, | |
| /// Represents the source of an agent request message. | |
| /// </summary> | |
| /// <remarks> | |
| /// Input messages for a specific agent run can originate from various sources. | |
| /// This type helps to identify whether a message came from outside the agent pipeline, |
| // Copyright (c) Microsoft. All rights reserved. | ||
|
|
||
| using System; | ||
| using Microsoft.Extensions.AI; |
There was a problem hiding this comment.
The using Microsoft.Extensions.AI; directive appears to be unused in this file and may trigger build warnings; consider removing it to keep the file warning-free and consistent with the rest of the codebase.
Suggested change
| using Microsoft.Extensions.AI; |
| AgentRequestMessageSource source = new("Test"); | ||
|
|
||
| // Act | ||
| bool result = source.Equals(source); |
There was a problem hiding this comment.
Comparison of identical values.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Motivation and Context
This is addressing feedback that messages should be annotated with their source, rather than split by type when passed to the AIContextProvider and ChatHistoryProvider.
#2542
Description
Contribution Checklist