Enable AIFunctionFactory to resolve parameters from an IServiceProvider

src/Libraries/Microsoft.Extensions.AI.Abstractions/Utilities/AIJsonSchemaCreateOptions.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Reflection;
 using System.Text.Json.Nodes;
 
 #pragma warning disable S1067 // Expressions should not be too complex
@@ -23,6 +24,17 @@ public sealed class AIJsonSchemaCreateOptions : IEquatable<AIJsonSchemaCreateOpt
     /// </summary>
     public Func<AIJsonSchemaCreateContext, JsonNode, JsonNode>? TransformSchemaNode { get; init; }
 
+    /// <summary>
+    /// Gets a callback that is invoked for every parameter in the <see cref="MethodBase"/> provided to
+    /// <see cref="AIJsonUtilities.CreateFunctionJsonSchema"/> in order to determine whether it should
+    /// be included in the generated schema.
+    /// </summary>
+    /// <remarks>
+    /// By default, when <see cref="IncludeParameter"/> is <see langword="null"/>,
+    /// all parameters are included in the generated schema.
+    /// </remarks>
+    public Func<ParameterInfo, bool>? IncludeParameter { get; init; }
+
     /// <summary>
     /// Gets a value indicating whether to include the type keyword in inferred schemas for .NET enums.
     /// </summary>
@@ -44,19 +56,24 @@ public sealed class AIJsonSchemaCreateOptions : IEquatable<AIJsonSchemaCreateOpt
     public bool RequireAllProperties { get; init; } = true;
 
     /// <inheritdoc/>
-    public bool Equals(AIJsonSchemaCreateOptions? other)
-    {
-        return other is not null &&
-            TransformSchemaNode == other.TransformSchemaNode &&
-            IncludeTypeInEnumSchemas == other.IncludeTypeInEnumSchemas &&
-            DisallowAdditionalProperties == other.DisallowAdditionalProperties &&
-            IncludeSchemaKeyword == other.IncludeSchemaKeyword &&
-            RequireAllProperties == other.RequireAllProperties;
-    }
+    public bool Equals(AIJsonSchemaCreateOptions? other) =>
+        other is not null &&
+        TransformSchemaNode == other.TransformSchemaNode &&
+        IncludeParameter == other.IncludeParameter &&
+        IncludeTypeInEnumSchemas == other.IncludeTypeInEnumSchemas &&
+        DisallowAdditionalProperties == other.DisallowAdditionalProperties &&
+        IncludeSchemaKeyword == other.IncludeSchemaKeyword &&
+        RequireAllProperties == other.RequireAllProperties;
 
     /// <inheritdoc />
     public override bool Equals(object? obj) => obj is AIJsonSchemaCreateOptions other && Equals(other);
 
     /// <inheritdoc />
-    public override int GetHashCode() => (TransformSchemaNode, IncludeTypeInEnumSchemas, DisallowAdditionalProperties, IncludeSchemaKeyword, RequireAllProperties).GetHashCode();
+    public override int GetHashCode() =>
+        (TransformSchemaNode,
+         IncludeParameter,
+         IncludeTypeInEnumSchemas,
+         DisallowAdditionalProperties,
+         IncludeSchemaKeyword,
+         RequireAllProperties).GetHashCode();
 }

src/Libraries/Microsoft.Extensions.AI.Abstractions/Utilities/AIJsonUtilities.Schema.cs

@@ -12,7 +12,6 @@
 using System.Text.Json.Nodes;
 using System.Text.Json.Schema;
 using System.Text.Json.Serialization;
-using System.Threading;
 using Microsoft.Shared.Diagnostics;
 
 #pragma warning disable S1121 // Assignments should not be made from within sub-expressions
@@ -77,11 +76,11 @@ public static JsonElement CreateFunctionJsonSchema(
                 Throw.ArgumentException(nameof(parameter), "Parameter is missing a name.");
             }
 
-            if (parameter.ParameterType == typeof(CancellationToken))
+            if (inferenceOptions.IncludeParameter is { } includeParameter &&
+                !includeParameter(parameter))
             {
-                // CancellationToken is a special case that, by convention, we don't want to include in the schema.
-                // Invocations of methods that include a CancellationToken argument should also special-case CancellationToken
-                // to pass along what relevant token into the method's invocation.
+                // Skip parameters that should not be included in the schema.
+                // By default, all parameters are included.
                 continue;
             }
 

