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; }
|
||||
}
|
||||
Reference in New Issue
Block a user