.NET: [BREAKING] Add consistent message filtering to all providers. (#3851)

* Add consistent message filtering to all providers.

* Remove old chat history filtering classes

* Fix merge issues

* Fix unit test

* Enforce non-nullable property

* Fix merging bug and make troubleshooting source info easier by adding tostring implementation
This commit is contained in:
westey
2026-02-12 10:50:13 +00:00
committed by GitHub
Unverified
parent c99df98547
commit de82ffd40a
26 changed files with 918 additions and 502 deletions
@@ -390,6 +390,49 @@ public sealed class AgentRequestMessageSourceAttributionTests
#endregion
#region ToString Tests
[Fact]
public void ToString_WithSourceId_ReturnsTypeColonId()
{
// Arrange
AgentRequestMessageSourceAttribution attribution = new(AgentRequestMessageSourceType.AIContextProvider, "MyProvider");
// Act
string result = attribution.ToString();
// Assert
Assert.Equal("AIContextProvider:MyProvider", result);
}
[Fact]
public void ToString_WithNullSourceId_ReturnsTypeOnly()
{
// Arrange
AgentRequestMessageSourceAttribution attribution = new(AgentRequestMessageSourceType.ChatHistory, null);
// Act
string result = attribution.ToString();
// Assert
Assert.Equal("ChatHistory", result);
}
[Fact]
public void ToString_Default_ReturnsExternalOnly()
{
// Arrange
AgentRequestMessageSourceAttribution attribution = default;
// Act
string result = attribution.ToString();
// Assert
Assert.Equal("External", result);
}
#endregion
#region Inequality Operator Tests
[Fact]
@@ -414,6 +414,46 @@ public sealed class AgentRequestMessageSourceTypeTests
#endregion
#region ToString Tests
[Fact]
public void ToString_ReturnsValue()
{
// Arrange
AgentRequestMessageSourceType source = new("CustomSource");
// Act
string result = source.ToString();
// Assert
Assert.Equal("CustomSource", result);
}
[Fact]
public void ToString_StaticExternal_ReturnsExternal()
{
// Arrange & Act
string result = AgentRequestMessageSourceType.External.ToString();
// Assert
Assert.Equal("External", result);
}
[Fact]
public void ToString_Default_ReturnsExternal()
{
// Arrange
AgentRequestMessageSourceType source = default;
// Act
string result = source.ToString();
// Assert
Assert.Equal("External", result);
}
#endregion
#region IEquatable Tests
[Fact]
@@ -1,141 +0,0 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
using Moq;
using Moq.Protected;
namespace Microsoft.Agents.AI.Abstractions.UnitTests;
/// <summary>
/// Contains tests for the <see cref="ChatHistoryProviderExtensions"/> class.
/// </summary>
public sealed class ChatHistoryProviderExtensionsTests
{
private static readonly AIAgent s_mockAgent = new Mock<AIAgent>().Object;
private static readonly AgentSession s_mockSession = new Mock<AgentSession>().Object;
[Fact]
public void WithMessageFilters_ReturnsChatHistoryProviderMessageFilter()
{
// Arrange
Mock<ChatHistoryProvider> providerMock = new();
// Act
ChatHistoryProvider result = providerMock.Object.WithMessageFilters(
invokingMessagesFilter: msgs => msgs,
invokedMessagesFilter: ctx => ctx);
// Assert
Assert.IsType<ChatHistoryProviderMessageFilter>(result);
}
[Fact]
public async Task WithMessageFilters_InvokingFilter_IsAppliedAsync()
{
// Arrange
Mock<ChatHistoryProvider> providerMock = new();
List<ChatMessage> innerMessages = [new(ChatRole.User, "Hello"), new(ChatRole.Assistant, "Hi")];
ChatHistoryProvider.InvokingContext context = new(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Test")]);
providerMock
.Protected()
.Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(innerMessages);
ChatHistoryProvider filtered = providerMock.Object.WithMessageFilters(
invokingMessagesFilter: msgs => msgs.Where(m => m.Role == ChatRole.User));
// Act
List<ChatMessage> result = (await filtered.InvokingAsync(context, CancellationToken.None)).ToList();
// Assert
Assert.Single(result);
Assert.Equal(ChatRole.User, result[0].Role);
}
[Fact]
public async Task WithMessageFilters_InvokedFilter_IsAppliedAsync()
{
// Arrange
Mock<ChatHistoryProvider> providerMock = new();
List<ChatMessage> requestMessages =
[
new(ChatRole.System, "System") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "TestSource") } } },
new(ChatRole.User, "Hello")
];
ChatHistoryProvider.InvokedContext context = new(s_mockAgent, s_mockSession, requestMessages)
{
ResponseMessages = [new ChatMessage(ChatRole.Assistant, "Response")]
};
ChatHistoryProvider.InvokedContext? capturedContext = null;
providerMock
.Protected()
.Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
.Callback<ChatHistoryProvider.InvokedContext, CancellationToken>((ctx, _) => capturedContext = ctx)
.Returns(default(ValueTask));
ChatHistoryProvider filtered = providerMock.Object.WithMessageFilters(
invokedMessagesFilter: ctx =>
{
ctx.ResponseMessages = null;
return ctx;
});
// Act
await filtered.InvokedAsync(context, CancellationToken.None);
// Assert
Assert.NotNull(capturedContext);
Assert.Null(capturedContext.ResponseMessages);
}
[Fact]
public void WithAIContextProviderMessageRemoval_ReturnsChatHistoryProviderMessageFilter()
{
// Arrange
Mock<ChatHistoryProvider> providerMock = new();
// Act
ChatHistoryProvider result = providerMock.Object.WithAIContextProviderMessageRemoval();
// Assert
Assert.IsType<ChatHistoryProviderMessageFilter>(result);
}
[Fact]
public async Task WithAIContextProviderMessageRemoval_RemovesAIContextProviderMessagesAsync()
{
// Arrange
Mock<ChatHistoryProvider> providerMock = new();
List<ChatMessage> requestMessages =
[
new(ChatRole.System, "System") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "TestSource") } } },
new(ChatRole.User, "Hello"),
new(ChatRole.System, "Context") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "TestContextSource") } } }
];
ChatHistoryProvider.InvokedContext context = new(s_mockAgent, s_mockSession, requestMessages);
ChatHistoryProvider.InvokedContext? capturedContext = null;
providerMock
.Protected()
.Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
.Callback<ChatHistoryProvider.InvokedContext, CancellationToken>((ctx, _) => capturedContext = ctx)
.Returns(default(ValueTask));
ChatHistoryProvider filtered = providerMock.Object.WithAIContextProviderMessageRemoval();
// Act
await filtered.InvokedAsync(context, CancellationToken.None);
// Assert
Assert.NotNull(capturedContext);
Assert.Equal(2, capturedContext.RequestMessages.Count());
Assert.Contains("System", capturedContext.RequestMessages.Select(x => x.Text));
Assert.Contains("Hello", capturedContext.RequestMessages.Select(x => x.Text));
}
}
@@ -1,213 +0,0 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
using Moq;
using Moq.Protected;
namespace Microsoft.Agents.AI.Abstractions.UnitTests;
/// <summary>
/// Contains tests for the <see cref="ChatHistoryProviderMessageFilter"/> class.
/// </summary>
public sealed class ChatHistoryProviderMessageFilterTests
{
private static readonly AIAgent s_mockAgent = new Mock<AIAgent>().Object;
private static readonly AgentSession s_mockSession = new Mock<AgentSession>().Object;
[Fact]
public void Constructor_WithNullInnerProvider_ThrowsArgumentNullException()
{
// Arrange, Act & Assert
Assert.Throws<ArgumentNullException>(() => new ChatHistoryProviderMessageFilter(null!));
}
[Fact]
public void Constructor_WithOnlyInnerProvider_Throws()
{
// Arrange
var innerProviderMock = new Mock<ChatHistoryProvider>();
// Act & Assert
Assert.Throws<ArgumentException>(() => new ChatHistoryProviderMessageFilter(innerProviderMock.Object));
}
[Fact]
public void Constructor_WithAllParameters_CreatesInstance()
{
// Arrange
var innerProviderMock = new Mock<ChatHistoryProvider>();
IEnumerable<ChatMessage> InvokingFilter(IEnumerable<ChatMessage> msgs) => msgs;
ChatHistoryProvider.InvokedContext InvokedFilter(ChatHistoryProvider.InvokedContext ctx) => ctx;
// Act
var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, InvokingFilter, InvokedFilter);
// Assert
Assert.NotNull(filter);
}
[Fact]
public void StateKey_DelegatesToInnerProvider()
{
// Arrange
var innerProviderMock = new Mock<ChatHistoryProvider>();
innerProviderMock.Setup(p => p.StateKey).Returns("inner-state-key");
// Act
var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, x => x);
// Assert
Assert.Equal("inner-state-key", filter.StateKey);
}
[Fact]
public async Task InvokingAsync_WithNoOpFilters_ReturnsInnerProviderMessagesAsync()
{
// Arrange
var innerProviderMock = new Mock<ChatHistoryProvider>();
var expectedMessages = new List<ChatMessage>
{
new(ChatRole.User, "Hello"),
new(ChatRole.Assistant, "Hi there!")
};
var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Test")]);
innerProviderMock
.Protected()
.Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(expectedMessages);
var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, x => x, x => x);
// Act
var result = (await filter.InvokingAsync(context, CancellationToken.None)).ToList();
// Assert
Assert.Equal(2, result.Count);
Assert.Equal("Hello", result[0].Text);
Assert.Equal("Hi there!", result[1].Text);
innerProviderMock
.Protected()
.Verify<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", Times.Once(), ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>());
}
[Fact]
public async Task InvokingAsync_WithInvokingFilter_AppliesFilterAsync()
{
// Arrange
var innerProviderMock = new Mock<ChatHistoryProvider>();
var innerMessages = new List<ChatMessage>
{
new(ChatRole.User, "Hello"),
new(ChatRole.Assistant, "Hi there!"),
new(ChatRole.User, "How are you?")
};
var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Test")]);
innerProviderMock
.Protected()
.Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(innerMessages);
// Filter to only user messages
IEnumerable<ChatMessage> InvokingFilter(IEnumerable<ChatMessage> msgs) => msgs.Where(m => m.Role == ChatRole.User);
var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, InvokingFilter);
// Act
var result = (await filter.InvokingAsync(context, CancellationToken.None)).ToList();
// Assert
Assert.Equal(2, result.Count);
Assert.All(result, msg => Assert.Equal(ChatRole.User, msg.Role));
innerProviderMock
.Protected()
.Verify<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", Times.Once(), ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>());
}
[Fact]
public async Task InvokingAsync_WithInvokingFilter_CanModifyMessagesAsync()
{
// Arrange
var innerProviderMock = new Mock<ChatHistoryProvider>();
var innerMessages = new List<ChatMessage>
{
new(ChatRole.User, "Hello"),
new(ChatRole.Assistant, "Hi there!")
};
var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Test")]);
innerProviderMock
.Protected()
.Setup<ValueTask<IEnumerable<ChatMessage>>>("InvokingCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokingContext>(), ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(innerMessages);
// Filter that transforms messages
IEnumerable<ChatMessage> InvokingFilter(IEnumerable<ChatMessage> msgs) =>
msgs.Select(m => new ChatMessage(m.Role, $"[FILTERED] {m.Text}"));
var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, InvokingFilter);
// Act
var result = (await filter.InvokingAsync(context, CancellationToken.None)).ToList();
// Assert
Assert.Equal(2, result.Count);
Assert.Equal("[FILTERED] Hello", result[0].Text);
Assert.Equal("[FILTERED] Hi there!", result[1].Text);
}
[Fact]
public async Task InvokedAsync_WithInvokedFilter_AppliesFilterAsync()
{
// Arrange
var innerProviderMock = new Mock<ChatHistoryProvider>();
List<ChatMessage> requestMessages =
[
new(ChatRole.System, "System") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "TestSource") } } },
new(ChatRole.User, "Hello"),
];
var responseMessages = new List<ChatMessage> { new(ChatRole.Assistant, "Response") };
var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages)
{
ResponseMessages = responseMessages
};
ChatHistoryProvider.InvokedContext? capturedContext = null;
innerProviderMock
.Protected()
.Setup<ValueTask>("InvokedCoreAsync", ItExpr.IsAny<ChatHistoryProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>())
.Callback<ChatHistoryProvider.InvokedContext, CancellationToken>((ctx, ct) => capturedContext = ctx)
.Returns(default(ValueTask));
// Filter that modifies the context
ChatHistoryProvider.InvokedContext InvokedFilter(ChatHistoryProvider.InvokedContext ctx)
{
var modifiedRequestMessages = ctx.RequestMessages.Where(x => x.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External).Select(m => new ChatMessage(m.Role, $"[FILTERED] {m.Text}")).ToList();
return new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, modifiedRequestMessages)
{
ResponseMessages = ctx.ResponseMessages,
InvokeException = ctx.InvokeException
};
}
var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, invokedMessagesFilter: InvokedFilter);
// Act
await filter.InvokedAsync(context, CancellationToken.None);
// Assert
Assert.NotNull(capturedContext);
Assert.Single(capturedContext.RequestMessages);
Assert.Equal("[FILTERED] Hello", capturedContext.RequestMessages.First().Text);
innerProviderMock
.Protected()
.Verify<ValueTask>("InvokedCoreAsync", Times.Once(), ItExpr.IsAny<ChatHistoryProvider.InvokedContext>(), ItExpr.IsAny<CancellationToken>());
}
}
@@ -83,7 +83,7 @@ public class InMemoryChatHistoryProviderTests
};
var provider = new InMemoryChatHistoryProvider();
provider.SetMessages(session, [providerMessages[0]]);
provider.SetMessages(session, providerMessages);
var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, requestMessages)
{
ResponseMessages = responseMessages
@@ -397,6 +397,91 @@ public class InMemoryChatHistoryProviderTests
await Assert.ThrowsAsync<ArgumentNullException>(() => provider.InvokingAsync(null!, CancellationToken.None).AsTask());
}
[Fact]
public async Task InvokedAsync_DefaultFilter_ExcludesChatHistoryMessagesAsync()
{
// Arrange
var session = CreateMockSession();
var provider = new InMemoryChatHistoryProvider();
var requestMessages = new List<ChatMessage>
{
new(ChatRole.User, "External message"),
new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } },
new(ChatRole.System, "From context provider") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "ContextSource") } } },
};
var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, requestMessages)
{
ResponseMessages = [new ChatMessage(ChatRole.Assistant, "Response")]
};
// Act
await provider.InvokedAsync(context, CancellationToken.None);
// Assert - ChatHistory message excluded, AIContextProvider message included
var messages = provider.GetMessages(session);
Assert.Equal(3, messages.Count);
Assert.Equal("External message", messages[0].Text);
Assert.Equal("From context provider", messages[1].Text);
Assert.Equal("Response", messages[2].Text);
}
[Fact]
public async Task InvokedAsync_CustomFilter_OverridesDefaultAsync()
{
// Arrange
var session = CreateMockSession();
var provider = new InMemoryChatHistoryProvider(new InMemoryChatHistoryProviderOptions
{
StorageInputMessageFilter = messages => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External)
});
var requestMessages = new List<ChatMessage>
{
new(ChatRole.User, "External message"),
new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } },
new(ChatRole.System, "From context provider") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "ContextSource") } } },
};
var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, requestMessages)
{
ResponseMessages = [new ChatMessage(ChatRole.Assistant, "Response")]
};
// Act
await provider.InvokedAsync(context, CancellationToken.None);
// Assert - Custom filter keeps only External messages (both ChatHistory and AIContextProvider excluded)
var messages = provider.GetMessages(session);
Assert.Equal(2, messages.Count);
Assert.Equal("External message", messages[0].Text);
Assert.Equal("Response", messages[1].Text);
}
[Fact]
public async Task InvokingAsync_OutputFilter_FiltersOutputMessagesAsync()
{
// Arrange
var session = CreateMockSession();
var provider = new InMemoryChatHistoryProvider(new InMemoryChatHistoryProviderOptions
{
RetrievalOutputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.User)
});
provider.SetMessages(session,
[
new ChatMessage(ChatRole.User, "User message"),
new ChatMessage(ChatRole.Assistant, "Assistant message"),
new ChatMessage(ChatRole.System, "System message")
]);
// Act
var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []);
var result = (await provider.InvokingAsync(context, CancellationToken.None)).ToList();
// Assert - Only user messages pass through the output filter
Assert.Single(result);
Assert.Equal("User message", result[0].Text);
}
public class TestAIContent(string testData) : AIContent
{
public string TestData => testData;