Support [FromKeyedServices] in AIFunctionFactory (#6310)

Author: stephentoub

src/Libraries/Microsoft.Extensions.AI/Functions/AIFunctionFactory.cs

@@ -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
@@ -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>
@@ -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>
@@ -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>
@@ -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
@@ -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>
@@ -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>
@@ -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>
@@ -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
@@ -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)
@@ -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);

test/Libraries/Microsoft.Extensions.AI.Tests/Functions/AIFunctionFactoryTest.cs

@@ -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()
     {