src/Libraries/Microsoft.Extensions.AI/ChatCompletion/AIFunctionArguments.cs

@@ -0,0 +1,42 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections;
+using System.Collections.Generic;
+using Microsoft.Shared.Diagnostics;
+
+#pragma warning disable CA1710 // Identifiers should have correct suffix
+
+namespace Microsoft.Extensions.AI;
+
+/// <summary>Represents arguments to be used with <see cref="AIFunction.InvokeAsync"/>.</summary>
+/// <remarks>
+/// <see cref="AIFunction.InvokeAsync"/> may be invoked with arbitary <see cref="IEnumerable{T}"/>
+/// implementations. However, some <see cref="AIFunction"/> implementations may dynamically check
+/// the type of the arguments, and if it's an <see cref="AIFunctionArguments"/>, use it to access
+/// an <see cref="IServiceProvider"/> that's passed in separately from the arguments enumeration.
+/// </remarks>
+public class AIFunctionArguments : IEnumerable<KeyValuePair<string, object?>>
+{
+    /// <summary>The arguments represented by this instance.</summary>
+    private readonly IEnumerable<KeyValuePair<string, object?>> _arguments;
+
+    /// <summary>Initializes a new instance of the <see cref="AIFunctionArguments"/> class.</summary>
+    /// <param name="arguments">The arguments represented by this instance.</param>
+    /// <param name="serviceProvider">Options services associated with these arguments.</param>
+    public AIFunctionArguments(IEnumerable<KeyValuePair<string, object?>>? arguments, IServiceProvider? serviceProvider = null)
+    {
+        _arguments = Throw.IfNull(arguments);
+        ServiceProvider = serviceProvider;
+    }
+
+    /// <summary>Gets the services associated with these arguments.</summary>
+    public IServiceProvider? ServiceProvider { get; }
+
+    /// <inheritdoc />
+    public IEnumerator<KeyValuePair<string, object?>> GetEnumerator() => _arguments.GetEnumerator();
+
+    /// <inheritdoc />
+    IEnumerator IEnumerable.GetEnumerator() => ((IEnumerable)_arguments).GetEnumerator();
+}

src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FromServicesAttribute.cs

@@ -0,0 +1,16 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+
+namespace Microsoft.Extensions.AI;
+
+/// <summary>Indicates that a parameter to an <see cref="AIFunction"/> should be sourced from an associated <see cref="IServiceProvider"/>.</summary>
+[AttributeUsage(AttributeTargets.Parameter)]
+public sealed class FromServicesAttribute : Attribute
+{
+    /// <summary>Initializes a new instance of the <see cref="FromServicesAttribute"/> class.</summary>
+    public FromServicesAttribute()
+    {
+    }
+}

src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs

@@ -8,6 +8,7 @@
 using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 using Microsoft.Shared.Diagnostics;
@@ -58,11 +59,13 @@ public partial class FunctionInvokingChatClient : DelegatingChatClient
     /// </summary>
     /// <param name="innerClient">The underlying <see cref="IChatClient"/>, or the next instance in a chain of clients.</param>
     /// <param name="logger">An <see cref="ILogger"/> to use for logging information about function invocation.</param>
-    public FunctionInvokingChatClient(IChatClient innerClient, ILogger? logger = null)
+    /// <param name="services">An optional <see cref="IServiceProvider"/> to use for resolving services required by the <see cref="AIFunction"/> instances being invoked.</param>
+    public FunctionInvokingChatClient(IChatClient innerClient, ILogger? logger = null, IServiceProvider? services = null)
         : base(innerClient)
     {
-        _logger = logger ?? NullLogger.Instance;
+        _logger = logger ?? (ILogger?)services?.GetService<ILogger<FunctionInvokingChatClient>>() ?? NullLogger.Instance;
         _activitySource = innerClient.GetService<ActivitySource>();
+        Services = services;
     }
 
     /// <summary>
