Skip to content

Support [FromKeyedServices] in AIFunctionFactory #6310

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Apr 17, 2025
Merged

Support [FromKeyedServices] in AIFunctionFactory #6310

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 73 additions & 8 deletions src/Libraries/Microsoft.Extensions.AI/Functions/AIFunctionFactory.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,15 @@ public static partial class AIFunctionFactory
/// <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/>.
/// </description>
/// </item>
/// <item>
/// <description>
/// By default, parameters attributed with <see cref="FromKeyedServicesAttribute"/> are resolved from the <see cref="AIFunctionArguments.Services"/>
/// property and are not included in the JSON schema. If the parameter is optional, such that a default value is provided,
/// <see cref="AIFunctionArguments.Services"/> is allowed to be <see langword="null"/>; otherwise, <see cref="AIFunctionArguments.Services"/>
/// must be non-<see langword="null"/>, or else the invocation will fail with an exception due to the required nature of the parameter.
/// The handling of such parameters may be overridden via <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/>.
/// </description>
/// </item>
/// </list>
/// All other parameter types are, by default, bound from the <see cref="AIFunctionArguments"/> dictionary passed into <see cref="AIFunction.InvokeAsync"/>
/// and are included in the generated JSON schema. This may be overridden by the <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/> provided
Expand Down Expand Up @@ -131,7 +140,7 @@ public static AIFunction Create(Delegate method, AIFunctionFactoryOptions? optio
/// <description>
/// <see cref="CancellationToken"/> parameters are automatically bound to the <see cref="CancellationToken"/> passed into
/// the invocation via <see cref="AIFunction.InvokeAsync"/>'s <see cref="CancellationToken"/> parameter. The parameter is
/// not included in the generated JSON schema. The behavior of <see cref="CancellationToken"/> parameters may not be overridden.
/// not included in the generated JSON schema.
/// </description>
/// </item>
/// <item>
Expand All @@ -140,7 +149,6 @@ public static AIFunction Create(Delegate method, AIFunctionFactoryOptions? optio
/// and are not included in the JSON schema. If the parameter is optional, such that a default value is provided,
/// <see cref="AIFunctionArguments.Services"/> is allowed to be <see langword="null"/>; otherwise, <see cref="AIFunctionArguments.Services"/>
/// must be non-<see langword="null"/>, or else the invocation will fail with an exception due to the required nature of the parameter.
/// The handling of <see cref="IServiceProvider"/> parameters may be overridden via <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/>.
/// </description>
/// </item>
/// <item>
Expand All @@ -149,8 +157,15 @@ public static AIFunction Create(Delegate method, AIFunctionFactoryOptions? optio
/// passed into <see cref="AIFunction.InvokeAsync"/> and are not included in the JSON schema. If the <see cref="AIFunctionArguments"/>
/// instance passed to <see cref="AIFunction.InvokeAsync"/> is <see langword="null"/>, the <see cref="AIFunction"/> implementation
/// manufactures an empty instance, such that parameters of type <see cref="AIFunctionArguments"/> may always be satisfied, whether
/// optional or not. The handling of <see cref="AIFunctionArguments"/> parameters may be overridden via
/// <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/>.
/// optional or not.
/// </description>
/// </item>
/// <item>
/// <description>
/// By default, parameters attributed with <see cref="FromKeyedServicesAttribute"/> are resolved from the <see cref="AIFunctionArguments.Services"/>
/// property and are not included in the JSON schema. If the parameter is optional, such that a default value is provided,
/// <see cref="AIFunctionArguments.Services"/> is allowed to be <see langword="null"/>; otherwise, <see cref="AIFunctionArguments.Services"/>
/// must be non-<see langword="null"/>, or else the invocation will fail with an exception due to the required nature of the parameter.
/// </description>
/// </item>
/// </list>
Expand Down Expand Up @@ -236,6 +251,15 @@ public static AIFunction Create(Delegate method, string? name = null, string? de
/// <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/>.
/// </description>
/// </item>
/// <item>
/// <description>
/// By default, parameters attributed with <see cref="FromKeyedServicesAttribute"/> are resolved from the <see cref="AIFunctionArguments.Services"/>
/// property and are not included in the JSON schema. If the parameter is optional, such that a default value is provided,
/// <see cref="AIFunctionArguments.Services"/> is allowed to be <see langword="null"/>; otherwise, <see cref="AIFunctionArguments.Services"/>
/// must be non-<see langword="null"/>, or else the invocation will fail with an exception due to the required nature of the parameter.
/// The handling of such parameters may be overridden via <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/>.
/// </description>
/// </item>
/// </list>
/// All other parameter types are, by default, bound from the <see cref="AIFunctionArguments"/> dictionary passed into <see cref="AIFunction.InvokeAsync"/>
/// and are included in the generated JSON schema. This may be overridden by the <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/> provided
Expand Down Expand Up @@ -306,7 +330,7 @@ public static AIFunction Create(MethodInfo method, object? target, AIFunctionFac
/// <description>
/// <see cref="CancellationToken"/> parameters are automatically bound to the <see cref="CancellationToken"/> passed into
/// the invocation via <see cref="AIFunction.InvokeAsync"/>'s <see cref="CancellationToken"/> parameter. The parameter is
/// not included in the generated JSON schema. The behavior of <see cref="CancellationToken"/> parameters may not be overridden.
/// not included in the generated JSON schema.
/// </description>
/// </item>
/// <item>
Expand All @@ -315,7 +339,6 @@ public static AIFunction Create(MethodInfo method, object? target, AIFunctionFac
/// and are not included in the JSON schema. If the parameter is optional, such that a default value is provided,
/// <see cref="AIFunctionArguments.Services"/> is allowed to be <see langword="null"/>; otherwise, <see cref="AIFunctionArguments.Services"/>
/// must be non-<see langword="null"/>, or else the invocation will fail with an exception due to the required nature of the parameter.
/// The handling of <see cref="IServiceProvider"/> parameters may be overridden via <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/>.
/// </description>
/// </item>
/// <item>
Expand All @@ -324,8 +347,15 @@ public static AIFunction Create(MethodInfo method, object? target, AIFunctionFac
/// passed into <see cref="AIFunction.InvokeAsync"/> and are not included in the JSON schema. If the <see cref="AIFunctionArguments"/>
/// instance passed to <see cref="AIFunction.InvokeAsync"/> is <see langword="null"/>, the <see cref="AIFunction"/> implementation
/// manufactures an empty instance, such that parameters of type <see cref="AIFunctionArguments"/> may always be satisfied, whether
/// optional or not. The handling of <see cref="AIFunctionArguments"/> parameters may be overridden via
/// <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/>.
/// optional or not.
/// </description>
/// </item>
/// <item>
/// <description>
/// By default, parameters attributed with <see cref="FromKeyedServicesAttribute"/> are resolved from the <see cref="AIFunctionArguments.Services"/>
/// property and are not included in the JSON schema. If the parameter is optional, such that a default value is provided,
/// <see cref="AIFunctionArguments.Services"/> is allowed to be <see langword="null"/>; otherwise, <see cref="AIFunctionArguments.Services"/>
/// must be non-<see langword="null"/>, or else the invocation will fail with an exception due to the required nature of the parameter.
/// </description>
/// </item>
/// </list>
Expand Down Expand Up @@ -426,6 +456,15 @@ public static AIFunction Create(MethodInfo method, object? target, string? name
/// <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/>.
/// </description>
/// </item>
/// <item>
/// <description>
/// By default, parameters attributed with <see cref="FromKeyedServicesAttribute"/> are resolved from the <see cref="AIFunctionArguments.Services"/>
/// property and are not included in the JSON schema. If the parameter is optional, such that a default value is provided,
/// <see cref="AIFunctionArguments.Services"/> is allowed to be <see langword="null"/>; otherwise, <see cref="AIFunctionArguments.Services"/>
/// must be non-<see langword="null"/>, or else the invocation will fail with an exception due to the required nature of the parameter.
/// The handling of such parameters may be overridden via <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/>.
/// </description>
/// </item>
/// </list>
/// All other parameter types are, by default, bound from the <see cref="AIFunctionArguments"/> dictionary passed into <see cref="AIFunction.InvokeAsync"/>
/// and are included in the generated JSON schema. This may be overridden by the <see cref="AIFunctionFactoryOptions.ConfigureParameterBinding"/> provided
Expand Down Expand Up @@ -668,6 +707,13 @@ private ReflectionAIFunctionDescriptor(DescriptorKey key, JsonSerializerOptions
return false;
}

// If the parameter is attributed as [FromKeyedServices], exclude it, as we'll instead
// get its value from the IServiceProvider.
if (parameterInfo.GetCustomAttribute<FromKeyedServicesAttribute>(inherit: true) is not null)
{
return false;
}

// If there was an existing IncludeParameter delegate, now defer to it as we've
// excluded everything we need to exclude.
if (key.SchemaOptions.IncludeParameter is { } existingIncludeParameter)
Expand Down Expand Up @@ -806,6 +852,25 @@ static bool IsAsyncMethod(MethodInfo method)
};
}

// For [FromKeyedServices] parameters, we bind to the services passed directly to InvokeAsync via AIFunctionArguments.
if (parameter.GetCustomAttribute<FromKeyedServicesAttribute>(inherit: true) is { } keyedAttr)
{
return (arguments, _) =>
{
if ((arguments.Services as IKeyedServiceProvider)?.GetKeyedService(parameterType, keyedAttr.Key) is { } service)
{
return service;
}

if (!parameter.HasDefaultValue)
{
Throw.ArgumentException(nameof(arguments), $"No service of type '{parameterType}' with key '{keyedAttr.Key}' was found.");
}

return parameter.DefaultValue;
};
}

// For all other parameters, create a marshaller that tries to extract the value from the arguments dictionary.
// Resolve the contract used to marshal the value from JSON -- can throw if not supported or not found.
JsonTypeInfo typeInfo = serializerOptions.GetTypeInfo(parameterType);
Expand Down
67 changes: 45 additions & 22 deletions test/Libraries/Microsoft.Extensions.AI.Tests/Functions/AIFunctionFactoryTest.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -427,46 +427,69 @@ public async Task Create_NoInstance_DisposableAndAsyncDisposableInstanceCreatedD
}

[Fact]
public async Task ConfigureParameterBinding_CanBeUsedToSupportFromKeyedServices()
public async Task FromKeyedServices_ResolvesFromServiceProvider()
{
MyService service = new(42);

ServiceCollection sc = new();
sc.AddKeyedSingleton("key", service);
IServiceProvider sp = sc.BuildServiceProvider();

AIFunction f = AIFunctionFactory.Create(
([FromKeyedServices("key")] MyService service, int myInteger) => service.Value + myInteger,
new AIFunctionFactoryOptions
{
ConfigureParameterBinding = p =>
{
if (p.GetCustomAttribute<FromKeyedServicesAttribute>() is { } attr)
{
return new()
{
BindParameter = (p, a) =>
(a.Services as IKeyedServiceProvider)?.GetKeyedService(p.ParameterType, attr.Key) is { } s ? s :
p.HasDefaultValue ? p.DefaultValue :
throw new ArgumentException($"Unable to resolve argument for '{p.Name}'."),
ExcludeFromSchema = true
};
}
AIFunction f = AIFunctionFactory.Create(([FromKeyedServices("key")] MyService service, int myInteger) => service.Value + myInteger);

return default;
},
});
Assert.Contains("myInteger", f.JsonSchema.ToString());
Assert.DoesNotContain("service", f.JsonSchema.ToString());

Exception e = await Assert.ThrowsAsync<ArgumentException>(() => f.InvokeAsync(new() { ["myInteger"] = 1 }).AsTask());
Assert.Contains("No service of type", e.Message);

var result = await f.InvokeAsync(new() { ["myInteger"] = 1, Services = sp });
Assert.Contains("43", result?.ToString());
}

[Fact]
public async Task FromKeyedServices_NullKeysBindToNonKeyedServices()
{
MyService service = new(42);

ServiceCollection sc = new();
sc.AddSingleton(service);
IServiceProvider sp = sc.BuildServiceProvider();

AIFunction f = AIFunctionFactory.Create(([FromKeyedServices(null!)] MyService service, int myInteger) => service.Value + myInteger);

Assert.Contains("myInteger", f.JsonSchema.ToString());
Assert.DoesNotContain("service", f.JsonSchema.ToString());

Exception e = await Assert.ThrowsAsync<ArgumentException>(() => f.InvokeAsync(new() { ["myInteger"] = 1 }).AsTask());
Assert.Contains("Unable to resolve", e.Message);
Assert.Contains("No service of type", e.Message);

var result = await f.InvokeAsync(new() { ["myInteger"] = 1, Services = sp });
Assert.Contains("43", result?.ToString());
}

[Fact]
public async Task FromKeyedServices_OptionalDefaultsToNull()
{
MyService service = new(42);

ServiceCollection sc = new();
sc.AddKeyedSingleton("key", service);
IServiceProvider sp = sc.BuildServiceProvider();

AIFunction f = AIFunctionFactory.Create(([FromKeyedServices("key")] MyService? service = null, int myInteger = 0) =>
service is null ? "null " + 1 : (service.Value + myInteger).ToString());

Assert.Contains("myInteger", f.JsonSchema.ToString());
Assert.DoesNotContain("service", f.JsonSchema.ToString());

var result = await f.InvokeAsync(new() { ["myInteger"] = 1 });
Assert.Contains("null 1", result?.ToString());

result = await f.InvokeAsync(new() { ["myInteger"] = 1, Services = sp });
Assert.Contains("43", result?.ToString());
}

[Fact]
public async Task ConfigureParameterBinding_CanBeUsedToSupportFromContext()
{
Expand Down
Loading