.Net: Bugfix FunctionInvocation Middleware not working when no ChatOptions provided (#1030)

* Bugfix FunctionInvocation Middleware

* Throw when specialized RunOptions are used for Function Invocation middleware + UT
This commit is contained in:
Roger Barreto
2025-10-01 08:10:12 +01:00
committed by GitHub
Unverified
parent fb51d917fd
commit 3d04517877
2 changed files with 159 additions and 18 deletions
@@ -30,27 +30,34 @@ internal sealed class FunctionInvocationDelegatingAgent : DelegatingAIAgent
// Decorate options to add the middleware function
private AgentRunOptions? AgentRunOptionsWithFunctionMiddleware(AgentRunOptions? options)
{
if (options is ChatClientAgentRunOptions aco)
if (options is null || options.GetType() == typeof(AgentRunOptions))
{
var originalFactory = aco.ChatClientFactory;
aco.ChatClientFactory = chatClient =>
{
var builder = chatClient.AsBuilder();
if (originalFactory is not null)
{
builder.Use(originalFactory);
}
return builder.ConfigureOptions(co
=> co.Tools = co.Tools?.Select(tool => tool is AIFunction aiFunction
? new MiddlewareEnabledFunction(this.InnerAgent, aiFunction, this._delegateFunc)
: tool)
.ToList())
.Build();
};
options = new ChatClientAgentRunOptions();
}
if (options is not ChatClientAgentRunOptions aco)
{
throw new NotSupportedException($"Function Invocation Middleware is only supported without options or with {nameof(ChatClientAgentRunOptions)}.");
}
var originalFactory = aco.ChatClientFactory;
aco.ChatClientFactory = chatClient =>
{
var builder = chatClient.AsBuilder();
if (originalFactory is not null)
{
builder.Use(originalFactory);
}
return builder.ConfigureOptions(co
=> co.Tools = co.Tools?.Select(tool => tool is AIFunction aiFunction
? new MiddlewareEnabledFunction(this.InnerAgent, aiFunction, this._delegateFunc)
: tool)
.ToList())
.Build();
};
return options;
}
@@ -56,6 +56,135 @@ public sealed class FunctionInvocationDelegatingAgentTests
#region Function Invocation Tests
/// <summary>
/// Tests that middleware is invoked when functions are called during agent execution without options.
/// </summary>
[Fact]
public async Task RunAsync_WithFunctionCall_NoOptions_InvokesMiddlewareAsync()
{
// Arrange
var executionOrder = new List<string>();
var testFunction = AIFunctionFactory.Create(() =>
{
executionOrder.Add("Function-Executed");
return "Function result";
}, "TestFunction", "A test function");
var functionCall = new FunctionCallContent("call_123", "TestFunction", new Dictionary<string, object?>());
var mockChatClient = CreateMockChatClientWithFunctionCalls(functionCall);
var innerAgent = new ChatClientAgent(mockChatClient.Object, tools: [testFunction]);
var messages = new List<ChatMessage> { new(ChatRole.User, "Test message") };
async ValueTask<object?> MiddlewareCallbackAsync(AIAgent agent, FunctionInvocationContext context, Func<FunctionInvocationContext, CancellationToken, ValueTask<object?>> next, CancellationToken cancellationToken)
{
executionOrder.Add("Middleware-Pre");
var result = await next(context, cancellationToken);
executionOrder.Add("Middleware-Post");
return result;
}
var middleware = new FunctionInvocationDelegatingAgent(innerAgent, MiddlewareCallbackAsync);
// Act
await middleware.RunAsync(messages, null, null, CancellationToken.None);
// Assert
Assert.Contains("Middleware-Pre", executionOrder);
Assert.Contains("Function-Executed", executionOrder);
Assert.Contains("Middleware-Post", executionOrder);
// Verify execution order
var middlewarePreIndex = executionOrder.IndexOf("Middleware-Pre");
var functionIndex = executionOrder.IndexOf("Function-Executed");
var middlewarePostIndex = executionOrder.IndexOf("Middleware-Post");
Assert.True(middlewarePreIndex < functionIndex);
Assert.True(functionIndex < middlewarePostIndex);
}
/// <summary>
/// Tests that middleware is invoked when functions are called during agent execution without options.
/// </summary>
[Fact]
public async Task RunAsync_WithFunctionCall_AgentRunOptions_InvokesMiddlewareAsync()
{
// Arrange
var executionOrder = new List<string>();
var testFunction = AIFunctionFactory.Create(() =>
{
executionOrder.Add("Function-Executed");
return "Function result";
}, "TestFunction", "A test function");
var functionCall = new FunctionCallContent("call_123", "TestFunction", new Dictionary<string, object?>());
var mockChatClient = CreateMockChatClientWithFunctionCalls(functionCall);
var innerAgent = new ChatClientAgent(mockChatClient.Object, tools: [testFunction]);
var messages = new List<ChatMessage> { new(ChatRole.User, "Test message") };
async ValueTask<object?> MiddlewareCallbackAsync(AIAgent agent, FunctionInvocationContext context, Func<FunctionInvocationContext, CancellationToken, ValueTask<object?>> next, CancellationToken cancellationToken)
{
executionOrder.Add("Middleware-Pre");
var result = await next(context, cancellationToken);
executionOrder.Add("Middleware-Post");
return result;
}
var middleware = new FunctionInvocationDelegatingAgent(innerAgent, MiddlewareCallbackAsync);
// Act
await middleware.RunAsync(messages, null, new AgentRunOptions(), CancellationToken.None);
// Assert
Assert.Contains("Middleware-Pre", executionOrder);
Assert.Contains("Function-Executed", executionOrder);
Assert.Contains("Middleware-Post", executionOrder);
// Verify execution order
var middlewarePreIndex = executionOrder.IndexOf("Middleware-Pre");
var functionIndex = executionOrder.IndexOf("Function-Executed");
var middlewarePostIndex = executionOrder.IndexOf("Middleware-Post");
Assert.True(middlewarePreIndex < functionIndex);
Assert.True(functionIndex < middlewarePostIndex);
}
/// <summary>
/// Tests that middleware is invoked when functions are called during agent execution without options.
/// </summary>
[Fact]
public async Task RunAsync_WithFunctionCall_CustomAgentRunOptions_ThrowsNotSupportedAsync()
{
// Arrange
var executionOrder = new List<string>();
var testFunction = AIFunctionFactory.Create(() =>
{
executionOrder.Add("Function-Executed");
return "Function result";
}, "TestFunction", "A test function");
var functionCall = new FunctionCallContent("call_123", "TestFunction", new Dictionary<string, object?>());
var mockChatClient = CreateMockChatClientWithFunctionCalls(functionCall);
var innerAgent = new ChatClientAgent(mockChatClient.Object, tools: [testFunction]);
var messages = new List<ChatMessage> { new(ChatRole.User, "Test message") };
async ValueTask<object?> MiddlewareCallbackAsync(AIAgent agent, FunctionInvocationContext context, Func<FunctionInvocationContext, CancellationToken, ValueTask<object?>> next, CancellationToken cancellationToken)
{
executionOrder.Add("Middleware-Pre");
var result = await next(context, cancellationToken);
executionOrder.Add("Middleware-Post");
return result;
}
var middleware = new FunctionInvocationDelegatingAgent(innerAgent, MiddlewareCallbackAsync);
// Act
await Assert.ThrowsAsync<NotSupportedException>(() =>
middleware.RunAsync(messages, null, new CustomAgentRunOptions(), CancellationToken.None));
}
/// <summary>
/// Tests that middleware is invoked when functions are called during agent execution.
/// </summary>
@@ -848,4 +977,9 @@ public sealed class FunctionInvocationDelegatingAgentTests
{
return new ChatResponse([new ChatMessage(ChatRole.Assistant, "Default response")]);
}
/// <summary>
/// Custom AgentRunOptions class for testing
/// </summary>
private sealed class CustomAgentRunOptions : AgentRunOptions;
}