@@ -77,6 +80,9 @@ public static FunctionInvocationContext? CurrentContext
         protected set => _currentContext.Value = value;
     }
 
+    /// <summary>Gets the <see cref="IServiceProvider"/> used for resolving services required by the <see cref="AIFunction"/> instances being invoked.</summary>
+    public IServiceProvider? Services { get; }
+
     /// <summary>
     /// Gets or sets a value indicating whether to handle exceptions that occur during function calls.
     /// </summary>
@@ -687,8 +693,14 @@ FunctionResultContent CreateFunctionResultContent(FunctionInvocationResult resul
         object? result = null;
         try
         {
+            IEnumerable<KeyValuePair<string, object?>>? arguments = context.CallContent.Arguments;
+            if (Services is not null)
+            {
+                arguments = new AIFunctionArguments(arguments, Services);
+            }
+
             CurrentContext = context;
-            result = await context.Function.InvokeAsync(context.CallContent.Arguments, cancellationToken).ConfigureAwait(false);
+            result = await context.Function.InvokeAsync(arguments, cancellationToken).ConfigureAwait(false);
         }
         catch (Exception e)
         {

src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClientBuilderExtensions.cs

@@ -32,7 +32,7 @@ public static ChatClientBuilder UseFunctionInvocation(
         {
             loggerFactory ??= services.GetService<ILoggerFactory>();
 
-            var chatClient = new FunctionInvokingChatClient(innerClient, loggerFactory?.CreateLogger(typeof(FunctionInvokingChatClient)));
+            var chatClient = new FunctionInvokingChatClient(innerClient, loggerFactory?.CreateLogger(typeof(FunctionInvokingChatClient)), services);
             configure?.Invoke(chatClient);
             return chatClient;
         });

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

@@ -14,9 +14,14 @@
 using System.Text.Json.Serialization.Metadata;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Shared.Collections;
 using Microsoft.Shared.Diagnostics;
 
+#pragma warning disable CA1031 // Do not catch general exception types
+#pragma warning disable S2302 // "nameof" should be used
+#pragma warning disable S3011 // Reflection should not be used to increase accessibility of classes, methods, or fields
+
 namespace Microsoft.Extensions.AI;
 
 /// <summary>Provides factory methods for creating commonly used implementations of <see cref="AIFunction"/>.</summary>
@@ -244,6 +249,32 @@ public static ReflectionAIFunctionDescriptor GetOrCreate(MethodInfo method, AIFu
 
         private ReflectionAIFunctionDescriptor(DescriptorKey key, JsonSerializerOptions serializerOptions)
         {
+            AIJsonSchemaCreateOptions schemaOptions = new()
+            {
+                // This needs to be kept in sync with the shape of AIJsonSchemaCreateOptions.
+                TransformSchemaNode = key.SchemaOptions.TransformSchemaNode,
+                IncludeParameter = parameterInfo =>
+                {
+                    // Explicitly exclude from the schema CancellationToken parameters as well
+                    // as those annotated as [FromServices] or [FromKeyedServices]. These will be satisfied
+                    // from sources other than arguments to InvokeAsync.
+                    if (parameterInfo.ParameterType == typeof(CancellationToken) ||
+                        parameterInfo.GetCustomAttribute<FromServicesAttribute>(inherit: true) is not null ||
+                        parameterInfo.GetCustomAttribute<FromKeyedServicesAttribute>(inherit: true) is not null)
+                    {
+                        return false;
+                    }
+
+                    // For all other parameters, delegate to whatever behavior is specified in the options.
+                    // If none is specified, include the parameter.
+                    return key.SchemaOptions.IncludeParameter?.Invoke(parameterInfo) ?? true;
+                },
+                IncludeTypeInEnumSchemas = key.SchemaOptions.IncludeTypeInEnumSchemas,
+                DisallowAdditionalProperties = key.SchemaOptions.DisallowAdditionalProperties,
+                IncludeSchemaKeyword = key.SchemaOptions.IncludeSchemaKeyword,
+                RequireAllProperties = key.SchemaOptions.RequireAllProperties,
+            };
+
             // Get marshaling delegates for parameters.
             ParameterInfo[] parameters = key.Method.GetParameters();
             ParameterMarshallers = new Func<IReadOnlyDictionary<string, object?>, CancellationToken, object?>[parameters.Length];
@@ -264,7 +295,7 @@ private ReflectionAIFunctionDescriptor(DescriptorKey key, JsonSerializerOptions
                 Name,
                 Description,
                 serializerOptions,
-                key.SchemaOptions);
+                schemaOptions);
         }
 
         public string Name { get; }
@@ -337,6 +368,40 @@ static bool IsAsyncMethod(MethodInfo method)
                     cancellationToken;
             }
 
+            // For DI-based parameters, try to resolve from the service provider.
+            if (parameter.GetCustomAttribute<FromServicesAttribute>(inherit: true) is { } fsAttr)
+            {
+                return (arguments, _) =>
+                {
+                    if ((arguments as AIFunctionArguments)?.ServiceProvider is IServiceProvider services &&
+                        services.GetService(parameterType) is object service)
+                    {
+                        return service;
+                    }
+
+                    // No service could be resolved. Return a default value if it's optional, otherwise throw.
+                    return parameter.HasDefaultValue ?
+                        parameter.DefaultValue :
+                        throw new InvalidOperationException($"Unable to resolve service of type '{parameterType}' for parameter '{parameter.Name}'.");
+                };
+            }
+            else if (parameter.GetCustomAttribute<FromKeyedServicesAttribute>(inherit: true) is { } fksAttr)
+            {
+                return (arguments, _) =>
+                {
+                    if ((arguments as AIFunctionArguments)?.ServiceProvider is IKeyedServiceProvider services &&
+                        services.GetKeyedService(parameterType, fksAttr.Key) is object service)
+                    {
+                        return service;
+                    }
+
+                    // No service could be resolved. Return a default value if it's optional, otherwise throw.
+                    return parameter.HasDefaultValue ?
+                        parameter.DefaultValue :
+                        throw new InvalidOperationException($"Unable to resolve service of type '{parameterType}' with key '{fksAttr.Key}' for parameter '{parameter.Name}'.");
+                };
+            }
+
             // For all other parameters, create a marshaller that tries to extract the value from the arguments dictionary.
             return (arguments, _) =>
             {
@@ -355,7 +420,6 @@ static bool IsAsyncMethod(MethodInfo method)
 
                     object? MarshallViaJsonRoundtrip(object value)
                     {
-#pragma warning disable CA1031 // Do not catch general exception types
                         try
                         {
                             string json = JsonSerializer.Serialize(value, serializerOptions.GetTypeInfo(value.GetType()));
@@ -366,7 +430,6 @@ static bool IsAsyncMethod(MethodInfo method)
                             // Eat any exceptions and fall back to the original value to force a cast exception later on.
                             return value;
                         }
-#pragma warning restore CA1031
                     }
                 }
 
@@ -479,9 +542,7 @@ private static MethodInfo GetMethodFromGenericMethodDefinition(Type specializedT
 #if NET
             return (MethodInfo)specializedType.GetMemberWithSameMetadataDefinitionAs(genericMethodDefinition);
 #else
-#pragma warning disable S3011 // Reflection should not be used to increase accessibility of classes, methods, or fields
             const BindingFlags All = BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Static | BindingFlags.Instance;
-#pragma warning restore S3011 // Reflection should not be used to increase accessibility of classes, methods, or fields
             return specializedType.GetMethods(All).First(m => m.MetadataToken == genericMethodDefinition.MetadataToken);
 #endif
         }

test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/Utilities/AIJsonUtilitiesTests.cs

@@ -86,6 +86,12 @@ public static void AIJsonSchemaCreateOptions_UsesStructuralEquality()
                     property.SetValue(options2, transformer);
                     break;
 
+                case null when property.PropertyType == typeof(Func<ParameterInfo, bool>):
+                    Func<ParameterInfo, bool> includeParameter = static (parameter) => true;
+                    property.SetValue(options1, includeParameter);
+                    property.SetValue(options2, includeParameter);
+                    break;
+
                 default:
                     Assert.Fail($"Unexpected property type: {property.PropertyType}");
                     break;