Skip to content
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
// Copyright (c) Microsoft. All rights reserved.

using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
using Microsoft.Shared.DiagnosticIds;
using Microsoft.Shared.Diagnostics;

namespace Microsoft.Agents.AI.Compaction;

/// <summary>
/// Provides extension methods for <see cref="CompactionStrategy"/>.
/// </summary>
[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
public static class ChatStrategyExtensions
{
/// <summary>
/// Returns an <see cref="IChatReducer"/> that applies this <see cref="CompactionStrategy"/> to reduce a list of messages.
/// </summary>
/// <param name="strategy">The compaction strategy to wrap as an <see cref="IChatReducer"/>.</param>
/// <returns>
/// An <see cref="IChatReducer"/> that, on each call to <see cref="IChatReducer.ReduceAsync"/>, builds a
/// <see cref="CompactionMessageIndex"/> from the supplied messages and applies the strategy's compaction logic,
/// returning the resulting included messages.
/// </returns>
/// <remarks>
/// This allows any <see cref="CompactionStrategy"/> to be used wherever an <see cref="IChatReducer"/> is expected,
/// bridging the compaction pipeline into systems bound to the <c>Microsoft.Extensions.AI</c> <see cref="IChatReducer"/> contract.
/// </remarks>
public static IChatReducer AsChatReducer(this CompactionStrategy strategy)
{
Throw.IfNull(strategy);

return new CompactionStrategyChatReducer(strategy);
}

/// <summary>
/// An <see cref="IChatReducer"/> adapter that delegates to a <see cref="CompactionStrategy"/>.
/// </summary>
private sealed class CompactionStrategyChatReducer : IChatReducer
{
private readonly CompactionStrategy _strategy;

public CompactionStrategyChatReducer(CompactionStrategy strategy)
{
this._strategy = strategy;
}

/// <inheritdoc/>
public async Task<IEnumerable<ChatMessage>> ReduceAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)
{
CompactionMessageIndex index = CompactionMessageIndex.Create([.. messages]);
await this._strategy.CompactAsync(index, cancellationToken: cancellationToken).ConfigureAwait(false);
return index.GetIncludedMessages();
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,159 @@
// Copyright (c) Microsoft. All rights reserved.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Agents.AI.Compaction;
using Microsoft.Extensions.AI;

namespace Microsoft.Agents.AI.UnitTests.Compaction;

/// <summary>
/// Contains tests for the <see cref="ChatStrategyExtensions"/> class.
/// </summary>
public class ChatStrategyExtensionsTests
{
[Fact]
public void AsChatReducerNullStrategyThrows()
{
// Act & Assert
Assert.Throws<ArgumentNullException>(() => ((CompactionStrategy)null!).AsChatReducer());
}

[Fact]
public void AsChatReducerReturnsIChatReducer()
{
// Arrange
ChatReducerCompactionStrategy strategy = new(new IdentityReducer(), CompactionTriggers.Always);

// Act
IChatReducer reducer = strategy.AsChatReducer();

// Assert
Assert.NotNull(reducer);
}

[Fact]
public async Task ReduceAsyncReturnsAllMessagesWhenStrategyDoesNotCompactAsync()
{
// Arrange — trigger never fires, so no compaction occurs
ChatReducerCompactionStrategy strategy = new(new IdentityReducer(), CompactionTriggers.Never);
IChatReducer reducer = strategy.AsChatReducer();

List<ChatMessage> messages =
[
new(ChatRole.User, "Hello"),
new(ChatRole.Assistant, "Hi!"),
];

// Act
IEnumerable<ChatMessage> result = await reducer.ReduceAsync(messages, CancellationToken.None);

// Assert
Assert.Equal(messages, result);
}

[Fact]
public async Task ReduceAsyncCompactsMessagesWhenStrategyFiresAsync()
{
// Arrange — reducer keeps only the last message
ChatReducerCompactionStrategy strategy = new(
new TakeLastReducer(1),
CompactionTriggers.Always);
IChatReducer reducer = strategy.AsChatReducer();

List<ChatMessage> messages =
[
new(ChatRole.User, "First"),
new(ChatRole.Assistant, "Response 1"),
new(ChatRole.User, "Second"),
];

// Act
IEnumerable<ChatMessage> result = await reducer.ReduceAsync(messages, CancellationToken.None);

// Assert
List<ChatMessage> resultList = [.. result];
Assert.Single(resultList);
Assert.Equal("Second", resultList[0].Text);
}

[Fact]
public async Task ReduceAsyncPassesCancellationTokenToStrategyAsync()
{
// Arrange
using CancellationTokenSource cts = new();
CancellationToken capturedToken = default;

CapturingReducer capturingReducer = new(token => capturedToken = token);
ChatReducerCompactionStrategy strategy = new(capturingReducer, CompactionTriggers.Always);
IChatReducer reducer = strategy.AsChatReducer();

List<ChatMessage> messages =
[
new(ChatRole.User, "Hello"),
new(ChatRole.User, "World"),
];

// Act
await reducer.ReduceAsync(messages, cts.Token);

// Assert
Assert.Equal(cts.Token, capturedToken);
}

[Fact]
public async Task ReduceAsyncEmptyMessagesReturnsEmptyAsync()
{
// Arrange
ChatReducerCompactionStrategy strategy = new(new IdentityReducer(), CompactionTriggers.Always);
IChatReducer reducer = strategy.AsChatReducer();

// Act
IEnumerable<ChatMessage> result = await reducer.ReduceAsync([], CancellationToken.None);

// Assert
Assert.Empty(result);
}

/// <summary>
/// An <see cref="IChatReducer"/> that returns messages unchanged.
/// </summary>
private sealed class IdentityReducer : IChatReducer
{
public Task<IEnumerable<ChatMessage>> ReduceAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)
=> Task.FromResult(messages);
}

/// <summary>
/// An <see cref="IChatReducer"/> that keeps only the last <c>n</c> messages.
/// </summary>
private sealed class TakeLastReducer : IChatReducer
{
private readonly int _count;

public TakeLastReducer(int count) => this._count = count;

public Task<IEnumerable<ChatMessage>> ReduceAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)
=> Task.FromResult(messages.Reverse().Take(this._count));
}

/// <summary>
/// An <see cref="IChatReducer"/> that captures the <see cref="CancellationToken"/> passed to <see cref="ReduceAsync"/>.
/// </summary>
private sealed class CapturingReducer : IChatReducer
{
private readonly Action<CancellationToken> _capture;

public CapturingReducer(Action<CancellationToken> capture) => this._capture = capture;

public Task<IEnumerable<ChatMessage>> ReduceAsync(IEnumerable<ChatMessage> messages, CancellationToken cancellationToken = default)
{
this._capture(cancellationToken);
IEnumerable<ChatMessage> reducedMessages = [messages.Reverse().First()];
return Task.FromResult(reducedMessages);
}
}
}
Loading