Allow caching AIFunction

[raffaeler]

Description

The default concrete implementation of AIFunction is ReflectionAIFunction.
This object captures the reference of the object containing the function to be executed on-demand by the LLM.
This reference is held as private, therefore each AIFunction instance points to a unique instance of the "agent".

I would like to cache this ReflectionAIFunction for all the instances of my agents. This means that ReflectionAIFunction should get rid of (or allow to override/replace) the field called object? _target which is used to invoke the method.

Reproduction Steps

In the "agent" default constructor I create the AIFunction

var function = AIFunctionFactory.Create(MyMethod);

later on, I can invoke the function using:

function.InvokeAsync(arguments, cancellationToken);

But InvokeAsync does not allow to specify the target instance.
Using UnderlyingMethod would not help as I could not reuse the internal marshalling transforms.

Expected behavior

I need to be able to manage the lifecycle of my "agents" without recreating the AIFunction for each instance.

If an LLM needs to invoke an agents multiple times and this agent is stateful, reusing the same agent instance would be wrong.
There are many examples of stateful agents like caching data or dbcontext, text to augment prompts, correlation ids, etc.

Actual behavior

Currently, I need to re-create the AIFunction for each instance of my agent.

Regression?

No

Known Workarounds

I am not aware of any workarounds

Configuration

• ASP.NET Core 9
• Microsoft.Extensions.AI version: 9.3.0-preview.1.25161.3

Other information

Nothing relevant

[stephentoub]

Can you use a static method and just pass your agent as the first argument?

[raffaeler]

I am not sure to undertand how would you do it.
The InvokeAsync is called with the following arguments:

function.InvokeAsync(arguments, cancellationToken);

But the _target object is privately held inside AIFunction (ReflectionAIFunction).
How can I override the use of _target?

[stephentoub]

_target isn't used for statics

[raffaeler]

Sure, but I need to access the agent state from the function, otherwise I would not have the problem.

The point is that I want to statically cache AIFunction but I need to invoke a specific method instance to make my agents stateful.
In certain case they will persist the state because they need to wait for human intervention,

[stephentoub]

Can you show me the code you want to write? Are you manually invoking function? How!/where? From where do you get the agent instance you'd bind to? Etc

[raffaeler]

Yes, I manually invoke the functions requested by the LLM.
Currently I have to create a new instance every time because I have no other way to pass a state to the tool/agent.

When I process the requests, I see there are tool calls:

// ...
                else if (response.FinishReason == ChatFinishReason.ToolCalls)
                {
                    // The model requested to invoke a tool
                    // Tools are the new name for Functions
                    await ProcessToolRequest(conversation, assistantMessage, _prompts);
                    lastWasTool = true;
                }
// ...

Then I process each of them manually, creating the ChatMessage for each response.
Please note that I will add more code here to eventually repeat the execution when the agent response fails the validation (for example it is not a json as requested).

    private async Task ProcessToolRequest(
        Conversation conversation,
        ChatMessage completion,
        IList<ChatMessage> prompts)
    {
        foreach (var toolCall in completion.Contents.OfType<FunctionCallContent>())
        {
            var functionName = toolCall.Name;
            var arguments = toolCall.Arguments;
            var result = await _toolInvoker.Invoke(
                conversation, functionName, arguments);

            ChatMessage responseMessage = new(
                ChatRole.Tool, [new FunctionResultContent(toolCall.CallId, result)]);

            prompts.Add(responseMessage);
        }
    }

Curently each invocation is done by creating a new instance of the tool, but now I need to make some of these agents stateful for the reasons explained in the previous messages.

This is where InvokeAsync is called.

    public async Task<object?> Invoke(
        Conversation conversation,
        string functionName,
        IDictionary<string, object?>? values)
    {
        var tool = _serviceProvider.GetKeyedService<ITool>(functionName);
        if (tool == null)
        {
            return $"The tool {functionName} is not part of the system yet.";
        }

        // the only workaround is to pass the state to the tool here
        // but I would like to load/save the tool instance in a cache

        var result = await tool.AIFunction.InvokeAsync(values) ?? string.Empty;
        return result;
    }

Everything about AIFunction can be cached out of the _target object which is the instance of the object.
What I would like to do is keep the function metadata separate from the instance in order to reuse the metadata.

[stephentoub]

Once a new build is published with #6158 and #6162, you should be able to make your scenario work well. Those don't let you substitute a this reference, but you shouldn't need to. You'd just have your AIFunction be for a static method that takes the equivalent of your receiver type as a regular argument, either using the built-in support for IServiceProvider to fetch the data from DI, or using a custom binder to thread the value through from the call site where you manually invoke your function.

Thanks for the feedback.

[raffaeler]

Thanks for the update @stephentoub
I need some time to put my hands on it as I am heading to the MVP Summit now.