Streaming Function Filters?

[Tyler-R-Kendrick]

I'm trying to emulate some of the behavior in the python guidance library around constraining llm outputs, but for the dotnet semantic kernel sdk.

I'd like to create a pluggable component that is idiomatic with the design of SK, and not shim behavior around it; but, I'm struggling to find the appropriate points of extension.

Ideally, I'd have a StreamingFunctionInvocationFilter that I could use to intercept and cancel the output stream from OpenAI, but I don't see anything that looks like this.

The closest approximation I can find is to do something like this:

public record struct FunctionStreamInvocationContext(
    FunctionInvocationContext Context,
    int TotalTokenCount,
    string ResultAggregate,
    string CurrentToken
);
public class ConstraintFunctionFilter(
    Func<FunctionStreamInvocationContext, bool> streamPredicate)
    : IFunctionInvocationFilter
{
    public async Task OnFunctionInvocationAsync(
        FunctionInvocationContext context,
        Func<FunctionInvocationContext, Task> next)
    {
        CancellationTokenSource tokenSource = new();
        var token = tokenSource.Token;
        var tokenCount = 0;
        var resultAggregate = string.Empty;
        await next(context);
        var result = context.Result;
        switch(result.GetValue<object>())
        {
            case IAsyncEnumerable<string> stream:
                await foreach (var item in stream.WithCancellation(token))
                {
                    resultAggregate += item;
                    var shouldContinue = streamPredicate(new(context, ++tokenCount, resultAggregate, item));
                    if (!shouldContinue)
                    {
                        tokenSource.Cancel();
                        break;
                    }
                }
                context.Result = new(context.Function, resultAggregate);
                break;
        }
    }
}

The idea here is to grab the stream, cancel it if criteria are met, and override the stream with a different result so subsequent filters don't continue enumeration of the original OpenAI request.

The problem is, I don't seem to be able to grab the byte data for tokens using this code - and I don't seem to be able to cancel the actual request (the stream continues generating tokens, since the cancellation only appears to apply to my enumeration of the stream and not the stream itself).

Any tips for intercepting this correctly in an idiomatic way?

[dmytrostruk]

Hi @Tyler-R-Kendrick , I think your approach is correct, here is an example how to use filters in streaming scenario:

semantic-kernel/dotnet/samples/Concepts/Filtering/FunctionInvocationFiltering.cs

Lines 167 to 187 in fe38a5f

	private sealed class StreamingFunctionFilterExample : IFunctionInvocationFilter
	{
	public async Task OnFunctionInvocationAsync(FunctionInvocationContext context, Func<FunctionInvocationContext, Task> next)
	{
	await next(context);
	// In streaming scenario, async enumerable is available in context result object.
	// To override data: get async enumerable from function result, override data and set new async enumerable in context result:
	var enumerable = context.Result.GetValue<IAsyncEnumerable<int>>();
	context.Result = new FunctionResult(context.Result, OverrideStreamingDataAsync(enumerable!));
	}
	private async IAsyncEnumerable<int> OverrideStreamingDataAsync(IAsyncEnumerable<int> data)
	{
	await foreach (var item in data)
	{
	// Example: override streaming data
	yield return item * 2;
	}
	}
	}

But if you want to get byte data for tokens or cancel actual HTTP request, you can define your interception and cancellation logic by extending HttpClient with DelegatingHandler:

semantic-kernel/dotnet/src/Plugins/Plugins.MsGraph/Connectors/Client/MsGraphClientLoggingHandler.cs

Lines 52 to 59 in fe38a5f

	protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
	{
	request.Headers.Add(ClientRequestIdHeaderName, Guid.NewGuid().ToString());
	this.LogHttpMessage(request.Headers, request.RequestUri, "REQUEST");
	HttpResponseMessage response = await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
	this.LogHttpMessage(response.Headers, response.RequestMessage?.RequestUri, "RESPONSE");
	return response;
	}

You can then pass your HttpClient to chat completion service:

semantic-kernel/dotnet/samples/Concepts/ChatCompletion/OpenAI_FunctionCalling.cs

Lines 99 to 108 in fe38a5f

	// Create a logging handler to output HTTP requests and responses
	var handler = new LoggingHandler(new HttpClientHandler(), this.Output);
	HttpClient httpClient = new(handler);
	// Create a kernel with OpenAI chat completion and WeatherPlugin
	IKernelBuilder kernelBuilder = Kernel.CreateBuilder();
	kernelBuilder.AddOpenAIChatCompletion(
	modelId: TestConfiguration.OpenAI.ChatModelId!,
	apiKey: TestConfiguration.OpenAI.ApiKey!,
	httpClient: httpClient);

I think at some point we can add CancellationToken to filter context, so it will be possible to monitor a cancellation, but I'm not sure that filters will get access to byte data from HTTP request, and I don't think they should have this capability, since it's HTTP-related logic, and there are existing mechanisms how to intercept it, like with example above. Please let me know if this resolves your scenario. Thanks!

[Tyler-R-Kendrick]

Just so I'm sure I understand the approach here, I would likely capture the cancellation token from the client request and request cancellation in the function filter. Is that correct?

I think that should be a sufficient workaround in the interim, but it seems that the cancellation token really should be passed down through function invocation. Should I make a feature request and PR for this behavior? Or is there something already in the pipeline for it?

[dmytrostruk]

I would likely capture the cancellation token from the client request and request cancellation in the function filter. Is that correct?

That's correct.

it seems that the cancellation token really should be passed down through function invocation. Should I make a feature request and PR for this behavior? Or is there something already in the pipeline for it?

I'm currently working on removing experimental flags from Filters feature, so I can create a PR for cancellation token in the same scope of work. If you could create a feature request, that would be great. Thanks!

[Tyler-R-Kendrick]

I've opened a request here. I'll close out this discussion. Thanks for your quick replies!

[dmytrostruk]

I would likely capture the cancellation token from the client request and request cancellation in the function filter. Is that correct?
That's correct.

@Tyler-R-Kendrick A little clarification here, cancellation token would allow to see if any cancellation happened above the stack so it will be possible to stop the operation inside the filter. It would also allow to pass the original cancellation token from function invocation in filter to subsequent async operations when needed. But cancellation token itself is not capable of cancelling the operation, it is a responsibility of CancellationTokenSource. But even today, if for some reason you want to stop function execution, you don't need a cancellation token, you can throw an exception with details, like in this example:

semantic-kernel/dotnet/samples/Concepts/Filtering/FunctionInvocationFiltering.cs

Lines 196 to 208 in 775994e

	try
	{
	await next(context);
	}
	catch (Exception exception)
	{
	this._logger.LogError(exception, "Something went wrong during function invocation");
	// Example: override function result value
	context.Result = new FunctionResult(context.Result, "Friendly message instead of exception");
	// Example: Rethrow another type of exception if needed
	// throw new InvalidOperationException("New exception");