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
@@ -0,0 +1,163 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
using System;
|
||||
using System.Collections.Concurrent;
|
||||
using System.Collections.Generic;
|
||||
using System.Linq;
|
||||
using System.Threading;
|
||||
using System.Threading.Tasks;
|
||||
using Microsoft.Extensions.AI;
|
||||
using Microsoft.Extensions.Logging;
|
||||
using Microsoft.Extensions.VectorData;
|
||||
using Microsoft.Shared.Diagnostics;
|
||||
|
||||
namespace Microsoft.Agents.AI.Functions;
|
||||
|
||||
/// <summary>
|
||||
/// Represents a contextual function provider that performs RAG (Retrieval-Augmented Generation) on the provided functions to identify
|
||||
/// the most relevant functions for the current context. The provider vectorizes the provided function names and descriptions
|
||||
/// and stores them in the specified vector store, allowing for a vector search to find the most relevant
|
||||
/// functions for a given context and provide the functions to the AI model/agent.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// <list type="bullet">
|
||||
/// <item>
|
||||
/// The provider is designed to work with in-memory vector stores. Using other vector stores
|
||||
/// will require the data synchronization and data lifetime management to be done by the caller.
|
||||
/// </item>
|
||||
/// <item>
|
||||
/// The in-memory vector store is supposed to be created per provider and not shared between providers
|
||||
/// unless each provider uses a different collection name. Not following this may lead to a situation
|
||||
/// where one provider identifies a function belonging to another provider as relevant and, as a result,
|
||||
/// an attempt to access it by the first provider will fail because the function is not registered with it.
|
||||
/// </item>
|
||||
/// <item>
|
||||
/// The provider uses function name as a key for the records and as such the specified vector store
|
||||
/// should support record keys of string type.
|
||||
/// </item>
|
||||
/// </list>
|
||||
/// </remarks>
|
||||
public sealed class ContextualFunctionProvider : AIContextProvider
|
||||
{
|
||||
private readonly FunctionStore _functionStore;
|
||||
private readonly ConcurrentQueue<ChatMessage> _recentMessages = [];
|
||||
private readonly ContextualFunctionProviderOptions _options;
|
||||
private bool _areFunctionsVectorized;
|
||||
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="ContextualFunctionProvider"/> class.
|
||||
/// </summary>
|
||||
/// <param name="vectorStore">An instance of a vector store.</param>
|
||||
/// <param name="vectorDimensions">The number of dimensions to use for the memory embeddings.</param>
|
||||
/// <param name="functions">The functions to vectorize and store for searching related functions.</param>
|
||||
/// <param name="maxNumberOfFunctions">The maximum number of relevant functions to retrieve from the vector store.</param>
|
||||
/// <param name="options">Further optional settings for configuring the provider.</param>
|
||||
/// <param name="loggerFactory">The logger factory to use for logging. If not provided, no logging will be performed.</param>
|
||||
public ContextualFunctionProvider(
|
||||
VectorStore vectorStore,
|
||||
int vectorDimensions,
|
||||
IEnumerable<AIFunction> functions,
|
||||
int maxNumberOfFunctions,
|
||||
ContextualFunctionProviderOptions? options = null,
|
||||
ILoggerFactory? loggerFactory = null)
|
||||
{
|
||||
Throw.IfNull(vectorStore);
|
||||
Throw.IfLessThan(vectorDimensions, 1, "Vector dimensions must be greater than 0");
|
||||
Throw.IfNull(functions);
|
||||
Throw.IfLessThan(maxNumberOfFunctions, 1, "Max number of functions must be greater than 0");
|
||||
|
||||
this._options = options ?? new ContextualFunctionProviderOptions();
|
||||
Throw.IfLessThan(this._options.NumberOfRecentMessagesInContext, 1, "Number of recent messages to include into context must be greater than 0");
|
||||
|
||||
this._functionStore = new FunctionStore(
|
||||
vectorStore,
|
||||
string.IsNullOrWhiteSpace(this._options.CollectionName) ? "functions" : this._options.CollectionName,
|
||||
vectorDimensions,
|
||||
functions,
|
||||
maxNumberOfFunctions,
|
||||
loggerFactory,
|
||||
options: new()
|
||||
{
|
||||
EmbeddingValueProvider = this._options.EmbeddingValueProvider,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
/// <inheritdoc />
|
||||
public override async ValueTask<AIContext> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default)
|
||||
{
|
||||
Throw.IfNull(context);
|
||||
|
||||
// Vectorize the functions if they are not already vectorized
|
||||
if (!this._areFunctionsVectorized)
|
||||
{
|
||||
await this._functionStore.SaveAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
this._areFunctionsVectorized = true;
|
||||
}
|
||||
|
||||
// Build the search context
|
||||
var searchContext = await this.BuildContextAsync(context.RequestMessages, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// Get the function relevant to the context
|
||||
var functions = await this._functionStore
|
||||
.SearchAsync(searchContext, cancellationToken: cancellationToken)
|
||||
.ConfigureAwait(false);
|
||||
|
||||
return new AIContext { Tools = [.. functions] };
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public override ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default)
|
||||
{
|
||||
Throw.IfNull(context);
|
||||
|
||||
// Add the request and response messages to the recent messages queue
|
||||
foreach (var message in context.RequestMessages)
|
||||
{
|
||||
this._recentMessages.Enqueue(message);
|
||||
}
|
||||
|
||||
if (context.ResponseMessages is not null)
|
||||
{
|
||||
foreach (var message in context.ResponseMessages)
|
||||
{
|
||||
this._recentMessages.Enqueue(message);
|
||||
}
|
||||
}
|
||||
|
||||
// If there are more messages than the configured limit, remove the oldest ones
|
||||
while (this._recentMessages.Count > this._options.NumberOfRecentMessagesInContext)
|
||||
{
|
||||
this._recentMessages.TryDequeue(out _);
|
||||
}
|
||||
|
||||
return default;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Builds the context from chat messages.
|
||||
/// </summary>
|
||||
/// <param name="newMessages">The new messages.</param>
|
||||
/// <param name="cancellationToken">The cancellation token to use for cancellation.</param>
|
||||
private async Task<string> BuildContextAsync(IEnumerable<ChatMessage> newMessages, CancellationToken cancellationToken)
|
||||
{
|
||||
if (this._options.ContextEmbeddingValueProvider is not null)
|
||||
{
|
||||
// Ensure we only take the recent messages up to the configured limit
|
||||
var recentMessages = this._recentMessages
|
||||
.Skip(Math.Max(0, this._recentMessages.Count - this._options.NumberOfRecentMessagesInContext));
|
||||
|
||||
return await this._options.ContextEmbeddingValueProvider.Invoke(recentMessages, newMessages, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
// Build context from the recent messages that already include the new messages
|
||||
return string.Join(
|
||||
Environment.NewLine,
|
||||
this._recentMessages
|
||||
.Skip(Math.Max(0, this._recentMessages.Count - this._options.NumberOfRecentMessagesInContext))
|
||||
.Concat(newMessages)
|
||||
.Where(m => !string.IsNullOrWhiteSpace(m?.Text))
|
||||
.Select(m => m.Text));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
using System;
|
||||
using System.Collections.Generic;
|
||||
using System.Threading;
|
||||
using System.Threading.Tasks;
|
||||
using Microsoft.Extensions.AI;
|
||||
|
||||
namespace Microsoft.Agents.AI.Functions;
|
||||
|
||||
/// <summary>
|
||||
/// Options for the <see cref="ContextualFunctionProvider"/>.
|
||||
/// </summary>
|
||||
public sealed class ContextualFunctionProviderOptions
|
||||
{
|
||||
/// <summary>
|
||||
/// Gets or sets the collection name to use for storing and retrieving functions.
|
||||
/// </summary>
|
||||
/// <value>If not set, the default value "functions" will be used.</value>
|
||||
public string? CollectionName { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the number of recent messages(messages from previous model/agent invocations) the provider uses to form a context.
|
||||
/// The provider collects all messages from all model/agent invocations, up to this number,
|
||||
/// and prepends them to the new messages of the current model/agent invocation to build a context.
|
||||
/// While collecting new messages, the provider will remove the oldest messages
|
||||
/// to keep the number of recent messages within the specified limit.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Using the recent messages together with the new messages can be very useful
|
||||
/// in cases where the model/agent is prompted to perform a task that requires details from
|
||||
/// previous invocation(s). For example, if the agent is asked to provision an Azure resource in the first
|
||||
/// invocation and deploy the resource in the second invocation, the second invocation will need
|
||||
/// information about the provisioned resource in the first invocation to deploy it.
|
||||
/// </remarks>
|
||||
public int NumberOfRecentMessagesInContext { get; set; } = 2;
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets a callback function that returns a value used to create a context embedding. The value is vectorized,
|
||||
/// and the resulting vector is used to perform vector searches for functions relevant to the context.
|
||||
/// If not provided, the default behavior is to concatenate the non-empty messages into a single string,
|
||||
/// separated by a new line.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The callback receives three parameters:
|
||||
/// `recentMessages` - messages from the previous model/agent invocations.
|
||||
/// `newMessages` - the new messages of the current model/agent invocation.
|
||||
/// `cancellationToken` - a cancellation token that can be used to cancel the operation.
|
||||
/// </remarks>
|
||||
public Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>, CancellationToken, Task<string>>? ContextEmbeddingValueProvider { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets a callback function that returns a value used to create a function embedding. The value is vectorized,
|
||||
/// and the resulting vector is stored in the vector store for use in vector searches for functions relevant
|
||||
/// to the context.
|
||||
/// If not provided, the default behavior is to concatenate the function name and description into a single string.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The callback receives two parameters:
|
||||
/// `function` - the function to get embedding value for.
|
||||
/// `cancellationToken` - a cancellation token that can be used to cancel the operation.
|
||||
/// </remarks>
|
||||
public Func<AIFunction, CancellationToken, Task<string>>? EmbeddingValueProvider { get; set; }
|
||||
}
|
||||
@@ -0,0 +1,185 @@
|
||||
// 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 Microsoft.Extensions.Logging;
|
||||
using Microsoft.Extensions.Logging.Abstractions;
|
||||
using Microsoft.Extensions.VectorData;
|
||||
using Microsoft.Shared.Diagnostics;
|
||||
|
||||
namespace Microsoft.Agents.AI.Functions;
|
||||
|
||||
/// <summary>
|
||||
/// Represents a vector store for <see cref="AIFunction"/> objects where the function name and description can be used for similarity searches.
|
||||
/// </summary>
|
||||
internal sealed class FunctionStore
|
||||
{
|
||||
private readonly VectorStore _vectorStore;
|
||||
private readonly Dictionary<string, AIFunction> _functionByName;
|
||||
private readonly string _collectionName;
|
||||
private readonly int _maxNumberOfFunctions;
|
||||
private readonly ILogger _logger;
|
||||
private readonly FunctionStoreOptions _options;
|
||||
private readonly VectorStoreCollection<object, Dictionary<string, object?>> _collection;
|
||||
private bool _isCollectionExistenceAsserted;
|
||||
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="FunctionStore"/> class.
|
||||
/// </summary>
|
||||
/// <param name="vectorStore">The vector store to use for storing functions.</param>
|
||||
/// <param name="collectionName">The name of the collection to use for storing and retrieving functions.</param>
|
||||
/// <param name="vectorDimensions">The number of dimensions to use for the memory embeddings.</param>
|
||||
/// <param name="functions">The functions to vectorize and store for searching related functions.</param>
|
||||
/// <param name="maxNumberOfFunctions">The maximum number of relevant functions to retrieve from the vector store.</param>
|
||||
/// <param name="loggerFactory">The logger factory to use for logging. If not provided, no logging will be performed.</param>
|
||||
/// <param name="options">The options to use for the function store.</param>
|
||||
internal FunctionStore(
|
||||
VectorStore vectorStore,
|
||||
string collectionName,
|
||||
int vectorDimensions,
|
||||
IEnumerable<AIFunction> functions,
|
||||
int maxNumberOfFunctions,
|
||||
ILoggerFactory? loggerFactory = default,
|
||||
FunctionStoreOptions? options = null)
|
||||
{
|
||||
Throw.IfNull(vectorStore);
|
||||
Throw.IfNullOrWhitespace(collectionName);
|
||||
Throw.IfLessThan(vectorDimensions, 1, "Vector dimensions must be greater than 0");
|
||||
Throw.IfNull(functions);
|
||||
Throw.IfLessThan(maxNumberOfFunctions, 1, "Max number of functions must be greater than 0");
|
||||
|
||||
this._vectorStore = vectorStore;
|
||||
this._collectionName = collectionName;
|
||||
this._functionByName = functions.ToDictionary(function => function.Name);
|
||||
this._maxNumberOfFunctions = maxNumberOfFunctions;
|
||||
this._logger = (loggerFactory ?? NullLoggerFactory.Instance).CreateLogger<FunctionStore>();
|
||||
this._options = options ?? new FunctionStoreOptions();
|
||||
|
||||
// Create and assert the collection support record keys of string type
|
||||
this._collection = this._vectorStore.GetDynamicCollection(collectionName, new VectorStoreCollectionDefinition()
|
||||
{
|
||||
Properties = [
|
||||
new VectorStoreKeyProperty("Name", typeof(string)),
|
||||
new VectorStoreVectorProperty("Embedding", typeof(string), dimensions: vectorDimensions)
|
||||
]
|
||||
});
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Saves the functions to the vector store.
|
||||
/// </summary>
|
||||
/// <param name="cancellationToken">The cancellation token to use for cancellation.</param>
|
||||
public async Task SaveAsync(CancellationToken cancellationToken = default)
|
||||
{
|
||||
// Get function data to vectorize
|
||||
var nameSourcePairs = await this.GetFunctionsVectorizationInfoAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
var functionRecords = new List<Dictionary<string, object?>>(nameSourcePairs.Count);
|
||||
|
||||
// Create vector store records
|
||||
for (var i = 0; i < nameSourcePairs.Count; i++)
|
||||
{
|
||||
var (name, vectorizationSource) = nameSourcePairs[i];
|
||||
|
||||
functionRecords.Add(new Dictionary<string, object?>()
|
||||
{
|
||||
["Name"] = name,
|
||||
["Embedding"] = vectorizationSource
|
||||
});
|
||||
}
|
||||
|
||||
// Create collection and upsert all vector store records
|
||||
await this._collection.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
await this._collection.UpsertAsync(functionRecords, cancellationToken: cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Searches for functions based on the provided context.
|
||||
/// </summary>
|
||||
/// <param name="context">The context to search for functions.</param>
|
||||
/// <param name="cancellationToken">The cancellation token to use for cancellation.</param>
|
||||
public async Task<IEnumerable<AIFunction>> SearchAsync(string context, CancellationToken cancellationToken = default)
|
||||
{
|
||||
await this.AssertCollectionExistsAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
List<VectorSearchResult<Dictionary<string, object?>>> results = new();
|
||||
|
||||
await foreach (var result in this._collection
|
||||
.SearchAsync(context, top: this._maxNumberOfFunctions, cancellationToken: cancellationToken).ConfigureAwait(false))
|
||||
{
|
||||
results.Add(result);
|
||||
}
|
||||
|
||||
this._logger.LogFunctionsSearchResults(context, this._maxNumberOfFunctions, results);
|
||||
|
||||
return results.Select(result => this._functionByName[(string)result.Record["Name"]!]);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Get the function vectorization information, which includes the function name and the source used for vectorization.
|
||||
/// </summary>
|
||||
/// <param name="cancellationToken">The cancellation token to use for cancellation.</param>
|
||||
/// <returns>The function name and vectorization source pairs.</returns>
|
||||
private async Task<List<FunctionVectorizationInfo>> GetFunctionsVectorizationInfoAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
List<FunctionVectorizationInfo> nameSourcePairs = new(this._functionByName.Count);
|
||||
|
||||
var provider = this._options.EmbeddingValueProvider ?? ((function, _) =>
|
||||
{
|
||||
string descriptionPart = string.IsNullOrEmpty(function.Description) ? string.Empty : $", description: {function.Description}";
|
||||
return Task.FromResult($"Function name: {function.Name}{descriptionPart}");
|
||||
});
|
||||
|
||||
foreach (KeyValuePair<string, AIFunction> pair in this._functionByName)
|
||||
{
|
||||
var vectorizationSource = await provider.Invoke(pair.Value, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
nameSourcePairs.Add(new FunctionVectorizationInfo(pair.Key, vectorizationSource));
|
||||
}
|
||||
|
||||
this._logger.LogFunctionsVectorizationInfo(nameSourcePairs);
|
||||
|
||||
return nameSourcePairs;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Asserts that the collection exists in the vector store.
|
||||
/// </summary>
|
||||
/// <param name="cancellationToken">The cancellation token to use for cancellation.</param>
|
||||
private async Task AssertCollectionExistsAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
if (!this._isCollectionExistenceAsserted)
|
||||
{
|
||||
if (!await this._collection.CollectionExistsAsync(cancellationToken).ConfigureAwait(false))
|
||||
{
|
||||
throw new InvalidOperationException($"Collection '{this._collectionName}' does not exist.");
|
||||
}
|
||||
|
||||
this._isCollectionExistenceAsserted = true;
|
||||
}
|
||||
}
|
||||
|
||||
internal readonly struct FunctionVectorizationInfo
|
||||
{
|
||||
public string Name { get; }
|
||||
|
||||
public string VectorizationSource { get; }
|
||||
|
||||
public FunctionVectorizationInfo(string name, string vectorizationSource)
|
||||
{
|
||||
this.Name = name;
|
||||
this.VectorizationSource = vectorizationSource;
|
||||
}
|
||||
|
||||
public void Deconstruct(out string name, out string vectorizationSource)
|
||||
{
|
||||
name = this.Name;
|
||||
vectorizationSource = this.VectorizationSource;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
using System.Collections.Generic;
|
||||
using System.Diagnostics.CodeAnalysis;
|
||||
using System.Linq;
|
||||
using Microsoft.Extensions.Logging;
|
||||
using Microsoft.Extensions.VectorData;
|
||||
|
||||
namespace Microsoft.Agents.AI.Functions;
|
||||
|
||||
[ExcludeFromCodeCoverage]
|
||||
internal static class FunctionStoreLoggingExtensions
|
||||
{
|
||||
internal static void LogFunctionsVectorizationInfo(this ILogger logger, IList<FunctionStore.FunctionVectorizationInfo> vectorizationInfo)
|
||||
{
|
||||
logger.LogInformation("ContextualFunctionProvider: Number of function to vectorize: {Count}", vectorizationInfo.Count);
|
||||
|
||||
if (logger.IsEnabled(LogLevel.Trace))
|
||||
{
|
||||
logger.LogTrace("ContextualFunctionProvider: Functions vectorization info: {VectorizationInfo}",
|
||||
string.Join(", ", vectorizationInfo.Select(info => $"\"Function: {info.Name}, VectorizationSource: {info.VectorizationSource}\"")));
|
||||
}
|
||||
}
|
||||
|
||||
internal static void LogFunctionsSearchResults(this ILogger logger, string context, int maxNumberOfFunctionsToReturn, IList<VectorSearchResult<Dictionary<string, object?>>> results)
|
||||
{
|
||||
logger.LogInformation("ContextualFunctionProvider: Search returned {Count} functions, with a maximum limit of {MaxCount}", results.Count, maxNumberOfFunctionsToReturn);
|
||||
|
||||
if (logger.IsEnabled(LogLevel.Trace))
|
||||
{
|
||||
logger.LogTrace("ContextualFunctionProvider: Functions search results for context {Context} with a maximum limit of {MaxCount}: {Results}",
|
||||
$"\"{context}\"",
|
||||
maxNumberOfFunctionsToReturn,
|
||||
string.Join(", ", results.Select(result => $"\"Function: {result.Record["Name"]}, Score: {result.Score}\"")));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
using System;
|
||||
using System.Threading;
|
||||
using System.Threading.Tasks;
|
||||
using Microsoft.Extensions.AI;
|
||||
|
||||
namespace Microsoft.Agents.AI.Functions;
|
||||
|
||||
/// <summary>
|
||||
/// Options for the <see cref="FunctionStore"/>
|
||||
/// </summary>
|
||||
internal sealed class FunctionStoreOptions
|
||||
{
|
||||
/// <summary>
|
||||
/// A callback function that returns a value used to create a function embedding. The value is vectorized,
|
||||
/// and the resulting vector is stored in the vector store for use in vector searches for functions relevant
|
||||
/// to the context.
|
||||
/// If not provided, the default behavior is to concatenate the function name and description into a single string.
|
||||
/// </summary>
|
||||
public Func<AIFunction, CancellationToken, Task<string>>? EmbeddingValueProvider { get; set; }
|
||||
}
|
||||
+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