[API Proposal]: FunctionInvocationContext.IterationState

[verdie-g]

Background and motivation

I have several use-cases where I would like to keep a state for the duration of a FunctionInvokingChatClient iteration.

Retries

FunctionInvokingChatClient.FunctionInvoker can be used to implement a retry mechanism but without an easy access to a state object, it's hard to keep track of how many retries were already performed.

Caching

One tool can return user data X and another tool data Y for the same user. Both tools need to fetch the same data which contains X and Y. X and Y are not returned by the same tool to avoid unnecessarily flood the context. Though, when both tools are called in the same iteration I'd like to fetch the data once. The iteration state could help with that.

API Proposal

namespace System.Collections.Generic;

public class FunctionInvocationContext
{
    public IDictionary<object,object?> IterationState { get; }
}

With FunctionInvokingChatClient.AllowConcurrentInvocation=true, we might need to deal with concurrent accesses. Should it be ConcurrentDictionary?

API Usage

In this example, tools can throw a custom AIFunctionRetryException to indicate to the LLM to retry the call. If too many retries are performed, a different message is sent to the LLM to try to stop the loop.

chatClient = new FunctionInvokingChatClient(chatClient)
{
    FunctionInvoker = async (ctx, ct) =>
    {
        try
        {
            return await ctx.Function.InvokeAsync(ctx.Arguments, ct);
        }
        catch (AIFunctionRetryException e)
        {
            string retryKey = $"retry.{ctx.Function.Name}";
            int retries;
            if (ctx.IterationState.TryGetValue(retryKey, out object? retriesObj))
            {
                retries = (int)retriesObj!;
            }
            else
            {
                retries = 1;
                ctx.IterationState[retryKey] = retries;
            }

            if (retries > 3)
            {
                return "Stop or something";
            }

            return $"""
                {e.Message}

                Fix the errors and try again.
                """;
        }
    },
};

Alternative Designs

DI Iteration Scope

Alternatively, if there was a DI scope created for the iteration, I could create a DI service by iteration. Not super familiar with DI scopes so I'm not sure that's possible.

AsyncLocal

The user could probably still create its own AsyncLocal but in my experience it's a feature that it beyond most C# developers.

Risks

No response