.NET: IChatClient's service provider lost when using ChatAgentClient, breaking custom AITool implementations

[cgillum]

Context

I have a custom AITool implementation that needs to resolve method parameters from an IServiceProvider (using the AIFunctionFactoryOptions.ParameterBindingOptions extensibility in Microsoft.Extensions.AI.Abstractions). I won't go into the details in case it distracts from the core issue.

This scenario used to work well in earlier builds (last one I tried was commit 3ee9ddd, about 1 month old). However, starting today, I noticed that my custom AITool implementation is no longer working because AIFunctionArguments.Services always returns EmptyServiceProvider. Before it would return the IServiceProvider that gets passed to ChatClientBuilder.Build(IServiceProvider).

Repro Steps

Here's the code I had before, where the IChatClient infrastructure correctly used my provided IServiceProvider for binding to AITool parameters.

IServiceProvider sp = // ...

IChatClient chatClient = new AzureOpenAIClient(new Uri(azureOpenAiEndpoint), credential)
    .GetChatClient(azureOpenAiDeploymentName)
    .AsIChatClient()
    .UseFunctionInvocation()
    .Build(sp); // <-- add my service provider

ChatClientAgent agentClient = new(chatClient, agentOptions);
await agentClient.RunStreamingAsync(...); // <-- invokes AITool infrastructure

Investigation

The issue seems to be here:

agent-framework/dotnet/src/Microsoft.Extensions.AI.Agents/ChatCompletion/ChatClientExtensions.cs

Lines 22 to 27 in 78125f0

	chatBuilder.Use((IChatClient innerClient, IServiceProvider services) =>
	{
	var loggerFactory = services.GetService<ILoggerFactory>();
	return new NewFunctionInvokingChatClient(innerClient, loggerFactory, services);
	});

For some reason, the NewFunctionInvokingChatClient is always given EmptyServiceProvider as the services parameter. It somehow loses the IServiceProvider that I gave to my IChatClient originally.

Workaround

I worked around this by creating a custom IChatClient that implements NewFunctionInvokingChatClient and passing in my IServiceProvider there. That causes the code I highlighted above to get skipped, allowing my scenario to start working again. It took quite a bit of debugging to figure this out, though and it feels very fragile.

Hoping there is a simple fix to ensure that I don't have to rely on fragile workarounds. But it seems that the Agent Framework wrappers are breaking existing stable Microsoft.Extensions.AI behavior/expectations (custom parameter binding, in my case), so probably worth fixing.

[rogerbarreto]

Hi @cgillum thanks for providing this information, since you didn't provide a full reproduction code using the ParameterBinding code, I tried using the default path and had no issues resolving the ServiceProvider at the function invocation.

using System.ComponentModel;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.AI.Agents;
using Microsoft.Extensions.DependencyInjection;
using OpenAI;

var apiKey = Environment.GetEnvironmentVariable("OPENAI_API_KEY") ?? throw new InvalidOperationException("OPENAI_API_KEY is not set.");
var modelId = System.Environment.GetEnvironmentVariable("OPENAI_MODELID") ?? "gpt-4o";

IServiceProvider sp = new ServiceCollection()
    .AddSingleton<Location>()
    .BuildServiceProvider();

IChatClient chatClient = new OpenAIClient(apiKey)
    .GetChatClient(modelId)
    .AsIChatClient()
    .AsBuilder()
    .UseFunctionInvocation()
    .Build(sp); // <-- add my service provider

[Description("Get the weather for a given location.")]
static string GetWeather(IServiceProvider sp)
    => $"The weather in {sp.GetService<Location>()} is cloudy with a high of 15°C.";

var myFunction = AIFunctionFactory.Create(GetWeather);

ChatClientAgentOptions agentOptions = new() { ChatOptions = new ChatOptions() { Tools = [myFunction] } };
ChatClientAgent agentClient = new(chatClient, agentOptions);

var thread = agentClient.GetNewThread();

var result = await agentClient.RunAsync("What's the weather?", thread);

Console.WriteLine(result);

public class Location
{
    public override string ToString() => "Rio de Janeiro";
}

Please I appreciate if you can provide a more detailed code-base where the error does arise.

Thanks

[cgillum]

Hi @rogerbarreto , thanks for responding to the issue and trying to reproduce. Here is the full code that I was using to generate a custom AITool implementation.

/// <summary>
/// A factory for creating Durable Task agent tools.
/// </summary>
public static class DurableAgentToolFactory
{
    /// <summary>
    /// Creates a Durable Task agent tool.
    /// </summary>
    /// <remarks>
    /// This method is used to create tools that have the <see cref="IDurableAgentContext"/> parameter.
    /// </remarks>
    /// <param name="method">The method to create the tool from.</param>
    /// <param name="options">The options to create the tool with.</param>
    /// <returns>The created tool.</returns>
    public static AITool Create(Delegate method, AIFunctionFactoryOptions? options = null)
    {
        // This is a helper method to bind the IDurableAgentContext parameter in agent tool calls.
        static AIFunctionFactoryOptions.ParameterBindingOptions AgentContextBinding(ParameterInfo parameter)
        {
            if (parameter.ParameterType == typeof(IDurableAgentContext))
            {
                return new AIFunctionFactoryOptions.ParameterBindingOptions()
                {
                    BindParameter = (p, a) => a.Services?.GetRequiredService<IDurableAgentContext>(),
                    ExcludeFromSchema = true,
                };
            }

            return default;
        }

        // Merge the options with the default options. Mutating the provided options is not ideal,
        // but it's the simplest way to support the case where the user wants to provide a custom
        // ConfigureParameterBinding delegate. We can revisit this if AIFunctionFactoryOptions
        // supports a more explicit way to merge options.
        options ??= new AIFunctionFactoryOptions();
        if (options.ConfigureParameterBinding is null)
        {
            options.ConfigureParameterBinding = AgentContextBinding;
        }
        else
        {
            // Wrap the user-provided binding to add the default binding.
            Func<ParameterInfo, AIFunctionFactoryOptions.ParameterBindingOptions> original = options.ConfigureParameterBinding;
            options.ConfigureParameterBinding = (p) =>
            {
                // User-provided binding takes precedence. If it returns a non-default value, use that.
                AIFunctionFactoryOptions.ParameterBindingOptions result = original(p);
                if (result != default)
                {
                    return result;
                }

                // Fall back to the default binding.
                return AgentContextBinding(p);
            };
        }

        return AIFunctionFactory.Create(method, options);
    }
}

