mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Port ContextualFunctionProvider from SK
This commit is contained in:
westey
+311
@@ -0,0 +1,311 @@
|
||||
// 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.Functions;
|
||||
using Microsoft.Extensions.AI;
|
||||
using Microsoft.Extensions.VectorData;
|
||||
using Moq;
|
||||
|
||||
namespace Microsoft.Agents.AI.UnitTests.Functions;
|
||||
|
||||
/// <summary>
|
||||
/// Contains unit tests for the <see cref="ContextualFunctionProvider"/> class.
|
||||
/// </summary>
|
||||
public sealed class ContextualFunctionProviderTests
|
||||
{
|
||||
private readonly Mock<VectorStore> _vectorStoreMock;
|
||||
private readonly Mock<VectorStoreCollection<object, Dictionary<string, object?>>> _collectionMock;
|
||||
|
||||
public ContextualFunctionProviderTests()
|
||||
{
|
||||
this._vectorStoreMock = new Mock<VectorStore>(MockBehavior.Strict);
|
||||
this._collectionMock = new Mock<VectorStoreCollection<object, Dictionary<string, object?>>>(MockBehavior.Strict);
|
||||
|
||||
this._vectorStoreMock
|
||||
.Setup(vs => vs.GetDynamicCollection(It.IsAny<string>(), It.IsAny<VectorStoreCollectionDefinition>()))
|
||||
.Returns(this._collectionMock.Object);
|
||||
|
||||
this._collectionMock
|
||||
.Setup(c => c.CollectionExistsAsync(It.IsAny<CancellationToken>()))
|
||||
.ReturnsAsync(true);
|
||||
|
||||
this._collectionMock
|
||||
.Setup(c => c.EnsureCollectionExistsAsync(It.IsAny<CancellationToken>()))
|
||||
.Returns(Task.CompletedTask);
|
||||
|
||||
this._collectionMock
|
||||
.Setup(c => c.UpsertAsync(It.IsAny<IEnumerable<Dictionary<string, object?>>>(), It.IsAny<CancellationToken>()))
|
||||
.Returns(Task.CompletedTask);
|
||||
|
||||
this._collectionMock
|
||||
.Setup(c => c.SearchAsync(It.IsAny<string>(), It.IsAny<int>(), null, It.IsAny<CancellationToken>()))
|
||||
.Returns(AsyncEnumerable.Empty<VectorSearchResult<Dictionary<string, object?>>>());
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public void Constructor_ShouldThrow_OnInvalidArguments()
|
||||
{
|
||||
// Arrange
|
||||
var vectorStore = new Mock<VectorStore>().Object;
|
||||
var functions = new List<AIFunction> { CreateFunction("f1") };
|
||||
|
||||
// Act & Assert
|
||||
Assert.Throws<ArgumentNullException>(() => new ContextualFunctionProvider(null!, 1, functions, 3));
|
||||
Assert.Throws<ArgumentOutOfRangeException>(() => new ContextualFunctionProvider(vectorStore, 0, functions, 3));
|
||||
Assert.Throws<ArgumentNullException>(() => new ContextualFunctionProvider(vectorStore, 1, null!, 3));
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task Invoking_ShouldVectorizeFunctions_Once_Async()
|
||||
{
|
||||
// Arrange
|
||||
var function = CreateFunction("f1", "desc");
|
||||
var functions = new List<AIFunction> { function };
|
||||
|
||||
this._collectionMock
|
||||
.Setup(c => c.UpsertAsync(It.IsAny<IEnumerable<Dictionary<string, object?>>>(), It.IsAny<CancellationToken>()))
|
||||
.Returns(Task.CompletedTask);
|
||||
|
||||
var provider = new ContextualFunctionProvider(
|
||||
vectorStore: this._vectorStoreMock.Object,
|
||||
vectorDimensions: 1536,
|
||||
functions: functions,
|
||||
maxNumberOfFunctions: 5);
|
||||
|
||||
var messages = new List<ChatMessage> { new() { Contents = [new TextContent("hello")] } };
|
||||
var context = new AIContextProvider.InvokingContext(messages);
|
||||
|
||||
// Act
|
||||
await provider.InvokingAsync(context);
|
||||
await provider.InvokingAsync(context);
|
||||
|
||||
// Assert
|
||||
this._collectionMock.Verify(
|
||||
c => c.UpsertAsync(It.IsAny<IEnumerable<Dictionary<string, object?>>>(), It.IsAny<CancellationToken>()),
|
||||
Times.Once);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task Invoking_ShouldReturnRelevantFunctions_Async()
|
||||
{
|
||||
// Arrange
|
||||
var function = CreateFunction("f1", "desc");
|
||||
var functions = new List<AIFunction> { function };
|
||||
|
||||
var searchResult = new VectorSearchResult<Dictionary<string, object?>>(
|
||||
new Dictionary<string, object?>
|
||||
{
|
||||
["Name"] = function.Name,
|
||||
["Description"] = function.Description
|
||||
},
|
||||
0.99f
|
||||
);
|
||||
|
||||
this._collectionMock
|
||||
.Setup(c => c.SearchAsync(It.IsAny<string>(), It.IsAny<int>(), null, It.IsAny<CancellationToken>()))
|
||||
.Returns(new[] { searchResult }.ToAsyncEnumerable());
|
||||
|
||||
var provider = new ContextualFunctionProvider(
|
||||
vectorStore: this._vectorStoreMock.Object,
|
||||
vectorDimensions: 1536,
|
||||
functions: functions,
|
||||
maxNumberOfFunctions: 5);
|
||||
|
||||
var messages = new List<ChatMessage> { new() { Contents = [new TextContent("context")] } };
|
||||
var context = new AIContextProvider.InvokingContext(messages);
|
||||
|
||||
// Act
|
||||
var result = await provider.InvokingAsync(context);
|
||||
|
||||
// Assert
|
||||
Assert.NotNull(result);
|
||||
Assert.NotNull(result.Tools);
|
||||
Assert.Single(result.Tools);
|
||||
Assert.Equal("f1", result.Tools[0].Name);
|
||||
this._collectionMock.Verify(
|
||||
c => c.SearchAsync("context", 5, null, It.IsAny<CancellationToken>()),
|
||||
Times.Once);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task BuildContext_ShouldUseContextEmbeddingValueProvider_Async()
|
||||
{
|
||||
// Arrange
|
||||
var functions = new List<AIFunction> { CreateFunction("f1") };
|
||||
var options = new ContextualFunctionProviderOptions
|
||||
{
|
||||
NumberOfRecentMessagesInContext = 3,
|
||||
ContextEmbeddingValueProvider = (recentMessages, newMessages, _) =>
|
||||
{
|
||||
Assert.Equal(3, recentMessages.Count());
|
||||
Assert.Single(newMessages);
|
||||
return Task.FromResult("custom context");
|
||||
}
|
||||
};
|
||||
|
||||
var provider = new ContextualFunctionProvider(
|
||||
vectorStore: this._vectorStoreMock.Object,
|
||||
vectorDimensions: 1536,
|
||||
functions: functions,
|
||||
maxNumberOfFunctions: 5,
|
||||
options: options);
|
||||
|
||||
var message1 = new ChatMessage() { Contents = [new TextContent("msg1")] };
|
||||
var message2 = new ChatMessage() { Contents = [new TextContent("msg2")] };
|
||||
var message3 = new ChatMessage() { Contents = [new TextContent("msg3")] };
|
||||
var message4 = new ChatMessage() { Contents = [new TextContent("msg4")] };
|
||||
var message5 = new ChatMessage() { Contents = [new TextContent("msg5")] };
|
||||
|
||||
// Simulate previous invocations to populate recent messages
|
||||
await provider.InvokedAsync(new AIContextProvider.InvokedContext([message1], null) { ResponseMessages = [] });
|
||||
await provider.InvokedAsync(new AIContextProvider.InvokedContext([message2], null) { ResponseMessages = [] });
|
||||
await provider.InvokedAsync(new AIContextProvider.InvokedContext([message3], null) { ResponseMessages = [] });
|
||||
await provider.InvokedAsync(new AIContextProvider.InvokedContext([message4], null) { ResponseMessages = [] });
|
||||
|
||||
var messages = new List<ChatMessage> { message5 };
|
||||
var context = new AIContextProvider.InvokingContext(messages);
|
||||
|
||||
// Act
|
||||
await provider.InvokingAsync(context);
|
||||
|
||||
// Assert
|
||||
this._collectionMock.Verify(
|
||||
c => c.SearchAsync("custom context", It.IsAny<int>(), null, It.IsAny<CancellationToken>()),
|
||||
Times.Once);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task BuildContext_ShouldConcatenateMessages_Async()
|
||||
{
|
||||
// Arrange
|
||||
var functions = new List<AIFunction> { CreateFunction("f1") };
|
||||
var options = new ContextualFunctionProviderOptions
|
||||
{
|
||||
NumberOfRecentMessagesInContext = 3
|
||||
};
|
||||
|
||||
var provider = new ContextualFunctionProvider(
|
||||
vectorStore: this._vectorStoreMock.Object,
|
||||
vectorDimensions: 1536,
|
||||
functions: functions,
|
||||
maxNumberOfFunctions: 5,
|
||||
options: options);
|
||||
|
||||
var message1 = new ChatMessage() { Contents = [new TextContent("msg1")] };
|
||||
var message2 = new ChatMessage() { Contents = [new TextContent("msg2")] };
|
||||
var message3 = new ChatMessage() { Contents = [new TextContent("msg3")] };
|
||||
var message4 = new ChatMessage() { Contents = [new TextContent("msg4")] };
|
||||
var message5 = new ChatMessage() { Contents = [new TextContent("msg5")] };
|
||||
|
||||
// Simulate previous invocations to populate recent messages
|
||||
await provider.InvokedAsync(new AIContextProvider.InvokedContext([message1], null) { ResponseMessages = [] });
|
||||
await provider.InvokedAsync(new AIContextProvider.InvokedContext([message2], null) { ResponseMessages = [] });
|
||||
await provider.InvokedAsync(new AIContextProvider.InvokedContext([message3], null) { ResponseMessages = [] });
|
||||
await provider.InvokedAsync(new AIContextProvider.InvokedContext([message4], null) { ResponseMessages = [] });
|
||||
|
||||
// Act
|
||||
var invokingContext = new AIContextProvider.InvokingContext([message5]);
|
||||
var context = await provider.InvokingAsync(invokingContext);
|
||||
|
||||
// Assert
|
||||
var expected = string.Join(Environment.NewLine, ["msg2", "msg3", "msg4", "msg5"]);
|
||||
this._collectionMock.Verify(c => c.SearchAsync(expected, It.IsAny<int>(), null, It.IsAny<CancellationToken>()), Times.Once);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task BuildContext_ShouldUseEmbeddingValueProvider_Async()
|
||||
{
|
||||
// Arrange
|
||||
List<Dictionary<string, object?>>? upsertedRecords = null;
|
||||
this._collectionMock
|
||||
.Setup(c => c.UpsertAsync(It.IsAny<IEnumerable<Dictionary<string, object?>>>(), It.IsAny<CancellationToken>()))
|
||||
.Callback<IEnumerable<Dictionary<string, object?>>, CancellationToken>((records, _) => upsertedRecords = records.ToList())
|
||||
.Returns(Task.CompletedTask);
|
||||
|
||||
var functions = new List<AIFunction> { CreateFunction("f1", "desc1") };
|
||||
var options = new ContextualFunctionProviderOptions
|
||||
{
|
||||
EmbeddingValueProvider = (func, ct) => Task.FromResult($"custom embedding for {func.Name}:{func.Description}")
|
||||
};
|
||||
|
||||
var provider = new ContextualFunctionProvider(
|
||||
vectorStore: this._vectorStoreMock.Object,
|
||||
vectorDimensions: 1536,
|
||||
functions: functions,
|
||||
maxNumberOfFunctions: 5,
|
||||
options: options);
|
||||
|
||||
var messages = new List<ChatMessage>
|
||||
{
|
||||
new() { Contents = [new TextContent("ignored")] }
|
||||
};
|
||||
var context = new AIContextProvider.InvokingContext(messages);
|
||||
|
||||
// Act
|
||||
await provider.InvokingAsync(context);
|
||||
|
||||
// Assert
|
||||
Assert.NotNull(upsertedRecords);
|
||||
var embeddingSource = upsertedRecords!.SelectMany(r => r).FirstOrDefault(kv => kv.Key == "Embedding").Value as string;
|
||||
Assert.Equal("custom embedding for f1:desc1", embeddingSource);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task ContextEmbeddingValueProvider_ReceivesRecentAndNewMessages_Async()
|
||||
{
|
||||
// Arrange
|
||||
var functions = new List<AIFunction> { CreateFunction("f1") };
|
||||
|
||||
IEnumerable<ChatMessage>? capturedRecentMessages = null;
|
||||
IEnumerable<ChatMessage>? capturedNewMessages = null;
|
||||
|
||||
var options = new ContextualFunctionProviderOptions
|
||||
{
|
||||
NumberOfRecentMessagesInContext = 2,
|
||||
ContextEmbeddingValueProvider = (recentMessages, newMessages, ct) =>
|
||||
{
|
||||
capturedRecentMessages = recentMessages;
|
||||
capturedNewMessages = newMessages;
|
||||
|
||||
return Task.FromResult("context");
|
||||
}
|
||||
};
|
||||
|
||||
var provider = new ContextualFunctionProvider(
|
||||
vectorStore: this._vectorStoreMock.Object,
|
||||
vectorDimensions: 1536,
|
||||
functions: functions,
|
||||
maxNumberOfFunctions: 5,
|
||||
options: options);
|
||||
|
||||
// Add more messages than the number of messages to keep
|
||||
await provider.InvokedAsync(new AIContextProvider.InvokedContext([new() { Contents = [new TextContent("msg1")] }], null) { ResponseMessages = [] });
|
||||
await provider.InvokedAsync(new AIContextProvider.InvokedContext([new() { Contents = [new TextContent("msg2")] }], null) { ResponseMessages = [] });
|
||||
await provider.InvokedAsync(new AIContextProvider.InvokedContext([new() { Contents = [new TextContent("msg3")] }], null) { ResponseMessages = [] });
|
||||
|
||||
// Act
|
||||
var invokingContext = new AIContextProvider.InvokingContext([
|
||||
new() { Contents = [new TextContent("msg4")] },
|
||||
new() { Contents = [new TextContent("msg5")] }
|
||||
]);
|
||||
await provider.InvokingAsync(invokingContext);
|
||||
|
||||
// Assert
|
||||
Assert.NotNull(capturedRecentMessages);
|
||||
Assert.Equal("msg2", capturedRecentMessages.ElementAt(0).Text);
|
||||
Assert.Equal("msg3", capturedRecentMessages.ElementAt(1).Text);
|
||||
|
||||
Assert.NotNull(capturedNewMessages);
|
||||
Assert.Equal("msg4", capturedNewMessages.ElementAt(0).Text);
|
||||
Assert.Equal("msg5", capturedNewMessages.ElementAt(1).Text);
|
||||
}
|
||||
|
||||
private static AIFunction CreateFunction(string name, string description = "")
|
||||
{
|
||||
return AIFunctionFactory.Create(() => { }, name, description);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
// 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.Functions;
|
||||
using Microsoft.Extensions.AI;
|
||||
using Microsoft.Extensions.VectorData;
|
||||
using Moq;
|
||||
|
||||
namespace Microsoft.Agents.AI.UnitTests.Functions;
|
||||
|
||||
/// <summary>
|
||||
/// Contains unit tests for the <see cref="FunctionStore"/> class.
|
||||
/// </summary>
|
||||
public sealed class FunctionStoreTests
|
||||
{
|
||||
private readonly Mock<VectorStore> _vectorStoreMock;
|
||||
private readonly Mock<VectorStoreCollection<object, Dictionary<string, object?>>> _collectionMock;
|
||||
|
||||
public FunctionStoreTests()
|
||||
{
|
||||
this._vectorStoreMock = new Mock<VectorStore>(MockBehavior.Strict);
|
||||
this._collectionMock = new Mock<VectorStoreCollection<object, Dictionary<string, object?>>>(MockBehavior.Strict);
|
||||
|
||||
this._vectorStoreMock
|
||||
.Setup(vs => vs.GetDynamicCollection(It.IsAny<string>(), It.IsAny<VectorStoreCollectionDefinition>()))
|
||||
.Returns(this._collectionMock.Object);
|
||||
|
||||
this._collectionMock
|
||||
.Setup(c => c.CollectionExistsAsync(It.IsAny<CancellationToken>()))
|
||||
.ReturnsAsync(true);
|
||||
|
||||
this._collectionMock
|
||||
.Setup(c => c.EnsureCollectionExistsAsync(It.IsAny<CancellationToken>()))
|
||||
.Returns(Task.CompletedTask);
|
||||
|
||||
this._collectionMock
|
||||
.Setup(c => c.UpsertAsync(It.IsAny<IEnumerable<Dictionary<string, object?>>>(), It.IsAny<CancellationToken>()))
|
||||
.Returns(Task.CompletedTask);
|
||||
|
||||
this._collectionMock
|
||||
.Setup(c => c.SearchAsync(It.IsAny<string>(), It.IsAny<int>(), null, It.IsAny<CancellationToken>()))
|
||||
.Returns(AsyncEnumerable.Empty<VectorSearchResult<Dictionary<string, object?>>>());
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public void Constructor_ShouldThrowOnInvalidArguments()
|
||||
{
|
||||
var functions = new List<AIFunction> { CreateFunction("f1") };
|
||||
|
||||
Assert.Throws<ArgumentNullException>(() => new FunctionStore(null!, "col", 1, functions, 3));
|
||||
Assert.Throws<ArgumentException>(() => new FunctionStore(this._vectorStoreMock.Object, "", 1, functions, 3));
|
||||
Assert.Throws<ArgumentOutOfRangeException>(() => new FunctionStore(this._vectorStoreMock.Object, "col", 0, functions, 3));
|
||||
Assert.Throws<ArgumentNullException>(() => new FunctionStore(this._vectorStoreMock.Object, "col", 1, null!, 3));
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task SaveAsync_ShouldUpsertFunctions_Async()
|
||||
{
|
||||
// Arrange
|
||||
var functions = new List<AIFunction>
|
||||
{
|
||||
CreateFunction("f1", "desc1"),
|
||||
CreateFunction("f2", "desc2")
|
||||
};
|
||||
|
||||
this._collectionMock.Setup(c => c.UpsertAsync(It.IsAny<IEnumerable<Dictionary<string, object?>>>(), It.IsAny<CancellationToken>()))
|
||||
.Returns(Task.CompletedTask)
|
||||
.Verifiable();
|
||||
|
||||
var store = new FunctionStore(this._vectorStoreMock.Object, "col", 3, functions, 3);
|
||||
|
||||
// Act
|
||||
await store.SaveAsync();
|
||||
|
||||
// Assert
|
||||
this._collectionMock.Verify(c => c.EnsureCollectionExistsAsync(It.IsAny<CancellationToken>()), Times.Once);
|
||||
this._collectionMock.Verify(c => c.UpsertAsync(It.Is<IEnumerable<Dictionary<string, object?>>>(records =>
|
||||
records.Count() == 2 &&
|
||||
records.Any(r => (r["Name"] as string) == "f1") &&
|
||||
records.Any(r => (r["Name"] as string) == "f2")
|
||||
), It.IsAny<CancellationToken>()), Times.Once);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task SearchAsync_ShouldReturnMatchingFunctionsAsync()
|
||||
{
|
||||
// Arrange
|
||||
var functions = new List<AIFunction>
|
||||
{
|
||||
CreateFunction("f1", "desc1"),
|
||||
CreateFunction("f2", "desc2"),
|
||||
CreateFunction("f3", "desc3")
|
||||
};
|
||||
|
||||
var searchResults = new List<VectorSearchResult<Dictionary<string, object?>>>
|
||||
{
|
||||
new(new Dictionary<string, object?> { ["Name"] = "f3" }, 0.3),
|
||||
new(new Dictionary<string, object?> { ["Name"] = "f2" }, 0.2),
|
||||
new(new Dictionary<string, object?> { ["Name"] = "f1" }, 0.1)
|
||||
};
|
||||
|
||||
this._collectionMock.Setup(c => c.SearchAsync(It.IsAny<string>(), It.IsAny<int>(), null, It.IsAny<CancellationToken>()))
|
||||
.Returns(searchResults.ToAsyncEnumerable());
|
||||
|
||||
var store = new FunctionStore(this._vectorStoreMock.Object, "col", 3, functions, 3);
|
||||
|
||||
// Act
|
||||
var result = await store.SearchAsync("desc3");
|
||||
|
||||
// Assert
|
||||
var resultList = result.ToList();
|
||||
Assert.Equal(3, resultList.Count);
|
||||
Assert.Equal("f3", resultList[0].Name);
|
||||
Assert.Equal("f2", resultList[1].Name);
|
||||
Assert.Equal("f1", resultList[2].Name);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task SearchAsync_ShouldThrowIfCollectionDoesNotExistAsync()
|
||||
{
|
||||
// Arrange
|
||||
var functions = new List<AIFunction> { CreateFunction("f1") };
|
||||
|
||||
this._collectionMock.Setup(c => c.CollectionExistsAsync(It.IsAny<CancellationToken>()))
|
||||
.ReturnsAsync(false);
|
||||
|
||||
var store = new FunctionStore(this._vectorStoreMock.Object, "col", 3, functions, 3);
|
||||
|
||||
// Act & Assert
|
||||
await Assert.ThrowsAsync<InvalidOperationException>(() => store.SearchAsync("query"));
|
||||
}
|
||||
|
||||
private static AIFunction CreateFunction(string name, string description = "desc")
|
||||
{
|
||||
return AIFunctionFactory.Create(() => { }, name, description);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user