My app then used DurableAgentToolFactory.Create to create an AITool that had the custom parameter binding:

DocumentationTools tools = new();
AITool tool = DurableAgentToolFactory.Create(tools.StartDocumentGeneration);

The problem was with this line of code, which does the custom parameter binding:

BindParameter = (p, a) => a.Services?.GetRequiredService<IDurableAgentContext>(),

The a.Services was always returning a EmptyServiceProvider, so no services were correctly discovered.

[rogerbarreto]

Many of the extensions already support the IServiceProvider support, I added the missing existing ones that cover fully the option for those.

Using the IServiceProvider extensions, the provided sample works as expected. Here's a full reproduction code that I used for checking.

using Microsoft.Agents.AI;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;
using OpenAI;
using System;
using System.ClientModel;
using System.ClientModel.Primitives;
using System.ComponentModel;
using System.Net.Http;
using System.Reflection;
using System.Threading.Tasks;

internal class Program
{
    private static async Task Main(string[] args)
    {
        IServiceProvider sp = new ServiceCollection()
            .AddSingleton<IDurableAgentContext, MyContext>()
            .BuildServiceProvider();

        using var httpHandler = new MyHttpHandler();
        using var httpClient = new HttpClient(httpHandler);

        const string AgentName = "Bing-Agent-Framework";
        const string AgentInstructions = "You are an AI assistant that helps people find information.";

        string apiKey = Environment.GetEnvironmentVariable("OPENAI_API_KEY") ?? throw new InvalidOperationException("OPENAI_API_KEY is not set.");
        AITool tool = DurableAgentToolFactory.Create(TimePerCity.GetCurrentTime);

        AIAgent agent = new OpenAIClient(new ApiKeyCredential(apiKey), new() { Transport = new HttpClientPipelineTransport(httpClient) })
            .GetChatClient("gpt-4o-mini")
            .CreateAIAgent(name: AgentName, instructions: AgentInstructions, tools: [tool], services: sp);

        AgentThread thread = agent.GetNewThread();

        ChatMessage userMessage = new ChatMessage(ChatRole.User, "What is the time in Dublin?");

        var result = await agent.RunAsync(userMessage, thread);

        Console.WriteLine(result);
    }
}

public interface IDurableAgentContext
{
    string Name { get; }
}

public class TimePerCity
{
    [Description("Gets the current time in given city.")]
    public static string GetCurrentTime(IDurableAgentContext context, string city) => $"Context: {context.Name}, Time: 12:00 PM";
}

public class MyContext : IDurableAgentContext
{
    public string Name { get; set; } = "MyContextInstance";
}

/// <summary>
/// A factory for creating Durable Task agent tools.
/// </summary>
public static class DurableAgentToolFactory
{
    /// <summary>
    /// Creates a Durable Task agent tool.
    /// </summary>
    /// <remarks>
    /// This method is used to create tools that have the <see cref="IDurableAgentContext"/> parameter.
    /// </remarks>
    /// <param name="method">The method to create the tool from.</param>
    /// <param name="options">The options to create the tool with.</param>
    /// <returns>The created tool.</returns>
    public static AITool Create(Delegate method, AIFunctionFactoryOptions? options = null)
    {
        // This is a helper method to bind the IDurableAgentContext parameter in agent tool calls.
        static AIFunctionFactoryOptions.ParameterBindingOptions AgentContextBinding(ParameterInfo parameter)
        {
            if (parameter.ParameterType == typeof(IDurableAgentContext))
            {
                return new AIFunctionFactoryOptions.ParameterBindingOptions()
                {
                    BindParameter = (p, a)
                        => a.Services?.GetRequiredService<IDurableAgentContext>(),
                    ExcludeFromSchema = true,
                };
            }

            return default;
        }

        // Merge the options with the default options. Mutating the provided options is not ideal,
        // but it's the simplest way to support the case where the user wants to provide a custom
        // ConfigureParameterBinding delegate. We can revisit this if AIFunctionFactoryOptions
        // supports a more explicit way to merge options.
        options ??= new AIFunctionFactoryOptions();
        if (options.ConfigureParameterBinding is null)
        {
            options.ConfigureParameterBinding = AgentContextBinding;
        }
        else
        {
            // Wrap the user-provided binding to add the default binding.
            Func<ParameterInfo, AIFunctionFactoryOptions.ParameterBindingOptions> original = options.ConfigureParameterBinding;
            options.ConfigureParameterBinding = (p) =>
            {
                // User-provided binding takes precedence. If it returns a non-default value, use that.
                AIFunctionFactoryOptions.ParameterBindingOptions result = original(p);
                if (result != default)
                {
                    return result;
                }

                // Fall back to the default binding.
                return AgentContextBinding(p);
            };
        }

        return AIFunctionFactory.Create(method, options);
    }
}